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.

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',
+    );
+  });
+});

query-planner-js/src/buildPlan.ts

@@ -105,6 +105,7 @@
 import { enforceQueryPlannerConfigDefaults, QueryPlannerConfig, validateQueryPlannerConfig } from "./config";
 import { generateAllPlansAndFindBest } from "./generateAllPlans";
 import { QueryPlan, ResponsePath, SequenceNode, PlanNode, ParallelNode, FetchNode, SubscriptionNode, trimSelectionNodes } from "./QueryPlan";
+import { performance } from 'perf_hooks';
 
 
 const debug = newDebugLogger('plan');
@@ -394,6 +395,9 @@
     readonly costFunction: CostFunction,
     initialContext: PathContext,
     typeConditionedFetching: boolean,
+
+    // time (calculated relative from performance.now) after which query plan calculation may be aborted
+    readonly deadline: number,
     excludedDestinations: ExcludedDestinations = [],
     excludedConditions: ExcludedConditions = [],
   ) {
@@ -445,6 +449,9 @@
   }
 
   private handleOpenBranch(selection: Selection, options: SimultaneousPathsWithLazyIndirectPaths<RV>[]) {
+    if (performance.now() > this.deadline) {
+      throw new Error('Query plan took too long to calculate');
+    }
     const operation = selection.element;
     debug.group(() => `Handling open branch: ${operation}`);
     let newOptions: SimultaneousPathsWithLazyIndirectPaths<RV>[] = [];
@@ -771,6 +778,7 @@
       this.costFunction,
       context,
       this.typeConditionedFetching,
+      this.deadline,
       excludedDestinations,
       addConditionExclusion(excludedConditions, edge.conditions),
     ).findBestPlan();
@@ -3531,10 +3539,10 @@
   let main: PlanNode | undefined = undefined;
   let primarySelection: MutableSelectionSet | undefined = undefined;
   let deferred: DeferredNode[] = [];
-
-  const { operation, processor } = parameters;
+  const { operation, processor, config } = parameters;
+  const deadline = performance.now() + config.maxQueryPlanningTime;
   if (operation.rootKind === 'mutation') {
-    const dependencyGraphs = computeRootSerialDependencyGraph(parameters, hasDefers);
+    const dependencyGraphs = computeRootSerialDependencyGraph(parameters, hasDefers, deadline);
     for (const dependencyGraph of dependencyGraphs) {
       const { main: localMain, deferred: localDeferred } = dependencyGraph.process(processor, operation.rootKind);
       // Note that `reduceSequence` "flatten" sequence if needs be.
@@ -3554,6 +3562,7 @@
       parameters,
       0,
       hasDefers,
+      deadline,
     );
     ({ main, deferred } = dependencyGraph.process(processor, operation.rootKind));
     primarySelection = dependencyGraph.deferTracking.primarySelection;
@@ -3675,12 +3684,14 @@
   parameters: PlanningParameters<RootVertex>,
   startFetchIdGen: number,
   hasDefer: boolean,
+  deadline: number,
 ): FetchDependencyGraph {
   return computeRootParallelBestPlan(
     parameters,
     parameters.operation.selectionSet,
     startFetchIdGen,
     hasDefer,
+    deadline,
   )[0];
 }
 
@@ -3689,6 +3700,7 @@
   selection: SelectionSet,
   startFetchIdGen: number,
   hasDefers: boolean,
+  deadline: number,
 ): [FetchDependencyGraph, OpPathTree<RootVertex>, number] {
   const planningTraversal = new QueryPlanningTraversal(
     parameters,
@@ -3699,6 +3711,7 @@
     defaultCostFunction,
     emptyContext,
     parameters.config.typeConditionedFetching,
+    deadline,
   );
   const plan = planningTraversal.findBestPlan();
   // Getting no plan means the query is essentially unsatisfiable (it's a valid query, but we can prove it will never return a result),
@@ -3726,17 +3739,18 @@
 function computeRootSerialDependencyGraph(
   parameters: PlanningParameters<RootVertex>,
   hasDefers: boolean,
+  deadline: number,
 ): FetchDependencyGraph[] {
   const { supergraphSchema, federatedQueryGraph, operation, root } = parameters;
   const rootType = hasDefers ? supergraphSchema.schemaDefinition.rootType(root.rootKind) : undefined;
   // We have to serially compute a plan for each top-level selection.
   const splittedRoots = splitTopLevelFields(operation.selectionSet);
   const graphs: FetchDependencyGraph[] = [];
   let startingFetchId = 0;
-  let [prevDepGraph, prevPaths] = computeRootParallelBestPlan(parameters, splittedRoots[0], startingFetchId, hasDefers);
+  let [prevDepGraph, prevPaths] = computeRootParallelBestPlan(parameters, splittedRoots[0], startingFetchId, hasDefers, deadline);
   let prevSubgraph = onlyRootSubgraph(prevDepGraph);
   for (let i = 1; i < splittedRoots.length; i++) {
-    const [newDepGraph, newPaths] = computeRootParallelBestPlan(parameters, splittedRoots[i], prevDepGraph.nextFetchId(), hasDefers);
+    const [newDepGraph, newPaths] = computeRootParallelBestPlan(parameters, splittedRoots[i], prevDepGraph.nextFetchId(), hasDefers, deadline);
     const newSubgraph = onlyRootSubgraph(newDepGraph);
     if (prevSubgraph === newSubgraph) {
       // The new operation (think 'mutation' operation) is on the same subgraph than the previous one, so we can concat them in a single fetch

query-planner-js/src/config.ts

@@ -118,6 +118,11 @@ export type QueryPlannerConfig = {
    * If you aren't aware of this flag, you probably don't need it.
    */
     typeConditionedFetching?: boolean,
+
+    /**
+     * The maximum time to wait for query planning to complete before giving up and throwing an exception. Defaults to 2 minutes
+     */
+    maxQueryPlanningTime?: number,
 }
 
 export function enforceQueryPlannerConfigDefaults(
@@ -128,6 +133,7 @@ export function enforceQueryPlannerConfigDefaults(
     reuseQueryFragments: true,
     generateQueryFragments: false,
     cache: new InMemoryLRUCache<QueryPlan>({maxSize: Math.pow(2, 20) * 50 }),
+    maxQueryPlanningTime: 2 * 60 * 1000, // two minutes
     ...config,
     incrementalDelivery: {
       enableDefer: false,