simplify

Author: joshua-spacetime

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

@@ -158,11 +158,11 @@ mod jemalloc_profiling {
 }
 
 // The internal router is for things that are not meant to be exposed to the public API.
-pub fn router<S>(task_dumps: TaskDumpRegistry) -> axum::Router<S>
+pub fn router<S>() -> axum::Router<S>
 where
     S: NodeDelegate + Clone + 'static,
 {
     axum::Router::new()
         .nest("/heap", jemalloc_profiling::jemalloc_router())
-        .nest("/task-dump", task_dump::router(task_dumps))
+        .nest("/task-dump", task_dump::router())
 }

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

@@ -1,124 +1,73 @@
-use axum::{extract::Extension, routing::get};
+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 _, sync::Arc, time::Duration};
+    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, sync::Mutex};
-
-    const MAX_TIMEOUT_MS: u64 = 30_000;
+    use tokio::runtime::Handle;
 
-    #[derive(Clone)]
-    struct Runtime {
-        handle: Handle,
-        dump_lock: Arc<Mutex<()>>,
-    }
+    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, Runtime>>,
+        runtimes: Arc<BTreeMap<&'static str, Handle>>,
     }
 
     impl TaskDumpRegistry {
         pub fn new(runtimes: impl IntoIterator<Item = (&'static str, Handle)>) -> Self {
-            let runtimes = runtimes
-                .into_iter()
-                .map(|(name, handle)| {
-                    (
-                        name,
-                        Runtime {
-                            handle,
-                            dump_lock: Arc::new(Mutex::new(())),
-                        },
-                    )
-                })
-                .collect();
             Self {
-                runtimes: Arc::new(runtimes),
+                runtimes: Arc::new(runtimes.into_iter().collect()),
             }
         }
-
-        fn get(&self, name: &str) -> Option<&Runtime> {
-            self.runtimes.get(name)
-        }
-
-        fn names(&self) -> impl Iterator<Item = &'static str> + '_ {
-            self.runtimes.keys().copied()
-        }
     }
 
     #[derive(Deserialize)]
     pub(super) struct TaskDumpQuery {
         runtime: String,
-        timeout_ms: u64,
-    }
-
-    impl TaskDumpQuery {
-        fn validate(&self) -> Result<(), (StatusCode, String)> {
-            if !(1..=MAX_TIMEOUT_MS).contains(&self.timeout_ms) {
-                return Err((
-                    StatusCode::BAD_REQUEST,
-                    format!("timeout_ms must be between 1 and {MAX_TIMEOUT_MS}"),
-                ));
-            }
-            Ok(())
-        }
+        timeout_ms: Option<NonZeroU64>,
     }
 
     pub(super) async fn handle_get_task_dump(
-        Extension(registry): Extension<TaskDumpRegistry>,
+        registry: Option<Extension<TaskDumpRegistry>>,
         Query(query): Query<TaskDumpQuery>,
     ) -> Result<Response, (StatusCode, String)> {
-        query.validate()?;
+        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.get(&query.runtime).cloned() else {
-            let valid = registry.names().collect::<Vec<_>>().join(", ");
+        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 permit = runtime.dump_lock.try_lock_owned().map_err(|_| {
-            (
-                StatusCode::CONFLICT,
-                format!("a task dump is already in progress for runtime {:?}", query.runtime),
-            )
-        })?;
-
-        // Keep the permit in the spawned task so that a timed-out dump continues
-        // to exclude new requests until Tokio's dump future actually finishes.
-        let dump_task = tokio::spawn(async move {
-            let dump = runtime.handle.dump().await;
-            (permit, dump)
-        });
-        let (_permit, dump) = tokio::time::timeout(Duration::from_millis(query.timeout_ms), dump_task)
+        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 {}ms while dumping Tokio runtime {:?}",
-                        query.timeout_ms, query.runtime
+                        "timed out after {timeout_ms}ms while dumping Tokio runtime {:?}",
+                        query.runtime
                     ),
                 )
-            })?
-            .map_err(|err| {
-                (
-                    StatusCode::INTERNAL_SERVER_ERROR,
-                    format!("task dump worker failed for runtime {:?}: {err}", query.runtime),
-                )
             })?;
 
         let runtime_name = query.runtime;
