diff --git a/CHANGELOG.md b/CHANGELOG.md index 72cc5109..1a368422 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,6 +25,36 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `--experimental-per-shard-rewrite` is deprecated (warns, no-op). ### Fixed +- **Graceful shutdown now drains connections (#438 F1, conn#8).** On + SIGTERM/SIGINT each shard previously tore its runtime down the moment its + event loop observed the cancellation, dropping every pending connection + task mid-poll: in-flight replies were truncated and the blocking/subscriber + shutdown arms (`-ERR server shutting down`) never ran — measured 19/50 + BLPOP-blocked clients losing their shutdown reply on Linux io_uring + (macOS passed only by scheduler luck). Shards now run a bounded drain + before persistence teardown: parked stage-1/2 reads are woken through + their cancellers (re-fired per tick to close the re-park race), token arms + cover the tracking/subscriber/blocking parks, and the loop waits for the + shard's live connection tasks to exit through the normal flush+FIN + epilogue, up to a 5 s ceiling so a wedged peer cannot hold up shutdown. + Writing the drain's red test surfaced a second, tokio-only leak in the + same class: connections accepted by the CENTRAL listener were io-bound to + the MAIN runtime's driver (tokio io resources bind at creation), so their + reads/writes died with main's runtime no matter what the shard drained — + with SO_REUSEPORT splitting accepts roughly evenly, about half of all + tokio connections failed their final writes with "A Tokio 1.x context was + found, but it is being shutdown". Forwarded streams are now re-registered + with the owning shard's runtime driver at spawn (the monoio path already + did this by forwarding std streams). +- **Parked-connection teardown can no longer close a reused fd out from + under a kill scan (#438 F2, conn#9 / sec L1).** `kill_clients`' + fd-liveness invariant (registry deregister strictly before fd close) held + in handler tasks by local-before-parameter drop order, but a task-parked + connection's watcher co-owned guard and stream as future upvars, whose + drop order is merely capture order — and the F1 drain makes dropping that + future a routine path. Both are now wrapped in a `ParkedSession` whose + hand-written `Drop` deregisters before closing, with the invariant + restated at the kill site. - **Remotely-triggerable shard-thread crash on pipelined batch tails (#438 follow-on).** On any `--shards >= 2` deployment, one pipelined write shaped `[…, SUBSCRIBE|BLPOP|…]` aborted the whole server: diff --git a/src/client_registry.rs b/src/client_registry.rs index 50704843..52ed631a 100644 --- a/src/client_registry.rs +++ b/src/client_registry.rs @@ -410,14 +410,23 @@ pub fn client_info(id: u64) -> Option { /// self-kill, then closes). pub fn kill_clients(filter: &KillFilter, self_id: Option) -> u64 { // Lock ordering / fd-liveness invariant, per stripe (c10k W5): the raw-fd - // `shutdown` is race-free because a connection's `RegistryGuard` (a local) - // drops — calling `deregister`, which needs its OWN stripe's WRITE lock — - // strictly before its `stream` (a parameter) drops and closes the fd. So - // while we hold a stripe's READ lock and an entry is present in it, - // `deregister` for that entry is blocked, its `stream` has not dropped, - // and `kill_fd` is still an open socket. The invariant only ever involves - // one entry and its own stripe, so striping preserves it; entries in - // other stripes are simply not visited while their lock is free. + // `shutdown` is race-free because a connection's `RegistryGuard` drops — + // calling `deregister`, which needs its OWN stripe's WRITE lock — + // strictly before its `stream` drops and closes the fd. So while we hold + // a stripe's READ lock and an entry is present in it, `deregister` for + // that entry is blocked, its `stream` has not dropped, and `kill_fd` is + // still an open socket. The invariant only ever involves one entry and + // its own stripe, so striping preserves it; entries in other stripes are + // simply not visited while their lock is free. + // + // WHO upholds guard-before-stream (F2, #438): in a handler task it falls + // out of drop order — the guard is a local, the stream a parameter, and + // locals drop first. A task-parked connection has NO handler task; its + // watcher co-owns both as future upvars (drop order = capture order, no + // language guarantee that helps), so they are wrapped in + // `conn_accept::ParkedSession`, whose hand-written Drop deregisters + // before closing. Any future owner of a {guard, stream} pair must + // preserve this order or wrap in ParkedSession. let mut count = 0u64; let kill_entry = |entry: &ClientEntry, count: &mut u64| { entry.live.kill_flag.store(true, Ordering::Relaxed); diff --git a/src/server/conn/handler_monoio/idle_park.rs b/src/server/conn/handler_monoio/idle_park.rs index 160ec4ad..33f79c3e 100644 --- a/src/server/conn/handler_monoio/idle_park.rs +++ b/src/server/conn/handler_monoio/idle_park.rs @@ -164,6 +164,32 @@ pub(crate) fn sweep(now_ms: u64) -> usize { }) } +/// F1 (#438): cancel EVERY currently-parked stage-1/2 read, regardless of +/// age. Called from the shard event loop's shutdown drain — the parked +/// reads are plain awaits (not selects), so the shutdown token alone cannot +/// wake them; the canceller is the only same-thread wake the design has. +/// The woken handler sees the sweep-cancel error, checks the shutdown token +/// and exits through the normal flush+FIN epilogue instead of re-parking. +/// +/// Re-fired every drain tick: a connection that was mid-batch at the first +/// call parks again only if it raced the token check, and the next tick +/// catches it. Idempotent on unparked slots. +pub(crate) fn cancel_all_parked() -> usize { + REGISTRY.with(|r| { + let mut cancelled = 0usize; + for slot in r.borrow().values() { + if slot.parked_since_ms.get() != 0 { + let old = slot.canceller.replace(Canceller::new()); + let fresh = old.cancel(); + slot.canceller.replace(fresh); + slot.parked_since_ms.set(0); + cancelled += 1; + } + } + cancelled + }) +} + /// Release the parked working set. The rent buffer is dropped outright (the /// pre-park sizing reallocates the probe size next iteration); the scratch /// buffers are only released when empty — a non-empty `read_buf` holds a @@ -355,6 +381,21 @@ mod idle_park_tests { assert_eq!(sweep(300_000 + IDLE_DOWNSHIFT_MS), 1); } + #[test] + fn cancel_all_parked_ignores_age_and_unparked() { + let reg_young = register(90_010); + let reg_old = register(90_011); + let reg_idle = register(90_012); + reg_young.slot.mark_parked(10_000); // just parked — sweep would skip + reg_old.slot.mark_parked_stage2(1); // ancient stage-2 park + // reg_idle: not parked at all. + assert_eq!(cancel_all_parked(), 2, "both parked slots, any age/stage"); + assert_eq!(reg_young.slot.parked_since_ms.get(), 0); + assert_eq!(reg_old.slot.parked_since_ms.get(), 0); + assert_eq!(cancel_all_parked(), 0, "idempotent once unparked"); + drop(reg_idle); + } + #[test] fn registration_drop_deregisters() { { diff --git a/src/server/conn/handler_monoio/mod.rs b/src/server/conn/handler_monoio/mod.rs index 90b40d50..ff34698c 100644 --- a/src/server/conn/handler_monoio/mod.rs +++ b/src/server/conn/handler_monoio/mod.rs @@ -847,6 +847,11 @@ pub(crate) async fn handle_connection_sharded_monoio< let track_buf = std::mem::take(&mut tmp_buf); let mut push_frame: Option = None; monoio::select! { + // F1 (#438): tracking conns park in this select (not the + // cancel-registered reads below), so the shutdown drain + // reaches them via the token. Losing the read future drops + // `track_buf` — acceptable, the connection is exiting. + _ = shutdown.cancelled() => { break; } read_result = stream.read(track_buf) => { let (result, returned_buf) = read_result; tmp_buf = returned_buf; @@ -937,6 +942,11 @@ pub(crate) async fn handle_connection_sharded_monoio< // ECONNRESET, …): terminate. Parking instead would spin // park→wake→park forever — the dead fd stays readable. Err(ref e) if !idle_park::is_sweep_cancel(e) => break, + // F1 (#438): cancelled by the shutdown drain, not the + // stage-2 sweep — exit through the flush+FIN epilogue + // instead of task-parking into a watcher that would just + // be dropped. + Err(_) if shutdown.is_cancelled() => break, Err(_) => { // Cancelled by the stage-2 sweep: exit the task. // read_buf holds at most MAX_PARKED_REMAINDER bytes @@ -997,6 +1007,10 @@ pub(crate) async fn handle_connection_sharded_monoio< // task-park path parks WITHOUT reading, so a mistaken // downshift here would feed the park→wake→park spin.) Err(ref e) if !idle_park::is_sweep_cancel(e) => break, + // F1 (#438): cancelled by the shutdown drain, not the idle + // sweep — exit through the flush+FIN epilogue instead of + // re-parking a read nothing will ever complete. + Err(_) if shutdown.is_cancelled() => break, Err(_) => { // Cancelled by the idle sweep: shed the working set, // re-park small. diff --git a/src/shard/conn_accept.rs b/src/shard/conn_accept.rs index 2df85fb0..e8113cca 100644 --- a/src/shard/conn_accept.rs +++ b/src/shard/conn_accept.rs @@ -40,6 +40,43 @@ type StdRwLock = std::sync::RwLock; /// and release the slot via `record_connection_closed()`. const TLS_HANDSHAKE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); +thread_local! { + /// F1 (#438): connection-scoped tasks alive on THIS shard thread — + /// handler tasks, park watchers, resumed and migrated handlers. Both + /// runtimes pin connection tasks to their shard thread (`monoio::spawn` + /// / `tokio::task::spawn_local`), so a thread-local suffices and each + /// shard's shutdown drain waits only on its own connections. Excludes + /// deliberately short-lived tasks (maxclients-reject writer) and PSYNC + /// hijack tasks (replication links have their own shutdown handling and + /// must not be able to pin a shard's drain past its deadline… they can + /// still eat the bounded deadline — see `SHUTDOWN_DRAIN` in event_loop). + static LIVE_CONN_TASKS: std::cell::Cell = const { std::cell::Cell::new(0) }; +} + +/// Number of live connection tasks on this shard thread (shutdown drain). +pub(crate) fn live_conn_tasks() -> usize { + LIVE_CONN_TASKS.with(|c| c.get()) +} + +/// RAII increment of [`LIVE_CONN_TASKS`]. Created OUTSIDE the spawned +/// future and moved into it, so the count is visible the moment the spawn +/// call returns — a drain that runs between spawn and first poll must see +/// the task, not conclude the shard is idle. +pub(crate) struct ConnTaskGuard(()); + +impl ConnTaskGuard { + pub(crate) fn enter() -> Self { + LIVE_CONN_TASKS.with(|c| c.set(c.get() + 1)); + ConnTaskGuard(()) + } +} + +impl Drop for ConnTaskGuard { + fn drop(&mut self) { + LIVE_CONN_TASKS.with(|c| c.set(c.get().saturating_sub(1))); + } +} + /// Process-wide TCP listen backlog (S-5, `--tcp-backlog`). Set once at /// startup before any listener binds; 1024 is the historical hardcoded /// value. A process-wide static (rather than threading a param through the @@ -164,6 +201,30 @@ pub(crate) fn spawn_tokio_connection( use crate::server::connection::handle_connection_sharded; use crate::server::connection::handle_connection_sharded_inner; + // F1 (#438): re-register the stream with THIS shard's runtime io driver. + // A connection accepted by the central listener arrives bound to the MAIN + // runtime's driver (tokio io resources bind at creation); its io then + // dies the moment main's runtime drops — during shutdown every drained + // reply write on such a connection failed with "A Tokio 1.x context was + // found, but it is being shutdown" while the shard's drain was still + // running, and no shard-side ordering can save a foreign-driver stream. + // (SO_REUSEPORT splits accepts roughly evenly between the central and + // per-shard listeners on Linux, so about half of all connections were + // affected.) The monoio path is immune: it forwards STD streams and + // registers them on the shard's ring right here. `from_std` in this + // function's (shard-thread) runtime context is the tokio equivalent. + let tcp_stream = match tcp_stream + .into_std() + .and_then(tokio::net::TcpStream::from_std) + { + Ok(s) => s, + Err(e) => { + tracing::warn!( + "Shard {shard_id}: shard-driver re-registration failed: {e}; dropping connection" + ); + return; + } + }; let aff = affinity_tracker.clone(); let rsm = remote_subscriber_map.clone(); let sdbs = shard_databases.clone(); @@ -265,7 +326,9 @@ pub(crate) fn spawn_tokio_connection( .map(|a| a.to_string()) .unwrap_or_else(|_| "unknown".to_string()); let maxclients = maxclients_tokio; + let conn_task = ConnTaskGuard::enter(); tokio::task::spawn_local(async move { + let _conn_task = conn_task; // F1: counted until task exit // maxclients check for TLS connections (plain TCP checks in handle_connection_sharded) if !crate::admin::metrics_setup::try_accept_connection(maxclients) { // c10k C3: rate-limited (see the plaintext site below). @@ -334,7 +397,9 @@ pub(crate) fn spawn_tokio_connection( }); } else { // Plain TCP connection + let conn_task = ConnTaskGuard::enter(); tokio::task::spawn_local(async move { + let _conn_task = conn_task; // F1: counted until task exit handle_connection_sharded(tcp_stream, &conn_ctx, sd, cid).await; }); } @@ -517,7 +582,9 @@ pub(crate) fn spawn_migrated_tokio_connection( }; #[cfg(not(unix))] let kill_fd = -1; + let conn_task = ConnTaskGuard::enter(); tokio::task::spawn_local(async move { + let _conn_task = conn_task; // F1: counted until task exit let _ = handle_connection_sharded_inner( tcp_stream, peer_addr, @@ -736,7 +803,9 @@ pub(crate) fn spawn_monoio_connection( if let (true, Some(tls_swap)) = (is_tls, tls_config.as_ref()) { // Load current TLS config from ArcSwap — new connections see reloaded certs let tls_cfg = tls_swap.load_full(); + let conn_task = ConnTaskGuard::enter(); monoio::spawn(async move { + let _conn_task = conn_task; // F1: counted until task exit // maxclients slot already taken by the pre-spawn gate above // (still BEFORE the handshake, preserving #17's stalled- // ClientHello bound); released by record_connection_closed @@ -809,7 +878,9 @@ pub(crate) fn spawn_monoio_connection( let dtx2 = dispatch_tx.clone(); #[cfg(target_os = "linux")] let notifiers2 = all_notifiers.to_vec(); + let conn_task = ConnTaskGuard::enter(); monoio::spawn(async move { + let _conn_task = conn_task; // F1: counted until task exit // maxclients slot already taken by the pre-spawn gate above; // released by the migration-aware decrement at task end. // R-6 fail-open: keep what we need to re-serve the client @@ -1084,6 +1155,56 @@ impl Drop for ParkedCount { } } +/// F2 (#438, conn#9 / sec L1): a parked connection's registry guard and +/// stream, co-owned by one watcher future. `kill_clients`' per-stripe +/// fd-liveness invariant requires the guard (whose drop deregisters under +/// the stripe WRITE lock) to drop STRICTLY BEFORE the stream (whose drop +/// closes the fd): while a kill scan holds a stripe READ lock and sees an +/// entry, that entry's fd must still be an open socket, or the scan's +/// `shutdown(2)` lands on a reused fd of an unrelated connection. Inside a +/// handler task the ordering falls out of locals (guard) dropping before +/// parameters (stream); in this watcher both were co-captured upvars, whose +/// drop order is capture order — never violated in practice, but one +/// refactor away from it, and the F1 shutdown drain makes dropping this +/// future a routine path rather than a teardown-only one. The hand-written +/// Drop makes the order explicit and refactor-proof. +/// Generic over the guard type so the drop-order contract is unit-testable +/// under BOTH runtimes (monoio-only unit tests are invisible to CI — every +/// CI test job runs the tokio feature set). Production use instantiates +/// `G = RegistryGuard`. +/// +/// FIELD ORDER IS THE CONTRACT: struct fields drop in declaration order +/// (RFC 1857), so `guard` (deregister, serializing against any in-flight +/// kill scan on this entry's stripe) drops strictly before `stream` (fd +/// close). Do not reorder these fields — the +/// `parked_session_drops_guard_before_stream` unit test pins the order, +/// and it lets the wake path destructure without `Option`s or a manual +/// `Drop` (which would forbid moving the parts back out). +pub(crate) struct ParkedSession { + guard: G, + stream: S, +} + +#[cfg_attr(not(all(feature = "runtime-monoio", unix)), allow(dead_code))] +impl ParkedSession { + pub(crate) fn new(guard: G, stream: S) -> Self { + ParkedSession { guard, stream } + } + + /// Wake path: hand both back to the resumed handler, where the + /// local-before-parameter drop ordering holds again. + pub(crate) fn into_parts(self) -> (G, S) { + (self.guard, self.stream) + } +} + +#[cfg(all(feature = "runtime-monoio", unix))] +impl ParkedSession { + fn park_readable(&self) -> impl std::future::Future> { + self.stream.park_readable() + } +} + #[cfg(all(feature = "runtime-monoio", unix))] fn spawn_parked_idle_watcher( stream: S, @@ -1100,13 +1221,19 @@ fn spawn_parked_idle_watcher( + ParkWatchable + 'static, { + // F1: counted so the shutdown drain waits for this watcher to run its + // shutdown arm (clean close + accounting) instead of dropping it cold. + let conn_task = ConnTaskGuard::enter(); + // F2: guard-first drop order, whichever way this future ends. + let session = ParkedSession::new(registry_guard, stream); let fut: std::pin::Pin>> = Box::pin(async move { + let _conn_task = conn_task; // Observability for `INFO clients.parked_clients`. RAII so the gauge // stays accurate even when this future is dropped outright (runtime // teardown), where neither select arm runs. let _parked = ParkedCount::enter(); let woke = monoio::select! { - res = stream.park_readable() => { + res = session.park_readable() => { // Err (fd error) also resumes: the handler's first read // surfaces the real error and tears down cleanly. let _ = res; @@ -1116,8 +1243,7 @@ fn spawn_parked_idle_watcher( }; if !woke { // Server shutdown: this watcher owns the close accounting. - drop(registry_guard); - drop(stream); + drop(session); crate::admin::metrics_setup::record_connection_closed(); return; } @@ -1128,8 +1254,7 @@ fn spawn_parked_idle_watcher( // fleet-wide `timeout` expiry wakes them in a burst. Close here // instead — this watcher already owns the close accounting. if crate::client_registry::is_killed(client_id) { - drop(registry_guard); - drop(stream); + drop(session); crate::admin::metrics_setup::record_connection_closed(); return; } @@ -1137,6 +1262,7 @@ fn spawn_parked_idle_watcher( // handoff) — the entry is never dropped, so there is no window where // a racing CLIENT LIST misses the conn or CLIENT KILL ID returns 0, // and name/connected_at/kill state persist across the wake. + let (registry_guard, stream) = session.into_parts(); spawn_resumed_parked_conn( stream, state, @@ -1172,7 +1298,9 @@ fn spawn_resumed_parked_conn( + 'static, { use crate::server::connection::handle_connection_sharded_monoio; + let conn_task = ConnTaskGuard::enter(); monoio::spawn(async move { + let _conn_task = conn_task; // F1: counted until task exit let initial = std::mem::take(&mut state.read_buf_remainder); let peer = state.peer_addr.clone(); let (outcome, stream_back) = handle_connection_sharded_monoio( @@ -1397,7 +1525,9 @@ pub(crate) fn spawn_migrated_monoio_connection( }; #[cfg(not(unix))] let kill_fd = -1; + let conn_task = ConnTaskGuard::enter(); monoio::spawn(async move { + let _conn_task = conn_task; // F1: counted until task exit // c1M P1: kept for the ParkIdle routing below (`sd` moves // into the handler call). let sd_park = sd.clone(); @@ -1456,3 +1586,67 @@ pub(crate) fn spawn_migrated_monoio_connection( } } } + +#[cfg(test)] +mod conn_accept_tests { + use super::*; + use std::cell::RefCell; + use std::rc::Rc; + + /// F2 (#438): the whole point of ParkedSession — guard drops strictly + /// before stream, however the session ends. + struct DropRecorder(&'static str, Rc>>); + impl Drop for DropRecorder { + fn drop(&mut self) { + self.1.borrow_mut().push(self.0); + } + } + + #[test] + fn parked_session_drops_guard_before_stream() { + let order = Rc::new(RefCell::new(Vec::new())); + let session = ParkedSession::new( + DropRecorder("guard", order.clone()), + DropRecorder("stream", order.clone()), + ); + drop(session); + assert_eq!( + &*order.borrow(), + &["guard", "stream"], + "kill_clients fd-liveness invariant: deregister before fd close" + ); + } + + #[test] + fn parked_session_into_parts_transfers_without_dropping() { + let order = Rc::new(RefCell::new(Vec::new())); + let session = ParkedSession::new( + DropRecorder("guard", order.clone()), + DropRecorder("stream", order.clone()), + ); + let (guard, stream) = session.into_parts(); + assert!( + order.borrow().is_empty(), + "into_parts must not drop either part (registration handoff)" + ); + // The caller re-establishes local-before-parameter ordering; here we + // just verify both came back live. + drop(guard); + drop(stream); + assert_eq!(&*order.borrow(), &["guard", "stream"]); + } + + /// F1 (#438): the drain's stop condition — live count tracks guard + /// lifetimes, visible immediately at enter (pre-spawn). + #[test] + fn conn_task_guard_counts_immediately_and_saturates() { + let base = live_conn_tasks(); + let g1 = ConnTaskGuard::enter(); + let g2 = ConnTaskGuard::enter(); + assert_eq!(live_conn_tasks(), base + 2); + drop(g1); + assert_eq!(live_conn_tasks(), base + 1); + drop(g2); + assert_eq!(live_conn_tasks(), base); + } +} diff --git a/src/shard/event_loop.rs b/src/shard/event_loop.rs index 336105d6..0e54fba1 100644 --- a/src/shard/event_loop.rs +++ b/src/shard/event_loop.rs @@ -36,6 +36,15 @@ use super::shared_databases::ShardDatabases; use super::uring_handler; use super::{conn_accept, persistence_tick, spsc_handler, timers}; +/// F1 (#438): ceiling on the graceful-shutdown connection drain. Normal +/// drains finish in single-digit milliseconds (parked reads wake instantly, +/// in-flight batches are already bounded by the C1 write timeout); the +/// ceiling exists so a wedged peer — zero-window reader with +/// `client_write_timeout_ms 0`, or a hung cross-shard leg — cannot hold up +/// process shutdown. On expiry the remaining tasks are dropped, which is +/// exactly the pre-F1 behaviour for every task. +const SHUTDOWN_DRAIN_MAX: Duration = Duration::from_secs(5); + /// c10k hardening B4: may the experimental io_uring bridge be armed? /// /// The bridge binds a SECOND `SO_REUSEPORT` listener on the server's own port @@ -1932,6 +1941,33 @@ impl super::Shard { } _ = shutdown.cancelled() => { info!("Shard {} shutting down", self.id); + // F1 (#438): bounded connection drain BEFORE persistence + // teardown — the `break` below returns from `run`, and + // dropping the LocalSet/runtime kills every connection + // task still pending, truncating in-flight replies and + // skipping the blocking/subscriber shutdown arms. The + // token (already cancelled) wakes the tokio selects' + // shutdown arms; this loop just keeps the thread polling + // until they have all exited through the flush+FIN + // epilogue, or the deadline expires (a wedged peer must + // not hold up shutdown — its task is then dropped, the + // pre-F1 behaviour). + { + let drain_deadline = std::time::Instant::now() + SHUTDOWN_DRAIN_MAX; + loop { + let live = conn_accept::live_conn_tasks(); + if live == 0 { + break; + } + if std::time::Instant::now() >= drain_deadline { + tracing::warn!( + "Shard {shard_id}: shutdown drain timed out with {live} connection task(s) still live; dropping them" + ); + break; + } + tokio::time::sleep(std::time::Duration::from_millis(2)).await; + } + } persistence_tick::drain_and_shutdown_spill( &mut spill_thread, &mut shard_manifest, @@ -2064,6 +2100,51 @@ impl super::Shard { // Check shutdown before awaiting (non-blocking) if shutdown.is_cancelled() { info!("Shard {} shutting down (monoio)", self.id); + // F1 (#438): bounded connection drain BEFORE persistence + // teardown — the `break` below returns from `run`, and + // dropping the monoio runtime kills every connection task + // still pending, truncating in-flight replies and + // skipping the blocking/subscriber shutdown arms (found + // live: 19/50 BLPOP clients lost their shutdown reply on + // Linux io_uring). Stage-1/2 idle-park reads are plain + // awaits the token cannot wake, so fire their cancellers; + // the woken handlers see the cancelled token and exit + // through the flush+FIN epilogue. Re-fired every tick to + // close the mid-batch re-park race. Deadline-bounded: a + // wedged peer must not hold up shutdown — its task is + // then dropped, the pre-F1 behaviour. + { + let drain_deadline = std::time::Instant::now() + SHUTDOWN_DRAIN_MAX; + let mut ticks = 0u32; + loop { + crate::server::conn::handler_monoio::idle_park::cancel_all_parked(); + let live = conn_accept::live_conn_tasks(); + if live == 0 { + break; + } + if std::time::Instant::now() >= drain_deadline { + tracing::warn!( + "Shard {shard_id}: shutdown drain timed out with {live} connection task(s) still live; dropping them" + ); + break; + } + // Fast ticks catch the one legal re-park race (a + // task that passed its post-batch shutdown check + // before the token cancelled, then parked after + // the first canceller sweep); after that no task + // can park again, so back off — the O(registry) + // canceller scan every 2 ms for the full 5 s + // ceiling would be a shutdown-only CPU spike at + // high connection counts. + let tick = if ticks < 5 { + std::time::Duration::from_millis(2) + } else { + std::time::Duration::from_millis(50) + }; + ticks += 1; + monoio::time::sleep(tick).await; + } + } persistence_tick::drain_and_shutdown_spill( &mut spill_thread, &mut shard_manifest, diff --git a/tests/shutdown_drain.rs b/tests/shutdown_drain.rs new file mode 100644 index 00000000..9948e7b3 --- /dev/null +++ b/tests/shutdown_drain.rs @@ -0,0 +1,255 @@ +//! F1 (#438, conn#8): graceful shutdown must drain connection tasks. +//! +//! Before the fix, the shard event loop's `shutdown.cancelled()` arm ran its +//! persistence teardown and immediately `break` — `run` returned, the monoio +//! runtime was dropped, and every spawned connection task was dropped +//! mid-poll. A client blocked in BLPOP never received the +//! `-ERR server shutting down` reply its shutdown arm exists to send; a +//! client with an in-flight batch could see its reply truncated. Clients saw +//! a bare socket close instead of reply-then-FIN. +//! +//! The fix adds a bounded drain phase to the shutdown arm: fire the idle-park +//! cancellers so stage-1/2 parked reads wake (they check the shutdown token +//! and exit through the normal flush+FIN epilogue), then keep polling the +//! runtime until the shard's live connection tasks reach zero or the drain +//! deadline expires, and only then tear down persistence and drop the +//! runtime. +//! +//! All tests are `#[cfg(unix)]`: they signal the server with SIGTERM. + +#![cfg(unix)] + +mod common; + +use std::io::{Read, Write}; +use std::net::TcpStream; +use std::process::{Child, Command, Stdio}; +use std::time::{Duration, Instant}; + +fn moon_binary() -> Option { + if let Ok(p) = std::env::var("MOON_BIN") { + return Some(std::path::PathBuf::from(p)); + } + let cargo_bin = std::path::PathBuf::from(env!("CARGO_BIN_EXE_moon")); + if cargo_bin.exists() { + return Some(cargo_bin); + } + None +} + +struct Moon { + child: Child, + port: u16, + _tmp_dir: tempfile::TempDir, +} + +impl Drop for Moon { + fn drop(&mut self) { + let _ = self.child.kill(); + let _ = self.child.wait(); + } +} + +fn spawn_moon(shards: &str) -> Option { + let bin = moon_binary()?; + let tmp_dir = tempfile::tempdir().expect("tempdir"); + let dir_str = tmp_dir.path().to_str().unwrap().to_string(); + let shards = shards.to_string(); + let (child, port) = common::spawn_listening(|port| { + Command::new(&bin) + .args([ + "--port", + &port.to_string(), + "--shards", + &shards, + "--admin-port", + "0", + "--appendonly", + "no", + "--disk-free-min-pct", + "0", + "--dir", + &dir_str, + ]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .expect("spawn moon") + }); + let moon = Moon { + child, + port, + _tmp_dir: tmp_dir, + }; + let deadline = Instant::now() + Duration::from_secs(10); + while Instant::now() < deadline { + if let Ok(mut c) = TcpStream::connect(("127.0.0.1", moon.port)) { + let _ = c.set_read_timeout(Some(Duration::from_millis(500))); + if c.write_all(b"*1\r\n$4\r\nPING\r\n").is_ok() { + let mut buf = [0u8; 64]; + if let Ok(n) = c.read(&mut buf) + && n > 0 + && buf.starts_with(b"+PONG") + { + return Some(moon); + } + } + } + std::thread::sleep(Duration::from_millis(100)); + } + eprintln!("skipping: moon did not become ready on port {}", moon.port); + None +} + +fn sigterm(moon: &Moon) { + let status = Command::new("kill") + .args(["-TERM", &moon.child.id().to_string()]) + .status() + .expect("send SIGTERM"); + assert!(status.success(), "kill -TERM failed"); +} + +/// Wait for the server process to exit, panicking past `secs`. Every test +/// calls this so a drain phase that hangs shutdown is itself a failure. +fn await_exit(moon: &mut Moon, secs: u64) { + let deadline = Instant::now() + Duration::from_secs(secs); + loop { + match moon.child.try_wait().expect("try_wait") { + Some(_) => return, + None if Instant::now() >= deadline => { + panic!("server did not exit within {secs}s of SIGTERM (drain must be bounded)") + } + None => std::thread::sleep(Duration::from_millis(50)), + } + } +} + +/// Read everything until EOF (or panic on timeout). Returns the bytes the +/// server flushed before closing. +fn read_to_eof(stream: &mut TcpStream, timeout: Duration) -> Vec { + stream.set_read_timeout(Some(timeout)).unwrap(); + let mut out = Vec::new(); + let mut chunk = [0u8; 4096]; + loop { + match stream.read(&mut chunk) { + Ok(0) => return out, + Ok(n) => out.extend_from_slice(&chunk[..n]), + Err(e) => panic!( + "no EOF from server within {timeout:?} (got {} bytes so far: {:?}): {e}", + out.len(), + String::from_utf8_lossy(&out) + ), + } + } +} + +/// Clients blocked in BLPOP must receive `-ERR server shutting down` and a +/// clean FIN when the server is SIGTERMed — their handler tasks have a +/// shutdown arm that sends exactly that reply, but pre-fix the tasks were +/// dropped unpolled when the shard runtime was torn down. +/// +/// Many connections, ALL must get the reply: a single conn can win the +/// teardown race by scheduling luck (macOS/kqueue reliably does; Linux +/// io_uring measured 31/50 pre-fix), so one-conn asserts under-test. With 50 +/// conns the pre-fix loss is statistically certain on the losing platforms +/// and the post-fix drain must be exhaustive anyway. +fn blocked_clients_drain(shards: &str) { + let Some(mut moon) = spawn_moon(shards) else { + return; + }; + const N: usize = 50; + let mut conns = Vec::with_capacity(N); + for i in 0..N { + let mut c = TcpStream::connect(("127.0.0.1", moon.port)).expect("connect"); + let key = format!("drainkey{i}"); + let cmd = format!( + "*3\r\n$5\r\nBLPOP\r\n${}\r\n{key}\r\n$1\r\n0\r\n", + key.len() + ); + c.write_all(cmd.as_bytes()).expect("send BLPOP"); + conns.push(c); + } + // Let the waits register (including the remote leg at shards>=2). + std::thread::sleep(Duration::from_millis(600)); + sigterm(&moon); + let mut lost = Vec::new(); + for (i, mut c) in conns.into_iter().enumerate() { + let bytes = read_to_eof(&mut c, Duration::from_secs(8)); + let text = String::from_utf8_lossy(&bytes).into_owned(); + let ok = text.contains("shutting down") + || text.starts_with("*-1\r\n") + || text.starts_with("$-1"); + if !ok { + lost.push((i, text)); + } + } + assert!( + lost.is_empty(), + "{}/{N} blocked clients lost their shutdown reply (task dropped before its shutdown arm ran); first: {:?}", + lost.len(), + lost.first() + ); + await_exit(&mut moon, 15); +} + +#[test] +fn blocked_clients_get_shutdown_reply_shards1() { + blocked_clients_drain("1"); +} + +#[test] +fn blocked_clients_get_shutdown_reply_shards2() { + blocked_clients_drain("2"); +} + +/// An idle connection parked in a stage-1/2 read must see a clean FIN (read +/// == 0, no ECONNRESET) shortly after SIGTERM, and the server must still +/// exit promptly — the drain phase is bounded. +#[test] +fn parked_idle_conn_clean_fin_on_sigterm() { + let Some(mut moon) = spawn_moon("1") else { + return; + }; + let mut idle = TcpStream::connect(("127.0.0.1", moon.port)).expect("connect"); + idle.write_all(b"*1\r\n$4\r\nPING\r\n").expect("ping"); + let mut buf = [0u8; 32]; + idle.set_read_timeout(Some(Duration::from_secs(2))).unwrap(); + let n = idle.read(&mut buf).expect("pong"); + assert!(buf[..n].starts_with(b"+PONG")); + // Sit idle past the 1 s stage-1 sweep so the read is parked/downshifted. + std::thread::sleep(Duration::from_millis(1600)); + sigterm(&moon); + // Runtime-behavior split: monoio's cancelled park exits with a bare FIN; + // the tokio main select's shutdown arm sends `-ERR server shutting down` + // first. Both are clean drains — what must NOT happen is a reset or a + // hang (read_to_eof panics on timeout, await_exit bounds the process). + let bytes = read_to_eof(&mut idle, Duration::from_secs(5)); + let text = String::from_utf8_lossy(&bytes); + assert!( + bytes.is_empty() || text.contains("shutting down"), + "idle conn should see a bare FIN or the shutdown error, got: {text:?}" + ); + await_exit(&mut moon, 10); +} + +/// A subscriber must also drain: its select loop has a shutdown arm today, +/// but pre-fix the task was never polled after cancellation. After the fix +/// the subscriber's socket closes with a clean FIN and the process exits. +#[test] +fn subscriber_conn_clean_fin_on_sigterm() { + let Some(mut moon) = spawn_moon("2") else { + return; + }; + let mut sub = TcpStream::connect(("127.0.0.1", moon.port)).expect("connect"); + sub.write_all(b"*2\r\n$9\r\nSUBSCRIBE\r\n$7\r\ndrainch\r\n") + .expect("subscribe"); + let mut buf = [0u8; 128]; + sub.set_read_timeout(Some(Duration::from_secs(2))).unwrap(); + let n = sub.read(&mut buf).expect("subscribe reply"); + assert!(n > 0, "subscribe must be acknowledged"); + std::thread::sleep(Duration::from_millis(300)); + sigterm(&moon); + // Everything already acknowledged; the drain just needs to close cleanly. + let _ = read_to_eof(&mut sub, Duration::from_secs(5)); + await_exit(&mut moon, 10); +}