allow forwarding of cancellation reasons within computations/queues (#4636)

This is currently not directly testable (I think!) without dipping into Executor internals because resolvers currently use a shared abort signal, and so at this point it's just for polish.

Although that may change in the future and this hardens the shape of our abort semantics.

Author: yaacovCR

src/execution/incremental/Computation.ts

@@ -9,9 +9,12 @@ type MaybePromise<T> =
 /** @internal **/
 export class Computation<T> {
   private _fn: () => PromiseOrValue<T>;
-  private _onCancel: (() => void) | undefined;
+  private _onCancel: ((reason?: unknown) => void) | undefined;
   private _maybePromise?: MaybePromise<T>;
-  constructor(fn: () => PromiseOrValue<T>, onCancel?: () => void) {
+  constructor(
+    fn: () => PromiseOrValue<T>,
+    onCancel?: (reason?: unknown) => void,
+  ) {
     this._fn = fn;
     this._onCancel = onCancel;
   }
@@ -51,21 +54,21 @@ export class Computation<T> {
       }
     }
   }
-  cancel(): void {
+  cancel(reason?: unknown): void {
     const maybePromise = this._maybePromise;
     if (!maybePromise) {
       this._maybePromise = {
         status: 'rejected',
-        reason: new Error('Cancelled!'),
+        reason,
       };
       return;
     }
     const status = maybePromise.status;
     if (status === 'pending' && this._onCancel) {
-      this._onCancel();
+      this._onCancel(reason);
       this._maybePromise = {
         status: 'rejected',
-        reason: new Error('Cancelled!'),
+        reason,
       };
     }
   }

src/execution/incremental/IncrementalExecutor.ts

@@ -321,10 +321,10 @@ export class IncrementalExecutor<
   override cancel(reason?: unknown): void {
     super.cancel(reason);
     for (const task of this.tasks) {
-      task.computation.cancel();
+      task.computation.cancel(reason);
     }
     for (const stream of this.streams) {
-      stream.queue.abort();
+      stream.queue.abort(reason);
     }
   }
 
@@ -503,7 +503,7 @@ export class IncrementalExecutor<
               groupedFieldSet,
               deliveryGroupMap,
             ),
-          () => executor.cancel(),
+          (reason) => executor.cancel(reason),
         ),
       };
 
@@ -581,10 +581,14 @@ export class IncrementalExecutor<
       return { groups, tasks, streams };
     }
 
