Skip deferred fragments in collectFields only when same label is used

Author: robrichard

src/execution/__tests__/collectFields-test.ts

@@ -106,5 +106,26 @@ describe('collectFields', () => {
 
       expect(fieldDetailsList).to.have.lengthOf(1);
     });
+
+    it('prevents infinite loop when deferred fragment is used with same label', () => {
+      const { newDeferUsages } = collectRootFields(`
+        query {
+          ...FragmentOne @defer(label: "Foo")
+        }
+        fragment FragmentOne on Query {
+          ...FragmentTwo @defer(label: "Bar")
+          field
+        }
+        fragment FragmentTwo on Query {
+          ...FragmentThree @defer(label: "Baz")
+          }
+        fragment FragmentThree on Query {
+          ...FragmentOne @defer(label: "Qux")
+          field
+        }
+      `);
+
+      expect(newDeferUsages).to.have.lengthOf(4);
+    });
   });
 });

src/execution/collectFields.ts

@@ -69,12 +69,19 @@ export interface FragmentDetails {
   variableSignatures?: ObjMap<GraphQLVariableSignature> | undefined;
 }
 
+interface VisitedFragmentNames {
+  // Fragments that have been visited without @defer.
+  immediateFragments: Set<string>;
+  // Map of fragment name to a set of labels that have been visited with @defer.
+  deferredFragments: Map<string, Set<string | null>>;
+}
+
 interface CollectFieldsContext {
   schema: GraphQLSchema;
   fragments: ObjMap<FragmentDetails>;
   variableValues: VariableValues;
   runtimeType: GraphQLObjectType;
-  visitedFragmentNames: Map<string, boolean>;
+  visitedFragmentNames: VisitedFragmentNames;
   hideSuggestions: boolean;
   forbiddenDirectiveInstances: Array<DirectiveNode>;
   forbidSkipAndInclude: boolean;
@@ -110,7 +117,10 @@ export function collectFields(
     fragments,
     variableValues,
     runtimeType,
-    visitedFragmentNames: new Map(),
+    visitedFragmentNames: {
+      immediateFragments: new Set(),
+      deferredFragments: new Map(),
+    },
     hideSuggestions,
     forbiddenDirectiveInstances: [],
     forbidSkipAndInclude,
@@ -151,7 +161,10 @@ export function collectSubfields(
     fragments,
     variableValues,
     runtimeType: returnType,
-    visitedFragmentNames: new Map(),
+    visitedFragmentNames: {
+      immediateFragments: new Set(),
+      deferredFragments: new Map(),
+    },
     hideSuggestions,
     forbiddenDirectiveInstances: [],
     forbidSkipAndInclude: false,
@@ -291,26 +304,38 @@ function collectFieldsImpl(
           deferUsage,
         );
 
-        const visitedAsDeferred = visitedFragmentNames.get(fragName);
+        const visitedAsImmediate =
+          visitedFragmentNames.immediateFragments.has(fragName);
+
+        if (visitedAsImmediate) {
+          // Even if the fragment is deferred, we don't need to
+          // collect it again if it has already been visited without @defer.
+          continue;
+        }
 
         let maybeNewDeferUsage: DeferUsage | undefined;
-        if (!newDeferUsage) {
-          // If this spread is not deferred, it may be skipped when already visited
-          // as a non-deferred spread. If it was previously visited as a deferred spread,
-          // it must be revisited.
-          if (visitedAsDeferred === false) {
-            continue;
+
+        if (newDeferUsage) {
+          let visitedAsDeferredByLabelSet =
+            visitedFragmentNames.deferredFragments.get(fragName);
+          if (!visitedAsDeferredByLabelSet) {
+            visitedAsDeferredByLabelSet = new Set();
+            visitedFragmentNames.deferredFragments.set(
+              fragName,
+              visitedAsDeferredByLabelSet,
+            );
           }
-          visitedFragmentNames.set(fragName, false);
-          maybeNewDeferUsage = deferUsage;
-        } else {
-          // If this spread is deferred, it can be skipped if it has already been visited.
-          if (visitedAsDeferred !== undefined) {
+          if (visitedAsDeferredByLabelSet.has(newDeferUsage.label ?? null)) {
+            // If the fragment has already been visited as deferred by the same
+            // label, skip it.
             continue;
           }
-          visitedFragmentNames.set(fragName, true);
           newDeferUsages.push(newDeferUsage);
           maybeNewDeferUsage = newDeferUsage;
+          visitedAsDeferredByLabelSet.add(newDeferUsage.label ?? null);
+        } else {
+          maybeNewDeferUsage = deferUsage;
+          visitedFragmentNames.immediateFragments.add(fragName);
         }
 
         const fragmentVariableSignatures = fragment.variableSignatures;

src/execution/incremental/__tests__/defer-test.ts

@@ -673,6 +673,90 @@ describe('Execute: defer directive', () => {
     });
   });
 
+  it('Can defer same fragment with different labels', async () => {
+    const document = parse(`
+      query HeroNameQuery {
+        hero {
+          ...TopFragment @defer(label: "DeferTop1")
+          ...TopFragment @defer(label: "DeferTop2")
+        }
+      }
+      fragment TopFragment on Hero {
+        name
+      }
+    `);
+    const result = await complete(document);
+    expectJSON(result).toDeepEqual([
+      {
+        data: { hero: {} },
+        pending: [
+          { id: '0', path: ['hero'], label: 'DeferTop1' },
+          { id: '1', path: ['hero'], label: 'DeferTop2' },
+        ],
+        hasNext: true,
+      },
+      {
+        hasNext: false,
+        incremental: [{ id: '0', data: { name: 'Luke' } }],
+        completed: [{ id: '0' }, { id: '1' }],
+      },
+    ]);
+  });
+
+  it('Can skip deferred fragment if same label is used', async () => {
+    const document = parse(`
+      query HeroNameQuery {
+        hero {
+          ...TopFragment @defer(label: "DeferTop")
+          ...TopFragment @defer(label: "DeferTop")
+        }
+      }
+      fragment TopFragment on Hero {
+        name
+      }
+    `);
+    const result = await complete(document);
+    expectJSON(result).toDeepEqual([
+      {
+        data: { hero: {} },
+        pending: [{ id: '0', path: ['hero'], label: 'DeferTop' }],
+        hasNext: true,
+      },
+      {
+        hasNext: false,
+        incremental: [{ id: '0', data: { name: 'Luke' } }],
+        completed: [{ id: '0' }],
+      },
+    ]);
+  });
+
+  it('Can skip deferred fragment if no label is used', async () => {
+    const document = parse(`
+      query HeroNameQuery {
+        hero {
+          ...TopFragment @defer
+          ...TopFragment @defer
+        }
+      }
+      fragment TopFragment on Hero {
+        name
+      }
+    `);
+    const result = await complete(document);
+    expectJSON(result).toDeepEqual([
+      {
+        data: { hero: {} },
+        pending: [{ id: '0', path: ['hero'] }],
+        hasNext: true,
+      },
+      {
+        hasNext: false,
+        incremental: [{ id: '0', data: { name: 'Luke' } }],
+        completed: [{ id: '0' }],
+      },
+    ]);
+  });
+
   it('Can defer an inline fragment', async () => {
     const document = parse(`
       query HeroNameQuery {