refactor: extract Queue from IncrementalPublisher and IncrementalGraph

Queue is a batching async-aware iterator-like protocol meant to be consumed in the following (and only the following) pattern:

>  let batch;
>  if ((batch = queue.currentBatch()) !== undefined) {
>    doSomethingWithBatch(batch);
>  }
>  while ((batch = await queue.nextBatch) !== undefined) {
>    doSomethingWithBatch(batch);
>  }

A `push()` methods is provided in the Queue constructor (a la repeaters, see https://repeater.js.org/) so that `push()` can be private to the code that constructs the Queue (although it can be saved to be passed to other code).

A `stop()` method is also provided as a the second argument to the constructor for convienence, but it is also available on the queue object itself so that it can be called by the executor or by the consumer. So this works:

>  const queue = new Queue(
>    async (push, stop) => {
>      push(1);
>      push(2)
>      await Promise.resolve();
>      push(3);
>      stop();
>    },
>  );
>
>  const batch1 = Array.from(queue.nextBatch()); // batch1 = [1, 2]
>  const batch2 = Array.from(await queue.nextBatchAsync()); // batch2 = [3]

as does this:

>  let push;
>  const queue = new Queue(
>    (_push) => {
>      push = _push;
>    },
>  );
>
>  push(1);
>  push(2);
>
>  const batch1 = Array.from(queue.nextBatch()); // batch1 = [1, 2]
>
>  const batch2Promise = queue.nextBatchAsync();
>
>  await Promise.resolve();
>  push(3);
>  queue.stop();
>
>  const batch2 = await batch2Promise; // batch2 = [3]

Note: concurrent calls to `currentBatch()` and `nextBatch` will return the same batch and are not encouraged.

A `toAsyncIterable(mapFn)` method transforms coalesces each batch of items into a single value (or undefined if not value is to be emitted for the batch).

Using queues, we are able to remove all logic for handling the implicit queue from IncrementalPublisher, retaining only the `_handleCompletedBatch()`, while adding only the required `push()` and `stop()` calls within the IncrementalGraph.

Tests do not change, except that `.return()` and `.throw()` (but not `next()` have an extra tick secondary to the additional layers of `withCleanup()`, so that the tests required slight adjustment.

Author: yaacovCR

src/execution/IncrementalGraph.ts

@@ -1,10 +1,10 @@
 import { BoxedPromiseOrValue } from '../jsutils/BoxedPromiseOrValue.js';
 import { invariant } from '../jsutils/invariant.js';
 import { isPromise } from '../jsutils/isPromise.js';
-import { promiseWithResolvers } from '../jsutils/promiseWithResolvers.js';
 
 import type { GraphQLError } from '../error/GraphQLError.js';
 
+import { Queue } from './Queue.js';
 import type {
   DeferredFragmentRecord,
   DeliveryGroup,
@@ -22,16 +22,18 @@ import { isDeferredFragmentRecord, isPendingExecutionGroup } from './types.js';
  */
 export class IncrementalGraph {
   private _rootNodes: Set<DeliveryGroup>;
-
-  private _completedQueue: Array<IncrementalDataRecordResult>;
-  private _nextQueue: Array<
-    (iterable: Iterable<IncrementalDataRecordResult> | undefined) => void
-  >;
+  private _completed: Queue<IncrementalDataRecordResult>;
+  // _push and _stop are assigned in the executor which is executed
+  // synchronously by the Queue constructor.
+  private _push!: (item: IncrementalDataRecordResult) => void;
+  private _stop!: () => void;
 
   constructor() {
     this._rootNodes = new Set();
-    this._completedQueue = [];
-    this._nextQueue = [];
+    this._completed = new Queue<IncrementalDataRecordResult>((push, stop) => {
+      this._push = push;
+      this._stop = stop;
+    });
   }
 
   getNewRootNodes(
@@ -92,29 +94,6 @@ export class IncrementalGraph {
     }
   }
 
-  *currentCompletedBatch(): Generator<IncrementalDataRecordResult> {
-    let completed;
-    while ((completed = this._completedQueue.shift()) !== undefined) {
-      yield completed;
-    }
-  }
-
-  nextCompletedBatch(): Promise<
-    Iterable<IncrementalDataRecordResult> | undefined
-  > {
-    const { promise, resolve } = promiseWithResolvers<
-      Iterable<IncrementalDataRecordResult> | undefined
-    >();
-    this._nextQueue.push(resolve);
-    return promise;
-  }
-
-  abort(): void {
-    for (const resolve of this._nextQueue) {
-      resolve(undefined);
-    }
-  }
-
   hasNext(): boolean {
     return this._rootNodes.size > 0;
   }
@@ -146,17 +125,28 @@ export class IncrementalGraph {
     const newRootNodes = this._promoteNonEmptyToRoot(
       deferredFragmentRecord.children,
     );
+    this._maybeStop();
     return { newRootNodes, successfulExecutionGroups };
   }
 
   removeDeferredFragment(
     deferredFragmentRecord: DeferredFragmentRecord,
   ): boolean {
-    return this._rootNodes.delete(deferredFragmentRecord);
+    const deleted = this._rootNodes.delete(deferredFragmentRecord);
+    if (!deleted) {
+      return false;
+    }
+    this._maybeStop();
+    return true;
   }
 
   removeStream(streamRecord: StreamRecord): void {
     this._rootNodes.delete(streamRecord);
+    this._maybeStop();
+  }
+
+  subscribe(): Queue<IncrementalDataRecordResult> {
+    return this._completed;
   }
 
   private _addIncrementalDataRecords(
@@ -246,9 +236,9 @@ export class IncrementalGraph {
     const value = completedExecutionGroup.value;
     if (isPromise(value)) {
       // eslint-disable-next-line @typescript-eslint/no-floating-promises
-      value.then((resolved) => this._enqueue(resolved));
+      value.then((resolved) => this._push(resolved));
     } else {
-      this._enqueue(value);
+      this._push(value);
     }
   }
 
@@ -266,7 +256,7 @@ export class IncrementalGraph {
           : streamItemRecord().value;
       if (isPromise(result)) {
         if (items.length > 0) {
-          this._enqueue({
+          this._push({
             streamRecord,
             result:
               // TODO add additional test case or rework for coverage
@@ -290,14 +280,14 @@ export class IncrementalGraph {
       }
       if (result.item === undefined) {
         if (items.length > 0) {
-          this._enqueue({
+          this._push({
             streamRecord,
             result: errors.length > 0 ? { items, errors } : { items },
             newDeferredFragmentRecords,
             incrementalDataRecords,
           });
         }
-        this._enqueue(
+        this._push(
           result.errors === undefined
             ? { streamRecord }
             : {
@@ -320,12 +310,9 @@ export class IncrementalGraph {
     }
   }
 
-  private _enqueue(completed: IncrementalDataRecordResult): void {
-    this._completedQueue.push(completed);
-    const next = this._nextQueue.shift();
-    if (next === undefined) {
-      return;
+  private _maybeStop(): void {
+    if (!this.hasNext()) {
+      this._stop();
     }
-    next(this.currentCompletedBatch());
   }
 }

src/execution/IncrementalPublisher.ts

@@ -64,13 +64,11 @@ 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();
@@ -95,14 +93,17 @@ class IncrementalPublisher {
       ? { errors, data, pending, hasNext: true }
       : { data, pending, hasNext: true };
 
-    const subsequentResults = withCleanup(this._subscribe(), async () => {
-      this._isDone = true;
-      this._context.abortSignalListener?.disconnect();
-      this._incrementalGraph.abort();
-      await this._returnAsyncIteratorsIgnoringErrors();
-    });
+    const subsequentResults = this._incrementalGraph
+      .subscribe()
+      .toAsyncGenerator((batch) => this._handleCompletedBatch(batch));
 
-    return { initialResult, subsequentResults };
+    return {
+      initialResult,
+      subsequentResults: withCleanup(subsequentResults, async () => {
+        this._context.abortSignalListener?.disconnect();
+        await this._returnAsyncIteratorsIgnoringErrors();
+      }),
+    };
   }
 
   private _toPendingResults(
@@ -128,55 +129,39 @@ class IncrementalPublisher {
     return String(this._nextId++);
   }
 
-  private async *_subscribe(): AsyncGenerator<
-    SubsequentIncrementalExecutionResult,
-    void,
-    void
-  > {
-    while (!this._isDone) {
-      const context: SubsequentIncrementalExecutionResultContext = {
-        pending: [],
-        incremental: [],
-        completed: [],
-      };
-
-      let batch: Iterable<IncrementalDataRecordResult> | undefined =
-        this._incrementalGraph.currentCompletedBatch();
-      do {
-        for (const completedResult of batch) {
-          this._handleCompletedIncrementalData(completedResult, context);
-        }
-
-        const { incremental, completed } = context;
-        if (incremental.length > 0 || completed.length > 0) {
-          const hasNext = this._incrementalGraph.hasNext();
-
-          if (!hasNext) {
-            this._isDone = true;
-          }
+  private _handleCompletedBatch(
+    batch: Iterable<IncrementalDataRecordResult>,
+  ): SubsequentIncrementalExecutionResult | undefined {
+    const context: SubsequentIncrementalExecutionResultContext = {
+      pending: [],
+      incremental: [],
+      completed: [],
+    };
 
-          const subsequentIncrementalExecutionResult: SubsequentIncrementalExecutionResult =
-            { hasNext };
+    for (const completedResult of batch) {
+      this._handleCompletedIncrementalData(completedResult, context);
+    }
 
-          const pending = context.pending;
-          if (pending.length > 0) {
-            subsequentIncrementalExecutionResult.pending = pending;
-          }
-          if (incremental.length > 0) {
-            subsequentIncrementalExecutionResult.incremental = incremental;
-          }
-          if (completed.length > 0) {
-            subsequentIncrementalExecutionResult.completed = completed;
-          }
+    const { incremental, completed } = context;
+    if (incremental.length === 0 && completed.length === 0) {
+      return;
+    }
 
-          yield subsequentIncrementalExecutionResult;
-          break;
-        }
+    const hasNext = this._incrementalGraph.hasNext();
 
-        // eslint-disable-next-line no-await-in-loop
-        batch = await this._incrementalGraph.nextCompletedBatch();
-      } while (batch !== undefined);
+    const subsequentIncrementalExecutionResult: SubsequentIncrementalExecutionResult =
+      { hasNext };
+    const pending = context.pending;
+    if (pending.length > 0) {
+      subsequentIncrementalExecutionResult.pending = pending;
+    }
+    if (incremental.length > 0) {
+      subsequentIncrementalExecutionResult.incremental = incremental;
+    }
+    if (completed.length > 0) {
+      subsequentIncrementalExecutionResult.completed = completed;
     }
+    return subsequentIncrementalExecutionResult;
   }
 
   private _handleCompletedIncrementalData(

src/execution/Queue.ts

@@ -0,0 +1,113 @@
+import { isPromise } from '../jsutils/isPromise.js';
+import type { PromiseOrValue } from '../jsutils/PromiseOrValue.js';
+import { promiseWithResolvers } from '../jsutils/promiseWithResolvers.js';
+
+import { withCleanup } from './withCleanup.js';
+
+/**
+ * @internal
+ */
+export class Queue<T> {
+  private _items: Array<T>;
+  private _stopped: boolean;
+  private _resolvers: Array<(iterable: Generator<T> | undefined) => void>;
+
+  constructor(
+    executor: (
+      push: (item: T) => void,
+      stop: () => void,
+    ) => PromiseOrValue<void>,
+  ) {
+    this._items = [];
+    this._stopped = false;
+    this._resolvers = [];
+    let result;
+    try {
+      result = executor(this._push.bind(this), this.stop.bind(this));
+    } catch {
+      // Ignore errors
+    }
+    if (isPromise(result)) {
+      result.catch(() => {
+        // Ignore errors
+      });
+    }
+  }
+
+  currentBatch(): Generator<T> | undefined {
+    if (this._items.length > 0) {
+      return this.batch();
+    }
+  }
+
+  nextBatch(): Promise<Generator<T> | undefined> {
+    if (this._items.length > 0) {
+      return Promise.resolve(this.batch());
+    }
+    if (this._stopped) {
+      return Promise.resolve(undefined);
+    }
+    const { promise, resolve } = promiseWithResolvers<
+      Generator<T> | undefined
+    >();
+    this._resolvers.push(resolve);
+    return promise;
+  }
+
+  stop(): void {
+    this._stopped = true;
+    this._resolve(undefined);
+  }
+
+  toAsyncGenerator<U>(
+    mapFn: (generator: Generator<T>) => U | undefined,
+  ): AsyncGenerator<U, void, void> {
+    return withCleanup(this.toAsyncGeneratorImpl(mapFn), () => this.stop());
+  }
+
+  private async *toAsyncGeneratorImpl<U>(
+    mapFn: (generator: Generator<T, void, void>) => U | undefined,
+  ): AsyncGenerator<U> {
+    let batch = this.currentBatch();
+    if (batch !== undefined) {
+      const maybe = mapFn(batch);
+      if (maybe !== undefined) {
+        yield maybe;
+      }
+    }
+
+    if (this._stopped) {
+      return;
+    }
+
+    // eslint-disable-next-line no-await-in-loop
+    while ((batch = await this.nextBatch()) !== undefined) {
+      const maybe = mapFn(batch);
+      if (maybe !== undefined) {
+        yield maybe;
+      }
+    }
+  }
+
+  private _push(item: T): void {
+    if (this._stopped) {
+      this._resolve(undefined);
+    }
+    this._items.push(item);
+    this._resolve(this.batch());
+  }
+
+  private _resolve(maybeIterable: Generator<T> | undefined): void {
+    for (const resolve of this._resolvers) {
+      resolve(maybeIterable);
+    }
+    this._resolvers = [];
+  }
+
+  private *batch(): Generator<T> {
+    let item: T | undefined;
+    while ((item = this._items.shift()) !== undefined) {
+      yield item;
+    }
+  }
+}