feat(apiserver): Add support for tracing and graceful shutdown

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"

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 => {},
+    }
+}

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.