use withCleanup to wrap actual generators

Author: yaacovCR

src/execution/IncrementalGraph.ts

@@ -97,11 +97,6 @@ export class IncrementalGraph {
     while ((completed = this._completedQueue.shift()) !== undefined) {
       yield completed;
     }
-    if (this._rootNodes.size === 0) {
-      for (const resolve of this._nextQueue) {
-        resolve(undefined);
-      }
-    }
   }
 
   nextCompletedBatch(): Promise<

src/execution/IncrementalPublisher.ts

@@ -28,6 +28,7 @@ import {
   isCompletedExecutionGroup,
   isFailedExecutionGroup,
 } from './types.js';
+import { withCleanup } from './withCleanup.js';
 
 export function buildIncrementalResponse(
   context: IncrementalPublisherContext,
@@ -63,11 +64,13 @@ interface SubsequentIncrementalExecutionResultContext {
  * @internal
  */
 class IncrementalPublisher {
+  private _isDone: boolean;
   private _context: IncrementalPublisherContext;
   private _nextId: number;
   private _incrementalGraph: IncrementalGraph;
 
   constructor(context: IncrementalPublisherContext) {
+    this._isDone = false;
     this._context = context;
     this._nextId = 0;
     this._incrementalGraph = new IncrementalGraph();
@@ -92,10 +95,14 @@ class IncrementalPublisher {
       ? { errors, data, pending, hasNext: true }
       : { data, pending, hasNext: true };
 
-    return {
-      initialResult,
-      subsequentResults: this._subscribe(),
-    };
+    const subsequentResults = withCleanup(this._subscribe(), async () => {
+      this._isDone = true;
+      this._context.abortSignalListener?.disconnect();
+      this._incrementalGraph.abort();
+      await this._returnAsyncIteratorsIgnoringErrors();
+    });
+
+    return { initialResult, subsequentResults };
   }
 
   private _toPendingResults(
@@ -121,22 +128,12 @@ class IncrementalPublisher {
     return String(this._nextId++);
   }
 
-  private _subscribe(): AsyncGenerator<
+  private async *_subscribe(): AsyncGenerator<
     SubsequentIncrementalExecutionResult,
     void,
     void
   > {
-    let isDone = false;
-
-    const _next = async (): Promise<
-      IteratorResult<SubsequentIncrementalExecutionResult, void>
-    > => {
-      if (isDone) {
-        this._context.abortSignalListener?.disconnect();
-        await this._returnAsyncIteratorsIgnoringErrors();
-        return { value: undefined, done: true };
-      }
-
+    while (!this._isDone) {
       const context: SubsequentIncrementalExecutionResultContext = {
         pending: [],
         incremental: [],
@@ -155,7 +152,7 @@ class IncrementalPublisher {
           const hasNext = this._incrementalGraph.hasNext();
 
           if (!hasNext) {
-            isDone = true;
+            this._isDone = true;
           }
 
           const subsequentIncrementalExecutionResult: SubsequentIncrementalExecutionResult =
@@ -172,50 +169,14 @@ class IncrementalPublisher {
             subsequentIncrementalExecutionResult.completed = completed;
           }
 
-          return { value: subsequentIncrementalExecutionResult, done: false };
+          yield subsequentIncrementalExecutionResult;
+          break;
         }
 
         // eslint-disable-next-line no-await-in-loop
         batch = await this._incrementalGraph.nextCompletedBatch();
       } while (batch !== undefined);
-
-      // TODO: add test for this case
-      /* c8 ignore next */
-      this._context.abortSignalListener?.disconnect();
-      await this._returnAsyncIteratorsIgnoringErrors();
-      return { value: undefined, done: true };
-    };
-
-    const _return = async (): Promise<
-      IteratorResult<SubsequentIncrementalExecutionResult, void>
-    > => {
-      isDone = true;
-      this._incrementalGraph.abort();
-      await this._returnAsyncIterators();
-      return { value: undefined, done: true };
-    };
-
-    const _throw = async (
-      error?: unknown,
-    ): Promise<IteratorResult<SubsequentIncrementalExecutionResult, void>> => {
-      isDone = true;
-      this._incrementalGraph.abort();
-      await this._returnAsyncIterators();
-      // eslint-disable-next-line @typescript-eslint/prefer-promise-reject-errors
-      return Promise.reject(error);
-    };
-
-    return {
-      [Symbol.asyncIterator]() {
-        return this;
-      },
-      next: _next,
-      return: _return,
-      throw: _throw,
-      async [Symbol.asyncDispose]() {
-        await _return();
-      },
-    };
+    }
   }
 
   private _handleCompletedIncrementalData(

src/execution/__tests__/simplePubSub.ts

@@ -1,4 +1,6 @@
-import { assert } from 'chai';
+import { promiseWithResolvers } from '../../jsutils/promiseWithResolvers.js';
+
+import { withCleanup } from '../withCleanup.js';
 
 /**
  * Create an AsyncIterator from an EventEmitter. Useful for mocking a
@@ -21,57 +23,49 @@ export class SimplePubSub<T> {
   }
 
   getSubscriber<R>(transform: (value: T) => R): AsyncGenerator<R, void, void> {
-    const pullQueue: Array<(result: IteratorResult<R, void>) => void> = [];
+    let pendingNext: ((result: R) => void) | undefined;
     const pushQueue: Array<R> = [];
     let listening = true;
     this._subscribers.add(pushValue);
 
     const emptyQueue = () => {
       listening = false;
       this._subscribers.delete(pushValue);
-      for (const resolve of pullQueue) {
-        resolve({ value: undefined, done: true });
+      if (pendingNext) {
+        pendingNext(undefined as R);
       }
-      pullQueue.length = 0;
+      pendingNext = undefined;
       pushQueue.length = 0;
     };
 
-    return {
-      next(): Promise<IteratorResult<R, void>> {
-        if (!listening) {
-          return Promise.resolve({ value: undefined, done: true });
-        }
-
+    async function* getSubscriberImpl(): AsyncGenerator<R, void, void> {
+      // eslint-disable-next-line no-unmodified-loop-condition
+      while (listening) {
         if (pushQueue.length > 0) {
           const value = pushQueue[0];
           pushQueue.shift();
-          return Promise.resolve({ value, done: false });
+          yield value;
+          continue;
         }
-        return new Promise((resolve) => pullQueue.push(resolve));
-      },
-      return(): Promise<IteratorResult<R, void>> {
-        emptyQueue();
-        return Promise.resolve({ value: undefined, done: true });
-      },
-      throw(error: unknown) {
-        emptyQueue();
-        // eslint-disable-next-line @typescript-eslint/prefer-promise-reject-errors
-        return Promise.reject(error);
-      },
-      [Symbol.asyncIterator]() {
-        return this;
-      },
-      async [Symbol.asyncDispose]() {
-        await this.return();
-      },
-    };
+
+        const { promise, resolve } = promiseWithResolvers<R>();
+        pendingNext = resolve;
+        // eslint-disable-next-line no-await-in-loop
+        const value = await promise;
+        if (!listening) {
+          return;
+        }
+        yield value;
+      }
+    }
+
+    return withCleanup(getSubscriberImpl(), emptyQueue);
 
     function pushValue(event: T): void {
       const value: R = transform(event);
-      if (pullQueue.length > 0) {
-        const receiver = pullQueue.shift();
-        assert(receiver != null);
-        receiver({ value, done: false });
+      if (pendingNext) {
+        pendingNext(value);
+        pendingNext = undefined;
       } else {
         pushQueue.push(value);
       }