[8.14] [Reporting/CSV] Resolve max_concurrent_shard_requests issue (#…

…182536) (#182742) # Backport This will backport the following commits from `main` to `8.14`: - [[Reporting/CSV] Resolve max_concurrent_shard_requests issue (#182536)](#182536)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport)
elastic · May 7, 2024 · 03da0c9 · 03da0c9
1 parent ed86cd0
commit 03da0c9
Show file tree

Hide file tree

Showing 5 changed files with 260 additions and 73 deletions.
diff --git a/config/serverless.yml b/config/serverless.yml
@@ -145,6 +145,7 @@ xpack.screenshotting.enabled: false
 xpack.reporting.queue.pollInterval: 3m
 xpack.reporting.roles.enabled: false
 xpack.reporting.statefulSettings.enabled: false
+xpack.reporting.csv.maxConcurrentShardRequests: 0
 
 # Disabled Observability plugins
 xpack.ux.enabled: false

diff --git a/packages/kbn-generate-csv/src/lib/search_cursor_pit.test.ts b/packages/kbn-generate-csv/src/lib/search_cursor_pit.test.ts
@@ -6,22 +6,53 @@
  * Side Public License, v 1.
  */
 
+import * as Rx from 'rxjs';
+
+import type { estypes } from '@elastic/elasticsearch';
 import type { IScopedClusterClient, Logger } from '@kbn/core/server';
 import { elasticsearchServiceMock, loggingSystemMock } from '@kbn/core/server/mocks';
 import type { ISearchClient } from '@kbn/data-plugin/common';
 import { createSearchSourceMock } from '@kbn/data-plugin/common/search/search_source/mocks';
 import { createSearchRequestHandlerContext } from '@kbn/data-plugin/server/search/mocks';
-import type { SearchCursor, SearchCursorSettings } from './search_cursor';
+import type { SearchCursorSettings } from './search_cursor';
 import { SearchCursorPit } from './search_cursor_pit';
 
+class TestSearchCursorPit extends SearchCursorPit {
+  constructor(...args: ConstructorParameters<typeof SearchCursorPit>) {
+    super(...args);
+  }
+
+  public getCursorId() {
+    return this.cursorId;
+  }
+
+  public openPointInTime() {
+    return super.openPointInTime();
+  }
+
+  public searchWithPit(...args: Parameters<SearchCursorPit['searchWithPit']>) {
+    return super.searchWithPit(...args);
+  }
+
+  public setSearchAfter(...args: Parameters<SearchCursorPit['setSearchAfter']>) {
+    return super.setSearchAfter(...args);
+  }
+
+  public getSearchAfter() {
+    return super.getSearchAfter();
+  }
+}
+
 describe('CSV Export Search Cursor', () => {
   let settings: SearchCursorSettings;
   let es: IScopedClusterClient;
   let data: ISearchClient;
   let logger: Logger;
-  let cursor: SearchCursor;
+  let cursor: TestSearchCursorPit;
+
+  let openPointInTimeSpy: jest.SpyInstance<Promise<estypes.OpenPointInTimeResponse>>;
 
-  beforeEach(async () => {
+  beforeEach(() => {
     settings = {
       scroll: {
         duration: jest.fn(() => '10m'),
@@ -34,56 +65,116 @@ describe('CSV Export Search Cursor', () => {
 
     es = elasticsearchServiceMock.createScopedClusterClient();
     data = createSearchRequestHandlerContext();
-    jest.spyOn(es.asCurrentUser, 'openPointInTime').mockResolvedValue({ id: 'somewhat-pit-id' });
+
+    openPointInTimeSpy = jest
+      .spyOn(es.asCurrentUser, 'openPointInTime')
+      .mockResolvedValue({ id: 'somewhat-pit-id' });
 
     logger = loggingSystemMock.createLogger();
+  });
 
-    cursor = new SearchCursorPit(
-      'test-index-pattern-string',
-      settings,
-      { data, es },
-      new AbortController(),
-      logger
-    );
+  describe('with default settings', () => {
+    beforeEach(async () => {
+      cursor = new TestSearchCursorPit(
+        'test-index-pattern-string',
+        settings,
+        { data, es },
+        new AbortController(),
+        logger
+      );
 
-    const openPointInTimeSpy = jest
-      // @ts-expect-error create spy on private method
-      .spyOn(cursor, 'openPointInTime');
+      await cursor.initialize();
 
-    await cursor.initialize();
+      expect(openPointInTimeSpy).toBeCalledTimes(1);
+    });
 
-    expect(openPointInTimeSpy).toBeCalledTimes(1);
-  });
+    it('supports pit and max_concurrent_shard_requests', async () => {
+      const dataSearchSpy = jest
+        .spyOn(data, 'search')
+        .mockReturnValue(Rx.of({ rawResponse: { hits: { hits: [] } } }));
 
-  it('supports point-in-time', async () => {
-    const searchWithPitSpy = jest
-      // @ts-expect-error create spy on private method
-      .spyOn(cursor, 'searchWithPit')
-      // @ts-expect-error mock resolved value for spy on private method
-      .mockResolvedValueOnce({ rawResponse: { hits: [] } });
+      const searchSource = createSearchSourceMock();
+      await cursor.getPage(searchSource);
 
-    const searchSource = createSearchSourceMock();
-    await cursor.getPage(searchSource);
-    expect(searchWithPitSpy).toBeCalledTimes(1);
-  });
+      expect(dataSearchSpy).toBeCalledTimes(1);
+      expect(dataSearchSpy).toBeCalledWith(
+        {
+          params: {
+            body: expect.objectContaining({ pit: { id: 'somewhat-pit-id', keep_alive: '10m' } }),
+            max_concurrent_shard_requests: 5,
+          },
+        },
+        expect.objectContaining({
+          strategy: 'es',
+          transport: { maxRetries: 0, requestTimeout: '10m' },
+        })
+      );
+    });
+
+    it('can update internal cursor ID', () => {
+      cursor.updateIdFromResults({ pit_id: 'very-typical-pit-id', hits: { hits: [] } });
+      expect(cursor.getCursorId()).toBe('very-typical-pit-id');
+    });
 
-  it('can update internal cursor ID', () => {
-    cursor.updateIdFromResults({ pit_id: 'very-typical-pit-id', hits: { hits: [] } });
-    // @ts-expect-error private field
-    expect(cursor.cursorId).toBe('very-typical-pit-id');
+    it('manages search_after', () => {
+      cursor.setSearchAfter([
+        {
+          _index: 'test-index',
+          _id: 'test-doc-id',
+          sort: ['Wed Jan 17 15:35:47 MST 2024', 42],
+        },
+      ]);
+
+      expect(cursor.getSearchAfter()).toEqual(['Wed Jan 17 15:35:47 MST 2024', 42]);
+    });
   });
 
-  it('manages search_after', () => {
-    // @ts-expect-error access private method
-    cursor.setSearchAfter([
-      {
-        _index: 'test-index',
-        _id: 'test-doc-id',
-        sort: ['Wed Jan 17 15:35:47 MST 2024', 42],
-      },
-    ]);
+  describe('with max_concurrent_shard_requests=0', () => {
+    beforeEach(async () => {
+      settings.maxConcurrentShardRequests = 0;
+
+      cursor = new TestSearchCursorPit(
+        'test-index-pattern-string',
+        settings,
+        { data, es },
+        new AbortController(),
+        logger
+      );
+
+      await cursor.initialize();
+
+      expect(openPointInTimeSpy).toBeCalledTimes(1);
+    });
+
+    it('suppresses max_concurrent_shard_requests from search body', async () => {
+      const dataSearchSpy = jest
+        .spyOn(data, 'search')
+        .mockReturnValue(Rx.of({ rawResponse: { hits: { hits: [] } } }));
+
+      const searchSource = createSearchSourceMock();
+      await cursor.getPage(searchSource);
 
-    // @ts-expect-error access private method
-    expect(cursor.getSearchAfter()).toEqual(['Wed Jan 17 15:35:47 MST 2024', 42]);
+      expect(dataSearchSpy).toBeCalledTimes(1);
+      expect(dataSearchSpy).toBeCalledWith(
+        {
+          params: {
+            body: {
+              fields: [],
+              pit: { id: 'somewhat-pit-id', keep_alive: '10m' },
+              query: { bool: { filter: [], must: [], must_not: [], should: [] } },
+              runtime_mappings: {},
+              script_fields: {},
+              stored_fields: ['*'],
+            },
+            max_concurrent_shard_requests: undefined,
+          },
+        },
+        {
+          abortSignal: expect.any(AbortSignal),
+          strategy: 'es',
+          transport: { maxRetries: 0, requestTimeout: '10m' },
+        }
+      );
+    });
   });
 });
diff --git a/packages/kbn-generate-csv/src/lib/search_cursor_pit.ts b/packages/kbn-generate-csv/src/lib/search_cursor_pit.ts
@@ -37,7 +37,7 @@ export class SearchCursorPit extends SearchCursor {
     this.cursorId = await this.openPointInTime();
   }
 
-  private async openPointInTime() {
+  protected async openPointInTime() {
     const { includeFrozen, maxConcurrentShardRequests, scroll, taskInstanceFields } = this.settings;
 
     let pitId: string | undefined;
@@ -74,13 +74,17 @@ export class SearchCursorPit extends SearchCursor {
     return pitId;
   }
 
-  private async searchWithPit(searchBody: SearchRequest) {
+  protected async searchWithPit(searchBody: SearchRequest) {
     const { maxConcurrentShardRequests, scroll, taskInstanceFields } = this.settings;
 
+    // maxConcurrentShardRequests=0 is not supported
+    const effectiveMaxConcurrentShardRequests =
+      maxConcurrentShardRequests > 0 ? maxConcurrentShardRequests : undefined;
+
     const searchParamsPit = {
       params: {
         body: searchBody,
-        max_concurrent_shard_requests: maxConcurrentShardRequests,
+        max_concurrent_shard_requests: effectiveMaxConcurrentShardRequests,
       },
     };
 
@@ -146,14 +150,14 @@ export class SearchCursorPit extends SearchCursor {
     this.setSearchAfter(hits); // for pit only
   }
 
-  private getSearchAfter() {
+  protected getSearchAfter() {
     return this.searchAfter;
   }
 
   /**
    * For managing the search_after parameter, needed for paging using point-in-time
    */
-  private setSearchAfter(hits: Array<estypes.SearchHit<unknown>>) {
+  protected setSearchAfter(hits: Array<estypes.SearchHit<unknown>>) {
     // Update last sort results for next query. PIT is used, so the sort results
     // automatically include _shard_doc as a tiebreaker
     this.searchAfter = hits[hits.length - 1]?.sort as estypes.SortResults | undefined;