perf(ecstore): retain remote shard HTTP chunks (#5991)

* perf(ecstore): retain remote shard HTTP chunks

* fix(ecstore): bound remote shard chunk retention

* fix(rio): persist empty chunk limit across polls
This commit is contained in:
GatewayJ
2026-08-13 15:00:44 +08:00
committed by GitHub
parent e11fcfbd08
commit 36deab8670
8 changed files with 1183 additions and 30 deletions
@@ -31,7 +31,7 @@ use rustfs_config::{
DEFAULT_INTERNODE_DATA_TRANSPORT, ENV_RUSTFS_INTERNODE_DATA_TRANSPORT, INTERNODE_DATA_TRANSPORT_TCP,
KNOWN_INTERNODE_DATA_TRANSPORT_BACKENDS,
};
use rustfs_rio::{HttpReader, HttpWriter};
use rustfs_rio::{ChunkReaderBox, HttpChunkReader, HttpReader, HttpWriter};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::future::Future;
@@ -221,6 +221,11 @@ pub struct NsScannerCapabilityRequest {
#[async_trait]
pub trait InternodeDataTransport: Send + Sync + std::fmt::Debug {
async fn open_read(&self, request: ReadStreamRequest) -> Result<FileReader>;
/// Opens an owned-chunk stream when this transport can retain receive-buffer
/// ownership. `None` preserves the established `open_read` fallback.
async fn open_read_chunks(&self, _request: ReadStreamRequest) -> Result<Option<ChunkReaderBox>> {
Ok(None)
}
async fn open_write(&self, request: WriteStreamRequest) -> Result<FileWriter>;
async fn open_walk_dir(&self, request: WalkDirStreamRequest) -> Result<FileReader>;
async fn open_ns_scanner(&self, _request: NsScannerStreamRequest) -> Result<FileReader> {
@@ -247,6 +252,15 @@ impl InternodeDataTransport for TcpHttpInternodeDataTransport {
))
}
async fn open_read_chunks(&self, request: ReadStreamRequest) -> Result<Option<ChunkReaderBox>> {
let url = build_read_file_stream_url(&request);
let mut headers = json_headers();
build_auth_headers(&url, &Method::GET, &mut headers)?;
Ok(Some(Box::new(
HttpChunkReader::new_with_stall_timeout(url, Method::GET, headers, None, request.stall_timeout).await?,
)))
}
async fn open_write(&self, request: WriteStreamRequest) -> Result<FileWriter> {
let server_epoch = self.put_file_auth_capability(&request.endpoint).await?;
let nonce = server_epoch.map(|_| Uuid::new_v4());
@@ -522,6 +522,33 @@ impl RemoteDisk {
}
}
async fn open_read_chunks_with_retry(&self, request: ReadStreamRequest) -> Result<Option<rustfs_rio::ChunkReaderBox>> {
let mut attempt = 1;
let mut last_retry_classification = None;
loop {
match self.data_transport.open_read_chunks(request.clone()).await {
Ok(reader) => {
if attempt > 1
&& let Some(classification) = last_retry_classification
{
crate::cluster::rpc::runtime_sources::record_remote_disk_open_read_retry_success(classification);
}
return Ok(reader);
}
Err(err) if attempt < REMOTE_DISK_OPEN_READ_MAX_ATTEMPTS && Self::is_retryable_open_read_error(&err) => {
if let Some(classification) = err.internode_http_error_kind() {
let classification = classification.metric_label();
crate::cluster::rpc::runtime_sources::record_remote_disk_open_read_retry(classification);
last_retry_classification = Some(classification);
}
tokio::time::sleep(REMOTE_DISK_OPEN_READ_RETRY_BACKOFF).await;
attempt += 1;
}
Err(err) => return Err(err),
}
}
}
pub fn record_capacity_probe(&self, total: u64, used: u64, free: u64) {
self.health.record_capacity_probe(total, used, free);
}
@@ -2454,6 +2481,30 @@ impl DiskAPI for RemoteDisk {
.await
}
async fn read_file_stream_chunks(
&self,
volume: &str,
path: &str,
offset: usize,
length: usize,
) -> Result<Option<rustfs_rio::ChunkReaderBox>> {
if self.health.is_faulty() {
return Err(DiskError::FaultyDisk);
}
let disk = self.disk_ref().await;
let stall_timeout = get_object_disk_read_timeout();
self.open_read_chunks_with_retry(ReadStreamRequest {
endpoint: self.endpoint.grid_host(),
disk,
volume: volume.to_string(),
path: path.to_string(),
offset,
length,
stall_timeout: (!stall_timeout.is_zero()).then_some(stall_timeout),
})
.await
}
/// Buffered read for remote disks.
/// The transport stream is collected into owned Bytes for caller sharing.
#[tracing::instrument(level = "trace", skip_all)]
+15
View File
@@ -2022,6 +2022,21 @@ impl DiskAPI for LocalDiskWrapper {
.await
}
async fn read_file_stream_chunks(
&self,
volume: &str,
path: &str,
offset: usize,
length: usize,
) -> Result<Option<rustfs_rio::ChunkReaderBox>> {
self.track_disk_health_with_op(
"read_file_stream_chunks",
|| async { self.disk.read_file_stream_chunks(volume, path, offset, length).await },
get_max_timeout_duration(),
)
.await
}
async fn read_file_mmap_copy(&self, volume: &str, path: &str, offset: usize, length: usize) -> Result<bytes::Bytes> {
self.track_disk_health_with_op(
"read_file_mmap_copy",
+26
View File
@@ -65,6 +65,7 @@ use error::{Error, Result};
use local::LocalDisk;
use rustfs_filemeta::{FileInfo, ObjectPartInfo, RawFileInfo};
use rustfs_madmin::info_commands::DiskMetrics;
use rustfs_rio::ChunkReaderBox;
use serde::{Deserialize, Serialize};
use std::{fmt::Debug, path::PathBuf, sync::Arc, time::Duration};
use time::OffsetDateTime;
@@ -427,6 +428,19 @@ impl DiskAPI for Disk {
}
}
async fn read_file_stream_chunks(
&self,
volume: &str,
path: &str,
offset: usize,
length: usize,
) -> Result<Option<ChunkReaderBox>> {
match self {
Disk::Local(_) => Ok(None),
Disk::Remote(remote_disk) => remote_disk.read_file_stream_chunks(volume, path, offset, length).await,
}
}
#[tracing::instrument(level = "trace", skip_all)]
async fn read_file_mmap_copy(&self, volume: &str, path: &str, offset: usize, length: usize) -> Result<Bytes> {
match self {
@@ -865,6 +879,18 @@ pub trait DiskAPI: Debug + Send + Sync + 'static {
async fn read_file(&self, volume: &str, path: &str) -> Result<FileReader>;
async fn read_file_stream(&self, volume: &str, path: &str, offset: usize, length: usize) -> Result<FileReader>;
/// Returns an owned-chunk stream when the backing transport can preserve
/// receive-buffer ownership. `None` retains the ordinary reader path.
async fn read_file_stream_chunks(
&self,
_volume: &str,
_path: &str,
_offset: usize,
_length: usize,
) -> Result<Option<ChunkReaderBox>> {
Ok(None)
}
/// File read using mmap-then-copy on Unix or an efficient read on non-Unix.
async fn read_file_mmap_copy(&self, volume: &str, path: &str, offset: usize, length: usize) -> Result<Bytes>;
+521 -9
View File
@@ -14,7 +14,10 @@
use pin_project_lite::pin_project;
use rustfs_utils::HashAlgorithm;
use std::future::poll_fn;
use std::io::IoSlice;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tracing::error;
@@ -23,6 +26,18 @@ const LOG_COMPONENT_ECSTORE: &str = "ecstore";
const LOG_SUBSYSTEM_ERASURE: &str = "erasure";
const EVENT_BITROT_SHORT_SHARD_READ: &str = "bitrot_short_shard_read";
const EVENT_BITROT_HASH_MISMATCH: &str = "bitrot_hash_mismatch";
const MAX_RETAINED_CHUNKS_PER_BLOCK: usize = 64;
const MAX_CHUNK_POLLS_PER_YIELD: usize = MAX_RETAINED_CHUNKS_PER_BLOCK + 1;
/// Result of polling an optional owned-chunk handoff.
pub enum ShardChunkRead {
/// The source does not support owned-chunk handoff and remains untouched.
Unsupported,
/// The source reached EOF.
Eof,
/// A non-empty chunk containing at most the requested number of bytes.
Chunk(bytes::Bytes),
}
/// A shard source that may already hold its bytes in memory.
///
@@ -42,6 +57,12 @@ pub trait ShardSource: AsyncRead + Send + Sync + Unpin {
fn try_take_block(&mut self, _n: usize) -> Option<bytes::Bytes> {
None
}
/// Polls one owned chunk when the source supports chunk handoff.
/// `Unsupported` must leave the source untouched.
fn poll_read_chunk(self: Pin<&mut Self>, _cx: &mut Context<'_>, _max: usize) -> Poll<std::io::Result<ShardChunkRead>> {
Poll::Ready(Ok(ShardChunkRead::Unsupported))
}
}
/// Borrowed and owned byte slices are ordinary streaming sources: they carry no
@@ -75,6 +96,9 @@ pin_project! {
// contiguous on-disk `[hash][data]` block so both are pulled in a single
// pass; grown lazily and never shrunk.
buf: Vec<u8>,
// Reused owned chunk vector for the remote HTTP fast path. Keeping the
// allocation with the reader avoids allocating once per bitrot block.
chunks: Vec<bytes::Bytes>,
skip_verify: bool,
last_verify_duration: Duration,
}
@@ -91,6 +115,7 @@ where
hash_algo: algo,
shard_size,
buf: Vec::new(),
chunks: Vec::new(),
skip_verify,
last_verify_duration: Duration::ZERO,
}
@@ -260,11 +285,6 @@ where
let need = hash_size + want;
// In-memory fast path: the block is already resident, so slice it instead
// of copying it into the scratch buffer first (rustfs/backlog#1159). One
// copy (`extend_from_slice`) instead of two. A source that cannot serve
// `need` bytes returns `None` and falls through to the scratch path,
// keeping the short-read contract.
if let Some(block) = self.inner.try_take_block(need) {
let (data, verify) = split_and_verify(&self.hash_algo, self.skip_verify, &block)?;
out.extend_from_slice(data);
@@ -272,6 +292,126 @@ where
return Ok(want);
}
self.chunks.clear();
let handed_off = {
let inner = &mut self.inner;
let chunks = &mut self.chunks;
let tail_buf = &mut self.buf;
let mut received = 0usize;
poll_fn(|cx| {
for _ in 0..MAX_CHUNK_POLLS_PER_YIELD {
let next = match Pin::new(&mut *inner).poll_read_chunk(cx, need - received) {
Poll::Ready(Ok(next)) => next,
Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
Poll::Pending => return Poll::Pending,
};
let chunk = match next {
ShardChunkRead::Unsupported if received == 0 => return Poll::Ready(Ok(false)),
ShardChunkRead::Unsupported => {
return Poll::Ready(Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"chunk handoff became unavailable after transferring data",
)));
}
ShardChunkRead::Eof => {
return Poll::Ready(Err(short_shard_read(received.saturating_sub(hash_size), want)));
}
ShardChunkRead::Chunk(chunk) => chunk,
};
if received == 0 {
tail_buf.clear();
}
if chunk.is_empty() {
return Poll::Ready(Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"chunk handoff returned an empty chunk",
)));
}
let remaining = need - received;
if chunk.len() > remaining {
return Poll::Ready(Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"chunk handoff exceeded its requested boundary",
)));
}
received += chunk.len();
if chunks.len() == MAX_RETAINED_CHUNKS_PER_BLOCK {
if tail_buf.is_empty() {
tail_buf.reserve_exact(need - (received - chunk.len()));
}
tail_buf.extend_from_slice(&chunk);
} else {
chunks.push(chunk);
}
if received == need {
return Poll::Ready(Ok(true));
}
}
cx.waker().wake_by_ref();
Poll::Pending
})
.await?
};
if handed_off {
if self.chunks.len() == 1 && self.buf.is_empty() {
let block = &self.chunks[0];
let (data, verify) = split_and_verify(&self.hash_algo, self.skip_verify, block)?;
out.extend_from_slice(data);
self.last_verify_duration = verify;
return Ok(want);
}
let block_chunks = || {
self.chunks
.iter()
.map(|chunk| chunk.as_ref())
.chain((!self.buf.is_empty()).then_some(self.buf.as_slice()))
};
if !self.skip_verify {
let verify_start = std::time::Instant::now();
let actual_hash = self
.hash_algo
.hash_encode_slices(block_chunks().scan(hash_size, |skip, chunk| {
let start = (*skip).min(chunk.len());
*skip -= start;
Some(&chunk[start..])
}));
let verify = verify_start.elapsed();
let mut hash_offset = 0;
let mut remaining = hash_size;
for chunk in block_chunks() {
let take = remaining.min(chunk.len());
if actual_hash.as_ref()[hash_offset..hash_offset + take] != chunk[..take] {
error!(
event = EVENT_BITROT_HASH_MISMATCH,
component = LOG_COMPONENT_ECSTORE,
subsystem = LOG_SUBSYSTEM_ERASURE,
state = "failed",
data_len = want,
"bitrot hash mismatch"
);
return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "bitrot hash mismatch"));
}
hash_offset += take;
remaining -= take;
if remaining == 0 {
break;
}
}
self.last_verify_duration = verify;
}
let mut skip = hash_size;
for chunk in block_chunks() {
let start = skip.min(chunk.len());
skip -= start;
out.extend_from_slice(&chunk[start..]);
}
return Ok(want);
}
// Streaming path: same single pass and same verification as `read`; only
// the sink differs (`extend_from_slice` into `out` instead of
// `copy_from_slice` into a pre-zeroed buffer).
@@ -677,18 +817,167 @@ impl BitrotWriterWrapper {
#[cfg(test)]
mod tests {
use super::ShardSource;
use super::{
BitrotReader, BitrotWriter, BitrotWriterWrapper, CustomWriter, bitrot_shard_file_size, bitrot_verify, write_all_vectored,
};
use super::{MAX_RETAINED_CHUNKS_PER_BLOCK, ShardChunkRead, ShardSource};
use bytes::Bytes;
use rustfs_utils::HashAlgorithm;
use std::io::{Cursor, IoSlice};
use std::collections::VecDeque;
use std::io::{self, Cursor, IoSlice};
use std::pin::Pin;
use std::sync::{
Arc,
atomic::{AtomicUsize, Ordering},
};
use std::task::{Context, Poll};
use tokio::io::{AsyncWrite, AsyncWriteExt};
use std::time::Duration;
use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, ReadBuf};
struct FragmentedSource {
chunks: VecDeque<Bytes>,
}
impl FragmentedSource {
fn new(bytes: Vec<u8>, fragment_sizes: &[usize]) -> Self {
let mut chunks = VecDeque::new();
let mut offset = 0;
for &size in fragment_sizes {
let end = (offset + size).min(bytes.len());
if offset < end {
chunks.push_back(Bytes::copy_from_slice(&bytes[offset..end]));
}
offset = end;
}
if offset < bytes.len() {
chunks.push_back(Bytes::copy_from_slice(&bytes[offset..]));
}
Self { chunks }
}
}
impl AsyncRead for FragmentedSource {
fn poll_read(self: Pin<&mut Self>, _cx: &mut Context<'_>, _buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
Poll::Ready(Err(io::Error::other("fragmented source must use chunk handoff")))
}
}
impl ShardSource for FragmentedSource {
fn poll_read_chunk(mut self: Pin<&mut Self>, _cx: &mut Context<'_>, max: usize) -> Poll<io::Result<ShardChunkRead>> {
let Some(mut chunk) = self.chunks.pop_front() else {
return Poll::Ready(Ok(ShardChunkRead::Eof));
};
if chunk.len() > max {
self.chunks.push_front(chunk.split_off(max));
chunk.truncate(max);
}
Poll::Ready(Ok(ShardChunkRead::Chunk(chunk)))
}
}
struct GeneratedChunkSource {
bytes: Bytes,
offset: usize,
fragment_size: usize,
fail_at: Option<usize>,
}
impl GeneratedChunkSource {
fn new(bytes: Vec<u8>, fragment_size: usize) -> Self {
assert!(fragment_size > 0);
Self {
bytes: Bytes::from(bytes),
offset: 0,
fragment_size,
fail_at: None,
}
}
fn failing(bytes: Vec<u8>, fragment_size: usize, fail_at: usize) -> Self {
Self {
fail_at: Some(fail_at),
..Self::new(bytes, fragment_size)
}
}
}
impl AsyncRead for GeneratedChunkSource {
fn poll_read(self: Pin<&mut Self>, _cx: &mut Context<'_>, _buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
Poll::Ready(Err(io::Error::other("generated source must use chunk handoff")))
}
}
impl ShardSource for GeneratedChunkSource {
fn poll_read_chunk(mut self: Pin<&mut Self>, _cx: &mut Context<'_>, max: usize) -> Poll<io::Result<ShardChunkRead>> {
if self.fail_at == Some(self.offset) {
return Poll::Ready(Err(rustfs_rio::new_test_internode_http_io_error(
rustfs_rio::InternodeHttpErrorKind::BodyStreamAborted,
)));
}
if self.offset == self.bytes.len() {
return Poll::Ready(Ok(ShardChunkRead::Eof));
}
let error_limit = self.fail_at.unwrap_or(self.bytes.len());
let take = self
.fragment_size
.min(max)
.min(error_limit - self.offset)
.min(self.bytes.len() - self.offset);
let start = self.offset;
self.offset += take;
Poll::Ready(Ok(ShardChunkRead::Chunk(self.bytes.slice(start..start + take))))
}
}
struct InvalidChunkSource {
mode: InvalidChunkMode,
}
#[derive(Clone, Copy)]
enum InvalidChunkMode {
Empty,
Oversized,
UnsupportedAfterChunk,
Unsupported,
}
impl AsyncRead for InvalidChunkSource {
fn poll_read(self: Pin<&mut Self>, _cx: &mut Context<'_>, _buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
Poll::Ready(Err(io::Error::other("invalid source must use chunk handoff")))
}
}
impl ShardSource for InvalidChunkSource {
fn poll_read_chunk(mut self: Pin<&mut Self>, _cx: &mut Context<'_>, max: usize) -> Poll<io::Result<ShardChunkRead>> {
match self.mode {
InvalidChunkMode::Empty => Poll::Ready(Ok(ShardChunkRead::Chunk(Bytes::new()))),
InvalidChunkMode::Oversized => Poll::Ready(Ok(ShardChunkRead::Chunk(Bytes::from(vec![0; max + 1])))),
InvalidChunkMode::UnsupportedAfterChunk => {
self.mode = InvalidChunkMode::Unsupported;
Poll::Ready(Ok(ShardChunkRead::Chunk(Bytes::from_static(b"x"))))
}
InvalidChunkMode::Unsupported => Poll::Ready(Ok(ShardChunkRead::Unsupported)),
}
}
}
struct ScratchReuseSource {
block: Option<Bytes>,
saw_reused_scratch: bool,
}
impl AsyncRead for ScratchReuseSource {
fn poll_read(mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
let Some(block) = self.block.take() else {
return Poll::Ready(Ok(()));
};
self.saw_reused_scratch = buf.initialize_unfilled()[..block.len()].iter().all(|byte| *byte == 0xa5);
buf.put_slice(&block);
Poll::Ready(Ok(()))
}
}
impl ShardSource for ScratchReuseSource {}
#[derive(Default)]
struct VectoredCountingWriter {
@@ -1446,6 +1735,70 @@ mod tests {
assert!(out.is_empty(), "corrupt bytes must never reach the caller's buffer");
}
#[tokio::test]
async fn chunked_handoff_verifies_data_split_across_hash_boundaries() {
const SHARD: usize = 4096;
let algo = HashAlgorithm::HighwayHash256S;
let data: Vec<u8> = (0..SHARD).map(|index| (index % 251) as u8).collect();
let mut encoded = Vec::new();
BitrotWriter::new(&mut encoded, SHARD, algo.clone())
.write(&data)
.await
.expect("write shard");
let mut output = Vec::with_capacity(SHARD);
BitrotReader::new(FragmentedSource::new(encoded, &[3, 11, 19, 37, 128]), SHARD, algo, false)
.read_appending(&mut output, SHARD)
.await
.expect("fragmented shard must verify");
assert_eq!(output, data);
}
#[tokio::test]
async fn chunked_handoff_never_appends_a_corrupt_shard() {
const SHARD: usize = 4096;
let algo = HashAlgorithm::HighwayHash256S;
let mut encoded = Vec::new();
BitrotWriter::new(&mut encoded, SHARD, algo.clone())
.write(&vec![9u8; SHARD])
.await
.expect("write shard");
let last = encoded.len() - 1;
encoded[last] ^= 0xff;
let mut output = Vec::with_capacity(SHARD);
let err = BitrotReader::new(FragmentedSource::new(encoded, &[7, 17, 31]), SHARD, algo, false)
.read_appending(&mut output, SHARD)
.await
.expect_err("corrupt fragmented shard must fail");
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
assert!(output.is_empty());
}
#[tokio::test]
async fn chunked_handoff_does_not_hash_when_verification_is_skipped() {
const SHARD: usize = 4096;
let algo = HashAlgorithm::HighwayHash256S;
let mut encoded = Vec::new();
BitrotWriter::new(&mut encoded, SHARD, algo.clone())
.write(&vec![9u8; SHARD])
.await
.expect("write shard");
encoded[0] ^= 0xff;
let mut output = Vec::with_capacity(SHARD);
let mut reader = BitrotReader::new(FragmentedSource::new(encoded, &[7, 17, 31]), SHARD, algo, true);
reader
.read_appending(&mut output, SHARD)
.await
.expect("skipped verification must accept fragmented shard bytes");
assert_eq!(reader.last_verify_duration(), Duration::ZERO);
assert_eq!(output, vec![9u8; SHARD]);
}
#[tokio::test]
async fn read_appending_rejects_a_want_larger_than_the_shard() {
let algo = HashAlgorithm::HighwayHash256;
@@ -1497,10 +1850,21 @@ mod tests {
// Equivalence: same bytes out of both paths.
let mut via_mem: Vec<u8> = Vec::with_capacity(SHARD);
BitrotReader::new(Cursor::new(Bytes::from(encoded.clone())), SHARD, algo.clone(), false)
let mut memory_reader = BitrotReader::new(Cursor::new(Bytes::from(encoded.clone())), SHARD, algo.clone(), false);
memory_reader
.read_appending(&mut via_mem, SHARD)
.await
.expect("in-memory read");
assert_eq!(
memory_reader.chunks.capacity(),
0,
"the synchronous fast path must not allocate chunk storage"
);
assert_eq!(
memory_reader.buf.capacity(),
0,
"the synchronous fast path must not allocate scratch storage"
);
let mut via_stream: Vec<u8> = Vec::with_capacity(SHARD);
BitrotReader::new(Cursor::new(encoded), SHARD, algo, false)
@@ -1537,4 +1901,152 @@ mod tests {
assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
assert!(out.is_empty(), "corrupt bytes must never reach the caller's buffer");
}
#[tokio::test]
async fn streaming_fallback_reuses_initialized_scratch() {
const SHARD: usize = 4096;
let algo = HashAlgorithm::HighwayHash256S;
let data = vec![7u8; SHARD];
let encoded = encode_one_block(&data, SHARD, algo.clone()).await;
let source = ScratchReuseSource {
block: Some(Bytes::copy_from_slice(&encoded)),
saw_reused_scratch: false,
};
let mut reader = BitrotReader::new(source, SHARD, algo, false);
reader.buf = vec![0xa5; encoded.len()];
let mut output = Vec::new();
reader
.read_appending(&mut output, SHARD)
.await
.expect("streaming fallback should verify");
assert!(reader.inner.saw_reused_scratch, "capability probing must not clear reusable scratch");
assert_eq!(output, data);
}
#[tokio::test]
async fn chunked_handoff_bounds_production_sized_one_byte_fragments() {
const SHARD: usize = 1024 * 1024 / 4;
let algo = HashAlgorithm::HighwayHash256S;
let data: Vec<u8> = (0..SHARD).map(|index| (index % 251) as u8).collect();
let encoded = encode_one_block(&data, SHARD, algo.clone()).await;
let encoded_len = encoded.len();
let mut reader = BitrotReader::new(GeneratedChunkSource::new(encoded, 1), SHARD, algo, false);
let mut output = Vec::with_capacity(SHARD);
reader
.read_appending(&mut output, SHARD)
.await
.expect("one-byte fragments should verify with bounded retained state");
assert_eq!(output, data);
assert_eq!(reader.chunks.len(), MAX_RETAINED_CHUNKS_PER_BLOCK);
assert!(reader.chunks.capacity() <= MAX_RETAINED_CHUNKS_PER_BLOCK);
assert_eq!(reader.buf.len(), encoded_len - MAX_RETAINED_CHUNKS_PER_BLOCK);
}
#[tokio::test]
async fn chunked_handoff_keeps_sixty_four_frames_zero_copy_and_respects_poll_budget() {
const SHARD: usize = 1024 * 1024;
const FRAME: usize = 16 * 1024;
let algo = HashAlgorithm::HighwayHash256S;
let small_data = vec![3u8; 4096];
let small_encoded = encode_one_block(&small_data, 4096, algo.clone()).await;
let mut exact_reader =
BitrotReader::new(FragmentedSource::new(small_encoded.clone(), &[1; 63]), 4096, algo.clone(), false);
let mut exact_output = Vec::new();
exact_reader
.read_appending(&mut exact_output, 4096)
.await
.expect("exactly sixty-four frames should verify");
assert_eq!(exact_output, small_data);
assert_eq!(exact_reader.chunks.len(), MAX_RETAINED_CHUNKS_PER_BLOCK);
assert!(exact_reader.buf.is_empty(), "the threshold itself must remain zero-copy");
let mut yielded_reader = BitrotReader::new(FragmentedSource::new(small_encoded, &[1; 65]), 4096, algo.clone(), false);
let mut yielded_output = Vec::new();
let mut yielded_read = Box::pin(yielded_reader.read_appending(&mut yielded_output, 4096));
let mut cx = Context::from_waker(std::task::Waker::noop());
assert!(std::future::Future::poll(yielded_read.as_mut(), &mut cx).is_pending());
assert!(matches!(std::future::Future::poll(yielded_read.as_mut(), &mut cx), Poll::Ready(Ok(4096))));
drop(yielded_read);
assert_eq!(yielded_output, small_data);
let data = vec![7u8; SHARD];
let encoded = encode_one_block(&data, SHARD, algo.clone()).await;
let mut reader = BitrotReader::new(FragmentedSource::new(encoded, &[FRAME; 64]), SHARD, algo, false);
let mut output = Vec::with_capacity(SHARD);
let mut read = Box::pin(reader.read_appending(&mut output, SHARD));
assert!(
matches!(std::future::Future::poll(read.as_mut(), &mut cx), Poll::Ready(Ok(SHARD))),
"sixty-five normal HTTP frames should complete without a cooperative yield"
);
drop(read);
assert_eq!(output, data);
assert_eq!(reader.chunks.len(), MAX_RETAINED_CHUNKS_PER_BLOCK);
assert_eq!(reader.buf.len(), HashAlgorithm::HighwayHash256S.size());
}
#[tokio::test]
async fn chunked_tail_failures_preserve_errors_and_output() {
const SHARD: usize = 4096;
let algo = HashAlgorithm::HighwayHash256S;
let data = vec![7u8; SHARD];
let encoded = encode_one_block(&data, SHARD, algo.clone()).await;
let sentinel = vec![1u8, 2, 3];
let mut short_output = sentinel.clone();
let short_err = BitrotReader::new(GeneratedChunkSource::new(encoded[..100].to_vec(), 1), SHARD, algo.clone(), false)
.read_appending(&mut short_output, SHARD)
.await
.expect_err("EOF after the retention threshold must stay a short read");
assert_eq!(short_err.kind(), io::ErrorKind::UnexpectedEof);
assert_eq!(short_output, sentinel);
let mut corrupt = encoded.clone();
let last = corrupt.len() - 1;
corrupt[last] ^= 0xff;
let mut corrupt_output = sentinel.clone();
let corrupt_err = BitrotReader::new(GeneratedChunkSource::new(corrupt, 1), SHARD, algo.clone(), false)
.read_appending(&mut corrupt_output, SHARD)
.await
.expect_err("corrupt coalesced tail must fail verification");
assert_eq!(corrupt_err.kind(), io::ErrorKind::InvalidData);
assert_eq!(corrupt_output, sentinel);
let mut failed_output = sentinel.clone();
let body_err = BitrotReader::new(GeneratedChunkSource::failing(encoded, 1, 65), SHARD, algo, false)
.read_appending(&mut failed_output, SHARD)
.await
.expect_err("a terminal body error must not become EOF");
let source = body_err
.get_ref()
.and_then(|source| source.downcast_ref::<rustfs_rio::InternodeHttpError>())
.expect("body error should retain internode classification");
assert_eq!(source.kind(), rustfs_rio::InternodeHttpErrorKind::BodyStreamAborted);
assert_eq!(failed_output, sentinel);
}
#[tokio::test]
async fn chunked_handoff_rejects_invalid_source_contracts() {
const SHARD: usize = 64;
for mode in [
InvalidChunkMode::Empty,
InvalidChunkMode::Oversized,
InvalidChunkMode::UnsupportedAfterChunk,
] {
let source = InvalidChunkSource { mode };
let mut output = vec![9u8];
let err = BitrotReader::new(source, SHARD, HashAlgorithm::HighwayHash256S, false)
.read_appending(&mut output, SHARD)
.await
.expect_err("invalid chunk contracts must fail closed");
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
assert_eq!(output, vec![9u8]);
}
}
}
+117 -2
View File
@@ -22,12 +22,13 @@ use crate::diagnostics::get::{
#[cfg(feature = "hotpath")]
use crate::disk::FileWriter;
use crate::disk::{self, DiskAPI as _, DiskStore, FileReader, MmapCopyStageMetrics, error::DiskError};
use crate::erasure::coding::{BitrotReader, BitrotWriterWrapper, CustomWriter};
use crate::erasure::coding::{BitrotReader, BitrotWriterWrapper, CustomWriter, ShardChunkRead};
use bytes::Bytes;
use rustfs_config::{
DEFAULT_OBJECT_MMAP_READ_ENABLE, DEFAULT_OBJECT_MMAP_READ_MAX_LENGTH, ENV_OBJECT_MMAP_READ_ENABLE,
ENV_OBJECT_MMAP_READ_MAX_LENGTH, ENV_OBJECT_ZERO_COPY_ENABLE,
};
use rustfs_rio::ChunkReaderBox;
use rustfs_utils::HashAlgorithm;
use std::future::Future;
use std::io::{self, Cursor};
@@ -51,6 +52,7 @@ tokio::task_local! {
/// (rustfs/backlog#1159). Everything else is a stream and keeps the old path.
pub enum ShardReader {
InMemory(Cursor<Bytes>),
Chunked(ChunkReaderBox),
Stream(Box<dyn AsyncRead + Send + Sync + Unpin>),
}
@@ -58,6 +60,7 @@ impl AsyncRead for ShardReader {
fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut tokio::io::ReadBuf<'_>) -> Poll<std::io::Result<()>> {
match self.get_mut() {
Self::InMemory(cursor) => Pin::new(cursor).poll_read(cx, buf),
Self::Chunked(reader) => Pin::new(&mut **reader).poll_read(cx, buf),
Self::Stream(reader) => Pin::new(reader).poll_read(cx, buf),
}
}
@@ -67,7 +70,19 @@ impl crate::erasure::coding::ShardSource for ShardReader {
fn try_take_block(&mut self, n: usize) -> Option<Bytes> {
match self {
Self::InMemory(cursor) => cursor.try_take_block(n),
Self::Stream(_) => None,
Self::Chunked(_) | Self::Stream(_) => None,
}
}
fn poll_read_chunk(self: Pin<&mut Self>, cx: &mut Context<'_>, max: usize) -> Poll<io::Result<ShardChunkRead>> {
let Self::Chunked(reader) = self.get_mut() else {
return Poll::Ready(Ok(ShardChunkRead::Unsupported));
};
match Pin::new(&mut **reader).poll_read_chunk(cx, max) {
Poll::Ready(Ok(Some(chunk))) => Poll::Ready(Ok(ShardChunkRead::Chunk(chunk))),
Poll::Ready(Ok(None)) => Poll::Ready(Ok(ShardChunkRead::Eof)),
Poll::Ready(Err(err)) => Poll::Ready(Err(err)),
Poll::Pending => Poll::Pending,
}
}
}
@@ -345,6 +360,17 @@ async fn open_disk_reader(
let metrics_path = metrics_path.filter(|_| rustfs_io_metrics::get_stage_metrics_enabled());
let stage_metrics_enabled = metrics_path.is_some();
// Preserve HTTP body ownership only on healthy remote reads. Instrumented
// and local paths retain their existing AsyncRead wrappers.
if use_mmap_read
&& !disk.is_local()
&& !stage_metrics_enabled
&& !cfg!(feature = "hotpath")
&& let Some(reader) = disk.read_file_stream_chunks(bucket, path, offset, length).await?
{
return Ok(ShardReader::Chunked(reader));
}
// Mmap-copy materializes the whole `offset..offset+length` range as one
// owned allocation before any byte is served, and GET/heal shard reads
// request the entire part span in one call. Over-cap reads (e.g. a huge
@@ -780,6 +806,50 @@ pub async fn create_bitrot_writer(
#[cfg(test)]
mod tests {
use super::*;
use rustfs_rio::ChunkReader;
use std::collections::VecDeque;
struct TestChunkReader {
chunks: VecDeque<Bytes>,
}
impl TestChunkReader {
fn new(bytes: Bytes, fragment_sizes: &[usize]) -> Self {
let mut chunks = VecDeque::new();
let mut offset = 0;
for &size in fragment_sizes {
let end = (offset + size).min(bytes.len());
if offset < end {
chunks.push_back(bytes.slice(offset..end));
}
offset = end;
}
if offset < bytes.len() {
chunks.push_back(bytes.slice(offset..));
}
Self { chunks }
}
}
impl AsyncRead for TestChunkReader {
fn poll_read(self: Pin<&mut Self>, _cx: &mut Context<'_>, _buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
Poll::Ready(Err(io::Error::other("test chunk reader must use chunk handoff")))
}
}
impl ChunkReader for TestChunkReader {
fn poll_read_chunk(mut self: Pin<&mut Self>, _cx: &mut Context<'_>, max: usize) -> Poll<io::Result<Option<Bytes>>> {
let Some(mut chunk) = self.chunks.pop_front() else {
return Poll::Ready(Ok(None));
};
let take = chunk.len().min(max);
if take < chunk.len() {
self.chunks.push_front(chunk.split_off(take));
}
chunk.truncate(take);
Poll::Ready(Ok(Some(chunk)))
}
}
#[cfg(feature = "hotpath")]
use crate::cluster::rpc::RemoteDisk;
@@ -1669,4 +1739,49 @@ mod tests {
println!("error: {error:?}");
assert_eq!(error, DiskError::DiskNotFound);
}
#[tokio::test]
async fn shard_reader_chunked_path_verifies_fragmented_remote_block() {
const SHARD_SIZE: usize = 1024;
let algo = HashAlgorithm::HighwayHash256S;
let data = vec![42u8; SHARD_SIZE];
let mut encoded = Vec::new();
crate::erasure::coding::BitrotWriter::new(&mut encoded, SHARD_SIZE, algo.clone())
.write(&data)
.await
.expect("test shard should encode");
let source = TestChunkReader::new(Bytes::from(encoded), &[3, 7, 17, 31]);
let mut reader = BitrotReader::new(ShardReader::Chunked(Box::new(source)), SHARD_SIZE, algo, false);
let mut output = Vec::with_capacity(SHARD_SIZE);
reader
.read_appending(&mut output, SHARD_SIZE)
.await
.expect("fragmented remote shard should verify");
assert_eq!(output, data);
}
#[tokio::test]
async fn shard_reader_chunked_path_handles_more_than_one_poll_budget() {
const SHARD_SIZE: usize = 1024;
let algo = HashAlgorithm::HighwayHash256S;
let data = vec![42u8; SHARD_SIZE];
let mut encoded = Vec::new();
crate::erasure::coding::BitrotWriter::new(&mut encoded, SHARD_SIZE, algo.clone())
.write(&data)
.await
.expect("test shard should encode");
let fragment_sizes = vec![1; encoded.len()];
let source = TestChunkReader::new(Bytes::from(encoded), &fragment_sizes);
let mut reader = BitrotReader::new(ShardReader::Chunked(Box::new(source)), SHARD_SIZE, algo, false);
let mut output = Vec::with_capacity(SHARD_SIZE);
reader
.read_appending(&mut output, SHARD_SIZE)
.await
.expect("fragmented remote shard should verify after multiple polls");
assert_eq!(output, data);
}
}
+358 -16
View File
@@ -52,6 +52,8 @@ const HTTP_VERSION_10_LABEL: &str = "http/1.0";
const HTTP_VERSION_11_LABEL: &str = "http/1.1";
const HTTP_VERSION_2_LABEL: &str = "h2";
const HTTP_VERSION_UNKNOWN_LABEL: &str = "unknown";
const MAX_CONSECUTIVE_EMPTY_CHUNKS: usize = 64;
const EXCESSIVE_EMPTY_CHUNKS_ERROR: &str = "HTTP body returned too many empty chunks";
pub const INTERNODE_DISK_ERROR_HEADER: &str = "x-rustfs-disk-error";
pub const INTERNODE_FILE_NOT_FOUND: &str = "file-not-found";
pub const INTERNODE_VOLUME_NOT_FOUND: &str = "volume-not-found";
@@ -883,6 +885,26 @@ fn internode_status_error(method: &Method, url: &str, operation: Option<&'static
InternodeHttpError::new(classified, context).into_io_error()
}
type HttpByteStream = Pin<Box<dyn Stream<Item = std::io::Result<Bytes>> + Send + Sync>>;
/// An async reader that can also transfer received HTTP body chunks without
/// copying their contents into an intermediate caller buffer.
pub trait ChunkReader: AsyncRead + Send + Sync + Unpin {
/// Returns the next non-empty owned chunk, limited to `max` bytes.
/// `None` is EOF.
fn poll_read_chunk(self: Pin<&mut Self>, cx: &mut Context<'_>, max: usize) -> Poll<io::Result<Option<Bytes>>>;
}
pub type ChunkReaderBox = Box<dyn ChunkReader>;
struct HttpReaderInit {
stream: HttpByteStream,
track_internode_metrics: bool,
internode_operation: Option<&'static str>,
stall_timeout: Option<Duration>,
request_started: Instant,
}
pin_project! {
pub struct HttpReader {
url:String,
@@ -895,7 +917,22 @@ pin_project! {
request_started: Instant,
duration_recorded: bool,
#[pin]
inner: StreamReader<Pin<Box<dyn Stream<Item=std::io::Result<Bytes>>+Send+Sync>>, Bytes>,
inner: StreamReader<HttpByteStream, Bytes>,
}
}
pin_project! {
pub struct HttpChunkReader {
track_internode_metrics: bool,
internode_operation: Option<&'static str>,
stall_timeout: Option<Duration>,
stall_timer: Option<Pin<Box<Sleep>>>,
request_started: Instant,
duration_recorded: bool,
consecutive_empty_chunks: usize,
#[pin]
inner: HttpByteStream,
current: Option<Bytes>,
}
}
@@ -934,12 +971,34 @@ impl HttpReader {
_read_buf_size: usize,
stall_timeout: Option<Duration>,
) -> io::Result<Self> {
let track_internode_metrics = is_internode_rpc_url(&url);
let internode_operation = internode_rpc_operation(&url);
let client = get_http_client(&url).await.inspect_err(|_| {
let init = Self::open(&url, &method, &headers, body, stall_timeout).await?;
Ok(Self {
inner: StreamReader::new(init.stream),
url,
method,
headers,
track_internode_metrics: init.track_internode_metrics,
internode_operation: init.internode_operation,
stall_timer: None,
stall_timeout: init.stall_timeout,
request_started: init.request_started,
duration_recorded: false,
})
}
async fn open(
url: &str,
method: &Method,
headers: &HeaderMap,
body: Option<Vec<u8>>,
stall_timeout: Option<Duration>,
) -> io::Result<HttpReaderInit> {
let track_internode_metrics = is_internode_rpc_url(url);
let internode_operation = internode_rpc_operation(url);
let client = get_http_client(url).await.inspect_err(|_| {
record_internode_error(track_internode_metrics, internode_operation);
})?;
let mut request: RequestBuilder = client.request(method.clone(), url.clone()).headers(headers.clone());
let mut request: RequestBuilder = client.request(method.clone(), url).headers(headers.clone());
if let Some(body) = body {
request = request.body(body);
}
@@ -949,7 +1008,7 @@ impl HttpReader {
record_internode_operation_duration(track_internode_metrics, internode_operation, request_started.elapsed());
record_internode_error(track_internode_metrics, internode_operation);
record_internode_classified_error(track_internode_metrics, internode_operation, classify_reqwest_error(&e));
internode_reqwest_error(&method, &url, internode_operation, e)
internode_reqwest_error(method, url, internode_operation, e)
})?;
record_internode_http_version(track_internode_metrics, internode_operation, http_version_metric_label(resp.version()));
@@ -959,12 +1018,12 @@ impl HttpReader {
record_internode_operation_duration(track_internode_metrics, internode_operation, request_started.elapsed());
record_internode_error(track_internode_metrics, internode_operation);
record_internode_classified_error(track_internode_metrics, internode_operation, classified.kind);
return Err(internode_classified_error(&method, &url, internode_operation, classified));
return Err(internode_classified_error(method, url, internode_operation, classified));
}
record_internode_outgoing_request(track_internode_metrics, internode_operation);
let stream_error_url = url.clone();
let stream_error_url = url.to_owned();
let stream_error_method = method.clone();
let stream = resp.bytes_stream().map_err(move |e| {
record_internode_error(track_internode_metrics, internode_operation);
@@ -973,17 +1032,12 @@ impl HttpReader {
internode_reqwest_body_error(&stream_error_method, &stream_error_url, internode_operation, e)
});
Ok(Self {
inner: StreamReader::new(Box::pin(stream)),
url,
method,
headers,
Ok(HttpReaderInit {
stream: Box::pin(stream),
track_internode_metrics,
internode_operation,
stall_timer: None,
stall_timeout,
request_started,
duration_recorded: false,
})
}
pub fn url(&self) -> &str {
@@ -1000,7 +1054,6 @@ impl HttpReader {
impl AsyncRead for HttpReader {
fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
let mut this = self.project();
let filled_before = buf.filled().len();
match this.inner.as_mut().poll_read(cx, buf) {
Poll::Ready(Ok(())) => {
@@ -1053,6 +1106,129 @@ impl AsyncRead for HttpReader {
}
}
impl HttpChunkReader {
pub async fn new_with_stall_timeout(
url: String,
method: Method,
headers: HeaderMap,
body: Option<Vec<u8>>,
stall_timeout: Option<Duration>,
) -> io::Result<Self> {
let init = HttpReader::open(&url, &method, &headers, body, stall_timeout).await?;
Ok(Self {
inner: init.stream,
current: None,
track_internode_metrics: init.track_internode_metrics,
internode_operation: init.internode_operation,
stall_timer: None,
stall_timeout: init.stall_timeout,
request_started: init.request_started,
duration_recorded: false,
consecutive_empty_chunks: 0,
})
}
}
fn excessive_empty_chunks_error() -> Error {
Error::new(io::ErrorKind::InvalidData, EXCESSIVE_EMPTY_CHUNKS_ERROR)
}
impl AsyncRead for HttpChunkReader {
fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
if buf.remaining() == 0 {
return Poll::Ready(Ok(()));
}
match ChunkReader::poll_read_chunk(self.as_mut(), cx, buf.remaining()) {
Poll::Ready(Ok(Some(chunk))) => {
buf.put_slice(&chunk);
Poll::Ready(Ok(()))
}
Poll::Ready(Ok(None)) => Poll::Ready(Ok(())),
Poll::Ready(Err(err)) => Poll::Ready(Err(err)),
Poll::Pending => Poll::Pending,
}
}
}
impl ChunkReader for HttpChunkReader {
fn poll_read_chunk(self: Pin<&mut Self>, cx: &mut Context<'_>, max: usize) -> Poll<io::Result<Option<Bytes>>> {
if max == 0 {
return Poll::Ready(Err(Error::new(io::ErrorKind::InvalidInput, "chunk read limit must be non-zero")));
}
let mut this = self.project();
if *this.consecutive_empty_chunks >= MAX_CONSECUTIVE_EMPTY_CHUNKS {
return Poll::Ready(Err(excessive_empty_chunks_error()));
}
loop {
if let Some(mut current) = this.current.take() {
let take = current.len().min(max);
let chunk = current.split_to(take);
if !current.is_empty() {
*this.current = Some(current);
}
record_internode_recv_bytes(*this.track_internode_metrics, *this.internode_operation, take);
*this.stall_timer = None;
return Poll::Ready(Ok(Some(chunk)));
}
match this.inner.as_mut().poll_next(cx) {
Poll::Ready(Some(Ok(bytes))) if bytes.is_empty() => {
*this.consecutive_empty_chunks += 1;
if *this.consecutive_empty_chunks == MAX_CONSECUTIVE_EMPTY_CHUNKS {
record_internode_error(*this.track_internode_metrics, *this.internode_operation);
return Poll::Ready(Err(excessive_empty_chunks_error()));
}
}
Poll::Ready(Some(Ok(bytes))) => {
*this.consecutive_empty_chunks = 0;
*this.current = Some(bytes);
}
Poll::Ready(Some(Err(err))) => {
record_internode_operation_duration_once(
*this.track_internode_metrics,
*this.internode_operation,
*this.request_started,
this.duration_recorded,
);
return Poll::Ready(Err(err));
}
Poll::Ready(None) => {
record_internode_operation_duration_once(
*this.track_internode_metrics,
*this.internode_operation,
*this.request_started,
this.duration_recorded,
);
*this.stall_timer = None;
return Poll::Ready(Ok(None));
}
Poll::Pending => {
let Some(stall_timeout) = *this.stall_timeout else {
return Poll::Pending;
};
let timer = this.stall_timer.get_or_insert_with(|| Box::pin(time::sleep(stall_timeout)));
if timer.as_mut().poll(cx).is_ready() {
record_internode_operation_duration_once(
*this.track_internode_metrics,
*this.internode_operation,
*this.request_started,
this.duration_recorded,
);
record_internode_stall_timeout(*this.track_internode_metrics, *this.internode_operation);
record_internode_error(*this.track_internode_metrics, *this.internode_operation);
return Poll::Ready(Err(Error::new(
io::ErrorKind::TimedOut,
"HttpReader stall timeout: no data received before deadline",
)));
}
return Poll::Pending;
}
}
}
}
}
impl EtagResolvable for HttpReader {
fn is_etag_reader(&self) -> bool {
false
@@ -2012,6 +2188,144 @@ mod tests {
handle.abort();
}
#[tokio::test]
async fn http_chunk_reader_handoff_preserves_boundaries_and_eof() {
let state = TestState::default();
let Some((url, handle)) = start_test_server(state.clone()).await else {
return;
};
let mut reader = HttpChunkReader::new_with_stall_timeout(url, Method::GET, HeaderMap::new(), None, None)
.await
.expect("reader should open");
assert_eq!(reader.consecutive_empty_chunks, 0);
let zero = std::future::poll_fn(|cx| Pin::new(&mut reader).poll_read_chunk(cx, 0))
.await
.expect_err("zero chunk bound is invalid");
assert_eq!(zero.kind(), io::ErrorKind::InvalidInput);
let first = std::future::poll_fn(|cx| Pin::new(&mut reader).poll_read_chunk(cx, 2))
.await
.expect("first chunk read should succeed")
.expect("first chunk should not be EOF");
assert_eq!(first, b"he"[..]);
let second = std::future::poll_fn(|cx| Pin::new(&mut reader).poll_read_chunk(cx, 8))
.await
.expect("second chunk read should succeed")
.expect("second chunk should not be EOF");
assert_eq!(second, b"llo"[..]);
let eof = std::future::poll_fn(|cx| Pin::new(&mut reader).poll_read_chunk(cx, 8))
.await
.expect("EOF should not be an error");
assert!(eof.is_none());
assert_eq!(state.get_count.load(Ordering::SeqCst), 1);
handle.abort();
}
#[test]
fn http_chunk_reader_rejects_excessive_empty_chunks_on_both_interfaces() {
let make_reader = |inner: HttpByteStream| HttpChunkReader {
track_internode_metrics: false,
internode_operation: None,
stall_timeout: None,
stall_timer: None,
request_started: Instant::now(),
duration_recorded: false,
consecutive_empty_chunks: 0,
inner,
current: None,
};
let empty_chunks_then_data = |empty_chunks| {
let items = (0..empty_chunks)
.map(|_| Ok(Bytes::new()))
.chain(std::iter::once(Ok(Bytes::from_static(b"data"))));
make_reader(Box::pin(stream::iter(items)))
};
let mut cx = Context::from_waker(std::task::Waker::noop());
let mut chunk_reader = empty_chunks_then_data(MAX_CONSECUTIVE_EMPTY_CHUNKS - 1);
let Poll::Ready(Ok(Some(chunk))) = Pin::new(&mut chunk_reader).poll_read_chunk(&mut cx, 4) else {
panic!("data after fewer than the maximum empty chunks should be returned");
};
assert_eq!(chunk, b"data"[..]);
let mut async_reader = empty_chunks_then_data(MAX_CONSECUTIVE_EMPTY_CHUNKS - 1);
let mut storage = [0; 4];
let mut read_buf = ReadBuf::new(&mut storage);
let Poll::Ready(Ok(())) = Pin::new(&mut async_reader).poll_read(&mut cx, &mut read_buf) else {
panic!("AsyncRead should return data after fewer than the maximum empty chunks");
};
assert_eq!(read_buf.filled(), b"data");
let mut chunk_reader = empty_chunks_then_data(MAX_CONSECUTIVE_EMPTY_CHUNKS);
let Poll::Ready(Err(err)) = Pin::new(&mut chunk_reader).poll_read_chunk(&mut cx, 4) else {
panic!("excessive empty chunks should fail closed");
};
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
let Poll::Ready(Err(err)) = Pin::new(&mut chunk_reader).poll_read_chunk(&mut cx, 4) else {
panic!("empty chunk limit failure should remain sticky");
};
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
let mut async_reader = empty_chunks_then_data(MAX_CONSECUTIVE_EMPTY_CHUNKS);
let mut storage = [0; 4];
let mut read_buf = ReadBuf::new(&mut storage);
let Poll::Ready(Err(err)) = Pin::new(&mut async_reader).poll_read(&mut cx, &mut read_buf) else {
panic!("excessive empty chunks should fail closed through AsyncRead");
};
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
let mut read_buf = ReadBuf::new(&mut storage);
let Poll::Ready(Err(err)) = Pin::new(&mut async_reader).poll_read(&mut cx, &mut read_buf) else {
panic!("AsyncRead empty chunk limit failure should remain sticky");
};
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
let make_pending_reader = || {
let mut empty_chunks = 0;
let mut returned_pending = false;
let items = stream::poll_fn(move |cx| {
if empty_chunks < MAX_CONSECUTIVE_EMPTY_CHUNKS - 1 {
empty_chunks += 1;
return Poll::Ready(Some(Ok(Bytes::new())));
}
if !returned_pending {
returned_pending = true;
cx.waker().wake_by_ref();
return Poll::Pending;
}
if empty_chunks < MAX_CONSECUTIVE_EMPTY_CHUNKS {
empty_chunks += 1;
return Poll::Ready(Some(Ok(Bytes::new())));
}
Poll::Ready(None)
});
make_reader(Box::pin(items))
};
let mut chunk_reader = make_pending_reader();
assert!(Pin::new(&mut chunk_reader).poll_read_chunk(&mut cx, 4).is_pending());
let Poll::Ready(Err(err)) = Pin::new(&mut chunk_reader).poll_read_chunk(&mut cx, 4) else {
panic!("the consecutive empty chunk limit must survive Pending");
};
assert_eq!(err.kind(), io::ErrorKind::InvalidData);
let items = (0..MAX_CONSECUTIVE_EMPTY_CHUNKS - 1)
.map(|_| Ok(Bytes::new()))
.chain(std::iter::once(Ok(Bytes::from_static(b"one"))))
.chain((0..MAX_CONSECUTIVE_EMPTY_CHUNKS - 1).map(|_| Ok(Bytes::new())))
.chain(std::iter::once(Ok(Bytes::from_static(b"two"))));
let mut chunk_reader = make_reader(Box::pin(stream::iter(items)));
let Poll::Ready(Ok(Some(first))) = Pin::new(&mut chunk_reader).poll_read_chunk(&mut cx, 3) else {
panic!("data should reset the consecutive empty chunk count");
};
assert_eq!(first, b"one"[..]);
let Poll::Ready(Ok(Some(second))) = Pin::new(&mut chunk_reader).poll_read_chunk(&mut cx, 3) else {
panic!("empty chunks after data should start a new sequence");
};
assert_eq!(second, b"two"[..]);
}
#[tokio::test]
async fn http_reader_records_walk_dir_recv_bytes() {
let state = TestState::default();
@@ -2347,6 +2661,34 @@ mod tests {
handle.abort();
}
#[tokio::test]
async fn http_chunk_reader_surfaces_body_error_after_partial_data() {
let state = TestState::default();
let Some((base_url, handle)) = start_test_server(state).await else {
return;
};
let url = base_url.replace("/stream", "/fail-after-partial");
let mut reader = HttpChunkReader::new_with_stall_timeout(url, Method::GET, HeaderMap::new(), None, None)
.await
.expect("chunk reader should accept the successful response headers");
let chunk = std::future::poll_fn(|cx| Pin::new(&mut reader).poll_read_chunk(cx, 64))
.await
.expect("partial response bytes should arrive before the terminal error")
.expect("partial response should not be EOF");
let err = std::future::poll_fn(|cx| Pin::new(&mut reader).poll_read_chunk(cx, 64))
.await
.expect_err("terminal body errors must not become clean EOF");
assert_eq!(chunk, b"partial"[..]);
let source = err
.get_ref()
.and_then(|source| source.downcast_ref::<InternodeHttpError>())
.expect("body error should retain internode classification");
assert_eq!(source.kind(), InternodeHttpErrorKind::BodyStreamAborted);
handle.abort();
}
#[test]
fn classify_http_status_marks_retryable_gateway_errors() {
let unavailable = classify_http_status(reqwest::StatusCode::SERVICE_UNAVAILABLE);
+80 -2
View File
@@ -12,11 +12,11 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use blake2::{Blake2b512, Digest as Blake2Digest};
use blake2::Blake2b512;
use highway::{HighwayHash, HighwayHasher, Key};
use md5::Md5;
use serde::{Deserialize, Serialize};
use sha2::Sha256;
use sha2::{Digest, Sha256};
/// Magic HH-256 key: HH-256 hash of first 100 decimals of π as utf-8 with zero key.
const MAGIC_HIGHWAY_HASH256_KEY: [u8; 32] = [
@@ -140,6 +140,62 @@ impl HashAlgorithm {
}
}
/// Hash byte slices as one logical byte stream without concatenating them.
#[inline]
pub fn hash_encode_slices<'a, I>(&self, slices: I) -> impl AsRef<[u8]>
where
I: IntoIterator<Item = &'a [u8]>,
{
match self {
HashAlgorithm::Md5 => {
let mut hasher = Md5::new();
for slice in slices {
hasher.update(slice);
}
HashEncoded::Md5(hasher.finalize().into())
}
HashAlgorithm::HighwayHash256 => {
let mut hasher = HighwayHasher::new(MAGIC_HIGHWAY_HASH256_PARSED_KEY);
for slice in slices {
hasher.append(slice);
}
HashEncoded::HighwayHash256(u8x32_from_u64x4(hasher.finalize256()))
}
HashAlgorithm::SHA256 => {
let mut hasher = Sha256::new();
for slice in slices {
hasher.update(slice);
}
HashEncoded::Sha256(hasher.finalize().into())
}
HashAlgorithm::HighwayHash256S => {
let mut hasher = HighwayHasher::new(MAGIC_HIGHWAY_HASH256_PARSED_KEY);
for slice in slices {
hasher.append(slice);
}
HashEncoded::HighwayHash256S(u8x32_from_u64x4(hasher.finalize256()))
}
HashAlgorithm::HighwayHash256SLegacy => {
let mut hasher = HighwayHasher::new(LEGACY_HIGHWAY_HASH256_PARSED_KEY);
for slice in slices {
hasher.append(slice);
}
HashEncoded::HighwayHash256SLegacy(u8x32_from_u64x4(hasher.finalize256()))
}
HashAlgorithm::BLAKE2b512 => {
let mut hasher = Blake2b512::new();
for slice in slices {
hasher.update(slice);
}
let hash = hasher.finalize();
let mut out = [0u8; 64];
out.copy_from_slice(hash.as_ref());
HashEncoded::Blake2b512(out)
}
HashAlgorithm::None => HashEncoded::None,
}
}
/// Return the output size in bytes for the hash algorithm.
///
/// # Returns
@@ -222,6 +278,28 @@ mod tests {
assert_eq!(hash.len(), 0);
}
#[test]
fn hash_encode_slices_matches_contiguous_for_all_algorithms() {
let data = b"fragmented bitrot hash input";
let slices = [&data[..3], &data[3..11], &data[11..], &[]];
for algo in [
HashAlgorithm::Md5,
HashAlgorithm::SHA256,
HashAlgorithm::HighwayHash256,
HashAlgorithm::HighwayHash256S,
HashAlgorithm::HighwayHash256SLegacy,
HashAlgorithm::BLAKE2b512,
HashAlgorithm::None,
] {
assert_eq!(
algo.hash_encode_slices(slices).as_ref(),
algo.hash_encode(data).as_ref(),
"fragmented hash must match contiguous hash for {algo:?}"
);
}
}
#[test]
fn test_hash_encode_md5() {
let data = b"test data";