polish(incremental): replace iterator with current and next accessors

-- we do not need a formal iterator as we are the only consumer.

Author: yaacovCR

src/execution/IncrementalGraph.ts

@@ -28,7 +28,7 @@ export class IncrementalGraph {
 
   private _completedQueue: Array<IncrementalDataRecordResult>;
   private _nextQueue: Array<
-    (iterable: IteratorResult<Iterable<IncrementalDataRecordResult>>) => void
+    (iterable: Iterable<IncrementalDataRecordResult> | undefined) => void
   >;
 
   constructor() {
@@ -70,37 +70,32 @@ export class IncrementalGraph {
     }
   }
 
-  completedIncrementalData() {
-    return {
-      [Symbol.asyncIterator]() {
-        return this;
-      },
-      next: (): Promise<
-        IteratorResult<Iterable<IncrementalDataRecordResult>>
-      > => {
-        const firstResult = this._completedQueue.shift();
-        if (firstResult !== undefined) {
-          return Promise.resolve({
-            value: this._yieldCurrentCompletedIncrementalData(firstResult),
-            done: false,
-          });
-        }
-        const { promise, resolve } =
-          promiseWithResolvers<
-            IteratorResult<Iterable<IncrementalDataRecordResult>>
-          >();
-        this._nextQueue.push(resolve);
-        return promise;
-      },
-      return: (): Promise<
-        IteratorResult<Iterable<IncrementalDataRecordResult>>
-      > => {
-        for (const resolve of this._nextQueue) {
-          resolve({ value: undefined, done: true });
-        }
-        return Promise.resolve({ value: undefined, done: true });
-      },
-    };
+  *currentCompletedBatch(): Generator<IncrementalDataRecordResult> {
+    let completed;
+    while ((completed = this._completedQueue.shift()) !== undefined) {
+      yield completed;
+    }
+    if (this._rootNodes.size === 0) {
+      for (const resolve of this._nextQueue) {
+        resolve(undefined);
+      }
+    }
+  }
+
+  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 {
@@ -327,24 +322,13 @@ export class IncrementalGraph {
     first: IncrementalDataRecordResult,
   ): Generator<IncrementalDataRecordResult> {
     yield first;
-    let completed;
-    while ((completed = this._completedQueue.shift()) !== undefined) {
-      yield completed;
-    }
-    if (this._rootNodes.size === 0) {
-      for (const resolve of this._nextQueue) {
-        resolve({ value: undefined, done: true });
-      }
-    }
+    yield* this.currentCompletedBatch();
   }
 
   private _enqueue(completed: IncrementalDataRecordResult): void {
     const next = this._nextQueue.shift();
     if (next !== undefined) {
-      next({
-        value: this._yieldCurrentCompletedIncrementalData(completed),
-        done: false,
-      });
+      next(this._yieldCurrentCompletedIncrementalData(completed));
       return;
     }
     this._completedQueue.push(completed);

src/execution/IncrementalPublisher.ts

@@ -135,14 +135,10 @@ class IncrementalPublisher {
         completed: [],
       };
 
-      const completedIncrementalData =
-        this._incrementalGraph.completedIncrementalData();
-      // use the raw iterator rather than 'for await ... of' so as not to trigger the
-      // '.return()' method on the iterator when exiting the loop with the next value
-      const asyncIterator = completedIncrementalData[Symbol.asyncIterator]();
-      let iteration = await asyncIterator.next();
-      while (!iteration.done) {
-        for (const completedResult of iteration.value) {
+      let batch: Iterable<IncrementalDataRecordResult> | undefined =
+        this._incrementalGraph.currentCompletedBatch();
+      do {
+        for (const completedResult of batch) {
           this._handleCompletedIncrementalData(completedResult, context);
         }
 
@@ -151,7 +147,6 @@ class IncrementalPublisher {
           const hasNext = this._incrementalGraph.hasNext();
 
           if (!hasNext) {
-            // eslint-disable-next-line require-atomic-updates
             isDone = true;
           }
 
@@ -173,8 +168,8 @@ class IncrementalPublisher {
         }
 
         // eslint-disable-next-line no-await-in-loop
-        iteration = await asyncIterator.next();
-      }
+        batch = await this._incrementalGraph.nextCompletedBatch();
+      } while (batch !== undefined);
 
       await this._returnAsyncIteratorsIgnoringErrors();
       return { value: undefined, done: true };
@@ -184,6 +179,7 @@ class IncrementalPublisher {
       IteratorResult<SubsequentIncrementalExecutionResult, void>
     > => {
       isDone = true;
+      this._incrementalGraph.abort();
       await this._returnAsyncIterators();
       return { value: undefined, done: true };
     };
@@ -192,6 +188,7 @@ class IncrementalPublisher {
       error?: unknown,
     ): Promise<IteratorResult<SubsequentIncrementalExecutionResult, void>> => {
       isDone = true;
+      this._incrementalGraph.abort();
       await this._returnAsyncIterators();
       return Promise.reject(error);
     };
@@ -363,8 +360,6 @@ class IncrementalPublisher {
   }
 
   private async _returnAsyncIterators(): Promise<void> {
-    await this._incrementalGraph.completedIncrementalData().return();
-
     const cancellableStreams = this._context.cancellableStreams;
     if (cancellableStreams === undefined) {
       return;