From 20f56af09c3cc9db16dc511b7827d533d0be4c11 Mon Sep 17 00:00:00 2001 From: houseme Date: Sat, 27 Jun 2026 21:46:38 +0800 Subject: [PATCH] feat(get): tune output response handoff (#3956) --- .../src/erasure/coding/decode_reader.rs | 16 ++- crates/io-metrics/src/lib.rs | 67 +++++++++++ rustfs/src/app/object_usecase.rs | 110 ++++++++++++++++-- rustfs/src/storage/sse.rs | 15 ++- scripts/run_get_codec_streaming_smoke.sh | 5 + 5 files changed, 203 insertions(+), 10 deletions(-) diff --git a/crates/ecstore/src/erasure/coding/decode_reader.rs b/crates/ecstore/src/erasure/coding/decode_reader.rs index c197b9007..ba4e6b734 100644 --- a/crates/ecstore/src/erasure/coding/decode_reader.rs +++ b/crates/ecstore/src/erasure/coding/decode_reader.rs @@ -405,6 +405,8 @@ where E: ErasureDecodeEngine + Clone + Send + Sync + 'static, { fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { + let filled_before_poll = buf.filled().len(); + loop { if self.output_pos < self.output_buf.len() { if self.prefetched_bufs.len() < self.fill_policy.max_inflight() @@ -431,7 +433,10 @@ where copy_start.elapsed().as_secs_f64(), ); } - return Poll::Ready(Ok(())); + if buf.remaining() == 0 { + return Poll::Ready(Ok(())); + } + continue; } if let Some(next_buf) = self.prefetched_bufs.pop_front() { @@ -441,6 +446,9 @@ where continue; } + if self.prefetch_error.is_some() && buf.filled().len() > filled_before_poll { + return Poll::Ready(Ok(())); + } if let Some(err) = self.prefetch_error.take() { return Poll::Ready(Err(err)); } @@ -452,7 +460,11 @@ where if self.output_wait_started_at.is_none() { self.output_wait_started_at = Some(Instant::now()); } - let prefetch = ready!(self.poll_prefetch(cx)); + let prefetch = match self.poll_prefetch(cx) { + Poll::Ready(result) => result, + Poll::Pending if buf.filled().len() > filled_before_poll => return Poll::Ready(Ok(())), + Poll::Pending => return Poll::Pending, + }; if let Some(started_at) = self.output_wait_started_at.take() { rustfs_io_metrics::record_get_object_fill_waited_by_output( self.metrics_path, diff --git a/crates/io-metrics/src/lib.rs b/crates/io-metrics/src/lib.rs index efba48a27..74f94d077 100644 --- a/crates/io-metrics/src/lib.rs +++ b/crates/io-metrics/src/lib.rs @@ -339,6 +339,71 @@ pub fn record_get_object_response_handoff( record_get_object_response_handoff_duration("s3_handler", duration_secs); } +/// Record ReaderStream capacity chosen for GetObject handoff. +#[inline(always)] +pub fn record_get_object_reader_stream_buffer_size(strategy: &str, buffer_source: &str, buffer_size_bytes: usize) { + if !get_stage_metrics_enabled() { + return; + } + histogram!( + "rustfs_io_get_object_reader_stream_buffer_size_bytes", + "strategy" => strategy.to_string(), + "buffer_source" => buffer_source.to_string() + ) + .record(usize_to_f64(buffer_size_bytes)); +} + +/// Record ReaderStream poll outcomes for GetObject handoff attribution. +#[inline(always)] +pub fn record_get_object_reader_stream_poll( + strategy: &str, + buffer_source: &str, + outcome: &'static str, + remaining_before: usize, + bytes: usize, + duration_secs: f64, +) { + if !get_stage_metrics_enabled() { + return; + } + let bytes = u64::try_from(bytes).unwrap_or(u64::MAX); + counter!( + "rustfs_io_get_object_reader_stream_poll_total", + "strategy" => strategy.to_string(), + "buffer_source" => buffer_source.to_string(), + "outcome" => outcome + ) + .increment(1); + counter!( + "rustfs_io_get_object_reader_stream_poll_bytes_total", + "strategy" => strategy.to_string(), + "buffer_source" => buffer_source.to_string(), + "outcome" => outcome + ) + .increment(bytes); + histogram!( + "rustfs_io_get_object_reader_stream_poll_remaining_bytes", + "strategy" => strategy.to_string(), + "buffer_source" => buffer_source.to_string(), + "outcome" => outcome + ) + .record(usize_to_f64(remaining_before)); + histogram!( + "rustfs_io_get_object_reader_stream_poll_bytes", + "strategy" => strategy.to_string(), + "buffer_source" => buffer_source.to_string(), + "outcome" => outcome + ) + .record(usize_to_f64(bytes as usize)); + histogram!( + "rustfs_io_get_object_reader_stream_poll_duration_seconds", + "strategy" => strategy.to_string(), + "buffer_source" => buffer_source.to_string(), + "outcome" => outcome + ) + .record(duration_secs); +} + /// Record I/O queue congestion observation. #[inline(always)] pub fn record_io_queue_congestion() { @@ -1548,6 +1613,8 @@ mod tests { record_get_object_fill_waited_by_output("codec_streaming", "single_inflight", 0.0003); record_get_object_fill_cancelled_on_drop("codec_streaming", "single_inflight"); record_get_object_reader_prefetch_bytes("codec_streaming", "single_inflight", 4096); + record_get_object_reader_stream_buffer_size("standard", "selected", 131072); + record_get_object_reader_stream_poll("standard", "selected", "ready_data", 8192, 4096, 0.0002); assert!(0.0003_f64.is_sign_positive()); } diff --git a/rustfs/src/app/object_usecase.rs b/rustfs/src/app/object_usecase.rs index b1ce47e24..24fe2bf3f 100644 --- a/rustfs/src/app/object_usecase.rs +++ b/rustfs/src/app/object_usecase.rs @@ -325,9 +325,16 @@ impl GetObjectStreamStrategy { const LARGE_SEQUENTIAL_GET_THRESHOLD_BYTES: i64 = 1024 * 1024 * 1024; const LARGE_SEQUENTIAL_GET_STREAM_BUFFER_CAP_BYTES: usize = 4 * MI_B; const LARGE_SEQUENTIAL_GET_READAHEAD_MULTIPLIER: usize = 2; +const LARGE_BODY_READER_STREAM_BUFFER_FLOOR_BYTES: usize = MI_B; +const LARGE_BODY_READER_STREAM_BUFFER_THRESHOLD_BYTES: i64 = 4 * MI_B as i64; const ENV_RUSTFS_GET_READER_STREAM_BUFFER_SIZE: &str = "RUSTFS_GET_READER_STREAM_BUFFER_SIZE"; +const ENV_RUSTFS_GET_OUTPUT_HANDOFF_ATTRIBUTION_ENABLE: &str = "RUSTFS_GET_OUTPUT_HANDOFF_ATTRIBUTION_ENABLE"; const GET_READER_STREAM_BUFFER_SOURCE_SELECTED: &str = "selected"; const GET_READER_STREAM_BUFFER_SOURCE_ENV_OVERRIDE: &str = "env_override"; +const GET_READER_STREAM_POLL_PENDING: &str = "pending"; +const GET_READER_STREAM_POLL_READY_DATA: &str = "ready_data"; +const GET_READER_STREAM_POLL_READY_EMPTY: &str = "ready_empty"; +const GET_READER_STREAM_POLL_READY_ERROR: &str = "ready_error"; fn get_reader_stream_buffer_size_override() -> Option { static GET_READER_STREAM_BUFFER_SIZE_OVERRIDE: OnceLock> = OnceLock::new(); @@ -339,6 +346,11 @@ fn get_reader_stream_buffer_size_override() -> Option { }) } +fn is_get_output_handoff_attribution_enabled() -> bool { + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| rustfs_utils::get_env_bool(ENV_RUSTFS_GET_OUTPUT_HANDOFF_ATTRIBUTION_ENABLE, false)) +} + fn resolve_reader_stream_buffer_size(selected_size: usize, override_size: Option) -> (usize, &'static str) { if let Some(override_size) = override_size.filter(|value| *value > 0) { return (override_size, GET_READER_STREAM_BUFFER_SOURCE_ENV_OVERRIDE); @@ -347,6 +359,20 @@ fn resolve_reader_stream_buffer_size(selected_size: usize, override_size: Option (selected_size.max(1), GET_READER_STREAM_BUFFER_SOURCE_SELECTED) } +fn tune_reader_stream_buffer_size( + selected_size: usize, + response_content_length: i64, + stream_strategy: GetObjectStreamStrategy, +) -> usize { + if stream_strategy == GetObjectStreamStrategy::Standard + && response_content_length >= LARGE_BODY_READER_STREAM_BUFFER_THRESHOLD_BYTES + { + return selected_size.max(LARGE_BODY_READER_STREAM_BUFFER_FLOOR_BYTES); + } + + selected_size +} + async fn enqueue_transitioned_delete_cleanup( store: Arc, bucket: &str, @@ -414,6 +440,8 @@ pin_project! { struct GetObjectReaderStream { #[pin] inner: ReaderStream, + strategy: &'static str, + buffer_source: &'static str, remaining: usize, emitted: usize, expected: usize, @@ -434,9 +462,14 @@ impl GetObjectReaderStream where R: AsyncRead, { - fn new(reader: R, capacity: usize, remaining: usize) -> Self { + fn new(reader: R, capacity: usize, remaining: usize, strategy: &'static str, buffer_source: &'static str) -> Self { + if is_get_output_handoff_attribution_enabled() { + rustfs_io_metrics::record_get_object_reader_stream_buffer_size(strategy, buffer_source, capacity); + } Self { inner: ReaderStream::with_capacity(reader, capacity), + strategy, + buffer_source, remaining, emitted: 0, expected: remaining, @@ -470,7 +503,9 @@ where return Poll::Ready(None); } - match this.inner.as_mut().poll_next(cx) { + let remaining_before = *this.remaining; + let poll_start = std::time::Instant::now(); + let result: Poll> = match this.inner.as_mut().poll_next(cx) { Poll::Ready(Some(Ok(mut bytes))) => { if bytes.len() > *this.remaining { bytes.truncate(*this.remaining); @@ -500,11 +535,34 @@ where error = %err, "GetObject ReaderStream returned error" ); - Poll::Ready(Some(Err(Box::new(err)))) + Poll::Ready(Some(Err(Box::new(err) as S3StdError))) } Poll::Ready(None) => Poll::Ready(None), Poll::Pending => Poll::Pending, + }; + + let emitted_bytes = match &result { + Poll::Ready(Some(Ok(bytes))) => bytes.len(), + _ => 0, + }; + let outcome = match &result { + Poll::Ready(Some(Ok(bytes))) if !bytes.is_empty() => GET_READER_STREAM_POLL_READY_DATA, + Poll::Ready(Some(Ok(_))) | Poll::Ready(None) => GET_READER_STREAM_POLL_READY_EMPTY, + Poll::Ready(Some(Err(_))) => GET_READER_STREAM_POLL_READY_ERROR, + Poll::Pending => GET_READER_STREAM_POLL_PENDING, + }; + if is_get_output_handoff_attribution_enabled() { + rustfs_io_metrics::record_get_object_reader_stream_poll( + this.strategy, + this.buffer_source, + outcome, + remaining_before, + emitted_bytes, + poll_start.elapsed().as_secs_f64(), + ); } + + result } fn size_hint(&self) -> (usize, Option) { @@ -1916,8 +1974,10 @@ impl DefaultObjectUsecase { R: AsyncRead + Send + Sync + Unpin + 'static, { let expected = usize::try_from(response_content_length.max(0)).unwrap_or(usize::MAX); + let tuned_stream_buffer_size = + tune_reader_stream_buffer_size(stream_buffer_size, response_content_length, stream_strategy); let (stream_buffer_size, buffer_source) = - resolve_reader_stream_buffer_size(stream_buffer_size, get_reader_stream_buffer_size_override()); + resolve_reader_stream_buffer_size(tuned_stream_buffer_size, get_reader_stream_buffer_size_override()); let get_stage_metrics_enabled = rustfs_io_metrics::get_stage_metrics_enabled(); if get_stage_metrics_enabled { rustfs_io_metrics::record_get_object_stream_strategy( @@ -1928,7 +1988,7 @@ impl DefaultObjectUsecase { } let handoff_start = get_stage_metrics_enabled.then(std::time::Instant::now); let reader = GetObjectStreamingReader::new(reader, bucket, key, expected, get_object_disk_read_timeout()); - let stream = GetObjectReaderStream::new(reader, stream_buffer_size, expected); + let stream = GetObjectReaderStream::new(reader, stream_buffer_size, expected, stream_strategy.as_str(), buffer_source); let blob = StreamingBlob::new(stream); if let Some(handoff_start) = handoff_start { rustfs_io_metrics::record_get_object_response_handoff( @@ -6106,6 +6166,30 @@ mod tests { assert_eq!(small_buffer_size, 512 * 1024); } + #[test] + fn tune_reader_stream_buffer_size_raises_large_standard_streams_only() { + assert_eq!( + tune_reader_stream_buffer_size(128 * 1024, 10 * MI_B as i64, GetObjectStreamStrategy::Standard), + LARGE_BODY_READER_STREAM_BUFFER_FLOOR_BYTES + ); + assert_eq!( + tune_reader_stream_buffer_size(512 * 1024, 10 * MI_B as i64, GetObjectStreamStrategy::Standard), + LARGE_BODY_READER_STREAM_BUFFER_FLOOR_BYTES + ); + assert_eq!( + tune_reader_stream_buffer_size(2 * MI_B, 10 * MI_B as i64, GetObjectStreamStrategy::Standard), + 2 * MI_B + ); + assert_eq!( + tune_reader_stream_buffer_size(128 * 1024, MI_B as i64, GetObjectStreamStrategy::Standard), + 128 * 1024 + ); + assert_eq!( + tune_reader_stream_buffer_size(128 * 1024, 10 * MI_B as i64, GetObjectStreamStrategy::LargeSequentialReadahead), + 128 * 1024 + ); + } + #[test] fn resolve_reader_stream_buffer_size_keeps_selected_default() { let (buffer_size, source) = resolve_reader_stream_buffer_size(128 * 1024, None); @@ -6273,7 +6357,13 @@ mod tests { #[tokio::test] async fn get_object_reader_stream_tracks_remaining_length() { - let mut stream = GetObjectReaderStream::new(std::io::Cursor::new(b"hello".to_vec()), 2, 5); + let mut stream = GetObjectReaderStream::new( + std::io::Cursor::new(b"hello".to_vec()), + 2, + 5, + GetObjectStreamStrategy::Standard.as_str(), + GET_READER_STREAM_BUFFER_SOURCE_SELECTED, + ); assert_eq!(stream.remaining_length().exact(), Some(5)); @@ -6289,7 +6379,13 @@ mod tests { #[tokio::test] async fn get_object_reader_stream_truncates_to_expected_length() { - let stream = GetObjectReaderStream::new(std::io::Cursor::new(b"hello!".to_vec()), 64, 5); + let stream = GetObjectReaderStream::new( + std::io::Cursor::new(b"hello!".to_vec()), + 64, + 5, + GetObjectStreamStrategy::Standard.as_str(), + GET_READER_STREAM_BUFFER_SOURCE_SELECTED, + ); let chunks = stream .collect::>() diff --git a/rustfs/src/storage/sse.rs b/rustfs/src/storage/sse.rs index edc529d41..b2df5ae13 100644 --- a/rustfs/src/storage/sse.rs +++ b/rustfs/src/storage/sse.rs @@ -2157,6 +2157,14 @@ pub fn reset_sse_dek_provider() { } } +#[cfg(test)] +#[cfg(feature = "rio-v2")] +pub fn set_sse_dek_provider_for_test(provider: Arc) { + if let Ok(mut slot) = GLOBAL_SSE_DEK_PROVIDER.write() { + *slot = Some(provider); + } +} + // ============================================================================ // Legacy Functions (SSE-S3 / SSE-KMS) // ============================================================================ @@ -2462,7 +2470,7 @@ mod tests { use s3s::S3ErrorCode; use s3s::dto::{SSECustomerAlgorithm, SSECustomerKey, SSECustomerKeyMD5, ServerSideEncryption}; use std::collections::HashMap; - use std::sync::OnceLock; + use std::sync::{Arc, OnceLock}; use temp_env::async_with_vars; use tokio::sync::Mutex; @@ -2994,6 +3002,11 @@ mod tests { .await .expect("kms test key should be created"); + let provider = KmsSseDekProvider::new_with_service_manager(manager.clone()) + .await + .expect("kms provider should initialize from the configured test manager"); + super::set_sse_dek_provider_for_test(Arc::new(provider)); + let client_context = HashMap::from([("tenant".to_string(), "alpha".to_string())]); let request = EncryptionRequest { bucket: "bucket", diff --git a/scripts/run_get_codec_streaming_smoke.sh b/scripts/run_get_codec_streaming_smoke.sh index d67448e68..a552e9606 100755 --- a/scripts/run_get_codec_streaming_smoke.sh +++ b/scripts/run_get_codec_streaming_smoke.sh @@ -26,6 +26,7 @@ ROUND_COOLDOWN_SECS=0 MODE="both" CODEC_ENGINES="legacy" CODEC_MAX_INFLIGHT=1 +OUTPUT_HANDOFF_ATTRIBUTION=false OUT_DIR="" RUSTFS_BIN="${PROJECT_ROOT}/target/release/rustfs" WARP_BIN="warp" @@ -53,6 +54,7 @@ Core options: --mode Which profile(s) to run (default: both) --codec-engine Codec engine(s) for codec profiles (default: legacy) --codec-max-inflight RUSTFS_GET_CODEC_STREAMING_MAX_INFLIGHT (default: 1) + --handoff-attribution Enable output handoff attribution metrics --address RustFS listen address (default: 127.0.0.1:19030) --bucket Benchmark bucket (default: rustfs-get-codec-smoke) --sizes Object sizes (default: 1MiB,4MiB,10MiB) @@ -141,6 +143,7 @@ parse_args() { --mode) MODE="$2"; shift 2 ;; --codec-engine) CODEC_ENGINES="$2"; shift 2 ;; --codec-max-inflight) CODEC_MAX_INFLIGHT="$2"; shift 2 ;; + --handoff-attribution) OUTPUT_HANDOFF_ATTRIBUTION=true; shift ;; --address) ADDRESS="$2"; shift 2 ;; --bucket) BUCKET="$2"; shift 2 ;; --sizes) SIZES="$2"; shift 2 ;; @@ -306,6 +309,7 @@ rustfs_volumes=${volumes} RUSTFS_GET_CODEC_STREAMING_ENABLE=${codec_enabled} RUSTFS_GET_CODEC_STREAMING_ENGINE=${codec_engine} RUSTFS_GET_CODEC_STREAMING_MAX_INFLIGHT=${CODEC_MAX_INFLIGHT} +RUSTFS_GET_OUTPUT_HANDOFF_ATTRIBUTION_ENABLE=${OUTPUT_HANDOFF_ATTRIBUTION} RUSTFS_GET_CODEC_STREAMING_MIN_SIZE=${CODEC_MIN_SIZE} metrics_path=${metrics_path} RUSTFS_SCANNER_ENABLED=false @@ -384,6 +388,7 @@ start_server() { export RUSTFS_GET_CODEC_STREAMING_ENABLE="$codec_enabled" export RUSTFS_GET_CODEC_STREAMING_ENGINE="$codec_engine" export RUSTFS_GET_CODEC_STREAMING_MAX_INFLIGHT="$CODEC_MAX_INFLIGHT" + export RUSTFS_GET_OUTPUT_HANDOFF_ATTRIBUTION_ENABLE="$OUTPUT_HANDOFF_ATTRIBUTION" export RUSTFS_GET_CODEC_STREAMING_MIN_SIZE="$CODEC_MIN_SIZE" export RUSTFS_REGION="$REGION" export RUSTFS_RPC_SECRET="rustfs-get-codec-smoke-rpc-secret"