coalesce multiple stream items into single incremental entry

Author: yaacovCR

src/execution/IncrementalPublisher.ts

@@ -224,7 +224,7 @@ class IncrementalPublisher {
   ): ExperimentalIncrementalExecutionResults {
     this._addIncrementalDataRecords(incrementalDataRecords);
     this._pruneEmpty();
-    this._enqueueNewRootDeferredFragments();
+    this._enqueueNew();
 
     const pending = this._pendingSourcesToResults();
 
@@ -257,22 +257,17 @@ class IncrementalPublisher {
         }
 
         if (pending) {
-          this._enqueueCompletedIncrementalDataRecordResult(
+          this._enqueueCompletedDeferredGroupedFieldSetResult(
             incrementalDataRecord.result,
           );
         }
 
         continue;
       }
 
-      const streamRecord = incrementalDataRecord.streamRecord;
-      if (streamRecord.id === undefined) {
-        this._newPending.add(streamRecord);
+      if (incrementalDataRecord.id === undefined) {
+        this._newPending.add(incrementalDataRecord);
       }
-
-      this._enqueueCompletedIncrementalDataRecordResult(
-        incrementalDataRecord.result,
-      );
     }
   }
 
@@ -333,8 +328,9 @@ class IncrementalPublisher {
     }
   }
 
