Add internal endpoint for tokio task dumps (#5594)

# Description of Changes

Introduces another `/internal` route for capturing tokio task dumps.

The endpoint is available under `/internal/task-dump` and returns a
plain-text dump of every task on the selected Tokio runtime.

### Usage Examples

#### Dump the main runtime

```bash
curl -fsS \
  'http://127.0.0.1:3000/internal/task-dump?runtime=main' \
  > task-dump-main.txt
```

`timeout_ms` defaults to 2 seconds.

#### Dump the replication runtime with a longer timeout

```bash
curl -fsS \
  'http://127.0.0.1:80/internal/task-dump?runtime=replication&timeout_ms=5000' \
  > task-dump-replication.txt
```

An unknown runtime returns `400 Bad Request` with the valid names. If
the runtime cannot complete the dump within the requested timeout, the
endpoint returns `504 Gateway Timeout`.

# API and ABI breaking changes

None

# Expected complexity level and risk

2

# Testing

- [x] Manual test on supported platform (not macos)

Author: joshua-spacetime

Cargo.lock

@@ -17,7 +17,7 @@ spacetimedb-schema.workspace = true
 
 base64.workspace = true
 http-body-util.workspace = true
-tokio = { version = "1.2", features = ["full"] }
+tokio = { workspace = true, features = ["full"] }
 lazy_static = "1.4.0"
 log = "0.4.4"
 serde = "1.0.136"
@@ -60,6 +60,9 @@ thiserror.workspace = true
 [target.'cfg(not(target_env = "msvc"))'.dependencies]
 jemalloc_pprof.workspace = true
 
+[target.'cfg(all(target_os = "linux", any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64")))'.dependencies]
+tokio = { workspace = true, features = ["taskdump"] }
+
 [dev-dependencies]
  tower = "0.5"
 jsonwebtoken.workspace = true

crates/client-api/Cargo.toml

@@ -1,5 +1,9 @@
 use crate::NodeDelegate;
 
+mod task_dump;
+
+pub use task_dump::TaskDumpRegistry;
+
 #[cfg(not(target_env = "msvc"))]
 mod jemalloc_profiling {
     use axum::body::Body;
@@ -158,5 +162,7 @@ pub fn router<S>() -> axum::Router<S>
 where
     S: NodeDelegate + Clone + 'static,
 {
-    axum::Router::new().nest("/heap", jemalloc_profiling::jemalloc_router())
+    axum::Router::new()
+        .nest("/heap", jemalloc_profiling::jemalloc_router())
+        .nest("/task-dump", task_dump::router())
 }

crates/client-api/src/routes/internal.rs

@@ -0,0 +1,174 @@
+use axum::routing::get;
+
+#[cfg(all(
+    tokio_unstable,
+    target_os = "linux",
+    any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64")
+))]
+mod imp {
+    use std::{collections::BTreeMap, fmt::Write as _, num::NonZeroU64, sync::Arc, time::Duration};
+
+    use axum::{
+        extract::{Extension, Query},
+        response::Response,
+    };
+    use http::{header::CONTENT_TYPE, StatusCode};
+    use serde::Deserialize;
+    use tokio::runtime::Handle;
+
+    const DEFAULT_TIMEOUT_MS: u64 = 2_000;
+
+    /// The Tokio runtimes which can be inspected by the internal task dump endpoint.
+    #[derive(Clone, Default)]
+    pub struct TaskDumpRegistry {
+        runtimes: Arc<BTreeMap<&'static str, Handle>>,
+    }
+
+    impl TaskDumpRegistry {
+        pub fn new(runtimes: impl IntoIterator<Item = (&'static str, Handle)>) -> Self {
+            Self {
+                runtimes: Arc::new(runtimes.into_iter().collect()),
+            }
+        }
+    }
+
+    #[derive(Deserialize)]
+    pub(super) struct TaskDumpQuery {
+        runtime: String,
+        timeout_ms: Option<NonZeroU64>,
+    }
+
+    pub(super) async fn handle_get_task_dump(
+        registry: Option<Extension<TaskDumpRegistry>>,
+        Query(query): Query<TaskDumpQuery>,
+    ) -> Result<Response, (StatusCode, String)> {
+        let Some(Extension(registry)) = registry else {
+            return Err((
+                StatusCode::NOT_IMPLEMENTED,
+                "Tokio task dumps are not configured for this server".into(),
+            ));
+        };
+
+        let Some(runtime) = registry.runtimes.get(query.runtime.as_str()).cloned() else {
+            let valid = registry.runtimes.keys().copied().collect::<Vec<_>>().join(", ");
+            return Err((
+                StatusCode::BAD_REQUEST,
+                format!("unknown Tokio runtime {:?}; valid runtimes: {valid}", query.runtime),
+            ));
+        };
+
+        let timeout_ms = query.timeout_ms.map_or(DEFAULT_TIMEOUT_MS, NonZeroU64::get);
+        let dump = tokio::time::timeout(Duration::from_millis(timeout_ms), runtime.dump())
+            .await
+            .map_err(|_| {
+                (
+                    StatusCode::GATEWAY_TIMEOUT,
+                    format!(
+                        "timed out after {timeout_ms}ms while dumping Tokio runtime {:?}",
+                        query.runtime
+                    ),
+                )
+            })?;
+
+        let runtime_name = query.runtime;
+        let body = tokio::task::spawn_blocking(move || format_dump(&runtime_name, dump))
+            .await
+            .map_err(|err| {
+                (
+                    StatusCode::INTERNAL_SERVER_ERROR,
+                    format!("task dump formatting failed: {err}"),
+                )
+            })?;
+
+        Response::builder()
+            .header(CONTENT_TYPE, "text/plain; charset=utf-8")
+            .body(body.into())
+            .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()))
+    }
+
+    fn format_dump(runtime_name: &str, dump: tokio::runtime::Dump) -> String {
+        let tasks = dump.tasks();
+        let mut output = String::new();
+        writeln!(output, "runtime: {runtime_name}").unwrap();
+        writeln!(output, "tasks: {}", tasks.iter().count()).unwrap();
+
+        for task in tasks.iter() {
+            writeln!(output, "\nTASK {}:", task.id()).unwrap();
+            writeln!(output, "{}", task.trace()).unwrap();
+        }
+
+        output
+    }
+
+    #[cfg(test)]
+    mod tests {
+        use super::*;
+        use http_body_util::BodyExt as _;
+        use tokio::runtime::Handle;
+
+        #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+        async fn dumps_registered_runtime_with_default_timeout() {
+            let registry = TaskDumpRegistry::new([("main", Handle::current())]);
+            let response = handle_get_task_dump(
+                Some(Extension(registry)),
+                Query(TaskDumpQuery {
+                    runtime: "main".into(),
+                    timeout_ms: None,
+                }),
+            )
+            .await
+            .unwrap();
+
+            assert_eq!(response.status(), StatusCode::OK);
+            let body = response.into_body().collect().await.unwrap().to_bytes();
+            let body = std::str::from_utf8(&body).unwrap();
+            assert!(body.starts_with("runtime: main\ntasks: "));
+        }
+    }
+}
+
+#[cfg(not(all(
+    tokio_unstable,
+    target_os = "linux",
+    any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64")
+)))]
+mod imp {
+    use axum::response::{IntoResponse as _, Response};
+    use http::StatusCode;
+    use tokio::runtime::Handle;
+
+    /// The Tokio runtimes which can be inspected by the internal task dump endpoint.
+    #[derive(Clone, Default)]
+    pub struct TaskDumpRegistry;
+
+    impl TaskDumpRegistry {
+        pub fn new(_: impl IntoIterator<Item = (&'static str, Handle)>) -> Self {
+            Self
+        }
+    }
+
+    pub(super) async fn handle_get_task_dump() -> Response {
+        (
+            StatusCode::NOT_IMPLEMENTED,
+            "Tokio task dumps require a Linux aarch64, x86, or x86_64 build with tokio_unstable enabled",
+        )
+            .into_response()
+    }
+
+    #[cfg(test)]
+    mod tests {
+        use super::*;
+
+        #[tokio::test]
+        async fn reports_unsupported_platform() {
+            assert_eq!(handle_get_task_dump().await.status(), StatusCode::NOT_IMPLEMENTED);
+        }
+    }
+}
+
+use imp::handle_get_task_dump;
+pub use imp::TaskDumpRegistry;
+
+pub fn router<S: Clone + Send + Sync + 'static>() -> axum::Router<S> {
+    axum::Router::new().route("/", get(handle_get_task_dump))
+}

