feat(execution): track async work for rejection safety

canonize async work tracking so promises started during execution are still observed after early errors or abort paths

route defaultTypeResolver through tracked async helpers to avoid dropping pending isTypeOf promise rejections

Author: yaacovCR

src/execution/AsyncWorkTracker.ts

@@ -0,0 +1,52 @@
+import { isPromise } from '../jsutils/isPromise.js';
+import type { PromiseOrValue } from '../jsutils/PromiseOrValue.js';
+
+/** @internal */
+export class AsyncWorkTracker {
+  pendingAsyncWork: Set<Promise<void>>;
+
+  constructor() {
+    this.pendingAsyncWork = new Set<Promise<void>>();
+  }
+
+  add(promise: Promise<unknown>): void {
+    const pendingAsyncWork = this.pendingAsyncWork;
+    const promiseToSettle = promise.then(
+      () => {
+        pendingAsyncWork.delete(promiseToSettle);
+      },
+      () => {
+        pendingAsyncWork.delete(promiseToSettle);
+      },
+    );
+    pendingAsyncWork.add(promiseToSettle);
+  }
+
+  addValues(values: ReadonlyArray<PromiseOrValue<unknown>>): void {
+    for (const value of values) {
+      if (isPromise(value)) {
+        this.add(value);
+      }
+    }
+  }
+
+  promiseAllTrackOnReject<T>(
+    values: ReadonlyArray<PromiseOrValue<T>>,
+  ): Promise<Array<T>> {
+    const promise = Promise.all(values);
+    promise.then(undefined, () => {
+      this.addValues(values);
+    });
+    return promise;
+  }
+
+  promiseCombinatorWithTracking<T, TResult>(
+    values: ReadonlyArray<PromiseOrValue<T>>,
+    combinator: (
+      promises: ReadonlyArray<PromiseOrValue<T>>,
+    ) => Promise<TResult>,
+  ): Promise<TResult> {
+    this.add(Promise.allSettled(values));
+    return combinator(values);
+  }
+}

src/execution/Executor.ts

@@ -32,6 +32,7 @@ import type {
   GraphQLObjectType,
   GraphQLOutputType,
   GraphQLResolveInfo,
+  GraphQLResolveInfoHelpers,
   GraphQLTypeResolver,
 } from '../type/definition.js';
 import {
@@ -231,6 +232,11 @@ export class Executor<
   abortResultPromise: ((reason?: unknown) => void) | undefined;
   resolverAbortController: AbortController | undefined;
   getAbortSignal: () => AbortSignal | undefined;
+  getAsyncHelpers: () => GraphQLResolveInfoHelpers;
+  promiseAll: <T>(
+    values: ReadonlyArray<PromiseOrValue<T>>,
+  ) => Promise<Array<T>>;
+  trackPromise: (promise: Promise<unknown>) => void;
 
   constructor(
     validatedExecutionArgs: ValidatedExecutionArgs,
@@ -249,8 +255,12 @@ export class Executor<
     } else {
       this.sharedExecutionContext = sharedExecutionContext;
     }
-    const { getAbortSignal } = this.sharedExecutionContext;
+    const { getAbortSignal, getAsyncHelpers, promiseAll, trackPromise } =
+      this.sharedExecutionContext;
     this.getAbortSignal = getAbortSignal;
+    this.getAsyncHelpers = getAsyncHelpers;
+    this.promiseAll = promiseAll;
+    this.trackPromise = trackPromise;
   }
 
   executeQueryOrMutationOrSubscriptionEvent(): PromiseOrValue<
@@ -261,10 +271,7 @@ export class Executor<
     if (externalAbortSignal) {
       externalAbortSignal.throwIfAborted();
       const onExternalAbort = () => {
-        const aborted = this.abort(externalAbortSignal.reason);
-        if (isPromise(aborted)) {
-          aborted.catch(() => undefined);
-        }
+        this.abort(externalAbortSignal.reason);
       };
       removeExternalAbortListener = () =>
         externalAbortSignal.removeEventListener('abort', onExternalAbort);
@@ -326,6 +333,7 @@ export class Executor<
             return this.buildResponse(null);
           },
         );
+        this.sharedExecutionContext.asyncWorkTracker.add(promise);
         const { promise: cancellablePromise, abort: abortResultPromise } =
           withCancellation(promise);
         this.abortResultPromise = abortResultPromise;
@@ -349,7 +357,7 @@ export class Executor<
     }
   }
 
-  abort(reason?: unknown): PromiseOrValue<void> {
+  abort(reason?: unknown): void {
     if (this.aborted) {
       return;
     }
@@ -508,8 +516,9 @@ export class Executor<
       }
     } catch (error) {
       if (containsPromise) {
-        // Ensure that any promises returned by other fields are handled, as they may also reject.
-        promiseForObject(results).catch(() => undefined);
+        this.sharedExecutionContext.asyncWorkTracker.addValues(
+          Object.values(results),
+        );
       }
       throw error;
     }
