[Search Sessions] Apply awaits to avoid unhandled errors (#90593)

* Apply awaits to avoid unhandled errors

* catch and ignore tracking error

* added reject test for search service

* Improve search service api test coverage

Co-authored-by: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>

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

@@ -7,7 +7,7 @@
  */
 
 import type { MockedKeys } from '@kbn/utility-types/jest';
-import { CoreSetup, CoreStart } from '../../../../core/server';
+import { CoreSetup, CoreStart, SavedObject } from '../../../../core/server';
 import { coreMock } from '../../../../core/server/mocks';
 
 import { DataPluginStart } from '../plugin';
@@ -86,13 +86,22 @@ describe('Search service', () => {
   describe('asScopedProvider', () => {
     let mockScopedClient: IScopedSearchClient;
     let searcPluginStart: ISearchStart<IEsSearchRequest, IEsSearchResponse<any>>;
-    let mockStrategy: jest.Mocked<ISearchStrategy>;
+    let mockStrategy: any;
+    let mockStrategyNoCancel: jest.Mocked<ISearchStrategy>;
     let mockSessionService: ISearchSessionService<any>;
     let mockSessionClient: jest.Mocked<IScopedSearchSessionsClient>;
     const sessionId = '1234';
 
     beforeEach(() => {
-      mockStrategy = { search: jest.fn().mockReturnValue(of({})) };
+      mockStrategy = {
+        search: jest.fn().mockReturnValue(of({})),
+        cancel: jest.fn(),
+        extend: jest.fn(),
+      };
+
+      mockStrategyNoCancel = {
+        search: jest.fn().mockReturnValue(of({})),
+      };
 
       mockSessionClient = createSearchSessionsClientMock();
       mockSessionService = {
@@ -104,6 +113,7 @@ describe('Search service', () => {
         expressions: expressionsPluginMock.createSetupContract(),
       });
       pluginSetup.registerSearchStrategy('es', mockStrategy);
+      pluginSetup.registerSearchStrategy('nocancel', mockStrategyNoCancel);
       pluginSetup.__enhance({
         defaultStrategy: 'es',
         sessionService: mockSessionService,
@@ -123,7 +133,7 @@ describe('Search service', () => {
       it('searches using the original request if not restoring, trackId is not called if there is no id in the response', async () => {
         const searchRequest = { params: {} };
         const options = { sessionId, isStored: false, isRestore: false };
-        mockSessionClient.trackId = jest.fn();
+        mockSessionClient.trackId = jest.fn().mockResolvedValue(undefined);
 
         mockStrategy.search.mockReturnValue(
           of({
@@ -165,10 +175,27 @@ describe('Search service', () => {
         expect(request).toStrictEqual({ ...searchRequest, id: 'my_id' });
       });
 
+      it('does not fail if `trackId` throws', async () => {
+        const searchRequest = { params: {} };
+        const options = { sessionId, isStored: false, isRestore: false };
+        mockSessionClient.trackId = jest.fn().mockRejectedValue(undefined);
+
+        mockStrategy.search.mockReturnValue(
+          of({
+            id: 'my_id',
+            rawResponse: {} as any,
+          })
+        );
+
+        await mockScopedClient.search(searchRequest, options).toPromise();
+
+        expect(mockSessionClient.trackId).toBeCalledTimes(1);
+      });
+
       it('calls `trackId` for every response, if the response contains an `id` and not restoring', async () => {
         const searchRequest = { params: {} };
         const options = { sessionId, isStored: false, isRestore: false };
-        mockSessionClient.trackId = jest.fn();
+        mockSessionClient.trackId = jest.fn().mockResolvedValue(undefined);
 
         mockStrategy.search.mockReturnValue(
           of(
@@ -195,7 +222,7 @@ describe('Search service', () => {
         const searchRequest = { params: {} };
         const options = { sessionId, isStored: true, isRestore: true };
         mockSessionClient.getId = jest.fn().mockResolvedValueOnce('my_id');
-        mockSessionClient.trackId = jest.fn();
+        mockSessionClient.trackId = jest.fn().mockResolvedValue(undefined);
 
         await mockScopedClient.search(searchRequest, options).toPromise();
 
@@ -206,12 +233,258 @@ describe('Search service', () => {
         const searchRequest = { params: {} };
         const options = {};
         mockSessionClient.getId = jest.fn().mockResolvedValueOnce('my_id');
-        mockSessionClient.trackId = jest.fn();
+        mockSessionClient.trackId = jest.fn().mockResolvedValue(undefined);
 
         await mockScopedClient.search(searchRequest, options).toPromise();
 
         expect(mockSessionClient.trackId).not.toBeCalled();
       });
     });
+
+    describe('cancelSession', () => {
+      const mockSavedObject: SavedObject = {
+        id: 'd7170a35-7e2c-48d6-8dec-9a056721b489',
+        type: 'search-session',
+        attributes: {
+          name: 'my_name',
+          appId: 'my_app_id',
+          urlGeneratorId: 'my_url_generator_id',
+          idMapping: {},
+        },
+        references: [],
+      };
+
+      it('cancels a saved object with no search ids', async () => {
+        mockSessionClient.getSearchIdMapping = jest
+          .fn()
+          .mockResolvedValue(new Map<string, string>());
+        mockSessionClient.cancel = jest.fn().mockResolvedValue(mockSavedObject);
+        const cancelSpy = jest.spyOn(mockScopedClient, 'cancel');
+
+        await mockScopedClient.cancelSession('123');
+
+        expect(mockSessionClient.cancel).toHaveBeenCalledTimes(1);
+        expect(cancelSpy).not.toHaveBeenCalled();
+      });
+
+      it('cancels a saved object and search ids', async () => {
+        const mockMap = new Map<string, string>();
+        mockMap.set('abc', 'es');
+        mockSessionClient.getSearchIdMapping = jest.fn().mockResolvedValue(mockMap);
+        mockStrategy.cancel = jest.fn();
+        mockSessionClient.cancel = jest.fn().mockResolvedValue(mockSavedObject);
+
+        await mockScopedClient.cancelSession('123');
+
+        expect(mockSessionClient.cancel).toHaveBeenCalledTimes(1);
+
+        const [searchId, options] = mockStrategy.cancel.mock.calls[0];
+        expect(mockStrategy.cancel).toHaveBeenCalledTimes(1);
+        expect(searchId).toBe('abc');
+        expect(options).toHaveProperty('strategy', 'es');
+      });
+
+      it('cancels a saved object with some strategies that dont support cancellation, dont throw an error', async () => {
+        const mockMap = new Map<string, string>();
+        mockMap.set('abc', 'nocancel');
+        mockMap.set('def', 'es');
+        mockSessionClient.getSearchIdMapping = jest.fn().mockResolvedValue(mockMap);
+        mockStrategy.cancel = jest.fn();
+        mockSessionClient.cancel = jest.fn().mockResolvedValue(mockSavedObject);
+
+        await mockScopedClient.cancelSession('123');
+
+        expect(mockSessionClient.cancel).toHaveBeenCalledTimes(1);
+
+        const [searchId, options] = mockStrategy.cancel.mock.calls[0];
+        expect(mockStrategy.cancel).toHaveBeenCalledTimes(1);
+        expect(searchId).toBe('def');
+        expect(options).toHaveProperty('strategy', 'es');
+      });
+
+      it('cancels a saved object with some strategies that dont exist, dont throw an error', async () => {
+        const mockMap = new Map<string, string>();
+        mockMap.set('abc', 'notsupported');
+        mockMap.set('def', 'es');
+        mockSessionClient.getSearchIdMapping = jest.fn().mockResolvedValue(mockMap);
+        mockStrategy.cancel = jest.fn();
+        mockSessionClient.cancel = jest.fn().mockResolvedValue(mockSavedObject);
+
+        await mockScopedClient.cancelSession('123');
+
+        expect(mockSessionClient.cancel).toHaveBeenCalledTimes(1);
+
+        const [searchId, options] = mockStrategy.cancel.mock.calls[0];
+        expect(mockStrategy.cancel).toHaveBeenCalledTimes(1);
+        expect(searchId).toBe('def');
+        expect(options).toHaveProperty('strategy', 'es');
+      });
+    });
+
+    describe('deleteSession', () => {
+      const mockSavedObject: SavedObject = {
+        id: 'd7170a35-7e2c-48d6-8dec-9a056721b489',
+        type: 'search-session',
+        attributes: {
+          name: 'my_name',
+          appId: 'my_app_id',
+          urlGeneratorId: 'my_url_generator_id',
+          idMapping: {},
+        },
+        references: [],
+      };
+
+      it('deletes a saved object with no search ids', async () => {
+        mockSessionClient.getSearchIdMapping = jest
+          .fn()
+          .mockResolvedValue(new Map<string, string>());
+        mockSessionClient.delete = jest.fn().mockResolvedValue(mockSavedObject);
+        const cancelSpy = jest.spyOn(mockScopedClient, 'cancel');
+
+        await mockScopedClient.deleteSession('123');
+
+        expect(mockSessionClient.delete).toHaveBeenCalledTimes(1);
+        expect(cancelSpy).not.toHaveBeenCalled();
+      });
+
+      it('deletes a saved object and search ids', async () => {
+        const mockMap = new Map<string, string>();
+        mockMap.set('abc', 'es');
+        mockSessionClient.getSearchIdMapping = jest.fn().mockResolvedValue(mockMap);
+        mockSessionClient.delete = jest.fn().mockResolvedValue(mockSavedObject);
+        mockStrategy.cancel = jest.fn();
+
+        await mockScopedClient.deleteSession('123');
+
+        expect(mockSessionClient.delete).toHaveBeenCalledTimes(1);
+
+        const [searchId, options] = mockStrategy.cancel.mock.calls[0];
+        expect(mockStrategy.cancel).toHaveBeenCalledTimes(1);
+        expect(searchId).toBe('abc');
+        expect(options).toHaveProperty('strategy', 'es');
+      });
+
+      it('deletes a saved object with some strategies that dont support cancellation, dont throw an error', async () => {
+        const mockMap = new Map<string, string>();
+        mockMap.set('abc', 'nocancel');
+        mockMap.set('def', 'es');
+        mockSessionClient.getSearchIdMapping = jest.fn().mockResolvedValue(mockMap);
+        mockSessionClient.delete = jest.fn().mockResolvedValue(mockSavedObject);
+        mockStrategy.cancel = jest.fn();
+
+        await mockScopedClient.deleteSession('123');
+
+        expect(mockSessionClient.delete).toHaveBeenCalledTimes(1);
+
+        const [searchId, options] = mockStrategy.cancel.mock.calls[0];
+        expect(mockStrategy.cancel).toHaveBeenCalledTimes(1);
+        expect(searchId).toBe('def');
+        expect(options).toHaveProperty('strategy', 'es');
+      });
+
+      it('deletes a saved object with some strategies that dont exist, dont throw an error', async () => {
+        const mockMap = new Map<string, string>();
+        mockMap.set('abc', 'notsupported');
+        mockMap.set('def', 'es');
+        mockSessionClient.getSearchIdMapping = jest.fn().mockResolvedValue(mockMap);
+        mockStrategy.cancel = jest.fn();
+        mockSessionClient.delete = jest.fn().mockResolvedValue(mockSavedObject);
+
+        await mockScopedClient.deleteSession('123');
+
+        expect(mockSessionClient.delete).toHaveBeenCalledTimes(1);
+
+        const [searchId, options] = mockStrategy.cancel.mock.calls[0];
+        expect(mockStrategy.cancel).toHaveBeenCalledTimes(1);
+        expect(searchId).toBe('def');
+        expect(options).toHaveProperty('strategy', 'es');
+      });
+    });
+
+    describe('extendSession', () => {
+      const mockSavedObject: SavedObject = {
+        id: 'd7170a35-7e2c-48d6-8dec-9a056721b489',
+        type: 'search-session',
+        attributes: {
+          name: 'my_name',
+          appId: 'my_app_id',
+          urlGeneratorId: 'my_url_generator_id',
+          idMapping: {},
+        },
+        references: [],
+      };
+
+      it('extends a saved object with no search ids', async () => {
+        mockSessionClient.getSearchIdMapping = jest
+          .fn()
+          .mockResolvedValue(new Map<string, string>());
+        mockSessionClient.extend = jest.fn().mockResolvedValue(mockSavedObject);
+        mockStrategy.extend = jest.fn();
+
+        await mockScopedClient.extendSession('123', new Date('2020-01-01'));
+
+        expect(mockSessionClient.extend).toHaveBeenCalledTimes(1);
+        expect(mockStrategy.extend).not.toHaveBeenCalled();
+      });
+
+      it('extends a saved object and search ids', async () => {
+        const mockMap = new Map<string, string>();
+        mockMap.set('abc', 'es');
+        mockSessionClient.getSearchIdMapping = jest.fn().mockResolvedValue(mockMap);
+        mockSessionClient.extend = jest.fn().mockResolvedValue(mockSavedObject);
+        mockStrategy.extend = jest.fn();
+
+        await mockScopedClient.extendSession('123', new Date('2020-01-01'));
+
+        expect(mockSessionClient.extend).toHaveBeenCalledTimes(1);
+        expect(mockStrategy.extend).toHaveBeenCalledTimes(1);
+        const [searchId, keepAlive, options] = mockStrategy.extend.mock.calls[0];
+        expect(searchId).toBe('abc');
+        expect(keepAlive).toContain('ms');
+        expect(options).toHaveProperty('strategy', 'es');
+      });
+
+      it('doesnt extend the saved object with some strategies that dont support cancellation, throws an error', async () => {
+        const mockMap = new Map<string, string>();
+        mockMap.set('abc', 'nocancel');
+        mockMap.set('def', 'es');
+        mockSessionClient.getSearchIdMapping = jest.fn().mockResolvedValue(mockMap);
+        mockSessionClient.extend = jest.fn().mockResolvedValue(mockSavedObject);
+        mockStrategy.extend = jest.fn().mockResolvedValue({});
+
+        const extendRes = mockScopedClient.extendSession('123', new Date('2020-01-01'));
+
+        await expect(extendRes).rejects.toThrowError(
+          'Failed to extend the expiration of some searches'
+        );
+
+        expect(mockSessionClient.extend).not.toHaveBeenCalled();
+        const [searchId, keepAlive, options] = mockStrategy.extend.mock.calls[0];
+        expect(searchId).toBe('def');
+        expect(keepAlive).toContain('ms');
+        expect(options).toHaveProperty('strategy', 'es');
+      });
+
+      it('doesnt extend the saved object with some strategies that dont exist, throws an error', async () => {
+        const mockMap = new Map<string, string>();
+        mockMap.set('abc', 'notsupported');
+        mockMap.set('def', 'es');
+        mockSessionClient.getSearchIdMapping = jest.fn().mockResolvedValue(mockMap);
+        mockSessionClient.extend = jest.fn().mockResolvedValue(mockSavedObject);
+        mockStrategy.extend = jest.fn().mockResolvedValue({});
+
+        const extendRes = mockScopedClient.extendSession('123', new Date('2020-01-01'));
+
+        await expect(extendRes).rejects.toThrowError(
+          'Failed to extend the expiration of some searches'
+        );
+
+        expect(mockSessionClient.extend).not.toHaveBeenCalled();
+        const [searchId, keepAlive, options] = mockStrategy.extend.mock.calls[0];
+        expect(searchId).toBe('def');
+        expect(keepAlive).toContain('ms');
+        expect(options).toHaveProperty('strategy', 'es');
+      });
+    });
   });
 });

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

@@ -275,7 +275,10 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
         switchMap((searchRequest) => strategy.search(searchRequest, options, deps)),
         tap((response) => {
           if (!options.sessionId || !response.id || options.isRestore) return;
-          deps.searchSessionsClient.trackId(request, response.id, options);
+          // intentionally swallow tracking error, as it shouldn't fail the search
+          deps.searchSessionsClient.trackId(request, response.id, options).catch((trackErr) => {
+            this.logger.error(trackErr);
+          });
         })
       );
     } catch (e) {
@@ -283,7 +286,11 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
     }
   };
 
-  private cancel = (deps: SearchStrategyDependencies, id: string, options: ISearchOptions = {}) => {
+  private cancel = async (
+    deps: SearchStrategyDependencies,
+    id: string,
+    options: ISearchOptions = {}
+  ) => {
     const strategy = this.getSearchStrategy(options.strategy);
     if (!strategy.cancel) {
       throw new KbnServerError(
@@ -294,7 +301,7 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
     return strategy.cancel(id, options, deps);
   };
 
-  private extend = (
+  private extend = async (
     deps: SearchStrategyDependencies,
     id: string,
     keepAlive: string,
@@ -309,25 +316,26 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
 
   private cancelSessionSearches = async (deps: SearchStrategyDependencies, sessionId: string) => {
     const searchIdMapping = await deps.searchSessionsClient.getSearchIdMapping(sessionId);
-
-    for (const [searchId, strategyName] of searchIdMapping.entries()) {
-      const searchOptions = {
-        sessionId,
-        strategy: strategyName,
-        isStored: true,
-      };
-      this.cancel(deps, searchId, searchOptions);
-    }
+    await Promise.allSettled(
+      Array.from(searchIdMapping).map(([searchId, strategyName]) => {
+        const searchOptions = {
+          sessionId,
+          strategy: strategyName,
+          isStored: true,
+        };
+        return this.cancel(deps, searchId, searchOptions);
+      })
+    );
   };
 
   private cancelSession = async (deps: SearchStrategyDependencies, sessionId: string) => {
     const response = await deps.searchSessionsClient.cancel(sessionId);
-    this.cancelSessionSearches(deps, sessionId);
+    await this.cancelSessionSearches(deps, sessionId);
     return response;
   };
 
   private deleteSession = async (deps: SearchStrategyDependencies, sessionId: string) => {
-    this.cancelSessionSearches(deps, sessionId);
+    await this.cancelSessionSearches(deps, sessionId);
     return deps.searchSessionsClient.delete(sessionId);
   };
 
@@ -339,13 +347,19 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
     const searchIdMapping = await deps.searchSessionsClient.getSearchIdMapping(sessionId);
     const keepAlive = `${moment(expires).diff(moment())}ms`;
 
-    for (const [searchId, strategyName] of searchIdMapping.entries()) {
-      const searchOptions = {
-        sessionId,
-        strategy: strategyName,
-        isStored: true,
-      };
-      await this.extend(deps, searchId, keepAlive, searchOptions);
+    const result = await Promise.allSettled(
+      Array.from(searchIdMapping).map(([searchId, strategyName]) => {
+        const searchOptions = {
+          sessionId,
+          strategy: strategyName,
+          isStored: true,
+        };
+        return this.extend(deps, searchId, keepAlive, searchOptions);
+      })
+    );
+
+    if (result.some((extRes) => extRes.status === 'rejected')) {
+      throw new Error('Failed to extend the expiration of some searches');
     }
 
     return deps.searchSessionsClient.extend(sessionId, expires);