From c7c4f0bd9446187272ddbe7967889a95a2f48729 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Fri, 7 Aug 2026 20:18:34 +0700 Subject: [PATCH 1/2] fix(server): accept-loop HOL block, migration fd ownership, park auth gates (#438 F3-F6 + conn-secondary) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit F3 (routing F10) — the central accept loop delivered every connection with a blocking send on the routed shard's bounded (4096) conn channel: one wedged shard froze accepts for ALL shards. The monoio legs were worse than reviewed — `flume::Sender::send` there is synchronous, so a full channel stalled the entire listener THREAD (plain + TLS legs alike). All five delivery sites (tokio plain/TLS, monoio plain/TLS x2 select shapes) now `try_send` and rotate to the next shard with room (`try_route_conn`); the affinity hint is best-effort and any live shard serves the client. Only when EVERY channel refuses does the loop fall back to the pre-F3 blocking send — server-wide saturation, where back-pressure is the correct behaviour — with a warn. Unit-tested (rotation, disconnected-skip, all-full payload return, target-first). F4 (routing F12) — MigrateConnectionPayload carried the migrated socket as a raw i32 with no Drop: a migration message still queued in an SPSC ring at shutdown, or drained into pending_migrations but never spawned, leaked the fd and stranded the client on a connection no task would ever serve (silent, until the fd table filled). The payload now owns the socket as an OwnedFd end to end — producer (tokio into_std; monoio dup wrapped at birth) → ring → pending_migrations → spawn_migrated_* — so every undelivered path closes the socket and the client gets a FIN. This also discharges the F1-drain interaction flagged on PR #445: queued migrations dropped by a shutting-down shard now close deterministically. Three `unsafe from_raw_fd` blocks became safe ownership conversions (net -3 unsafe). Unit test pins the drop semantics; a revert to a raw fd no longer compiles. Resumed-parked conns stay deliberately can_migrate:false — the resume wrapper has no MigrateConnection arm, and migration sampling exists for HOT conns while this path is only reachable by idling past --conn-park-secs; rationale documented at the spawn site. F5 (sec L3) — both migrated-connection spawn sites built the target shard's ConnectionContext with requirepass: None. Session auth was unaffected (it travels in MigratedConnectionState), but AUTH issued on a migrated conn wrongly answered "no password is set", and any future auth-re-derivation from the context would have failed open. The real requirepass is now read from runtime config at spawn. F6 (sec L2) — unauthenticated connections could task-park: on an auth-enabled server an attacker could open maxclients silent sockets, never AUTH, and hold each at the parked ~3.3 KB watcher footprint indefinitely and invisibly. The stage-2 parkable predicate now requires conn.authenticated — pre-AUTH conns keep their full handler task (costly enough to surface in monitoring, still reaped by `timeout N`); no-auth servers see no change (authenticated=true from accept). Red/green on macOS monoio: parked_clients=2 pre-fix, =1 post-fix (authed idle sibling doubles as the positive control that the park machinery stayed live). conn-secondary — * is_sweep_cancel matched any bare errno 125, so an ECANCELED from a non-sweep source read as "sweep cancel -> re-park": for a dead fd that is a permanent park->wake->park spin at 100% CPU (the exact E11 failure class the errno check was built to prevent). The idle-park slot now records provenance (`swept`, set by sweep/drain before firing, re-armed false at every park) and the park arms require it via consume-once `was_swept_cancel`; blocking.rs keeps the plain errno check where a stray 125 only re-arms a level-triggered poll (no spin possible). Unit-tested (provenance required, consume-once, stale-flag cleared on re-park, drain marks too). * client_registry::register double-counted TOTAL_CLIENTS and the shard gauges when inserting over a still-present id (the kept_registration miss-arm fail-safe racing an entry back into existence): permanent +1 skew feeding shard_overloaded routing. register() now balances against the replaced entry, and the miss arm (an invariant violation, not a code path) warns loudly. Unit-tested. * `timeout` read-once: already fixed by the D1 chore-sweep rework — both shard chores re-read runtime_config.timeout every second and CONFIG SET timeout is live (verified, comment at the monoio chore says exactly this). No change. Gates: fmt + clippy -D warnings (default and tokio,jemalloc), new unit tests green under both feature sets, parked_idle_parity 7/7 + shutdown_drain 4/4 + tls_park_keyupdate on macOS monoio AND tokio (F6 leg degenerates to parked=0 under tokio as documented), full suites on macOS monoio (28 suites; one cross_shard_consistency_red load flake, solo-green — same documented family as cdg6e), Linux VM monoio (lib 4519/4519; client_tracking_invalidation push-delivery flake fires on UNMODIFIED main too — 3x-vs-3x merge-base A/B, 1/3 main vs 2/3 branch, same missing-second-key signature; pre-existing, queued for the test-hygiene sweep) and Linux VM tokio CI-parity (exit 0, 190 suites). Bench gates waived — no hot-path cost: try_route_conn replaces a blocking send with try_send attempts on the accept path (not the command path); the OwnedFd is layout-identical to i32 with a Drop that only runs on the cold migration/teardown path; the park-arm provenance check is two Cell ops on the already-cold cancelled-read path; the parkable predicate gains one bool test evaluated once per park decision. Refs #438 (F3 = routing F10, F4 = routing F12, F5 = sec L3, F6 = sec L2, conn-secondary) author: Tin Dang --- CHANGELOG.md | 45 +++++++ src/client_registry.rs | 47 ++++++- src/server/conn/handler_monoio/idle_park.rs | 69 ++++++++++ src/server/conn/handler_monoio/mod.rs | 27 +++- src/server/conn/handler_sharded/mod.rs | 25 ++-- src/server/listener.rs | 140 +++++++++++++++++--- src/shard/conn_accept.rs | 116 ++++++++-------- src/shard/dispatch.rs | 64 ++++++++- src/shard/event_loop.rs | 11 +- src/shard/spsc_handler.rs | 3 +- tests/parked_idle_parity.rs | 65 ++++++++- 11 files changed, 520 insertions(+), 92 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1a368422..9488c0fe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,6 +25,51 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `--experimental-per-shard-rewrite` is deprecated (warns, no-op). ### Fixed +- **Central accept loop no longer head-of-line blocks on one wedged shard + (#438 F3).** Every central-listener delivery (tokio plain/TLS, monoio + plain/TLS — the monoio sends were *synchronous*, stalling the whole + listener thread) previously blocked on the routed shard's bounded conn + channel; one shard wedged at its 4096-conn cap froze accepts for every + other shard. Deliveries now `try_send` and rotate to the next shard with + room; only when every shard's channel is full does the loop fall back to + the blocking send (server-wide saturation, where back-pressure is + correct). +- **Connection-migration fds can no longer leak (#438 F4).** + `MigrateConnectionPayload` carried the socket as a raw `i32` with no drop + semantics: a migration message still queued when a shard shut down — or + drained but never spawned — leaked the fd and stranded the client on a + connection no task would ever serve. The payload now owns the socket as an + `OwnedFd` end to end (producer → SPSC ring → pending-migrations queue → + target-shard spawn), so every undelivered path closes the socket and the + client sees a FIN. Three `unsafe from_raw_fd` blocks became safe + ownership conversions in the process. (Resumed-parked connections remain + deliberately `can_migrate:false`; rationale documented at the spawn site.) +- **Migrated connections now carry the real `requirepass` (#438 F5, sec + L3).** The target-shard `ConnectionContext` was built with + `requirepass: None`; the session's auth state was unaffected (it travels + in `MigratedConnectionState`), but a later `AUTH` on a migrated + connection wrongly answered "no password is set", and any future code + deriving auth from the context would have failed open. +- **Unauthenticated connections never task-park (#438 F6, sec L2).** On an + auth-enabled server, a client could open sockets and never authenticate + nor speak; each one downshifted and task-parked into the ~3.3 KB watcher + state, letting an attacker hold a maxclients-worth of silent connections + indefinitely at near-zero cost. Pre-AUTH connections now stay un-parked + (full handler task — visible in monitoring, still reaped by `timeout N`); + no-auth servers are unaffected. +- **Idle-park cancel provenance + registry counter hygiene (#438 + conn-secondary).** A bare `ECANCELED` (errno 125) from any non-sweep + source was indistinguishable from the idle sweep's cancel and re-parked + the connection — for a dead fd that is a permanent park→wake→park spin at + 100% CPU. Park arms now require the sweep/drain to have marked the slot + (`was_swept_cancel`, consume-once) before treating 125 as a park signal. + Separately, a re-register over a still-present client id double-counted + `TOTAL_CLIENTS`/shard gauges permanently (skewing `shard_overloaded` + routing); `register` now balances against the replaced entry and the + `kept_registration` miss-arm logs the invariant violation loudly. (The + third secondary finding — `timeout` read once at connection setup — was + already fixed by the D1 chore-sweep rework, which re-reads the config + every second.) - **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 diff --git a/src/client_registry.rs b/src/client_registry.rs index 52ed631a..f8c8a5af 100644 --- a/src/client_registry.rs +++ b/src/client_registry.rs @@ -312,8 +312,25 @@ pub fn register( shard, live: Arc::clone(&live), }; - stripe(id).write().insert(id, entry); - TOTAL_CLIENTS.fetch_add(1, Ordering::Relaxed); + // #438 conn-secondary: a re-register over a still-present id (the + // `kept_registration` miss-arm fail-safe firing while the entry somehow + // survived, or any future double-register bug) must not double-count — + // the single eventual deregister would then leave TOTAL_CLIENTS and the + // shard gauges permanently high, skewing `shard_overloaded` routing. + // Balance against whatever the insert replaced. + let replaced = stripe(id).write().insert(id, entry); + if let Some(old) = replaced { + tracing::warn!( + "client registry: id {} re-registered while present (shard {} -> {})", + id, + old.shard, + shard + ); + shard_conns_delta(old.shard, false); + crate::admin::metrics_setup::record_shard_connection_delta(old.shard, -1.0); + } else { + TOTAL_CLIENTS.fetch_add(1, Ordering::Relaxed); + } shard_conns_delta(shard, true); crate::admin::metrics_setup::record_shard_connection_delta(shard, 1.0); live @@ -1282,4 +1299,30 @@ mod idle_timeout_tests { parked_delta(false); assert_eq!(parked_clients(), before); } + + /// #438 conn-secondary: a re-register of a still-present id (the + /// kept_registration miss-arm fail-safe racing an entry back into + /// existence) must NOT increment TOTAL_CLIENTS again — the single + /// eventual deregister would then leave the count permanently +1, + /// skewing `shard_overloaded` routing. (TOTAL_CLIENTS is process-global + /// and other tests register concurrently, so the assert brackets only + /// the one replacing insert — the narrowest possible window.) + #[test] + fn replacing_register_does_not_double_count() { + use std::sync::atomic::Ordering; + const ID: u64 = 9_060; + let _a = register(ID, "t:1".into(), "default".into(), 908, -1); + let t1 = TOTAL_CLIENTS.load(Ordering::Relaxed); + let _b = register(ID, "t:1".into(), "default".into(), 909, -1); + let t2 = TOTAL_CLIENTS.load(Ordering::Relaxed); + // Pre-fix, the unconditional fetch_add makes this read t1+1 even + // with zero concurrency. (An unrelated test registering inside this + // two-instruction window could false-fail; the window is nanoseconds + // and the suite's other global-counter tests accept the same odds.) + assert_eq!(t2, t1, "replacing insert must not double-count"); + // The entry itself was replaced, not duplicated, and one deregister + // fully clears it. + deregister(ID); + assert!(live_handle(ID).is_none(), "single deregister must clear"); + } } diff --git a/src/server/conn/handler_monoio/idle_park.rs b/src/server/conn/handler_monoio/idle_park.rs index 33f79c3e..2de22975 100644 --- a/src/server/conn/handler_monoio/idle_park.rs +++ b/src/server/conn/handler_monoio/idle_park.rs @@ -66,6 +66,12 @@ pub(super) struct IdleSlot { /// read — the sweep then applies [`park_after_ms`] instead of /// [`IDLE_DOWNSHIFT_MS`], and the woken handler exits its task. stage2: Cell, + /// #438 conn-secondary: provenance for the cancel errno. Set by + /// [`sweep`]/[`cancel_all_parked`] when THEY fire this slot's canceller, + /// consumed by the handler's error check ([`IdleSlot::was_swept_cancel`]). + /// A bare `ECANCELED` (125) that no sweep produced is a REAL error and + /// must terminate the connection, not park it. + swept: Cell, } impl IdleSlot { @@ -78,6 +84,7 @@ impl IdleSlot { /// cached clock; clamped to ≥1 so 0 stays the "not parked" sentinel. pub(super) fn mark_parked(&self, now_ms: u64) { self.stage2.set(false); + self.swept.set(false); self.parked_since_ms.set(now_ms.max(1)); } @@ -85,6 +92,7 @@ impl IdleSlot { /// verified parkable, cancel = task exit. pub(super) fn mark_parked_stage2(&self, now_ms: u64) { self.stage2.set(true); + self.swept.set(false); self.parked_since_ms.set(now_ms.max(1)); } @@ -92,6 +100,18 @@ impl IdleSlot { pub(super) fn mark_unparked(&self) { self.parked_since_ms.set(0); } + + /// #438 conn-secondary: was this error produced by a sweep/drain cancel + /// of THIS slot? Checks the errno shape AND consumes the provenance flag + /// the cancelling chore set. A raw 125 without the flag (some non-sweep + /// `ECANCELED` source) reads as a real error — the caller terminates + /// instead of parking a connection nothing will ever cancel-wake again. + /// The flag is re-armed `false` at every park, so a raced data delivery + /// (cancel fired, bytes won) cannot leave a stale `true` behind past the + /// next park. + pub(super) fn was_swept_cancel(&self, e: &std::io::Error) -> bool { + is_sweep_cancel(e) && self.swept.replace(false) + } } thread_local! { @@ -118,6 +138,7 @@ pub(super) fn register(client_id: u64) -> IdleParkRegistration { canceller: RefCell::new(Canceller::new()), parked_since_ms: Cell::new(0), stage2: Cell::new(false), + swept: Cell::new(false), }); REGISTRY.with(|r| r.borrow_mut().insert(client_id, slot.clone())); IdleParkRegistration { slot, client_id } @@ -148,6 +169,9 @@ pub(crate) fn sweep(now_ms: u64) -> usize { IDLE_DOWNSHIFT_MS }; if parked != 0 && now_ms.saturating_sub(parked) >= threshold { + // Provenance BEFORE the cancel fires: the woken handler may + // check `was_swept_cancel` on this same thread's next poll. + slot.swept.set(true); // Consume-and-replace: cancel() fires every associated op and // returns a fresh canceller for the connection's next park. let old = slot.canceller.replace(Canceller::new()); @@ -179,6 +203,10 @@ pub(crate) fn cancel_all_parked() -> usize { let mut cancelled = 0usize; for slot in r.borrow().values() { if slot.parked_since_ms.get() != 0 { + // Same provenance mark as `sweep` — the drained handler's + // shutdown-token arm exits either way, but the errno check + // runs first and must classify this 125 as sweep-produced. + slot.swept.set(true); let old = slot.canceller.replace(Canceller::new()); let fresh = old.cancel(); slot.canceller.replace(fresh); @@ -396,6 +424,47 @@ mod idle_park_tests { drop(reg_idle); } + /// #438 conn-secondary: a bare errno 125 the sweep did not produce must + /// read as a REAL error; a swept cancel reads as sweep exactly once. + #[test] + fn was_swept_cancel_requires_sweep_provenance() { + let e125 = std::io::Error::from_raw_os_error(125); + let reg = register(90_020); + reg.slot.mark_parked(10_000); + // No sweep ran: a raw ECANCELED is NOT a sweep cancel. + assert!(!reg.slot.was_swept_cancel(&e125)); + // Sweep fires → provenance set → consumed exactly once. + reg.slot.mark_parked(20_000); + assert_eq!(sweep(20_000 + IDLE_DOWNSHIFT_MS), 1); + assert!(reg.slot.was_swept_cancel(&e125)); + assert!(!reg.slot.was_swept_cancel(&e125), "flag is consume-once"); + // Non-125 errors never match even right after a sweep. + reg.slot.mark_parked(40_000); + assert_eq!(sweep(40_000 + IDLE_DOWNSHIFT_MS), 1); + let reset = std::io::Error::from(std::io::ErrorKind::ConnectionReset); + assert!(!reg.slot.was_swept_cancel(&reset)); + } + + /// A sweep whose cancel raced a data delivery leaves the flag set; the + /// next park must re-arm it to false so it cannot mask a later real 125. + #[test] + fn stale_sweep_flag_cleared_on_next_park() { + let e125 = std::io::Error::from_raw_os_error(125); + let reg = register(90_021); + reg.slot.mark_parked_stage2(10_000); + assert_eq!(sweep(10_000 + park_after_ms()), 1); + // Handler woke with DATA (raced) — flag never consumed. Re-park: + reg.slot.mark_parked(50_000); + assert!( + !reg.slot.was_swept_cancel(&e125), + "re-park must clear stale provenance" + ); + // Drain cancel marks provenance too. + reg.slot.mark_parked(60_000); + assert_eq!(cancel_all_parked(), 1); + assert!(reg.slot.was_swept_cancel(&e125)); + } + #[test] fn registration_drop_deregisters() { { diff --git a/src/server/conn/handler_monoio/mod.rs b/src/server/conn/handler_monoio/mod.rs index ff34698c..4b3f4d00 100644 --- a/src/server/conn/handler_monoio/mod.rs +++ b/src/server/conn/handler_monoio/mod.rs @@ -382,6 +382,15 @@ pub(crate) async fn handle_connection_sharded_monoio< let (client_live, registry_guard) = match park.kept_registration { Some(guard) => { let live = crate::client_registry::live_handle(client_id).unwrap_or_else(|| { + // #438 conn-secondary: reaching here means the entry vanished + // while its guard was alive — an invariant violation, not a + // code path. Fail open (re-register so the conn stays served + // and killable) but LOUDLY; register() itself now balances + // the counters if the entry raced back into existence. + tracing::warn!( + "client {}: registry entry missing on park resume despite live guard; re-registering", + client_id + ); crate::client_registry::register( client_id, peer_addr.clone(), @@ -919,6 +928,16 @@ pub(crate) async fn handle_connection_sharded_monoio< && !conn.in_multi && conn.command_queue.is_empty() && conn.active_cross_txn.is_none() + // F6 (#438, sec L2): unauthenticated conns never task-park. + // On an auth-enabled server, parking a pre-AUTH conn would + // let an attacker hold a maxclients-worth of silent sockets + // at ~3.3 KB each, indefinitely and invisibly cheap. Keeping + // them un-parked leaves each one pinned to a full handler + // task — costly enough to surface in CPU/RSS monitoring — + // and `timeout N` (the slowloris knob) still reaps them. On + // no-auth servers `authenticated` is true from accept, so + // this changes nothing there. + && conn.authenticated // Replica-handshake conns (sent REPLCONF) never park: PSYNC // on a resumed parked conn is unsupported (warn+close). && !conn.saw_replconf @@ -941,7 +960,10 @@ pub(crate) async fn handle_connection_sharded_monoio< // Real socket/TLS error (EOF without close_notify, // 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, + // #438 conn-secondary: `was_swept_cancel` also demands + // sweep provenance — a bare errno 125 that no sweep of + // ours produced is a real error too. + Err(ref e) if !reg.slot.was_swept_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 @@ -1006,7 +1028,8 @@ pub(crate) async fn handle_connection_sharded_monoio< // always performed a read that re-surfaced them; the // 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, + // #438 conn-secondary: provenance-checked — see stage 2. + Err(ref e) if !reg.slot.was_swept_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. diff --git a/src/server/conn/handler_sharded/mod.rs b/src/server/conn/handler_sharded/mod.rs index 0be9804a..04ea4a5e 100644 --- a/src/server/conn/handler_sharded/mod.rs +++ b/src/server/conn/handler_sharded/mod.rs @@ -192,12 +192,17 @@ pub(crate) async fn handle_connection_sharded( Some(stream), ) = (result.0, result.1) { - use std::os::unix::io::IntoRawFd; match stream.into_std() { Ok(std_stream) => { - let raw_fd = std_stream.into_raw_fd(); + // F4 (#438): the fd rides the channel as an OwnedFd — an + // undelivered message (ring torn down at shutdown) closes the + // socket on drop instead of leaking it with a stranded client. + let owned_fd = std::os::fd::OwnedFd::from(std_stream); let msg = ShardMessage::MigrateConnection(Box::new( - crate::shard::dispatch::MigrateConnectionPayload { fd: raw_fd, state }, + crate::shard::dispatch::MigrateConnectionPayload { + fd: owned_fd, + state, + }, )); let target_idx = ChannelMesh::target_index(ctx.shard_id, target_shard); let push_result = { @@ -250,12 +255,12 @@ pub(crate) async fn handle_connection_sharded( // serving the client on this shard with the state // recovered from the undelivered message, instead of // dropping the connection. Migration disabled so the - // affinity tracker cannot loop. - use std::os::unix::io::FromRawFd; - // SAFETY: fd is a valid, uniquely-owned descriptor from - // into_raw_fd(); the failed push left ownership with us. - let std_stream = - unsafe { std::net::TcpStream::from_raw_fd(payload.fd) }; + // affinity tracker cannot loop. F4 (#438): the fd + // comes back as an OwnedFd; the safe From conversion + // replaces the old unsafe from_raw_fd. + use std::os::fd::AsRawFd; + let raw_kill_fd = payload.fd.as_raw_fd(); + let std_stream = std::net::TcpStream::from(payload.fd); match tokio::net::TcpStream::from_std(std_stream) { Ok(tcp_stream) => { let _ = handle_connection_sharded_inner( @@ -267,7 +272,7 @@ pub(crate) async fn handle_connection_sharded( false, // can_migrate BytesMut::new(), Some(&payload.state), - payload.fd, + raw_kill_fd, ) .await; } diff --git a/src/server/listener.rs b/src/server/listener.rs index 0f462a4a..afb53927 100644 --- a/src/server/listener.rs +++ b/src/server/listener.rs @@ -7,7 +7,7 @@ use parking_lot::Mutex; use std::path::PathBuf; use std::sync::atomic::AtomicU64; use std::sync::{Arc, RwLock}; -use tracing::{debug, error, info}; +use tracing::{debug, error, info, warn}; #[cfg(feature = "runtime-tokio")] use crate::command::connection as conn_cmd; @@ -29,6 +29,33 @@ use crate::tracking::TrackingTable; #[cfg(feature = "runtime-tokio")] type SharedDatabases = Arc>>; +/// #438 F3: deliver an accepted connection to a shard WITHOUT head-of-line +/// blocking the central accept loop on one wedged shard's full conn channel. +/// +/// Tries the routed shard first, then every other shard once (rotating from +/// the routed one). A full — or disconnected — channel just moves on to the +/// next shard; the affinity hint is best-effort, and any live shard serves +/// the client correctly. Returns `Ok(shard)` with the shard that actually +/// took the connection, or `Err(payload)` when EVERY shard's channel refused +/// it — at that point the caller decides between back-pressure (blocking +/// send, the pre-F3 behaviour, correct when the whole server is saturated) +/// and dropping (correct during shutdown when the receivers are gone). +fn try_route_conn( + txs: &[channel::MpscSender], + target: usize, + payload: T, +) -> Result { + let mut payload = payload; + for i in 0..txs.len() { + let shard = (target + i) % txs.len(); + match txs[shard].try_send(payload) { + Ok(()) => return Ok(shard), + Err(e) => payload = e.into_inner(), + } + } + Err(payload) +} + #[cfg(feature = "runtime-tokio")] use super::connection; #[cfg(feature = "runtime-tokio")] @@ -395,9 +422,13 @@ pub async fn run_sharded( continue; } debug!("New TLS connection from {} -> shard {}", addr, tls_next_shard); - let tx = &tls_txs[tls_next_shard]; - if tx.send_async((stream, true)).await.is_err() { - error!("Failed to send TLS connection to shard {}", tls_next_shard); + // F3 (#438): rotate past a full shard channel instead of + // head-of-line blocking every other shard's accepts. + if let Err(payload) = try_route_conn(&tls_txs, tls_next_shard, (stream, true)) { + warn!("All shard conn channels full; back-pressuring TLS accept on shard {}", tls_next_shard); + if tls_txs[tls_next_shard].send_async(payload).await.is_err() { + error!("Failed to send TLS connection to shard {}", tls_next_shard); + } } tls_next_shard = (tls_next_shard + 1) % tls_num_shards; } @@ -467,9 +498,15 @@ pub async fn run_sharded( } }; debug!("New connection from {} -> shard {}", addr, target_shard); - let tx = &conn_txs[target_shard]; - if tx.send_async((stream, false)).await.is_err() { - error!("Failed to send connection to shard {}", target_shard); + // F3 (#438): rotate past a full shard channel instead of + // head-of-line blocking every other shard's accepts. Only + // when EVERY channel is full does the accept loop block — + // server-wide saturation, where back-pressure is correct. + if let Err(payload) = try_route_conn(&conn_txs, target_shard, (stream, false)) { + warn!("All shard conn channels full; back-pressuring accept on shard {}", target_shard); + if conn_txs[target_shard].send_async(payload).await.is_err() { + error!("Failed to send connection to shard {}", target_shard); + } } } Err(e) => { @@ -609,9 +646,16 @@ pub async fn run_sharded( // which relinquished ownership. We take sole ownership here. unsafe { std::net::TcpStream::from_raw_fd(fd) } }; - let tx = &conn_txs[target_shard]; - if tx.send((std_stream, false)).is_err() { - error!("Failed to send connection to shard {}", target_shard); + // F3 (#438): rotate past a full shard channel; a + // blocking `send` here stalls the whole listener + // THREAD (this is a sync flume send on the monoio + // main loop), freezing TLS accepts too. Block only + // when every channel is full (server saturated). + if let Err(payload) = try_route_conn(&conn_txs, target_shard, (std_stream, false)) { + warn!("All shard conn channels full; back-pressuring accept on shard {}", target_shard); + if conn_txs[target_shard].send(payload).is_err() { + error!("Failed to send connection to shard {}", target_shard); + } } } Err(e) => { accept_backoff.record_error("Accept error", &e).await; } @@ -639,9 +683,12 @@ pub async fn run_sharded( // which relinquished ownership. We take sole ownership here. unsafe { std::net::TcpStream::from_raw_fd(fd) } }; - let tx = &conn_txs[tls_next_shard]; - if tx.send((std_stream, true)).is_err() { - error!("Failed to send TLS connection to shard {}", tls_next_shard); + // F3 (#438): same rotation as the plain leg above. + if let Err(payload) = try_route_conn(&conn_txs, tls_next_shard, (std_stream, true)) { + warn!("All shard conn channels full; back-pressuring TLS accept on shard {}", tls_next_shard); + if conn_txs[tls_next_shard].send(payload).is_err() { + error!("Failed to send TLS connection to shard {}", tls_next_shard); + } } tls_next_shard = (tls_next_shard + 1) % num_shards; } @@ -703,9 +750,16 @@ pub async fn run_sharded( // which relinquished ownership. We take sole ownership here. unsafe { std::net::TcpStream::from_raw_fd(fd) } }; - let tx = &conn_txs[target_shard]; - if tx.send((std_stream, false)).is_err() { - error!("Failed to send connection to shard {}", target_shard); + // F3 (#438): rotate past a full shard channel; a + // blocking `send` here stalls the whole listener + // THREAD (this is a sync flume send on the monoio + // main loop), freezing TLS accepts too. Block only + // when every channel is full (server saturated). + if let Err(payload) = try_route_conn(&conn_txs, target_shard, (std_stream, false)) { + warn!("All shard conn channels full; back-pressuring accept on shard {}", target_shard); + if conn_txs[target_shard].send(payload).is_err() { + error!("Failed to send connection to shard {}", target_shard); + } } } Err(e) => { @@ -723,3 +777,57 @@ pub async fn run_sharded( Ok(()) } + +#[cfg(test)] +mod listener_tests { + use super::try_route_conn; + use crate::runtime::channel; + + /// F3 (#438): one wedged shard's full channel must not stop delivery — + /// the connection rotates to the next shard. + #[test] + fn full_target_rotates_to_next_shard() { + let (tx0, _rx0) = channel::mpsc_bounded::(1); + let (tx1, rx1) = channel::mpsc_bounded::(1); + let (tx2, _rx2) = channel::mpsc_bounded::(1); + tx0.try_send(99).unwrap(); // wedge shard 0 + let txs = vec![tx0, tx1, tx2]; + assert_eq!(try_route_conn(&txs, 0, 7), Ok(1)); + assert_eq!(rx1.try_recv(), Ok(7)); + } + + /// A disconnected shard (receiver dropped — shard thread gone) is skipped + /// the same way a full one is. + #[test] + fn disconnected_shard_is_skipped() { + let (tx0, rx0) = channel::mpsc_bounded::(1); + let (tx1, rx1) = channel::mpsc_bounded::(1); + drop(rx0); + let txs = vec![tx0, tx1]; + assert_eq!(try_route_conn(&txs, 0, 7), Ok(1)); + assert_eq!(rx1.try_recv(), Ok(7)); + } + + /// Every channel refusing returns the payload for the caller's blocking + /// back-pressure fallback — the connection is never silently dropped. + #[test] + fn all_full_returns_payload() { + let (tx0, _rx0) = channel::mpsc_bounded::(1); + let (tx1, _rx1) = channel::mpsc_bounded::(1); + tx0.try_send(1).unwrap(); + tx1.try_send(2).unwrap(); + let txs = vec![tx0, tx1]; + assert_eq!(try_route_conn(&txs, 1, 7), Err(7)); + } + + /// Rotation starts AT the routed shard, preserving affinity/round-robin + /// placement whenever the target has room. + #[test] + fn target_with_room_keeps_the_connection() { + let (tx0, rx0) = channel::mpsc_bounded::(1); + let (tx1, _rx1) = channel::mpsc_bounded::(1); + let txs = vec![tx0, tx1]; + assert_eq!(try_route_conn(&txs, 0, 7), Ok(0)); + assert_eq!(rx0.try_recv(), Ok(7)); + } +} diff --git a/src/shard/conn_accept.rs b/src/shard/conn_accept.rs index e8113cca..d6508f1a 100644 --- a/src/shard/conn_accept.rs +++ b/src/shard/conn_accept.rs @@ -407,29 +407,24 @@ pub(crate) fn spawn_tokio_connection( /// Spawn a migrated connection handler on the target shard (Tokio runtime). /// -/// Reconstructs a `TcpStream` from a raw FD transferred via `ShardMessage::MigrateConnection`, -/// prepends synthetic RESP commands for state restoration (SELECT, CLIENT SETNAME), and -/// spawns a handler with `requirepass = None` (pre-authenticated). -/// -/// # Safety -/// -/// The caller must ensure `fd` is a valid, open file descriptor representing a connected -/// TCP socket. Ownership is transferred: this function consumes the FD (via `from_raw_fd`). +/// Reconstructs a `TcpStream` from the `OwnedFd` transferred via +/// `ShardMessage::MigrateConnection` (#438 F4: ownership is in the type — the +/// conversion chain is entirely safe, and an undelivered fd closes on drop), +/// prepends synthetic RESP commands for state restoration (SELECT, CLIENT +/// SETNAME), and spawns the handler. /// /// # Limitations /// /// TLS connections cannot be migrated because TLS session state lives in userspace and /// cannot be reconstructed from a raw FD. Only plain TCP connections should be migrated. -// `unix`-gated alongside `runtime-tokio`: the signature takes a -// `std::os::unix::io::RawFd` and reconstructs the socket via `from_raw_fd`, -// both of which only exist on Unix. Moon targets Linux + macOS (both Unix), so -// on every supported build `unix` is always true; the extra predicate just makes -// the platform coupling explicit (CodeRabbit PR #144). Full non-Unix support -// would also require gating `MigrateConnectionPayload.fd` — out of scope. +// `unix`-gated alongside `runtime-tokio`: the signature takes an `OwnedFd` +// (`MigrateFd` aliases it on Unix only). Moon targets Linux + macOS (both +// Unix), so on every supported build `unix` is always true; the extra +// predicate just makes the platform coupling explicit (CodeRabbit PR #144). #[cfg(all(feature = "runtime-tokio", unix))] #[allow(clippy::too_many_arguments)] pub(crate) fn spawn_migrated_tokio_connection( - fd: std::os::unix::io::RawFd, + fd: crate::shard::dispatch::MigrateFd, mut state: MigratedConnectionState, shard_databases: &Arc, dispatch_tx: &Rc>>>, @@ -459,24 +454,15 @@ pub(crate) fn spawn_migrated_tokio_connection( spill_file_id: &Rc>, disk_offload_dir: &Option, ) { - use std::os::unix::io::FromRawFd; - use crate::server::connection::handle_connection_sharded_inner; - // `fd` was produced by `libc::dup()` on the source shard before being pushed - // through the SPSC channel. That dup is a fresh, owned kernel fd, distinct from - // any other open fd. Ownership is transferred exactly once through the channel — - // the source shard drops the original stream immediately after `dup`, and on - // SPSC push failure the producer reconstructs an `OwnedFd` to close the dup. - // Here on the consumer side we take ownership by wrapping it in `TcpStream`, - // whose `Drop` closes the fd exactly once. No aliasing, no double-close. - // SAFETY: fd is a valid, uniquely-owned dup'd socket transferred via SPSC. - let std_stream = unsafe { std::net::TcpStream::from_raw_fd(fd) }; + // F4 (#438): ownership rode the channel as an OwnedFd, so the conversion + // to TcpStream is the safe From impl — no raw-fd reasoning left here. + let std_stream = std::net::TcpStream::from(fd); if let Err(e) = std_stream.set_nonblocking(true) { tracing::warn!( - "Shard {}: migrated fd {} set_nonblocking failed: {}", + "Shard {}: migrated fd set_nonblocking failed: {}", shard_id, - fd, e ); return; // std_stream Drop closes FD @@ -525,6 +511,8 @@ pub(crate) fn spawn_migrated_tokio_connection( let sc = script_cache_rc.clone(); let acl = acl_table.clone(); let rtcfg = runtime_config.clone(); + // F5 (#438): read before rtcfg moves into ConnectionContext::new. + let reqpass = rtcfg.read().requirepass.clone(); let scfg = server_config.clone(); let notifiers = all_notifiers.to_vec(); let snap_tx = snapshot_trigger_tx.clone(); @@ -548,7 +536,14 @@ pub(crate) fn spawn_migrated_tokio_connection( num_shards, psr, blk, - None, // requirepass: None = pre-authenticated + // F5 (#438, sec L3): carry the REAL requirepass. The conn's + // auth state comes from MigratedConnectionState (a migrated + // conn is already authenticated), so this is inert for the + // session itself — but a later AUTH on the migrated conn now + // validates against the actual password instead of erroring + // with "no password is set", and any future code that + // re-derives auth from ctx.requirepass sees the truth. + reqpass, pool_for_ctx, trk, rs, @@ -1023,9 +1018,15 @@ pub(crate) fn spawn_monoio_connection( ) .await; } else { + // F4 (#438): wrap the dup immediately — from here on + // the fd is closed by Drop on every path (undelivered + // message, ring torn down at shutdown), never leaked. + // SAFETY: dup_fd is the fresh, uniquely-owned fd + // libc::dup returned above; nothing else closes it. + let owned_fd = unsafe { std::os::unix::io::OwnedFd::from_raw_fd(dup_fd) }; let mut pending_msg = Some(ShardMessage::MigrateConnection(Box::new( crate::shard::dispatch::MigrateConnectionPayload { - fd: dup_fd, + fd: owned_fd, state, }, ))); @@ -1060,11 +1061,8 @@ pub(crate) fn spawn_monoio_connection( if _migrated { drop(stream); // hand-off confirmed: target owns dup_fd } else if let Some(ShardMessage::MigrateConnection(payload)) = pending_msg { - // SAFETY: dup_fd is a valid dup'd socket from libc::dup - // above; the push never succeeded so we still own it. - drop(unsafe { - std::os::unix::io::OwnedFd::from_raw_fd(payload.fd) - }); + // F4 (#438): dropping the payload closes the + // dup'd fd (OwnedFd) — no manual reclaim needed. tracing::warn!( "Shard {}: migration SPSC full, keeping connection {} local (monoio)", shard_id, cid @@ -1309,6 +1307,13 @@ fn spawn_resumed_parked_conn( &conn_ctx, shutdown.clone(), client_id, + // can_migrate=false is DELIBERATE (#438 F4 review): this wrapper + // has no MigrateConnection arm — enabling migration here would + // need the full dup/SPSC/fail-open plumbing of the primary spawn + // site. The affinity value is also marginal: a conn only reaches + // this path by idling past --conn-park-secs, and migration + // sampling exists for HOT connections. A resumed conn that turns + // hot simply stays on the shard the affinity router first chose. false, // can_migrate initial, Some(&state), @@ -1361,25 +1366,21 @@ fn spawn_resumed_parked_conn( /// Spawn a migrated connection handler on the target shard (monoio runtime). /// -/// Reconstructs a `monoio::net::TcpStream` from a raw FD transferred via -/// `ShardMessage::MigrateConnection`, prepends synthetic RESP commands for state -/// restoration, and spawns a handler with `requirepass = None` (pre-authenticated). -/// -/// # Safety -/// -/// Same safety requirements as `spawn_migrated_tokio_connection`: the caller must -/// ensure `fd` is a valid, open file descriptor for a connected TCP socket. +/// Reconstructs a `monoio::net::TcpStream` from the `OwnedFd` transferred via +/// `ShardMessage::MigrateConnection` (#438 F4: safe ownership transfer, drop +/// closes), prepends synthetic RESP commands for state restoration, and +/// spawns the handler. /// /// # Limitations /// /// TLS connections cannot be migrated (TLS session state is in userspace). // `unix`-gated alongside `runtime-monoio` for the same reason as the tokio twin -// above: the `RawFd` signature + `from_raw_fd` are Unix-only. Always true on -// supported targets (Linux + macOS). +// above: `MigrateFd` is `OwnedFd` on Unix only. Always true on supported +// targets (Linux + macOS). #[cfg(all(feature = "runtime-monoio", unix))] #[allow(clippy::too_many_arguments)] pub(crate) fn spawn_migrated_monoio_connection( - fd: std::os::unix::io::RawFd, + fd: crate::shard::dispatch::MigrateFd, mut state: MigratedConnectionState, shard_databases: &Arc, dispatch_tx: &Rc>>>, @@ -1409,19 +1410,15 @@ pub(crate) fn spawn_migrated_monoio_connection( spill_file_id: &Rc>, disk_offload_dir: &Option, ) { - use std::os::unix::io::FromRawFd; - use crate::server::connection::handle_connection_sharded_monoio; - // Same ownership chain as `spawn_migrated_tokio_connection`: `fd` is a dup'd - // socket transferred exactly once through SPSC, source already dropped its handle. - // SAFETY: fd is a valid, uniquely-owned dup'd socket; TcpStream is sole close-owner. - let std_stream = unsafe { std::net::TcpStream::from_raw_fd(fd) }; + // F4 (#438): same as the tokio twin — OwnedFd in, safe From conversion, + // TcpStream is sole close-owner from here. + let std_stream = std::net::TcpStream::from(fd); if let Err(e) = std_stream.set_nonblocking(true) { tracing::warn!( - "Shard {}: migrated fd {} set_nonblocking failed: {}", + "Shard {}: migrated fd set_nonblocking failed: {}", shard_id, - fd, e ); // Source shard's `try_accept_connection` already counted this @@ -1470,6 +1467,8 @@ pub(crate) fn spawn_migrated_monoio_connection( let sc = script_cache_rc.clone(); let acl = acl_table.clone(); let rtcfg = runtime_config.clone(); + // F5 (#438): read before rtcfg moves into ConnectionContext::new. + let reqpass = rtcfg.read().requirepass.clone(); let scfg = server_config.clone(); let notifiers = all_notifiers.to_vec(); let snap_tx = snapshot_trigger_tx.clone(); @@ -1493,7 +1492,14 @@ pub(crate) fn spawn_migrated_monoio_connection( num_shards, psr, blk, - None, // requirepass: None = pre-authenticated + // F5 (#438, sec L3): carry the REAL requirepass. The conn's + // auth state comes from MigratedConnectionState (a migrated + // conn is already authenticated), so this is inert for the + // session itself — but a later AUTH on the migrated conn now + // validates against the actual password instead of erroring + // with "no password is set", and any future code that + // re-derives auth from ctx.requirepass sees the truth. + reqpass, pool_for_ctx, trk, rs, diff --git a/src/shard/dispatch.rs b/src/shard/dispatch.rs index f89f0f6a..3c714cde 100644 --- a/src/shard/dispatch.rs +++ b/src/shard/dispatch.rs @@ -363,13 +363,27 @@ pub type RawSocketFd = std::os::unix::io::RawFd; #[cfg(not(unix))] pub type RawSocketFd = i32; +/// Owned socket fd carried by a connection-migration message (#438 F4). +/// +/// On Unix this is `std::os::fd::OwnedFd`, so an undelivered or dropped +/// `MigrateConnectionPayload` — SPSC ring torn down at shutdown, message +/// drained but never spawned — CLOSES the socket instead of leaking the fd +/// and stranding the client on a connection no task will ever serve again. +/// On non-unix targets (where the migration path is compiled out and the +/// payload is never constructed) it stays the raw alias so the types compile. +#[cfg(unix)] +pub type MigrateFd = std::os::fd::OwnedFd; +#[cfg(not(unix))] +pub type MigrateFd = RawSocketFd; + /// Boxed payload for `ShardMessage::MigrateConnection` (Phase 177, hot-path split). /// /// `MigratedConnectionState` already holds heap-backed strings/bytes but still /// exceeds 120 B inline. Moving it behind a Box keeps the enum slot in the /// cache-line budget set by the slotted variants. pub struct MigrateConnectionPayload { - pub fd: RawSocketFd, + /// Owned by the message (#438 F4): dropping the payload closes the socket. + pub fd: MigrateFd, pub state: crate::server::conn::affinity::MigratedConnectionState, } @@ -1619,3 +1633,51 @@ mod tests { ); } } + +#[cfg(all(test, unix))] +mod migrate_fd_tests { + use super::MigrateConnectionPayload; + use std::os::fd::{AsRawFd, OwnedFd}; + + /// F4 (#438): an undelivered `MigrateConnectionPayload` (SPSC ring torn + /// down at shutdown, drain exited before spawn) must CLOSE its socket on + /// drop — the pre-F4 raw-i32 field leaked the fd and stranded the client. + /// The `OwnedFd` field makes this a type-level guarantee; this test pins + /// the type so a revert to a raw fd cannot land silently. + #[test] + fn dropped_payload_closes_socket() { + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind"); + let stream = + std::net::TcpStream::connect(listener.local_addr().expect("addr")).expect("connect"); + let fd = OwnedFd::from(stream); + let raw = fd.as_raw_fd(); + let payload = MigrateConnectionPayload { + fd, + state: crate::server::conn::affinity::MigratedConnectionState { + selected_db: 0, + authenticated: true, + client_name: None, + protocol_version: 2, + current_user: "default".to_string(), + flags: 0, + read_buf_remainder: bytes::BytesMut::new(), + client_id: 1, + peer_addr: "t".to_string(), + workspace_id: None, + }, + }; + // SAFETY: fcntl(F_GETFD) on any integer is memory-safe; an invalid fd + // just returns -1/EBADF. + assert!( + unsafe { libc::fcntl(raw, libc::F_GETFD) } >= 0, + "fd must be open while the payload owns it" + ); + drop(payload); + // SAFETY: as above. + assert_eq!( + unsafe { libc::fcntl(raw, libc::F_GETFD) }, + -1, + "dropping the payload must close the socket" + ); + } +} diff --git a/src/shard/event_loop.rs b/src/shard/event_loop.rs index 0e54fba1..7d01128c 100644 --- a/src/shard/event_loop.rs +++ b/src/shard/event_loop.rs @@ -999,8 +999,11 @@ impl super::Shard { let cached_clock = CachedClock::new(); // Pending FD migrations collected from SPSC drain (spawn wired in Plan 50-02). + // F4 (#438): the fd is an OwnedFd — entries still queued when the event + // loop exits (shutdown) drop here and CLOSE their sockets, giving the + // client a FIN instead of a permanently stranded silent connection. let mut pending_migrations: Vec<( - crate::shard::dispatch::RawSocketFd, + crate::shard::dispatch::MigrateFd, crate::server::conn::affinity::MigratedConnectionState, )> = Vec::new(); @@ -1497,7 +1500,7 @@ impl super::Shard { { tracing::info!( "Shard {}: accepting migrated connection (fd={}, client_id={}, from={})", - shard_id, fd, state.client_id, state.peer_addr + shard_id, std::os::fd::AsRawFd::as_raw_fd(&fd), state.client_id, state.peer_addr ); #[cfg(feature = "runtime-tokio")] conn_accept::spawn_migrated_tokio_connection( @@ -1602,7 +1605,7 @@ impl super::Shard { { tracing::info!( "Shard {}: accepting migrated connection (fd={}, client_id={}, from={})", - shard_id, fd, state.client_id, state.peer_addr + shard_id, std::os::fd::AsRawFd::as_raw_fd(&fd), state.client_id, state.peer_addr ); #[cfg(feature = "runtime-tokio")] conn_accept::spawn_migrated_tokio_connection( @@ -2323,7 +2326,7 @@ impl super::Shard { tracing::info!( "Shard {}: accepting migrated connection (fd={}, client_id={}, from={})", shard_id, - fd, + std::os::fd::AsRawFd::as_raw_fd(&fd), state.client_id, state.peer_addr ); diff --git a/src/shard/spsc_handler.rs b/src/shard/spsc_handler.rs index 4d15de86..7bd9c198 100644 --- a/src/shard/spsc_handler.rs +++ b/src/shard/spsc_handler.rs @@ -131,7 +131,7 @@ pub(crate) fn drain_spsc_shared( script_cache: &Rc>, cached_clock: &CachedClock, pending_migrations: &mut Vec<( - crate::shard::dispatch::RawSocketFd, + crate::shard::dispatch::MigrateFd, crate::server::conn::affinity::MigratedConnectionState, )>, pending_cdc_subscribes: &mut Vec, @@ -291,6 +291,7 @@ pub(crate) fn drain_spsc_shared( execute_batch.push(msg); } ShardMessage::MigrateConnection(payload) => { + let payload = *payload; pending_migrations.push((payload.fd, payload.state)); } ShardMessage::CdcSubscribe(payload) => { diff --git a/tests/parked_idle_parity.rs b/tests/parked_idle_parity.rs index c14a8974..8d4b3bbd 100644 --- a/tests/parked_idle_parity.rs +++ b/tests/parked_idle_parity.rs @@ -18,7 +18,7 @@ use std::time::Duration; /// Past downshift (1s) + park threshold (2s) + sweep cadence (1s) + margin. const PARK_WAIT: Duration = Duration::from_millis(4600); -fn spawn_moon(dir: &std::path::Path, port: u16) -> std::process::Child { +fn spawn_moon_with(dir: &std::path::Path, port: u16, extra: &[&str]) -> std::process::Child { Command::new(common::find_moon_binary()) .args([ "--port", @@ -32,12 +32,17 @@ fn spawn_moon(dir: &std::path::Path, port: u16) -> std::process::Child { "--conn-park-secs", "2", ]) + .args(extra) .stdout(Stdio::null()) .stderr(Stdio::null()) .spawn() .expect("spawn moon") } +fn spawn_moon(dir: &std::path::Path, port: u16) -> std::process::Child { + spawn_moon_with(dir, port, &[]) +} + fn read_exact_deadline(stream: &mut TcpStream, want: usize) -> Vec { stream .set_read_timeout(Some(Duration::from_secs(10))) @@ -404,3 +409,61 @@ fn partial_frame_survives_a_park() { let _ = child.kill(); let _ = child.wait(); } + +/// F6 (#438, sec L2): an UNAUTHENTICATED connection on an auth-enabled server +/// must never task-park. Parking a pre-AUTH conn would let an attacker hold a +/// maxclients-worth of silent sockets at ~3.3 KB each, indefinitely and +/// invisibly cheap; keeping them un-parked leaves each pinned to a full +/// handler task, costly enough to surface in monitoring. +/// +/// The authed idle sibling doubles as the positive control: it MUST park +/// (proving the park machinery is live, so the unauth conn's absence from the +/// gauge is the F6 exclusion, not a broken park path). Task parking is +/// monoio-only; under tokio both counts are 0 and the test degenerates to +/// "unauth conn is never parked" (vacuously green, same caveat as the rest of +/// this suite). +#[test] +fn unauthenticated_conn_never_task_parks() { + let dir = tempfile::tempdir().expect("tempdir"); + let (mut child, port) = + common::spawn_listening(|p| spawn_moon_with(dir.path(), p, &["--requirepass", "park-f6"])); + + // Positive control: authenticated, then idle past the park threshold. + let mut authed = TcpStream::connect(("127.0.0.1", port)).expect("connect authed"); + let r = command_reply(&mut authed, "AUTH park-f6\r\n"); + assert!(r.starts_with("+OK"), "AUTH failed: {r}"); + + // Victim: connects and never speaks — no AUTH, no bytes. + let silent = TcpStream::connect(("127.0.0.1", port)).expect("connect silent"); + + std::thread::sleep(PARK_WAIT); + + // Control conn (freshly active, not parked) reads the gauge. + let mut control = TcpStream::connect(("127.0.0.1", port)).expect("connect control"); + let r = command_reply(&mut control, "AUTH park-f6\r\n"); + assert!(r.starts_with("+OK"), "control AUTH failed: {r}"); + let info = command_reply(&mut control, "INFO clients\r\n"); + let parked: u64 = info + .lines() + .find_map(|l| l.strip_prefix("parked_clients:")) + .unwrap_or_else(|| panic!("parked_clients missing from INFO clients:\n{info}")) + .trim() + .parse() + .expect("parked_clients value"); + + let expected = if cfg!(feature = "runtime-monoio") { + 1 + } else { + 0 + }; + assert_eq!( + parked, expected, + "exactly the authed idle conn may park; an unauthenticated conn must \ + never task-park (parked_clients={parked}, expected {expected})" + ); + + drop(silent); + drop(authed); + let _ = child.kill(); + let _ = child.wait(); +} From 44f9df58fe60c48ae96e950a40f5b8e4b8040b7f Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Fri, 7 Aug 2026 20:39:15 +0700 Subject: [PATCH 2/2] =?UTF-8?q?fix(server):=20review=20round=20for=20#446?= =?UTF-8?q?=20=E2=80=94=20live-shard=20fallback,=20deterministic=20tests,?= =?UTF-8?q?=20F6=20identity=20check?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review findings (Qodo + CodeRabbit, converging on one real defect): 1. Fallback could drop a conn a live shard would take (both bots, CONFIRMED): when try_route_conn exhausted the rotation, the blocking fallback targeted the ORIGINALLY routed shard — if that receiver was disconnected while another shard was merely full, the fallback errored instantly and the accepted conn was dropped. The helper now reports the first LIVE-but-full shard from the rotation (Err((Option, payload))); callers back-pressure on THAT shard, and only drop (with an error log) when every receiver is gone — i.e. shutdown, where dropping is correct. New unit tests: disconnected-target-with-live-full-shard, all-disconnected. 2. Flaky global-state unit tests (both bots): the TOTAL_CLIENTS bracket could false-fail if an unrelated test registered between its two loads, and a failed assert skipped cleanup; it now retries with a fresh id (bug reads +1 on EVERY attempt; fix needs one undisturbed attempt) and deregisters before the verdict. The F4 fd-drop test probed the raw fd number, which the parallel runner could reuse between close and probe; it now observes the close through the PEER's read() -> EOF — no fd-number reasoning at all. 3. F6 test identified WHICH conn parked (CodeRabbit): parked_clients=1 also holds if the silent conn parked and the authed one didn't. The test now closes the authed conn and requires the gauge to drain to 0 while the silent conn stays open. Skipped with rationale (replied on-thread): - "Split oversized files" — pre-existing sizes, out of scope for a hardening PR; noted for a future refactor wave. - New-unsafe flag on the monoio dup wrap — consolidates two prior unsafe sites at the same location; net unsafe count for the PR is -3, SAFETY comment present. Gates: fmt + clippy -D warnings (both feature sets), listener/registry/ migrate-fd unit tests green both feature sets, parked_idle_parity monoio 3/4 full-green runs (the one bad run failed on connect/read setup races in OTHER tests — the documented harness flake class queued for the test-hygiene sweep; the new assertions never tripped) and tokio 7/7. Refs #438, #446 author: Tin Dang --- src/client_registry.rs | 44 ++++++++----- src/server/listener.rs | 124 ++++++++++++++++++++++++++---------- src/shard/dispatch.rs | 31 ++++----- tests/parked_idle_parity.rs | 27 +++++++- 4 files changed, 160 insertions(+), 66 deletions(-) diff --git a/src/client_registry.rs b/src/client_registry.rs index f8c8a5af..2d9987fe 100644 --- a/src/client_registry.rs +++ b/src/client_registry.rs @@ -1304,25 +1304,35 @@ mod idle_timeout_tests { /// kept_registration miss-arm fail-safe racing an entry back into /// existence) must NOT increment TOTAL_CLIENTS again — the single /// eventual deregister would then leave the count permanently +1, - /// skewing `shard_overloaded` routing. (TOTAL_CLIENTS is process-global - /// and other tests register concurrently, so the assert brackets only - /// the one replacing insert — the narrowest possible window.) + /// skewing `shard_overloaded` routing. + /// + /// TOTAL_CLIENTS is process-global and other tests mutate it + /// concurrently, so the bracket retries with a fresh id when an + /// unrelated registration lands inside the two-load window (#446 + /// review): with the fix, one undisturbed attempt reads t2 == t1; with + /// the pre-fix unconditional fetch_add, EVERY attempt reads t1+1. + /// Cleanup runs before the verdict so a failure cannot leak entries + /// into later tests. #[test] fn replacing_register_does_not_double_count() { use std::sync::atomic::Ordering; - const ID: u64 = 9_060; - let _a = register(ID, "t:1".into(), "default".into(), 908, -1); - let t1 = TOTAL_CLIENTS.load(Ordering::Relaxed); - let _b = register(ID, "t:1".into(), "default".into(), 909, -1); - let t2 = TOTAL_CLIENTS.load(Ordering::Relaxed); - // Pre-fix, the unconditional fetch_add makes this read t1+1 even - // with zero concurrency. (An unrelated test registering inside this - // two-instruction window could false-fail; the window is nanoseconds - // and the suite's other global-counter tests accept the same odds.) - assert_eq!(t2, t1, "replacing insert must not double-count"); - // The entry itself was replaced, not duplicated, and one deregister - // fully clears it. - deregister(ID); - assert!(live_handle(ID).is_none(), "single deregister must clear"); + let mut clean_attempt_seen = false; + for attempt in 0..5u64 { + let id = 9_060 + attempt; + let _a = register(id, "t:1".into(), "default".into(), 908, -1); + let t1 = TOTAL_CLIENTS.load(Ordering::Relaxed); + let _b = register(id, "t:1".into(), "default".into(), 909, -1); + let t2 = TOTAL_CLIENTS.load(Ordering::Relaxed); + deregister(id); + assert!(live_handle(id).is_none(), "single deregister must clear"); + if t2 == t1 { + clean_attempt_seen = true; + break; + } + } + assert!( + clean_attempt_seen, + "replacing insert added to TOTAL_CLIENTS on every attempt — double-count" + ); } } diff --git a/src/server/listener.rs b/src/server/listener.rs index afb53927..8029ece5 100644 --- a/src/server/listener.rs +++ b/src/server/listener.rs @@ -36,24 +36,32 @@ type SharedDatabases = Arc>>; /// the routed one). A full — or disconnected — channel just moves on to the /// next shard; the affinity hint is best-effort, and any live shard serves /// the client correctly. Returns `Ok(shard)` with the shard that actually -/// took the connection, or `Err(payload)` when EVERY shard's channel refused -/// it — at that point the caller decides between back-pressure (blocking -/// send, the pre-F3 behaviour, correct when the whole server is saturated) -/// and dropping (correct during shutdown when the receivers are gone). +/// took the connection. When every channel refused, returns the payload plus +/// `Some(shard)` for the first LIVE-but-full shard in rotation order — the +/// caller back-pressures with a blocking send on THAT shard (review of #446: +/// blocking on the routed shard would drop the connection whenever the +/// routed receiver is gone while another shard is merely full) — or `None` +/// when every receiver is disconnected (shutdown; dropping is correct). fn try_route_conn( txs: &[channel::MpscSender], target: usize, payload: T, -) -> Result { +) -> Result, T)> { let mut payload = payload; + let mut first_full: Option = None; for i in 0..txs.len() { let shard = (target + i) % txs.len(); match txs[shard].try_send(payload) { Ok(()) => return Ok(shard), - Err(e) => payload = e.into_inner(), + Err(e) => { + if first_full.is_none() && matches!(e, flume::TrySendError::Full(_)) { + first_full = Some(shard); + } + payload = e.into_inner(); + } } } - Err(payload) + Err((first_full, payload)) } #[cfg(feature = "runtime-tokio")] @@ -424,10 +432,14 @@ pub async fn run_sharded( debug!("New TLS connection from {} -> shard {}", addr, tls_next_shard); // F3 (#438): rotate past a full shard channel instead of // head-of-line blocking every other shard's accepts. - if let Err(payload) = try_route_conn(&tls_txs, tls_next_shard, (stream, true)) { - warn!("All shard conn channels full; back-pressuring TLS accept on shard {}", tls_next_shard); - if tls_txs[tls_next_shard].send_async(payload).await.is_err() { - error!("Failed to send TLS connection to shard {}", tls_next_shard); + if let Err((live_full, payload)) = try_route_conn(&tls_txs, tls_next_shard, (stream, true)) { + if let Some(shard) = live_full { + warn!("All shard conn channels full; back-pressuring TLS accept on shard {}", shard); + if tls_txs[shard].send_async(payload).await.is_err() { + error!("Failed to send TLS connection to shard {}", shard); + } + } else { + error!("All shard conn channels disconnected; dropping accepted TLS connection"); } } tls_next_shard = (tls_next_shard + 1) % tls_num_shards; @@ -501,11 +513,16 @@ pub async fn run_sharded( // F3 (#438): rotate past a full shard channel instead of // head-of-line blocking every other shard's accepts. Only // when EVERY channel is full does the accept loop block — - // server-wide saturation, where back-pressure is correct. - if let Err(payload) = try_route_conn(&conn_txs, target_shard, (stream, false)) { - warn!("All shard conn channels full; back-pressuring accept on shard {}", target_shard); - if conn_txs[target_shard].send_async(payload).await.is_err() { - error!("Failed to send connection to shard {}", target_shard); + // server-wide saturation, where back-pressure is correct — + // and it blocks on a shard verified LIVE by the rotation. + if let Err((live_full, payload)) = try_route_conn(&conn_txs, target_shard, (stream, false)) { + if let Some(shard) = live_full { + warn!("All shard conn channels full; back-pressuring accept on shard {}", shard); + if conn_txs[shard].send_async(payload).await.is_err() { + error!("Failed to send connection to shard {}", shard); + } + } else { + error!("All shard conn channels disconnected; dropping accepted connection"); } } } @@ -650,11 +667,16 @@ pub async fn run_sharded( // blocking `send` here stalls the whole listener // THREAD (this is a sync flume send on the monoio // main loop), freezing TLS accepts too. Block only - // when every channel is full (server saturated). - if let Err(payload) = try_route_conn(&conn_txs, target_shard, (std_stream, false)) { - warn!("All shard conn channels full; back-pressuring accept on shard {}", target_shard); - if conn_txs[target_shard].send(payload).is_err() { - error!("Failed to send connection to shard {}", target_shard); + // when every channel is full (server saturated), + // on a shard the rotation verified is still live. + if let Err((live_full, payload)) = try_route_conn(&conn_txs, target_shard, (std_stream, false)) { + if let Some(shard) = live_full { + warn!("All shard conn channels full; back-pressuring accept on shard {}", shard); + if conn_txs[shard].send(payload).is_err() { + error!("Failed to send connection to shard {}", shard); + } + } else { + error!("All shard conn channels disconnected; dropping accepted connection"); } } } @@ -684,10 +706,14 @@ pub async fn run_sharded( unsafe { std::net::TcpStream::from_raw_fd(fd) } }; // F3 (#438): same rotation as the plain leg above. - if let Err(payload) = try_route_conn(&conn_txs, tls_next_shard, (std_stream, true)) { - warn!("All shard conn channels full; back-pressuring TLS accept on shard {}", tls_next_shard); - if conn_txs[tls_next_shard].send(payload).is_err() { - error!("Failed to send TLS connection to shard {}", tls_next_shard); + if let Err((live_full, payload)) = try_route_conn(&conn_txs, tls_next_shard, (std_stream, true)) { + if let Some(shard) = live_full { + warn!("All shard conn channels full; back-pressuring TLS accept on shard {}", shard); + if conn_txs[shard].send(payload).is_err() { + error!("Failed to send TLS connection to shard {}", shard); + } + } else { + error!("All shard conn channels disconnected; dropping accepted TLS connection"); } } tls_next_shard = (tls_next_shard + 1) % num_shards; @@ -754,11 +780,16 @@ pub async fn run_sharded( // blocking `send` here stalls the whole listener // THREAD (this is a sync flume send on the monoio // main loop), freezing TLS accepts too. Block only - // when every channel is full (server saturated). - if let Err(payload) = try_route_conn(&conn_txs, target_shard, (std_stream, false)) { - warn!("All shard conn channels full; back-pressuring accept on shard {}", target_shard); - if conn_txs[target_shard].send(payload).is_err() { - error!("Failed to send connection to shard {}", target_shard); + // when every channel is full (server saturated), + // on a shard the rotation verified is still live. + if let Err((live_full, payload)) = try_route_conn(&conn_txs, target_shard, (std_stream, false)) { + if let Some(shard) = live_full { + warn!("All shard conn channels full; back-pressuring accept on shard {}", shard); + if conn_txs[shard].send(payload).is_err() { + error!("Failed to send connection to shard {}", shard); + } + } else { + error!("All shard conn channels disconnected; dropping accepted connection"); } } } @@ -809,15 +840,42 @@ mod listener_tests { } /// Every channel refusing returns the payload for the caller's blocking - /// back-pressure fallback — the connection is never silently dropped. + /// back-pressure fallback — the connection is never silently dropped — + /// together with a LIVE full shard to block on (rotation order from the + /// target). #[test] - fn all_full_returns_payload() { + fn all_full_returns_payload_and_live_shard() { let (tx0, _rx0) = channel::mpsc_bounded::(1); let (tx1, _rx1) = channel::mpsc_bounded::(1); tx0.try_send(1).unwrap(); tx1.try_send(2).unwrap(); let txs = vec![tx0, tx1]; - assert_eq!(try_route_conn(&txs, 1, 7), Err(7)); + assert_eq!(try_route_conn(&txs, 1, 7), Err((Some(1), 7))); + } + + /// #446 review: a DISCONNECTED routed shard must not become the blocking + /// fallback while another shard is merely full — the fallback shard must + /// be one whose receiver is alive. + #[test] + fn disconnected_target_falls_back_to_live_full_shard() { + let (tx0, rx0) = channel::mpsc_bounded::(1); + let (tx1, _rx1) = channel::mpsc_bounded::(1); + drop(rx0); // routed shard's receiver is gone + tx1.try_send(1).unwrap(); // other shard alive but full + let txs = vec![tx0, tx1]; + assert_eq!(try_route_conn(&txs, 0, 7), Err((Some(1), 7))); + } + + /// Every receiver gone (shutdown): no live shard to block on — the + /// caller drops the connection instead of blocking forever. + #[test] + fn all_disconnected_reports_no_live_shard() { + let (tx0, rx0) = channel::mpsc_bounded::(1); + let (tx1, rx1) = channel::mpsc_bounded::(1); + drop(rx0); + drop(rx1); + let txs = vec![tx0, tx1]; + assert_eq!(try_route_conn(&txs, 0, 7), Err((None, 7))); } /// Rotation starts AT the routed shard, preserving affinity/round-robin diff --git a/src/shard/dispatch.rs b/src/shard/dispatch.rs index 3c714cde..8da9448c 100644 --- a/src/shard/dispatch.rs +++ b/src/shard/dispatch.rs @@ -1637,22 +1637,26 @@ mod tests { #[cfg(all(test, unix))] mod migrate_fd_tests { use super::MigrateConnectionPayload; - use std::os::fd::{AsRawFd, OwnedFd}; + use std::os::fd::OwnedFd; /// F4 (#438): an undelivered `MigrateConnectionPayload` (SPSC ring torn /// down at shutdown, drain exited before spawn) must CLOSE its socket on /// drop — the pre-F4 raw-i32 field leaked the fd and stranded the client. /// The `OwnedFd` field makes this a type-level guarantee; this test pins /// the type so a revert to a raw fd cannot land silently. + /// + /// #446 review: verified through the PEER's eyes (read → EOF) instead of + /// probing the raw fd number, which another thread could reuse between + /// close and probe under the parallel test runner. #[test] fn dropped_payload_closes_socket() { + use std::io::Read; let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind"); - let stream = + let mut client = std::net::TcpStream::connect(listener.local_addr().expect("addr")).expect("connect"); - let fd = OwnedFd::from(stream); - let raw = fd.as_raw_fd(); + let (server_side, _) = listener.accept().expect("accept"); let payload = MigrateConnectionPayload { - fd, + fd: OwnedFd::from(server_side), state: crate::server::conn::affinity::MigratedConnectionState { selected_db: 0, authenticated: true, @@ -1666,18 +1670,15 @@ mod migrate_fd_tests { workspace_id: None, }, }; - // SAFETY: fcntl(F_GETFD) on any integer is memory-safe; an invalid fd - // just returns -1/EBADF. - assert!( - unsafe { libc::fcntl(raw, libc::F_GETFD) } >= 0, - "fd must be open while the payload owns it" - ); drop(payload); - // SAFETY: as above. + client + .set_read_timeout(Some(std::time::Duration::from_secs(10))) + .expect("set timeout"); + let mut buf = [0u8; 1]; + let n = client.read(&mut buf).expect("read after payload drop"); assert_eq!( - unsafe { libc::fcntl(raw, libc::F_GETFD) }, - -1, - "dropping the payload must close the socket" + n, 0, + "peer must observe EOF once dropping the payload closes the socket" ); } } diff --git a/tests/parked_idle_parity.rs b/tests/parked_idle_parity.rs index 8d4b3bbd..c807fc68 100644 --- a/tests/parked_idle_parity.rs +++ b/tests/parked_idle_parity.rs @@ -462,8 +462,33 @@ fn unauthenticated_conn_never_task_parks() { never task-park (parked_clients={parked}, expected {expected})" ); - drop(silent); + // #446 review: the aggregate gauge alone can't say WHICH conn parked. + // Close the authed conn and require the gauge to drain to 0 while the + // silent conn stays open — if the parked one had been `silent`, the + // gauge would stay at 1. drop(authed); + let deadline = std::time::Instant::now() + Duration::from_secs(10); + loop { + let info = command_reply(&mut control, "INFO clients\r\n"); + let parked_now: u64 = info + .lines() + .find_map(|l| l.strip_prefix("parked_clients:")) + .unwrap_or_else(|| panic!("parked_clients missing from INFO clients:\n{info}")) + .trim() + .parse() + .expect("parked_clients value"); + if parked_now == 0 { + break; + } + assert!( + std::time::Instant::now() < deadline, + "parked_clients stuck at {parked_now} after closing the authed \ + conn — the parked connection was not the authenticated one" + ); + std::thread::sleep(Duration::from_millis(200)); + } + + drop(silent); let _ = child.kill(); let _ = child.wait(); }