experiment: lazily create DeferredFragments

src/execution/IncrementalGraph.ts

@@ -1,12 +1,14 @@
 import { BoxedPromiseOrValue } from '../jsutils/BoxedPromiseOrValue.js';
 import { invariant } from '../jsutils/invariant.js';
 import { isPromise } from '../jsutils/isPromise.js';
+import type { Path } from '../jsutils/Path.js';
+import { pathAtFieldDepth } from '../jsutils/Path.js';
 import { promiseWithResolvers } from '../jsutils/promiseWithResolvers.js';
 
 import type { GraphQLError } from '../error/GraphQLError.js';
 
+import type { DeferUsage } from './collectFields.js';
 import type {
-  DeferredFragmentRecord,
   DeferredGroupedFieldSetRecord,
   IncrementalDataRecord,
   IncrementalDataRecordResult,
@@ -16,22 +18,69 @@ import type {
   SubsequentResultRecord,
 } from './types.js';
 import {
+  DeferredFragmentRecord,
   isDeferredFragmentRecord,
   isDeferredGroupedFieldSetRecord,
 } from './types.js';
 
+/**
+ * @internal
+ */
+class DeferredFragmentFactory {
+  private _rootDeferredFragments = new Map<
+    DeferUsage,
+    DeferredFragmentRecord
+  >();
+
+  get(deferUsage: DeferUsage, path: Path | undefined): DeferredFragmentRecord {
+    const deferUsagePath = pathAtFieldDepth(path, deferUsage.fieldDepth);
+    let deferredFragmentRecords:
+      | Map<DeferUsage, DeferredFragmentRecord>
+      | undefined;
+    if (deferUsagePath === undefined) {
+      deferredFragmentRecords = this._rootDeferredFragments;
+    } else {
+      deferredFragmentRecords = (
+        deferUsagePath as unknown as {
+          deferredFragmentRecords: Map<DeferUsage, DeferredFragmentRecord>;
+        }
+      ).deferredFragmentRecords;
+      if (deferredFragmentRecords === undefined) {
+        deferredFragmentRecords = new Map();
+        (
+          deferUsagePath as unknown as {
+            deferredFragmentRecords: Map<DeferUsage, DeferredFragmentRecord>;
+          }
+        ).deferredFragmentRecords = deferredFragmentRecords;
+      }
+    }
+    let deferredFragmentRecord = deferredFragmentRecords.get(deferUsage);
+    if (deferredFragmentRecord === undefined) {
+      const { label, parentDeferUsage } = deferUsage;
+      deferredFragmentRecord = new DeferredFragmentRecord(
+        deferUsagePath,
+        label,
+        parentDeferUsage,
+      );
+      deferredFragmentRecords.set(deferUsage, deferredFragmentRecord);
+    }
+    return deferredFragmentRecord;
+  }
+}
+
 /**
  * @internal
  */
 export class IncrementalGraph {
   private _rootNodes: Set<SubsequentResultRecord>;
-
+  private _deferredFragmentFactory: DeferredFragmentFactory;
   private _completedQueue: Array<IncrementalDataRecordResult>;
   private _nextQueue: Array<
     (iterable: Iterable<IncrementalDataRecordResult> | undefined) => void
   >;
 
   constructor() {
+    this._deferredFragmentFactory = new DeferredFragmentFactory();
     this._rootNodes = new Set();
     this._completedQueue = [];
     this._nextQueue = [];
@@ -52,8 +101,15 @@ export class IncrementalGraph {
   addCompletedReconcilableDeferredGroupedFieldSet(
     reconcilableResult: ReconcilableDeferredGroupedFieldSetResult,
   ): void {
-    for (const deferredFragmentRecord of reconcilableResult
-      .deferredGroupedFieldSetRecord.deferredFragmentRecords) {
+    const { deferUsages, path } =
+      reconcilableResult.deferredGroupedFieldSetRecord;
+    const deferredFragmentRecords: Array<DeferredFragmentRecord> = [];
+    for (const deferUsage of deferUsages) {
+      const deferredFragmentRecord = this._deferredFragmentFactory.get(
+        deferUsage,
+        path,
+      );
+      deferredFragmentRecords.push(deferredFragmentRecord);
       deferredFragmentRecord.deferredGroupedFieldSetRecords.delete(
         reconcilableResult.deferredGroupedFieldSetRecord,
       );
@@ -64,12 +120,31 @@ export class IncrementalGraph {
     if (incrementalDataRecords !== undefined) {
       this._addIncrementalDataRecords(
         incrementalDataRecords,
-        reconcilableResult.deferredGroupedFieldSetRecord
-          .deferredFragmentRecords,
+        deferredFragmentRecords,
       );
     }
   }
 
+  getDeepestDeferredFragmentRecord(
+    initialDeferUsage: DeferUsage,
+    deferUsages: ReadonlySet<DeferUsage>,
+    path: Path | undefined,
+  ): DeferredFragmentRecord {
+    let bestDeferUsage = initialDeferUsage;
+    let maxFieldDepth = initialDeferUsage.fieldDepth;
+    for (const deferUsage of deferUsages) {
+      if (deferUsage === initialDeferUsage) {
+        continue;
+      }
+      const fieldDepth = deferUsage.fieldDepth;
+      if (fieldDepth > maxFieldDepth) {
+        maxFieldDepth = fieldDepth;
+        bestDeferUsage = deferUsage;
+      }
+    }
+    return this._deferredFragmentFactory.get(bestDeferUsage, path);
+  }
+
   *currentCompletedBatch(): Generator<IncrementalDataRecordResult> {
     let completed;
     while ((completed = this._completedQueue.shift()) !== undefined) {
@@ -102,12 +177,20 @@ export class IncrementalGraph {
     return this._rootNodes.size > 0;
   }
 
-  completeDeferredFragment(deferredFragmentRecord: DeferredFragmentRecord):
+  completeDeferredFragment(
+    deferUsage: DeferUsage,
+    path: Path | undefined,
+  ):
     | {
+        deferredFragmentRecord: DeferredFragmentRecord;
         newRootNodes: ReadonlyArray<SubsequentResultRecord>;
         reconcilableResults: ReadonlyArray<ReconcilableDeferredGroupedFieldSetResult>;
       }
     | undefined {
+    const deferredFragmentRecord = this._deferredFragmentFactory.get(
+      deferUsage,
+      path,
+    );
     if (
       !this._rootNodes.has(deferredFragmentRecord) ||
       deferredFragmentRecord.deferredGroupedFieldSetRecords.size > 0
@@ -119,8 +202,13 @@ export class IncrementalGraph {
     );
     this._removeRootNode(deferredFragmentRecord);
     for (const reconcilableResult of reconcilableResults) {
-      for (const otherDeferredFragmentRecord of reconcilableResult
-        .deferredGroupedFieldSetRecord.deferredFragmentRecords) {
+      const { deferUsages, path: resultPath } =
+        reconcilableResult.deferredGroupedFieldSetRecord;
+      for (const otherDeferUsage of deferUsages) {
+        const otherDeferredFragmentRecord = this._deferredFragmentFactory.get(
+          otherDeferUsage,
+          resultPath,
+        );
         otherDeferredFragmentRecord.reconcilableResults.delete(
           reconcilableResult,
         );
@@ -129,17 +217,22 @@ export class IncrementalGraph {
     const newRootNodes = this._promoteNonEmptyToRoot(
       deferredFragmentRecord.children,
     );
-    return { newRootNodes, reconcilableResults };
+    return { deferredFragmentRecord, newRootNodes, reconcilableResults };
   }
 
   removeDeferredFragment(
-    deferredFragmentRecord: DeferredFragmentRecord,
-  ): boolean {
+    deferUsage: DeferUsage,
+    path: Path | undefined,
+  ): DeferredFragmentRecord | undefined {
+    const deferredFragmentRecord = this._deferredFragmentFactory.get(
+      deferUsage,
+      path,
+    );
     if (!this._rootNodes.has(deferredFragmentRecord)) {
-      return false;
+      return;
     }
     this._removeRootNode(deferredFragmentRecord);
-    return true;
+    return deferredFragmentRecord;
   }
 
   removeStream(streamRecord: StreamRecord): void {
@@ -159,7 +252,12 @@ export class IncrementalGraph {
   ): void {
     for (const incrementalDataRecord of incrementalDataRecords) {
       if (isDeferredGroupedFieldSetRecord(incrementalDataRecord)) {
-        for (const deferredFragmentRecord of incrementalDataRecord.deferredFragmentRecords) {
+        const { deferUsages, path } = incrementalDataRecord;
+        for (const deferUsage of deferUsages) {
+          const deferredFragmentRecord = this._deferredFragmentFactory.get(
+            deferUsage,
+            path,
+          );
           this._addDeferredFragment(
             deferredFragmentRecord,
             initialResultChildren,
@@ -216,9 +314,17 @@ export class IncrementalGraph {
   private _completesRootNode(
     deferredGroupedFieldSetRecord: DeferredGroupedFieldSetRecord,
   ): boolean {
-    return deferredGroupedFieldSetRecord.deferredFragmentRecords.some(
-      (deferredFragmentRecord) => this._rootNodes.has(deferredFragmentRecord),
-    );
+    const { deferUsages, path } = deferredGroupedFieldSetRecord;
+    for (const deferUsage of deferUsages) {
+      const deferredFragmentRecord = this._deferredFragmentFactory.get(
+        deferUsage,
+        path,
+      );
+      if (this._rootNodes.has(deferredFragmentRecord)) {
+        return true;
+      }
+    }
+    return false;
   }
 
   private _addDeferredFragment(
@@ -228,12 +334,16 @@ export class IncrementalGraph {
     if (this._rootNodes.has(deferredFragmentRecord)) {
       return;
     }
-    const parent = deferredFragmentRecord.parent;
-    if (parent === undefined) {
+    const parentDeferUsage = deferredFragmentRecord.parentDeferUsage;
+    if (parentDeferUsage === undefined) {
       invariant(initialResultChildren !== undefined);
       initialResultChildren.add(deferredFragmentRecord);
       return;
     }
+    const parent = this._deferredFragmentFactory.get(
+      parentDeferUsage,
+      deferredFragmentRecord.path,
+    );
     parent.children.add(deferredFragmentRecord);
     this._addDeferredFragment(parent, initialResultChildren);
   }

src/execution/IncrementalPublisher.ts

@@ -8,7 +8,6 @@ import { IncrementalGraph } from './IncrementalGraph.js';
 import type {
   CancellableStreamRecord,
   CompletedResult,
-  DeferredFragmentRecord,
   DeferredGroupedFieldSetResult,
   ExperimentalIncrementalExecutionResults,
   IncrementalDataRecord,
@@ -221,20 +220,21 @@ class IncrementalPublisher {
     deferredGroupedFieldSetResult: DeferredGroupedFieldSetResult,
     context: SubsequentIncrementalExecutionResultContext,
   ): void {
+    const { deferUsages, path } =
+      deferredGroupedFieldSetResult.deferredGroupedFieldSetRecord;
     if (
       isNonReconcilableDeferredGroupedFieldSetResult(
         deferredGroupedFieldSetResult,
       )
     ) {
-      for (const deferredFragmentRecord of deferredGroupedFieldSetResult
-        .deferredGroupedFieldSetRecord.deferredFragmentRecords) {
-        const id = deferredFragmentRecord.id;
-        if (
-          !this._incrementalGraph.removeDeferredFragment(deferredFragmentRecord)
-        ) {
+      for (const deferUsage of deferUsages) {
+        const deferredFragmentRecord =
+          this._incrementalGraph.removeDeferredFragment(deferUsage, path);
+        if (deferredFragmentRecord === undefined) {
           // This can occur if multiple deferred grouped field sets error for a fragment.
           continue;
         }
+        const id = deferredFragmentRecord.id;
         invariant(id !== undefined);
         context.completed.push({
           id,
@@ -248,34 +248,43 @@ class IncrementalPublisher {
       deferredGroupedFieldSetResult,
     );
 
-    for (const deferredFragmentRecord of deferredGroupedFieldSetResult
-      .deferredGroupedFieldSetRecord.deferredFragmentRecords) {
+    for (const deferUsage of deferUsages) {
       const completion = this._incrementalGraph.completeDeferredFragment(
-        deferredFragmentRecord,
+        deferUsage,
+        path,
       );
       if (completion === undefined) {
         continue;
       }
-      const id = deferredFragmentRecord.id;
-      invariant(id !== undefined);
-      const incremental = context.incremental;
-      const { newRootNodes, reconcilableResults } = completion;
+      const { deferredFragmentRecord, newRootNodes, reconcilableResults } =
+        completion;
       context.pending.push(...this._toPendingResults(newRootNodes));
+      const incremental = context.incremental;
       for (const reconcilableResult of reconcilableResults) {
-        const { bestId, subPath } = this._getBestIdAndSubPath(
-          id,
-          deferredFragmentRecord,
-          reconcilableResult,
-        );
+        const { deferUsages: resultDeferUsages, path: resultPath } =
+          reconcilableResult.deferredGroupedFieldSetRecord;
+        const bestDeferredFragmentRecord =
+          this._incrementalGraph.getDeepestDeferredFragmentRecord(
+            deferUsage,
+            resultDeferUsages,
+            resultPath,
+          );
+        const bestId = bestDeferredFragmentRecord.id;
+        invariant(bestId !== undefined);
         const incrementalEntry: IncrementalDeferResult = {
           ...reconcilableResult.result,
           id: bestId,
         };
-        if (subPath !== undefined) {
+        const subPath = pathToArray(resultPath).slice(
+          pathToArray(bestDeferredFragmentRecord.path).length,
+        );
+        if (subPath.length > 0) {
           incrementalEntry.subPath = subPath;
         }
         incremental.push(incrementalEntry);
       }
+      const id = deferredFragmentRecord.id;
+      invariant(id !== undefined);
       context.completed.push({ id });
     }
   }
@@ -326,39 +335,6 @@ class IncrementalPublisher {
     }
   }
 
-  private _getBestIdAndSubPath(
-    initialId: string,
-    initialDeferredFragmentRecord: DeferredFragmentRecord,
-    deferredGroupedFieldSetResult: DeferredGroupedFieldSetResult,
-  ): { bestId: string; subPath: ReadonlyArray<string | number> | undefined } {
-    let maxLength = pathToArray(initialDeferredFragmentRecord.path).length;
-    let bestId = initialId;
-
-    for (const deferredFragmentRecord of deferredGroupedFieldSetResult
-      .deferredGroupedFieldSetRecord.deferredFragmentRecords) {
-      if (deferredFragmentRecord === initialDeferredFragmentRecord) {
-        continue;
-      }
-      const id = deferredFragmentRecord.id;
-      // TODO: add test case for when an fragment has not been released, but might be processed for the shortest path.
-      /* c8 ignore next 3 */
-      if (id === undefined) {
-        continue;
-      }
-      const fragmentPath = pathToArray(deferredFragmentRecord.path);
-      const length = fragmentPath.length;
-      if (length > maxLength) {
-        maxLength = length;
-        bestId = id;
-      }
-    }
-    const subPath = deferredGroupedFieldSetResult.path.slice(maxLength);
-    return {
-      bestId,
-      subPath: subPath.length > 0 ? subPath : undefined,
-    };
-  }
-
   private async _returnAsyncIterators(): Promise<void> {
     const cancellableStreams = this._context.cancellableStreams;
     if (cancellableStreams === undefined) {