feat(execution): expose async-work-finished execution hook

Expose an execution hook that runs once tracked async work has finished,
including work that can outlive a returned execution result.

The hook is available through `hooks.asyncWorkFinished` on execution args.

An `execute` wrapper that waits for all tracked async work can be written as:

```ts
function executeAndWaitForAsyncWorkFinished(
  args: ExecutionArgs,
): PromiseOrValue<ExecutionResult> {
  let hookHasFired = false;
  const { promise: hookFinished, resolve: resolveHookFinished } =
    Promise.withResolvers<void>();

  const userAsyncWorkFinishedHook = args.hooks?.asyncWorkFinished;
  const result = execute({
    ...args,
    hooks: {
      ...args.hooks,
      asyncWorkFinished(info) {
        try {
          userAsyncWorkFinishedHook?.(info);
        } finally {
          hookHasFired = true;
          resolveHookFinished();
        }
      },
    },
  });

  return hookHasFired ? result : hookFinished.then(() => result);
}
```

To ensure resolver-side async work is also tracked and awaited, use
`info.getAsyncHelpers().track(...)` or `info.getAsyncHelpers().promiseAll(...)`.

`promiseAll(...)` is optimized for the common case where the returned promise is
awaited (or returned) from the resolver. It only starts tracking on rejection,
and does so as a side-effect. Un-awaited async side effects are an anti-pattern:

```ts
resolve(_source, _args, _context, info) {
  const { promiseAll } = info.getAsyncHelpers();
  promiseAll([Promise.reject(new Error('bad')), pendingCleanup]).catch(
    () => undefined,
  );
  return 'ok';
}
```

In that anti-pattern, tracking starts only after rejection (on a later
microtask), so this work is not guaranteed to delay
`hooks.asyncWorkFinished`.

Use `track(...)` for un-awaited async side effects:

```ts
resolve(_source, _args, _context, info) {
  const { track } = info.getAsyncHelpers();
  track([doCleanupAsync().catch(() => undefined)]);
  return 'ok';
}
```

Author: yaacovCR

src/execution/AsyncWorkTracker.ts

@@ -30,6 +30,15 @@ export class AsyncWorkTracker {
     }
   }
 
+  wait(): Promise<void> | void {
+    // wait can complete synchronously when there is no tracked async work,
+    // which allows synchronous execution paths to remain synchronous.
+    if (this.pendingAsyncWork.size === 0) {
+      return;
+    }
+    return this.waitForPendingAsyncWork();
+  }
+
   promiseAllTrackOnReject<T>(
     values: ReadonlyArray<PromiseOrValue<T>>,
   ): Promise<Array<T>> {
@@ -39,4 +48,11 @@ export class AsyncWorkTracker {
     });
     return promise;
   }
+
+  private async waitForPendingAsyncWork(): Promise<void> {
+    while (this.pendingAsyncWork.size > 0) {
+      // eslint-disable-next-line no-await-in-loop
+      await Promise.allSettled(Array.from(this.pendingAsyncWork));
+    }
+  }
 }

src/execution/Executor.ts

@@ -61,6 +61,8 @@ import { createSharedExecutionContext } from './createSharedExecutionContext.js'
 import { buildResolveInfo } from './execute.js';
 import type { StreamUsage } from './getStreamUsage.js';
 import { getStreamUsage as _getStreamUsage } from './getStreamUsage.js';
+import type { ExecutionHooks } from './hooks.js';
+import { runAsyncWorkFinishedHook } from './hooks.js';
 import { returnIteratorCatchingErrors } from './returnIteratorCatchingErrors.js';
 import type { VariableValues } from './values.js';
 import { getArgumentValues } from './values.js';
@@ -116,6 +118,7 @@ export interface ValidatedExecutionArgs {
   errorPropagation: boolean;
   externalAbortSignal: AbortSignal | undefined;
   enableEarlyExecution: boolean;
+  hooks: ExecutionHooks | undefined;
 }
 
 /**
@@ -370,17 +373,31 @@ export class Executor<
     this.aborted = true;
   }
 
+  finishSharedExecution(): void {
+    this.resolverAbortController?.abort();
+    const asyncWorkFinishedHook =
+      this.validatedExecutionArgs.hooks?.asyncWorkFinished;
+    if (asyncWorkFinishedHook === undefined) {
+      return;
+    }
+    runAsyncWorkFinishedHook(
+      this.validatedExecutionArgs,
+      this.sharedExecutionContext,
+      asyncWorkFinishedHook,
+    );
+  }
+
   /**
    * Given a completed execution context and data, build the `{ errors, data }`
    * response defined by the "Response" section of the GraphQL specification.
    */
   buildResponse(
     data: ObjMap<unknown> | null,
   ): ExecutionResult | TAlternativeInitialResponse {
+    this.finishSharedExecution();
     this.finish();
     const errors = this.collectedErrors.errors;
     const result = errors.length ? { errors, data } : { data };
-    this.resolverAbortController?.abort();
     return result;
   }
 

