diff --git a/rustfs/src/storage/rpc/http_service.rs b/rustfs/src/storage/rpc/http_service.rs index 0e2cfe61e..bdd54fe12 100644 --- a/rustfs/src/storage/rpc/http_service.rs +++ b/rustfs/src/storage/rpc/http_service.rs @@ -39,7 +39,6 @@ use rustfs_io_metrics::internode_metrics::{ INTERNODE_OPERATION_NS_SCANNER, INTERNODE_OPERATION_PUT_FILE_CAPABILITY, INTERNODE_OPERATION_PUT_FILE_STREAM, INTERNODE_OPERATION_READ_FILE_STREAM, INTERNODE_OPERATION_WALK_DIR, INTERNODE_TRANSPORT_BACKEND_TCP_HTTP, }; -use rustfs_utils::net::bytes_stream; use s3s::Body; use s3s::dto::StreamingBlob; use serde::de::DeserializeOwned; @@ -51,7 +50,7 @@ use std::pin::Pin; use std::sync::{Arc, LazyLock, Weak}; use std::task::{Context, Poll}; use std::time::{Duration, Instant}; -use tokio::io::{self, AsyncWriteExt}; +use tokio::io::{self, AsyncReadExt, AsyncWriteExt}; use tokio::sync::{Mutex, oneshot}; use tokio_util::{io::ReaderStream, sync::CancellationToken}; use tower::Service; @@ -663,15 +662,24 @@ where R: tokio::io::AsyncRead + Unpin + Send + Sync + 'static, { let metrics = runtime_sources::current_internode_metrics(); - let stream = ReaderStream::with_capacity(reader, DEFAULT_READ_BUFFER_SIZE).map_ok(move |bytes| { + let read_buffer_size = read_file_stream_buffer_size(length); + let read_limit = if length == 0 { + u64::MAX + } else { + u64::try_from(length).unwrap_or(u64::MAX) + }; + let stream = ReaderStream::with_capacity(reader.take(read_limit), read_buffer_size).map_ok(move |bytes| { metrics.record_sent_bytes_for_operation_and_backend(operation, INTERNODE_TRANSPORT_BACKEND_TCP_HTTP, bytes.len()); bytes }); + Box::pin(stream) +} +fn read_file_stream_buffer_size(length: usize) -> usize { if length == 0 { - Box::pin(stream) + DEFAULT_READ_BUFFER_SIZE } else { - Box::pin(bytes_stream(stream, length)) + length.min(DEFAULT_READ_BUFFER_SIZE) } } @@ -1664,7 +1672,7 @@ fn put_file_stage_error_message(stage: &str, query: &PutFileQuery, err: &dyn std #[cfg(test)] mod tests { use super::{ - DiskError, InternodeRpcService, LOG_SUBSYSTEM_DIRECTORY_WALK, LOG_SUBSYSTEM_FILE_TRANSFER, + DEFAULT_READ_BUFFER_SIZE, DiskError, InternodeRpcService, LOG_SUBSYSTEM_DIRECTORY_WALK, LOG_SUBSYSTEM_FILE_TRANSFER, LOG_SUBSYSTEM_NAMESPACE_SCANNER, LOG_SUBSYSTEM_ROUTING, NS_SCANNER_BODY_SHA256_QUERY, NS_SCANNER_CAPABILITY_CHALLENGE_QUERY, NS_SCANNER_CYCLE_QUERY, NS_SCANNER_LEADER_EPOCH_QUERY, NS_SCANNER_PATH, NS_SCANNER_REQUEST_ID_QUERY, NS_SCANNER_SERVER_EPOCH_QUERY, NS_SCANNER_SESSION_ID_QUERY, @@ -1673,10 +1681,10 @@ mod tests { append_walk_dir_completion, internode_http_operation, internode_rpc_subsystem, is_internode_rpc_path, ns_scanner_response_body, ns_scanner_server_epoch_matches, put_body_size_mismatch, put_file_auth_nonce, put_file_capability_response, put_file_server_epoch_matches, put_file_stage_error_message, put_file_target_lock, - read_file_body_stream, remote_scanner_claim_rejection, response_with_disk_error, supports_walk_dir_stream_completion, - validate_walk_dir_completion_request, verify_internode_rpc_signature, verify_ns_scanner_body_digest, - verify_walk_dir_body_digest, walk_dir_response_body, write_authenticated_put_file, write_body_chunks_to_writer, - write_put_file_body_chunks_to_writer, + read_file_body_stream, read_file_stream_buffer_size, remote_scanner_claim_rejection, response_with_disk_error, + supports_walk_dir_stream_completion, validate_walk_dir_completion_request, verify_internode_rpc_signature, + verify_ns_scanner_body_digest, verify_walk_dir_body_digest, walk_dir_response_body, write_authenticated_put_file, + write_body_chunks_to_writer, write_put_file_body_chunks_to_writer, }; use crate::storage::storage_api::ecstore_rpc::{build_put_file_auth_trailer, gen_signature_headers}; use crate::storage::storage_api::rpc_consumer::http_service::{DiskAPI as _, DiskOption, DiskStore, Endpoint, new_disk}; @@ -1704,6 +1712,23 @@ mod tests { use tokio_stream::StreamExt; use tokio_stream::iter; + struct RejectExtraPollReader { + emitted: bool, + } + + impl io::AsyncRead for RejectExtraPollReader { + fn poll_read(mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut io::ReadBuf<'_>) -> Poll> { + if self.emitted { + return Poll::Ready(Err(io::Error::other("reader polled past the requested length"))); + } + + let bytes = b"hello world"; + buf.put_slice(&bytes[..buf.remaining().min(bytes.len())]); + self.emitted = true; + Poll::Ready(Ok(())) + } + } + async fn new_put_file_test_disk() -> (DiskStore, tempfile::TempDir) { let dir = tempfile::tempdir().expect("temp directory should be created"); let endpoint = @@ -2842,11 +2867,7 @@ mod tests { #[tokio::test] async fn read_file_body_stream_truncates_to_requested_length() { - let (reader, mut writer) = tokio::io::duplex(64); - tokio::spawn(async move { - writer.write_all(b"hello world").await.expect("write succeeds"); - }); - + let reader = RejectExtraPollReader { emitted: false }; let mut stream = read_file_body_stream(reader, 5, INTERNODE_OPERATION_READ_FILE_STREAM); let mut out = Vec::new(); while let Some(chunk) = stream.next().await { @@ -2856,6 +2877,18 @@ mod tests { assert_eq!(out, b"hello"); } + #[test] + fn read_file_body_stream_sizes_buffer_to_requested_length() { + for (length, expected_capacity) in [ + (0, DEFAULT_READ_BUFFER_SIZE), + (40 * 1024, 40 * 1024), + (DEFAULT_READ_BUFFER_SIZE, DEFAULT_READ_BUFFER_SIZE), + (DEFAULT_READ_BUFFER_SIZE + 1, DEFAULT_READ_BUFFER_SIZE), + ] { + assert_eq!(read_file_stream_buffer_size(length), expected_capacity); + } + } + #[test] fn read_file_error_response_marks_only_missing_disk_errors() { for (error, expected) in [