crates/client-api/src/routes/internal/task_dump.rs

@@ -13,6 +13,7 @@ pub mod metrics;
 pub mod prometheus;
 pub mod subscribe;
 
+pub use self::internal::TaskDumpRegistry;
 use self::{database::DatabaseRoutes, identity::IdentityRoutes};
 
 /// This API call is just designed to allow clients to determine whether or not they can

crates/client-api/src/routes/mod.rs

@@ -7,7 +7,7 @@ use std::sync::Arc;
 
 use crate::{StandaloneEnv, StandaloneOptions};
 use anyhow::Context;
-use axum::extract::DefaultBodyLimit;
+use axum::extract::{DefaultBodyLimit, Extension};
 use clap::ArgAction::SetTrue;
 use clap::{Arg, ArgMatches};
 use spacetimedb::config::{parse_config, CertificateAuthority};
@@ -19,6 +19,7 @@ use spacetimedb::worker_metrics;
 use spacetimedb_client_api::routes::database::DatabaseRoutes;
 use spacetimedb_client_api::routes::router;
 use spacetimedb_client_api::routes::subscribe::WebSocketOptions;
+use spacetimedb_client_api::routes::TaskDumpRegistry;
 use spacetimedb_paths::cli::{PrivKeyPath, PubKeyPath};
 use spacetimedb_paths::server::{ConfigToml, ServerDataDir};
 use tokio::net::TcpListener;
