diff --git a/CHANGELOG.md b/CHANGELOG.md index 7b6397da3..524dcd164 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,25 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Fixed +- **A client that disconnects while blocked is now reaped (c10k hardening + A1).** The infinite-wait `select!` behind `BLPOP`/`BRPOP`/`BLMOVE`/ + `BZPOPMIN`/`BZPOPMAX`/`BLMPOP`/`BZMPOP`/`BRPOPLPUSH` had exactly two arms, + `reply_rx` and `shutdown` — nothing watched the socket. `BLPOP key 0` + followed by a disconnect therefore leaked the handler task, the + `WaitEntry`, the client-registry entry and the maxclients slot *forever*: + infinite waiters carry `deadline: None` so the deadline sweep skips them, + and `timeout` exempts blocked clients by design (Redis parity). A few + thousand throwaway connections wedged the server until restart, using one + unauthenticated command. The wait now also watches the peer, and a + vanished client tears down every registration it held (local and remote) + before closing. `CLIENT KILL` of a blocked client works through the same + path, since it closes the fd with `shutdown(2)`. Bytes a client legally + pipelines behind its blocking command are carried into the parse stream + rather than consumed and dropped — they used to wait in the kernel, so + the handler now skips exactly one read to drain them. **Known gap:** TLS + connections keep the old behaviour; the vendored monoio-rustls read loops + internally until rustls yields plaintext, so a post-readiness read can + park inside a partial record, which the cancel-free design must never do. - **A timed-out cross-shard `BLPOP` no longer eats the next push to that key (c10k hardening A2/A3).** Two defects compounded into silent data loss at `--shards > 1`. First, the tokio single-key cleanup condition diff --git a/src/server/conn/blocking.rs b/src/server/conn/blocking.rs index 3195eb2d0..157fc2688 100644 --- a/src/server/conn/blocking.rs +++ b/src/server/conn/blocking.rs @@ -4,7 +4,6 @@ use std::rc::Rc; use std::sync::Arc; use bytes::Bytes; -#[cfg(feature = "runtime-monoio")] use bytes::BytesMut; use futures::StreamExt; use ringbuf::HeapProd; @@ -28,6 +27,140 @@ use super::util::extract_bytes; const BLOCK_REGISTER_FAILED: &[u8] = b"MOONERR blocking registration failed: owning shard not draining"; +/// What a blocking command produced. +/// +/// c10k hardening A1: the wait used to have only two exits — a reply or a +/// timeout — so a client that vanished mid-`BLPOP key 0` was unreachable +/// forever. [`PeerGone`](BlockingOutcome::PeerGone) is the third, and it is +/// deliberately NOT a `Frame`: there is nobody left to send one to, and the +/// caller must close the connection rather than write into a dead socket. +pub(crate) enum BlockingOutcome { + /// Normal completion: wake-up value, timeout nil, or an error frame. + Reply(Frame), + /// The peer went away while blocked (EOF, RST, or a `CLIENT KILL` + /// `shutdown(2)`). Every registration this waiter held has already been + /// torn down; the caller must close the connection without replying. + PeerGone, +} + +/// Scratch size for one drain of a blocked client's socket. Only ever used by +/// a client that pipelines behind a blocking command, which is legal but rare +/// — the buffer is allocated once per such connection and reused across the +/// wait loop, never on any dispatch path. +#[cfg(feature = "runtime-monoio")] +const PEER_DRAIN_BUF: usize = 512; + +/// Ceiling on bytes a blocked client may pipeline behind its blocking command +/// before the connection is dropped. +/// +/// The carry buffer is the one place A1 accumulates client input outside the +/// handler's own read loop, so it gets its own bound rather than inheriting +/// the read loop's (which has none — that is finding C2, a separate fix). The +/// value is far above any legitimate pipeline and far below a memory problem +/// at c10k. +const PEER_CARRY_LIMIT: usize = 64 * 1024 * 1024; + +/// c10k A1 (tokio): classify one peer read taken while a client is blocked. +/// +/// `true` means the wait is over because the peer is gone: EOF, a socket +/// error, or a client that pipelined past [`PEER_CARRY_LIMIT`]. `false` means +/// the bytes were carried and the client stays blocked. +#[cfg(feature = "runtime-tokio")] +fn peer_wake_ends_wait(read: std::io::Result, carry: &BytesMut) -> bool { + match read { + Ok(0) => true, + Ok(_) => carry.len() > PEER_CARRY_LIMIT, + // `Interrupted` is a retry signal, not a dead peer (tokio never + // surfaces `WouldBlock` here — it returns `Pending` instead). + Err(ref e) if e.kind() == std::io::ErrorKind::Interrupted => false, + Err(_) => true, + } +} + +/// c10k A1: result of one peer-readiness wake while a client is blocked. +#[cfg(feature = "runtime-monoio")] +enum PeerWake { + /// Bytes arrived and were carried forward; the client is alive and stays + /// blocked (Redis executes pipelined commands only after the block ends). + Alive, + /// EOF or a fatal socket error — the connection is finished. + Gone, +} + +/// c10k A1 (monoio): consume one readiness worth of bytes from a blocked +/// client's socket into `carry`. +/// +/// MUST be called only after [`IdleParkRead::peer_readable`] has fired, and +/// MUST NOT be called from inside a `select!`. Both rules exist for the same +/// reason: level-triggered readiness guarantees this read completes at once +/// (data, EOF, or error), so it can never be cancelled mid-flight and lose +/// client bytes — the failure mode that rules out a naive read arm. +/// +/// Racing the reply is harmless: `reply_tx` is a `flume` bounded(1), so a +/// wake-up delivered while this read is in flight simply sits in the channel +/// and the next loop turn picks it up. +#[cfg(feature = "runtime-monoio")] +async fn drain_peer_bytes( + stream: &mut S, + scratch: &mut Vec, + carry: &mut BytesMut, +) -> PeerWake { + if scratch.len() != PEER_DRAIN_BUF { + scratch.resize(PEER_DRAIN_BUF, 0); + } + let buf = std::mem::take(scratch); + let (res, buf) = stream.read(buf).await; + let wake = match res { + Ok(0) => PeerWake::Gone, + Ok(n) => { + carry.extend_from_slice(&buf[..n]); + PeerWake::Alive + } + // NOT death. `readable(false)` does a real `poll(2)`, but the legacy + // (epoll/kqueue) driver can still hand back a readiness that yields + // nothing, and a cancel is a control signal, not a peer event. + // Misreading either as EOF would drop a perfectly healthy blocked + // client — the exact inverse of the bug being fixed. Re-arming is + // safe: the next poll only resolves on a real event. + Err(ref e) + if matches!( + e.kind(), + std::io::ErrorKind::WouldBlock | std::io::ErrorKind::Interrupted + ) || super::handler_monoio::idle_park::is_sweep_cancel(e) => + { + PeerWake::Alive + } + Err(_) => PeerWake::Gone, + }; + *scratch = buf; + wake +} + +/// c10k A1 (monoio): handle one resolved peer-readiness poll. +/// +/// `true` means the wait is over because the peer is gone. A readiness error +/// counts as gone; readiness success means bytes or EOF are waiting, so the +/// (uncancellable, immediately-completing) drain decides which. +#[cfg(feature = "runtime-monoio")] +async fn peer_wake_ends_wait_monoio( + ready: std::io::Result<()>, + stream: &mut S, + scratch: &mut Vec, + carry: &mut BytesMut, +) -> bool { + if let Err(e) = ready { + // Same rule as the drain below: only a real error is death. + return !(matches!( + e.kind(), + std::io::ErrorKind::WouldBlock | std::io::ErrorKind::Interrupted + ) || super::handler_monoio::idle_park::is_sweep_cancel(&e)); + } + match drain_peer_bytes(stream, scratch, carry).await { + PeerWake::Gone => true, + PeerWake::Alive => carry.len() > PEER_CARRY_LIMIT, + } +} + /// Push a blocking-registry control message (`BlockRegister` / `BlockCancel`) /// to the shard that owns the key. /// @@ -226,8 +359,13 @@ pub(crate) fn convert_blocking_to_nonblocking(cmd: &[u8], args: &[Frame]) -> Fra /// first-wakeup-wins, BlockCancel cleanup on completion/timeout/shutdown. /// /// CRITICAL: RefCell borrows MUST be dropped before any .await point. +/// +/// c10k A1: `stream` is watched for EOF throughout the wait. `tokio`'s +/// `AsyncReadExt::read_buf` is documented cancel-safe (a losing `select!` +/// branch reads nothing), so here the watch is a plain read arm and any bytes +/// it does pick up are carried into `carry` for the handler's read loop. #[cfg(feature = "runtime-tokio")] -pub(crate) async fn handle_blocking_command( +pub(crate) async fn handle_blocking_command( cmd: &[u8], args: &[Frame], selected_db: usize, @@ -237,19 +375,25 @@ pub(crate) async fn handle_blocking_command( num_shards: usize, dispatch_tx: &Rc>>>, shutdown: &CancellationToken, -) -> Frame { + stream: &mut S, + carry: &mut BytesMut, +) -> BlockingOutcome +where + S: tokio::io::AsyncRead + Unpin, +{ use futures::stream::FuturesUnordered; + use tokio::io::AsyncReadExt; // Parse timeout (last argument for all blocking commands) let timeout_secs = match parse_blocking_timeout(cmd, args) { Ok(t) => t, - Err(e) => return e, + Err(e) => return BlockingOutcome::Reply(e), }; // Parse keys and command-specific args let (keys, blocked_cmd_factory) = match parse_blocking_args(cmd, args) { Ok(v) => v, - Err(e) => return e, + Err(e) => return BlockingOutcome::Reply(e), }; // --- Non-blocking fast path: try to get data immediately --- @@ -265,7 +409,7 @@ pub(crate) async fn handle_blocking_command( None }); if let Some(frame) = maybe_frame { - return frame; + return BlockingOutcome::Reply(frame); } // Borrow released at with_shard_db boundary — safe before await. } @@ -280,7 +424,7 @@ pub(crate) async fn handle_blocking_command( if keys.len() == 1 { let target = key_to_shard(&keys[0], num_shards); let is_remote = target != shard_id; - let (reply_tx, reply_rx) = channel::oneshot::>(); + let (reply_tx, mut reply_rx) = channel::oneshot::>(); let wait_id = blocking_registry.borrow_mut().next_wait_id(); if is_remote { // Remote registration via SPSC — bounded and shutdown-aware @@ -295,7 +439,9 @@ pub(crate) async fn handle_blocking_command( }, )); if !push_block_msg(shutdown, dispatch_tx, shard_id, target, msg, None).await { - return Frame::Error(Bytes::from_static(BLOCK_REGISTER_FAILED)); + return BlockingOutcome::Reply(Frame::Error(Bytes::from_static( + BLOCK_REGISTER_FAILED, + ))); } } else { // Local registration @@ -310,34 +456,58 @@ pub(crate) async fn handle_blocking_command( .register(selected_db, keys[0].clone(), entry); } + // A1: every arm below is a `loop` because the peer-watch arm can fire + // repeatedly without ending the wait — a client is allowed to pipeline + // behind a blocking command, and Redis runs those commands only after + // the block resolves. `&mut reply_rx` is re-selected across turns + // safely: `OneshotReceiver` caches its inner `flume` future, so the + // waker registration survives and no wake-up is lost. let result = if let Some(dl) = deadline { - tokio::select! { - res = reply_rx => { - match res { - Ok(Some(frame)) => frame, - Ok(None) | Err(_) => Frame::Null, + let sleep = tokio::time::sleep(dl.saturating_duration_since(std::time::Instant::now())); + tokio::pin!(sleep); + loop { + tokio::select! { + res = &mut reply_rx => { + break match res { + Ok(Some(frame)) => Some(frame), + Ok(None) | Err(_) => Some(Frame::Null), + }; + } + _ = &mut sleep => { + blocking_registry.borrow_mut().remove_wait(wait_id); + break Some(Frame::Null); + } + _ = shutdown.cancelled() => { + blocking_registry.borrow_mut().remove_wait(wait_id); + break Some(Frame::Error(Bytes::from_static(b"ERR server shutting down"))); + } + read = stream.read_buf(carry) => { + if peer_wake_ends_wait(read, carry) { + blocking_registry.borrow_mut().remove_wait(wait_id); + break None; + } } - } - _ = tokio::time::sleep(dl.saturating_duration_since(std::time::Instant::now())) => { - blocking_registry.borrow_mut().remove_wait(wait_id); - Frame::Null - } - _ = shutdown.cancelled() => { - blocking_registry.borrow_mut().remove_wait(wait_id); - Frame::Error(Bytes::from_static(b"ERR server shutting down")) } } } else { - tokio::select! { - res = reply_rx => { - match res { - Ok(Some(frame)) => frame, - Ok(None) | Err(_) => Frame::Null, + loop { + tokio::select! { + res = &mut reply_rx => { + break match res { + Ok(Some(frame)) => Some(frame), + Ok(None) | Err(_) => Some(Frame::Null), + }; + } + _ = shutdown.cancelled() => { + blocking_registry.borrow_mut().remove_wait(wait_id); + break Some(Frame::Error(Bytes::from_static(b"ERR server shutting down"))); + } + read = stream.read_buf(carry) => { + if peer_wake_ends_wait(read, carry) { + blocking_registry.borrow_mut().remove_wait(wait_id); + break None; + } } - } - _ = shutdown.cancelled() => { - blocking_registry.borrow_mut().remove_wait(wait_id); - Frame::Error(Bytes::from_static(b"ERR server shutting down")) } } }; @@ -370,7 +540,13 @@ pub(crate) async fn handle_blocking_command( ) .await; } - return result; + // A1: the cancel above runs for a vanished peer too — that is the + // whole point. Leaving the owner-side `WaitEntry` behind is what made + // a disconnected `BLPOP key 0` unreapable. + return match result { + Some(frame) => BlockingOutcome::Reply(frame), + None => BlockingOutcome::PeerGone, + }; } // --- Multi-key coordinator: register on ALL keys across local + remote shards --- @@ -444,12 +620,17 @@ pub(crate) async fn handle_blocking_command( ) .await; drop(receivers); - return Frame::Error(Bytes::from_static(BLOCK_REGISTER_FAILED)); + return BlockingOutcome::Reply(Frame::Error(Bytes::from_static(BLOCK_REGISTER_FAILED))); } // Await first successful result from any key/shard. // FuturesUnordered may return Err (sender dropped by remove_wait cleanup) before // returning the successful Ok. We must skip Err/None results and keep polling. + // + // A1: `peer_gone` breaks the same loop as every other terminal condition, + // so the shared cleanup below runs identically — a vanished multi-key + // waiter unwinds ALL of its registrations, local and remote. + let mut peer_gone = false; let frame = if let Some(dl) = deadline { let sleep = tokio::time::sleep(dl.saturating_duration_since(std::time::Instant::now())); tokio::pin!(sleep); @@ -468,6 +649,9 @@ pub(crate) async fn handle_blocking_command( result_frame = Frame::Error(Bytes::from_static(b"ERR server shutting down")); break; } + read = stream.read_buf(carry) => { + if peer_wake_ends_wait(read, carry) { peer_gone = true; break; } + } } } result_frame @@ -486,6 +670,9 @@ pub(crate) async fn handle_blocking_command( result_frame = Frame::Error(Bytes::from_static(b"ERR server shutting down")); break; } + read = stream.read_buf(carry) => { + if peer_wake_ends_wait(read, carry) { peer_gone = true; break; } + } } } result_frame @@ -505,7 +692,11 @@ pub(crate) async fn handle_blocking_command( // Drop remaining receivers; remote senders get Err on send -- harmless drop(receivers); - frame + if peer_gone { + BlockingOutcome::PeerGone + } else { + BlockingOutcome::Reply(frame) + } } /// Monoio version of handle_blocking_command. @@ -517,9 +708,20 @@ pub(crate) async fn handle_blocking_command( /// - `monoio::time::sleep(Duration::from_micros(10))` for SPSC backpressure /// /// CRITICAL: RefCell borrows MUST be dropped before any .await point. +/// +/// c10k A1: `stream` is watched for EOF throughout the wait, but unlike the +/// tokio twin the watch is a two-step — [`IdleParkRead::peer_readable`] inside +/// the `select!`, [`drain_peer_bytes`] outside it. A monoio read owns its +/// buffer, so a read arm losing the race could take client bytes down with the +/// cancelled op; a readiness poll owns nothing and is safe to drop. +/// +/// Streams whose `peer_readable` is the never-resolving default (TLS — see the +/// trait docs) keep the pre-A1 behaviour: their blocked clients are still not +/// reapable by disconnect. TLS at least requires a completed handshake, so it +/// is not the unauthenticated one-command DoS this fixes. #[cfg(feature = "runtime-monoio")] #[allow(clippy::await_holding_refcell_ref)] -pub(crate) async fn handle_blocking_command_monoio( +pub(crate) async fn handle_blocking_command_monoio( cmd: &[u8], args: &[Frame], selected_db: usize, @@ -530,21 +732,30 @@ pub(crate) async fn handle_blocking_command_monoio( dispatch_tx: &Rc>>>, shutdown: &CancellationToken, spsc_notifiers: &[Arc], -) -> Frame { + stream: &mut S, + carry: &mut BytesMut, +) -> BlockingOutcome +where + S: super::handler_monoio::idle_park::IdleParkRead, +{ use futures::stream::FuturesUnordered; // Parse timeout (last argument for all blocking commands) let timeout_secs = match parse_blocking_timeout(cmd, args) { Ok(t) => t, - Err(e) => return e, + Err(e) => return BlockingOutcome::Reply(e), }; // Parse keys and command-specific args let (keys, blocked_cmd_factory) = match parse_blocking_args(cmd, args) { Ok(v) => v, - Err(e) => return e, + Err(e) => return BlockingOutcome::Reply(e), }; + // Reused across every drain in this wait — one allocation for a client + // that actually pipelines, none for the overwhelming majority that do not. + let mut peer_scratch: Vec = Vec::new(); + // --- Non-blocking fast path: try to get data immediately --- // Use with_shard_db (thread-local ShardSlice) — no RwLock guard needed. { @@ -557,7 +768,7 @@ pub(crate) async fn handle_blocking_command_monoio( None }); if let Some(frame) = immediate_result { - return frame; + return BlockingOutcome::Reply(frame); } } @@ -571,7 +782,7 @@ pub(crate) async fn handle_blocking_command_monoio( if keys.len() == 1 { let target = key_to_shard(&keys[0], num_shards); let is_remote = target != shard_id; - let (reply_tx, reply_rx) = channel::oneshot::>(); + let (reply_tx, mut reply_rx) = channel::oneshot::>(); let wait_id = blocking_registry.borrow_mut().next_wait_id(); if is_remote { // A4: the old `loop { try_push; sleep(10µs) }` had no retry cap @@ -596,7 +807,9 @@ pub(crate) async fn handle_blocking_command_monoio( ) .await { - return Frame::Error(Bytes::from_static(BLOCK_REGISTER_FAILED)); + return BlockingOutcome::Reply(Frame::Error(Bytes::from_static( + BLOCK_REGISTER_FAILED, + ))); } } else { let entry = crate::blocking::WaitEntry { @@ -610,34 +823,54 @@ pub(crate) async fn handle_blocking_command_monoio( .register(selected_db, keys[0].clone(), entry); } + // A1: see the tokio twin for why these are loops. The peer arm only + // awaits READINESS; the read that follows happens after the select! + // has already resolved, where no cancellation can reach it. let result = if let Some(dl) = deadline { - monoio::select! { - res = reply_rx => { - match res { - Ok(Some(frame)) => frame, - Ok(None) | Err(_) => Frame::Null, + let mut sleep = std::pin::pin!(monoio::time::sleep( + dl.saturating_duration_since(std::time::Instant::now()) + )); + loop { + let ready = monoio::select! { + res = &mut reply_rx => { + break match res { + Ok(Some(frame)) => Some(frame), + Ok(None) | Err(_) => Some(Frame::Null), + }; } - } - _ = monoio::time::sleep(dl.saturating_duration_since(std::time::Instant::now())) => { - blocking_registry.borrow_mut().remove_wait(wait_id); - Frame::Null - } - _ = shutdown.cancelled() => { + _ = &mut sleep => { + blocking_registry.borrow_mut().remove_wait(wait_id); + break Some(Frame::Null); + } + _ = shutdown.cancelled() => { + blocking_registry.borrow_mut().remove_wait(wait_id); + break Some(Frame::Error(Bytes::from_static(b"ERR server shutting down"))); + } + r = stream.peer_readable() => r, + }; + if peer_wake_ends_wait_monoio(ready, stream, &mut peer_scratch, carry).await { blocking_registry.borrow_mut().remove_wait(wait_id); - Frame::Error(Bytes::from_static(b"ERR server shutting down")) + break None; } } } else { - monoio::select! { - res = reply_rx => { - match res { - Ok(Some(frame)) => frame, - Ok(None) | Err(_) => Frame::Null, + loop { + let ready = monoio::select! { + res = &mut reply_rx => { + break match res { + Ok(Some(frame)) => Some(frame), + Ok(None) | Err(_) => Some(Frame::Null), + }; } - } - _ = shutdown.cancelled() => { + _ = shutdown.cancelled() => { + blocking_registry.borrow_mut().remove_wait(wait_id); + break Some(Frame::Error(Bytes::from_static(b"ERR server shutting down"))); + } + r = stream.peer_readable() => r, + }; + if peer_wake_ends_wait_monoio(ready, stream, &mut peer_scratch, carry).await { blocking_registry.borrow_mut().remove_wait(wait_id); - Frame::Error(Bytes::from_static(b"ERR server shutting down")) + break None; } } }; @@ -654,7 +887,10 @@ pub(crate) async fn handle_blocking_command_monoio( ) .await; } - return result; + return match result { + Some(frame) => BlockingOutcome::Reply(frame), + None => BlockingOutcome::PeerGone, + }; } // --- Multi-key coordinator: register on ALL keys across local + remote shards --- @@ -732,19 +968,20 @@ pub(crate) async fn handle_blocking_command_monoio( ) .await; drop(receivers); - return Frame::Error(Bytes::from_static(BLOCK_REGISTER_FAILED)); + return BlockingOutcome::Reply(Frame::Error(Bytes::from_static(BLOCK_REGISTER_FAILED))); } // Await first successful result from any key/shard. // FuturesUnordered may return Err (sender dropped by remove_wait cleanup) before // returning the successful Ok. We must skip Err/None results and keep polling. + let mut peer_gone = false; let frame = if let Some(dl) = deadline { let mut sleep = std::pin::pin!(monoio::time::sleep( dl.saturating_duration_since(std::time::Instant::now()) )); let mut result_frame = Frame::Null; loop { - monoio::select! { + let ready = monoio::select! { result = receivers.next() => { match result { Some(Ok(Some(frame))) => { result_frame = frame; break; } @@ -757,13 +994,18 @@ pub(crate) async fn handle_blocking_command_monoio( result_frame = Frame::Error(Bytes::from_static(b"ERR server shutting down")); break; } + r = stream.peer_readable() => r, + }; + if peer_wake_ends_wait_monoio(ready, stream, &mut peer_scratch, carry).await { + peer_gone = true; + break; } } result_frame } else { let mut result_frame = Frame::Null; loop { - monoio::select! { + let ready = monoio::select! { result = receivers.next() => { match result { Some(Ok(Some(frame))) => { result_frame = frame; break; } @@ -775,6 +1017,11 @@ pub(crate) async fn handle_blocking_command_monoio( result_frame = Frame::Error(Bytes::from_static(b"ERR server shutting down")); break; } + r = stream.peer_readable() => r, + }; + if peer_wake_ends_wait_monoio(ready, stream, &mut peer_scratch, carry).await { + peer_gone = true; + break; } } result_frame @@ -794,7 +1041,11 @@ pub(crate) async fn handle_blocking_command_monoio( // Drop remaining receivers; remote senders get Err on send -- harmless drop(receivers); - frame + if peer_gone { + BlockingOutcome::PeerGone + } else { + BlockingOutcome::Reply(frame) + } } /// Parse timeout from a blocking command. diff --git a/src/server/conn/handler_monoio/dispatch.rs b/src/server/conn/handler_monoio/dispatch.rs index ad005048a..e3d607645 100644 --- a/src/server/conn/handler_monoio/dispatch.rs +++ b/src/server/conn/handler_monoio/dispatch.rs @@ -1589,10 +1589,20 @@ pub(super) enum BlockingResult { Handled, /// Write error during flush. Caller should return Done. WriteError, + /// c10k A1: the client vanished while blocked. Its registrations are gone; + /// the caller must close the connection WITHOUT writing a reply. + PeerGone, } /// Handle blocking commands (BLPOP, BRPOP, BLMOVE, etc.). -pub(super) async fn try_handle_blocking( +/// +/// c10k A1: `read_buf` is the handler's own accumulation buffer, passed in so +/// the peer watch can carry any bytes the client pipelines behind its blocking +/// command straight back into the parse stream. It holds only the unparsed +/// tail of the current batch here, so appending preserves wire order. +pub(super) async fn try_handle_blocking< + S: monoio::io::AsyncWriteRent + super::idle_park::IdleParkRead, +>( cmd: &[u8], cmd_args: &[Frame], conn: &mut ConnectionState, @@ -1601,6 +1611,7 @@ pub(super) async fn try_handle_blocking( local_leg_write_idxs: &mut Vec, codec: &mut crate::server::codec::RespCodec, write_buf: &mut bytes::BytesMut, + read_buf: &mut bytes::BytesMut, stream: &mut S, shutdown: &CancellationToken, ) -> BlockingResult { @@ -1648,7 +1659,7 @@ pub(super) async fn try_handle_blocking( } } - let blocking_response = handle_blocking_command_monoio( + let outcome = handle_blocking_command_monoio( cmd, cmd_args, conn.selected_db, @@ -1659,9 +1670,19 @@ pub(super) async fn try_handle_blocking( &ctx.dispatch_tx, shutdown, &ctx.spsc_notifiers, + stream, + read_buf, ) .await; + let blocking_response = match outcome { + crate::server::conn::blocking::BlockingOutcome::Reply(frame) => frame, + crate::server::conn::blocking::BlockingOutcome::PeerGone => { + responses.clear(); + return BlockingResult::PeerGone; + } + }; + // Encode blocking response directly codec.encode_frame(&blocking_response, write_buf); responses.clear(); diff --git a/src/server/conn/handler_monoio/idle_park.rs b/src/server/conn/handler_monoio/idle_park.rs index 8620ca0ed..160ec4ada 100644 --- a/src/server/conn/handler_monoio/idle_park.rs +++ b/src/server/conn/handler_monoio/idle_park.rs @@ -234,6 +234,27 @@ pub(crate) trait IdleParkRead: AsyncReadRent { /// stage-1 read. Default: nothing (plain TCP has no stream-owned /// buffers); TLS releases its drained wrapper buffers here. fn on_idle_downshift(&mut self) {} + + /// c10k A1: await read-readiness WITHOUT consuming anything. + /// + /// Two properties make this the right primitive for watching a blocked + /// client's socket: + /// * it is cancel-safe by construction — a readiness poll owns no + /// buffer, so losing the `select!` race cannot lose client bytes + /// (`spawn_parked_idle_watcher` relies on the same property); + /// * once it fires, the fd is level-triggered ready, so the follow-up + /// read completes immediately with data, EOF or an error and can be + /// issued OUTSIDE the `select!`, where nothing can cancel it. + /// + /// The default never resolves, which leaves streams without a race-free + /// standalone readiness await on exactly the pre-A1 behaviour. TLS + /// deliberately keeps that default: `Stream::read_inner` loops on + /// `read_io` until rustls yields plaintext, so a post-readiness read can + /// park inside a partial record — the one thing the design above must + /// never do. See the module docs on `blocking.rs` for the consequence. + fn peer_readable(&self) -> impl std::future::Future> { + std::future::pending() + } } impl IdleParkRead for monoio::net::TcpStream { @@ -247,6 +268,13 @@ impl IdleParkRead for monoio::net::TcpStream { ) -> impl std::future::Future>> { monoio::io::CancelableAsyncReadRent::cancelable_read(self, buf, c) } + + /// `relaxed = false`: a real `poll(2)`, no false positives. A spurious + /// wake would cost a wasted read, not correctness, but the blocked-client + /// watcher must never mistake one for EOF. + fn peer_readable(&self) -> impl std::future::Future> { + self.readable(false) + } } /// TLS streams downshift too (c10k P4b): the vendored monoio-rustls diff --git a/src/server/conn/handler_monoio/mod.rs b/src/server/conn/handler_monoio/mod.rs index 25f7d32f2..e8a2ba776 100644 --- a/src/server/conn/handler_monoio/mod.rs +++ b/src/server/conn/handler_monoio/mod.rs @@ -265,6 +265,9 @@ pub(crate) async fn handle_connection_sharded_monoio< buf }; let mut write_buf = BytesMut::with_capacity(8192); + // c10k A1: set when a blocking command leaves unparsed input in + // `read_buf` (see the read-skip guard in the main loop). + let mut carried_input = false; let mut codec = RespCodec::default(); let mut conn = super::core::ConnectionState::new( client_id, @@ -686,7 +689,17 @@ pub(crate) async fn handle_connection_sharded_monoio< if tmp_buf.len() != park_len { tmp_buf.resize(park_len, 0); } - if let Some(dur) = idle_timeout { + // c10k A1: a blocking command's peer watch may have pulled bytes the + // client pipelined behind it out of the kernel and into `read_buf`. + // Pre-A1 those bytes stayed in the socket, so the read below returned + // them at once; now they are already here, and parking in read() first + // would hang the pipelined command until the client happened to send + // more. Skip exactly one read and let the parser drain them. A carry + // that is only a partial frame parses to nothing, `frames.is_empty()` + // sends us straight back here, and the flag is already cleared. + if std::mem::take(&mut carried_input) { + // Nothing to read: `read_buf` already holds unparsed input. + } else if let Some(dur) = idle_timeout { // Timeout-aware read: select between read and sleep. // monoio::select! drops the losing future, so tmp_buf ownership transfers. // We allocate a fresh buffer when timeout is enabled (safety feature, not hot path). @@ -1400,6 +1413,7 @@ pub(crate) async fn handle_connection_sharded_monoio< &mut local_leg_write_idxs, &mut codec, &mut write_buf, + &mut read_buf, &mut stream, &shutdown, ) @@ -1407,8 +1421,17 @@ pub(crate) async fn handle_connection_sharded_monoio< { dispatch::BlockingResult::NotBlocking => {} dispatch::BlockingResult::Queued => continue, - dispatch::BlockingResult::Handled => break, + dispatch::BlockingResult::Handled => { + // A1: anything left in read_buf is either this batch's + // unparsed tail or bytes the peer watch carried; parse + // before the next read either way. + carried_input = !read_buf.is_empty(); + break; + } dispatch::BlockingResult::WriteError => return (MonoioHandlerResult::Done, None), + // c10k A1: peer vanished mid-block. Nothing to write; the + // registry entry and maxclients slot are released by returning. + dispatch::BlockingResult::PeerGone => return (MonoioHandlerResult::Done, None), } // --- MULTI queue mode: queue commands when in transaction --- diff --git a/src/server/conn/handler_sharded/mod.rs b/src/server/conn/handler_sharded/mod.rs index 13b46bf70..b7a55fcb9 100644 --- a/src/server/conn/handler_sharded/mod.rs +++ b/src/server/conn/handler_sharded/mod.rs @@ -299,6 +299,9 @@ pub(crate) async fn handle_connection_sharded_inner< buf }; let mut write_buf = BytesMut::with_capacity(8192); + // c10k A1: set when a blocking command leaves unparsed input in + // `read_buf` (see the read-skip guard in the main select). + let mut carried_input = false; let parse_config = crate::protocol::ParseConfig::default(); let mut conn = super::core::ConnectionState::new( client_id, @@ -387,7 +390,17 @@ pub(crate) async fn handle_connection_sharded_inner< } tokio::select! { result = async { - if let Some(dur) = idle_timeout { + // c10k A1: a blocking command's peer watch may have pulled + // bytes the client pipelined behind it out of the kernel and + // into `read_buf`. Pre-A1 they stayed in the socket, so this + // read returned them at once; now they are already here and + // awaiting a read would hang the pipelined command. Report + // them as if just read — a carry that is only a partial frame + // parses to nothing and the loop comes straight back with the + // flag already cleared. + if std::mem::take(&mut carried_input) && !read_buf.is_empty() { + Ok(read_buf.len()) + } else if let Some(dur) = idle_timeout { match tokio::time::timeout(dur, stream.read_buf(&mut read_buf)).await { Ok(r) => r, Err(_) => Err(std::io::Error::new(std::io::ErrorKind::TimedOut, "idle timeout")), @@ -1009,13 +1022,31 @@ pub(crate) async fn handle_connection_sharded_inner< } } if stream.write_all(&write_buf).await.is_err() { arena.reset(); return (HandlerResult::Done, None); } - let blocking_response = handle_blocking_command( + // c10k A1: `read_buf` doubles as the carry buffer — it + // holds only the unparsed tail of this batch here, so + // bytes the client pipelines while blocked append in + // wire order. + let blocking_outcome = handle_blocking_command( cmd, cmd_args, conn.selected_db, &ctx.shard_databases, &ctx.blocking_registry, ctx.shard_id, ctx.num_shards, &ctx.dispatch_tx, &shutdown, + &mut stream, &mut read_buf, ).await; + let blocking_response = match blocking_outcome { + crate::server::conn::blocking::BlockingOutcome::Reply(frame) => frame, + // Peer vanished mid-block: registrations are torn + // down, nothing to reply to, close the connection. + crate::server::conn::blocking::BlockingOutcome::PeerGone => { + arena.reset(); + return (HandlerResult::Done, None); + } + }; let blocking_response = apply_resp3_conversion(cmd, blocking_response, conn.protocol_version); responses = Vec::with_capacity(1); responses.push(blocking_response); + // A1: anything left in read_buf is either this batch's + // unparsed tail or bytes the peer watch carried; parse + // before awaiting the next read either way. + carried_input = !read_buf.is_empty(); break; } diff --git a/tests/blocking_peer_eof.rs b/tests/blocking_peer_eof.rs new file mode 100644 index 000000000..fd773a719 --- /dev/null +++ b/tests/blocking_peer_eof.rs @@ -0,0 +1,340 @@ +//! c10k hardening A1 — a client that vanishes while blocked must be reaped. +//! +//! The defect: the infinite-wait `select!` behind BLPOP/BRPOP/BLMOVE/BZPOP* +//! had exactly two arms, `reply_rx` and `shutdown`. Nothing watched the +//! socket, so `BLPOP key 0` followed by a disconnect left the handler task, +//! the `WaitEntry`, the client-registry entry and the maxclients slot alive +//! forever — infinite waiters carry `deadline: None`, so the deadline sweep +//! never reaps them either, and `timeout` exempts blocked clients by design +//! (Redis parity). A few thousand throwaway connections wedged the server +//! until restart, using one unauthenticated command. +//! +//! The fix adds a peer-watch arm. These tests pin both halves: the slot is +//! released when the peer goes away, and a client that legitimately pipelines +//! behind its blocking command still gets every reply (the watch consumes +//! from the socket, so those bytes have to be carried into the parse stream — +//! pre-A1 they simply waited in the kernel). +//! +//! Run with: +//! cargo test --release --test blocking_peer_eof + +mod common; + +use std::io::{Read, Write}; +use std::net::TcpStream; +use std::process::{Child, Command, Stdio}; +use std::time::{Duration, Instant}; + +fn spawn(dir: &std::path::Path, port: u16, shards: &str) -> Child { + Command::new(common::find_moon_binary()) + .args([ + "--port", + &port.to_string(), + "--shards", + shards, + "--dir", + dir.to_str().unwrap(), + "--disk-free-min-pct", + "0", + // Without an explicit cap moon auto-sizes maxmemory to ~80% of + // host RAM and provisions a multi-GB per-shard page cache. Four + // of these starting at once makes startup itself flaky; the test + // stores a handful of tiny keys, so cap it small. + "--maxmemory", + "268435456", + ]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .expect("spawn moon") +} + +fn connect(port: u16) -> TcpStream { + let s = TcpStream::connect(("127.0.0.1", port)).expect("connect"); + s.set_read_timeout(Some(Duration::from_secs(10))) + .expect("read timeout"); + s.set_nodelay(true).expect("nodelay"); + s +} + +fn send(stream: &mut TcpStream, parts: &[&str]) { + let mut out = format!("*{}\r\n", parts.len()); + for p in parts { + out.push_str(&format!("${}\r\n{}\r\n", p.len(), p)); + } + stream.write_all(out.as_bytes()).expect("write"); +} + +fn read_some(stream: &mut TcpStream) -> std::io::Result { + let mut buf = [0u8; 8192]; + let n = stream.read(&mut buf)?; + Ok(String::from_utf8_lossy(&buf[..n]).into_owned()) +} + +fn connected_clients(port: u16) -> u64 { + let mut c = connect(port); + send(&mut c, &["INFO", "clients"]); + let body = read_some(&mut c).expect("INFO reply"); + for line in body.lines() { + if let Some(rest) = line.strip_prefix("connected_clients:") { + return rest.trim().parse().unwrap_or(0); + } + } + panic!("INFO clients has no connected_clients; got:\n{body}"); +} + +/// Poll until `connected_clients` drops to `at_most`, or give up. Returns the +/// last value seen so failures report the real number. +fn wait_for_clients(port: u16, at_most: u64, budget: Duration) -> u64 { + let start = Instant::now(); + let mut last = u64::MAX; + while start.elapsed() < budget { + last = connected_clients(port); + // The probe connection above is itself counted while it is open. + if last <= at_most { + return last; + } + std::thread::sleep(Duration::from_millis(200)); + } + last +} + +struct Server { + child: Child, + port: u16, + dir: std::path::PathBuf, +} + +impl Drop for Server { + fn drop(&mut self) { + common::sigkill(&mut self.child); + let _ = std::fs::remove_dir_all(&self.dir); + } +} + +fn server(tag: &str, shards: &str) -> Option { + let bin = common::find_moon_binary(); + if !bin.exists() { + eprintln!("skipping: no moon binary; build with `cargo build --release`"); + return None; + } + let (child, port) = common::spawn_listening(|port| { + let dir = std::env::temp_dir().join(format!("moon-{tag}-{port}")); + let _ = std::fs::create_dir_all(&dir); + spawn(&dir, port, shards) + }); + let dir = std::env::temp_dir().join(format!("moon-{tag}-{port}")); + Some(Server { child, port, dir }) +} + +/// THE A1 REGRESSION TEST (single-key). Pre-fix every one of these +/// connections stayed connected forever. +#[test] +fn disconnected_single_key_blocked_clients_are_reaped() { + let Some(srv) = server("a1-single", "1") else { + return; + }; + const N: usize = 24; + + for i in 0..N { + let mut c = connect(srv.port); + // Block forever on a key nobody will ever push to. + send(&mut c, &["BLPOP", &format!("a1:single:{i}"), "0"]); + // Give the server a moment to actually register the wait before the + // disconnect — otherwise the test could pass by never blocking. + std::thread::sleep(Duration::from_millis(20)); + drop(c); + } + + // Only the INFO probe connection itself should remain. + let remaining = wait_for_clients(srv.port, 2, Duration::from_secs(15)); + assert!( + remaining <= 2, + "{N} disconnected `BLPOP key 0` clients must release their slots; \ + connected_clients is still {remaining} — this is the A1 leak" + ); + + // And the server must still be fully serviceable. + let mut probe = connect(srv.port); + send(&mut probe, &["PING"]); + assert!( + read_some(&mut probe).expect("pong").starts_with("+PONG"), + "server must still serve after the disconnect storm" + ); +} + +/// Same leak via the multi-key coordinator, across shards so the waiter holds +/// REMOTE registrations too — those are torn down by `BlockCancel`, a +/// different cleanup path from the single-key one. +#[test] +fn disconnected_multi_key_blocked_clients_are_reaped() { + let Some(srv) = server("a1-multi", "4") else { + return; + }; + const N: usize = 16; + + for i in 0..N { + let mut c = connect(srv.port); + // Several keys with different hash slots — at 4 shards this fans out + // to remote owners, exercising the BlockCancel unwind. + send( + &mut c, + &[ + "BLPOP", + &format!("a1:m:{i}:alpha"), + &format!("a1:m:{i}:beta"), + &format!("a1:m:{i}:gamma"), + &format!("a1:m:{i}:delta"), + "0", + ], + ); + std::thread::sleep(Duration::from_millis(20)); + drop(c); + } + + let remaining = wait_for_clients(srv.port, 2, Duration::from_secs(15)); + assert!( + remaining <= 2, + "{N} disconnected multi-key blocked clients must release their slots; \ + connected_clients is still {remaining}" + ); + + // The registrations must be gone too, not just the sockets: a leaked + // ghost waiter would swallow this push instead of leaving it on the list. + let mut pusher = connect(srv.port); + send(&mut pusher, &["RPUSH", "a1:m:0:alpha", "payload"]); + let _ = read_some(&mut pusher); + send(&mut pusher, &["LLEN", "a1:m:0:alpha"]); + let llen = read_some(&mut pusher).unwrap_or_default(); + assert!( + llen.starts_with(":1"), + "the pushed element must stay on the list — a ghost waiter swallowed \ + it instead (LLEN said {llen:?})" + ); +} + +/// A live blocked client must NOT be mistaken for a dead one. This is the +/// inverse failure mode of the fix: the peer watch reads from the socket, so +/// a client that pipelines behind its blocking command must still be woken, +/// and must still get the replies to what it pipelined. +#[test] +fn pipelining_behind_a_blocking_command_still_works() { + let Some(srv) = server("a1-carry", "1") else { + return; + }; + let mut blocker = connect(srv.port); + + send(&mut blocker, &["BLPOP", "a1:carry", "0"]); + std::thread::sleep(Duration::from_millis(200)); + + // Legal RESP: more commands arrive while the client is blocked. Redis + // runs them after the block resolves; the peer watch must carry them, not + // eat them and not treat them as a disconnect. + send(&mut blocker, &["PING"]); + std::thread::sleep(Duration::from_millis(200)); + + // The blocked client must still be considered alive and wakeable. + let mut pusher = connect(srv.port); + send(&mut pusher, &["RPUSH", "a1:carry", "value"]); + let _ = read_some(&mut pusher); + + // Expect the BLPOP reply, then the pipelined PONG. They may or may not + // land in the same read, so accumulate. + blocker + .set_read_timeout(Some(Duration::from_secs(5))) + .expect("read timeout"); + let mut got = String::new(); + let deadline = Instant::now() + Duration::from_secs(5); + while Instant::now() < deadline && !(got.contains("value") && got.contains("PONG")) { + match read_some(&mut blocker) { + Ok(s) if s.is_empty() => break, + Ok(s) => got.push_str(&s), + Err(_) => break, + } + } + assert!( + got.contains("value"), + "blocked client must still be woken by the push; got {got:?}" + ); + assert!( + got.contains("PONG"), + "the command pipelined while blocked must not be swallowed by the \ + peer watch — it has to be carried into the parse stream; got {got:?}" + ); +} + +/// The carry may be a PARTIAL frame. The read-skip guard must then fall +/// through to a real read instead of spinning on unparseable bytes or +/// stalling until the client happens to send something else. +#[test] +fn partial_frame_carried_while_blocked_completes_later() { + let Some(srv) = server("a1-partial", "1") else { + return; + }; + let mut blocker = connect(srv.port); + send(&mut blocker, &["BLPOP", "a1:partial", "0"]); + std::thread::sleep(Duration::from_millis(200)); + + // Half of `PING`, sent while blocked — the peer watch will carry these + // bytes, and they cannot be parsed on their own. + blocker.write_all(b"*1\r\n$4\r\nPI").expect("write partial"); + std::thread::sleep(Duration::from_millis(200)); + + let mut pusher = connect(srv.port); + send(&mut pusher, &["RPUSH", "a1:partial", "value"]); + let _ = read_some(&mut pusher); + + blocker + .set_read_timeout(Some(Duration::from_secs(5))) + .expect("read timeout"); + let mut got = read_some(&mut blocker).unwrap_or_default(); + assert!( + got.contains("value"), + "blocked client must be woken even with a partial frame carried; got {got:?}" + ); + + // Now complete the frame; the server must have kept the partial bytes. + blocker.write_all(b"NG\r\n").expect("write rest"); + let deadline = Instant::now() + Duration::from_secs(5); + while Instant::now() < deadline && !got.contains("PONG") { + match read_some(&mut blocker) { + Ok(s) if s.is_empty() => break, + Ok(s) => got.push_str(&s), + Err(_) => break, + } + } + assert!( + got.contains("PONG"), + "the carried partial frame must complete once the rest arrives — the \ + carried bytes were dropped instead; got {got:?}" + ); +} + +/// A blocked client that stays connected and silent must stay blocked. Guards +/// against the peer watch spinning or self-triggering on an idle socket. +#[test] +fn silent_blocked_client_stays_blocked() { + let Some(srv) = server("a1-idle", "1") else { + return; + }; + let mut blocker = connect(srv.port); + send(&mut blocker, &["BLPOP", "a1:idle", "0"]); + + // Sit still well past any internal sweep tick. + std::thread::sleep(Duration::from_secs(4)); + + // Still registered? A push must reach it. + let mut pusher = connect(srv.port); + send(&mut pusher, &["RPUSH", "a1:idle", "late"]); + let _ = read_some(&mut pusher); + + blocker + .set_read_timeout(Some(Duration::from_secs(5))) + .expect("read timeout"); + let woken = read_some(&mut blocker).unwrap_or_default(); + assert!( + woken.contains("late"), + "an idle-but-connected blocked client must stay blocked and wakeable; \ + got {woken:?}" + ); +}