diff --git a/Cargo.lock b/Cargo.lock index 2fca17a44f6..d1f8130fd09 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9582,6 +9582,7 @@ version = "1.48.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ff360e02eab121e0bc37a2d3b4d4dc622e6eda3a8e5253d5435ecf5bd4c68408" dependencies = [ + "backtrace", "bytes", "libc", "mio 1.1.0", diff --git a/crates/client-api/Cargo.toml b/crates/client-api/Cargo.toml index afbb5a2a340..09fcb563300 100644 --- a/crates/client-api/Cargo.toml +++ b/crates/client-api/Cargo.toml @@ -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 diff --git a/crates/client-api/src/routes/internal.rs b/crates/client-api/src/routes/internal.rs index b9faf10d391..79bae5ce5d4 100644 --- a/crates/client-api/src/routes/internal.rs +++ b/crates/client-api/src/routes/internal.rs @@ -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() -> axum::Router 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()) } diff --git a/crates/client-api/src/routes/internal/task_dump.rs b/crates/client-api/src/routes/internal/task_dump.rs new file mode 100644 index 00000000000..bd616658183 --- /dev/null +++ b/crates/client-api/src/routes/internal/task_dump.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>, + } + + impl TaskDumpRegistry { + pub fn new(runtimes: impl IntoIterator) -> Self { + Self { + runtimes: Arc::new(runtimes.into_iter().collect()), + } + } + } + + #[derive(Deserialize)] + pub(super) struct TaskDumpQuery { + runtime: String, + timeout_ms: Option, + } + + pub(super) async fn handle_get_task_dump( + registry: Option>, + Query(query): Query, + ) -> Result { + 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::>().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) -> 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() -> axum::Router { + axum::Router::new().route("/", get(handle_get_task_dump)) +} diff --git a/crates/client-api/src/routes/mod.rs b/crates/client-api/src/routes/mod.rs index 0a649db75b4..c7bbdeeaf2e 100644 --- a/crates/client-api/src/routes/mod.rs +++ b/crates/client-api/src/routes/mod.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 diff --git a/crates/standalone/src/subcommands/start.rs b/crates/standalone/src/subcommands/start.rs index 2cc1ede03c3..8961d55bc42 100644 --- a/crates/standalone/src/subcommands/start.rs +++ b/crates/standalone/src/subcommands/start.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,11 +198,8 @@ 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(); @@ -209,7 +207,10 @@ pub async fn exec(args: &ArgMatches, db_cores: JobCores) -> anyhow::Result<()> { 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).