[Search] Search batching using bfetch

docs/development/plugins/data/public/kibana-plugin-plugins-data-public.plugin.setup.md

@@ -7,15 +7,15 @@
 <b>Signature:</b>
 
 ```typescript
-setup(core: CoreSetup<DataStartDependencies, DataPublicPluginStart>, { expressions, uiActions, usageCollection }: DataSetupDependencies): DataPublicPluginSetup;
+setup(core: CoreSetup<DataStartDependencies, DataPublicPluginStart>, { bfetch, expressions, uiActions, usageCollection }: DataSetupDependencies): DataPublicPluginSetup;
 ```
 
 ## Parameters
 
 |  Parameter | Type | Description |
 |  --- | --- | --- |
 |  core | <code>CoreSetup&lt;DataStartDependencies, DataPublicPluginStart&gt;</code> |  |
-|  { expressions, uiActions, usageCollection } | <code>DataSetupDependencies</code> |  |
+|  { bfetch, expressions, uiActions, usageCollection } | <code>DataSetupDependencies</code> |  |
 
 <b>Returns:</b>
 

docs/development/plugins/data/public/kibana-plugin-plugins-data-public.searchinterceptordeps.bfetch.md

@@ -0,0 +1,11 @@
+<!-- Do not edit this file. It is automatically generated by API Documenter. -->
+
+[Home](./index.md) &gt; [kibana-plugin-plugins-data-public](./kibana-plugin-plugins-data-public.md) &gt; [SearchInterceptorDeps](./kibana-plugin-plugins-data-public.searchinterceptordeps.md) &gt; [bfetch](./kibana-plugin-plugins-data-public.searchinterceptordeps.bfetch.md)
+
+## SearchInterceptorDeps.bfetch property
+
+<b>Signature:</b>
+
+```typescript
+bfetch: BfetchPublicSetup;
+```

docs/development/plugins/data/public/kibana-plugin-plugins-data-public.searchinterceptordeps.md

@@ -14,6 +14,7 @@ export interface SearchInterceptorDeps
 
 |  Property | Type | Description |
 |  --- | --- | --- |
+|  [bfetch](./kibana-plugin-plugins-data-public.searchinterceptordeps.bfetch.md) | <code>BfetchPublicSetup</code> |  |
 |  [http](./kibana-plugin-plugins-data-public.searchinterceptordeps.http.md) | <code>CoreSetup['http']</code> |  |
 |  [session](./kibana-plugin-plugins-data-public.searchinterceptordeps.session.md) | <code>ISessionService</code> |  |
 |  [startServices](./kibana-plugin-plugins-data-public.searchinterceptordeps.startservices.md) | <code>Promise&lt;[CoreStart, any, unknown]&gt;</code> |  |

docs/development/plugins/data/server/kibana-plugin-plugins-data-server.plugin.setup.md

