From f980b26e108e0cf66ecc8f491ad33f1ba7bd450f Mon Sep 17 00:00:00 2001 From: Hubert Gruszecki Date: Fri, 18 Sep 2026 11:29:48 +0200 Subject: [PATCH] feat(message_bus): read replica sockets ahead, bypassing large frames The link shard that owns a peer's replica socket burns a full core under sustained replicated writes. Nothing batches its reads: read_message issues one io_uring read for the 256-byte header and a second for the body. Bytes already sitting in the kernel receive queue wait for the next call. A read-ahead buffer sized by message_bus.replica_read_buffer_size lets one socket read serve a whole burst. It adds no latency, because the reader never waits for the buffer to fill. One fill delivers at most one buffer, so a body above that size crosses the buffer in buffer-sized pieces and pays one extra pass over the payload and one read per piece. A read that large goes straight into the frame's own buffer instead. Zero keeps the unbuffered path, so the A/B baseline is a config change and not a separate build. The plaintext writer already coalesces into one writev, so it is left alone. replica_socket_reads_total and replica_inbound_frames_total carry the evidence on the real workload: their ratio is the batching factor, and only link shards bump them. A read counts once it completes, so a link torn down mid-read cannot inflate the ratio. --- core/configs/src/server_config/defaults.rs | 1 + core/configs/src/server_config/displays.rs | 8 +- core/configs/src/server_config/message_bus.rs | 104 +++++- core/integration/tests/cluster/mod.rs | 1 + .../tests/cluster/replica_read_batching.rs | 150 +++++++++ core/message_bus/src/config.rs | 8 + core/message_bus/src/framing.rs | 94 +++++- core/message_bus/src/installer/replica.rs | 12 +- core/message_bus/src/lib.rs | 82 +++++ core/message_bus/src/transports/tcp.rs | 296 +++++++++++++++++- core/server/config.toml | 13 + core/shard/src/lib.rs | 2 + core/shard/src/metrics.rs | 27 ++ 13 files changed, 774 insertions(+), 24 deletions(-) create mode 100644 core/integration/tests/cluster/replica_read_batching.rs diff --git a/core/configs/src/server_config/defaults.rs b/core/configs/src/server_config/defaults.rs index 6019a8b0d9..d96b274ac5 100644 --- a/core/configs/src/server_config/defaults.rs +++ b/core/configs/src/server_config/defaults.rs @@ -296,6 +296,7 @@ impl Default for MessageBusConfig { MessageBusConfig { max_batch: bus.max_batch as usize, max_message_size: bus.max_message_size.parse().unwrap(), + replica_read_buffer_size: bus.replica_read_buffer_size.parse().unwrap(), peer_queue_capacity: bus.peer_queue_capacity as usize, client_queue_capacity: bus.client_queue_capacity as usize, reconnect_period: bus.reconnect_period.parse().unwrap(), diff --git a/core/configs/src/server_config/displays.rs b/core/configs/src/server_config/displays.rs index ee96a0bbe4..06e2b50adc 100644 --- a/core/configs/src/server_config/displays.rs +++ b/core/configs/src/server_config/displays.rs @@ -96,12 +96,14 @@ impl Display for MessageBusConfig { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { write!( f, - "{{ max_batch: {}, max_message_size: {}, peer_queue_capacity: {}, \ - reconnect_period: {}, close_peer_timeout: {}, close_grace: {}, \ - handshake_grace: {} }}", + "{{ max_batch: {}, max_message_size: {}, replica_read_buffer_size: {}, \ + peer_queue_capacity: {}, client_queue_capacity: {}, reconnect_period: {}, \ + close_peer_timeout: {}, close_grace: {}, handshake_grace: {} }}", self.max_batch, self.max_message_size, + self.replica_read_buffer_size, self.peer_queue_capacity, + self.client_queue_capacity, self.reconnect_period, self.close_peer_timeout, self.close_grace, diff --git a/core/configs/src/server_config/message_bus.rs b/core/configs/src/server_config/message_bus.rs index 8ce3ea57c3..5a008c4dc3 100644 --- a/core/configs/src/server_config/message_bus.rs +++ b/core/configs/src/server_config/message_bus.rs @@ -52,6 +52,7 @@ //! [`super::server::ServerConfig::load`]. use super::COMPONENT; +use super::defaults::SERVER_CONFIG; use crate::ConfigurationError; use configs::ConfigEnv; use iggy_common::{IggyByteSize, IggyDuration, MAX_MESSAGE_SIZE_UPPER_BYTES, Validatable}; @@ -70,7 +71,9 @@ use serde_with::{DisplayFromStr, serde_as}; /// failure until both are reconciled. pub const IOV_MAX_LIMIT: usize = 512; -const DEFAULT_CLIENT_QUEUE_CAPACITY: usize = 256; +/// Floor for a nonzero [`MessageBusConfig::replica_read_buffer_size`]. +/// A buffer below one page costs more reads and copies than no buffer. +const MIN_REPLICA_READ_BUFFER_BYTES: u64 = 4 * 1024; /// Tunables for the message bus that ships consensus traffic between /// replicas and SDK-client traffic between shards. @@ -88,6 +91,18 @@ pub struct MessageBusConfig { #[config_env(leaf)] pub max_message_size: IggyByteSize, + /// Read-ahead buffer per plaintext replica link. Costs one buffer of + /// this size per installed connection. The framing layer decodes + /// every complete frame a fill delivered, so a burst costs one read + /// instead of one or two per frame. A body at least this size is + /// read straight into its frame rather than through the buffer. Zero + /// keeps the unbuffered path. No effect under `cluster.tls`, whose + /// reader buffers inside compio's `SyncStream` adapter, nor on the + /// client plane, which never wraps its read half. + #[serde(default = "default_replica_read_buffer_size")] + #[config_env(leaf)] + pub replica_read_buffer_size: IggyByteSize, + /// Bound on each replica peer's mpsc queue. Writer task drains; the /// `send_to_*` path enqueues. Too small drops under burst; too /// large delays backpressure signalling. @@ -145,6 +160,17 @@ impl Validatable for MessageBusConfig { ); return Err(ConfigurationError::InvalidConfigurationValue); } + // The one key in this section where zero is legal: it selects the + // unbuffered read path, so the A/B baseline is a config change + // rather than a separate build. + let replica_read_buffer = self.replica_read_buffer_size.as_bytes_u64(); + if replica_read_buffer != 0 && replica_read_buffer < MIN_REPLICA_READ_BUFFER_BYTES { + eprintln!( + "{COMPONENT} message_bus.replica_read_buffer_size ({replica_read_buffer}) must be \ + 0 (unbuffered) or at least {MIN_REPLICA_READ_BUFFER_BYTES} bytes" + ); + return Err(ConfigurationError::InvalidConfigurationValue); + } if self.peer_queue_capacity == 0 { eprintln!("{COMPONENT} message_bus.peer_queue_capacity must be > 0"); return Err(ConfigurationError::InvalidConfigurationValue); @@ -187,8 +213,19 @@ impl Validatable for MessageBusConfig { } } -const fn default_client_queue_capacity() -> usize { - DEFAULT_CLIENT_QUEUE_CAPACITY +/// Zero fails validation, so the value comes from the embedded `config.toml`. +fn default_client_queue_capacity() -> usize { + SERVER_CONFIG.message_bus.client_queue_capacity as usize +} + +/// [`IggyByteSize`]'s own `Default` is 0 bytes, which turns the read-ahead +/// buffer off, so the value comes from the embedded `config.toml`. +fn default_replica_read_buffer_size() -> IggyByteSize { + SERVER_CONFIG + .message_bus + .replica_read_buffer_size + .parse() + .expect("message_bus.replica_read_buffer_size is a byte size") } #[cfg(test)] @@ -241,7 +278,10 @@ mod tests { .remove("client_queue_capacity"); config["peer_queue_capacity"] = serde_json::json!(8192); let decoded: MessageBusConfig = serde_json::from_value(config).unwrap(); - assert_eq!(decoded.client_queue_capacity, DEFAULT_CLIENT_QUEUE_CAPACITY); + assert_eq!( + decoded.client_queue_capacity, + baseline().client_queue_capacity + ); assert_eq!(decoded.peer_queue_capacity, 8192); decoded.validate().unwrap(); } @@ -253,6 +293,62 @@ mod tests { assert!(config.validate().is_err()); } + #[test] + fn default_replica_read_buffer_is_256_kib() { + assert_eq!( + baseline().replica_read_buffer_size.as_bytes_u64(), + 256 * 1024 + ); + } + + /// Zero is the unbuffered path, not a misconfiguration: it is the + /// A/B baseline arm and must survive validation. + #[test] + fn accepts_zero_replica_read_buffer() { + let mut c = baseline(); + c.replica_read_buffer_size = IggyByteSize::from(0_u64); + assert!(c.validate().is_ok()); + } + + #[test] + fn accepts_replica_read_buffer_at_floor() { + let mut c = baseline(); + c.replica_read_buffer_size = IggyByteSize::from(MIN_REPLICA_READ_BUFFER_BYTES); + assert!(c.validate().is_ok()); + } + + #[test] + fn rejects_replica_read_buffer_below_floor() { + let mut c = baseline(); + c.replica_read_buffer_size = IggyByteSize::from(MIN_REPLICA_READ_BUFFER_BYTES - 1); + assert!(c.validate().is_err()); + } + + #[test] + fn replica_read_buffer_parses_byte_size_strings() { + for (text, expected) in [("0", 0_u64), ("64 KiB", 64 * 1024)] { + let mut config = serde_json::to_value(baseline()).unwrap(); + config["replica_read_buffer_size"] = serde_json::json!(text); + let decoded: MessageBusConfig = serde_json::from_value(config).unwrap(); + assert_eq!(decoded.replica_read_buffer_size.as_bytes_u64(), expected); + decoded.validate().unwrap(); + } + } + + #[test] + fn missing_replica_read_buffer_keeps_config_default() { + let mut config = serde_json::to_value(baseline()).unwrap(); + config + .as_object_mut() + .unwrap() + .remove("replica_read_buffer_size"); + let decoded: MessageBusConfig = serde_json::from_value(config).unwrap(); + assert_eq!( + decoded.replica_read_buffer_size, + baseline().replica_read_buffer_size + ); + } + #[test] fn rejects_zero_max_message_size() { let mut c = baseline(); diff --git a/core/integration/tests/cluster/mod.rs b/core/integration/tests/cluster/mod.rs index e9f81fb936..c4d183fb07 100644 --- a/core/integration/tests/cluster/mod.rs +++ b/core/integration/tests/cluster/mod.rs @@ -31,4 +31,5 @@ mod partition_dedup; mod partition_primary_routing; mod partition_state_transfer; mod register_forwarding; +mod replica_read_batching; mod staggered_bootstrap; diff --git a/core/integration/tests/cluster/replica_read_batching.rs b/core/integration/tests/cluster/replica_read_batching.rs new file mode 100644 index 0000000000..1f212ca377 --- /dev/null +++ b/core/integration/tests/cluster/replica_read_batching.rs @@ -0,0 +1,150 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The plaintext replica plane counts what it reads. +//! +//! `replica_socket_reads_total` and `replica_inbound_frames_total` are the +//! direct evidence for `message_bus.replica_read_buffer_size`: their ratio is +//! the batching factor the read-ahead buffer exists to raise. Only a shard +//! that owns a replica socket bumps them, so a nonzero pair also names the +//! link shard. +//! +//! One shard per node, so shard 0 is that link shard and the assertion needs +//! no search. What is pinned here is that both counters reach the scrape at +//! all: the `Rc` reaches the reader task, the take in `tick_partitions` runs, +//! and the two `ShardMetrics` counters are registered. The ratio itself is +//! only logged. Frame arrival timing on a live link is not controlled, so a +//! threshold on it would be a flaky test; +//! `framing::tests::buffered_read_batches_many_frames_per_socket_read` asserts +//! the factor where the traffic shape is fixed. +//! +//! `frames >= reads` is deliberately not asserted: a body split across two +//! segments, or one straddling the end of the buffer, costs two reads for one +//! frame. + +use iggy::prelude::*; +use integration::iggy_harness; +use std::time::Duration; +use tokio::time::{Instant, sleep}; + +use crate::server::http_client::HttpClient; + +const STREAM: &str = "replica-read-stream"; +const TOPIC: &str = "replica-read-topic"; +const PARTITION_ID: u32 = 0; +const MESSAGES: u32 = 64; + +/// The counters land on the scrape through the shard's partition tick, so the +/// first scrape after a produce can still read zero. +const COUNTER_BUDGET: Duration = Duration::from_secs(10); +const COUNTER_POLL: Duration = Duration::from_millis(250); + +/// One shard per node, so this is the only shard and it owns both replica +/// links. +const LINK_SHARD: u16 = 0; + +/// Read one `shard`-labelled counter out of the Prometheus text exposition. +/// +/// The sub-registry label comes first in the label set, and these two counters +/// carry no others, so the series name plus the label is an exact line prefix. +fn shard_counter(metrics: &str, name: &str, shard: u16) -> Option { + let prefix = format!("{name}{{shard=\"{shard}\"}} "); + metrics + .lines() + .find_map(|line| line.strip_prefix(&prefix)?.trim().parse().ok()) +} + +async fn scrape(http: &HttpClient) -> String { + http.client + .get(http.url("/metrics")) + .bearer_auth(&http.token) + .send() + .await + .expect("metrics response") + .text() + .await + .expect("metrics text") +} + +#[iggy_harness(cluster_nodes = 3, server(sharding.cpu_allocation = "0..1"))] +async fn given_a_replicating_cluster_when_scraping_should_report_replica_reads_and_frames( + harness: &TestHarness, +) { + let client = harness.new_client().await.unwrap(); + client + .login_user(DEFAULT_ROOT_USERNAME, DEFAULT_ROOT_PASSWORD) + .await + .unwrap(); + let stream = Identifier::named(STREAM).unwrap(); + let topic = Identifier::named(TOPIC).unwrap(); + client.create_stream(STREAM).await.unwrap(); + client + .create_topic( + &stream, + TOPIC, + &TopicCreateOptions { + partitions_count: Some(1), + message_expiry: Some(IggyExpiry::NeverExpire), + ..TopicCreateOptions::default() + }, + ) + .await + .expect("create topic"); + + for index in 0..MESSAGES { + let mut messages = vec![ + IggyMessage::builder() + .payload(format!("payload-{index}").into()) + .build() + .expect("message build"), + ]; + client + .send_messages( + &stream, + &topic, + &Partitioning::partition_id(PARTITION_ID), + &mut messages, + ) + .await + .unwrap_or_else(|error| panic!("send_messages {index}: {error}")); + } + + let http = HttpClient::login_root(harness).await; + let deadline = Instant::now() + COUNTER_BUDGET; + loop { + let metrics = scrape(&http).await; + let reads = shard_counter(&metrics, "replica_socket_reads_total", LINK_SHARD); + let frames = shard_counter(&metrics, "replica_inbound_frames_total", LINK_SHARD); + if let (Some(reads), Some(frames)) = (reads, frames) + && reads > 0 + && frames > 0 + { + println!( + "link shard {LINK_SHARD}: {frames} frames over {reads} socket reads, \ + {:.2} frames per read", + frames as f64 / reads as f64 + ); + return; + } + assert!( + Instant::now() < deadline, + "link shard {LINK_SHARD} reported reads={reads:?} frames={frames:?} within \ + {COUNTER_BUDGET:?}; both must be nonzero once replica traffic has flowed" + ); + sleep(COUNTER_POLL).await; + } +} diff --git a/core/message_bus/src/config.rs b/core/message_bus/src/config.rs index 19731b4f77..aefaf5624f 100644 --- a/core/message_bus/src/config.rs +++ b/core/message_bus/src/config.rs @@ -138,6 +138,12 @@ pub struct MessageBusConfig { /// validator; undersize or oversize frames are rejected. pub max_message_size: usize, + /// Read-ahead buffer per plaintext replica link, in bytes. Threaded + /// into `TcpTransportConn::with_replica_read` by the replica + /// installer; zero selects the unbuffered read path. The client + /// plane and the TLS-family transports ignore it. + pub replica_read_buffer_size: usize, + /// Bound on each replica peer's mpsc queue. The writer task drains; the /// `send_to_*` path enqueues. Too small drops under burst; too /// large delays backpressure signalling. @@ -216,6 +222,8 @@ impl From<&ServerConfig> for MessageBusConfig { max_batch: bus.max_batch, max_message_size: usize::try_from(bus.max_message_size.as_bytes_u64()) .expect("message_bus.max_message_size fits usize on supported targets"), + replica_read_buffer_size: usize::try_from(bus.replica_read_buffer_size.as_bytes_u64()) + .expect("message_bus.replica_read_buffer_size fits usize on supported targets"), peer_queue_capacity: bus.peer_queue_capacity, client_queue_capacity: bus.client_queue_capacity, reconnect_period: bus.reconnect_period.get_duration(), diff --git a/core/message_bus/src/framing.rs b/core/message_bus/src/framing.rs index ab626c65ed..3eab855150 100644 --- a/core/message_bus/src/framing.rs +++ b/core/message_bus/src/framing.rs @@ -130,7 +130,9 @@ pub async fn read_message( // Stage 2: grow the same `Owned` in place and fill the tail via a // slice read. Total allocations for a body frame: ONE // (`Owned::with_capacity(HEADER_SIZE)` plus one in-place realloc of - // the backing AVec). Zero memcpys of the data. + // the backing AVec). The body lands in this `Owned` in one pass; a + // buffered read half copies only what it read ahead (at most its + // buffer size) before the rest goes straight in. let mut owned = owned; // The `.map_err` arm is unreachable today: `Owned` is // the only `IoBufMut` ever fed here, and its `reserve_exact` impl @@ -174,8 +176,11 @@ fn to_read_error(e: &std::io::Error) -> IggyError { #[cfg(test)] mod tests { use super::*; + use crate::ReplicaReadStats; + use crate::transports::tcp::{CountingRead, ReadAhead}; use compio::net::{TcpListener, TcpStream}; use iggy_binary_protocol::{Command, SIZE_FIELD_OFFSET}; + use std::rc::Rc; #[allow(clippy::cast_possible_truncation)] fn make_header_only(command: Command) -> Message { @@ -224,6 +229,93 @@ mod tests { assert!(matches!(res, Err(IggyError::InvalidCommand))); } + /// The one place the batching factor is asserted: here the traffic + /// shape is controlled, so many frames provably land in one socket + /// read. On a live link the ratio depends on arrival timing. + #[compio::test] + #[allow(clippy::future_not_send)] + async fn buffered_read_batches_many_frames_per_socket_read() { + const FRAMES: usize = 64; + const CAPACITY: usize = 64 * 1024; + + let (mut a, b) = local_pair().await; + let batch: Vec<_> = (0..FRAMES) + .map(|_| make_header_only(Command::Ping).into_frozen()) + .collect(); + a.write_vectored_all(batch).await.0.unwrap(); + + let stats = Rc::new(ReplicaReadStats::default()); + let mut reader = ReadAhead::new(CAPACITY, CountingRead::new(b, Rc::clone(&stats))); + for _ in 0..FRAMES { + let read = read_message(&mut reader, MAX_MESSAGE_SIZE).await.unwrap(); + assert_eq!(read.header().command, Command::Ping); + } + + let reads = stats.take().reads; + assert!( + (1..FRAMES as u64 / 8).contains(&reads), + "{FRAMES} frames of one write cost {reads} reads; unbuffered would cost {FRAMES}" + ); + } + + /// A frame several times the buffer decodes intact, and does not + /// cross the buffer in buffer-sized pieces: a read at least one + /// buffer long goes straight to the socket. + #[compio::test] + #[allow(clippy::future_not_send, clippy::cast_possible_truncation)] + async fn buffered_read_decodes_frame_larger_than_buffer() { + const CAPACITY: usize = 4 * 1024; + const MULTIPLE: usize = 8; + let frame_len = MULTIPLE * CAPACITY; + + let mut msg = Message::::new(frame_len).transmute_header( + |_, h: &mut GenericHeader| { + h.command = Command::Ping; + h.size = frame_len as u32; + }, + ); + for (at, byte) in msg.as_mut_slice()[HEADER_SIZE..].iter_mut().enumerate() { + *byte = (at % 251) as u8; + } + let expected = msg.as_slice().to_vec(); + + let (mut a, b) = local_pair().await; + let stats = Rc::new(ReplicaReadStats::default()); + let mut reader = ReadAhead::new(CAPACITY, CountingRead::new(b, Rc::clone(&stats))); + write_message(&mut a, msg).await.unwrap(); + + let read = read_message(&mut reader, MAX_MESSAGE_SIZE).await.unwrap(); + assert_eq!(read.as_slice(), expected.as_slice()); + let reads = stats.take().reads; + assert!( + reads >= 2, + "a frame {MULTIPLE}x the buffer cannot arrive in one read" + ); + assert!( + reads < MULTIPLE as u64, + "{MULTIPLE} buffers of body cost {reads} reads; copied through the buffer \ + it cannot cost fewer than {MULTIPLE}" + ); + } + + /// EOF partway through a header stays a closed connection rather than + /// a framing error once the read half is buffered. + #[compio::test] + #[allow(clippy::future_not_send)] + async fn buffered_read_reports_eof_after_partial_header() { + use compio::io::AsyncWriteExt; + let (mut a, b) = local_pair().await; + a.write_all(vec![0u8; HEADER_SIZE / 2]).await.0.unwrap(); + drop(a); + + let mut reader = ReadAhead::new( + 4 * 1024, + CountingRead::new(b, Rc::new(ReplicaReadStats::default())), + ); + let res = read_message(&mut reader, MAX_MESSAGE_SIZE).await; + assert!(matches!(res, Err(IggyError::ConnectionClosed))); + } + #[compio::test] #[allow(clippy::future_not_send)] async fn read_rejects_undersize_size_field() { diff --git a/core/message_bus/src/installer/replica.rs b/core/message_bus/src/installer/replica.rs index 5bb4556fbe..172c521733 100644 --- a/core/message_bus/src/installer/replica.rs +++ b/core/message_bus/src/installer/replica.rs @@ -123,12 +123,16 @@ pub fn install_replica_inbound( async move { match ctx.tls.clone() { None => { + let read_buffer = bus.config().replica_read_buffer_size; + let read_stats = bus.replica_read_stats(); accept_and_install( &bus, stream, &ctx, ChannelBinding::Plaintext, - TcpTransportConn::new, + move |stream| { + TcpTransportConn::new(stream).with_replica_read(read_buffer, read_stats) + }, on_message, &peer, ) @@ -232,13 +236,17 @@ pub fn install_replica_outbound( async move { match ctx.tls.clone() { None => { + let read_buffer = bus.config().replica_read_buffer_size; + let read_stats = bus.replica_read_stats(); dial_and_install( &bus, stream, &ctx, peer_id, ChannelBinding::Plaintext, - TcpTransportConn::new, + move |stream| { + TcpTransportConn::new(stream).with_replica_read(read_buffer, read_stats) + }, on_message, ) .await diff --git a/core/message_bus/src/lib.rs b/core/message_bus/src/lib.rs index dc2348daaa..1471469b79 100644 --- a/core/message_bus/src/lib.rs +++ b/core/message_bus/src/lib.rs @@ -432,6 +432,58 @@ pub type AcceptedWssClientFn = std::rc::Rc; +/// Socket-read accounting for one plaintext replica link. +/// +/// The link shard reads every inbound consensus frame off one socket, so +/// the ratio of decoded frames to completed socket reads is the batching +/// factor the read-ahead buffer exists to raise. Only link shards read +/// replica sockets, so a nonzero value also identifies the link shard. +/// +/// `Cell` rather than atomics: the reader task that bumps these and the +/// shard sweep that drains them run on the same compio thread. Shared +/// with the reader task through an `Rc`, on the precedent of +/// `installer::replica`'s `install_aborted`. +#[derive(Debug, Default)] +pub struct ReplicaReadStats { + reads: Cell, + frames: Cell, +} + +impl ReplicaReadStats { + /// One completed socket read on the replica read half. + pub fn record_read(&self) { + self.reads.set(self.reads.get() + 1); + } + + /// One frame decoded out of the replica read half. + pub fn record_frame(&self) { + self.frames.set(self.frames.get() + 1); + } + + /// Drain both counts, resetting them to zero. + /// + /// Take semantics because the caller feeds the result to a cumulative + /// Prometheus counter's `inc_by`: handing back running totals would + /// double count every sweep. + pub const fn take(&self) -> ReplicaReadMetrics { + ReplicaReadMetrics { + reads: self.reads.replace(0), + frames: self.frames.replace(0), + } + } +} + +/// Deltas since the previous [`ReplicaReadStats::take`]. +/// +/// `frames >= reads` is not an invariant: a body that arrives in two +/// segments, or one that straddles the end of the read-ahead buffer, +/// costs two reads for one frame. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub struct ReplicaReadMetrics { + pub reads: u64, + pub frames: u64, +} + /// Point-to-point message delivery between consensus participants. /// /// `Ok(())` means "accepted for delivery" - NOT "delivered to peer". @@ -627,6 +679,16 @@ pub trait MessageBus { fn realtime_micros(&self) -> u64 { iggy_common::IggyTimestamp::now().as_micros() } + + /// Drain this bus's plaintext replica socket-read counters. + /// + /// The provided default returns zeros so the simulator bus and the + /// test buses need no override: they own no replica socket. The + /// production bus takes from its own [`ReplicaReadStats`], which the + /// reader task of every plaintext replica link shares. + fn take_replica_read_stats(&self) -> ReplicaReadMetrics { + ReplicaReadMetrics::default() + } } /// Production message bus backed by real TCP connections. @@ -707,6 +769,11 @@ pub struct IggyMessageBus { /// pending; entries expire lazily inside [`Self::check_dial_pending`]. /// Only shard 0's bus ever populates this. pending_dials: RefCell>, + /// Socket-read counters shared by the reader task of every plaintext + /// replica link this shard owns. Drained through + /// [`MessageBus::take_replica_read_stats`] by the shard's tick; + /// stays at zero on a shard that owns no such link. + replica_read_stats: Rc, } impl IggyMessageBus { @@ -819,6 +886,7 @@ impl IggyMessageBus { replica_handshake_slots: RefCell::new(ahash::AHashMap::new()), replica_slot_seq: Cell::new(0), pending_dials: RefCell::new(ahash::AHashMap::new()), + replica_read_stats: Rc::new(ReplicaReadStats::default()), } } @@ -1109,6 +1177,12 @@ impl IggyMessageBus { &self.config } + /// Handle every plaintext replica reader task on this shard bumps. + #[must_use] + pub fn replica_read_stats(&self) -> Rc { + Rc::clone(&self.replica_read_stats) + } + /// Cheap clone of the root shutdown token. /// /// Handed to accept loops, read tasks, writer tasks, and periodic tasks @@ -1296,6 +1370,10 @@ impl MessageBus for std::rc::Rc { fn realtime_micros(&self) -> u64 { (**self).realtime_micros() } + + fn take_replica_read_stats(&self) -> ReplicaReadMetrics { + (**self).take_replica_read_stats() + } } #[allow(clippy::future_not_send)] @@ -1394,6 +1472,10 @@ impl MessageBus for IggyMessageBus { fn track_background(&self, handle: JoinHandle<()>) { Self::track_background(self, handle); } + + fn take_replica_read_stats(&self) -> ReplicaReadMetrics { + self.replica_read_stats.take() + } } /// Extract the owning shard from a client id. diff --git a/core/message_bus/src/transports/tcp.rs b/core/message_bus/src/transports/tcp.rs index 53845b41e1..bec39cae69 100644 --- a/core/message_bus/src/transports/tcp.rs +++ b/core/message_bus/src/transports/tcp.rs @@ -27,11 +27,20 @@ //! and emits them via `compio::io::AsyncWriteExt::write_vectored_all` //! (one writev in the common case, looping on short writes; zero //! intermediate copies of `Frozen`). +//! +//! The replica plane wraps its read half through +//! `TcpTransportConn::with_replica_read`: socket reads are counted on +//! every link, and above a nonzero `replica_read_buffer_size` one read +//! fills a buffer the framing layer then decodes many frames out of. The +//! client plane keeps the bare read half. use super::{ActorContext, TransportConn, TransportListener}; +use crate::ReplicaReadStats; use crate::framing; use crate::lifecycle::BusMessage; -use compio::io::AsyncWriteExt; +use compio::BufResult; +use compio::buf::IoBufMut; +use compio::io::{AsyncRead, AsyncWriteExt}; use compio::net::{TcpListener, TcpStream}; use compio::runtime::fd::PollFd; use futures::FutureExt; @@ -41,6 +50,7 @@ use std::io; use std::mem; use std::net::SocketAddr; use std::os::fd::AsRawFd; +use std::rc::Rc; use tracing::{debug, error, trace}; /// Inbound TCP listener wrapper. @@ -88,18 +98,44 @@ impl TransportListener for TcpTransportListener { /// reader and writer tasks internally. pub(crate) struct TcpTransportConn { stream: TcpStream, + /// Replica plane only: read-ahead capacity in bytes (0 selects the + /// unbuffered path) paired with the link's socket-read counters. + /// `None` on the client plane, which never wraps its read half. + replica_read: Option<(usize, Rc)>, } impl TcpTransportConn { #[must_use] pub(crate) const fn new(stream: TcpStream) -> Self { - Self { stream } + Self { + stream, + replica_read: None, + } + } + + /// Count socket reads on this link, and read ahead into a buffer of + /// `capacity` bytes when that is nonzero. + /// + /// Called only from the replica installer. The client installer + /// stays on [`Self::new`], so no client connection pays for either. + #[must_use] + pub(crate) fn with_replica_read( + mut self, + capacity: usize, + stats: Rc, + ) -> Self { + self.replica_read = Some((capacity, stats)); + self } } impl TransportConn for TcpTransportConn { #[allow(clippy::future_not_send)] async fn run(self, ctx: ActorContext) { + let Self { + stream, + replica_read, + } = self; // Capture a refcounted poll fd BEFORE `into_split` so the // shutdown watchdog can wake the reader's parked io_uring read // SQE via `libc::shutdown(SHUT_RD)`. No `select!` over a TCP @@ -108,8 +144,8 @@ impl TransportConn for TcpTransportConn { // io_uring read is not cancel-safe in the protocol sense). // compio 0.19 dropped `TcpStream::to_shared_fd`; `to_poll_fd` // gives an equivalently refcounted fd handle (PollFd: AsRawFd). - let poll_fd = self.stream.to_poll_fd().ok(); - let (read_half, write_half) = self.stream.into_split(); + let poll_fd = stream.to_poll_fd().ok(); + let (read_half, write_half) = stream.into_split(); let ActorContext { in_tx, rx, @@ -127,13 +163,24 @@ impl TransportConn for TcpTransportConn { let writer_shutdown = shutdown; let reader_peer = peer.clone(); - let reader_handle = compio::runtime::spawn(reader_loop( - read_half, - in_tx, - max_message_size, - label, - reader_peer, - )); + let reader_handle = match replica_read { + Some((capacity, stats)) => compio::runtime::spawn(reader_loop( + ReadAhead::new(capacity, CountingRead::new(read_half, Rc::clone(&stats))), + in_tx, + max_message_size, + Some(stats), + label, + reader_peer, + )), + None => compio::runtime::spawn(reader_loop( + read_half, + in_tx, + max_message_size, + None, + label, + reader_peer, + )), + }; let writer_handle = compio::runtime::spawn(writer_loop( write_half, rx, @@ -187,26 +234,140 @@ fn spawn_shutdown_watchdog( .detach(); } +/// [`AsyncRead`] adapter that counts completed reads of the inner +/// source. +/// +/// Sits under the [`ReadAhead`] buffer, so it sees socket reads rather +/// than the buffer hits the framing layer makes against the buffer. +/// Wraps the replica read half even when read-ahead is off, so the +/// unbuffered arm of an A/B run reports a ratio rather than zeros. +pub(crate) struct CountingRead { + inner: R, + stats: Rc, +} + +impl CountingRead { + pub(crate) const fn new(inner: R, stats: Rc) -> Self { + Self { inner, stats } + } +} + +impl AsyncRead for CountingRead { + #[allow(clippy::future_not_send)] + async fn read(&mut self, buf: B) -> BufResult { + let res = self.inner.read(buf).await; + // Counted regardless of outcome: an error and an `Ok(0)` each + // cost the same submit and completion as a full read. Recorded + // after the read so a link force-cancelled at teardown does not + // count a read that never completed. + self.stats.record_read(); + res + } +} + +/// Replica read half that reads ahead into a buffer while keeping large +/// reads off it. +/// +/// Socket reads land in the buffer, so the framing layer decodes every +/// complete frame one fill delivered: a burst of small frames costs one +/// read instead of one or two per frame. A read at least one buffer long +/// goes straight to the socket instead, because a fill cannot deliver +/// more than `capacity`: a larger frame would cross the buffer in +/// buffer-sized pieces, which costs one extra pass over the payload and +/// one socket read per piece. `capacity == 0` buffers nothing, because +/// every read is then at least one buffer long. +/// +/// A fill moves the buffer into the inner read and restores it when that +/// read completes, so a fill future dropped mid-read leaves the buffer +/// empty rather than half-filled. Only tearing the peer down drops a +/// fill, and that drops the whole read half with it. +pub(crate) struct ReadAhead { + inner: R, + capacity: usize, + buf: Vec, + /// Bytes of `buf` already handed to the caller. + pos: usize, +} + +impl ReadAhead { + pub(crate) fn new(capacity: usize, inner: R) -> Self { + Self { + inner, + capacity, + buf: Vec::with_capacity(capacity), + pos: 0, + } + } +} + +impl ReadAhead { + /// Refill the buffer from the socket. A short read is normal and + /// means the socket held nothing more. + #[allow(clippy::future_not_send)] + async fn fill(&mut self) -> io::Result<()> { + self.buf.clear(); + self.pos = 0; + // No-op unless a dropped fill left the buffer without capacity. + self.buf.reserve_exact(self.capacity); + let buf = mem::take(&mut self.buf); + let BufResult(res, buf) = self.inner.read(buf).await; + self.buf = buf; + res.map(|_| ()) + } +} + +impl AsyncRead for ReadAhead { + #[allow(clippy::future_not_send)] + async fn read(&mut self, mut buf: B) -> BufResult { + if self.pos == self.buf.len() { + // The caller left room for a whole buffer or more, so the + // buffer would only add a copy on the way through. + if buf.buf_capacity() >= self.capacity { + return self.inner.read(buf).await; + } + if let Err(e) = self.fill().await { + return BufResult(Err(e), buf); + } + } + let mut src: &[u8] = &self.buf[self.pos..]; + let BufResult(res, buf) = src.read(buf).await; + if let Ok(n) = res { + self.pos += n; + } + BufResult(res, buf) + } +} + /// Read framed consensus messages off the wire and forward each to /// [`ActorContext::in_tx`]. Exits on EOF, framing error, or send-side /// closure. /// +/// Generic over the read half so the replica plane can hand in a +/// [`ReadAhead`] wrapping [`CountingRead`] while the client plane hands +/// in the bare `TcpStream`. `stats` is `Some` on every replica link and +/// counts decoded frames against the socket reads `CountingRead` counts. +/// /// No `select!` over the TCP read. Cooperative shutdown is delivered /// via [`spawn_shutdown_watchdog`], which calls `libc::shutdown(SHUT_RD)` /// when the bus token fires; the in-flight read returns `Ok(0)` and /// `framing::read_message` surfaces it as an EOF error on the next -/// iteration. +/// iteration. Frames already complete in the read-ahead buffer are still +/// decoded and delivered before that happens. #[allow(clippy::future_not_send)] -async fn reader_loop( - mut read_half: TcpStream, +async fn reader_loop( + mut read_half: R, in_tx: async_channel::Sender>, max_message_size: usize, + stats: Option>, label: &'static str, peer: String, ) { loop { match framing::read_message(&mut read_half, max_message_size).await { Ok(msg) => { + if let Some(stats) = &stats { + stats.record_frame(); + } if in_tx.send(msg).await.is_err() { debug!(%label, %peer, "tcp reader: inbound queue dropped"); return; @@ -365,6 +526,12 @@ mod tests { (client_res.unwrap(), server) } + /// The shipped default, so a bump in `config.toml` is exercised here + /// rather than pinned to a stale literal. + fn default_replica_read_buffer() -> usize { + crate::MessageBusConfig::default().replica_read_buffer_size + } + #[allow(clippy::future_not_send)] fn drive( conn: TcpTransportConn, @@ -439,6 +606,91 @@ mod tests { let _ = server_handle.await; } + /// Same contract as the plaintext twin above, with the replica + /// plane's read-ahead buffer at the configured default, plus the + /// counters the A/B run reads. + #[compio::test] + #[allow(clippy::future_not_send)] + async fn run_pumps_three_frames_through_buffered_reader() { + let (client, server) = local_pair().await; + let stats = Rc::new(ReplicaReadStats::default()); + let (client_out, _client_in, client_shutdown, client_handle) = + drive(TcpTransportConn::new(client)); + let (_server_out, server_in, server_shutdown, server_handle) = drive( + TcpTransportConn::new(server) + .with_replica_read(default_replica_read_buffer(), Rc::clone(&stats)), + ); + + for cmd in [Command::Ping, Command::Prepare, Command::Request] { + client_out.send(header_only(cmd).into()).await.unwrap(); + } + + let recv_with_timeout = |rx: &async_channel::Receiver>| { + let rx = rx.clone(); + async move { + compio::time::timeout(Duration::from_secs(2), rx.recv()) + .await + .expect("recv within 2s") + .expect("ok") + } + }; + let a = recv_with_timeout(&server_in).await; + let b = recv_with_timeout(&server_in).await; + let c = recv_with_timeout(&server_in).await; + assert_eq!(a.header().command, Command::Ping); + assert_eq!(b.header().command, Command::Prepare); + assert_eq!(c.header().command, Command::Request); + + let counts = stats.take(); + assert_eq!(counts.frames, 3); + assert!( + counts.reads >= 1, + "three frames cannot arrive in zero reads" + ); + + client_shutdown.trigger(); + server_shutdown.trigger(); + let _ = client_handle.await; + let _ = server_handle.await; + } + + /// Capacity 0 keeps the unbuffered read path and still counts, so the + /// baseline arm of an A/B run reports the same ratio as the buffered + /// arms instead of zeros. + #[compio::test] + #[allow(clippy::future_not_send)] + async fn unbuffered_replica_read_still_counts() { + let (client, server) = local_pair().await; + let stats = Rc::new(ReplicaReadStats::default()); + let (client_out, _client_in, client_shutdown, client_handle) = + drive(TcpTransportConn::new(client)); + let (_server_out, server_in, server_shutdown, server_handle) = + drive(TcpTransportConn::new(server).with_replica_read(0, Rc::clone(&stats))); + + for cmd in [Command::Ping, Command::Prepare, Command::Request] { + client_out.send(header_only(cmd).into()).await.unwrap(); + let read = compio::time::timeout(Duration::from_secs(2), server_in.recv()) + .await + .expect("recv within 2s") + .expect("ok"); + assert_eq!(read.header().command, cmd); + } + + let counts = stats.take(); + assert_eq!(counts.frames, 3); + assert!( + counts.reads >= counts.frames, + "the unbuffered path costs at least one read per frame, got {} for {}", + counts.reads, + counts.frames + ); + + client_shutdown.trigger(); + server_shutdown.trigger(); + let _ = client_handle.await; + let _ = server_handle.await; + } + /// A frame fragmented to `IOV_MAX` or past it must reach the peer whole /// and in order: `sendmsg` rejects more than `IOV_MAX` iovecs with /// `EMSGSIZE`, so the writer has to chunk the flattened batch instead of @@ -534,6 +786,22 @@ mod tests { assert!(res.is_ok(), "run must exit within 2s of shutdown"); } + /// The `SHUT_RD` watchdog still tears the reader down when its read + /// half sits under a `ReadAhead`. + #[compio::test] + #[allow(clippy::future_not_send)] + async fn buffered_run_exits_on_shutdown_signal() { + let (client, _server) = local_pair().await; + let conn = TcpTransportConn::new(client).with_replica_read( + default_replica_read_buffer(), + Rc::new(ReplicaReadStats::default()), + ); + let (_out_tx, _in_rx, shutdown, handle) = drive(conn); + shutdown.trigger(); + let res = compio::time::timeout(Duration::from_secs(2), handle).await; + assert!(res.is_ok(), "run must exit within 2s of shutdown"); + } + #[compio::test] #[allow(clippy::future_not_send)] async fn run_reports_oversize_frame_and_exits() { diff --git a/core/server/config.toml b/core/server/config.toml index b32988b0fb..84a36c7405 100644 --- a/core/server/config.toml +++ b/core/server/config.toml @@ -1057,6 +1057,19 @@ max_batch = 256 # also breaks Go clients, whose frame cap is a hard 64 MiB constant. max_message_size = "64 MiB" +# Read-ahead buffer per plaintext replica link. The framing layer decodes every +# complete frame a fill delivered, and a fill takes what the socket holds rather +# than waiting for a full buffer. A body at least this size is read straight into +# its frame. 0 keeps the unbuffered path (one read per header, one per body). +# No-op under cluster.tls. +# +# 256 KiB holds 64 average Prepares or 1024 PrepareOks of one burst, and is +# twice the kernel's default receive buffer (tcp_rmem[1], 128 KiB). Kernel +# autotuning grows the receive queue past that default once the reader falls +# behind, so under a CPU-bound reader this value alone sets the batching +# factor. Boot rejects a nonzero value below 4 KiB. +replica_read_buffer_size = "256 KiB" + # Bound on each replica peer's mpsc queue. The writer task drains; the # send_to_* path enqueues. peer_queue_capacity = 4096 diff --git a/core/shard/src/lib.rs b/core/shard/src/lib.rs index 803547fba2..287b747b8b 100644 --- a/core/shard/src/lib.rs +++ b/core/shard/src/lib.rs @@ -7635,6 +7635,8 @@ where } } self.metrics.record_persistence(&persistence_metrics); + self.metrics + .record_replica_reads(&self.bus.take_replica_read_stats()); self.metrics .set_repair_ring(repair_ring_entries, repair_ring_bytes); diff --git a/core/shard/src/metrics.rs b/core/shard/src/metrics.rs index 4caf701f20..f16de0c2d7 100644 --- a/core/shard/src/metrics.rs +++ b/core/shard/src/metrics.rs @@ -44,6 +44,7 @@ use prometheus_client::registry::Registry; use std::sync::{Arc, OnceLock}; use iggy_common::ConsumerKind; +use message_bus::ReplicaReadMetrics; /// Label for `frame_drops_total`. /// @@ -246,6 +247,8 @@ pub struct ShardMetrics { metadata_prepare_gap_drops_total: Counter, metadata_read_frontier_refusals_total: Counter, client_requests_denied_queue_full_total: Counter, + replica_socket_reads_total: Counter, + replica_inbound_frames_total: Counter, partition_consumer_offsets_denied_total: Family, consumer_offset_denied_counters: [Counter; 2], partition_consumer_offsets_stranded: Family, @@ -320,6 +323,8 @@ impl ShardMetrics { metadata_prepare_gap_drops_total: Counter::default(), metadata_read_frontier_refusals_total: Counter::default(), client_requests_denied_queue_full_total: Counter::default(), + replica_socket_reads_total: Counter::default(), + replica_inbound_frames_total: Counter::default(), partition_consumer_offsets_denied_total, consumer_offset_denied_counters, partition_consumer_offsets_stranded, @@ -347,6 +352,18 @@ impl ShardMetrics { self.partition_wal_errors.inc_by(metrics.failed_writes); } + /// Fold one sweep's worth of replica socket-read deltas in. + /// + /// Deltas, never running totals: [`ReplicaReadMetrics`] comes from a + /// take that resets the source, so feeding cumulative values here + /// would double count every sweep. Both counters stay at zero on a + /// shard that owns no plaintext replica link, which is what makes + /// them identify the link shard. + pub fn record_replica_reads(&self, metrics: &ReplicaReadMetrics) { + self.replica_socket_reads_total.inc_by(metrics.reads); + self.replica_inbound_frames_total.inc_by(metrics.frames); + } + fn register_persistence(&self, registry: &mut Registry) { registry.register( "partition_wal_disk_bytes", @@ -823,6 +840,16 @@ impl ShardMetrics { "replicated prepares dropped out of order by a backup's gap check", self.partition_prepare_gap_drops_total.clone(), ); + registry.register( + "replica_socket_reads", + "completed socket reads on this shard's plaintext replica links", + self.replica_socket_reads_total.clone(), + ); + registry.register( + "replica_inbound_frames", + "frames decoded off this shard's plaintext replica links", + self.replica_inbound_frames_total.clone(), + ); registry.register( "metadata_prepare_gap_drops", "replicated metadata prepares dropped out of order by a backup's gap check",