feat: add optional shutdown timeout to servers (#43)

hseeberger · Nov 22, 2023 · 5f098c6 · 5f098c6
1 parent 1f36e4b
commit 5f098c6
Show file tree

Hide file tree

Showing 9 changed files with 63 additions and 26 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -20,6 +20,7 @@ publish       = false
 anyhow                = { version = "1.0", features = [ "backtrace" ] }
 axum                  = { version = "0.6", features = [ "http2", "json" ] }
 configured            = { version = "0.7" }
+humantime-serde       = { version = "1.1" }
 opentelemetry         = { version = "0.21" }
 opentelemetry_sdk     = { version = "0.21", features = [ "rt-tokio" ] }
 opentelemetry-http    = { version = "0.10", features = [ "reqwest" ] }

diff --git a/hello-tracing-backend/Cargo.toml b/hello-tracing-backend/Cargo.toml
@@ -16,6 +16,7 @@ hello-tracing-common  = { path = "../hello-tracing-common" }
 anyhow                = { workspace = true }
 axum                  = { workspace = true }
 configured            = { workspace = true }
+humantime-serde       = { workspace = true }
 opentelemetry         = { workspace = true }
 opentelemetry-http    = { workspace = true }
 opentelemetry-otlp    = { workspace = true }

diff --git a/hello-tracing-backend/config/default.yaml b/hello-tracing-backend/config/default.yaml
@@ -1,6 +1,7 @@
 api:
   addr: 0.0.0.0
   port: 8080
+  shutdown-timeout: 3s
 
 tracing:
   service-name: hello-tracing-backend

diff --git a/hello-tracing-backend/src/api/mod.rs b/hello-tracing-backend/src/api/mod.rs
@@ -4,8 +4,11 @@ use anyhow::{Context, Result};
 use axum::{body::Body, http::Request};
 use hello_tracing_common::otel::http::{accept_trace, record_trace_id};
 use serde::Deserialize;
-use std::net::{IpAddr, SocketAddr};
-use tokio::signal::unix::{signal, SignalKind};
+use std::{net::IpAddr, time::Duration};
+use tokio::{
+    signal::unix::{signal, SignalKind},
+    time,
+};
 use tonic::transport::Server;
 use tower::ServiceBuilder;
 use tower_http::trace::TraceLayer;
@@ -16,16 +19,16 @@ use tracing::{field, info_span, Span};
 pub struct Config {
     addr: IpAddr,
     port: u16,
-}
-
-impl Config {
-    fn socket_addr(&self) -> SocketAddr {
-        SocketAddr::new(self.addr, self.port)
-    }
+    #[serde(with = "humantime_serde")]
+    shutdown_timeout: Option<Duration>,
 }
 
 pub async fn serve(config: Config) -> Result<()> {
-    let socket_addr = config.socket_addr();
+    let Config {
+        addr,
+        port,
+        shutdown_timeout,
+    } = config;
 
     let app = Server::builder()
         .layer(
@@ -36,19 +39,22 @@ pub async fn serve(config: Config) -> Result<()> {
         )
         .add_service(v0::hello());
 
-    app.serve_with_shutdown(socket_addr, shutdown_signal())
+    app.serve_with_shutdown((addr, port).into(), shutdown_signal(shutdown_timeout))
         .await
-        .context("serving the api")
+        .context("run server")
 }
 
 fn make_span(request: &Request<Body>) -> Span {
     let headers = request.headers();
     info_span!("incoming request", ?headers, trace_id = field::Empty)
 }
 
-async fn shutdown_signal() {
+async fn shutdown_signal(shutdown_timeout: Option<Duration>) {
     signal(SignalKind::terminate())
         .expect("install SIGTERM handler")
         .recv()
         .await;
+    if let Some(shutdown_timeout) = shutdown_timeout {
+        time::sleep(shutdown_timeout).await;
+    }
 }
diff --git a/hello-tracing-gateway/Cargo.toml b/hello-tracing-gateway/Cargo.toml
@@ -16,6 +16,7 @@ hello-tracing-common  = { path = "../hello-tracing-common" }
 anyhow                = { workspace = true }
 axum                  = { workspace = true }
 configured            = { workspace = true }
+humantime-serde       = { workspace = true }
 opentelemetry         = { workspace = true }
 opentelemetry-http    = { workspace = true }
 opentelemetry-otlp    = { workspace = true }

diff --git a/hello-tracing-gateway/config/default.yaml b/hello-tracing-gateway/config/default.yaml
@@ -1,6 +1,7 @@
 api:
   addr: 0.0.0.0
   port: 8080
+  shutdown-timeout: 3s
 
 # backend:
 #   endpoint: ...

diff --git a/hello-tracing-gateway/src/api/mod.rs b/hello-tracing-gateway/src/api/mod.rs
@@ -11,8 +11,11 @@ use axum::{
 };
 use hello_tracing_common::otel::http::{accept_trace, record_trace_id};
 use serde::Deserialize;
-use std::net::{IpAddr, SocketAddr};
-use tokio::signal::unix::{signal, SignalKind};
+use std::{net::IpAddr, time::Duration};
+use tokio::{
+    signal::unix::{signal, SignalKind},
+    time,
+};
 use tower::ServiceBuilder;
 use tower_http::trace::TraceLayer;
 use tracing::{field, info_span, Span};
@@ -22,15 +25,17 @@ use tracing::{field, info_span, Span};
 pub struct Config {
     addr: IpAddr,
     port: u16,
-}
-
-impl Config {
-    fn socket_addr(&self) -> SocketAddr {
-        SocketAddr::new(self.addr, self.port)
-    }
+    #[serde(with = "humantime_serde")]
+    shutdown_timeout: Option<Duration>,
 }
 
 pub async fn serve(config: Config, backend: Backend) -> Result<()> {
+    let Config {
+        addr,
+        port,
+        shutdown_timeout,
+    } = config;
+
     let app = Router::new()
         .route("/", get(ready))
         .nest("/v0", v0::app(backend))
@@ -41,11 +46,11 @@ pub async fn serve(config: Config, backend: Backend) -> Result<()> {
                 .map_request(record_trace_id),
         );
 
-    Server::bind(&config.socket_addr())
+    Server::bind(&(addr, port).into())
         .serve(app.into_make_service())
-        .with_graceful_shutdown(shutdown_signal())
+        .with_graceful_shutdown(shutdown_signal(shutdown_timeout))
         .await
-        .context("serving api")
+        .context("run server")
 }
 
 async fn ready() -> impl IntoResponse {
@@ -57,9 +62,12 @@ fn make_span(request: &Request<Body>) -> Span {
     info_span!("incoming request", ?headers, trace_id = field::Empty)
 }
 
-async fn shutdown_signal() {
+async fn shutdown_signal(shutdown_timeout: Option<Duration>) {
     signal(SignalKind::terminate())
         .expect("install SIGTERM handler")
         .recv()
         .await;
+    if let Some(shutdown_timeout) = shutdown_timeout {
+        time::sleep(shutdown_timeout).await;
+    }
 }
diff --git a/justfile b/justfile
@@ -32,14 +32,14 @@ run-gateway port="8080" backend_port="8081":
 		APP__API__PORT={{port}} \
 		APP__BACKEND__ENDPOINT=http://localhost:{{backend_port}} \
 		cargo run -p hello-tracing-gateway \
-		> $HOME/tmp/hello-tracing-gateway.log
+		| jq
 
 run-backend port="8081":
 	RUST_LOG=hello_tracing_backend=debug,info \
 		CONFIG_DIR=hello-tracing-backend/config \
 		APP__API__PORT={{port}} \
 		cargo run -p hello-tracing-backend \
-		> $HOME/tmp/hello-tracing-backend.log
+		| jq
 
 docker tag="latest":
 	docker build \