@@ -7,7 +7,7 @@
 <b>Signature:</b>
 
 ```typescript
-setup(core: CoreSetup<DataPluginStartDependencies, DataPluginStart>, { expressions, usageCollection }: DataPluginSetupDependencies): {
+setup(core: CoreSetup<DataPluginStartDependencies, DataPluginStart>, { bfetch, expressions, usageCollection }: DataPluginSetupDependencies): {
         __enhance: (enhancements: DataEnhancements) => void;
         search: ISearchSetup;
         fieldFormats: {
@@ -21,7 +21,7 @@ setup(core: CoreSetup<DataPluginStartDependencies, DataPluginStart>, { expressio
 |  Parameter | Type | Description |
 |  --- | --- | --- |
 |  core | <code>CoreSetup&lt;DataPluginStartDependencies, DataPluginStart&gt;</code> |  |
-|  { expressions, usageCollection } | <code>DataPluginSetupDependencies</code> |  |
+|  { bfetch, expressions, usageCollection } | <code>DataPluginSetupDependencies</code> |  |
 
 <b>Returns:</b>
 

src/plugins/bfetch/public/batching/batch_utils.test.ts

@@ -0,0 +1,86 @@
+/*
+ * 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 { BatchItem } from './types';
+import { getDonePromise } from './batch_utils';
+import { defer } from 'src/plugins/kibana_utils/common';
+
+const tick = () => new Promise((resolve) => setTimeout(resolve, 1));
+
+describe('getDonePromise()', () => {
+  test('Triggers when aborted', async () => {
+    const abortController = new AbortController();
+    const item: BatchItem<any, any> = {
+      future: defer<any>(),
+      payload: null,
+      done: false,
+      signal: abortController.signal,
+    };
+    const b = getDonePromise(item);
+
+    const spy = jest.fn();
+
+    b.then(spy);
+
+    abortController.abort();
+    await tick();
+    expect(spy).toHaveBeenCalledTimes(1);
+    expect(item.done).toBeTruthy();
+  });
+
+  test('Triggers when resolved', async () => {
+    const abortController = new AbortController();
+    const item: BatchItem<any, any> = {
+      future: defer<any>(),
+      payload: null,
+      done: false,
+      signal: abortController.signal,
+    };
+    const b = getDonePromise(item);
+
+    const spy = jest.fn();
+
+    b.then(spy);
+
+    item.future.resolve(null);
+    await tick();
+    expect(spy).toHaveBeenCalledTimes(1);
+    expect(item.done).toBeTruthy();
+  });
+
+  test('Triggers when rejected', async () => {
+    const abortController = new AbortController();
+    const item: BatchItem<any, any> = {
+      future: defer<any>(),
+      payload: null,
+      done: false,
+      signal: abortController.signal,
+    };
+    const b = getDonePromise(item);
+
+    const spy = jest.fn();
+
+    b.then(spy);
+
+    item.future.reject(null);
+    await tick();
+    expect(spy).toHaveBeenCalledTimes(1);
+    expect(item.done).toBeTruthy();
+  });
+});

src/plugins/bfetch/public/batching/batch_utils.ts

@@ -0,0 +1,50 @@
+/*
+ * 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 { BatchItem } from './types';
+import { AbortError } from '../../../kibana_utils/public';
+
+export function isBatchDone(items: Array<BatchItem<any, any>>): boolean {
+  return items.every((item) => item.done);
+}
+
+export function getDonePromise(item: BatchItem<any, any>) {
+  return new Promise<void>((resolve) => {
+    const onDone = () => {
+      item.done = true;
+      if (item.signal) item.signal.removeEventListener('abort', onDone);
+      resolve();
+    };
+    item.future.promise.then(onDone, onDone);
+    if (item.signal) item.signal.addEventListener('abort', onDone);
+  });
+}
+
+export function rejectOnAbort(item: BatchItem<any, any>) {
+  const cleanup = () => {
+    if (item.signal) item.signal.removeEventListener('abort', rejectAborted);
+  };
+  const rejectAborted = () => {
+    item.future.reject(new AbortError());
+    cleanup();
+  };
+
+  if (item.signal) item.signal.addEventListener('abort', rejectAborted);
+  item.future.promise.then(cleanup, cleanup);
+}

src/plugins/bfetch/public/batching/create_streaming_batched_function.test.ts

@@ -19,7 +19,7 @@
 
 import { createStreamingBatchedFunction } from './create_streaming_batched_function';
 import { fetchStreaming as fetchStreamingReal } from '../streaming/fetch_streaming';
-import { defer, of } from '../../../kibana_utils/public';
+import { AbortError, defer, of } from '../../../kibana_utils/public';
 import { Subject } from 'rxjs';
 
 const getPromiseState = (promise: Promise<unknown>): Promise<'resolved' | 'rejected' | 'pending'> =>
@@ -168,6 +168,28 @@ describe('createStreamingBatchedFunction()', () => {
       expect(fetchStreaming).toHaveBeenCalledTimes(1);
     });
 
+    test('ignores a request with an aborted signal', async () => {
+      const { fetchStreaming } = setup();
+      const fn = createStreamingBatchedFunction({
+        url: '/test',
+        fetchStreaming,
+        maxItemAge: 5,
+        flushOnMaxItems: 3,
+      });
+
+      const abortController = new AbortController();
+      abortController.abort();
+
+      of(fn({ foo: 'bar' }, abortController.signal));
+      fn({ baz: 'quix' });
+
+      await new Promise((r) => setTimeout(r, 6));
+      const { body } = fetchStreaming.mock.calls[0][0];
+      expect(JSON.parse(body)).toEqual({
+        batch: [{ baz: 'quix' }],
+      });
+    });
+
     test('sends POST request to correct endpoint with items in array batched sorted in call order', async () => {
       const { fetchStreaming } = setup();
       const fn = createStreamingBatchedFunction({
@@ -423,6 +445,44 @@ describe('createStreamingBatchedFunction()', () => {
       expect(result3).toEqual({ b: '3' });
     });
 
+    describe('when requests are aborted', () => {
+      test('rejects promise on abort and lets others continue', async () => {
+        const { fetchStreaming, stream } = setup();
+        const fn = createStreamingBatchedFunction({
+          url: '/test',
+          fetchStreaming,
+          maxItemAge: 5,
+          flushOnMaxItems: 3,
+        });
+
+        const abortController = new AbortController();
+        const promise = fn({ a: '1' }, abortController.signal);
+        const promise2 = fn({ a: '2' });
+        await new Promise((r) => setTimeout(r, 6));
+
+        expect(await isPending(promise)).toBe(true);
+
+        abortController.abort();
+        await new Promise((r) => setTimeout(r, 6));
+
+        expect(await isPending(promise)).toBe(false);
+        const [, error] = await of(promise);
+        expect(error).toBeInstanceOf(AbortError);
+
+        stream.next(
+          JSON.stringify({
+            id: 1,
+            result: { b: '2' },
+          }) + '\n'
+        );
+
+        await new Promise((r) => setTimeout(r, 1));
+
+        const [result2] = await of(promise2);
+        expect(result2).toEqual({ b: '2' });
+      });
+    });
+
     describe('when stream closes prematurely', () => {
       test('rejects pending promises with CONNECTION error code', async () => {
         const { fetchStreaming, stream } = setup();

src/plugins/bfetch/public/batching/create_streaming_batched_function.ts

@@ -17,7 +17,7 @@
  * under the License.
  */
 
-import { defer, Defer } from '../../../kibana_utils/public';
+import { AbortError, defer } from '../../../kibana_utils/public';
 import {
   ItemBufferParams,
   TimedItemBufferParams,
@@ -27,13 +27,8 @@ import {
 } from '../../common';
 import { fetchStreaming, split } from '../streaming';
 import { normalizeError } from '../../common';
-
-export interface BatchItem<Payload, Result> {
-  payload: Payload;
-  future: Defer<Result>;
-}
-
-export type BatchedFunc<Payload, Result> = (payload: Payload) => Promise<Result>;
+import { BatchedFunc, BatchItem } from './types';
+import { isBatchDone, getDonePromise, rejectOnAbort } from './batch_utils';
 
 export interface BatchedFunctionProtocolError extends ErrorLike {
   code: string;
@@ -83,31 +78,63 @@ export const createStreamingBatchedFunction = <Payload, Result extends object>(
     maxItemAge = 10,
   } = params;
   const [fn] = createBatchedFunction<BatchedFunc<Payload, Result>, BatchItem<Payload, Result>>({
-    onCall: (payload: Payload) => {
+    onCall: (payload: Payload, signal?: AbortSignal) => {
       const future = defer<Result>();
       const entry: BatchItem<Payload, Result> = {
         payload,
         future,
+        signal,
+        done: false,
       };
       return [future.promise, entry];
     },
     onBatch: async (items) => {
       try {
-        let responsesReceived = 0;
-        const batch = items.map(({ payload }) => payload);
+        const promises: Array<Promise<void>> = [];
+        const abortController = new AbortController();
+
+        // Filter out and reject any items who's signal is already aborted
+        items = items.filter((item) => {
+          if (item.signal?.aborted) {
+            item.future.reject(new AbortError());
+          }
+          return !item.signal?.aborted;
+        });
+
+        // Prepare batch
+        const batch = items
+          .filter((item) => {
+            if (item.signal?.aborted) {
+              item.future.reject(new AbortError());
+            }
+            return !item.signal?.aborted;
+          })
+          .map((item) => {
+            // Subscribe to reject promise on abort
+            rejectOnAbort(item);
+
+            // Track batch progress
+            promises.push(getDonePromise(item));
+
+            // Return payload to be sent
+            return item.payload;
+          });
+
+        // abort when all items were either resolved, rejected or aborted
+        Promise.all(promises).then(() => abortController.abort());
+
         const { stream } = fetchStreamingInjected({
           url,
           body: JSON.stringify({ batch }),
           method: 'POST',
+          signal: abortController.signal,
         });
         stream.pipe(split('\n')).subscribe({
           next: (json: string) => {
             const response = JSON.parse(json) as BatchResponseItem<Result, ErrorLike>;
             if (response.error) {
-              responsesReceived++;
               items[response.id].future.reject(response.error);
             } else if (response.result !== undefined) {
-              responsesReceived++;
               items[response.id].future.resolve(response.result);
             }
           },
@@ -117,7 +144,7 @@ export const createStreamingBatchedFunction = <Payload, Result extends object>(
             for (const { future } of items) future.reject(normalizedError);
           },
           complete: () => {
-            const streamTerminatedPrematurely = responsesReceived !== items.length;
+            const streamTerminatedPrematurely = !isBatchDone(items);
             if (streamTerminatedPrematurely) {
               const error: BatchedFunctionProtocolError = {
                 message: 'Connection terminated prematurely.',