From 4815d608a217e350a1f6a80b15c108171159c937 Mon Sep 17 00:00:00 2001 From: GatewayJ <835269233@qq.com> Date: Tue, 30 Jun 2026 11:19:46 +0800 Subject: [PATCH] fix(ecstore): preserve recovery quorum identity (#3937) * fix(ecstore): preserve recovery quorum identity * fix(get): hold read permits through response streaming * fix(heal): queue read repair at low priority * fix(ecstore): harden recovery read validation --- .../ecstore/src/cache_value/metacache_set.rs | 150 +++++++++++++++++- crates/ecstore/src/set_disk/read.rs | 105 ++++++++++-- rustfs/src/app/object_usecase.rs | 87 ++++++++-- rustfs/src/storage/concurrency/manager.rs | 26 +++ 4 files changed, 333 insertions(+), 35 deletions(-) diff --git a/crates/ecstore/src/cache_value/metacache_set.rs b/crates/ecstore/src/cache_value/metacache_set.rs index 5e13ee8c2..5c4b1474f 100644 --- a/crates/ecstore/src/cache_value/metacache_set.rs +++ b/crates/ecstore/src/cache_value/metacache_set.rs @@ -27,6 +27,7 @@ use std::{ }; use tokio::io::AsyncRead; use tokio::spawn; +use tokio::sync::Mutex as TokioMutex; use tokio::time::timeout; use tokio_util::sync::CancellationToken; use tracing::{debug, error, warn}; @@ -59,6 +60,10 @@ fn is_missing_path_error(err: &DiskError) -> bool { matches!(err, DiskError::FileNotFound | DiskError::FileVersionNotFound | DiskError::VolumeNotFound) } +async fn take_fallback_candidate(fallback_items: &Arc>>) -> Option { + fallback_items.lock().await.pop_front() +} + #[cfg(test)] #[derive(Clone)] pub(crate) enum TestReaderBehavior { @@ -67,6 +72,7 @@ pub(crate) enum TestReaderBehavior { Stall, IgnoreCancel, ProducerError(DiskError), + PrimaryErrorThenFallback(DiskError), PartialThenTimeout(Vec), } @@ -89,6 +95,8 @@ pub struct ListPathRawOptions { #[cfg(test)] pub(crate) test_reader_behaviors: Vec, #[cfg(test)] + pub(crate) test_fallback_reader_behaviors: Vec, + #[cfg(test)] pub(crate) peek_timeout: Option, // pub agreed: Option>, // pub partial: Option]) + Send + Sync>>, @@ -112,6 +120,8 @@ impl Clone for ListPathRawOptions { #[cfg(test)] test_reader_behaviors: self.test_reader_behaviors.clone(), #[cfg(test)] + test_fallback_reader_behaviors: self.test_fallback_reader_behaviors.clone(), + #[cfg(test)] peek_timeout: self.peek_timeout, ..Default::default() } @@ -128,7 +138,11 @@ pub async fn list_path_raw(rx: CancellationToken, opts: ListPathRawOptions) -> d let mut jobs: Vec>> = Vec::new(); let mut readers = Vec::with_capacity(opts.disks.len()); - let fds = opts.fallback_disks.iter().flatten().cloned().collect::>(); + let fds = Arc::new(TokioMutex::new(opts.fallback_disks.iter().flatten().cloned().collect::>())); + #[cfg(test)] + let test_fallbacks = Arc::new(TokioMutex::new( + opts.test_fallback_reader_behaviors.iter().cloned().collect::>(), + )); let max_disk_failures = opts.disks.len().saturating_sub(opts.min_disks); let producer_errs: Arc<[OnceLock]> = (0..opts.disks.len()).map(|_| OnceLock::new()).collect::>().into(); @@ -137,14 +151,16 @@ pub async fn list_path_raw(rx: CancellationToken, opts: ListPathRawOptions) -> d for (disk_idx, disk) in opts.disks.iter().enumerate() { let opdisk = disk.clone(); let opts_clone = opts.clone(); - let mut fds_clone = fds.clone(); + let fds_clone = fds.clone(); + #[cfg(test)] + let test_fallbacks_clone = test_fallbacks.clone(); let cancel_rx_clone = cancel_rx.clone(); let producer_errs_clone = producer_errs.clone(); let (rd, wr) = tokio::io::duplex(64); readers.push(MetacacheReader::new(rd)); jobs.push(spawn(async move { #[cfg(test)] - if let Some(behavior) = opts_clone.test_reader_behaviors.get(disk_idx).cloned() { + let test_primary_error = if let Some(behavior) = opts_clone.test_reader_behaviors.get(disk_idx).cloned() { match behavior { TestReaderBehavior::Eof => return Ok(()), TestReaderBehavior::Entries(entries) => { @@ -168,6 +184,7 @@ pub async fn list_path_raw(rx: CancellationToken, opts: ListPathRawOptions) -> d record_producer_error(&producer_errs_clone, disk_idx, &err); return Err(err); } + TestReaderBehavior::PrimaryErrorThenFallback(err) => Some(err), TestReaderBehavior::PartialThenTimeout(entries) => { let mut wr = wr; let mut out = rustfs_filemeta::MetacacheWriter::new(&mut wr); @@ -178,7 +195,9 @@ pub async fn list_path_raw(rx: CancellationToken, opts: ListPathRawOptions) -> d return Err(err); } } - } + } else { + None + }; let mut wr = wr; let wakl_opts = WalkDirOptions { @@ -195,7 +214,13 @@ pub async fn list_path_raw(rx: CancellationToken, opts: ListPathRawOptions) -> d let mut need_fallback = false; let mut last_err = None; - if let Some(disk) = opdisk { + #[cfg(test)] + if let Some(err) = test_primary_error { + last_err = Some(err); + need_fallback = true; + } + + if !need_fallback && let Some(disk) = opdisk { let primary_walk_started = std::time::Instant::now(); match disk.walk_dir(wakl_opts, &mut wr).await { Ok(_res) => { @@ -238,7 +263,7 @@ pub async fn list_path_raw(rx: CancellationToken, opts: ListPathRawOptions) -> d need_fallback = true; } } - } else { + } else if !need_fallback { last_err = Some(DiskError::DiskNotFound); need_fallback = true; } @@ -249,8 +274,37 @@ pub async fn list_path_raw(rx: CancellationToken, opts: ListPathRawOptions) -> d } while need_fallback { + #[cfg(test)] + if let Some(behavior) = take_fallback_candidate(&test_fallbacks_clone).await { + match behavior { + TestReaderBehavior::Eof => { + need_fallback = false; + last_err = None; + continue; + } + TestReaderBehavior::Entries(entries) => { + let mut out = rustfs_filemeta::MetacacheWriter::new(&mut wr); + out.write(&entries).await.expect("test fallback entries should be written"); + out.close().await.expect("test fallback entries should close"); + need_fallback = false; + last_err = None; + continue; + } + TestReaderBehavior::ProducerError(err) | TestReaderBehavior::PrimaryErrorThenFallback(err) => { + last_err = Some(err); + continue; + } + TestReaderBehavior::Stall + | TestReaderBehavior::IgnoreCancel + | TestReaderBehavior::PartialThenTimeout(_) => { + last_err = Some(DiskError::Timeout); + continue; + } + } + } + let mut disk_op = None; - while let Some(disk) = fds_clone.pop_front() { + while let Some(disk) = take_fallback_candidate(&fds_clone).await { if disk.is_online().await { disk_op = Some(disk); break; @@ -753,6 +807,88 @@ mod tests { assert!(!is_missing_path_error(&DiskError::FileAccessDenied)); } + #[tokio::test] + async fn fallback_candidates_are_claimed_once_across_producers() { + let mut queue = VecDeque::new(); + queue.push_back(1usize); + let candidates = Arc::new(TokioMutex::new(queue)); + + let first = take_fallback_candidate(&candidates).await; + let second = take_fallback_candidate(&candidates).await; + + assert_eq!(first, Some(1)); + assert_eq!(second, None); + } + + fn fallback_test_entry() -> MetaCacheEntry { + MetaCacheEntry { + name: "bucket/object".to_string(), + metadata: vec![1, 2, 3], + cached: None, + reusable: false, + } + } + + #[tokio::test] + async fn list_path_raw_does_not_reuse_one_fallback_for_multiple_failed_producers() { + let err = list_path_raw( + CancellationToken::new(), + ListPathRawOptions { + disks: vec![None, None], + min_disks: 2, + test_reader_behaviors: vec![ + TestReaderBehavior::PrimaryErrorThenFallback(DiskError::DiskNotFound), + TestReaderBehavior::PrimaryErrorThenFallback(DiskError::DiskNotFound), + ], + test_fallback_reader_behaviors: vec![TestReaderBehavior::Entries(vec![fallback_test_entry()])], + ..Default::default() + }, + ) + .await + .expect_err("one fallback must not be counted as two failed primary producers"); + + assert_eq!(err, DiskError::DiskNotFound); + } + + #[tokio::test] + async fn list_path_raw_uses_distinct_fallbacks_to_restore_quorum() { + let seen = Arc::new(Mutex::new(Vec::new())); + let seen_clone = seen.clone(); + + list_path_raw( + CancellationToken::new(), + ListPathRawOptions { + disks: vec![None, None], + min_disks: 2, + test_reader_behaviors: vec![ + TestReaderBehavior::PrimaryErrorThenFallback(DiskError::DiskNotFound), + TestReaderBehavior::PrimaryErrorThenFallback(DiskError::DiskNotFound), + ], + test_fallback_reader_behaviors: vec![ + TestReaderBehavior::Entries(vec![fallback_test_entry()]), + TestReaderBehavior::Entries(vec![fallback_test_entry()]), + ], + partial: Some(Box::new(move |entries: MetaCacheEntries, _: &[Option]| { + let seen = seen_clone.clone(); + Box::pin(async move { + let matches = entries + .0 + .iter() + .flatten() + .filter(|entry| entry.name == "bucket/object") + .count(); + seen.lock().expect("seen mutex poisoned").push(matches); + }) + })), + ..Default::default() + }, + ) + .await + .expect("distinct fallback producers should restore listing quorum"); + + assert_eq!(seen.lock().expect("seen mutex poisoned").as_slice(), &[2]); + } + #[tokio::test] async fn list_path_raw_returns_timeout_when_reader_stalls_before_completion() { let err = list_path_raw( diff --git a/crates/ecstore/src/set_disk/read.rs b/crates/ecstore/src/set_disk/read.rs index 29b00bf0e..98c3c6042 100644 --- a/crates/ecstore/src/set_disk/read.rs +++ b/crates/ecstore/src/set_disk/read.rs @@ -358,6 +358,14 @@ impl MetadataQuorumAccumulator { if self.requested_version_id.is_empty() { return None; } + if self.conflicting_metadata + || self.delete_marker_seen + || self.not_found_responses > 0 + || self.version_not_found_responses > 0 + || self.hard_errors > 0 + { + return None; + } if self.matching_version_votes >= self.read_quorum_for_version() { return Some(MetadataEarlyStopDecision { reason: GET_METADATA_EARLY_STOP_REASON_VERSION_MATCH_QUORUM, @@ -432,6 +440,7 @@ fn metadata_early_stop_candidate_matches(left: &FileInfo, right: &FileInfo) -> b && left.parts == right.parts && left.checksum == right.checksum && left.versioned == right.versioned + && left.data_dir == right.data_dir && left.erasure.algorithm == right.erasure.algorithm && left.erasure.data_blocks == right.erasure.data_blocks && left.erasure.parity_blocks == right.erasure.parity_blocks @@ -728,7 +737,7 @@ async fn submit_read_repair_heal_with_submitter( bucket.to_string(), Some(object.to_string()), false, - Some(HealChannelPriority::Normal), + Some(HealChannelPriority::Low), Some(pool_index), Some(set_index), ); @@ -2468,10 +2477,13 @@ mod metadata_cache_tests { use super::*; use rustfs_common::heal_channel::HealAdmissionDropReason; use serial_test::serial; + use std::sync::Mutex; use std::sync::atomic::{AtomicUsize, Ordering}; static SLOW_READ_REPAIR_SUBMITTER_CALLS: AtomicUsize = AtomicUsize::new(0); static DROPPED_READ_REPAIR_SUBMITTER_CALLS: AtomicUsize = AtomicUsize::new(0); + static CAPTURED_READ_REPAIR_PRIORITY: Mutex> = Mutex::new(None); + static CAPTURED_READ_REPAIR_CALLS: AtomicUsize = AtomicUsize::new(0); fn slow_read_repair_submitter(_request: rustfs_common::heal_channel::HealChannelRequest) -> ReadRepairAdmissionFuture { SLOW_READ_REPAIR_SUBMITTER_CALLS.fetch_add(1, Ordering::Relaxed); @@ -2488,6 +2500,15 @@ mod metadata_cache_tests { }) } + fn capture_read_repair_submitter(request: rustfs_common::heal_channel::HealChannelRequest) -> ReadRepairAdmissionFuture { + CAPTURED_READ_REPAIR_CALLS.fetch_add(1, Ordering::Relaxed); + *CAPTURED_READ_REPAIR_PRIORITY.lock().expect("capture mutex poisoned") = Some(request.priority); + Box::pin(async { + tokio::time::sleep(Duration::from_millis(1)).await; + ReadRepairAdmissionOutcome::Response(HealAdmissionResult::Accepted) + }) + } + async fn new_metadata_cache_test_set() -> Arc { SetDisks::new( "metadata-cache-test".to_string(), @@ -2671,6 +2692,41 @@ mod metadata_cache_tests { release_read_repair_heal_reservation(&released_key).await; } + #[tokio::test] + #[serial] + async fn submit_read_repair_heal_uses_low_priority() { + CAPTURED_READ_REPAIR_CALLS.store(0, Ordering::Relaxed); + *CAPTURED_READ_REPAIR_PRIORITY.lock().expect("capture mutex poisoned") = None; + let bucket = format!("bucket-{}", Uuid::new_v4()); + + submit_read_repair_heal_with_submitter( + ReadRepairHealSubmission { + bucket: &bucket, + object: "object", + version_id: None, + pool_index: 0, + set_index: 0, + part_number: Some(1), + reason: "missing_shards", + }, + capture_read_repair_submitter, + ) + .await; + + tokio::time::timeout(Duration::from_secs(1), async { + while CAPTURED_READ_REPAIR_CALLS.load(Ordering::Relaxed) == 0 { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("background read-repair submitter should be called"); + + assert_eq!( + *CAPTURED_READ_REPAIR_PRIORITY.lock().expect("capture mutex poisoned"), + Some(HealChannelPriority::Low) + ); + } + #[test] fn resolved_read_repair_version_prefers_selected_fileinfo_version() { let mut fi = valid_test_fileinfo("object"); @@ -3165,6 +3221,21 @@ mod tests { assert_eq!(accumulator.final_miss_reason(), GET_METADATA_EARLY_STOP_REASON_CONFLICTING_METADATA); } + #[test] + fn metadata_quorum_accumulator_falls_back_on_split_data_dir() { + let mut accumulator = metadata_early_stop_accumulator(); + let mut first = metadata_early_stop_candidate("object", 1); + let mut second = metadata_early_stop_candidate("object", 2); + first.data_dir = Some(Uuid::parse_str("00000000-0000-0000-0000-000000000001").expect("static uuid should parse")); + second.data_dir = Some(Uuid::parse_str("00000000-0000-0000-0000-000000000002").expect("static uuid should parse")); + + accumulator.observe_file_info(&first); + accumulator.observe_file_info(&second); + + assert!(accumulator.early_stop_decision().is_none()); + assert_eq!(accumulator.final_miss_reason(), GET_METADATA_EARLY_STOP_REASON_CONFLICTING_METADATA); + } + #[test] fn metadata_quorum_accumulator_falls_back_on_explicit_version_quorum() { let mut accumulator = MetadataQuorumAccumulator::new(4, 2, false); @@ -3342,28 +3413,40 @@ mod tests { } #[test] - fn version_early_stop_tracks_matching_votes_independently_of_candidate() { + fn version_early_stop_falls_back_on_conflicting_metadata() { let requested_vid = Uuid::new_v4(); let mut accumulator = version_early_stop_accumulator(&requested_vid.to_string()); // Two valid responses with matching version_id but different erasure.index - // (so they conflict on the candidate path but still count for version votes) + // (so they conflict on the candidate path and must keep the fanout open) let mut fi1 = version_early_stop_candidate("object", 1, requested_vid); fi1.size = 100; let mut fi2 = version_early_stop_candidate("object", 2, requested_vid); - fi2.size = 200; // different size → conflicting metadata on candidate path + fi2.size = 200; accumulator.observe_file_info(&fi1); accumulator.observe_file_info(&fi2); - // Candidate path sees conflict, but version path sees quorum assert!(accumulator.early_stop_decision().is_none()); - assert_eq!( - accumulator.version_early_stop_decision(), - Some(MetadataEarlyStopDecision { - reason: GET_METADATA_EARLY_STOP_REASON_VERSION_MATCH_QUORUM - }) - ); + assert!(accumulator.version_early_stop_decision().is_none()); + assert_eq!(accumulator.final_miss_reason(), GET_METADATA_EARLY_STOP_REASON_CONFLICTING_METADATA); + } + + #[test] + fn version_early_stop_falls_back_on_split_data_dir() { + let requested_vid = Uuid::new_v4(); + let mut accumulator = version_early_stop_accumulator(&requested_vid.to_string()); + let mut first = version_early_stop_candidate("object", 1, requested_vid); + let mut second = version_early_stop_candidate("object", 2, requested_vid); + first.data_dir = Some(Uuid::parse_str("00000000-0000-0000-0000-000000000001").expect("static uuid should parse")); + second.data_dir = Some(Uuid::parse_str("00000000-0000-0000-0000-000000000002").expect("static uuid should parse")); + + accumulator.observe_file_info(&first); + accumulator.observe_file_info(&second); + + assert!(accumulator.early_stop_decision().is_none()); + assert!(accumulator.version_early_stop_decision().is_none()); + assert_eq!(accumulator.final_miss_reason(), GET_METADATA_EARLY_STOP_REASON_CONFLICTING_METADATA); } fn codec_streaming_test_object_info(fi: &FileInfo) -> ObjectInfo { diff --git a/rustfs/src/app/object_usecase.rs b/rustfs/src/app/object_usecase.rs index c44fbe7a0..5c8ff5c78 100644 --- a/rustfs/src/app/object_usecase.rs +++ b/rustfs/src/app/object_usecase.rs @@ -168,7 +168,7 @@ use std::sync::{Arc, Mutex, OnceLock}; use std::time::Duration; use time::{OffsetDateTime, format_description::well_known::Rfc3339}; use tokio::io::{AsyncRead, ReadBuf}; -use tokio::sync::RwLock; +use tokio::sync::{OwnedSemaphorePermit, RwLock}; use tokio_tar::Archive; use tokio_util::io::{ReaderStream, StreamReader}; use tracing::{debug, error, instrument, warn}; @@ -252,9 +252,9 @@ struct GetObjectBootstrap { concurrent_requests: usize, } -struct GetObjectIoPlanning<'a> { +struct GetObjectIoPlanning { /// `None` when inline fast path skips disk I/O semaphore. - _disk_permit: Option>, + disk_permit: Option, permit_wait_duration: Duration, queue_status: concurrency::IoQueueStatus, queue_utilization: f64, @@ -287,8 +287,8 @@ struct GetObjectReadSetup { is_inline_fast_path: bool, } -struct GetObjectPreparedRead<'a> { - io_planning: GetObjectIoPlanning<'a>, +struct GetObjectPreparedRead { + io_planning: GetObjectIoPlanning, read_setup: GetObjectReadSetup, } @@ -441,6 +441,35 @@ pin_project! { } } +pin_project! { + // Keep the disk-read admission permit tied to the response body. This is + // intentionally conservative backpressure: a streaming GET should occupy a + // read slot until the client drains or drops the body. + struct DiskReadPermitReader { + #[pin] + inner: R, + _disk_permit: OwnedSemaphorePermit, + } +} + +impl DiskReadPermitReader { + fn new(inner: R, disk_permit: OwnedSemaphorePermit) -> Self { + Self { + inner, + _disk_permit: disk_permit, + } + } +} + +impl AsyncRead for DiskReadPermitReader +where + R: AsyncRead, +{ + fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { + self.project().inner.poll_read(cx, buf) + } +} + pin_project! { struct GetObjectReaderStream { #[pin] @@ -2038,16 +2067,16 @@ impl DefaultObjectUsecase { }) } - async fn acquire_get_object_io_planning<'a>( - manager: &'a ConcurrencyManager, + async fn acquire_get_object_io_planning( + manager: &ConcurrencyManager, wrapper: &RequestTimeoutWrapper, timeout_config: &GetObjectTimeoutPolicy, bucket: &str, key: &str, - ) -> S3Result> { + ) -> S3Result { let permit_wait_start = std::time::Instant::now(); let disk_permit = manager - .acquire_disk_read_permit() + .acquire_owned_disk_read_permit() .await .map_err(|_| s3_error!(InternalError, "disk read semaphore closed"))?; let permit_wait_duration = permit_wait_start.elapsed(); @@ -2083,7 +2112,7 @@ impl DefaultObjectUsecase { Self::ensure_get_object_not_timed_out(wrapper, timeout_config, bucket, key, GetObjectTimeoutStage::BeforeRead)?; Ok(GetObjectIoPlanning { - _disk_permit: Some(disk_permit), + disk_permit: Some(disk_permit), permit_wait_duration, queue_status, queue_utilization, @@ -2141,9 +2170,9 @@ impl DefaultObjectUsecase { }) } #[allow(clippy::too_many_arguments)] - async fn prepare_get_object_read_execution<'a>( + async fn prepare_get_object_read_execution( req: &S3Request, - manager: &'a ConcurrencyManager, + manager: &ConcurrencyManager, wrapper: &RequestTimeoutWrapper, timeout_config: &GetObjectTimeoutPolicy, bucket: &str, @@ -2151,7 +2180,7 @@ impl DefaultObjectUsecase { rs: Option, opts: &ObjectOptions, part_number: Option, - ) -> S3Result> { + ) -> S3Result { let h = req.headers.clone(); // SF05: Store lookup first (cached via SF01 moka cache). @@ -2187,7 +2216,7 @@ impl DefaultObjectUsecase { // SF05: Skip disk I/O semaphore for inline fast path — data is already in memory. let io_planning = if read_setup.is_inline_fast_path { GetObjectIoPlanning { - _disk_permit: None, + disk_permit: None, permit_wait_duration: Duration::ZERO, queue_status: concurrency::IoQueueStatus::default(), queue_utilization: 0.0, @@ -3427,9 +3456,12 @@ impl DefaultObjectUsecase { ) .await?; let GetObjectPreparedRead { io_planning, read_setup } = prepared_read; - let permit_wait_duration = io_planning.permit_wait_duration; - let queue_status = io_planning.queue_status; - let queue_utilization = io_planning.queue_utilization; + let GetObjectIoPlanning { + disk_permit, + permit_wait_duration, + queue_status, + queue_utilization, + } = io_planning; let GetObjectReadSetup { info, @@ -3447,6 +3479,11 @@ impl DefaultObjectUsecase { encryption_applied, is_inline_fast_path: _, } = read_setup; + let final_stream = if let Some(disk_permit) = disk_permit { + wrap_reader(DiskReadPermitReader::new(final_stream, disk_permit)) + } else { + final_stream + }; let versioning_start = rustfs_io_metrics::get_stage_metrics_enabled().then(std::time::Instant::now); let versioned = BucketVersioningSys::prefix_enabled(&bucket, &key).await; @@ -6110,6 +6147,22 @@ mod tests { assert_eq!(err.kind(), std::io::ErrorKind::TimedOut); } + #[tokio::test] + async fn disk_read_permit_reader_holds_permit_until_reader_is_dropped() { + let semaphore = Arc::new(tokio::sync::Semaphore::new(1)); + let permit = semaphore + .clone() + .acquire_owned() + .await + .expect("test semaphore should grant owned permit"); + + let reader = DiskReadPermitReader::new(std::io::Cursor::new(Vec::::new()), permit); + assert_eq!(semaphore.available_permits(), 0); + + drop(reader); + assert_eq!(semaphore.available_permits(), 1); + } + #[tokio::test] async fn build_get_object_body_keeps_large_objects_on_streaming_path_without_preread() { let reads = Arc::new(AtomicUsize::new(0)); diff --git a/rustfs/src/storage/concurrency/manager.rs b/rustfs/src/storage/concurrency/manager.rs index 43499165c..71f7fa0df 100644 --- a/rustfs/src/storage/concurrency/manager.rs +++ b/rustfs/src/storage/concurrency/manager.rs @@ -169,6 +169,14 @@ impl ConcurrencyManager { self.disk_read_semaphore.acquire().await } + /// Acquire an owned permit to perform a disk read operation. + /// + /// Use this when the permit must outlive the borrow of the manager, such as + /// response body streams that continue after the S3 handler returns. + pub async fn acquire_owned_disk_read_permit(&self) -> Result { + self.disk_read_semaphore.clone().acquire_owned().await + } + // ============================================ // Adaptive I/O Strategy Methods // ============================================ @@ -726,6 +734,24 @@ mod integration_tests { assert_eq!(snapshot.limit, Some(manager.scheduler_config().max_concurrent_reads)); } + #[tokio::test] + #[serial] + async fn test_concurrency_manager_owned_disk_read_permit_tracks_queue_snapshot() { + let manager = ConcurrencyManager::new(); + + let permit = manager + .acquire_owned_disk_read_permit() + .await + .expect("owned disk read permit should be acquired"); + let snapshot = manager.get_object_admission_snapshot(); + + assert_eq!(snapshot.active, Some(1)); + assert_eq!(snapshot.limit, Some(manager.scheduler_config().max_concurrent_reads)); + + drop(permit); + assert_eq!(manager.get_object_admission_snapshot().active, Some(0)); + } + #[tokio::test] #[serial] async fn test_concurrency_manager_workload_admission_snapshot_tracks_put_requests() {