mirror of
https://github.com/rustfs/rustfs.git
synced 2026-09-06 03:59:14 +00:00
Zhengchao An
1dd81cf276
Add `tee_reader` / `tee_reader_with_options` in `rustfs-rio`: a `TeePrimary` that drives the source and a `TeeSecondary` that observes an identical copy of every chunk through a byte-bounded queue. The primary returns `Pending` when the queue is full, so both sides advance at the pace of the slowest consumer; it is meant for small objects only. Termination: source EOF and errors propagate to the secondary with the same `io::ErrorKind`; dropping the secondary turns the primary into a pass-through; dropping the primary early fails the secondary with `BrokenPipe` by default, or hands the remaining source to a background drain task bounded by `max_drain_bytes` when `TeeOptions::drain_on_primary_drop` is set. `TeeSecondary::into_stream` exposes the queued `Bytes` chunks without an extra copy. Includes a proptest equivalence test, backpressure, error, drop, drain-limit and cancel-safety tests, and a criterion bench comparing tee throughput against a direct read (64 MiB in 1 MiB chunks).
104 lines
3.3 KiB
Rust
104 lines
3.3 KiB
Rust
// Copyright 2024 RustFS Team
|
|
//
|
|
// Licensed 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.
|
|
|
|
//! Throughput of `tee_reader` versus reading the same source directly:
|
|
//! 64 MiB of data served in 1 MiB chunks, consumed with 1 MiB reads.
|
|
|
|
use bytes::Bytes;
|
|
use criterion::{Criterion, Throughput, criterion_group, criterion_main};
|
|
use rustfs_rio::tee_reader;
|
|
use std::hint::black_box;
|
|
use std::io;
|
|
use std::pin::Pin;
|
|
use std::task::{Context, Poll};
|
|
use tokio::io::{AsyncRead, AsyncReadExt, ReadBuf};
|
|
|
|
const CHUNK_BYTES: usize = 1024 * 1024;
|
|
const TOTAL_BYTES: usize = 64 * 1024 * 1024;
|
|
const TEE_BUFFER_BYTES: usize = 4 * CHUNK_BYTES;
|
|
|
|
/// In-memory source that serves at most `CHUNK_BYTES` per poll.
|
|
struct ChunkedSource {
|
|
data: Bytes,
|
|
pos: usize,
|
|
}
|
|
|
|
impl AsyncRead for ChunkedSource {
|
|
fn poll_read(mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
|
|
let remaining = self.data.len() - self.pos;
|
|
let n = CHUNK_BYTES.min(remaining).min(buf.remaining());
|
|
buf.put_slice(&self.data[self.pos..self.pos + n]);
|
|
self.pos += n;
|
|
Poll::Ready(Ok(()))
|
|
}
|
|
}
|
|
|
|
async fn consume<R: AsyncRead + Unpin>(mut reader: R) -> usize {
|
|
let mut buf = vec![0u8; CHUNK_BYTES];
|
|
let mut total = 0;
|
|
loop {
|
|
let n = reader.read(&mut buf).await.expect("read");
|
|
if n == 0 {
|
|
return total;
|
|
}
|
|
total += n;
|
|
}
|
|
}
|
|
|
|
fn bench_tee_reader(c: &mut Criterion) {
|
|
let runtime = tokio::runtime::Builder::new_multi_thread()
|
|
.worker_threads(2)
|
|
.enable_all()
|
|
.build()
|
|
.expect("build tokio runtime for tee_reader benchmark");
|
|
let data = Bytes::from(vec![0xA5u8; TOTAL_BYTES]);
|
|
|
|
let mut group = c.benchmark_group("tee_reader_64mib_1mib_chunks");
|
|
group.throughput(Throughput::Bytes(TOTAL_BYTES as u64));
|
|
group.sample_size(10);
|
|
|
|
group.bench_function("direct_read", |b| {
|
|
b.iter(|| {
|
|
let source = ChunkedSource {
|
|
data: data.clone(),
|
|
pos: 0,
|
|
};
|
|
let total = runtime.block_on(consume(source));
|
|
black_box(total)
|
|
})
|
|
});
|
|
|
|
group.bench_function("tee_primary_plus_secondary", |b| {
|
|
b.iter(|| {
|
|
let source = ChunkedSource {
|
|
data: data.clone(),
|
|
pos: 0,
|
|
};
|
|
let (primary, secondary) = tee_reader(source, TEE_BUFFER_BYTES);
|
|
let totals = runtime.block_on(async {
|
|
let secondary_task = tokio::spawn(consume(secondary));
|
|
let primary_total = consume(primary).await;
|
|
let secondary_total = secondary_task.await.expect("secondary task");
|
|
(primary_total, secondary_total)
|
|
});
|
|
black_box(totals)
|
|
})
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
criterion_group!(benches, bench_tee_reader);
|
|
criterion_main!(benches);
|