Fix potential issue with nested @defer in non-deferrable case (#2312)

The code handling `@defer` works roughly like this:
1. as we cross a query element (field or spread) with a `@defer`, we
  remember that we just "entered" a deferred section.
2. as we handle the initial elements within that `@defer`, we notice
  that we just entered a `@defer`, ignore in which subgraph we were,
  and instead attempt to "re-enter" subgraphs using a key.
3. if that work, we continue from whichever possible subgraphs we just
  re-entered.
4. Otherwise, if we can't handle the new element by re-entering a
  subgraph, then we fall back to "normal" operation and try to handle
  the element from the original subgraph in which we were before
  entering the @defer: this correspond to case where we cannot do a
  router-based-defer due to not having proper @key to do it.

The issue however is that we were trying to re-enter subgraphs even when
the next element was a fragment. This meant that we re-entered any
subgraphs that had a key, but that may exclude the one subgraph in which
we previously were. If we get a new `@defer` followed by a field that
was only in the "original" subgraph, then we could end up not being able
to find that field _even_ when falling back to step 4 above, because we
had already re-entered different subgraphs due to the 1st `@defer`.

The solution implemented by this commit consists in delaying the
subgraph re-entry until we actually get to a field, since this is where
we _have_ to re-enter.

gateway-js/CHANGELOG.md

@@ -7,6 +7,7 @@ This CHANGELOG pertains only to Apollo Federation packages in the 2.x range. The
 ## 2.2.3
 
 - Fix possible assertion error during query planning [PR #2299](https://github.com/apollographql/federation/pull/2299).
+- Fix potential issue with nested `@defer` in non-deferrable case [PR #2312](https://github.com/apollographql/federation/pull/2312).
 
 ## 2.2.2
 

query-graphs-js/src/graphPath.ts

@@ -335,7 +335,7 @@ export class GraphPath<TTrigger, RV extends Vertex = Vertex, TNullEdge extends n
               edgeConditions: withReplacedLastElement(this.props.edgeConditions, conditionsResolution.pathTree ?? null),
               edgeToTail: updatedEdge,
               runtimeTypesOfTail: runtimeTypesWithoutPreviousCast,
-              deferOnTail: defer,
+              deferOnTail: defer ?? this.props.deferOnTail,
             });
           }
         }
@@ -370,7 +370,10 @@ export class GraphPath<TTrigger, RV extends Vertex = Vertex, TNullEdge extends n
       edgeToTail: edge,
       runtimeTypesOfTail: updateRuntimeTypes(this.props.runtimeTypesOfTail, edge),
       runtimeTypesBeforeTailIfLastIsCast: edge?.transition?.kind === 'DownCast' ? this.props.runtimeTypesOfTail : undefined,
-      deferOnTail: defer,
+      // If there is no new `defer` taking precedence, and the edge is downcast, then we inherit the prior version. This
+      // is because we only try to re-enter subgraphs for @defer on concrete fields, and so as long as we add downcasts,
+      // we should remember that we still need to try re-entering the subgraph.
+      deferOnTail: defer ?? (edge && edge.transition.kind === 'DownCast' ? this.props.deferOnTail : undefined),
     });
   }
 
