abort without a pending next

Author: yaacovCR

src/execution/PromiseCanceller.ts

@@ -28,19 +28,14 @@ export class PromiseCanceller {
     this.abortSignal.removeEventListener('abort', this.abort);
   }
 
-  cancellablePromise<T>(
-    originalPromise: Promise<T>,
-    onCancel?: (() => unknown) | undefined,
-  ): Promise<T> {
+  cancellablePromise<T>(originalPromise: Promise<T>): 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 = () => {
-      onCancel?.();
       reject(this.abortSignal.reason);
     };
     this._aborts.add(abort);
@@ -62,20 +57,45 @@ export class PromiseCanceller {
     const iterator = iterable[Symbol.asyncIterator]();
 
     if (iterator.return) {
+      const _next = iterator.next.bind(iterator);
       const _return = iterator.return.bind(iterator);
-      const _returnIgnoringErrors = async (): Promise<IteratorResult<T>> => {
+
+      const abort = () => {
         _return().catch(() => {
           /* c8 ignore next */
           // ignore
         });
-        return Promise.resolve({ value: undefined, done: true });
       };
 
+      if (this.abortSignal.aborted) {
+        abort();
+        // eslint-disable-next-line @typescript-eslint/prefer-promise-reject-errors
+        const onMethod = () => Promise.reject(this.abortSignal.reason);
+        return {
+          [Symbol.asyncIterator]: () => ({
+            next: onMethod,
+            return: onMethod,
+          }),
+        };
+      }
+
+      this._aborts.add(abort);
+      const on = (method: () => Promise<IteratorResult<T>>) => async () => {
+        try {
+          const iteration = await this.cancellablePromise(method());
+          if (iteration.done) {
+            this._aborts.delete(abort);
+          }
+          return iteration;
+        } catch (error) {
+          this._aborts.delete(abort);
+          throw error;
+        }
+      };
       return {
         [Symbol.asyncIterator]: () => ({
-          next: () =>
-            this.cancellablePromise(iterator.next(), _returnIgnoringErrors),
-          return: () => this.cancellablePromise(_return()),
+          next: on(_next),
+          return: on(_return),
         }),
       };
     }

src/execution/__tests__/PromiseCanceller-test.ts

@@ -70,62 +70,6 @@ describe('PromiseCanceller', () => {
 
       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!');
-    });
-
-    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 */
-      });
-
-      let onCancelCalled = false;
-      const onCancel = () => {
-        onCancelCalled = true;
-      };
-
-      const withCancellation = promiseCanceller.cancellablePromise(
-        promise,
-        onCancel,
-      );
-
-      expect(onCancelCalled).to.equal(true);
-
-      await expectPromise(withCancellation).toRejectWith('Cancelled!');
-    });
   });
 
   describe('cancellableAsyncIterable', () => {
@@ -197,13 +141,11 @@ describe('PromiseCanceller', () => {
 
       abortController.abort(new Error('Cancelled!'));
 
-      expect(returned).to.equal(false);
+      expect(returned).to.equal(true);
 
       const nextPromise =
         cancellableAsyncIterable[Symbol.asyncIterator]().next();
 
-      expect(returned).to.equal(true);
-
       await expectPromise(nextPromise).toRejectWith('Cancelled!');
     });
 
@@ -218,7 +160,8 @@ describe('PromiseCanceller', () => {
       let returned = false;
       const asyncIterable = {
         [Symbol.asyncIterator]: () => ({
-          next: () => Promise.resolve({ value: 1, done: false }),
+          /* c8 ignore next */
+          next: () => Promise.resolve({ value: undefined, done: true }),
           return: () => {
             returned = true;
             return Promise.resolve({ value: undefined, done: true });
@@ -229,13 +172,11 @@ describe('PromiseCanceller', () => {
       const cancellableAsyncIterable =
         promiseCanceller.cancellableIterable(asyncIterable);
 
-      expect(returned).to.equal(false);
+      expect(returned).to.equal(true);
 
       const nextPromise =
         cancellableAsyncIterable[Symbol.asyncIterator]().next();
 
-      expect(returned).to.equal(true);
-
       await expectPromise(nextPromise).toRejectWith('Cancelled!');
     });
   });

src/execution/__tests__/abort-signal-test.ts

@@ -795,7 +795,7 @@ describe('Execute: Cancellation', () => {
     });
   });
 
-  it('should stop the execution when aborted during subscription', async () => {
+  it('should stop the execution when aborted prior to return of the subscription resolver', async () => {
     const abortController = new AbortController();
     const document = parse(`
       subscription {
@@ -830,7 +830,7 @@ describe('Execute: Cancellation', () => {
     });
   });
 
-  it('should stop the execution when aborted during subscription', async () => {
+  it('should successfully wrap the subscription', async () => {
     const abortController = new AbortController();
     const document = parse(`
       subscription {
@@ -845,6 +845,41 @@ describe('Execute: Cancellation', () => {
       rootValue: {
         async *foo() {
           yield await Promise.resolve({ foo: 'foo' });
+        },
+      },
+    });
+
+    assert(isAsyncIterable(subscription));
+
+    expectJSON(await subscription.next()).toDeepEqual({
+      value: {
+        data: {
+          foo: 'foo',
+        },
+      },
+      done: false,
+    });
+
+    expectJSON(await subscription.next()).toDeepEqual({
+      value: undefined,
+      done: true,
+    });
+  });
+
+  it('should stop the execution when aborted during subscription', async () => {
+    const abortController = new AbortController();
+    const document = parse(`
+      subscription {
+        foo
+      }
+    `);
+
+    const subscription = subscribe({
+      document,
+      schema,
+      abortSignal: abortController.signal,
+      rootValue: {
+        async *foo() {
           yield await Promise.resolve({ foo: 'foo' }); /* c8 ignore start */
         } /* c8 ignore stop */,
       },