Add support to automatically cancel query plan request if planning takes longer than a pre-configured amount of time.

.changeset/beige-suits-occur.md

@@ -0,0 +1,6 @@
+---
+"@apollo/query-planner": minor
+"@apollo/gateway": minor
+---
+
+Query planner now has support to automatically abort query plan requests that take longer than a configured amount of time. Default value is 2 minutes. Value is set by `maxQueryPlanningTime` value in `QueryPlannerConfig` options.

composition-js/src/compose.ts

@@ -17,6 +17,7 @@ import { buildFederatedQueryGraph, buildSupergraphAPIQueryGraph } from "@apollo/
 import { MergeResult, mergeSubgraphs } from "./merging";
 import { validateGraphComposition } from "./validate";
 import { CompositionHint } from "./hints";
+import { performance } from 'perf_hooks';
 
 export type CompositionResult = CompositionFailure | CompositionSuccess;
 
@@ -133,7 +134,10 @@ export function validateSatisfiability({ supergraphSchema, supergraphSdl} : Sati
   const supergraph = supergraphSchema ? new Supergraph(supergraphSchema, null) : Supergraph.build(supergraphSdl, { supportedFeatures: null });
   const supergraphQueryGraph = buildSupergraphAPIQueryGraph(supergraph);
   const federatedQueryGraph = buildFederatedQueryGraph(supergraph, false);
-  return validateGraphComposition(supergraph.schema, supergraph.subgraphNameToGraphEnumValue(), supergraphQueryGraph, federatedQueryGraph);
+
+  // for satisfiability, we don't want to really have a deadline here, so set it far into the future
+  const deadline = performance.now() + 1000 * 60 * 60 * 24; // 1 day
+  return validateGraphComposition(deadline, supergraph.schema, supergraph.subgraphNameToGraphEnumValue(), supergraphQueryGraph, federatedQueryGraph);
 }
 
 type ValidateSubgraphsAndMergeResult = MergeResult | { errors: GraphQLError[] };

composition-js/src/validate.ts

