diff --git a/Cargo.lock b/Cargo.lock index c223fc2..975c6cd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1709,6 +1709,7 @@ dependencies = [ "criterion", "filetime", "futures-util", + "libc", "portable-pty", "sandd-protocol", "serde", diff --git a/sandd/Cargo.toml b/sandd/Cargo.toml index e30254f..f366a61 100644 --- a/sandd/Cargo.toml +++ b/sandd/Cargo.toml @@ -78,6 +78,8 @@ filetime = "0.2" [dev-dependencies] tempfile = "3.8" criterion = { version = "0.5", features = ["async_tokio"] } +# Used only by shutdown tests to raise SIGTERM to our own process. +libc = "0.2" [[bench]] name = "snapshot_bench" diff --git a/sandd/src/main.rs b/sandd/src/main.rs index 4a31437..c5ea341 100644 --- a/sandd/src/main.rs +++ b/sandd/src/main.rs @@ -22,6 +22,45 @@ use tracing::{debug, error, info, warn}; /// tailscaled. Localhost-only: reachable solely by this container's daemon. const TUNNEL_SOCKS_PROXY: &str = "127.0.0.1:1055"; +/// Why the serve loop returned, so main() knows whether to reconnect or exit. +#[derive(Debug, PartialEq, Eq)] +enum ServeOutcome { + /// The connection dropped (server closed, socket error). main() reconnects. + Disconnected, + /// We received SIGTERM/SIGINT and closed the connection cleanly. main() exits + /// the process instead of reconnecting. + Shutdown, +} + +/// Resolve when the process is asked to terminate: SIGTERM (what `kubectl delete +/// pod` / `docker stop` send) or SIGINT (Ctrl-C). On non-unix, only Ctrl-C. +async fn shutdown_signal() { + #[cfg(unix)] + { + use tokio::signal::unix::{signal, SignalKind}; + // If we can't install a handler there is nothing sensible to do but keep + // running; a failed handler must not take the daemon down. + let mut sigterm = match signal(SignalKind::terminate()) { + Ok(s) => s, + Err(e) => { + warn!("Failed to install SIGTERM handler: {}", e); + // Fall back to only Ctrl-C. + let _ = tokio::signal::ctrl_c().await; + return; + } + }; + tokio::select! { + _ = sigterm.recv() => info!("Received SIGTERM, shutting down"), + _ = tokio::signal::ctrl_c() => info!("Received SIGINT, shutting down"), + } + } + #[cfg(not(unix))] + { + let _ = tokio::signal::ctrl_c().await; + info!("Received Ctrl-C, shutting down"); + } +} + #[derive(Parser, Debug)] #[command(name = "sandd")] #[command( @@ -111,7 +150,18 @@ async fn main() -> Result<()> { ) .await { - Ok(_) => info!("Connection closed gracefully"), + // A clean shutdown (SIGTERM/SIGINT) must NOT reconnect — exit the + // process so the pod terminates promptly and the controller sees the + // Close frame we just sent. + Ok(ServeOutcome::Shutdown) => { + info!("Shutdown complete"); + return Ok(()); + } + // The specific reason (server Close, socket error, stream end, + // registration failure) is already logged at the break site inside + // serve(); avoid claiming "gracefully" here since Disconnected also + // covers error paths. main() only needs to know: reconnect. + Ok(ServeOutcome::Disconnected) => info!("Connection closed, reconnecting"), Err(e) => error!("Connection error: {}", e), } @@ -126,7 +176,7 @@ async fn connect_and_serve( heartbeat_interval: u64, labels: HashMap, tunnel: bool, -) -> Result<()> { +) -> Result { info!("Connecting to server at {}", server_url); // Connect with subprotocol using client builder @@ -175,7 +225,7 @@ async fn connect_and_serve( .await .context("tunnel: WebSocket handshake over SOCKS5 failed")?; log_negotiated_protocol(&response); - return serve(ws_stream, daemon_id, heartbeat_interval, labels).await; + return serve(ws_stream, daemon_id, heartbeat_interval, labels, shutdown_signal()).await; } let (ws_stream, response) = match tokio_tungstenite::connect_async(request).await { @@ -186,7 +236,7 @@ async fn connect_and_serve( } }; log_negotiated_protocol(&response); - serve(ws_stream, daemon_id, heartbeat_interval, labels).await + serve(ws_stream, daemon_id, heartbeat_interval, labels, shutdown_signal()).await } /// Log the WebSocket subprotocol the server negotiated (shared by both transports). @@ -203,14 +253,21 @@ fn log_negotiated_protocol( /// Run the daemon session over an established WebSocket stream. Generic over the /// transport so the direct (connect_async) and tunnel (SOCKS5) paths share one /// implementation. -async fn serve( +/// +/// `shutdown` is the future that, once resolved, triggers a graceful close: in +/// production it is `shutdown_signal()` (SIGTERM/SIGINT); tests inject a future +/// they control so the shutdown path can be exercised without raising a real, +/// process-wide signal mid-connection. +async fn serve( ws_stream: tokio_tungstenite::WebSocketStream, daemon_id: &str, heartbeat_interval: u64, labels: HashMap, -) -> Result<()> + shutdown: F, +) -> Result where S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static, + F: std::future::Future, { info!("WebSocket connection established"); @@ -244,7 +301,7 @@ where info!("Registration successful: {}", message); } else { error!("Registration failed: {}", message); - return Ok(()); + return Ok(ServeOutcome::Disconnected); } } _ => { @@ -284,47 +341,82 @@ where } }); - // Message processing loop - while let Some(msg) = ws_rx.next().await { - let msg = match msg { - Ok(WsMessage::Text(text)) => text, - Ok(WsMessage::Close(_)) => { - info!("Server closed connection"); - break; - } - Err(e) => { - error!("WebSocket error: {}", e); - break; + // Message processing loop. Also races a shutdown signal so that on + // SIGTERM/SIGINT we send a WebSocket Close frame BEFORE exiting: the + // controller removes a daemon the moment it sees that Close (server.rs + // handle_websocket), so a graceful pod deletion deregisters immediately + // instead of waiting out the ~90s heartbeat-timeout reaper. + // + // Pin the shutdown future once so it can be polled across loop iterations + // without being moved (it may be `!Unpin`). + tokio::pin!(shutdown); + let outcome = loop { + tokio::select! { + // Poll shutdown FIRST. With `biased`, tokio checks branches top to + // bottom and only reaches a later branch when earlier ones are + // Pending; if the message branch were first and inbound messages were + // continuously ready (buffered / high throughput), the shutdown branch + // could starve and delay the Close frame + process exit. Shutdown is + // usually Pending, so in normal operation this falls straight through + // to draining messages; it can't starve them because shutdown fires + // once and breaks the loop. + biased; + + _ = &mut shutdown => { + // Best-effort clean close so the controller deregisters us now. + let mut tx = ws_tx_clone.lock().await; + if let Err(e) = tx.send(WsMessage::Close(None)).await { + warn!("Failed to send Close frame on shutdown: {}", e); + } + break ServeOutcome::Shutdown; } - _ => continue, - }; - let message: Message = match serde_json::from_str(&msg) { - Ok(m) => m, - Err(e) => { - error!("Failed to parse message: {}", e); - continue; + msg = ws_rx.next() => { + let msg = match msg { + Some(Ok(WsMessage::Text(text))) => text, + Some(Ok(WsMessage::Close(_))) => { + info!("Server closed connection"); + break ServeOutcome::Disconnected; + } + Some(Err(e)) => { + error!("WebSocket error: {}", e); + break ServeOutcome::Disconnected; + } + None => { + // Stream ended. + break ServeOutcome::Disconnected; + } + _ => continue, + }; + + let message: Message = match serde_json::from_str(&msg) { + Ok(m) => m, + Err(e) => { + error!("Failed to parse message: {}", e); + continue; + } + }; + + // Handle message inline + if let Err(e) = handle_message( + message, + ws_tx_clone.clone(), + executor.clone(), + session_manager.clone(), + snapshot_manager.clone(), + ) + .await + { + error!("Error handling message: {}", e); + } } - }; - - // Handle message inline - if let Err(e) = handle_message( - message, - ws_tx_clone.clone(), - executor.clone(), - session_manager.clone(), - snapshot_manager.clone(), - ) - .await - { - error!("Error handling message: {}", e); } - } + }; heartbeat_handle.abort(); info!("Disconnected from agent"); - Ok(()) + Ok(outcome) } async fn handle_message( @@ -737,3 +829,160 @@ async fn setup_tunnel(args: &Args) -> Result<()> { Err(anyhow::anyhow!("Timeout waiting for mesh IP assignment")) } + +#[cfg(test)] +mod shutdown_tests { + //! Tests for the graceful-shutdown path added to `serve`: on a shutdown + //! signal the daemon must send a WebSocket Close frame and return + //! `ServeOutcome::Shutdown` (so `main` exits instead of reconnecting). The + //! Close is what lets the controller deregister the daemon immediately + //! rather than waiting out its ~90s heartbeat-timeout reaper. + //! + //! `serve` takes the shutdown future as a parameter precisely so these tests + //! can trigger it deterministically, without raising a real, process-wide + //! SIGTERM in the middle of a test run. + + use super::*; + use futures_util::{SinkExt, StreamExt}; + use tokio::net::{TcpListener, TcpStream}; + use tokio_tungstenite::tungstenite::protocol::Message as WsMessage; + use tokio_tungstenite::{accept_async, WebSocketStream}; + + /// Stand up an in-process WebSocket server on localhost and connect a client + /// to it. Returns (client_stream_for_serve, accepted_server_stream). The + /// server side lets a test act as the controller: ack registration, then + /// observe what the daemon sends (e.g. the Close frame on shutdown). + async fn ws_pair() -> ( + WebSocketStream>, + WebSocketStream, + ) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + // Accept concurrently with the client dial so neither side blocks. + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + accept_async(stream).await.unwrap() + }); + + let url = format!("ws://{}/ws", addr); + let (client, _resp) = tokio_tungstenite::connect_async(url).await.unwrap(); + let server = server.await.unwrap(); + (client, server) + } + + /// Play the controller: read the daemon's Register and reply RegisterAck so + /// `serve` proceeds past registration into its main loop. + async fn ack_registration(server: &mut WebSocketStream) { + let reg = server.next().await.unwrap().unwrap(); + let text = reg.into_text().unwrap(); + let msg: Message = serde_json::from_str(&text).unwrap(); + assert!( + matches!(msg, Message::Register { .. }), + "expected Register first, got: {:?}", + msg + ); + let ack = Message::RegisterAck { + success: true, + message: "ok".to_string(), + }; + server + .send(WsMessage::Text(serde_json::to_string(&ack).unwrap())) + .await + .unwrap(); + } + + /// The core regression: when the shutdown future fires, `serve` returns + /// `Shutdown` AND the peer receives a Close frame. + /// + /// `serve`'s future is not `Send` (SessionManager is !Sync), so it can't be + /// `tokio::spawn`ed; instead we run it concurrently with the controller side + /// via `join!` on the current task. + #[tokio::test] + async fn shutdown_sends_close_and_returns_shutdown() { + let (client, mut server) = ws_pair().await; + + // Trigger shutdown via a oneshot fired AFTER registration is acked, so + // the ordering is deterministic (no wall-clock sleep racing the + // handshake). `serve` awaits the receiver; mapping away the RecvError + // matches its `Future` bound and means a dropped sender also + // resolves shutdown rather than hanging. + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>(); + let shutdown = async { + let _ = shutdown_rx.await; + }; + + let daemon = serve(client, "test-daemon", 3600, HashMap::new(), shutdown); + + let controller = async { + ack_registration(&mut server).await; + // Registration is complete and serve() is now in its select! loop; + // fire shutdown deterministically. + shutdown_tx.send(()).unwrap(); + // The controller side should observe a Close frame. Tolerate any + // pre-close traffic (e.g. a heartbeat), though the 3600s interval + // makes that unlikely in-test. + let mut saw_close = false; + while let Some(frame) = server.next().await { + match frame { + Ok(WsMessage::Close(_)) => { + saw_close = true; + break; + } + Ok(_) => continue, + Err(_) => break, + } + } + saw_close + }; + + let (outcome, saw_close) = tokio::join!(daemon, controller); + assert!(saw_close, "daemon did not send a Close frame on shutdown"); + assert_eq!(outcome.unwrap(), ServeOutcome::Shutdown); + } + + /// The counterpart: if the controller closes the connection, `serve` returns + /// `Disconnected` (so `main` reconnects) — NOT `Shutdown`. Guards against the + /// shutdown branch swallowing ordinary disconnects. + #[tokio::test] + async fn server_close_returns_disconnected() { + let (client, mut server) = ws_pair().await; + + // A shutdown future that never resolves: only the server-close path can + // end this session. + let shutdown = std::future::pending::<()>(); + + let daemon = serve(client, "test-daemon", 3600, HashMap::new(), shutdown); + + let controller = async { + ack_registration(&mut server).await; + // Controller closes the connection. + server.close(None).await.unwrap(); + }; + + let (outcome, ()) = tokio::join!(daemon, controller); + assert_eq!(outcome.unwrap(), ServeOutcome::Disconnected); + } + + /// `shutdown_signal()` must resolve when the process receives SIGTERM (what + /// `kubectl delete pod` / `docker stop` send). Uses a real self-signal; unix + /// only. This asserts the wiring, not the WebSocket behavior above. + #[cfg(unix)] + #[tokio::test] + async fn shutdown_signal_resolves_on_sigterm() { + // Raise SIGTERM to our own process after a short delay, then confirm the + // helper's future completes rather than hanging. + tokio::spawn(async { + tokio::time::sleep(Duration::from_millis(50)).await; + // SAFETY: raising a signal to our own PID is sound; kill(2) with a + // valid signal number has no memory-safety implications. + unsafe { + libc::raise(libc::SIGTERM); + } + }); + + tokio::time::timeout(Duration::from_secs(5), shutdown_signal()) + .await + .expect("shutdown_signal did not resolve on SIGTERM"); + } +}