publisher: iterate only through ready items

Author: yaacovCR

src/execution/__tests__/stream-test.ts

@@ -1286,6 +1286,10 @@ describe('Execute: stream directive', () => {
       },
       {
         incremental: [
+          {
+            items: [{ name: 'Luke' }],
+            path: ['nestedObject', 'nestedFriendList', 0],
+          },
           {
             data: { scalarField: null },
             path: ['otherNestedObject'],
@@ -1297,11 +1301,10 @@ describe('Execute: stream directive', () => {
               },
             ],
           },
-          {
-            items: [{ name: 'Luke' }],
-            path: ['nestedObject', 'nestedFriendList', 0],
-          },
         ],
+        hasNext: true,
+      },
+      {
         hasNext: false,
       },
     ]);

src/execution/execute.ts

@@ -2129,20 +2129,16 @@ function filterSubsequentPayloads(
 }
 
 function getIncrementalResult(
-  pending: ReadonlySet<IncrementalDataRecord>,
+  completedRecords: ReadonlySet<IncrementalDataRecord>,
   publisher: Publisher<
     IncrementalDataRecord,
     SubsequentIncrementalExecutionResult
   >,
 ): SubsequentIncrementalExecutionResult | undefined {
   const incrementalResults: Array<IncrementalResult> = [];
   let encounteredCompletedAsyncIterator = false;
-  for (const incrementalDataRecord of pending) {
+  for (const incrementalDataRecord of completedRecords) {
     const incrementalResult: IncrementalResult = {};
-    if (!incrementalDataRecord.isCompleted) {
-      continue;
-    }
-    publisher.delete(incrementalDataRecord);
     if (isStreamItemsRecord(incrementalDataRecord)) {
       const items = incrementalDataRecord.items;
       if (incrementalDataRecord.isCompletedAsyncIterator) {
@@ -2165,7 +2161,6 @@ function getIncrementalResult(
     }
     incrementalResults.push(incrementalResult);
   }
-
   return incrementalResults.length
     ? { incremental: incrementalResults, hasNext: publisher.hasNext() }
     : encounteredCompletedAsyncIterator && !publisher.hasNext()
@@ -2211,7 +2206,7 @@ class DeferredFragmentRecord {
     this.parentContext = opts.parentContext;
     this.errors = [];
     this._exeContext = opts.exeContext;
-    this._exeContext.publisher.add(this);
+    this._exeContext.publisher.introduce(this);
     this.isCompleted = false;
     this.data = null;
     this.promise = new Promise<ObjMap<unknown> | null>((resolve) => {
@@ -2221,6 +2216,7 @@ class DeferredFragmentRecord {
     }).then((data) => {
       this.data = data;
       this.isCompleted = true;
+      this._exeContext.publisher.release(this);
     });
   }
 
@@ -2262,7 +2258,7 @@ class StreamItemsRecord {
     this.asyncIterator = opts.asyncIterator;
     this.errors = [];
     this._exeContext = opts.exeContext;
-    this._exeContext.publisher.add(this);
+    this._exeContext.publisher.introduce(this);
     this.isCompleted = false;
     this.items = null;
     this.promise = new Promise<Array<unknown> | null>((resolve) => {
@@ -2272,6 +2268,7 @@ class StreamItemsRecord {
     }).then((items) => {
       this.items = items;
       this.isCompleted = true;
+      this._exeContext.publisher.release(this);
     });
   }
 

src/jsutils/Publisher.ts

@@ -1,9 +1,7 @@
-interface ContainsPromise {
-  promise: Promise<void>;
-}
-
 /** @internal */
-export class Publisher<I extends ContainsPromise, R> {
+export class Publisher<I, R> {
+  _unreleased: Set<I>;
+  _released: Set<I>;
   _pending: Set<I>;
   _update: (completed: Set<I>, publisher: Publisher<I, R>) => R | undefined;
   _onAbruptClose: (pending: ReadonlySet<I>) => Promise<void>;
@@ -19,6 +17,8 @@ export class Publisher<I extends ContainsPromise, R> {
     ) => R | undefined,
     onAbruptClose: (pending: ReadonlySet<I>) => Promise<void>,
   ) {
+    this._unreleased = new Set();
+    this._released = new Set();
     this._pending = new Set();
     this._update = update;
     this._onAbruptClose = onAbruptClose;
@@ -42,11 +42,22 @@ export class Publisher<I extends ContainsPromise, R> {
     return this._pending.size > 0;
   }
 
-  add(item: I) {
+  introduce(item: I) {
+    this._unreleased.add(item);
     this._pending.add(item);
   }
 
+  release(item: I): void {
+    if (this._pending.has(item)) {
+      this._unreleased.delete(item);
+      this._released.add(item);
+      this._trigger();
+    }
+  }
+
   delete(item: I) {
+    this._unreleased.delete(item);
+    this._released.delete(item);
     this._pending.delete(item);
     this._trigger();
   }
@@ -55,29 +66,31 @@ export class Publisher<I extends ContainsPromise, R> {
     let isDone = false;
 
     const _next = async (): Promise<IteratorResult<R, void>> => {
-      if (isDone) {
-        return { value: undefined, done: true };
-      }
-
-      await Promise.race(Array.from(this._pending).map((item) => item.promise));
-
-      if (isDone) {
-        // a different call to next has exhausted all payloads
-        return { value: undefined, done: true };
+      // eslint-disable-next-line no-constant-condition
+      while (true) {
+        if (isDone) {
+          return { value: undefined, done: true };
+        }
+
+        for (const item of this._released) {
+          this._pending.delete(item);
+        }
+        const released = this._released;
+        this._released = new Set();
+
+        const result = this._update(released, this);
+
+        if (!this.hasNext()) {
+          isDone = true;
+        }
+
+        if (result !== undefined) {
+          return { value: result, done: false };
+        }
+
+        // eslint-disable-next-line no-await-in-loop
+        await this._signalled;
       }
-
-      const result = this._update(this._pending, this);
-      const hasNext = this._pending.size > 0;
-
-      if (result === undefined && hasNext) {
-        return _next();
-      }
-
-      if (!hasNext) {
-        isDone = true;
-      }
-
-      return { value: result as R, done: false };
     };
 
     const _return = async (): Promise<IteratorResult<R, void>> => {