diff --git a/examples/concurrent_pread_bench.rs b/examples/concurrent_pread_bench.rs index 8b52fe3..4502af2 100644 --- a/examples/concurrent_pread_bench.rs +++ b/examples/concurrent_pread_bench.rs @@ -110,14 +110,16 @@ struct Config { read_size: usize, concurrency: usize, total_ops: usize, + shards: usize, verify: bool, } fn parse_args() -> Result { let a: Vec = std::env::args().collect(); - if a.len() != 7 { + if a.len() != 7 && a.len() != 8 { return Err( - "usage: concurrent_pread_bench ".to_string(), + "usage: concurrent_pread_bench [shards]" + .to_string(), ); } let parse = |name: &str, raw: &str| -> Result { @@ -131,6 +133,7 @@ fn parse_args() -> Result { read_size: parse("read_size", &a[4])? as usize, concurrency: parse("concurrency", &a[5])? as usize, total_ops: parse("total_ops", &a[6])? as usize, + shards: if a.len() == 8 { parse("shards", &a[7])? as usize } else { 1 }, verify: matches!(std::env::var("BENCH_VERIFY").as_deref(), Ok("1") | Ok("true")), }; validate(&cfg)?; @@ -291,7 +294,7 @@ async fn run(cfg: &Config) -> Result, String> { let driver = if cfg.strategy.uses_uring() { let depth = (cfg.concurrency.max(64) * 2).next_power_of_two() as u32; Some(Arc::new( - UringDriver::probe_and_start(depth).map_err(|e| format!("probe io_uring: {e:?}"))?, + UringDriver::probe_and_start_sharded(depth, cfg.shards).map_err(|e| format!("probe io_uring: {e:?}"))?, )) } else { None @@ -416,10 +419,11 @@ fn main() -> ExitCode { let ops = lats.len(); let iops = ops as f64 / secs; let mbps = (ops as f64 * cfg.read_size as f64 / (1024.0 * 1024.0)) / secs; - // CSV: strategy,file_size,read_size,concurrency,ops,secs,IOPS,MBps,p50_us,p99_us,p999_us + // CSV: strategy,shards,file_size,read_size,concurrency,ops,secs,IOPS,MBps,p50_us,p99_us,p999_us println!( - "{},{},{},{},{},{:.6},{:.0},{:.1},{},{},{}", + "{},{},{},{},{},{},{:.6},{:.0},{:.1},{},{},{}", cfg.strategy.name(), + cfg.shards, cfg.file_size, cfg.read_size, cfg.concurrency, diff --git a/src/driver.rs b/src/driver.rs index 25ad9ea..8cb1dca 100644 --- a/src/driver.rs +++ b/src/driver.rs @@ -20,7 +20,7 @@ use std::os::fd::{AsRawFd, FromRawFd}; use std::os::unix::ffi::OsStrExt; use std::pin::Pin; use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::mpsc::{self, TryRecvError}; use std::task::{Context, Poll}; use std::thread::JoinHandle; @@ -240,7 +240,7 @@ struct DriverStats { } /// Point-in-time copy of the driver counters. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub struct StatsSnapshot { /// Read ops handed to the kernel. pub submitted: u64, @@ -476,26 +476,109 @@ impl Drop for ReadHandle { } /// Process-level io_uring driver: one ring, one driver thread. -pub struct UringDriver { +/// One io_uring ring plus the thread that drives it. +/// +/// Every cancel-safety invariant holds *per shard*, exactly as it did when a +/// driver owned a single ring: this shard's pending table owns its buffers and +/// fds until their CQEs, its permits are released only when a pending entry is +/// dropped, and its bounded drain is what shutdown joins on. A `ReadHandle` +/// carries the `tx` and `wake` of the shard that accepted it, so a cancel or a +/// deferred submission always routes back to that same shard. +struct Shard { tx: mpsc::Sender, handle: Option>, stats: Arc, - next_id: AtomicU64, - /// Backpressure permits (one per allowed in-flight op). Closed when the - /// driver exits so any waiting `ReadHandle` resolves with a driver-gone - /// error instead of hanging (rustfs/backlog#1102). + /// Backpressure permits (one per allowed in-flight op on this ring). Closed + /// when the driver thread exits so any waiting `ReadHandle` resolves with a + /// driver-gone error instead of hanging (rustfs/backlog#1102). sem: Arc, - /// Signaled after every message send so the driver loop wakes immediately + /// Signaled after every message send so this shard's loop wakes immediately /// instead of waiting out the heartbeat (backlog#1102). wake_efd: Arc, } +impl Shard { + /// Ask the shard's thread to drain and exit, then join it. Idempotent: the + /// `JoinHandle` is taken, so a later `Drop` is a no-op. + fn join(&mut self) { + if let Some(h) = self.handle.take() { + let _ = self.tx.send(Msg::Shutdown); + self.wake_efd.signal(); + let _ = h.join(); + } + } +} + +impl Drop for Shard { + fn drop(&mut self) { + self.join(); + } +} + +pub struct UringDriver { + /// One or more independent rings. A cache-hit buffered read completes inline + /// inside `io_uring_enter`, so the thread driving a ring performs that + /// read's memcpy — which caps a single-ring driver at one core's memory + /// bandwidth (~5 GB/s measured, rustfs/backlog#1145). Sharding lifts that + /// ceiling roughly linearly while keeping the ring set per-disk, so a stalled + /// disk still cannot starve another disk's rings (rustfs/backlog#1055). + shards: Vec, + next_id: AtomicU64, + /// Round-robin cursor for shard selection. Relaxed: it only has to spread + /// ops, never to order them. + rr: AtomicUsize, +} + impl UringDriver { /// Create the ring AND verify a real `IORING_OP_READ` round-trip on a /// temp file before accepting work. `io_uring_setup` succeeding is not /// enough: gVisor/seccomp environments can create a ring whose ops then /// fail with ENOSYS/EINVAL (backlog#894 probe design). + /// Start a single-ring driver. Identical to `probe_and_start_sharded(entries, 1)`. pub fn probe_and_start(entries: u32) -> Result { + Self::probe_and_start_sharded(entries, 1) + } + + /// Start a driver backed by `shards` independent rings, each with `entries` + /// SQ slots and its own driver thread. + /// + /// Use more than one shard when the workload hits the page cache: such reads + /// complete inline in `io_uring_enter`, so a single driver thread performs + /// every one of their memcpys and caps the driver at one core's memory + /// bandwidth. Measured on a 16-core host (rustfs/backlog#1145): 1 ring → + /// 4890 MB/s, 2 → 8969 MB/s, 4 → 15806 MB/s, with per-ring throughput flat. + /// Reads that miss the cache are device-bound and do not need sharding. + /// + /// In-flight ops are capped at `entries` *per shard* (the invariant that + /// makes CQ overflow structurally unreachable holds per ring), so the whole + /// driver admits up to `shards * entries` concurrent reads. + /// + /// `shards` is clamped to at least 1. Probing happens on the first shard, so + /// a restricted environment fails exactly as it does for a single ring; if a + /// later shard fails to start, the ones already running are shut down and + /// joined before the error is returned. + pub fn probe_and_start_sharded(entries: u32, shards: usize) -> Result { + let mut started = Vec::with_capacity(shards.max(1)); + for _ in 0..shards.max(1) { + // `?` drops `started`, whose `Shard::drop` joins each running thread. + started.push(Self::start_shard(entries)?); + } + Ok(Self { + shards: started, + next_id: AtomicU64::new(1), + rr: AtomicUsize::new(0), + }) + } + + /// Pick the shard for the next op. Round-robin spreads the inline-completion + /// memcpy across driver threads; correctness does not depend on the choice, + /// because the handle remembers which shard took the op. + fn shard(&self) -> &Shard { + let n = self.shards.len(); + &self.shards[self.rr.fetch_add(1, Ordering::Relaxed) % n] + } + + fn start_shard(entries: u32) -> Result { let mut ring = IoUring::new(entries).map_err(ProbeFailure::Setup)?; // Require the NODROP feature (kernel >= 5.5). Without it, CQ overflow // silently drops CQEs, stranding pending entries forever and hanging @@ -530,11 +613,10 @@ impl UringDriver { .spawn(move || drive(ring, rx, thread_stats, thread_sem, cq_efd, thread_wake)) .expect("spawn driver thread"); - Ok(Self { + Ok(Shard { tx, handle: Some(handle), stats, - next_id: AtomicU64::new(1), sem, wake_efd, }) @@ -572,6 +654,13 @@ impl UringDriver { assert_eq!(id & CANCEL_BIT, 0, "op id overflowed into the cancel bit"); let (done, rx) = oneshot::channel(); + // Bind the op to one shard for its whole life: the permit, the message, + // the wake, and any later cancel all go to this ring. The handle holds + // clones of that shard's `tx`/`wake_efd`, so nothing can route a cancel + // to a ring whose pending table does not hold the op. The rejection paths + // below return an `Inert` handle that never sends, but still need a `tx`. + let shard = self.shard(); + // Reject an offset the kernel would answer with a runtime EINVAL that // must NOT be mistaken for an environment restriction (C7, // rustfs/backlog#1059). The kernel reads `off` as a signed loff_t, so @@ -587,7 +676,7 @@ impl UringDriver { return ReadHandle { id, rx, - tx: self.tx.clone(), + tx: shard.tx.clone(), finished: false, cancel_on_drop: false, state: HandleState::Inert, @@ -608,7 +697,7 @@ impl UringDriver { return ReadHandle { id, rx, - tx: self.tx.clone(), + tx: shard.tx.clone(), finished: false, cancel_on_drop: false, state: HandleState::Inert, @@ -628,7 +717,7 @@ impl UringDriver { return ReadHandle { id, rx, - tx: self.tx.clone(), + tx: shard.tx.clone(), finished: false, cancel_on_drop: false, state: HandleState::Inert, @@ -640,12 +729,12 @@ impl UringDriver { // released only when the pending entry is dropped at the CQE (C10, // rustfs/backlog#1060). Acquisition never blocks the caller's thread // (rustfs/backlog#1102). - match Arc::clone(&self.sem).try_acquire_owned() { + match Arc::clone(&shard.sem).try_acquire_owned() { // Fast path: a permit was free, so submit eagerly — no allocation, // no await, and the op is in flight the moment `submit` returns, // exactly as with the previous blocking implementation. Ok(permit) => { - if let Err(mpsc::SendError(msg)) = self.tx.send(Msg::Read { + if let Err(mpsc::SendError(msg)) = shard.tx.send(Msg::Read { id, file, offset, @@ -665,18 +754,18 @@ impl UringDriver { return ReadHandle { id, rx, - tx: self.tx.clone(), + tx: shard.tx.clone(), finished: false, cancel_on_drop: false, state: HandleState::Inert, }; } // Wake the driver loop so the read starts immediately. - self.wake_efd.signal(); + shard.wake_efd.signal(); ReadHandle { id, rx, - tx: self.tx.clone(), + tx: shard.tx.clone(), finished: false, cancel_on_drop: true, state: HandleState::Submitted, @@ -688,17 +777,17 @@ impl UringDriver { Err(TryAcquireError::NoPermits) => ReadHandle { id, rx, - tx: self.tx.clone(), + tx: shard.tx.clone(), finished: false, cancel_on_drop: true, state: HandleState::WaitingPermit { - acquire: Box::pin(Arc::clone(&self.sem).acquire_owned()), + acquire: Box::pin(Arc::clone(&shard.sem).acquire_owned()), file, offset, len, align, done, - wake: Arc::clone(&self.wake_efd), + wake: Arc::clone(&shard.wake_efd), }, }, // The driver has exited and closed the semaphore. @@ -707,7 +796,7 @@ impl UringDriver { ReadHandle { id, rx, - tx: self.tx.clone(), + tx: shard.tx.clone(), finished: false, cancel_on_drop: false, state: HandleState::Inert, @@ -716,34 +805,46 @@ impl UringDriver { } } + /// Counters summed across every shard. The conservation identities the + /// cancel-safety tests assert (`submitted == delivered + orphan_reclaimed`, + /// `in_flight == 0` after a clean drain) hold per shard, so they hold for + /// the sum. pub fn stats(&self) -> StatsSnapshot { - StatsSnapshot { - submitted: self.stats.submitted.load(Ordering::SeqCst), - delivered: self.stats.delivered.load(Ordering::SeqCst), - orphan_reclaimed: self.stats.orphan_reclaimed.load(Ordering::SeqCst), - in_flight: self.stats.in_flight.load(Ordering::SeqCst), - cancel_succeeded: self.stats.cancel_succeeded.load(Ordering::SeqCst), - cancel_not_found: self.stats.cancel_not_found.load(Ordering::SeqCst), - cancel_already: self.stats.cancel_already.load(Ordering::SeqCst), - cq_overflow: self.stats.cq_overflow.load(Ordering::SeqCst), + let mut snap = StatsSnapshot::default(); + for shard in &self.shards { + let s = &shard.stats; + snap.submitted += s.submitted.load(Ordering::SeqCst); + snap.delivered += s.delivered.load(Ordering::SeqCst); + snap.orphan_reclaimed += s.orphan_reclaimed.load(Ordering::SeqCst); + snap.in_flight += s.in_flight.load(Ordering::SeqCst); + snap.cancel_succeeded += s.cancel_succeeded.load(Ordering::SeqCst); + snap.cancel_not_found += s.cancel_not_found.load(Ordering::SeqCst); + snap.cancel_already += s.cancel_already.load(Ordering::SeqCst); + snap.cq_overflow += s.cq_overflow.load(Ordering::SeqCst); } + snap } - /// Stop accepting work, cancel all in-flight ops, drain the ring to - /// `in_flight == 0`, then join the driver thread. Only after that is the - /// ring dropped/unmapped — the shutdown ordering P2 requires. + /// Stop accepting work, cancel all in-flight ops, drain every ring to + /// `in_flight == 0`, then join each driver thread. Only after that is a ring + /// dropped/unmapped — the shutdown ordering P2 requires, per shard. + /// + /// Shards are asked to stop first and joined afterwards, so their bounded + /// drains overlap instead of serializing `shards * DRAIN_TIMEOUT`. pub fn shutdown(mut self) -> StatsSnapshot { - let _ = self.tx.send(Msg::Shutdown); - self.wake_efd.signal(); - if let Some(h) = self.handle.take() { - let _ = h.join(); + for shard in &self.shards { + let _ = shard.tx.send(Msg::Shutdown); + shard.wake_efd.signal(); + } + for shard in &mut self.shards { + shard.join(); } let snap = self.stats(); - // A clean drain leaves in_flight == 0. A non-zero count here means the - // bounded drain bailed out on a hung device and leaked the ring+buffers - // to stay memory-safe (C4, rustfs/backlog#1055) — a degraded but safe - // outcome, not a panic. Callers/tests that require a clean drain assert - // on the returned snapshot themselves. + // A clean drain leaves in_flight == 0. A non-zero count here means some + // shard's bounded drain bailed out on a hung device and leaked its + // ring+buffers to stay memory-safe (C4, rustfs/backlog#1055) — a degraded + // but safe outcome, not a panic. Callers/tests that require a clean drain + // assert on the returned snapshot themselves. if snap.in_flight != 0 { eprintln!( "uring-spike shutdown: {} ops still in flight (bounded-drain bailout on a hung device)", @@ -756,10 +857,16 @@ impl UringDriver { impl Drop for UringDriver { fn drop(&mut self) { - if let Some(h) = self.handle.take() { - let _ = self.tx.send(Msg::Shutdown); - self.wake_efd.signal(); - let _ = h.join(); + // Ask every shard to stop before joining any of them, so their bounded + // drains overlap. Dropping the `Vec` would instead run each + // `Shard::drop` in turn, serializing up to `shards * DRAIN_TIMEOUT` on a + // hung device. `Shard::join` is idempotent, so the later drops are no-ops. + for shard in &self.shards { + let _ = shard.tx.send(Msg::Shutdown); + shard.wake_efd.signal(); + } + for shard in &mut self.shards { + shard.join(); } } } diff --git a/tests/cancel.rs b/tests/cancel.rs index 75fe156..5ff104b 100644 --- a/tests/cancel.rs +++ b/tests/cancel.rs @@ -45,6 +45,21 @@ fn driver_or_skip(name: &str) -> Option { } } +/// Like `driver_or_skip` but with `shards` independent rings (backlog#1145). +fn sharded_driver_or_skip(name: &str, shards: usize) -> Option { + match UringDriver::probe_and_start_sharded(64, shards) { + Ok(d) => Some(d), + Err(e) => { + assert!( + e.is_expected_restriction(), + "probe failed OUTSIDE the expected restriction errno class: {e:?}" + ); + eprintln!("SKIP {name}: restricted environment, graceful degradation path taken ({e:?})"); + None + } + } +} + /// Deterministic pseudo-random content so reads are verifiable. fn make_content(len: usize) -> Vec { let mut state: u64 = 0x2545F4914F6CDD1D; @@ -627,3 +642,73 @@ async fn drop_without_shutdown_drains_and_cancels() { } drop(pipe_write); } + +/// Sharding the driver's rings (backlog#1145) must not weaken the cancel-safety +/// model: every op the driver accepted is still delivered to a live caller or +/// orphan-reclaimed at its CQE, and the aggregate counters conserve. Ops are +/// round-robined across shards, so a cancel routed to the wrong ring — or a +/// pending table that lost an entry across the split — shows up here as a +/// conservation violation or a content mismatch. +/// +/// Deliberately mixes dropped and kept handles, exactly as `cancel_stress` does, +/// but across 4 rings so the drop lands on a shard other than the one the next +/// submit picks. +#[tokio::test(flavor = "multi_thread")] +async fn sharded_driver_conserves_buffers_across_shards() { + const SHARDS: usize = 4; + let Some(driver) = sharded_driver_or_skip("sharded_driver_conserves_buffers_across_shards", SHARDS) else { + return; + }; + const LEN: usize = 8 << 20; + const OPS: usize = 256; + const READ_LEN: usize = 64 << 10; + let content = make_content(LEN); + let (path, file) = temp_file_with(&content, "sharded-stress"); + + let mut kept = Vec::new(); + for i in 0..OPS { + let offset = (i * 97_611) % (LEN - READ_LEN); + let handle = driver.read_at(Arc::clone(&file), offset as u64, READ_LEN); + if i % 2 == 0 { + drop(handle); + } else { + kept.push((offset, handle)); + } + } + + // Correctness across shards: a kept read must return its own bytes, not a + // neighbouring shard's. + for (offset, handle) in kept { + let got = handle.await.expect("kept read failed"); + assert_eq!(got, &content[offset..offset + READ_LEN], "mismatch at offset {offset}"); + } + + let snap = driver.shutdown(); + assert!(snap.submitted <= OPS as u64, "more ops submitted than requested: {snap:?}"); + assert!(snap.delivered >= (OPS / 2) as u64, "kept half not all delivered: {snap:?}"); + assert_eq!( + snap.delivered + snap.orphan_reclaimed, + snap.submitted, + "some buffers are unaccounted for across shards: {snap:?}" + ); + assert_eq!(snap.in_flight, 0, "shutdown left ops in flight: {snap:?}"); + assert_eq!(snap.cq_overflow, 0, "per-shard in-flight cap must keep CQ overflow unreachable: {snap:?}"); + let _ = std::fs::remove_file(path); +} + +/// A single-shard sharded driver must behave exactly like `probe_and_start`, +/// which is what every other test in this file exercises. +#[tokio::test(flavor = "multi_thread")] +async fn sharded_driver_with_one_shard_matches_single_ring() { + let Some(driver) = sharded_driver_or_skip("sharded_driver_with_one_shard_matches_single_ring", 1) else { + return; + }; + const LEN: usize = 1 << 20; + let content = make_content(LEN); + let (path, file) = temp_file_with(&content, "sharded-one"); + let got = driver.read_at(Arc::clone(&file), 4096, 8192).await.expect("read"); + assert_eq!(got, &content[4096..4096 + 8192]); + let snap = driver.shutdown(); + assert_eq!(snap.delivered + snap.orphan_reclaimed, snap.submitted, "{snap:?}"); + let _ = std::fs::remove_file(path); +}