Ensure `build_query_plan()` cancels when the request cancels #6840

SimonSapin · 2025-02-20T19:33:43Z

When a request is cancelled, resources (compute time/memory) allocated to that request should be released. However, query planning is an exception, and currently keeps running after the request is cancelled. This PR adds cooperative cancellation to build_query_plan(), to ensure that when the request is cancelled, query planning errors shortly after.

svc-apollo-docs · 2025-02-20T19:34:28Z

✅ Docs preview ready

The preview is ready to be viewed. View the preview

File Changes

0 new, 2 changed, 0 removed

* graphos/reference/upgrade/from-router-v1.mdx
* graphos/routing/(latest)/customization/rhai.mdx

Build ID: 6f6b4476b6c27277dc04f455

URL: https://www.apollographql.com/docs/deploy-preview/6f6b4476b6c27277dc04f455

SimonSapin · 2025-02-20T19:34:56Z

apollo-federation/src/query_plan/query_planning_traversal.rs

@@ -339,6 +354,7 @@ impl<'a: 'b, 'b> QueryPlanningTraversal<'a, 'b> {
    )]
    fn find_best_plan_inner(&mut self) -> Result<Option<&BestQueryPlanInfo>, FederationError> {
        while !self.open_branches.is_empty() {
+            self.parameters.check_cancellation()?;


Checks like this would need to be added to all loops where the planner can spend significant amount of time.

I think handle_open_branch is the one. Other functions won't take too long. Even if advance_with_operation_element created a million options, it returns reasonably quickly. It's the handle_open_branch that actually visits them. One decision point is how often we want to check (like every 100 or 1000 iterations?). Also, since handle_open_branch is called recursively, it may check either only at the top-level or always.

dariuszkuc · 2025-02-21T06:06:48Z

apollo-router/src/compute_job.rs

+    ///
+    /// In this case, a long-running job should try to cancel itself
+    /// to avoid needless resource consumption.
+    pub(crate) fn check_for_cooperative_cancellation(&self) -> ControlFlow<()> {


Pretty cool! While tying it to the client request (disconnect/time out) is a great start - if we were to expose this feature as a user specified config value it sounds like it should be as simple as wrapping those QP tasks with a timeout?

I don't think we should be exposing this specifically as a feature. Cancellation at the query planning level should work in the same way as backpressure in other parts of the router. That is, query planning should adhere to the router's existing traffic_shaping.router.timeout configuration.

I agree that for now this feature should be tied to the request lifecycle. That being said, sometime in the future, I think we should revisit whether it would make sense to further expand on this functionality as there are some potential additional usecases (eg folks might want to still plan the query even though their request timed out with the assumption that be subsequent request would hit the cached plan, similarly we might want to have some limits when doing qp cache prewarming etc).

sachindshinde · 2025-02-28T20:22:26Z

From what I understand, the cancellation check is relatively cheap, so as long as it's not in a tight loop it should be fine to run it. I've added cancellation checks to other potentially problematic loops/iterations in query planner, added a PR description and changelog, and marked this PR as ready for review.

duckki

My concern is calling check_for_cooperative_cancellation too often, since It tries access memory that is not in the cpu cache. But, I may be wrong because CPU architecture might have changed since the last time I studied it :)

duckki · 2025-02-28T20:49:51Z

apollo-federation/src/operation/tests/mod.rs

+        fragments,
+        schema,
+        &Default::default(),
+        &|| Ok(()),


I wish we declare a named thing in place of || Ok(()), so it's more readable.

I agree. A named function can be used directly like &never_cancel with:

fn never_cancel() -> Result<(), SingleFederationError> { Ok(()) }

Makes sense, I've shifted test files to using &never_cancel in 083ebc1

duckki · 2025-02-28T22:21:29Z

Ok, I've done some reading and I'm convinced this won't slowdown our QP. In the future, when we add a timeout in the cancelation check, that may add up some cpu cycles. But, that still seems reasonable.

duckki · 2025-02-28T22:28:09Z

apollo-federation/src/query_plan/query_planner.rs

+            &|| {
+                QueryPlanningParameters::check_cancellation_with(
+                    &options.check_for_cooperative_cancellation,
+                )


I was wondering why we needed the check_cancellation_with method. Then, I realized normalize_operation has to be called before QueryPlanningParameters is constructed. That's unfortunate.

duckki

Approved. Only a minor suggestion to refactor some obscure closures in tests.

SimonSapin · 2025-03-03T09:31:56Z

apollo-federation/src/operation/tests/mod.rs

+        fragments,
+        schema,
+        &Default::default(),
+        &|| Ok(()),


I agree. A named function can be used directly like &never_cancel with:

fn never_cancel() -> Result<(), SingleFederationError> { Ok(()) }

SimonSapin · 2025-03-03T09:35:13Z

apollo-federation/src/query_plan/fetch_dependency_graph.rs

@@ -3488,6 +3488,7 @@ pub(crate) fn compute_nodes_for_tree(
    initial_node_path: FetchDependencyGraphNodePath,
    initial_defer_context: DeferContext,
    initial_conditions: &OpGraphPathContext,
+    check_cancellation: &impl Fn() -> Result<(), SingleFederationError>,


Using impl Fn here makes the whole method generic and monomorphized separately for every distinct callback that can be passed. This has some compilation time and code size cost.

Assuming that check_cancellation isn’t called in very hot loop, would it be better to use &dyn Fn()? (Here and in other places with a similar argument)

It shouldn't be in any hot loops, I've switched it to using &dyn Fn() in 083ebc1

…ck cancellation)

When a request is cancelled, resources (compute time/memory) allocated to that request should be released. However, query planning is an exception, and currently keeps running after the request is cancelled. This PR adds cooperative cancellation to `build_query_plan()`, to ensure that when the request is cancelled, query planning errors shortly after. Co-authored-by: Sachin D. Shinde <sachin@apollographql.com> (cherry picked from commit 2c554fc) # Conflicts: # apollo-federation/src/query_graph/path_tree.rs # apollo-router/src/compute_job.rs # apollo-router/src/introspection.rs # apollo-router/src/query_planner/query_planner_service.rs # apollo-router/src/services/layers/query_analysis.rs

This comment has been minimized.

Sign in to view

SimonSapin commented Feb 20, 2025

View reviewed changes

dariuszkuc reviewed Feb 21, 2025

View reviewed changes

sachindshinde changed the title ~~Initial scaffolding for cooperative planner cancellation~~ Ensure build_query_plan() cancels when the request cancels Feb 28, 2025

sachindshinde force-pushed the simon/coop branch 2 times, most recently from d051bb8 to e3bb09b Compare February 28, 2025 18:53

sachindshinde marked this pull request as ready for review February 28, 2025 18:53

sachindshinde requested review from a team, sachindshinde, goto-bus-stop and TylerBloom as code owners February 28, 2025 18:53

sachindshinde force-pushed the simon/coop branch from e3bb09b to ba9877f Compare February 28, 2025 19:16

sachindshinde requested a review from a team as a code owner February 28, 2025 20:21

sachindshinde added the backport-1.x Backport this PR to 1.x label Feb 28, 2025

duckki reviewed Feb 28, 2025

View reviewed changes

duckki approved these changes Feb 28, 2025

View reviewed changes

apollographql deleted a comment from duckki Mar 3, 2025

SimonSapin commented Mar 3, 2025

View reviewed changes

sachindshinde force-pushed the simon/coop branch from 67fad15 to 083ebc1 Compare March 3, 2025 21:13

SimonSapin and others added 4 commits March 3, 2025 13:17

Initial scaffolding for cooperative planner cancellation

2090fc8

Add cancellation checks to potentially heavy loops in query planning

adc76b0

Add a changeset

a36b08e

Address feedback (switch to dyn, and use named function for no-op che…

4072854

…ck cancellation)

sachindshinde force-pushed the simon/coop branch from 083ebc1 to 4072854 Compare March 3, 2025 21:17

sachindshinde merged commit 2c554fc into dev Mar 3, 2025
15 checks passed

sachindshinde deleted the simon/coop branch March 3, 2025 22:57

SimonSapin mentioned this pull request Mar 3, 2025

Ensure build_query_plan() cancels when the request cancels (backport #6840) #6904

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Ensure `build_query_plan()` cancels when the request cancels #6840

Ensure `build_query_plan()` cancels when the request cancels #6840

SimonSapin commented Feb 20, 2025 •

edited by sachindshinde

Loading

This comment has been minimized.

svc-apollo-docs commented Feb 20, 2025 •

edited

Loading

SimonSapin Feb 20, 2025

duckki Feb 22, 2025 •

edited

Loading

dariuszkuc Feb 21, 2025

lrlna Feb 24, 2025

dariuszkuc Feb 25, 2025

sachindshinde commented Feb 28, 2025 •

edited

Loading

duckki left a comment

duckki Feb 28, 2025

SimonSapin Mar 3, 2025

sachindshinde Mar 3, 2025

duckki commented Feb 28, 2025 •

edited

Loading

duckki Feb 28, 2025 •

edited

Loading

duckki left a comment

SimonSapin Mar 3, 2025

SimonSapin Mar 3, 2025

sachindshinde Mar 3, 2025

Ensure build_query_plan() cancels when the request cancels #6840

Ensure build_query_plan() cancels when the request cancels #6840

Conversation

SimonSapin commented Feb 20, 2025 • edited by sachindshinde Loading

This comment has been minimized.

svc-apollo-docs commented Feb 20, 2025 • edited Loading

✅ Docs preview ready

Choose a reason for hiding this comment

duckki Feb 22, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sachindshinde commented Feb 28, 2025 • edited Loading

duckki left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

duckki commented Feb 28, 2025 • edited Loading

duckki Feb 28, 2025 • edited Loading

Choose a reason for hiding this comment

duckki left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Ensure `build_query_plan()` cancels when the request cancels #6840

Ensure `build_query_plan()` cancels when the request cancels #6840

SimonSapin commented Feb 20, 2025 •

edited by sachindshinde

Loading

svc-apollo-docs commented Feb 20, 2025 •

edited

Loading

duckki Feb 22, 2025 •

edited

Loading

sachindshinde commented Feb 28, 2025 •

edited

Loading

duckki commented Feb 28, 2025 •

edited

Loading

duckki Feb 28, 2025 •

edited

Loading