apollographql · garypen · Mar 26, 2024 · Mar 15, 2024 · Mar 18, 2024 · Mar 19, 2024
diff --git a/.changesets/fix_geal_request_completion.md b/.changesets/fix_geal_request_completion.md
@@ -0,0 +1,12 @@
+### Execute the entire request pipeline if the client closed the connection ([Issue #4569](https://github.com/apollographql/router/issues/4569)), [Issue #4576](https://github.com/apollographql/router/issues/4576)), ([Issue #4589](https://github.com/apollographql/router/issues/4589)), ([Issue #4590](https://github.com/apollographql/router/issues/4590)), ([Issue #4611](https://github.com/apollographql/router/issues/4611))
+
+The router is now making sure that the entire request handling pipeline is executed when the client closes the connection early, to let telemetry and any rhai scrit or coprocessor perform their tasks before canceling. Before that, when a client canceled a request, the entire execution was dropped and parts of the router, like telemetry, could not run properly. It now executes up to the first response event (in the case of subscription or `@defer` usage), adds a 499 status code to the response and skips the remaining subgraph requests.
+
+This change will report more requests to Studio and the configured telemetry, which will appear like a sudden increase in errors, because those failing requests were not reported before. To keep the previous behavior of immediately dropping execution for canceled requests, it is possible with the following option:
+
+```yaml
+supergraph:
+  early_cancel: true
+```
+
+By [@Geal](https://github.com/Geal) in https://github.com/apollographql/router/pull/4770
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -258,7 +258,7 @@ commands:
             - run:
                 name: Install CMake
                 command: |
-                  choco install cmake -y
+                  choco install cmake.install -y
                   echo 'export PATH="/c/Program Files/CMake/bin:$PATH"' >> "$BASH_ENV"
                   exit $LASTEXITCODE
       - when:

diff --git a/RELEASE_CHECKLIST.md b/RELEASE_CHECKLIST.md
diff --git a/apollo-router/src/axum_factory/axum_http_server_factory.rs b/apollo-router/src/axum_factory/axum_http_server_factory.rs
@@ -37,6 +37,8 @@ use tower::ServiceBuilder;
 use tower::ServiceExt;
 use tower_http::decompression::DecompressionBody;
 use tower_http::trace::TraceLayer;
+use tracing::instrument::WithSubscriber;
+use tracing::Instrument;
 
 use super::listeners::ensure_endpoints_consistency;
 use super::listeners::ensure_listenaddrs_consistency;
@@ -64,6 +66,7 @@ use crate::services::router;
 use crate::uplink::license_enforcement::LicenseState;
 use crate::uplink::license_enforcement::APOLLO_ROUTER_LICENSE_EXPIRED;
 use crate::uplink::license_enforcement::LICENSE_EXPIRED_SHORT_MESSAGE;
+use crate::Context;
 
 static ACTIVE_SESSION_COUNT: AtomicU64 = AtomicU64::new(0);
 
@@ -549,16 +552,17 @@ pub(super) fn main_router<RF>(
 where
     RF: RouterFactory,
 {
+    let early_cancel = configuration.supergraph.early_cancel;
     let mut router = Router::new().route(
         &configuration.supergraph.sanitized_path(),
         get({
             move |Extension(service): Extension<RF>, request: Request<DecompressionBody<Body>>| {
-                handle_graphql(service.create().boxed(), request)
+                handle_graphql(service.create().boxed(), early_cancel, request)
             }
         })
         .post({
             move |Extension(service): Extension<RF>, request: Request<DecompressionBody<Body>>| {
-                handle_graphql(service.create().boxed(), request)
+                handle_graphql(service.create().boxed(), early_cancel, request)
             }
         }),
     );
@@ -569,13 +573,13 @@ where
             get({
                 move |Extension(service): Extension<RF>,
                       request: Request<DecompressionBody<Body>>| {
-                    handle_graphql(service.create().boxed(), request)
+                    handle_graphql(service.create().boxed(), early_cancel, request)
                 }
             })
             .post({
                 move |Extension(service): Extension<RF>,
                       request: Request<DecompressionBody<Body>>| {
-                    handle_graphql(service.create().boxed(), request)
+                    handle_graphql(service.create().boxed(), early_cancel, request)
                 }
             }),
         );
@@ -586,6 +590,7 @@ where
 
 async fn handle_graphql(
     service: router::BoxService,
+    early_cancel: bool,
     http_request: Request<DecompressionBody<Body>>,
 ) -> impl IntoResponse {
     let _guard = SessionCountGuard::start();
@@ -602,7 +607,33 @@ async fn handle_graphql(
         .get(ACCEPT_ENCODING)
         .cloned();
 
-    let res = service.oneshot(request).await;
+    let res = if early_cancel {
+        service.oneshot(request).await
+    } else {
+        // to make sure we can record request handling when the client closes the connection prematurely,
+        // we execute the request in a separate task that will run until we get the first response, which
+        // means it went through the entire pipeline at least once (not looking at deferred responses or
+        // subscription events). This is a bit wasteful, so to avoid unneeded subgraph calls, we insert in
+        // the context a flag to indicate that the request is canceled and subgraph calls should not be made
+        let mut cancel_handler = CancelHandler::new(&context);
+        let task = service
+            .oneshot(request)
+            .with_current_subscriber()
+            .in_current_span();
+        let res = match tokio::task::spawn(task).await {
+            Ok(res) => res,
+            Err(_e) => {
+                return (
+                    StatusCode::INTERNAL_SERVER_ERROR,
+                    "router service call failed",
+                )
+                    .into_response();
+            }
+        };
+        cancel_handler.on_response();
+        res
+    };
+
     let dur = context.busy_time();
     let processing_seconds = dur.as_secs_f64();
 
@@ -654,12 +685,48 @@ async fn handle_graphql(
     }
 }
 
+struct CancelHandler<'a> {
+    context: &'a Context,
+    got_first_response: bool,
+    span: tracing::Span,
+}
+
+impl<'a> CancelHandler<'a> {
+    fn new(context: &'a Context) -> Self {
+        CancelHandler {
+            context,
+            got_first_response: false,
+            span: tracing::Span::current(),
+        }
+    }
+
+    fn on_response(&mut self) {
+        self.got_first_response = true;
+    }
+}
+
+impl<'a> Drop for CancelHandler<'a> {
+    fn drop(&mut self) {
+        if !self.got_first_response {
+            self.span
+                .in_scope(|| tracing::error!("broken pipe: the client closed the connection"));
+            self.context.extensions().lock().insert(CanceledRequest);
+        }
+    }
+}
+
+pub(crate) struct CanceledRequest;
+
 #[cfg(test)]
 mod tests {
     use std::str::FromStr;
 
-    use super::*;
+    use http::header::ACCEPT;
+    use http::header::CONTENT_TYPE;
+    use tower::Service;
 
+    use super::*;
+    use crate::assert_snapshot_subscriber;
     #[test]
     fn test_span_mode_default() {
         let config =
@@ -687,4 +754,35 @@ mod tests {
         let mode = span_mode(&config);
         assert_eq!(mode, SpanMode::Deprecated);
     }
+
+    #[tokio::test]
+    async fn request_cancel() {
+        let mut http_router = crate::TestHarness::builder()
+            .schema(include_str!("../testdata/supergraph.graphql"))
+            .build_http_service()
+            .await
+            .unwrap();
+
+        async {
+            let _res = tokio::time::timeout(
+                std::time::Duration::from_micros(100),
+                http_router.call(
+                    http::Request::builder()
+                        .method("POST")
+                        .uri("/")
+                        .header(ACCEPT, "application/json")
+                        .header(CONTENT_TYPE, "application/json")
+                        .body(hyper::Body::from(r#"{"query":"query { me { name }}"}"#))
+                        .unwrap(),
+                ),
+            )
+            .await;
+
+            tokio::time::sleep(std::time::Duration::from_millis(1000)).await;
+        }
+        .with_subscriber(assert_snapshot_subscriber!(
+            tracing_core::LevelFilter::ERROR
+        ))
+        .await
+    }
 }
diff --git a/apollo-router/src/axum_factory/mod.rs b/apollo-router/src/axum_factory/mod.rs
@@ -12,6 +12,7 @@ use std::sync::OnceLock;
 use axum::Router;
 pub(crate) use axum_http_server_factory::span_mode;
 pub(crate) use axum_http_server_factory::AxumHttpServerFactory;
+pub(crate) use axum_http_server_factory::CanceledRequest;
 pub(crate) use listeners::ListenAddrAndRouter;
 
 static ENDPOINT_CALLBACK: OnceLock<Arc<dyn Fn(Router) -> Router + Send + Sync>> = OnceLock::new();

diff --git a/...ts/apollo_router__axum_factory__axum_http_server_factory__tests__request_cancel@logs.snap b/...ts/apollo_router__axum_factory__axum_http_server_factory__tests__request_cancel@logs.snap
@@ -0,0 +1,7 @@
+---
+source: apollo-router/src/axum_factory/axum_http_server_factory.rs
+expression: yaml
+---
+- fields: {}
+  level: ERROR
+  message: "broken pipe: the client closed the connection"
diff --git a/apollo-router/src/configuration/mod.rs b/apollo-router/src/configuration/mod.rs
@@ -616,6 +616,12 @@ pub(crate) struct Supergraph {
 
     /// Query planning options
     pub(crate) query_planning: QueryPlanning,
+
+    /// abort request handling when the client drops the connection.
+    /// Default: false.
+    /// When set to true, some parts of the request pipeline like telemetry will not work properly,
+    /// but request handling will stop immediately when the client connection is closed.
+    pub(crate) early_cancel: bool,
 }
 
 fn default_defer_support() -> bool {
@@ -632,6 +638,7 @@ impl Supergraph {
         defer_support: Option<bool>,
         query_planning: Option<QueryPlanning>,
         reuse_query_fragments: Option<bool>,
+        early_cancel: Option<bool>,
     ) -> Self {
         Self {
             listen: listen.unwrap_or_else(default_graphql_listen),
@@ -640,6 +647,7 @@ impl Supergraph {
             defer_support: defer_support.unwrap_or_else(default_defer_support),
             query_planning: query_planning.unwrap_or_default(),
             reuse_query_fragments,
+            early_cancel: early_cancel.unwrap_or_default(),
         }
     }
 }
@@ -655,6 +663,7 @@ impl Supergraph {
         defer_support: Option<bool>,
         query_planning: Option<QueryPlanning>,
         reuse_query_fragments: Option<bool>,
+        early_cancel: Option<bool>,
     ) -> Self {
         Self {
             listen: listen.unwrap_or_else(test_listen),
@@ -663,6 +672,7 @@ impl Supergraph {
             defer_support: defer_support.unwrap_or_else(default_defer_support),
             query_planning: query_planning.unwrap_or_default(),
             reuse_query_fragments,
+            early_cancel: early_cancel.unwrap_or_default(),
         }
     }
 }

diff --git a/...r/src/configuration/snapshots/apollo_router__configuration__tests__schema_generation.snap b/...r/src/configuration/snapshots/apollo_router__configuration__tests__schema_generation.snap
@@ -2686,7 +2686,8 @@ expression: "&schema"
           "warmed_up_queries": null,
           "experimental_plans_limit": null,
           "experimental_paths_limit": null
-        }
+        },
+        "early_cancel": false
       },
       "type": "object",
       "properties": {
@@ -2695,6 +2696,11 @@ expression: "&schema"
           "default": true,
           "type": "boolean"
         },
+        "early_cancel": {
+          "description": "abort request handling when the client drops the connection. Default: false. When set to true, some parts of the request pipeline like telemetry will not work properly, but request handling will stop immediately when the client connection is closed.",
+          "default": false,
+          "type": "boolean"
+        },
         "experimental_reuse_query_fragments": {
           "description": "Enable reuse of query fragments Default: depends on the federation version",
           "default": null,

diff --git a/apollo-router/src/plugins/coprocessor/execution.rs b/apollo-router/src/plugins/coprocessor/execution.rs
@@ -12,6 +12,7 @@ use tower_service::Service;
 
 use super::externalize_header_map;
 use super::*;
+use crate::batching::BatchQuery;
 use crate::graphql;
 use crate::layers::async_checkpoint::OneShotAsyncCheckpointLayer;
 use crate::layers::ServiceBuilderExt;
@@ -289,6 +290,18 @@ where
 
             execution_response
         };
+
+        // Handle cancelled batch queries
+        // FIXME: This should be way higher up the call chain so that custom plugins / rhai / etc. can
+        // automatically work with batched queries and cancellations.
+        let batch_query_opt = res.context.extensions().lock().remove::<BatchQuery>();
+        if let Some(mut batch_query) = batch_query_opt {
+            // TODO: How do we reliably get the reason for the coprocessor cancellation here?
+            batch_query
+                .signal_cancelled("coprocessor cancelled request at execution layer".to_string())
+                .await;
+        }
+
         return Ok(ControlFlow::Break(res));
     }
 

diff --git a/apollo-router/src/plugins/coprocessor/mod.rs b/apollo-router/src/plugins/coprocessor/mod.rs
@@ -30,6 +30,7 @@ use tower::Service;
 use tower::ServiceBuilder;
 use tower::ServiceExt;
 
+use crate::batching::BatchQuery;
 use crate::error::Error;
 use crate::layers::async_checkpoint::OneShotAsyncCheckpointLayer;
 use crate::layers::ServiceBuilderExt;
@@ -685,6 +686,17 @@ where
             }
         }
 