@@ -522,7 +531,7 @@ export class Executor<
     // Otherwise, results is a map from field name to the result of resolving that
     // field, which is possibly a promise. Return a promise that will return this
     // same map, but with any promises replaced with the values they resolved to.
-    return promiseForObject(results);
+    return promiseForObject(results, this.promiseAll);
   }
 
   /**
@@ -559,6 +568,7 @@ export class Executor<
       parentType,
       path,
       this.getAbortSignal,
+      this.getAsyncHelpers,
     );
 
     // Get the resolve function, regardless of if its result is normal or abrupt (error).
@@ -855,24 +865,26 @@ export class Executor<
         index++;
       }
     } catch (error) {
-      // eslint-disable-next-line @typescript-eslint/no-floating-promises
-      returnIteratorCatchingErrors(asyncIterator);
+      this.trackPromise(returnIteratorCatchingErrors(asyncIterator));
       if (containsPromise) {
-        Promise.all(completedResults).catch(() => undefined);
+        this.sharedExecutionContext.asyncWorkTracker.addValues(
+          completedResults,
+        );
       }
       throw error;
     }
 
     // Throwing on completion outside of the loop may allow engines to better optimize
     if (this.aborted) {
       if (!iteration?.done) {
-        // eslint-disable-next-line @typescript-eslint/no-floating-promises
-        returnIteratorCatchingErrors(asyncIterator);
+        this.trackPromise(returnIteratorCatchingErrors(asyncIterator));
       }
       throw new Error('Aborted!');
     }
 
-    return containsPromise ? Promise.all(completedResults) : completedResults;
+    return containsPromise
+      ? this.promiseAll(completedResults)
+      : completedResults;
   }
 
   /* c8 ignore next 12 */
