Move computation-heavy tasks to a dedicated thread pool

.changesets/fix_simon_compute_task.md

@@ -0,0 +1,13 @@
+### Move computation-heavy tasks to a dedicated thread pool ([PR #6122](https://github.com/apollographql/router/pull/6122))
+
+These components can take non-trivial amounts of CPU time:
+* GraphQL parsing
+* GraphQL validation
+* Query planning
+* Schema introspection
+
+In order to avoid blocking threads that execute asynchronous code, they are now run (in their respective Rust implementations) in a new pool of as many threads as CPU cores are available. Previously we used Tokio’s [`spawn_blocking`] for this purpose, but it is appears to be intended for blocking I/O and uses up to 512 threads so it isn’t a great fit for computation tasks.
+
+[`spawn_blocking`]: https://docs.rs/tokio/latest/tokio/task/fn.spawn_blocking.html
+
+By [@SimonSapin](https://github.com/SimonSapin) in https://github.com/apollographql/router/pull/6122

Cargo.lock

@@ -361,6 +361,7 @@ dependencies = [
  "tempfile",
  "test-log",
  "thiserror",
+ "threadpool",
  "tikv-jemallocator",
  "time",
  "tokio",

apollo-router/Cargo.toml

@@ -209,6 +209,7 @@ static_assertions = "1.1.0"
 strum_macros = "0.26.0"
 sys-info = "0.9.1"
 thiserror = "1.0.61"
+threadpool = "1.8.1"
 tokio.workspace = true
 tokio-stream = { version = "0.1.15", features = ["sync", "net"] }
 tokio-util = { version = "0.7.11", features = ["net", "codec", "time"] }

apollo-router/src/compute_task.rs

@@ -0,0 +1,48 @@
+use std::sync::OnceLock;
+
+use tokio::sync::oneshot;
+
+static POOL: OnceLock<threadpool::ThreadPool> = OnceLock::new();
+
+/// Drop-in replacement for Tokio’s [`spawn_blocking`]
+/// but intended for for tasks that keep a CPU core active: it uses a thread pool
+/// limited to the number of available CPU cores.
+///
+/// `spawn_blocking` on the other hand appears to be intended for tasks that pause a thread
+/// in a blocking syscall: `max_blocking_threads` is configurable but defaults to 512,
+/// which is too high for tasks where the CPU is actively running.
+/// Configuring it to $NUM_CPUS feels risky in case a library we use relies on `spawn_blocking`
+/// in the originally intended way.
+///
+/// [`spawn_blocking`]: https://docs.rs/tokio/latest/tokio/task/fn.spawn_blocking.html
+///
+/// ## Example
+///
+/// Doing this in an async function can work but isn't great for Tokio performance:
+///
+/// ```no_run
+/// let output = some_heavy_computation();
+/// ```
+///
+/// It can be replaced with:
+///
+/// ```no_run
+/// use crate::compute_task;
+///
+/// let output = compute_task::execute(|| some_heavy_computation()).await;
+/// ```
+pub(crate) fn execute<T, F>(f: F) -> oneshot::Receiver<T>
+where
+    F: FnOnce() -> T + Send + 'static,
+    T: Send + 'static,
+{
+    // https://docs.rs/threadpool/latest/threadpool/struct.Builder.html#method.num_threads
+    // > defaults the number of threads to the number of CPUs
+    let pool = POOL.get_or_init(|| threadpool::Builder::new().build());
+
+    let (tx, rx) = oneshot::channel();
+    pool.execute(move || {
+        let _ = tx.send(f());
+    });
+    rx
+}

apollo-router/src/lib.rs

@@ -56,6 +56,7 @@ mod apollo_studio_interop;
 pub(crate) mod axum_factory;
 mod batching;
 mod cache;
+mod compute_task;
 mod configuration;
 mod context;
 mod error;

apollo-router/src/query_planner/bridge_query_planner.rs

@@ -31,6 +31,7 @@ use super::PlanNode;
 use super::QueryKey;
 use crate::apollo_studio_interop::generate_usage_reporting;
 use crate::cache::storage::CacheStorage;
+use crate::compute_task;
 use crate::configuration::IntrospectionMode as IntrospectionConfig;
 use crate::configuration::QueryPlannerMode;
 use crate::error::PlanErrors;
@@ -292,7 +293,7 @@ impl PlannerMode {
             PlannerMode::Rust(rust_planner) => {
                 let doc = doc.clone();
                 let rust_planner = rust_planner.clone();
-                let (plan, mut root_node) = tokio::task::spawn_blocking(move || {
+                let (plan, mut root_node) = compute_task::execute(move || {
                     let start = Instant::now();
 
                     let query_plan_options = QueryPlanOptions {
@@ -543,11 +544,9 @@ impl BridgeQueryPlanner {
             IntrospectionMode::Rust => {
                 let schema = self.schema.clone();
                 let response = Box::new(
-                    tokio::task::spawn_blocking(move || {
-                        Self::rust_introspection(&schema, &key, &doc)
-                    })
-                    .await
-                    .expect("Introspection panicked")?,
+                    compute_task::execute(move || Self::rust_introspection(&schema, &key, &doc))
+                        .await
+                        .expect("Introspection panicked")?,
                 );
                 return Ok(QueryPlannerContent::Response { response });
             }
@@ -581,7 +580,7 @@ impl BridgeQueryPlanner {
                     .map_err(QueryPlannerError::Introspection);
                 let schema = self.schema.clone();
                 let js_result_clone = js_result.clone();
-                tokio::task::spawn_blocking(move || {
+                compute_task::execute(move || {
                     let rust_result = match Self::rust_introspection(&schema, &key, &doc) {
                         Ok(response) => {
                             if response.errors.is_empty() {

apollo-router/src/services/layers/query_analysis.rs

@@ -13,10 +13,10 @@ use http::StatusCode;
 use lru::LruCache;
 use router_bridge::planner::UsageReporting;
 use tokio::sync::Mutex;
-use tokio::task;
 
 use crate::apollo_studio_interop::generate_extended_references;
 use crate::apollo_studio_interop::ExtendedReferenceStats;
+use crate::compute_task;
 use crate::context::OPERATION_KIND;
 use crate::context::OPERATION_NAME;
 use crate::graphql::Error;
@@ -85,11 +85,11 @@ impl QueryAnalysisLayer {
         let schema = self.schema.clone();
         let conf = self.configuration.clone();
 
-        // Must be created *outside* of the spawn_blocking or the span is not connected to the
+        // Must be created *outside* of the compute_task or the span is not connected to the
         // parent
         let span = tracing::info_span!(QUERY_PARSING_SPAN_NAME, "otel.kind" = "INTERNAL");
 
-        task::spawn_blocking(move || {
+        compute_task::execute(move || {
             span.in_scope(|| {
                 let doc = Query::parse_document(
                     &query,