emit all deferred fragments as soon as possible

yaacovCR · yaacovCR · commit 929cdc84cea6 · 2023-12-22T11:30:16.000+02:00
diff --git a/src/execution/IncrementalPublisher.ts b/src/execution/IncrementalPublisher.ts
@@ -196,14 +196,11 @@ export class IncrementalPublisher {
     this._reset();
   }
 
-  reportNewDeferFragmentRecord(
-    deferredFragmentRecord: DeferredFragmentRecord,
-    parentIncrementalResultRecord:
-      | InitialResultRecord
-      | DeferredFragmentRecord
-      | StreamItemsRecord,
+  reportNewSubsequentResultRecord(
+    subsequentResultRecord: SubsequentResultRecord,
+    parentIncrementalDataRecord: IncrementalDataRecord,
   ): void {
-    parentIncrementalResultRecord.children.add(deferredFragmentRecord);
+    parentIncrementalDataRecord.children.add(subsequentResultRecord);
   }
 
   reportNewDeferredGroupedFieldSetRecord(
@@ -217,19 +214,6 @@ export class IncrementalPublisher {
     }
   }
 
-  reportNewStreamItemsRecord(
-    streamItemsRecord: StreamItemsRecord,
-    parentIncrementalDataRecord: IncrementalDataRecord,
-  ): void {
-    if (isDeferredGroupedFieldSetRecord(parentIncrementalDataRecord)) {
-      for (const parent of parentIncrementalDataRecord.deferredFragmentRecords) {
-        parent.children.add(streamItemsRecord);
-      }
-    } else {
-      parentIncrementalDataRecord.children.add(streamItemsRecord);
-    }
-  }
-
   completeDeferredGroupedFieldSet(
     deferredGroupedFieldSetRecord: DeferredGroupedFieldSetRecord,
     data: ObjMap<unknown>,
@@ -349,8 +333,9 @@ export class IncrementalPublisher {
 
     const streams = new Set<StreamRecord>();
 
-    const children = this._getChildren(erroringIncrementalDataRecord);
-    const descendants = this._getDescendants(children);
+    const descendants = this._getDescendants(
+      erroringIncrementalDataRecord.children,
+    );
 
     for (const child of descendants) {
       if (!this._nullsChildSubsequentResultRecord(child, nullPathArray)) {
@@ -538,19 +523,19 @@ export class IncrementalPublisher {
     const incrementalResults: Array<IncrementalResult> = [];
     const completedResults: Array<CompletedResult> = [];
     for (const subsequentResultRecord of completedRecords) {
-      for (const child of subsequentResultRecord.children) {
-        if (child.filtered) {
-          continue;
-        }
-        const pendingSource = isStreamItemsRecord(child)
-          ? child.streamRecord
-          : child;
-        if (!pendingSource.pendingSent) {
-          newPendingSources.add(pendingSource);
-        }
-        this._publish(child);
-      }
       if (isStreamItemsRecord(subsequentResultRecord)) {
+        for (const child of subsequentResultRecord.children) {
+          if (child.filtered) {
+            continue;
+          }
+          const pendingSource = isStreamItemsRecord(child)
+            ? child.streamRecord
+            : child;
+          if (!pendingSource.pendingSent) {
+            newPendingSources.add(pendingSource);
+          }
+          this._publish(child);
+        }
         if (subsequentResultRecord.isFinalRecord) {
           newPendingSources.delete(subsequentResultRecord.streamRecord);
           completedResults.push(
@@ -585,6 +570,18 @@ export class IncrementalPublisher {
           continue;
         }
         for (const deferredGroupedFieldSetRecord of subsequentResultRecord.deferredGroupedFieldSetRecords) {
+          for (const child of deferredGroupedFieldSetRecord.children) {
+            if (child.filtered) {
+              continue;
+            }
+            const pendingSource = isStreamItemsRecord(child)
+              ? child.streamRecord
+              : child;
+            if (!pendingSource.pendingSent) {
+              newPendingSources.add(pendingSource);
+            }
+            this._publish(child);
+          }
           if (!deferredGroupedFieldSetRecord.sent) {
             deferredGroupedFieldSetRecord.sent = true;
             const incrementalResult: IncrementalDeferResult =
@@ -613,6 +610,8 @@ export class IncrementalPublisher {
     let idWithLongestPath: string | undefined;
     for (const deferredFragmentRecord of deferredFragmentRecords) {
       const id = deferredFragmentRecord.id;
+      // TODO add test
+      /* c8 ignore next 3 */
       if (id === undefined) {
         continue;
       }
@@ -668,39 +667,27 @@ export class IncrementalPublisher {
 
     if (subsequentResultRecord._pending.size > 0) {
       this._introduce(subsequentResultRecord);
-    } else if (
-      subsequentResultRecord.deferredGroupedFieldSetRecords.size > 0 ||
-      subsequentResultRecord.children.size > 0
-    ) {
+    } else if (subsequentResultRecord.deferredGroupedFieldSetRecords.size > 0) {
       this._push(subsequentResultRecord);
     }
   }
 
-  private _getChildren(
-    erroringIncrementalDataRecord: IncrementalDataRecord,
-  ): ReadonlySet<SubsequentResultRecord> {
-    const children = new Set<SubsequentResultRecord>();
-    if (isDeferredGroupedFieldSetRecord(erroringIncrementalDataRecord)) {
-      for (const erroringIncrementalResultRecord of erroringIncrementalDataRecord.deferredFragmentRecords) {
-        for (const child of erroringIncrementalResultRecord.children) {
-          children.add(child);
-        }
-      }
-    } else {
-      for (const child of erroringIncrementalDataRecord.children) {
-        children.add(child);
-      }
-    }
-    return children;
-  }
-
   private _getDescendants(
     children: ReadonlySet<SubsequentResultRecord>,
     descendants = new Set<SubsequentResultRecord>(),
   ): ReadonlySet<SubsequentResultRecord> {
     for (const child of children) {
       descendants.add(child);
-      this._getDescendants(child.children, descendants);
+      if (isStreamItemsRecord(child)) {
+        this._getDescendants(child.children, descendants);
+      } else {
+        for (const deferredGroupedFieldSetRecord of child.deferredGroupedFieldSetRecords) {
+          this._getDescendants(
+            deferredGroupedFieldSetRecord.children,
+            descendants,
+          );
+        }
+      }
     }
     return descendants;
   }
@@ -736,12 +723,6 @@ export class IncrementalPublisher {
   }
 }
 
-function isDeferredGroupedFieldSetRecord(
-  incrementalDataRecord: unknown,
-): incrementalDataRecord is DeferredGroupedFieldSetRecord {
-  return incrementalDataRecord instanceof DeferredGroupedFieldSetRecord;
-}
-
 function isStreamItemsRecord(
   subsequentResultRecord: unknown,
 ): subsequentResultRecord is StreamItemsRecord {
@@ -764,6 +745,7 @@ export class DeferredGroupedFieldSetRecord {
   deferredFragmentRecords: ReadonlyArray<DeferredFragmentRecord>;
   groupedFieldSet: GroupedFieldSet;
   shouldInitiateDefer: boolean;
+  children: Set<SubsequentResultRecord>;
   errors: Array<GraphQLError>;
   data: ObjMap<unknown> | undefined;
   sent: boolean;
@@ -778,6 +760,7 @@ export class DeferredGroupedFieldSetRecord {
     this.deferredFragmentRecords = opts.deferredFragmentRecords;
     this.groupedFieldSet = opts.groupedFieldSet;
     this.shouldInitiateDefer = opts.shouldInitiateDefer;
+    this.children = new Set();
     this.errors = [];
     this.sent = false;
   }
@@ -788,7 +771,6 @@ export class DeferredFragmentRecord {
   path: ReadonlyArray<string | number>;
   label: string | undefined;
   id: string | undefined;
-  children: Set<SubsequentResultRecord>;
   deferredGroupedFieldSetRecords: Set<DeferredGroupedFieldSetRecord>;
   errors: Array<GraphQLError>;
   filtered: boolean;
@@ -798,7 +780,6 @@ export class DeferredFragmentRecord {
   constructor(opts: { path: Path | undefined; label: string | undefined }) {
     this.path = pathToArray(opts.path);
     this.label = opts.label;
-    this.children = new Set();
     this.filtered = false;
     this.deferredGroupedFieldSetRecords = new Set();
     this.errors = [];
diff --git a/src/execution/__tests__/defer-test.ts b/src/execution/__tests__/defer-test.ts
@@ -946,41 +946,29 @@ describe('Execute: defer directive', () => {
             },
           },
         },
-        pending: [{ id: '0', path: ['hero'] }],
+        pending: [
+          { id: '0', path: ['hero'] },
+          { id: '1', path: ['hero', 'nestedObject'] },
+          { id: '2', path: ['hero', 'nestedObject', 'deeperObject'] },
+        ],
         hasNext: true,
       },
       {
-        pending: [{ id: '1', path: ['hero', 'nestedObject'] }],
         incremental: [
           {
             data: { bar: 'bar' },
-            id: '0',
-            subPath: ['nestedObject', 'deeperObject'],
+            id: '2',
           },
-        ],
-        completed: [{ id: '0' }],
-        hasNext: true,
-      },
-      {
-        pending: [{ id: '2', path: ['hero', 'nestedObject', 'deeperObject'] }],
-        incremental: [
           {
             data: { baz: 'baz' },
-            id: '1',
-            subPath: ['deeperObject'],
+            id: '2',
           },
-        ],
-        hasNext: true,
-        completed: [{ id: '1' }],
-      },
-      {
-        incremental: [
           {
             data: { bak: 'bak' },
             id: '2',
           },
         ],
-        completed: [{ id: '2' }],
+        completed: [{ id: '0' }, { id: '1' }, { id: '2' }],
         hasNext: false,
       },
     ]);
@@ -1023,34 +1011,27 @@ describe('Execute: defer directive', () => {
           },
         },
         pending: [
-          { id: '0', path: ['hero'] },
+          { id: '0', path: ['hero', 'nestedObject', 'deeperObject'] },
           { id: '1', path: ['hero', 'nestedObject', 'deeperObject'] },
         ],
         hasNext: true,
       },
       {
-        pending: [{ id: '2', path: ['hero', 'nestedObject', 'deeperObject'] }],
         incremental: [
           {
             data: {
               foo: 'foo',
             },
-            id: '1',
+            id: '0',
           },
-        ],
-        completed: [{ id: '0' }, { id: '1' }],
-        hasNext: true,
-      },
-      {
-        incremental: [
           {
             data: {
               bar: 'bar',
             },
-            id: '2',
+            id: '1',
           },
         ],
-        completed: [{ id: '2' }],
+        completed: [{ id: '0' }, { id: '1' }],
         hasNext: false,
       },
     ]);
diff --git a/src/execution/buildFieldPlan.ts b/src/execution/buildFieldPlan.ts
@@ -65,11 +65,14 @@ export function buildFieldPlan(
   for (const [responseKey, fieldDetailsList] of fields) {
     const deferUsageSet = new Set<DeferUsage>();
     let inOriginalResult = false;
+    let inParentResult = false;
     for (const fieldDetails of fieldDetailsList) {
       const deferUsage = fieldDetails.deferUsage;
       if (deferUsage === undefined) {
         inOriginalResult = true;
         continue;
+      } else if (parentDeferUsages.has(deferUsage)) {
+        inParentResult = true;
       }
       deferUsageSet.add(deferUsage);
       if (!knownDeferUsages.has(deferUsage)) {
@@ -79,13 +82,10 @@ export function buildFieldPlan(
     }
     if (inOriginalResult) {
       deferUsageSet.clear();
-    } else {
+    } else if (inParentResult) {
       deferUsageSet.forEach((deferUsage) => {
-        const ancestors = getAncestors(deferUsage);
-        for (const ancestor of ancestors) {
-          if (deferUsageSet.has(ancestor)) {
-            deferUsageSet.delete(deferUsage);
-          }
+        if (!parentDeferUsages.has(deferUsage)) {
+          deferUsageSet.delete(deferUsage);
         }
       });
     }
@@ -153,13 +153,3 @@ export function buildFieldPlan(
     newDeferUsages: Array.from(newDeferUsages),
   };
 }
-
-function getAncestors(deferUsage: DeferUsage): ReadonlyArray<DeferUsage> {
-  const ancestors: Array<DeferUsage> = [];
-  let parentDeferUsage: DeferUsage | undefined = deferUsage.parentDeferUsage;
-  while (parentDeferUsage !== undefined) {
-    ancestors.unshift(parentDeferUsage);
-    parentDeferUsage = parentDeferUsage.parentDeferUsage;
-  }
-  return ancestors;
-}
diff --git a/src/execution/collectFields.ts b/src/execution/collectFields.ts
@@ -28,7 +28,6 @@ import { getDirectiveValues } from './values.js';
 
 export interface DeferUsage {
   label: string | undefined;
-  parentDeferUsage: DeferUsage | undefined;
 }
 
 export interface FieldDetails {
@@ -159,7 +158,6 @@ function collectFieldsImpl(
           operation,
           variableValues,
           selection,
-          parentDeferUsage,
         );
 
         collectFieldsImpl(
@@ -179,7 +177,6 @@ function collectFieldsImpl(
           operation,
           variableValues,
           selection,
-          parentDeferUsage,
         );
 
         if (
@@ -223,7 +220,6 @@ function getDeferUsage(
   operation: OperationDefinitionNode,
   variableValues: { [variable: string]: unknown },
   node: FragmentSpreadNode | InlineFragmentNode,
-  parentDeferUsage: DeferUsage | undefined,
 ): DeferUsage | undefined {
   const defer = getDirectiveValues(GraphQLDeferDirective, node, variableValues);
 
@@ -242,7 +238,6 @@ function getDeferUsage(
 
   return {
     label: typeof defer.label === 'string' ? defer.label : undefined,
-    parentDeferUsage,
   };
 }
 
diff --git a/src/execution/execute.ts b/src/execution/execute.ts
