perf(ecstore): avoid per-block shard vector allocation (#6037)

Keep encoded shards in one contiguous Bytes buffer while they cross the streaming write queue, and materialize Vec<Bytes> only for the existing public APIs.

Co-authored-by: heihutu <heihutu@gmail.com>
This commit is contained in:
houseme
2026-08-13 09:55:24 +08:00
committed by GitHub
parent a5594c3d89
commit f21e88b112
2 changed files with 195 additions and 60 deletions
+92 -40
View File
@@ -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<T>(
sender.send(InflightEntry::new(entry, bytes)).await
}
fn queued_batch_bytes(batch: &[Vec<Bytes>]) -> 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<BitrotWriterWrapper>, err: &mut Option<Error>, shard: &Bytes) {
async fn write_shard(writer_opt: &mut Option<BitrotWriterWrapper>, err: &mut Option<Error>, 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<Bytes>) -> 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<Item = &'b [u8]>) -> 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<Self>, encode_buf: Vec<u8>, len: usize) -> std::io::Result<(Vec<Bytes>, Vec<u8>)> {
async fn encode_block(self: Arc<Self>, encode_buf: Vec<u8>, len: usize) -> std::io::Result<(EncodedBlock, Vec<u8>)> {
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<Self>, encode_buf: BytesMut, len: usize) -> std::io::Result<Vec<Bytes>> {
async fn encode_block_bytes_mut(self: Arc<Self>, encode_buf: BytesMut, len: usize) -> std::io::Result<EncodedBlock> {
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::<InflightEntry<Vec<Bytes>>>(inflight_blocks);
let (tx, mut rx) = mpsc::channel::<InflightEntry<EncodedBlock>>(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::<InflightEntry<Vec<Vec<Bytes>>>>(channel_capacity);
let (tx, mut rx) = mpsc::channel::<InflightEntry<Vec<EncodedBlock>>>(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<Vec<u8>> {
async fn committed_shards_for_pipeline(pipeline: EncodePipeline, uses_legacy: bool, payload: &[u8]) -> Vec<Vec<u8>> {
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<Vec<u8>> {
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<Arc<Mutex<Vec<u8>>>> = (0..TOTAL_SHARDS).map(|_| Arc::new(Mutex::new(Vec::new()))).collect();
let mut writers: Vec<BitrotWriterWrapper> = 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<u8>> = 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()
);
}
}
}
}
+103 -20
View File
@@ -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<Item = &[u8]> {
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<Bytes> {
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 = <reed_solomon_erasure::galois_8::Field as reed_solomon_erasure::Field>::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<Vec<Bytes>> {
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<EncodedBlock> {
self.encode_data_block_inner(data)
}
fn encode_data_block_inner(&self, data: &[u8]) -> io::Result<EncodedBlock> {
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<Vec<Bytes>> {
pub fn encode_data_bytes_mut(&self, data_buffer: BytesMut, data_len: usize) -> io::Result<Vec<Bytes>> {
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<EncodedBlock> {
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<EncodedBlock> {
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::<Vec<_>>();
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::<EncodedBlock>(),
std::mem::size_of::<Bytes>() + std::mem::size_of::<usize>(),
"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