From b9a4a2dcb844c502adf44c816318be79a77b4696 Mon Sep 17 00:00:00 2001
From: Eguo Wang <wangeguo@gmail.com>
Date: Tue, 2 Jan 2024 10:47:48 +0800
Subject: [PATCH] feat(apiserver): Add support for tracing and graceful
 shutdown

---
 Cargo.lock            | 12 ++++++------
 Cargo.toml            |  2 +-
 apiserver/src/app.rs  | 44 +++++++++++++++++++++++++++++++++++++------
 apiserver/src/main.rs |  8 ++++++--
 4 files changed, 51 insertions(+), 15 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index d6e199c..8327bde 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -61,7 +61,7 @@ checksum = "0942ffc6dcaadf03badf6e6a2d0228460359d5e34b57ccdc720b7382dfbd5ec5"
 
 [[package]]
 name = "amp-apiserver"
-version = "0.6.15"
+version = "0.6.16"
 dependencies = [
  "amp-common",
  "amp-resources",
@@ -121,7 +121,7 @@ dependencies = [
 
 [[package]]
 name = "amp-controllers"
-version = "0.6.15"
+version = "0.6.16"
 dependencies = [
  "amp-common",
  "amp-resolver",
@@ -146,7 +146,7 @@ dependencies = [
 
 [[package]]
 name = "amp-crdgen"
-version = "0.6.15"
+version = "0.6.16"
 dependencies = [
  "amp-common",
  "clap",
@@ -158,7 +158,7 @@ dependencies = [
 
 [[package]]
 name = "amp-resolver"
-version = "0.6.15"
+version = "0.6.16"
 dependencies = [
  "amp-common",
  "amp-resources",
@@ -172,7 +172,7 @@ dependencies = [
 
 [[package]]
 name = "amp-resources"
-version = "0.6.15"
+version = "0.6.16"
 dependencies = [
  "amp-common",
  "anyhow",
@@ -192,7 +192,7 @@ dependencies = [
 
 [[package]]
 name = "amp-syncer"
-version = "0.6.15"
+version = "0.6.16"
 dependencies = [
  "amp-common",
  "async-nats",
diff --git a/Cargo.toml b/Cargo.toml
index afa4e9a..ac20237 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,5 +1,5 @@
 [workspace.package]
-version = "0.6.15"
+version = "0.6.16"
 edition = "2021"
 license = "Apache-2.0"
 repository = "https://github.com/amphitheatre-app/amphitheatre"
diff --git a/apiserver/src/app.rs b/apiserver/src/app.rs
index e7e2fd2..a171c46 100644
--- a/apiserver/src/app.rs
+++ b/apiserver/src/app.rs
@@ -12,25 +12,57 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::net::SocketAddr;
-use std::sync::Arc;
-
 use crate::context::Context;
 use crate::{routes, swagger};
 
+use std::net::SocketAddr;
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::signal;
+use tower_http::timeout::TimeoutLayer;
+use tower_http::trace::TraceLayer;
+
 pub async fn run(ctx: Arc<Context>) {
     let port = ctx.config.port;
 
     // build our application with a route
-    let app = routes::build().merge(swagger::build()).with_state(ctx);
+    let app = routes::build().merge(swagger::build()).with_state(ctx).layer((
+        TraceLayer::new_for_http(),
+        // Graceful shutdown will wait for outstanding requests to complete. Add a timeout so
+        // requests don't hang forever.
+        TimeoutLayer::new(Duration::from_secs(10)),
+    ));
 
     // run our app with hyper, and serve it over HTTP
     let addr = SocketAddr::from(([0, 0, 0, 0], port));
     let listener = tokio::net::TcpListener::bind(addr).await.unwrap();
 
-    // Run this server for ... forever!
-    if let Err(err) = axum::serve(listener, app).await {
+    // Run the server with graceful shutdown
+    let server = axum::serve(listener, app).with_graceful_shutdown(shutdown_signal()).await;
+    if let Err(err) = server {
         tracing::error!("Server error: {}", err);
         std::process::exit(1)
     }
 }
+
+async fn shutdown_signal() {
+    let ctrl_c = async {
+        signal::ctrl_c().await.expect("failed to install Ctrl+C handler");
+    };
+
+    #[cfg(unix)]
+    let terminate = async {
+        signal::unix::signal(signal::unix::SignalKind::terminate())
+            .expect("failed to install signal handler")
+            .recv()
+            .await;
+    };
+
+    #[cfg(not(unix))]
+    let terminate = std::future::pending::<()>();
+
+    tokio::select! {
+        _ = ctrl_c => {},
+        _ = terminate => {},
+    }
+}
diff --git a/apiserver/src/main.rs b/apiserver/src/main.rs
index 250e991..9b5cbd9 100644
--- a/apiserver/src/main.rs
+++ b/apiserver/src/main.rs
@@ -22,11 +22,15 @@ use amphitheatre::context::Context;
 use clap::Parser;
 use tracing::metadata::LevelFilter;
 use tracing_subscriber::EnvFilter;
+use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
-    let filter = EnvFilter::builder().with_default_directive(LevelFilter::INFO.into()).from_env_lossy();
-    tracing_subscriber::fmt().with_env_filter(filter).init();
+    // Enable tracing.
+    tracing_subscriber::registry()
+        .with(EnvFilter::builder().with_default_directive(LevelFilter::INFO.into()).from_env_lossy())
+        .with(tracing_subscriber::fmt::layer())
+        .init();
 
     // This returns an error if the `.env` file doesn't exist, but that's not what we want
     // since we're not going to use a `.env` file if we deploy this application.