feat(execution): expose async work finished execution hook

Author: yaacovCR

src/execution/Executor.ts

@@ -62,6 +62,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 { runPostExecutionHooks } from './hooks.js';
 import { returnIteratorCatchingErrors } from './returnIteratorCatchingErrors.js';
 import type { VariableValues } from './values.js';
 import { getArgumentValues } from './values.js';
@@ -117,6 +119,7 @@ export interface ValidatedExecutionArgs {
   errorPropagation: boolean;
   externalAbortSignal: AbortSignal | undefined;
   enableEarlyExecution: boolean;
+  experimentalHooks: ExecutionHooks | undefined;
 }
 
 /**
@@ -368,9 +371,23 @@ export class Executor<
     }
     this.abortResultPromise?.(this.abortReason);
     this.resolverAbortController?.abort(this.abortReason);
+    const hooks = this.validatedExecutionArgs.experimentalHooks;
+    if (hooks !== undefined) {
+      runPostExecutionHooks(
+        this.validatedExecutionArgs,
+        this.sharedExecutionContext,
+      );
+    }
   }
 
   finish(): void {
+    const hooks = this.validatedExecutionArgs.experimentalHooks;
+    if (hooks !== undefined) {
+      runPostExecutionHooks(
+        this.validatedExecutionArgs,
+        this.sharedExecutionContext,
+      );
+    }
     this.throwIfAborted();
     this.aborted = true;
   }

src/execution/__tests__/hooks-test.ts

@@ -0,0 +1,293 @@
+import { assert, expect } from 'chai';
+import { describe, it } from 'mocha';
+
+import { expectPromise } from '../../__testUtils__/expectPromise.js';
+import { resolveOnNextTick } from '../../__testUtils__/resolveOnNextTick.js';
+
+import { promiseWithResolvers } from '../../jsutils/promiseWithResolvers.js';
+
+import { parse } from '../../language/parser.js';
+
+import type { GraphQLResolveInfo } from '../../type/definition.js';
+import { GraphQLObjectType } from '../../type/definition.js';
+import { GraphQLString } from '../../type/scalars.js';
+import { GraphQLSchema } from '../../type/schema.js';
+
+import { buildSchema } from '../../utilities/buildASTSchema.js';
+
+import { createSharedExecutionContext } from '../createSharedExecutionContext.js';
+import { execute } from '../execute.js';
+import type { ValidatedExecutionArgs } from '../Executor.js';
+import { runPostExecutionHooks } from '../hooks.js';
+
+const executeHookSchema = new GraphQLSchema({
+  query: new GraphQLObjectType({
+    name: 'Query',
+    fields: {
+      test: {
+        type: GraphQLString,
+        resolve: () => 'ok',
+      },
+    },
+  }),
+});
+
+const cancellationHookSchema = buildSchema(`
+  type Todo {
+    id: ID
+    items: [String]
+  }
+
+  type Query {
+    todo: Todo
+  }
+`);
+
+describe('Execute: Hooks', () => {
+  it('ignores errors thrown by hooks', async () => {
+    const calls: Array<string> = [];
+    const { promise: hooksFinished, resolve: resolveHooksFinished } =
+      promiseWithResolvers<undefined>();
+
+    const result = execute({
+      schema: executeHookSchema,
+      document: parse('{ test }'),
+      experimentalHooks: {
+        queryOrMutationOrSubscriptionEventAsyncWorkFinished() {
+          calls.push('asyncWork');
+          resolveHooksFinished(undefined);
+          throw new Error(
+            'queryOrMutationOrSubscriptionEventAsyncWorkFinished failed',
+          );
+        },
+      },
+    });
+
+    expect(result).to.deep.equal({
+      data: {
+        test: 'ok',
+      },
+    });
+    await hooksFinished;
+    expect(calls).to.deep.equal(['asyncWork']);
+  });
+
+  it('runs post execution hooks for synchronous execution', async () => {
+    const calls: Array<string> = [];
+    const { promise: hooksFinished, resolve: resolveHooksFinished } =
+      promiseWithResolvers<undefined>();
+
+    const result = execute({
+      schema: executeHookSchema,
+      document: parse('{ test }'),
+      experimentalHooks: {
+        queryOrMutationOrSubscriptionEventAsyncWorkFinished() {
+          calls.push('asyncWork');
+          resolveHooksFinished(undefined);
+        },
+      },
+    });
+
+    expect(result).to.deep.equal({
+      data: {
+        test: 'ok',
+      },
+    });
+    await hooksFinished;
+    expect(calls).to.deep.equal(['asyncWork']);
+  });
+
+  it('runs post execution hooks for asynchronous execution', async () => {
+    const { promise: resolvedValue, resolve } = promiseWithResolvers<string>();
+    const calls: Array<string> = [];
+    const { promise: hooksFinished, resolve: resolveHooksFinished } =
+      promiseWithResolvers<undefined>();
+    const asyncSchema = new GraphQLSchema({
+      query: new GraphQLObjectType({
+        name: 'Query',
+        fields: {
+          test: {
+            type: GraphQLString,
+            resolve: () => resolvedValue,
+          },
+        },
+      }),
+    });
+
+    const resultPromise = execute({
+      schema: asyncSchema,
+      document: parse('{ test }'),
+      experimentalHooks: {
+        queryOrMutationOrSubscriptionEventAsyncWorkFinished() {
+          calls.push('asyncWork');
+          resolveHooksFinished(undefined);
+        },
+      },
+    });
+
+    expect(calls).to.deep.equal([]);
+    resolve('ok');
+
+    const result = await resultPromise;
+    expect(result).to.deep.equal({
+      data: {
+        test: 'ok',
+      },
+    });
+    await hooksFinished;
+    expect(calls).to.deep.equal(['asyncWork']);
+  });
+
+  it('runs post execution hooks for aborted execution', async () => {
+    const abortController = new AbortController();
+    const { promise: pendingCleanup, resolve: resolveCleanup } =
+      promiseWithResolvers<string>();
+    const { promise: asyncWorkFinished, resolve: resolveAsyncWorkFinished } =
+      promiseWithResolvers<undefined>();
+    const calls: Array<string> = [];
+    const document = parse(`
+      query {
+        todo {
+          id
+        }
+      }
+    `);
+
+    const resultPromise = execute({
+      document,
+      schema: cancellationHookSchema,
+      abortSignal: abortController.signal,
+      experimentalHooks: {
+        queryOrMutationOrSubscriptionEventAsyncWorkFinished() {
+          calls.push('asyncWork');
+          resolveAsyncWorkFinished(undefined);
+        },
+      },
+      rootValue: {
+        todo: (_args: any, _context: any, info: GraphQLResolveInfo) => {
+          const abortSignal = info.getAbortSignal();
+          assert(abortSignal instanceof AbortSignal);
+          const abortPromise = new Promise<never>((_resolve, reject) => {
+            abortSignal.addEventListener('abort', () => {
+              reject(new Error('This operation was aborted'));
+            });
+          });
+          return info
+            .getAsyncHelpers()
+            .promiseAll([abortPromise, pendingCleanup]);
+        },
+      },
+    });
+
+    abortController.abort();
+    await expectPromise(resultPromise).toRejectWith(
+      'This operation was aborted',
+    );
+    expect(calls).to.deep.equal([]);
+
+    resolveCleanup('done');
+    await asyncWorkFinished;
+
+    expect(calls).to.deep.equal(['asyncWork']);
+  });
+
+  it('fires queryOrMutationOrSubscriptionEventAsyncWorkFinished after async iterator return cleanup', async () => {
+    const abortController = new AbortController();
+    const document = parse(`
+      query {
+        todo {
+          items
+        }
+      }
+    `);
+    const { promise: nextReturned, resolve: resolveNextReturned } =
+      promiseWithResolvers<IteratorResult<string>>();
+    const { promise: nextStarted, resolve: resolveNextStarted } =
+      // eslint-disable-next-line @typescript-eslint/no-invalid-void-type
+      promiseWithResolvers<void>();
+    const { promise: returnStarted, resolve: resolveReturnStarted } =
+      // eslint-disable-next-line @typescript-eslint/no-invalid-void-type
+      promiseWithResolvers<void>();
+    const { promise: returnFinished, resolve: resolveReturnFinished } =
+      promiseWithResolvers<IteratorResult<string>>();
+    const { promise: asyncWorkFinished, resolve: resolveAsyncWorkFinished } =
+      promiseWithResolvers<undefined>();
+    const asyncIterator = {
+      [Symbol.asyncIterator]() {
+        return this;
+      },
+      next() {
+        resolveNextStarted(undefined);
+        return nextReturned;
+      },
+      return() {
+        resolveReturnStarted(undefined);
+        return returnFinished;
+      },
+    };
+
+    const resultPromise = execute({
+      schema: cancellationHookSchema,
+      document,
+      abortSignal: abortController.signal,
+      experimentalHooks: {
+        queryOrMutationOrSubscriptionEventAsyncWorkFinished() {
+          resolveAsyncWorkFinished(undefined);
+        },
+      },
+      rootValue: {
+        todo: {
+          items: asyncIterator,
+        },
+      },
+    });
+
+    await nextStarted;
+    abortController.abort();
+    resolveNextReturned({ value: 'value', done: false });
+
+    await expectPromise(resultPromise).toRejectWith(
+      'This operation was aborted',
+    );
+    await returnStarted;
+
+    let asyncWorkHookCalled = false;
+    const asyncWorkObserved = asyncWorkFinished.then(() => {
+      asyncWorkHookCalled = true;
+    });
+    await resolveOnNextTick();
+    expect(asyncWorkHookCalled).to.equal(false);
+
+    resolveReturnFinished({ value: undefined, done: true });
+    await asyncWorkFinished;
+    await asyncWorkObserved;
+  });
+
+  it('runs post execution hooks only once', async () => {
+    const calls: Array<string> = [];
+    const { promise: hooksFinished, resolve: resolveHooksFinished } =
+      promiseWithResolvers<undefined>();
+
+    const validatedExecutionArgs = {
+      experimentalHooks: {
+        queryOrMutationOrSubscriptionEventAsyncWorkFinished() {
+          calls.push('asyncWork');
+          resolveHooksFinished(undefined);
+        },
+      },
+    } as Pick<ValidatedExecutionArgs, 'experimentalHooks'>;
+
+    const sharedExecutionContext = createSharedExecutionContext(undefined);
+    runPostExecutionHooks(
+      validatedExecutionArgs as ValidatedExecutionArgs,
+      sharedExecutionContext,
+    );
+    runPostExecutionHooks(
+      validatedExecutionArgs as ValidatedExecutionArgs,
+      sharedExecutionContext,
+    );
+
+    await hooksFinished;
+    expect(calls).to.deep.equal(['asyncWork']);
+  });
+});

src/execution/createSharedExecutionContext.ts

@@ -7,6 +7,7 @@ import { AsyncWorkTracker } from './AsyncWorkTracker.js';
 /** @internal */
 export interface SharedExecutionContext {
   asyncWorkTracker: AsyncWorkTracker;
+  postExecutionHooksStarted: boolean;
   getAbortSignal: () => AbortSignal | undefined;
   promiseAll: <T>(
     values: ReadonlyArray<PromiseOrValue<T>>,
@@ -75,6 +76,7 @@ export function createSharedExecutionContext(
 
   return {
     asyncWorkTracker,
+    postExecutionHooksStarted: false,
     getAbortSignal: () => abortSignal,
     promiseAll,
     getAsyncHelpers,

src/execution/execute.ts

@@ -40,6 +40,7 @@ import type { ExecutionResult, ValidatedExecutionArgs } from './Executor.js';
 import { Executor } from './Executor.js';
 import { ExecutorThrowingOnIncremental } from './ExecutorThrowingOnIncremental.js';
 import { getVariableSignature } from './getVariableSignature.js';
+import type { ExecutionHooks } from './hooks.js';
 import type { ExperimentalIncrementalExecutionResults } from './incremental/IncrementalExecutor.js';
 import { IncrementalExecutor } from './incremental/IncrementalExecutor.js';
 import { mapAsyncIterable } from './mapAsyncIterable.js';
@@ -298,6 +299,7 @@ export interface ExecutionArgs {
   hideSuggestions?: Maybe<boolean>;
   abortSignal?: Maybe<AbortSignal>;
   enableEarlyExecution?: Maybe<boolean>;
+  experimentalHooks?: Maybe<ExecutionHooks>;
   /** Additional execution options. */
   options?: {
     /** Set the maximum number of errors allowed for coercing (defaults to 50). */
@@ -330,6 +332,7 @@ export function validateExecutionArgs(
     perEventExecutor,
     abortSignal: externalAbortSignal,
     enableEarlyExecution,
+    experimentalHooks,
     options,
   } = args;
 
@@ -419,6 +422,7 @@ export function validateExecutionArgs(
     errorPropagation,
     externalAbortSignal: externalAbortSignal ?? undefined,
     enableEarlyExecution: enableEarlyExecution === true,
+    experimentalHooks: experimentalHooks ?? undefined,
   };
 }