feat: Native axum executor integration

examples/metrics-service/Cargo.toml

@@ -7,8 +7,9 @@ publish = false
 [dependencies]
 tokio = { version = "1.51.0", 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"

examples/metrics-service/src/main.rs

@@ -1,4 +1,3 @@
-mod axum_traced;
 mod buffer;
 mod ddb;
 mod routes;
@@ -10,7 +9,7 @@ use aws_config::BehaviorVersion;
 use clap::Parser;
 #[cfg(target_os = "linux")]
 use dial9_tokio_telemetry::telemetry::{CpuProfilingConfig, SchedEventConfig};
-use dial9_tokio_telemetry::telemetry::{RotatingWriter, TracedRuntime};
+use dial9_tokio_telemetry::telemetry::{RotatingWriter, TelemetryHandle, TracedRuntime};
 use tokio::runtime::Builder;
 use tokio_util::sync::CancellationToken;
 
@@ -253,10 +252,11 @@ fn main() -> std::io::Result<()> {
             }
         });
 
-        axum_traced::serve(listener, app.into_make_service(), handle.clone())
+        let executor = TelemetryExecutor(handle);
+        axum::serve(listener, app.into_make_service())
+            .with_executor(executor)
             .with_graceful_shutdown(async move { shutdown.cancelled().await })
-            .await
-            .unwrap();
+            .await;
     });
 
     // Drop the runtime first so worker threads exit and flush their
@@ -267,3 +267,16 @@ fn main() -> std::io::Result<()> {
 
     Ok(())
 }
+
+/// Adapter to allow `TelemetryHandle` to be used as an executor for axum's serve.
+#[derive(Clone)]
+struct TelemetryExecutor(TelemetryHandle);
+
+impl axum::serve::Executor for TelemetryExecutor {
+    fn execute<Fut>(&self, fut: Fut)
+    where
+        Fut: std::future::Future<Output = ()> + Send + 'static,
+    {
+        self.0.spawn(fut);
+    }
+}