diff --git a/Cargo.lock b/Cargo.lock index d1443a1..da371bc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1814,6 +1814,15 @@ dependencies = [ "lazy_static", ] +[[package]] +name = "signal-hook-registry" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9e9e0b4211b72e7b8b6e85c807d36c212bdb33ea8587f7569562a84df5465b1" +dependencies = [ + "libc", +] + [[package]] name = "signature" version = "2.2.0" @@ -2255,6 +2264,7 @@ dependencies = [ "mio", "parking_lot", "pin-project-lite", + "signal-hook-registry", "socket2", "tokio-macros", "windows-sys 0.52.0", diff --git a/Cargo.toml b/Cargo.toml index 0fcb12c..1e83826 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -38,7 +38,7 @@ sqlx = { version = "0.8", features = [ "uuid", ] } subtle = "2" -tokio = { version = "1", features = ["parking_lot", "rt-multi-thread"] } +tokio = { version = "1", features = ["parking_lot", "rt-multi-thread", "signal"] } tokio-util = { version = "0.7" } tower = { version = "0.4" } tower-http = { version = "0.5", features = [ diff --git a/src/main.rs b/src/main.rs index 09c504e..43b7029 100644 --- a/src/main.rs +++ b/src/main.rs @@ -10,7 +10,7 @@ use std::{env, net::Ipv4Addr}; use auth::HashedBytes; use axum::{extract::FromRef, routing::post, Router}; use file_storage::FileStorage; -use tokio::net::TcpListener; +use tokio::{net::TcpListener, signal}; type Pool = sqlx::postgres::PgPool; @@ -66,11 +66,37 @@ async fn main() -> anyhow::Result<()> { let addr = (Ipv4Addr::UNSPECIFIED, 3000); let listener = TcpListener::bind(addr).await?; - axum::serve(listener, router).await?; + axum::serve(listener, router) + .with_graceful_shutdown(shutdown_signal()) + .await?; Ok(()) } +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 => {}, + } +} + fn app(state: AppState) -> Router { use axum::{http::header, routing::get}; use endpoints::{