[data.search.bsearch] Forward request abortSignal to search strategy

src/plugins/data/common/search/strategies/es_search/types.ts

@@ -15,7 +15,8 @@ export type ISearchRequestParams = {
   trackTotalHits?: boolean;
 } & estypes.SearchRequest;
 
-export interface IEsSearchRequest extends IKibanaSearchRequest<ISearchRequestParams> {
+export interface IEsSearchRequest<T extends ISearchRequestParams = ISearchRequestParams>
+  extends IKibanaSearchRequest<T> {
   indexType?: string;
 }
 

src/plugins/data/server/search/routes/bsearch.ts

@@ -11,6 +11,7 @@ import { catchError } from 'rxjs/operators';
 import { BfetchServerSetup } from '@kbn/bfetch-plugin/server';
 import type { ExecutionContextSetup } from '@kbn/core/server';
 import apm from 'elastic-apm-node';
+import { getRequestAbortedSignal } from '../..';
 import {
   IKibanaSearchRequest,
   IKibanaSearchResponse,
@@ -28,6 +29,7 @@ export function registerBsearchRoute(
     IKibanaSearchResponse
   >('/internal/bsearch', (request) => {
     const search = getScoped(request);
+    const abortSignal = getRequestAbortedSignal(request.events.aborted$);
     return {
       /**
        * @param requestOptions
@@ -39,7 +41,7 @@ export function registerBsearchRoute(
           apm.addLabels(executionContextService.getAsLabels());
 
           return firstValueFrom(
-            search.search(requestData, restOptions).pipe(
+            search.search(requestData, { ...restOptions, abortSignal }).pipe(
               catchError((err) => {
                 // Re-throw as object, to get attributes passed to the client
                 // eslint-disable-next-line no-throw-literal

src/plugins/data/server/search/search_service.ts

@@ -437,11 +437,7 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
     }
   };
 
-  private cancel = async (
-    deps: SearchStrategyDependencies,
-    id: string,
-    options: ISearchOptions = {}
-  ) => {
+  private cancel = (deps: SearchStrategyDependencies, id: string, options: ISearchOptions = {}) => {
     const strategy = this.getSearchStrategy(options.strategy);
     if (!strategy.cancel) {
       throw new KbnServerError(
@@ -468,14 +464,18 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
   private cancelSessionSearches = async (deps: SearchStrategyDependencies, sessionId: string) => {
     const searchIdMapping = await deps.searchSessionsClient.getSearchIdMapping(sessionId);
     await Promise.allSettled(
-      Array.from(searchIdMapping).map(([searchId, strategyName]) => {
+      Array.from(searchIdMapping).map(async ([searchId, strategyName]) => {
         const searchOptions = {
           sessionId,
           strategy: strategyName,
           isStored: true,
         };
 
-        return this.cancel(deps, searchId, searchOptions);
+        try {
+          await this.cancel(deps, searchId, searchOptions);
+        } catch (e) {
+          this.logger.error(`cancelSessionSearches error: ${e.message}`);
+        }
       })
     );
   };

src/plugins/data/server/search/strategies/eql_search/eql_search_strategy.test.ts

@@ -15,6 +15,7 @@ import { getMockSearchConfig } from '../../../../config.mock';
 
 const getMockEqlResponse = () => ({
   body: {
+    id: 'my-search-id',
     is_partial: false,
     is_running: false,
     took: 162,
@@ -54,13 +55,16 @@ describe('EQL search strategy', () => {
   describe('search()', () => {
     let mockEqlSearch: jest.Mock;
     let mockEqlGet: jest.Mock;
+    let mockEqlDelete: jest.Mock;
     let mockDeps: SearchStrategyDependencies;
     let params: Required<EqlSearchStrategyRequest>['params'];
     let options: Required<EqlSearchStrategyRequest>['options'];
 
     beforeEach(() => {
       mockEqlSearch = jest.fn().mockResolvedValueOnce(getMockEqlResponse());
       mockEqlGet = jest.fn().mockResolvedValueOnce(getMockEqlResponse());
+      mockEqlDelete = jest.fn();
+
       mockDeps = {
         uiSettingsClient: {
           get: jest.fn(),
@@ -70,6 +74,7 @@ describe('EQL search strategy', () => {
             eql: {
               get: mockEqlGet,
               search: mockEqlSearch,
+              delete: mockEqlDelete,
             },
           },
         },
@@ -124,6 +129,34 @@ describe('EQL search strategy', () => {
       });
     });
 
+    it('should delete when aborted', async () => {
+      const response = getMockEqlResponse();
+      mockEqlSearch.mockReset().mockResolvedValueOnce({
+        ...response,
+        body: {
+          ...response.body,
+          is_running: true,
+        },
+      });
+      const eqlSearch = await eqlSearchStrategyProvider(mockSearchConfig, mockLogger);
+      const abortController = new AbortController();
+      const abortSignal = abortController.signal;
+
+      // Abort after an incomplete first response is returned
+      setTimeout(() => abortController.abort(), 100);
+
+      let err: any;
+      try {
+        await eqlSearch.search({ options, params }, { abortSignal }, mockDeps).toPromise();
+      } catch (e) {
+        err = e;
+      }
+
+      expect(mockEqlSearch).toBeCalled();
+      expect(err).not.toBeUndefined();
+      expect(mockEqlDelete).toBeCalled();
+    });
+
     describe('arguments', () => {
       it('sends along async search options', async () => {
         const eqlSearch = await eqlSearchStrategyProvider(mockSearchConfig, mockLogger);

src/plugins/data/server/search/strategies/eql_search/eql_search_strategy.ts

@@ -9,6 +9,7 @@
 import type { TransportResult } from '@elastic/elasticsearch';
 import { tap } from 'rxjs/operators';
 import type { IScopedClusterClient, Logger } from '@kbn/core/server';
+import { getKbnServerError } from '@kbn/kibana-utils-plugin/server';
 import { SearchConfigSchema } from '../../../../config';
 import {
   EqlSearchStrategyRequest,
@@ -27,15 +28,19 @@ export const eqlSearchStrategyProvider = (
   searchConfig: SearchConfigSchema,
   logger: Logger
 ): ISearchStrategy<EqlSearchStrategyRequest, EqlSearchStrategyResponse> => {
-  async function cancelAsyncSearch(id: string, esClient: IScopedClusterClient) {
+  function cancelAsyncSearch(id: string, esClient: IScopedClusterClient) {
     const client = esClient.asCurrentUser.eql;
-    await client.delete({ id });
+    return client.delete({ id });
   }
 
   return {
     cancel: async (id, options, { esClient }) => {
       logger.debug(`_eql/delete ${id}`);
-      await cancelAsyncSearch(id, esClient);
+      try {
+        await cancelAsyncSearch(id, esClient);
+      } catch (e) {
+        throw getKbnServerError(e);
+      }
     },
 
     search: ({ id, ...request }, options: IAsyncSearchOptions, { esClient, uiSettingsClient }) => {
@@ -85,8 +90,15 @@ export const eqlSearchStrategyProvider = (
       };
 
       const cancel = async () => {
-        if (id) {
+        if (!id) return;
+        try {
           await cancelAsyncSearch(id, esClient);
+        } catch (e) {
+          // A 404 means either this search request does not exist, or that it is already cancelled
+          if (e.meta?.statusCode === 404) return;
+
+          // Log all other (unexpected) error messages
+          logger.error(`cancelEqlSearch error: ${e.message}`);
         }
       };
 

src/plugins/data/server/search/strategies/ese_search/ese_search_strategy.test.ts

@@ -136,6 +136,30 @@ describe('ES search strategy', () => {
         expect(request).toHaveProperty('keep_alive', '60000ms');
       });
 
+      it('allows overriding keep_alive and wait_for_completion_timeout', async () => {
+        mockGetCaller.mockResolvedValueOnce(mockAsyncResponse);
+
+        const params = {
+          index: 'logstash-*',
+          body: { query: {} },
+          wait_for_completion_timeout: '10s',
+          keep_alive: '5m',
+        };
+        const esSearch = await enhancedEsSearchStrategyProvider(
+          mockLegacyConfig$,
+          mockSearchConfig,
+          mockLogger
+        );
+
+        await esSearch.search({ id: 'foo', params }, {}, mockDeps).toPromise();
+
+        expect(mockGetCaller).toBeCalled();
+        const request = mockGetCaller.mock.calls[0][0];
+        expect(request.id).toEqual('foo');
+        expect(request).toHaveProperty('wait_for_completion_timeout', '10s');
+        expect(request).toHaveProperty('keep_alive', '5m');
+      });
+
       it('sets transport options on POST requests', async () => {
         const transportOptions = { maxRetries: 1 };
         mockSubmitCaller.mockResolvedValueOnce(mockAsyncResponse);
@@ -260,6 +284,38 @@ describe('ES search strategy', () => {
 
         expect(mockApiCaller).toBeCalledTimes(0);
       });
+
+      it('should delete when aborted', async () => {
+        mockSubmitCaller.mockResolvedValueOnce({
+          ...mockAsyncResponse,
+          body: {
+            ...mockAsyncResponse.body,
+            is_running: true,
+          },
+        });
+
+        const params = { index: 'logstash-*', body: { query: {} } };
+        const esSearch = await enhancedEsSearchStrategyProvider(
+          mockLegacyConfig$,
+          mockSearchConfig,
+          mockLogger
+        );
+        const abortController = new AbortController();
+        const abortSignal = abortController.signal;
+
+        // Abort after an incomplete first response is returned
+        setTimeout(() => abortController.abort(), 100);
+
+        let err: KbnServerError | undefined;
+        try {
+          await esSearch.search({ params }, { abortSignal }, mockDeps).toPromise();
+        } catch (e) {
+          err = e;
+        }
+        expect(mockSubmitCaller).toBeCalled();
+        expect(err).not.toBeUndefined();
+        expect(mockDeleteCaller).toBeCalled();
+      });
     });
 
     describe('with sessionId', () => {
@@ -367,6 +423,44 @@ describe('ES search strategy', () => {
         expect(request).toHaveProperty('wait_for_completion_timeout');
         expect(request).not.toHaveProperty('keep_alive');
       });
+
+      it('should not delete a saved session when aborted', async () => {
+        mockSubmitCaller.mockResolvedValueOnce({
+          ...mockAsyncResponse,
+          body: {
+            ...mockAsyncResponse.body,
+            is_running: true,
+          },
+        });
+
+        const params = { index: 'logstash-*', body: { query: {} } };
+        const esSearch = await enhancedEsSearchStrategyProvider(
+          mockLegacyConfig$,
+          mockSearchConfig,
+          mockLogger
+        );
+        const abortController = new AbortController();
+        const abortSignal = abortController.signal;
+
+        // Abort after an incomplete first response is returned
+        setTimeout(() => abortController.abort(), 100);
+
+        let err: KbnServerError | undefined;
+        try {
+          await esSearch
+            .search(
+              { params },
+              { abortSignal, sessionId: '1', isSearchStored: true, isStored: true },
+              mockDeps
+            )
+            .toPromise();
+        } catch (e) {
+          err = e;
+        }
+        expect(mockSubmitCaller).toBeCalled();
+        expect(err).not.toBeUndefined();
+        expect(mockDeleteCaller).not.toBeCalled();
+      });
     });
 
     it('throws normalized error if ResponseError is thrown', async () => {

src/plugins/data/server/search/strategies/ese_search/ese_search_strategy.ts

@@ -12,6 +12,7 @@ import { catchError, tap } from 'rxjs/operators';
 import type * as estypes from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
 import { firstValueFrom, from } from 'rxjs';
 import { getKbnServerError } from '@kbn/kibana-utils-plugin/server';
+import { IAsyncSearchRequestParams } from '../..';
 import { getKbnSearchError, KbnSearchError } from '../../report_search_error';
 import type { ISearchStrategy, SearchStrategyDependencies } from '../../types';
 import type {
@@ -43,26 +44,28 @@ export const enhancedEsSearchStrategyProvider = (
   logger: Logger,
   usage?: SearchUsage,
   useInternalUser: boolean = false
-): ISearchStrategy => {
-  async function cancelAsyncSearch(id: string, esClient: IScopedClusterClient) {
-    try {
-      const client = useInternalUser ? esClient.asInternalUser : esClient.asCurrentUser;
-      await client.asyncSearch.delete({ id });
-    } catch (e) {
-      throw getKbnServerError(e);
-    }
+): ISearchStrategy<IEsSearchRequest<IAsyncSearchRequestParams>> => {
+  function cancelAsyncSearch(id: string, esClient: IScopedClusterClient) {
+    const client = useInternalUser ? esClient.asInternalUser : esClient.asCurrentUser;
+    return client.asyncSearch.delete({ id });
   }
 
   function asyncSearch(
-    { id, ...request }: IEsSearchRequest,
+    { id, ...request }: IEsSearchRequest<IAsyncSearchRequestParams>,
     options: IAsyncSearchOptions,
     { esClient, uiSettingsClient }: SearchStrategyDependencies
   ) {
     const client = useInternalUser ? esClient.asInternalUser : esClient.asCurrentUser;
 
     const search = async () => {
       const params = id
-        ? getDefaultAsyncGetParams(searchConfig, options)
+        ? {
+            ...getDefaultAsyncGetParams(searchConfig, options),
+            ...(request.params?.keep_alive ? { keep_alive: request.params.keep_alive } : {}),
+            ...(request.params?.wait_for_completion_timeout
+              ? { wait_for_completion_timeout: request.params.wait_for_completion_timeout }
+              : {}),
+          }
         : {
             ...(await getDefaultAsyncSubmitParams(uiSettingsClient, searchConfig, options)),
             ...request.params,
@@ -89,8 +92,15 @@ export const enhancedEsSearchStrategyProvider = (
     };
 
     const cancel = async () => {
-      if (id) {
+      if (!id || options.isStored) return;
+      try {
         await cancelAsyncSearch(id, esClient);
+      } catch (e) {
+        // A 404 means either this search request does not exist, or that it is already cancelled
+        if (e.meta?.statusCode === 404) return;
+
+        // Log all other (unexpected) error messages
+        logger.error(`cancelAsyncSearch error: ${e.message}`);
       }
     };
 
@@ -179,7 +189,11 @@ export const enhancedEsSearchStrategyProvider = (
      */
     cancel: async (id, options, { esClient }) => {
       logger.debug(`cancel ${id}`);
-      await cancelAsyncSearch(id, esClient);
+      try {
+        await cancelAsyncSearch(id, esClient);
+      } catch (e) {
+        throw getKbnServerError(e);
+      }
     },
     /**
      *