use single phase of resolver abortSignal cancellation

Previously, we used a lazily created abortSignal per resolver so as to shift triggering the abortSignal for a resolver that returned a streamed asyncIterable from the original executor to the stream. This is expensive and difficult to reason about and has been replaced by a single phase of cancellation when the entire execution finishes.

Note: when executing a subscription, we are still able to cancel resolver abortSignals after execution of each subscription event.

Author: yaacovCR

src/execution/Executor.ts

@@ -217,42 +217,41 @@ export class Executor<
 > {
   validatedExecutionArgs: ValidatedExecutionArgs;
   finished: boolean;
-  initialResponseAbortController: AbortController | undefined;
-  resolverAbortControllers: Map<Path, AbortController>;
   collectedErrors: CollectedErrors;
+  internalAbortController: AbortController;
+  removeExternalAbortListener: (() => void) | undefined;
+  resolverAbortController: AbortController | undefined;
+  sharedResolverAbortSignal: AbortSignal;
 
-  constructor(validatedExecutionArgs: ValidatedExecutionArgs) {
+  constructor(
+    validatedExecutionArgs: ValidatedExecutionArgs,
+    sharedResolverAbortSignal?: AbortSignal,
+  ) {
     this.validatedExecutionArgs = validatedExecutionArgs;
     this.finished = false;
-    this.resolverAbortControllers = new Map();
     this.collectedErrors = new CollectedErrors();
+    this.internalAbortController = new AbortController();
+
+    if (sharedResolverAbortSignal === undefined) {
+      this.resolverAbortController = new AbortController();
+      this.sharedResolverAbortSignal = this.resolverAbortController.signal;
+    } else {
+      this.sharedResolverAbortSignal = sharedResolverAbortSignal;
+    }
   }
 
   executeQueryOrMutationOrSubscriptionEvent(): PromiseOrValue<
     ExecutionResult | TAlternativeInitialResponse
   > {
-    const abortController = (this.initialResponseAbortController =
-      new AbortController());
-
-    const validatedExecutionArgs = this.validatedExecutionArgs;
-    const externalAbortSignal = validatedExecutionArgs.externalAbortSignal;
-    let removeAbortListener: (() => void) | undefined;
+    const externalAbortSignal = this.validatedExecutionArgs.externalAbortSignal;
     if (externalAbortSignal) {
-      if (externalAbortSignal.aborted) {
-        throw new Error(externalAbortSignal.reason);
-      }
+      externalAbortSignal.throwIfAborted();
       const onExternalAbort = () => this.cancel(externalAbortSignal.reason);
-      removeAbortListener = () =>
+      this.removeExternalAbortListener = () =>
         externalAbortSignal.removeEventListener('abort', onExternalAbort);
       externalAbortSignal.addEventListener('abort', onExternalAbort);
     }
 
-    const onFinish = () => {
-      removeAbortListener?.();
-      this.finish();
-      abortController.signal.throwIfAborted();
-    };
-
     try {
       const {
         schema,
@@ -261,7 +260,7 @@ export class Executor<
         operation,
         variableValues,
         hideSuggestions,
-      } = validatedExecutionArgs;
+      } = this.validatedExecutionArgs;
 
       const { operation: operationType, selectionSet } = operation;
 
@@ -293,62 +292,40 @@ export class Executor<
       if (isPromise(result)) {
         const promise = result.then(
           (data) => {
-            onFinish();
+            this.finish();
             return this.buildResponse(data);
           },
           (error: unknown) => {
-            onFinish();
+            this.finish();
             this.collectedErrors.add(error as GraphQLError, undefined);
             return this.buildResponse(null);
           },
         );
-        return externalAbortSignal
-          ? cancellablePromise(promise, abortController.signal)
-          : promise;
+        return cancellablePromise(promise, this.internalAbortController.signal);
       }
-      onFinish();
+      this.finish();
       return this.buildResponse(result);
     } catch (error) {
+      this.finish();
       this.collectedErrors.add(error as GraphQLError, undefined);
-      onFinish();
       return this.buildResponse(null);
     }
   }
 
-  cancel(reason: unknown): void {
+  cancel(reason?: unknown): void {
     if (!this.finished) {
-      this.initialResponseAbortController?.abort(reason);
-      this.finish(reason);
+      this.finish();
+      this.internalAbortController.abort(reason);
+      this.resolverAbortController?.abort(reason);
     }
   }
 
-  finish(reason?: unknown): void {
+  finish(): void {
     if (!this.finished) {
       this.finished = true;
-      this.triggerResolverAbortSignals(reason);
-    }
-  }
-
-  triggerResolverAbortSignals(reason?: unknown): void {
-    const { resolverAbortControllers } = this;
-    const finishReason =
-      reason ?? new Error('Execution has already completed.');
-    for (const abortController of resolverAbortControllers.values()) {
-      abortController.abort(finishReason);
+      this.removeExternalAbortListener?.();
     }
-  }
-
-  getAbortSignal(path: Path): AbortSignal {
-    const resolverAbortSignal = this.resolverAbortControllers.get(path)?.signal;
-    if (resolverAbortSignal !== undefined) {
-      return resolverAbortSignal;
-    }
-    const abortController = new AbortController();
-    this.resolverAbortControllers.set(path, abortController);
-    if (this.finished) {
-      abortController.abort(new Error('Execution has already completed.'));
-    }
-    return abortController.signal;
+    this.internalAbortController.signal.throwIfAborted();
   }
 
   /**
@@ -358,6 +335,7 @@ export class Executor<
   buildResponse(
     data: ObjMap<unknown> | null,
   ): ExecutionResult | TAlternativeInitialResponse {
+    this.resolverAbortController?.abort();
     const errors = this.collectedErrors.errors;
     return errors.length ? { errors, data } : { data };
   }
@@ -542,7 +520,7 @@ export class Executor<
       toNodes(fieldDetailsList),
       parentType,
       path,
-      () => this.getAbortSignal(path),
+      () => this.sharedResolverAbortSignal,
     );
 
     // Get the resolve function, regardless of if its result is normal or abrupt (error).
@@ -571,7 +549,6 @@ export class Executor<
           path,
           result,
           positionContext,
-          true,
         );
       }
 
@@ -587,22 +564,13 @@ export class Executor<
       if (isPromise(completed)) {
         // Note: we don't rely on a `catch` method, but we do expect "thenable"
         // to take a second callback for the error case.
-        return completed.then(
-          (resolved) => {
-            this.resolverAbortControllers.delete(path);
-            return resolved;
-          },
-          (rawError: unknown) => {
-            this.resolverAbortControllers.delete(path);
-            this.handleFieldError(rawError, returnType, fieldDetailsList, path);
-            return null;
-          },
-        );
+        return completed.then(undefined, (rawError: unknown) => {
+          this.handleFieldError(rawError, returnType, fieldDetailsList, path);
+          return null;
+        });
       }
-      this.resolverAbortControllers.delete(path);
       return completed;
     } catch (rawError) {
-      this.resolverAbortControllers.delete(path);
       this.handleFieldError(rawError, returnType, fieldDetailsList, path);
       return null;
     }
@@ -755,7 +723,6 @@ export class Executor<
     path: Path,
     result: Promise<unknown>,
     positionContext: TPositionContext | undefined,
-    isFieldValue?: boolean,
   ): Promise<unknown> {
     try {
       const resolved = await result;
@@ -774,14 +741,8 @@ export class Executor<
       if (isPromise(completed)) {
         completed = await completed;
       }
-      if (isFieldValue) {
-        this.resolverAbortControllers.delete(path);
-      }
       return completed;
     } catch (rawError) {
-      if (isFieldValue) {
-        this.resolverAbortControllers.delete(path);
-      }
       this.handleFieldError(rawError, returnType, fieldDetailsList, path);
       return null;
     }

src/execution/incremental/IncrementalExecutor.ts

@@ -264,9 +264,10 @@ export class IncrementalExecutor<
 
   constructor(
     validatedExecutionArgs: ValidatedExecutionArgs,
+    sharedResolverAbortSignal?: AbortSignal,
     deferUsageSet?: DeferUsageSet,
   ) {
-    super(validatedExecutionArgs);
+    super(validatedExecutionArgs, sharedResolverAbortSignal);
     this.deferUsageSet = deferUsageSet;
     this.groups = [];
     this.tasks = [];
@@ -276,7 +277,11 @@ export class IncrementalExecutor<
   createSubExecutor(
     deferUsageSet?: DeferUsageSet,
   ): IncrementalExecutor<TExperimental> {
-    return new IncrementalExecutor(this.validatedExecutionArgs, deferUsageSet);
+    return new IncrementalExecutor(
+      this.validatedExecutionArgs,
+      this.sharedResolverAbortSignal,
+      deferUsageSet,
+    );
   }
 
   override cancel(reason?: unknown): void {
@@ -296,20 +301,21 @@ export class IncrementalExecutor<
   override buildResponse(
     data: ObjMap<unknown> | null,
   ): ExecutionResult | TExperimental {
-    const errors = this.collectedErrors.errors;
     const work = this.getIncrementalWork();
     const { tasks, streams } = work;
     if (tasks?.length === 0 && streams?.length === 0) {
-      return errors.length ? { errors, data } : { data };
+      return super.buildResponse(data);
     }
 
+    const errors = this.collectedErrors.errors;
     invariant(data !== null);
     const incrementalPublisher = new IncrementalPublisher();
     return incrementalPublisher.buildResponse(
       data,
       errors,
       work,
       this.validatedExecutionArgs.externalAbortSignal,
+      () => this.resolverAbortController?.abort(),
     ) as TExperimental;
   }
 
@@ -509,6 +515,7 @@ export class IncrementalExecutor<
         (resolved) =>
           this.buildExecutionGroupResult(deliveryGroups, path, resolved),
         (error: unknown) => {
+          this.cancel();
           throw error;
         },
       );

src/execution/incremental/IncrementalPublisher.ts

@@ -47,6 +47,7 @@ export class IncrementalPublisher {
     errors: ReadonlyArray<GraphQLError>,
     work: IncrementalWork,
     abortSignal: AbortSignal | undefined,
+    onFinished: () => void,
   ): ExperimentalIncrementalExecutionResults {
     const { initialGroups, initialStreams, events } = createWorkQueue<
       ExecutionGroupValue,
@@ -61,12 +62,13 @@ export class IncrementalPublisher {
       });
     }
 
-    let onWorkQueueFinished: (() => void) | undefined;
     if (abortSignal) {
       abortSignal.addEventListener('abort', abort);
-      onWorkQueueFinished = () =>
-        abortSignal.removeEventListener('abort', abort);
     }
+    const onWorkQueueFinished = () => {
+      onFinished();
+      abortSignal?.removeEventListener('abort', abort);
+    };
 
     const pending = this._toPendingResults(initialGroups, initialStreams);
 
@@ -78,7 +80,7 @@ export class IncrementalPublisher {
       mapAsyncIterable(events, (batch) =>
         this._handleBatch(batch, onWorkQueueFinished),
       ),
-      () => onWorkQueueFinished?.(),
+      () => onWorkQueueFinished(),
     );
 
     return {

src/execution/legacyIncremental/BranchingIncrementalExecutor.ts

@@ -89,27 +89,29 @@ export class BranchingIncrementalExecutor extends IncrementalExecutor<Experiment
   ): IncrementalExecutor<ExperimentalIncrementalExecutionResults> {
     return new BranchingIncrementalExecutor(
       this.validatedExecutionArgs,
+      this.sharedResolverAbortSignal,
       deferUsageSet,
     );
   }
 
   override buildResponse(
     data: ObjMap<unknown> | null,
   ): ExecutionResult | ExperimentalIncrementalExecutionResults {
-    const errors = this.collectedErrors.errors;
     const work = this.getIncrementalWork();
     const { tasks, streams } = work;
     if (tasks?.length === 0 && streams?.length === 0) {
-      return errors.length ? { errors, data } : { data };
+      return super.buildResponse(data);
     }
 
+    const errors = this.collectedErrors.errors;
     invariant(data !== null);
     const incrementalPublisher = new BranchingIncrementalPublisher();
     return incrementalPublisher.buildResponse(
       data,
       errors,
       work,
       this.validatedExecutionArgs.externalAbortSignal,
+      () => this.resolverAbortController?.abort(),
     );
   }
 

src/execution/legacyIncremental/BranchingIncrementalPublisher.ts

@@ -44,6 +44,7 @@ export class BranchingIncrementalPublisher {
     errors: ReadonlyArray<GraphQLError>,
     work: IncrementalWork,
     abortSignal: AbortSignal | undefined,
+    onFinished: () => void,
   ): ExperimentalIncrementalExecutionResults {
     const { initialStreams, events } = createWorkQueue<
       ExecutionGroupValue,
@@ -62,12 +63,13 @@ export class BranchingIncrementalPublisher {
       });
     }
 
-    let onWorkQueueFinished: (() => void) | undefined;
     if (abortSignal) {
       abortSignal.addEventListener('abort', abort);
-      onWorkQueueFinished = () =>
-        abortSignal.removeEventListener('abort', abort);
     }
+    const onWorkQueueFinished = () => {
+      onFinished();
+      abortSignal?.removeEventListener('abort', abort);
+    };
 
     const initialResult: InitialIncrementalExecutionResult = errors.length
       ? { errors, data, hasNext: true }
@@ -77,7 +79,7 @@ export class BranchingIncrementalPublisher {
       mapAsyncIterable(events, (batch) =>
         this._handleBatch(batch, onWorkQueueFinished),
       ),
-      () => onWorkQueueFinished?.(),
+      () => onWorkQueueFinished(),
     );
 
     return {