@@ -197,19 +198,19 @@ pub async fn exec(args: &ArgMatches, db_cores: JobCores) -> anyhow::Result<()> {
     )
     .await?;
     worker_metrics::spawn_jemalloc_stats(listen_addr.clone());
-    worker_metrics::spawn_tokio_stats(
-        listen_addr.clone(),
-        "main".to_string(),
-        tokio::runtime::Handle::current(),
-    );
+    let main_rt = tokio::runtime::Handle::current();
+    worker_metrics::spawn_tokio_stats(listen_addr.clone(), "main".to_string(), main_rt.clone());
     worker_metrics::spawn_page_pool_stats(listen_addr.clone(), ctx.page_pool().clone());
     worker_metrics::spawn_bsatn_rlb_pool_stats(listen_addr.clone(), ctx.bsatn_rlb_pool().clone());
     let mut db_routes = DatabaseRoutes::default();
     db_routes.root_post = db_routes.root_post.layer(DefaultBodyLimit::disable());
     db_routes.db_put = db_routes.db_put.layer(DefaultBodyLimit::disable());
     db_routes.pre_publish = db_routes.pre_publish.layer(DefaultBodyLimit::disable());
     let extra = axum::Router::new().nest("/health", spacetimedb_client_api::routes::health::router());
-    let service = router(&ctx, db_routes, IdentityRoutes::default(), extra).with_state(ctx.clone());
+    let task_dumps = TaskDumpRegistry::new([("main", main_rt)]);
+    let service = router(&ctx, db_routes, IdentityRoutes::default(), extra)
+        .layer(Extension(task_dumps))
+        .with_state(ctx.clone());
 
     // Check if the requested port is available on both IPv4 and IPv6.
     // If not, offer to find an available port by incrementing (unless non-interactive).