reduce ticks stemming from withCleanup (#4526)

replace withCleanup() with more targeted withConcurrentAbruptClose().

next() calls are no longer wrapped, reducing extra ticks

 = refactors mapAsyncIterable to use the more targeted function
= refactors IncrementalPublisher to also perform required clean up on a
normal close.

* this change also alters the behavior of a stream test to no longer
pass spuriously and allows it to be further streamlined

Author: yaacovCR

src/execution/IncrementalPublisher.ts

@@ -23,7 +23,7 @@ import type {
   SubsequentIncrementalExecutionResult,
 } from './types.js';
 import { isCompletedExecutionGroup, isFailedExecutionGroup } from './types.js';
-import { withCleanup } from './withCleanup.js';
+import { withConcurrentAbruptClose } from './withConcurrentAbruptClose.js';
 
 // eslint-disable-next-line max-params
 export function buildIncrementalResponse(
@@ -99,13 +99,17 @@ class IncrementalPublisher {
 
     return {
       initialResult,
-      subsequentResults: withCleanup(subsequentResults, async () => {
-        this._abortSignalListener?.disconnect();
-        await this._returnAsyncIteratorsIgnoringErrors();
-      }),
+      subsequentResults: withConcurrentAbruptClose(subsequentResults, () =>
+        this._cleanUp(),
+      ),
     };
   }
 
+  private async _cleanUp(): Promise<void> {
+    this._abortSignalListener?.disconnect();
+    await this._returnAsyncIteratorsIgnoringErrors();
+  }
+
   private _ensureId(deliveryGroup: DeliveryGroup): string {
     return (deliveryGroup.id ??= String(this._nextId++));
   }
@@ -160,6 +164,14 @@ class IncrementalPublisher {
     if (completed.length > 0) {
       subsequentIncrementalExecutionResult.completed = completed;
     }
+
+    if (!hasNext) {
+      this._cleanUp().catch(() => {
+        /* c8 ignore next 2 */
+        // ignore error
+      });
+    }
+
     return subsequentIncrementalExecutionResult;
   }
 

src/execution/Queue.ts

@@ -2,7 +2,7 @@ import { isPromise } from '../jsutils/isPromise.js';
 import type { PromiseOrValue } from '../jsutils/PromiseOrValue.js';
 import { promiseWithResolvers } from '../jsutils/promiseWithResolvers.js';
 
-import { withCleanup } from './withCleanup.js';
+import { withConcurrentAbruptClose } from './withConcurrentAbruptClose.js';
 
 /**
  * @internal
@@ -42,7 +42,9 @@ export class Queue<T> {
   subscribe<U>(
     mapFn: (generator: Generator<T>) => U | undefined,
   ): AsyncGenerator<U, void, void> {
-    return withCleanup(this.subscribeImpl(mapFn), () => this.stop());
+    return withConcurrentAbruptClose(this.subscribeImpl(mapFn), () =>
+      this.stop(),
+    );
   }
 
   private async *subscribeImpl<U>(

src/execution/__tests__/Queue-test.ts

@@ -84,8 +84,8 @@ describe('Queue', () => {
     const sub = queue.subscribe((batch) => Array.from(batch)[0]);
     expect(await sub.next()).to.deep.equal({ done: false, value: 1 });
     expect(await sub.next()).to.deep.equal({ done: false, value: 4 });
-    expect(await sub.next()).to.deep.equal({ done: false, value: 16 });
-    expect(await sub.next()).to.deep.equal({ done: false, value: 28 });
+    expect(await sub.next()).to.deep.equal({ done: false, value: 13 });
+    expect(await sub.next()).to.deep.equal({ done: false, value: 22 });
   });
 
   it('should allow the executor to indicate completion', async () => {
@@ -180,8 +180,8 @@ describe('Queue', () => {
       }
     });
     expect(await sub.next()).to.deep.equal({ done: false, value: [2] });
-    expect(await sub.next()).to.deep.equal({ done: false, value: [8] });
-    expect(await sub.next()).to.deep.equal({ done: false, value: [14] });
+    expect(await sub.next()).to.deep.equal({ done: false, value: [6] });
+    expect(await sub.next()).to.deep.equal({ done: false, value: [10] });
   });
 
   it('should condense pushes during map into the same batch', async () => {

src/execution/__tests__/simplePubSub.ts

@@ -1,6 +1,6 @@
 import { promiseWithResolvers } from '../../jsutils/promiseWithResolvers.js';
 
-import { withCleanup } from '../withCleanup.js';
+import { withConcurrentAbruptClose } from '../withConcurrentAbruptClose.js';
 
 /**
  * Create an AsyncIterator from an EventEmitter. Useful for mocking a
@@ -59,7 +59,7 @@ export class SimplePubSub<T> {
       }
     }
 
-    return withCleanup(getSubscriberImpl(), emptyQueue);
+    return withConcurrentAbruptClose(getSubscriberImpl(), emptyQueue);
 
     function pushValue(event: T): void {
       const value: R = transform(event);

src/execution/__tests__/stream-test.ts

@@ -1795,23 +1795,18 @@ describe('Execute: stream directive', () => {
 
   it('Returns iterator and ignores errors when stream payloads are filtered', async () => {
     let returned = false;
-    let requested = false;
     const iterable = {
       [Symbol.asyncIterator]: () => ({
         next: () => {
-          /* c8 ignore start */
-          if (requested) {
-            // stream is filtered, next is not called, and so this is not reached.
-            return Promise.reject(new Error('Oops'));
-          } /* c8 ignore stop */
-          requested = true;
-          const friend = friends[0];
+          if (returned) {
+            return Promise.resolve({
+              done: true,
+              value: undefined,
+            });
+          }
           return Promise.resolve({
             done: false,
-            value: {
-              name: friend.name,
-              nonNullName: null,
-            },
+            value: friends[0],
           });
         },
         return: () => {

…/execution/__tests__/withCleanup-test.ts …ests__/withConcurrentAbruptClose-test.tssrc/execution/__tests__/withCleanup-test.ts renamed to src/execution/__tests__/withConcurrentAbruptClose-test.ts src/execution/__tests__/withCleanup-test.ts renamed to src/execution/__tests__/withConcurrentAbruptClose-test.ts

@@ -3,81 +3,122 @@ import { describe, it } from 'mocha';
 
 import { expectPromise } from '../../__testUtils__/expectPromise.js';
 
-import { withCleanup } from '../withCleanup.js';
+import { withConcurrentAbruptClose } from '../withConcurrentAbruptClose.js';
 
 /* eslint-disable @typescript-eslint/require-await */
-describe('withCleanup', () => {
-  it('calls cleanup function when completes', async () => {
+describe('withConcurrentAbruptClose', () => {
+  it('calls function when returned', async () => {
     async function* source() {
       yield 1;
     }
 
     let done = false;
-    const generator = withCleanup(source(), () => {
+
+    const generator = withConcurrentAbruptClose(source(), () => {
       done = true;
     });
 
     expect(await generator.next()).to.deep.equal({ value: 1, done: false });
     expect(done).to.equal(false);
-    expect(await generator.next()).to.deep.equal({
+    expect(await generator.return()).to.deep.equal({
       value: undefined,
       done: true,
     });
     expect(done).to.equal(true);
   });
 
-  it('calls cleanup function when completes with error', async () => {
+  it('ignores sync errors when returned', async () => {
     async function* source() {
       yield 1;
-      throw new Error('Oops');
     }
 
-    let done = false;
-    const generator = withCleanup(source(), () => {
-      done = true;
+    const generator = withConcurrentAbruptClose(source(), () => {
+      throw new Error('Oops');
     });
 
     expect(await generator.next()).to.deep.equal({ value: 1, done: false });
-    expect(done).to.equal(false);
-    await expectPromise(generator.next()).toRejectWith('Oops');
-    expect(done).to.equal(true);
+    expect(await generator.return()).to.deep.equal({
+      value: undefined,
+      done: true,
+    });
   });
 
-  it('calls cleanup function when returned', async () => {
+  it('ignores async errors when returned', async () => {
     async function* source() {
       yield 1;
     }
 
-    let done = false;
-    const generator = withCleanup(source(), () => {
-      done = true;
-    });
+    const generator = withConcurrentAbruptClose(source(), () =>
+      Promise.reject(new Error('Oops')),
+    );
 
     expect(await generator.next()).to.deep.equal({ value: 1, done: false });
-    expect(done).to.equal(false);
     expect(await generator.return()).to.deep.equal({
       value: undefined,
       done: true,
     });
-    expect(done).to.equal(true);
   });
 
-  it('calls cleanup function when thrown', async () => {
+  it('calls function when thrown', async () => {
     async function* source() {
       yield 1;
     }
 
     let done = false;
-    const generator = withCleanup(source(), () => {
-      done = true;
-    });
+    let error;
+    const generator = withConcurrentAbruptClose(
+      source(),
+      () => {
+        done = true;
+      },
+      (err) => {
+        done = true;
+        error = err;
+      },
+    );
 
     expect(await generator.next()).to.deep.equal({ value: 1, done: false });
     expect(done).to.equal(false);
-    await expectPromise(generator.throw(new Error('Oops'))).toRejectWith(
-      'Oops',
-    );
+    const oops = new Error('Oops');
+    await expectPromise(generator.throw(oops)).toRejectWith('Oops');
     expect(done).to.equal(true);
+    expect(error).to.equal(oops);
+  });
+
+  it('ignores sync errors when thrown', async () => {
+    async function* source() {
+      yield 1;
+    }
+
+    const generator = withConcurrentAbruptClose(
+      source(),
+      () => {
+        throw new Error('Ignored');
+      },
+      () => {
+        throw new Error('Ignored');
+      },
+    );
+
+    expect(await generator.next()).to.deep.equal({ value: 1, done: false });
+    const oops = new Error('Oops');
+    await expectPromise(generator.throw(oops)).toRejectWith('Oops');
+  });
+
+  it('ignores async errors when thrown', async () => {
+    async function* source() {
+      yield 1;
+    }
+
+    const generator = withConcurrentAbruptClose(
+      source(),
+      () => Promise.reject(new Error('Ignored')),
+      () => Promise.reject(new Error('Ignored')),
+    );
+
+    expect(await generator.next()).to.deep.equal({ value: 1, done: false });
+    const oops = new Error('Oops');
+    await expectPromise(generator.throw(oops)).toRejectWith('Oops');
   });
 
   it('calls cleanup function when disposed', async () => {
@@ -109,17 +150,17 @@ describe('withCleanup', () => {
       },
     };
 
-    let cleanedUp = false;
+    let called = false;
     {
-      await using generator = withCleanup(source, () => {
-        cleanedUp = true;
+      await using generator = withConcurrentAbruptClose(source, () => {
+        called = true;
       });
 
       expect(await generator.next()).to.deep.equal({ value: 1, done: false });
       expect(await generator.next()).to.deep.equal({ value: 2, done: false });
     }
 
-    expect(cleanedUp).to.equal(true);
+    expect(called).to.equal(true);
     expect(returned).to.equal(true);
   });
 
@@ -128,7 +169,7 @@ describe('withCleanup', () => {
       yield 1;
     }
 
-    const generator = withCleanup(source(), () => {
+    const generator = withConcurrentAbruptClose(source(), () => {
       /* noop */
     });
 

src/execution/execute.ts

@@ -88,7 +88,7 @@ import {
   getDirectiveValues,
   getVariableValues,
 } from './values.js';
-import { withCleanup } from './withCleanup.js';
+import { withConcurrentAbruptClose } from './withConcurrentAbruptClose.js';
 
 /* eslint-disable max-params */
 // This file contains a lot of such errors but we plan to refactor it anyway
@@ -2264,7 +2264,7 @@ function mapSourceToResponse(
   }
 
   return abortSignalListener
-    ? withCleanup(
+    ? withConcurrentAbruptClose(
         mapAsyncIterable(
           cancellableIterable(resultOrStream, abortSignalListener),
           mapFn,

src/execution/mapAsyncIterable.ts

@@ -1,7 +1,7 @@
 import { isPromise } from '../jsutils/isPromise.js';
 import type { PromiseOrValue } from '../jsutils/PromiseOrValue.js';
 
-import { withCleanup } from './withCleanup.js';
+import { withConcurrentAbruptClose } from './withConcurrentAbruptClose.js';
 
 /**
  * Given an AsyncIterable and a callback function, return an AsyncIterator
@@ -11,17 +11,35 @@ export function mapAsyncIterable<T, U>(
   iterable: AsyncGenerator<T> | AsyncIterable<T>,
   callback: (value: T) => PromiseOrValue<U>,
 ): AsyncGenerator<U, void, void> {
-  return withCleanup(mapAsyncIterableImpl(iterable, callback), async () => {
-    const iterator = iterable[Symbol.asyncIterator]();
-    if (typeof iterator.return === 'function') {
-      try {
-        await iterator.return(); /* c8 ignore start */
-      } catch (_error) {
-        // FIXME: add test case
-        /* ignore error */
-      } /* c8 ignore stop */
-    }
-  });
+  const iterator = iterable[Symbol.asyncIterator]();
+  const returnFn = iterator.return?.bind(iterator);
+  const throwFn = iterator.throw?.bind(iterator);
+
+  const onReturn = returnFn
+    ? async () => {
+        await callIgnoringErrors(returnFn);
+      }
+    : () => Promise.resolve();
+
+  const onThrow = throwFn
+    ? async (reason?: unknown) => {
+        await callIgnoringErrors(() => throwFn(reason));
+      }
+    : onReturn;
+
+  return withConcurrentAbruptClose(
+    mapAsyncIterableImpl(iterable, callback),
+    onReturn,
+    onThrow,
+  );
+}
+
+async function callIgnoringErrors(fn: () => Promise<unknown>): Promise<void> {
+  try {
+    await fn();
+  } catch {
+    // ignore error
+  }
 }
 
 async function* mapAsyncIterableImpl<T, U, R = undefined>(