feat(rio): add bounded tee reader for one-read two-consumer paths (#7066)

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).

Cargo.lock

@@ -10449,6 +10449,7 @@ dependencies = [
  "base64-simd",
  "bytes",
  "crc-fast",
+ "criterion",
  "faster-hex",
  "futures",
  "hex-simd",
@@ -10460,6 +10461,7 @@ dependencies = [
  "md-5 0.11.0",
  "minlz",
  "pin-project-lite",
+ "proptest",
  "rand 0.10.2",
  "reqwest",
  "rustfs-config",

crates/rio/Cargo.toml

@@ -90,8 +90,15 @@ s3s = { workspace = true, features = ["minio"] }
 hex-simd.workspace = true
 
 [dev-dependencies]
+tokio = { workspace = true, features = ["test-util"] }
 tokio-test = { workspace = true }
+criterion = { workspace = true, features = ["html_reports"] }
+proptest = { workspace = true }
 axum = { workspace = true }
 hyper = { workspace = true, features = ["http2", "server"] }
 hyper-util = { workspace = true, features = ["tokio"] }
 http-body-util = { workspace = true }
+
+[[bench]]
+name = "tee_reader"
+harness = false

crates/rio/benches/tee_reader.rs

@@ -0,0 +1,103 @@
+// 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);

crates/rio/src/lib.rs

@@ -118,6 +118,12 @@ pub use hardlimit_reader::HardLimitReader;
 
 mod hash_reader;
 pub use hash_reader::*;
+
+mod tee_reader;
+pub use tee_reader::{
+    DEFAULT_TEE_MAX_DRAIN_BYTES, TeeDrainLimitExceeded, TeeOptions, TeePrimary, TeeSecondary, TeeStream, tee_reader,
+    tee_reader_with_options,
+};
 mod checksum;
 pub use checksum::*;