-  private _enqueueNewRootDeferredFragments(): void {
-    const incrementalDataRecordsToWatch = new Set<IncrementalDataRecord>();
+  private _enqueueNew(): void {
+    const deferredGroupedFieldSetRecords =
+      new Set<DeferredGroupedFieldSetRecord>();
     for (const subsequentResultRecord of this._newPending) {
       if (isDeferredFragmentRecord(subsequentResultRecord)) {
         for (const deferredGroupedFieldSetRecord of subsequentResultRecord.deferredGroupedFieldSetRecords) {
@@ -351,19 +347,21 @@ class IncrementalPublisher {
               continue;
             }
           }
-          incrementalDataRecordsToWatch.add(deferredGroupedFieldSetRecord);
+          deferredGroupedFieldSetRecords.add(deferredGroupedFieldSetRecord);
         }
+      } else if (isStreamRecord(subsequentResultRecord)) {
+        this._enqueueStreamRecord(subsequentResultRecord);
       }
     }
 
-    for (const incrementalDataRecord of incrementalDataRecordsToWatch) {
-      this._enqueueCompletedIncrementalDataRecordResult(
-        incrementalDataRecord.result,
+    for (const deferredGroupedFieldSetRecord of deferredGroupedFieldSetRecords) {
+      this._enqueueCompletedDeferredGroupedFieldSetResult(
+        deferredGroupedFieldSetRecord.result,
       );
     }
   }
 
-  private _enqueueCompletedIncrementalDataRecordResult(
+  private _enqueueCompletedDeferredGroupedFieldSetResult(
     result: PromiseOrValue<IncrementalDataRecordResult>,
   ): void {
     if (isPromise(result)) {
@@ -378,6 +376,81 @@ class IncrementalPublisher {
     }
   }
 
+  private _enqueueStreamRecord(record: StreamRecord): void {
+    const streamItemRecords = record.streamItemRecords;
+    if (isPromise(streamItemRecords)) {
+      // eslint-disable-next-line @typescript-eslint/no-floating-promises
+      streamItemRecords.then((resolved) =>
+        this._enqueueStreamItems(record, resolved),
+      );
+    } else {
+      // eslint-disable-next-line @typescript-eslint/no-floating-promises
+      this._enqueueStreamItems(record, streamItemRecords);
+    }
+  }
+
+  private async _enqueueStreamItems(
+    streamRecord: StreamRecord,
+    streamItemRecords: Array<StreamItemRecord>,
+  ): Promise<void> {
+    let items: Array<unknown> = [];
+    let errors: Array<GraphQLError> = [];
+    let incrementalDataRecords: Array<IncrementalDataRecord> = [];
+    let streamItemRecord: StreamItemRecord | undefined;
+    while ((streamItemRecord = streamItemRecords.shift()) !== undefined) {
+      let result = streamItemRecord.result;
+      if (isPromise(result)) {
+        if (items.length > 0) {
+          this._completedResultQueue.push({
+            streamRecord,
+            result:
+              // TODO: add test case or rework for coverage
+              errors.length > 0 /* c8 ignore start */
+                ? { items, errors } /* c8 ignore stop */
+                : { items },
+            incrementalDataRecords,
+          });
+          this._trigger();
+          items = [];
+          errors = [];
+          incrementalDataRecords = [];
+        }
+        // eslint-disable-next-line no-await-in-loop
+        result = await result;
+        // wait an additional tick to coalesce resolving additional promises
+        // within the queue
+        // eslint-disable-next-line no-await-in-loop
+        await Promise.resolve();
+      }
+      if (result.item === undefined) {
+        if (items.length > 0) {
+          this._completedResultQueue.push({
+            streamRecord,
+            result: errors.length > 0 ? { items, errors } : { items },
+            incrementalDataRecords,
+          });
+        }
+        this._completedResultQueue.push(
+          result.errors === undefined
+            ? { streamRecord }
+            : {
+                streamRecord,
+                errors: result.errors,
+              },
+        );
+        this._trigger();
+        return;
+      }
+      items.push(result.item);
+      if (result.errors !== undefined) {
+        errors.push(...result.errors);
+      }
+      if (result.incrementalDataRecords !== undefined) {
+        incrementalDataRecords.push(...result.incrementalDataRecords);
+      }
+    }
+  }
+
   private _pendingSourcesToResults(): Array<PendingResult> {
     const pendingResults: Array<PendingResult> = [];
     for (const pendingSource of this._newPending) {
@@ -424,7 +497,7 @@ class IncrementalPublisher {
           }
 
           this._pruneEmpty();
-          this._enqueueNewRootDeferredFragments();
+          this._enqueueNew();
 
           pending = [...pending, ...this._pendingSourcesToResults()];
         }
@@ -720,9 +793,25 @@ function isNonReconcilableDeferredGroupedFieldSetResult(
   return deferredGroupedFieldSetResult.errors !== undefined;
 }
 
-export interface DeferredGroupedFieldSetRecord {
+/** @internal */
+export class DeferredGroupedFieldSetRecord {
   deferredFragmentRecords: ReadonlyArray<DeferredFragmentRecord>;
   result: PromiseOrValue<DeferredGroupedFieldSetResult>;
+
+  constructor(opts: {
+    deferredFragmentRecords: ReadonlyArray<DeferredFragmentRecord>;
+    result: PromiseOrValue<DeferredGroupedFieldSetResult>;
+  }) {
+    this.deferredFragmentRecords = opts.deferredFragmentRecords;
+    const result = opts.result;
+    this.result = result;
+    if (isPromise(result)) {
+      // eslint-disable-next-line @typescript-eslint/no-floating-promises
+      result.then((resolved) => {
+        this.result = resolved;
+      });
+    }
+  }
 }
 
 export interface SubsequentResultRecord {
@@ -755,55 +844,93 @@ export class DeferredFragmentRecord {
   }
 }
 
-export interface CancellableStreamRecord extends SubsequentResultRecord {
+/** @internal */
+export class StreamRecord {
+  path: Path | undefined;
+  label: string | undefined;
+  id?: string | undefined;
+  streamItemRecords: PromiseOrValue<Array<StreamItemRecord>>;
+
+  constructor(opts: {
+    path: Path | undefined;
+    label: string | undefined;
+    streamItemRecords: PromiseOrValue<Array<StreamItemRecord>>;
+  }) {
+    this.path = opts.path;
+    this.label = opts.label;
+    const streamItemRecords = opts.streamItemRecords;
+    this.streamItemRecords = streamItemRecords;
+    if (isPromise(streamItemRecords)) {
+      // eslint-disable-next-line @typescript-eslint/no-floating-promises
+      streamItemRecords.then((resolved) => {
+        this.streamItemRecords = resolved;
+      });
+    }
+  }
+}
+
+/** @internal */
+export class CancellableStreamRecord extends StreamRecord {
   earlyReturn: () => Promise<unknown>;
+
+  constructor(opts: {
+    path: Path | undefined;
+    label: string | undefined;
+    streamItemRecords: PromiseOrValue<Array<StreamItemRecord>>;
+    earlyReturn: () => Promise<unknown>;
+  }) {
+    super(opts);
+    this.earlyReturn = opts.earlyReturn;
+  }
+}
+
+function isStreamRecord(
+  record: SubsequentResultRecord | IncrementalDataRecord,
+): record is StreamRecord {
+  return record instanceof StreamRecord;
 }
 
 function isCancellableStreamRecord(
   subsequentResultRecord: SubsequentResultRecord,
 ): subsequentResultRecord is CancellableStreamRecord {
-  return 'earlyReturn' in subsequentResultRecord;
-}
-
-interface ReconcilableStreamItemsResult {
-  streamRecord: SubsequentResultRecord;
-  result: BareStreamItemsResult;
-  incrementalDataRecords: ReadonlyArray<IncrementalDataRecord> | undefined;
-  errors?: never;
+  return subsequentResultRecord instanceof CancellableStreamRecord;
 }
 
-export function isReconcilableStreamItemsResult(
-  streamItemsResult: StreamItemsResult,
-): streamItemsResult is ReconcilableStreamItemsResult {
-  return streamItemsResult.result !== undefined;
+/** @internal */
+export class StreamItemRecord {
+  result: PromiseOrValue<StreamItemResult>;
+  constructor(result: PromiseOrValue<StreamItemResult>) {
+    this.result = result;
+    if (isPromise(result)) {
+      // eslint-disable-next-line @typescript-eslint/no-floating-promises
+      result.then((resolved) => {
+        this.result = resolved;
+      });
+    }
+  }
 }
 
-interface TerminatingStreamItemsResult {
-  streamRecord: SubsequentResultRecord;
-  result?: never;
-  incrementalDataRecords?: never;
-  errors?: never;
+export interface StreamItemResult {
+  item?: unknown;
+  incrementalDataRecords?: ReadonlyArray<IncrementalDataRecord> | undefined;
+  errors?: ReadonlyArray<GraphQLError> | undefined;
 }
 
-interface NonReconcilableStreamItemsResult {
+interface StreamItemsResult {
   streamRecord: SubsequentResultRecord;
-  errors: ReadonlyArray<GraphQLError>;
-  result?: never;
+  result?: BareStreamItemsResult | undefined;
+  incrementalDataRecords?: ReadonlyArray<IncrementalDataRecord> | undefined;
+  errors?: ReadonlyArray<GraphQLError> | undefined;
 }
 
-export type StreamItemsResult =
-  | ReconcilableStreamItemsResult
-  | TerminatingStreamItemsResult
-  | NonReconcilableStreamItemsResult;
-
 export interface StreamItemsRecord {
   streamRecord: SubsequentResultRecord;
   result: PromiseOrValue<StreamItemsResult>;
 }
 
 export type IncrementalDataRecord =
   | DeferredGroupedFieldSetRecord
-  | StreamItemsRecord;
+  | StreamRecord;
 
 export type IncrementalDataRecordResult =
   | DeferredGroupedFieldSetResult