@@ -993,15 +1005,17 @@ export class Executor<
         index++;
       }
     } catch (error) {
-      const maybePromises = containsPromise ? completedResults : [];
-      maybePromises.push(...collectIteratorPromises(iterator));
-      if (maybePromises.length) {
-        Promise.all(maybePromises).catch(() => undefined);
+      const asyncWorkTracker = this.sharedExecutionContext.asyncWorkTracker;
+      if (containsPromise) {
+        asyncWorkTracker.addValues(completedResults);
       }
+      asyncWorkTracker.addValues(collectIteratorPromises(iterator));
       throw error;
     }
 
-    return containsPromise ? Promise.all(completedResults) : completedResults;
+    return containsPromise
+      ? this.promiseAll(completedResults)
+      : completedResults;
   }
 
   completeMaybePromisedListItemValue(

src/execution/__tests__/AsyncWorkTracker-test.ts

@@ -0,0 +1,147 @@
+import { expect } from 'chai';
+import { describe, it } from 'mocha';
+
+import { expectEqualPromisesOrValues } from '../../__testUtils__/expectEqualPromisesOrValues.js';
+import { resolveOnNextTick } from '../../__testUtils__/resolveOnNextTick.js';
+
+import { promiseWithResolvers } from '../../jsutils/promiseWithResolvers.js';
+
+import { AsyncWorkTracker } from '../AsyncWorkTracker.js';
+
+describe('AsyncWorkTracker', () => {
+  it('works to track promises', async () => {
+    const tracker = new AsyncWorkTracker();
+    const delayed = promiseWithResolvers<number>();
+
+    tracker.add(delayed.promise);
+    expect(tracker.pendingAsyncWork.size).to.equal(1);
+    delayed.resolve(1);
+    await resolveOnNextTick();
+    expect(tracker.pendingAsyncWork.size).to.equal(0);
+  });
+});
+
+describe('promiseAllTrackOnReject', () => {
+  it('resolves like Promise.all', async () => {
+    const tracker = new AsyncWorkTracker();
+
+    const values = [Promise.resolve(1), Promise.resolve(2), Promise.resolve(3)];
+
+    await expectEqualPromisesOrValues([
+      tracker.promiseAllTrackOnReject(values),
+      Promise.all(values),
+    ]);
+  });
+
+  it('resolves synchronous values without tracking', async () => {
+    const tracker = new AsyncWorkTracker();
+
+    const result = await tracker.promiseAllTrackOnReject([1, 2, 3]);
+
+    expect(result).to.deep.equal([1, 2, 3]);
+    expect(tracker.pendingAsyncWork.size).to.equal(0);
+  });
+
+  it('does not add an extra microtask on fulfilled promiseAll results', async () => {
+    const tracker = new AsyncWorkTracker();
+    let settled = false;
+
+    const promise = Promise.resolve(1);
+    const trackedPromise = tracker.promiseAllTrackOnReject([promise]);
+    trackedPromise.then(
+      () => {
+        settled = true;
+      },
+      () => undefined,
+    );
+    await Promise.all([promise]);
+    expect(settled).to.equal(true);
+  });
+
+  it('tracks all promises only after rejection', async () => {
+    const delayed = promiseWithResolvers<undefined>();
+    const tracker = new AsyncWorkTracker();
+    const result = tracker.promiseAllTrackOnReject([
+      Promise.reject(new Error('bad')),
+      delayed.promise,
+    ] as const);
+    expect(tracker.pendingAsyncWork.size).to.equal(0);
+
+    await result.catch(() => undefined);
+    expect(tracker.pendingAsyncWork.size).to.equal(1);
+    delayed.resolve(undefined);
+
+    await resolveOnNextTick();
+    expect(tracker.pendingAsyncWork.size).to.equal(0);
+  });
+
+  it('tracks promises until they settle and catches later rejections', async () => {
+    let unhandledRejection: unknown = null;
+    const unhandledRejectionListener = (reason: unknown) => {
+      unhandledRejection = reason;
+    };
+    // eslint-disable-next-line no-undef
+    process.on('unhandledRejection', unhandledRejectionListener);
+
+    const tracker = new AsyncWorkTracker();
+    const delayed = promiseWithResolvers<undefined>();
+    const result = tracker.promiseAllTrackOnReject([
+      Promise.reject(new Error('bad')),
+      delayed.promise,
+    ] as const);
+
+    await result.catch(() => undefined);
+    expect(tracker.pendingAsyncWork.size).to.equal(1);
+
+    delayed.reject(new Error('late bad'));
+    await new Promise((resolve) => setTimeout(resolve, 20));
+
+    // eslint-disable-next-line no-undef
+    process.removeListener('unhandledRejection', unhandledRejectionListener);
+
+    expect(tracker.pendingAsyncWork.size).to.equal(0);
+    expect(unhandledRejection).to.equal(null);
+  });
+});
+
+describe('promiseCombinatorWithTracking', () => {
+  it('tracks pending promises after Promise.any resolves', async () => {
+    const tracker = new AsyncWorkTracker();
+
+    const delayed = promiseWithResolvers<number>();
+    const result = await tracker.promiseCombinatorWithTracking(
+      [1, delayed.promise],
+      // eslint-disable-next-line @typescript-eslint/await-thenable
+      (promises) => Promise.any(promises),
+    );
+    expect(result).to.equal(1);
+
+    await resolveOnNextTick();
+    expect(tracker.pendingAsyncWork.size).to.equal(1);
+
+    delayed.resolve(2);
+    await Promise.allSettled(Array.from(tracker.pendingAsyncWork));
+    await resolveOnNextTick();
+    expect(tracker.pendingAsyncWork.size).to.equal(0);
+  });
+
+  it('tracks pending promises after Promise.race rejects', async () => {
+    const tracker = new AsyncWorkTracker();
+
+    const delayed = promiseWithResolvers<number>();
+    const result = tracker.promiseCombinatorWithTracking(
+      [Promise.reject(new Error('bad')), delayed.promise],
+      // eslint-disable-next-line @typescript-eslint/await-thenable
+      (promises) => Promise.race(promises),
+    );
+
+    await result.catch(() => undefined);
+    await resolveOnNextTick();
+    expect(tracker.pendingAsyncWork.size).to.equal(1);
+
+    delayed.resolve(1);
+    await Promise.allSettled(Array.from(tracker.pendingAsyncWork));
+    await resolveOnNextTick();
+    expect(tracker.pendingAsyncWork.size).to.equal(0);
+  });
+});

src/execution/__tests__/executor-test.ts

@@ -258,6 +258,14 @@ describe('Execute: Handles basic execution tasks', () => {
       'operation',
       'variableValues',
       'getAbortSignal',
+      'getAsyncHelpers',
+    );
+    const asyncHelpers = resolvedInfo?.getAsyncHelpers();
+    expect(asyncHelpers).to.have.all.keys(
+      'promiseAll',
+      'promiseAny',
+      'promiseRace',
+      'trackPromise',
     );
 
     const operation = document.definitions[0];
@@ -295,6 +303,25 @@ describe('Execute: Handles basic execution tasks', () => {
     expect(abortSignal).to.be.instanceOf(AbortSignal);
     expect(resolvedInfo?.getAbortSignal()).to.equal(abortSignal);
 
+    expect(resolvedInfo?.getAsyncHelpers()).to.equal(asyncHelpers);
+
+    const promiseAll = asyncHelpers?.promiseAll;
+    expect(promiseAll).to.be.a('function');
+    expect(resolvedInfo?.getAsyncHelpers().promiseAll).to.equal(promiseAll);
+
+    const promiseAny = asyncHelpers?.promiseAny;
+    expect(promiseAny).to.be.a('function');
+    expect(resolvedInfo?.getAsyncHelpers().promiseAny).to.equal(promiseAny);
+
+    const promiseRace = asyncHelpers?.promiseRace;
+    expect(promiseRace).to.be.a('function');
+    expect(resolvedInfo?.getAsyncHelpers().promiseRace).to.equal(promiseRace);
+
+    const trackPromise = asyncHelpers?.trackPromise;
+    expect(trackPromise).to.be.a('function');
+    expect(resolvedInfo?.getAsyncHelpers().trackPromise).to.equal(trackPromise);
+    trackPromise?.(Promise.resolve());
+
     resolve();
 
     await result;