+    const cancellationReason = new Error(
+      'Cancelled secondary to null within original result',
+    );
+
     const filteredTasks: Array<ExecutionGroup> = [];
     for (const task of tasks) {
       if (collectedErrors.hasNulledPosition(task.path)) {
-        task.computation.cancel();
+        task.computation.cancel(cancellationReason);
       } else {
         filteredTasks.push(task);
       }
@@ -593,7 +597,7 @@ export class IncrementalExecutor<
     const filteredStreams: Array<ItemStream> = [];
     for (const stream of streams) {
       if (collectedErrors.hasNulledPosition(stream.path)) {
-        stream.queue.cancel();
+        stream.queue.abort(cancellationReason);
       } else {
         filteredStreams.push(stream);
       }
@@ -713,11 +717,13 @@ export class IncrementalExecutor<
     const { enableEarlyExecution } = this.validatedExecutionArgs;
     const queue = new Queue<StreamItemResult>(
       async ({ push, stop, started, stopped }) => {
-        const cancelStreamItems = new Set<() => void>();
+        const cancelStreamItems = new Set<(reason?: unknown) => void>();
 
         // eslint-disable-next-line @typescript-eslint/no-floating-promises
-        stopped.then(() => {
-          cancelStreamItems.forEach((cancelStreamItem) => cancelStreamItem());
+        stopped.then((reason) => {
+          cancelStreamItems.forEach((cancelStreamItem) =>
+            cancelStreamItem(reason),
+          );
           returnIteratorCatchingErrors(iterator);
         });
         await (enableEarlyExecution ? Promise.resolve() : started);
@@ -763,7 +769,8 @@ export class IncrementalExecutor<
           );
           if (isPromise(streamItemResult)) {
             if (enableEarlyExecution) {
-              const cancelStreamItem = () => executor.cancel();
+              const cancelStreamItem = (reason?: unknown) =>
+                executor.cancel(reason);
               cancelStreamItems.add(cancelStreamItem);
               streamItemResult = streamItemResult.finally(() => {
                 cancelStreamItems.delete(cancelStreamItem);

src/execution/incremental/Queue.ts

@@ -87,7 +87,7 @@ export class Queue<T> {
   private _batchRequests = new Set<BatchRequest<T>>();
 
   private _resolveStarted: () => void;
-  private _resolveStopped: () => void;
+  private _resolveStopped: (reason?: unknown) => void;
 
   constructor(
     executor: ({
@@ -106,8 +106,7 @@ export class Queue<T> {
 
     this._resolveStarted = resolveStarted;
     const { promise: stopped, resolve: resolveStopped } =
-      // eslint-disable-next-line @typescript-eslint/no-invalid-void-type
-      promiseWithResolvers<void>();
+      promiseWithResolvers<unknown>();
     this._resolveStopped = resolveStopped;
 
     try {
@@ -152,7 +151,7 @@ export class Queue<T> {
     if (this._isStopped) {
       return;
     }
-    this._terminate();
+    this._terminate(reason);
     if (this._batchRequests.size) {
       this._batchRequests.forEach((request) => request.reject(reason));
       this._batchRequests.clear();
@@ -285,7 +284,7 @@ export class Queue<T> {
     return maybePushPromise;
   }
 
-  private _terminate(): void {
+  private _terminate(reason?: unknown): void {
     for (const entry of this._entries) {
       if (entry.kind === 'item') {
         this._release();
@@ -294,7 +293,7 @@ export class Queue<T> {
     this._entries.length = 0;
     this._stopRequested = true;
     this._isStopped = true;
-    this._resolveStopped();
+    this._resolveStopped(reason);
   }
 
   private _stop(reason?: unknown): void {
@@ -343,7 +342,7 @@ export class Queue<T> {
         this._entries.shift();
         this._release();
         this._isStopped = true;
-        this._resolveStopped();
+        this._resolveStopped(settled.reason);
         this._batchRequests = new Set();
         requests.forEach((request) => request.reject(settled.reason));
       }

src/execution/incremental/WorkQueue.ts

@@ -238,7 +238,7 @@ export function createWorkQueue<
         }
       });
       // eslint-disable-next-line @typescript-eslint/no-floating-promises
-      stopped.then(() => cancel());
+      stopped.then((reason) => cancel(reason));
     },
     1,
   ).subscribe((graphEvents) => handleGraphEvents(graphEvents));
@@ -249,39 +249,39 @@ export function createWorkQueue<
     events,
   };
 
-  function cancel(): void {
+  function cancel(reason?: unknown): void {
     for (const group of rootGroups) {
-      cancelGroup(group);
+      cancelGroup(group, reason);
     }
     for (const stream of rootStreams) {
-      cancelStream(stream);
+      cancelStream(stream, reason);
     }
   }
 
-  function cancelGroup(group: G): void {
+  function cancelGroup(group: G, reason?: unknown): void {
     const groupNode = groupNodes.get(group);
     if (groupNode) {
       for (const task of groupNode.tasks) {
-        cancelTask(task);
+        cancelTask(task, reason);
       }
       for (const childGroup of groupNode.childGroups) {
-        cancelGroup(childGroup);
+        cancelGroup(childGroup, reason);
       }
     }
   }
 
-  function cancelTask(task: Task<T, I, G, S>): void {
-    task.computation.cancel();
+  function cancelTask(task: Task<T, I, G, S>, reason?: unknown): void {
+    task.computation.cancel(reason);
     const taskNode = taskNodes.get(task);
     if (taskNode) {
       for (const childStream of taskNode.childStreams) {
-        cancelStream(childStream);
+        cancelStream(childStream, reason);
       }
     }
   }
 
-  function cancelStream(stream: S): void {
-    stream.queue.cancel();
+  function cancelStream(stream: S, reason?: unknown): void {
+    stream.queue.abort(reason);
   }
 
   function maybeIntegrateWork(

src/execution/incremental/__tests__/Computation-test.ts

@@ -168,4 +168,31 @@ describe('Computation', () => {
     expect(onCancelRan).to.equal(true);
     expect(() => computation.result()).to.throw('Cancelled!');
   });
+
+  it('can be cancelled with a provided reason before running', () => {
+    const abortReason = new Error('aborted');
+    const computation = new Computation(() => ({ value: 123 }));
+
+    computation.cancel(abortReason);
+    expect(() => computation.result()).to.throw('aborted');
+  });
+
+  it('forwards cancellation reason to onCancel while running asynchronously', () => {
+    const abortReason = new Error('aborted');
+    let onCancelReason: unknown;
+    const computation = new Computation(
+      () =>
+        new Promise(() => {
+          // Never resolves.
+        }),
+      (reason) => {
+        onCancelReason = reason;
+      },
+    );
+
+    computation.prime();
+    computation.cancel(abortReason);
+    expect(onCancelReason).to.equal(abortReason);
+    expect(() => computation.result()).to.throw('aborted');
+  });
 });

src/execution/incremental/__tests__/Queue-test.ts

@@ -157,6 +157,19 @@ describe('Queue', () => {
     expect(await sub.next()).to.deep.equal({ done: true, value: undefined });
   });
 
+  it('cancel is a no-op after stopping', async () => {
+    const queue = new Queue(({ stop }) => {
+      stop();
+    });
+
+    const sub = queue.subscribe();
+
+    expect(queue.isStopped()).to.equal(true);
+    queue.cancel();
+
+    expect(await sub.next()).to.deep.equal({ done: true, value: undefined });
+  });
+
   it('abort is a no-op after stopping', async () => {
     const queue = new Queue(({ stop }) => {
       stop();
@@ -643,14 +656,17 @@ describe('Queue', () => {
   it('stops in an error state when calling stopped with a reason, i.e. the last call to next to reject with that reason', async () => {
     let stoppedPromise!: Promise<unknown>;
     let stopped = false;
+    let stoppedReason: unknown;
+    const stopReason = new Error('Oops');
     const sub = new Queue(({ push, stop, stopped: _stoppedPromise }) => {
       stoppedPromise = _stoppedPromise;
 
-      stoppedPromise.then(() => {
+      stoppedPromise.then((reason) => {
         stopped = true;
+        stoppedReason = reason;
       });
       push(1);
-      stop(new Error('Oops'));
+      stop(stopReason);
     }).subscribe();
 
     expect(stopped).to.equal(false);
@@ -659,6 +675,7 @@ describe('Queue', () => {
     await expectPromise(sub.next()).toRejectWith('Oops');
 
     expect(stopped).to.equal(true);
+    expect(stoppedReason).to.equal(stopReason);
   });
 
   it('cancels existing requests when calling cancel', async () => {

src/execution/incremental/__tests__/WorkQueue-test.ts

@@ -1,6 +1,7 @@
 import { expect } from 'chai';
 import { describe, it } from 'mocha';
 
+import { expectPromise } from '../../../__testUtils__/expectPromise.js';
 import { resolveOnNextTick } from '../../../__testUtils__/resolveOnNextTick.js';
 
 import { isPromise } from '../../../jsutils/isPromise.js';
@@ -1266,4 +1267,61 @@ describe('WorkQueue', () => {
     expect(childTaskCancelled).to.equal(true);
     expect(childStreamCancelled).to.equal(true);
   });
+
+  it('forwards throw reason when cancelling nested work', async () => {
+    const root: TestGroup = { parent: undefined };
+    const rootStream = streamFrom([{ value: 1 }]);
+    const abortReason = new Error('Abort nested work');
+
+    let childStreamCancelReason: unknown;
+    const childStreamQueue = new Queue<TestStreamItem>(({ stopped }) => {
+      // eslint-disable-next-line @typescript-eslint/no-floating-promises
+      stopped.then((reason) => {
+        childStreamCancelReason = reason;
+      });
+    });
+    const childStream: TestStream = { queue: childStreamQueue };
+
+    let childTaskCancelReason: unknown;
+    const childTask: Task<
+      TestTaskValue,
+      TestStreamItemValue,
+      TestGroup,
+      TestStream
+    > = {
+      groups: [root],
+      computation: new Computation(
+        () =>
+          new Promise(() => {
+            // never resolves
+          }),
+        (reason) => {
+          childTaskCancelReason = reason;
+        },
+      ),
+    };
+
+    const rootTask = makeTask([root], () => ({
+      value: 'root',
+      work: { streams: [childStream] },
+    }));
+
+    const workQueue = createWorkQueue({
+      groups: [root],
+      tasks: [rootTask, childTask],
+      streams: [rootStream],
+    });
+
+    const iterator = workQueue.events[Symbol.asyncIterator]();
+    await iterator.next();
+
+    await expectPromise(iterator.throw(abortReason)).toRejectWith(
+      'Abort nested work',
+    );
+
+    await resolveOnNextTick();
+
+    expect(childTaskCancelReason).to.equal(abortReason);
+    expect(childStreamCancelReason).to.equal(abortReason);
+  });
 });