diff --git a/crates/rutster/src/config.rs b/crates/rutster/src/config.rs index fe5d8bd..96af499 100644 --- a/crates/rutster/src/config.rs +++ b/crates/rutster/src/config.rs @@ -79,6 +79,20 @@ pub fn max_sessions(raw: Option) -> Result { } } +/// `RUTSTER_DRAIN_DEADLINE_SECS` — how long shutdown waits for in-flight +/// calls before the hard stop. Default 0 = today's instant shutdown (dev +/// loop); production sets minutes. Calls are non-migratable by design, so +/// drain-then-terminate is the ONLY graceful scale-in shape. +pub fn drain_deadline(raw: Option) -> Result { + match raw { + None => Ok(std::time::Duration::ZERO), + Some(s) => s + .parse::() + .map(std::time::Duration::from_secs) + .map_err(|e| format!("RUTSTER_DRAIN_DEADLINE_SECS {s:?}: {e}")), + } +} + #[cfg(test)] mod tests { use super::*; @@ -141,6 +155,25 @@ mod tests { assert!(err.contains("RUTSTER_MAX_SESSIONS")); } + #[test] + fn drain_deadline_defaults_to_zero_when_unset() { + assert_eq!(drain_deadline(None).unwrap(), std::time::Duration::ZERO); + } + + #[test] + fn drain_deadline_parses_seconds() { + assert_eq!( + drain_deadline(Some("300".into())).unwrap(), + std::time::Duration::from_secs(300) + ); + } + + #[test] + fn drain_deadline_rejects_garbage_with_var_name_in_error() { + let err = drain_deadline(Some("soon".into())).unwrap_err(); + assert!(err.contains("RUTSTER_DRAIN_DEADLINE_SECS")); + } + #[test] fn media_address_config_parses_all_three() { let cfg = media_address_config( diff --git a/crates/rutster/src/main.rs b/crates/rutster/src/main.rs index fc01f39..93332be 100644 --- a/crates/rutster/src/main.rs +++ b/crates/rutster/src/main.rs @@ -58,21 +58,55 @@ async fn main() { // RUTSTER_TAP_BIND pattern in rutster-brain-realtime). let addr: SocketAddr = rutster::config::http_bind(std::env::var("RUTSTER_HTTP_BIND").ok()) .expect("RUTSTER_HTTP_BIND must be host:port"); + let drain_deadline = + rutster::config::drain_deadline(std::env::var("RUTSTER_DRAIN_DEADLINE_SECS").ok()) + .expect("RUTSTER_DRAIN_DEADLINE_SECS must be seconds"); info!(%addr, "listening"); let listener = tokio::net::TcpListener::bind(addr).await.unwrap(); + + // Two-phase shutdown (slice-5, review M1): + // SIGTERM → Drain (stop admitting; keep serving HTTP for in-flight + // calls' DELETE/offer; readyz flips so the LB pulls us) → deadline + // → HTTP quiesce → media_thread.shutdown() (the hard stop). + let (http_stop_tx, http_stop_rx) = tokio::sync::oneshot::channel::<()>(); + let drain_cmd_tx = app_state.cmd_tx.clone(); + tokio::spawn(async move { + shutdown_signal().await; + if !drain_deadline.is_zero() { + let (reply, drained) = tokio::sync::oneshot::channel(); + if drain_cmd_tx + .send(rutster::media_thread::MediaCmd::Drain { reply }) + .await + .is_ok() + { + match tokio::time::timeout(drain_deadline, drained).await { + Ok(_) => info!("drain complete; proceeding to shutdown"), + Err(_) => info!(?drain_deadline, "drain deadline hit; shutting down anyway"), + } + } + } + let _ = http_stop_tx.send(()); + }); + axum::serve(listener, router(app_state)) - .with_graceful_shutdown(shutdown_signal()) + .with_graceful_shutdown(async { + let _ = http_stop_rx.await; + }) .await .unwrap(); media_thread.shutdown(); } -/// Ctrl-C / SIGTERM handler (spec §4.5). After the signal, `main` returns -/// from `axum::serve(...)` and calls `media_thread.shutdown()`, which sends -/// `MediaCmd::Shutdown` to the dedicated thread. The thread drains its -/// session map, exits its loop, and is joined. No in-flight call -/// preservation story in the dev loop. +/// Ctrl-C / SIGTERM handler (spec §4.5). This future resolves the instant +/// the signal arrives — it does NOT wait for the drain. The caller (`main`, +/// slice-5 review M1) is the two-phase shutdown: after this resolves, it +/// sends `MediaCmd::Drain` and keeps `axum::serve` running so in-flight +/// calls can still hit DELETE/offer while readyz flips the LB away from +/// this node; once the map empties or `RUTSTER_DRAIN_DEADLINE_SECS` elapses +/// (whichever first — default 0 skips this wait entirely, preserving the +/// old instant-shutdown dev loop), HTTP quiesces and `media_thread.shutdown()` +/// sends `MediaCmd::Shutdown` for the hard stop. async fn shutdown_signal() { let ctrl_c = async { tokio::signal::ctrl_c() diff --git a/crates/rutster/src/media_thread.rs b/crates/rutster/src/media_thread.rs index 15b2fb0..84f7913 100644 --- a/crates/rutster/src/media_thread.rs +++ b/crates/rutster/src/media_thread.rs @@ -91,6 +91,12 @@ pub enum MediaCmd { }, /// Graceful shutdown — drain + drop + join. Shutdown { reply: oneshot::Sender<()> }, + /// Enter drain (review M1): reject new Registers ("draining" → 503), + /// keep ticking existing sessions, fire `reply` when the map empties. + /// The deadline lives with the CALLER (main.rs times out and proceeds + /// to Shutdown) — the thread just reports emptiness; that keeps + /// "how long is too long" an operator policy, not thread logic. + Drain { reply: oneshot::Sender<()> }, /// Cold-path capacity/health snapshot (readyz + the autoscaling /// signal; review M2). Cheap: counters the loop already maintains. Stats { reply: oneshot::Sender }, @@ -216,6 +222,8 @@ fn run_media_thread( let max_sessions = opts.max_sessions; let mut tick_overruns: u64 = 0; let mut last_tick_micros: u64 = 0; + let mut draining = false; + let mut drain_done: Option> = None; info!("media thread started"); loop { @@ -225,6 +233,14 @@ fn run_media_thread( while let Ok(cmd) = cmd_rx.try_recv() { match cmd { MediaCmd::Register { tap_url, reply } => { + // Drain (review M1) takes priority over the capacity + // check below: once draining, EVERY new Register is + // shed regardless of headroom — we're bleeding out, + // not admitting. + if draining { + let _ = reply.send(Err("draining: not accepting new sessions".into())); + continue; + } // Admission control (review M2): shed the marginal // call with a crisp 503 instead of degrading every // in-flight call's 20ms budget. "node full" is the @@ -303,11 +319,19 @@ fn run_media_thread( let _ = reply.send(MediaStats { sessions: sessions.len(), max_sessions, - draining: false, // Task 4 wires the real flag + draining, tick_overruns, last_tick_micros, }); } + MediaCmd::Drain { reply } => { + draining = true; + if sessions.is_empty() { + let _ = reply.send(()); + } else { + drain_done = Some(reply); + } + } } } @@ -377,6 +401,14 @@ fn run_media_thread( debug!(channel_id = %id, "session evicted after close"); } + // Drain completion check — after eviction so a tick that closes + // the last session completes the drain in the same iteration. + if draining && sessions.is_empty() { + if let Some(done) = drain_done.take() { + let _ = done.send(()); + } + } + // === Step 3: compensated sleep (review M2). === // The old fixed sleep(META_TICK) silently stretched the effective // tick to 10ms + work; under load that pushes every session past @@ -435,6 +467,64 @@ mod tests { thread.shutdown(); } + #[test] + fn drain_rejects_new_sessions_and_completes_when_map_empties() { + let rt = tokio::runtime::Runtime::new().unwrap(); + let handle = rt.handle().clone(); + let url = url::Url::parse("ws://127.0.0.1:8081/echo").unwrap(); + let thread = + MediaThread::spawn(url.clone(), MediaThreadOpts::default(), handle).expect("spawn"); + + // one live session + let (reply, rx) = oneshot::channel(); + thread + .cmd_tx() + .blocking_send(MediaCmd::Register { + tap_url: url::Url::parse("ws://127.0.0.1:1/echo").unwrap(), + reply, + }) + .unwrap(); + let id = rx.blocking_recv().unwrap().unwrap(); + + // drain: must NOT complete while the session lives + let (drain_reply, mut drain_rx) = oneshot::channel(); + thread + .cmd_tx() + .blocking_send(MediaCmd::Drain { reply: drain_reply }) + .unwrap(); + std::thread::sleep(Duration::from_millis(50)); + assert!( + drain_rx.try_recv().is_err(), + "drain must wait for in-flight sessions" + ); + + // new registers are shed with the contract string + let (reply, rx) = oneshot::channel(); + thread + .cmd_tx() + .blocking_send(MediaCmd::Register { + tap_url: url::Url::parse("ws://127.0.0.1:1/echo").unwrap(), + reply, + }) + .unwrap(); + let err = rx.blocking_recv().unwrap().expect_err("draining rejects"); + assert!(err.contains("draining"), "{err}"); + + // deleting the last session completes the drain + let (reply, rx) = oneshot::channel(); + thread + .cmd_tx() + .blocking_send(MediaCmd::Delete { id, reply }) + .unwrap(); + rx.blocking_recv().unwrap(); + assert!( + drain_rx.blocking_recv().is_ok(), + "drain completes once the map is empty" + ); + + thread.shutdown(); + } + #[test] fn media_thread_register_and_shutdown_round_trips() { let rt = tokio::runtime::Runtime::new().unwrap();