diff --git a/rustfs/src/app/object/extract.rs b/rustfs/src/app/object/extract.rs index 12f367a13..b773e37a5 100644 --- a/rustfs/src/app/object/extract.rs +++ b/rustfs/src/app/object/extract.rs @@ -63,6 +63,14 @@ pin_project! { } } +pin_project! { + struct ExtractArchiveDecoderReader { + #[pin] + inner: R, + _permit: OwnedSemaphorePermit, + } +} + #[derive(Debug, Default)] struct ExtractArchiveUploadState { etag: Option, @@ -100,6 +108,21 @@ impl ExtractArchiveEtagReader { } } +impl ExtractArchiveDecoderReader { + fn new(inner: R, permit: OwnedSemaphorePermit) -> Self { + Self { inner, _permit: permit } + } +} + +impl AsyncRead for ExtractArchiveDecoderReader +where + R: AsyncRead, +{ + fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { + self.project().inner.poll_read(cx, buf) + } +} + fn extract_archive_incomplete_body(remaining: u64) -> std::io::Error { let Ok(remaining) = i64::try_from(remaining) else { return std::io::Error::new(std::io::ErrorKind::InvalidData, "archive remaining body length exceeds i64"); @@ -510,6 +533,36 @@ fn try_acquire_extract_staging_permit(manager: &ConcurrencyManager, staging_weig }) } +async fn build_admitted_extract_archive_decoder( + manager: &ConcurrencyManager, + key: &str, + tracked_archive: R, +) -> S3Result>> +where + R: AsyncRead + Send + Unpin + 'static, +{ + // Admission precedes stream inspection so saturation cannot allocate or + // drive another codec. The returned reader owns the permit through archive + // finalization and transport-length validation. + let permit = manager.try_acquire_snowball_archive_decoder().ok_or_else(|| { + object_s3_error( + S3ErrorCode::SlowDown, + "Snowball archive decoder limit reached, please reduce your request rate", + ) + })?; + let (detected_archive_format, sniffed_archive) = + CompressionFormat::sniff(tracked_archive).await.map_err(|err| match err { + ZipError::InspectStream(source) => map_extract_archive_error(source), + _ => s3_error!(InvalidArgument, "Failed to detect archive compression"), + })?; + let archive_format = resolve_extract_archive_format(key, detected_archive_format); + let decoder = archive_format.get_decoder(sniffed_archive).map_err(|e| { + error!(error = ?e, "Archive decoder creation failed"); + s3_error!(InvalidArgument, "get_decoder err") + })?; + Ok(ExtractArchiveDecoderReader::new(decoder, permit)) +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum ExtractBatchAction { Stage, @@ -2067,16 +2120,7 @@ impl DefaultObjectUsecase { let extract_limits = put_object_extract_limits(); let tracked_archive = ExtractArchiveEtagReader::new(archive_reader, expected_archive_length, archive_upload_state.clone()); - let (detected_archive_format, sniffed_archive) = - CompressionFormat::sniff(tracked_archive).await.map_err(|err| match err { - ZipError::InspectStream(source) => map_extract_archive_error(source), - _ => s3_error!(InvalidArgument, "Failed to detect archive compression"), - })?; - let archive_format = resolve_extract_archive_format(&key, detected_archive_format); - let decoder = archive_format.get_decoder(sniffed_archive).map_err(|e| { - error!(error = ?e, "Archive decoder creation failed"); - s3_error!(InvalidArgument, "get_decoder err") - })?; + let decoder = build_admitted_extract_archive_decoder(get_concurrency_manager(), &key, tracked_archive).await?; let decoder = ExtractDecodedLimitReader::new(decoder, extract_limits.max_decoded_size); let mut ar = build_put_object_extract_archive(decoder, extract_limits); @@ -2621,6 +2665,9 @@ impl DefaultObjectUsecase { } state.etag.as_ref().map(|etag| to_s3s_etag(etag)) }; + // Keep decoder admission through body-complete validation, then release + // it before response checksum and completion bookkeeping. + drop(decoder); apply_trailing_checksums( input.checksum_algorithm.as_ref().map(|a| a.as_str()), &req.trailing_headers, @@ -2693,6 +2740,120 @@ mod tests { }) } + #[tokio::test] + async fn snowball_archive_decoder_admission_is_global_and_lifetime_bound() { + struct PanicOnRead; + struct ErrorOnRead; + + impl AsyncRead for PanicOnRead { + fn poll_read(self: Pin<&mut Self>, _cx: &mut Context<'_>, _buf: &mut ReadBuf<'_>) -> Poll> { + panic!("a saturated decoder admission must not inspect the archive body") + } + } + + impl AsyncRead for ErrorOnRead { + fn poll_read(self: Pin<&mut Self>, _cx: &mut Context<'_>, _buf: &mut ReadBuf<'_>) -> Poll> { + Poll::Ready(Err(std::io::Error::other("injected decoder source failure"))) + } + } + + let manager = ConcurrencyManager::new(); + let clone = manager.clone(); + let mut held = Vec::new(); + while let Some(permit) = manager.try_acquire_snowball_archive_decoder() { + held.push(permit); + } + assert!(!held.is_empty(), "the global decoder gate must admit at least one archive"); + + let error = match build_admitted_extract_archive_decoder(&clone, "archive.tar", PanicOnRead).await { + Ok(_) => panic!("a saturated decoder gate must reject without constructing another decoder"), + Err(error) => error, + }; + assert_eq!(error.code(), &S3ErrorCode::SlowDown); + + drop( + held.pop() + .expect("one decoder permit must be available for the lifetime test"), + ); + let error = match build_admitted_extract_archive_decoder(&clone, "archive.tar", ErrorOnRead).await { + Ok(_) => panic!("archive inspection failure must remain an error"), + Err(error) => error, + }; + assert_eq!(error.code(), &S3ErrorCode::InvalidArgument); + let released_after_error = clone + .try_acquire_snowball_archive_decoder() + .expect("archive inspection failure must release decoder capacity"); + drop(released_after_error); + + let mut builder = Builder::new(Vec::new()); + let mut header = Header::new_gnu(); + header.set_size(0); + header.set_cksum(); + builder + .append_data(&mut header, "member.txt", &b""[..]) + .await + .expect("decoder lifetime fixture should append its member"); + let archive_bytes = builder.into_inner().await.expect("decoder lifetime fixture should finalize"); + let expected_length = u64::try_from(archive_bytes.len()).expect("fixture length must fit u64"); + let upload_state = Arc::new(Mutex::new(ExtractArchiveUploadState::default())); + let tracked_archive = + ExtractArchiveEtagReader::new(std::io::Cursor::new(archive_bytes), expected_length, upload_state.clone()); + let decoder = build_admitted_extract_archive_decoder(&clone, "archive.tar", tracked_archive) + .await + .expect("released decoder capacity must be reusable"); + assert!( + manager.try_acquire_snowball_archive_decoder().is_none(), + "the decoder reader must retain admission while it is active" + ); + + let extract_limits = put_object_extract_limits(); + let decoder = ExtractDecodedLimitReader::new(decoder, extract_limits.max_decoded_size); + let mut archive = build_put_object_extract_archive(decoder, extract_limits); + let mut entries = archive.entries().expect("admitted archive entries should be readable"); + let entry = entries + .next() + .await + .expect("admitted archive should contain its member") + .expect("admitted archive member should parse"); + assert_eq!(entry.path_bytes().expect("archive member path should parse").as_ref(), b"member.txt"); + drop(entry); + assert!(entries.next().await.is_none(), "admitted archive should contain one member"); + drop(entries); + let mut decoder = match archive.into_inner() { + Ok(decoder) => decoder, + Err(_) => panic!("admitted archive should finalize"), + }; + tokio::io::copy(&mut decoder, &mut tokio::io::sink()) + .await + .expect("admitted archive should consume its remaining transport body"); + assert!( + upload_state + .lock() + .expect("archive upload state lock must remain healthy") + .body_complete, + "transport-length validation must complete while decoder admission is held" + ); + assert!( + manager.try_acquire_snowball_archive_decoder().is_none(), + "archive finalization and transport validation must retain decoder admission" + ); + drop(decoder); + assert!( + clone.try_acquire_snowball_archive_decoder().is_some(), + "dropping the finalized decoder must release its global slot" + ); + + let cancelled = + build_admitted_extract_archive_decoder(&manager, "archive.tar", std::io::Cursor::new(b"cancelled".to_vec())) + .await + .expect("the decoder gate must remain reusable"); + drop(cancelled); + assert!( + clone.try_acquire_snowball_archive_decoder().is_some(), + "dropping an unfinished decoder must release admission for cancellation" + ); + } + #[test] fn snowball_max_inflight_has_a_serial_compatibility_floor_and_bounded_ceiling() { assert_eq!(EXTRACT_DEFAULT_MAX_INFLIGHT, 1); diff --git a/rustfs/src/storage/concurrency/manager.rs b/rustfs/src/storage/concurrency/manager.rs index 779f9b5c0..68d450d65 100644 --- a/rustfs/src/storage/concurrency/manager.rs +++ b/rustfs/src/storage/concurrency/manager.rs @@ -35,6 +35,10 @@ use tokio::sync::{OwnedSemaphorePermit, Semaphore}; use tracing::debug; const DERIVED_LARGE_PUT_ADMISSION_LIMIT_MAX: usize = 32; +// Framed S2 alone can retain one encoded and one decoded block of roughly +// 4 MiB each, while other codecs have their own larger windows. Four keeps +// useful request parallelism without scaling codec memory and CPU with clients. +const SNOWBALL_ARCHIVE_DECODER_LIMIT: usize = 4; pub(crate) const SNOWBALL_MEMBER_COMMIT_LIMIT: usize = 32; pub(crate) const SNOWBALL_STAGING_BYTES_LIMIT: usize = 4 * MI_B; @@ -71,6 +75,8 @@ pub struct ConcurrencyManager { metrics_collector: Arc, /// Foreground write admission policy, resolved once at startup. foreground_write_admission_policy: ForegroundWriteAdmissionPolicy, + /// Bounds active Snowball archive inspection and decoding across requests. + snowball_archive_decoder_semaphore: Arc, /// Snowball members are internal PUTs, so they use a separate global gate /// from preparation through the independently owned post-commit tail. snowball_member_commit_semaphore: Arc, @@ -425,6 +431,7 @@ impl ConcurrencyManager { bandwidth_monitor, metrics_collector, foreground_write_admission_policy, + snowball_archive_decoder_semaphore: Arc::new(Semaphore::new(SNOWBALL_ARCHIVE_DECODER_LIMIT)), snowball_member_commit_semaphore: Arc::new(Semaphore::new(SNOWBALL_MEMBER_COMMIT_LIMIT)), snowball_staging_bytes_semaphore: Arc::new(Semaphore::new(SNOWBALL_STAGING_BYTES_LIMIT)), } @@ -578,6 +585,11 @@ impl ConcurrencyManager { .await } + /// Try to acquire one global Snowball archive decoder slot. + pub(crate) fn try_acquire_snowball_archive_decoder(&self) -> Option { + self.snowball_archive_decoder_semaphore.clone().try_acquire_owned().ok() + } + /// Acquire one global Snowball member lifecycle slot. pub(crate) async fn acquire_snowball_member_commit(&self) -> Result { self.snowball_member_commit_semaphore.clone().acquire_owned().await @@ -1095,8 +1107,8 @@ mod integration_tests { use super::super::io_schedule::{IoLoadLevel, IoPriority}; use super::super::request_guard::GetObjectGuard; use super::{ - ConcurrencyManager, ForegroundWriteAdmission, SNOWBALL_MEMBER_COMMIT_LIMIT, SNOWBALL_STAGING_BYTES_LIMIT, - derive_large_put_admission_limit, + ConcurrencyManager, ForegroundWriteAdmission, SNOWBALL_ARCHIVE_DECODER_LIMIT, SNOWBALL_MEMBER_COMMIT_LIMIT, + SNOWBALL_STAGING_BYTES_LIMIT, derive_large_put_admission_limit, }; use crate::storage::storage_api::concurrency_consumer::PutObjectGuard; use rustfs_concurrency::{AdmissionState, WorkloadAdmissionSnapshotProvider, WorkloadClass}; @@ -1109,6 +1121,20 @@ mod integration_tests { let manager = ConcurrencyManager::new(); let clone = manager.clone(); + let decoder_permits = manager + .snowball_archive_decoder_semaphore + .clone() + .try_acquire_many_owned( + u32::try_from(SNOWBALL_ARCHIVE_DECODER_LIMIT).expect("Snowball decoder limit must fit into u32"), + ) + .expect("the exact Snowball decoder limit must be available"); + assert!( + clone.try_acquire_snowball_archive_decoder().is_none(), + "a cloned manager must share the global decoder gate" + ); + drop(decoder_permits); + assert!(clone.try_acquire_snowball_archive_decoder().is_some()); + let commit_permits = manager .snowball_member_commit_semaphore .clone()