+        // Handle cancelled batch queries
+        // FIXME: This should be way higher up the call chain so that custom plugins / rhai / etc. can
+        // automatically work with batched queries and cancellations.
+        let batch_query_opt = res.context.extensions().lock().remove::<BatchQuery>();
+        if let Some(mut batch_query) = batch_query_opt {
+            // TODO: How do we reliably get the reason for the coprocessor cancellation here?
+            batch_query
+                .signal_cancelled("coprocessor cancelled request at router layer".to_string())
+                .await;
+        }
+
         return Ok(ControlFlow::Break(res));
     }
 
@@ -1014,6 +1026,18 @@ where
 
             subgraph_response
         };
+
+        // Handle cancelled batch queries
+        // FIXME: This should be way higher up the call chain so that custom plugins / rhai / etc. can
+        // automatically work with batched queries and cancellations.
+        let batch_query_opt = res.context.extensions().lock().remove::<BatchQuery>();
+        if let Some(mut batch_query) = batch_query_opt {
+            // TODO: How do we reliably get the reason for the coprocessor cancellation here?
+            batch_query
+                .signal_cancelled("coprocessor cancelled request at subgraph layer".to_string())
+                .await;
+        }
+
         return Ok(ControlFlow::Break(res));
     }
 

diff --git a/apollo-router/src/plugins/coprocessor/supergraph.rs b/apollo-router/src/plugins/coprocessor/supergraph.rs
@@ -279,6 +279,18 @@ where
 
             supergraph_response
         };
+
+        // Handle cancelled batch queries
+        // FIXME: This should be way higher up the call chain so that custom plugins / rhai / etc. can
+        // automatically work with batched queries and cancellations.
+        let batch_query_opt = res.context.extensions().lock().remove::<BatchQuery>();
+        if let Some(mut batch_query) = batch_query_opt {
+            // TODO: How do we reliably get the reason for the coprocessor cancellation here?
+            batch_query
+                .signal_cancelled("coprocessor cancelled request at supergraph layer".to_string())
+                .await;
+        }
+
         return Ok(ControlFlow::Break(res));
     }