UI to stop async searches

src/plugins/data/public/mocks.ts

@@ -17,13 +17,10 @@
  * under the License.
  */
 
-// eslint-disable-next-line @kbn/eslint/no-restricted-paths
-import { coreMock } from '../../../../src/core/public/mocks';
 import { Plugin, DataPublicPluginSetup, DataPublicPluginStart, IndexPatternsContract } from '.';
 import { fieldFormatsMock } from '../common/field_formats/mocks';
-import { searchSetupMock } from './search/mocks';
+import { searchSetupMock, searchStartMock } from './search/mocks';
 import { queryServiceMock } from './query/mocks';
-import { getCalculateAutoTimeExpression } from './search/aggs/buckets/lib/date_utils';
 
 export type Setup = jest.Mocked<ReturnType<Plugin['setup']>>;
 export type Start = jest.Mocked<ReturnType<Plugin['start']>>;
@@ -36,52 +33,25 @@ const autocompleteMock: any = {
 
 const createSetupContract = (): Setup => {
   const querySetupMock = queryServiceMock.createSetupContract();
-  const setupContract = {
+  return {
     autocomplete: autocompleteMock,
     search: searchSetupMock,
     fieldFormats: fieldFormatsMock as DataPublicPluginSetup['fieldFormats'],
     query: querySetupMock,
-    __LEGACY: {
-      esClient: {
-        search: jest.fn(),
-        msearch: jest.fn(),
-      },
-    },
   };
-
-  return setupContract;
 };
 
 const createStartContract = (): Start => {
-  const coreStart = coreMock.createStart();
   const queryStartMock = queryServiceMock.createStartContract();
-  const startContract = {
+  return {
     autocomplete: autocompleteMock,
-    getSuggestions: jest.fn(),
-    search: {
-      aggs: {
-        calculateAutoTimeExpression: getCalculateAutoTimeExpression(coreStart.uiSettings),
-      },
-      search: jest.fn(),
-      __LEGACY: {
-        esClient: {
-          search: jest.fn(),
-          msearch: jest.fn(),
-        },
-      },
-    },
+    search: searchStartMock,
     fieldFormats: fieldFormatsMock as DataPublicPluginStart['fieldFormats'],
     query: queryStartMock,
     ui: {
       IndexPatternSelect: jest.fn(),
       SearchBar: jest.fn(),
     },
-    __LEGACY: {
-      esClient: {
-        search: jest.fn(),
-        msearch: jest.fn(),
-      },
-    },
     indexPatterns: ({
       make: () => ({
         fieldsFetcher: {
@@ -91,7 +61,6 @@ const createStartContract = (): Start => {
       get: jest.fn().mockReturnValue(Promise.resolve({})),
     } as unknown) as IndexPatternsContract,
   };
-  return startContract;
 };
 
 export { searchSourceMock } from './search/mocks';

src/plugins/data/public/search/mocks.ts

@@ -23,10 +23,27 @@ import { getCalculateAutoTimeExpression } from './aggs/buckets/lib/date_utils';
 
 export * from './search_source/mocks';
 
+const coreStart = coreMock.createStart();
+
 export const searchSetupMock = {
   aggs: {
     calculateAutoTimeExpression: getCalculateAutoTimeExpression(coreMock.createSetup().uiSettings),
   },
   registerSearchStrategyContext: jest.fn(),
   registerSearchStrategyProvider: jest.fn(),
 };
+
+export const searchStartMock = {
+  cancelPendingSearches: jest.fn(),
+  getPendingSearchesCount$: jest.fn(),
+  aggs: {
+    calculateAutoTimeExpression: getCalculateAutoTimeExpression(coreStart.uiSettings),
+  },
+  search: jest.fn(),
+  __LEGACY: {
+    esClient: {
+      search: jest.fn(),
+      msearch: jest.fn(),
+    },
+  },
+};

src/plugins/data/public/search/search_interceptor.ts

@@ -0,0 +1,82 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { Subject } from 'rxjs';
+import { finalize } from 'rxjs/operators';
+import { IKibanaSearchRequest } from '../../common/search';
+import { ISearch, ISearchOptions } from './i_search';
+
+export class SearchInterceptor {
+  /**
+   * `abortController` used to signal all searches to abort.
+   */
+  private abortController = new AbortController();
+
+  /**
+   * Number of in-progress search requests.
+   */
+  private pendingCount: number = 0;
+
+  /**
+   * Observable that emits when the number of pending requests changes.
+   */
+  private pendingCount$: Subject<number> = new Subject();
+
+  /**
+   * Abort our `AbortController`, which in turn aborts any intercepted searches.
+   */
+  cancelPending = () => {
+    this.abortController.abort();
+    this.abortController = new AbortController();
+  };
+
+  /**
+   * Searches using the given `search` method. Overrides the `AbortSignal` with one that will abort
+   * either when `cancelPending` is called, or when the original `AbortSignal` is aborted. Updates
+   * the `pendingCount` when the request is started/finalized.
+   */
+  search = (search: ISearch<any>, request: IKibanaSearchRequest, options?: ISearchOptions) => {
+    this.pendingCount$.next(++this.pendingCount);
+
+    // Create a new `AbortController` that will abort when our either our private `AbortController`
+    // aborts, or the given `AbortSignal` aborts.
+    const abortController = new AbortController();
+    if (options?.signal) {
+      options.signal.addEventListener('abort', () => {
+        abortController.abort();
+      });
+    }
+    this.abortController.signal.addEventListener('abort', () => {
+      abortController.abort();
+    });
+    const { signal } = abortController;
+
+    return search(request as any, { ...options, signal }).pipe(
+      finalize(() => this.pendingCount$.next(--this.pendingCount))
+    );
+  };
+
+  /**
+   * Returns the current number of pending searches. This could mean either one of the search
+   * requests is still in flight, or that it has only received partial responses.
+   */
+  getPendingCount$ = () => {
+    return this.pendingCount$.asObservable();
+  };
+}

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

@@ -18,14 +18,14 @@
  */
 
 import { Plugin, CoreSetup, CoreStart, PackageInfo } from '../../../../core/public';
-
 import { getCalculateAutoTimeExpression } from './aggs/buckets/lib/date_utils';
 import { SYNC_SEARCH_STRATEGY, syncSearchStrategyProvider } from './sync_search_strategy';
 import { ISearchSetup, ISearchStart, TSearchStrategyProvider, TSearchStrategiesMap } from './types';
 import { TStrategyTypes } from './strategy_types';
 import { getEsClient, LegacyApiCaller } from './es_client';
 import { ES_SEARCH_STRATEGY, DEFAULT_SEARCH_STRATEGY } from '../../common/search';
-import { esSearchStrategyProvider } from './es_search/es_search_strategy';
+import { esSearchStrategyProvider } from './es_search';
+import { SearchInterceptor } from './search_interceptor';
 
 /**
  * The search plugin exposes two registration methods for other plugins:
@@ -74,7 +74,17 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
   }
 
   public start(core: CoreStart): ISearchStart {
+    /**
+     * A global object that intercepts all searches and provides convenience methods for cancelling
+     * all pending search requests, as well as getting the number of pending search requests.
+     * TODO: Make this modular so that apps can opt in/out of search collection, or even provide
+     * their own search collector instances
+     */
+    const searchInterceptor = new SearchInterceptor();
+
     return {
+      cancelPendingSearches: () => searchInterceptor.cancelPending(),
+      getPendingSearchesCount$: () => searchInterceptor.getPendingCount$(),
       aggs: {
         calculateAutoTimeExpression: getCalculateAutoTimeExpression(core.uiSettings),
       },
@@ -84,7 +94,7 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
           core,
           getSearchStrategy: this.getSearchStrategy,
         });
-        return search(request as any, options);
+        return searchInterceptor.search(search, request, options);
       },
       __LEGACY: {
         esClient: this.esClient!,

src/plugins/data/public/search/search_source/search_source.test.ts

@@ -17,29 +17,19 @@
  * under the License.
  */
 
-import { SearchSource } from '../search_source';
+import { SearchSource } from './search_source';
 import { IndexPattern } from '../..';
 import { setSearchService, setUiSettings, setInjectedMetadata } from '../../services';
 
 import {
   injectedMetadataServiceMock,
   uiSettingsServiceMock,
 } from '../../../../../core/public/mocks';
+import { searchStartMock } from '../mocks';
 
 setUiSettings(uiSettingsServiceMock.createStartContract());
 setInjectedMetadata(injectedMetadataServiceMock.createSetupContract());
-setSearchService({
-  aggs: {
-    calculateAutoTimeExpression: jest.fn().mockReturnValue('1d'),
-  },
-  search: jest.fn(),
-  __LEGACY: {
-    esClient: {
-      search: jest.fn(),
-      msearch: jest.fn(),
-    },
-  },
-});
+setSearchService(searchStartMock);
 
 jest.mock('../fetch', () => ({
   fetchSoon: jest.fn().mockResolvedValue({}),

src/plugins/data/public/search/search_strategy/default_search_strategy.test.ts

@@ -20,6 +20,7 @@
 import { IUiSettingsClient } from '../../../../../core/public';
 import { SearchStrategySearchParams } from './types';
 import { defaultSearchStrategy } from './default_search_strategy';
+import { searchStartMock } from '../mocks';
 
 const { search } = defaultSearchStrategy;
 
@@ -55,25 +56,24 @@ describe('defaultSearchStrategy', function() {
       searchMockResponse.abort.mockClear();
       searchMock.mockClear();
 
+      const searchService = searchStartMock;
+      searchService.aggs.calculateAutoTimeExpression = jest.fn().mockReturnValue('1d');
+      searchService.search = newSearchMock;
+      searchService.__LEGACY = {
+        esClient: {
+          search: searchMock,
+          msearch: msearchMock,
+        },
+      };
+
       searchArgs = {
         searchRequests: [
           {
             index: { title: 'foo' },
           },
         ],
         esShardTimeout: 0,
-        searchService: {
-          aggs: {
-            calculateAutoTimeExpression: jest.fn().mockReturnValue('1d'),
-          },
-          search: newSearchMock,
-          __LEGACY: {
-            esClient: {
-              search: searchMock,
-              msearch: msearchMock,
-            },
-          },
-        },
+        searchService,
       };
 
       es = searchArgs.searchService.__LEGACY.esClient;

src/plugins/data/public/search/types.ts

@@ -18,6 +18,7 @@
  */
 
 import { CoreStart } from 'kibana/public';
+import { Observable } from 'rxjs';
 import { TimeRange } from '../../common';
 import { ISearch, ISearchGeneric } from './i_search';
 import { TStrategyTypes } from './strategy_types';
@@ -91,6 +92,8 @@ export interface ISearchSetup {
 export interface ISearchStart {
   aggs: SearchAggsStart;
   search: ISearchGeneric;
+  cancelPendingSearches: () => void;
+  getPendingSearchesCount$: () => Observable<number>;
   __LEGACY: {
     esClient: LegacyApiCaller;
   };