@@ -306,6 +306,7 @@ function generateWitnessValue(type: InputType): any {
  * @param federatedQueryGraph the (federated) `QueryGraph` corresponding the subgraphs having been composed to obtain `supergraphSchema`.
  */
 export function validateGraphComposition(
+  deadline: number,
   supergraphSchema: Schema,
   subgraphNameToGraphEnumValue: Map<string, string>,
   supergraphAPI: QueryGraph,
@@ -315,6 +316,7 @@ export function validateGraphComposition(
   hints? : CompositionHint[],
 } {
   const { errors, hints } = new ValidationTraversal(
+    deadline,
     supergraphSchema,
     subgraphNameToGraphEnumValue,
     supergraphAPI,
@@ -697,12 +699,14 @@ class ValidationTraversal {
   private readonly context: ValidationContext;
 
   constructor(
+    deadline: number,
     supergraphSchema: Schema,
     subgraphNameToGraphEnumValue: Map<string, string>,
     supergraphAPI: QueryGraph,
     federatedQueryGraph: QueryGraph,
   ) {
     this.conditionResolver = simpleValidationConditionResolver({
+      deadline,
       supergraph: supergraphSchema,
       queryGraph: federatedQueryGraph,
       withCaching: true,

query-graphs-js/src/__tests__/graphPath.test.ts

@@ -15,6 +15,9 @@ import {
 import { QueryGraph, Vertex, buildFederatedQueryGraph } from "../querygraph";
 import { emptyContext } from "../pathContext";
 import { simpleValidationConditionResolver } from "../conditionsValidation";
+import { performance } from 'perf_hooks';
+
+const FAR_FUTURE_DEADLINE = performance.now() + 1000 * 60 * 60 * 24;
 
 function parseSupergraph(subgraphs: number, schema: string): { supergraph: Schema, api: Schema, queryGraph: QueryGraph } {
   assert(subgraphs >= 1, 'Should have at least 1 subgraph');
@@ -68,7 +71,7 @@ function createOptions(supergraph: Schema, queryGraph: QueryGraph): Simultaneous
   return createInitialOptions(
     initialPath,
     emptyContext,
-    simpleValidationConditionResolver({ supergraph, queryGraph }),
+    simpleValidationConditionResolver({ deadline: FAR_FUTURE_DEADLINE, supergraph, queryGraph }),
     [],
     [],
     new Map(),
@@ -104,7 +107,7 @@ describe("advanceSimultaneousPathsWithOperation", () => {
     const initial = createOptions(supergraph, queryGraph)[0];
 
     // Then picking `t`, which should be just the one option of picking it in S1 at this point.
-    const allAfterT = advanceSimultaneousPathsWithOperation(supergraph, initial, field(api, "Query.t"), new Map());
+    const allAfterT = advanceSimultaneousPathsWithOperation(FAR_FUTURE_DEADLINE, supergraph, initial, field(api, "Query.t"), new Map());
     assert(allAfterT, 'Should have advanced correctly');
     expect(allAfterT).toHaveLength(1);
     const afterT = allAfterT[0];
@@ -118,7 +121,7 @@ describe("advanceSimultaneousPathsWithOperation", () => {
     expect(indirect.paths[0].toString()).toBe(`Query(S1) --[t]--> T(S1) --[{ otherId } ⊢ key()]--> T(S2) (types: [T])`);
     expect(indirect.paths[1].toString()).toBe(`Query(S1) --[t]--> T(S1) --[{ id } ⊢ key()]--> T(S3) (types: [T])`);
 
-    const allForId = advanceSimultaneousPathsWithOperation(supergraph, afterT, field(api, "T.id"), new Map());
+    const allForId = advanceSimultaneousPathsWithOperation(FAR_FUTURE_DEADLINE, supergraph, afterT, field(api, "T.id"), new Map());
     assert(allForId, 'Should have advanced correctly');
 
     // Here, `id` is a direct path from both of our indirect paths. However, it makes no sense to use the 2nd
@@ -156,7 +159,7 @@ describe("advanceSimultaneousPathsWithOperation", () => {
     const initial = createOptions(supergraph, queryGraph)[0];
 
     // Then picking `t`, which should be just the one option of picking it in S1 at this point.
-    const allAfterT = advanceSimultaneousPathsWithOperation(supergraph, initial, field(api, "Query.t"), new Map());
+    const allAfterT = advanceSimultaneousPathsWithOperation(FAR_FUTURE_DEADLINE, supergraph, initial, field(api, "Query.t"), new Map());
     assert(allAfterT, 'Should have advanced correctly');
     expect(allAfterT).toHaveLength(1);
     const afterT = allAfterT[0];
@@ -170,7 +173,7 @@ describe("advanceSimultaneousPathsWithOperation", () => {
     expect(indirect.paths[0].toString()).toBe(`Query(S1) --[t]--> T(S1) --[{ otherId } ⊢ key()]--> T(S2) (types: [T])`);
     expect(indirect.paths[1].toString()).toBe(`Query(S1) --[t]--> T(S1) --[{ id1 id2 } ⊢ key()]--> T(S3) (types: [T])`);
 
-    const allForId = advanceSimultaneousPathsWithOperation(supergraph, afterT, field(api, "T.id1"), new Map());
+    const allForId = advanceSimultaneousPathsWithOperation(FAR_FUTURE_DEADLINE, supergraph, afterT, field(api, "T.id1"), new Map());
     assert(allForId, 'Should have advanced correctly');
 
     // Here, `id1` is a direct path from both of our indirect paths. However, it makes no sense to use the 2nd
@@ -202,7 +205,7 @@ describe("advanceSimultaneousPathsWithOperation", () => {
     const initial = createOptions(supergraph, queryGraph)[0];
 
     // Then picking `t`, which should be just the one option of picking it in S1 at this point.
-    const allAfterT = advanceSimultaneousPathsWithOperation(supergraph, initial, field(api, "Query.t"), new Map());
+    const allAfterT = advanceSimultaneousPathsWithOperation(FAR_FUTURE_DEADLINE, supergraph, initial, field(api, "Query.t"), new Map());
     assert(allAfterT, 'Should have advanced correctly');
     expect(allAfterT).toHaveLength(1);
     const afterT = allAfterT[0];

query-graphs-js/src/conditionsValidation.ts

@@ -18,6 +18,8 @@ import { cachingConditionResolver } from "./conditionsCaching";
 
 class ConditionValidationState {
   constructor(
+    // deadline by which time we want the computation to abort
+    readonly deadline: number,
     // Selection that belongs to the condition we're validating.
     readonly selection: Selection,
     // All the possible "simultaneous paths" we could be in the subgraph when we reach this state selection.
@@ -28,6 +30,7 @@ class ConditionValidationState {
     const newOptions: SimultaneousPathsWithLazyIndirectPaths[] = [];
     for (const paths of this.subgraphOptions) {
       const pathsOptions = advanceSimultaneousPathsWithOperation(
+        this.deadline,
         supergraph,
         paths,
         this.selection.element,
@@ -50,6 +53,7 @@ class ConditionValidationState {
     }
     return this.selection.selectionSet ? this.selection.selectionSet.selections().map(
       s => new ConditionValidationState(
+        this.deadline,
         s,
         newOptions,
       )
@@ -69,10 +73,12 @@ class ConditionValidationState {
  * conditions are satisfied.
  */
 export function simpleValidationConditionResolver({
+  deadline,
   supergraph,
   queryGraph,
   withCaching,
 }: {
+  deadline: number,
   supergraph: Schema,
   queryGraph: QueryGraph,
   withCaching?: boolean,
@@ -92,7 +98,7 @@ export function simpleValidationConditionResolver({
       new SimultaneousPathsWithLazyIndirectPaths(
         [initialPath],
         context,
-        simpleValidationConditionResolver({ supergraph, queryGraph, withCaching }),
+        simpleValidationConditionResolver({ deadline, supergraph, queryGraph, withCaching }),
         excludedDestinations,
         excludedConditions,
         new Map(),
@@ -103,6 +109,7 @@ export function simpleValidationConditionResolver({
     for (const selection of conditions.selections()) {
       stack.push(
         new ConditionValidationState(
+          deadline,
           selection,
           initialOptions,
         ),

query-graphs-js/src/graphPath.ts

@@ -42,6 +42,7 @@
 import { DownCast, Transition } from "./transition";
 import { PathContext, emptyContext, isPathContext } from "./pathContext";
 import { v4 as uuidv4 } from 'uuid';
+import { performance } from 'perf_hooks';
 
 const debug = newDebugLogger('path');
 
@@ -2210,6 +2211,7 @@
 // The lists of options can be empty, which has the special meaning that the operation is guaranteed to have no results (it corresponds to unsatisfiable conditions),
 // meaning that as far as query planning goes, we can just ignore the operation but otherwise continue.
 export function advanceSimultaneousPathsWithOperation<V extends Vertex>(
+  deadline: number,
   supergraphSchema: Schema,
   subgraphSimultaneousPaths: SimultaneousPathsWithLazyIndirectPaths<V>,
   operation: OperationElement,
@@ -2228,6 +2230,7 @@
     if (!shouldReenterSubgraph) {
       debug.group(() => `Direct options`);
       const { options: advanceOptions, hasOnlyTypeExplodedResults } = advanceWithOperation(
+        deadline,
         supergraphSchema,
         path,
         operation,
@@ -2275,6 +2278,7 @@
         for (const pathWithNonCollecting of pathsWithNonCollecting.paths) {
           debug.group(() => `For indirect path ${pathWithNonCollecting}:`);
           const { options: pathWithOperation } = advanceWithOperation(
+            deadline,
             supergraphSchema,
             pathWithNonCollecting,
             operation,
@@ -2334,6 +2338,7 @@
     if (options.length === 0 && shouldReenterSubgraph) {
       debug.group(() => `Cannot defer (no indirect options); falling back to direct options`);
       const { options: advanceOptions } = advanceWithOperation(
+        deadline,
         supergraphSchema,
         path,
         operation,
@@ -2601,6 +2606,7 @@
 // We also actually need to return a set of options of simultaneous paths. Cause when we type explode, we create simultaneous paths, but
 // as a field might be resolve by multiple subgraphs, we may have options created.
 function advanceWithOperation<V extends Vertex>(
+  deadline: number,
   supergraphSchema: Schema,
   path: OpGraphPath<V>,
   operation: OperationElement,
@@ -2613,6 +2619,10 @@
 } {
   debug.group(() => `Trying to advance ${path} directly with ${operation}`);
 
+  if (performance.now() > deadline) {
+    throw new Error('Query plan took too long to calculate');
+  }
+
   const currentType = path.tail.type;
   if (isFederatedGraphRootType(currentType)) {
     // We cannot advance any operation from there: we need to take the initial non-collecting edges first.
@@ -2725,6 +2735,7 @@
           const castOp = new FragmentElement(currentType, implemType.name);
           debug.group(() => `Handling implementation ${implemType}`);
           const implemOptions = advanceSimultaneousPathsWithOperation(
+            deadline,
             supergraphSchema,
             new SimultaneousPathsWithLazyIndirectPaths([path], context, conditionResolver, [], [], overrideConditions),
             castOp,
@@ -2748,6 +2759,7 @@
           for (const optPaths of implemOptions) {
             debug.group(() => `For ${simultaneousPathsToString(optPaths)}`);
             const withFieldOptions = advanceSimultaneousPathsWithOperation(
+              deadline,
               supergraphSchema,
               optPaths,
               operation,
@@ -2824,6 +2836,7 @@
           debug.group(() => `Trying ${tName}`);
           const castOp = new FragmentElement(currentType, tName, operation.appliedDirectives);
           const implemOptions = advanceSimultaneousPathsWithOperation(
+            deadline,
             supergraphSchema,
             new SimultaneousPathsWithLazyIndirectPaths([path], context, conditionResolver, [], [], overrideConditions),
             castOp,

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

@@ -10721,3 +10721,99 @@ test('ensure we are removing all useless groups', () => {
       }
     `);
 });
+
+describe('`maxQueryPlanningTime` configuration', () => {
+  // Simple schema that should still take longer the a millisecond to plan
+  const typeDefs = gql`
+    type Query {
+      t: T @shareable
+    }
+
+    type T @key(fields: "id") @shareable {
+      id: ID!
+      v1: Int
+      v2: Int
+      v3: Int
+      v4: Int
+    }
+  `;
+
+  const subgraphs = [
+    {
+      name: 'Subgraph1',
+      typeDefs,
+    },
+    {
+      name: 'Subgraph2',
+      typeDefs,
+    },
+  ];
+
+  test('maxQueryPlanningTime works when not set', () => {
+    const config = { debug: { maxEvaluatedPlans: undefined } };
+    const [api, queryPlanner] = composeAndCreatePlannerWithOptions(
+      subgraphs,
+      config,
+    );
+    const operation = operationFromDocument(
+      api,
+      gql`
+        {
+          t {
+            v1
+            v2
+            v3
+            v4
+          }
+        }
+      `,
+    );
+
+    const plan = queryPlanner.buildQueryPlan(operation);
+    expect(plan).toMatchInlineSnapshot(`
+      QueryPlan {
+        Fetch(service: "Subgraph1") {
+          {
+            t {
+              v1
+              v2
+              v3
+              v4
+            }
+          }
+        },
+      }
+    `);
+
+    const stats = queryPlanner.lastGeneratedPlanStatistics();
+    expect(stats?.evaluatedPlanCount).toBe(16);
+  });
+
+  test('maxQueryPlanningTime exceeded (10 microseconds)', () => {
+    const config = {
+      debug: { maxEvaluatedPlans: undefined },
+      maxQueryPlanningTime: 0.01,
+    };
+    const [api, queryPlanner] = composeAndCreatePlannerWithOptions(
+      subgraphs,
+      config,
+    );
+    const operation = operationFromDocument(
+      api,
+      gql`
+        {
+          t {
+            v1
+            v2
+            v3
+            v4
+          }
+        }
+      `,
+    );
+
+    expect(() => queryPlanner.buildQueryPlan(operation)).toThrow(
+      'Query plan took too long to calculate',
+    );
+  });
+});