@@ -1586,9 +1589,10 @@ export function advanceSimultaneousPathsWithOperation<V extends Vertex>(
     let options: SimultaneousPaths<V>[] | undefined = undefined;
 
     debug.group(() => `Computing options for ${path}`);
+    const shouldReenterSubgraph = path.deferOnTail && operation.kind === 'Field';
     // If we're just entering a deferred section, then we will need to re-enter subgraphs, so we should not consider
-    // direct options and instead force and indirect path.
-    if (!path.deferOnTail) {
+    // direct options and instead force an indirect path.
+    if (!shouldReenterSubgraph) {
       debug.group(() => `Direct options`);
       options = advanceWithOperation(
         supergraphSchema,
@@ -1620,41 +1624,43 @@ export function advanceSimultaneousPathsWithOperation<V extends Vertex>(
 
     // If there was not valid direct path (or we didn't check those because we enter a defer), that's ok, we'll just try with non-collecting edges.
     options = options ?? [];
-    debug.group(`Computing indirect paths:`);
-    // Then adds whatever options can be obtained by taking some non-collecting edges first.
-    const pathsWithNonCollecting = subgraphSimultaneousPaths.indirectOptions(updatedContext, i);
-    debug.groupEnd(() => pathsWithNonCollecting.paths.length == 0 ? `no indirect paths` : `${pathsWithNonCollecting.paths.length} indirect paths`);
-    if (pathsWithNonCollecting.paths.length > 0) {
-      debug.group('Validating indirect options:');
-      for (const pathWithNonCollecting of pathsWithNonCollecting.paths) {
-        debug.group(() => `For indirect path ${pathWithNonCollecting}:`);
-        const pathWithOperation = advanceWithOperation(
-          supergraphSchema,
-          pathWithNonCollecting,
-          operation,
-          updatedContext,
-          subgraphSimultaneousPaths.conditionResolver
-        );
-        // If we can't advance the operation after that path, ignore it, it's just not an option.
-        if (!pathWithOperation) {
-          debug.groupEnd(() => `Ignoring: cannot be advanced with ${operation}`);
-          continue;
+    if (operation.kind === 'Field') {
+      debug.group(`Computing indirect paths:`);
+      // Then adds whatever options can be obtained by taking some non-collecting edges first.
+      const pathsWithNonCollecting = subgraphSimultaneousPaths.indirectOptions(updatedContext, i);
+      debug.groupEnd(() => pathsWithNonCollecting.paths.length == 0 ? `no indirect paths` : `${pathsWithNonCollecting.paths.length} indirect paths`);
+      if (pathsWithNonCollecting.paths.length > 0) {
+        debug.group('Validating indirect options:');
+        for (const pathWithNonCollecting of pathsWithNonCollecting.paths) {
+          debug.group(() => `For indirect path ${pathWithNonCollecting}:`);
+          const pathWithOperation = advanceWithOperation(
+            supergraphSchema,
+            pathWithNonCollecting,
+            operation,
+            updatedContext,
+            subgraphSimultaneousPaths.conditionResolver
+          );
+          // If we can't advance the operation after that path, ignore it, it's just not an option.
+          if (!pathWithOperation) {
+            debug.groupEnd(() => `Ignoring: cannot be advanced with ${operation}`);
+            continue;
+          }
+          debug.groupEnd(() => `Adding valid option: ${pathWithOperation}`);
+          // advancedWithOperation can return an empty list only if the operation if a fragment with a condition that, on top of the "current" type
+          // is unsatisfiable. But as we've only taken type preserving transitions, we cannot get an empty results at this point if we haven't
+          // had one when testing direct transitions above (in which case we have exited the method early).
+          assert(pathWithOperation.length > 0, () => `Unexpected empty options after non-collecting path ${pathWithNonCollecting} for ${operation}`);
+          options = options.concat(pathWithOperation);
         }
-        debug.groupEnd(() => `Adding valid option: ${pathWithOperation}`);
-        // advancedWithOperation can return an empty list only if the operation if a fragment with a condition that, on top of the "current" type
-        // is unsatisfiable. But as we've only taken type preserving transitions, we cannot get an empty results at this point if we haven't
-        // had one when testing direct transitions above (in which case we have exited the method early).
-        assert(pathWithOperation.length > 0, () => `Unexpected empty options after non-collecting path ${pathWithNonCollecting} for ${operation}`);
-        options = options.concat(pathWithOperation);
+        debug.groupEnd();
       }
-      debug.groupEnd();
     }
 
     // If we were entering a @defer, we've skipped the potential "direct" options because we need an "indirect" one (a key/root query)
     // to be able to actualy defer. But in rare cases, it's possible we actually couldn't resolve the key fields needed to take a key
     // but could still find a direct path. If so, it means it's a corner case where we cannot do query-planner-based-@defer and have
     // to fall back on not deferring.
-    if (options.length === 0 && path.deferOnTail) {
+    if (options.length === 0 && shouldReenterSubgraph) {
       debug.group(() => `Cannot defer (no indirect options); falling back to direct options`);
       options = advanceWithOperation(
         supergraphSchema,

query-planner-js/CHANGELOG.md

@@ -4,6 +4,7 @@ This CHANGELOG pertains only to Apollo Federation packages in the 2.x range. The
 
 ## vNext
 
+- Fix potential issue with nested `@defer` in non-deferrable case [PR #2312](https://github.com/apollographql/federation/pull/2312).
 - Fix possible assertion error during query planning [PR #2299](https://github.com/apollographql/federation/pull/2299).
 
 ## 2.2.2

query-planner-js/src/__tests__/buildPlan.defer.test.ts

@@ -1314,6 +1314,84 @@ describe('nested @defer', () => {
       }
     `);
   });
+
+  test('on entity but with unuseful key', () => {
+    const subgraph1 = {
+      name: 'Subgraph1',
+      typeDefs: gql`
+        type Query {
+          t : T
+        }
+
+        type T {
+          id: ID! @shareable
+          a: Int
+          b: Int
+        }
+      `
+    }
+
+    const subgraph2 = {
+      name: 'Subgraph2',
+      typeDefs: gql`
+        type T @key(fields: "id") {
+          id: ID!
+        }
+      `
+    }
+
+    const [api, queryPlanner] = composeAndCreatePlannerWithDefer(subgraph1, subgraph2);
+    const operation = operationFromDocument(api, gql`
+      {
+        t {
+          ... @defer {
+            a
+            ... @defer {
+              b
+            }
+          }
+        }
+      }
+    `);
+
+    const plan = queryPlanner.buildQueryPlan(operation);
+    // Note that nothing can effectively be deferred, so everything is fetched in the very first fetch, but
+    // we then have deferred sections that just sub-select what each defer should return to "simulate" the
+    // deferred responses.
+    expect(plan).toMatchInlineSnapshot(`
+      QueryPlan {
+        Defer {
+          Primary {
+            :
+            Fetch(service: "Subgraph1") {
+              {
+                t {
+                  a
+                  b
+                }
+              }
+            }
+          }, [
+            Deferred(depends: [], path: "t") {
+              Defer {
+                Primary {
+                  {
+                    a
+                  }:
+                }, [
+                  Deferred(depends: [], path: "t") {
+                    {
+                      b
+                    }:
+                  },
+                ]
+              }
+            },
+          ]
+        },
+      }
+    `);
+  });
 });
 
 describe('@defer on mutation', () => {

query-planner-js/src/buildPlan.ts

@@ -268,6 +268,7 @@ class QueryPlanningTaversal<RV extends Vertex> {
 
   private handleOpenBranch(selection: Selection, options: SimultaneousPathsWithLazyIndirectPaths<RV>[]) {
     const operation = selection.element();
+    debug.group(() => `Handling open branch: ${operation}`);
     let newOptions: SimultaneousPathsWithLazyIndirectPaths<RV>[] = [];
     for (const option of options) {
       const followupForOption = advanceSimultaneousPathsWithOperation(this.supergraphSchema, option, operation);
@@ -305,6 +306,7 @@ class QueryPlanningTaversal<RV extends Vertex> {
         if (operation.kind === 'FragmentElement') {
           this.closedBranches.push([option.paths.map(p => terminateWithNonRequestedTypenameField(p))]);
         }
+        debug.groupEnd(() => `Terminating branch with no possible results`);
         return;
       }
       newOptions = newOptions.concat(followupForOption);
@@ -316,13 +318,14 @@ class QueryPlanningTaversal<RV extends Vertex> {
       // This should never happen for a top-level query planning (unless the supergraph has *not* been
       // validated), but can happen when computing sub-plans for a key condition.
       if (this.isTopLevel) {
-        debug.log(`No valid options to advance ${selection} from ${advanceOptionsToString(options)}`);
+        debug.groupEnd(() => `No valid options to advance ${selection} from ${advanceOptionsToString(options)}`);
         throw new Error(`Was not able to find any options for ${selection}: This shouldn't have happened.`);
       } else {
         // We clear both open branches and closed ones as a mean to terminate the plan computation with
         // no plan
         this.stack.splice(0, this.stack.length);
         this.closedBranches.splice(0, this.closedBranches.length);
+        debug.groupEnd(() => `No possible plan for ${selection} from ${advanceOptionsToString(options)}; terminating condition`);
         return;
       }
     }
@@ -331,9 +334,11 @@ class QueryPlanningTaversal<RV extends Vertex> {
       for (const branch of mapOptionsToSelections(selection.selectionSet, newOptions)) {
         this.stack.push(branch);
       }
+      debug.groupEnd();
     } else {
       const updated = this.maybeEliminateStrictlyMoreCostlyPaths(newOptions);
       this.closedBranches.push(updated);
+      debug.groupEnd(() => `Branch finished with ${updated.length} options`);
     }
   }