feat(get): tune output response handoff (#3956)

This commit is contained in:
houseme
2026-06-27 21:46:38 +08:00
committed by GitHub
parent 68d5d1d41d
commit 20f56af09c
5 changed files with 203 additions and 10 deletions
@@ -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<io::Result<()>> {
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,
+67
View File
@@ -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());
}
+103 -7
View File
@@ -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<usize> {
static GET_READER_STREAM_BUFFER_SIZE_OVERRIDE: OnceLock<Option<usize>> = OnceLock::new();
@@ -339,6 +346,11 @@ fn get_reader_stream_buffer_size_override() -> Option<usize> {
})
}
fn is_get_output_handoff_attribution_enabled() -> bool {
static ENABLED: OnceLock<bool> = 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>) -> (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<ECStore>,
bucket: &str,
@@ -414,6 +440,8 @@ pin_project! {
struct GetObjectReaderStream<R> {
#[pin]
inner: ReaderStream<R>,
strategy: &'static str,
buffer_source: &'static str,
remaining: usize,
emitted: usize,
expected: usize,
@@ -434,9 +462,14 @@ impl<R> GetObjectReaderStream<R>
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<Option<Self::Item>> = 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<usize>) {
@@ -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::<Vec<_>>()
+14 -1
View File
@@ -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<dyn SseDekProvider>) {
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",
+5
View File
@@ -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 <legacy|codec|both> Which profile(s) to run (default: both)
--codec-engine <csv> Codec engine(s) for codec profiles (default: legacy)
--codec-max-inflight <n> RUSTFS_GET_CODEC_STREAMING_MAX_INFLIGHT (default: 1)
--handoff-attribution Enable output handoff attribution metrics
--address <host:port> RustFS listen address (default: 127.0.0.1:19030)
--bucket <name> Benchmark bucket (default: rustfs-get-codec-smoke)
--sizes <csv> 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"