refactor(workqueue): propagate async computation abort cleanup (#4640)

Author: yaacovCR

src/execution/incremental/Computation.ts

@@ -9,11 +9,11 @@ type MaybePromise<T> =
 /** @internal **/
 export class Computation<T> {
   private _fn: () => PromiseOrValue<T>;
-  private _onAbort: ((reason?: unknown) => void) | undefined;
+  private _onAbort: ((reason?: unknown) => PromiseOrValue<void>) | undefined;
   private _maybePromise?: MaybePromise<T>;
   constructor(
     fn: () => PromiseOrValue<T>,
-    onAbort?: (reason?: unknown) => void,
+    onAbort?: (reason?: unknown) => PromiseOrValue<void>,
   ) {
     this._fn = fn;
     this._onAbort = onAbort;
@@ -54,7 +54,7 @@ export class Computation<T> {
       }
     }
   }
-  abort(reason?: unknown): void {
+  abort(reason?: unknown): PromiseOrValue<void> {
     const maybePromise = this._maybePromise;
     if (!maybePromise) {
       this._maybePromise = {
@@ -64,12 +64,14 @@ export class Computation<T> {
       return;
     }
     const status = maybePromise.status;
-    if (status === 'pending' && this._onAbort) {
-      this._onAbort(reason);
+    if (status === 'pending') {
       this._maybePromise = {
         status: 'rejected',
         reason,
       };
+      if (this._onAbort) {
+        return this._onAbort(reason);
+      }
     }
   }
 }

src/execution/incremental/IncrementalExecutor.ts

@@ -321,7 +321,13 @@ export class IncrementalExecutor<
   override cancel(reason?: unknown): void {
     super.cancel(reason);
     for (const task of this.tasks) {
-      task.computation.abort(reason);
+      const aborted = task.computation.abort(reason);
+      /* c8 ignore start */
+      // TODO: add coverage
+      if (isPromise(aborted)) {
+        aborted.catch(() => undefined);
+      }
+      /* c8 ignore stop */
     }
     for (const stream of this.streams) {
       const aborted = stream.queue.abort(reason);
@@ -594,7 +600,13 @@ export class IncrementalExecutor<
     const filteredTasks: Array<ExecutionGroup> = [];
     for (const task of tasks) {
       if (collectedErrors.hasNulledPosition(task.path)) {
-        task.computation.abort(cancellationReason);
+        const aborted = task.computation.abort(cancellationReason);
+        /* c8 ignore start */
+        // TODO: add coverage
+        if (isPromise(aborted)) {
+          aborted.catch(() => undefined);
+        }
+        /* c8 ignore stop */
       } else {
         filteredTasks.push(task);
       }

src/execution/incremental/WorkQueue.ts

@@ -282,7 +282,10 @@ export function createWorkQueue<
     reason: unknown,
     cancelPromises: Array<Promise<unknown>>,
   ): void {
-    task.computation.abort(reason);
+    const abortResult = task.computation.abort(reason);
+    if (isPromise(abortResult)) {
+      cancelPromises.push(abortResult);
+    }
     const taskNode = taskNodes.get(task);
     if (taskNode) {
       for (const childStream of taskNode.childStreams) {
@@ -296,10 +299,9 @@ export function createWorkQueue<
     reason: unknown,
     cancelPromises: Array<Promise<unknown>>,
   ): void {
-    const cancelResult =
-      reason === undefined ? stream.queue.cancel() : stream.queue.abort(reason);
-    if (isPromise(cancelResult)) {
-      cancelPromises.push(cancelResult);
+    const abortResult = stream.queue.abort(reason);
+    if (isPromise(abortResult)) {
+      cancelPromises.push(abortResult);
     }
   }
 

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

@@ -4,8 +4,20 @@ import { describe, it } from 'mocha';
 import { expectPromise } from '../../../__testUtils__/expectPromise.js';
 import { resolveOnNextTick } from '../../../__testUtils__/resolveOnNextTick.js';
 
+import { isPromise } from '../../../jsutils/isPromise.js';
+
 import { Computation } from '../Computation.js';
 
+function abortIgnoringCleanup(
+  computation: Computation<unknown>,
+  reason?: unknown,
+): void {
+  const aborted = computation.abort(reason);
+  if (isPromise(aborted)) {
+    aborted.catch(() => undefined);
+  }
+}
+
 describe('Computation', () => {
   it('can return a result', () => {
     const computation = new Computation(() => ({ value: 123 }));
@@ -114,7 +126,7 @@ describe('Computation', () => {
         onAbortRan = true;
       },
     );
-    computation.abort();
+    abortIgnoringCleanup(computation);
     expect(() => computation.result()).to.throw('Cancelled!');
     expect(onAbortRan).to.equal(false);
   });
@@ -129,7 +141,7 @@ describe('Computation', () => {
     );
 
     computation.prime();
-    computation.abort();
+    abortIgnoringCleanup(computation);
     expect(computation.result()).to.deep.equal({ value: 123 });
     expect(onAbortRan).to.equal(false);
   });
@@ -146,7 +158,7 @@ describe('Computation', () => {
     );
 
     computation.prime();
-    computation.abort();
+    abortIgnoringCleanup(computation);
     expect(() => computation.result()).to.throw('failure');
     expect(onAbortRan).to.equal(false);
   });
@@ -164,16 +176,52 @@ describe('Computation', () => {
     );
 
     computation.prime();
-    computation.abort();
+    abortIgnoringCleanup(computation);
     expect(onAbortRan).to.equal(true);
     expect(() => computation.result()).to.throw('Cancelled!');
   });
 
+  it('returns async abort cleanup while running', async () => {
+    let resolveCleanup!: () => void;
+    const cleanupPromise = new Promise<void>((resolve) => {
+      resolveCleanup = resolve;
+    });
+    const computation = new Computation(
+      () =>
+        new Promise(() => {
+          // Never resolves.
+        }),
+      () => cleanupPromise,
+    );
+
+    computation.prime();
+    const abortResult = computation.abort();
+    expect(abortResult).to.equal(cleanupPromise);
+    expect(isPromise(abortResult)).to.equal(true);
+    if (!isPromise(abortResult)) {
+      throw new Error('Expected async abort cleanup promise.');
+    }
+
+    let abortSettled = false;
+    abortResult.then(
+      () => {
+        abortSettled = true;
+      },
+      () => {
+        abortSettled = true;
+      },
+    );
+    expect(abortSettled).to.equal(false);
+    resolveCleanup();
+    await abortResult;
+    expect(abortSettled).to.equal(true);
+  });
+
   it('can be aborted with a provided reason before running', () => {
     const abortReason = new Error('aborted');
     const computation = new Computation(() => ({ value: 123 }));
 
-    computation.abort(abortReason);
+    abortIgnoringCleanup(computation, abortReason);
     expect(() => computation.result()).to.throw('aborted');
   });
 
@@ -191,7 +239,7 @@ describe('Computation', () => {
     );
 
     computation.prime();
-    computation.abort(abortReason);
+    abortIgnoringCleanup(computation, abortReason);
     expect(onAbortReason).to.equal(abortReason);
     expect(() => computation.result()).to.throw('aborted');
   });

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

@@ -1228,6 +1228,9 @@ describe('WorkQueue', () => {
     const childStream: TestStream = { queue: childStreamQueue };
 
     let childTaskCancelled = false;
+    const { promise: childTaskCleanup, resolve: resolveChildTaskCleanup } =
+      // eslint-disable-next-line @typescript-eslint/no-invalid-void-type
+      promiseWithResolvers<void>();
     const childTask: Task<
       TestTaskValue,
       TestStreamItemValue,
@@ -1242,6 +1245,7 @@ describe('WorkQueue', () => {
           }),
         () => {
           childTaskCancelled = true;
+          return childTaskCleanup;
         },
       ),
     };
@@ -1278,6 +1282,7 @@ describe('WorkQueue', () => {
     await resolveOnNextTick();
 
     resolveChildStreamCleanup();
+    resolveChildTaskCleanup();
 
     expect(await returnPromise).to.deep.equal({
       value: undefined,