clean up PromiseCanceller when used with subscriptions

yaacovCR · yaacovCR · commit a521d6928a03 · 2024-11-04T14:19:19.000+02:00
diff --git a/src/execution/PromiseCanceller.ts b/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);
@@ -61,28 +56,22 @@ export class PromiseCanceller {
   cancellableIterable<T>(iterable: AsyncIterable<T>): AsyncIterable<T> {
     const iterator = iterable[Symbol.asyncIterator]();
 
+    const _next = iterator.next.bind(iterator);
+
     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),
+          next: () => this.cancellablePromise(_next()),
           return: () => this.cancellablePromise(_return()),
         }),
       };
     }
 
     return {
       [Symbol.asyncIterator]: () => ({
-        next: () => this.cancellablePromise(iterator.next()),
+        next: () => this.cancellablePromise(_next()),
       }),
     };
   }
diff --git a/src/execution/__tests__/PromiseCanceller-test.ts b/src/execution/__tests__/PromiseCanceller-test.ts
@@ -1,4 +1,3 @@
-import { expect } from 'chai';
 import { describe, it } from 'mocha';
 
 import { expectPromise } from '../../__testUtils__/expectPromise.js';
@@ -70,62 +69,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', () => {
@@ -174,69 +117,5 @@ describe('PromiseCanceller', () => {
 
       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();
-
-      expect(returned).to.equal(true);
-
-      await expectPromise(nextPromise).toRejectWith('Cancelled!');
-    });
-
-    it('works to call return when already aborted', async () => {
-      const abortController = new AbortController();
-      const abortSignal = abortController.signal;
-
-      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!');
-    });
   });
 });
diff --git a/src/execution/__tests__/abort-signal-test.ts b/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 a subscription resolver', async () => {
     const abortController = new AbortController();
     const document = parse(`
       subscription {
@@ -830,6 +830,44 @@ describe('Execute: Cancellation', () => {
     });
   });
 
+  it('should successfully wrap the subscription', async () => {
+    const abortController = new AbortController();
+    const document = parse(`
+      subscription {
+        foo
+      }
+    `);
+
+    async function* foo() {
+      yield await Promise.resolve({ foo: 'foo' });
+    }
+
+    const subscription = await subscribe({
+      document,
+      schema,
+      abortSignal: abortController.signal,
+      rootValue: {
+        foo: Promise.resolve(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(`
@@ -838,15 +876,16 @@ describe('Execute: Cancellation', () => {
       }
     `);
 
+    async function* foo() {
+      yield await Promise.resolve({ foo: 'foo' });
+    }
+
     const subscription = subscribe({
       document,
       schema,
       abortSignal: abortController.signal,
       rootValue: {
-        async *foo() {
-          yield await Promise.resolve({ foo: 'foo' });
-          yield await Promise.resolve({ foo: 'foo' }); /* c8 ignore start */
-        } /* c8 ignore stop */,
+        foo: foo(),
       },
     });
 
diff --git a/src/execution/__tests__/mapAsyncIterable-test.ts b/src/execution/__tests__/mapAsyncIterable-test.ts
@@ -89,6 +89,54 @@ describe('mapAsyncIterable', () => {
     });
   });
 
+  it('calls done when completes', async () => {
+    async function* source() {
+      yield 1;
+      yield 2;
+      yield 3;
+    }
+
+    let done = false;
+    const doubles = mapAsyncIterable(
+      source(),
+      (x) => Promise.resolve(x + x),
+      () => {
+        done = true;
+      },
+    );
+
+    expect(await doubles.next()).to.deep.equal({ value: 2, done: false });
+    expect(await doubles.next()).to.deep.equal({ value: 4, done: false });
+    expect(await doubles.next()).to.deep.equal({ value: 6, done: false });
+    expect(done).to.equal(false);
+    expect(await doubles.next()).to.deep.equal({
+      value: undefined,
+      done: true,
+    });
+    expect(done).to.equal(true);
+  });
+
+  it('calls done when completes with error', async () => {
+    async function* source() {
+      yield 1;
+      throw new Error('Oops');
+    }
+
+    let done = false;
+    const doubles = mapAsyncIterable(
+      source(),
+      (x) => Promise.resolve(x + x),
+      () => {
+        done = true;
+      },
+    );
+
+    expect(await doubles.next()).to.deep.equal({ value: 2, done: false });
+    expect(done).to.equal(false);
+    await expectPromise(doubles.next()).toRejectWith('Oops');
+    expect(done).to.equal(true);
+  });
+
   it('allows returning early from mapped async generator', async () => {
     async function* source() {
       try {
diff --git a/src/execution/execute.ts b/src/execution/execute.ts
@@ -2099,6 +2099,13 @@ export function subscribe(
   return mapSourceToResponse(validatedExecutionArgs, resultOrStream);
 }
 
+/**
+ *
+ * For each payload yielded from a subscription, map it over the normal
+ * GraphQL `execute` function, with `payload` as the rootValue.
+ * This implements the "MapSourceToResponseEvent" algorithm described in
+ * the GraphQL specification..
+ */
 function mapSourceToResponse(
   validatedExecutionArgs: ValidatedExecutionArgs,
   resultOrStream: ExecutionResult | AsyncIterable<unknown>,
@@ -2107,10 +2114,22 @@ function mapSourceToResponse(
     return resultOrStream;
   }
 
-  // For each payload yielded from a subscription, map it over the normal
-  // GraphQL `execute` function, with `payload` as the rootValue.
-  // This implements the "MapSourceToResponseEvent" algorithm described in
-  // the GraphQL specification..
+  const abortSignal = validatedExecutionArgs.abortSignal;
+  if (abortSignal) {
+    const promiseCanceller = new PromiseCanceller(abortSignal);
+    return mapAsyncIterable(
+      promiseCanceller?.cancellableIterable(resultOrStream),
+      (payload: unknown) => {
+        const perEventExecutionArgs: ValidatedExecutionArgs = {
+          ...validatedExecutionArgs,
+          rootValue: payload,
+        };
+        return validatedExecutionArgs.perEventExecutor(perEventExecutionArgs);
+      },
+      () => promiseCanceller.disconnect(),
+    );
+  }
+
   return mapAsyncIterable(resultOrStream, (payload: unknown) => {
     const perEventExecutionArgs: ValidatedExecutionArgs = {
       ...validatedExecutionArgs,
@@ -2265,16 +2284,16 @@ function executeSubscription(
     // used to represent an authenticated user, or request-specific caches.
     const result = resolveFn(rootValue, args, contextValue, info, abortSignal);
 
-    const promiseCanceller = abortSignal
-      ? new PromiseCanceller(abortSignal)
-      : undefined;
-
     if (isPromise(result)) {
+      const promiseCanceller = abortSignal
+        ? new PromiseCanceller(abortSignal)
+        : undefined;
+
       const promise = promiseCanceller?.cancellablePromise(result) ?? result;
       return promise.then(assertEventStream).then(
         (resolved) => {
           promiseCanceller?.disconnect();
-          return promiseCanceller?.cancellableIterable(resolved) ?? resolved;
+          return resolved;
         },
         (error: unknown) => {
           promiseCanceller?.disconnect();
@@ -2284,7 +2303,7 @@ function executeSubscription(
     }
 
     const eventStream = assertEventStream(result);
-    return promiseCanceller?.cancellableIterable(eventStream) ?? eventStream;
+    return eventStream;
   } catch (error) {
     throw locatedError(error, fieldNodes, pathToArray(path));
   }
diff --git a/src/execution/mapAsyncIterable.ts b/src/execution/mapAsyncIterable.ts
