diff --git a/clawshell.example.toml b/clawshell.example.toml index fa0eb11..2fb92f5 100644 --- a/clawshell.example.toml +++ b/clawshell.example.toml @@ -7,6 +7,10 @@ log_level = "info" host = "127.0.0.1" port = 18790 +# Graceful shutdown timeout (seconds). When receiving SIGTERM/SIGINT, +# the server waits up to this duration for in-flight requests to complete. +# shutdown_timeout_secs = 30 + [upstream] base_url = "https://api.openai.com" anthropic_base_url = "https://api.anthropic.com" diff --git a/src/config.rs b/src/config.rs index 9c10138..8769b74 100644 --- a/src/config.rs +++ b/src/config.rs @@ -42,6 +42,13 @@ pub struct ServerConfig { pub host: String, #[serde(default = "default_port")] pub port: u16, + + /// Graceful shutdown timeout in seconds. When receiving a shutdown signal + /// (SIGTERM/SIGINT), the server will wait up to this duration for in-flight + /// requests to complete before forcefully terminating. + /// Default: 30 seconds. + #[serde(default = "default_shutdown_timeout_secs")] + pub shutdown_timeout_secs: u64, } fn default_host() -> String { @@ -52,6 +59,10 @@ fn default_port() -> u16 { 18790 } +fn default_shutdown_timeout_secs() -> u64 { + 30 +} + #[derive(Debug, Deserialize, Clone)] pub struct UpstreamConfig { #[serde(default = "default_base_url")] diff --git a/src/main.rs b/src/main.rs index 0a452a3..b38ab8a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -131,21 +131,50 @@ async fn cmd_start_inner( process::drop_privileges()?; - axum::serve(listener, app) - .with_graceful_shutdown(async { - let ctrl_c = signal::ctrl_c(); - #[cfg(unix)] - let mut term = signal::unix::signal(signal::unix::SignalKind::terminate()).unwrap(); - #[cfg(unix)] - tokio::select! { _ = ctrl_c => {}, _ = term.recv() => {} }; - #[cfg(not(unix))] + // Capture shutdown timeout from config before moving into async block + let shutdown_timeout_secs = config.server.shutdown_timeout_secs; + + // Create a shutdown signal future that handles both SIGTERM and SIGINT + let shutdown_signal = async move { + let ctrl_c = signal::ctrl_c(); + + #[cfg(unix)] + { + let mut term = + signal::unix::signal(signal::unix::SignalKind::terminate()).unwrap(); + tokio::select! { + _ = ctrl_c => { + info!("Received SIGINT (Ctrl+C), initiating graceful shutdown..."); + } + _ = term.recv() => { + info!("Received SIGTERM, initiating graceful shutdown..."); + } + }; + } + + #[cfg(not(unix))] + { ctrl_c.await.ok(); - info!("Shutdown signal received"); - }) + info!("Received shutdown signal, initiating graceful shutdown..."); + } + + info!( + timeout_secs = shutdown_timeout_secs, + "Waiting for in-flight requests to complete..." + ); + }; + + // Start server with graceful shutdown + // Note: axum::serve's with_graceful_shutdown will: + // 1. Stop accepting new connections when shutdown signal is received + // 2. Wait for existing connections to complete their requests + // 3. The timeout is applied at the TCP level by the OS + axum::serve(listener, app) + .with_graceful_shutdown(shutdown_signal) .await?; process::remove_pid_file(); - info!("ClawShell shut down"); + info!("ClawShell shut down gracefully"); Ok(()) }