@@ -155,44 +104,16 @@ mod imp {
     mod tests {
         use super::*;
         use http_body_util::BodyExt as _;
-
-        #[tokio::test]
-        async fn registry_names_are_sorted() {
-            let handle = Handle::current();
-            let registry = TaskDumpRegistry::new([("replication", handle.clone()), ("main", handle)]);
-
-            assert_eq!(registry.names().collect::<Vec<_>>(), ["main", "replication"]);
-        }
-
-        #[test]
-        fn timeout_must_be_in_range() {
-            for timeout_ms in [1, MAX_TIMEOUT_MS] {
-                assert!(TaskDumpQuery {
-                    runtime: "main".into(),
-                    timeout_ms,
-                }
-                .validate()
-                .is_ok());
-            }
-
-            for timeout_ms in [0, MAX_TIMEOUT_MS + 1] {
-                assert!(TaskDumpQuery {
-                    runtime: "main".into(),
-                    timeout_ms,
-                }
-                .validate()
-                .is_err());
-            }
-        }
+        use tokio::runtime::Handle;
 
         #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-        async fn dumps_registered_runtime() {
+        async fn dumps_registered_runtime_with_default_timeout() {
             let registry = TaskDumpRegistry::new([("main", Handle::current())]);
             let response = handle_get_task_dump(
-                Extension(registry),
+                Some(Extension(registry)),
                 Query(TaskDumpQuery {
                     runtime: "main".into(),
-                    timeout_ms: 10_000,
+                    timeout_ms: None,
                 }),
             )
             .await
@@ -248,8 +169,6 @@ mod imp {
 use imp::handle_get_task_dump;
 pub use imp::TaskDumpRegistry;
 
-pub fn router<S: Clone + Send + Sync + 'static>(registry: TaskDumpRegistry) -> axum::Router<S> {
-    axum::Router::new()
-        .route("/", get(handle_get_task_dump))
-        .layer(Extension(registry))
+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/mod.rs

@@ -27,25 +27,6 @@ pub fn router<S>(
     identity_routes: IdentityRoutes<S>,
     extra: axum::Router<S>,
 ) -> axum::Router<S>
-where
-    S: NodeDelegate + ControlStateDelegate + Authorization + Clone + 'static,
-{
-    router_with_task_dumps(
-        ctx,
-        database_routes,
-        identity_routes,
-        extra,
-        TaskDumpRegistry::default(),
-    )
-}
-
-pub fn router_with_task_dumps<S>(
-    ctx: &S,
-    database_routes: DatabaseRoutes<S>,
-    identity_routes: IdentityRoutes<S>,
-    extra: axum::Router<S>,
-    task_dumps: TaskDumpRegistry,
-) -> axum::Router<S>
 where
     S: NodeDelegate + ControlStateDelegate + Authorization + Clone + 'static,
 {
@@ -66,5 +47,5 @@ where
 
     axum::Router::new()
         .nest("/v1", router.layer(cors))
-        .nest("/internal", internal::router(task_dumps))
+        .nest("/internal", internal::router())
 }

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};
@@ -17,7 +17,7 @@ use spacetimedb::startup::{self, TracingOptions};
 use spacetimedb::util::jobs::JobCores;
 use spacetimedb::worker_metrics;
 use spacetimedb_client_api::routes::database::DatabaseRoutes;
-use spacetimedb_client_api::routes::router_with_task_dumps;
+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};
@@ -208,8 +208,9 @@ pub async fn exec(args: &ArgMatches, db_cores: JobCores) -> anyhow::Result<()> {
     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 task_dumps = TaskDumpRegistry::new([("main", main_rt)]);
-    let service =
-        router_with_task_dumps(&ctx, db_routes, IdentityRoutes::default(), extra, task_dumps).with_state(ctx.clone());
+    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).