From 030c6b1bb03dd13b15acc0c83b74d13e2156eed3 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Wed, 29 Jul 2026 12:28:10 +0700 Subject: [PATCH 01/17] refactor(server): delete vestigial pending_wakers relay + dead conn_state module (c10k W4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The pending_wakers waker relay (Rc>>) had ZERO registrants anywhere in the tree — the only waker push in the codebase is runtime/cancel.rs on a different structure. It was retired as the reply-wake mechanism when swf0 proved cross-thread oneshot wakes work under monoio's `sync` feature (M2), but the plumbing survived: allocated per shard, cloned per accepted connection (2 spawn paths + 2 migration fail-open paths), threaded through 4 call sites, and drained twice per event-loop iteration for nothing. Also deletes src/server/conn_state.rs (209 lines): a duplicate ConnectionState/ConnectionContext pair marked "Phase 44: Defined only", never adopted, zero references. Found by the 2026-07-29 c10k connection-plane review (tmp/C10K-REVIEW.md); this is workstream W4 of tmp/C10K-PLAN.md. Both runtimes compile clean (cargo check default + runtime-tokio,jemalloc). CLAUDE.md monoio-waker note updated to match. author: Tin Dang --- CLAUDE.md | 2 +- src/server/conn/handler_monoio/mod.rs | 5 - src/server/conn_state.rs | 209 -------------------------- src/server/mod.rs | 1 - src/shard/conn_accept.rs | 13 +- src/shard/event_loop.rs | 31 +--- 6 files changed, 11 insertions(+), 250 deletions(-) delete mode 100644 src/server/conn_state.rs diff --git a/CLAUDE.md b/CLAUDE.md index fb9103533..17c39e67b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -145,7 +145,7 @@ orb run -m moon-dev bash -c 'sudo apt-get update -qq && sudo apt-get install -y - Never hold a lock across `.await` points. - Replace `.read().unwrap()` / `.write().unwrap()` with `.read()` / `.write()` (parking_lot doesn't poison). - Per-shard locks only — no global locks on the write path. -- **monoio cross-thread wakers:** `monoio::spawn` creates `!Send` tasks; `Waker::wake()` from another OS thread does NOT reach them. The cross-shard **reply** path therefore has the connection **await a `flume` oneshot directly** (the target shard sends the reply on it after executing). A `pending_wakers: Rc>>` relay is still swept each event-loop iteration (pub/sub + backpressure paths thread it), but it is **no longer the reply-wake mechanism** — the old "register your waker, event loop drains it after SPSC" design was retired when test `swf0` disproved its premise (see `event_loop.rs` ~1730). For cross-thread signalling, prefer `flume::bounded(1)` over custom atomic oneshots. +- **monoio cross-thread wakers:** `monoio::spawn` creates `!Send` tasks; `Waker::wake()` from another OS thread does NOT reach them. The cross-shard **reply** path therefore has the connection **await a `flume` oneshot directly** (the target shard sends the reply on it after executing). The old `pending_wakers` relay ("register your waker, event loop drains it after SPSC") was retired when test `swf0` disproved its premise, and the vestigial plumbing was **deleted entirely** in the c10k wave (it had zero registrants). For cross-thread signalling, prefer `flume::bounded(1)` over custom atomic oneshots. ### Error Handling - All command errors return `Frame::Error(Bytes)` — no `Result` types in dispatch paths. diff --git a/src/server/conn/handler_monoio/mod.rs b/src/server/conn/handler_monoio/mod.rs index 164b8895d..1aa94dd68 100644 --- a/src/server/conn/handler_monoio/mod.rs +++ b/src/server/conn/handler_monoio/mod.rs @@ -177,11 +177,6 @@ pub(crate) async fn handle_connection_sharded_monoio< client_id: u64, can_migrate: bool, initial_read_buf: BytesMut, - // Event-loop waker relay. No longer used by this handler — the cross-shard - // reply path awaits its oneshot directly (M2, spsc-wake-floor; cross-thread - // wake proven by swf0). The relay plumbing is kept for future registrants - // and the event loop still sweeps it every iteration. - _pending_wakers: Rc>>, migrated_state: Option<&MigratedConnectionState>, // Raw socket fd for CLIENT KILL force-close (R-3), or -1 if unavailable // (non-unix). Threaded from the concrete spawn site; the generic `S` here diff --git a/src/server/conn_state.rs b/src/server/conn_state.rs deleted file mode 100644 index 10e121b70..000000000 --- a/src/server/conn_state.rs +++ /dev/null @@ -1,209 +0,0 @@ -use std::cell::RefCell; -use std::rc::Rc; -use std::sync::{Arc, RwLock}; - -use bytes::Bytes; -use ringbuf::HeapProd; - -use crate::blocking::BlockingRegistry; -use crate::cluster::ClusterState; -use crate::config::{RuntimeConfig, ServerConfig}; -use crate::persistence::aof::AofWriterPool; -use crate::protocol::Frame; -use crate::pubsub::PubSubRegistry; -use crate::replication::state::ReplicationState; -use crate::runtime::cancel::CancellationToken; -use crate::runtime::channel; -use crate::scripting::ScriptCache; -use crate::shard::dispatch::ShardMessage; -use crate::shard::shared_databases::ShardDatabases; -use crate::storage::entry::CachedClock; -use crate::tracking::{TrackingState, TrackingTable}; -use crate::transaction::CrossStoreTxn; - -/// Shared immutable state provided to each connection handler by the shard. -/// -/// Bundles the ~20 parameters currently passed individually to -/// `handle_connection_sharded_inner`. All fields use shared ownership -/// (`Rc`, `Arc`) so cloning is cheap. -/// -/// **Phase 44:** Defined only. Adoption in connection handlers deferred to Phase 48. -#[allow(dead_code)] -pub struct ConnectionContext { - pub shard_databases: Arc, - pub shard_id: usize, - pub num_shards: usize, - pub dispatch_tx: Rc>>>, - pub pubsub_registry: Rc>, - pub blocking_registry: Rc>, - pub shutdown: CancellationToken, - pub requirepass: Option, - /// Per-shard AOF writer pool. **Sole AOF interface** after the - /// 2d/2e migration sequence. Built by spawn sites in `shard/conn_accept.rs` - /// from the manifest layout — TopLevel wraps a single sender, - /// PerShard owns one sender per shard. Step 2f flips spawn sites - /// to construct PerShard pools when the on-disk manifest demands it. - pub aof_pool: Option>, - pub tracking_table: Rc>, - pub repl_state: Option>>, - pub cluster_state: Option>>, - pub lua: Rc, - pub script_cache: Rc>, - pub func_registry: Rc>, - pub config_port: u16, - pub acl_table: Arc>, - pub runtime_config: Arc>, - pub config: Arc, - pub spsc_notifiers: Vec>, - pub snapshot_trigger_tx: channel::WatchSender, - pub cached_clock: CachedClock, -} - -/// Per-connection mutable state capturing auth, transaction, pub/sub, and -/// protocol modes as a flat struct with explicit fields. -/// -/// Created by `ConnectionState::new()` at connection accept time. -/// Mirrors the local variables currently declared at the top of each handler body. -/// -/// **Phase 44:** Defined only. Adoption in connection handlers deferred to Phase 48. -pub struct ConnectionState { - /// RESP protocol version (2 or 3, set by HELLO command) - pub protocol_version: u8, - /// Currently selected database index - pub selected_db: usize, - /// Whether the client has authenticated (true if no password required) - pub authenticated: bool, - /// ACL username for the current session - pub current_user: String, - /// Whether the client is in a MULTI transaction block - pub in_multi: bool, - /// Active cross-store transaction (None if not in transaction). - /// Mutually exclusive with in_multi (MULTI/EXEC is KV-only). - pub active_cross_txn: Option, - /// Queued commands for MULTI/EXEC - pub command_queue: Vec, - /// Client-side tracking state - pub tracking_state: TrackingState, - /// Receiver for tracking invalidation messages - pub tracking_rx: Option>, - /// Client name set by CLIENT SETNAME - pub client_name: Option, - /// Cluster ASKING flag (set by ASKING, cleared before routing check) - pub asking: bool, - /// Pub/sub subscription count (>0 means client is in subscriber mode) - pub subscription_count: usize, - /// Subscriber ID for pub/sub (0 means not subscribed) - pub subscriber_id: u64, - /// Pub/sub message sender - pub pubsub_tx: Option>, - /// Pub/sub message receiver - pub pubsub_rx: Option>, - /// Unique client ID assigned at connection accept - pub client_id: u64, - /// Peer address string for logging - pub peer_addr: String, -} - -impl ConnectionState { - /// Create a new ConnectionState for a freshly accepted connection. - /// - /// `authenticated` is set to `true` if no requirepass is configured. - pub fn new(client_id: u64, peer_addr: String, requirepass_set: bool) -> Self { - Self { - protocol_version: 2, - selected_db: 0, - authenticated: !requirepass_set, - current_user: "default".to_string(), - in_multi: false, - active_cross_txn: None, - command_queue: Vec::new(), - tracking_state: TrackingState::default(), - tracking_rx: None, - client_name: None, - asking: false, - subscription_count: 0, - subscriber_id: 0, - pubsub_tx: None, - pubsub_rx: None, - client_id, - peer_addr, - } - } - - /// Returns true if the client is in pub/sub subscriber mode. - pub fn is_subscriber(&self) -> bool { - self.subscription_count > 0 - } - - /// Returns true if the client is in a MULTI transaction block. - pub fn is_in_transaction(&self) -> bool { - self.in_multi - } - - /// Check if connection is in a cross-store transaction. - #[inline] - pub fn in_cross_txn(&self) -> bool { - self.active_cross_txn.is_some() - } - - /// Get the active transaction's ID, if any. - #[inline] - #[allow(dead_code)] // API reserved for future handler-level TXN integration - pub fn cross_txn_id(&self) -> Option { - self.active_cross_txn.as_ref().map(|t| t.txn_id) - } - - /// Get the active transaction's snapshot LSN, if any. - #[inline] - #[allow(dead_code)] // API reserved for future handler-level TXN integration - pub fn cross_txn_snapshot(&self) -> Option { - self.active_cross_txn.as_ref().map(|t| t.snapshot_lsn) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn new_connection_state_defaults() { - let state = ConnectionState::new(42, "127.0.0.1:9999".to_string(), false); - assert_eq!(state.protocol_version, 2); - assert_eq!(state.selected_db, 0); - assert!(state.authenticated); // no requirepass - assert_eq!(state.current_user, "default"); - assert!(!state.in_multi); - assert!(state.command_queue.is_empty()); - assert!(!state.asking); - assert_eq!(state.subscription_count, 0); - assert_eq!(state.client_id, 42); - assert_eq!(state.peer_addr, "127.0.0.1:9999"); - } - - #[test] - fn new_connection_state_with_requirepass() { - let state = ConnectionState::new(1, "10.0.0.1:1234".to_string(), true); - assert!(!state.authenticated); // requirepass is set - } - - #[test] - fn is_subscriber_false_by_default() { - let state = ConnectionState::new(1, "x".to_string(), false); - assert!(!state.is_subscriber()); - } - - #[test] - fn is_subscriber_true_when_subscribed() { - let mut state = ConnectionState::new(1, "x".to_string(), false); - state.subscription_count = 3; - assert!(state.is_subscriber()); - } - - #[test] - fn is_in_transaction_reflects_multi() { - let mut state = ConnectionState::new(1, "x".to_string(), false); - assert!(!state.is_in_transaction()); - state.in_multi = true; - assert!(state.is_in_transaction()); - } -} diff --git a/src/server/mod.rs b/src/server/mod.rs index b93c9d608..562666685 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -1,7 +1,6 @@ pub mod accept_backoff; pub mod codec; pub mod conn; -pub mod conn_state; #[cfg(feature = "runtime-tokio")] pub mod embedded; pub mod expiration; diff --git a/src/shard/conn_accept.rs b/src/shard/conn_accept.rs index 61a6304ea..7813de54c 100644 --- a/src/shard/conn_accept.rs +++ b/src/shard/conn_accept.rs @@ -554,7 +554,6 @@ pub(crate) fn spawn_monoio_connection( shard_id: usize, num_shards: usize, config_port: u16, - pending_wakers: &Rc>>, spill_sender: &Option>, spill_file_id: &Rc>, disk_offload_dir: &Option, @@ -626,7 +625,6 @@ pub(crate) fn spawn_monoio_connection( let notifiers = all_notifiers.to_vec(); let snap_tx = snapshot_trigger_tx.clone(); let clk = cached_clock.clone(); - let pw = pending_wakers.clone(); let all_regs = all_pubsub_registries.to_vec(); let all_rsm = all_remote_sub_maps.to_vec(); @@ -696,7 +694,6 @@ pub(crate) fn spawn_monoio_connection( cid, false, // can_migrate: TLS connections cannot transfer session state BytesMut::new(), - pw, None, // fresh connection kill_fd, ) @@ -736,7 +733,7 @@ pub(crate) fn spawn_monoio_connection( // R-6 fail-open: keep what we need to re-serve the client // locally if a migration hand-off cannot be delivered. #[cfg(target_os = "linux")] - let (sd2, pw2, peer2) = (sd.clone(), pw.clone(), peer_addr.clone()); + let (sd2, peer2) = (sd.clone(), peer_addr.clone()); let _result = handle_connection_sharded_monoio( tcp_stream, peer_addr, @@ -745,7 +742,6 @@ pub(crate) fn spawn_monoio_connection( cid, cfg!(target_os = "linux"), // can_migrate: FD dup requires libc (Linux only) BytesMut::new(), - pw, None, // fresh connection kill_fd, ) @@ -838,7 +834,7 @@ pub(crate) fn spawn_monoio_connection( let _ = handle_connection_sharded_monoio( stream, peer2, &conn_ctx, sd2, cid, false, // can_migrate - BytesMut::new(), pw2, + BytesMut::new(), Some(&state), kill_fd, ) .await; @@ -894,7 +890,7 @@ pub(crate) fn spawn_monoio_connection( let _ = handle_connection_sharded_monoio( stream, peer2, &conn_ctx, sd2, cid, false, // can_migrate: pin locally, no retry loop - BytesMut::new(), pw2, + BytesMut::new(), Some(&payload.state), kill_fd, ) .await; @@ -964,7 +960,6 @@ pub(crate) fn spawn_migrated_monoio_connection( shard_id: usize, num_shards: usize, config_port: u16, - pending_wakers: &Rc>>, spill_sender: &Option>, spill_file_id: &Rc>, disk_offload_dir: &Option, @@ -1038,7 +1033,6 @@ pub(crate) fn spawn_migrated_monoio_connection( let all_regs = all_pubsub_registries.to_vec(); let all_rsm = all_remote_sub_maps.to_vec(); let aff = pubsub_affinity.clone(); - let pw = pending_wakers.clone(); let spill_tx = spill_sender.clone(); let spill_fid = spill_file_id.clone(); let do_dir = disk_offload_dir.clone(); @@ -1095,7 +1089,6 @@ pub(crate) fn spawn_migrated_monoio_connection( cid, false, // can_migrate: already-migrated connections skip re-migration sampling migration_buf, - pw, Some(&state), kill_fd, ) diff --git a/src/shard/event_loop.rs b/src/shard/event_loop.rs index 6b5bd70ce..4e5662b41 100644 --- a/src/shard/event_loop.rs +++ b/src/shard/event_loop.rs @@ -1238,16 +1238,12 @@ impl super::Shard { } } - // Waker relay, swept after every drain cycle. HISTORICAL NOTE: this was - // built on the belief that monoio's !Send executor cannot receive - // cross-thread Waker::wake() — that is FALSE with the `sync` feature - // (enabled in Cargo.toml): a remote wake rides a per-thread waker - // channel + driver unpark (eventfd/kqueue), proven at runtime by - // tests/spsc_wake_floor_red.rs::swf0. The hot cross-shard reply path - // now awaits its oneshot directly (handler_monoio, M2); the relay is - // kept for future registrants and stays correct either way. - #[cfg(feature = "runtime-monoio")] - let pending_wakers: Rc>> = Rc::new(RefCell::new(Vec::new())); + // NOTE: the old `pending_wakers` relay (register-waker, event-loop + // sweeps after SPSC drain) was deleted in the c10k wave — it had zero + // registrants since M2 made the cross-shard reply path await its + // oneshot directly (cross-thread wake via monoio's `sync` feature, + // proven by tests/spsc_wake_floor_red.rs::swf0). For cross-thread + // signalling prefer flume oneshots, not a waker relay. // R-4: backoff for the tokio per-shard SO_REUSEPORT accept branch so an // fd-exhaustion storm can't hot-spin this shard's select loop. @@ -1468,7 +1464,6 @@ impl super::Shard { &cached_clock, &remote_sub_map_arc, &all_pubsub_registries, &all_remote_sub_maps, &affinity_tracker, shard_id, num_shards, config_port, - &pending_wakers, &spill_sender, &spill_file_id, &disk_offload_dir, ); } @@ -1574,7 +1569,6 @@ impl super::Shard { &cached_clock, &remote_sub_map_arc, &all_pubsub_registries, &all_remote_sub_maps, &affinity_tracker, shard_id, num_shards, config_port, - &pending_wakers, &spill_sender, &spill_file_id, &disk_offload_dir, ); } @@ -1955,7 +1949,6 @@ impl super::Shard { shard_id, num_shards, config_port, - &pending_wakers, &spill_sender, &spill_file_id, &disk_offload_dir, @@ -1997,17 +1990,11 @@ impl super::Shard { shard_id, num_shards, config_port, - &pending_wakers, &spill_sender, &spill_file_id, &disk_offload_dir, ); } - // Wake cross-shard response tasks that registered during the previous iteration. - #[cfg(feature = "runtime-monoio")] - for waker in pending_wakers.borrow_mut().drain(..) { - waker.wake(); - } // Monoio runtime: direct-await on 1ms periodic tick. // AVOID monoio::select! — it leaks ~100 bytes per re-entry (internal future @@ -2131,7 +2118,7 @@ impl super::Shard { } // --- Every-wake body (mirrors the tokio notify arm): drain SPSC, - // handle drain outputs, sweep the pending_wakers relay --- + // handle drain outputs --- let mut pending_snapshot = None; // No outer with_shard — each arm takes its own flat borrow. let hit_cap = spsc_handler::drain_spsc_shared( @@ -2179,9 +2166,6 @@ impl super::Shard { let wal_dir = wal_writer.as_ref().map(|w| w.wal_dir()); cdc_registry.register_pending(pending_cdc_subscribes.drain(..), wal_dir); } - for waker in pending_wakers.borrow_mut().drain(..) { - waker.wake(); - } persistence_tick::handle_pending_snapshot( pending_snapshot, &mut snapshot_state, @@ -2231,7 +2215,6 @@ impl super::Shard { shard_id, num_shards, config_port, - &pending_wakers, &spill_sender, &spill_file_id, &disk_offload_dir, From 01491f09170c3d2a21bed9b73646f3a73def2e7b Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Wed, 29 Jul 2026 12:33:36 +0700 Subject: [PATCH 02/17] perf(server): fix per-connection pipeline memory ratchet (c10k W1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Empirically confirmed (tmp/C10K-REVIEW.md E5): one 1024-deep pipeline permanently ratcheted a connection from 56.5 KB to ~217 KB RSS until disconnect — 5000 such conns pinned 1.09 GB, stable after 35s idle. ~147 KB of that is the `responses`/`frames` Vec pair growing to the 1024-frame batch cap (72 B/Frame) and only ever `.clear()`ing. Fixes, monoio handler: - Clear + shrink both batch scratch vecs at END of the batch iteration (before parking in read()), via new `util::shrink_batch_vec` — shrink fires only above 256-frame capacity, so p99 small batches never pay a realloc; a burst-then-idle conn now parks at steady-state capacity. (Shrinking at the top-of-loop clear sites is too late: an idle conn parks in read() before ever reaching them.) - I/O buffer shrink floor lowered 64 KiB → 16 KiB (util:: IO_BUF_SHRINK_TRIGGER): 16–64 KiB high-waters previously ratcheted until disconnect. - Subscriber mode: rent the loop-level tmp_buf via mem::take instead of allocating + zeroing a fresh 8 KiB per select iteration (the `__memset_zva64` flame entry, 2.5% at c10k). The buffer is only lost when the pubsub arm wins (dropped read op owns it — io_uring cancel semantics); command traffic now allocates nothing. Fixes, tokio handler: - Same 16 KiB shrink floor for read/write buffers. - Bump arena capped: reset() retains the largest chunk, so one huge batch pinned its high-water until disconnect; >64 KiB arenas are now rebuilt at 4 KiB. Red/green: 3 new unit tests for shrink_batch_vec thresholds; conn suite 42/42 green; both runtimes compile clean. RSS A/B re-run of E5/E1 gates the branch before merge. Workstream W1 of tmp/C10K-PLAN.md. author: Tin Dang --- src/server/conn/handler_monoio/mod.rs | 34 +++++++++++--- src/server/conn/handler_sharded/mod.rs | 18 ++++++-- src/server/conn/util.rs | 62 ++++++++++++++++++++++++++ 3 files changed, 106 insertions(+), 8 deletions(-) diff --git a/src/server/conn/handler_monoio/mod.rs b/src/server/conn/handler_monoio/mod.rs index 1aa94dd68..d6d7fa67c 100644 --- a/src/server/conn/handler_monoio/mod.rs +++ b/src/server/conn/handler_monoio/mod.rs @@ -300,10 +300,21 @@ pub(crate) async fn handle_connection_sharded_monoio< #[allow(clippy::unwrap_used)] // conn.pubsub_rx is always Some when conn.subscription_count > 0 let rx = conn.pubsub_rx.as_ref().unwrap(); - let sub_tmp_buf = vec![0u8; 8192]; + // c10k W1: rent the main loop's tmp_buf instead of allocating a + // fresh zeroed 8 KiB per select iteration. The buffer is only + // LOST when the pubsub arm wins (the dropped read op owns it — + // io_uring cancel semantics); the refill below re-arms it. A + // read-win restores ownership, so command traffic allocates + // nothing. Stale bytes past `n` are never read (`&buf[..n]`). + if tmp_buf.len() < 8192 { + tmp_buf.resize(8192, 0); + } + let sub_tmp_buf = std::mem::take(&mut tmp_buf); monoio::select! { read_result = stream.read(sub_tmp_buf) => { let (result, buf) = read_result; + tmp_buf = buf; + let buf = &tmp_buf; match result { Ok(0) => { // Client half-closed — break out of loop. @@ -2433,20 +2444,33 @@ pub(crate) async fn handle_connection_sharded_monoio< break; } - // Shrink buffers if they grew too large - if read_buf.capacity() > 65536 { + // Shrink buffers if they grew too large (c10k W1: floor lowered + // 64 KiB → 16 KiB — the old floor let 16–64 KiB high-waters ratchet + // until disconnect; see tmp/C10K-REVIEW.md). + if read_buf.capacity() > super::util::IO_BUF_SHRINK_TRIGGER { let remaining = read_buf.split(); read_buf = BytesMut::with_capacity(8192); if !remaining.is_empty() { read_buf.extend_from_slice(&remaining); } } - if write_buf.capacity() > 65536 { + if write_buf.capacity() > super::util::IO_BUF_SHRINK_TRIGGER { write_buf = BytesMut::with_capacity(8192); } - if tmp_buf.capacity() > 65536 { + if tmp_buf.capacity() > super::util::IO_BUF_SHRINK_TRIGGER { tmp_buf = vec![0u8; 8192]; } + // c10k W1: drop the batch scratch capacity BEFORE parking in read(). + // Both vecs are dead scratch here (responses fully serialized to + // write_buf, frames fully dispatched); clearing at the top of the + // next iteration is too late for a burst-then-idle connection, which + // parks holding the 1024-frame high-water (~74 KB each — the E5 + // permanent-ratchet finding). Sustained >256-frame pipelines pay one + // shrink+regrow per batch (~sub-µs vs a 1024-command batch). + responses.clear(); + super::util::shrink_batch_vec(&mut responses); + frames.clear(); + super::util::shrink_batch_vec(&mut frames); } // --- Graceful TCP shutdown: send FIN to client to avoid CLOSE_WAIT --- diff --git a/src/server/conn/handler_sharded/mod.rs b/src/server/conn/handler_sharded/mod.rs index 219d06896..2ad389519 100644 --- a/src/server/conn/handler_sharded/mod.rs +++ b/src/server/conn/handler_sharded/mod.rs @@ -2136,7 +2136,15 @@ pub(crate) async fn handle_connection_sharded_inner< } } - arena.reset(); + // c10k W1: cap the bump arena before reuse. `reset()` retains + // the largest chunk, so one huge batch would otherwise pin its + // high-water allocation until disconnect (the tokio-side twin + // of the monoio batch-vec ratchet, tmp/C10K-REVIEW.md E5). + if arena.allocated_bytes() > 65536 { + arena = bumpalo::Bump::with_capacity(4096); + } else { + arena.reset(); + } // AUTH rate limiting: delay response to slow down brute-force attacks if auth_delay_ms > 0 { @@ -2190,8 +2198,12 @@ pub(crate) async fn handle_connection_sharded_inner< ); } - if write_buf.capacity() > 65536 { write_buf = BytesMut::with_capacity(8192); } - if read_buf.capacity() > 65536 { + // c10k W1: shrink floor lowered 64 KiB → 16 KiB (see + // super::util::IO_BUF_SHRINK_TRIGGER). + if write_buf.capacity() > super::util::IO_BUF_SHRINK_TRIGGER { + write_buf = BytesMut::with_capacity(8192); + } + if read_buf.capacity() > super::util::IO_BUF_SHRINK_TRIGGER { let remaining = read_buf.split(); read_buf = BytesMut::with_capacity(8192); if !remaining.is_empty() { read_buf.extend_from_slice(&remaining); } diff --git a/src/server/conn/util.rs b/src/server/conn/util.rs index b76fa6360..dd2ae9e22 100644 --- a/src/server/conn/util.rs +++ b/src/server/conn/util.rs @@ -103,3 +103,65 @@ pub(crate) fn unpropagate_subscription( .remove(channel, shard_id, is_pattern); } } + +/// Post-batch capacity governor for the per-connection batch scratch vectors +/// (c10k W1). `responses`/`frames` are cleared and reused across batches; one +/// deep pipeline grows them to the 1024-frame batch cap (~74 KB each at +/// 72 B/Frame) and `.clear()` never returns that capacity — measured as a +/// permanent ~160 KB/conn RSS ratchet until disconnect (tmp/C10K-REVIEW.md, +/// experiment E5). Shrinking only above the trigger keeps every realloc off +/// the small-batch path: a p99 batch (≤64 frames) never grows past the +/// trigger, so it never pays a shrink. +pub(crate) const BATCH_VEC_STEADY_CAP: usize = 64; +pub(crate) const BATCH_VEC_SHRINK_TRIGGER: usize = 256; + +/// Shrink an emptied batch scratch vector back to steady-state capacity. +/// Call only after `clear()` — shrinking an empty Vec moves no elements. +#[inline] +pub(crate) fn shrink_batch_vec(v: &mut Vec) { + debug_assert!(v.is_empty(), "shrink_batch_vec expects a cleared vec"); + if v.capacity() > BATCH_VEC_SHRINK_TRIGGER { + v.shrink_to(BATCH_VEC_STEADY_CAP); + } +} + +/// I/O buffer shrink floor (c10k W1): a connection that once carried a large +/// value keeps its high-water read/write/rent buffer until disconnect. The +/// old floor (64 KiB) let 16–64 KiB high-waters ratchet forever; 16 KiB +/// bounds steady-state at 2× the 8 KiB working size while still amortizing +/// growth for genuinely large frames in flight. +pub(crate) const IO_BUF_SHRINK_TRIGGER: usize = 16384; + +#[cfg(test)] +mod shrink_tests { + use super::*; + + #[test] + fn oversized_batch_vec_shrinks_to_steady_cap() { + let mut v: Vec = Vec::with_capacity(1024); + v.clear(); + shrink_batch_vec(&mut v); + assert!( + v.capacity() <= BATCH_VEC_SHRINK_TRIGGER, + "capacity {} must shrink below trigger", + v.capacity() + ); + assert!(v.capacity() >= BATCH_VEC_STEADY_CAP); + } + + #[test] + fn steady_state_batch_vec_is_untouched() { + let mut v: Vec = Vec::with_capacity(BATCH_VEC_SHRINK_TRIGGER); + let cap_before = v.capacity(); + shrink_batch_vec(&mut v); + assert_eq!(v.capacity(), cap_before, "at-trigger capacity must not shrink"); + } + + #[test] + fn small_batch_vec_is_untouched() { + let mut v: Vec = Vec::with_capacity(BATCH_VEC_STEADY_CAP); + let cap_before = v.capacity(); + shrink_batch_vec(&mut v); + assert_eq!(v.capacity(), cap_before); + } +} From 85d5c590ad4a4ba99aa86a1fa60f273df661db2d Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Wed, 29 Jul 2026 12:37:25 +0700 Subject: [PATCH 03/17] perf(server): box active_cross_txn out of every connection future (c10k W2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CrossStoreTxn is ~2.2 KB of inline SmallVec capacity (kv_undo SmallVec<[UndoRecord;16]> alone is 1160 B). Stored unboxed as Option in ConnectionState, it sat inside EVERY connection's ~21 KB task future — transacting or not — as part of the measured 56.5 KB/idle-conn footprint (tmp/C10K-REVIEW.md §2). Option> moves that to the heap for the rare transacting connection only; TXN.BEGIN (cold path, not in the no-alloc dispatch list) pays one allocation. Call-site churn is minimal: field access auto-derefs, Option<&CrossStoreTxn> consumers switch as_ref() → as_deref(), by-value abort/commit sites deref the box. New regression guard: connection_state_stays_small asserts size_of::() ≤ 768 B (was ~2.7 KB). txn (46) + transaction (39) suites green; both runtimes compile clean. Workstream W2 of tmp/C10K-PLAN.md. author: Tin Dang --- src/server/conn/core.rs | 5 ++++- src/server/conn/handler_monoio/ft.rs | 6 +++--- src/server/conn/handler_monoio/mod.rs | 2 +- src/server/conn/handler_monoio/txn.rs | 6 +++--- src/server/conn/handler_sharded/ft.rs | 6 +++--- src/server/conn/handler_sharded/mod.rs | 2 +- src/server/conn/handler_sharded/txn.rs | 6 +++--- src/server/conn/tests.rs | 13 +++++++++++++ 8 files changed, 31 insertions(+), 15 deletions(-) diff --git a/src/server/conn/core.rs b/src/server/conn/core.rs index ae195618b..21350d7e6 100644 --- a/src/server/conn/core.rs +++ b/src/server/conn/core.rs @@ -255,7 +255,10 @@ pub(crate) struct ConnectionState { pub in_multi: bool, /// Active cross-store transaction (None if not in transaction). /// Mutually exclusive with in_multi (MULTI/EXEC is KV-only). - pub active_cross_txn: Option, + /// Boxed (c10k W2): CrossStoreTxn is ~2.2 KB of inline SmallVecs which + /// otherwise sits in EVERY connection's task future, transacting or not + /// (tmp/C10K-REVIEW.md §2). TXN.BEGIN (cold path) pays one heap alloc. + pub active_cross_txn: Option>, /// Active workspace binding for this connection (None = no workspace context). /// Set by WS.AUTH, cleared on connection drop. pub workspace_id: Option, diff --git a/src/server/conn/handler_monoio/ft.rs b/src/server/conn/handler_monoio/ft.rs index fb7021e0d..b73c445c4 100644 --- a/src/server/conn/handler_monoio/ft.rs +++ b/src/server/conn/handler_monoio/ft.rs @@ -297,7 +297,7 @@ pub(super) async fn try_handle_ft_command( let as_of_lsn = match resolve_ft_search_as_of_lsn( cmd_args, Some(&ctx.shard_databases), - conn.active_cross_txn.as_ref(), + conn.active_cross_txn.as_deref(), ) { Ok(lsn) => lsn, Err(err_frame) => { @@ -405,7 +405,7 @@ pub(super) async fn try_handle_ft_command( match resolve_ft_search_as_of_lsn( cmd_args, Some(&ctx.shard_databases), - conn.active_cross_txn.as_ref(), + conn.active_cross_txn.as_deref(), ) { Err(err_frame) => err_frame, Ok(as_of_lsn) => { @@ -760,7 +760,7 @@ pub(super) async fn try_handle_ft_command( resolve_ft_search_as_of_lsn( cmd_args, Some(&ctx.shard_databases), - conn.active_cross_txn.as_ref(), + conn.active_cross_txn.as_deref(), ) } else { Ok(0u64) diff --git a/src/server/conn/handler_monoio/mod.rs b/src/server/conn/handler_monoio/mod.rs index d6d7fa67c..a8e0a869b 100644 --- a/src/server/conn/handler_monoio/mod.rs +++ b/src/server/conn/handler_monoio/mod.rs @@ -2492,7 +2492,7 @@ pub(crate) async fn handle_connection_sharded_monoio< ctx.num_shards, &ctx.dispatch_tx, &ctx.spsc_notifiers, - txn, + *txn, ) .await; } diff --git a/src/server/conn/handler_monoio/txn.rs b/src/server/conn/handler_monoio/txn.rs index acad7821c..1a19f9237 100644 --- a/src/server/conn/handler_monoio/txn.rs +++ b/src/server/conn/handler_monoio/txn.rs @@ -33,11 +33,11 @@ pub(super) fn try_handle_txn_begin( // Get next txn_id and snapshot_lsn from vector store's transaction manager let active = crate::shard::slice::with_shard(|s| s.vector_store.txn_manager_mut().begin()); - conn.active_cross_txn = Some(CrossStoreTxn::new( + conn.active_cross_txn = Some(Box::new(CrossStoreTxn::new( active.txn_id, active.snapshot_lsn, conn.selected_db, - )); + ))); responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); } Err(e) => responses.push(e), @@ -280,7 +280,7 @@ pub(super) async fn try_handle_txn_abort( ctx.num_shards, &ctx.dispatch_tx, &ctx.spsc_notifiers, - txn, + *txn, ) .await; responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); diff --git a/src/server/conn/handler_sharded/ft.rs b/src/server/conn/handler_sharded/ft.rs index 11d5d6f6a..6a794bc85 100644 --- a/src/server/conn/handler_sharded/ft.rs +++ b/src/server/conn/handler_sharded/ft.rs @@ -126,7 +126,7 @@ pub(super) async fn try_handle_ft_command( let as_of_lsn = match resolve_ft_search_as_of_lsn( cmd_args, Some(&ctx.shard_databases), - conn.active_cross_txn.as_ref(), + conn.active_cross_txn.as_deref(), ) { Ok(lsn) => lsn, Err(err_frame) => { @@ -231,7 +231,7 @@ pub(super) async fn try_handle_ft_command( match resolve_ft_search_as_of_lsn( cmd_args, Some(&ctx.shard_databases), - conn.active_cross_txn.as_ref(), + conn.active_cross_txn.as_deref(), ) { Err(err_frame) => err_frame, Ok(as_of_lsn) => { @@ -461,7 +461,7 @@ pub(super) async fn try_handle_ft_command( Some(resolve_ft_search_as_of_lsn( cmd_args, Some(&ctx.shard_databases), - conn.active_cross_txn.as_ref(), + conn.active_cross_txn.as_deref(), )) } else { None diff --git a/src/server/conn/handler_sharded/mod.rs b/src/server/conn/handler_sharded/mod.rs index 2ad389519..4e97d9169 100644 --- a/src/server/conn/handler_sharded/mod.rs +++ b/src/server/conn/handler_sharded/mod.rs @@ -2254,7 +2254,7 @@ pub(crate) async fn handle_connection_sharded_inner< ctx.num_shards, &ctx.dispatch_tx, &ctx.spsc_notifiers, - txn, + *txn, ) .await; } diff --git a/src/server/conn/handler_sharded/txn.rs b/src/server/conn/handler_sharded/txn.rs index c04ee9e1a..879c76acc 100644 --- a/src/server/conn/handler_sharded/txn.rs +++ b/src/server/conn/handler_sharded/txn.rs @@ -34,11 +34,11 @@ pub(super) fn try_handle_txn_begin( // Unconditional slice path: ShardSlice is always initialized. let active = crate::shard::slice::with_shard(|s| s.vector_store.txn_manager_mut().begin()); - conn.active_cross_txn = Some(CrossStoreTxn::new( + conn.active_cross_txn = Some(Box::new(CrossStoreTxn::new( active.txn_id, active.snapshot_lsn, conn.selected_db, - )); + ))); responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); } Err(e) => responses.push(e), @@ -275,7 +275,7 @@ pub(super) async fn try_handle_txn_abort( ctx.num_shards, &ctx.dispatch_tx, &ctx.spsc_notifiers, - txn, + *txn, ) .await; responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); diff --git a/src/server/conn/tests.rs b/src/server/conn/tests.rs index d9617b70a..bde58493e 100644 --- a/src/server/conn/tests.rs +++ b/src/server/conn/tests.rs @@ -562,3 +562,16 @@ fn test_inline_get_genuine_miss_still_answers_inline() { assert!(read_buf.is_empty()); assert_eq!(&write_buf[..], b"$-1\r\n"); } + +/// c10k W2 regression guard: `active_cross_txn` must stay boxed. Unboxed, +/// CrossStoreTxn's inline SmallVecs put ~2.2 KB into EVERY connection's task +/// future (tmp/C10K-REVIEW.md §2). If this assert fires, something re-inlined +/// large state into ConnectionState — box it instead. +#[test] +fn connection_state_stays_small() { + let sz = std::mem::size_of::(); + assert!( + sz <= 768, + "ConnectionState is {sz} B — keep bulky fields boxed (was 2.7 KB before c10k W2)" + ); +} From 50dfefe81ec6edfa3caff5c65e303eeac4dcf72b Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Wed, 29 Jul 2026 12:46:15 +0700 Subject: [PATCH 04/17] fix(server): loud maxclients rejection + early gate + RLIMIT_NOFILE check (c10k W3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Redis parity break found by the 2026-07-29 c10k review: the monoio accept paths rejected over-cap connections with a warn! and a silent close — a client at the cap (default: exactly 10000) saw an unexplained EOF. The tokio plain-TCP path already wrote the error; monoio (default runtime) and both TLS paths did not. - monoio: the maxclients CAS now runs BEFORE the per-connection clone-fest and ConnectionContext construction — a rejected connection previously paid 3 O(num_shards) Vec clones, ~25 Arc refcount bumps and a spawned handler task before the gate ran. On rejection moon now writes `-ERR max number of clients reached\r\n` best-effort, then closes. TLS keeps slot-before-handshake semantics (#17); the in-task gates are removed (slot taken pre-spawn, released at task end as before — rejected conns were never counted, so accounting is unchanged). - tokio TLS: same best-effort error write on rejection (plaintext write surfaces client-side as a loud handshake failure, matching Redis's raw-fd behavior). - startup (unix): getrlimit(RLIMIT_NOFILE) vs maxclients + reserved fds (rlimit_reserved_fds: 64 + 16/shard for listeners/WAL/AOF/spill/uring). Soft limit is raised toward the hard limit when short (logged); if the hard limit still cannot fit, a loud warning states the real client ceiling. Previously moon never looked at RLIMIT_NOFILE and a 10k target on a 1024-fd shell died in silent EMFILE accept-backoff loops. Red/green: new tests/maxclients_reject_parity.rs — rejected_connection_receives_err_max_clients (red before: silent EOF; green after) and slot_frees_on_disconnect (guards the accounting restructure). Unit test for rlimit_reserved_fds. Both runtimes compile. Workstream W3 of tmp/C10K-PLAN.md. author: Tin Dang --- src/main.rs | 59 ++++++++++++++ src/shard/conn_accept.rs | 56 ++++++++++---- tests/maxclients_reject_parity.rs | 124 ++++++++++++++++++++++++++++++ 3 files changed, 224 insertions(+), 15 deletions(-) create mode 100644 tests/maxclients_reject_parity.rs diff --git a/src/main.rs b/src/main.rs index d5f4d9437..f753e6da8 100644 --- a/src/main.rs +++ b/src/main.rs @@ -617,6 +617,47 @@ fn main() -> anyhow::Result<()> { tracing::warn!("{msg}"); } + // c10k W3: make maxclients honest against RLIMIT_NOFILE (Redis parity — + // Redis raises the soft limit to fit maxclients or shrinks maxclients + // loudly; moon previously never looked, so a 10k-conn target on a 1024-fd + // shell died in silent EMFILE accept-backoff loops). Non-unix: no-op. + #[cfg(unix)] + { + let reserved = rlimit_reserved_fds(num_shards); + // SAFETY: getrlimit/setrlimit with a valid resource constant and a + // pointer to a properly initialized rlimit struct — plain libc FFI + // with no aliasing or lifetime concerns. + unsafe { + let mut rl = libc::rlimit { + rlim_cur: 0, + rlim_max: 0, + }; + if libc::getrlimit(libc::RLIMIT_NOFILE, &mut rl) == 0 { + let need = config.maxclients as u64 + reserved; + if config.maxclients > 0 && u64::from(rl.rlim_cur) < need { + let want = need.min(u64::from(rl.rlim_max)); + let mut raised = rl; + raised.rlim_cur = want as libc::rlim_t; + if libc::setrlimit(libc::RLIMIT_NOFILE, &raised) == 0 { + tracing::info!( + "Raised RLIMIT_NOFILE soft limit {} -> {} to fit maxclients {} (+{} reserved fds)", + rl.rlim_cur, want, config.maxclients, reserved + ); + } + if want < need { + tracing::warn!( + "RLIMIT_NOFILE hard limit {} cannot fit maxclients {} (+{} reserved fds): \ + accepts beyond ~{} clients will fail with EMFILE. Raise `ulimit -n` \ + or lower --maxclients.", + rl.rlim_max, config.maxclients, reserved, + want.saturating_sub(reserved) + ); + } + } + } + } + } + // Create channel mesh for inter-shard communication let mut mesh = ChannelMesh::new(num_shards, CHANNEL_BUFFER_SIZE); @@ -2042,10 +2083,28 @@ pub fn should_warn_undersubscription(maxclients: usize, num_shards: usize) -> Op } } +/// Fds moon needs beyond client sockets (c10k W3 rlimit check): listeners +/// (per-shard SO_REUSEPORT + central), per-shard WAL/AOF segments + spill + +/// manifest handles, io_uring fds, logs. Deliberately generous — the cost of +/// over-reserving is a slightly earlier warning, the cost of under-reserving +/// is EMFILE on the persistence path mid-flight. +pub fn rlimit_reserved_fds(num_shards: usize) -> u64 { + 64 + (num_shards as u64) * 16 +} + #[cfg(test)] mod tests { + use super::rlimit_reserved_fds; use super::should_warn_undersubscription; + #[test] + fn reserved_fds_scale_with_shards() { + assert_eq!(rlimit_reserved_fds(1), 80); + assert_eq!(rlimit_reserved_fds(16), 320); + // Never zero — listeners + logs always exist. + assert!(rlimit_reserved_fds(0) >= 64); + } + #[test] fn no_warn_single_shard() { assert!(should_warn_undersubscription(100, 1).is_none()); diff --git a/src/shard/conn_accept.rs b/src/shard/conn_accept.rs index 7813de54c..e12cae513 100644 --- a/src/shard/conn_accept.rs +++ b/src/shard/conn_accept.rs @@ -272,6 +272,13 @@ pub(crate) fn spawn_tokio_connection( "Shard {}: TLS connection rejected: maxclients reached", shard_id ); + // c10k W3 Redis parity: loud rejection. The plaintext write + // surfaces client-side as a handshake failure (best effort). + use tokio::io::AsyncWriteExt; + let mut rejected = tcp_stream; + let _ = rejected + .write_all(b"-ERR max number of clients reached\r\n") + .await; return; } // R-3: capture the underlying TCP fd before the handshake moves the @@ -579,6 +586,34 @@ pub(crate) fn spawn_monoio_connection( }; #[cfg(not(unix))] let kill_fd = -1; + + // c10k W3: maxclients gate BEFORE the per-connection clone-fest + + // ConnectionContext construction. A rejected connection previously + // paid 3 O(num_shards) Vec clones, ~25 Arc bumps and a spawned + // handler task before the CAS ran — and was then closed SILENTLY. + // Redis parity: write `-ERR max number of clients reached` (best + // effort; for TLS clients the plaintext write surfaces as a loud + // handshake failure, matching Redis's raw-fd behavior), then + // close. The slot was never taken, so no record_connection_closed. + let maxclients = runtime_config.read().maxclients; + if !crate::admin::metrics_setup::try_accept_connection(maxclients) { + tracing::warn!( + "Shard {}: connection rejected: maxclients reached", + shard_id + ); + monoio::spawn(async move { + use monoio::io::AsyncWriteRentExt; + let mut rejected = tcp_stream; + let (_res, _buf) = rejected + .write_all(bytes::Bytes::from_static( + b"-ERR max number of clients reached\r\n", + )) + .await; + // drop closes the fd + }); + return; + } + let aff = affinity_tracker.clone(); let rsm = remote_subscriber_map.clone(); let sdbs = shard_databases.clone(); @@ -667,18 +702,14 @@ pub(crate) fn spawn_monoio_connection( do_dir, ); - let maxclients = conn_ctx.runtime_config.read().maxclients; if let (true, Some(tls_swap)) = (is_tls, tls_config.as_ref()) { // Load current TLS config from ArcSwap — new connections see reloaded certs let tls_cfg = tls_swap.load_full(); monoio::spawn(async move { - if !crate::admin::metrics_setup::try_accept_connection(maxclients) { - tracing::warn!( - "Shard {}: TLS connection rejected: maxclients reached", - shard_id - ); - return; - } + // maxclients slot already taken by the pre-spawn gate above + // (still BEFORE the handshake, preserving #17's stalled- + // ClientHello bound); released by record_connection_closed + // at the end of this task. let acceptor = monoio_rustls::TlsAcceptor::from(tls_cfg); // #17: bound the handshake so a stalled ClientHello cannot // pin the fd + maxclients slot forever. @@ -723,13 +754,8 @@ pub(crate) fn spawn_monoio_connection( #[cfg(target_os = "linux")] let notifiers2 = all_notifiers.to_vec(); monoio::spawn(async move { - if !crate::admin::metrics_setup::try_accept_connection(maxclients) { - tracing::warn!( - "Shard {}: connection rejected: maxclients reached", - shard_id - ); - return; - } + // maxclients slot already taken by the pre-spawn gate above; + // released by the migration-aware decrement at task end. // R-6 fail-open: keep what we need to re-serve the client // locally if a migration hand-off cannot be delivered. #[cfg(target_os = "linux")] diff --git a/tests/maxclients_reject_parity.rs b/tests/maxclients_reject_parity.rs new file mode 100644 index 000000000..4a42b2b5e --- /dev/null +++ b/tests/maxclients_reject_parity.rs @@ -0,0 +1,124 @@ +//! c10k W3: maxclients rejection must be LOUD (Redis parity). +//! +//! Redis writes `-ERR max number of clients reached\r\n` before closing a +//! connection that exceeds `maxclients`. Moon used to log a warn! and close +//! silently — a client at the cap saw an unexplained EOF (found by the +//! 2026-07-29 c10k review, tmp/C10K-REVIEW.md defect #1). + +mod common; + +use std::io::{Read, Write}; +use std::net::TcpStream; +use std::process::{Command, Stdio}; +use std::time::Duration; + +fn spawn_moon(dir: &std::path::Path, port: u16, maxclients: u32) -> std::process::Child { + Command::new(common::find_moon_binary()) + .args([ + "--port", + &port.to_string(), + "--shards", + "1", + "--maxclients", + &maxclients.to_string(), + "--dir", + dir.to_str().unwrap(), + "--disk-free-min-pct", + "0", + ]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .expect("spawn moon") +} + +fn ping_ok(stream: &mut TcpStream) -> bool { + stream.set_read_timeout(Some(Duration::from_secs(5))).ok(); + if stream.write_all(b"PING\r\n").is_err() { + return false; + } + let mut buf = [0u8; 16]; + matches!(stream.read(&mut buf), Ok(n) if buf[..n].starts_with(b"+PONG")) +} + +/// Acquire the single maxclients slot, retrying while transient holders +/// (spawn_listening's liveness probe, a prior conn's async deregistration) +/// release it. +fn acquire_slot(port: u16) -> TcpStream { + let deadline = std::time::Instant::now() + Duration::from_secs(10); + loop { + let mut conn = TcpStream::connect(("127.0.0.1", port)).expect("connect"); + if ping_ok(&mut conn) { + return conn; + } + assert!( + std::time::Instant::now() < deadline, + "could not acquire the maxclients slot within 10s" + ); + std::thread::sleep(Duration::from_millis(100)); + } +} + +#[test] +fn rejected_connection_receives_err_max_clients() { + let dir = tempfile::tempdir().expect("tempdir"); + let (mut child, port) = common::spawn_listening(|p| spawn_moon(dir.path(), p, 1)); + + // Conn A takes the single slot and must be fully served. + let mut conn_a = acquire_slot(port); + + // Conn B exceeds maxclients: it must receive the Redis-parity error + // before the server closes it — not a silent EOF. + let mut conn_b = + TcpStream::connect(("127.0.0.1", port)).expect("conn B TCP-connects (accept then reject)"); + conn_b + .set_read_timeout(Some(Duration::from_secs(10))) + .expect("set timeout"); + let mut buf = Vec::new(); + let mut chunk = [0u8; 256]; + loop { + match conn_b.read(&mut chunk) { + Ok(0) => break, + Ok(n) => buf.extend_from_slice(&chunk[..n]), + Err(e) => panic!( + "conn B read failed before any reply (silent rejection): {e}; got {:?}", + String::from_utf8_lossy(&buf) + ), + } + } + let reply = String::from_utf8_lossy(&buf); + assert!( + reply.starts_with("-ERR max number of clients reached"), + "over-cap connection must get the Redis parity error, got: {reply:?}" + ); + + // The slot-holder must be unaffected by the rejection. + assert!(ping_ok(&mut conn_a), "conn A must survive conn B's rejection"); + + drop(conn_a); + common::sigkill(&mut child); +} + +#[test] +fn slot_frees_on_disconnect() { + let dir = tempfile::tempdir().expect("tempdir"); + let (mut child, port) = common::spawn_listening(|p| spawn_moon(dir.path(), p, 1)); + + let conn_a = acquire_slot(port); + drop(conn_a); + + // After A disconnects, the slot must free (poll: deregistration is async). + let deadline = std::time::Instant::now() + Duration::from_secs(10); + loop { + let mut conn_c = TcpStream::connect(("127.0.0.1", port)).expect("conn C"); + if ping_ok(&mut conn_c) { + break; + } + assert!( + std::time::Instant::now() < deadline, + "slot never freed after disconnect" + ); + std::thread::sleep(Duration::from_millis(100)); + } + common::sigkill(&mut child); +} From e3821e4c6c451bcf499d498d79cfb1b875c6d62b Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Wed, 29 Jul 2026 12:49:24 +0700 Subject: [PATCH 05/17] perf(blocking): deadline-heap sweep replaces full-registry walk (c10k W6) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit BlockingRegistry::expire_timed_out ran at 100 Hz per shard and walked EVERY queue of EVERY blocked waiter — two full passes (scan + retain) — even when nothing was due: ~2M entry visits/s/shard at 10k blocked BLPOP clients, quadratic when many waiters shared a key (tmp/C10K-REVIEW.md defect #3). There was no deadline-ordered structure at all. Now a min-heap of (deadline, db_index, key) indexes every waiter that HAS a deadline. The sweep pops due entries and scans only the queues they name; block-forever waiters (BLPOP key 0) never enter the heap and are never visited. Stale heap entries (waiter served or cancelled before its deadline) pop to a no-op at their original deadline — bounded by registration rate x timeout, the same envelope as the queues themselves. Steady state with nothing due is a single peek. expire_timed_out now returns the visited-entry count as the test/ observability surface. Red/green: 4 new unit tests — sweep_skips_undeadlined_waiters (100 forever-waiters + 1 due => exactly 1 visit), shared-key partial expiry preserves FIFO, served-waiter stale-entry no-op, multi-key waiter expires from all queues. Full blocking suite 8/8. Workstream W6 of tmp/C10K-PLAN.md. author: Tin Dang --- src/blocking/mod.rs | 153 ++++++++++++++++++++++++++++++++++++++++++-- src/shard/timers.rs | 7 +- 2 files changed, 152 insertions(+), 8 deletions(-) diff --git a/src/blocking/mod.rs b/src/blocking/mod.rs index 0891cb1f7..2dd437962 100644 --- a/src/blocking/mod.rs +++ b/src/blocking/mod.rs @@ -66,6 +66,17 @@ pub struct BlockingRegistry { waiters: HashMap<(usize, Bytes), VecDeque>, /// wait_id -> list of (db_index, key) for cross-key cleanup on wakeup/timeout. wait_keys: HashMap>, + /// Deadline index (c10k W6): min-heap of (deadline, db_index, key) for + /// every registered waiter that HAS a deadline. `expire_timed_out` used + /// to walk EVERY queue of EVERY blocked waiter at its 100 Hz cadence + /// (~2M entry visits/s/shard at 10k blocked clients, tmp/C10K-REVIEW.md + /// defect #3); with the heap it touches only queues holding an actually + /// -due candidate. Entries are lazily invalidated: a waiter served or + /// cancelled before its deadline leaves a stale heap entry that pops to + /// a no-op at its original deadline — bounded by registration rate × + /// timeout, the same envelope as the queues themselves. Zero-timeout + /// (block-forever) waiters never enter the heap. + deadlines: std::collections::BinaryHeap>, /// Monotonically increasing wait_id counter (lower 48 bits). next_id: u64, /// Shard ID encoded in upper 16 bits of wait_id for global uniqueness. @@ -81,6 +92,7 @@ impl BlockingRegistry { BlockingRegistry { waiters: HashMap::new(), wait_keys: HashMap::new(), + deadlines: std::collections::BinaryHeap::new(), next_id: 0, shard_id, } @@ -100,6 +112,11 @@ impl BlockingRegistry { let wait_id = entry.wait_id; let queue_key = (db_index, key.clone()); + if let Some(deadline) = entry.deadline { + self.deadlines + .push(std::cmp::Reverse((deadline, db_index, key))); + } + self.waiters .entry(queue_key.clone()) .or_insert_with(VecDeque::new) @@ -151,18 +168,36 @@ impl BlockingRegistry { /// Expire all timed-out waiters. Sends None through their reply channels. /// - /// Two-pass approach: first collect timed-out entries, then clean up. - pub fn expire_timed_out(&mut self, now: std::time::Instant) { - // Pass 1: collect timed-out (wait_id, reply_tx) pairs + /// Heap-driven (c10k W6): pops due deadline-index entries and scans ONLY + /// the queues they name; every other blocked waiter is untouched. Runs at + /// 100 Hz from the shard timer, so the no-expiry steady state must stay + /// O(1). Returns the number of queue entries visited (observability + + /// test surface — the old implementation visited every blocked waiter). + pub fn expire_timed_out(&mut self, now: std::time::Instant) -> usize { + let mut visited = 0usize; let mut timed_out: Vec<(u64, crate::runtime::channel::OneshotSender>)> = Vec::new(); let mut timed_out_ids: Vec = Vec::new(); - for queue in self.waiters.values_mut() { + while self + .deadlines + .peek() + .is_some_and(|std::cmp::Reverse((d, _, _))| *d <= now) + { + #[allow(clippy::unwrap_used)] // peek above just proved non-empty + let std::cmp::Reverse((_, db_index, key)) = self.deadlines.pop().unwrap(); + let queue_key = (db_index, key); + // A missing queue is a STALE heap entry (waiter served/cancelled + // before its deadline) — the lazy-invalidation no-op. + let Some(queue) = self.waiters.get_mut(&queue_key) else { + continue; + }; let mut i = 0; while i < queue.len() { + visited += 1; let is_expired = queue[i].deadline.map_or(false, |d| d <= now); if is_expired { + #[allow(clippy::unwrap_used)] // i < queue.len() by loop guard let entry = queue.remove(i).unwrap(); timed_out_ids.push(entry.wait_id); timed_out.push((entry.wait_id, entry.reply_tx)); @@ -170,11 +205,11 @@ impl BlockingRegistry { i += 1; } } + if queue.is_empty() { + self.waiters.remove(&queue_key); + } } - // Remove empty queues - self.waiters.retain(|_, q| !q.is_empty()); - // Send None (timeout) to all timed-out waiters for (_wait_id, reply_tx) in timed_out { let _ = reply_tx.send(None); @@ -187,6 +222,7 @@ impl BlockingRegistry { for id in timed_out_ids { self.wait_keys.remove(&id); } + visited } } @@ -303,3 +339,106 @@ mod tests { assert!(!reg.has_waiters(0, &Bytes::from_static(b"nokey"))); } } + +#[cfg(test)] +mod deadline_heap_tests { + use super::*; + use std::time::{Duration, Instant}; + + fn entry(reg: &mut BlockingRegistry, deadline: Option) -> (u64, WaitEntry) { + let id = reg.next_wait_id(); + let (tx, _rx) = crate::runtime::channel::oneshot(); + ( + id, + WaitEntry { + wait_id: id, + cmd: BlockedCommand::BLPop, + reply_tx: tx, + deadline, + }, + ) + } + + /// c10k W6: the 100 Hz sweep must not touch block-forever waiters. The + /// old implementation walked every queue of every blocked client. + #[test] + fn sweep_skips_undeadlined_waiters() { + let mut reg = BlockingRegistry::new(0); + let now = Instant::now(); + for i in 0..100u32 { + let (_, e) = entry(&mut reg, None); + reg.register(0, Bytes::from(format!("k{i}")), e); + } + let (_, due) = entry(&mut reg, Some(now - Duration::from_millis(1))); + reg.register(0, Bytes::from_static(b"due-key"), due); + + let visited = reg.expire_timed_out(now); + assert_eq!(visited, 1, "only the due queue's single entry is visited"); + assert!(!reg.has_waiters(0, &Bytes::from_static(b"due-key"))); + assert!(reg.has_waiters(0, &Bytes::from_static(b"k0"))); + assert!(reg.has_waiters(0, &Bytes::from_static(b"k99"))); + + // Steady state after the sweep: nothing due, zero visits. + assert_eq!(reg.expire_timed_out(now), 0); + } + + /// Shared key: expire only due entries, preserve FIFO of the rest. + #[test] + fn shared_key_partial_expiry_preserves_fifo() { + let mut reg = BlockingRegistry::new(0); + let now = Instant::now(); + let key = Bytes::from_static(b"shared"); + let (_, e1) = entry(&mut reg, Some(now - Duration::from_millis(5))); + reg.register(0, key.clone(), e1); + let (id2, e2) = entry(&mut reg, Some(now + Duration::from_secs(60))); + reg.register(0, key.clone(), e2); + let (id3, e3) = entry(&mut reg, None); + reg.register(0, key.clone(), e3); + + reg.expire_timed_out(now); + assert_eq!(reg.pop_front(0, &key).unwrap().wait_id, id2); + assert_eq!(reg.pop_front(0, &key).unwrap().wait_id, id3); + assert!(reg.pop_front(0, &key).is_none()); + } + + /// A waiter served before its deadline leaves a stale heap entry — the + /// sweep at its original deadline must be a no-op, not a panic or a + /// spurious timeout reply. + #[test] + fn served_waiter_stale_heap_entry_is_noop() { + let mut reg = BlockingRegistry::new(0); + let now = Instant::now(); + let key = Bytes::from_static(b"served"); + let (id, e) = entry(&mut reg, Some(now + Duration::from_millis(10))); + reg.register(0, key.clone(), e); + + let popped = reg.pop_front(0, &key).expect("served"); + assert_eq!(popped.wait_id, id); + reg.remove_wait(id); + + let visited = reg.expire_timed_out(now + Duration::from_secs(1)); + assert_eq!(visited, 0, "stale heap entry must no-op"); + } + + /// Multi-key waiter (BLPOP k1 k2) with one shared deadline: both queue + /// entries removed in one sweep, one visit each. + #[test] + fn multikey_waiter_expires_from_all_queues() { + let mut reg = BlockingRegistry::new(0); + let now = Instant::now(); + let id = reg.next_wait_id(); + let deadline = Some(now - Duration::from_millis(1)); + for k in [&b"mk1"[..], &b"mk2"[..]] { + let (tx, _rx) = crate::runtime::channel::oneshot(); + reg.register( + 0, + Bytes::copy_from_slice(k), + WaitEntry { wait_id: id, cmd: BlockedCommand::BLPop, reply_tx: tx, deadline }, + ); + } + let visited = reg.expire_timed_out(now); + assert_eq!(visited, 2); + assert!(!reg.has_waiters(0, &Bytes::from_static(b"mk1"))); + assert!(!reg.has_waiters(0, &Bytes::from_static(b"mk2"))); + } +} diff --git a/src/shard/timers.rs b/src/shard/timers.rs index 91e9b1630..9a6a954b4 100644 --- a/src/shard/timers.rs +++ b/src/shard/timers.rs @@ -209,9 +209,14 @@ pub(crate) fn run_eviction( } /// Expire timed-out blocked clients. +/// +/// Heap-driven since c10k W6: cost is O(due + stale-heap-entries), not +/// O(all blocked waiters), so the 100 Hz cadence is safe at 10k+ blocked +/// clients. The visit count is intentionally dropped here — it exists for +/// tests/observability at the registry layer. pub(crate) fn expire_blocked_clients(blocking_rc: &Rc>) { let now = std::time::Instant::now(); - blocking_rc.borrow_mut().expire_timed_out(now); + let _ = blocking_rc.borrow_mut().expire_timed_out(now); } /// Checkpoint tick interval in milliseconds. From e59c93e6b2402c2ebef147c09ff5fc5ed9eeb3de Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Wed, 29 Jul 2026 12:51:24 +0700 Subject: [PATCH 06/17] fix(shard): rotate SPSC drain start to prevent peer-shard starvation (c10k W7) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit drain_spsc_shared always began at consumers[0] with a shared 256-message budget per cycle. Whenever the budget ran out before the tail was reached — sustained cross-shard pressure from a hot peer — low-index peers monopolized every cycle and high-index peers starved indefinitely (R-5 from the 2026-07-06 conn-plane review, re-confirmed by the 2026-07-29 c10k review, tmp/C10K-REVIEW.md defect #4). The drain now starts at a per-thread rotating index (thread-local Cell, same pattern as the existing DRAIN_SCRATCH), so the budget cut-off truncates a DIFFERENT tail each cycle and every peer is head-most once per n cycles. SnapshotBegin early-exit semantics unchanged. The visit order is extracted as rotated_indices(start, n) with unit tests: full coverage per cycle, rotating head, n=0 safety. spsc suite 20/20, shard suite 191/191. Workstream W7 of tmp/C10K-PLAN.md. (The give-up error path was already clean — both slotted dispatch paths return `ERR cross-shard dispatch backpressure` to the client — so W7 is rotation only.) author: Tin Dang --- src/shard/spsc_handler.rs | 57 ++++++++++++++++++++++++++++++++++++++- 1 file changed, 56 insertions(+), 1 deletion(-) diff --git a/src/shard/spsc_handler.rs b/src/shard/spsc_handler.rs index ae30d5f21..be48e235f 100644 --- a/src/shard/spsc_handler.rs +++ b/src/shard/spsc_handler.rs @@ -221,8 +221,27 @@ pub(crate) fn drain_spsc_shared( other_messages.push(msg); } + // c10k W7: rotate the consumer start index per drain call. The loop used + // to start at consumers[0] every cycle, so whenever the 256-message + // budget ran out before the tail was reached, a hot low-index peer shard + // starved high-index peers indefinitely (tmp/C10K-REVIEW.md defect #4). + thread_local! { + static DRAIN_START: Cell = const { Cell::new(0) }; + } + let n = consumers.len(); + let start = if n > 0 { + DRAIN_START.with(|c| { + let s = c.get() % n; + c.set(s + 1); + s + }) + } else { + 0 + }; + let mut snapshot_seen = false; - for consumer in consumers.iter_mut() { + for idx in rotated_indices(start, n) { + let consumer = &mut consumers[idx]; if snapshot_seen { break; } @@ -4456,3 +4475,39 @@ mod drain_cap_tests { assert!(consumers[0].is_empty(), "all 300 messages must be consumed"); } } + +/// c10k W7: visit order for the SPSC consumer drain — a rotation of +/// `0..n` beginning at `start`. Extracted so the fairness property is unit +/// -testable without standing up the full drain call. +#[inline] +pub(crate) fn rotated_indices(start: usize, n: usize) -> impl Iterator { + (0..n).map(move |off| (start + off) % n) +} + +#[cfg(test)] +mod drain_rotation_tests { + use super::rotated_indices; + + #[test] + fn every_cycle_covers_all_consumers_exactly_once() { + for start in 0..5 { + let mut seen: Vec = rotated_indices(start, 5).collect(); + assert_eq!(seen.len(), 5); + seen.sort_unstable(); + assert_eq!(seen, vec![0, 1, 2, 3, 4]); + } + } + + #[test] + fn successive_starts_change_the_first_visited_consumer() { + // The budget cut-off truncates the TAIL of the visit order; rotating + // the head is what guarantees no consumer is permanently tail-most. + let firsts: Vec = (0..4).map(|s| rotated_indices(s, 4).next().unwrap()).collect(); + assert_eq!(firsts, vec![0, 1, 2, 3]); + } + + #[test] + fn zero_consumers_is_empty_not_panic() { + assert_eq!(rotated_indices(0, 0).count(), 0); + } +} From bed94a341ba86b9407398b8b57d8a5bc83fbd0c9 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Wed, 29 Jul 2026 12:55:23 +0700 Subject: [PATCH 07/17] perf(server): stripe the global client registry 16 ways (c10k W5) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The registry was one process-wide write-preferring RwLock: every accept and every close serialized on it across all shard threads (the only cross-shard serialization point on the connect path), and a CLIENT LIST formatted all N clients under the read lock — blocking every pending register/deregister on all shards for the duration (at 1M conns that is a ~128 MB format under a global lock). CLIENT KILL scanned O(N) even for KILL ID on an id-keyed map. (tmp/C10K-REVIEW.md defect #2.) Now 16 stripes keyed by id % 16: - register/deregister/update/is_killed: single-stripe locks (contention /16, and never blocked by a listing of the other 15 stripes). - CLIENT LIST: walks one stripe at a time, pre-sized from a lock-free total counter. Cross-stripe ordering is not insertion order — Redis makes no CLIENT LIST ordering guarantee. - CLIENT KILL ID: O(1) single-stripe lookup. ADDR/USER filters walk stripes one at a time; a conn registering into an already-visited stripe mid-walk is missed — the same inherent race the single-lock version had at command granularity. - The R-3 fd-shutdown liveness invariant is preserved per stripe: an entry visible under its stripe's READ lock blocks that entry's deregister (same stripe's WRITE lock), so kill_fd is still open. New cross-stripe test (list sees all stripes; KILL ID exactly one victim; KILL USER reaches every stripe). client_registry suite 12/12. Workstream W5 of tmp/C10K-PLAN.md. author: Tin Dang --- src/client_registry.rs | 178 ++++++++++++++++++++++++++++++++--------- 1 file changed, 138 insertions(+), 40 deletions(-) diff --git a/src/client_registry.rs b/src/client_registry.rs index 7f004fe2e..6d10bd653 100644 --- a/src/client_registry.rs +++ b/src/client_registry.rs @@ -1,12 +1,14 @@ //! Global client connection registry for CLIENT LIST/INFO/KILL. //! //! Every connection registers on accept and deregisters on close. -//! The registry is a global `parking_lot::RwLock` touched only on -//! connect/disconnect and CLIENT commands. Per-batch state (db, idle time, -//! flags, kill checks) flows through the lock-free [`ClientLiveState`] handle -//! that `register` returns (QW8, 2026-06 review finding 1.3 — previously the -//! steady-state loop took the global write lock after every pipeline batch -//! and the global read lock for every kill check). +//! The registry is a 16-way striped `parking_lot::RwLock` (c10k W5; +//! stripe = `id % 16`) touched only on connect/disconnect and CLIENT +//! commands — accepts/closes contend per-stripe, and CLIENT LIST/KILL walk +//! one stripe at a time instead of freezing the whole registry. Per-batch +//! state (db, idle time, flags, kill checks) flows through the lock-free +//! [`ClientLiveState`] handle that `register` returns (QW8, 2026-06 review +//! finding 1.3 — previously the steady-state loop took the global write lock +//! after every pipeline batch and the global read lock for every kill check). use parking_lot::RwLock; use std::collections::HashMap; @@ -14,9 +16,26 @@ use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, LazyLock}; use std::time::Instant; -/// Global client registry. -static REGISTRY: LazyLock>> = - LazyLock::new(|| RwLock::new(HashMap::new())); +/// Stripe count for the global registry (c10k W5). Power of two, sized so +/// that at 10k conns a stripe holds ~625 entries: register/deregister +/// contention drops 16×, and CLIENT LIST/KILL walks lock ONE stripe at a +/// time instead of stalling every accept/close on all shards behind a +/// single write-preferring RwLock (tmp/C10K-REVIEW.md defect #2). +const STRIPES: usize = 16; + +/// Global client registry, striped by `id % STRIPES`. +static REGISTRY: LazyLock<[RwLock>; STRIPES]> = + LazyLock::new(|| std::array::from_fn(|_| RwLock::new(HashMap::new()))); + +/// Live client count across all stripes — pre-sizes CLIENT LIST output +/// without locking anything. +static TOTAL_CLIENTS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0); + +/// The stripe owning client `id`. +#[inline] +fn stripe(id: u64) -> &'static RwLock> { + ®ISTRY[(id as usize) % STRIPES] +} /// Lock-free per-connection state, shared between the connection task /// (writer, once per batch) and CLIENT LIST/INFO/KILL (occasional readers). @@ -141,14 +160,16 @@ pub fn register( shard, live: Arc::clone(&live), }; - REGISTRY.write().insert(id, entry); + stripe(id).write().insert(id, entry); + TOTAL_CLIENTS.fetch_add(1, Ordering::Relaxed); crate::admin::metrics_setup::record_shard_connection_delta(shard, 1.0); live } /// Deregister a client connection. pub fn deregister(id: u64) { - if let Some(entry) = REGISTRY.write().remove(&id) { + if let Some(entry) = stripe(id).write().remove(&id) { + TOTAL_CLIENTS.fetch_sub(1, Ordering::Relaxed); crate::admin::metrics_setup::record_shard_connection_delta(entry.shard, -1.0); } } @@ -157,7 +178,7 @@ pub fn deregister(id: u64) { /// never the steady-state batch loop; batch state goes through the /// [`ClientLiveState`] handle instead). pub fn update(id: u64, f: F) { - if let Some(entry) = REGISTRY.write().get_mut(&id) { + if let Some(entry) = stripe(id).write().get_mut(&id) { f(entry); } } @@ -167,7 +188,7 @@ pub fn update(id: u64, f: F) { /// Registry-lookup variant for code without the live handle; connection /// loops use `ClientLiveState::is_killed` (lock-free) instead. pub fn is_killed(id: u64) -> bool { - REGISTRY.read().get(&id).is_some_and(|e| e.live.is_killed()) + stripe(id).read().get(&id).is_some_and(|e| e.live.is_killed()) } /// Format all clients as a CLIENT LIST string. @@ -175,11 +196,19 @@ pub fn is_killed(id: u64) -> bool { /// Each line: `id=N addr=... fd=0 name=... db=N ...` /// Returns the full response string. pub fn client_list() -> String { - let registry = REGISTRY.read(); let now = Instant::now(); - let mut result = String::with_capacity(registry.len() * 128); - for entry in registry.values() { - format_client_line(&mut result, entry, now); + // Pre-size from the lock-free total; walk one stripe at a time so a big + // listing never blocks register/deregister on the other 15 stripes + // (c10k W5 — the old single lock held every accept/close hostage for the + // whole O(N) format). Ordering across stripes is not insertion order; + // Redis makes no CLIENT LIST ordering guarantee. + let mut result = + String::with_capacity(TOTAL_CLIENTS.load(Ordering::Relaxed).saturating_mul(128)); + for lock in REGISTRY.iter() { + let stripe = lock.read(); + for entry in stripe.values() { + format_client_line(&mut result, entry, now); + } } // Remove trailing newline if present if result.ends_with('\n') { @@ -190,9 +219,8 @@ pub fn client_list() -> String { /// Format a single client's info (for CLIENT INFO). pub fn client_info(id: u64) -> Option { - let registry = REGISTRY.read(); let now = Instant::now(); - registry.get(&id).map(|entry| { + stripe(id).read().get(&id).map(|entry| { let mut result = String::with_capacity(128); format_client_line(&mut result, entry, now); if result.ends_with('\n') { @@ -216,29 +244,53 @@ pub fn client_info(id: u64) -> Option { /// per-batch kill check closes the connection (Redis replies first on /// self-kill, then closes). pub fn kill_clients(filter: &KillFilter, self_id: Option) -> u64 { - // Hold the read lock across the whole loop. This is what makes the raw-fd - // `shutdown` free of a use-after-reuse race: a connection's `RegistryGuard` - // (a local) drops — and calls `deregister`, which needs the registry WRITE - // lock — strictly before its `stream` (a parameter) drops and closes the - // fd. So while we hold the READ lock and an entry is present, `deregister` - // for it is blocked, its `stream` has not dropped, and `kill_fd` is still - // an open socket. Once an entry is gone, we never touch its fd. - let registry = REGISTRY.read(); + // Lock ordering / fd-liveness invariant, per stripe (c10k W5): the raw-fd + // `shutdown` is race-free because a connection's `RegistryGuard` (a local) + // drops — calling `deregister`, which needs its OWN stripe's WRITE lock — + // strictly before its `stream` (a parameter) drops and closes the fd. So + // while we hold a stripe's READ lock and an entry is present in it, + // `deregister` for that entry is blocked, its `stream` has not dropped, + // and `kill_fd` is still an open socket. The invariant only ever involves + // one entry and its own stripe, so striping preserves it; entries in + // other stripes are simply not visited while their lock is free. let mut count = 0u64; - for entry in registry.values() { - let matches = match filter { - KillFilter::Id(target_id) => entry.id == *target_id, - KillFilter::Addr(addr) => entry.addr == *addr, - KillFilter::User(user) => entry.user == *user, - }; - if matches { - entry.live.kill_flag.store(true, Ordering::Relaxed); - // Self-kill stays cooperative: shutting down our own socket here - // would happen mid-command, before the reply is flushed. - if self_id != Some(entry.id) { - force_close_fd(entry.live.kill_fd); + let kill_entry = |entry: &ClientEntry, count: &mut u64| { + entry.live.kill_flag.store(true, Ordering::Relaxed); + // Self-kill stays cooperative: shutting down our own socket here + // would happen mid-command, before the reply is flushed. + if self_id != Some(entry.id) { + force_close_fd(entry.live.kill_fd); + } + *count += 1; + }; + match filter { + // Id filter: O(1) — direct stripe lookup instead of the old full + // registry scan (the map was always id-keyed). + KillFilter::Id(target_id) => { + let guard = stripe(*target_id).read(); + if let Some(entry) = guard.get(target_id) { + kill_entry(entry, &mut count); + } + } + // Addr/User filters: walk stripes ONE at a time — a broad kill no + // longer stalls every accept/close globally. A connection that + // registers into an already-visited stripe mid-walk is missed, the + // same inherent race the single-lock version had at command + // granularity. + KillFilter::Addr(_) | KillFilter::User(_) => { + for lock in REGISTRY.iter() { + let guard = lock.read(); + for entry in guard.values() { + let matches = match filter { + KillFilter::Id(_) => false, // handled by the arm above + KillFilter::Addr(addr) => entry.addr == *addr, + KillFilter::User(user) => entry.user == *user, + }; + if matches { + kill_entry(entry, &mut count); + } + } } - count += 1; } } count @@ -571,3 +623,49 @@ mod tests { assert!(matches!(filter, KillFilter::User(u) if u == "alice")); } } + +#[cfg(test)] +mod striping_tests { + use super::*; + + /// Ids chosen to span every stripe; registry state is process-global, so + /// use a distinct id range from the legacy tests (which use small ids). + const BASE: u64 = 91_000; + + fn register_n(n: u64, user: &str) -> Vec { + (0..n) + .map(|i| { + let id = BASE + i; + let _ = register(id, format!("10.0.0.{i}:1234"), user.to_string(), 0, -1); + id + }) + .collect() + } + + #[test] + fn cross_stripe_list_and_targeted_kill() { + let ids = register_n(40, "stripe-user"); + + // CLIENT LIST must see every entry regardless of stripe. + let list = client_list(); + for id in &ids { + assert!(list.contains(&format!("id={id} ")), "id {id} missing from list"); + } + + // Kill-by-id is a single-stripe O(1) lookup — exactly one victim. + assert_eq!(kill_clients(&KillFilter::Id(BASE + 17), None), 1); + assert!(is_killed(BASE + 17)); + assert!(!is_killed(BASE + 18)); + + // Kill-by-user walks all stripes and reaches every entry. + assert_eq!( + kill_clients(&KillFilter::User("stripe-user".to_string()), None), + 40 + ); + for id in &ids { + assert!(is_killed(*id)); + deregister(*id); + } + assert!(client_info(BASE).is_none(), "deregistered id must be gone"); + } +} From 112f8bf745fabd720e616b87feb0309195150789 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Wed, 29 Jul 2026 12:59:33 +0700 Subject: [PATCH 08/17] fix(server): load-gate the IP-affinity funnel (c10k W8) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The AffinityTracker pins ALL central-accepted connections from an IP to the shard where one of its connections migrated (key locality) or first subscribed (pub/sub locality) — with zero load feedback. A saturated shard was never removed from rotation, giving ~2x worst-case connection skew from few client IPs, and FD migration compounded it (tmp/C10K-REVIEW.md defect #5). - client_registry now maintains readable per-shard connection counters (OnceLock>, initialized from startup with the resolved shard count; inert when uninitialized — tests/embedded). The metrics gauge could not serve this: gauges are write-only and gated on METRICS_INITIALIZED. - All three listener affinity lookups filter the hint through shard_overloaded(): a shard above 2x the mean conn count (+16 floor, so tiny fleets never flap) falls back to round-robin instead of receiving further funneled connections. Affinity resumes once the shard drains below threshold. - maybe_evict's sort stays as-is — review found it amortized ~60 ops/insert past cap (fires once per 4096 inserts), not a real cost. 4 unit tests on the pure threshold predicate (balanced, at-threshold, funneled, small-fleet floor, zero-shard safety). Workstream W8 of tmp/C10K-PLAN.md. author: Tin Dang --- src/client_registry.rs | 89 ++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 4 ++ src/server/listener.rs | 26 +++++++++--- 3 files changed, 113 insertions(+), 6 deletions(-) diff --git a/src/client_registry.rs b/src/client_registry.rs index 6d10bd653..c2261e6e7 100644 --- a/src/client_registry.rs +++ b/src/client_registry.rs @@ -37,6 +37,62 @@ fn stripe(id: u64) -> &'static RwLock> { ®ISTRY[(id as usize) % STRIPES] } +/// Per-shard live connection counts (c10k W8). Readable — unlike the +/// metrics gauge — so the listener's IP-affinity funnel can refuse to pile +/// further connections onto an already-overloaded shard. Initialized once +/// at startup; when unset (unit tests, embedded), the load gate is inert. +static SHARD_CONNS: std::sync::OnceLock> = std::sync::OnceLock::new(); + +/// Initialize the per-shard connection counters. Called once from startup +/// with the resolved shard count; later calls are no-ops. +pub fn init_shard_conn_counts(num_shards: usize) { + let _ = SHARD_CONNS.set( + (0..num_shards) + .map(|_| AtomicUsize::new(0)) + .collect::>() + .into_boxed_slice(), + ); +} + +#[inline] +fn shard_conns_delta(shard: usize, add: bool) { + if let Some(counts) = SHARD_CONNS.get() + && let Some(c) = counts.get(shard) + { + if add { + c.fetch_add(1, Ordering::Relaxed); + } else { + c.fetch_sub(1, Ordering::Relaxed); + } + } +} + +/// Is `shard` carrying disproportionate connection load? Used by the +/// listener to gate the IP-affinity hint: one migrated/subscribed client +/// used to pin ALL subsequent central-accepted connections from its IP onto +/// one shard with no load feedback (~2× worst-case skew, +/// tmp/C10K-REVIEW.md defect #5). Threshold: 2× the mean + a small +/// absolute floor so tiny fleets never flap. +pub fn shard_overloaded(shard: usize, num_shards: usize) -> bool { + let Some(counts) = SHARD_CONNS.get() else { + return false; + }; + let Some(c) = counts.get(shard) else { + return false; + }; + overload_threshold_exceeded( + c.load(Ordering::Relaxed), + TOTAL_CLIENTS.load(Ordering::Relaxed), + num_shards, + ) +} + +/// Pure overload predicate (unit-test surface for the W8 gate). +#[inline] +pub(crate) fn overload_threshold_exceeded(count: usize, total: usize, num_shards: usize) -> bool { + count > 2 * (total / num_shards.max(1)) + 16 +} + /// Lock-free per-connection state, shared between the connection task /// (writer, once per batch) and CLIENT LIST/INFO/KILL (occasional readers). pub struct ClientLiveState { @@ -162,6 +218,7 @@ pub fn register( }; stripe(id).write().insert(id, entry); TOTAL_CLIENTS.fetch_add(1, Ordering::Relaxed); + shard_conns_delta(shard, true); crate::admin::metrics_setup::record_shard_connection_delta(shard, 1.0); live } @@ -170,6 +227,7 @@ pub fn register( pub fn deregister(id: u64) { if let Some(entry) = stripe(id).write().remove(&id) { TOTAL_CLIENTS.fetch_sub(1, Ordering::Relaxed); + shard_conns_delta(entry.shard, false); crate::admin::metrics_setup::record_shard_connection_delta(entry.shard, -1.0); } } @@ -669,3 +727,34 @@ mod striping_tests { assert!(client_info(BASE).is_none(), "deregistered id must be gone"); } } + +#[cfg(test)] +mod overload_gate_tests { + use super::overload_threshold_exceeded; + + #[test] + fn balanced_load_is_not_overloaded() { + // 4 shards, 1000 total, this shard at the mean. + assert!(!overload_threshold_exceeded(250, 1000, 4)); + // Right at 2x mean + 16 is still allowed (strictly-greater gate). + assert!(!overload_threshold_exceeded(516, 1000, 4)); + } + + #[test] + fn funneled_shard_trips_the_gate() { + assert!(overload_threshold_exceeded(517, 1000, 4)); + assert!(overload_threshold_exceeded(900, 1000, 4)); + } + + #[test] + fn small_fleets_never_flap() { + // Absolute floor: with 10 total conns, even a 10/0/0/0 split stays + // under mean*2+16 — affinity wins for tiny fleets. + assert!(!overload_threshold_exceeded(10, 10, 4)); + } + + #[test] + fn zero_shards_is_safe() { + assert!(!overload_threshold_exceeded(0, 0, 0)); + } +} diff --git a/src/main.rs b/src/main.rs index f753e6da8..aa6d79448 100644 --- a/src/main.rs +++ b/src/main.rs @@ -658,6 +658,10 @@ fn main() -> anyhow::Result<()> { } } + // c10k W8: readable per-shard connection counters — feeds the listener's + // IP-affinity load gate. + moon::client_registry::init_shard_conn_counts(num_shards); + // Create channel mesh for inter-shard communication let mut mesh = ChannelMesh::new(num_shards, CHANNEL_BUFFER_SIZE); diff --git a/src/server/listener.rs b/src/server/listener.rs index 8c2e6ed43..0f462a4aa 100644 --- a/src/server/listener.rs +++ b/src/server/listener.rs @@ -443,10 +443,16 @@ pub async fn run_sharded( continue; } - // Affinity-aware routing: check if this IP has a preferred shard + // Affinity-aware routing: check if this IP has a preferred shard. + // c10k W8: the hint is load-gated — an overloaded shard falls + // back to round-robin instead of funneling every same-IP conn. let target_shard = { let peer_ip = addr.ip(); - if let Some(preferred) = affinity_tracker.read().lookup(&peer_ip) { + if let Some(preferred) = affinity_tracker + .read() + .lookup(&peer_ip) + .filter(|&s| !crate::client_registry::shard_overloaded(s, num_shards)) + { if preferred < num_shards { preferred } else { @@ -576,10 +582,14 @@ pub async fn run_sharded( let _ = monoio::io::AsyncWriteRentExt::write_all(&mut stream, err_msg).await; continue; } - // Affinity-aware routing + // Affinity-aware routing (c10k W8: load-gated hint) let target_shard = { let peer_ip = addr.ip(); - if let Some(preferred) = affinity_tracker.read().lookup(&peer_ip) { + if let Some(preferred) = affinity_tracker + .read() + .lookup(&peer_ip) + .filter(|&s| !crate::client_registry::shard_overloaded(s, num_shards)) + { if preferred < num_shards { preferred } else { let s = next_shard; next_shard = (next_shard + 1) % num_shards; @@ -666,10 +676,14 @@ pub async fn run_sharded( continue; } - // Affinity-aware routing + // Affinity-aware routing (c10k W8: load-gated hint) let target_shard = { let peer_ip = addr.ip(); - if let Some(preferred) = affinity_tracker.read().lookup(&peer_ip) { + if let Some(preferred) = affinity_tracker + .read() + .lookup(&peer_ip) + .filter(|&s| !crate::client_registry::shard_overloaded(s, num_shards)) + { if preferred < num_shards { preferred } else { let s = next_shard; next_shard = (next_shard + 1) % num_shards; From 7a3f921a8ad8c740fd9e06c9e95440bdd17afd0a Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Wed, 29 Jul 2026 13:02:00 +0700 Subject: [PATCH 09/17] docs(server): CHANGELOG for c10k wave + loud MOON_URING bypass warning (c10k T3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The experimental tokio->io_uring bridge accepts connections that bypass maxclients and the client registry (CLIENT LIST/KILL blind). Full accounting integration is deliberately unscheduled — the bridge is documented as broken under sustained load — but the bypass now announces itself at startup instead of being silent. c1M roadmap (parked-idle conns, task-future diet, uring-entries knob, TLS caps) captured in .planning/rfcs/c1m-connection-plane.md. author: Tin Dang --- CHANGELOG.md | 43 +++++++++++++++++++++++++++++++++++++++++ src/shard/event_loop.rs | 10 ++++++++++ 2 files changed, 53 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index dd2aabea2..b81fa1ea3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,49 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed +- **c10k connection-plane wave (PR #TBD)** — from the 2026-07-29 empirical + review (`tmp/C10K-REVIEW.md`: 10k/25k live-connection ramps, idle-CPU + measurement, symbolized flames, pipeline-ratchet A/B on the Linux VM): + - **Pipeline memory ratchet fixed (W1).** One 1024-deep pipeline + permanently ratcheted a connection from 56.5 KB to ~217 KB RSS until + disconnect (measured: 5000 such conns pinned 1.09 GB). Batch scratch + vecs now shrink after oversized batches, I/O buffer shrink floor + lowered 64 KiB → 16 KiB, tokio bump arena capped, subscriber-mode + per-message 8 KiB alloc+zero removed. + - **maxclients rejection is loud (W3).** Over-cap connections now receive + `-ERR max number of clients reached` (Redis parity) instead of a silent + EOF on the monoio + TLS paths; the gate runs before per-connection + context construction; startup now checks `RLIMIT_NOFILE` against + maxclients, raises the soft limit when possible, and warns with the + real client ceiling otherwise. + - **Blocked-client sweep is deadline-heap driven (W6)** — O(due) instead + of walking every blocked waiter at 100 Hz (~2M entry visits/s/shard at + 10k BLPOP clients); block-forever waiters are never visited. + - **SPSC drain starvation fixed (W7).** The cross-shard drain rotates its + start consumer, so the 256-message budget no longer lets a hot + low-index peer starve high-index peers indefinitely. + - **IP-affinity funnel load-gated (W8).** A shard above 2× the mean + connection count no longer receives affinity-funneled connections + (falls back to round-robin) — previously one migrated/subscribed client + pinned all same-IP connections to one shard with no load feedback. + +### Changed +- **Client registry striped 16 ways (W5).** Accepts/closes no longer + serialize on one global lock; `CLIENT LIST` locks one stripe at a time + (output is no longer insertion-ordered — Redis makes no ordering + guarantee); `CLIENT KILL ID` is an O(1) lookup instead of a full scan. +- **`active_cross_txn` boxed (W2)** — ~2.2 KB of inline transaction + SmallVecs no longer ride in every connection's task future; + `ConnectionState` is guarded ≤768 B by a regression test. +- **Dead plumbing deleted (W4):** the zero-registrant `pending_wakers` + waker relay (per-conn Rc clone + twice-per-iteration sweep) and the + never-adopted duplicate `server/conn_state.rs` module (209 lines). +- The experimental `MOON_URING=1` tokio bridge now logs a loud startup + warning that its connections bypass maxclients and CLIENT LIST/KILL + (T3; full accounting integration tracked in + `.planning/rfcs/c1m-connection-plane.md`). + ## [0.8.2] — 2026-07-20 ### Fixed diff --git a/src/shard/event_loop.rs b/src/shard/event_loop.rs index 4e5662b41..3dcdae4e4 100644 --- a/src/shard/event_loop.rs +++ b/src/shard/event_loop.rs @@ -112,6 +112,16 @@ impl super::Shard { Ok(mut d) => match d.init() { Ok(()) => { info!("Shard {} started (io_uring mode)", self.id); + // c10k T3: known limitation, stated loudly — bridge + // connections bypass maxclients AND the client + // registry (CLIENT LIST/KILL blind); capacity is + // bounded only by the driver's FdTable. + tracing::warn!( + "Shard {} io_uring bridge: connections accepted here bypass \ + maxclients and CLIENT LIST/KILL (experimental path; see \ + .planning/rfcs/c1m-connection-plane.md)", + self.id + ); Some(d) } Err(e) => { From 6fe18446ff822e6fb7f9d4b7a9912c4753cd66fa Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Wed, 29 Jul 2026 13:05:09 +0700 Subject: [PATCH 10/17] style(server): fmt + clippy zero-warnings for the c10k wave rlimit math kept in rlim_t domain (kills useless u64::from lints); batch-vec shrink helpers cfg-gated monoio-or-test (tokio-only build saw them as dead code); cargo fmt over the new test/impl code. author: Tin Dang --- src/blocking/mod.rs | 7 ++++++- src/client_registry.rs | 10 ++++++++-- src/main.rs | 18 ++++++++++++------ src/server/conn/util.rs | 11 ++++++++++- src/shard/spsc_handler.rs | 4 +++- tests/maxclients_reject_parity.rs | 5 ++++- 6 files changed, 43 insertions(+), 12 deletions(-) diff --git a/src/blocking/mod.rs b/src/blocking/mod.rs index 2dd437962..2568899e9 100644 --- a/src/blocking/mod.rs +++ b/src/blocking/mod.rs @@ -433,7 +433,12 @@ mod deadline_heap_tests { reg.register( 0, Bytes::copy_from_slice(k), - WaitEntry { wait_id: id, cmd: BlockedCommand::BLPop, reply_tx: tx, deadline }, + WaitEntry { + wait_id: id, + cmd: BlockedCommand::BLPop, + reply_tx: tx, + deadline, + }, ); } let visited = reg.expire_timed_out(now); diff --git a/src/client_registry.rs b/src/client_registry.rs index c2261e6e7..a739174bf 100644 --- a/src/client_registry.rs +++ b/src/client_registry.rs @@ -246,7 +246,10 @@ pub fn update(id: u64, f: F) { /// Registry-lookup variant for code without the live handle; connection /// loops use `ClientLiveState::is_killed` (lock-free) instead. pub fn is_killed(id: u64) -> bool { - stripe(id).read().get(&id).is_some_and(|e| e.live.is_killed()) + stripe(id) + .read() + .get(&id) + .is_some_and(|e| e.live.is_killed()) } /// Format all clients as a CLIENT LIST string. @@ -707,7 +710,10 @@ mod striping_tests { // CLIENT LIST must see every entry regardless of stripe. let list = client_list(); for id in &ids { - assert!(list.contains(&format!("id={id} ")), "id {id} missing from list"); + assert!( + list.contains(&format!("id={id} ")), + "id {id} missing from list" + ); } // Kill-by-id is a single-stripe O(1) lookup — exactly one victim. diff --git a/src/main.rs b/src/main.rs index aa6d79448..d32bf1a8e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -633,15 +633,19 @@ fn main() -> anyhow::Result<()> { rlim_max: 0, }; if libc::getrlimit(libc::RLIMIT_NOFILE, &mut rl) == 0 { - let need = config.maxclients as u64 + reserved; - if config.maxclients > 0 && u64::from(rl.rlim_cur) < need { - let want = need.min(u64::from(rl.rlim_max)); + // rlim_t is u64 on every supported unix; `reserved` is u64. + let need: libc::rlim_t = config.maxclients as libc::rlim_t + reserved; + if config.maxclients > 0 && rl.rlim_cur < need { + let want = need.min(rl.rlim_max); let mut raised = rl; - raised.rlim_cur = want as libc::rlim_t; + raised.rlim_cur = want; if libc::setrlimit(libc::RLIMIT_NOFILE, &raised) == 0 { tracing::info!( "Raised RLIMIT_NOFILE soft limit {} -> {} to fit maxclients {} (+{} reserved fds)", - rl.rlim_cur, want, config.maxclients, reserved + rl.rlim_cur, + want, + config.maxclients, + reserved ); } if want < need { @@ -649,7 +653,9 @@ fn main() -> anyhow::Result<()> { "RLIMIT_NOFILE hard limit {} cannot fit maxclients {} (+{} reserved fds): \ accepts beyond ~{} clients will fail with EMFILE. Raise `ulimit -n` \ or lower --maxclients.", - rl.rlim_max, config.maxclients, reserved, + rl.rlim_max, + config.maxclients, + reserved, want.saturating_sub(reserved) ); } diff --git a/src/server/conn/util.rs b/src/server/conn/util.rs index dd2ae9e22..76aa0f709 100644 --- a/src/server/conn/util.rs +++ b/src/server/conn/util.rs @@ -112,12 +112,17 @@ pub(crate) fn unpropagate_subscription( /// experiment E5). Shrinking only above the trigger keeps every realloc off /// the small-batch path: a p99 batch (≤64 frames) never grows past the /// trigger, so it never pays a shrink. +/// (Consumed by the monoio handler only — the tokio handler allocates its +/// batch vecs per batch; cfg-gated so the tokio-only build stays warning-free.) +#[cfg(any(feature = "runtime-monoio", test))] pub(crate) const BATCH_VEC_STEADY_CAP: usize = 64; +#[cfg(any(feature = "runtime-monoio", test))] pub(crate) const BATCH_VEC_SHRINK_TRIGGER: usize = 256; /// Shrink an emptied batch scratch vector back to steady-state capacity. /// Call only after `clear()` — shrinking an empty Vec moves no elements. #[inline] +#[cfg(any(feature = "runtime-monoio", test))] pub(crate) fn shrink_batch_vec(v: &mut Vec) { debug_assert!(v.is_empty(), "shrink_batch_vec expects a cleared vec"); if v.capacity() > BATCH_VEC_SHRINK_TRIGGER { @@ -154,7 +159,11 @@ mod shrink_tests { let mut v: Vec = Vec::with_capacity(BATCH_VEC_SHRINK_TRIGGER); let cap_before = v.capacity(); shrink_batch_vec(&mut v); - assert_eq!(v.capacity(), cap_before, "at-trigger capacity must not shrink"); + assert_eq!( + v.capacity(), + cap_before, + "at-trigger capacity must not shrink" + ); } #[test] diff --git a/src/shard/spsc_handler.rs b/src/shard/spsc_handler.rs index be48e235f..4ff6474b5 100644 --- a/src/shard/spsc_handler.rs +++ b/src/shard/spsc_handler.rs @@ -4502,7 +4502,9 @@ mod drain_rotation_tests { fn successive_starts_change_the_first_visited_consumer() { // The budget cut-off truncates the TAIL of the visit order; rotating // the head is what guarantees no consumer is permanently tail-most. - let firsts: Vec = (0..4).map(|s| rotated_indices(s, 4).next().unwrap()).collect(); + let firsts: Vec = (0..4) + .map(|s| rotated_indices(s, 4).next().unwrap()) + .collect(); assert_eq!(firsts, vec![0, 1, 2, 3]); } diff --git a/tests/maxclients_reject_parity.rs b/tests/maxclients_reject_parity.rs index 4a42b2b5e..dbf5c5150 100644 --- a/tests/maxclients_reject_parity.rs +++ b/tests/maxclients_reject_parity.rs @@ -93,7 +93,10 @@ fn rejected_connection_receives_err_max_clients() { ); // The slot-holder must be unaffected by the rejection. - assert!(ping_ok(&mut conn_a), "conn A must survive conn B's rejection"); + assert!( + ping_ok(&mut conn_a), + "conn A must survive conn B's rejection" + ); drop(conn_a); common::sigkill(&mut child); From fdcd5d904793ed7dcb9ab0e1b864b82887070b35 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Wed, 29 Jul 2026 13:54:05 +0700 Subject: [PATCH 11/17] feat(config): --uring-entries knob for per-shard ring sizing (c10k P4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Expose monoio's already-public RuntimeBuilder::with_entries through moon config: `--uring-entries N` sets the io_uring submission-queue size per shard ring (monoio default 1024; CQ = 2x by kernel), and sizes the legacy epoll/kqueue driver's mio event batch on that path. 1024 in-flight ops per shard is small for 10k+ conns/shard under bursty pipelines (.planning/rfcs/c1m-connection-plane.md Proposal 4). Plumbing mirrors the --io-driver epoll precedent: a process-global AtomicU32 in runtime/mod.rs set once in main BEFORE any shard thread spawns, consumed by all three monoio build paths (legacy, tuned io_uring, fusion fallback). Values below monoio's 256 floor clamp up with a warning; unset keeps monoio's default; no effect under the tokio runtime (loud warn). Zero vendored-monoio changes needed — with_entries() was already public and forwarded on every builder path. Tests: config parse (default None / value / reject non-numeric) + runtime static default/set/clamp. author: Tin Dang --- src/config.rs | 21 ++++++++++++++++++ src/main.rs | 22 +++++++++++++++++++ src/runtime/mod.rs | 45 ++++++++++++++++++++++++++++++++++++++ src/runtime/monoio_impl.rs | 33 ++++++++++++++++++---------- 4 files changed, 109 insertions(+), 12 deletions(-) diff --git a/src/config.rs b/src/config.rs index f9f439c32..b3577cadc 100644 --- a/src/config.rs +++ b/src/config.rs @@ -295,6 +295,15 @@ pub struct ServerConfig { #[arg(long = "uring-sqpoll")] pub uring_sqpoll_ms: Option, + /// io_uring submission-queue entries per shard ring (default: monoio's + /// 1024; CQ is sized 2x by the kernel). Raise for high-connection shards + /// under bursty pipelines — 1024 in-flight ops per shard is small at 10k+ + /// conns/shard. Values below 256 are raised to 256 (monoio's floor). + /// Also sizes the legacy (epoll/kqueue) driver's event batch. Applies to + /// the monoio runtime only. + #[arg(long = "uring-entries")] + pub uring_entries: Option, + /// I/O driver for the monoio runtime. "auto" lets FusionDriver pick /// (io_uring on Linux when available, else epoll/kqueue); "epoll" forces /// the legacy poller. Measured on GCE ARM (c4a Axion, 2026-07): epoll is @@ -2001,6 +2010,18 @@ mod tests { assert!(ServerConfig::try_parse_from(["moon", "--io-driver", "iouring"]).is_err()); } + #[test] + fn test_uring_entries_flag() { + let config = ServerConfig::parse_from::<[&str; 0], &str>([]); + assert_eq!( + config.uring_entries, None, + "default must be None (monoio's built-in 1024)" + ); + let config = ServerConfig::parse_from(["moon", "--uring-entries", "4096"]); + assert_eq!(config.uring_entries, Some(4096)); + assert!(ServerConfig::try_parse_from(["moon", "--uring-entries", "x"]).is_err()); + } + #[test] fn test_ft_search_workers_flag() { let config = ServerConfig::parse_from::<[&str; 0], &str>([]); diff --git a/src/main.rs b/src/main.rs index d32bf1a8e..856dd3eb3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1074,6 +1074,28 @@ fn main() -> anyhow::Result<()> { #[cfg(not(feature = "runtime-monoio"))] tracing::warn!("--io-busy-poll-us has no effect under the tokio runtime"); } + if let Some(entries) = config.uring_entries { + // Same before-shard-spawn contract as force_legacy_driver above. + moon::runtime::set_uring_entries(entries); + #[cfg(feature = "runtime-monoio")] + { + let effective = entries.max(moon::runtime::MIN_URING_ENTRIES); + if effective != entries { + tracing::warn!( + "--uring-entries {} below the {} floor; using {}", + entries, + moon::runtime::MIN_URING_ENTRIES, + effective + ); + } + tracing::info!( + "I/O ring size: {} SQ entries per shard (--uring-entries)", + effective + ); + } + #[cfg(not(feature = "runtime-monoio"))] + tracing::warn!("--uring-entries has no effect under the tokio runtime"); + } // Graph traversal timeout default — set BEFORE shard threads spawn so every // TraversalGuard::with_default_timeout observes it (per-query TIMEOUT overrides). diff --git a/src/runtime/mod.rs b/src/runtime/mod.rs index bef1b48f8..a32667458 100644 --- a/src/runtime/mod.rs +++ b/src/runtime/mod.rs @@ -45,6 +45,33 @@ pub fn legacy_driver_forced() -> bool { || std::env::var_os("MOON_NO_URING").is_some() } +/// Process-wide io_uring SQ-entries override (`--uring-entries N`), set ONCE +/// from main BEFORE any shard thread spawns — same contract as +/// [`force_legacy_driver`]. 0 = unset (monoio's default 1024). Also sizes the +/// legacy driver's mio event batch. +static URING_ENTRIES: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(0); + +/// monoio's `RuntimeBuilder::with_entries` floor; values below are raised so +/// the operator-visible behavior matches what the ring actually gets. +pub const MIN_URING_ENTRIES: u32 = 256; + +/// Set the per-shard ring size. Values below [`MIN_URING_ENTRIES`] are clamped +/// up (mirrors monoio's internal floor). Call before shard threads spawn. +pub fn set_uring_entries(entries: u32) { + URING_ENTRIES.store( + entries.max(MIN_URING_ENTRIES), + std::sync::atomic::Ordering::Release, + ); +} + +/// The configured ring size, or `None` when the operator left the default. +pub fn uring_entries() -> Option { + match URING_ENTRIES.load(std::sync::atomic::Ordering::Acquire) { + 0 => None, + n => Some(n), + } +} + /// True when the epoll busy-poll park is configured — via the /// `--io-busy-poll-us` flag (the caller passes the config value) or the /// `MOON_EPOLL_SPIN_US` env fallback the vendored driver also honors. Gates @@ -360,6 +387,24 @@ mod yield_costfree_tests { } } +#[cfg(test)] +mod uring_entries_tests { + // Shared process-global: run assertions in one test to avoid ordering + // races between parallel test threads. + #[test] + fn uring_entries_default_set_and_clamp() { + assert_eq!(super::uring_entries(), None, "unset must read as None"); + super::set_uring_entries(4096); + assert_eq!(super::uring_entries(), Some(4096)); + super::set_uring_entries(64); + assert_eq!( + super::uring_entries(), + Some(super::MIN_URING_ENTRIES), + "sub-floor values clamp up to monoio's minimum" + ); + } +} + // --- Runtime-specific type aliases for TCP networking --- // These allow subsystem code to use a single type name regardless of runtime. diff --git a/src/runtime/monoio_impl.rs b/src/runtime/monoio_impl.rs index 4ef3514c1..260babafc 100644 --- a/src/runtime/monoio_impl.rs +++ b/src/runtime/monoio_impl.rs @@ -65,13 +65,17 @@ impl RuntimeFactory for MonoioRuntimeFactory { // explicitly. (Previously MOON_NO_URING only gated the tokio bridge // + uring_active(); the monoio driver choice ignored it — discovered // during the c4a p=1 driver A/B.) + // --uring-entries: per-shard ring size (SQ entries; also the legacy + // driver's mio event-batch capacity). monoio's default is 1024. + let ring_entries = crate::runtime::uring_entries(); if crate::runtime::legacy_driver_forced() { - let mut rt = monoio::RuntimeBuilder::::new() - .enable_timer() - .build() - .unwrap_or_else(|e| { - panic!("failed to build monoio legacy runtime '{}': {}", name, e) - }); + let mut builder = monoio::RuntimeBuilder::::new(); + if let Some(n) = ring_entries { + builder = builder.with_entries(n); + } + let mut rt = builder.enable_timer().build().unwrap_or_else(|e| { + panic!("failed to build monoio legacy runtime '{}': {}", name, e) + }); rt.block_on(f); return; } @@ -121,11 +125,12 @@ impl RuntimeFactory for MonoioRuntimeFactory { .setup_single_issuer() .setup_defer_taskrun(); } - match monoio::RuntimeBuilder::::new() - .uring_builder(urb) - .enable_timer() - .build() - { + let mut builder = + monoio::RuntimeBuilder::::new().uring_builder(urb); + if let Some(n) = ring_entries { + builder = builder.with_entries(n); + } + match builder.enable_timer().build() { Ok(mut rt) => { rt.block_on(f); return; @@ -141,7 +146,11 @@ impl RuntimeFactory for MonoioRuntimeFactory { } } - let mut rt = monoio::RuntimeBuilder::::new() + let mut builder = monoio::RuntimeBuilder::::new(); + if let Some(n) = ring_entries { + builder = builder.with_entries(n); + } + let mut rt = builder .enable_timer() .build() .unwrap_or_else(|e| panic!("failed to build monoio runtime '{}': {}", name, e)); From 39203484d14031214e19fb09557422c9948785cc Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Wed, 29 Jul 2026 13:54:19 +0700 Subject: [PATCH 12/17] perf(server): box cold txn-abort + FT.* awaitees out of the conn future (c10k P3) -Zprint-type-sizes on the 58-variant connection state machine showed its size is set entirely by two COLD suspend points: awaiting try_handle_txn_abort (5568 B awaitee -> 8319 B variant) and the disconnect-teardown abort_cross_store_txn_routed (5376 B -> 8127 B), with try_handle_ft_command third (2816 B -> 5567 B). Box::pin those three behind their existing name-check fast paths, in both the monoio and sharded handlers: - TXN.ABORT rollback + leaked-txn disconnect teardown: alloc only when an abort actually executes - FT.* body split into ft_command_inner and boxed after the "FT." prefix check: alloc only per executed FT command (noise next to the search itself), never on the per-command filter chain Measured (nightly -Zprint-type-sizes A/B, macOS aarch64 release-fast): per-conn monoio task allocation 9280 -> 5952 B plain TCP (-36%), 13248 -> 9920 B TLS, 9472 -> 6144 B migrated; handler future 8640 -> 5312 B. Meets the RFC Proposal 3 target (<=8 KB) without touching the hot KV dispatch path (the next-largest variant is the hot cross-shard dispatch await, deliberately left inline). author: Tin Dang --- src/server/conn/handler_monoio/ft.rs | 14 ++++++++++++++ src/server/conn/handler_monoio/mod.rs | 7 +++++-- src/server/conn/handler_monoio/txn.rs | 7 +++++-- src/server/conn/handler_sharded/ft.rs | 14 ++++++++++++++ src/server/conn/handler_sharded/mod.rs | 7 +++++-- src/server/conn/handler_sharded/txn.rs | 7 +++++-- 6 files changed, 48 insertions(+), 8 deletions(-) diff --git a/src/server/conn/handler_monoio/ft.rs b/src/server/conn/handler_monoio/ft.rs index b73c445c4..90e4dadb6 100644 --- a/src/server/conn/handler_monoio/ft.rs +++ b/src/server/conn/handler_monoio/ft.rs @@ -200,7 +200,21 @@ pub(super) async fn try_handle_ft_command( if cmd.len() <= 3 || !cmd[..3].eq_ignore_ascii_case(b"FT.") { return false; } + // Box::pin (c10k future diet): the FT.* body is a ~2.8 KB state machine + // that would otherwise live inline in every connection future. The alloc + // happens only when an FT.* command actually executes — never on the + // name-mismatch fast path above — and is noise next to the search itself. + Box::pin(ft_command_inner(cmd, cmd_args, frame, conn, ctx, responses)).await +} +async fn ft_command_inner( + cmd: &[u8], + cmd_args: &[Frame], + frame: &Frame, + conn: &ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { if ctx.num_shards > 1 { // Multi-shard: dispatch via SPSC #[cfg(feature = "text-index")] diff --git a/src/server/conn/handler_monoio/mod.rs b/src/server/conn/handler_monoio/mod.rs index a8e0a869b..e2b534adc 100644 --- a/src/server/conn/handler_monoio/mod.rs +++ b/src/server/conn/handler_monoio/mod.rs @@ -2485,7 +2485,10 @@ pub(crate) async fn handle_connection_sharded_monoio< // sharded runtime block in handler_sharded.rs so both paths delegate to the same // shared helper. FIN has already been sent; shard state is still intact. if let Some(txn) = conn.active_cross_txn.take() { - crate::transaction::abort::abort_cross_store_txn_routed( + // Box::pin (c10k future diet): this ~5.4 KB rollback state machine + // otherwise sits inline in EVERY connection future; boxing costs one + // alloc on the leaked-txn teardown path only. + Box::pin(crate::transaction::abort::abort_cross_store_txn_routed( &ctx.shard_databases, ctx.shard_id, conn.selected_db, @@ -2493,7 +2496,7 @@ pub(crate) async fn handle_connection_sharded_monoio< &ctx.dispatch_tx, &ctx.spsc_notifiers, *txn, - ) + )) .await; } diff --git a/src/server/conn/handler_monoio/txn.rs b/src/server/conn/handler_monoio/txn.rs index 1a19f9237..319c63090 100644 --- a/src/server/conn/handler_monoio/txn.rs +++ b/src/server/conn/handler_monoio/txn.rs @@ -273,7 +273,10 @@ pub(super) async fn try_handle_txn_abort( // src/transaction/abort.rs for lock ordering. // Multi-shard: graph legs route to the shards owning // each graph name via ShardMessage::GraphRollback. - crate::transaction::abort::abort_cross_store_txn_routed( + // Box::pin (c10k future diet): keeps the ~5.4 KB rollback + // state machine out of the per-connection future; the alloc + // only happens when TXN.ABORT actually executes. + Box::pin(crate::transaction::abort::abort_cross_store_txn_routed( &ctx.shard_databases, ctx.shard_id, conn.selected_db, @@ -281,7 +284,7 @@ pub(super) async fn try_handle_txn_abort( &ctx.dispatch_tx, &ctx.spsc_notifiers, *txn, - ) + )) .await; responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); } else { diff --git a/src/server/conn/handler_sharded/ft.rs b/src/server/conn/handler_sharded/ft.rs index 6a794bc85..84025cfbe 100644 --- a/src/server/conn/handler_sharded/ft.rs +++ b/src/server/conn/handler_sharded/ft.rs @@ -25,7 +25,21 @@ pub(super) async fn try_handle_ft_command( if cmd.len() <= 3 || !cmd[..3].eq_ignore_ascii_case(b"FT.") { return false; } + // Box::pin (c10k future diet): the FT.* body is a ~2.8 KB state machine + // that would otherwise live inline in every connection future. The alloc + // happens only when an FT.* command actually executes — never on the + // name-mismatch fast path above — and is noise next to the search itself. + Box::pin(ft_command_inner(cmd, cmd_args, frame, conn, ctx, responses)).await +} +async fn ft_command_inner( + cmd: &[u8], + cmd_args: &[Frame], + frame: &Frame, + conn: &ConnectionState, + ctx: &ConnectionContext, + responses: &mut Vec, +) -> bool { if ctx.num_shards > 1 { // Multi-shard: dispatch via SPSC #[cfg(feature = "text-index")] diff --git a/src/server/conn/handler_sharded/mod.rs b/src/server/conn/handler_sharded/mod.rs index 4e97d9169..13b46bf70 100644 --- a/src/server/conn/handler_sharded/mod.rs +++ b/src/server/conn/handler_sharded/mod.rs @@ -2247,7 +2247,10 @@ pub(crate) async fn handle_connection_sharded_inner< // Closes T-161-05 — without this, a disconnect after TXN.BEGIN + SET would leak // kv_intents and pin the key invisible for all subsequent readers. if let Some(txn) = conn.active_cross_txn.take() { - crate::transaction::abort::abort_cross_store_txn_routed( + // Box::pin (c10k future diet): this ~5.4 KB rollback state machine + // otherwise sits inline in EVERY connection future; boxing costs one + // alloc on the leaked-txn teardown path only. + Box::pin(crate::transaction::abort::abort_cross_store_txn_routed( &ctx.shard_databases, ctx.shard_id, conn.selected_db, @@ -2255,7 +2258,7 @@ pub(crate) async fn handle_connection_sharded_inner< &ctx.dispatch_tx, &ctx.spsc_notifiers, *txn, - ) + )) .await; } diff --git a/src/server/conn/handler_sharded/txn.rs b/src/server/conn/handler_sharded/txn.rs index 879c76acc..2e1521040 100644 --- a/src/server/conn/handler_sharded/txn.rs +++ b/src/server/conn/handler_sharded/txn.rs @@ -268,7 +268,10 @@ pub(super) async fn try_handle_txn_abort( // src/transaction/abort.rs for lock ordering. // Multi-shard: graph legs route to the shards owning // each graph name via ShardMessage::GraphRollback. - crate::transaction::abort::abort_cross_store_txn_routed( + // Box::pin (c10k future diet): keeps the ~5.4 KB rollback + // state machine out of the per-connection future; the alloc + // only happens when TXN.ABORT actually executes. + Box::pin(crate::transaction::abort::abort_cross_store_txn_routed( &ctx.shard_databases, ctx.shard_id, conn.selected_db, @@ -276,7 +279,7 @@ pub(super) async fn try_handle_txn_abort( &ctx.dispatch_tx, &ctx.spsc_notifiers, *txn, - ) + )) .await; responses.push(Frame::SimpleString(Bytes::from_static(b"OK"))); } else { From dc94c4b2c50776ec7fb19b25efd01cc5ee086721 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Wed, 29 Jul 2026 13:56:23 +0700 Subject: [PATCH 13/17] =?UTF-8?q?docs(changelog):=20c10k=20round=202=20?= =?UTF-8?q?=E2=80=94=20--uring-entries=20+=20cold-future=20boxing?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit author: Tin Dang --- CHANGELOG.md | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index b81fa1ea3..8ef5b2d40 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -33,7 +33,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 (falls back to round-robin) — previously one migrated/subscribed client pinned all same-IP connections to one shard with no load feedback. +### Added +- **`--uring-entries N` (c10k P4)** — per-shard io_uring submission-queue + size (monoio default 1024; also the legacy driver's event-batch + capacity). Values below monoio's 256 floor clamp up with a warning; + monoio runtime only. + ### Changed +- **Cold command futures boxed out of the connection state machine + (c10k P3).** TXN.ABORT rollback, leaked-txn disconnect teardown, and + the FT.* body now `Box::pin` behind their name-check fast paths; the + per-connection monoio task allocation drops 9.3 KB → 5.9 KB plain TCP + (13.2 → 9.9 KB TLS) with zero allocation on the KV dispatch path. - **Client registry striped 16 ways (W5).** Accepts/closes no longer serialize on one global lock; `CLIENT LIST` locks one stripe at a time (output is no longer insertion-ordered — Redis makes no ordering From b80c2bf933641acad4bbd7a72604a2b906167d79 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Wed, 29 Jul 2026 15:02:08 +0700 Subject: [PATCH 14/17] =?UTF-8?q?perf(server):=20two-stage=20idle=20park?= =?UTF-8?q?=20=E2=80=94=20downshift=20idle=20conns=20to=20a=20512B=20worki?= =?UTF-8?q?ng=20set=20(c10k=20W11)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A connection parked in read() pinned its full working set forever: an 8 KiB rent buffer submitted to the kernel plus the empty 8 KiB read/write scratch — ~24 KiB of the measured 43.5 KB idle RSS/conn, and the floor the c1M RFC said only connection parking could reclaim. Two-stage park, no timers or allocs on the hot path: - Stage 1 (full buffers): cancel-capable streams read via monoio's cancelable_read, with the park timestamp (shard cached clock) in a shard-thread-local IdleSlot registry. One Rc alloc per CONNECTION at setup; per park just two Cell stores + a CancelHandle Rc clone. - The shard event loop's existing 1s chore sweeps the registry and fires the Canceller of any read parked >=1s. Cancel-and-await is loss-free on both drivers (io_uring AsyncCancel; legacy READ_CANCELED): a completion that raced the cancel still delivers its bytes. - Stage 2: the woken handler drops the rent buffer, releases the empty scratch buffers, and re-parks on a 512 B probe buffer with no further chore involvement — one cancellation per idle period, not per second. - Real data restores the full working set lazily (pre-park sizing) and re-arms stage 1. Bursts >512 B spill into the next full-size read. TLS streams opt out via IdleParkRead::SUPPORTS_IDLE_PARK (monoio_rustls has no cancelable read) and keep the pre-W11 path byte-for-byte; the tokio runtime is untouched. Measured (moon-dev VM, 2 shards, 10k idle conns): 43.5 -> 19.9 KB/conn after 30s idle (-54%; campaign total 56.5 -> 19.9, -65%); post-wake steady state 27.8 KB/conn (lazy regrowth also sheds the old always-8KiB read/write scratch); PING sweep over all 10k downshifted conns 0.07s, 0 failures; RSS returns on close. Parity suite (probe-size wake, >512B pipeline, 4 KiB value round trip, undisturbed active sibling) green on BOTH cancel paths: macOS kqueue legacy driver and Linux io_uring. author: Tin Dang --- src/server/conn/handler_monoio/idle_park.rs | 244 ++++++++++++++++++++ src/server/conn/handler_monoio/mod.rs | 60 ++++- src/shard/event_loop.rs | 5 + tests/idle_downshift_parity.rs | 145 ++++++++++++ 4 files changed, 450 insertions(+), 4 deletions(-) create mode 100644 src/server/conn/handler_monoio/idle_park.rs create mode 100644 tests/idle_downshift_parity.rs diff --git a/src/server/conn/handler_monoio/idle_park.rs b/src/server/conn/handler_monoio/idle_park.rs new file mode 100644 index 000000000..24f76a9a3 --- /dev/null +++ b/src/server/conn/handler_monoio/idle_park.rs @@ -0,0 +1,244 @@ +//! c10k W11 — idle-connection buffer downshift. +//! +//! A connection parked in `read()` pins its full working set: an 8 KiB rent +//! buffer plus the (empty) 8 KiB read-accumulation and write buffers — +//! ~24 KiB of the measured ~43 KiB idle RSS/conn. This module shrinks that +//! to ~0.5 KiB for connections idle ≥ [`IDLE_DOWNSHIFT_MS`] without putting +//! a timer or an allocation on the hot path: +//! +//! 1. Each cancel-capable connection registers one [`IdleSlot`] in a +//! shard-thread-local registry (one `Rc` alloc per connection, at setup). +//! 2. Stage-1 park: the handler records the park timestamp (shard cached +//! clock — no syscall) and reads via monoio's `cancelable_read`. +//! 3. The shard event loop's EXISTING 1 s chore calls [`sweep`]: any read +//! parked ≥ the threshold gets its `Canceller` fired. Cancel-and-await +//! is loss-free: if the op had already completed with data, awaiting the +//! same future delivers those bytes; otherwise it returns `ECANCELED`. +//! 4. On a cancelled (data-less) wake the handler releases its buffers and +//! re-parks in stage 2 with a [`IDLE_PROBE_BUF`]-byte buffer and NO +//! further chore involvement — a pure-idle connection is cancelled +//! exactly once, then sits at the small footprint until real data. +//! 5. Real data (stage 1 or 2) restores the full working set lazily. +//! +//! Hot connections never meet the sweep: their parks are microseconds, the +//! chore only sees `parked_since_ms` ≠ 0 entries older than the threshold. +//! Per-park hot-path cost: two `Cell` stores + one `Rc` clone of the +//! `CancelHandle` (no allocation). +//! +//! TLS connections (`monoio_rustls` streams) do not implement cancelable +//! reads; they opt out via [`IdleParkRead::SUPPORTS_IDLE_PARK`] and keep +//! the pre-W11 behavior byte-for-byte. + +use std::cell::{Cell, RefCell}; +use std::collections::HashMap; +use std::rc::Rc; + +use monoio::io::{AsyncReadRent, CancelHandle, Canceller}; + +/// Park duration after which the sweep cancels a stage-1 read (ms). +pub(super) const IDLE_DOWNSHIFT_MS: u64 = 1000; +/// Rent-buffer size while downshifted. A typical idle→active wake (one +/// command) fits; larger bursts spill into the next full-size read. +pub(super) const IDLE_PROBE_BUF: usize = 512; +/// Full rent-buffer size (the pre-W11 constant working size). +pub(super) const PARK_BUF_FULL: usize = 8192; + +/// Per-connection shared state between the handler task and the shard chore. +pub(super) struct IdleSlot { + /// Canceller for the connection's in-flight stage-1 read. The sweep + /// consumes it (`Canceller::cancel(self) -> Self`) and stores the fresh + /// replacement; the handler takes a fresh `CancelHandle` from whatever + /// canceller is current at every park. + canceller: RefCell, + /// Shard-cached-clock ms when the connection parked its stage-1 read; + /// 0 = not parked in stage 1 (processing, downshifted, or gone). + parked_since_ms: Cell, +} + +impl IdleSlot { + /// Take a cancel handle for the current park attempt. + pub(super) fn handle(&self) -> CancelHandle { + self.canceller.borrow().handle() + } + + /// Mark the connection parked (stage 1). `now_ms` comes from the shard + /// cached clock; clamped to ≥1 so 0 stays the "not parked" sentinel. + pub(super) fn mark_parked(&self, now_ms: u64) { + self.parked_since_ms.set(now_ms.max(1)); + } + + /// Mark the connection no longer parked (woke or is processing). + pub(super) fn mark_unparked(&self) { + self.parked_since_ms.set(0); + } +} + +thread_local! { + /// All cancel-capable connections on this shard thread, by client id. + static REGISTRY: RefCell>> = RefCell::new(HashMap::new()); +} + +/// RAII registration: removes the slot from the thread registry on drop +/// (handler return, migration hand-off, or task cancellation). +pub(super) struct IdleParkRegistration { + pub(super) slot: Rc, + client_id: u64, +} + +impl Drop for IdleParkRegistration { + fn drop(&mut self) { + REGISTRY.with(|r| r.borrow_mut().remove(&self.client_id)); + } +} + +/// Register a connection for idle-park sweeping on this shard thread. +pub(super) fn register(client_id: u64) -> IdleParkRegistration { + let slot = Rc::new(IdleSlot { + canceller: RefCell::new(Canceller::new()), + parked_since_ms: Cell::new(0), + }); + REGISTRY.with(|r| r.borrow_mut().insert(client_id, slot.clone())); + IdleParkRegistration { slot, client_id } +} + +/// Cancel every stage-1 read parked ≥ [`IDLE_DOWNSHIFT_MS`]. Called from the +/// shard event loop's 1 s chore (same thread as the connections). Returns the +/// number of reads cancelled (observability + tests). +/// +/// The cancelled connection wakes with `ECANCELED` (or its raced data), +/// downshifts, and re-parks in stage 2 with `parked_since_ms == 0`, so each +/// idle connection is cancelled exactly once per idle period. +pub(crate) fn sweep(now_ms: u64) -> usize { + REGISTRY.with(|r| { + let mut cancelled = 0usize; + for slot in r.borrow().values() { + let parked = slot.parked_since_ms.get(); + if parked != 0 && now_ms.saturating_sub(parked) >= IDLE_DOWNSHIFT_MS { + // 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()); + let fresh = old.cancel(); + slot.canceller.replace(fresh); + // One-shot: the waking handler also clears this, but resetting + // here keeps a slow-to-wake connection from being re-cancelled + // by the next sweep. + slot.parked_since_ms.set(0); + cancelled += 1; + } + } + cancelled + }) +} + +/// Release the parked working set. The rent buffer is dropped outright (the +/// pre-park sizing reallocates the probe size next iteration); the scratch +/// buffers are only released when empty — a non-empty `read_buf` holds a +/// partial frame that must survive the re-park. +pub(super) fn downshift_idle_buffers( + tmp_buf: &mut Vec, + read_buf: &mut bytes::BytesMut, + write_buf: &mut bytes::BytesMut, +) { + *tmp_buf = Vec::new(); + if read_buf.is_empty() && read_buf.capacity() > 0 { + *read_buf = bytes::BytesMut::new(); + } + if write_buf.is_empty() && write_buf.capacity() > 0 { + *write_buf = bytes::BytesMut::new(); + } +} + +/// Stream capability trait for the two-stage idle park. +/// +/// The default is a plain read that ignores the cancel handle and reports +/// `SUPPORTS_IDLE_PARK = false`, keeping non-cancelable streams (TLS) on the +/// exact pre-W11 path. Only streams whose `cancelable_read` is loss-free +/// under cancel-and-await may set `SUPPORTS_IDLE_PARK = true`. +pub(crate) trait IdleParkRead: AsyncReadRent { + const SUPPORTS_IDLE_PARK: bool = false; + + fn idle_park_read( + &mut self, + buf: Vec, + _c: CancelHandle, + ) -> impl std::future::Future>> { + self.read(buf) + } +} + +impl IdleParkRead for monoio::net::TcpStream { + const SUPPORTS_IDLE_PARK: bool = true; + + fn idle_park_read( + &mut self, + buf: Vec, + c: CancelHandle, + ) -> impl std::future::Future>> { + monoio::io::CancelableAsyncReadRent::cancelable_read(self, buf, c) + } +} + +/// TLS streams keep the pre-W11 behavior (no cancelable read in +/// monoio_rustls; dropping its read future mid-handshake-record would also +/// desync the TLS session state). +impl IdleParkRead for monoio_rustls::ServerTlsStream {} + +#[cfg(test)] +mod idle_park_tests { + use super::*; + + #[test] + fn sweep_cancels_only_overdue_parks() { + let reg = register(90_001); + // Not parked: never cancelled. + assert_eq!(sweep(10_000), 0, "unparked slot must not be swept"); + // Parked, under threshold: untouched. + reg.slot.mark_parked(10_000); + assert_eq!(sweep(10_000 + IDLE_DOWNSHIFT_MS - 1), 0); + // Over threshold: cancelled exactly once, then the one-shot reset + // keeps later sweeps away until the next park. + assert_eq!(sweep(10_000 + IDLE_DOWNSHIFT_MS), 1); + assert_eq!( + reg.slot.parked_since_ms.get(), + 0, + "sweep must reset park mark" + ); + assert_eq!( + sweep(10_000 + 10 * IDLE_DOWNSHIFT_MS), + 0, + "one-shot per park" + ); + // Re-park arms it again with the REPLACED canceller. + reg.slot.mark_parked(60_000); + assert_eq!(sweep(60_000 + IDLE_DOWNSHIFT_MS), 1); + } + + #[test] + fn registration_drop_deregisters() { + { + let reg = register(90_002); + reg.slot.mark_parked(5_000); + } + assert_eq!( + sweep(1_000_000), + 0, + "dropped registration must leave the registry" + ); + } + + #[test] + fn downshift_releases_only_empty_scratch() { + let mut tmp = vec![0u8; PARK_BUF_FULL]; + let mut rb = bytes::BytesMut::with_capacity(PARK_BUF_FULL); + let mut wb = bytes::BytesMut::with_capacity(PARK_BUF_FULL); + rb.extend_from_slice(b"partial-frame"); // must survive + downshift_idle_buffers(&mut tmp, &mut rb, &mut wb); + assert_eq!(tmp.capacity(), 0, "rent buffer must be dropped"); + assert_eq!(wb.capacity(), 0, "empty write buffer must be released"); + assert_eq!(&rb[..], b"partial-frame", "partial frame must be kept"); + assert!( + rb.capacity() >= 13, + "non-empty read buffer keeps its storage" + ); + } +} diff --git a/src/server/conn/handler_monoio/mod.rs b/src/server/conn/handler_monoio/mod.rs index e2b534adc..1cd1ef643 100644 --- a/src/server/conn/handler_monoio/mod.rs +++ b/src/server/conn/handler_monoio/mod.rs @@ -5,6 +5,7 @@ mod dispatch; mod ft; +pub(crate) mod idle_park; mod pubsub; mod read; mod txn; @@ -168,7 +169,7 @@ fn run_write_eviction_gate( #[cfg(feature = "runtime-monoio")] #[tracing::instrument(skip_all, level = "debug")] pub(crate) async fn handle_connection_sharded_monoio< - S: monoio::io::AsyncReadRent + monoio::io::AsyncWriteRent, + S: monoio::io::AsyncReadRent + monoio::io::AsyncWriteRent + idle_park::IdleParkRead, >( mut stream: S, peer_addr: String, @@ -245,6 +246,12 @@ pub(crate) async fn handle_connection_sharded_monoio< // Monoio's ownership I/O takes ownership and returns the buffer, so we reassign. let mut tmp_buf = vec![0u8; 8192]; + // c10k W11: two-stage idle park (see idle_park.rs). Cancel-capable + // streams register for the shard chore's ≥1s sweep; `downshifted` tracks + // whether this connection currently holds the probe-sized working set. + let idle_reg = S::SUPPORTS_IDLE_PARK.then(|| idle_park::register(client_id)); + let mut downshifted = false; + // Client idle timeout: 0 = disabled (read once, avoid lock on hot path) let idle_timeout_secs = ctx.runtime_config.read().timeout; let idle_timeout = if idle_timeout_secs > 0 { @@ -590,9 +597,16 @@ pub(crate) async fn handle_connection_sharded_monoio< } // Read data from stream using monoio ownership I/O. - // Reuse pre-allocated buffer; restore length to 8192 for the read. - if tmp_buf.len() < 8192 { - tmp_buf.resize(8192, 0); + // Reuse pre-allocated buffer; restore length for the read. While + // downshifted (c10k W11) the park size is the probe buffer; the + // resize below re-inflates lazily on the first post-idle iteration. + let park_len = if downshifted { + idle_park::IDLE_PROBE_BUF + } else { + idle_park::PARK_BUF_FULL + }; + if tmp_buf.len() != park_len { + tmp_buf.resize(park_len, 0); } if let Some(dur) = idle_timeout { // Timeout-aware read: select between read and sleep. @@ -650,6 +664,44 @@ pub(crate) async fn handle_connection_sharded_monoio< } continue; } + } else if downshifted { + // c10k W11 stage 2: park with the probe buffer, no chore state. + // Real data restores the full working set (lazily, via the + // pre-park sizing above) and re-arms stage 1. + let (result, returned_buf) = stream.read(tmp_buf).await; + tmp_buf = returned_buf; + match result { + Ok(0) => break, + Ok(n) => { + read_buf.extend_from_slice(&tmp_buf[..n]); + downshifted = false; + } + Err(_) => break, + } + } else if let Some(reg) = idle_reg.as_ref() { + // c10k W11 stage 1: full-size read, registered for the shard + // chore's ≥1s idle sweep. Cancel-and-await is loss-free: a + // completion that raced the cancel still delivers its bytes. + let handle = reg.slot.handle(); + reg.slot.mark_parked(ctx.cached_clock.ms()); + let (result, returned_buf) = stream.idle_park_read(tmp_buf, handle).await; + reg.slot.mark_unparked(); + tmp_buf = returned_buf; + match result { + Ok(0) => break, + Ok(n) => { + read_buf.extend_from_slice(&tmp_buf[..n]); + } + Err(_) => { + // Cancelled by the idle sweep — or a real socket error, + // which the stage-2 re-park's own read surfaces + // immediately. Either way: shed the working set, re-park + // small. No error-kind matching needed. + idle_park::downshift_idle_buffers(&mut tmp_buf, &mut read_buf, &mut write_buf); + downshifted = true; + continue; + } + } } else { let (result, returned_buf) = stream.read(tmp_buf).await; tmp_buf = returned_buf; diff --git a/src/shard/event_loop.rs b/src/shard/event_loop.rs index 3dcdae4e4..7ab0c1c57 100644 --- a/src/shard/event_loop.rs +++ b/src/shard/event_loop.rs @@ -2447,6 +2447,11 @@ impl super::Shard { } } timers::sync_wal_v3(&mut wal_writer); + // c10k W11: cancel reads parked ≥1s so idle connections + // downshift to the probe-sized working set. Same thread + // as the connection tasks (thread-per-core), no locks. + let _ = + crate::server::conn::handler_monoio::idle_park::sweep(cached_clock.ms()); // P3+MA1+MA2: MVCC committed prune + zombie sweep + kill old snapshots // + RECL_* + segment-stall. crate::shard::slice::with_shard(|s| { diff --git a/tests/idle_downshift_parity.rs b/tests/idle_downshift_parity.rs new file mode 100644 index 000000000..fdb238800 --- /dev/null +++ b/tests/idle_downshift_parity.rs @@ -0,0 +1,145 @@ +//! c10k W11: the idle-connection buffer downshift must be invisible on the +//! wire. A connection idle past the ~1s downshift threshold sheds its 8 KiB +//! working set and re-parks on a 512 B probe buffer — every byte sent +//! afterwards (single commands, >512 B pipelines, multi-KB values) must +//! still be served exactly as before, and an interleaved active connection +//! must never be disturbed by its idle sibling's downshift. +//! +//! (Runtime note: the downshift is monoio-only; under the tokio runtime the +//! sweep never fires and this suite degenerates to a plain parity check — +//! green either way.) + +mod common; + +use std::io::{Read, Write}; +use std::net::TcpStream; +use std::process::{Command, Stdio}; +use std::time::Duration; + +fn spawn_moon(dir: &std::path::Path, port: u16) -> std::process::Child { + Command::new(common::find_moon_binary()) + .args([ + "--port", + &port.to_string(), + "--shards", + "1", + "--dir", + dir.to_str().unwrap(), + "--disk-free-min-pct", + "0", + ]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .expect("spawn moon") +} + +fn read_exact_deadline(stream: &mut TcpStream, want: usize) -> Vec { + stream + .set_read_timeout(Some(Duration::from_secs(10))) + .expect("set timeout"); + let mut out = Vec::with_capacity(want); + let mut chunk = [0u8; 4096]; + while out.len() < want { + match stream.read(&mut chunk) { + Ok(0) => break, + Ok(n) => out.extend_from_slice(&chunk[..n]), + Err(e) => panic!("read failed after {} of {} bytes: {}", out.len(), want, e), + } + } + out +} + +fn ping(stream: &mut TcpStream) { + stream.write_all(b"PING\r\n").expect("write PING"); + let r = read_exact_deadline(stream, 7); + assert_eq!(&r, b"+PONG\r\n"); +} + +/// Idle past the downshift threshold, then drive traffic shaped to cross +/// every downshift boundary: a probe-sized wake, a >512 B pipeline (forces +/// the probe read to spill into a full-size follow-up), and a multi-KB +/// value (exercises read-buffer regrowth from released capacity). +#[test] +fn idle_connection_serves_all_traffic_after_downshift() { + let dir = tempfile::tempdir().expect("tempdir"); + let (mut child, port) = common::spawn_listening(|p| spawn_moon(dir.path(), p)); + + let mut conn = TcpStream::connect(("127.0.0.1", port)).expect("connect"); + conn.set_nodelay(true).ok(); + ping(&mut conn); + + // Cross the downshift threshold (1s) with margin for sweep cadence. + std::thread::sleep(Duration::from_millis(2600)); + + // 1. Small wake: fits the 512 B probe buffer. + ping(&mut conn); + + // 2. Idle again, then a 100-deep PING pipeline (700 B > probe size): + // the first probe read must deliver its prefix and the remainder must + // arrive via the re-inflated read path with zero loss or reorder. + std::thread::sleep(Duration::from_millis(2600)); + let pipeline = b"PING\r\n".repeat(100); + conn.write_all(&pipeline).expect("write pipeline"); + let replies = read_exact_deadline(&mut conn, 7 * 100); + assert_eq!(replies.len(), 700, "all 100 pipelined replies must arrive"); + assert!( + replies.chunks(7).all(|c| c == b"+PONG\r\n"), + "every pipelined reply must be +PONG" + ); + + // 3. Idle again, then a 4 KiB SET + GET round trip: the released + // read/write buffers must regrow transparently. + std::thread::sleep(Duration::from_millis(2600)); + let payload = vec![b'x'; 4096]; + let mut set_cmd = Vec::new(); + set_cmd.extend_from_slice(b"*3\r\n$3\r\nSET\r\n$8\r\nbigvalue\r\n$4096\r\n"); + set_cmd.extend_from_slice(&payload); + set_cmd.extend_from_slice(b"\r\n"); + conn.write_all(&set_cmd).expect("write SET"); + let r = read_exact_deadline(&mut conn, 5); + assert_eq!(&r, b"+OK\r\n"); + conn.write_all(b"*2\r\n$3\r\nGET\r\n$8\r\nbigvalue\r\n") + .expect("write GET"); + let want = format!("$4096\r\n{}\r\n", String::from_utf8(payload).unwrap()); + let r = read_exact_deadline(&mut conn, want.len()); + assert_eq!( + r, + want.as_bytes(), + "4 KiB value must survive the downshifted round trip" + ); + + let _ = child.kill(); + let _ = child.wait(); +} + +/// An active connection must be completely unaffected while an idle sibling +/// on the same shard is downshifted by the sweep. +#[test] +fn active_sibling_is_undisturbed_by_idle_downshift() { + let dir = tempfile::tempdir().expect("tempdir"); + let (mut child, port) = common::spawn_listening(|p| spawn_moon(dir.path(), p)); + + let mut idle = TcpStream::connect(("127.0.0.1", port)).expect("connect idle"); + ping(&mut idle); + + let mut active = TcpStream::connect(("127.0.0.1", port)).expect("connect active"); + active.set_nodelay(true).ok(); + + // ~3s of continuous activity spanning multiple sweep firings while the + // sibling sits idle and gets downshifted. + let deadline = std::time::Instant::now() + Duration::from_millis(3000); + let mut rounds = 0u32; + while std::time::Instant::now() < deadline { + ping(&mut active); + rounds += 1; + std::thread::sleep(Duration::from_millis(5)); + } + assert!(rounds > 100, "active connection must keep full throughput"); + + // The downshifted idle sibling still answers. + ping(&mut idle); + + let _ = child.kill(); + let _ = child.wait(); +} From 6ef45e527959084a5ca0fb3952ff762b649db261 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Wed, 29 Jul 2026 15:16:37 +0700 Subject: [PATCH 15/17] docs(changelog): c10k W11 idle-conn downshift entry author: Tin Dang --- CHANGELOG.md | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8ef5b2d40..d000b0d3e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -40,6 +40,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 monoio runtime only. ### Changed +- **Idle connections downshift to a ~0.5 KB working set (c10k W11).** + A connection parked in `read()` ≥1s has its read cancelled by the shard's + 1s chore (loss-free cancel-and-await on both io_uring and epoll/kqueue), + sheds its 8 KiB rent buffer and empty scratch buffers, and re-parks on a + 512 B probe buffer. Idle RSS/conn 43.5 → 19.9 KB (10k-conn VM A/B); + GCE-pinned p=1/p=16 A/B perf-neutral. TLS conns and the tokio runtime + keep the previous behavior. - **Cold command futures boxed out of the connection state machine (c10k P3).** TXN.ABORT rollback, leaked-txn disconnect teardown, and the FT.* body now `Box::pin` behind their name-check fast paths; the From d2f3e7cfffdb720393f77fd873b11c5f23621567 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Wed, 29 Jul 2026 16:48:39 +0700 Subject: [PATCH 16/17] =?UTF-8?q?fix(tests):=20add=20uring=5Fentries=20to?= =?UTF-8?q?=20test=20ServerConfig=20literals=20=E2=80=94=20repairs=20red?= =?UTF-8?q?=20tokio=20Check?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-2 fallout (--uring-entries, fdcd5d90): three tokio-only integration suites build ServerConfig struct literals and were never rebuilt by the lib-suite gates, so the missing `uring_entries` field only surfaced in CI's full-matrix Check job (E0063 in mq_integration, txn_kv_wiring, workspace_integration ×2). Adds `uring_entries: None` to all four literals. Also moves the round-2 `uring_entries_tests` module to the end of src/runtime/mod.rs (clippy `items_after_test_module`, visible only under --all-targets). Verified: all three suites compile and pass under --no-default-features --features runtime-tokio,jemalloc (VM; txn_kv_wiring's crash-recovery leg needs disk headroom — the macOS host root volume sits under the 5% diskfull floor). Lesson for the gate checklist: a config-struct field addition must be followed by `cargo clippy --all-targets` on BOTH matrices — lib suites do not compile integration tests. author: Tin Dang --- src/runtime/mod.rs | 36 +++++++++++++++++----------------- tests/mq_integration.rs | 1 + tests/txn_kv_wiring.rs | 1 + tests/workspace_integration.rs | 2 ++ 4 files changed, 22 insertions(+), 18 deletions(-) diff --git a/src/runtime/mod.rs b/src/runtime/mod.rs index a32667458..36bdb5350 100644 --- a/src/runtime/mod.rs +++ b/src/runtime/mod.rs @@ -387,24 +387,6 @@ mod yield_costfree_tests { } } -#[cfg(test)] -mod uring_entries_tests { - // Shared process-global: run assertions in one test to avoid ordering - // races between parallel test threads. - #[test] - fn uring_entries_default_set_and_clamp() { - assert_eq!(super::uring_entries(), None, "unset must read as None"); - super::set_uring_entries(4096); - assert_eq!(super::uring_entries(), Some(4096)); - super::set_uring_entries(64); - assert_eq!( - super::uring_entries(), - Some(super::MIN_URING_ENTRIES), - "sub-floor values clamp up to monoio's minimum" - ); - } -} - // --- Runtime-specific type aliases for TCP networking --- // These allow subsystem code to use a single type name regardless of runtime. @@ -455,3 +437,21 @@ pub type RuntimeFactoryImpl = TokioRuntimeFactory; #[cfg(feature = "runtime-monoio")] pub type RuntimeFactoryImpl = MonoioRuntimeFactory; + +#[cfg(test)] +mod uring_entries_tests { + // Shared process-global: run assertions in one test to avoid ordering + // races between parallel test threads. + #[test] + fn uring_entries_default_set_and_clamp() { + assert_eq!(super::uring_entries(), None, "unset must read as None"); + super::set_uring_entries(4096); + assert_eq!(super::uring_entries(), Some(4096)); + super::set_uring_entries(64); + assert_eq!( + super::uring_entries(), + Some(super::MIN_URING_ENTRIES), + "sub-floor values clamp up to monoio's minimum" + ); + } +} diff --git a/tests/mq_integration.rs b/tests/mq_integration.rs index f7119829c..ea6ac01cf 100644 --- a/tests/mq_integration.rs +++ b/tests/mq_integration.rs @@ -90,6 +90,7 @@ async fn start_mq_server(num_shards: usize) -> (u16, CancellationToken) { uring_sqpoll_ms: None, io_driver: "auto".to_string(), io_busy_poll_us: 0, + uring_entries: None, ft_search_workers: None, admin_port: 0, slowlog_log_slower_than: 10000, diff --git a/tests/txn_kv_wiring.rs b/tests/txn_kv_wiring.rs index 70d223590..a46eeb408 100644 --- a/tests/txn_kv_wiring.rs +++ b/tests/txn_kv_wiring.rs @@ -102,6 +102,7 @@ async fn start_txn_server(num_shards: usize, persistence_dir: &str) -> (u16, Can uring_sqpoll_ms: None, io_driver: "auto".to_string(), io_busy_poll_us: 0, + uring_entries: None, ft_search_workers: None, admin_port: 0, slowlog_log_slower_than: 10000, diff --git a/tests/workspace_integration.rs b/tests/workspace_integration.rs index 4dc7acf22..8e9eba567 100644 --- a/tests/workspace_integration.rs +++ b/tests/workspace_integration.rs @@ -83,6 +83,7 @@ async fn start_workspace_server(num_shards: usize) -> (u16, CancellationToken) { uring_sqpoll_ms: None, io_driver: "auto".to_string(), io_busy_poll_us: 0, + uring_entries: None, ft_search_workers: None, admin_port: 0, slowlog_log_slower_than: 10000, @@ -325,6 +326,7 @@ async fn start_workspace_server_with_auth( uring_sqpoll_ms: None, io_driver: "auto".to_string(), io_busy_poll_us: 0, + uring_entries: None, ft_search_workers: None, admin_port: 0, slowlog_log_slower_than: 10000, From d2d01449050fc61e77a66a6e107a2067973281b5 Mon Sep 17 00:00:00 2001 From: Tin Dang Date: Wed, 29 Jul 2026 16:48:49 +0700 Subject: [PATCH 17/17] =?UTF-8?q?feat(server):=20c10k=20P4b=20=E2=80=94=20?= =?UTF-8?q?vendor=20TLS=20wrapper=20stack,=20idle=20TLS=20conns=20shed=204?= =?UTF-8?q?0=20KB=20(round=204)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round 4 of the c10k campaign, per the c1M RFC sequencing: the TLS diet. Root cause (round 3): monoio-io-wrapper eagerly allocates two fixed 16 KiB Box<[u8]> buffers per TLS stream and never frees them, and monoio-rustls has no cancelable read, so W11's idle downshift could not cover TLS connections at all — idle TLS RSS sat flat at ~87.5 KB/conn. Changes (vendored crates follow the existing vendor/monoio "moon patch" comment style; wired via [patch.crates-io]): - vendor/monoio-io-wrapper: Buffer is lazy (allocated on first use) and releasable (release_if_empty, only while drained; reallocates transparently on next use). New do_io_cancelable whose errors — including ECANCELED — are returned directly and NOT stashed for status replay: a stashed cancel error would resurface on the next plain read and tear down a healthy connection (unit-tested in-crate, 5 tests). - vendor/monoio-rustls: Stream::release_idle_buffers() and a CancelableAsyncReadRent impl that forwards the cancel handle into the wrapper's socket read. Loss-free under cancel-and-await: bytes from a completion racing the cancel persist in the wrapper buffer + rustls deframer and are consumed by the next read. - src/server/conn/handler_monoio/idle_park.rs: TLS streams now set SUPPORTS_IDLE_PARK = true and implement the new on_idle_downshift hook (releases both wrapper buffers). Lazy-alloc alone would be worthless — the handshake itself fills both buffers, so the win requires the idle-park cancel to fire the release. - tests/tls_idle_downshift_parity.rs: wire parity on one continuous TLS session across three downshift/wake cycles (probe wake, 700 B pipeline, 4 KiB value round trip), rustls client with throwaway cert; green on kqueue (macOS) and io_uring (VM). E9 A/B (moon-dev VM, 3000 idle TLS conns, shards=2): pre-P4b 87.5 KB/conn flat forever → P4b 47.4 KB/conn after 30 s idle (−46%), re-downshifts to 48.7 after a wake-sweep; sweeps bad=0; RSS returns on close. Matches the RFC's −30–40 KB estimate. Remaining TLS floor (rustls session ~15–20 KB + 9.9 KB task future) is Proposal 1 territory. Gates: tokio lib suite 3609 pass; VM monoio lib suite pass; fmt + both clippy matrices clean; unsafe/unwrap audits pass. refs: tmp/C10K-REVIEW.md round 4, .planning/rfcs/c1m-connection-plane.md author: Tin Dang --- CHANGELOG.md | 12 + Cargo.lock | 4 - Cargo.toml | 2 + src/server/conn/handler_monoio/idle_park.rs | 39 +- src/server/conn/handler_monoio/mod.rs | 1 + tests/tls_idle_downshift_parity.rs | 232 +++++++++ vendor/monoio-io-wrapper/Cargo.toml | 35 ++ vendor/monoio-io-wrapper/Cargo.toml.orig | 19 + vendor/monoio-io-wrapper/README.adoc | 13 + vendor/monoio-io-wrapper/src/lib.rs | 191 +++++++ vendor/monoio-io-wrapper/src/safe_io.rs | 524 ++++++++++++++++++++ vendor/monoio-io-wrapper/src/unsafe_io.rs | 191 +++++++ vendor/monoio-rustls/Cargo.toml | 74 +++ vendor/monoio-rustls/Cargo.toml.orig | 32 ++ vendor/monoio-rustls/README.md | 1 + vendor/monoio-rustls/examples/connect.rs | 36 ++ vendor/monoio-rustls/src/client.rs | 78 +++ vendor/monoio-rustls/src/error.rs | 20 + vendor/monoio-rustls/src/lib.rs | 31 ++ vendor/monoio-rustls/src/server.rs | 74 +++ vendor/monoio-rustls/src/stream.rs | 383 ++++++++++++++ 21 files changed, 1981 insertions(+), 11 deletions(-) create mode 100644 tests/tls_idle_downshift_parity.rs create mode 100644 vendor/monoio-io-wrapper/Cargo.toml create mode 100644 vendor/monoio-io-wrapper/Cargo.toml.orig create mode 100644 vendor/monoio-io-wrapper/README.adoc create mode 100644 vendor/monoio-io-wrapper/src/lib.rs create mode 100644 vendor/monoio-io-wrapper/src/safe_io.rs create mode 100644 vendor/monoio-io-wrapper/src/unsafe_io.rs create mode 100644 vendor/monoio-rustls/Cargo.toml create mode 100644 vendor/monoio-rustls/Cargo.toml.orig create mode 100644 vendor/monoio-rustls/README.md create mode 100644 vendor/monoio-rustls/examples/connect.rs create mode 100644 vendor/monoio-rustls/src/client.rs create mode 100644 vendor/monoio-rustls/src/error.rs create mode 100644 vendor/monoio-rustls/src/lib.rs create mode 100644 vendor/monoio-rustls/src/server.rs create mode 100644 vendor/monoio-rustls/src/stream.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index d000b0d3e..867756bfa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -47,6 +47,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 512 B probe buffer. Idle RSS/conn 43.5 → 19.9 KB (10k-conn VM A/B); GCE-pinned p=1/p=16 A/B perf-neutral. TLS conns and the tokio runtime keep the previous behavior. +- **Idle TLS connections shed their 32 KB wrapper buffers (c10k P4b).** + `monoio-rustls` + `monoio-io-wrapper` are now vendored (moon-patch style, + like `vendor/monoio`): the wrapper's two eagerly-allocated, never-freed + 16 KiB per-stream buffers are lazy + releasable, and the TLS stream + implements a loss-free cancelable read, so W11's idle downshift now + covers TLS connections too — a parked TLS conn drops moon's working set + AND both wrapper buffers (they reallocate lazily on the next I/O). + Cancel errors are deliberately not stashed for replay in the wrapper + (a stashed ECANCELED would resurface on the next read and tear down a + healthy connection). Wire parity across downshift/wake cycles is + regression-tested on one continuous TLS session + (`tests/tls_idle_downshift_parity.rs`). - **Cold command futures boxed out of the connection state machine (c10k P3).** TXN.ABORT rollback, leaked-txn disconnect teardown, and the FT.* body now `Box::pin` behind their name-check fast paths; the diff --git a/Cargo.lock b/Cargo.lock index 323d2e6ad..ea90fe2c9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1788,8 +1788,6 @@ dependencies = [ [[package]] name = "monoio-io-wrapper" version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4bcfaa76e5daf87cc4d31b4d1b6bc93c12db59c19df50b9200afdbde42077655" dependencies = [ "monoio", ] @@ -1808,8 +1806,6 @@ dependencies = [ [[package]] name = "monoio-rustls" version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6e31f422825bd7fb19957af6eaf89d7234ba143fcc0e515f5a2f526e332d1875" dependencies = [ "bytes", "monoio", diff --git a/Cargo.toml b/Cargo.toml index 98b3fff08..d7f2177d3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -135,6 +135,8 @@ text-index = [ # when the MOON_* spin knobs are unset. [patch.crates-io] monoio = { path = "vendor/monoio" } +monoio-rustls = { path = "vendor/monoio-rustls" } +monoio-io-wrapper = { path = "vendor/monoio-io-wrapper" } [target.'cfg(target_os = "linux")'.dependencies] io-uring = "0.7" diff --git a/src/server/conn/handler_monoio/idle_park.rs b/src/server/conn/handler_monoio/idle_park.rs index 24f76a9a3..ad06a8521 100644 --- a/src/server/conn/handler_monoio/idle_park.rs +++ b/src/server/conn/handler_monoio/idle_park.rs @@ -25,9 +25,11 @@ //! Per-park hot-path cost: two `Cell` stores + one `Rc` clone of the //! `CancelHandle` (no allocation). //! -//! TLS connections (`monoio_rustls` streams) do not implement cancelable -//! reads; they opt out via [`IdleParkRead::SUPPORTS_IDLE_PARK`] and keep -//! the pre-W11 behavior byte-for-byte. +//! TLS connections participate too (c10k P4b): the vendored monoio-rustls +//! implements a loss-free cancelable read, and their +//! [`IdleParkRead::on_idle_downshift`] additionally releases the wrapper's +//! two 16 KiB stream buffers (vendored monoio-io-wrapper reallocates them +//! lazily on the next I/O). use std::cell::{Cell, RefCell}; use std::collections::HashMap; @@ -164,6 +166,11 @@ pub(crate) trait IdleParkRead: AsyncReadRent { ) -> impl std::future::Future>> { self.read(buf) } + + /// Stream-side hook fired when the handler downshifts after a cancelled + /// 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) {} } impl IdleParkRead for monoio::net::TcpStream { @@ -178,10 +185,28 @@ impl IdleParkRead for monoio::net::TcpStream { } } -/// TLS streams keep the pre-W11 behavior (no cancelable read in -/// monoio_rustls; dropping its read future mid-handshake-record would also -/// desync the TLS session state). -impl IdleParkRead for monoio_rustls::ServerTlsStream {} +/// TLS streams downshift too (c10k P4b): the vendored monoio-rustls +/// cancelable read forwards the cancel to the inner socket read; raced +/// bytes persist in the wrapper buffer + rustls deframer, and the cancel +/// error is deliberately NOT stashed for replay (vendored +/// `SafeRead::do_io_cancelable`), so the stage-2 probe read resumes the +/// session cleanly. The downshift hook additionally sheds the wrapper's +/// two 16 KiB stream buffers. +impl IdleParkRead for monoio_rustls::ServerTlsStream { + const SUPPORTS_IDLE_PARK: bool = true; + + fn idle_park_read( + &mut self, + buf: Vec, + c: CancelHandle, + ) -> impl std::future::Future>> { + monoio::io::CancelableAsyncReadRent::cancelable_read(self, buf, c) + } + + fn on_idle_downshift(&mut self) { + self.release_idle_buffers(); + } +} #[cfg(test)] mod idle_park_tests { diff --git a/src/server/conn/handler_monoio/mod.rs b/src/server/conn/handler_monoio/mod.rs index 1cd1ef643..5fdc74ee5 100644 --- a/src/server/conn/handler_monoio/mod.rs +++ b/src/server/conn/handler_monoio/mod.rs @@ -698,6 +698,7 @@ pub(crate) async fn handle_connection_sharded_monoio< // immediately. Either way: shed the working set, re-park // small. No error-kind matching needed. idle_park::downshift_idle_buffers(&mut tmp_buf, &mut read_buf, &mut write_buf); + stream.on_idle_downshift(); downshifted = true; continue; } diff --git a/tests/tls_idle_downshift_parity.rs b/tests/tls_idle_downshift_parity.rs new file mode 100644 index 000000000..434225388 --- /dev/null +++ b/tests/tls_idle_downshift_parity.rs @@ -0,0 +1,232 @@ +//! c10k P4b: the TLS idle-connection downshift must be invisible on the +//! wire. A TLS connection idle past the ~1s threshold gets its parked read +//! cancelled (vendored monoio-rustls cancelable read), sheds moon's working +//! set AND the wrapper's two 16 KiB stream buffers, and re-parks on a probe +//! buffer — every byte sent afterwards must still be served exactly as +//! before over the SAME TLS session (a cancel that desyncs the session or +//! replays a stale error surfaces here as a broken read). +//! +//! (Runtime note: the downshift is monoio-only; under the tokio runtime the +//! sweep never fires and this suite degenerates to a plain TLS parity +//! check — green either way.) + +mod common; + +use std::io::{Read, Write}; +use std::net::TcpStream; +use std::process::{Command, Stdio}; +use std::sync::Arc; +use std::time::Duration; + +/// Accept any server certificate — the server uses a throwaway self-signed +/// cert generated per test run; transport privacy is irrelevant here. +#[derive(Debug)] +struct AcceptAnyCert(rustls::crypto::CryptoProvider); + +impl rustls::client::danger::ServerCertVerifier for AcceptAnyCert { + fn verify_server_cert( + &self, + _end_entity: &rustls::pki_types::CertificateDer<'_>, + _intermediates: &[rustls::pki_types::CertificateDer<'_>], + _server_name: &rustls::pki_types::ServerName<'_>, + _ocsp_response: &[u8], + _now: rustls::pki_types::UnixTime, + ) -> Result { + Ok(rustls::client::danger::ServerCertVerified::assertion()) + } + + fn verify_tls12_signature( + &self, + _message: &[u8], + _cert: &rustls::pki_types::CertificateDer<'_>, + _dss: &rustls::DigitallySignedStruct, + ) -> Result { + Ok(rustls::client::danger::HandshakeSignatureValid::assertion()) + } + + fn verify_tls13_signature( + &self, + _message: &[u8], + _cert: &rustls::pki_types::CertificateDer<'_>, + _dss: &rustls::DigitallySignedStruct, + ) -> Result { + Ok(rustls::client::danger::HandshakeSignatureValid::assertion()) + } + + fn supported_verify_schemes(&self) -> Vec { + self.0.signature_verification_algorithms.supported_schemes() + } +} + +struct TlsConn { + session: rustls::ClientConnection, + sock: TcpStream, +} + +impl TlsConn { + fn connect(port: u16) -> Self { + let provider = rustls::crypto::aws_lc_rs::default_provider(); + let config = rustls::ClientConfig::builder_with_provider(Arc::new(provider.clone())) + .with_safe_default_protocol_versions() + .expect("protocol versions") + .dangerous() + .with_custom_certificate_verifier(Arc::new(AcceptAnyCert(provider))) + .with_no_client_auth(); + let server_name = rustls::pki_types::ServerName::try_from("localhost").unwrap(); + let session = + rustls::ClientConnection::new(Arc::new(config), server_name).expect("client conn"); + let sock = TcpStream::connect(("127.0.0.1", port)).expect("tcp connect"); + sock.set_nodelay(true).ok(); + sock.set_read_timeout(Some(Duration::from_secs(10))).ok(); + Self { session, sock } + } + + fn write_all(&mut self, data: &[u8]) { + rustls::Stream::new(&mut self.session, &mut self.sock) + .write_all(data) + .expect("tls write"); + } + + fn read_exact_deadline(&mut self, want: usize) -> Vec { + let mut tls = rustls::Stream::new(&mut self.session, &mut self.sock); + let mut out = Vec::with_capacity(want); + let mut chunk = [0u8; 4096]; + while out.len() < want { + match tls.read(&mut chunk) { + Ok(0) => break, + Ok(n) => out.extend_from_slice(&chunk[..n]), + Err(e) => panic!( + "tls read failed after {} of {} bytes: {}", + out.len(), + want, + e + ), + } + } + out + } + + fn ping(&mut self) { + self.write_all(b"PING\r\n"); + let r = self.read_exact_deadline(7); + assert_eq!(&r, b"+PONG\r\n"); + } +} + +/// Generate a throwaway self-signed cert. Returns false (test skips) if the +/// `openssl` CLI is unavailable on this machine. +fn generate_cert(dir: &std::path::Path) -> bool { + let status = Command::new("openssl") + .args([ + "req", + "-x509", + "-newkey", + "rsa:2048", + "-nodes", + "-days", + "1", + "-subj", + "/CN=localhost", + ]) + .arg("-keyout") + .arg(dir.join("key.pem")) + .arg("-out") + .arg(dir.join("cert.pem")) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status(); + matches!(status, Ok(s) if s.success()) +} + +fn spawn_moon_tls(dir: &std::path::Path, port: u16, tls_port: u16) -> std::process::Child { + Command::new(common::find_moon_binary()) + .args([ + "--port", + &port.to_string(), + "--shards", + "1", + "--dir", + dir.join("data").to_str().unwrap(), + "--disk-free-min-pct", + "0", + "--tls-port", + &tls_port.to_string(), + ]) + .arg("--tls-cert-file") + .arg(dir.join("cert.pem")) + .arg("--tls-key-file") + .arg(dir.join("key.pem")) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .expect("spawn moon with TLS") +} + +/// Idle a TLS connection across multiple sweep firings, then drive traffic +/// shaped to cross every downshift boundary: a probe-sized wake, a >512 B +/// pipeline, and a multi-KB value (regrows both moon's buffers and the +/// released TLS wrapper buffers). All on one TLS session — any session +/// desync from the cancel path fails loudly. +#[test] +fn tls_idle_connection_serves_all_traffic_after_downshift() { + let dir = tempfile::tempdir().expect("tempdir"); + if !generate_cert(dir.path()) { + eprintln!("SKIP: openssl CLI not available for cert generation"); + return; + } + let tls_port = common::reserve_port(); + let (mut child, _port) = common::spawn_listening(|p| spawn_moon_tls(dir.path(), p, tls_port)); + // spawn_listening waits on the plain port; give the TLS listener a + // moment if it comes up second. + let mut conn = None; + for _ in 0..50 { + match std::panic::catch_unwind(|| TlsConn::connect(tls_port)) { + Ok(c) => { + conn = Some(c); + break; + } + Err(_) => std::thread::sleep(Duration::from_millis(100)), + } + } + let mut conn = conn.expect("TLS connect"); + conn.ping(); + + // Cross the downshift threshold (1s) with margin for sweep cadence. + std::thread::sleep(Duration::from_millis(2600)); + + // 1. Small wake over the downshifted TLS session. + conn.ping(); + + // 2. Idle again, then a 100-deep PING pipeline (700 B plaintext). + std::thread::sleep(Duration::from_millis(2600)); + let pipeline = b"PING\r\n".repeat(100); + conn.write_all(&pipeline); + let replies = conn.read_exact_deadline(7 * 100); + assert_eq!(replies.len(), 700, "all 100 pipelined replies must arrive"); + assert!( + replies.chunks(7).all(|c| c == b"+PONG\r\n"), + "every pipelined reply must be +PONG" + ); + + // 3. Idle again, then a 4 KiB SET + GET round trip. + std::thread::sleep(Duration::from_millis(2600)); + let payload = vec![b'x'; 4096]; + let mut set_cmd = Vec::new(); + set_cmd.extend_from_slice(b"*3\r\n$3\r\nSET\r\n$8\r\nbigvalue\r\n$4096\r\n"); + set_cmd.extend_from_slice(&payload); + set_cmd.extend_from_slice(b"\r\n"); + conn.write_all(&set_cmd); + let r = conn.read_exact_deadline(5); + assert_eq!(&r, b"+OK\r\n"); + conn.write_all(b"*2\r\n$3\r\nGET\r\n$8\r\nbigvalue\r\n"); + let want = format!("$4096\r\n{}\r\n", String::from_utf8(payload).unwrap()); + let r = conn.read_exact_deadline(want.len()); + assert_eq!( + r, + want.as_bytes(), + "4 KiB value must survive the downshifted TLS round trip" + ); + + let _ = child.kill(); + let _ = child.wait(); +} diff --git a/vendor/monoio-io-wrapper/Cargo.toml b/vendor/monoio-io-wrapper/Cargo.toml new file mode 100644 index 000000000..9e54e20d8 --- /dev/null +++ b/vendor/monoio-io-wrapper/Cargo.toml @@ -0,0 +1,35 @@ +# THIS FILE IS AUTOMATICALLY GENERATED BY CARGO +# +# When uploading crates to the registry Cargo will automatically +# "normalize" Cargo.toml files for maximal compatibility +# with all versions of Cargo and also rewrite `path` dependencies +# to registry (e.g., crates.io) dependencies. +# +# If you are reading this file be aware that the original Cargo.toml +# will likely look very different (and much more reasonable). +# See Cargo.toml.orig for the original contents. + +[package] +edition = "2021" +name = "monoio-io-wrapper" +version = "0.1.1" +authors = [ + "ChiHai ", + "Rain Jiang ", +] +description = "A read/write wrapper to bridge sync and async io" +homepage = "https://github.com/monoio-rs/monoio-tls" +readme = "README.adoc" +categories = [ + "asynchronous", + "network-programming", +] +license = "MIT/Apache-2.0" +repository = "https://github.com/monoio-rs/monoio-tls" + +[dependencies.monoio] +version = "0.2.0" + +[features] +default = [] +unsafe_io = [] diff --git a/vendor/monoio-io-wrapper/Cargo.toml.orig b/vendor/monoio-io-wrapper/Cargo.toml.orig new file mode 100644 index 000000000..dd0090e6a --- /dev/null +++ b/vendor/monoio-io-wrapper/Cargo.toml.orig @@ -0,0 +1,19 @@ +[package] +name = "monoio-io-wrapper" +version = "0.1.1" + +authors = ["ChiHai ", "Rain Jiang "] +categories = ["asynchronous", "network-programming"] +description = "A read/write wrapper to bridge sync and async io" +edition = "2021" +homepage = "https://github.com/monoio-rs/monoio-tls" +license = "MIT/Apache-2.0" +readme = "README.adoc" +repository = "https://github.com/monoio-rs/monoio-tls" + +[dependencies] +monoio = { workspace = true } + +[features] +default = [] +unsafe_io = [] diff --git a/vendor/monoio-io-wrapper/README.adoc b/vendor/monoio-io-wrapper/README.adoc new file mode 100644 index 000000000..d81d05a4f --- /dev/null +++ b/vendor/monoio-io-wrapper/README.adoc @@ -0,0 +1,13 @@ += Monoio IO Wrapper +An io wrapper to bind std io and monoio async io. + +|=== +|Return |do_io |read / write + +.2+|SafeIO .2+| Current async r/w result | 1. WouldBlock: Empty(r) or Full(w) and need calling `do_io` +|2. Other: success io or last `do_io` error + +.2+|UnsafeIO | 1. WouldBlock: Not capture mem block info, need calling `read` / `write` | 1. WouldBlock: mem block info is captured and need calling `do_io` +|2. Other: current async r/w result | 2. Other: success io or last `do_io` error + +|=== \ No newline at end of file diff --git a/vendor/monoio-io-wrapper/src/lib.rs b/vendor/monoio-io-wrapper/src/lib.rs new file mode 100644 index 000000000..173e65cc9 --- /dev/null +++ b/vendor/monoio-io-wrapper/src/lib.rs @@ -0,0 +1,191 @@ +#![allow(clippy::unsafe_removed_from_name)] + +use monoio::io::{AsyncReadRent, AsyncWriteRent}; + +mod safe_io; +#[cfg(feature = "unsafe_io")] +mod unsafe_io; + +#[derive(Debug)] +pub enum ReadBuffer { + Safe(safe_io::SafeRead), + #[cfg(feature = "unsafe_io")] + Unsafe(unsafe_io::UnsafeRead), +} + +#[derive(Debug)] +pub enum WriteBuffer { + Safe(safe_io::SafeWrite), + #[cfg(feature = "unsafe_io")] + Unsafe(unsafe_io::UnsafeWrite), +} + +impl ReadBuffer { + /// Create a new ReadBuffer with given buffer size. + #[inline] + pub fn new(buffer_size: usize) -> Self { + Self::Safe(safe_io::SafeRead::new(buffer_size)) + } + + /// Create a new ReadBuffer that uses unsafe I/O. + /// # Safety + /// Users must make sure the buffer ptr and len is valid until io finished. + /// So the Future cannot be dropped directly. Consider using CancellableIO. + #[inline] + #[cfg(feature = "unsafe_io")] + pub const unsafe fn new_unsafe() -> Self { + Self::Unsafe(unsafe_io::UnsafeRead::new()) + } + + #[inline] + pub async fn do_io(&mut self, mut io: IO) -> std::io::Result { + match self { + Self::Safe(b) => b.do_io(&mut io).await, + #[cfg(feature = "unsafe_io")] + Self::Unsafe(b) => unsafe { b.do_io(&mut io).await }, + } + } + + /// moon patch (c10k P4b): cancelable variant of `do_io`. Cancel errors + /// surface directly without being stashed for replay (see + /// `SafeRead::do_io_cancelable`). The unsafe-io variant has no cancel + /// support and falls back to a plain read. + #[inline] + pub async fn do_io_cancelable( + &mut self, + mut io: IO, + c: monoio::io::CancelHandle, + ) -> std::io::Result { + match self { + Self::Safe(b) => b.do_io_cancelable(&mut io, c).await, + #[cfg(feature = "unsafe_io")] + Self::Unsafe(b) => unsafe { b.do_io(&mut io).await }, + } + } + + /// moon patch (c10k P4b): drop the backing storage while drained; it + /// reallocates lazily on next use. No-op (false) for unsafe-io buffers. + #[inline] + pub fn release_if_empty(&mut self) -> bool { + match self { + Self::Safe(b) => b.release_if_empty(), + #[cfg(feature = "unsafe_io")] + Self::Unsafe(_) => false, + } + } + + #[inline] + #[cfg(feature = "unsafe_io")] + pub fn is_safe(&self) -> bool { + match self { + Self::Safe(_) => true, + #[cfg(feature = "unsafe_io")] + Self::Unsafe(_) => false, + } + } + + #[inline] + #[cfg(not(feature = "unsafe_io"))] + pub const fn is_safe(&self) -> bool { + true + } +} + +impl Default for ReadBuffer { + #[inline] + fn default() -> Self { + Self::Safe(Default::default()) + } +} + +impl std::io::Read for ReadBuffer { + #[inline] + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + match self { + Self::Safe(b) => b.read(buf), + #[cfg(feature = "unsafe_io")] + Self::Unsafe(b) => b.read(buf), + } + } +} + +impl WriteBuffer { + /// Create a new WriteBuffer with given buffer size. + #[inline] + pub fn new(buffer_size: usize) -> Self { + Self::Safe(safe_io::SafeWrite::new(buffer_size)) + } + + /// Create a new WriteBuffer that uses unsafe I/O. + /// # Safety + /// Users must make sure the buffer ptr and len is valid until io finished. + /// So the Future cannot be dropped directly. Consider using CancellableIO. + #[inline] + #[cfg(feature = "unsafe_io")] + pub const unsafe fn new_unsafe() -> Self { + Self::Unsafe(unsafe_io::UnsafeWrite::new()) + } + + #[inline] + pub async fn do_io(&mut self, mut io: IO) -> std::io::Result { + match self { + Self::Safe(buf) => buf.do_io(&mut io).await, + #[cfg(feature = "unsafe_io")] + Self::Unsafe(buf) => unsafe { buf.do_io(&mut io).await }, + } + } + + /// moon patch (c10k P4b): drop the backing storage while drained; it + /// reallocates lazily on next use. No-op (false) for unsafe-io buffers. + #[inline] + pub fn release_if_empty(&mut self) -> bool { + match self { + Self::Safe(b) => b.release_if_empty(), + #[cfg(feature = "unsafe_io")] + Self::Unsafe(_) => false, + } + } + + #[inline] + #[cfg(feature = "unsafe_io")] + pub fn is_safe(&self) -> bool { + match self { + Self::Safe(_) => true, + #[cfg(feature = "unsafe_io")] + Self::Unsafe(_) => false, + } + } + + #[inline] + #[cfg(not(feature = "unsafe_io"))] + pub const fn is_safe(&self) -> bool { + true + } +} + +impl Default for WriteBuffer { + #[inline] + fn default() -> Self { + Self::Safe(Default::default()) + } +} + +impl std::io::Write for WriteBuffer { + #[inline] + fn write(&mut self, buf: &[u8]) -> std::io::Result { + match self { + Self::Safe(b) => b.write(buf), + #[cfg(feature = "unsafe_io")] + Self::Unsafe(b) => b.write(buf), + } + } + + #[inline] + fn flush(&mut self) -> std::io::Result<()> { + match self { + Self::Safe(b) => b.flush(), + #[cfg(feature = "unsafe_io")] + Self::Unsafe(b) => b.flush(), + } + } +} diff --git a/vendor/monoio-io-wrapper/src/safe_io.rs b/vendor/monoio-io-wrapper/src/safe_io.rs new file mode 100644 index 000000000..072e67112 --- /dev/null +++ b/vendor/monoio-io-wrapper/src/safe_io.rs @@ -0,0 +1,524 @@ +use std::{fmt::Debug, io, mem}; + +use monoio::{ + buf::{IoBuf, IoBufMut}, + io::{AsyncReadRent, AsyncWriteRent, AsyncWriteRentExt, CancelHandle, CancelableAsyncReadRent}, +}; + +const BUFFER_SIZE: usize = 16 * 1024; + +struct Buffer { + read: usize, + write: usize, + // moon patch (c10k P4b): target capacity. `buf` is either exactly this + // size or released (len 0, only legal while drained); storage is + // materialized on first use and can be dropped again via + // `release_if_empty`, so an idle stream holds no backing allocation. + cap: usize, + buf: Box<[u8]>, +} + +impl Default for Buffer { + fn default() -> Self { + Self::new(BUFFER_SIZE) + } +} + +impl Buffer { + fn new(size: usize) -> Self { + Self { + read: 0, + write: 0, + cap: size, + // moon patch (c10k P4b): lazy — allocated on first use. + buf: Box::default(), + } + } + + // moon patch (c10k P4b): materialize backing storage before an I/O or + // copy-in touches it. Only legal transitions: released <-> allocated, + // both while drained. + fn ensure_allocated(&mut self) { + if self.buf.len() != self.cap { + debug_assert!(self.is_empty()); + self.buf = vec![0; self.cap].into_boxed_slice(); + } + } + + // moon patch (c10k P4b): drop the backing storage while drained. + fn release_if_empty(&mut self) -> bool { + if self.is_empty() && !self.buf.is_empty() { + self.read = 0; + self.write = 0; + self.buf = Box::default(); + true + } else { + false + } + } + + fn len(&self) -> usize { + self.write - self.read + } + + fn is_empty(&self) -> bool { + self.len() == 0 + } + + fn available(&self) -> usize { + self.buf.len() - self.write + } + + fn is_full(&self) -> bool { + self.available() == 0 + } + + fn advance(&mut self, n: usize) { + assert!(self.read + n <= self.write); + self.read += n; + if self.read == self.write { + self.read = 0; + self.write = 0; + } + } +} + +unsafe impl monoio::buf::IoBuf for Buffer { + fn read_ptr(&self) -> *const u8 { + unsafe { self.buf.as_ptr().add(self.read) } + } + + fn bytes_init(&self) -> usize { + self.write - self.read + } +} + +unsafe impl monoio::buf::IoBufMut for Buffer { + fn write_ptr(&mut self) -> *mut u8 { + unsafe { self.buf.as_mut_ptr().add(self.write) } + } + + fn bytes_total(&mut self) -> usize { + self.buf.len() - self.write + } + + unsafe fn set_init(&mut self, pos: usize) { + self.write += pos; + } +} + +pub struct SafeRead { + // the option is only meant for temporary take, it always should be some + buffer: Option, + status: ReadStatus, +} + +impl Debug for SafeRead { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("SafeRead") + .field("status", &self.status) + .finish() + } +} + +#[derive(Debug)] +enum ReadStatus { + Eof, + Err(io::Error), + Ok, +} + +impl Default for SafeRead { + fn default() -> Self { + Self { + buffer: Some(Buffer::default()), + status: ReadStatus::Ok, + } + } +} + +impl SafeRead { + /// Create a new SafeRead with given buffer size. + pub fn new(buffer_size: usize) -> Self { + Self { + buffer: Some(Buffer::new(buffer_size)), + status: ReadStatus::Ok, + } + } + + /// `do_io` do async read from io to inner buffer. + /// # Handle return value + /// _: the read result. + pub async fn do_io(&mut self, mut io: IO) -> io::Result { + // if there are some data inside the buffer, just return. + let buffer = self.buffer.as_mut().expect("buffer mut expected"); + if !buffer.is_empty() { + return Ok(buffer.len()); + } + // moon patch (c10k P4b): lazily (re)materialize released storage. + buffer.ensure_allocated(); + + // read from raw io + // # Safety + // We have already checked it is not None. + let buffer = unsafe { self.buffer.take().unwrap_unchecked() }; + let (result, buf) = io.read(buffer).await; + self.buffer = Some(buf); + match result { + Ok(0) => { + self.status = ReadStatus::Eof; + result + } + Ok(_) => { + self.status = ReadStatus::Ok; + result + } + Err(e) => { + let rerr = e.kind().into(); + self.status = ReadStatus::Err(e); + Err(rerr) + } + } + } + + /// moon patch (c10k P4b): cancelable twin of `do_io`, used by the + /// idle-park sweep path. Errors from the cancelable read — including + /// ECANCELED from a fired cancel — are returned directly and NOT stashed + /// in `status`: stashing would replay a spurious error into the next + /// plain read after an idle-park cancel and tear down a healthy + /// connection. EOF is genuine state and is still recorded. + pub async fn do_io_cancelable( + &mut self, + mut io: IO, + c: CancelHandle, + ) -> io::Result { + let buffer = self.buffer.as_mut().expect("buffer mut expected"); + if !buffer.is_empty() { + return Ok(buffer.len()); + } + buffer.ensure_allocated(); + + // # Safety + // We have already checked it is not None. + let buffer = unsafe { self.buffer.take().unwrap_unchecked() }; + let (result, buf) = io.cancelable_read(buffer, c).await; + self.buffer = Some(buf); + match &result { + Ok(0) => self.status = ReadStatus::Eof, + Ok(_) => self.status = ReadStatus::Ok, + Err(_) => {} + } + result + } + + /// moon patch (c10k P4b): drop the backing storage while drained; it + /// reallocates lazily on the next `do_io`. Returns whether storage was + /// released. + pub fn release_if_empty(&mut self) -> bool { + self.buffer + .as_mut() + .expect("buffer mut expected") + .release_if_empty() + } +} + +impl io::Read for SafeRead { + /// `read` from buffer. + /// # Handle return value + /// 1. Err(WouldBlock): the buffer is empty, call do_io to fetch more. + /// 2. _: handle it. + fn read(&mut self, buf: &mut [u8]) -> io::Result { + // if buffer is empty, return WoundBlock. + let buffer = self.buffer.as_mut().expect("buffer mut expected"); + if buffer.is_empty() { + return match mem::replace(&mut self.status, ReadStatus::Ok) { + ReadStatus::Eof => Ok(0), + ReadStatus::Err(e) => Err(e), + ReadStatus::Ok => Err(io::ErrorKind::WouldBlock.into()), + }; + } + + // now buffer is not empty. copy it. + let to_copy = buffer.len().min(buf.len()); + unsafe { std::ptr::copy_nonoverlapping(buffer.read_ptr(), buf.as_mut_ptr(), to_copy) }; + buffer.advance(to_copy); + + Ok(to_copy) + } +} + +pub struct SafeWrite { + // the option is only meant for temporary take, it always should be some + buffer: Option, + status: WriteStatus, +} + +impl Debug for SafeWrite { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("SafeWrite") + .field("status", &self.status) + .finish() + } +} + +#[derive(Debug)] +enum WriteStatus { + Err(io::Error), + Ok, +} + +impl Default for SafeWrite { + fn default() -> Self { + Self { + buffer: Some(Buffer::default()), + status: WriteStatus::Ok, + } + } +} + +impl SafeWrite { + /// Create a new SafeWrite with given buffer size. + pub fn new(buffer_size: usize) -> Self { + Self { + buffer: Some(Buffer::new(buffer_size)), + status: WriteStatus::Ok, + } + } + + /// `do_io` do async write from inner buffer to io. + /// # Handle return value + /// _: the write_all result(note: the data may have been written even when error). + pub async fn do_io(&mut self, mut io: IO) -> io::Result { + // if the buffer is empty, just return. + let buffer = self.buffer.as_ref().expect("buffer ref expected"); + if buffer.is_empty() { + return Ok(0); + } + + // buffer is not empty now. write it. + // # Safety + // We have already checked it is not None. + let buffer = unsafe { self.buffer.take().unwrap_unchecked() }; + let (result, buffer) = io.write_all(buffer).await; + self.buffer = Some(buffer); + match result { + Ok(written_len) => { + unsafe { self.buffer.as_mut().unwrap_unchecked().advance(written_len) }; + Ok(written_len) + } + Err(e) => { + let rerr = e.kind().into(); + self.status = WriteStatus::Err(e); + Err(rerr) + } + } + } + + /// moon patch (c10k P4b): drop the backing storage while drained; it + /// reallocates lazily on the next copy-in. Returns whether storage was + /// released. + pub fn release_if_empty(&mut self) -> bool { + self.buffer + .as_mut() + .expect("buffer mut expected") + .release_if_empty() + } +} + +impl io::Write for SafeWrite { + /// `write` to buffer. + /// # Handle return value + /// 1. Err(WouldBlock): the buffer is full, call do_io to flush it. + /// 2. _: handle it. + fn write(&mut self, buf: &[u8]) -> io::Result { + // if there is too much data inside the buffer, return WoundBlock + let buffer = self.buffer.as_mut().expect("buffer mut expected"); + // moon patch (c10k P4b): lazily (re)materialize released storage + // before the fullness check and copy-in. + buffer.ensure_allocated(); + match mem::replace(&mut self.status, WriteStatus::Ok) { + WriteStatus::Err(e) => return Err(e), + WriteStatus::Ok if buffer.is_full() => return Err(io::ErrorKind::WouldBlock.into()), + _ => (), + } + + // there is space inside the buffer, copy to it. + let to_copy = buf.len().min(buffer.available()); + unsafe { std::ptr::copy_nonoverlapping(buf.as_ptr(), buffer.write_ptr(), to_copy) }; + unsafe { buffer.set_init(to_copy) }; + Ok(to_copy) + } + + /// `flush` to buffer. + /// # Handle return value + /// 1. Err(WouldBlock): the buffer is full, call do_io to flush it. + /// 2. _: handle it. + fn flush(&mut self) -> io::Result<()> { + let buffer = self.buffer.as_mut().expect("buffer mut expected"); + match mem::replace(&mut self.status, WriteStatus::Ok) { + WriteStatus::Err(e) => Err(e), + WriteStatus::Ok if !buffer.is_empty() => Err(io::ErrorKind::WouldBlock.into()), + _ => Ok(()), + } + } +} + +// moon patch (c10k P4b): tests for lazy allocation, release, and the +// no-stash cancel-error contract. +#[cfg(test)] +mod moon_patch_tests { + use std::future::Future; + use std::io::{Read, Write}; + use std::pin::pin; + use std::task::{Context, Poll, Waker}; + + use monoio::buf::{IoBufMut, IoVecBufMut}; + use monoio::io::{AsyncReadRent, CancelHandle, CancelableAsyncReadRent, Canceller}; + use monoio::BufResult; + + use super::*; + + /// Drive an immediately-ready future to completion without a runtime. + fn now(fut: F) -> F::Output { + let mut fut = pin!(fut); + let waker = Waker::noop(); + let mut cx = Context::from_waker(waker); + match fut.as_mut().poll(&mut cx) { + Poll::Ready(v) => v, + Poll::Pending => panic!("mock future must be immediately ready"), + } + } + + /// Mock io: serves one scripted result per read call. + struct MockIo { + results: Vec, io::Error>>, + } + + impl AsyncReadRent for MockIo { + async fn read(&mut self, mut buf: T) -> BufResult { + match self.results.remove(0) { + Ok(data) => { + let n = data.len().min(buf.bytes_total()); + unsafe { + std::ptr::copy_nonoverlapping(data.as_ptr(), buf.write_ptr(), n); + buf.set_init(n); + } + (Ok(n), buf) + } + Err(e) => (Err(e), buf), + } + } + + async fn readv(&mut self, buf: T) -> BufResult { + (Ok(0), buf) + } + } + + impl CancelableAsyncReadRent for MockIo { + async fn cancelable_read( + &mut self, + buf: T, + _c: CancelHandle, + ) -> BufResult { + // The scripted results decide the outcome; a fired cancel is + // modelled by scripting Err(ECANCELED), matching what monoio's + // real drivers return for a cancelled op. + self.read(buf).await + } + + async fn cancelable_readv( + &mut self, + buf: T, + _c: CancelHandle, + ) -> BufResult { + (Ok(0), buf) + } + } + + fn allocated(read: &SafeRead) -> usize { + read.buffer.as_ref().unwrap().buf.len() + } + + #[test] + fn read_buffer_starts_released_and_allocates_on_do_io() { + let mut r = SafeRead::new(64); + assert_eq!(allocated(&r), 0, "must start lazy"); + let mut io = MockIo { + results: vec![Ok(b"hello".to_vec())], + }; + assert_eq!(now(r.do_io(&mut io)).unwrap(), 5); + assert_eq!(allocated(&r), 64, "do_io must materialize storage"); + let mut out = [0u8; 8]; + assert_eq!(r.read(&mut out).unwrap(), 5); + assert_eq!(&out[..5], b"hello"); + assert!(r.release_if_empty(), "drained buffer must release"); + assert_eq!(allocated(&r), 0); + // Released buffer reallocates transparently on the next do_io. + let mut io = MockIo { + results: vec![Ok(b"again".to_vec())], + }; + assert_eq!(now(r.do_io(&mut io)).unwrap(), 5); + assert_eq!(r.read(&mut out).unwrap(), 5); + assert_eq!(&out[..5], b"again"); + } + + #[test] + fn release_refuses_undrained_buffer() { + let mut r = SafeRead::new(64); + let mut io = MockIo { + results: vec![Ok(b"pending".to_vec())], + }; + now(r.do_io(&mut io)).unwrap(); + assert!(!r.release_if_empty(), "undrained data must be kept"); + assert_eq!(allocated(&r), 64); + } + + #[test] + fn write_buffer_lazy_alloc_and_release() { + let mut w = SafeWrite::new(64); + assert_eq!(w.buffer.as_ref().unwrap().buf.len(), 0, "must start lazy"); + assert_eq!(w.write(b"abc").unwrap(), 3); + assert_eq!(w.buffer.as_ref().unwrap().buf.len(), 64); + assert!(!w.release_if_empty(), "unflushed data must be kept"); + // Drain via advance (as write_all's do_io would). + w.buffer.as_mut().unwrap().advance(3); + assert!(w.release_if_empty()); + assert_eq!(w.buffer.as_ref().unwrap().buf.len(), 0); + // And it comes back on the next copy-in. + assert_eq!(w.write(b"xyz").unwrap(), 3); + assert_eq!(w.buffer.as_ref().unwrap().buf.len(), 64); + } + + #[test] + fn cancel_error_is_not_stashed_for_replay() { + let mut r = SafeRead::new(64); + let canceller = Canceller::new(); + let mut io = MockIo { + results: vec![Err(io::Error::from_raw_os_error(125))], + }; + let err = now(r.do_io_cancelable(&mut io, canceller.handle())).unwrap_err(); + assert_eq!(err.raw_os_error(), Some(125), "ECANCELED surfaces once"); + // The next sync read must see a clean WouldBlock, not a replayed + // cancel error (which would tear down a healthy connection). + let mut out = [0u8; 8]; + let replay = r.read(&mut out).unwrap_err(); + assert_eq!(replay.kind(), io::ErrorKind::WouldBlock); + // And a later cancelable read with a fresh handle works. + let mut io = MockIo { + results: vec![Ok(b"ok".to_vec())], + }; + assert_eq!(now(r.do_io_cancelable(&mut io, canceller.handle())).unwrap(), 2); + } + + #[test] + fn real_error_is_still_stashed_on_plain_do_io() { + let mut r = SafeRead::new(64); + let mut io = MockIo { + results: vec![Err(io::Error::new(io::ErrorKind::ConnectionReset, "boom"))], + }; + assert!(now(r.do_io(&mut io)).is_err()); + let mut out = [0u8; 8]; + let replay = r.read(&mut out).unwrap_err(); + assert_eq!(replay.kind(), io::ErrorKind::ConnectionReset); + } +} diff --git a/vendor/monoio-io-wrapper/src/unsafe_io.rs b/vendor/monoio-io-wrapper/src/unsafe_io.rs new file mode 100644 index 000000000..ed737ab16 --- /dev/null +++ b/vendor/monoio-io-wrapper/src/unsafe_io.rs @@ -0,0 +1,191 @@ +use std::{io, mem}; + +use monoio::{ + buf::{IoBuf, IoBufMut}, + io::{AsyncReadRent, AsyncWriteRent}, +}; + +/// Used by both UnsafeRead and UnsafeWrite. +#[derive(Debug)] +enum Status { + /// We haven't do real io, and maybe the dest is recorded. + WaitFill(Option<(*const u8, usize)>), + /// We have already do real io. The length maybe zero or non-zero. + Filled(Result), +} + +impl Default for Status { + fn default() -> Self { + Status::WaitFill(None) + } +} + +/// UnsafeRead is a wrapper of some meta data. +/// It implements std::io::Read trait. But it do real io in an async way. +/// On the first read, it may returns WouldBlock error, which means the +/// `fullfill` should be called to do real io. +/// The data is read directly into the dest that last std read passes. +/// Note that this action is an unsafe hack to avoid data copy. +/// You can only use this wrapper when you make sure the read dest is always +/// a valid buffer. +#[derive(Default, Debug)] +pub struct UnsafeRead { + status: Status, +} + +impl UnsafeRead { + pub const fn new() -> Self { + Self { + status: Status::WaitFill(None), + } + } + + /// `do_io` must be called after calling to io::Read::read. + /// # Handle return value + /// 1. Ok(n): previous read data length. + /// 2. Err(e): previous error. + /// 3. Err(WouldBlock): need calling read to capture ptr and len. + /// # Safety + /// User must make sure the former buffer is still valid until io finished. + pub async unsafe fn do_io(&mut self, mut io: IO) -> io::Result { + match self.status { + Status::WaitFill(Some((ptr, len))) => { + let buf = RawBuf { ptr, len }; + let read_ret = io.read(buf).await.0; + let rret: io::Result = match &read_ret { + Ok(n) => Ok(*n), + Err(e) => Err(e.kind().into()), + }; + self.status = Status::Filled(read_ret); + rret + } + Status::WaitFill(None) => Err(io::ErrorKind::WouldBlock.into()), + Status::Filled(ref prev_ret) => match prev_ret { + Ok(n) => Ok(*n), + Err(e) => Err(e.kind().into()), + }, + } + } + + /// Clear status. + pub fn reset(&mut self) { + self.status = Status::WaitFill(None); + } +} + +impl io::Read for UnsafeRead { + /// `read` from buffer(not really a buffer). + /// # Handle return value + /// 1. Err(WouldBlock): ptr and len captured, need call do_io and then retry read. + /// 2. _: handle it. + fn read(&mut self, buf: &mut [u8]) -> io::Result { + match mem::replace(&mut self.status, Status::WaitFill(None)) { + Status::WaitFill(_) => { + let ptr = buf.as_ptr(); + let len = buf.len(); + self.status = Status::WaitFill(Some((ptr, len))); + Err(io::ErrorKind::WouldBlock.into()) + } + Status::Filled(ret) => ret, + } + } +} + +/// UnsafeWrite behaves like `UnsafeRead`. +#[derive(Default, Debug)] +pub struct UnsafeWrite { + status: Status, +} + +impl UnsafeWrite { + pub const fn new() -> Self { + Self { + status: Status::WaitFill(None), + } + } + + /// `do_io` must be called after calling to io::Write::write. + /// # Handle return value + /// 1. Ok(n): previous written data length. + /// 2. Err(e): previous error. + /// 3. Err(WouldBlock): need calling write to capture ptr and len. + /// # Safety + /// User must make sure the former buffer is still valid until io finished. + pub async unsafe fn do_io(&mut self, mut io: IO) -> io::Result { + match self.status { + Status::WaitFill(Some((ptr, len))) => { + let buf = RawBuf { ptr, len }; + let write_ret = io.write(buf).await.0; + let rret: io::Result = match &write_ret { + Ok(n) => Ok(*n), + Err(e) => Err(e.kind().into()), + }; + self.status = Status::Filled(write_ret); + rret + } + Status::WaitFill(None) => Err(io::ErrorKind::WouldBlock.into()), + Status::Filled(ref prev_ret) => match prev_ret { + Ok(n) => Ok(*n), + Err(e) => Err(e.kind().into()), + }, + } + } + + /// Clear status. + pub fn reset(&mut self) { + self.status = Status::WaitFill(None); + } +} + +impl io::Write for UnsafeWrite { + /// `read` to buffer. + /// # Handle return value + /// 1. Err(WouldBlock): ptr and len captured, need call do_io and then retry write. + /// 2. _: handle it. + fn write(&mut self, buf: &[u8]) -> io::Result { + match mem::replace(&mut self.status, Status::WaitFill(None)) { + Status::WaitFill(_) => { + let ptr = buf.as_ptr(); + let len = buf.len(); + self.status = Status::WaitFill(Some((ptr, len))); + Err(io::ErrorKind::WouldBlock.into()) + } + Status::Filled(ret) => ret, + } + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } +} + +/// RawBuf is a wrapper of buffer ptr and len. +/// It seems that it hold the ownership of the buffer, which infact not. +/// Use this wrapper only when you can make sure the buffer ptr lives +/// longer than the wrapper. +struct RawBuf { + ptr: *const u8, + len: usize, +} + +unsafe impl IoBuf for RawBuf { + fn read_ptr(&self) -> *const u8 { + self.ptr + } + + fn bytes_init(&self) -> usize { + self.len + } +} + +unsafe impl IoBufMut for RawBuf { + fn write_ptr(&mut self) -> *mut u8 { + self.ptr as *mut u8 + } + + fn bytes_total(&mut self) -> usize { + self.len + } + + unsafe fn set_init(&mut self, _pos: usize) {} +} diff --git a/vendor/monoio-rustls/Cargo.toml b/vendor/monoio-rustls/Cargo.toml new file mode 100644 index 000000000..fe70a0543 --- /dev/null +++ b/vendor/monoio-rustls/Cargo.toml @@ -0,0 +1,74 @@ +# THIS FILE IS AUTOMATICALLY GENERATED BY CARGO +# +# When uploading crates to the registry Cargo will automatically +# "normalize" Cargo.toml files for maximal compatibility +# with all versions of Cargo and also rewrite `path` dependencies +# to registry (e.g., crates.io) dependencies. +# +# If you are reading this file be aware that the original Cargo.toml +# will likely look very different (and much more reasonable). +# See Cargo.toml.orig for the original contents. + +[package] +edition = "2021" +name = "monoio-rustls" +version = "0.4.0" +authors = [ + "ChiHai ", + "Rain Jiang ", +] +build = false +autobins = false +autoexamples = false +autotests = false +autobenches = false +description = "Asynchronous TLS streams wrapper for Monoio based on Rustls." +homepage = "https://github.com/monoio-rs/monoio-tls" +readme = "README.md" +categories = [ + "asynchronous", + "cryptography", + "network-programming", +] +license = "MIT/Apache-2.0" +repository = "https://github.com/monoio-rs/monoio-tls" + +[lib] +name = "monoio_rustls" +path = "src/lib.rs" + +[[example]] +name = "connect" +path = "examples/connect.rs" + +[dependencies.bytes] +version = "1" + +[dependencies.monoio] +version = "0.2.0" + +[dependencies.monoio-io-wrapper] +version = "0.1.1" + +[dependencies.rustls] +version = "~0.23.4" +features = ["std"] +default-features = false + +[dependencies.thiserror] +version = "1" + +[dev-dependencies.monoio] +version = "0.2.0" + +[dev-dependencies.webpki-roots] +version = "~0.26.1" + +[features] +default = [ + "logging", + "tls12", +] +logging = ["rustls/logging"] +tls12 = ["rustls/tls12"] +unsafe_io = ["monoio-io-wrapper/unsafe_io"] diff --git a/vendor/monoio-rustls/Cargo.toml.orig b/vendor/monoio-rustls/Cargo.toml.orig new file mode 100644 index 000000000..944a5104b --- /dev/null +++ b/vendor/monoio-rustls/Cargo.toml.orig @@ -0,0 +1,32 @@ +[package] +name = "monoio-rustls" +version = "0.4.0" + +authors = ["ChiHai ", "Rain Jiang "] +categories = ["asynchronous", "cryptography", "network-programming"] +description = "Asynchronous TLS streams wrapper for Monoio based on Rustls." +edition = "2021" +homepage = "https://github.com/monoio-rs/monoio-tls" +license = "MIT/Apache-2.0" +readme = "README.md" +repository = "https://github.com/monoio-rs/monoio-tls" + +[dependencies] +monoio = { workspace = true } +bytes = { workspace = true } +thiserror = { workspace = true } + +monoio-io-wrapper = { version = "0.1.1", path = "../monoio-io-wrapper" } +rustls = { version = "~0.23.4", default-features = false, features = ["std"] } + +[features] +default = ["logging", "tls12"] +logging = ["rustls/logging"] +tls12 = ["rustls/tls12"] +# Once unsafe_io is enabled, you may not drop the future before it returns ready. +# It saves one buffer copy than disabled. +unsafe_io = ["monoio-io-wrapper/unsafe_io"] + +[dev-dependencies] +monoio = { workspace = true } +webpki-roots = "~0.26.1" diff --git a/vendor/monoio-rustls/README.md b/vendor/monoio-rustls/README.md new file mode 100644 index 000000000..fc045d0db --- /dev/null +++ b/vendor/monoio-rustls/README.md @@ -0,0 +1 @@ +# Monoio-rustls diff --git a/vendor/monoio-rustls/examples/connect.rs b/vendor/monoio-rustls/examples/connect.rs new file mode 100644 index 000000000..114cdd04c --- /dev/null +++ b/vendor/monoio-rustls/examples/connect.rs @@ -0,0 +1,36 @@ +use std::sync::Arc; + +use monoio::{ + io::{AsyncReadRent, AsyncWriteRentExt}, + net::TcpStream, +}; +use monoio_rustls::TlsConnector; +use rustls::{RootCertStore}; + +#[monoio::main] +async fn main() { + let mut root_store = RootCertStore::empty(); + root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned()); + let config = rustls::ClientConfig::builder() + .with_root_certificates(root_store) + .with_no_client_auth(); + + let connector = TlsConnector::from(Arc::new(config)); + let stream = TcpStream::connect("rsproxy.cn:443").await.unwrap(); + println!("rsproxy.cn:443 connected"); + + let domain = rustls::pki_types::ServerName::try_from("rsproxy.cn").unwrap(); + let mut stream = connector.connect(domain, stream).await.unwrap(); + println!("handshake success"); + + let content = b"GET / HTTP/1.0\r\nHost: rsproxy.cn\r\n\r\n"; + let (r, _) = stream.write_all(content).await; + r.expect("unable to write http request"); + println!("http request sent"); + + let buf = vec![0_u8; 64]; + let (r, buf) = stream.read(buf).await; + r.expect("unable to read http response"); + let resp = String::from_utf8(buf).unwrap(); + println!("http response recv: \n\n{}", resp); +} diff --git a/vendor/monoio-rustls/src/client.rs b/vendor/monoio-rustls/src/client.rs new file mode 100644 index 000000000..b32e82119 --- /dev/null +++ b/vendor/monoio-rustls/src/client.rs @@ -0,0 +1,78 @@ +use std::sync::Arc; + +use monoio::io::{AsyncReadRent, AsyncWriteRent, OwnedReadHalf, OwnedWriteHalf}; +use rustls::{pki_types::ServerName, ClientConfig, ClientConnection}; + +use crate::{stream::Stream, TlsError}; + +/// A wrapper around an underlying raw stream which implements the TLS protocol. +pub type TlsStream = Stream; +/// TlsStream for read only. +pub type TlsStreamReadHalf = OwnedReadHalf>; +/// TlsStream for write only. +pub type TlsStreamWriteHalf = OwnedWriteHalf>; + +/// A wrapper around a `rustls::ClientConfig`, providing an async `connect` method. +#[derive(Clone)] +pub struct TlsConnector { + inner: Arc, + #[cfg(feature = "unsafe_io")] + unsafe_io: bool, +} + +impl From> for TlsConnector { + fn from(inner: Arc) -> TlsConnector { + TlsConnector { + inner, + #[cfg(feature = "unsafe_io")] + unsafe_io: false, + } + } +} + +impl From for TlsConnector { + fn from(inner: ClientConfig) -> TlsConnector { + TlsConnector { + inner: Arc::new(inner), + #[cfg(feature = "unsafe_io")] + unsafe_io: false, + } + } +} + +impl TlsConnector { + /// Enable unsafe-io. + /// # Safety + /// Users must make sure the buffer ptr and len is valid until io finished. + /// So the Future cannot be dropped directly. Consider using CancellableIO. + #[cfg(feature = "unsafe_io")] + pub unsafe fn unsafe_io(self, enabled: bool) -> Self { + Self { + unsafe_io: enabled, + ..self + } + } + + pub async fn connect( + &self, + domain: ServerName<'static>, + stream: IO, + ) -> Result, TlsError> + where + IO: AsyncReadRent + AsyncWriteRent, + { + let session = ClientConnection::new(self.inner.clone(), domain)?; + #[cfg(feature = "unsafe_io")] + let mut stream = if self.unsafe_io { + // # Safety + // Users already maked unsafe io. + unsafe { Stream::new_unsafe(stream, session) } + } else { + Stream::new(stream, session) + }; + #[cfg(not(feature = "unsafe_io"))] + let mut stream = Stream::new(stream, session); + stream.handshake().await?; + Ok(stream) + } +} diff --git a/vendor/monoio-rustls/src/error.rs b/vendor/monoio-rustls/src/error.rs new file mode 100644 index 000000000..0c4f8a60e --- /dev/null +++ b/vendor/monoio-rustls/src/error.rs @@ -0,0 +1,20 @@ +use std::io; + +use thiserror::Error; + +#[derive(Error, Debug)] +pub enum TlsError { + #[error("io error")] + Io(#[from] std::io::Error), + #[error("rustls error")] + Rustls(#[from] rustls::Error), +} + +impl From for io::Error { + fn from(e: TlsError) -> Self { + match e { + TlsError::Io(e) => e, + TlsError::Rustls(e) => io::Error::new(io::ErrorKind::Other, e), + } + } +} diff --git a/vendor/monoio-rustls/src/lib.rs b/vendor/monoio-rustls/src/lib.rs new file mode 100644 index 000000000..0524338f4 --- /dev/null +++ b/vendor/monoio-rustls/src/lib.rs @@ -0,0 +1,31 @@ +#![allow(stable_features)] + +mod client; +mod error; +mod server; +mod stream; + +pub use client::{ + TlsConnector, TlsStream as ClientTlsStream, TlsStreamReadHalf as ClientTlsStreamReadHalf, + TlsStreamWriteHalf as ClientTlsStreamWriteHalf, +}; +pub use error::TlsError; +pub use server::{ + TlsAcceptor, TlsStream as ServerTlsStream, TlsStreamReadHalf as ServerTlsStreamReadHalf, + TlsStreamWriteHalf as ServerTlsStreamWriteHalf, +}; + +/// A wrapper around an underlying raw stream which implements the TLS protocol. +pub type TlsStream = stream::Stream; + +impl From> for TlsStream { + fn from(value: ClientTlsStream) -> Self { + value.map_conn(|c| c.into()) + } +} + +impl From> for TlsStream { + fn from(value: ServerTlsStream) -> Self { + value.map_conn(|c| c.into()) + } +} diff --git a/vendor/monoio-rustls/src/server.rs b/vendor/monoio-rustls/src/server.rs new file mode 100644 index 000000000..0d0ea2909 --- /dev/null +++ b/vendor/monoio-rustls/src/server.rs @@ -0,0 +1,74 @@ +use std::sync::Arc; + +use monoio::io::{AsyncReadRent, AsyncWriteRent, OwnedReadHalf, OwnedWriteHalf}; +use rustls::{ServerConfig, ServerConnection}; + +use crate::{stream::Stream, TlsError}; + +/// A wrapper around an underlying raw stream which implements the TLS protocol. +pub type TlsStream = Stream; +/// TlsStream for read only. +pub type TlsStreamReadHalf = OwnedReadHalf>; +/// TlsStream for write only. +pub type TlsStreamWriteHalf = OwnedWriteHalf>; + +/// A wrapper around a `rustls::ServerConfig`, providing an async `accept` method. +#[derive(Clone)] +pub struct TlsAcceptor { + inner: Arc, + #[cfg(feature = "unsafe_io")] + unsafe_io: bool, +} + +impl From> for TlsAcceptor { + fn from(inner: Arc) -> TlsAcceptor { + TlsAcceptor { + inner, + #[cfg(feature = "unsafe_io")] + unsafe_io: false, + } + } +} + +impl From for TlsAcceptor { + fn from(inner: ServerConfig) -> TlsAcceptor { + TlsAcceptor { + inner: Arc::new(inner), + #[cfg(feature = "unsafe_io")] + unsafe_io: false, + } + } +} + +impl TlsAcceptor { + /// Enable unsafe-io. + /// # Safety + /// Users must make sure the buffer ptr and len is valid until io finished. + /// So the Future cannot be dropped directly. Consider using CancellableIO. + #[cfg(feature = "unsafe_io")] + pub unsafe fn unsafe_io(self, enabled: bool) -> Self { + Self { + unsafe_io: enabled, + ..self + } + } + + pub async fn accept(&self, stream: IO) -> Result, TlsError> + where + IO: AsyncReadRent + AsyncWriteRent, + { + let session = ServerConnection::new(self.inner.clone())?; + #[cfg(feature = "unsafe_io")] + let mut stream = if self.unsafe_io { + // # Safety + // Users already maked unsafe io. + unsafe { Stream::new_unsafe(stream, session) } + } else { + Stream::new(stream, session) + }; + #[cfg(not(feature = "unsafe_io"))] + let mut stream = Stream::new(stream, session); + stream.handshake().await?; + Ok(stream) + } +} diff --git a/vendor/monoio-rustls/src/stream.rs b/vendor/monoio-rustls/src/stream.rs new file mode 100644 index 000000000..978c410c8 --- /dev/null +++ b/vendor/monoio-rustls/src/stream.rs @@ -0,0 +1,383 @@ +use std::{ + io::{self, Read, Write}, + ops::{Deref, DerefMut}, +}; + +use monoio::{ + buf::{IoBuf, IoBufMut, IoVecBuf, IoVecBufMut, RawBuf}, + io::{AsyncReadRent, AsyncWriteRent, CancelHandle, CancelableAsyncReadRent, Split}, + BufResult, +}; +use monoio_io_wrapper::{ReadBuffer, WriteBuffer}; +use rustls::{ClientConnection, ConnectionCommon, ServerConnection, SideData}; + +#[derive(Debug)] +pub struct Stream { + pub(crate) io: IO, + pub(crate) session: C, + r_buffer: ReadBuffer, + w_buffer: WriteBuffer, +} + +impl Stream { + #[inline] + pub fn alpn_protocol(&self) -> Option> { + self.session.alpn_protocol().map(|s| s.to_vec()) + } +} + +impl Stream { + #[inline] + pub fn alpn_protocol(&self) -> Option> { + self.session.alpn_protocol().map(|s| s.to_vec()) + } +} + +unsafe impl Split for Stream {} + +impl Stream { + pub fn new(io: IO, session: C) -> Self { + Self { + io, + session, + r_buffer: Default::default(), + w_buffer: Default::default(), + } + } + + /// Enable unsafe-io. + /// # Safety + /// Users must make sure the buffer ptr and len is valid until io finished. + /// So the Future cannot be dropped directly. Consider using CancellableIO. + #[cfg(feature = "unsafe_io")] + pub unsafe fn new_unsafe(io: IO, session: C) -> Self { + Self { + io, + session, + r_buffer: ReadBuffer::new_unsafe(), + w_buffer: WriteBuffer::new_unsafe(), + } + } + + pub fn into_parts(self) -> (IO, C) { + (self.io, self.session) + } + + /// moon patch (c10k P4b): release both wrapper buffers if drained. + /// Called by moon's idle-park downshift so a parked TLS connection + /// sheds its 2×16 KiB wrapper allocations; they reallocate lazily on + /// the next I/O. Returns true if anything was released. + pub fn release_idle_buffers(&mut self) -> bool { + let r = self.r_buffer.release_if_empty(); + let w = self.w_buffer.release_if_empty(); + r || w + } + + pub(crate) fn map_conn C2>(self, f: F) -> Stream { + Stream { + io: self.io, + session: f(self.session), + r_buffer: self.r_buffer, + w_buffer: self.w_buffer, + } + } +} + +impl Stream +where + C: DerefMut + Deref>, +{ + pub(crate) async fn read_io(&mut self, splitted: bool) -> io::Result { + let n = loop { + match self.session.read_tls(&mut self.r_buffer) { + Ok(n) => { + break n; + } + Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => { + #[allow(unused_unsafe)] + unsafe { + self.r_buffer.do_io(&mut self.io).await? + }; + continue; + } + Err(err) => return Err(err), + } + }; + + let state = match self.session.process_new_packets() { + Ok(state) => state, + Err(err) => { + // When to write_io? If we do this in read call, the UnsafeWrite may crash + // when we impl split in an UnsafeCell way. + // Here we choose not to do write when read. + // User should manually shutdown it on error. + if !splitted { + let _ = self.write_io().await; + } + return Err(io::Error::new(io::ErrorKind::InvalidData, err)); + } + }; + + if state.peer_has_closed() && self.session.is_handshaking() { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "tls handshake alert", + )); + } + + Ok(n) + } + + pub(crate) async fn write_io(&mut self) -> io::Result { + let n = loop { + match self.session.write_tls(&mut self.w_buffer) { + Ok(n) => { + if self.w_buffer.is_safe() { + self.w_buffer.do_io(&mut self.io).await?; + } + break n; + } + Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => { + // here we don't have to check WouldBlock since we already captured the + // mem block info under unsafe-io. + #[allow(unused_unsafe)] + unsafe { + self.w_buffer.do_io(&mut self.io).await? + }; + continue; + } + Err(err) => return Err(err), + } + }; + + Ok(n) + } + + pub(crate) async fn handshake(&mut self) -> io::Result<(usize, usize)> { + let mut wrlen = 0; + let mut rdlen = 0; + let mut eof = false; + + loop { + while self.session.wants_write() && self.session.is_handshaking() { + wrlen += self.write_io().await?; + } + while !eof && self.session.wants_read() && self.session.is_handshaking() { + let n = self.read_io(false).await?; + rdlen += n; + if n == 0 { + eof = true; + } + } + + match (eof, self.session.is_handshaking()) { + (true, true) => { + let err = io::Error::new(io::ErrorKind::UnexpectedEof, "tls handshake eof"); + return Err(err); + } + (false, true) => (), + (_, false) => { + break; + } + }; + } + + // flush buffer + while self.session.wants_write() { + wrlen += self.write_io().await?; + } + + Ok((rdlen, wrlen)) + } + + /// moon patch (c10k P4b): cancelable twin of `read_io`. A cancel fired + /// mid-read surfaces here as ECANCELED WITHOUT being stashed in the + /// wrapper-buffer status (see `SafeRead::do_io_cancelable`), so the next + /// plain read resumes the TLS session cleanly. Unsplit streams only. + pub(crate) async fn read_io_cancelable(&mut self, c: CancelHandle) -> io::Result + where + IO: CancelableAsyncReadRent, + { + let n = loop { + match self.session.read_tls(&mut self.r_buffer) { + Ok(n) => { + break n; + } + Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => { + self.r_buffer.do_io_cancelable(&mut self.io, c.clone()).await?; + continue; + } + Err(err) => return Err(err), + } + }; + + let state = match self.session.process_new_packets() { + Ok(state) => state, + Err(err) => { + let _ = self.write_io().await; + return Err(io::Error::new(io::ErrorKind::InvalidData, err)); + } + }; + + if state.peer_has_closed() && self.session.is_handshaking() { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "tls handshake alert", + )); + } + + Ok(n) + } + + pub(crate) async fn read_inner( + &mut self, + mut buf: T, + splitted: bool, + ) -> BufResult { + let slice = unsafe { std::slice::from_raw_parts_mut(buf.write_ptr(), buf.bytes_total()) }; + loop { + // read from rustls to buffer + match self.session.reader().read(slice) { + Ok(n) => { + unsafe { buf.set_init(n) }; + return (Ok(n), buf); + } + // we need more data, read something. + Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => (), + Err(e) => { + return (Err(e), buf); + } + } + + // now we need data, read something into rustls + if let Err(e) = self.read_io(splitted).await { + return (Err(e), buf); + } + } + } +} + +impl AsyncReadRent for Stream +where + C: DerefMut + Deref>, +{ + async fn read(&mut self, buf: T) -> BufResult { + self.read_inner(buf, false).await + } + + async fn readv(&mut self, mut buf: T) -> BufResult { + let n = match unsafe { RawBuf::new_from_iovec_mut(&mut buf) } { + Some(raw_buf) => self.read(raw_buf).await.0, + None => Ok(0), + }; + if let Ok(n) = n { + unsafe { buf.set_init(n) }; + } + (n, buf) + } +} + +// moon patch (c10k P4b): cancelable read so moon's W11 idle-park sweep can +// cancel a TLS read parked on the underlying socket. Loss-free under +// cancel-and-await: bytes from a completion that raced the cancel persist in +// the wrapper buffer / rustls deframer and are consumed by the next read. +impl CancelableAsyncReadRent for Stream +where + IO: AsyncReadRent + AsyncWriteRent + CancelableAsyncReadRent, + C: DerefMut + Deref>, +{ + async fn cancelable_read(&mut self, mut buf: T, c: CancelHandle) -> BufResult { + let slice = unsafe { std::slice::from_raw_parts_mut(buf.write_ptr(), buf.bytes_total()) }; + loop { + // read from rustls to buffer + match self.session.reader().read(slice) { + Ok(n) => { + unsafe { buf.set_init(n) }; + return (Ok(n), buf); + } + // we need more data, read something. + Err(ref err) if err.kind() == io::ErrorKind::WouldBlock => (), + Err(e) => { + return (Err(e), buf); + } + } + + // now we need data, read something into rustls + if let Err(e) = self.read_io_cancelable(c.clone()).await { + return (Err(e), buf); + } + } + } + + async fn cancelable_readv(&mut self, mut buf: T, c: CancelHandle) -> BufResult { + let n = match unsafe { RawBuf::new_from_iovec_mut(&mut buf) } { + Some(raw_buf) => self.cancelable_read(raw_buf, c).await.0, + None => Ok(0), + }; + if let Ok(n) = n { + unsafe { buf.set_init(n) }; + } + (n, buf) + } +} + +impl AsyncWriteRent for Stream +where + C: DerefMut + Deref>, +{ + async fn write(&mut self, buf: T) -> BufResult { + // construct slice + let slice = unsafe { std::slice::from_raw_parts(buf.read_ptr(), buf.bytes_init()) }; + + // flush rustls inner write buffer to make sure there is space for new data + if self.session.wants_write() { + if let Err(e) = self.write_io().await { + return (Err(e), buf); + } + } + + // write slice to rustls + let n = match self.session.writer().write(slice) { + Ok(n) => n, + Err(e) => return (Err(e), buf), + }; + + // write from rustls to connection + while self.session.wants_write() { + match self.write_io().await { + Ok(0) => { + break; + } + Ok(_) => (), + Err(e) => return (Err(e), buf), + } + } + (Ok(n), buf) + } + + // TODO: use real writev + async fn writev(&mut self, buf_vec: T) -> BufResult { + let n = match unsafe { RawBuf::new_from_iovec(&buf_vec) } { + Some(raw_buf) => self.write(raw_buf).await.0, + None => Ok(0), + }; + (n, buf_vec) + } + + async fn flush(&mut self) -> io::Result<()> { + self.session.writer().flush()?; + while self.session.wants_write() { + self.write_io().await?; + } + self.io.flush().await + } + + async fn shutdown(&mut self) -> io::Result<()> { + self.session.send_close_notify(); + + while self.session.wants_write() { + self.write_io().await?; + } + self.io.shutdown().await + } +}