refactor(execution): track pending work (#4657)

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

route defaultTypeResolver promise handling through tracked async helpers

motivation:

- #4658

Author: yaacovCR

src/execution/AsyncWorkTracker.ts

@@ -0,0 +1,42 @@
+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;
+  }
+}

src/execution/Executor.ts

@@ -32,6 +32,7 @@ import type {
   GraphQLObjectType,
   GraphQLOutputType,
   GraphQLResolveInfo,
+  GraphQLResolveInfoHelpers,
   GraphQLTypeResolver,
 } from '../type/definition.js';
 import {
@@ -231,6 +232,10 @@ 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>>;
 
   constructor(
     validatedExecutionArgs: ValidatedExecutionArgs,
@@ -249,8 +254,11 @@ export class Executor<
     } else {
       this.sharedExecutionContext = sharedExecutionContext;
     }
-    const { getAbortSignal } = this.sharedExecutionContext;
+    const { getAbortSignal, getAsyncHelpers, promiseAll } =
+      this.sharedExecutionContext;
     this.getAbortSignal = getAbortSignal;
+    this.getAsyncHelpers = getAsyncHelpers;
+    this.promiseAll = promiseAll;
   }
 
   executeQueryOrMutationOrSubscriptionEvent(): PromiseOrValue<
@@ -261,10 +269,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);
@@ -324,6 +329,7 @@ export class Executor<
             return this.buildResponse(null);
           },
         );
+        this.sharedExecutionContext.asyncWorkTracker.add(promise);
         const { promise: cancellablePromise, abort: abortResultPromise } =
           withCancellation(promise);
         this.abortResultPromise = abortResultPromise;
@@ -347,7 +353,7 @@ export class Executor<
     }
   }
 
-  abort(reason?: unknown): PromiseOrValue<void> {
+  abort(reason?: unknown): void {
     if (this.aborted) {
       return;
     }
@@ -506,8 +512,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;
     }
@@ -520,7 +527,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);
   }
 
   /**
@@ -557,6 +564,7 @@ export class Executor<
       parentType,
       path,
       this.getAbortSignal,
+      this.getAsyncHelpers,
     );
 
     // Get the resolve function, regardless of if its result is normal or abrupt (error).
@@ -853,24 +861,30 @@ export class Executor<
         index++;
       }
     } catch (error) {
-      // eslint-disable-next-line @typescript-eslint/no-floating-promises
-      returnIteratorCatchingErrors(asyncIterator);
+      this.sharedExecutionContext.asyncWorkTracker.add(
+        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.sharedExecutionContext.asyncWorkTracker.add(
+          returnIteratorCatchingErrors(asyncIterator),
+        );
       }
       throw new Error('Aborted!');
     }
 
-    return containsPromise ? Promise.all(completedResults) : completedResults;
+    return containsPromise
+      ? this.promiseAll(completedResults)
+      : completedResults;
   }
 
   /* c8 ignore next 12 */
@@ -991,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,105 @@
+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);
+  });
+});

src/execution/__tests__/executor-test.ts

@@ -258,7 +258,10 @@ describe('Execute: Handles basic execution tasks', () => {
       'operation',
       'variableValues',
       'getAbortSignal',
+      'getAsyncHelpers',
     );
+    const asyncHelpers = resolvedInfo?.getAsyncHelpers();
+    expect(asyncHelpers).to.have.all.keys('track');
 
     const operation = document.definitions[0];
     assert(operation.kind === Kind.OPERATION_DEFINITION);
@@ -295,6 +298,13 @@ 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 track = asyncHelpers?.track;
+    expect(track).to.be.a('function');
+    expect(resolvedInfo?.getAsyncHelpers().track).to.equal(track);
+    track?.([Promise.resolve()]);
+
     resolve();
 
     await result;

src/execution/createSharedExecutionContext.ts

@@ -1,12 +1,38 @@
+import type { PromiseOrValue } from '../jsutils/PromiseOrValue.js';
+
+import type { GraphQLResolveInfoHelpers } from '../type/index.js';
+
+import { AsyncWorkTracker } from './AsyncWorkTracker.js';
+
 /** @internal */
 export interface SharedExecutionContext {
+  asyncWorkTracker: AsyncWorkTracker;
   getAbortSignal: () => AbortSignal | undefined;
+  getAsyncHelpers: () => GraphQLResolveInfoHelpers;
+  promiseAll: <T>(
+    values: ReadonlyArray<PromiseOrValue<T>>,
+  ) => Promise<Array<T>>;
 }
 
 export function createSharedExecutionContext(
   abortSignal: AbortSignal | undefined,
 ): SharedExecutionContext {
+  const asyncWorkTracker = new AsyncWorkTracker();
+  let resolveInfoHelpers: GraphQLResolveInfoHelpers | undefined;
+
+  const promiseAll = <T>(
+    values: ReadonlyArray<PromiseOrValue<T>>,
+  ): Promise<Array<T>> => asyncWorkTracker.promiseAllTrackOnReject(values);
+
+  const getAsyncHelpers = (): GraphQLResolveInfoHelpers =>
+    (resolveInfoHelpers ??= {
+      track: (maybePromises) => asyncWorkTracker.addValues(maybePromises),
+    });
+
   return {
+    asyncWorkTracker,
     getAbortSignal: () => abortSignal,
+    getAsyncHelpers,
+    promiseAll,
   };
 }