From 8aba025bd6b69e3394db46f8bd2abf170978d20d Mon Sep 17 00:00:00 2001 From: Chris Date: Sat, 24 Oct 2026 22:19:06 +0900 Subject: [PATCH] feat(concurrency): retain resource leases through native I/O (#8886) feat(concurrency): retain resource leases across owned disk work --- crates/concurrency/src/adaptive.rs | 832 +++++++++++++++++++++ crates/concurrency/src/lib.rs | 1 + crates/ecstore/src/disk/cleanup_runtime.rs | 66 +- crates/ecstore/src/disk/disk_store.rs | 40 +- crates/ecstore/src/disk/io_admission.rs | 250 ++++++- crates/ecstore/src/disk/local.rs | 55 +- crates/ecstore/src/disk/mod.rs | 7 +- crates/ecstore/src/disk/os.rs | 509 ++++++++++--- crates/ecstore/src/lib.rs | 2 +- crates/ecstore/src/store/heal.rs | 42 +- 10 files changed, 1675 insertions(+), 129 deletions(-) create mode 100644 crates/concurrency/src/adaptive.rs diff --git a/crates/concurrency/src/adaptive.rs b/crates/concurrency/src/adaptive.rs new file mode 100644 index 000000000..943caf749 --- /dev/null +++ b/crates/concurrency/src/adaptive.rs @@ -0,0 +1,832 @@ +// Copyright 2026 RustFS Team +// Licensed under the Apache License, Version 2.0. + +//! Mutable admission capacity without replacing or revoking held permits. + +use std::num::NonZeroU16; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::time::Duration; +use tokio::sync::{Notify, Semaphore}; +use tokio::time::Instant; + +/// A resource owner publishes a capacity; running units retain their permits +/// after a reduction. Lease slots are reserved once, with no per-unit Arc +/// allocation. Tokio wait queues may lazily initialize platform mutex storage +/// on first contention. It does not discover devices, start +/// workers, or decide resource policy. Cancellation drops the acquire future. +/// +/// Uncontended admission is atomic. Parked work takes a FIFO admission turn, +/// so repeated fast acquisitions cannot starve an older waiter. Source/age +/// priority still belongs to the caller's scheduler, before joining this queue. +pub struct AdaptiveAdmission { + // One CAS observes both capacity and active count, so a concurrent shrink + // cannot admit a new unit using an already superseded capacity. + state: AtomicU64, + origin: Instant, + valid_until_millis: AtomicU64, + // Only the FIFO turn holder subscribes. A stored notification is harmless: + // capacity and freshness are always checked again after waking. + changed: Notify, + queued: AtomicUsize, + // FIFO admission turn, not a mutable-state lock. No storage lock or + // admitted permit may be held while joining. This private semaphore is + // never closed; release and resize do not acquire it. + turn: Semaphore, + maximum: u32, + slots: Box<[PermitSlots]>, + next_chunk: AtomicUsize, +} + +struct PermitSlots { + free: AtomicU64, + references: [AtomicUsize; 64], +} + +struct QueuedAdmission<'a>(&'a AdaptiveAdmission); + +impl Drop for QueuedAdmission<'_> { + fn drop(&mut self) { + self.0.queued.fetch_sub(1, Ordering::AcqRel); + } +} + +/// One active unit. Clones share that unit without allocating; the final drop +/// returns it to the same owner, including after resize or waiter cancellation. +#[must_use = "hold the permit for the admitted unit's complete lifetime"] +pub struct AdaptivePermit { + owner: Arc, + slot: usize, +} + +impl Clone for AdaptivePermit { + fn clone(&self) -> Self { + // Every slot reference owns a strong owner reference. Clone the Arc + // first, so its overflow protection also bounds the slot refcount. + let owner = Arc::clone(&self.owner); + owner.slots[self.slot / 64].references[self.slot % 64].fetch_add(1, Ordering::Relaxed); + Self { owner, slot: self.slot } + } +} + +impl std::fmt::Debug for AdaptivePermit { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AdaptivePermit") + .field("slot", &self.slot) + .field("owner", &self.owner.snapshot()) + .finish() + } +} + +/// One coherent view of configured capacity and actual active units. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct AdaptiveAdmissionSnapshot { + /// Current cap on new admissions. Zero pauses all new units. + pub limit: u32, + /// Existing units, which may temporarily exceed a reduced limit. + pub active: u32, +} + +/// Confidence-qualified pressure for one complete sampling window. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ResourcePressure { + /// Missing, stale, reset, or insufficient observations; never evidence of idle. + Unknown, + /// A resource or foreground-protection threshold was reached. + Busy, + /// All required observations were fresh and below their low watermarks. + Stable, +} + +/// Resource-owner input to the controller, after applying operator and kernel +/// caps. This describes one owner; callers must not duplicate a device's budget +/// for aliases, pools, or remote requesters. +#[derive(Debug, Clone, Copy)] +pub struct FeedbackObservation { + /// Identity of the resource/telemetry generation. Replacement resets probes. + pub generation: u64, + /// Upper bound imposed by resources and explicit configuration; zero blocks. + pub max_concurrency: u32, + /// Complete-window pressure, including the confidence of its sources. + pub pressure: ResourcePressure, + /// Cumulative completed work in this generation, not just scheduled work. + pub completed_units: Option, +} + +/// Why the latest concurrency limit was selected. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum FeedbackReason { + /// The resource cap is unavailable; no new unit can be bounded safely. + Unavailable, + /// A changed owner or reset work counter requires a new baseline. + NewBaseline, + /// The resource or operator cap reduced the permitted concurrency. + CapReduced, + /// Pressure caused an immediate multiplicative reduction. + Pressure, + /// Incomplete telemetry retains at most one maintenance unit. + LowConfidence, + /// Waiting for consecutive stable windows, work progress and cooldown. + Recovering, + /// One additional unit is being probed after a stable recovery period. + Probe, +} + +/// One controller decision, suitable for metrics without dynamic labels. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct FeedbackDecision { + /// New admissions use this cap; existing permits remain valid. + pub limit: u32, + /// Decision cause, distinct from a proof of work completion or storage health. + pub reason: FeedbackReason, +} + +/// Single-owner additive probe/multiplicative decrease controller. It adjusts +/// only concurrency. Every source must be fresh before declaring `Stable`. +/// +/// The sampler runs every second. A gap over two seconds resets recovery; eight +/// consecutive stable seconds and observed completed work allow one extra unit. +/// Pressure or lower caps apply immediately. This policy is not a performance +/// claim; its parameters require the same workload/foreground-SLO acceptance. +pub struct FeedbackController { + previous: Option<(Instant, FeedbackObservation)>, + stable_since: Option, + last_change: Instant, + limit: u32, +} + +impl Default for FeedbackController { + fn default() -> Self { + Self::new() + } +} + +impl FeedbackController { + /// Start without a resource observation or an admission grant. + pub fn new() -> Self { + Self { + previous: None, + stable_since: None, + last_change: Instant::now(), + limit: 0, + } + } + + /// Consume one sample and publish a three-second lease to the same gate. + /// If this function stops being called, admission expires independently. + pub fn update(&mut self, gate: &AdaptiveAdmission, observation: Option) -> FeedbackDecision { + let observation = observation.map(|mut observation| { + observation.max_concurrency = observation.max_concurrency.min(gate.maximum); + observation + }); + let decision = self.decide(Instant::now(), observation); + gate.set_sampled_limit(decision.limit, Duration::from_secs(3)); + decision + } + + fn decide(&mut self, now: Instant, observation: Option) -> FeedbackDecision { + let Some(observation) = observation else { + self.previous = None; + self.stable_since = None; + self.last_change = now; + self.limit = 0; + return FeedbackDecision { + limit: 0, + reason: FeedbackReason::Unavailable, + }; + }; + let previous = self.previous.replace((now, observation)); + let contiguous = previous.is_some_and(|(at, old)| { + now.saturating_duration_since(at) <= Duration::from_secs(2) + && old.generation == observation.generation + && matches!((old.completed_units, observation.completed_units), (Some(a), Some(b)) if b >= a) + }); + let advanced = previous + .is_some_and(|(_, old)| matches!((old.completed_units, observation.completed_units), (Some(a), Some(b)) if b > a)); + let old_limit = self.limit; + let reason = if observation.max_concurrency == 0 { + self.limit = 0; + self.stable_since = None; + FeedbackReason::CapReduced + } else if observation.pressure == ResourcePressure::Busy { + self.limit = (self.limit / 2).max(1).min(observation.max_concurrency); + self.stable_since = None; + // Repeated busy windows restart cooldown, even at the minimum cap. + self.last_change = now; + FeedbackReason::Pressure + } else if observation.max_concurrency < self.limit { + self.limit = observation.max_concurrency; + self.stable_since = None; + FeedbackReason::CapReduced + } else if observation.pressure == ResourcePressure::Unknown || observation.completed_units.is_none() { + self.limit = 1.min(observation.max_concurrency); + self.stable_since = None; + FeedbackReason::LowConfidence + } else if !contiguous || self.limit == 0 { + self.limit = 1.min(observation.max_concurrency); + self.stable_since = None; + self.last_change = now; + FeedbackReason::NewBaseline + } else { + let since = *self.stable_since.get_or_insert(now); + if advanced + && now.saturating_duration_since(since) >= Duration::from_secs(8) + && now.saturating_duration_since(self.last_change) >= Duration::from_secs(8) + && self.limit < observation.max_concurrency + { + self.limit += 1; + FeedbackReason::Probe + } else { + FeedbackReason::Recovering + } + }; + if old_limit != self.limit { + self.last_change = now; + } + FeedbackDecision { + limit: self.limit, + reason, + } + } +} + +impl AdaptiveAdmission { + /// Construct a bounded owner, with zero initial capacity allowed for unknown + /// input. The immutable maximum reserves at most 65,535 lease slots once; + /// acquiring or cloning a permit does not allocate a per-unit Arc. + pub fn new(maximum: NonZeroU16, limit: u32) -> Arc { + let maximum_u32 = u32::from(maximum.get()); + let maximum = usize::from(maximum.get()); + let slots = (0..maximum.div_ceil(64)) + .map(|chunk| { + let count = (maximum - chunk * 64).min(64); + PermitSlots { + free: AtomicU64::new(if count == 64 { u64::MAX } else { (1_u64 << count) - 1 }), + references: std::array::from_fn(|_| AtomicUsize::new(0)), + } + }) + .collect::>() + .into_boxed_slice(); + Arc::new(Self { + state: AtomicU64::new(u64::from(limit.min(maximum_u32)) << 32), + origin: Instant::now(), + valid_until_millis: AtomicU64::new(u64::MAX), + changed: Notify::new(), + queued: AtomicUsize::new(0), + turn: Semaphore::new(1), + maximum: maximum_u32, + slots, + next_chunk: AtomicUsize::new(0), + }) + } + + /// Publish a cap bounded by the immutable maximum, without losing held work + /// or changing owner identity. + pub fn set_limit(&self, limit: u32) { + let limit = limit.min(self.maximum); + let _ = self.state.try_update(Ordering::AcqRel, Ordering::Acquire, |state| { + Some((u64::from(limit) << 32) | (state & u64::from(u32::MAX))) + }); + self.changed.notify_one(); + } + + /// Publish a cap with a finite observation lifetime. Expired observations + /// block new admissions even if the sampling worker has stopped. Existing + /// permits are retained. A single resource owner must serialize publishers. + /// + /// Capacity is published before freshness, so a concurrent reader can + /// temporarily reject a fresh update but cannot use an increased lifetime + /// with the previous, potentially larger capacity. + pub fn set_sampled_limit(&self, limit: u32, valid_for: Duration) { + self.set_limit(limit); + let deadline = self + .elapsed_millis() + .saturating_add(u64::try_from(valid_for.as_millis()).unwrap_or(u64::MAX)); + self.valid_until_millis.store(deadline, Ordering::Release); + self.changed.notify_one(); + } + + fn elapsed_millis(&self) -> u64 { + u64::try_from(self.origin.elapsed().as_millis()).unwrap_or(u64::MAX) + } + + fn fresh(&self) -> bool { + self.elapsed_millis() < self.valid_until_millis.load(Ordering::Acquire) + } + + /// Read the cap and active count from the same atomic observation. + pub fn snapshot(&self) -> AdaptiveAdmissionSnapshot { + // Acquire freshness first: a newly published lease must not be paired + // with capacity loaded before its preceding reduction. + let fresh = self.fresh(); + let state = self.state.load(Ordering::Acquire); + AdaptiveAdmissionSnapshot { + limit: if fresh { + u32::try_from(state >> 32).unwrap_or(u32::MAX) + } else { + 0 + }, + active: u32::try_from(state & u64::from(u32::MAX)).unwrap_or(u32::MAX), + } + } + + /// Admit immediately if capacity remains, without a per-unit allocation. + /// The uncontended path uses atomics; refunding a racing acquisition may + /// wake the FIFO head through Tokio's notification queue. + pub fn try_acquire(self: &Arc) -> Option { + if self.queued.load(Ordering::Acquire) != 0 { + return None; + } + let permit = self.acquire_capacity()?; + // A waiter can register between the first queue check and the capacity + // CAS. Refund before exposing the permit rather than stealing its turn. + (self.queued.load(Ordering::Acquire) == 0).then_some(permit) + } + + fn acquire_capacity(self: &Arc) -> Option { + if !self.fresh() { + return None; + } + let state = self.state.load(Ordering::Acquire); + if state & u64::from(u32::MAX) >= state >> 32 { + return None; + } + let slot = self.reserve_slot()?; + if self + .state + .try_update(Ordering::AcqRel, Ordering::Acquire, |state| { + let active = state & u64::from(u32::MAX); + (active < state >> 32).then(|| state + 1) + }) + .is_err() + { + self.release_slot(slot); + self.changed.notify_one(); + return None; + } + self.slots[slot / 64].references[slot % 64].store(1, Ordering::Release); + let permit = AdaptivePermit { + owner: self.clone(), + slot, + }; + // The capacity CAS can retry across an observation expiry. Recheck + // before giving ownership to the caller; dropping refunds this unit. + self.fresh().then_some(permit) + } + + fn reserve_slot(&self) -> Option { + let first = self.next_chunk.fetch_add(1, Ordering::Relaxed) % self.slots.len(); + for offset in 0..self.slots.len() { + let chunk_index = (first + offset) % self.slots.len(); + let free = &self.slots[chunk_index].free; + if let Ok(bits) = free.try_update(Ordering::AcqRel, Ordering::Acquire, |bits| (bits != 0).then(|| bits & (bits - 1))) + { + // trailing_zeros is in 0..64 because the zero word is rejected. + return Some(chunk_index * 64 + bits.trailing_zeros() as usize); + } + } + None + } + + fn release_slot(&self, slot: usize) { + self.slots[slot / 64].free.fetch_or(1_u64 << (slot % 64), Ordering::Release); + } + + /// Wait for capacity. The future owns no permit until it returns; dropping + /// a pending future leaves capacity untouched. Wrap in the caller's existing + /// cancellation/deadline select before acquiring any storage commit locks. + pub async fn acquire(self: &Arc) -> AdaptivePermit { + if let Some(permit) = self.try_acquire() { + return permit; + } + self.queued.fetch_add(1, Ordering::AcqRel); + let _queued = QueuedAdmission(self); + // Retain the successful permit inside its Result. No code can close + // this private semaphore, so acquisition cannot return an error. + let _turn = self.turn.acquire().await; + loop { + let notified = self.changed.notified(); + tokio::pin!(notified); + // Register before checking capacity to avoid a lost release/resize. + notified.as_mut().enable(); + if let Some(permit) = self.acquire_capacity() { + return permit; + } + notified.await; + } + } +} + +impl Drop for AdaptivePermit { + fn drop(&mut self) { + if self.owner.slots[self.slot / 64].references[self.slot % 64].fetch_sub(1, Ordering::AcqRel) != 1 { + return; + } + // Publish the free slot before exposing another admission capacity. + self.owner.release_slot(self.slot); + self.owner.state.fetch_sub(1, Ordering::AcqRel); + self.owner.changed.notify_one(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn shared_unit_stays_admitted_until_native_worker_releases_its_last_clone() { + let gate = AdaptiveAdmission::new(NonZeroU16::new(1).expect("one unit"), 1); + let outer = gate.acquire().await; + let native = outer.clone(); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let worker = tokio::task::spawn_blocking(move || { + started_tx.send(()).expect("worker started"); + release_rx.blocking_recv().expect("release native worker"); + drop(native); + }); + started_rx.await.expect("native worker is live"); + drop(outer); + assert_eq!(gate.snapshot().active, 1); + assert!(gate.try_acquire().is_none()); + gate.set_limit(0); + assert_eq!(gate.snapshot().active, 1); + release_tx.send(()).expect("release tail"); + worker.await.expect("native tail completed"); + assert_eq!(gate.snapshot().active, 0); + gate.set_limit(1); + assert!(gate.try_acquire().is_some()); + } + + #[tokio::test] + async fn bounded_slots_cross_word_boundaries_and_recycle_only_after_final_drop() { + let maximum = NonZeroU16::new(65).expect("partial second chunk"); + let gate = AdaptiveAdmission::new(maximum, u32::MAX); + assert_eq!(gate.snapshot().limit, 65); + for _ in 0..4 { + let units: Vec<_> = (0..65).map(|_| gate.try_acquire().expect("free slot")).collect(); + let tails = units.clone(); + assert_eq!(gate.snapshot().active, 65); + assert!(gate.try_acquire().is_none()); + drop(units); + assert_eq!(gate.snapshot().active, 65); + gate.set_limit(0); + for (remaining, tail) in (0..65).rev().zip(tails) { + drop(tail); + assert_eq!(gate.snapshot().active, remaining); + } + gate.set_limit(u32::MAX); + assert_eq!(gate.snapshot().limit, 65); + } + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn shared_units_survive_concurrent_clone_drop_reuse_and_resize() { + let gate = AdaptiveAdmission::new(NonZeroU16::new(8).expect("bounded owner"), 8); + let mut workers = Vec::new(); + for _ in 0..16 { + let gate = gate.clone(); + workers.push(tokio::spawn(async move { + for _ in 0..64 { + let root = gate.acquire().await; + let one = root.clone(); + let two = one.clone(); + drop(root); + tokio::task::yield_now().await; + drop(two); + tokio::task::yield_now().await; + drop(one); + } + })); + } + for limit in [0, 1, 8, 2, 0, 8] { + gate.set_limit(limit); + tokio::task::yield_now().await; + } + tokio::time::timeout(Duration::from_secs(5), async { + for worker in workers { + worker.await.expect("lease worker"); + } + }) + .await + .expect("all lease workers complete"); + assert_eq!(gate.snapshot(), AdaptiveAdmissionSnapshot { limit: 8, active: 0 }); + let full: Vec<_> = (0..8).map(|_| gate.try_acquire().expect("all slots returned")).collect(); + assert!(gate.try_acquire().is_none()); + drop(full); + } + + #[tokio::test] + async fn queued_work_cannot_be_starved_by_new_fast_acquisitions() { + use std::task::{Context, Poll, Waker}; + let gate = AdaptiveAdmission::new(NonZeroU16::new(8).expect("bounded test owner"), 1); + let held = gate.acquire().await; + let mut first = Box::pin(gate.acquire()); + let mut second = Box::pin(gate.acquire()); + let mut context = Context::from_waker(Waker::noop()); + assert!(first.as_mut().poll(&mut context).is_pending()); + assert!(second.as_mut().poll(&mut context).is_pending()); + drop(held); + assert!(gate.try_acquire().is_none(), "new work must not steal a parked worker's turn"); + assert!(second.as_mut().poll(&mut context).is_pending(), "later work cannot pass the queue head"); + let Poll::Ready(first_permit) = first.as_mut().poll(&mut context) else { + panic!("queue head must run") + }; + drop(first_permit); + let Poll::Ready(second_permit) = second.as_mut().poll(&mut context) else { + panic!("second waiter must run") + }; + drop(second_permit); + assert!(gate.try_acquire().is_some(), "drained queue restores the atomic fast path"); + } + + #[tokio::test] + async fn cancelling_queued_work_preserves_fifo_capacity_and_fast_path() { + use std::task::{Context, Poll, Waker}; + let gate = AdaptiveAdmission::new(NonZeroU16::new(8).expect("bounded test owner"), 0); + let mut head = Box::pin(gate.acquire()); + let mut middle = Box::pin(gate.acquire()); + let mut tail = Box::pin(gate.acquire()); + let mut context = Context::from_waker(Waker::noop()); + assert!(head.as_mut().poll(&mut context).is_pending()); + assert!(middle.as_mut().poll(&mut context).is_pending()); + assert!(tail.as_mut().poll(&mut context).is_pending()); + drop(middle); + drop(head); + gate.set_limit(1); + assert!( + gate.try_acquire().is_none(), + "cancelled head must not expose the surviving waiter's capacity" + ); + let Poll::Ready(permit) = tail.as_mut().poll(&mut context) else { + panic!("surviving waiter must run") + }; + drop(permit); + assert!(gate.try_acquire().is_some()); + + gate.set_limit(0); + let mut cancelled = Box::pin(gate.acquire()); + assert!(cancelled.as_mut().poll(&mut context).is_pending()); + drop(cancelled); + gate.set_limit(1); + assert!( + gate.try_acquire().is_some(), + "cancelling the final waiter must not strand an admission turn" + ); + } + + #[tokio::test] + async fn cancelling_a_notified_head_hands_capacity_to_the_next_waiter() { + use std::task::{Context, Poll, Waker}; + let gate = AdaptiveAdmission::new(NonZeroU16::new(2).expect("two units"), 0); + let mut head = Box::pin(gate.acquire()); + let mut next = Box::pin(gate.acquire()); + let mut context = Context::from_waker(Waker::noop()); + assert!(head.as_mut().poll(&mut context).is_pending()); + assert!(next.as_mut().poll(&mut context).is_pending()); + gate.set_limit(1); + // Cancel after notification, before the head consumes its capacity. + drop(head); + let Poll::Ready(permit) = next.as_mut().poll(&mut context) else { + panic!("successor must observe capacity without another notification") + }; + assert_eq!(gate.snapshot(), AdaptiveAdmissionSnapshot { limit: 1, active: 1 }); + drop(permit); + } + + #[tokio::test(start_paused = true)] + async fn one_resize_serves_multiple_fifo_turns_but_stored_wakes_cannot_bypass_expiry() { + use std::task::{Context, Poll, Waker}; + let gate = AdaptiveAdmission::new(NonZeroU16::new(2).expect("two units"), 0); + let mut first = Box::pin(gate.acquire()); + let mut second = Box::pin(gate.acquire()); + let mut context = Context::from_waker(Waker::noop()); + assert!(first.as_mut().poll(&mut context).is_pending()); + assert!(second.as_mut().poll(&mut context).is_pending()); + gate.set_sampled_limit(2, Duration::from_secs(1)); + let Poll::Ready(one) = first.as_mut().poll(&mut context) else { + panic!("first turn must observe resize") + }; + let Poll::Ready(two) = second.as_mut().poll(&mut context) else { + panic!("second turn must observe resize without releasing the first unit") + }; + tokio::time::advance(Duration::from_secs(1)).await; + drop(one); + drop(two); + let mut expired = Box::pin(gate.acquire()); + assert!(expired.as_mut().poll(&mut context).is_pending(), "stored wake is not a fresh lease"); + gate.set_sampled_limit(1, Duration::from_secs(1)); + let Poll::Ready(renewed) = expired.as_mut().poll(&mut context) else { + panic!("fresh observation must wake expired admission") + }; + drop(renewed); + assert_eq!(gate.snapshot().active, 0); + } + use std::future::Future; + use std::task::Poll; + + #[tokio::test] + async fn shrink_preserves_all_held_units_and_blocks_until_below_new_limit() { + let owner = AdaptiveAdmission::new(NonZeroU16::new(8).expect("bounded test owner"), 3); + let first = owner.acquire().await; + let second = owner.acquire().await; + let third = owner.acquire().await; + owner.set_limit(1); + assert_eq!(owner.snapshot(), AdaptiveAdmissionSnapshot { limit: 1, active: 3 }); + drop(first); + assert!(owner.try_acquire().is_none()); + drop(second); + assert!(owner.try_acquire().is_none()); + drop(third); + let fourth = owner.try_acquire().expect("last old permit returned to original owner"); + assert_eq!(owner.snapshot(), AdaptiveAdmissionSnapshot { limit: 1, active: 1 }); + drop(fourth); + assert_eq!(owner.snapshot().active, 0); + } + + #[tokio::test] + async fn zero_capacity_and_cancelled_waiter_do_not_leak_or_bypass() { + let owner = AdaptiveAdmission::new(NonZeroU16::new(8).expect("bounded test owner"), 0); + let mut waiting = Box::pin(owner.acquire()); + std::future::poll_fn(|cx| { + assert!(waiting.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + drop(waiting); + assert_eq!(owner.snapshot(), AdaptiveAdmissionSnapshot { limit: 0, active: 0 }); + assert!(owner.try_acquire().is_none()); + owner.set_limit(1); + let permit = owner.acquire().await; + assert!(owner.try_acquire().is_none()); + drop(permit); + assert_eq!(owner.snapshot().active, 0); + } + + #[tokio::test] + async fn resize_wakes_a_registered_waiter_without_an_unrelated_release() { + let owner = AdaptiveAdmission::new(NonZeroU16::new(8).expect("bounded test owner"), 0); + let mut waiting = Box::pin(owner.acquire()); + std::future::poll_fn(|cx| { + assert!(waiting.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + owner.set_limit(1); + let permit = tokio::time::timeout(std::time::Duration::from_secs(1), waiting) + .await + .expect("resize notification"); + assert_eq!(owner.snapshot().active, 1); + drop(permit); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn concurrent_resize_release_and_acquire_keep_the_owner_count_exact() { + let owner = AdaptiveAdmission::new(NonZeroU16::new(8).expect("bounded test owner"), 8); + let mut tasks = Vec::new(); + for _ in 0..16 { + let owner = owner.clone(); + tasks.push(tokio::spawn(async move { + for _ in 0..100 { + let permit = owner.acquire().await; + tokio::task::yield_now().await; + drop(permit); + } + })); + } + for limit in [1, 0, 4, 2, 8, 0, u32::MAX, 1] { + owner.set_limit(limit); + tokio::task::yield_now().await; + } + owner.set_limit(8); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + for task in tasks { + task.await.expect("admission task"); + } + }) + .await + .expect("all waiters make progress"); + assert_eq!(owner.snapshot(), AdaptiveAdmissionSnapshot { limit: 8, active: 0 }); + } + + #[tokio::test(start_paused = true)] + async fn expired_observation_blocks_new_work_without_revoking_held_permits() { + let owner = AdaptiveAdmission::new(NonZeroU16::new(8).expect("bounded test owner"), 0); + owner.set_sampled_limit(2, Duration::from_secs(2)); + let held = owner.acquire().await; + tokio::time::advance(Duration::from_secs(2)).await; + assert!(owner.try_acquire().is_none()); + assert_eq!(owner.snapshot(), AdaptiveAdmissionSnapshot { limit: 0, active: 1 }); + drop(held); + assert_eq!(owner.snapshot().active, 0); + owner.set_sampled_limit(1, Duration::from_secs(2)); + let fresh = owner.try_acquire().expect("fresh observation reopens same owner"); + owner.set_sampled_limit(8, Duration::ZERO); + assert!(owner.try_acquire().is_none()); + drop(fresh); + assert_eq!(owner.snapshot().active, 0); + } + + fn healthy(completed: u64) -> FeedbackObservation { + FeedbackObservation { + generation: 1, + max_concurrency: 8, + pressure: ResourcePressure::Stable, + completed_units: Some(completed), + } + } + + #[tokio::test(start_paused = true)] + async fn controller_grows_one_unit_per_cooldown_and_cuts_pressure_immediately() { + let gate = AdaptiveAdmission::new(NonZeroU16::new(8).expect("bounded test owner"), 0); + let mut controller = FeedbackController::new(); + assert_eq!(controller.update(&gate, Some(healthy(0))).limit, 1); + for completed in 1..=25 { + tokio::time::advance(Duration::from_secs(1)).await; + let limit = controller.update(&gate, Some(healthy(completed))).limit; + assert_eq!( + limit, + if completed < 9 { + 1 + } else if completed < 17 { + 2 + } else if completed < 25 { + 3 + } else { + 4 + } + ); + } + let held = gate.try_acquire().expect("existing work"); + let mut pressure = healthy(26); + pressure.pressure = ResourcePressure::Busy; + assert_eq!( + controller.update(&gate, Some(pressure)), + FeedbackDecision { + limit: 2, + reason: FeedbackReason::Pressure + } + ); + pressure.max_concurrency = 0; + assert_eq!(controller.update(&gate, Some(pressure)).limit, 0); + assert_eq!(gate.snapshot().active, 1); + assert!(gate.try_acquire().is_none()); + drop(held); + } + + #[tokio::test(start_paused = true)] + async fn controller_missing_low_sample_stalled_or_changed_sources_never_expand() { + for kind in 0..5 { + let gate = AdaptiveAdmission::new(NonZeroU16::new(8).expect("bounded test owner"), 0); + let mut controller = FeedbackController::new(); + for tick in 0..32 { + let mut observation = healthy(tick); + match kind { + 0 => observation.pressure = ResourcePressure::Unknown, + 1 => observation.completed_units = None, + 2 => observation.completed_units = Some(0), + 3 => observation.generation = tick, + _ => {} + } + tokio::time::advance(Duration::from_secs(if kind == 4 { 3 } else { 1 })).await; + assert!(controller.update(&gate, Some(observation)).limit <= 1, "kind {kind}"); + } + assert_eq!( + controller.update(&gate, None), + FeedbackDecision { + limit: 0, + reason: FeedbackReason::Unavailable + } + ); + assert!(gate.try_acquire().is_none()); + } + } + + #[tokio::test(start_paused = true)] + async fn controller_resumes_after_zero_cap_without_requiring_impossible_progress() { + let gate = AdaptiveAdmission::new(NonZeroU16::new(8).expect("bounded test owner"), 0); + let mut controller = FeedbackController::new(); + let mut sample = healthy(0); + sample.max_concurrency = 0; + assert_eq!(controller.update(&gate, Some(sample)).limit, 0); + tokio::time::advance(Duration::from_secs(1)).await; + sample.max_concurrency = 8; + assert_eq!(controller.update(&gate, Some(sample)).limit, 1); + assert!(gate.try_acquire().is_some(), "minimum progress restarts inside the restored cap"); + } + + #[tokio::test(start_paused = true)] + async fn controller_applies_pressure_reduction_and_lower_cap_in_the_same_window() { + let gate = AdaptiveAdmission::new(NonZeroU16::new(8).expect("bounded test owner"), 0); + let mut controller = FeedbackController::new(); + for completed in 0..=25 { + controller.update(&gate, Some(healthy(completed))); + tokio::time::advance(Duration::from_secs(1)).await; + } + assert_eq!(gate.snapshot().limit, 4); + let mut sample = healthy(26); + sample.max_concurrency = 3; + sample.pressure = ResourcePressure::Busy; + assert_eq!(controller.update(&gate, Some(sample)).limit, 2); + } +} diff --git a/crates/concurrency/src/lib.rs b/crates/concurrency/src/lib.rs index 7e6ba9706..e070e5127 100644 --- a/crates/concurrency/src/lib.rs +++ b/crates/concurrency/src/lib.rs @@ -36,6 +36,7 @@ // Re-exported io-core type used by downstream timeout implementations. pub use rustfs_io_core::OperationProgress; +pub mod adaptive; mod backpressure; mod deadlock; mod queue; diff --git a/crates/ecstore/src/disk/cleanup_runtime.rs b/crates/ecstore/src/disk/cleanup_runtime.rs index 8293e577a..c6f62a398 100644 --- a/crates/ecstore/src/disk/cleanup_runtime.rs +++ b/crates/ecstore/src/disk/cleanup_runtime.rs @@ -16,6 +16,7 @@ //! PUT admission precedes staging; receiver and GC disk work share a separate //! budget so coordinators waiting for peers cannot exhaust receiver capacity. +use crate::disk::io_admission::{current_resource_unit, with_optional_resource_unit}; use crate::disk::{LOG_COMPONENT_ECSTORE, LOG_SUBSYSTEM_DISK}; use std::collections::BTreeSet; use std::future::Future; @@ -422,7 +423,7 @@ impl CleanupRuntime { // the result receiver's lifetime to discard work cancelled during // runtime scheduling, before its first poll can start filesystem I/O. let (result_tx, result_rx) = tokio::sync::oneshot::channel(); - let task = runtime.spawn(async move { + let task = runtime.spawn(with_optional_resource_unit(current_resource_unit(), async move { if result_tx.is_closed() { return; } @@ -430,7 +431,7 @@ impl CleanupRuntime { // through I/O even if the waiter cancels and closes result_rx. let result = execute_disk(pending, Some(active), execution, started, future).await; let _ = result_tx.send(result); - }); + })); match result_rx.await { Ok(result) => result, Err(err) => { @@ -469,11 +470,11 @@ impl CleanupRuntime { gates: Some((self.disk_admission.clone(), self.disk_execution.clone())), }; handle - .spawn(async move { + .spawn(with_optional_resource_unit(current_resource_unit(), async move { let active = execution.acquire_owned().await.map_err(super::error::DiskError::other)?; let _active = Admission::new(active, "rustfs_cleanup_gc_active"); scan(budget).await - }) + })) .await .map_err(super::error::DiskError::other)? } @@ -582,6 +583,7 @@ where F::Output: Send + 'static, { let started = rustfs_io_metrics::put_stage_metrics_enabled().then(Instant::now); + let future = with_optional_resource_unit(current_resource_unit(), future); match reservation { Some(reservation) => reservation.handle.clone().spawn(async move { // This semaphore is never closed. Both queued work and overflow @@ -620,6 +622,62 @@ mod tests { use std::time::Duration; use tokio::sync::oneshot; + #[tokio::test] + async fn cleanup_dispatch_retains_resource_units_without_inheriting_disk_policy() { + use crate::disk::io_admission::{current_disk_io_context, with_background_disk_io, with_background_resource_unit}; + use rustfs_concurrency::adaptive::AdaptiveAdmission; + use std::num::NonZeroU16; + use tokio_util::sync::CancellationToken; + + for mode in 0..4 { + let runtime = Arc::new(CleanupRuntime::new(config(1, 1, 1)).expect("cleanup runtime")); + let worker_runtime = runtime.clone(); + let gate = AdaptiveAdmission::new(NonZeroU16::new(1).expect("one unit"), 1); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let work = async move { + started_tx + .send((current_resource_unit().is_some(), current_disk_io_context().is_none())) + .expect("report worker context"); + release_rx.await.expect("release owned worker"); + Ok::<_, super::super::error::DiskError>(()) + }; + let task = tokio::spawn(with_background_disk_io( + CancellationToken::new(), + with_background_resource_unit(gate.acquire().await, async move { + match mode { + 0 => worker_runtime + .disk_operation(work, DiskAdmission::Request) + .await + .expect("disk operation"), + 1 => worker_runtime.gc_operation(move |_| work).await.expect("GC operation"), + 2 => { + let reservation = worker_runtime.reserve().expect("cleanup reservation"); + run(Some(reservation), async { work.await.expect("cleanup operation") }).await; + } + _ => spawn(None, work) + .await + .expect("unisolated worker") + .expect("unisolated operation"), + } + }), + )); + let (inherited, original_policy) = deadline(started_rx).await.expect("owned work starts"); + task.abort(); + assert!(task.await.expect_err("waiter cancelled").is_cancelled()); + let held = gate.snapshot().active; + release_tx.send(()).expect("owned work survives cancellation"); + deadline(async { + while gate.snapshot().active != 0 { + tokio::task::yield_now().await; + } + }) + .await; + assert!(inherited, "mode {mode} lost resource ownership"); + assert!(original_policy, "mode {mode} changed the detached worker's disk policy"); + assert_eq!(held, 1, "mode {mode} released a live worker's unit"); + } + } fn config(pending: usize, workers: usize, blocking_threads: usize) -> Config { Config { pending, diff --git a/crates/ecstore/src/disk/disk_store.rs b/crates/ecstore/src/disk/disk_store.rs index 90dc5ed78..c510c970a 100644 --- a/crates/ecstore/src/disk/disk_store.rs +++ b/crates/ecstore/src/disk/disk_store.rs @@ -15,7 +15,7 @@ use crate::disk::cleanup_runtime; use crate::disk::io_admission::{ AdmissionReader, BoxedChunkReader, DiskIoAdmission, DiskIoClass, DiskIoContext, current_disk_io_context, - with_disk_io_context, with_disk_io_permit, + current_resource_unit, with_disk_io_context, with_disk_io_permit, with_optional_resource_unit, }; use crate::disk::{ CheckPartsResp, DataDirDeleteStatus, DeleteOptions, DiskAPI, DiskError, DiskInfo, DiskInfoOptions, DiskLocation, Endpoint, @@ -449,10 +449,10 @@ where if external_guard.is_none() { return operation().await; } - tokio::spawn(async move { + tokio::spawn(with_optional_resource_unit(current_resource_unit(), async move { let _external_guard = external_guard; operation().await - }) + })) .await .map_err(|_| Error::other("owned mutation task failed"))? } @@ -1648,6 +1648,7 @@ impl LocalDiskWrapper { cancellation: CancellationToken::new(), admission: None, permit: None, + resource_unit: current_resource_unit(), }) } @@ -3785,32 +3786,49 @@ mod tests { #[tokio::test] async fn owned_mutation_keeps_publication_guard_after_waiter_cancellation() { + use rustfs_concurrency::adaptive::AdaptiveAdmission; + use std::num::NonZeroU16; + + let resources = AdaptiveAdmission::new(NonZeroU16::new(1).expect("one resource unit"), 1); + let resource_unit = resources.acquire().await; let drops = Arc::new(std::sync::atomic::AtomicUsize::new(0)); let guard: Arc = Arc::new(DropProbe(Arc::clone(&drops))); let (started_tx, started_rx) = tokio::sync::oneshot::channel(); let (release_tx, release_rx) = tokio::sync::oneshot::channel(); let (finished_tx, finished_rx) = tokio::sync::oneshot::channel(); - let waiter = tokio::spawn(run_owned_mutation(Some(guard), move || async move { - started_tx.send(()).expect("mutation should signal start"); - release_rx.await.expect("mutation should be released"); - finished_tx.send(()).expect("mutation should signal completion"); - Ok::<_, Error>(()) - })); + let waiter = tokio::spawn(crate::disk::with_background_resource_unit( + resource_unit, + run_owned_mutation(Some(guard), move || async move { + let inherited = current_resource_unit().is_some(); + assert!(current_disk_io_context().is_none(), "preserve the owned task's original disk policy"); + started_tx.send(inherited).expect("mutation should signal start"); + release_rx.await.expect("mutation should be released"); + finished_tx.send(()).expect("mutation should signal completion"); + Ok::<_, Error>(()) + }), + )); - started_rx.await.expect("mutation owner should start"); + let inherited = started_rx.await.expect("mutation owner should start"); waiter.abort(); + assert!(waiter.await.expect_err("caller was cancelled").is_cancelled()); + let held_after_cancel = resources.snapshot().active; assert_eq!(drops.load(std::sync::atomic::Ordering::SeqCst), 0); release_tx.send(()).expect("mutation owner should still be alive"); finished_rx.await.expect("mutation owner should finish"); tokio::time::timeout(Duration::from_secs(1), async { - while drops.load(std::sync::atomic::Ordering::SeqCst) == 0 { + while drops.load(std::sync::atomic::Ordering::SeqCst) == 0 || resources.snapshot().active != 0 { tokio::task::yield_now().await; } }) .await .expect("publication guard should be released after mutation completion"); + assert!(inherited, "owned mutation must retain its resource unit"); + assert_eq!( + held_after_cancel, 1, + "cancelling the waiter must not release a live mutation's resource unit" + ); } struct PendingWriter; diff --git a/crates/ecstore/src/disk/io_admission.rs b/crates/ecstore/src/disk/io_admission.rs index 3a9b63d04..611e5d144 100644 --- a/crates/ecstore/src/disk/io_admission.rs +++ b/crates/ecstore/src/disk/io_admission.rs @@ -24,6 +24,7 @@ use std::task::{Context, Poll}; use std::time::Duration; use metrics::{counter, gauge, histogram}; +use rustfs_concurrency::adaptive::AdaptivePermit; use rustfs_rio::ChunkReader; use tokio::io::{AsyncRead, ReadBuf}; use tokio::sync::Notify; @@ -35,6 +36,7 @@ const BACKGROUND_MAX_WAIT: Duration = Duration::from_millis(100); tokio::task_local! { static DISK_IO_CONTEXT: DiskIoContext; + static RESOURCE_IO_UNIT: AdaptivePermit; } #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -49,6 +51,7 @@ pub(crate) struct DiskIoContext { pub(crate) cancellation: CancellationToken, pub(crate) admission: Option>, pub(crate) permit: Option, + pub(crate) resource_unit: Option, } /// Mark disk I/O issued by `future` as background work for local-disk admission. @@ -60,12 +63,43 @@ pub async fn with_background_disk_io(cancellation: CancellationToken, cancellation, admission: None, permit: None, + resource_unit: current_resource_unit(), }, future, ) .await } +/// Retain an already admitted resource unit without changing disk classification, +/// cancellation or admission. Acquire it before namespace, cache or commit locks. +/// Existing background scopes and disk leases propagate it into native I/O. +/// Use this for physical I/O ownership, not node-wide scheduling capacity that +/// must return with a cancelled waiter so healthy disks can progress independently +/// of a stalled native operation. Escaping buffers need their own memory leases. +#[doc(hidden)] +pub async fn with_background_resource_unit(unit: AdaptivePermit, future: F) -> F::Output { + RESOURCE_IO_UNIT.scope(unit, future).await +} + +pub(crate) fn current_resource_unit() -> Option { + RESOURCE_IO_UNIT.try_with(Clone::clone).ok().or_else(|| { + DISK_IO_CONTEXT + .try_with(|context| context.resource_unit.clone()) + .ok() + .flatten() + }) +} + +/// Transfer only resource ownership into an existing spawned task. Disk +/// classification/permits are not inherited: changing those can introduce a +/// wait behind a disk permit still held by the parent. +pub(crate) async fn with_optional_resource_unit(unit: Option, future: F) -> F::Output { + match unit { + Some(unit) => with_background_resource_unit(unit, future).await, + None => future.await, + } +} + pub(crate) async fn with_disk_io_context(context: DiskIoContext, future: F) -> F::Output { DISK_IO_CONTEXT.scope(context, future).await } @@ -79,7 +113,11 @@ pub(crate) async fn with_disk_io_permit(permit: Option, cancellation: CancellationToken::new(), admission: None, permit: None, + resource_unit: None, }); + if context.resource_unit.is_none() { + context.resource_unit = permit.lease.resource_unit.clone(); + } context.admission = Some(permit.admission()); with_disk_io_context( DiskIoContext { @@ -91,15 +129,32 @@ pub(crate) async fn with_disk_io_permit(permit: Option, .await } -pub(crate) async fn with_optional_disk_io_context(context: Option, future: F) -> F::Output { - match context { - Some(context) => with_disk_io_context(context, future).await, - None => future.await, +pub(crate) fn with_optional_disk_io_context( + context: Option, + future: F, +) -> impl Future { + // Capture at future construction, before JoinSet/spawn moves polling to a + // different task. A resource-only scope has no DiskIoContext to carry it. + let unit = current_resource_unit(); + async move { + let future = with_optional_resource_unit(unit, future); + match context { + Some(context) => with_disk_io_context(context, future).await, + None => future.await, + } } } pub(crate) fn current_disk_io_context() -> Option { - DISK_IO_CONTEXT.try_with(Clone::clone).ok() + DISK_IO_CONTEXT + .try_with(|context| { + let mut context = context.clone(); + if let Ok(unit) = RESOURCE_IO_UNIT.try_with(Clone::clone) { + context.resource_unit = Some(unit); + } + context + }) + .ok() } /// Acquire one low-level permit only when the current walker explicitly owns a gate. @@ -175,6 +230,10 @@ impl DiskIoAdmission { } pub(crate) fn foreground(self: &Arc) -> DiskIoPermit { + self.foreground_with_resource(current_resource_unit()) + } + + fn foreground_with_resource(self: &Arc, resource_unit: Option) -> DiskIoPermit { self.foreground_active.fetch_add(1, Ordering::AcqRel); if self.background_waiters.load(Ordering::Acquire) > 0 { self.record_state(); @@ -183,11 +242,20 @@ impl DiskIoAdmission { lease: Arc::new(DiskIoPermitLease { admission: self.clone(), class: DiskIoClass::Foreground, + resource_unit, }), } } pub(crate) async fn background(self: &Arc, cancellation: &CancellationToken) -> std::io::Result { + self.background_with_resource(cancellation, current_resource_unit()).await + } + + async fn background_with_resource( + self: &Arc, + cancellation: &CancellationToken, + resource_unit: Option, + ) -> std::io::Result { let started_at = Instant::now(); let ticket = { let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); @@ -248,6 +316,7 @@ impl DiskIoAdmission { lease: Arc::new(DiskIoPermitLease { admission: self.clone(), class: DiskIoClass::Background, + resource_unit, }), }); } @@ -360,6 +429,7 @@ impl DiskIoPermit { struct DiskIoPermitLease { admission: Arc, class: DiskIoClass, + resource_unit: Option, } pub(crate) fn spawn_blocking_with_disk_io_permit( @@ -370,8 +440,12 @@ where T: Send + 'static, F: FnOnce() -> T + Send + 'static, { + // A supplied disk permit can predate this resource scope or carry another + // owner's unit. Preserve the current unit independently until native I/O ends. + let resource_unit = current_resource_unit(); tokio::task::spawn_blocking(move || { let _io_permit = io_permit; + let _resource_unit = resource_unit; operation() }) } @@ -436,7 +510,7 @@ impl AdmissionReader { } if self.context.class == DiskIoClass::Foreground { - self.permit = Some(self.admission.foreground()); + self.permit = Some(self.admission.foreground_with_resource(self.context.resource_unit.clone())); return Poll::Ready(Ok(())); } @@ -444,7 +518,10 @@ impl AdmissionReader { if wait.is_none() { let admission = self.admission.clone(); let cancellation = self.context.cancellation.clone(); - *wait = Some(Box::pin(async move { admission.background(&cancellation).await })); + let resource_unit = self.context.resource_unit.clone(); + *wait = Some(Box::pin( + async move { admission.background_with_resource(&cancellation, resource_unit).await }, + )); } let Some(wait_future) = wait.as_mut() else { return Poll::Ready(Err(io::Error::other("background disk I/O admission waiter is missing"))); @@ -514,8 +591,166 @@ impl ChunkReader for BoxedChunkReader { #[cfg(test)] mod tests { use super::*; + use rustfs_concurrency::adaptive::AdaptiveAdmission; + use std::num::NonZeroU16; use tokio::io::{AsyncReadExt, AsyncWriteExt}; + #[tokio::test] + async fn resource_unit_outlives_cancelled_scope_until_native_io_finishes() { + for permit_mode in 0..4 { + let gate = AdaptiveAdmission::new(NonZeroU16::new(1).expect("one unit"), 1); + let previous_gate = AdaptiveAdmission::new(NonZeroU16::new(1).expect("previous owner"), 1); + let unit = gate.acquire().await; + let disk = DiskIoAdmission::new("resource-unit-native"); + let existing_permit = match permit_mode { + 2 => Some(disk.foreground()), + 3 => Some(with_background_resource_unit(previous_gate.acquire().await, async { disk.foreground() }).await), + _ => None, + }; + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let task = tokio::spawn(with_background_resource_unit(unit, async move { + with_background_disk_io(CancellationToken::new(), async { + let permit = if existing_permit.is_some() { + existing_permit + } else if permit_mode == 1 { + Some(disk.background(&CancellationToken::new()).await.expect("disk permit")) + } else { + None + }; + spawn_blocking_with_disk_io_permit(permit, move || { + started_tx.send(()).expect("native I/O started"); + release_rx.blocking_recv().expect("native I/O release"); + }) + .await + .expect("native I/O completed"); + }) + .await; + })); + started_rx.await.expect("native worker owns unit"); + task.abort(); + assert!(task.await.expect_err("scope was cancelled").is_cancelled()); + let held = gate.snapshot().active; + let previous_held = previous_gate.snapshot().active; + let rejected = gate.try_acquire().is_none(); + gate.set_limit(0); + release_tx.send(()).expect("release native I/O"); + tokio::time::timeout(Duration::from_secs(5), async { + while gate.snapshot().active != 0 || previous_gate.snapshot().active != 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("native completion releases the unit"); + gate.set_limit(1); + assert!(gate.try_acquire().is_some()); + assert_eq!(held, 1, "mode {permit_mode}: native I/O still owns the current unit"); + assert!(rejected, "mode {permit_mode}: a live native worker must retain admission"); + assert_eq!(previous_held, u32::from(permit_mode == 3), "preserve any older disk lease's unit"); + } + } + + #[tokio::test] + async fn escaped_disk_permit_restores_resource_context_without_a_new_unit() { + let gate = AdaptiveAdmission::new(NonZeroU16::new(1).expect("one unit"), 1); + let disk = DiskIoAdmission::new("resource-unit-escape"); + let permit = with_background_resource_unit(gate.acquire().await, async { + disk.background(&CancellationToken::new()).await.expect("disk permit") + }) + .await; + assert_eq!(gate.snapshot().active, 1); + with_background_disk_io(CancellationToken::new(), async { + with_disk_io_permit(Some(permit), async { + let unit = current_resource_unit().expect("escaped disk lease restores unit into an existing context"); + assert_eq!(gate.snapshot().active, 1, "cloning retains one unit rather than charging twice"); + drop(unit); + }) + .await; + }) + .await; + assert_eq!(gate.snapshot().active, 0); + } + + #[tokio::test] + async fn escaped_stream_retains_resource_unit_outside_task_local_scope() { + let gate = AdaptiveAdmission::new(NonZeroU16::new(1).expect("one unit"), 1); + let disk = DiskIoAdmission::new("resource-unit-stream"); + let (mut writer, inner) = tokio::io::duplex(8); + let mut reader = with_background_resource_unit(gate.acquire().await, async { + with_background_disk_io(CancellationToken::new(), async { + AdmissionReader::new(inner, disk, current_disk_io_context().expect("resource context")) + }) + .await + }) + .await; + assert_eq!(gate.snapshot().active, 1); + writer.write_all(b"rust").await.expect("write complete stream"); + drop(writer); + let mut bytes = Vec::new(); + reader.read_to_end(&mut bytes).await.expect("read outside original scope"); + assert_eq!(bytes, b"rust"); + assert_eq!(gate.snapshot().active, 1, "reader still owns its context"); + drop(reader); + assert_eq!(gate.snapshot().active, 0); + } + + #[tokio::test] + async fn resource_scope_preserves_existing_disk_policy_and_cancellation() { + let gate = AdaptiveAdmission::new(NonZeroU16::new(1).expect("one unit"), 1); + with_background_resource_unit(gate.acquire().await, async { + assert!(current_disk_io_context().is_none(), "resource ownership does not classify disk I/O"); + assert!(current_resource_unit().is_some()); + }) + .await; + let cancellation = CancellationToken::new(); + let existing = DiskIoContext { + class: DiskIoClass::Foreground, + cancellation: cancellation.clone(), + admission: None, + permit: None, + resource_unit: None, + }; + with_disk_io_context(existing, async { + with_background_resource_unit(gate.acquire().await, async { + cancellation.cancel(); + let context = current_disk_io_context().expect("existing disk context"); + assert_eq!(context.class, DiskIoClass::Foreground); + assert!( + context.resource_unit.is_some(), + "captured disk context must retain the current resource unit" + ); + assert!(context.cancellation.is_cancelled(), "retain the original cancellation token"); + assert!(current_resource_unit().is_some()); + }) + .await; + }) + .await; + assert_eq!(gate.snapshot().active, 0); + } + + #[tokio::test] + async fn resource_only_scope_survives_existing_optional_context_spawn_boundary() { + let gate = AdaptiveAdmission::new(NonZeroU16::new(1).expect("one unit"), 1); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let mut worker = None; + with_background_resource_unit(gate.acquire().await, async { + worker = Some(tokio::spawn(with_optional_disk_io_context(current_disk_io_context(), async { + release_rx.await.expect("release child after caller scope ends"); + assert!(current_disk_io_context().is_none()); + assert!(current_resource_unit().is_some()); + }))); + }) + .await; + let held = gate.snapshot().active; + release_tx.send(()).expect("child remains alive"); + worker + .expect("scoped child was spawned") + .await + .expect("child preserves resource-only scope"); + assert_eq!(held, 1); + assert_eq!(gate.snapshot().active, 0); + } + struct BlockingWorker(tokio::task::JoinHandle<()>); async fn wait_for_background_waiters(admission: &DiskIoAdmission, expected: usize) { @@ -636,6 +871,7 @@ mod tests { cancellation: CancellationToken::new(), admission: None, permit: None, + resource_unit: None, }, ); let read = tokio::spawn(async move { diff --git a/crates/ecstore/src/disk/local.rs b/crates/ecstore/src/disk/local.rs index 5e60c9d3d..febf6765f 100644 --- a/crates/ecstore/src/disk/local.rs +++ b/crates/ecstore/src/disk/local.rs @@ -50,7 +50,7 @@ use crate::diagnostics::get::{ #[cfg(test)] use crate::disk::HEALING_MARKER_PATH; use crate::disk::disk_store::{get_drive_walkdir_stall_timeout, get_object_disk_read_timeout}; -use crate::disk::io_admission::{DiskIoPermit, spawn_blocking_with_disk_io_permit}; +use crate::disk::io_admission::{DiskIoPermit, current_resource_unit, spawn_blocking_with_disk_io_permit}; use crate::disk::{ BUCKET_META_PREFIX, CHECK_PART_FILE_CORRUPT, CHECK_PART_FILE_NOT_FOUND, CHECK_PART_SUCCESS, CHECK_PART_UNKNOWN, CHECK_PART_VOLUME_NOT_FOUND, CheckPartsResp, ConditionalFileUpdate, CurrentDeleteGeneration, DataDirDeleteStatus, @@ -79,6 +79,7 @@ use metrics::counter; #[cfg(target_os = "linux")] use metrics::gauge; use parking_lot::{Mutex as ParkingLotMutex, RwLock as ParkingLotRwLock}; +use rustfs_concurrency::adaptive::AdaptivePermit; use rustfs_filemeta::{ Cache, FileInfo, FileInfoOpts, FileMeta, MetaCacheEntry, MetacacheWriter, ObjectPartInfo, Opts, RawFileInfo, UpdateFn, ValidationMode, get_file_info, read_xl_meta_no_data_sync, @@ -314,6 +315,7 @@ enum MetadataReadOp { } struct MetadataReadJob { + resource_unit: Option, path: PathBuf, io_permit: Option, op: MetadataReadOp, @@ -335,6 +337,7 @@ impl MetadataReadJob { if self.is_cancelled() { return; } + let _resource_unit = self.resource_unit; let _io_permit = self.io_permit; match self.op { MetadataReadOp::ReadAll { tx } => { @@ -394,7 +397,8 @@ impl MetadataReadQueue { }) } - fn enqueue(self: &Arc, job: MetadataReadJob) { + fn enqueue(self: &Arc, mut job: MetadataReadJob) { + job.resource_unit = current_resource_unit(); let should_spawn = { let mut state = self.state.lock(); state.jobs.push_back(job); @@ -8252,6 +8256,7 @@ impl LocalDisk { let path = file_path.as_ref().to_path_buf(); let (tx, rx) = tokio::sync::oneshot::channel(); self.metadata_read_queue.enqueue(MetadataReadJob { + resource_unit: None, path, io_permit, op: MetadataReadOp::ReadMetadataOnly { tx }, @@ -8429,6 +8434,7 @@ impl LocalDisk { let path = file_path.as_ref().to_path_buf(); let (tx, rx) = tokio::sync::oneshot::channel(); self.metadata_read_queue.enqueue(MetadataReadJob { + resource_unit: None, path, io_permit, op: MetadataReadOp::ReadAll { tx }, @@ -25858,6 +25864,46 @@ mod test { } } + #[tokio::test] + async fn metadata_read_job_retains_resource_unit_until_native_completion() { + use crate::disk::with_background_resource_unit; + use rustfs_concurrency::adaptive::AdaptiveAdmission; + use std::num::NonZeroU16; + + let dir = tempfile::tempdir().expect("metadata root"); + let queue = MetadataReadQueue::new(os::PublicationRoot::new(dir.path()).expect("publication root")); + let gate = AdaptiveAdmission::new(NonZeroU16::new(1).expect("one unit"), 1); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + with_background_resource_unit(gate.acquire().await, async { + queue.enqueue(MetadataReadJob { + resource_unit: None, + path: PathBuf::new(), + io_permit: None, + op: MetadataReadOp::TestCustom { + is_closed: Box::new(|| false), + run: Box::new(move || { + started_tx.send(()).expect("native metadata worker started"); + release_rx.blocking_recv().expect("release native metadata work"); + }), + }, + }); + started_rx.await.expect("worker owns queued job"); + }) + .await; + let held = gate.snapshot().active; + release_tx.send(()).expect("release native job after caller scope ends"); + tokio::time::timeout(Duration::from_secs(5), async { + while gate.snapshot().active != 0 || queue.state.lock().active_workers != 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("native metadata job and worker retire"); + assert_eq!(held, 1, "leaving the caller scope must not release a live metadata job's unit"); + assert!(gate.try_acquire().is_some()); + } + #[tokio::test] async fn metadata_read_queue_bounds_concurrency_skips_cancelled_and_recovers_from_panic() { let dir = tempfile::tempdir().expect("tempdir should be created"); @@ -25878,6 +25924,7 @@ mod test { let entered_tx = entered_tx.clone(); let release_rx = Arc::clone(&release_rx); queue.enqueue(MetadataReadJob { + resource_unit: None, path: PathBuf::new(), io_permit: None, op: MetadataReadOp::TestCustom { @@ -25905,12 +25952,14 @@ mod test { let (cancelled_tx, cancelled_rx) = tokio::sync::oneshot::channel(); drop(cancelled_rx); queue.enqueue(MetadataReadJob { + resource_unit: None, path: dir.path().join("cancelled-never-opened"), io_permit: None, op: MetadataReadOp::ReadMetadataOnly { tx: cancelled_tx }, }); let executed_cancelled_clone = Arc::clone(&executed_cancelled); queue.enqueue(MetadataReadJob { + resource_unit: None, path: PathBuf::new(), io_permit: None, op: MetadataReadOp::TestCustom { @@ -25923,6 +25972,7 @@ mod test { // Enqueue a panicking job followed by a normal completion job to verify panic respawn. queue.enqueue(MetadataReadJob { + resource_unit: None, path: PathBuf::new(), io_permit: None, op: MetadataReadOp::TestCustom { @@ -25934,6 +25984,7 @@ mod test { let active_trailing = Arc::clone(&active); let max_active_trailing = Arc::clone(&max_active); queue.enqueue(MetadataReadJob { + resource_unit: None, path: PathBuf::new(), io_permit: None, op: MetadataReadOp::TestCustom { diff --git a/crates/ecstore/src/disk/mod.rs b/crates/ecstore/src/disk/mod.rs index ae7cbdb05..b4219971d 100644 --- a/crates/ecstore/src/disk/mod.rs +++ b/crates/ecstore/src/disk/mod.rs @@ -31,8 +31,11 @@ pub mod os; #[cfg(test)] pub(crate) use io_admission::DiskIoClass; -pub use io_admission::with_background_disk_io; -pub(crate) use io_admission::{acquire_context_disk_io_permit, current_disk_io_context, with_optional_disk_io_context}; +pub(crate) use io_admission::{ + acquire_context_disk_io_permit, current_disk_io_context, current_resource_unit, with_optional_disk_io_context, + with_optional_resource_unit, +}; +pub use io_admission::{with_background_disk_io, with_background_resource_unit}; pub const RUSTFS_META_BUCKET: &str = ".rustfs.sys"; pub const MIGRATING_META_BUCKET: &str = ".minio.sys"; diff --git a/crates/ecstore/src/disk/os.rs b/crates/ecstore/src/disk/os.rs index d5742fd56..4dcb5757a 100644 --- a/crates/ecstore/src/disk/os.rs +++ b/crates/ecstore/src/disk/os.rs @@ -759,9 +759,15 @@ impl OpenedDstDirFsyncGroup { struct DstDirFsyncWaiter { namespace_owner: Option>, + resource_unit: Option, result_tx: oneshot::Sender, } +struct DstDirFsyncOwner { + _namespace: Option>, + _resource: Option, +} + #[derive(Clone)] struct SharedDstDirFsyncError { kind: io::ErrorKind, @@ -851,6 +857,7 @@ impl DstDirFsyncGroupCommit { group_state.pending.push_back(DstDirFsyncWaiter { result_tx, namespace_owner, + resource_unit: super::io_admission::current_resource_unit(), }); let start_worker = !group_state.worker_running; if start_worker { @@ -907,7 +914,7 @@ impl DstDirFsyncGroupCommit { } #[cfg(unix)] -async fn fsync_open_dst_dir_group(group: &DstDirFsyncGroup, namespace_owners: Vec>) -> io::Result<()> { +async fn fsync_open_dst_dir_group(group: &DstDirFsyncGroup, namespace_owners: Vec) -> io::Result<()> { #[cfg(test)] let dir = group.dir.clone(); let dir_file = group.dir_file.clone(); @@ -930,7 +937,7 @@ async fn fsync_open_dst_dir_group(group: &DstDirFsyncGroup, namespace_owners: Ve } #[cfg(not(unix))] -async fn fsync_open_dst_dir_group(group: &DstDirFsyncGroup, namespace_owners: Vec>) -> io::Result<()> { +async fn fsync_open_dst_dir_group(group: &DstDirFsyncGroup, namespace_owners: Vec) -> io::Result<()> { let _namespace_owners = namespace_owners; fsync_dir(&group.dir).await } @@ -997,7 +1004,17 @@ fn run_dst_dir_fsync_group_worker(group: Arc) -> impl std::fut #[cfg(test)] fsync_dir_recorder::record_grouped(&group.dir, batch.len()); - let namespace_owners = batch.iter_mut().filter_map(|waiter| waiter.namespace_owner.take()).collect(); + let namespace_owners = batch + .iter_mut() + .filter_map(|waiter| { + let namespace = waiter.namespace_owner.take(); + let resource = waiter.resource_unit.take(); + (namespace.is_some() || resource.is_some()).then_some(DstDirFsyncOwner { + _namespace: namespace, + _resource: resource, + }) + }) + .collect(); let result = fsync_open_dst_dir_group(&group, namespace_owners) .await .map_err(SharedDstDirFsyncError::from_error); @@ -1056,7 +1073,7 @@ async fn fsync_dst_dir_group_commit_opened( } let dir = dir.to_path_buf(); - let opened = tokio::task::spawn_blocking(move || { + let opened = super::io_admission::spawn_blocking_with_disk_io_permit(None, move || { #[cfg(target_os = "linux")] if let Some(file) = opened { return OpenedDstDirFsyncGroup::from_file(&dir, file); @@ -1190,6 +1207,7 @@ type FileFdatasyncGroupKey = (usize, bool); struct FileFdatasyncWaiter { files: Vec, + resource_unit: Option, enqueued_at: Option, wait_role: &'static str, result_tx: oneshot::Sender, @@ -1260,7 +1278,8 @@ impl FileFdatasyncGroupCommit { let (result_tx, result_rx) = oneshot::channel(); let key = (Arc::as_ptr(&disk_permits) as usize, FSYNC_ON_CURRENT_RUNTIME.get()); let mut registry = self.inner.lock(); - registry.groups.retain(|_, group| group.disk_permits.strong_count() > 0); + // A retired limiter can still have queued or native work. Its worker + // owns the registry entry and returns that budget on completion/cancel. if registry.total_waiters >= MAX_FILE_FDATASYNC_WAITERS { return Err(io::Error::new( io::ErrorKind::WouldBlock, @@ -1300,6 +1319,7 @@ impl FileFdatasyncGroupCommit { }; group_state.pending.push_back(FileFdatasyncWaiter { files, + resource_unit: super::io_admission::current_resource_unit(), enqueued_at: rustfs_io_metrics::put_stage_timer(), wait_role, result_tx, @@ -1332,7 +1352,13 @@ impl FileFdatasyncGroupCommit { fn remove_idle_group(&self, group: &Arc) { let mut registry = self.inner.lock(); let group_state = group.inner.lock(); - if !group_state.worker_running && group_state.pending.is_empty() { + if !group_state.worker_running + && group_state.pending.is_empty() + && registry + .groups + .get(&group.key) + .is_some_and(|registered| Arc::ptr_eq(registered, group)) + { registry.groups.remove(&group.key); } } @@ -1352,80 +1378,142 @@ impl FileFdatasyncGroupCommit { } } -async fn run_file_fdatasync_group_worker(group: Arc) { - loop { - #[cfg(test)] - file_sync_probe::run_before_group_batch(); - tokio::task::yield_now().await; - let wait = file_fdatasync_group_commit_wait(); - if !wait.is_zero() { - tokio::time::sleep(wait).await; - } - let (batch, batch_file_count): (Vec, usize) = { - let mut group_state = group.inner.lock(); - let batch_file_count = group_state.pending_files; - group_state.pending_files = 0; - (group_state.pending.drain(..).collect(), batch_file_count) - }; - if batch.is_empty() { - let mut group_state = group.inner.lock(); - group_state.worker_running = false; - drop(group_state); - FILE_FDATASYNC_GROUP_COMMIT.remove_idle_group(&group); +struct FileFdatasyncWorkerGuard { + group: Arc, + in_flight_waiters: usize, + in_flight_files: usize, + armed: bool, +} + +impl Drop for FileFdatasyncWorkerGuard { + fn drop(&mut self) { + if !self.armed { return; } - - rustfs_io_metrics::record_put_rename_fdatasync_group_outstanding( - rustfs_io_metrics::PUT_RENAME_FDATASYNC_GROUP_OUTSTANDING_STATE_BATCH_WAITERS, - batch.len(), - ); - rustfs_io_metrics::record_put_rename_fdatasync_group_outstanding( - rustfs_io_metrics::PUT_RENAME_FDATASYNC_GROUP_OUTSTANDING_STATE_BATCH_FILES, - batch_file_count, - ); - for waiter in &batch { - if let Some(enqueued_at) = waiter.enqueued_at { - rustfs_io_metrics::record_put_rename_fdatasync_group_wait( - waiter.wait_role, - enqueued_at.elapsed().as_secs_f64() * 1000.0, - ); + let pending = { + // Same registry -> group order as enqueue and normal completion. + let mut registry = FILE_FDATASYNC_GROUP_COMMIT.inner.lock(); + let mut state = self.group.inner.lock(); + let pending = std::mem::take(&mut state.pending); + let pending_files = std::mem::take(&mut state.pending_files); + state.worker_running = false; + if registry + .groups + .get(&self.group.key) + .is_some_and(|group| Arc::ptr_eq(group, &self.group)) + { + registry.total_waiters = registry.total_waiters.saturating_sub(pending.len() + self.in_flight_waiters); + registry.total_files = registry.total_files.saturating_sub(pending_files + self.in_flight_files); + registry.groups.remove(&self.group.key); } - } - let batch_files: Vec = batch.iter().flat_map(|waiter| waiter.files.iter().cloned()).collect(); - #[cfg(test)] - file_sync_probe::record_group_batch(batch_file_count); - rustfs_io_metrics::record_put_rename_fdatasync_batch( - rustfs_io_metrics::PUT_RENAME_FDATASYNC_BATCH_MODE_PARALLEL, - batch_file_count, - ); - let result = if let Some(disk_permits) = group.disk_permits.upgrade() { - run_file_sync_blocking(disk_permits, move || sync_files(&batch_files)) + pending + }; + // Queued work never reached the disk; native work retains its own units. + // Destructors and notification wakeups run outside the registry locks. + drop(pending); + } +} + +fn run_file_fdatasync_group_worker(group: Arc) -> impl std::future::Future { + // Capture before spawning so even an unpolled worker releases its queue. + let worker_guard = FileFdatasyncWorkerGuard { + group: group.clone(), + in_flight_waiters: 0, + in_flight_files: 0, + armed: true, + }; + async move { + let mut worker_guard = worker_guard; + loop { + #[cfg(test)] + file_sync_probe::run_before_group_batch(); + tokio::task::yield_now().await; + let wait = file_fdatasync_group_commit_wait(); + if !wait.is_zero() { + tokio::time::sleep(wait).await; + } + let (mut batch, batch_file_count): (Vec, usize) = { + let mut group_state = group.inner.lock(); + let batch_file_count = group_state.pending_files; + group_state.pending_files = 0; + (group_state.pending.drain(..).collect(), batch_file_count) + }; + if batch.is_empty() { + let mut group_state = group.inner.lock(); + worker_guard.armed = false; + group_state.worker_running = false; + drop(group_state); + FILE_FDATASYNC_GROUP_COMMIT.remove_idle_group(&group); + return; + } + + worker_guard.in_flight_waiters = batch.len(); + worker_guard.in_flight_files = batch_file_count; + rustfs_io_metrics::record_put_rename_fdatasync_group_outstanding( + rustfs_io_metrics::PUT_RENAME_FDATASYNC_GROUP_OUTSTANDING_STATE_BATCH_WAITERS, + batch.len(), + ); + rustfs_io_metrics::record_put_rename_fdatasync_group_outstanding( + rustfs_io_metrics::PUT_RENAME_FDATASYNC_GROUP_OUTSTANDING_STATE_BATCH_FILES, + batch_file_count, + ); + for waiter in &batch { + if let Some(enqueued_at) = waiter.enqueued_at { + rustfs_io_metrics::record_put_rename_fdatasync_group_wait( + waiter.wait_role, + enqueued_at.elapsed().as_secs_f64() * 1000.0, + ); + } + } + // Move each producer's existing path buffer and unit into the physical + // batch. No path cloning or separate resource-owner allocation is needed. + let batch_files: Vec<_> = batch + .iter_mut() + .map(|waiter| (std::mem::take(&mut waiter.files), waiter.resource_unit.take())) + .collect(); + #[cfg(test)] + file_sync_probe::record_group_batch(batch_file_count); + rustfs_io_metrics::record_put_rename_fdatasync_batch( + rustfs_io_metrics::PUT_RENAME_FDATASYNC_BATCH_MODE_PARALLEL, + batch_file_count, + ); + let result = if let Some(disk_permits) = group.disk_permits.upgrade() { + run_file_sync_blocking(disk_permits, move || { + for (files, _resource_unit) in &batch_files { + sync_files(files)?; + } + Ok(()) + }) .await .map_err(SharedFileFdatasyncError::from_error) - } else { - Err(SharedFileFdatasyncError::from_error(io::Error::other( - "file fdatasync group commit limiter dropped", - ))) - }; - FILE_FDATASYNC_GROUP_COMMIT.complete_batch(batch.len(), batch_file_count); - - let should_stop = { - let mut group_state = group.inner.lock(); - if group_state.pending.is_empty() { - group_state.worker_running = false; - true } else { - false + Err(SharedFileFdatasyncError::from_error(io::Error::other( + "file fdatasync group commit limiter dropped", + ))) + }; + FILE_FDATASYNC_GROUP_COMMIT.complete_batch(batch.len(), batch_file_count); + worker_guard.in_flight_waiters = 0; + worker_guard.in_flight_files = 0; + + let should_stop = { + let mut group_state = group.inner.lock(); + if group_state.pending.is_empty() { + worker_guard.armed = false; + group_state.worker_running = false; + true + } else { + false + } + }; + if should_stop { + FILE_FDATASYNC_GROUP_COMMIT.remove_idle_group(&group); + } + for waiter in batch { + let _ = waiter.result_tx.send(result.clone()); + } + if should_stop { + return; } - }; - if should_stop { - FILE_FDATASYNC_GROUP_COMMIT.remove_idle_group(&group); - } - for waiter in batch { - let _ = waiter.result_tx.send(result.clone()); - } - if should_stop { - return; } } } @@ -1513,6 +1601,11 @@ pub(super) fn is_isolated_io_thread() -> bool { /// Spawn a blocking task on the fsync-dedicated runtime if configured, /// otherwise fall back to the main tokio blocking pool. fn fsync_spawn_blocking(f: impl FnOnce() -> T + Send + 'static) -> tokio::task::JoinHandle { + let resource_unit = super::io_admission::current_resource_unit(); + let f = move || { + let _resource_unit = resource_unit; + f() + }; if FSYNC_ON_CURRENT_RUNTIME.get() { return tokio::task::spawn_blocking(f); } @@ -1693,7 +1786,10 @@ impl DiskBlockingQueue { operation: impl FnOnce(&dyn Fn() -> bool) -> io::Result + Send + 'static, ) -> io::Result { let (tx, rx) = oneshot::channel(); + let resource_unit = super::io_admission::current_resource_unit(); self.enqueue(Box::new(move || { + // Retain each producer separately, never on the shared worker. + let _resource_unit = resource_unit; let result = operation(&|| tx.is_closed()); let _ = tx.send(result); })); @@ -7734,6 +7830,10 @@ mod tests { #[tokio::test] #[serial_test::serial(dst_dir_fsync_group_commit)] async fn grouped_fsync_physical_batch_keeps_all_owners_after_worker_cancellation() { + use crate::disk::io_admission::with_background_resource_unit; + use rustfs_concurrency::adaptive::AdaptiveAdmission; + use std::num::NonZeroU16; + let resources: [_; 3] = std::array::from_fn(|_| AdaptiveAdmission::new(NonZeroU16::new(1).expect("one unit"), 1)); let temp_dir = tempdir().expect("fixture directory"); let dir = temp_dir.path().canonicalize().expect("canonical fsync path"); let first_ctx = Arc::new(crate::runtime::instance::InstanceContext::new()); @@ -7742,19 +7842,25 @@ mod tests { let second_owner = second_ctx.begin_namespace_commit(); let first_probe = Arc::downgrade(&first_owner); let second_probe = Arc::downgrade(&second_owner); - let (first_rx, group) = DST_DIR_FSYNC_GROUP_COMMIT - .enqueue_opened( - OpenedDstDirFsyncGroup::open(&dir).expect("open first waiter directory"), - Some(first_owner), - ) - .expect("queue first real waiter"); + let (first_rx, group) = with_background_resource_unit(resources[0].acquire().await, async { + DST_DIR_FSYNC_GROUP_COMMIT + .enqueue_opened( + OpenedDstDirFsyncGroup::open(&dir).expect("open first waiter directory"), + Some(first_owner), + ) + .expect("queue first real waiter") + }) + .await; let group = group.expect("first waiter starts the group"); - let (second_rx, second_worker) = DST_DIR_FSYNC_GROUP_COMMIT - .enqueue_opened( - OpenedDstDirFsyncGroup::open(&dir).expect("open second waiter directory"), - Some(second_owner), - ) - .expect("queue second real waiter"); + let (second_rx, second_worker) = with_background_resource_unit(resources[1].acquire().await, async { + DST_DIR_FSYNC_GROUP_COMMIT + .enqueue_opened( + OpenedDstDirFsyncGroup::open(&dir).expect("open second waiter directory"), + Some(second_owner), + ) + .expect("queue second real waiter") + }) + .await; assert!(second_worker.is_none(), "same directory must join the same batch"); assert_eq!(group.inner.lock().pending.len(), 2); let (entered_tx, entered_rx) = tokio::sync::oneshot::channel(); @@ -7778,12 +7884,15 @@ mod tests { let queued_owner = queued_ctx.begin_namespace_commit(); let queued_probe = Arc::downgrade(&queued_owner); let queued_generation = queued_ctx.namespace_commit_generation(); - let (queued_rx, queued_worker) = DST_DIR_FSYNC_GROUP_COMMIT - .enqueue_opened( - OpenedDstDirFsyncGroup::open(&dir).expect("open queued waiter directory"), - Some(queued_owner), - ) - .expect("queue a waiter after the physical batch was frozen"); + let (queued_rx, queued_worker) = with_background_resource_unit(resources[2].acquire().await, async { + DST_DIR_FSYNC_GROUP_COMMIT + .enqueue_opened( + OpenedDstDirFsyncGroup::open(&dir).expect("open queued waiter directory"), + Some(queued_owner), + ) + .expect("queue a waiter after the physical batch was frozen") + }) + .await; assert!(queued_worker.is_none()); assert_eq!(group.inner.lock().pending.len(), 1); drop((first_rx, second_rx)); @@ -7798,6 +7907,7 @@ mod tests { assert_eq!(DST_DIR_FSYNC_GROUP_COMMIT.counts_for_test(), (0, 0)); let first_pending = first_ctx.namespace_commits_pending() && first_probe.upgrade().is_some(); let second_pending = second_ctx.namespace_commits_pending() && second_probe.upgrade().is_some(); + let retained_resources = resources.each_ref().map(|gate| gate.snapshot().active); let generations = (first_ctx.namespace_commit_generation(), second_ctx.namespace_commit_generation()); drop(release_tx); tokio::time::timeout(Duration::from_secs(5), async { @@ -7812,6 +7922,12 @@ mod tests { }) .await .expect("physical fsync must release every batch owner"); + assert_eq!( + retained_resources, + [1, 1, 0], + "physical batch retains each producer; queued cancellation releases its unit" + ); + assert_eq!(resources.each_ref().map(|gate| gate.snapshot().active), [0; 3]); assert!(fsync_dir_recorder::was_fsynced(&dir), "the detached syscall must really execute"); assert!( first_pending && second_pending, @@ -8794,6 +8910,217 @@ mod tests { assert_eq!(err.kind(), io::ErrorKind::NotFound); } + #[tokio::test] + #[serial_test::serial(file_sync_probe)] + async fn resource_grouped_fdatasync_retired_limiter_keeps_its_budget_until_worker_cleanup() { + use crate::disk::io_admission::with_background_resource_unit; + use rustfs_concurrency::adaptive::AdaptiveAdmission; + use std::num::NonZeroU16; + + let gate = AdaptiveAdmission::new(NonZeroU16::new(1).expect("one unit"), 1); + let retired = Arc::new(Semaphore::new(1)); + let (old_rx, old_group) = with_background_resource_unit(gate.acquire().await, async { + FILE_FDATASYNC_GROUP_COMMIT.enqueue(retired.clone(), vec![PathBuf::from("retired")]) + }) + .await + .expect("enqueue retiring producer"); + let old_worker = run_file_fdatasync_group_worker(old_group.expect("retiring group")); + drop(retired); + + let live = Arc::new(Semaphore::new(1)); + let (new_rx, new_group) = FILE_FDATASYNC_GROUP_COMMIT + .enqueue(live, vec![PathBuf::from("successor")]) + .expect("enqueue new producer"); + let new_worker = run_file_fdatasync_group_worker(new_group.expect("new group")); + let before_cleanup = FILE_FDATASYNC_GROUP_COMMIT.counts_for_test(); + drop(old_worker); + let after_old = FILE_FDATASYNC_GROUP_COMMIT.counts_for_test(); + drop(new_worker); + let after_both = FILE_FDATASYNC_GROUP_COMMIT.counts_for_test(); + clear_file_fdatasync_group_commit_for_test(); + assert!(old_rx.await.is_err()); + assert!(new_rx.await.is_err()); + assert_eq!(gate.snapshot().active, 0); + assert_eq!(before_cleanup, (2, 2, 2), "retiring queued work still owns a registry budget"); + assert_eq!(after_old, (1, 1, 1), "old cleanup must return exactly its own budget"); + assert_eq!(after_both, (0, 0, 0), "no phantom waiters or files after cancellation"); + } + + #[tokio::test] + #[serial_test::serial(file_sync_probe)] + async fn resource_grouped_fdatasync_unpolled_worker_releases_its_producer() { + use crate::disk::io_admission::with_background_resource_unit; + use rustfs_concurrency::adaptive::AdaptiveAdmission; + use std::num::NonZeroU16; + + let gate = AdaptiveAdmission::new(NonZeroU16::new(1).expect("one unit"), 1); + let limiter = Arc::new(Semaphore::new(1)); + let (rx, group) = with_background_resource_unit(gate.acquire().await, async { + FILE_FDATASYNC_GROUP_COMMIT.enqueue(limiter, vec![PathBuf::from("never-dispatched")]) + }) + .await + .expect("enqueue producer"); + let group = group.expect("first producer starts group"); + let held_before_drop = gate.snapshot().active; + drop(run_file_fdatasync_group_worker(group.clone())); + let counts = FILE_FDATASYNC_GROUP_COMMIT.counts_for_test(); + let held_after_drop = gate.snapshot().active; + let pending = group.inner.lock().pending.len(); + let running = group.inner.lock().worker_running; + // Also drain a failing implementation before making assertions. + drop(std::mem::take(&mut group.inner.lock().pending)); + clear_file_fdatasync_group_commit_for_test(); + assert!(rx.await.is_err()); + assert_eq!((held_before_drop, held_after_drop), (1, 0)); + assert_eq!(counts, (0, 0, 0)); + assert_eq!((pending, running), (0, false)); + } + + #[tokio::test] + async fn file_fdatasync_stale_idle_cleanup_preserves_the_successor_group() { + let registry = FileFdatasyncGroupCommit::default(); + let limiter = Arc::new(Semaphore::new(1)); + let (first_rx, old) = registry + .enqueue(limiter.clone(), vec![PathBuf::from("old")]) + .expect("first group"); + let old = old.expect("first group starts worker"); + let first_waiter = old.inner.lock().pending.pop_front().expect("freeze first batch"); + old.inner.lock().pending_files = 0; + registry.complete_batch(1, 1); + old.inner.lock().worker_running = false; + registry.remove_idle_group(&old); + assert_eq!(registry.counts_for_test(), (0, 0, 0)); + let (mut next_rx, next) = registry + .enqueue(limiter, vec![PathBuf::from("new")]) + .expect("successor group"); + let next = next.expect("successor starts worker"); + assert!(!Arc::ptr_eq(&old, &next)); + // An older worker can resume cleanup after a successor was installed. + registry.remove_idle_group(&old); + assert_eq!(registry.counts_for_test(), (1, 1, 1)); + assert!( + registry + .inner + .lock() + .groups + .get(&next.key) + .is_some_and(|group| Arc::ptr_eq(group, &next)) + ); + assert!(matches!(next_rx.try_recv(), Err(oneshot::error::TryRecvError::Empty))); + assert!(next.inner.lock().worker_running); + assert_eq!(next.inner.lock().pending.len(), 1); + assert!(first_waiter.result_tx.send(Ok(())).is_ok()); + assert!(first_rx.await.expect("original result").is_ok()); + } + + #[tokio::test] + #[serial_test::serial(file_sync_probe)] + async fn resource_grouped_fdatasync_keeps_native_producers_and_releases_cancelled_queue() { + use crate::disk::io_admission::with_background_resource_unit; + use rustfs_concurrency::adaptive::AdaptiveAdmission; + use std::num::NonZeroU16; + + let resources: [_; 3] = std::array::from_fn(|_| AdaptiveAdmission::new(NonZeroU16::new(1).expect("one unit"), 1)); + let root = tempdir().expect("file sync fixture"); + let files: [_; 3] = std::array::from_fn(|index| root.path().join(index.to_string())); + for file in &files { + std::fs::write(file, b"durable payload").expect("write real file"); + } + let _probe = file_sync_probe::set_blocking(root.path()); + let limiter = Arc::new(Semaphore::new(1)); + let mut receivers = Vec::new(); + let mut group = None; + for index in 0..2 { + let (rx, worker) = with_background_resource_unit(resources[index].acquire().await, async { + FILE_FDATASYNC_GROUP_COMMIT.enqueue(limiter.clone(), vec![files[index].clone()]) + }) + .await + .expect("enqueue producer"); + receivers.push(rx); + if index == 0 { + group = worker; + } else { + assert!(worker.is_none(), "second producer joins the same batch"); + } + } + let group = group.expect("first producer starts worker"); + let worker = tokio::spawn(run_file_fdatasync_group_worker(group.clone())); + file_sync_probe::wait_for_active(1).await; + let (queued, new_worker) = with_background_resource_unit(resources[2].acquire().await, async { + FILE_FDATASYNC_GROUP_COMMIT.enqueue(limiter.clone(), vec![files[2].clone()]) + }) + .await + .expect("enqueue after batch is frozen"); + assert!(new_worker.is_none()); + worker.abort(); + assert!(worker.await.expect_err("cancel group owner").is_cancelled()); + let retained = resources.each_ref().map(|gate| gate.snapshot().active); + let counts = FILE_FDATASYNC_GROUP_COMMIT.counts_for_test(); + let (pending, running) = { + let state = group.inner.lock(); + (state.pending.len(), state.worker_running) + }; + // Preserve failure observations, then clean up even the old implementation. + let remaining = std::mem::take(&mut group.inner.lock().pending); + drop(remaining); + clear_file_fdatasync_group_commit_for_test(); + assert!(queued.await.is_err()); + file_sync_probe::release(); + file_sync_probe::wait_for_idle().await; + let permit = limiter.acquire().await.expect("native disk permit returned"); + drop(permit); + drop(receivers); + assert_eq!(retained, [1, 1, 0], "native owners survive, queued producer is released"); + assert_eq!(resources.each_ref().map(|gate| gate.snapshot().active), [0; 3]); + assert_eq!(counts, (0, 0, 0), "cancelled worker clears its registry budget"); + assert_eq!((pending, running), (0, false)); + assert_eq!(file_sync_probe::attempts(), 2, "undispatched third file must not be synced"); + } + + #[tokio::test] + async fn resource_native_dispatch_retains_each_unit_after_caller_cancellation() { + use crate::disk::io_admission::with_background_resource_unit; + use rustfs_concurrency::adaptive::AdaptiveAdmission; + use std::num::NonZeroU16; + + let mut retained = Vec::new(); + for queued in [false, true] { + let gate = AdaptiveAdmission::new(NonZeroU16::new(1).expect("one unit"), 1); + let unit = gate.acquire().await; + let (started_tx, started_rx) = oneshot::channel(); + let (release_tx, release_rx) = oneshot::channel(); + let (finished_tx, finished_rx) = oneshot::channel(); + let worker = tokio::spawn(with_background_resource_unit(unit, async move { + let operation = move || { + started_tx.send(()).expect("native operation entered"); + let _ = release_rx.blocking_recv(); + let _ = finished_tx.send(()); + Ok(()) + }; + if queued { + DiskBlockingQueue::new().run(operation).await + } else { + fsync_spawn_blocking(operation).await.expect("native operation joined") + } + })); + started_rx.await.expect("native operation started"); + worker.abort(); + assert!(worker.await.expect_err("caller cancelled").is_cancelled()); + retained.push(gate.snapshot().active); + // Drain native work before assertions, including in the negative control. + release_tx.send(()).expect("release native operation"); + finished_rx.await.expect("native operation finished"); + tokio::time::timeout(Duration::from_secs(1), async { + while gate.snapshot().active != 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("native completion returns the unit"); + } + assert_eq!(retained, [1, 1], "fsync and namespace dispatch must retain their own producer's unit"); + } + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn disk_blocking_queue_bounds_concurrency_and_recovers_from_panic() { use std::sync::atomic::{AtomicUsize, Ordering}; diff --git a/crates/ecstore/src/lib.rs b/crates/ecstore/src/lib.rs index ff2844d6d..2e44df108 100644 --- a/crates/ecstore/src/lib.rs +++ b/crates/ecstore/src/lib.rs @@ -55,7 +55,7 @@ mod set_disk; mod storage_api_contracts; mod store; -pub use disk::with_background_disk_io; +pub use disk::{with_background_disk_io, with_background_resource_unit}; pub use store::PoolMetaWriteGateStatus; // pub mod checksum; diff --git a/crates/ecstore/src/store/heal.rs b/crates/ecstore/src/store/heal.rs index 83d542e30..03bb192d5 100644 --- a/crates/ecstore/src/store/heal.rs +++ b/crates/ecstore/src/store/heal.rs @@ -14,7 +14,7 @@ use super::*; use crate::core::pools::{POOL_META_NAME, load_pool_meta_identity_observing}; -use crate::disk::{self, error::DiskError}; +use crate::disk::{self, current_resource_unit, error::DiskError, with_optional_resource_unit}; use crate::services::rebalance::{REBAL_META_NAME, RebalStatus}; use crate::set_disk::get_lock_acquire_timeout; use crate::storage_api_contracts::heal::HealOperations as _; @@ -301,12 +301,15 @@ impl ECStore { let opts = *opts; // Dropping a caller's cancellation/timeout waiter must not release the lifecycle // owner while a storage operation is still committing. - tokio::spawn(BUCKET_HEAL_SCOPE.scope(scope.clone(), async move { - scope.check().map_err(Error::from)?; - let result = operation(store, bucket, opts).await; - scope.check().map_err(Error::from)?; - result - })) + tokio::spawn(with_optional_resource_unit( + current_resource_unit(), + BUCKET_HEAL_SCOPE.scope(scope.clone(), async move { + scope.check().map_err(Error::from)?; + let result = operation(store, bucket, opts).await; + scope.check().map_err(Error::from)?; + result + }), + )) .await .map_err(|error| Error::other_with_context("bucket heal owner task failed", error))? } @@ -1953,6 +1956,11 @@ mod tests { #[tokio::test] #[serial_test::serial] async fn bucket_incarnation_heal_owner_survives_cancelled_waiter_and_queued_writer() { + use crate::disk::{current_resource_unit, with_background_resource_unit}; + use rustfs_concurrency::adaptive::AdaptiveAdmission; + use std::num::NonZeroU16; + + let resources = AdaptiveAdmission::new(NonZeroU16::new(1).expect("one unit"), 1); let (_root, store, shutdown) = multi_pool_heal_store().await; let bucket = format!("heal-owner-{}", Uuid::new_v4().simple()); store @@ -1965,14 +1973,16 @@ mod tests { let (drained_tx, drained_rx) = tokio::sync::oneshot::channel(); let worker_store = store.clone(); let worker_bucket = bucket.clone(); - let waiter = tokio::spawn(async move { + let waiter = tokio::spawn(with_background_resource_unit(resources.acquire().await, async move { worker_store .run_bucket_heal_at_incarnation( &worker_bucket, identity, &HealOpts::default(), move |store, bucket, opts| async move { - started_tx.send(()).expect("announce held generation"); + started_tx + .send(current_resource_unit().is_some()) + .expect("announce held generation"); release_rx.await.expect("release commit"); let result = store .heal_object_with_proof(&bucket, "history.txt", &Uuid::new_v4().to_string(), &opts) @@ -1985,10 +1995,11 @@ mod tests { }, ) .await - }); - started_rx.await.expect("operation acquired generation fence"); + })); + let inherited = started_rx.await.expect("operation acquired generation fence"); waiter.abort(); assert!(waiter.await.expect_err("waiter aborted").is_cancelled()); + let held = resources.snapshot().active; let writer_store = store.clone(); let writer_bucket = bucket.clone(); let mut writer = Box::pin(writer_store.acquire_bucket_lifecycle_write_lock(&writer_bucket)); @@ -2007,6 +2018,15 @@ mod tests { .expect("lifecycle writer"); drop(guard); shutdown.cancel(); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while resources.snapshot().active != 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("drained heal releases its resource unit"); + assert!(inherited, "owned bucket heal must inherit resource ownership"); + assert_eq!(held, 1, "cancelled waiter must not release a live heal's resource unit"); } async fn minimal_heal_store() -> ECStore {