Show partial results after search has been canceled

src/platform/packages/shared/kbn-unified-histogram/components/chart/hooks/use_total_hits.ts

@@ -74,8 +74,8 @@ export const useTotalHits = ({
     return () => subscription.unsubscribe();
   }, [fetch, fetch$]);
 
-  const onAbort = useCallback(() => {
-    abortController.current?.abort();
+  const onAbort = useCallback((e: Event) => {
+    abortController.current?.abort((e.target as AbortSignal)?.reason);
   }, []);
 
   useEffect(() => {

src/platform/plugins/shared/dashboard/public/dashboard_app/no_data/dashboard_app_no_data.tsx

@@ -19,6 +19,7 @@ import {
 import { withSuspense } from '@kbn/shared-ux-utility';
 import type { LensSerializedState } from '@kbn/lens-plugin/public';
 import { getLensAttributesFromSuggestion } from '@kbn/visualization-utils';
+import { AbortReason } from '@kbn/kibana-utils-plugin/common';
 import {
   coreServices,
   dataService,
@@ -59,12 +60,12 @@ export const DashboardAppNoDataPage = ({
 
   useEffect(() => {
     return () => {
-      abortController?.abort();
+      abortController?.abort(AbortReason.CLEANUP);
     };
   }, [abortController]);
 
   const onTryESQL = useCallback(async () => {
-    abortController?.abort();
+    abortController?.abort(AbortReason.REPLACED);
     if (lensHelpersAsync.value) {
       const abc = new AbortController();
       const { dataViews } = dataService;

src/platform/plugins/shared/data/common/search/poll_search.test.ts

@@ -8,7 +8,7 @@
  */
 
 import { pollSearch } from './poll_search';
-import { AbortError } from '@kbn/kibana-utils-plugin/common';
+import { AbortError, AbortReason } from '@kbn/kibana-utils-plugin/common';
 
 describe('pollSearch', () => {
   function getMockedSearch$(resolveOnI = 1) {
@@ -86,6 +86,29 @@ describe('pollSearch', () => {
     expect(cancelFn).toBeCalledTimes(1);
   });
 
+  test('Does not throw or cancel if abort reason is CANCELED', async () => {
+    const searchFn = jest.fn().mockResolvedValue({
+      isRunning: false,
+      isPartial: false,
+      rawResponse: {},
+    });
+    const cancelFn = jest.fn();
+
+    const abortController = new AbortController();
+    setTimeout(() => abortController.abort(AbortReason.CANCELED), 100);
+
+    // Poll should complete without throwing AbortError and without calling cancel
+    await expect(
+      pollSearch(searchFn, cancelFn, { abortSignal: abortController.signal }).toPromise()
+    ).resolves.toEqual({
+      isRunning: false,
+      isPartial: false,
+      rawResponse: {},
+    });
+
+    expect(cancelFn).not.toHaveBeenCalled();
+  });
+
   test('Does not leak unresolved promises on cancel', async () => {
     const searchFn = getMockedSearch$(20);
     const cancelFn = jest.fn().mockRejectedValueOnce({ error: 'Oh no!' });

src/platform/plugins/shared/data/common/search/poll_search.ts

@@ -8,9 +8,20 @@
  */
 
 import type { Observable } from 'rxjs';
-import { from, timer, defer, fromEvent, EMPTY } from 'rxjs';
-import { expand, map, switchMap, takeUntil, takeWhile, tap } from 'rxjs';
-import { AbortError } from '@kbn/kibana-utils-plugin/common';
+import {
+  defer,
+  EMPTY,
+  expand,
+  from,
+  fromEvent,
+  switchMap,
+  takeUntil,
+  takeWhile,
+  tap,
+  throwError,
+  timer,
+} from 'rxjs';
+import { AbortError, AbortReason } from '@kbn/kibana-utils-plugin/common';
 import type { IKibanaSearchResponse } from '@kbn/search-types';
 import type { IAsyncSearchOptions } from '..';
 import { isAbortResponse, isRunningResponse } from '..';
@@ -54,9 +65,11 @@ export const pollSearch = <Response extends IKibanaSearchResponse>(
     }
 
     const aborted$ = (abortSignal ? fromEvent(abortSignal, 'abort') : EMPTY).pipe(
-      map(() => {
-        throw new AbortError();
-      })
+      switchMap((e) =>
+        (e.target as AbortSignal)?.reason === AbortReason.CANCELED
+          ? EMPTY
+          : throwError(new AbortError())
+      )
     );
 
     return from(search()).pipe(

src/platform/plugins/shared/data/public/search/constants.ts

@@ -7,6 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
+export const EVENT_TYPE_DATA_SEARCH_CANCEL = 'data_search_cancel';
 export const EVENT_TYPE_DATA_SEARCH_TIMEOUT = 'data_search_timeout';
 export const EVENT_PROPERTY_SEARCH_TIMEOUT_MS = 'timeout_ms';
 export const EVENT_PROPERTY_EXECUTION_CONTEXT = 'execution_context';

src/platform/plugins/shared/data/public/search/search_interceptor/search_abort_controller.test.ts

@@ -7,6 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
+import { AbortReason } from '@kbn/kibana-utils-plugin/common';
 import { SearchAbortController } from './search_abort_controller';
 
 const timeTravel = (msToRun = 0) => {
@@ -45,10 +46,21 @@ describe('search abort controller', () => {
     const controller2 = new AbortController();
     sac.addAbortSignal(controller2.signal);
     expect(sac.getSignal().aborted).toBe(false);
-    controller.abort();
+    controller.abort(AbortReason.CANCELED);
     expect(sac.getSignal().aborted).toBe(false);
-    controller2.abort();
-    expect(sac.getSignal().aborted).toBe(true);
+    controller2.abort(AbortReason.CANCELED);
+    const signal = sac.getSignal();
+    expect(signal.aborted).toBe(true);
+    expect(signal.reason).toBe(AbortReason.CANCELED);
+  });
+
+  test('isCanceled', () => {
+    const sac1 = new SearchAbortController();
+    sac1.abort(AbortReason.CANCELED);
+    expect(sac1.isCanceled()).toBe(true);
+    const sac2 = new SearchAbortController();
+    sac2.abort(AbortReason.TIMEOUT);
+    expect(sac2.isCanceled()).toBe(false);
   });
 
   test('aborts explicitly even if all inputs are not aborted', () => {

src/platform/plugins/shared/data/public/search/search_interceptor/search_abort_controller.ts

@@ -9,23 +9,18 @@
 
 import type { Subscription } from 'rxjs';
 import { timer } from 'rxjs';
-
-export enum AbortReason {
-  Timeout = 'timeout',
-}
+import { AbortReason } from '@kbn/kibana-utils-plugin/common';
 
 export class SearchAbortController {
   private inputAbortSignals: AbortSignal[] = new Array();
   private abortController: AbortController = new AbortController();
   private timeoutSub?: Subscription;
   private destroyed = false;
-  private reason?: AbortReason;
 
   constructor(timeout?: number) {
     if (timeout) {
       this.timeoutSub = timer(timeout).subscribe(() => {
-        this.reason = AbortReason.Timeout;
-        this.abortController.abort();
+        this.abortController.abort(AbortReason.TIMEOUT);
         this.timeoutSub!.unsubscribe();
       });
     }
@@ -34,7 +29,7 @@ export class SearchAbortController {
   private abortHandler = () => {
     const allAborted = this.inputAbortSignals.every((signal) => signal.aborted);
     if (allAborted) {
-      this.abortController.abort();
+      this.abortController.abort(this.inputAbortSignals[0].reason);
       this.cleanup();
     }
   };
@@ -67,12 +62,16 @@ export class SearchAbortController {
     return this.abortController.signal;
   }
 
-  public abort() {
+  public abort(reason?: AbortReason) {
     this.cleanup();
-    this.abortController.abort();
+    this.abortController.abort(reason);
   }
 
   public isTimeout() {
-    return this.reason === AbortReason.Timeout;
+    return this.abortController.signal.reason === AbortReason.TIMEOUT;
+  }
+
+  public isCanceled() {
+    return this.abortController.signal.reason === AbortReason.CANCELED;
   }
 }

src/platform/plugins/shared/data/public/search/search_interceptor/search_interceptor.test.ts

@@ -17,16 +17,15 @@ import { AbortError } from '@kbn/kibana-utils-plugin/public';
 import { EsError, type IEsError } from '@kbn/search-errors';
 import type { ISessionService } from '..';
 import { SearchSessionState } from '..';
-
 import * as searchPhaseException from '../../../common/search/test_data/search_phase_execution_exception.json';
 import * as resourceNotFoundException from '../../../common/search/test_data/resource_not_found_exception.json';
 import { BehaviorSubject } from 'rxjs';
 import { dataPluginMock } from '../../mocks';
 import { UI_SETTINGS } from '../../../common';
+import { AbortReason } from '@kbn/kibana-utils-plugin/common';
 import type { SearchServiceStartDependencies } from '../search_service';
 import type { Start as InspectorStart } from '@kbn/inspector-plugin/public';
 import { SearchTimeoutError, TimeoutErrorMode } from './timeout_error';
-
 import { SearchSessionIncompleteWarning } from './search_session_incomplete_warning';
 import { getMockSearchConfig } from '../../../config.mock';
 
@@ -1159,7 +1158,7 @@ describe('SearchInterceptor', () => {
         const abort = sessionService.trackSearch.mock.calls[0][0].abort;
         expect(abort).toBeInstanceOf(Function);
 
-        abort();
+        abort(AbortReason.REPLACED);
 
         await timeTravel(10);
 

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

@@ -10,18 +10,23 @@
 import { v4 as uuidv4 } from 'uuid';
 import { memoize, once } from 'lodash';
 import type { Observable, Subscription } from 'rxjs';
-import { BehaviorSubject, EMPTY, from, fromEvent, of, throwError } from 'rxjs';
 import {
+  BehaviorSubject,
   catchError,
+  EMPTY,
   filter,
   finalize,
+  from,
+  fromEvent,
   map,
+  of,
   shareReplay,
   skip,
   switchMap,
   take,
   takeUntil,
   tap,
+  throwError,
 } from 'rxjs';
 import type { estypes } from '@elastic/elasticsearch';
 import type {
@@ -51,14 +56,16 @@ import { toMountPoint } from '@kbn/react-kibana-mount';
 import type { KibanaServerError } from '@kbn/kibana-utils-plugin/public';
 import { AbortError } from '@kbn/kibana-utils-plugin/public';
 import type {
-  SanitizedConnectionRequestParams,
   IKibanaSearchRequest,
+  IKibanaSearchResponse,
+  ISearchOptions,
   ISearchOptionsSerializable,
+  SanitizedConnectionRequestParams,
 } from '@kbn/search-types';
 import { createEsError, isEsError, renderSearchError } from '@kbn/search-errors';
-import type { IKibanaSearchResponse, ISearchOptions } from '@kbn/search-types';
-import { defaultFreeze } from '@kbn/kibana-utils-plugin/common';
+import { AbortReason, defaultFreeze } from '@kbn/kibana-utils-plugin/common';
 import {
+  EVENT_TYPE_DATA_SEARCH_CANCEL,
   EVENT_TYPE_DATA_SEARCH_TIMEOUT,
   EVENT_PROPERTY_SEARCH_TIMEOUT_MS,
   EVENT_PROPERTY_EXECUTION_CONTEXT,
@@ -318,7 +325,7 @@ export class SearchInterceptor {
 
     const searchTracker = this.deps.session.isCurrentSession(sessionId)
       ? this.deps.session.trackSearch({
-          abort: () => searchAbortController.abort(),
+          abort: (reason?: AbortReason) => searchAbortController.abort(reason),
           poll: async (abortSignal) => {
             if (id) {
               await search({ abortSignal });
@@ -356,8 +363,14 @@ export class SearchInterceptor {
     );
 
     const cancel = async () => {
-      // If the request times out, we handle cancellation after we make the last call to retrieve the results
-      if (!id || isSavedToBackground || searchAbortController.isTimeout()) return;
+      // If the request times out/is canceled, we handle cancellation after we make the last call to retrieve the results
+      if (
+        !id ||
+        isSavedToBackground ||
+        searchAbortController.isTimeout() ||
+        searchAbortController.isCanceled()
+      )
+        return;
       try {
         await sendCancelRequest();
       } catch (e) {
@@ -399,14 +412,24 @@ export class SearchInterceptor {
           : response;
       }),
       catchError((e: Error) => {
-        // If we aborted (search:timeout advanced setting) and there was a partial response, return it instead of just erroring out
-        if (searchAbortController.isTimeout()) {
-          this.startRenderServices.analytics.reportEvent(EVENT_TYPE_DATA_SEARCH_TIMEOUT, {
-            [EVENT_PROPERTY_SEARCH_TIMEOUT_MS]: this.searchTimeout,
-            [EVENT_PROPERTY_EXECUTION_CONTEXT]: options.executionContext,
-          });
+        // If we aborted (search:timeout advanced setting) or the user canceled and there was a partial response, return it instead of just erroring out
+        if (id && (searchAbortController.isTimeout() || searchAbortController.isCanceled())) {
+          this.startRenderServices.analytics.reportEvent(
+            searchAbortController.isTimeout()
+              ? EVENT_TYPE_DATA_SEARCH_TIMEOUT
+              : EVENT_TYPE_DATA_SEARCH_CANCEL,
+            {
+              ...(searchAbortController.isTimeout() && {
+                [EVENT_PROPERTY_SEARCH_TIMEOUT_MS]: this.searchTimeout,
+              }),
+              [EVENT_PROPERTY_EXECUTION_CONTEXT]: options.executionContext,
+            }
+          );
           return from(
-            this.runSearch({ id, ...request }, { ...options, retrieveResults: true })
+            this.runSearch(
+              { id, ...request },
+              { ...options, abortSignal: new AbortController().signal, retrieveResults: true }
+            )
           ).pipe(
             map((response) =>
               options.strategy === ENHANCED_ES_SEARCH_STRATEGY
@@ -415,7 +438,9 @@ export class SearchInterceptor {
             ),
             tap(async () => {
               await sendCancelRequest();
-              this.handleSearchError(e, request?.params?.body ?? {}, options, true);
+              if (searchAbortController.isTimeout()) {
+                this.handleSearchError(e, request?.params?.body ?? {}, options, true);
+              }
             })
           );
         } else {
@@ -615,9 +640,11 @@ export class SearchInterceptor {
         // Abort the replay if the abortSignal is aborted.
         // The underlaying search will not abort unless searchAbortController fires.
         const aborted$ = (abortSignal ? fromEvent(abortSignal, 'abort') : EMPTY).pipe(
-          map(() => {
-            throw new AbortError();
-          })
+          switchMap((e) =>
+            (e.target as AbortSignal)?.reason === AbortReason.CANCELED
+              ? EMPTY
+              : throwError(new AbortError())
+          )
         );
 
         return response$.pipe(