From e1fb5f063e3de55370b792a715606222b3cdb234 Mon Sep 17 00:00:00 2001
From: Lukas Olson <olson.lukas@gmail.com>
Date: Sun, 13 Sep 2020 19:50:58 -0700
Subject: [PATCH] [Search] Re-add support for aborting when a connection is
 closed (#76470) (#77296)

* [Search] Re-add support for aborting when a connection is closed

* Fix types

Co-authored-by: Elastic Machine <elasticmachine@users.noreply.github.com>

Co-authored-by: Elastic Machine <elasticmachine@users.noreply.github.com>
---
 .../search/es_search/es_search_strategy.ts    |  9 ++++---
 .../server/search/es_search_strategy.ts       | 26 +++++++++++++------
 2 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/src/plugins/data/server/search/es_search/es_search_strategy.ts b/src/plugins/data/server/search/es_search/es_search_strategy.ts
index 106f974ed3457..e2ed500689cfa 100644
--- a/src/plugins/data/server/search/es_search/es_search_strategy.ts
+++ b/src/plugins/data/server/search/es_search/es_search_strategy.ts
@@ -52,10 +52,11 @@ export const esSearchStrategyProvider = (
       });
 
       try {
-        const esResponse = (await context.core.elasticsearch.client.asCurrentUser.search(
-          params
-        )) as ApiResponse<SearchResponse<any>>;
-        const rawResponse = esResponse.body;
+        // Temporary workaround until https://github.com/elastic/elasticsearch-js/issues/1297
+        const promise = context.core.elasticsearch.client.asCurrentUser.search(params);
+        if (options?.abortSignal)
+          options.abortSignal.addEventListener('abort', () => promise.abort());
+        const { body: rawResponse } = (await promise) as ApiResponse<SearchResponse<any>>;
 
         if (usage) usage.trackSuccess(rawResponse.took);
 
diff --git a/x-pack/plugins/data_enhanced/server/search/es_search_strategy.ts b/x-pack/plugins/data_enhanced/server/search/es_search_strategy.ts
index eda6178dc8e5b..72ea1f096e8fb 100644
--- a/x-pack/plugins/data_enhanced/server/search/es_search_strategy.ts
+++ b/x-pack/plugins/data_enhanced/server/search/es_search_strategy.ts
@@ -7,6 +7,7 @@
 import { first } from 'rxjs/operators';
 import { SearchResponse } from 'elasticsearch';
 import { Observable } from 'rxjs';
+import { TransportRequestPromise } from '@elastic/elasticsearch/lib/Transport';
 import { SharedGlobalConfig, RequestHandlerContext, Logger } from '../../../../../src/core/server';
 import {
   getTotalLoaded,
@@ -40,8 +41,8 @@ export const enhancedEsSearchStrategyProvider = (
 
     try {
       const response = isAsync
-        ? await asyncSearch(context, request)
-        : await rollupSearch(context, request);
+        ? await asyncSearch(context, request, options)
+        : await rollupSearch(context, request, options);
 
       if (
         usage &&
@@ -69,9 +70,10 @@ export const enhancedEsSearchStrategyProvider = (
 
   async function asyncSearch(
     context: RequestHandlerContext,
-    request: IEnhancedEsSearchRequest
+    request: IEnhancedEsSearchRequest,
+    options?: ISearchOptions
   ): Promise<IEsSearchResponse> {
-    let esResponse;
+    let promise: TransportRequestPromise<any>;
     const esClient = context.core.elasticsearch.client.asCurrentUser;
     const uiSettingsClient = await context.core.uiSettings.client;
 
@@ -89,14 +91,17 @@ export const enhancedEsSearchStrategyProvider = (
         ...request.params,
       });
 
-      esResponse = await esClient.asyncSearch.submit(submitOptions);
+      promise = esClient.asyncSearch.submit(submitOptions);
     } else {
-      esResponse = await esClient.asyncSearch.get({
+      promise = esClient.asyncSearch.get({
         id: request.id,
         ...toSnakeCase(asyncOptions),
       });
     }
 
+    // Temporary workaround until https://github.com/elastic/elasticsearch-js/issues/1297
+    if (options?.abortSignal) options.abortSignal.addEventListener('abort', () => promise.abort());
+    const esResponse = await promise;
     const { id, response, is_partial: isPartial, is_running: isRunning } = esResponse.body;
     return {
       id,
@@ -109,7 +114,8 @@ export const enhancedEsSearchStrategyProvider = (
 
   const rollupSearch = async function (
     context: RequestHandlerContext,
-    request: IEnhancedEsSearchRequest
+    request: IEnhancedEsSearchRequest,
+    options?: ISearchOptions
   ): Promise<IEsSearchResponse> {
     const esClient = context.core.elasticsearch.client.asCurrentUser;
     const uiSettingsClient = await context.core.uiSettings.client;
@@ -123,13 +129,17 @@ export const enhancedEsSearchStrategyProvider = (
       ...params,
     });
 
-    const esResponse = await esClient.transport.request({
+    const promise = esClient.transport.request({
       method,
       path,
       body,
       querystring,
     });
 
+    // Temporary workaround until https://github.com/elastic/elasticsearch-js/issues/1297
+    if (options?.abortSignal) options.abortSignal.addEventListener('abort', () => promise.abort());
+    const esResponse = await promise;
+
     const response = esResponse.body as SearchResponse<any>;
     return {
       rawResponse: response,