src/execution/__tests__/AsyncWorkTracker-test.ts

@@ -103,3 +103,83 @@ describe('promiseAllTrackOnReject', () => {
     expect(unhandledRejection).to.equal(null);
   });
 });
+
+describe('wait', () => {
+  it('returns synchronously when there is no pending async work', () => {
+    const tracker = new AsyncWorkTracker();
+    expect(tracker.wait()).to.equal(undefined);
+  });
+
+  it('waits when tracked async work is present', async () => {
+    const tracker = new AsyncWorkTracker();
+
+    const delayed = promiseWithResolvers<undefined>();
+    tracker.add(delayed.promise);
+
+    let settled = false;
+    const maybeWait = tracker.wait();
+    expect(maybeWait).to.be.instanceOf(Promise);
+    const wait = Promise.resolve(maybeWait).then(() => {
+      settled = true;
+    });
+    await resolveOnNextTick();
+    expect(settled).to.equal(false);
+
+    delayed.resolve(undefined);
+    await wait;
+    expect(settled).to.equal(true);
+  });
+
+  it('keeps waiting when tracked async work is followed by more tracked async work', async () => {
+    const tracker = new AsyncWorkTracker();
+
+    const delayed = promiseWithResolvers<undefined>();
+    tracker.add(delayed.promise);
+
+    let settled = false;
+    const maybeWait = tracker.wait();
+    expect(maybeWait).to.be.instanceOf(Promise);
+    const wait = Promise.resolve(maybeWait).then(() => {
+      settled = true;
+    });
+    await resolveOnNextTick();
+    expect(settled).to.equal(false);
+
+    delayed.resolve(undefined);
+
+    const anotherDelayed = promiseWithResolvers<undefined>();
+    tracker.add(anotherDelayed.promise);
+    await resolveOnNextTick();
+    expect(settled).to.equal(false);
+
+    anotherDelayed.resolve(undefined);
+
+    await wait;
+    expect(settled).to.equal(true);
+  });
+
+  it('does not wait for side-effect promiseAll tracking added on next microtask', async () => {
+    const tracker = new AsyncWorkTracker();
+    const delayed = promiseWithResolvers<undefined>();
+
+    tracker
+      .promiseAllTrackOnReject([
+        Promise.reject(new Error('bad')),
+        delayed.promise,
+      ])
+      .catch(() => undefined);
+
+    const maybeWait = tracker.wait();
+    expect(maybeWait).to.equal(undefined);
+
+    await resolveOnNextTick();
+    expect(tracker.pendingAsyncWork.size).to.equal(0);
+    await resolveOnNextTick();
+    expect(tracker.pendingAsyncWork.size).to.be.greaterThan(0);
+
+    delayed.resolve(undefined);
+    await Promise.allSettled(Array.from(tracker.pendingAsyncWork));
+    await resolveOnNextTick();
+    expect(tracker.pendingAsyncWork.size).to.equal(0);
+  });
+});

src/execution/__tests__/executor-test.ts

@@ -261,7 +261,7 @@ describe('Execute: Handles basic execution tasks', () => {
       'getAsyncHelpers',
     );
     const asyncHelpers = resolvedInfo?.getAsyncHelpers();
-    expect(asyncHelpers).to.have.all.keys('track');
+    expect(asyncHelpers).to.have.all.keys('promiseAll', 'track');
 
     const operation = document.definitions[0];
     assert(operation.kind === Kind.OPERATION_DEFINITION);
@@ -300,6 +300,10 @@ describe('Execute: Handles basic execution tasks', () => {
 
     expect(resolvedInfo?.getAsyncHelpers()).to.equal(asyncHelpers);
 
+    const promiseAll = asyncHelpers?.promiseAll;
+    expect(promiseAll).to.be.a('function');
+    expect(resolvedInfo?.getAsyncHelpers().promiseAll).to.equal(promiseAll);
+
     const track = asyncHelpers?.track;
     expect(track).to.be.a('function');
     expect(resolvedInfo?.getAsyncHelpers().track).to.equal(track);