Use axum's native with_executor, remove hand-rolled serve loop

Now that axum supports Serve::with_executor(), implement axum::serve::Executor directly on TelemetryHandle behind an "axum" feature flag.

Author: yulnr

Cargo.lock

@@ -34,6 +34,8 @@ bon = "3"
 aws-sdk-s3-transfer-manager = { version = "0.1.3", optional = true, default-features = false }
 aws-sdk-s3 = { version = "1", default-features = false, optional = true, features = ["behavior-version-latest", "sigv4a", "rt-tokio", "default-https-client"] }
 aws-config = { version = "1", optional = true, default-features = false }
+# TODO: update this once this is merged upstream
+axum = { git = "https://github.com/yulnr/axum", branch = "feat/custom-serve-executor", default-features = false, features = ["tokio", "http1", "http2"], optional = true }
 flate2 = { version = "1" }
 time = { version = "0.3", features = ["formatting", "macros"], optional = true }
 metrique-writer = "0.1"
@@ -42,6 +44,7 @@ metrique = { version = "0.1.20", features = ["local-format"] }
 [features]
 cpu-profiling = ["dep:dial9-perf-self-profile"]
 task-dump = ["tokio/taskdump"]
+axum = ["dep:axum"]
 worker-s3 = ["dep:aws-sdk-s3-transfer-manager", "dep:aws-sdk-s3", "dep:aws-config", "dep:time"]
 
 [dev-dependencies]

dial9-tokio-telemetry/Cargo.toml

@@ -216,6 +216,18 @@ impl TelemetryHandle {
     }
 }
 
+/// Implement `axum::serve::Executor` for `TelemetryHandle` to allow it to be
+/// used as an executor for axum's serve.
+#[cfg(feature = "axum")]
+impl axum::serve::Executor for TelemetryHandle {
+    fn execute<Fut>(&self, fut: Fut)
+    where
+        Fut: std::future::Future<Output = ()> + Send + 'static,
+    {
+        self.spawn(fut);
+    }
+}
+
 /// Holds the background worker thread and its stop signal.
 struct WorkerHandle {
     shutdown: Option<tokio::sync::oneshot::Sender<Duration>>,

dial9-tokio-telemetry/src/telemetry/recorder/mod.rs

@@ -7,8 +7,9 @@ publish = false
 [dependencies]
 tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "sync", "process"] }
 tokio-util = { version = "0.7", features = ["rt"] }
-axum = "0.8"
-axum-core = "0.5"
+# TODO: update this once this is merged upstream
+axum = { git = "https://github.com/yulnr/axum", branch = "feat/custom-serve-executor" }
+axum-core = { git = "https://github.com/yulnr/axum", branch = "feat/custom-serve-executor" }
 hyper = "1"
 hyper-util = { version = "0.1", features = ["tokio", "server-auto"] }
 tower = "0.5"
@@ -20,12 +21,12 @@ aws-config = "1"
 aws-sdk-dynamodb = "1"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
-dial9-tokio-telemetry = { path = "../../dial9-tokio-telemetry", features = ["worker-s3"] }
+dial9-tokio-telemetry = { path = "../../dial9-tokio-telemetry", features = ["worker-s3", "axum"] }
 reqwest = { version = "0.12", default-features = false, features = ["json"] }
 clap = { version = "4", features = ["derive"] }
 uuid = { version = "1", features = ["v4"] }
 hostname = "0.4"
 libc = "0.2"
 
 [target.'cfg(target_os = "linux")'.dependencies]
-dial9-tokio-telemetry = { path = "../../dial9-tokio-telemetry", features = ["cpu-profiling", "worker-s3"] }
+dial9-tokio-telemetry = { path = "../../dial9-tokio-telemetry", features = ["cpu-profiling", "worker-s3", "axum"] }

examples/metrics-service/Cargo.toml

@@ -1,4 +1,3 @@
-mod axum_traced;
 mod buffer;
 mod ddb;
 mod routes;
@@ -253,10 +252,11 @@ fn main() -> std::io::Result<()> {
             }
         });
 
-        axum_traced::serve(listener, app.into_make_service(), handle.clone())
+        axum::serve(listener, app.into_make_service())
+            .with_executor(handle.clone())
             .with_graceful_shutdown(async move { shutdown.cancelled().await })
-            .await
-            .unwrap();
+            .await;
+
         let shutdown = guard.graceful_shutdown(Duration::from_secs(5)).await;
         tracing::info!("dial9 shutdown: {shutdown:?}");
     });