chore: update axum for graceful shutdown

ldclabs · Dec 30, 2023 · 1aa2a0e · 1aa2a0e
1 parent d3873b5
commit 1aa2a0e
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 18 deletions.
diff --git a/crates/ns-indexer/Cargo.toml b/crates/ns-indexer/Cargo.toml
@@ -61,6 +61,7 @@ tower-http = { version = "0.5", features = [
   "decompression-zstd",
   "propagate-header",
   "cors",
+  "timeout",
 ] }
 validator = { version = "0.16", features = ["derive"] }
 

diff --git a/crates/ns-indexer/src/bin/main.rs b/crates/ns-indexer/src/bin/main.rs
@@ -93,16 +93,12 @@ fn main() -> anyhow::Result<()> {
                 .await
                 .expect("failed to bind");
 
-            // wait for "graceful shutdown" https://github.com/tokio-rs/axum/pull/2398
-            // let res = axum::serve(listener, app).await;
-            tokio::select! {
-                _ = async {
-                    match axum::serve(listener, app).await {
-                        Ok(_) => log::info!(target: "server", "indexer api finished"),
-                        Err(err) => log::error!(target: "server", "indexer api error: {}", err),
-                    }
-                } => {},
-                _ = shutdown.clone() => {},
+            match axum::serve(listener, app)
+                .with_graceful_shutdown(shutdown.clone())
+                .await
+            {
+                Ok(_) => log::info!(target: "server", "indexer api finished"),
+                Err(err) => log::error!(target: "server", "indexer api error: {}", err),
             }
 
             Ok::<(), anyhow::Error>(())

diff --git a/crates/ns-indexer/src/router.rs b/crates/ns-indexer/src/router.rs
@@ -1,10 +1,11 @@
 use axum::{middleware, routing, Router};
 use std::sync::Arc;
-use tower::ServiceBuilder;
+use std::time::Duration;
 use tower_http::{
     catch_panic::CatchPanicLayer,
     compression::{predicate::SizeAbove, CompressionLayer},
     cors::CorsLayer,
+    timeout::TimeoutLayer,
 };
 
 use ns_axum_web::context;
@@ -13,12 +14,6 @@ use ns_axum_web::encoding;
 use crate::api;
 
 pub fn new(state: Arc<api::IndexerAPI>) -> Router {
-    let mds = ServiceBuilder::new()
-        .layer(CatchPanicLayer::new())
-        .layer(middleware::from_fn(context::middleware))
-        .layer(CorsLayer::very_permissive())
-        .layer(CompressionLayer::new().compress_when(SizeAbove::new(encoding::MIN_ENCODING_SIZE)));
-
     Router::new()
         .route("/", routing::get(api::version))
         .route("/healthz", routing::get(api::healthz))
@@ -81,6 +76,12 @@ pub fn new(state: Arc<api::IndexerAPI>) -> Router {
         //         .route("/list_by_code", routing::get(api::ServiceAPI::get))
         //         .route("/list_by_submitter", routing::get(api::ServiceAPI::get)),
         // )
-        .route_layer(mds)
+        .layer((
+            CatchPanicLayer::new(),
+            TimeoutLayer::new(Duration::from_secs(10)),
+            middleware::from_fn(context::middleware),
+            CorsLayer::very_permissive(),
+            CompressionLayer::new().compress_when(SizeAbove::new(encoding::MIN_ENCODING_SIZE)),
+        ))
         .with_state(state)
 }