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/__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);

…/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

@@ -66,7 +66,7 @@ import { getVariableSignature } from './getVariableSignature.js';
 import { mapAsyncIterable } from './mapAsyncIterable.js';
 import type { VariableValues } from './values.js';
 import { getArgumentValues, 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
@@ -1557,7 +1557,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>(

src/execution/withCleanup.ts

@@ -0,0 +1,62 @@
+import { isPromise } from '../jsutils/isPromise.js';
+import type { PromiseOrValue } from '../jsutils/PromiseOrValue.js';
+
+const asyncDispose: typeof Symbol.asyncDispose =
+  Symbol.asyncDispose /* c8 ignore start */ ??
+  Symbol.for('Symbol.asyncDispose'); /* c8 ignore stop */
+
+/**
+ * Given an AsyncGenerator and provided functions, return an AsyncGenerator
+ * which calls the given functions when the generator is abruptly closed,
+ * calling the functions immediately even if the generator is paused.
+ *
+ * This is useful for allowing return and throw to trigger logic even if the
+ * generator is paused on a pending await within a `next()` call (including
+ * if that logic can cause that hanging `next()` call to return early).
+ *
+ * Errors from the provided functions are ignored.
+ *
+ * The provided functions should be idempotent, as they may be called
+ * multiple times.
+ */
+export function withConcurrentAbruptClose<T>(
+  generator: AsyncGenerator<T, void, void>,
+  beforeReturn: () => PromiseOrValue<void>,
+  beforeThrow: (error?: unknown) => PromiseOrValue<void> = beforeReturn,
+): AsyncGenerator<T, void, void> {
+  return {
+    [Symbol.asyncIterator]() {
+      return this;
+    },
+    next() {
+      return generator.next();
+    },
+    async return() {
+      ignoreErrors(beforeReturn);
+      return generator.return();
+    },
+    async throw(error?: unknown) {
+      ignoreErrors(() => beforeThrow(error));
+      return generator.throw(error);
+    },
+    async [asyncDispose]() {
+      ignoreErrors(beforeReturn);
+      if (typeof generator[asyncDispose] === 'function') {
+        await generator[asyncDispose]();
+      }
+    },
+  };
+}
+
+function ignoreErrors(fn: () => PromiseOrValue<unknown>): void {
+  try {
+    const result = fn();
+    if (isPromise(result)) {
+      result.catch(() => {
+        // ignore error
+      });
+    }
+  } catch {
+    // ignore error
+  }
+}