diff --git a/crates/ecstore/src/erasure/coding/encode.rs b/crates/ecstore/src/erasure/coding/encode.rs index 4b03ad60e..e600603d8 100644 --- a/crates/ecstore/src/erasure/coding/encode.rs +++ b/crates/ecstore/src/erasure/coding/encode.rs @@ -18,6 +18,7 @@ use crate::disk::error_reduce::{ }; use crate::erasure::coding::BitrotWriterWrapper; use crate::erasure::coding::Erasure; +use crate::erasure::coding::erasure::EncodedBlock; use crate::runtime::sources as runtime_sources; use bytes::{Bytes, BytesMut}; use futures::StreamExt; @@ -223,8 +224,8 @@ async fn send_queued( sender.send(InflightEntry::new(entry, bytes)).await } -fn queued_batch_bytes(batch: &[Vec]) -> usize { - batch.iter().map(|block| queued_block_bytes(block)).sum() +fn queued_batch_bytes(batch: &[EncodedBlock]) -> usize { + batch.iter().map(EncodedBlock::queued_bytes).sum() } fn dominant_error_summary_label(summary: &WriteQuorumFailureSummary) -> &'static str { @@ -336,7 +337,7 @@ impl<'a> MultiWriter<'a> { } } - async fn write_shard(writer_opt: &mut Option, err: &mut Option, shard: &Bytes) { + async fn write_shard(writer_opt: &mut Option, err: &mut Option, shard: &[u8]) { match writer_opt { Some(writer) => { match writer.write(shard).await { @@ -361,12 +362,20 @@ impl<'a> MultiWriter<'a> { } pub async fn write(&mut self, data: Vec) -> std::io::Result<()> { - assert_eq!(data.len(), self.writers.len()); + self.write_shards(data.iter().map(Bytes::as_ref)).await + } + + async fn write_block(&mut self, block: &EncodedBlock) -> std::io::Result<()> { + self.write_shards(block.shards()).await + } + + async fn write_shards<'b>(&mut self, shards: impl ExactSizeIterator) -> std::io::Result<()> { + assert_eq!(shards.len(), self.writers.len()); let budget = self.next_progress_budget(); { let mut futures = FuturesUnordered::new(); - for ((writer_opt, err), shard) in self.writers.iter_mut().zip(self.errs.iter_mut()).zip(data.iter()) { + for ((writer_opt, err), shard) in self.writers.iter_mut().zip(self.errs.iter_mut()).zip(shards) { if err.is_some() { continue; // Skip if we already have an error for this writer } @@ -490,10 +499,10 @@ impl<'a> MultiWriter<'a> { } impl Erasure { - async fn encode_block(self: Arc, encode_buf: Vec, len: usize) -> std::io::Result<(Vec, Vec)> { + async fn encode_block(self: Arc, encode_buf: Vec, len: usize) -> std::io::Result<(EncodedBlock, Vec)> { let encode_stage_start = stage_timer_if_enabled(); let encode_once = move || { - let res = self.encode_data(&encode_buf[..len]); + let res = self.encode_data_block(&encode_buf[..len]); (res, encode_buf) }; @@ -518,9 +527,9 @@ impl Erasure { Ok((res?, returned_buf)) } - async fn encode_block_bytes_mut(self: Arc, encode_buf: BytesMut, len: usize) -> std::io::Result> { + async fn encode_block_bytes_mut(self: Arc, encode_buf: BytesMut, len: usize) -> std::io::Result { let encode_stage_start = stage_timer_if_enabled(); - let encode_once = move || self.encode_data_bytes_mut(encode_buf, len); + let encode_once = move || self.encode_data_bytes_mut_block(encode_buf, len); let res = match tokio::runtime::Handle::current().runtime_flavor() { // Same rationale as encode_block: inline the short EC burst on the @@ -624,7 +633,7 @@ impl Erasure { let expanded_block_bytes = self.shard_size().saturating_mul(self.total_shard_count()); let max_inflight_bytes = erasure_encode_max_inflight_bytes(); let inflight_blocks = encode_channel_capacity(expanded_block_bytes, max_inflight_bytes); - let (tx, mut rx) = mpsc::channel::>>(inflight_blocks); + let (tx, mut rx) = mpsc::channel::>(inflight_blocks); let mut task = AbortOnDropTask::new(tokio::spawn(async move { let block_size = self.block_size; @@ -646,7 +655,7 @@ impl Erasure { let encode_buf = buf; let res = self.clone().encode_block_bytes_mut(encode_buf, n).await?; buf = BytesMut::with_capacity(ingest_capacity); - let queued_bytes = queued_block_bytes(&res); + let queued_bytes = res.queued_bytes(); let _producer_stage = rustfs_io_metrics::track_ec_encode_producer_bytes(queued_bytes); let send_wait_stage_start = stage_timer_if_enabled(); if let Err(err) = send_queued(&tx, res, queued_bytes).await { @@ -676,7 +685,7 @@ impl Erasure { let encode_buf = std::mem::take(&mut buf); let (res, returned_buf) = self.clone().encode_block(encode_buf, n).await?; buf = returned_buf; - let queued_bytes = queued_block_bytes(&res); + let queued_bytes = res.queued_bytes(); let _producer_stage = rustfs_io_metrics::track_ec_encode_producer_bytes(queued_bytes); let send_wait_stage_start = stage_timer_if_enabled(); if let Err(err) = send_queued(&tx, res, queued_bytes).await { @@ -720,9 +729,9 @@ impl Erasure { if block.is_empty() { break; } - let _writer_stage = rustfs_io_metrics::track_ec_encode_writer_bytes(queued_block_bytes(&block)); + let _writer_stage = rustfs_io_metrics::track_ec_encode_writer_bytes(block.queued_bytes()); let write_stage_start = stage_timer_if_enabled(); - if let Err(err) = writers.write(block).await { + if let Err(err) = writers.write_block(&block).await { write_err = Some(err); break; } @@ -769,7 +778,7 @@ impl Erasure { let inflight_blocks = encode_channel_capacity(expanded_block_bytes, max_inflight_bytes); let batch_blocks = encode_batch_block_count().min(inflight_blocks); let channel_capacity = inflight_blocks.div_ceil(batch_blocks).max(1); - let (tx, mut rx) = mpsc::channel::>>>(channel_capacity); + let (tx, mut rx) = mpsc::channel::>>(channel_capacity); let mut task = AbortOnDropTask::new(tokio::spawn(async move { let block_size = self.block_size; @@ -786,7 +795,7 @@ impl Erasure { let encode_buf = std::mem::take(&mut buf); let (res, returned_buf) = self.clone().encode_block(encode_buf, n).await?; buf = returned_buf; - let queued_bytes = queued_block_bytes(&res); + let queued_bytes = res.queued_bytes(); pending_batch_bytes = pending_batch_bytes.saturating_add(queued_bytes); pending_batch.push(res); drop(pending_batch_stage.take()); @@ -845,7 +854,7 @@ impl Erasure { let _writer_stage = rustfs_io_metrics::track_ec_encode_writer_bytes(queued_batch_bytes(&batch)); let write_stage_start = stage_timer_if_enabled(); for block in batch { - if let Err(err) = writers.write(block).await { + if let Err(err) = writers.write_block(&block).await { write_err = Some(err); break; } @@ -1895,7 +1904,11 @@ mod tests { let baseline = rustfs_io_metrics::current_ec_encode_inflight_bytes(); let (tx, rx) = mpsc::channel(2); let mut rx = rx; - let batch = vec![vec![Bytes::from_static(b"queued")], vec![Bytes::from_static(b"batch")]]; + let erasure = Erasure::new(1, 0, 16); + let batch = vec![ + erasure.encode_data_block(b"queued").expect("first block should encode"), + erasure.encode_data_block(b"batch").expect("second block should encode"), + ]; let batch_bytes = queued_batch_bytes(&batch); send_queued(&tx, batch, batch_bytes).await.expect("batch should be queued"); @@ -2236,11 +2249,11 @@ mod tests { .expect("bytesmut encode should succeed on current-thread runtime"); let expected_shard_size = payload.len().div_ceil(erasure.data_shards); - assert_eq!(shards.len(), erasure.total_shard_count()); - assert!(shards.iter().all(|shard| shard.len() == expected_shard_size)); + assert_eq!(shards.shards().len(), erasure.total_shard_count()); + assert!(shards.shards().all(|shard| shard.len() == expected_shard_size)); let mut restored = Vec::new(); - for shard in shards.iter().take(erasure.data_shards) { + for shard in shards.shards().take(erasure.data_shards) { restored.extend_from_slice(shard); } restored.truncate(payload.len()); @@ -2506,7 +2519,7 @@ mod tests { assert_eq!(&next[..], &data[16..]); } - async fn committed_shards_for_ingest_mode(use_bytesmut_ingest: bool, uses_legacy: bool, payload: &[u8]) -> Vec> { + async fn committed_shards_for_pipeline(pipeline: EncodePipeline, uses_legacy: bool, payload: &[u8]) -> Vec> { const DATA_SHARDS: usize = 2; const PARITY_SHARDS: usize = 2; const TOTAL_SHARDS: usize = DATA_SHARDS + PARITY_SHARDS; @@ -2520,10 +2533,16 @@ mod tests { 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"); + let (_reader, total) = match pipeline { + EncodePipeline::Vec => { + erasure + .encode_with_ingest_mode(reader, &mut writers, DATA_SHARDS, false) + .await + } + EncodePipeline::BytesMut => erasure.encode_with_ingest_mode(reader, &mut writers, DATA_SHARDS, true).await, + EncodePipeline::Batched => erasure.encode_batched(reader, &mut writers, DATA_SHARDS).await, + } + .expect("encode should succeed"); assert_eq!(total, payload.len()); committed @@ -2532,31 +2551,64 @@ mod tests { .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. + async fn expected_committed_shards(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| bitrot_writer(DeferredCommitWriter::new(c.clone()), BLOCK_SIZE / DATA_SHARDS)) + .collect(); + let erasure = Erasure::new_with_options(DATA_SHARDS, PARITY_SHARDS, BLOCK_SIZE, uses_legacy); + + for block in payload.chunks(BLOCK_SIZE) { + let shards = erasure.encode_data(block).expect("reference block should encode"); + for (writer, shard) in writers.iter_mut().zip(shards) { + let written = writer.write(&shard).await.expect("reference shard should write"); + assert_eq!(written, shard.len()); + } + } + for writer in &mut writers { + writer.shutdown().await.expect("reference writer should commit"); + } + + committed + .iter() + .map(|c| c.lock().expect("committed buffer should be lockable").clone()) + .collect() + } + + /// The streaming and batched paths must produce the same bitrot-wrapped shard + /// bytes as the public block encoder for both shard-size formulas and all block + /// boundary shapes. #[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(), + vec![1], + vec![2; BLOCK_SIZE - 1], (0..BLOCK_SIZE as u32).map(|i| i as u8).collect(), // exactly one full block - vec![3u8; BLOCK_SIZE * 4], // whole number of blocks + vec![4; BLOCK_SIZE + 1], + 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() - ); + let expected = expected_committed_shards(uses_legacy, payload).await; + for pipeline in [EncodePipeline::Vec, EncodePipeline::BytesMut, EncodePipeline::Batched] { + let actual = committed_shards_for_pipeline(pipeline, uses_legacy, payload).await; + assert_eq!( + actual, + expected, + "streaming shards must match the public block encoder (legacy={uses_legacy}, payload_len={})", + payload.len() + ); + } } } } diff --git a/crates/ecstore/src/erasure/coding/erasure.rs b/crates/ecstore/src/erasure/coding/erasure.rs index 8b4e213b2..23cac93c7 100644 --- a/crates/ecstore/src/erasure/coding/erasure.rs +++ b/crates/ecstore/src/erasure/coding/erasure.rs @@ -29,6 +29,46 @@ use tokio::io::AsyncRead; use tracing::warn; use uuid::Uuid; +pub(crate) struct EncodedBlock { + data: Bytes, + shard_size: usize, +} + +impl EncodedBlock { + fn empty() -> Self { + Self { + data: Bytes::new(), + shard_size: 0, + } + } + + pub(crate) fn is_empty(&self) -> bool { + self.data.is_empty() + } + + pub(crate) fn queued_bytes(&self) -> usize { + self.data.len() + } + + pub(crate) fn shards(&self) -> impl ExactSizeIterator { + debug_assert!(self.shard_size > 0, "only non-empty encoded blocks reach shard writers"); + debug_assert_eq!(self.data.len() % self.shard_size, 0); + self.data.chunks_exact(self.shard_size) + } + + fn into_shards(mut self, shard_count: usize) -> Vec { + if self.shard_size == 0 { + return vec![Bytes::new(); shard_count]; + } + + let mut shards = Vec::with_capacity(shard_count); + for _ in 0..shard_count { + shards.push(self.data.split_to(self.shard_size)); + } + shards + } +} + const MODERN_MAX_TOTAL_SHARDS: usize = ::ORDER; const MODERN_REED_SOLOMON_CACHE_MAX_ENTRIES: usize = 64; @@ -675,6 +715,17 @@ impl Erasure { #[tracing::instrument(level = "debug", skip_all, fields(data_len=data.len()))] #[hotpath::measure(impl_type = "Erasure")] pub fn encode_data(&self, data: &[u8]) -> io::Result> { + self.encode_data_block_inner(data) + .map(|block| block.into_shards(self.total_shard_count())) + } + + #[tracing::instrument(level = "debug", skip_all, fields(data_len=data.len()))] + #[hotpath::measure(label = "Erasure::encode_data", impl_type = "Erasure")] + pub(crate) fn encode_data_block(&self, data: &[u8]) -> io::Result { + self.encode_data_block_inner(data) + } + + fn encode_data_block_inner(&self, data: &[u8]) -> io::Result { let shard_size_fn = if self.uses_legacy { calc_shard_size_legacy } else { @@ -682,7 +733,7 @@ impl Erasure { }; let per_shard_size = shard_size_fn(data.len(), self.data_shards); if per_shard_size == 0 { - return Ok(vec![Bytes::new(); self.total_shard_count()]); + return Ok(EncodedBlock::empty()); } let need_total_size = per_shard_size * self.total_shard_count(); @@ -708,15 +759,10 @@ impl Erasure { } } - // Zero-copy split, all shards reference data_buffer - let mut data_buffer = data_buffer.freeze(); - let mut shards = Vec::with_capacity(self.total_shard_count()); - for _ in 0..self.total_shard_count() { - let shard = data_buffer.split_to(per_shard_size); - shards.push(shard); - } - - Ok(shards) + Ok(EncodedBlock { + data: data_buffer.freeze(), + shard_size: per_shard_size, + }) } /// Encode owned data, avoiding a copy when the caller already has a heap buffer. @@ -786,7 +832,17 @@ impl Erasure { /// `data_len <= block_size` — both shard-size formulas are monotone in /// `data_len` — so this function never reallocates the buffer. #[hotpath::measure(impl_type = "Erasure")] - pub fn encode_data_bytes_mut(&self, mut data_buffer: BytesMut, data_len: usize) -> io::Result> { + pub fn encode_data_bytes_mut(&self, data_buffer: BytesMut, data_len: usize) -> io::Result> { + self.encode_data_bytes_mut_block_inner(data_buffer, data_len) + .map(|block| block.into_shards(self.total_shard_count())) + } + + #[hotpath::measure(label = "Erasure::encode_data_bytes_mut", impl_type = "Erasure")] + pub(crate) fn encode_data_bytes_mut_block(&self, data_buffer: BytesMut, data_len: usize) -> io::Result { + self.encode_data_bytes_mut_block_inner(data_buffer, data_len) + } + + fn encode_data_bytes_mut_block_inner(&self, mut data_buffer: BytesMut, data_len: usize) -> io::Result { let shard_size_fn = if self.uses_legacy { calc_shard_size_legacy } else { @@ -794,7 +850,7 @@ impl Erasure { }; let per_shard_size = shard_size_fn(data_len, self.data_shards); if per_shard_size == 0 { - return Ok(vec![Bytes::new(); self.total_shard_count()]); + return Ok(EncodedBlock::empty()); } let need_total_size = per_shard_size * self.total_shard_count(); @@ -821,14 +877,10 @@ impl Erasure { } } - let mut data_buffer = data_buffer.freeze(); - let mut shards = Vec::with_capacity(self.total_shard_count()); - for _ in 0..self.total_shard_count() { - let shard = data_buffer.split_to(per_shard_size); - shards.push(shard); - } - - Ok(shards) + Ok(EncodedBlock { + data: data_buffer.freeze(), + shard_size: per_shard_size, + }) } /// Decode and reconstruct missing data shards in-place. @@ -1547,6 +1599,37 @@ mod tests { } } + #[test] + fn streaming_encoded_block_uses_one_contiguous_backing_buffer() { + let erasure = Erasure::new(8, 8, 64); + + for data_len in [1, 63, 64] { + let data = (0..data_len).map(|i| i as u8).collect::>(); + let expected = erasure.encode_data(&data).expect("public encode should succeed"); + let borrowed = erasure + .encode_data_block(&data) + .expect("borrowed streaming encode should succeed"); + let owned = erasure + .encode_data_bytes_mut_block(BytesMut::from(&data[..]), data.len()) + .expect("BytesMut streaming encode should succeed"); + + assert!(borrowed.shards().eq(expected.iter().map(Bytes::as_ref))); + assert!(owned.shards().eq(expected.iter().map(Bytes::as_ref))); + assert_eq!(borrowed.shards().len(), 16); + assert_eq!(borrowed.queued_bytes(), owned.queued_bytes()); + + let first = borrowed.shards().next().expect("encoded block should have shards").as_ptr(); + for (index, shard) in borrowed.shards().enumerate() { + assert_eq!(shard.as_ptr(), first.wrapping_add(index * shard.len())); + } + } + assert_eq!( + std::mem::size_of::(), + std::mem::size_of::() + std::mem::size_of::(), + "queue entries must contain one backing buffer handle, not per-shard handles" + ); + } + /// 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