mirror of
https://github.com/rustfs/rustfs.git
synced 2026-08-13 16:46:55 +00:00
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:
@@ -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)]
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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>;
|
||||
|
||||
|
||||
@@ -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]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
@@ -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);
|
||||
|
||||
@@ -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";
|
||||
|
||||
Reference in New Issue
Block a user