diff --git a/crates/ecstore/src/cluster/rpc/internode_data_transport.rs b/crates/ecstore/src/cluster/rpc/internode_data_transport.rs index b57e6e4ae..a9be2f95e 100644 --- a/crates/ecstore/src/cluster/rpc/internode_data_transport.rs +++ b/crates/ecstore/src/cluster/rpc/internode_data_transport.rs @@ -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; + /// 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> { + Ok(None) + } async fn open_write(&self, request: WriteStreamRequest) -> Result; async fn open_walk_dir(&self, request: WalkDirStreamRequest) -> Result; async fn open_ns_scanner(&self, _request: NsScannerStreamRequest) -> Result { @@ -247,6 +252,15 @@ impl InternodeDataTransport for TcpHttpInternodeDataTransport { )) } + async fn open_read_chunks(&self, request: ReadStreamRequest) -> Result> { + 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 { let server_epoch = self.put_file_auth_capability(&request.endpoint).await?; let nonce = server_epoch.map(|_| Uuid::new_v4()); diff --git a/crates/ecstore/src/cluster/rpc/remote_disk.rs b/crates/ecstore/src/cluster/rpc/remote_disk.rs index 7a2c1a3a8..9b78fc846 100644 --- a/crates/ecstore/src/cluster/rpc/remote_disk.rs +++ b/crates/ecstore/src/cluster/rpc/remote_disk.rs @@ -522,6 +522,33 @@ impl RemoteDisk { } } + async fn open_read_chunks_with_retry(&self, request: ReadStreamRequest) -> Result> { + 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> { + 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)] diff --git a/crates/ecstore/src/disk/disk_store.rs b/crates/ecstore/src/disk/disk_store.rs index fee70bc97..deb40a5cb 100644 --- a/crates/ecstore/src/disk/disk_store.rs +++ b/crates/ecstore/src/disk/disk_store.rs @@ -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> { + 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 { self.track_disk_health_with_op( "read_file_mmap_copy", diff --git a/crates/ecstore/src/disk/mod.rs b/crates/ecstore/src/disk/mod.rs index 0f9c28bee..26047fdfc 100644 --- a/crates/ecstore/src/disk/mod.rs +++ b/crates/ecstore/src/disk/mod.rs @@ -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> { + 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 { match self { @@ -865,6 +879,18 @@ pub trait DiskAPI: Debug + Send + Sync + 'static { async fn read_file(&self, volume: &str, path: &str) -> Result; async fn read_file_stream(&self, volume: &str, path: &str, offset: usize, length: usize) -> Result; + /// 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> { + 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; diff --git a/crates/ecstore/src/erasure/coding/bitrot.rs b/crates/ecstore/src/erasure/coding/bitrot.rs index a12529588..e1cfe3971 100644 --- a/crates/ecstore/src/erasure/coding/bitrot.rs +++ b/crates/ecstore/src/erasure/coding/bitrot.rs @@ -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 { 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> { + 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, + // Reused owned chunk vector for the remote HTTP fast path. Keeping the + // allocation with the reader avoids allocating once per bitrot block. + chunks: Vec, 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, + } + + impl FragmentedSource { + fn new(bytes: Vec, 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> { + 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> { + 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, + } + + impl GeneratedChunkSource { + fn new(bytes: Vec, fragment_size: usize) -> Self { + assert!(fragment_size > 0); + Self { + bytes: Bytes::from(bytes), + offset: 0, + fragment_size, + fail_at: None, + } + } + + fn failing(bytes: Vec, 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> { + 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> { + 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> { + 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> { + 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, + saw_reused_scratch: bool, + } + + impl AsyncRead for ScratchReuseSource { + fn poll_read(mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { + 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 = (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 = 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 = 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 = (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::()) + .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]); + } + } } diff --git a/crates/ecstore/src/io_support/bitrot.rs b/crates/ecstore/src/io_support/bitrot.rs index 2e046807d..5a3120760 100644 --- a/crates/ecstore/src/io_support/bitrot.rs +++ b/crates/ecstore/src/io_support/bitrot.rs @@ -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), + Chunked(ChunkReaderBox), Stream(Box), } @@ -58,6 +60,7 @@ impl AsyncRead for ShardReader { fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut tokio::io::ReadBuf<'_>) -> Poll> { 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 { 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> { + 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, + } + + 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> { + 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>> { + 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); + } } diff --git a/crates/rio/src/http_reader.rs b/crates/rio/src/http_reader.rs index 531c4edfb..8be68cef4 100644 --- a/crates/rio/src/http_reader.rs +++ b/crates/rio/src/http_reader.rs @@ -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> + 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>>; +} + +pub type ChunkReaderBox = Box; + +struct HttpReaderInit { + stream: HttpByteStream, + track_internode_metrics: bool, + internode_operation: Option<&'static str>, + stall_timeout: Option, + 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>+Send+Sync>>, Bytes>, + inner: StreamReader, + } +} + +pin_project! { + pub struct HttpChunkReader { + track_internode_metrics: bool, + internode_operation: Option<&'static str>, + stall_timeout: Option, + stall_timer: Option>>, + request_started: Instant, + duration_recorded: bool, + consecutive_empty_chunks: usize, + #[pin] + inner: HttpByteStream, + current: Option, } } @@ -934,12 +971,34 @@ impl HttpReader { _read_buf_size: usize, stall_timeout: Option, ) -> io::Result { - 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>, + stall_timeout: Option, + ) -> io::Result { + 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> { 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>, + stall_timeout: Option, + ) -> io::Result { + 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> { + 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>> { + 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::()) + .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); diff --git a/crates/utils/src/hash.rs b/crates/utils/src/hash.rs index b2f6f4eba..6720a1cf3 100644 --- a/crates/utils/src/hash.rs +++ b/crates/utils/src/hash.rs @@ -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, + { + 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";