add cancellation support to async iterable iteration

src/execution/PromiseCanceller.ts

@@ -1,8 +1,8 @@
 import { promiseWithResolvers } from '../jsutils/promiseWithResolvers.js';
 
 /**
- * A PromiseCanceller object can be used to cancel multiple promises
- * using a single AbortSignal.
+ * A PromiseCanceller object can be used to trigger multiple responses
+ * in response to a single AbortSignal.
  *
  * @internal
  */
@@ -28,14 +28,21 @@ export class PromiseCanceller {
     this.abortSignal.removeEventListener('abort', this.abort);
   }
 
-  withCancellation<T>(originalPromise: Promise<T>): Promise<T> {
+  cancellablePromise<T>(
+    originalPromise: Promise<T>,
+    onCancel?: (() => unknown) | undefined,
+  ): Promise<T> {
     if (this.abortSignal.aborted) {
+      onCancel?.();
       // eslint-disable-next-line @typescript-eslint/prefer-promise-reject-errors
       return Promise.reject(this.abortSignal.reason);
     }
 
     const { promise, resolve, reject } = promiseWithResolvers<T>();
-    const abort = () => reject(this.abortSignal.reason);
+    const abort = () => {
+      onCancel?.();
+      reject(this.abortSignal.reason);
+    };
     this._aborts.add(abort);
     originalPromise.then(
       (resolved) => {
@@ -50,4 +57,33 @@ export class PromiseCanceller {
 
     return promise;
   }
+
+  cancellableIterable<T>(iterable: AsyncIterable<T>): AsyncIterable<T> {
+    const iterator = iterable[Symbol.asyncIterator]();
+
+    if (iterator.return) {
+      const _return = iterator.return.bind(iterator);
+      const _returnIgnoringErrors = async (): Promise<IteratorResult<T>> => {
+        _return().catch(() => {
+          /* c8 ignore next */
+          // ignore
+        });
+        return Promise.resolve({ value: undefined, done: true });
+      };
+
+      return {
+        [Symbol.asyncIterator]: () => ({
+          next: () =>
+            this.cancellablePromise(iterator.next(), _returnIgnoringErrors),
+          return: () => this.cancellablePromise(_return()),
+        }),
+      };
+    }
+
+    return {
+      [Symbol.asyncIterator]: () => ({
+        next: () => this.cancellablePromise(iterator.next()),
+      }),
+    };
+  }
 }

src/execution/__tests__/PromiseCanceller-test.ts

@@ -1,56 +1,242 @@
+import { expect } from 'chai';
 import { describe, it } from 'mocha';
 
 import { expectPromise } from '../../__testUtils__/expectPromise.js';
 
 import { PromiseCanceller } from '../PromiseCanceller.js';
 
 describe('PromiseCanceller', () => {
-  it('works to cancel an already resolved promise', async () => {
-    const abortController = new AbortController();
-    const abortSignal = abortController.signal;
+  describe('cancellablePromise', () => {
+    it('works to cancel an already resolved promise', async () => {
+      const abortController = new AbortController();
+      const abortSignal = abortController.signal;
 
-    const promiseCanceller = new PromiseCanceller(abortSignal);
+      const promiseCanceller = new PromiseCanceller(abortSignal);
 
-    const promise = Promise.resolve(1);
+      const promise = Promise.resolve(1);
 
-    const withCancellation = promiseCanceller.withCancellation(promise);
+      const withCancellation = promiseCanceller.cancellablePromise(promise);
 
-    abortController.abort(new Error('Cancelled!'));
+      abortController.abort(new Error('Cancelled!'));
 
-    await expectPromise(withCancellation).toRejectWith('Cancelled!');
-  });
+      await expectPromise(withCancellation).toRejectWith('Cancelled!');
+    });
+
+    it('works to cancel an already resolved promise after abort signal triggered', async () => {
+      const abortController = new AbortController();
+      const abortSignal = abortController.signal;
+
+      abortController.abort(new Error('Cancelled!'));
+
+      const promiseCanceller = new PromiseCanceller(abortSignal);
+
+      const promise = Promise.resolve(1);
+
+      const withCancellation = promiseCanceller.cancellablePromise(promise);
+
+      await expectPromise(withCancellation).toRejectWith('Cancelled!');
+    });
+
+    it('works to cancel a hanging promise', async () => {
+      const abortController = new AbortController();
+      const abortSignal = abortController.signal;
+
+      const promiseCanceller = new PromiseCanceller(abortSignal);
+
+      const promise = new Promise(() => {
+        /* never resolves */
+      });
+
+      const withCancellation = promiseCanceller.cancellablePromise(promise);
+
+      abortController.abort(new Error('Cancelled!'));
+
+      await expectPromise(withCancellation).toRejectWith('Cancelled!');
+    });
 
-  it('works to cancel a hanging promise', async () => {
-    const abortController = new AbortController();
-    const abortSignal = abortController.signal;
+    it('works to cancel a hanging promise created after abort signal triggered', async () => {
+      const abortController = new AbortController();
+      const abortSignal = abortController.signal;
 
-    const promiseCanceller = new PromiseCanceller(abortSignal);
+      abortController.abort(new Error('Cancelled!'));
 
-    const promise = new Promise(() => {
-      /* never resolves */
+      const promiseCanceller = new PromiseCanceller(abortSignal);
+
+      const promise = new Promise(() => {
+        /* never resolves */
+      });
+
+      const withCancellation = promiseCanceller.cancellablePromise(promise);
+
+      await expectPromise(withCancellation).toRejectWith('Cancelled!');
+    });
+
+    it('works to trigger onCancel when cancelling a hanging promise', async () => {
+      const abortController = new AbortController();
+      const abortSignal = abortController.signal;
+
+      const promiseCanceller = new PromiseCanceller(abortSignal);
+
+      const promise = new Promise(() => {
+        /* never resolves */
+      });
+
+      let onCancelCalled = false;
+      const onCancel = () => {
+        onCancelCalled = true;
+      };
+
+      const withCancellation = promiseCanceller.cancellablePromise(
+        promise,
+        onCancel,
+      );
+
+      expect(onCancelCalled).to.equal(false);
+
+      abortController.abort(new Error('Cancelled!'));
+
+      expect(onCancelCalled).to.equal(true);
+
+      await expectPromise(withCancellation).toRejectWith('Cancelled!');
     });
 
-    const withCancellation = promiseCanceller.withCancellation(promise);
+    it('works to trigger onCancel when cancelling a hanging promise created after abort signal triggered', async () => {
+      const abortController = new AbortController();
+      const abortSignal = abortController.signal;
+
+      abortController.abort(new Error('Cancelled!'));
+
+      const promiseCanceller = new PromiseCanceller(abortSignal);
+
+      const promise = new Promise(() => {
+        /* never resolves */
+      });
 
-    abortController.abort(new Error('Cancelled!'));
+      let onCancelCalled = false;
+      const onCancel = () => {
+        onCancelCalled = true;
+      };
 
-    await expectPromise(withCancellation).toRejectWith('Cancelled!');
+      const withCancellation = promiseCanceller.cancellablePromise(
+        promise,
+        onCancel,
+      );
+
+      expect(onCancelCalled).to.equal(true);
+
+      await expectPromise(withCancellation).toRejectWith('Cancelled!');
+    });
   });
 
-  it('works to cancel a hanging promise created after abort signal triggered', async () => {
-    const abortController = new AbortController();
-    const abortSignal = abortController.signal;
+  describe('cancellableAsyncIterable', () => {
+    it('works to abort a next call', async () => {
+      const abortController = new AbortController();
+      const abortSignal = abortController.signal;
+
+      const promiseCanceller = new PromiseCanceller(abortSignal);
+
+      const asyncIterable = {
+        [Symbol.asyncIterator]: () => ({
+          next: () => Promise.resolve({ value: 1, done: false }),
+        }),
+      };
+
+      const cancellableAsyncIterable =
+        promiseCanceller.cancellableIterable(asyncIterable);
+
+      const nextPromise =
+        cancellableAsyncIterable[Symbol.asyncIterator]().next();
+
+      abortController.abort(new Error('Cancelled!'));
+
+      await expectPromise(nextPromise).toRejectWith('Cancelled!');
+    });
+
+    it('works to abort a next call when already aborted', async () => {
+      const abortController = new AbortController();
+      const abortSignal = abortController.signal;
+
+      abortController.abort(new Error('Cancelled!'));
+
+      const promiseCanceller = new PromiseCanceller(abortSignal);
+
+      const asyncIterable = {
+        [Symbol.asyncIterator]: () => ({
+          next: () => Promise.resolve({ value: 1, done: false }),
+        }),
+      };
+
+      const cancellableAsyncIterable =
+        promiseCanceller.cancellableIterable(asyncIterable);
+
+      const nextPromise =
+        cancellableAsyncIterable[Symbol.asyncIterator]().next();
 
-    abortController.abort(new Error('Cancelled!'));
+      await expectPromise(nextPromise).toRejectWith('Cancelled!');
+    });
+
+    it('works to call return', async () => {
+      const abortController = new AbortController();
+      const abortSignal = abortController.signal;
+
+      const promiseCanceller = new PromiseCanceller(abortSignal);
+
+      let returned = false;
+      const asyncIterable = {
+        [Symbol.asyncIterator]: () => ({
+          next: () => Promise.resolve({ value: 1, done: false }),
+          return: () => {
+            returned = true;
+            return Promise.resolve({ value: undefined, done: true });
+          },
+        }),
+      };
+
+      const cancellableAsyncIterable =
+        promiseCanceller.cancellableIterable(asyncIterable);
+
+      abortController.abort(new Error('Cancelled!'));
+
+      expect(returned).to.equal(false);
+
+      const nextPromise =
+        cancellableAsyncIterable[Symbol.asyncIterator]().next();
 
-    const promiseCanceller = new PromiseCanceller(abortSignal);
+      expect(returned).to.equal(true);
 
-    const promise = new Promise(() => {
-      /* never resolves */
+      await expectPromise(nextPromise).toRejectWith('Cancelled!');
     });
 
-    const withCancellation = promiseCanceller.withCancellation(promise);
+    it('works to call return when already aborted', async () => {
+      const abortController = new AbortController();
+      const abortSignal = abortController.signal;
 
-    await expectPromise(withCancellation).toRejectWith('Cancelled!');
+      abortController.abort(new Error('Cancelled!'));
+
+      const promiseCanceller = new PromiseCanceller(abortSignal);
+
+      let returned = false;
+      const asyncIterable = {
+        [Symbol.asyncIterator]: () => ({
+          next: () => Promise.resolve({ value: 1, done: false }),
+          return: () => {
+            returned = true;
+            return Promise.resolve({ value: undefined, done: true });
+          },
+        }),
+      };
+
+      const cancellableAsyncIterable =
+        promiseCanceller.cancellableIterable(asyncIterable);
+
+      expect(returned).to.equal(false);
+
+      const nextPromise =
+        cancellableAsyncIterable[Symbol.asyncIterator]().next();
+
+      expect(returned).to.equal(true);
+
+      await expectPromise(nextPromise).toRejectWith('Cancelled!');
+    });
   });
 });