From 92bf55ce62ce3412dc4e875636c329571d727131 Mon Sep 17 00:00:00 2001 From: houseme Date: Wed, 8 Jul 2026 08:45:21 +0800 Subject: [PATCH] perf(ecstore): right-size BytesMut encode ingest capacity to the EC-expanded block (#4396) --- crates/ecstore/src/erasure/coding/encode.rs | 173 ++++++++++++++++++- crates/ecstore/src/erasure/coding/erasure.rs | 67 ++++++- 2 files changed, 231 insertions(+), 9 deletions(-) diff --git a/crates/ecstore/src/erasure/coding/encode.rs b/crates/ecstore/src/erasure/coding/encode.rs index 6a51a17da..b0ce598c0 100644 --- a/crates/ecstore/src/erasure/coding/encode.rs +++ b/crates/ecstore/src/erasure/coding/encode.rs @@ -87,6 +87,48 @@ fn use_bytesmut_ingest() -> bool { rustfs_utils::get_env_bool(ENV_RUSTFS_ERASURE_ENCODE_BYTESMUT_INGEST, DEFAULT_RUSTFS_ERASURE_ENCODE_BYTESMUT_INGEST) }) } +/// Read up to `limit` bytes into `buf`'s uninitialized spare capacity, appending after its +/// current length, and distinguish a clean EOF from a short read. +/// +/// Mirrors `rustfs_utils::read_full_or_eof` semantics: returns `Ok(None)` when EOF is reached +/// before any byte is read, `Ok(Some(n))` once at least one byte is read, preserves +/// `InvalidData` errors (e.g. checksum mismatches) raised after a partial fill, and wraps other +/// partial-fill errors as `UnexpectedEof`. Unlike `resize` followed by a slice read, the spare +/// capacity is never zero-filled first, so full-block ingest skips a per-block memset. +async fn read_full_buf_or_eof(reader: &mut R, buf: &mut BytesMut, limit: usize) -> std::io::Result> +where + R: AsyncRead + Unpin, +{ + use bytes::BufMut as _; + use tokio::io::AsyncReadExt as _; + + debug_assert!(buf.is_empty(), "block ingest buffer must start empty"); + debug_assert!(limit > 0, "limit must be non-zero (block_size is validated upstream)"); + let mut total = 0; + while total < limit { + let mut limited = (&mut *buf).limit(limit - total); + let n = match reader.read_buf(&mut limited).await { + Ok(n) => n, + Err(e) => { + if total == 0 { + return Err(e); + } + // Preserve InvalidData (e.g. checksum mismatch) instead of wrapping it as + // UnexpectedEof, so proper error handling can occur upstream. + if e.kind() == std::io::ErrorKind::InvalidData { + return Err(e); + } + return Err(std::io::Error::new(std::io::ErrorKind::UnexpectedEof, e)); + } + }; + if n == 0 { + break; + } + total += n; + } + if total == 0 { Ok(None) } else { Ok(Some(total)) } +} + fn queued_block_bytes(block: &[Bytes]) -> usize { block.iter().map(Bytes::len).sum() } @@ -370,10 +412,28 @@ impl Erasure { #[cfg_attr(feature = "hotpath", hotpath::measure)] pub async fn encode( + self: Arc, + reader: R, + writers: &mut [Option], + quorum: usize, + ) -> std::io::Result<(R, usize)> + where + R: AsyncRead + Send + Sync + Unpin + 'static, + { + let use_bytesmut_ingest = use_bytesmut_ingest(); + self.encode_with_ingest_mode(reader, writers, quorum, use_bytesmut_ingest) + .await + } + + /// Streaming encode with an explicit ingest-buffer strategy. `encode` resolves the + /// strategy from `RUSTFS_ERASURE_ENCODE_BYTESMUT_INGEST`; tests call this directly to + /// exercise both paths regardless of the cached environment value. + async fn encode_with_ingest_mode( self: Arc, mut reader: R, writers: &mut [Option], quorum: usize, + use_bytesmut_ingest: bool, ) -> std::io::Result<(R, usize)> where R: AsyncRead + Send + Sync + Unpin + 'static, @@ -393,20 +453,24 @@ impl Erasure { let task = tokio::spawn(async move { let block_size = self.block_size; - let use_bytesmut_ingest = use_bytesmut_ingest(); let mut total = 0; if use_bytesmut_ingest { - let mut buf = BytesMut::with_capacity(block_size); - buf.resize(block_size, 0); + // HP-10 (rustfs/backlog#931): reserve the EC-expanded block size up front. + // Both shard-size formulas are monotone in data_len, so the resize to + // `need_total_size` inside `encode_data_bytes_mut` always stays within this + // capacity and never reallocates. Reading into uninitialized spare capacity + // (instead of resize + slice read) also skips zero-filling each fresh buffer. + let ingest_capacity = expanded_block_bytes.max(block_size); + let mut buf = BytesMut::with_capacity(ingest_capacity); loop { - match rustfs_utils::read_full_or_eof(&mut reader, &mut buf[..]).await { + match read_full_buf_or_eof(&mut reader, &mut buf, block_size).await { Ok(Some(n)) => { debug_assert!(n > 0, "non-zero block_size prevents zero-length reads"); + debug_assert_eq!(buf.len(), n, "ingest buffer length must equal bytes read"); total += n; let encode_buf = buf; let res = self.clone().encode_block_bytes_mut(encode_buf, n).await?; - buf = BytesMut::with_capacity(block_size); - buf.resize(block_size, 0); + buf = BytesMut::with_capacity(ingest_capacity); let queued_bytes = queued_block_bytes(&res); rustfs_io_metrics::add_ec_encode_inflight_bytes(queued_bytes); let send_wait_stage_start = stage_timer_if_enabled(); @@ -1157,6 +1221,103 @@ mod tests { } } + #[tokio::test] + async fn read_full_buf_or_eof_returns_none_on_empty_reader() { + let mut reader = Cursor::new(Vec::::new()); + let mut buf = BytesMut::with_capacity(16); + + let res = read_full_buf_or_eof(&mut reader, &mut buf, 8).await.unwrap(); + + assert_eq!(res, None); + assert!(buf.is_empty()); + } + + #[tokio::test] + async fn read_full_buf_or_eof_reads_partial_tail() { + let data = b"tail".to_vec(); + let mut reader = Cursor::new(data.clone()); + let mut buf = BytesMut::with_capacity(64); + + let res = read_full_buf_or_eof(&mut reader, &mut buf, 16).await.unwrap(); + + assert_eq!(res, Some(data.len())); + assert_eq!(&buf[..], &data[..]); + assert!(buf.capacity() >= 64, "pre-reserved spare capacity must be kept"); + } + + #[tokio::test] + async fn read_full_buf_or_eof_stops_at_limit() { + let data = vec![7u8; 32]; + let mut reader = Cursor::new(data.clone()); + let mut buf = BytesMut::with_capacity(64); + + let res = read_full_buf_or_eof(&mut reader, &mut buf, 16).await.unwrap(); + + assert_eq!(res, Some(16)); + assert_eq!(&buf[..], &data[..16]); + + // The remaining bytes are still readable as the next block. + let mut next = BytesMut::with_capacity(64); + let res = read_full_buf_or_eof(&mut reader, &mut next, 16).await.unwrap(); + assert_eq!(res, Some(16)); + assert_eq!(&next[..], &data[16..]); + } + + async fn committed_shards_for_ingest_mode(use_bytesmut_ingest: bool, uses_legacy: bool, payload: &[u8]) -> Vec> { + const DATA_SHARDS: usize = 2; + const PARITY_SHARDS: usize = 2; + const TOTAL_SHARDS: usize = DATA_SHARDS + PARITY_SHARDS; + const BLOCK_SIZE: usize = 64; + + let committed: Vec>>> = (0..TOTAL_SHARDS).map(|_| Arc::new(Mutex::new(Vec::new()))).collect(); + let mut writers: Vec> = committed + .iter() + .map(|c| Some(bitrot_writer(DeferredCommitWriter::new(c.clone()), BLOCK_SIZE / DATA_SHARDS))) + .collect(); + + let erasure = Arc::new(Erasure::new_with_options(DATA_SHARDS, PARITY_SHARDS, BLOCK_SIZE, uses_legacy)); + let reader = tokio::io::BufReader::new(Cursor::new(payload.to_vec())); + let (_reader, total) = erasure + .encode_with_ingest_mode(reader, &mut writers, DATA_SHARDS, use_bytesmut_ingest) + .await + .expect("encode should succeed"); + assert_eq!(total, payload.len()); + + committed + .iter() + .map(|c| c.lock().expect("committed buffer should be lockable").clone()) + .collect() + } + + /// HP-10 (rustfs/backlog#931) merge gate: the BytesMut ingest path must produce + /// byte-for-byte identical shard streams to the default Vec ingest path, for both + /// legacy-aware shard-size formulas, across empty, sub-block, exactly-full-block, + /// and multi-block-with-partial-tail payloads. + #[tokio::test] + async fn bytesmut_ingest_matches_vec_ingest_byte_for_byte() { + const BLOCK_SIZE: usize = 64; + let payloads: Vec> = vec![ + Vec::new(), + b"tiny".to_vec(), + (0..BLOCK_SIZE as u32).map(|i| i as u8).collect(), // exactly one full block + vec![3u8; BLOCK_SIZE * 4], // whole number of blocks + (0..(BLOCK_SIZE * 3 + 7) as u32).map(|i| (i % 251) as u8).collect(), // partial tail + ]; + + for uses_legacy in [false, true] { + for payload in &payloads { + let vec_path = committed_shards_for_ingest_mode(false, uses_legacy, payload).await; + let bytesmut_path = committed_shards_for_ingest_mode(true, uses_legacy, payload).await; + assert_eq!( + vec_path, + bytesmut_path, + "ingest paths must be byte-identical (legacy={uses_legacy}, payload_len={})", + payload.len() + ); + } + } + } + #[test] fn encode_channel_capacity_never_returns_zero() { assert_eq!(encode_channel_capacity(0, 1024), 1); diff --git a/crates/ecstore/src/erasure/coding/erasure.rs b/crates/ecstore/src/erasure/coding/erasure.rs index 5894010d1..93e5c9455 100644 --- a/crates/ecstore/src/erasure/coding/erasure.rs +++ b/crates/ecstore/src/erasure/coding/erasure.rs @@ -610,6 +610,12 @@ impl Erasure { /// Encode data from an owned `BytesMut` buffer, avoiding the initial copy /// from a borrowed slice into a fresh `BytesMut`. + /// + /// Capacity contract: when the caller pre-reserves + /// `shard_size() * total_shard_count()` bytes (the EC-expanded size of a full + /// block), the `resize(need_total_size)` below stays within capacity for every + /// `data_len <= block_size` — both shard-size formulas are monotone in + /// `data_len` — so this function never reallocates the buffer. #[cfg_attr(feature = "hotpath", hotpath::measure)] pub fn encode_data_bytes_mut(&self, mut data_buffer: BytesMut, data_len: usize) -> io::Result> { let shard_size_fn = if self.uses_legacy { @@ -1008,13 +1014,68 @@ mod tests { fn encode_data_bytes_mut_matches_borrowed_path() { for uses_legacy in [false, true] { let erasure = Erasure::new_with_options(4, 2, 64, uses_legacy); - for data in [Vec::new(), b"small payload".to_vec(), (0_u8..37).collect::>()] { + let block_size = erasure.block_size; + let expanded_block_bytes = erasure.shard_size() * erasure.total_shard_count(); + let cases: Vec> = vec![ + Vec::new(), + b"small payload".to_vec(), + (0_u8..37).collect(), + vec![0xA5; block_size - 1], // last block one byte short of full + vec![0x5A; block_size], // exactly one full block + ]; + for data in cases { let borrowed = erasure.encode_data(&data).expect("borrowed encode should succeed"); - let bytes_mut = BytesMut::from(&data[..]); + let owned = erasure - .encode_data_bytes_mut(bytes_mut, data.len()) + .encode_data_bytes_mut(BytesMut::from(&data[..]), data.len()) .expect("bytesmut encode should succeed"); assert_eq!(owned, borrowed); + + // Ingest-shaped buffer: spare capacity pre-reserved for the EC-expanded + // block, exactly what the HP-10 streaming ingest path hands over. + let mut ingest = BytesMut::with_capacity(expanded_block_bytes.max(block_size)); + ingest.extend_from_slice(&data); + let preallocated = erasure + .encode_data_bytes_mut(ingest, data.len()) + .expect("preallocated bytesmut encode should succeed"); + assert_eq!(preallocated, borrowed); + + // Buffer longer than data_len (stale bytes past the logical block) + // must be truncated before padding. + let mut oversized = BytesMut::from(&data[..]); + oversized.extend_from_slice(&[0xFF; 8]); + let truncated = erasure + .encode_data_bytes_mut(oversized, data.len()) + .expect("oversized bytesmut encode should succeed"); + assert_eq!(truncated, borrowed); + } + } + } + + /// HP-10 capacity invariant: both shard-size formulas are monotone in `data_len`, + /// so pre-reserving `shard_size(block_size) * total_shard_count` covers the + /// `need_total_size` of every block-or-smaller payload and the ingest buffer + /// never reallocates inside `encode_data_bytes_mut`. + #[test] + fn shard_size_monotonicity_bounds_expanded_block_capacity() { + for shard_fn in [calc_shard_size, calc_shard_size_legacy] { + for data_shards in 1usize..=16 { + for block_size in [1usize, 2, 63, 64, 65, 1024, 4096] { + let full = shard_fn(block_size, data_shards); + let mut prev = shard_fn(0, data_shards); + for data_len in 1..=block_size { + let cur = shard_fn(data_len, data_shards); + assert!(cur >= prev, "shard size must be monotone (len {data_len}, shards {data_shards})"); + assert!(cur <= full, "per-block shard size must not exceed the full-block bound"); + prev = cur; + } + } + // Spot-check the production-scale block size at its boundaries. + let block_size = 1usize << 20; + let full = shard_fn(block_size, data_shards); + for data_len in [0usize, 1, block_size / 2, block_size - 1, block_size] { + assert!(shard_fn(data_len, data_shards) <= full); + } } } }