Compare commits

...

58 Commits

Author SHA1 Message Date
overtrue ba1386a0fc test(multipart): preserve failures before fault barriers 2026-10-27 06:59:54 +08:00
overtrue 6c42c94f17 fix(shutdown): join actual TLS reload work before retirement 2026-10-27 06:39:13 +08:00
overtrue 17ce7097c4 fix(shutdown): retain HTTP child tasks through real retirement 2026-10-27 04:25:00 +08:00
overtrue 1fbda624c9 fix(shutdown): preserve maintenance and frontend join failures 2026-10-27 02:30:45 +08:00
overtrue 41befd4978 fix(heal): admit inventory probes and retain replacement work 2026-10-27 02:30:45 +08:00
overtrue 0856c5a3cb fix(startup): retain shutdown work and preserve deadline failures 2026-10-27 01:14:42 +08:00
overtrue 3b4a35a4b6 fix(heal): yield node admission while accepted native work drains 2026-10-27 01:14:42 +08:00
overtrue 057d05964b build: pin profiling sources and trust fixed ownership dependencies 2026-10-27 01:14:42 +08:00
overtrue d749a845d2 fix(heal): admit resume scans and retain replacement workers 2026-10-26 19:57:31 +08:00
overtrue 4a200357a6 fix(heal): drain channel commands before stopping MRF 2026-10-26 19:18:59 +08:00
overtrue c398c45ae5 test(heal): verify shutdown at native read boundaries 2026-10-26 18:06:44 +08:00
overtrue 7c5ec6fb66 fix(heal): join manager workers during shutdown 2026-10-26 17:22:39 +08:00
overtrue f139c19f8a test(heal): isolate process-global MRF consumer fixture 2026-10-26 15:10:12 +08:00
overtrue d49309b117 fix(heal): retain and join the MRF consumer on shutdown 2026-10-26 13:29:39 +08:00
overtrue 2188ad2574 fix(scanner): account namespace envelope decode storage 2026-10-26 13:17:13 +08:00
overtrue 4f78887475 fix(scanner): budget namespace frame read staging 2026-10-26 12:34:34 +08:00
overtrue 7b1ff0da45 fix(scanner): retain the enabled supervisor task handle 2026-10-26 12:15:33 +08:00
overtrue 94f344a7d4 fix(rpc): budget namespace scanner responses 2026-10-26 11:15:13 +08:00
overtrue ff834ac98f feat(scanner): measure cache copy allocation requests 2026-10-26 08:55:50 +08:00
overtrue 3f0e592264 perf(rpc): reuse reserved response reader storage 2026-10-26 08:05:09 +08:00
overtrue c15af8e16d fix(rpc): reserve background scanner pipe storage 2026-10-26 07:37:15 +08:00
overtrue cc05fd0e37 perf(rio): reuse reserved HTTP writer storage 2026-10-26 06:58:11 +08:00
overtrue 083e7e7350 fix(storage): retain inline and HTTP writer memory reservations 2026-10-26 06:34:37 +08:00
overtrue 6f1875e021 test(scanner): verify samples and clean concurrent collectors 2026-10-26 05:52:11 +08:00
overtrue 9a03007726 fix(ecstore): account for heal and RPC reader buffers 2026-10-26 05:52:11 +08:00
overtrue 28cad76ce3 fix(storage): retain memory charges for direct-write buffers 2026-10-26 01:20:14 +08:00
overtrue d4d071869c fix(storage): account for native io_uring read allocations 2026-10-26 01:05:40 +08:00
overtrue 88cc1a3f54 fix(storage): retain memory charges for background read results 2026-10-26 00:44:26 +08:00
overtrue d6e8973e59 fix(storage): account for retained background stream buffers 2026-10-26 00:01:37 +08:00
overtrue 6845775668 fix(ecstore): retain physical ownership during durable rename sync 2026-10-25 22:05:26 +08:00
overtrue 0d28dcaf37 fix(ecstore): retain resources through whole-file native I/O 2026-10-25 21:35:12 +08:00
overtrue fbd2bd2ca1 fix(ecstore): enforce receiver-side background admission 2026-10-25 20:56:51 +08:00
overtrue 39ac7321aa fix(ecstore): retain native io_uring read ownership 2026-10-25 18:53:04 +08:00
overtrue e161e65147 fix(ecstore): retain staged metadata write ownership 2026-10-25 17:26:39 +08:00
overtrue 9bcc86f49f fix(ecstore): retain resources during positioned reads 2026-10-25 16:57:16 +08:00
overtrue 564c4354c9 fix(storage): retain physical ownership during bitrot reads 2026-10-25 16:32:08 +08:00
overtrue e07caee90c fix(storage): retain resources through native filesystem jobs 2026-10-25 16:12:13 +08:00
overtrue 4a41482a18 chore: align physical admission with the updated node base 2026-10-25 15:42:30 +08:00
overtrue 040720aa58 chore: sync node admission with current main 2026-10-25 15:42:11 +08:00
overtrue 136296c76b chore: integrate main disk dispatch with physical admission 2026-10-25 15:35:27 +08:00
overtrue 5158b16917 fix(storage): retain physical resources through streamed reads 2026-10-25 14:57:34 +08:00
overtrue 3520821376 fix(storage): retain resources through buffered writer I/O 2026-10-25 14:32:29 +08:00
overtrue 2924e37351 fix(storage): retain resources through direct writer I/O 2026-10-25 14:04:56 +08:00
overtrue cedaa66673 fix(storage): retain resources during native format probes 2026-10-25 13:30:07 +08:00
overtrue 723e6fa58e feat(heal): admit physical resources before object repair 2026-10-25 13:04:18 +08:00
overtrue 058699c57a feat(storage): admit auxiliary physical resources without waiting 2026-10-25 12:35:48 +08:00
overtrue 1c7d0d057e fix(storage): retain only native disk resource subsets 2026-10-25 12:12:41 +08:00
overtrue 23c62f5dd7 feat(scanner): share physical topology across disk plans 2026-10-25 11:19:36 +08:00
overtrue 95f21f0267 feat(concurrency): combine preallocated resource selections 2026-10-25 11:19:36 +08:00
overtrue 49eb74f209 feat(scanner): gate background reads by physical device 2026-10-25 11:19:31 +08:00
overtrue d954886cef feat(concurrency): acquire physical groups before node work 2026-10-25 11:19:31 +08:00
overtrue 1fd619992f fix(scanner): retain resource ownership in native filesystem jobs 2026-10-25 11:19:31 +08:00
overtrue 7947978a45 feat(concurrency): retain bounded physical resource subsets 2026-10-25 11:19:31 +08:00
overtrue b2006b78e3 fix(concurrency): declare Tokio time dependency 2026-10-25 11:19:24 +08:00
overtrue 1123029eda test(scanner): attest adaptive admission in process receipts 2026-10-25 11:19:24 +08:00
overtrue 529b440f67 fix(scanner): yield cache leases for shared node admission 2026-10-25 11:19:24 +08:00
overtrue 4ed33f68a8 fix(scanner): drain candidates after listing completion 2026-10-25 11:19:24 +08:00
overtrue 28fbe69909 feat(scanner): share node admission with background heal 2026-10-25 11:19:24 +08:00
118 changed files with 24390 additions and 1614 deletions
Generated
+4 -8
View File
@@ -5236,8 +5236,7 @@ dependencies = [
[[package]]
name = "hotpath"
version = "0.29.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "62f2351902fa259b758c7303505273393f6b08172ab5bdd5739a2d2bc6d554e8"
source = "git+https://github.com/rustfs/hotpath.git?rev=4d7a2f15758e65579c7c80f3c14b8c90b7760d3b#4d7a2f15758e65579c7c80f3c14b8c90b7760d3b"
dependencies = [
"arc-swap",
"async-channel",
@@ -5271,14 +5270,12 @@ dependencies = [
[[package]]
name = "hotpath-drain"
version = "0.29.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "83dab2340134d9976d62f9f5172f9c6584c7f521b092f258e0e58e412e160b53"
source = "git+https://github.com/rustfs/hotpath.git?rev=4d7a2f15758e65579c7c80f3c14b8c90b7760d3b#4d7a2f15758e65579c7c80f3c14b8c90b7760d3b"
[[package]]
name = "hotpath-macros"
version = "0.29.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a04bbe5701b097f224a08187a78b75bb9d9b4bbaff6a279c2a50ba73edbb7aac"
source = "git+https://github.com/rustfs/hotpath.git?rev=4d7a2f15758e65579c7c80f3c14b8c90b7760d3b#4d7a2f15758e65579c7c80f3c14b8c90b7760d3b"
dependencies = [
"proc-macro2",
"quote",
@@ -10831,8 +10828,7 @@ dependencies = [
[[package]]
name = "rustfs-uring"
version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f7deda770b7981ee3bfaa202ecd18625fc794c962b8d58cd4f436952b7e66dcf"
source = "git+https://github.com/rustfs/uring.git?rev=eb8fc160e71d804427de1e905f5c226d580f71bf#eb8fc160e71d804427de1e905f5c226d580f71bf"
dependencies = [
"io-uring",
"libc",
+4 -1
View File
@@ -384,7 +384,7 @@ dav-server = "0.11.0"
# Performance Analysis and Memory Profiling
rustfs-mimalloc = { version = "0.7.2" }
# Preserve Unicode focus filters until rustfs/backlog#2302 is resolved.
hotpath = { version = "0.29.1", default-features = false }
hotpath = { version = "=0.29.1", default-features = false }
# Snapshot testing for output format regression detection
insta = { version = "1.49" }
@@ -424,6 +424,9 @@ strip = "none"
# Pin upstream fixes until compatible crates.io releases are available.
[patch.crates-io]
# Keep the ordered drain API fixed while its upstream patch is reviewed.
hotpath = { git = "https://github.com/rustfs/hotpath.git", rev = "4d7a2f15758e65579c7c80f3c14b8c90b7760d3b" }
hotpath-drain = { git = "https://github.com/rustfs/hotpath.git", rev = "4d7a2f15758e65579c7c80f3c14b8c90b7760d3b" }
# Preserve RustFS OIDC flows while migrating the old elliptic-curve dependency chain.
# Upstream: openidconnect-rs#244 and oauth2-rs#352.
openidconnect = { git = "https://github.com/ramosbugs/openidconnect-rs", rev = "2f77ee73bc4554e1cb6279adb386005429fb4e41" }
+1 -1
View File
@@ -26,7 +26,7 @@ rustfs-io-core = { workspace = true }
serde = { workspace = true, features = ["derive"] }
# Async runtime
tokio = { workspace = true, features = ["sync", "fs", "rt-multi-thread"] }
tokio = { workspace = true, features = ["sync", "fs", "rt-multi-thread", "time"] }
# Logging
tracing = { workspace = true }
+83 -14
View File
@@ -10,6 +10,10 @@ use std::time::Duration;
use tokio::sync::{Notify, Semaphore};
use tokio::time::Instant;
pub mod permit_set;
const MAX_PERMIT_REFERENCES: usize = usize::MAX / 2;
/// 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
@@ -36,6 +40,7 @@ pub struct AdaptiveAdmission {
maximum: u32,
slots: Box<[PermitSlots]>,
next_chunk: AtomicUsize,
completed: AtomicU64,
}
struct PermitSlots {
@@ -61,10 +66,12 @@ pub struct AdaptivePermit {
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.
// Subset carriers can retain slot references through a shared owner,
// so the slot count needs its own bound in addition to Arc's guard.
let owner = Arc::clone(&self.owner);
owner.slots[self.slot / 64].references[self.slot % 64].fetch_add(1, Ordering::Relaxed);
if owner.slots[self.slot / 64].references[self.slot % 64].fetch_add(1, Ordering::Relaxed) >= MAX_PERMIT_REFERENCES {
std::process::abort();
}
Self { owner, slot: self.slot }
}
}
@@ -114,11 +121,22 @@ impl NodeAdmission {
/// Publish feedback with this owner's actual completed-work counter.
pub fn update(&self, controller: &mut FeedbackController, observation: Option<FeedbackObservation>) -> FeedbackDecision {
self.update_with_probe(controller, observation, true)
}
/// Coordinate growth with other resource controllers in the same window.
/// Protective reductions never depend on permission to probe.
pub fn update_with_probe(
&self,
controller: &mut FeedbackController,
observation: Option<FeedbackObservation>,
allow_probe: bool,
) -> FeedbackDecision {
let observation = observation.map(|mut observation| {
observation.completed_units = self.completed();
observation
});
controller.update(&self.capacity, observation)
controller.update_with_probe(&self.capacity, observation, allow_probe)
}
/// Current leased scheduling cap and held async units.
@@ -249,16 +267,31 @@ impl FeedbackController {
/// 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<FeedbackObservation>) -> FeedbackDecision {
self.update_with_probe(gate, observation, true)
}
/// Allow at most the sampling owner's selected controller to probe upward.
pub fn update_with_probe(
&mut self,
gate: &AdaptiveAdmission,
observation: Option<FeedbackObservation>,
allow_probe: bool,
) -> FeedbackDecision {
let observation = observation.map(|mut observation| {
observation.max_concurrency = observation.max_concurrency.min(gate.maximum);
observation
});
let decision = self.decide(Instant::now(), observation);
let decision = self.decide_with_probe(Instant::now(), observation, allow_probe);
gate.set_sampled_limit(decision.limit, Duration::from_secs(3));
decision
}
fn decide(&mut self, now: Instant, observation: Option<FeedbackObservation>) -> FeedbackDecision {
fn decide_with_probe(
&mut self,
now: Instant,
observation: Option<FeedbackObservation>,
allow_probe: bool,
) -> FeedbackDecision {
let Some(observation) = observation else {
self.previous = None;
self.stable_since = None;
@@ -303,7 +336,8 @@ impl FeedbackController {
FeedbackReason::NewBaseline
} else {
let since = *self.stable_since.get_or_insert(now);
if advanced
if allow_probe
&& 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
@@ -325,6 +359,13 @@ impl FeedbackController {
}
impl AdaptiveAdmission {
/// Successful async operations using this physical owner. Native clones
/// retain capacity but never count another completion.
pub fn completed(&self) -> Option<u64> {
let value = self.completed.load(Ordering::Acquire);
(value != u64::MAX).then_some(value)
}
/// 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.
@@ -351,6 +392,7 @@ impl AdaptiveAdmission {
maximum: maximum_u32,
slots,
next_chunk: AtomicUsize::new(0),
completed: AtomicU64::new(0),
})
}
@@ -466,6 +508,21 @@ impl AdaptiveAdmission {
self.slots[slot / 64].free.fetch_or(1_u64 << (slot % 64), Ordering::Release);
}
fn release_reference(&self, slot: usize) {
self.release_reference_with(slot, || {});
}
fn release_reference_with(&self, slot: usize, before_release: impl FnOnce()) {
if self.slots[slot / 64].references[slot % 64].fetch_sub(1, Ordering::AcqRel) != 1 {
return;
}
before_release();
// Publish the free slot before exposing another admission capacity.
self.release_slot(slot);
self.state.fetch_sub(1, Ordering::AcqRel);
self.changed.notify_one();
}
/// 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.
@@ -493,13 +550,7 @@ impl AdaptiveAdmission {
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();
self.owner.release_reference(self.slot);
}
}
@@ -861,6 +912,24 @@ mod tests {
}
}
#[tokio::test(start_paused = true)]
async fn shared_probe_permission_blocks_growth_but_never_protective_cuts() {
let gate = AdaptiveAdmission::new(NonZeroU16::new(8).expect("bound"), 0);
let mut controller = FeedbackController::new();
controller.update_with_probe(&gate, Some(healthy(0)), false);
for completed in 1..=20 {
tokio::time::advance(Duration::from_secs(1)).await;
assert_eq!(controller.update_with_probe(&gate, Some(healthy(completed)), false).limit, 1);
}
tokio::time::advance(Duration::from_secs(1)).await;
assert_eq!(controller.update_with_probe(&gate, Some(healthy(21)), true).reason, FeedbackReason::Probe);
assert_eq!(gate.snapshot().limit, 2);
let mut busy = healthy(22);
busy.pressure = ResourcePressure::Busy;
assert_eq!(controller.update_with_probe(&gate, Some(busy), false).limit, 1);
assert_eq!(controller.update_with_probe(&gate, None, false).limit, 0);
}
#[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);
File diff suppressed because it is too large Load Diff
+102 -53
View File
@@ -24,6 +24,14 @@ use std::{
time::{Duration, SystemTime},
};
/// Copy an existing cache-model value without including surrounding work.
#[inline]
pub fn clone_cache_value<T: ToOwned + ?Sized>(source: &T) -> T::Owned {
#[cfg(feature = "hotpath-alloc")]
let _copy_scope = hotpath::MeasurementGuardSync::new("cache_model.owned_copy", false, false);
source.to_owned()
}
/// Maximum amount a persisted `last_update` may lead the local wall clock before the
/// persisted timestamp is treated as untrustworthy.
///
@@ -342,7 +350,7 @@ impl UnknownTierStats {
self.diagnostics_dropped = checked_saturating_add(1, self.diagnostics_dropped, &mut self.counter_overflowed);
continue;
}
self.diagnostics.push(diagnostic.clone());
self.diagnostics.push(clone_cache_value(diagnostic));
}
}
if !self.diagnostics.is_empty() {
@@ -367,7 +375,7 @@ impl UnknownTierStats {
if !self.fits_add(other) {
return None;
}
let mut merged = self.clone();
let mut merged = clone_cache_value(self);
merged.merge(other);
Some(merged)
}
@@ -407,7 +415,7 @@ impl AllTierStats {
if st.is_empty() {
continue;
}
let entry = self.tiers.entry(tier.clone()).or_default();
let entry = self.tiers.entry(clone_cache_value(tier)).or_default();
*entry = entry.add(st);
}
}
@@ -437,7 +445,7 @@ impl AllTierStats {
.tiers
.keys()
.filter(|tier| tier.as_str() != "STANDARD" && tier.as_str() != "REDUCED_REDUNDANCY" && !known.contains(tier.as_str()))
.cloned()
.map(clone_cache_value)
.collect();
for tier in retired_tiers {
if let Some(stats) = self.tiers.remove(&tier) {
@@ -1234,10 +1242,10 @@ impl DataUsageEntry {
}
pub fn add_child(&mut self, hash: &DataUsageHash) {
if self.children.contains(&hash.key()) {
if self.children.contains(&clone_cache_value(&hash.0)) {
return;
}
self.children.insert(hash.key());
self.children.insert(clone_cache_value(&hash.0));
}
pub fn merge(&mut self, other: &DataUsageEntry) {
@@ -1254,7 +1262,7 @@ impl DataUsageEntry {
s_rep.replica_size += o_rep.replica_size;
s_rep.replica_count += o_rep.replica_count;
for (arn, stat) in o_rep.targets.iter() {
let st = s_rep.targets.entry(arn.clone()).or_default();
let st = s_rep.targets.entry(clone_cache_value(arn)).or_default();
st.pending_size += stat.pending_size;
st.replicated_size += stat.replicated_size;
st.failed_size += stat.failed_size;
@@ -1448,7 +1456,7 @@ pub fn prefix_usage_in_cache(
} else {
format!("{bucket}/{prefix}")
};
let entry = cache.get(&hash_path(&root).key())?.clone();
let entry = clone_cache_value(cache.get(&clone_cache_value(&hash_path(&root).0))?);
let usage = PrefixUsageSummary::from_entry(&flatten_entry(cache, &entry, 0)?);
@@ -1591,38 +1599,38 @@ pub struct DataUsageCache {
impl DataUsageCache {
pub fn replace(&mut self, path: &str, parent: &str, e: DataUsageEntry) {
let hash = hash_path(path);
self.cache.insert(hash.key(), e);
self.cache.insert(clone_cache_value(&hash.0), e);
if !parent.is_empty() {
let phash = hash_path(parent);
let p = {
let p = self.cache.entry(phash.key()).or_default();
let p = self.cache.entry(clone_cache_value(&phash.0)).or_default();
p.add_child(&hash);
p.clone()
clone_cache_value(p)
};
self.cache.insert(phash.key(), p);
self.cache.insert(clone_cache_value(&phash.0), p);
}
}
pub fn replace_hashed(&mut self, hash: &DataUsageHash, parent: &Option<DataUsageHash>, e: &DataUsageEntry) {
self.cache.insert(hash.key(), e.clone());
self.cache.insert(clone_cache_value(&hash.0), clone_cache_value(e));
if let Some(parent) = parent {
self.cache.entry(parent.key()).or_default().add_child(hash);
self.cache.entry(clone_cache_value(&parent.0)).or_default().add_child(hash);
}
}
pub fn find(&self, path: &str) -> Option<DataUsageEntry> {
self.cache.get(&hash_path(path).key()).cloned()
self.cache.get(&clone_cache_value(&hash_path(path).0)).map(clone_cache_value)
}
pub fn find_children_copy(&mut self, h: DataUsageHash) -> DataUsageHashMap {
self.cache.entry(h.string()).or_default().children.clone()
clone_cache_value(&self.cache.entry(clone_cache_value(&h.0)).or_default().children)
}
pub fn flatten(&self, root: &DataUsageEntry) -> DataUsageEntry {
let mut root = root.clone();
for id in root.children.clone().iter() {
let mut root = clone_cache_value(root);
for id in clone_cache_value(&root.children).iter() {
if let Some(e) = self.cache.get(id) {
let mut e = e.clone();
let mut e = clone_cache_value(e);
if !e.children.is_empty() {
e = self.flatten(&e);
}
@@ -1634,16 +1642,16 @@ impl DataUsageCache {
}
pub fn copy_with_children(&mut self, src: &DataUsageCache, hash: &DataUsageHash, parent: &Option<DataUsageHash>) {
if let Some(e) = src.cache.get(&hash.string()) {
self.cache.insert(hash.key(), e.clone());
if let Some(e) = src.cache.get(&clone_cache_value(&hash.0)) {
self.cache.insert(clone_cache_value(&hash.0), clone_cache_value(e));
for ch in e.children.iter() {
if *ch == hash.key() {
if *ch == clone_cache_value(&hash.0) {
return;
}
self.copy_with_children(src, &DataUsageHash(ch.to_string()), &Some(hash.clone()));
self.copy_with_children(src, &DataUsageHash(clone_cache_value(ch)), &Some(clone_cache_value(hash)));
}
if let Some(parent) = parent {
let p = self.cache.entry(parent.key()).or_default();
let p = self.cache.entry(clone_cache_value(&parent.0)).or_default();
p.add_child(hash);
}
}
@@ -1651,14 +1659,14 @@ impl DataUsageCache {
pub fn delete_recursive(&mut self, hash: &DataUsageHash) {
let mut need_remove = Vec::new();
if let Some(v) = self.cache.get(&hash.string()) {
if let Some(v) = self.cache.get(&clone_cache_value(&hash.0)) {
for child in v.children.iter() {
need_remove.push(child.clone());
need_remove.push(clone_cache_value(child));
}
}
self.cache.remove(&hash.string());
self.cache.remove(&clone_cache_value(&hash.0));
need_remove.iter().for_each(|child| {
self.delete_recursive(&DataUsageHash(child.to_string()));
self.delete_recursive(&DataUsageHash(clone_cache_value(child)));
});
}
@@ -1679,7 +1687,7 @@ impl DataUsageCache {
}
pub fn search_parent(&self, hash: &DataUsageHash) -> Option<DataUsageHash> {
let want = hash.key();
let want = clone_cache_value(&hash.0);
if let Some(last_index) = want.rfind('/')
&& let Some(v) = self.find(&want[0..last_index])
&& v.children.contains(&want)
@@ -1690,7 +1698,7 @@ impl DataUsageCache {
for (k, v) in self.cache.iter() {
if v.children.contains(&want) {
let found = DataUsageHash(k.clone());
let found = DataUsageHash(clone_cache_value(k));
return Some(found);
}
}
@@ -1698,7 +1706,7 @@ impl DataUsageCache {
}
pub fn is_compacted(&self, hash: &DataUsageHash) -> bool {
match self.cache.get(&hash.key()) {
match self.cache.get(&clone_cache_value(&hash.0)) {
Some(due) => due.compacted,
None => false,
}
@@ -1723,7 +1731,7 @@ impl DataUsageCache {
if self.cache.len() < limit {
return;
}
let top = hash_path(&self.info.name).key();
let top = clone_cache_value(&hash_path(&self.info.name).0);
let top_e = match self.find(&top) {
Some(e) => e,
None => return,
@@ -1750,7 +1758,7 @@ impl DataUsageCache {
}
pub fn reduce_children_of(&mut self, path: &DataUsageHash, limit: usize, compact_self: bool) {
let e = match self.cache.get(&path.key()) {
let e = match self.cache.get(&clone_cache_value(&path.0)) {
Some(e) => e,
None => return,
};
@@ -1760,13 +1768,13 @@ impl DataUsageCache {
}
if e.children.len() > limit && compact_self {
let mut flat = self.size_recursive(&path.key()).unwrap_or_default();
let mut flat = self.size_recursive(&clone_cache_value(&path.0)).unwrap_or_default();
flat.compacted = true;
self.delete_recursive(path);
self.replace_hashed(path, &None, &flat);
return;
}
let total = self.total_children_rec(&path.key());
let total = self.total_children_rec(&clone_cache_value(&path.0));
if total < limit {
return;
}
@@ -1779,12 +1787,12 @@ impl DataUsageCache {
let mut candidate_index = 0;
while remove > 0 && candidate_index < candidates.len() {
let e = &candidates[candidate_index];
let candidate = e.path.clone();
let candidate = clone_cache_value(&e.path);
if candidate == *path && !compact_self {
break;
}
let removing = self.total_children_rec(&candidate.key());
let mut flat = match self.size_recursive(&candidate.key()) {
let removing = self.total_children_rec(&clone_cache_value(&candidate.0));
let mut flat = match self.size_recursive(&clone_cache_value(&candidate.0)) {
Some(flat) => flat,
None => {
candidate_index += 1;
@@ -1821,7 +1829,7 @@ impl DataUsageCache {
if o.root().is_none() {
return;
}
*self = o.clone();
*self = clone_cache_value(o);
return;
};
@@ -1834,7 +1842,8 @@ impl DataUsageCache {
}
existing_root.merge(&other_root);
self.cache.insert(hash_path(&self.info.name).key(), existing_root);
self.cache
.insert(clone_cache_value(&hash_path(&self.info.name).0), existing_root);
let root_hash = self.root_hash();
for key in other_root.children.iter() {
@@ -1845,7 +1854,7 @@ impl DataUsageCache {
if let Some(existing) = self.cache.get_mut(key) {
existing.merge(&flat);
} else {
self.replace_hashed(&DataUsageHash(key.clone()), &Some(root_hash.clone()), &flat);
self.replace_hashed(&DataUsageHash(clone_cache_value(key)), &Some(clone_cache_value(&root_hash)), &flat);
}
}
}
@@ -1889,7 +1898,7 @@ impl DataUsageCache {
for (arn, stat) in rs.targets.iter() {
bui.replication_info.insert(
arn.clone(),
clone_cache_value(arn),
BucketTargetUsageInfo {
replication_pending_size: stat.pending_size,
replicated_size: stat.replicated_size,
@@ -1902,7 +1911,7 @@ impl DataUsageCache {
);
}
}
buckets_usage.insert(bucket_name.clone(), bui);
buckets_usage.insert(clone_cache_value(bucket_name), bui);
}
DataUsageInfo {
@@ -1934,20 +1943,20 @@ struct Inner {
}
fn add(data_usage_cache: &DataUsageCache, path: &DataUsageHash, candidates: &mut Vec<Inner>) -> usize {
let e = match data_usage_cache.cache.get(&path.key()) {
let e = match data_usage_cache.cache.get(&clone_cache_value(&path.0)) {
Some(e) => e,
None => return 0,
};
let mut objects = e.objects;
for ch in e.children.iter() {
objects += add(data_usage_cache, &DataUsageHash(ch.clone()), candidates);
objects += add(data_usage_cache, &DataUsageHash(clone_cache_value(ch)), candidates);
}
// Collect internal nodes (with children) as compaction candidates.
// Leaf nodes have no children to remove, so compacting them is a no-op.
if !e.children.is_empty() {
candidates.push(Inner {
objects,
path: path.clone(),
path: clone_cache_value(path),
});
}
objects
@@ -1955,7 +1964,7 @@ fn add(data_usage_cache: &DataUsageCache, path: &DataUsageHash, candidates: &mut
fn mark(duc: &DataUsageCache, entry: &DataUsageEntry, found: &mut HashSet<String>) {
for k in entry.children.iter() {
found.insert(k.to_string());
found.insert(clone_cache_value(k));
if let Some(ch) = duc.cache.get(k) {
mark(duc, ch, found);
}
@@ -2402,11 +2411,11 @@ impl SizeSummary {
}
for entry in &other.size_reconciliation {
self.record_size_reconciliation(entry.clone());
self.record_size_reconciliation(clone_cache_value(entry));
}
self.size_reconciliation_truncated |= other.size_reconciliation_truncated;
for scope in &other.reconciliation_scopes {
self.record_reconciliation_scope(&scope.bucket, &scope.object);
self.record_reconciliation_scope_with(&scope.bucket, &scope.object, || clone_cache_value(scope));
}
}
@@ -2431,6 +2440,13 @@ impl SizeSummary {
/// Mark one object scope as refreshed. Duplicate scopes are suppressed so
/// merging summaries remains bounded and deterministic.
pub fn record_reconciliation_scope(&mut self, bucket: &str, object: &str) {
self.record_reconciliation_scope_with(bucket, object, || SizeReconciliationScope {
bucket: bucket.to_string(),
object: object.to_string(),
});
}
fn record_reconciliation_scope_with(&mut self, bucket: &str, object: &str, build: impl FnOnce() -> SizeReconciliationScope) {
if !self
.reconciliation_scopes
.iter()
@@ -2440,10 +2456,7 @@ impl SizeSummary {
self.size_reconciliation_truncated = true;
return;
}
self.reconciliation_scopes.push(SizeReconciliationScope {
bucket: bucket.to_string(),
object: object.to_string(),
});
self.reconciliation_scopes.push(build());
}
}
}
@@ -3373,6 +3386,42 @@ mod tests {
assert_eq!(summary1.versions, 15);
}
#[test]
fn test_reconciliation_scope_builds_only_after_admission() {
let mut summary = SizeSummary::new();
let builds = std::cell::Cell::new(0);
let build = || {
builds.set(builds.get() + 1);
SizeReconciliationScope {
bucket: "bucket".to_string(),
object: "object".to_string(),
}
};
summary.record_reconciliation_scope_with("bucket", "object", build);
assert_eq!(builds.get(), 1, "an admitted scope is copied once");
assert_eq!(summary.reconciliation_scopes.len(), 1);
assert_eq!(summary.reconciliation_scopes[0].bucket, "bucket");
assert_eq!(summary.reconciliation_scopes[0].object, "object");
summary.record_reconciliation_scope_with("bucket", "object", build);
assert_eq!(builds.get(), 1, "duplicate scopes must not be copied");
assert_eq!(summary.reconciliation_scopes.len(), 1);
assert!(!summary.size_reconciliation_truncated);
for index in 1..1024 {
summary.record_reconciliation_scope("bucket", &index.to_string());
}
summary.record_reconciliation_scope_with("bucket", "object", build);
assert_eq!(builds.get(), 1, "a duplicate at capacity must not be copied");
assert!(!summary.size_reconciliation_truncated);
summary.record_reconciliation_scope_with("bucket", "overflow", build);
assert_eq!(builds.get(), 1, "a rejected scope must not be copied");
assert_eq!(summary.reconciliation_scopes.len(), 1024);
assert!(summary.size_reconciliation_truncated);
}
#[test]
fn test_size_histogram_compat_rollup_sums_all_sub_buckets() {
let mut hist = SizeHistogram::default();
@@ -35,6 +35,7 @@ const DEADLINE: Duration = Duration::from_secs(45);
async fn fixture(
drives_per_node: usize,
storage_class: &str,
lock_diagnostics: bool,
) -> TestResult<(RustFSTestClusterEnvironment, Vec<RpcFaultProxy>, Vec<Client>, String)> {
let mut cluster =
RustFSTestClusterEnvironment::with_topology(ClusterTopology::single_pool_multidrive(4, drives_per_node)).await?;
@@ -54,6 +55,13 @@ async fn fixture(
] {
cluster.set_env(name, value);
}
if lock_diagnostics {
cluster.set_env(
"RUST_LOG",
"warn,rustfs_lock::distributed_lock=debug,rustfs_ecstore::object_lock_diag=trace,rustfs_ecstore::cluster::rpc::remote_locker=info",
);
cluster.set_env("RUSTFS_OBJECT_LOCK_DIAG_ENABLE", "true");
}
let mut proxies = Vec::new();
for node in 0..4 {
proxies.push(cluster.start_rpc_proxy_for_node(node).await?);
@@ -316,26 +324,120 @@ fn upload_path_scope(bucket: &str, key: &str, upload_id: &str) -> TestResult<Str
))
}
async fn with_paused_peers<T: Send + 'static>(
async fn with_paused_peers<T, E>(
proxies: &[RpcFaultProxy],
methods: &[&'static str],
limits: &[usize],
scope: &str,
request: impl std::future::Future<Output = T> + Send + 'static,
) -> TestResult<T> {
request: impl std::future::Future<Output = Result<T, SdkError<E>>> + Send + 'static,
) -> TestResult<Result<T, SdkError<E>>>
where
T: Send + 'static,
E: std::error::Error + ProvideErrorMetadata + Send + Sync + 'static,
{
let deadline = tokio::time::Instant::now() + DEADLINE;
let pauses: Vec<_> = proxies[1..=limits.len()]
.iter()
.zip(limits)
.map(|(proxy, limit)| proxy.pause(methods, *limit, Some(scope)))
.collect::<Result<_, _>>()?;
let task = tokio::spawn(request);
wait_for_faults(&pauses, methods).await?;
let result = tokio::time::timeout(DEADLINE, task).await??;
let mut task = tokio::spawn(request);
let mut joined = None;
let mut deadline_expired = false;
let mut barrier_result = {
let faults = wait_for_faults(&pauses, methods);
tokio::pin!(faults);
match tokio::time::timeout_at(deadline, async {
tokio::select! {
result = &mut faults => result,
result = &mut task => {
joined = Some(result);
// A ready request is not evidence that every selected fault
// was reached. Poll the original barrier once for that proof.
match futures::poll!(faults.as_mut()) {
std::task::Poll::Ready(result) => result,
std::task::Poll::Pending => Err("SDK request completed before the selected RPC fault barrier".into()),
}
}
}
})
.await
{
Ok(result) => result,
Err(error) => {
deadline_expired = true;
Err(std::io::Error::new(std::io::ErrorKind::TimedOut, error).into())
}
}
};
// Tokio's timeout polls an already-ready inner future before its timer.
// A late observation cannot certify the original absolute deadline.
deadline_expired |= tokio::time::Instant::now() >= deadline;
if deadline_expired {
barrier_result = Err(std::io::Error::new(
std::io::ErrorKind::TimedOut,
"selected RPC fault proof observed after the absolute MPU deadline",
)
.into());
}
if barrier_result.is_err() {
for pause in &pauses {
pause.discard_pending();
}
}
let mut client_task_aborted = false;
let request_result = match joined {
Some(result) => result,
None => match tokio::time::timeout_at(deadline, &mut task).await {
Ok(result) => result,
Err(error) => {
deadline_expired = true;
for pause in &pauses {
pause.discard_pending();
}
barrier_result = Err(std::io::Error::new(std::io::ErrorKind::TimedOut, error).into());
// This owns only the test's SDK request, not server/native work.
// Retire the client task without granting another operation budget.
client_task_aborted = true;
task.abort();
task.await
}
},
};
deadline_expired |= tokio::time::Instant::now() >= deadline;
if deadline_expired {
barrier_result = Err(std::io::Error::new(
std::io::ErrorKind::TimedOut,
"SDK request result observed after the absolute MPU deadline",
)
.into());
}
let hits: Vec<_> = pauses.iter().map(RpcPause::hits).collect();
for pause in &pauses {
pause.discard_pending();
}
drop(pauses);
Ok(result)
if barrier_result.is_err() {
let sdk_error = request_result.as_ref().ok().and_then(|result| result.as_ref().err());
let response = sdk_error.and_then(SdkError::raw_response);
eprintln!(
"MPU fault barrier failed: methods={methods:?}, hits={hits:?}, client_task_aborted={client_task_aborted}, deadline_expired={deadline_expired}, sdk_code={:?}, cause={:?}, request_id={:?}",
sdk_error
.and_then(|error| error.as_service_error())
.and_then(|error| error.meta().code()),
response.and_then(|response| response.headers().get("x-rustfs-error-cause")),
response.and_then(|response| response.headers().get("x-amz-request-id")),
);
}
match (barrier_result, request_result) {
// Preserve the actual SDK source instead of replacing an early service
// failure with an unrelated fault-wait timeout. The outer Err stays red.
(Err(_), Ok(Err(error))) => Err(Box::new(error)),
(Err(_), Err(error)) if !client_task_aborted => Err(Box::new(error)),
(Err(error), _) => Err(error),
(Ok(()), Ok(result)) => Ok(result),
(Ok(()), Err(error)) => Err(Box::new(error)),
}
}
fn quorum_outcome<T: std::fmt::Debug, E: std::fmt::Debug + ProvideErrorMetadata>(
@@ -583,19 +685,19 @@ async fn phase_quorum_matrix(
#[tokio::test]
async fn metadata_delay_preserves_write_quorum_and_rejects_quorum_minus_one() -> TestResult {
let (_cluster, proxies, clients, bucket) = fixture(1, "EC:2").await?;
let (_cluster, proxies, clients, bucket) = fixture(1, "EC:2", true).await?;
phase_quorum_matrix(&proxies, &clients, &bucket, false, 1).await
}
#[tokio::test]
async fn lock_delay_preserves_admission_quorum_and_reports_timeout() -> TestResult {
let (_cluster, proxies, clients, bucket) = fixture(1, "EC:2").await?;
let (_cluster, proxies, clients, bucket) = fixture(1, "EC:2", false).await?;
phase_quorum_matrix(&proxies, &clients, &bucket, true, 1).await
}
#[tokio::test]
async fn sixteen_shard_multipart_metadata_requires_twelve_votes() -> TestResult {
let (cluster, proxies, clients, bucket) = fixture(4, "EC:4").await?;
let (cluster, proxies, clients, bucket) = fixture(4, "EC:4", false).await?;
let storage = cluster_admin_ok(&cluster, http::Method::GET, "/rustfs/admin/v3/storageinfo", None).await?;
let storage: serde_json::Value = serde_json::from_str(&storage)?;
let storage: rustfs_madmin::StorageInfo = serde_json::from_value(storage["info"].clone())?;
@@ -608,7 +710,7 @@ async fn sixteen_shard_multipart_metadata_requires_twelve_votes() -> TestResult
#[tokio::test]
async fn cancelled_metadata_wait_releases_fence_for_bucket_recreation() -> TestResult {
let (_cluster, proxies, clients, bucket) = fixture(1, "EC:2").await?;
let (_cluster, proxies, clients, bucket) = fixture(1, "EC:2", false).await?;
let pauses: Vec<_> = proxies[1..]
.iter()
.map(|proxy| proxy.pause(&["WriteMetadata"], usize::MAX, Some(&bucket)))
@@ -734,7 +836,7 @@ async fn resume_at_rpc<T: Send + 'static>(
#[tokio::test]
async fn each_multipart_phase_resumes_after_metadata_and_lock_delays() -> TestResult {
let (_cluster, proxies, clients, bucket) = fixture(1, "EC:2").await?;
let (_cluster, proxies, clients, bucket) = fixture(1, "EC:2", false).await?;
for locks in [false, true] {
let key = format!("phase-{locks}");
let client = clients[0].clone();
+5 -3
View File
@@ -228,9 +228,11 @@ metrics = { workspace = true }
# The guard scripts/check_no_tokio_io_uring.sh allows an explicit
# io-uring integration; only the tokio "io-uring" runtime feature is banned.
[target.'cfg(target_os = "linux")'.dependencies]
# SharedReadBudget and the bounded read APIs are available in the 0.3.0
# registry release; keep the application on the published contract.
rustfs-uring = "0.3.0"
# Native completion owners keep physical I/O leases alive after cancellation.
rustfs-uring = { git = "https://github.com/rustfs/uring.git", rev = "eb8fc160e71d804427de1e905f5c226d580f71bf" }
[target.'cfg(target_os = "linux")'.dev-dependencies]
rustfs-uring = { git = "https://github.com/rustfs/uring.git", rev = "eb8fc160e71d804427de1e905f5c226d580f71bf", features = ["fault-injection"] }
[target.'cfg(windows)'.dependencies]
winapi-util.workspace = true
+13 -4
View File
@@ -427,6 +427,9 @@ pub mod data_usage {
}
pub mod disk {
pub use crate::cluster::rpc::background_pipe::{BackgroundPipeReader, BackgroundPipeWriter, background_pipe};
pub use crate::cluster::rpc::background_reader_stream::BackgroundReaderStream;
pub use crate::disk::background_memory::BackgroundEncodeBuffer;
pub use crate::disk::disk_store::get_object_disk_read_timeout;
pub use crate::disk::local::{ReplacementExecutionLease, ScanGuard};
#[cfg(all(feature = "test-util", not(windows)))]
@@ -442,6 +445,10 @@ pub mod disk {
ReadOptions, RenameDataResp, STORAGE_FORMAT_FILE, SnapshotLeaseToken, UpdateMetadataOpts, VolumeInfo, WalkDirOptions,
is_quota_mutation_fence_path, new_disk, validate_batch_read_version_item_count,
};
pub use crate::disk::{
BackgroundMemoryBudget, PhysicalDiskBinding, PhysicalDiskPlan, PhysicalDiskSample, PhysicalDiskSampler,
PhysicalResourcePlan, RemoteBackgroundIo, with_background_resource_unit, with_remote_background_io,
};
pub use bytes::Bytes;
pub use endpoint::Endpoint;
pub use error::DiskError;
@@ -579,6 +586,7 @@ pub mod rio {
}
pub mod rpc {
pub use crate::cluster::rpc::http_auth::with_rpc_background_request;
pub use crate::cluster::rpc::{
AuthenticatedChannel, KMS_SIGNAL_SUBSYSTEM, LocalPeerS3Client, MAX_NETWORK_PROBE_BYTES, MAX_NETWORK_PROBE_DURATION,
NetworkPeerProbeClient, NetworkPeerProbeError, NetworkPeerProbeMeasurement, NetworkPeerTarget, PEER_RESTDRY_RUN,
@@ -588,8 +596,8 @@ pub mod rpc {
TONIC_RPC_PREFIX, TonicInterceptor, build_put_file_auth_trailer, check_and_record_signed_rpc_nonce,
decode_heal_bucket_rpc_options, encode_heal_bucket_rpc_options, gen_signature_headers, gen_tonic_replay_scope_headers,
gen_tonic_signature_headers, gen_tonic_signature_interceptor, node_service_time_out_client,
node_service_time_out_client_no_auth, normalize_tonic_rpc_audience, set_tonic_canonical_body_digest,
sign_ns_scanner_capability, sign_ns_scanner_capability_with_checkpoint_coverage,
node_service_time_out_client_no_auth, normalize_tonic_rpc_audience, rpc_background_admission,
set_tonic_canonical_body_digest, sign_ns_scanner_capability, sign_ns_scanner_capability_with_checkpoint_coverage,
sign_ns_scanner_capability_with_tier_registry_generation, sign_put_file_capability, sign_tonic_rpc_response_proof,
tonic_boot_epoch_challenge, tonic_boot_epoch_response_headers, tonic_rpc_auth_failure_reason,
verify_ns_scanner_capability, verify_ns_scanner_capability_with_checkpoint_coverage,
@@ -639,8 +647,9 @@ pub mod storage {
pub use crate::store::HealWalkVersion;
pub use crate::store::{
BootstrapLocalTarget, ECStore, SCANNER_PUBLICATION_LEASE_TTL_MS, ScannerDataMovementPauseStatus, all_local_disk,
all_local_disk_path, find_local_disk_by_ref, init_local_disks, init_local_disks_with_instance_ctx, init_lock_clients,
prewarm_local_disk_id_map, prewarm_local_disk_id_map_with_instance_ctx,
all_local_disk_path, find_local_disk_by_ref, find_local_disk_by_ref_for_rpc, init_local_disks,
init_local_disks_with_instance_ctx, init_lock_clients, prewarm_local_disk_id_map,
prewarm_local_disk_id_map_with_instance_ctx,
};
pub use crate::store::{HealObjectAbsenceProof, HealObjectStorageResult};
}
@@ -0,0 +1,682 @@
// Copyright 2026 RustFS Team
// Licensed under the Apache License, Version 2.0.
use crate::disk::background_memory::BackgroundBuffer;
use std::io::{self, IoSlice};
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll, Waker, ready};
use tokio::io::{AsyncRead, AsyncWrite, DuplexStream, ReadBuf};
pub struct BackgroundPipeReader {
inner: PipeEnd,
}
pub struct BackgroundPipeWriter {
inner: PipeEnd,
}
impl BackgroundPipeWriter {
/// Capture only the payload owner, without retaining node or physical units.
pub fn memory_budget(&self) -> Option<Arc<crate::disk::BackgroundMemoryBudget>> {
match &self.inner {
PipeEnd::Foreground(_) => None,
PipeEnd::Background(state) => state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.buffer
.memory_budget(),
}
}
}
enum PipeEnd {
Foreground(DuplexStream),
Background(Arc<Mutex<PipeState>>),
}
// The pipe lock may reserve memory. The memory budget never locks a pipe.
struct PipeState {
buffer: BackgroundBuffer,
read_offset: usize,
readable: usize,
capacity: usize,
reader_closed: bool,
writer_closed: bool,
admission_failure: Option<io::ErrorKind>,
read_waker: Option<Waker>,
write_waker: Option<Waker>,
}
/// Capture only memory ownership while the request is admitted. Payload storage
/// is reserved and initialized once, before the first nonempty background write.
pub fn background_pipe(capacity: usize) -> (BackgroundPipeReader, BackgroundPipeWriter) {
let capacity = capacity.max(1);
let memory = crate::disk::current_resource_unit().and_then(|resources| resources.memory_budget());
match memory {
Some(memory) => {
let state = Arc::new(Mutex::new(PipeState {
buffer: BackgroundBuffer::new(Some(memory)),
read_offset: 0,
readable: 0,
capacity,
reader_closed: false,
writer_closed: false,
admission_failure: None,
read_waker: None,
write_waker: None,
}));
(
BackgroundPipeReader {
inner: PipeEnd::Background(state.clone()),
},
BackgroundPipeWriter {
inner: PipeEnd::Background(state),
},
)
}
None => {
let (reader, writer) = tokio::io::duplex(capacity);
(
BackgroundPipeReader {
inner: PipeEnd::Foreground(reader),
},
BackgroundPipeWriter {
inner: PipeEnd::Foreground(writer),
},
)
}
}
}
fn advance_ring(offset: usize, count: usize, capacity: usize) -> usize {
let to_end = capacity - offset;
if count < to_end { offset + count } else { count - to_end }
}
fn poll_background_read(state: &Mutex<PipeState>, cx: &mut Context<'_>, output: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
if output.remaining() == 0 {
return Poll::Ready(Ok(()));
}
let mut state = state.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
if state.readable == 0 {
if let Some(kind) = state.admission_failure.take() {
drop(state);
return Poll::Ready(Err(io::Error::new(kind, "background pipe storage admission failed")));
}
if state.writer_closed {
return Poll::Ready(Ok(()));
}
let replaced = state.read_waker.replace(cx.waker().clone());
drop(state);
drop(replaced);
return Poll::Pending;
}
let length = state.readable.min(output.remaining());
let first = length.min(state.capacity - state.read_offset);
output.put_slice(&state.buffer.as_slice()[state.read_offset..state.read_offset + first]);
output.put_slice(&state.buffer.as_slice()[..length - first]);
state.read_offset = advance_ring(state.read_offset, length, state.capacity);
state.readable -= length;
let wake = state.write_waker.take();
drop(state);
if let Some(wake) = wake {
wake.wake();
}
Poll::Ready(Ok(()))
}
fn poll_background_write(state: &Mutex<PipeState>, cx: &mut Context<'_>, input: &[IoSlice<'_>]) -> Poll<io::Result<usize>> {
let mut state = state.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
if state.reader_closed || state.writer_closed {
return Poll::Ready(Err(io::ErrorKind::BrokenPipe.into()));
}
if input.iter().all(|part| part.is_empty()) {
return Poll::Ready(Ok(0));
}
let available = state.capacity - state.readable;
if available == 0 {
let replaced = state.write_waker.replace(cx.waker().clone());
drop(state);
drop(replaced);
return Poll::Pending;
}
let capacity = state.capacity;
if let Err(error) = state.buffer.ensure_len(capacity) {
// Legacy walk completion may suppress ordinary producer errors. A new
// pipe admission failure must also reach the reader instead of EOF.
state.admission_failure = Some(error.kind());
state.writer_closed = true;
let wake = state.read_waker.take();
drop(state);
if let Some(wake) = wake {
wake.wake();
}
return Poll::Ready(Err(error));
}
let mut offset = advance_ring(state.read_offset, state.readable, capacity);
let mut written = 0;
for part in input {
let length = part.len().min(available - written);
let first = length.min(capacity - offset);
state.buffer.as_mut_slice()[offset..offset + first].copy_from_slice(&part[..first]);
state.buffer.as_mut_slice()[..length - first].copy_from_slice(&part[first..length]);
offset = advance_ring(offset, length, capacity);
written += length;
if written == available {
break;
}
}
state.readable += written;
let wake = state.read_waker.take();
drop(state);
if let Some(wake) = wake {
wake.wake();
}
Poll::Ready(Ok(written))
}
impl AsyncRead for BackgroundPipeReader {
fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, output: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
match &mut self.get_mut().inner {
PipeEnd::Foreground(reader) => Pin::new(reader).poll_read(cx, output),
PipeEnd::Background(state) => {
let coop = ready!(tokio::task::coop::poll_proceed(cx));
let result = poll_background_read(state, cx, output);
if result.is_ready() {
coop.made_progress();
}
result
}
}
}
}
impl AsyncWrite for BackgroundPipeWriter {
fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, input: &[u8]) -> Poll<io::Result<usize>> {
match &mut self.get_mut().inner {
PipeEnd::Foreground(writer) => Pin::new(writer).poll_write(cx, input),
PipeEnd::Background(state) => {
let coop = ready!(tokio::task::coop::poll_proceed(cx));
let result = poll_background_write(state, cx, &[IoSlice::new(input)]);
if result.is_ready() {
coop.made_progress();
}
result
}
}
}
fn poll_write_vectored(self: Pin<&mut Self>, cx: &mut Context<'_>, input: &[IoSlice<'_>]) -> Poll<io::Result<usize>> {
match &mut self.get_mut().inner {
PipeEnd::Foreground(writer) => Pin::new(writer).poll_write_vectored(cx, input),
PipeEnd::Background(state) => {
let coop = ready!(tokio::task::coop::poll_proceed(cx));
let result = poll_background_write(state, cx, input);
if result.is_ready() {
coop.made_progress();
}
result
}
}
}
fn is_write_vectored(&self) -> bool {
match &self.inner {
PipeEnd::Foreground(writer) => writer.is_write_vectored(),
PipeEnd::Background(_) => true,
}
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
match &mut self.get_mut().inner {
PipeEnd::Foreground(writer) => Pin::new(writer).poll_flush(cx),
PipeEnd::Background(_) => Poll::Ready(Ok(())),
}
}
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
match &mut self.get_mut().inner {
PipeEnd::Foreground(writer) => Pin::new(writer).poll_shutdown(cx),
PipeEnd::Background(state) => {
let mut state = state.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
state.writer_closed = true;
let wake = state.read_waker.take();
drop(state);
if let Some(wake) = wake {
wake.wake();
}
Poll::Ready(Ok(()))
}
}
}
}
impl Drop for BackgroundPipeReader {
fn drop(&mut self) {
if let PipeEnd::Background(state) = &self.inner {
let mut state = state.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
state.reader_closed = true;
let wake = state.write_waker.take();
let cancelled = state.read_waker.take();
drop(state);
drop(cancelled);
if let Some(wake) = wake {
wake.wake();
}
}
}
}
impl Drop for BackgroundPipeWriter {
fn drop(&mut self) {
if let PipeEnd::Background(state) = &self.inner {
let mut state = state.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
state.writer_closed = true;
let wake = state.read_waker.take();
let cancelled = state.write_waker.take();
drop(state);
drop(cancelled);
if let Some(wake) = wake {
wake.wake();
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::disk::BackgroundMemoryBudget;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
pub(super) fn budget(capacity: usize) -> Arc<BackgroundMemoryBudget> {
let budget = BackgroundMemoryBudget::new();
budget.set_sampled_limit(u64::try_from(capacity).expect("test capacity fits u64"), Duration::from_secs(30));
budget
}
pub(super) fn local_pipe(
memory: &Arc<BackgroundMemoryBudget>,
capacity: usize,
) -> (BackgroundPipeReader, BackgroundPipeWriter) {
let state = Arc::new(Mutex::new(PipeState {
buffer: BackgroundBuffer::new(Some(memory.clone())),
read_offset: 0,
readable: 0,
capacity,
reader_closed: false,
writer_closed: false,
admission_failure: None,
read_waker: None,
write_waker: None,
}));
(
BackgroundPipeReader {
inner: PipeEnd::Background(state.clone()),
},
BackgroundPipeWriter {
inner: PipeEnd::Background(state),
},
)
}
pub(super) fn write_now(writer: &mut (impl AsyncWrite + Unpin), bytes: &[u8]) -> Poll<io::Result<usize>> {
Pin::new(writer).poll_write(&mut Context::from_waker(Waker::noop()), bytes)
}
pub(super) fn read_now(reader: &mut (impl AsyncRead + Unpin), bytes: &mut [u8]) -> Poll<io::Result<usize>> {
let mut output = ReadBuf::new(bytes);
match Pin::new(reader).poll_read(&mut Context::from_waker(Waker::noop()), &mut output) {
Poll::Ready(Ok(())) => Poll::Ready(Ok(output.filled().len())),
Poll::Ready(Err(error)) => Poll::Ready(Err(error)),
Poll::Pending => Poll::Pending,
}
}
#[test]
fn fixed_ring_allocates_lazily_once_and_keeps_last_owner_charged() {
for capacity in [1, 7, 64 * 1024, 1024 * 1024] {
let memory = budget(capacity);
let (mut reader, mut writer) = local_pipe(&memory, capacity);
assert_eq!(memory.snapshot().used, 0);
assert!(matches!(write_now(&mut writer, &[]), Poll::Ready(Ok(0))));
assert_eq!(memory.snapshot().used, 0);
assert!(matches!(write_now(&mut writer, b"x"), Poll::Ready(Ok(1))));
assert_eq!(memory.snapshot().used, capacity as u64);
memory.set_sampled_limit(0, Duration::ZERO);
let mut output = [0];
assert!(matches!(read_now(&mut reader, &mut output), Poll::Ready(Ok(1))));
assert_eq!(output, *b"x");
assert!(matches!(write_now(&mut writer, b"y"), Poll::Ready(Ok(1))));
drop(writer);
assert_eq!(memory.snapshot().used, capacity as u64, "reader still owns the real ring");
assert!(matches!(read_now(&mut reader, &mut output), Poll::Ready(Ok(1))));
assert_eq!(output, *b"y");
assert!(matches!(read_now(&mut reader, &mut output), Poll::Ready(Ok(0))));
assert_eq!(memory.snapshot().used, capacity as u64, "draining does not destroy allocation ownership");
drop(reader);
assert_eq!(memory.snapshot().used, 0);
}
}
#[test]
fn admission_error_reaches_writer_and_reader_once_without_payload() {
for sample in [None, Some((6, Duration::from_secs(30))), Some((7, Duration::ZERO))] {
let memory = BackgroundMemoryBudget::new();
if let Some((limit, validity)) = sample {
memory.set_sampled_limit(limit, validity);
}
let (mut reader, mut writer) = local_pipe(&memory, 7);
let error = match write_now(&mut writer, b"x") {
Poll::Ready(Err(e)) => e,
_ => panic!("reject before copying first payload"),
};
assert!(
crate::disk::background_memory::is_memory_pressure(&error),
"writer retains the original budget marker"
);
assert_eq!(memory.snapshot().used, 0);
let error = match read_now(&mut reader, &mut [0]) {
Poll::Ready(Err(e)) => e,
_ => panic!("reader must see a pipe admission failure"),
};
assert_eq!(error.kind(), io::ErrorKind::WouldBlock);
assert!(matches!(read_now(&mut reader, &mut [0]), Poll::Ready(Ok(0))), "failure is delivered once");
assert!(matches!(write_now(&mut writer,b"x"),Poll::Ready(Err(e)) if e.kind()==io::ErrorKind::BrokenPipe));
drop((reader, writer));
assert_eq!(memory.snapshot().used, 0);
}
}
#[test]
fn allocation_overflow_is_a_reader_failure_without_a_leaked_reservation() {
let memory = BackgroundMemoryBudget::new();
memory.set_sampled_limit(u64::MAX, Duration::from_secs(30));
let (mut reader, mut writer) = local_pipe(&memory, usize::MAX);
assert!(matches!(write_now(&mut writer,b"x"),Poll::Ready(Err(e)) if e.kind()==io::ErrorKind::OutOfMemory));
assert_eq!(memory.snapshot().used, 0);
assert!(matches!(read_now(&mut reader,&mut[0]),Poll::Ready(Err(e)) if e.kind()==io::ErrorKind::OutOfMemory));
}
#[test]
fn scalar_and_vectored_wrap_match_a_reference_queue() {
use std::collections::VecDeque;
for capacity in [1, 3, 7, 31] {
let memory = budget(capacity);
let (mut reader, mut writer) = local_pipe(&memory, capacity);
let mut expected = VecDeque::new();
let mut next = 0u8;
for turn in 0..300 {
let input: Vec<_> = (0..(turn % 11 + 1))
.map(|_| {
let byte = next;
next = next.wrapping_add(1);
byte
})
.collect();
let available = capacity - expected.len();
let written = if turn % 2 == 0 {
write_now(&mut writer, &input)
} else {
let split = input.len() / 2;
Pin::new(&mut writer).poll_write_vectored(
&mut Context::from_waker(Waker::noop()),
&[
io::IoSlice::new(&[]),
io::IoSlice::new(&input[..split]),
io::IoSlice::new(&input[split..]),
],
)
};
if available == 0 {
assert!(written.is_pending());
} else {
let accepted = input.len().min(available);
assert!(matches!(written,Poll::Ready(Ok(n)) if n==accepted));
expected.extend(&input[..accepted]);
}
let mut output = vec![0; turn % 5 + 1];
let count = output.len().min(expected.len());
assert!(matches!(read_now(&mut reader,&mut output),Poll::Ready(Ok(n)) if n==count));
for byte in &output[..count] {
assert_eq!(Some(*byte), expected.pop_front());
}
assert_eq!(memory.snapshot().used, capacity as u64);
}
drop(writer);
let mut output = vec![0; capacity];
let count = expected.len();
assert!(matches!(read_now(&mut reader,&mut output),Poll::Ready(Ok(n)) if n==count));
assert_eq!(&output[..count], expected.into_iter().collect::<Vec<_>>().as_slice());
assert!(matches!(read_now(&mut reader, &mut output), Poll::Ready(Ok(0))));
drop(reader);
assert_eq!(memory.snapshot().used, 0);
}
}
#[tokio::test]
async fn reader_drop_wakes_blocked_writer_and_writer_drop_wakes_empty_reader() {
let memory = budget(7);
let (reader, mut writer) = local_pipe(&memory, 7);
writer.write_all(b"1234567").await.expect("fill the ring");
let task = tokio::spawn(async move { writer.write_all(b"blocked").await });
tokio::task::yield_now().await;
assert!(!task.is_finished());
drop(reader);
let error = tokio::time::timeout(Duration::from_secs(1), task)
.await
.expect("reader drop wakes writer")
.expect("writer task joins")
.expect_err("closed reader is BrokenPipe");
assert_eq!(error.kind(), io::ErrorKind::BrokenPipe);
assert_eq!(memory.snapshot().used, 0);
let (mut reader, writer) = local_pipe(&memory, 7);
let task = tokio::spawn(async move { reader.read(&mut [0]).await });
tokio::task::yield_now().await;
assert!(!task.is_finished());
drop(writer);
assert_eq!(
tokio::time::timeout(Duration::from_secs(1), task)
.await
.expect("writer drop wakes reader")
.expect("reader task joins")
.expect("clean EOF"),
0
);
}
#[tokio::test]
async fn read_wakes_a_full_writer_without_changing_byte_limit() {
let memory = budget(7);
let (mut reader, mut writer) = local_pipe(&memory, 7);
writer.write_all(b"1234567").await.expect("fill initial byte limit");
let task = tokio::spawn(async move {
writer.write_all(b"x").await.expect("resume after one byte is read");
writer
});
tokio::task::yield_now().await;
assert!(!task.is_finished());
let mut first = [0];
reader.read_exact(&mut first).await.expect("read one byte");
let writer = tokio::time::timeout(Duration::from_secs(1), task)
.await
.expect("read wakes writer")
.expect("writer joins");
assert_eq!(&first, b"1");
drop(writer);
let mut all = Vec::new();
reader.read_to_end(&mut all).await.expect("drain wrapped bytes");
assert_eq!(all, b"234567x");
drop(reader);
assert_eq!(memory.snapshot().used, 0);
}
#[tokio::test]
async fn shutdown_drains_then_eof_and_cancelled_writes_accept_nothing() {
let memory = budget(7);
let (mut reader, mut writer) = local_pipe(&memory, 7);
writer
.write_all(b"1234567")
.await
.expect("fill before cancelling a blocked write");
{
let blocked = writer.write(b"lost");
tokio::pin!(blocked);
assert!(futures::poll!(&mut blocked).is_pending());
}
writer.shutdown().await.expect("shutdown does not discard buffered bytes");
assert!(matches!(write_now(&mut writer,b"x"),Poll::Ready(Err(e)) if e.kind()==io::ErrorKind::BrokenPipe));
let mut out = Vec::new();
reader.read_to_end(&mut out).await.expect("drain after shutdown");
assert_eq!(out, b"1234567");
drop(reader);
assert_eq!(memory.snapshot().used, 7, "writer still holds the allocation");
drop(writer);
assert_eq!(memory.snapshot().used, 0);
}
#[test]
fn foreground_keeps_native_duplex_semantics() {
let (reader, mut writer) = background_pipe(7);
assert!(matches!(write_now(&mut writer, b"12345678"), Poll::Ready(Ok(7))));
assert!(write_now(&mut writer, b"x").is_pending());
drop(reader);
assert!(matches!(write_now(&mut writer,b"x"),Poll::Ready(Err(e)) if e.kind()==io::ErrorKind::BrokenPipe));
let (mut reader, writer) = background_pipe(7);
drop(writer);
assert!(matches!(read_now(&mut reader, &mut [0]), Poll::Ready(Ok(0))));
}
#[test]
fn cancelled_reader_leaves_the_real_ring_charged_until_writer_drop() {
let memory = budget(7);
let (reader, mut writer) = local_pipe(&memory, 7);
assert!(matches!(write_now(&mut writer, b"x"), Poll::Ready(Ok(1))));
drop(reader);
assert_eq!(memory.snapshot().used, 7, "the surviving writer still owns the ring allocation");
assert!(matches!(write_now(&mut writer,b"y"),Poll::Ready(Err(e)) if e.kind()==io::ErrorKind::BrokenPipe));
drop(writer);
assert_eq!(memory.snapshot().used, 0);
}
#[test]
fn runtime_shutdown_releases_queued_writer_but_not_the_surviving_reader() {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("build isolated runtime");
let memory = budget(7);
let (reader, mut writer) = local_pipe(&memory, 7);
runtime.block_on(async {
writer.write_all(b"x").await.expect("allocate ring before runtime shutdown");
tokio::spawn(async move {
let _writer = writer;
std::future::pending::<()>().await;
});
tokio::task::yield_now().await;
});
drop(runtime);
assert_eq!(memory.snapshot().used, 7, "reader holds the actual allocation after runtime cancellation");
drop(reader);
assert_eq!(memory.snapshot().used, 0);
}
#[tokio::test]
async fn ready_operations_keep_tokio_cooperative_yields() {
let memory = budget(7);
let (_reader, mut writer) = local_pipe(&memory, 7);
let observed = Arc::new(AtomicUsize::new(0));
let flag = observed.clone();
let task = tokio::spawn(async move {
flag.store(1, Ordering::SeqCst);
});
for _ in 0..1000 {
assert_eq!(writer.write(&[]).await.expect("empty operation remains successful"), 0);
}
assert_eq!(
observed.load(Ordering::SeqCst),
1,
"ready pipe operations must not starve other ready tasks"
);
task.await.expect("observer task joins");
}
}
#[cfg(test)]
mod waker_tests {
use super::*;
use std::sync::{
Weak,
atomic::{AtomicUsize, Ordering},
};
use std::task::Wake;
struct LockCheck {
state: Weak<Mutex<PipeState>>,
wakes: Arc<AtomicUsize>,
drops: Arc<AtomicUsize>,
}
impl LockCheck {
fn unlocked(&self) {
if let Some(state) = self.state.upgrade() {
assert!(state.try_lock().is_ok(), "waker code must run outside the pipe lock");
}
}
}
impl Wake for LockCheck {
fn wake(self: Arc<Self>) {
self.unlocked();
self.wakes.fetch_add(1, Ordering::SeqCst);
}
}
impl Drop for LockCheck {
fn drop(&mut self) {
self.unlocked();
self.drops.fetch_add(1, Ordering::SeqCst);
}
}
#[test]
fn replaced_and_woken_reader_and_writer_wakers_run_outside_the_lock() {
for waiting_reader in [true, false] {
let memory = super::tests::budget(7);
let (mut reader, mut writer) = super::tests::local_pipe(&memory, 7);
let PipeEnd::Background(state) = &reader.inner else { panic!("background fixture") };
let weak = Arc::downgrade(state);
let wakes = Arc::new(AtomicUsize::new(0));
let drops = Arc::new(AtomicUsize::new(0));
if !waiting_reader {
assert!(matches!(super::tests::write_now(&mut writer, b"1234567"), Poll::Ready(Ok(7))));
}
for _ in 0..2 {
let waker = Waker::from(Arc::new(LockCheck {
state: weak.clone(),
wakes: wakes.clone(),
drops: drops.clone(),
}));
let mut cx = Context::from_waker(&waker);
if waiting_reader {
let mut bytes = [0];
assert!(
Pin::new(&mut reader)
.poll_read(&mut cx, &mut ReadBuf::new(&mut bytes))
.is_pending()
);
} else {
assert!(Pin::new(&mut writer).poll_write(&mut cx, b"x").is_pending());
}
drop(waker);
}
assert_eq!(drops.load(Ordering::SeqCst), 1, "the replaced waker has been dropped");
if waiting_reader {
assert!(matches!(super::tests::write_now(&mut writer, b"x"), Poll::Ready(Ok(1))));
} else {
assert!(matches!(super::tests::read_now(&mut reader, &mut [0]), Poll::Ready(Ok(1))));
}
assert_eq!(wakes.load(Ordering::SeqCst), 1);
assert_eq!(drops.load(Ordering::SeqCst), 2);
}
}
}
@@ -0,0 +1,586 @@
// Copyright 2026 RustFS Team
// Licensed under the Apache License, Version 2.0.
use crate::disk::background_memory::{BackgroundMemoryBudget, Reservation};
use bytes::{Bytes, BytesMut};
use futures::Stream;
use std::io;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, ReadBuf};
use tokio_util::io::ReaderStream;
struct ReservedReadBytes {
// Free each payload view before releasing its shared allocation receipt.
bytes: Bytes,
_reservation: Arc<Reservation>,
}
impl AsRef<[u8]> for ReservedReadBytes {
fn as_ref(&self) -> &[u8] {
&self.bytes
}
}
struct ReaderBuffer {
// These initialized bytes are spare storage, not unconsumed source data.
bytes: BytesMut,
reservation: Option<Arc<Reservation>>,
budget: Arc<BackgroundMemoryBudget>,
capacity: usize,
}
impl ReaderBuffer {
fn new(budget: Arc<BackgroundMemoryBudget>, capacity: usize) -> Self {
Self {
bytes: BytesMut::new(),
reservation: None,
budget,
capacity: capacity.max(1),
}
}
fn clear(&mut self) {
self.bytes = BytesMut::new();
self.reservation = None;
}
fn ensure_readable(&mut self) -> io::Result<()> {
if !self.bytes.is_empty() {
return Ok(());
}
if self.reservation.is_some() {
if self.bytes.try_reclaim(self.capacity) {
self.bytes.resize(self.capacity, 0);
return Ok(());
}
// A live output owns this allocation. Drop our empty view before
// its receipt, then reserve independent replacement storage.
self.clear();
}
let reservation = self.budget.reserve(self.capacity)?;
let mut bytes = Vec::new();
bytes
.try_reserve_exact(self.capacity)
.map_err(|_| io::ErrorKind::OutOfMemory)?;
bytes.resize(self.capacity, 0);
// Fresh ordinary Bytes is unique and spans the complete allocation.
// Unlike from_owner output, this conversion cannot require a copy.
let bytes = Bytes::from(bytes.into_boxed_slice())
.try_into_mut()
.map_err(|_| io::Error::other("fresh reader storage was not unique"))?;
self.bytes = bytes;
self.reservation = Some(Arc::new(reservation));
Ok(())
}
fn take_bytes(&mut self, length: usize) -> io::Result<Bytes> {
let reservation = self
.reservation
.as_ref()
.cloned()
.ok_or_else(|| io::Error::other("reader storage has no memory reservation"))?;
Ok(Bytes::from_owner(ReservedReadBytes {
bytes: self.bytes.split_to(length).freeze(),
_reservation: reservation,
}))
}
}
/// RPC response chunks retain their background memory reservation until the
/// last clone or slice is dropped. Foreground readers use the existing stream.
pub struct BackgroundReaderStream<R> {
inner: ReaderState<R>,
}
enum ReaderState<R> {
Foreground(ReaderStream<R>),
Background { reader: Option<R>, buffer: ReaderBuffer },
}
impl<R: AsyncRead> BackgroundReaderStream<R> {
/// Capture memory ownership inside the admitted scope, before moving the
/// response body out of it. This stream retains no node or device capacity.
pub fn with_capacity(reader: R, capacity: usize) -> Self {
let memory = crate::disk::current_resource_unit().and_then(|resources| resources.memory_budget());
let inner = match memory {
Some(memory) => ReaderState::Background {
reader: Some(reader),
buffer: ReaderBuffer::new(memory, capacity),
},
None => ReaderState::Foreground(ReaderStream::with_capacity(reader, capacity)),
};
Self { inner }
}
}
impl<R: AsyncRead + Unpin> Stream for BackgroundReaderStream<R> {
type Item = io::Result<Bytes>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
match &mut this.inner {
ReaderState::Foreground(reader) => Pin::new(reader).poll_next(cx),
ReaderState::Background { reader, buffer } => {
let Some(inner) = reader.as_mut() else {
return Poll::Ready(None);
};
if let Err(error) = buffer.ensure_readable() {
*reader = None;
return Poll::Ready(Some(Err(error)));
}
let mut output = ReadBuf::new(&mut buffer.bytes);
match std::task::ready!(Pin::new(inner).poll_read(cx, &mut output)) {
Ok(()) => {
let length = output.filled().len();
if length == 0 {
buffer.clear();
*reader = None;
Poll::Ready(None)
} else {
let bytes = buffer.take_bytes(length);
if bytes.is_err() {
buffer.clear();
*reader = None;
}
Poll::Ready(Some(bytes))
}
}
Err(error) => {
buffer.clear();
*reader = None;
Poll::Ready(Some(Err(error)))
}
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::disk::{BackgroundMemoryBudget, test_resource_units, with_optional_resource_unit};
use futures::StreamExt;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
use tokio::io::AsyncWriteExt;
struct ObservedReader {
polls: Arc<AtomicUsize>,
fail: bool,
}
impl AsyncRead for ObservedReader {
fn poll_read(self: Pin<&mut Self>, _cx: &mut Context<'_>, output: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
self.polls.fetch_add(1, Ordering::SeqCst);
if self.fail {
Poll::Ready(Err(io::Error::from(io::ErrorKind::PermissionDenied)))
} else {
output.put_slice(b"data");
Poll::Ready(Ok(()))
}
}
}
#[tokio::test]
async fn background_frames_keep_capacity_until_the_last_slice() {
let memory = BackgroundMemoryBudget::new();
memory.set_sampled_limit(16, Duration::from_secs(30));
let (resources, gates) = test_resource_units(2);
let mut stream = with_optional_resource_unit(Some(resources.with_memory_budget(memory.clone())), async {
BackgroundReaderStream::with_capacity(io::Cursor::new(b"data"), 16)
})
.await;
assert!(
gates.iter().all(|gate| gate.snapshot().active == 0),
"the stream captures only memory ownership"
);
assert_eq!(memory.snapshot().used, 0, "construction does not allocate payload storage");
let bytes = stream.next().await.expect("first frame").expect("read frame");
assert_eq!(bytes.as_ref(), b"data");
assert_eq!(memory.snapshot().used, 16, "charge full capacity rather than visible length");
let slice = bytes.slice(1..3);
memory.set_sampled_limit(0, Duration::ZERO);
drop(bytes);
drop(stream);
assert_eq!(slice.as_ref(), b"at");
assert_eq!(memory.snapshot().used, 16);
drop(slice);
assert_eq!(memory.snapshot().used, 0);
}
#[tokio::test]
async fn background_denial_happens_before_reading_and_never_falls_back() {
for sample in [None, Some((7, Duration::from_secs(30))), Some((8, Duration::ZERO))] {
let memory = BackgroundMemoryBudget::new();
if let Some((limit, validity)) = sample {
memory.set_sampled_limit(limit, validity);
}
let polls = Arc::new(AtomicUsize::new(0));
let (resources, _) = test_resource_units(1);
let mut stream = with_optional_resource_unit(Some(resources.with_memory_budget(memory.clone())), async {
BackgroundReaderStream::with_capacity(
ObservedReader {
polls: polls.clone(),
fail: false,
},
8,
)
})
.await;
let error = stream
.next()
.await
.expect("allocation failure")
.expect_err("background allocation denied");
assert_eq!(error.kind(), io::ErrorKind::WouldBlock);
assert_eq!(polls.load(Ordering::SeqCst), 0, "denial cannot consume source bytes");
assert!(stream.next().await.is_none());
assert_eq!(memory.snapshot().used, 0);
}
}
#[tokio::test]
async fn background_pending_read_keeps_storage_through_shrink_or_cancellation() {
for cancelled in [false, true] {
let memory = BackgroundMemoryBudget::new();
memory.set_sampled_limit(16, Duration::from_secs(30));
let (resources, _) = test_resource_units(1);
let (mut sender, reader) = tokio::io::duplex(4);
let mut stream = with_optional_resource_unit(Some(resources.with_memory_budget(memory.clone())), async {
BackgroundReaderStream::with_capacity(reader, 16)
})
.await;
assert!(futures::poll!(stream.next()).is_pending());
memory.set_sampled_limit(0, Duration::ZERO);
assert_eq!(memory.snapshot().used, 16);
let bytes = if cancelled {
None
} else {
sender.write_all(b"data").await.expect("release pending read");
Some(
stream
.next()
.await
.expect("pending frame")
.expect("existing allocation survives shrink"),
)
};
drop(stream);
if let Some(bytes) = bytes {
assert_eq!(bytes.as_ref(), b"data");
assert_eq!(memory.snapshot().used, 16);
drop(bytes);
}
assert_eq!(memory.snapshot().used, 0);
}
}
#[tokio::test]
async fn background_read_error_and_eof_release_storage() {
for fail in [false, true] {
let memory = BackgroundMemoryBudget::new();
memory.set_sampled_limit(8, Duration::from_secs(30));
let (resources, _) = test_resource_units(1);
let mut stream = with_optional_resource_unit(Some(resources.with_memory_budget(memory.clone())), async {
let reader: Box<dyn AsyncRead + Unpin> = if fail {
Box::new(ObservedReader {
polls: Arc::new(AtomicUsize::new(0)),
fail: true,
})
} else {
Box::new(tokio::io::empty())
};
BackgroundReaderStream::with_capacity(reader, 8)
})
.await;
if fail {
let error = stream.next().await.expect("read failure").expect_err("preserve reader error");
assert_eq!(error.kind(), io::ErrorKind::PermissionDenied);
}
assert!(stream.next().await.is_none());
assert_eq!(memory.snapshot().used, 0);
}
}
#[tokio::test]
async fn retained_frame_blocks_new_storage_until_released() {
let memory = BackgroundMemoryBudget::new();
memory.set_sampled_limit(8, Duration::from_secs(30));
let (resources, _) = test_resource_units(1);
let mut stream = with_optional_resource_unit(Some(resources.with_memory_budget(memory.clone())), async {
BackgroundReaderStream::with_capacity(io::Cursor::new(b"firstonefollowup"), 8)
})
.await;
let first = stream.next().await.expect("first frame").expect("fits exact budget");
assert_eq!(first.as_ref(), b"firstone");
let error = stream
.next()
.await
.expect("next allocation")
.expect_err("held frame exhausts budget");
assert_eq!(error.kind(), io::ErrorKind::WouldBlock);
assert_eq!(memory.snapshot().used, 8);
drop(stream);
drop(first);
assert_eq!(memory.snapshot().used, 0);
}
#[tokio::test]
async fn consumed_frames_keep_reusable_capacity_until_eof() {
let memory = BackgroundMemoryBudget::new();
memory.set_sampled_limit(8, Duration::from_secs(30));
let (resources, _) = test_resource_units(1);
let mut stream = with_optional_resource_unit(Some(resources.with_memory_budget(memory.clone())), async {
BackgroundReaderStream::with_capacity(io::Cursor::new(b"firstonefollowup"), 8)
})
.await;
for expected in [b"firstone", b"followup"] {
let bytes = stream.next().await.expect("next frame").expect("released previous capacity");
assert_eq!(bytes.as_ref(), expected);
assert_eq!(memory.snapshot().used, 8);
drop(bytes);
assert_eq!(memory.snapshot().used, 8, "the reusable allocation remains owned by the stream");
}
assert!(stream.next().await.is_none());
assert_eq!(memory.snapshot().used, 0);
}
#[tokio::test]
async fn foreground_stream_and_unpolled_cancellation_preserve_existing_behavior() {
let mut stream = BackgroundReaderStream::with_capacity(io::Cursor::new(b"foreground"), 4);
assert!(matches!(stream.inner, ReaderState::Foreground(_)));
let mut bytes = Vec::new();
while let Some(chunk) = stream.next().await {
bytes.extend_from_slice(&chunk.expect("foreground chunk"));
}
assert_eq!(bytes, b"foreground");
let memory = BackgroundMemoryBudget::new();
let (resources, _) = test_resource_units(1);
let stream = with_optional_resource_unit(Some(resources.with_memory_budget(memory.clone())), async {
BackgroundReaderStream::with_capacity(tokio::io::empty(), 8)
})
.await;
drop(stream);
assert_eq!(memory.snapshot().used, 0);
}
}
#[cfg(test)]
mod reuse_tests {
use super::*;
use bytes::Bytes;
use futures::StreamExt;
use std::sync::Arc;
use std::time::Duration;
use std::{
io,
pin::Pin,
task::{Context, Poll},
};
use tokio::io::{AsyncRead, ReadBuf};
static DATA: [u8; 1024] = [0x37; 1024];
struct ShortReader {
offset: usize,
length: usize,
chunk: usize,
fail_at: Option<usize>,
}
impl AsyncRead for ShortReader {
fn poll_read(mut self: Pin<&mut Self>, _: &mut Context<'_>, out: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
if self.fail_at == Some(self.offset) {
return Poll::Ready(Err(io::ErrorKind::PermissionDenied.into()));
}
let n = (self.length - self.offset).min(self.chunk).min(out.remaining());
out.put_slice(&DATA[self.offset..self.offset + n]);
self.offset += n;
Poll::Ready(Ok(()))
}
}
fn budget(limit: u64) -> Arc<crate::disk::BackgroundMemoryBudget> {
let b = crate::disk::BackgroundMemoryBudget::new();
b.set_sampled_limit(limit, Duration::from_secs(30));
b
}
async fn stream<R: AsyncRead>(
budget: &Arc<crate::disk::BackgroundMemoryBudget>,
reader: R,
capacity: usize,
) -> BackgroundReaderStream<R> {
let (resources, _) = crate::disk::test_resource_units(1);
crate::disk::with_optional_resource_unit(Some(resources.with_memory_budget(budget.clone())), async {
BackgroundReaderStream::with_capacity(reader, capacity)
})
.await
}
#[tokio::test]
async fn complete_consumed_frames_reclaim_the_same_storage_after_expiry() {
let b = budget(64);
let mut s = stream(&b, io::Cursor::new(&DATA[..]), 64).await;
let mut pointer = None;
for _ in 0..16 {
let frame = s.next().await.expect("next full frame").expect("reclaimed backing");
assert_eq!(frame.len(), 64);
assert_eq!(frame.as_ref(), &DATA[..64]);
if let Some(p) = pointer {
assert_eq!(frame.as_ptr(), p);
} else {
pointer = Some(frame.as_ptr());
}
b.set_sampled_limit(0, Duration::ZERO);
drop(frame);
assert_eq!(b.snapshot().used, 64);
}
assert!(s.next().await.is_none());
assert_eq!(b.snapshot().used, 0);
}
#[tokio::test]
async fn short_live_frames_use_disjoint_tail_then_reserve_separate_storage() {
let b = budget(32);
let mut s = stream(
&b,
ShortReader {
offset: 0,
length: 32,
chunk: 4,
fail_at: None,
},
16,
)
.await;
let mut frames: Vec<Bytes> = Vec::new();
for i in 0..4 {
let frame = s.next().await.expect("short frame").expect("remaining initialized tail");
if i != 0 {
assert_eq!(frame.as_ptr(), frames[0].as_ptr().wrapping_add(i * 4));
}
assert_eq!(frame.as_ref(), &DATA[..4]);
frames.push(frame);
assert_eq!(b.snapshot().used, 16);
}
let next = s.next().await.expect("independent frame").expect("second backing fits");
assert_ne!(next.as_ptr(), frames[0].as_ptr());
assert_eq!(b.snapshot().used, 32);
drop(frames);
assert_eq!(b.snapshot().used, 16);
drop(s);
assert_eq!(b.snapshot().used, 16);
drop(next);
assert_eq!(b.snapshot().used, 0);
}
#[tokio::test]
async fn final_slice_prevents_reclaim_and_keeps_old_plus_new_charged() {
let b = budget(32);
let mut s = stream(&b, io::Cursor::new(&DATA[..32]), 16).await;
let first = s.next().await.expect("old frame").expect("old backing");
let last = first.slice(15..);
drop(first);
let next = s.next().await.expect("new frame").expect("second backing fits");
assert_eq!(last.as_ref(), &DATA[..1]);
assert_ne!(next.as_ptr(), last.as_ptr());
assert_eq!(b.snapshot().used, 32);
drop(s);
assert_eq!(b.snapshot().used, 32);
drop(next);
assert_eq!(b.snapshot().used, 16);
drop(last);
assert_eq!(b.snapshot().used, 0);
}
#[tokio::test]
async fn expiry_uses_tail_but_rejects_new_storage_before_consuming_source() {
let b = budget(16);
let mut s = stream(
&b,
ShortReader {
offset: 0,
length: 32,
chunk: 4,
fail_at: None,
},
16,
)
.await;
let first = s.next().await.expect("first frame").expect("initial backing");
b.set_sampled_limit(0, Duration::ZERO);
for _ in 0..3 {
drop(s.next().await.expect("tail frame").expect("reuse charged storage"));
}
let error = s
.next()
.await
.expect("new backing denied")
.expect_err("live frame prevents reclaim");
assert!(crate::disk::background_memory::is_memory_pressure(&error));
assert!(s.next().await.is_none());
assert_eq!(b.snapshot().used, 16);
drop(first);
assert_eq!(b.snapshot().used, 0);
}
#[tokio::test]
async fn eof_and_error_release_spare_but_keep_live_output_charged() {
for fail_at in [None, Some(4)] {
let b = budget(16);
let mut s = stream(
&b,
ShortReader {
offset: 0,
length: 4,
chunk: 4,
fail_at,
},
16,
)
.await;
let frame = s.next().await.expect("first short frame").expect("payload");
match fail_at {
None => assert!(s.next().await.is_none()),
Some(_) => assert_eq!(
s.next()
.await
.expect("read error")
.expect_err("preserve source failure")
.kind(),
io::ErrorKind::PermissionDenied
),
}
assert!(s.next().await.is_none());
assert_eq!(b.snapshot().used, 16);
assert_eq!(frame.as_ref(), &DATA[..4]);
drop(frame);
assert_eq!(b.snapshot().used, 0);
}
}
#[test]
fn runtime_shutdown_drops_a_pending_stream_without_refunding_early() {
for abort in [false, true] {
let b = budget(16);
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("runtime");
let (_tx, reader) = tokio::io::duplex(4);
let (began, wait) = tokio::sync::oneshot::channel();
let memory = b.clone();
let handle = rt.spawn(async move {
let mut s = stream(&memory, reader, 16).await;
assert!(futures::poll!(s.next()).is_pending());
assert_eq!(memory.snapshot().used, 16);
began.send(()).expect("notify pending allocation");
let _ = s.next().await;
});
rt.block_on(wait).expect("stream pending");
assert_eq!(b.snapshot().used, 16);
if abort {
handle.abort();
let _ = rt.block_on(handle);
}
drop(rt);
assert_eq!(b.snapshot().used, 0);
}
}
}
+13 -4
View File
@@ -17,10 +17,10 @@ use crate::cluster::rpc::http_auth::RPC_REPLAY_SCOPE_VERSION_HEADER;
use crate::cluster::rpc::http_auth::{
AuthenticatedPeerReplayCapabilities, RPC_AUTH_VERSION_HEADER, RPC_AUTH_VERSION_V2, RPC_BOOT_EPOCH_CHALLENGE_HEADER,
RPC_CONTENT_SHA256_HEADER, RPC_REPLAY_CACHE_CAPABILITY_HEADER, RPC_REPLAY_CACHE_CAPABILITY_PROOF_HEADER,
RollingMutationBodyDigest, TIMESTAMP_HEADER, internode_rpc_body_digest_strict,
verify_tonic_peer_replay_capabilities_response,
RollingMutationBodyDigest, TIMESTAMP_HEADER, gen_tonic_replay_scope_headers_with_admission, internode_rpc_body_digest_strict,
rpc_background_admission, verify_tonic_peer_replay_capabilities_response,
};
use crate::cluster::rpc::{gen_tonic_replay_scope_headers, gen_tonic_signature_headers, normalize_tonic_rpc_audience};
use crate::cluster::rpc::{gen_tonic_signature_headers, normalize_tonic_rpc_audience};
#[cfg(test)]
use crate::cluster::rpc::{tonic_boot_epoch_challenge, tonic_boot_epoch_response_headers};
use crate::disk::error::{DiskError, Error as DiskErrorType, RpcStatusError};
@@ -413,7 +413,16 @@ where
.get(RPC_CONTENT_SHA256_HEADER)
.and_then(|value| value.to_str().ok()),
) {
match gen_tonic_replay_scope_headers(audience, request.uri().path(), timestamp, content_sha256, boot_epoch) {
match rpc_background_admission(request.headers()).and_then(|background| {
gen_tonic_replay_scope_headers_with_admission(
audience,
request.uri().path(),
timestamp,
content_sha256,
boot_epoch,
background,
)
}) {
Ok(headers) => request.headers_mut().extend(headers),
Err(error) => debug!(error = %error, "could not attach replay-scoped RPC signature"),
}
+195 -24
View File
@@ -61,6 +61,9 @@ use uuid::Uuid;
type HmacSha256 = Hmac<Sha256>;
const SIGNATURE_HEADER: &str = "x-rustfs-signature";
pub const RPC_BACKGROUND_ADMISSION_HEADER: &str = "x-rustfs-rpc-background-admission";
const RPC_BACKGROUND_ADMISSION_V1: &str = "v1";
const RPC_BACKGROUND_ADMISSION_DOMAIN: &[u8] = b"rustfs-rpc-background-admission-v1\0";
pub(crate) const TIMESTAMP_HEADER: &str = "x-rustfs-timestamp";
pub(crate) const RPC_AUTH_VERSION_HEADER: &str = "x-rustfs-rpc-auth-version";
const RPC_SIGNATURE_V2_HEADER: &str = "x-rustfs-rpc-signature-v2";
@@ -77,6 +80,46 @@ pub(crate) const RPC_REPLAY_CACHE_CAPABILITY_HEADER: &str = "x-rustfs-rpc-replay
pub(crate) const RPC_REPLAY_CACHE_CAPABILITY_PROOF_HEADER: &str = "x-rustfs-rpc-replay-cache-capability-proof";
const RPC_REPLAY_SCOPE_VERSION_V3: &str = "3";
const RPC_RESPONSE_PROOF_DOMAIN: &[u8] = b"rustfs-rpc-response-proof-v1\0";
tokio::task_local! {
static RPC_BACKGROUND_REQUEST: bool;
}
pub(crate) fn current_rpc_background_request() -> bool {
RPC_BACKGROUND_REQUEST
.try_with(|background| *background)
.unwrap_or_else(|_| crate::disk::has_resource_io_scope())
}
/// Preserve the authenticated wire class while polling locally admitted work.
/// This scope carries only the class, never node, disk, or memory ownership.
pub async fn with_rpc_background_request<F: std::future::Future>(background: bool, operation: F) -> F::Output {
RPC_BACKGROUND_REQUEST.scope(background, operation).await
}
/// Parse the request class only; callers must verify the signature before using it.
pub fn rpc_background_admission(headers: &HeaderMap) -> std::io::Result<bool> {
let mut values = headers.get_all(RPC_BACKGROUND_ADMISSION_HEADER).iter();
match (values.next(), values.next()) {
(None, None) => Ok(false),
(Some(value), None) if value == RPC_BACKGROUND_ADMISSION_V1 => Ok(true),
_ => Err(std::io::Error::other("Invalid RPC background admission version")),
}
}
fn bind_background_admission(mac: &mut HmacSha256, background: bool) {
if background {
// Every auth lane binds the class, so stripping the marker or falling
// back to an older signature cannot turn admitted work into foreground I/O.
mac.update(RPC_BACKGROUND_ADMISSION_DOMAIN);
}
}
fn insert_background_admission(headers: &mut HeaderMap, background: bool) {
if background {
headers.insert(RPC_BACKGROUND_ADMISSION_HEADER, HeaderValue::from_static(RPC_BACKGROUND_ADMISSION_V1));
}
}
const RPC_REPLAY_SCOPE_DOMAIN: &[u8] = b"rustfs-rpc-replay-scope-v3\0";
const RPC_BOOT_EPOCH_PROOF_DOMAIN: &[u8] = b"rustfs-rpc-boot-epoch-proof-v1\0";
const RPC_REPLAY_CACHE_CAPABILITY_PROOF_DOMAIN: &[u8] = b"rustfs-rpc-replay-cache-capability-proof-v1\0";
@@ -519,21 +562,23 @@ fn redacted_rpc_path(url: &str) -> String {
}
/// Generate HMAC-SHA256 signature for the given data
fn generate_signature(secret: &str, url: &str, method: &Method, timestamp: i64) -> String {
fn generate_signature(secret: &str, url: &str, method: &Method, timestamp: i64, background: bool) -> String {
let data = signature_payload(url, method, timestamp);
let mut mac = <HmacSha256 as KeyInit>::new_from_slice(secret.as_bytes()).expect("HMAC can take key of any size");
bind_background_admission(&mut mac, background);
mac.update(data.as_bytes());
let result = mac.finalize();
base64_simd::STANDARD.encode_to_string(result.into_bytes())
}
fn verify_signature(secret: &str, url: &str, method: &Method, timestamp: i64, signature: &str) -> bool {
fn verify_signature(secret: &str, url: &str, method: &Method, timestamp: i64, signature: &str, background: bool) -> bool {
let Ok(signature) = base64_simd::STANDARD.decode_to_vec(signature) else {
return false;
};
let data = signature_payload(url, method, timestamp);
let mut mac = <HmacSha256 as KeyInit>::new_from_slice(secret.as_bytes()).expect("HMAC can take key of any size");
bind_background_admission(&mut mac, background);
mac.update(data.as_bytes());
mac.verify_slice(&signature).is_ok()
}
@@ -771,6 +816,7 @@ pub fn verify_ns_scanner_capability_with_checkpoint_coverage(
#[derive(Clone, Copy)]
struct SignatureV2Scope<'a> {
background: bool,
audience: &'a str,
service: &'a str,
rpc_method: &'a str,
@@ -780,6 +826,7 @@ struct SignatureV2Scope<'a> {
}
fn update_signature_v2(mac: &mut HmacSha256, scope: SignatureV2Scope<'_>) {
bind_background_admission(mac, scope.background);
for part in [
b"rustfs-rpc-auth-v2|".as_slice(),
scope.audience.as_bytes(),
@@ -818,6 +865,7 @@ fn verify_signature_v2(secret: &str, scope: SignatureV2Scope<'_>, signature: &st
#[derive(Clone, Copy)]
struct ReplayScope<'a> {
background: bool,
audience: &'a str,
path: &'a str,
timestamp: &'a str,
@@ -827,6 +875,7 @@ struct ReplayScope<'a> {
}
fn update_replay_scope(mac: &mut HmacSha256, scope: ReplayScope<'_>) {
bind_background_admission(mac, scope.background);
mac.update(RPC_REPLAY_SCOPE_DOMAIN);
for part in [
scope.audience.as_bytes(),
@@ -970,6 +1019,24 @@ pub fn gen_tonic_replay_scope_headers(
timestamp: &str,
content_sha256: &str,
boot_epoch: Uuid,
) -> std::io::Result<HeaderMap> {
gen_tonic_replay_scope_headers_with_admission(
audience,
path,
timestamp,
content_sha256,
boot_epoch,
current_rpc_background_request(),
)
}
pub(crate) fn gen_tonic_replay_scope_headers_with_admission(
audience: &str,
path: &str,
timestamp: &str,
content_sha256: &str,
boot_epoch: Uuid,
background: bool,
) -> std::io::Result<HeaderMap> {
if audience.is_empty() || !path.starts_with('/') || !valid_content_sha256(content_sha256) || boot_epoch.is_nil() {
return Err(std::io::Error::other("Invalid replay-scoped RPC signing scope"));
@@ -983,6 +1050,7 @@ pub fn gen_tonic_replay_scope_headers(
let signature = generate_replay_scope_signature(
&get_shared_secret()?,
ReplayScope {
background,
audience,
path,
timestamp,
@@ -992,6 +1060,7 @@ pub fn gen_tonic_replay_scope_headers(
},
)?;
let mut headers = HeaderMap::new();
insert_background_admission(&mut headers, background);
headers.insert(RPC_REPLAY_SCOPE_VERSION_HEADER, HeaderValue::from_static(RPC_REPLAY_SCOPE_VERSION_V3));
headers.insert(
RPC_REPLAY_SCOPE_SIGNATURE_HEADER,
@@ -1226,10 +1295,12 @@ pub fn build_auth_headers(url: &str, method: &Method, headers: &mut HeaderMap) -
pub fn gen_signature_headers(url: &str, method: &Method) -> std::io::Result<HeaderMap> {
let secret = get_shared_secret()?;
let timestamp = OffsetDateTime::now_utc().unix_timestamp();
let background = current_rpc_background_request();
let signature = generate_signature(&secret, url, method, timestamp);
let signature = generate_signature(&secret, url, method, timestamp, background);
let mut headers = HeaderMap::new();
insert_background_admission(&mut headers, background);
headers.insert(SIGNATURE_HEADER, HeaderValue::from_str(&signature).expect("Invalid header value"));
headers.insert(
TIMESTAMP_HEADER,
@@ -1264,10 +1335,12 @@ pub fn gen_tonic_signature_headers(
let timestamp_header = timestamp.to_string();
let body_nonce = (content_sha256 != UNSIGNED_PAYLOAD).then(|| Uuid::new_v4().to_string());
let nonce = body_nonce.as_deref().unwrap_or(UNSIGNED_PAYLOAD_NONCE);
let legacy_signature = generate_signature(&secret, TONIC_RPC_PREFIX, &Method::GET, timestamp);
let background = current_rpc_background_request();
let legacy_signature = generate_signature(&secret, TONIC_RPC_PREFIX, &Method::GET, timestamp, background);
let signature_v2 = generate_signature_v2(
&secret,
SignatureV2Scope {
background,
audience,
service,
rpc_method,
@@ -1278,6 +1351,7 @@ pub fn gen_tonic_signature_headers(
)?;
let mut headers = HeaderMap::new();
insert_background_admission(&mut headers, background);
headers.insert(SIGNATURE_HEADER, header_value(&legacy_signature, SIGNATURE_HEADER)?);
headers.insert(TIMESTAMP_HEADER, header_value(&timestamp_header, TIMESTAMP_HEADER)?);
headers.insert(RPC_AUTH_VERSION_HEADER, HeaderValue::from_static(RPC_AUTH_VERSION_V2));
@@ -1501,6 +1575,7 @@ fn verify_tonic_replay_scope_signature(audience: &str, path: &str, headers: &Hea
if !verify_replay_scope_signature(
&secret,
ReplayScope {
background: rpc_background_admission(headers)?,
audience,
path,
timestamp,
@@ -1625,6 +1700,9 @@ fn verify_tonic_rpc_signature_with_strictness(
strict: bool,
) -> std::io::Result<()> {
if !has_v2_auth_headers(headers) {
if rpc_background_admission(headers)? {
return Err(std::io::Error::other("RPC background admission requires v2 authentication"));
}
// RUSTFS_COMPAT_TODO(heal-rpc-auth-v2): accept old peers during rolling upgrades. Remove after the minimum
// supported RustFS peer version sends v2 authentication on every internode gRPC request.
if strict {
@@ -1696,6 +1774,7 @@ fn verify_tonic_rpc_signature_with_strictness(
if !verify_signature_v2(
&secret,
SignatureV2Scope {
background: rpc_background_admission(headers)?,
audience,
service,
rpc_method,
@@ -1736,7 +1815,7 @@ pub fn verify_rpc_signature(url: &str, method: &Method, headers: &HeaderMap) ->
// Verify signature with constant-time HMAC comparison.
let secret = get_shared_secret()?;
if !verify_signature(&secret, url, method, timestamp, signature) {
if !verify_signature(&secret, url, method, timestamp, signature, rpc_background_admission(headers)?) {
let rpc_path = redacted_rpc_path(url);
error!(
rpc_path = %rpc_path,
@@ -1828,6 +1907,98 @@ mod tests {
runtime_sources::ensure_test_rpc_secret();
}
#[tokio::test]
async fn background_admission_is_bound_in_http_and_both_tonic_signature_lanes() {
ensure_test_rpc_secret();
let secret = get_shared_secret().expect("RPC secret");
let path = "/node_service.NodeService/ReadAll";
for background in [false, true] {
let (http, tonic) = with_rpc_background_request(background, async {
(
gen_signature_headers("/rustfs/rpc/read?disk=a", &Method::GET).expect("HTTP signature"),
gen_tonic_signature_headers("node-a:9000", "node_service.NodeService", "ReadAll", None)
.expect("gRPC signature"),
)
})
.await;
assert_eq!(rpc_background_admission(&http).expect("HTTP class"), background);
assert_eq!(rpc_background_admission(&tonic).expect("gRPC class"), background);
verify_rpc_signature("/rustfs/rpc/read?disk=a", &Method::GET, &http).expect("HTTP authentication");
verify_tonic_rpc_signature_with_strictness("node-a:9000", path, &tonic, true).expect("v2 authentication");
let timestamp = tonic[TIMESTAMP_HEADER].to_str().expect("timestamp");
let content = tonic[RPC_CONTENT_SHA256_HEADER].to_str().expect("content digest");
assert_eq!(
verify_signature_v2(
&secret,
SignatureV2Scope {
background: false,
audience: "node-a:9000",
service: "node_service.NodeService",
rpc_method: "ReadAll",
timestamp,
nonce: tonic[RPC_NONCE_HEADER].to_str().expect("nonce"),
content_sha256: content,
},
tonic[RPC_SIGNATURE_V2_HEADER].to_str().expect("signature"),
),
!background,
"an old peer cannot authenticate a background request as foreground"
);
let mut replay = tonic.clone();
// The channel can run outside the caller's task-local scope.
replay.extend(
gen_tonic_replay_scope_headers_with_admission(
"node-a:9000",
path,
timestamp,
content,
tonic_rpc_boot_epoch(),
background,
)
.expect("v3 signature"),
);
for mutation in ["strip-or-add", "unknown", "duplicate"] {
let alter = |headers: &HeaderMap| {
let mut changed = headers.clone();
match mutation {
"strip-or-add" if background => {
changed.remove(RPC_BACKGROUND_ADMISSION_HEADER);
}
"strip-or-add" => insert_background_admission(&mut changed, true),
"unknown" => {
changed.insert(RPC_BACKGROUND_ADMISSION_HEADER, HeaderValue::from_static("v2"));
}
_ => {
insert_background_admission(&mut changed, true);
changed.append(RPC_BACKGROUND_ADMISSION_HEADER, HeaderValue::from_static("v1"));
}
}
changed
};
assert!(verify_rpc_signature("/rustfs/rpc/read?disk=a", &Method::GET, &alter(&http)).is_err());
assert!(verify_tonic_rpc_signature_with_strictness("node-a:9000", path, &alter(&tonic), true).is_err());
assert!(verify_tonic_replay_scope_signature("node-a:9000", path, &alter(&replay)).is_err());
}
verify_tonic_replay_scope_signature("node-a:9000", path, &replay).expect("v3 authentication");
assert!(
verify_tonic_replay_scope_signature("node-a:9000", path, &replay).is_err(),
"replay rejected"
);
if background {
let mut downgraded = HeaderMap::new();
for name in [SIGNATURE_HEADER, TIMESTAMP_HEADER, RPC_BACKGROUND_ADMISSION_HEADER] {
downgraded.insert(name, tonic[name].clone());
}
assert!(verify_tonic_rpc_signature_with_strictness("node-a:9000", path, &downgraded, false).is_err());
downgraded.remove(RPC_BACKGROUND_ADMISSION_HEADER);
assert!(verify_tonic_rpc_signature_with_strictness("node-a:9000", path, &downgraded, false).is_err());
}
}
}
#[test]
fn namespace_scanner_capability_proof_binds_challenge_and_server_epoch() {
let secret = "test-scanner-capability-secret";
@@ -1965,8 +2136,8 @@ mod tests {
let method = Method::GET;
let timestamp = 1640995200; // Fixed timestamp
let signature1 = generate_signature(secret, url, &method, timestamp);
let signature2 = generate_signature(secret, url, &method, timestamp);
let signature1 = generate_signature(secret, url, &method, timestamp, false);
let signature2 = generate_signature(secret, url, &method, timestamp, false);
assert_eq!(signature1, signature2, "Same inputs should produce same signature");
assert!(!signature1.is_empty(), "Signature should not be empty");
@@ -1979,10 +2150,10 @@ mod tests {
let method = Method::GET;
let timestamp = 1640995200;
let signature1 = generate_signature(secret, url, &method, timestamp);
let signature2 = generate_signature(secret, "http://different.com/api/test2", &method, timestamp);
let signature3 = generate_signature(secret, url, &Method::POST, timestamp);
let signature4 = generate_signature(secret, url, &method, timestamp + 1);
let signature1 = generate_signature(secret, url, &method, timestamp, false);
let signature2 = generate_signature(secret, "http://different.com/api/test2", &method, timestamp, false);
let signature3 = generate_signature(secret, url, &Method::POST, timestamp, false);
let signature4 = generate_signature(secret, url, &method, timestamp + 1, false);
assert_ne!(signature1, signature2, "Different URLs should produce different signatures");
assert_ne!(signature1, signature3, "Different methods should produce different signatures");
@@ -2088,14 +2259,14 @@ mod tests {
let url = "http://example.com/api/test";
let method = Method::GET;
let timestamp = 1640995200;
let signature = generate_signature(secret, url, &method, timestamp);
let signature = generate_signature(secret, url, &method, timestamp, false);
let mut tampered = base64_simd::STANDARD.decode_to_vec(&signature).unwrap();
tampered[0] ^= 1;
let tampered_signature = base64_simd::STANDARD.encode_to_string(tampered);
assert!(verify_signature(secret, url, &method, timestamp, &signature));
assert!(!verify_signature(secret, url, &method, timestamp, &tampered_signature));
assert!(!verify_signature(secret, url, &method, timestamp, "invalid-signature"));
assert!(verify_signature(secret, url, &method, timestamp, &signature, false));
assert!(!verify_signature(secret, url, &method, timestamp, &tampered_signature, false));
assert!(!verify_signature(secret, url, &method, timestamp, "invalid-signature", false));
}
#[test]
@@ -2112,11 +2283,11 @@ mod tests {
);
let method = Method::GET;
let timestamp = 1_640_995_200;
let signature = generate_signature(secret, signed_url, &method, timestamp);
let signature = generate_signature(secret, signed_url, &method, timestamp, false);
assert!(verify_signature(secret, signed_url, &method, timestamp, &signature));
assert!(!verify_signature(secret, downgraded_url, &method, timestamp, &signature));
assert!(!verify_signature(secret, tampered_body_digest, &method, timestamp, &signature));
assert!(verify_signature(secret, signed_url, &method, timestamp, &signature, false));
assert!(!verify_signature(secret, downgraded_url, &method, timestamp, &signature, false));
assert!(!verify_signature(secret, tampered_body_digest, &method, timestamp, &signature, false));
}
#[test]
@@ -2126,7 +2297,7 @@ mod tests {
let method = Method::GET;
let timestamp = OffsetDateTime::now_utc().unix_timestamp();
let secret = get_shared_secret().expect("test RPC secret should resolve");
let expected_signature = generate_signature(&secret, url, &method, timestamp);
let expected_signature = generate_signature(&secret, url, &method, timestamp, false);
let invalid_signature = "invalid-signature";
let logs = CapturedLogs::default();
let subscriber = tracing_subscriber::fmt()
@@ -2164,7 +2335,7 @@ mod tests {
// Set expired timestamp (older than SIGNATURE_VALID_DURATION)
let expired_timestamp = OffsetDateTime::now_utc().unix_timestamp() - SIGNATURE_VALID_DURATION - 10;
let secret = get_shared_secret().expect("test RPC secret should resolve");
let signature = generate_signature(&secret, url, &method, expired_timestamp);
let signature = generate_signature(&secret, url, &method, expired_timestamp, false);
headers.insert(SIGNATURE_HEADER, HeaderValue::from_str(&signature).unwrap());
headers.insert(TIMESTAMP_HEADER, HeaderValue::from_str(&expired_timestamp.to_string()).unwrap());
@@ -2186,7 +2357,7 @@ mod tests {
let future_timestamp = OffsetDateTime::now_utc().unix_timestamp() + SIGNATURE_VALID_DURATION + 10;
let secret = get_shared_secret().expect("test RPC secret should resolve");
let signature = generate_signature(&secret, url, &method, future_timestamp);
let signature = generate_signature(&secret, url, &method, future_timestamp, false);
headers.insert(SIGNATURE_HEADER, HeaderValue::from_str(&signature).unwrap());
headers.insert(TIMESTAMP_HEADER, HeaderValue::from_str(&future_timestamp.to_string()).unwrap());
@@ -2300,7 +2471,7 @@ mod tests {
// Test timestamp just within valid duration
let valid_timestamp = current_time - SIGNATURE_VALID_DURATION + 1;
let signature = generate_signature(&secret, url, &method, valid_timestamp);
let signature = generate_signature(&secret, url, &method, valid_timestamp, false);
headers.insert(SIGNATURE_HEADER, HeaderValue::from_str(&signature).unwrap());
headers.insert(TIMESTAMP_HEADER, HeaderValue::from_str(&valid_timestamp.to_string()).unwrap());
@@ -2311,7 +2482,7 @@ mod tests {
// Test timestamp just outside valid duration
let mut headers = HeaderMap::new();
let invalid_timestamp = current_time - SIGNATURE_VALID_DURATION - 15;
let signature = generate_signature(&secret, url, &method, invalid_timestamp);
let signature = generate_signature(&secret, url, &method, invalid_timestamp, false);
headers.insert(SIGNATURE_HEADER, HeaderValue::from_str(&signature).unwrap());
headers.insert(TIMESTAMP_HEADER, HeaderValue::from_str(&invalid_timestamp.to_string()).unwrap());
@@ -33,7 +33,9 @@ use rustfs_config::{
DEFAULT_INTERNODE_DATA_TRANSPORT, ENV_RUSTFS_INTERNODE_DATA_TRANSPORT, INTERNODE_DATA_TRANSPORT_TCP,
KNOWN_INTERNODE_DATA_TRANSPORT_BACKENDS,
};
use rustfs_rio::{ChunkReaderBox, HttpChunkReader, HttpReader, HttpWriter};
use rustfs_rio::{
ChunkReaderBox, HttpChunkReader, HttpReader, HttpWriter, HttpWriterMemoryReservation, HttpWriterMemoryReserver,
};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::future::Future;
@@ -66,6 +68,19 @@ const PUT_FILE_CAPABILITY_PROBE_TIMEOUT: Duration = Duration::from_secs(5);
const CONTENT_TYPE_JSON: &str = "application/json";
const CONTENT_TYPE_MSGPACK: &str = "application/msgpack";
fn http_writer_memory_reserver() -> Option<HttpWriterMemoryReserver> {
crate::disk::current_resource_unit()
.and_then(|resources| resources.memory_budget())
.map(|budget| {
// Payloads can outlive the operation; never retain its node or disk units.
Arc::new(move |capacity| {
budget
.reserve(capacity)
.map(|reservation| Box::new(reservation) as HttpWriterMemoryReservation)
}) as HttpWriterMemoryReserver
})
}
fn unsupported_transport_message(transport: &str) -> String {
format!(
"invalid {ENV_RUSTFS_INTERNODE_DATA_TRANSPORT}={transport:?}; supported values: {}",
@@ -346,13 +361,17 @@ impl InternodeDataTransport for TcpHttpInternodeDataTransport {
}
async fn open_write(&self, request: WriteStreamRequest) -> Result<FileWriter> {
let memory_reserver = http_writer_memory_reserver();
let server_epoch = self.put_file_auth_capability(&request.endpoint).await?;
let auth_scope = server_epoch.map(|server_epoch| (Uuid::new_v4(), server_epoch));
let url = build_put_file_stream_url(&request, auth_scope);
let endpoint = request.endpoint;
let mut headers = json_headers();
build_auth_headers(&url, &Method::PUT, &mut headers)?;
let writer = HttpWriter::new(url.clone(), Method::PUT, headers).await?;
let writer = match memory_reserver {
Some(reserve) => HttpWriter::new_with_memory_reserver(url.clone(), Method::PUT, headers, reserve).await?,
None => HttpWriter::new(url.clone(), Method::PUT, headers).await?,
};
match auth_scope {
Some((nonce, server_epoch)) => Ok(Box::new(PutFileAuthWriter::new(writer, url, nonce, endpoint, server_epoch))),
None => Ok(Box::new(writer)),
@@ -948,6 +967,193 @@ mod tests {
use tokio::io::AsyncWriteExt;
use tokio::sync::{Barrier, Notify};
fn memory_test_write_request(authenticated: bool) -> WriteStreamRequest {
crate::runtime::sources::ensure_test_rpc_secret();
let endpoint = format!("http://writer-memory-{}.invalid", Uuid::new_v4());
put_file_capability_cache_entry(&endpoint).write().cached = Some(if authenticated {
PutFileCapabilityState::V1 {
server_epoch: Uuid::new_v4(),
revalidate_after: Instant::now() + PUT_FILE_V1_CAPABILITY_TTL,
}
} else {
PutFileCapabilityState::LegacyUntil(Instant::now() + PUT_FILE_LEGACY_CAPABILITY_TTL)
});
WriteStreamRequest {
endpoint,
disk: "disk-a".to_owned(),
volume: "bucket".to_owned(),
path: "part.1".to_owned(),
append: false,
size: 4,
}
}
#[tokio::test]
async fn http_write_memory_reserver_outlives_only_the_memory_scope() {
use crate::disk::{BackgroundMemoryBudget, test_resource_units, with_optional_resource_unit};
assert!(http_writer_memory_reserver().is_none(), "foreground must use the existing writer");
let memory = BackgroundMemoryBudget::new();
memory.set_sampled_limit(2 * 1024 * 1024, Duration::from_secs(30));
let (resources, gates) = test_resource_units(2);
let request = memory_test_write_request(true);
let endpoint = request.endpoint.clone();
let writer = with_optional_resource_unit(
Some(resources.with_memory_budget(memory.clone())),
TcpHttpInternodeDataTransport.open_write(request),
)
.await;
PUT_FILE_CAPABILITY_CACHE.write().remove(&endpoint);
let mut writer = writer.expect("admitted transport writer");
assert!(
gates.iter().all(|gate| gate.snapshot().active == 0),
"the callback must not retain disk units"
);
assert_eq!(memory.snapshot().used, 1024 * 1024);
writer.write_all(b"data").await.expect("buffered bytes");
writer.flush().await.expect("transfer charged allocation to queue");
memory.set_sampled_limit(2 * 1024 * 1024, Duration::ZERO);
writer
.write_all(&vec![0x5a; 1024 * 1024 - 4])
.await
.expect("reuse the already charged tail after observation expiry");
assert_eq!(memory.snapshot().used, 1024 * 1024, "reuse must not allocate separate storage");
let error = writer
.write(b"rejected")
.await
.expect_err("expired observation must reject new storage");
assert_eq!(error.kind(), io::ErrorKind::WouldBlock);
assert!(
crate::disk::background_memory::is_memory_pressure(&error),
"preserve the typed memory-pressure marker"
);
assert_eq!(memory.snapshot().used, 1024 * 1024, "queued storage remains charged after expiry");
drop(writer);
tokio::task::yield_now().await;
assert_eq!(memory.snapshot().used, 0);
}
#[tokio::test]
async fn http_write_memory_denial_preserves_the_budget_error() {
use crate::disk::{BackgroundMemoryBudget, test_resource_units, with_optional_resource_unit};
for (authenticated, sample) in [false, true].into_iter().flat_map(|authenticated| {
[
None,
Some((1024 * 1024 - 1, Duration::from_secs(30))),
Some((1024 * 1024, Duration::ZERO)),
]
.into_iter()
.map(move |sample| (authenticated, sample))
}) {
let memory = BackgroundMemoryBudget::new();
if let Some((limit, valid_for)) = sample {
memory.set_sampled_limit(limit, valid_for);
}
let (resources, _) = test_resource_units(1);
let request = memory_test_write_request(authenticated);
let endpoint = request.endpoint.clone();
let result = with_optional_resource_unit(
Some(resources.with_memory_budget(memory.clone())),
TcpHttpInternodeDataTransport.open_write(request),
)
.await;
PUT_FILE_CAPABILITY_CACHE.write().remove(&endpoint);
let error = match result {
Err(Error::Io(error)) => error,
Err(error) => panic!("unexpected transport error: {error}"),
Ok(_) => panic!("unavailable memory must reject construction"),
};
assert!(crate::disk::background_memory::is_memory_pressure(&error));
assert_eq!(error.kind(), io::ErrorKind::WouldBlock);
assert_eq!(memory.snapshot().used, 0);
}
}
#[tokio::test]
async fn admitted_remote_reads_bind_receiver_admission_to_the_http_signature() {
use crate::disk::with_background_resource_unit;
use rustfs_concurrency::adaptive::AdaptiveAdmission;
use std::num::NonZeroU16;
crate::runtime::sources::ensure_test_rpc_secret();
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("remote admission listener");
let address = listener.local_addr().expect("remote admission address");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("remote read connection");
let mut bytes = Vec::new();
while !bytes.windows(4).any(|window| window == b"\r\n\r\n") {
let mut chunk = [0; 1024];
let read = socket.read(&mut chunk).await.expect("remote read headers");
assert!(read > 0 && bytes.len() < 16 * 1024, "bounded complete request headers");
bytes.extend_from_slice(&chunk[..read]);
}
let request = std::str::from_utf8(&bytes).expect("HTTP request");
let mut lines = request.split("\r\n");
let target = lines
.next()
.expect("request line")
.split_whitespace()
.nth(1)
.expect("request target")
.to_owned();
let mut headers = HeaderMap::new();
for line in lines.take_while(|line| !line.is_empty()) {
let (name, value) = line.split_once(':').expect("header field");
headers.append(
http::header::HeaderName::from_bytes(name.as_bytes()).expect("header name"),
HeaderValue::from_str(value.trim()).expect("header value"),
);
}
socket
.write_all(b"HTTP/1.1 200 OK\r\ncontent-length: 5\r\nconnection: close\r\n\r\nhello")
.await
.expect("remote read response");
(target, headers)
});
let gate = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let resource = gate.try_acquire().expect("sender device admission");
let mut reader = tokio::time::timeout(
Duration::from_secs(5),
with_background_resource_unit(
resource,
TcpHttpInternodeDataTransport.open_read(ReadStreamRequest {
endpoint: format!("http://{address}"),
disk: "disk-a".to_owned(),
volume: "bucket".to_owned(),
path: "part.1".to_owned(),
offset: 0,
length: 5,
stall_timeout: Some(Duration::from_secs(5)),
}),
),
)
.await
.expect("bounded remote request")
.expect("remote reader");
let mut body = Vec::new();
reader.read_to_end(&mut body).await.expect("complete remote body");
assert_eq!(body, b"hello");
let (target, headers) = server.await.expect("remote request capture");
assert_eq!(gate.snapshot().active, 0, "remote body does not retain the sender's device");
assert_eq!(
headers
.get("x-rustfs-rpc-background-admission")
.and_then(|value| value.to_str().ok()),
Some("v1"),
"resource-controlled remote work must request receiver-side admission"
);
crate::cluster::rpc::verify_rpc_signature(&target, &Method::GET, &headers).expect("bound request signature");
let mut stripped = headers;
stripped.remove("x-rustfs-rpc-background-admission");
assert!(
crate::cluster::rpc::verify_rpc_signature(&target, &Method::GET, &stripped).is_err(),
"removing background admission must invalidate the signature"
);
}
async fn wait_for_capability_flight_waiters(entry: &PutFileCapabilityCacheEntry, waiters: usize) {
tokio::time::timeout(Duration::from_secs(5), async {
loop {
+13 -10
View File
@@ -13,6 +13,8 @@
// limitations under the License.
pub(crate) mod background_monitor;
pub(crate) mod background_pipe;
pub(crate) mod background_reader_stream;
pub(crate) mod client;
pub(crate) mod context_propagation;
pub(crate) mod http_auth;
@@ -33,16 +35,17 @@ pub use client::{
// Re-exported through `api::rpc`; not every item is consumed inside this crate.
#[allow(unused_imports)]
pub use http_auth::{
TONIC_RPC_PREFIX, build_auth_headers, build_put_file_auth_trailer, check_and_record_signed_rpc_nonce, gen_signature_headers,
gen_tonic_replay_scope_headers, gen_tonic_signature_headers, normalize_tonic_rpc_audience, set_tonic_canonical_body_digest,
set_tonic_mutation_body_digest, set_tonic_rolling_canonical_body_digest, set_tonic_rolling_mutation_body_digest,
sign_ns_scanner_capability, sign_ns_scanner_capability_with_checkpoint_coverage,
sign_ns_scanner_capability_with_tier_registry_generation, sign_put_file_capability, sign_tonic_rpc_response_proof,
tonic_boot_epoch_challenge, tonic_boot_epoch_response_headers, tonic_rpc_auth_failure_reason, verify_ns_scanner_capability,
verify_ns_scanner_capability_with_checkpoint_coverage, verify_ns_scanner_capability_with_tier_registry_generation,
verify_put_file_auth_trailer, verify_put_file_capability, verify_rpc_signature, verify_tonic_boot_epoch_response,
verify_tonic_canonical_body_digest, verify_tonic_mutation_body_digest, verify_tonic_mutation_body_digest_reject_unsigned,
verify_tonic_rpc_response_proof, verify_tonic_rpc_signature, verify_tonic_rpc_signature_with_bootstrap,
RPC_BACKGROUND_ADMISSION_HEADER, TONIC_RPC_PREFIX, build_auth_headers, build_put_file_auth_trailer,
check_and_record_signed_rpc_nonce, gen_signature_headers, gen_tonic_replay_scope_headers, gen_tonic_signature_headers,
normalize_tonic_rpc_audience, rpc_background_admission, set_tonic_canonical_body_digest, set_tonic_mutation_body_digest,
set_tonic_rolling_canonical_body_digest, set_tonic_rolling_mutation_body_digest, sign_ns_scanner_capability,
sign_ns_scanner_capability_with_checkpoint_coverage, sign_ns_scanner_capability_with_tier_registry_generation,
sign_put_file_capability, sign_tonic_rpc_response_proof, tonic_boot_epoch_challenge, tonic_boot_epoch_response_headers,
tonic_rpc_auth_failure_reason, verify_ns_scanner_capability, verify_ns_scanner_capability_with_checkpoint_coverage,
verify_ns_scanner_capability_with_tier_registry_generation, verify_put_file_auth_trailer, verify_put_file_capability,
verify_rpc_signature, verify_tonic_boot_epoch_response, verify_tonic_canonical_body_digest,
verify_tonic_mutation_body_digest, verify_tonic_mutation_body_digest_reject_unsigned, verify_tonic_rpc_response_proof,
verify_tonic_rpc_signature, verify_tonic_rpc_signature_with_bootstrap,
};
#[cfg(test)]
pub(crate) use internode_data_transport::TcpHttpInternodeDataTransport;
+105 -6
View File
@@ -17,6 +17,7 @@ use crate::cluster::rpc::client::{
node_service_delete_recovery_client, node_service_time_out_client, node_service_time_out_client_for_class,
node_service_time_out_client_no_auth,
};
use crate::cluster::rpc::http_auth::{current_rpc_background_request, with_rpc_background_request};
use crate::cluster::rpc::internode_data_transport::{
InternodeDataTransport, NsScannerCapabilityRequest, NsScannerStreamRequest, ReadStreamRequest, WalkDirStreamRequest,
WriteStreamRequest,
@@ -533,6 +534,7 @@ struct RetryingRemoteReader {
reader: Option<FileReader>,
transport: Arc<dyn InternodeDataTransport>,
request: ReadStreamRequest,
background: bool,
emitted: usize,
retried: bool,
initial_read_timeout: Option<Duration>,
@@ -553,6 +555,7 @@ impl RetryingRemoteReader {
reader: Some(reader),
transport,
request,
background: current_rpc_background_request(),
emitted: 0,
retried: false,
initial_read_timeout,
@@ -570,9 +573,10 @@ impl RetryingRemoteReader {
let request = resumed_read_request(&self.request, self.emitted)?;
let recovery_timeout = self.recovery_timeout;
let transport = Arc::clone(&self.transport);
self.resume = Some(AbortOnDropTask::new(tokio::spawn(async move {
with_remote_read_recovery_timeout(recovery_timeout, transport.open_read_fresh(request)).await
})));
self.resume = Some(AbortOnDropTask::new(tokio::spawn(with_rpc_background_request(
self.background,
async move { with_remote_read_recovery_timeout(recovery_timeout, transport.open_read_fresh(request)).await },
))));
Ok(())
}
@@ -680,6 +684,7 @@ struct RetryingRemoteChunkReader {
reader: Option<rustfs_rio::ChunkReaderBox>,
transport: Arc<dyn InternodeDataTransport>,
request: ReadStreamRequest,
background: bool,
emitted: usize,
retried: bool,
initial_read_timeout: Option<Duration>,
@@ -700,6 +705,7 @@ impl RetryingRemoteChunkReader {
reader: Some(reader),
transport,
request,
background: current_rpc_background_request(),
emitted: 0,
retried: false,
initial_read_timeout,
@@ -717,9 +723,10 @@ impl RetryingRemoteChunkReader {
let request = resumed_read_request(&self.request, self.emitted)?;
let recovery_timeout = self.recovery_timeout;
let transport = Arc::clone(&self.transport);
self.resume = Some(AbortOnDropTask::new(tokio::spawn(async move {
with_remote_read_recovery_timeout(recovery_timeout, transport.open_read_chunks_fresh(request)).await
})));
self.resume = Some(AbortOnDropTask::new(tokio::spawn(with_rpc_background_request(
self.background,
async move { with_remote_read_recovery_timeout(recovery_timeout, transport.open_read_chunks_fresh(request)).await },
))));
Ok(())
}
@@ -7071,6 +7078,98 @@ mod tests {
.await;
}
#[tokio::test]
async fn remote_read_retries_preserve_signed_admission_after_sender_scope_ends() {
use crate::cluster::rpc::{rpc_background_admission, verify_rpc_signature};
use crate::disk::with_background_resource_unit;
use rustfs_concurrency::adaptive::AdaptiveAdmission;
use std::num::NonZeroU16;
crate::runtime::sources::ensure_test_rpc_secret();
for background in [false, true] {
for chunked in [false, true] {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("retry listener");
let address = listener.local_addr().expect("retry address");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("fresh connection");
let mut bytes = Vec::new();
while !bytes.windows(4).any(|window| window == b"\r\n\r\n") {
let mut chunk = [0; 1024];
let read = socket.read(&mut chunk).await.expect("retry request");
assert!(read > 0 && bytes.len() < 16 * 1024, "bounded complete headers");
bytes.extend_from_slice(&chunk[..read]);
}
let request = std::str::from_utf8(&bytes).expect("HTTP request");
let mut lines = request.split("\r\n");
let target = lines
.next()
.expect("request line")
.split_whitespace()
.nth(1)
.expect("target")
.to_owned();
let mut headers = http::HeaderMap::new();
for line in lines.take_while(|line| !line.is_empty()) {
let (name, value) = line.split_once(':').expect("header");
headers.append(
http::header::HeaderName::from_bytes(name.as_bytes()).expect("header name"),
http::HeaderValue::from_str(value.trim()).expect("header value"),
);
}
socket
.write_all(b"HTTP/1.1 200 OK\r\ncontent-length: 6\r\nconnection: close\r\n\r\n456789")
.await
.expect("retry response");
(target, headers)
});
let gate = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let request = ReadStreamRequest {
endpoint: format!("http://{address}"),
..resume_request(10)
};
let create = async {
let transport = Arc::new(TcpHttpInternodeDataTransport);
if chunked {
Box::new(RetryingRemoteChunkReader::new_with_timeouts(
resume_step_chunk_reader(ResumeReadStep::PartialThenReset(b"0123".to_vec())),
transport,
request,
None,
None,
)) as FileReader
} else {
Box::new(RetryingRemoteReader::new_with_timeouts(
resume_step_reader(ResumeReadStep::PartialThenReset(b"0123".to_vec())),
transport,
request,
None,
None,
)) as FileReader
}
};
let mut reader = if background {
with_background_resource_unit(gate.try_acquire().expect("sender admission"), create).await
} else {
create.await
};
assert_eq!(gate.snapshot().active, 0, "retry reader cannot retain sender disk capacity");
let mut output = Vec::new();
tokio::time::timeout(
Duration::from_secs(5),
with_rpc_background_request(!background, reader.read_to_end(&mut output)),
)
.await
.expect("bounded retry")
.expect("resumed body");
assert_eq!(output, b"0123456789", "chunked={chunked}, background={background}");
let (target, headers) = server.await.expect("captured retry");
assert!(target.contains("offset=11") && target.contains("length=6"), "{target}");
assert_eq!(rpc_background_admission(&headers).expect("admission marker"), background);
verify_rpc_signature(&target, &http::Method::GET, &headers).expect("signed retry class");
}
}
}
#[tokio::test]
async fn remote_reader_resumes_from_emitted_bytes_without_duplicates() {
let transport = Arc::new(ResumeTransport::with_read_steps(vec![ResumeReadStep::Data(b"456789".to_vec())]));
+48 -48
View File
@@ -624,58 +624,58 @@ mod tests {
#[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 crate::disk::io_admission::{current_disk_io_context, test_resource_units, with_background_disk_io};
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;
for physical_count in [1, 2] {
for mode in 0..4 {
let runtime = Arc::new(CleanupRuntime::new(config(1, 1, 1)).expect("cleanup runtime"));
let worker_runtime = runtime.clone();
let (unit, gates) = test_resource_units(physical_count);
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_optional_resource_unit(Some(unit), 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"),
}
_ => 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: Vec<_> = gates.iter().map(|gate| gate.snapshot().active).collect();
release_tx.send(()).expect("owned work survives cancellation");
deadline(async {
while gates.iter().any(|gate| gate.snapshot().active != 0) {
tokio::task::yield_now().await;
}
}),
));
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");
})
.await;
assert!(inherited, "mode {mode} lost resource ownership");
assert!(original_policy, "mode {mode} changed the detached worker's disk policy");
assert_eq!(held, vec![1; physical_count], "mode {mode} released a live worker's unit");
}
}
}
fn config(pending: usize, workers: usize, blocking_threads: usize) -> Config {
+689 -32
View File
@@ -40,6 +40,7 @@ use std::sync::OnceLock;
use std::sync::atomic::AtomicBool;
use std::{
collections::HashMap,
io,
path::PathBuf,
pin::Pin,
sync::{
@@ -1493,6 +1494,8 @@ pub struct LocalDiskWrapper {
metrics: Arc<DiskHealthMetricEpoch>,
/// Per-drive foreground/background admission shared across wrapper reconnects.
io_admission: Arc<DiskIoAdmission>,
/// Physical leaves are refreshed by the node sampler, never by object I/O.
physical_admission: Arc<crate::disk::io_admission::PhysicalDiskBinding>,
/// Whether health checking is enabled
health_check: bool,
/// Cancellation token for monitoring tasks
@@ -1601,6 +1604,7 @@ impl LocalDiskWrapper {
health,
metrics,
io_admission,
physical_admission: Arc::new(crate::disk::io_admission::PhysicalDiskBinding::default()),
health_check: health_check && env_health_check,
cancel_token: CancellationToken::new(),
disk_id: Arc::new(RwLock::new(None)),
@@ -1618,6 +1622,10 @@ impl LocalDiskWrapper {
}
}
pub(crate) fn physical_admission(&self) -> Arc<crate::disk::io_admission::PhysicalDiskBinding> {
self.physical_admission.clone()
}
pub fn get_disk(&self) -> Arc<LocalDisk> {
self.disk.clone()
}
@@ -1659,6 +1667,33 @@ impl LocalDiskWrapper {
context
}
// Projection cannot wait: callers may already hold namespace or cache locks.
// Enter before identity probes, disk permits, or owned native jobs are created.
async fn with_native_resources<F: std::future::Future>(&self, operation: F) -> Result<F::Output> {
let Some(resources) = current_resource_unit() else {
if super::remote_io_admission::active() {
return Err(DiskError::Io(io::Error::from(io::ErrorKind::WouldBlock)));
}
return Ok(operation.await);
};
let resources = if let Some(permits) = resources.as_set() {
let unavailable = || DiskError::Io(io::Error::from(io::ErrorKind::WouldBlock));
let plan = self.physical_admission.plan().ok_or_else(unavailable)?;
let native = plan
.capture_or_acquire_native(permits)
.map_err(|error| DiskError::Io(io::Error::new(io::ErrorKind::WouldBlock, error)))?;
if !plan.is_current() {
return Err(unavailable());
}
resources.with_permits(native)
} else {
resources
};
#[cfg(target_os = "linux")]
let resources = resources.with_mount_lease(self.disk.native_mount_lease());
Ok(with_optional_resource_unit(Some(resources), operation).await)
}
async fn acquire_disk_io_permit(&self) -> Result<Option<crate::disk::io_admission::DiskIoPermit>> {
if let Some(context) = current_disk_io_context() {
if let Some(permit) = context.permit
@@ -2069,6 +2104,19 @@ impl LocalDiskWrapper {
versions: Vec<FileInfoVersions>,
opts: DeleteOptions,
namespace_owner: Option<Arc<dyn Send + Sync>>,
) -> Vec<Option<Error>> {
let count = versions.len();
self.with_native_resources(self.delete_versions_admitted(volume, versions, opts, namespace_owner))
.await
.unwrap_or_else(|error| vec![Some(error); count])
}
async fn delete_versions_admitted(
&self,
volume: &str,
versions: Vec<FileInfoVersions>,
opts: DeleteOptions,
namespace_owner: Option<Arc<dyn Send + Sync>>,
) -> Vec<Option<Error>> {
let version_count = versions.len();
self.metrics.record_operation_call("delete_versions");
@@ -2121,6 +2169,19 @@ impl LocalDiskWrapper {
versions: Vec<GuardedFileInfoVersions>,
opts: DeleteOptions,
namespace_owner: Option<Arc<dyn Send + Sync>>,
) -> Vec<Option<Error>> {
let count = versions.len();
self.with_native_resources(self.delete_versions_guarded_admitted(volume, versions, opts, namespace_owner))
.await
.unwrap_or_else(|error| vec![Some(error); count])
}
async fn delete_versions_guarded_admitted(
&self,
volume: &str,
versions: Vec<GuardedFileInfoVersions>,
opts: DeleteOptions,
namespace_owner: Option<Arc<dyn Send + Sync>>,
) -> Vec<Option<Error>> {
let version_count = versions.len();
self.metrics.record_operation_call("delete_versions");
@@ -2319,6 +2380,28 @@ impl LocalDiskWrapper {
timeout_health_action: TimeoutHealthAction,
mutation: DiskMetricMutation,
) -> Result<T>
where
F: FnOnce(Option<crate::disk::io_admission::DiskIoPermit>) -> Fut,
Fut: std::future::Future<Output = Result<T>>,
{
self.with_native_resources(self.track_admitted_disk_operation(
op,
operation,
timeout_duration,
timeout_health_action,
mutation,
))
.await?
}
async fn track_admitted_disk_operation<T, F, Fut>(
&self,
op: &'static str,
operation: F,
timeout_duration: Duration,
timeout_health_action: TimeoutHealthAction,
mutation: DiskMetricMutation,
) -> Result<T>
where
F: FnOnce(Option<crate::disk::io_admission::DiskIoPermit>) -> Fut,
Fut: std::future::Future<Output = Result<T>>,
@@ -2358,6 +2441,7 @@ impl LocalDiskWrapper {
self.metrics.record_operation_latency(op, started.elapsed());
self.record_result_error_metrics(&result);
if result.is_ok() {
super::remote_io_admission::record_progress();
self.health.record_operation_success(&self.endpoint(), "operation_success");
self.metrics.record_mutation_success(mutation);
}
@@ -2372,6 +2456,7 @@ impl LocalDiskWrapper {
self.record_result_error_metrics(&operation_result);
// Log success; the waiting guard balances every exit path.
if operation_result.is_ok() {
super::remote_io_admission::record_progress();
self.health.record_operation_success(&self.endpoint(), "operation_success");
self.metrics.record_mutation_success(mutation);
}
@@ -2538,7 +2623,7 @@ impl DiskAPI for LocalDiskWrapper {
}
async fn is_online(&self) -> bool {
let Ok(Some(disk_id)) = self.disk.get_disk_id().await else {
let Ok(Some(disk_id)) = self.get_disk_id().await else {
return false;
};
@@ -2571,7 +2656,7 @@ impl DiskAPI for LocalDiskWrapper {
}
async fn get_disk_id(&self) -> Result<Option<Uuid>> {
self.disk.get_disk_id().await
self.with_native_resources(self.disk.get_disk_id()).await?
}
async fn set_disk_id(&self, id: Option<Uuid>) -> Result<()> {
@@ -2929,25 +3014,29 @@ impl DiskAPI for LocalDiskWrapper {
}
async fn read_file(&self, volume: &str, path: &str) -> Result<crate::disk::FileReader> {
let reader = self
.track_disk_health_with_op(
"read_file",
|| async { self.disk.read_file(volume, path).await },
get_max_timeout_duration(),
)
.await?;
Ok(Box::new(AdmissionReader::new(reader, self.io_admission.clone(), self.disk_io_context())))
self.track_disk_health_with_op(
"read_file",
|| async {
let reader = self.disk.read_file(volume, path).await?;
Ok(Box::new(AdmissionReader::new(reader, self.io_admission.clone(), self.disk_io_context()))
as crate::disk::FileReader)
},
get_max_timeout_duration(),
)
.await
}
async fn read_file_stream(&self, volume: &str, path: &str, offset: usize, length: usize) -> Result<crate::disk::FileReader> {
let reader = self
.track_disk_health_with_op(
"read_file_stream",
|| async { self.disk.read_file_stream(volume, path, offset, length).await },
get_max_timeout_duration(),
)
.await?;
Ok(Box::new(AdmissionReader::new(reader, self.io_admission.clone(), self.disk_io_context())))
self.track_disk_health_with_op(
"read_file_stream",
|| async {
let reader = self.disk.read_file_stream(volume, path, offset, length).await?;
Ok(Box::new(AdmissionReader::new(reader, self.io_admission.clone(), self.disk_io_context()))
as crate::disk::FileReader)
},
get_max_timeout_duration(),
)
.await
}
async fn read_file_stream_chunks(
@@ -2957,20 +3046,21 @@ impl DiskAPI for LocalDiskWrapper {
offset: usize,
length: usize,
) -> Result<Option<rustfs_rio::ChunkReaderBox>> {
let reader = self
.track_disk_health_with_op(
"read_file_stream_chunks",
|| async { self.disk.read_file_stream_chunks(volume, path, offset, length).await },
get_max_timeout_duration(),
)
.await?;
Ok(reader.map(|reader| {
Box::new(AdmissionReader::new(
BoxedChunkReader(reader),
self.io_admission.clone(),
self.disk_io_context(),
)) as rustfs_rio::ChunkReaderBox
}))
self.track_disk_health_with_op(
"read_file_stream_chunks",
|| async {
let reader = self.disk.read_file_stream_chunks(volume, path, offset, length).await?;
Ok(reader.map(|reader| {
Box::new(AdmissionReader::new(
BoxedChunkReader(reader),
self.io_admission.clone(),
self.disk_io_context(),
)) as rustfs_rio::ChunkReaderBox
}))
},
get_max_timeout_duration(),
)
.await
}
async fn read_file_mmap_copy(&self, volume: &str, path: &str, offset: usize, length: usize) -> Result<bytes::Bytes> {
@@ -3831,6 +3921,573 @@ mod tests {
);
}
#[tokio::test]
async fn local_disk_resource_identity_probes_keep_only_their_native_device_after_cancel() {
assert_native_identity_probe_cancellation(false).await;
}
#[tokio::test]
async fn remote_background_native_identity_cancellation_keeps_receiver_device() {
assert_native_identity_probe_cancellation(true).await;
}
async fn assert_native_identity_probe_cancellation(remote: bool) {
use rustfs_concurrency::adaptive::{
AdaptiveAdmission, FeedbackController, FeedbackObservation, NodeAdmission, ResourcePressure,
permit_set::AdaptivePermitSetOwner,
};
use std::num::NonZeroU16;
for read in [false, true] {
for online in [false, true] {
let dir = tempfile::tempdir().expect("format probe fixture");
let mut endpoint = Endpoint::try_from(dir.path().to_str().expect("fixture path")).expect("local endpoint");
endpoint.set_pool_index(0);
endpoint.set_set_index(0);
endpoint.set_disk_index(0);
let meta = dir.path().join(RUSTFS_META_BUCKET);
tokio::fs::create_dir_all(&meta).await.expect("metadata directory");
let mut format = FormatV3::new(1, 1);
format.erasure.this = format.erasure.sets[0][0];
tokio::fs::write(meta.join(FORMAT_CONFIG_FILE), format.to_json().expect("format bytes"))
.await
.expect("write format");
let disk = Arc::new(LocalDisk::new(&endpoint, false).await.expect("formatted disk"));
{
let mut cache = disk.format_info.write().await;
cache.last_check = None;
if read {
cache.file_info = None;
}
}
let (started_tx, started_rx) = tokio::sync::oneshot::channel();
let (release_tx, release_rx) = tokio::sync::oneshot::channel();
disk.install_format_probe_hook(read, move || {
let _ = started_tx.send(());
let _ = release_rx.blocking_recv();
});
let wrapper = Arc::new(LocalDiskWrapper::new(disk, false));
let a = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let b = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let shared = AdaptivePermitSetOwner::new(NonZeroU16::MIN, vec![a.clone(), b.clone()]).expect("catalog");
wrapper
.physical_admission
.set_test_shared_capacity(shared.clone(), vec![a.clone()])
.expect("native device A");
let node = NodeAdmission::new(NonZeroU16::MIN);
node.update(
&mut FeedbackController::new(),
Some(FeedbackObservation {
generation: 1,
max_concurrency: 1,
pressure: ResourcePressure::Unknown,
completed_units: None,
}),
);
let parent = (!remote).then(|| shared.try_acquire_with_node(&node).expect("planned devices"));
let task = tokio::spawn({
let wrapper = wrapper.clone();
let node = node.clone();
async move {
let operation = async {
if online {
assert!(wrapper.is_online().await);
} else {
wrapper.get_disk_id().await.expect("disk identity");
}
};
if remote {
crate::disk::with_remote_background_io(Some(crate::disk::RemoteBackgroundIo::new(node)), async {
crate::disk::remote_io_admission::admit(&wrapper.physical_admission).expect("RPC disk admission");
operation.await;
})
.await;
} else {
let parent = parent.expect("local operation admission");
crate::disk::with_background_resource_set(parent.resources().clone(), operation).await;
parent.complete();
}
}
});
tokio::time::timeout(Duration::from_secs(5), started_rx)
.await
.expect("native probe starts")
.expect("native signal");
task.abort();
assert!(task.await.expect_err("cancel waiter").is_cancelled());
let held = [node.snapshot().active, a.snapshot().active, b.snapshot().active];
let healthy_available = b.try_acquire().is_some();
release_tx.send(()).expect("release native probe before assertions");
tokio::time::timeout(Duration::from_secs(5), async {
while a.snapshot().active != 0 || b.snapshot().active != 0 {
tokio::task::yield_now().await;
}
})
.await
.expect("native probe releases its device");
assert_eq!(held, [0, 1, 0], "read={read}, online={online}");
assert!(healthy_available);
assert_eq!([node.completed(), a.completed(), b.completed()], [Some(0), Some(0), Some(0)]);
assert_eq!(wrapper.get_disk_id().await.expect("format remains readable"), Some(format.erasure.this));
assert!(wrapper.is_online().await);
}
}
}
#[tokio::test]
async fn local_disk_resource_projection_releases_unrelated_devices_after_native_cancellation() {
use crate::disk::io_admission::spawn_blocking_with_disk_io_permit;
use rustfs_concurrency::adaptive::{
AdaptiveAdmission, FeedbackController, FeedbackObservation, NodeAdmission, ResourcePressure,
permit_set::AdaptivePermitSetOwner,
};
use std::num::NonZeroU16;
for auxiliary in [false, true] {
let dir = tempfile::tempdir().expect("local disk fixture");
let endpoint = Endpoint::try_from(dir.path().to_str().expect("fixture path")).expect("local endpoint");
let disk = Arc::new(LocalDisk::new(&endpoint, false).await.expect("local disk"));
let wrapper = LocalDiskWrapper::new(disk, false);
let first = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let second = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let shared = AdaptivePermitSetOwner::new(NonZeroU16::MIN, vec![first.clone(), second.clone()])
.expect("shared operation resources");
wrapper
.physical_admission
.set_test_shared_capacity(shared.clone(), vec![if auxiliary { second.clone() } else { first.clone() }])
.expect("only the first device backs this wrapper");
let node = NodeAdmission::new(NonZeroU16::MIN);
node.update(
&mut FeedbackController::new(),
Some(FeedbackObservation {
generation: 1,
max_concurrency: 1,
pressure: ResourcePressure::Unknown,
completed_units: None,
}),
);
let parent = if auxiliary {
shared
.select(std::slice::from_ref(&first))
.expect("initial object device")
.try_acquire_with_node(&node)
} else {
shared.try_acquire_with_node(&node)
}
.expect("initial object admission");
let (started_tx, started_rx) = tokio::sync::oneshot::channel();
let (release_tx, release_rx) = tokio::sync::oneshot::channel();
let task = tokio::spawn(async move {
crate::disk::with_background_resource_set(parent.resources().clone(), async {
wrapper
.track_disk_health_with_op_timeout_action_and_mutation_with_io_permit(
"read_all",
move |permit| async move {
spawn_blocking_with_disk_io_permit(permit, move || {
started_tx.send(()).expect("signal native start");
release_rx.blocking_recv().expect("release native I/O");
})
.await
.expect("native I/O joined");
Ok(())
},
Duration::ZERO,
TimeoutHealthAction::MarkFailure,
DiskMetricMutation::None,
)
.await
.expect("wrapped I/O");
})
.await;
parent.complete();
});
tokio::time::timeout(Duration::from_secs(5), started_rx)
.await
.expect("native work starts")
.expect("start notification");
task.abort();
assert!(task.await.expect_err("cancel waiter").is_cancelled());
let held = [node.snapshot().active, first.snapshot().active, second.snapshot().active];
let healthy_available = if auxiliary {
first.try_acquire().is_some()
} else {
second.try_acquire().is_some()
};
release_tx.send(()).expect("finish native work before asserting");
tokio::time::timeout(Duration::from_secs(5), async {
while first.snapshot().active != 0 || second.snapshot().active != 0 {
tokio::task::yield_now().await;
}
})
.await
.expect("native resource return");
assert_eq!(
held,
if auxiliary { [0, 0, 1] } else { [0, 1, 0] },
"only the actual native device remains charged"
);
assert!(healthy_available, "healthy work must resume before native I/O finishes");
assert_eq!(node.completed(), Some(0), "cancelled work is not progress");
assert_eq!([first.completed(), second.completed()], [Some(0), Some(0)]);
}
}
#[tokio::test]
async fn local_disk_resource_projection_escaped_readers_keep_only_their_device() {
use rustfs_concurrency::adaptive::{AdaptiveAdmission, permit_set::AdaptivePermitSetOwner};
use std::num::NonZeroU16;
use tokio::io::AsyncReadExt;
let dir = tempfile::tempdir().expect("reader fixture");
let endpoint = Endpoint::try_from(dir.path().to_str().expect("fixture path")).expect("local endpoint");
let disk = Arc::new(LocalDisk::new(&endpoint, false).await.expect("local disk"));
let wrapper = LocalDiskWrapper::new(disk, false);
tokio::fs::create_dir_all(dir.path().join("bucket"))
.await
.expect("bucket directory");
let bytes = b"physical disk reader payload";
tokio::fs::write(dir.path().join("bucket/object"), bytes)
.await
.expect("reader bytes");
let first = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let second = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let shared = AdaptivePermitSetOwner::new(NonZeroU16::MIN, vec![first.clone(), second.clone()]).expect("parent resources");
wrapper
.physical_admission
.set_test_shared_capacity(shared.clone(), vec![first.clone()])
.expect("reader device");
for stream in [false, true] {
let a = first.try_acquire().expect("first device");
let b = second.try_acquire().expect("second device");
let parent = shared.capture(&[&a, &b]).expect("admitted parent");
drop((a, b));
let mut reader = crate::disk::with_background_resource_set(parent, async {
if stream {
wrapper.read_file_stream("bucket", "object", 0, bytes.len()).await
} else {
wrapper.read_file("bucket", "object").await
}
})
.await
.expect("reader escapes the resource scope");
assert_eq!([first.snapshot().active, second.snapshot().active], [1, 0]);
assert!(second.try_acquire().is_some(), "unrelated work is available while the reader lives");
let mut actual = Vec::new();
reader.read_to_end(&mut actual).await.expect("consume escaped reader");
assert_eq!(actual, bytes);
drop(reader);
assert_eq!(first.snapshot().active, 0);
}
}
#[cfg(unix)]
#[test]
fn local_disk_resource_projection_cancelled_stream_read_retains_native_device() {
use rustfs_concurrency::adaptive::{AdaptiveAdmission, permit_set::AdaptivePermitSetOwner};
use std::num::NonZeroU16;
use tokio::io::{AsyncRead, ReadBuf};
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.max_blocking_threads(1)
.build()
.expect("one-thread native read pool");
runtime.block_on(async {
let dir = tempfile::tempdir().expect("native reader fixture");
let endpoint = Endpoint::try_from(dir.path().to_str().expect("fixture path")).expect("local endpoint");
let disk = Arc::new(LocalDisk::new(&endpoint, false).await.expect("local disk"));
let wrapper = LocalDiskWrapper::new(disk, false);
std::fs::create_dir(dir.path().join("bucket")).expect("bucket directory");
let bytes = b"physical disk native reader payload";
std::fs::write(dir.path().join("bucket/object"), bytes).expect("reader bytes");
let first = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let second = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let shared =
AdaptivePermitSetOwner::new(NonZeroU16::MIN, vec![first.clone(), second.clone()]).expect("parent resources");
wrapper
.physical_admission
.set_test_shared_capacity(shared.clone(), vec![first.clone()])
.expect("reader device");
let mut observations = Vec::new();
for stream in [false, true] {
let a = first.try_acquire().expect("first device");
let b = second.try_acquire().expect("second device");
let parent = shared.capture(&[&a, &b]).expect("admitted parent");
drop((a, b));
let mut reader = crate::disk::with_background_resource_set(
parent,
crate::disk::with_background_disk_io(CancellationToken::new(), async {
if stream {
wrapper.read_file_stream("bucket", "object", 3, bytes.len() - 3).await
} else {
wrapper.read_file("bucket", "object").await
}
}),
)
.await
.expect("background reader escapes its parent");
assert_eq!([first.snapshot().active, second.snapshot().active], [1, 0]);
let (started_tx, started_rx) = tokio::sync::oneshot::channel();
let (release_tx, release_rx) = std::sync::mpsc::channel::<()>();
let blocker = tokio::task::spawn_blocking(move || {
let _ = started_tx.send(());
let _ = release_rx.recv();
});
started_rx.await.expect("native pool occupied after open");
let (pending_tx, pending_rx) = tokio::sync::oneshot::channel();
let task = tokio::spawn(async move {
let mut data = [0; 1];
let mut pending_tx = Some(pending_tx);
std::future::poll_fn(|cx| {
let mut buffer = ReadBuf::new(&mut data);
let result = Pin::new(&mut reader).poll_read(cx, &mut buffer);
if result.is_pending()
&& let Some(signal) = pending_tx.take()
{
let _ = signal.send(());
}
result
})
.await
});
tokio::time::timeout(Duration::from_secs(5), pending_rx)
.await
.expect("read queues behind native worker")
.expect("pending read signal");
let admitted = wrapper.background_io_active_for_tests();
task.abort();
assert!(task.await.expect_err("cancel pending read").is_cancelled());
let held = [first.snapshot().active, second.snapshot().active];
release_tx.send(()).expect("release native pool before asserting");
blocker.await.expect("blocking worker released");
tokio::task::spawn_blocking(|| ())
.await
.expect("drain the single native queue");
assert_eq!([first.snapshot().active, second.snapshot().active], [0, 0]);
observations.push((stream, admitted, held));
}
for (stream, admitted, held) in observations {
assert!(admitted, "stream={stream}: cancellation follows disk admission, not its wait");
assert_eq!(held, [1, 0], "stream={stream}: only the pending native read's device stays charged");
}
});
}
#[cfg(unix)]
#[tokio::test]
async fn projected_stream_memory_survives_scope_exit_and_hot_shrink() {
use crate::disk::io_admission::{BackgroundMemoryBudget, ResourceIoUnit, with_optional_resource_unit};
use rustfs_concurrency::adaptive::{AdaptiveAdmission, permit_set::AdaptivePermitSetOwner};
use std::num::NonZeroU16;
use tokio::io::AsyncReadExt;
let dir = tempfile::tempdir().expect("memory admission fixture");
let endpoint = Endpoint::try_from(dir.path().to_str().expect("fixture path")).expect("local endpoint");
let disk = Arc::new(LocalDisk::new(&endpoint, false).await.expect("local disk"));
let wrapper = LocalDiskWrapper::new(disk, false);
std::fs::create_dir(dir.path().join("bucket")).expect("bucket directory");
let bytes = b"0123456789abcdefghijklmnopqrstuv";
std::fs::write(dir.path().join("bucket/object"), bytes).expect("reader bytes");
let first = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let second = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let shared = AdaptivePermitSetOwner::new(NonZeroU16::MIN, vec![first.clone(), second.clone()]).expect("parent resources");
wrapper
.physical_admission
.set_test_shared_capacity(shared.clone(), vec![first.clone()])
.expect("selected device");
let memory = BackgroundMemoryBudget::new();
for stream in [false, true] {
memory.set_sampled_limit(4, Duration::from_secs(30));
let a = first.try_acquire().expect("first device");
let b = second.try_acquire().expect("second device");
let parent =
ResourceIoUnit::from(shared.capture(&[&a, &b]).expect("admitted parent")).with_memory_budget(memory.clone());
drop((a, b));
let mut reader = with_optional_resource_unit(Some(parent), async {
if stream {
wrapper.read_file_stream("bucket", "object", 0, bytes.len()).await
} else {
wrapper.read_file("bucket", "object").await
}
})
.await
.expect("reader escapes its resource scope");
assert_eq!(memory.snapshot().used, 0, "opening does not allocate a read buffer");
let mut prefix = [0; 4];
reader.read_exact(&mut prefix).await.expect("initial read");
assert_eq!(&prefix, &bytes[..4]);
assert_eq!(memory.snapshot().used, 4, "projection retains the same memory owner");
assert_eq!([first.snapshot().active, second.snapshot().active], [1, 0]);
memory.set_sampled_limit(0, Duration::ZERO);
let mut reused = [0; 2];
reader
.read_exact(&mut reused)
.await
.expect("already allocated storage remains usable");
assert_eq!(&reused, &bytes[4..6]);
let mut next = [99; 8];
let error = reader.read(&mut next).await.expect_err("new storage is blocked after shrink");
assert_eq!(error.kind(), std::io::ErrorKind::WouldBlock);
assert_eq!(next, [99; 8], "denied allocation exposes no stale bytes");
assert_eq!(memory.snapshot().used, 4);
memory.set_sampled_limit(12, Duration::from_secs(30));
reader
.read_exact(&mut next)
.await
.expect("growth includes old and replacement storage");
assert_eq!(&next, &bytes[6..14], "denial did not advance the file");
assert_eq!(memory.snapshot().used, 8);
drop(reader);
assert_eq!(memory.snapshot().used, 0);
assert_eq!([first.snapshot().active, second.snapshot().active], [0, 0]);
}
}
#[tokio::test]
async fn local_disk_resource_projection_auxiliary_metadata_retries_without_parking_or_false_progress() {
use rustfs_concurrency::adaptive::{
AdaptiveAdmission, FeedbackController, FeedbackObservation, NodeAdmission, ResourcePressure,
permit_set::AdaptivePermitSetOwner,
};
use std::num::NonZeroU16;
let dir = tempfile::tempdir().expect("metadata fixture");
let endpoint = Endpoint::try_from(dir.path().to_str().expect("fixture path")).expect("local endpoint");
let disk = Arc::new(LocalDisk::new(&endpoint, false).await.expect("local disk"));
let wrapper = LocalDiskWrapper::new(disk, false);
tokio::fs::create_dir_all(dir.path().join("bucket"))
.await
.expect("metadata bucket");
tokio::fs::write(dir.path().join("bucket/metadata"), b"auxiliary metadata")
.await
.expect("metadata bytes");
let a = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let b = AdaptiveAdmission::new(NonZeroU16::MIN, 0);
let shared = AdaptivePermitSetOwner::new(NonZeroU16::MIN, vec![a.clone(), b.clone()]).expect("topology");
wrapper
.physical_admission
.set_test_shared_capacity(shared.clone(), vec![b.clone()])
.expect("metadata on B");
let node = NodeAdmission::new(NonZeroU16::MIN);
node.update(
&mut FeedbackController::new(),
Some(FeedbackObservation {
generation: 1,
max_concurrency: 1,
pressure: ResourcePressure::Unknown,
completed_units: None,
}),
);
let parent = shared
.select(std::slice::from_ref(&a))
.expect("object on A")
.try_acquire_with_node(&node)
.expect("object admission");
crate::disk::with_background_resource_set(parent.resources().clone(), async {
let mut read = Box::pin(wrapper.read_all("bucket", "metadata"));
let rejected = std::future::poll_fn(|cx| Poll::Ready(read.as_mut().poll(cx))).await;
assert!(
matches!(rejected, Poll::Ready(Err(DiskError::Io(error))) if error.kind() == io::ErrorKind::WouldBlock),
"auxiliary admission must return immediately under storage locks"
);
drop(read);
assert!(
matches!(wrapper.get_disk_id().await, Err(DiskError::Io(error)) if error.kind() == io::ErrorKind::WouldBlock)
);
assert!(!wrapper.is_online().await, "identity checks cannot bypass closed physical capacity");
assert_eq!([a.snapshot().active, b.snapshot().active, node.snapshot().active], [1, 0, 1]);
assert_eq!([a.completed(), b.completed(), node.completed()], [Some(0), Some(0), Some(0)]);
b.set_limit(1);
for _ in 0..2 {
assert_eq!(
wrapper.read_all("bucket", "metadata").await.expect("retry metadata read"),
Bytes::from_static(b"auxiliary metadata")
);
assert_eq!(b.snapshot().active, 1, "repeated I/O reuses the parent's physical unit");
}
})
.await;
assert_eq!(wrapper.metrics_snapshot().total_errors_availability, 0);
assert_eq!([a.completed(), b.completed(), node.completed()], [Some(0), Some(0), Some(0)]);
parent.complete();
assert_eq!([a.completed(), b.completed(), node.completed()], [Some(1), Some(1), Some(1)]);
assert_eq!([a.snapshot().active, b.snapshot().active, node.snapshot().active], [0, 0, 0]);
}
#[tokio::test]
async fn local_disk_resource_projection_defers_missing_or_foreign_owners_before_io() {
use rustfs_concurrency::adaptive::{AdaptiveAdmission, permit_set::AdaptivePermitSetOwner};
use std::num::NonZeroU16;
let dir = tempfile::tempdir().expect("unavailable resource fixture");
let endpoint = Endpoint::try_from(dir.path().to_str().expect("fixture path")).expect("local endpoint");
let disk = Arc::new(LocalDisk::new(&endpoint, false).await.expect("local disk"));
let wrapper = LocalDiskWrapper::new(disk, false);
let gate = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let shared = AdaptivePermitSetOwner::new(NonZeroU16::MIN, vec![gate.clone()]).expect("parent");
let permit = gate.try_acquire().expect("parent device");
let parent = shared.capture(&[&permit]).expect("parent ownership");
drop(permit);
for foreign in [false, true] {
if foreign {
wrapper
.physical_admission
.set_test_capacity(vec![AdaptiveAdmission::new(NonZeroU16::MIN, 1)])
.expect("replaced device");
}
crate::disk::with_background_resource_set(parent.clone(), async {
let polled = std::cell::Cell::new(false);
let result = wrapper
.track_disk_health_with_op(
"read_all",
|| async {
polled.set(true);
Ok(())
},
Duration::ZERO,
)
.await;
assert!(matches!(result, Err(DiskError::Io(error)) if error.kind() == io::ErrorKind::WouldBlock));
assert!(!polled.get(), "unavailable ownership cannot start native work");
let group = FileInfoVersions::default();
let plain = wrapper
.delete_versions("bucket", vec![group.clone(); 2], DeleteOptions::default())
.await;
let guarded = wrapper
.delete_versions_guarded_with_namespace_owner(
"bucket",
vec![
GuardedFileInfoVersions {
group,
expected_current: crate::disk::CurrentDeleteGeneration::Absent
};
2
],
DeleteOptions::default(),
None,
)
.await;
for errors in [plain, guarded] {
assert_eq!(errors.len(), 2);
assert!(errors.iter().all(|error| matches!(error,
Some(DiskError::Io(error)) if error.kind() == io::ErrorKind::WouldBlock)));
}
})
.await;
assert_eq!(wrapper.health.waiting_count(), 0);
assert_eq!(
wrapper.metrics_snapshot().total_errors_availability,
0,
"resource deferral is not a failed drive"
);
assert_eq!(gate.snapshot().active, 1, "rejection does not refund parent ownership");
}
}
struct PendingWriter;
#[test]
+132 -62
View File
@@ -16,40 +16,15 @@ use std::{
collections::HashMap,
fs::Metadata,
path::{Path, PathBuf},
sync::{Arc, Mutex, OnceLock},
sync::Mutex,
time::{Duration, Instant},
};
use tokio::{
fs::{self, File},
io,
};
use tokio::{fs::File, io};
static READONLY_OPTIONS: OnceLock<Arc<fs::OpenOptions>> = OnceLock::new();
static WRITEONLY_OPTIONS: OnceLock<Arc<fs::OpenOptions>> = OnceLock::new();
static READWRITE_OPTIONS: OnceLock<Arc<fs::OpenOptions>> = OnceLock::new();
fn get_readonly_options() -> &'static Arc<fs::OpenOptions> {
READONLY_OPTIONS.get_or_init(|| {
let mut opts = fs::OpenOptions::new();
opts.read(true);
Arc::new(opts)
})
}
fn get_writeonly_options() -> &'static Arc<fs::OpenOptions> {
WRITEONLY_OPTIONS.get_or_init(|| {
let mut opts = fs::OpenOptions::new();
opts.write(true);
Arc::new(opts)
})
}
fn get_readwrite_options() -> &'static Arc<fs::OpenOptions> {
READWRITE_OPTIONS.get_or_init(|| {
let mut opts = fs::OpenOptions::new();
opts.read(true).write(true);
Arc::new(opts)
})
async fn native_io<T: Send + 'static>(operation: impl FnOnce() -> io::Result<T> + Send + 'static) -> io::Result<T> {
super::io_admission::spawn_blocking_with_resource_ownership(operation)
.await
.map_err(|_| io::Error::other("background task failed"))?
}
#[cfg(not(windows))]
@@ -117,33 +92,32 @@ pub const O_APPEND: FileMode = 0x00400;
// create_new: bool,
pub async fn open_file(path: impl AsRef<Path>, mode: FileMode) -> io::Result<File> {
let base_opts = match mode & (O_RDONLY | O_WRONLY | O_RDWR) {
O_RDONLY => get_readonly_options(),
O_WRONLY => get_writeonly_options(),
O_RDWR => get_readwrite_options(),
_ => get_readonly_options(),
};
if (mode & (O_CREATE | O_APPEND | O_TRUNC)) != 0 {
let mut opts = (**base_opts).clone();
if mode & O_CREATE != 0 {
opts.create(true);
let path = path.as_ref().to_path_buf();
native_io(move || {
let mut options = std::fs::OpenOptions::new();
match mode & (O_RDONLY | O_WRONLY | O_RDWR) {
O_WRONLY => options.write(true),
O_RDWR => options.read(true).write(true),
_ => options.read(true),
}
if mode & O_APPEND != 0 {
opts.append(true);
.create(mode & O_CREATE != 0)
.append(mode & O_APPEND != 0)
.truncate(mode & O_TRUNC != 0);
#[cfg(all(any(test, feature = "test-util"), not(windows)))]
if mode == O_RDONLY {
super::os::prepared_publication_test_hooks::run(
super::os::prepared_publication_test_hooks::Stage::NativeReadOpen,
&path,
);
}
if mode & O_TRUNC != 0 {
opts.truncate(true);
}
opts.open(path.as_ref()).await
} else {
base_opts.open(path.as_ref()).await
}
options.open(path).map(File::from_std)
})
.await
}
pub async fn access(path: impl AsRef<Path>) -> io::Result<()> {
fs::metadata(path).await?;
Ok(())
let path = path.as_ref().to_path_buf();
native_io(move || std::fs::metadata(path).map(|_| ())).await
}
pub fn access_std(path: impl AsRef<Path>) -> io::Result<()> {
@@ -152,7 +126,8 @@ pub fn access_std(path: impl AsRef<Path>) -> io::Result<()> {
}
pub async fn lstat(path: impl AsRef<Path>) -> io::Result<Metadata> {
fs::symlink_metadata(path).await
let path = path.as_ref().to_path_buf();
native_io(move || std::fs::symlink_metadata(path)).await
}
pub fn lstat_std(path: impl AsRef<Path>) -> io::Result<Metadata> {
@@ -160,7 +135,8 @@ pub fn lstat_std(path: impl AsRef<Path>) -> io::Result<Metadata> {
}
pub async fn make_dir_all(path: impl AsRef<Path>) -> io::Result<()> {
fs::create_dir_all(path.as_ref()).await
let path = path.as_ref().to_path_buf();
native_io(move || std::fs::create_dir_all(path)).await
}
fn is_dir_error(e: &io::Error) -> bool {
@@ -173,9 +149,13 @@ fn is_dir_error(e: &io::Error) -> bool {
#[tracing::instrument(level = "debug", skip_all)]
pub async fn remove(path: impl AsRef<Path>) -> io::Result<()> {
// Try remove_file first; fall back to remove_dir if it's a directory
match fs::remove_file(path.as_ref()).await {
let file_path = path.as_ref().to_path_buf();
match native_io(move || std::fs::remove_file(file_path)).await {
Ok(()) => Ok(()),
Err(e) if is_dir_error(&e) => fs::remove_dir(path.as_ref()).await,
Err(e) if is_dir_error(&e) => {
let path = path.as_ref().to_path_buf();
native_io(move || std::fs::remove_dir(path)).await
}
Err(e) => Err(e),
}
}
@@ -183,9 +163,13 @@ pub async fn remove(path: impl AsRef<Path>) -> io::Result<()> {
#[allow(dead_code, reason = "asserted by this file's tests (backlog#1823)")]
pub async fn remove_all(path: impl AsRef<Path>) -> io::Result<()> {
// Try remove_file first; fall back to remove_dir_all if it's a directory
match fs::remove_file(path.as_ref()).await {
let file_path = path.as_ref().to_path_buf();
match native_io(move || std::fs::remove_file(file_path)).await {
Ok(()) => Ok(()),
Err(e) if is_dir_error(&e) => fs::remove_dir_all(path.as_ref()).await,
Err(e) if is_dir_error(&e) => {
let path = path.as_ref().to_path_buf();
native_io(move || std::fs::remove_dir_all(path)).await
}
Err(e) => Err(e),
}
}
@@ -211,11 +195,14 @@ pub fn remove_all_std(path: impl AsRef<Path>) -> io::Result<()> {
}
pub async fn mkdir(path: impl AsRef<Path>) -> io::Result<()> {
fs::create_dir(path.as_ref()).await
let path = path.as_ref().to_path_buf();
native_io(move || std::fs::create_dir(path)).await
}
pub async fn rename(from: impl AsRef<Path>, to: impl AsRef<Path>) -> io::Result<()> {
fs::rename(from, to).await
let from = from.as_ref().to_path_buf();
let to = to.as_ref().to_path_buf();
native_io(move || std::fs::rename(from, to)).await
}
#[cfg(any(not(windows), test))]
@@ -277,7 +264,7 @@ pub async fn cached_access(path: impl AsRef<Path>) -> io::Result<()> {
return Err(io::Error::new(io::ErrorKind::NotFound, "bucket not found (cached)"));
}
let result = fs::metadata(&path_buf).await;
let result = access(&path_buf).await;
match &result {
Ok(_) => BUCKET_EXISTENCE_CACHE.record(path_buf, true),
@@ -301,6 +288,89 @@ mod tests {
use tempfile::TempDir;
use tokio::io::AsyncWriteExt;
#[test]
fn cancelled_filesystem_jobs_retain_physical_resources() {
use crate::disk::io_admission::{test_resource_units, with_optional_resource_unit};
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.max_blocking_threads(1)
.build()
.expect("single native worker");
runtime.block_on(async {
let mut observations = Vec::new();
for count in 1..=2 {
for operation in [
"access",
"cached_access",
"lstat",
"mkdir",
"make_dir_all",
"remove",
"remove_all",
"rename",
"open",
] {
let dir = TempDir::new().expect("filesystem ownership fixture");
let source = dir.path().join("source");
let target = dir.path().join("target");
match operation {
"mkdir" | "make_dir_all" | "open" => {}
"remove_all" => {
std::fs::create_dir(&source).expect("recursive removal directory");
std::fs::write(source.join("child"), b"payload").expect("recursive removal child");
}
_ => std::fs::write(&source, b"payload").expect("existing file"),
}
let (resources, gates) = test_resource_units(count);
let (started_tx, started_rx) = tokio::sync::oneshot::channel();
let (release_tx, release_rx) = std::sync::mpsc::channel();
let blocker = tokio::task::spawn_blocking(move || {
let _ = started_tx.send(());
let _ = release_rx.recv();
});
started_rx.await.expect("native worker occupied");
let mut future = Box::pin(with_optional_resource_unit(Some(resources), async {
match operation {
"access" => access(&source).await,
"cached_access" => cached_access(&source).await,
"lstat" => lstat(&source).await.map(|_| ()),
"mkdir" => mkdir(&source).await,
"make_dir_all" => make_dir_all(source.join("child")).await,
"remove" => remove(&source).await,
"remove_all" => remove_all(&source).await,
"rename" => rename(&source, &target).await,
"open" => open_file(&source, O_CREATE | O_WRONLY).await.map(|_| ()),
_ => unreachable!("fixed operation list"),
}
}));
assert!(futures::poll!(&mut future).is_pending(), "{operation}: native job queued");
drop(future);
let held: Vec<_> = gates.iter().map(|gate| gate.snapshot().active).collect();
release_tx.send(()).expect("release native worker before assertions");
blocker.await.expect("native blocker completed");
tokio::task::spawn_blocking(|| ()).await.expect("native queue drained");
assert!(gates.iter().all(|gate| gate.snapshot().active == 0));
match operation {
"mkdir" => assert!(source.is_dir()),
"make_dir_all" => assert!(source.join("child").is_dir()),
"open" => assert!(source.is_file()),
"remove" => assert!(!source.exists()),
"rename" => {
assert!(!source.exists());
assert_eq!(std::fs::read(&target).expect("renamed file"), b"payload");
}
_ => {}
}
observations.push((count, operation, held));
}
}
for (count, operation, held) in observations {
assert_eq!(held, vec![1; count], "{operation}: queued native job keeps its physical resources");
}
});
}
#[tokio::test]
async fn test_file_mode_constants() {
assert_eq!(O_RDONLY, 0x00000);
+218 -15
View File
@@ -14,6 +14,18 @@
//! Per-local-disk foreground-priority admission for scanner I/O.
pub(crate) mod memory;
mod physical_identity;
mod physical_owners;
mod physical_sampler;
pub use memory::BackgroundMemoryBudget;
pub(crate) use physical_sampler::PhysicalHealBinding;
pub use physical_sampler::{
PhysicalDiskBinding, PhysicalDiskPlan, PhysicalDiskSample, PhysicalDiskSampler, PhysicalResourcePlan,
};
use std::collections::VecDeque;
use std::future::Future;
use std::io;
@@ -25,18 +37,118 @@ use std::time::Duration;
use metrics::{counter, gauge, histogram};
use rustfs_concurrency::adaptive::AdaptivePermit;
use rustfs_concurrency::adaptive::permit_set::AdaptivePermitSet;
use rustfs_rio::ChunkReader;
use tokio::io::{AsyncRead, ReadBuf};
use tokio::sync::Notify;
use tokio::time::Instant;
use tokio_util::sync::CancellationToken;
pub(crate) mod remote;
pub use remote::{RemoteBackgroundIo, with_remote_background_io};
const BACKGROUND_FOREGROUND_BURST: u32 = 8;
const BACKGROUND_MAX_WAIT: Duration = Duration::from_millis(100);
tokio::task_local! {
static DISK_IO_CONTEXT: DiskIoContext;
static RESOURCE_IO_UNIT: AdaptivePermit;
static RESOURCE_IO_UNIT: ResourceIoUnit;
}
/// Native I/O retains the physical subset selected before storage work. A
/// layered disk may touch several leaves; unrelated devices stay outside it.
#[derive(Clone, Debug)]
pub(crate) struct ResourceIoUnit {
permits: ResourceIoPermits,
memory_budget: Option<Arc<BackgroundMemoryBudget>>,
#[cfg(target_os = "linux")]
_mount_lease: Option<Arc<std::fs::File>>,
}
#[derive(Debug)]
enum ResourceIoPermits {
Single(AdaptivePermit),
Set(AdaptivePermitSet),
}
impl Clone for ResourceIoPermits {
fn clone(&self) -> Self {
match self {
Self::Single(permit) => Self::Single(permit.clone()),
Self::Set(permits) => Self::Set(permits.clone()),
}
}
}
impl From<AdaptivePermit> for ResourceIoUnit {
fn from(permit: AdaptivePermit) -> Self {
Self {
permits: ResourceIoPermits::Single(permit),
memory_budget: BackgroundMemoryBudget::current(),
#[cfg(target_os = "linux")]
_mount_lease: None,
}
}
}
impl From<AdaptivePermitSet> for ResourceIoUnit {
fn from(permits: AdaptivePermitSet) -> Self {
Self {
permits: ResourceIoPermits::Set(permits),
memory_budget: BackgroundMemoryBudget::current(),
#[cfg(target_os = "linux")]
_mount_lease: None,
}
}
}
impl ResourceIoUnit {
pub(crate) fn memory_budget(&self) -> Option<Arc<BackgroundMemoryBudget>> {
self.memory_budget.clone()
}
pub(crate) fn with_permits(mut self, permits: AdaptivePermitSet) -> Self {
self.permits = ResourceIoPermits::Set(permits);
self
}
#[cfg(test)]
pub(crate) fn with_memory_budget(mut self, budget: Arc<BackgroundMemoryBudget>) -> Self {
self.memory_budget = Some(budget);
self
}
pub(crate) fn as_set(&self) -> Option<&AdaptivePermitSet> {
match &self.permits {
ResourceIoPermits::Single(_) => None,
ResourceIoPermits::Set(permits) => Some(permits),
}
}
#[cfg(target_os = "linux")]
pub(crate) fn with_mount_lease(mut self, mount_lease: Arc<std::fs::File>) -> Self {
// The selected disk's proc-fd paths must survive native jobs after its
// async wrapper retires. Clones carry the same pinned mount generation.
self._mount_lease = Some(mount_lease);
self
}
}
#[cfg(test)]
pub(crate) fn test_resource_units(count: usize) -> (ResourceIoUnit, Vec<Arc<rustfs_concurrency::adaptive::AdaptiveAdmission>>) {
use rustfs_concurrency::adaptive::{AdaptiveAdmission, permit_set::AdaptivePermitSetOwner};
use std::num::NonZeroU16;
assert!((1..=2).contains(&count), "test single and layered disks");
let gates: Vec<_> = (0..count).map(|_| AdaptiveAdmission::new(NonZeroU16::MIN, 1)).collect();
let mut admitted: Vec<_> = gates.iter().map(|gate| gate.try_acquire().expect("test resource")).collect();
let unit = if count == 1 {
ResourceIoUnit::from(admitted.pop().expect("single resource"))
} else {
let owner = AdaptivePermitSetOwner::new(NonZeroU16::MIN, gates.clone()).expect("layered resource");
ResourceIoUnit::from(owner.capture(&admitted.iter().collect::<Vec<_>>()).expect("admitted subset"))
};
(unit, gates)
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
@@ -51,7 +163,7 @@ pub(crate) struct DiskIoContext {
pub(crate) cancellation: CancellationToken,
pub(crate) admission: Option<Arc<DiskIoAdmission>>,
pub(crate) permit: Option<DiskIoPermit>,
pub(crate) resource_unit: Option<AdaptivePermit>,
pub(crate) resource_unit: Option<ResourceIoUnit>,
}
/// Mark disk I/O issued by `future` as background work for local-disk admission.
@@ -78,24 +190,46 @@ pub async fn with_background_disk_io<F: Future>(cancellation: CancellationToken,
/// of a stalled native operation. Escaping buffers need their own memory leases.
#[doc(hidden)]
pub async fn with_background_resource_unit<F: Future>(unit: AdaptivePermit, future: F) -> F::Output {
RESOURCE_IO_UNIT.scope(unit, future).await
RESOURCE_IO_UNIT.scope(ResourceIoUnit::from(unit), future).await
}
pub(crate) fn current_resource_unit() -> Option<AdaptivePermit> {
RESOURCE_IO_UNIT.try_with(Clone::clone).ok().or_else(|| {
DISK_IO_CONTEXT
.try_with(|context| context.resource_unit.clone())
.ok()
.flatten()
})
/// Retain only the previously admitted physical leaves touched by this I/O.
/// Projection of owned units never waits and may run under storage locks.
/// Acquire the planned resources before locking; auxiliary leaves discovered
/// later must use nonblocking admission. Do not pass unrelated parent resources
/// to a native job that touches only some of its disks.
#[doc(hidden)]
pub async fn with_background_resource_set<F: Future>(units: AdaptivePermitSet, future: F) -> F::Output {
RESOURCE_IO_UNIT.scope(ResourceIoUnit::from(units), future).await
}
pub(crate) fn current_resource_unit() -> Option<ResourceIoUnit> {
RESOURCE_IO_UNIT
.try_with(Clone::clone)
.ok()
.or_else(|| {
DISK_IO_CONTEXT
.try_with(|context| context.resource_unit.clone())
.ok()
.flatten()
})
.or_else(remote::resources)
}
pub(crate) fn has_resource_io_scope() -> bool {
RESOURCE_IO_UNIT.try_with(|_| ()).is_ok()
|| DISK_IO_CONTEXT
.try_with(|context| context.resource_unit.is_some())
.unwrap_or(false)
|| remote::active()
}
/// 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<F: Future>(unit: Option<AdaptivePermit>, future: F) -> F::Output {
pub(crate) async fn with_optional_resource_unit<F: Future>(unit: Option<ResourceIoUnit>, future: F) -> F::Output {
match unit {
Some(unit) => with_background_resource_unit(unit, future).await,
Some(unit) => RESOURCE_IO_UNIT.scope(unit, future).await,
None => future.await,
}
}
@@ -233,7 +367,7 @@ impl DiskIoAdmission {
self.foreground_with_resource(current_resource_unit())
}
fn foreground_with_resource(self: &Arc<Self>, resource_unit: Option<AdaptivePermit>) -> DiskIoPermit {
fn foreground_with_resource(self: &Arc<Self>, resource_unit: Option<ResourceIoUnit>) -> DiskIoPermit {
self.foreground_active.fetch_add(1, Ordering::AcqRel);
if self.background_waiters.load(Ordering::Acquire) > 0 {
self.record_state();
@@ -254,7 +388,7 @@ impl DiskIoAdmission {
async fn background_with_resource(
self: &Arc<Self>,
cancellation: &CancellationToken,
resource_unit: Option<AdaptivePermit>,
resource_unit: Option<ResourceIoUnit>,
) -> std::io::Result<DiskIoPermit> {
let started_at = Instant::now();
let ticket = {
@@ -429,7 +563,7 @@ impl DiskIoPermit {
struct DiskIoPermitLease {
admission: Arc<DiskIoAdmission>,
class: DiskIoClass,
resource_unit: Option<AdaptivePermit>,
resource_unit: Option<ResourceIoUnit>,
}
pub(crate) fn spawn_blocking_with_disk_io_permit<T, F>(
@@ -450,6 +584,18 @@ where
})
}
/// Run native filesystem work with the current resource ownership. Cancelling
/// the async waiter cannot refund a unit while the native operation still runs.
/// This does not acquire capacity or change foreground/background policy.
#[doc(hidden)]
pub fn spawn_blocking_with_resource_ownership<T, F>(operation: F) -> tokio::task::JoinHandle<T>
where
T: Send + 'static,
F: FnOnce() -> T + Send + 'static,
{
spawn_blocking_with_disk_io_permit(None, operation)
}
type AdmissionWait = Pin<Box<dyn Future<Output = io::Result<DiskIoPermit>> + Send>>;
impl Drop for DiskIoPermitLease {
@@ -595,6 +741,63 @@ mod tests {
use std::num::NonZeroU16;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
#[tokio::test]
async fn physical_subset_survives_native_cancellation_without_retaining_unrelated_disks() {
use rustfs_concurrency::adaptive::permit_set::AdaptivePermitSetOwner;
for disk_permit in [false, true] {
let maximum = NonZeroU16::new(1).expect("one physical unit");
let first = AdaptiveAdmission::new(maximum, 1);
let second = AdaptiveAdmission::new(maximum, 1);
let unrelated = AdaptiveAdmission::new(maximum, 1);
let owner = AdaptivePermitSetOwner::new(maximum, vec![first.clone(), second.clone()]).expect("layered disk");
let a = first.try_acquire().expect("first backing device");
let b = second.try_acquire().expect("second backing device");
let healthy = unrelated.try_acquire().expect("another disk in the parent operation");
let subset = owner.capture(&[&a, &b]).expect("capture before storage work");
drop((a, b, owner));
let disk = DiskIoAdmission::new("physical-subset-native");
let (started_tx, started_rx) = tokio::sync::oneshot::channel();
let (release_tx, release_rx) = tokio::sync::oneshot::channel();
let task = tokio::spawn(async move {
let _healthy = healthy;
with_background_resource_set(subset, async {
let permit = if disk_permit {
Some(disk.background(&CancellationToken::new()).await.expect("logical 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("release native I/O");
})
.await
.expect("native I/O completed");
})
.await;
});
started_rx.await.expect("native worker retains subset");
task.abort();
assert!(task.await.expect_err("cancel async waiter").is_cancelled());
let active = [first.snapshot().active, second.snapshot().active, unrelated.snapshot().active];
let healthy_available = unrelated.try_acquire().is_some();
first.set_limit(0);
second.set_limit(0);
release_tx.send(()).expect("finish native work");
tokio::time::timeout(Duration::from_secs(5), async {
while first.snapshot().active != 0 || second.snapshot().active != 0 {
tokio::task::yield_now().await;
}
})
.await
.expect("both backing devices are returned after native completion");
assert_eq!(active, [1, 1, 0], "only the native I/O's physical subset remains charged");
assert!(healthy_available, "cancellation frees the independent disk immediately");
assert!(first.try_acquire().is_none());
assert!(second.try_acquire().is_none(), "old completions cannot reopen shrunken owners");
}
}
#[tokio::test]
async fn resource_unit_outlives_cancelled_scope_until_native_io_finishes() {
for permit_mode in 0..4 {
@@ -0,0 +1,547 @@
// Copyright 2026 RustFS Team
// Licensed under the Apache License, Version 2.0.
use bytes::Bytes;
use std::io;
use std::sync::{Arc, Mutex, OnceLock};
use std::time::Duration;
use tokio::time::Instant;
static BACKGROUND_MEMORY: OnceLock<Arc<BackgroundMemoryBudget>> = OnceLock::new();
/// Process-wide payload storage retained by admitted background I/O. The
/// existing resource sampler publishes limits; this owner starts no workers.
#[derive(Debug)]
pub struct BackgroundMemoryBudget {
// This short critical section only changes counters and a deadline. It
// never allocates, calls storage, or acquires another lock while held.
state: Mutex<State>,
}
#[derive(Debug)]
struct State {
limit: u64,
used: u64,
valid_until: Instant,
}
/// A coherent view of the leased cap and retained payload storage.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct BackgroundMemorySnapshot {
/// Zero when the latest resource observation has expired.
pub limit: u64,
/// Existing storage remains charged even after the cap shrinks below it.
pub used: u64,
}
impl BackgroundMemoryBudget {
/// Start closed until the existing resource sampler publishes a limit.
pub fn new() -> Arc<Self> {
Arc::new(Self {
state: Mutex::new(State {
limit: 0,
used: 0,
valid_until: Instant::now(),
}),
})
}
/// Install the sole process owner before starting adaptive background work.
/// Returns false if an owner was already installed.
pub fn install(self: &Arc<Self>) -> bool {
BACKGROUND_MEMORY.set(self.clone()).is_ok()
}
/// Return only the installed process memory policy; no node or disk ownership.
/// No policy is installed in the default, non-adaptive startup path.
pub fn current() -> Option<Arc<Self>> {
BACKGROUND_MEMORY.get().cloned()
}
/// Update admission without revoking storage retained by native tasks or
/// readers. Missing observations must publish zero capacity.
pub fn set_sampled_limit(&self, limit: u64, valid_for: Duration) {
let mut state = self.state.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
state.limit = limit;
let now = Instant::now();
state.valid_until = now.checked_add(valid_for).unwrap_or(now);
}
/// Read the cap and retained byte count under the same short lock.
pub fn snapshot(&self) -> BackgroundMemorySnapshot {
let state = self.state.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
BackgroundMemorySnapshot {
limit: if Instant::now() < state.valid_until { state.limit } else { 0 },
used: state.used,
}
}
pub(crate) fn reserve(self: &Arc<Self>, bytes: usize) -> io::Result<Reservation> {
let bytes = u64::try_from(bytes).map_err(|_| io::Error::from(io::ErrorKind::OutOfMemory))?;
let mut state = self.state.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
let used = state
.used
.checked_add(bytes)
.ok_or_else(|| io::Error::new(io::ErrorKind::WouldBlock, BackgroundMemoryExhausted))?;
if bytes != 0 && (Instant::now() >= state.valid_until || used > state.limit) {
return Err(io::Error::new(io::ErrorKind::WouldBlock, BackgroundMemoryExhausted));
}
state.used = used;
Ok(Reservation {
owner: self.clone(),
bytes,
})
}
}
#[derive(Debug)]
struct BackgroundMemoryExhausted;
impl std::fmt::Display for BackgroundMemoryExhausted {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str("background buffer memory budget exhausted")
}
}
impl std::error::Error for BackgroundMemoryExhausted {}
pub(crate) fn is_memory_pressure(error: &io::Error) -> bool {
error.kind() == io::ErrorKind::OutOfMemory || error.get_ref().is_some_and(|inner| inner.is::<BackgroundMemoryExhausted>())
}
pub(crate) struct Reservation {
owner: Arc<BackgroundMemoryBudget>,
bytes: u64,
}
impl Reservation {
pub(crate) fn belongs_to(&self, budget: &Arc<BackgroundMemoryBudget>) -> bool {
Arc::ptr_eq(&self.owner, budget)
}
pub(crate) fn wrap(bytes: Bytes, reservation: Option<Self>) -> Bytes {
if bytes.is_empty() {
drop(bytes);
drop(reservation);
return Bytes::new();
}
match reservation {
Some(reservation) => Bytes::from_owner(BackgroundBytes {
bytes,
_reservation: reservation,
}),
None => bytes,
}
}
}
struct BackgroundBytes {
// Free the payload before refunding it, including after the last slice.
bytes: Bytes,
_reservation: Reservation,
}
impl AsRef<[u8]> for BackgroundBytes {
fn as_ref(&self) -> &[u8] {
&self.bytes
}
}
impl Drop for Reservation {
fn drop(&mut self) {
self.owner
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.used -= self.bytes;
}
}
/// Reusable initialized payload storage. The reservation is embedded in the
/// existing native job state, with no per-read owner allocation or payload copy.
pub(crate) struct BackgroundBuffer {
// Field order releases the allocation before refunding its reservation.
bytes: Box<[u8]>,
reservation: Option<Reservation>,
budget: Option<Arc<BackgroundMemoryBudget>>,
}
impl BackgroundBuffer {
pub(crate) fn new(budget: Option<Arc<BackgroundMemoryBudget>>) -> Self {
Self {
bytes: Box::default(),
reservation: None,
budget,
}
}
/// Call only after the previous contents have been consumed. Growth charges
/// both allocations until the old buffer is freed, including after shrink.
pub(crate) fn ensure_len(&mut self, length: usize) -> io::Result<()> {
self.grow(length, 0)
}
/// Preserve an initialized prefix when replacing storage. Callers keep
/// `preserve` within both the existing allocation and the requested length.
pub(crate) fn grow(&mut self, length: usize, preserve: usize) -> io::Result<()> {
if self.bytes.len() >= length {
return Ok(());
}
let reservation = self.budget.as_ref().map(|budget| budget.reserve(length)).transpose()?;
let mut bytes = Vec::new();
bytes
.try_reserve_exact(length)
.map_err(|_| io::Error::from(io::ErrorKind::OutOfMemory))?;
bytes.resize(length, 0);
bytes[..preserve].copy_from_slice(&self.bytes[..preserve]);
self.bytes = bytes.into_boxed_slice();
self.reservation = reservation;
Ok(())
}
/// Metadata is only a size hint: concurrent growth still reads through EOF.
/// Every replacement is charged before copying the retained prefix.
pub(crate) fn read_to_end(mut self, reader: &mut impl io::Read, hint: usize) -> io::Result<Bytes> {
self.ensure_len(hint)?;
let mut filled = 0;
loop {
if filled == self.bytes.len() {
// Probe EOF without forcing a larger allocation for an exact fit.
let mut probe = [0];
match reader.read(&mut probe) {
Ok(0) => break,
Ok(_) => {
let length = self.bytes.len().checked_mul(2).ok_or(io::ErrorKind::OutOfMemory)?.max(32);
self.grow(length, filled)?;
self.bytes[filled] = probe[0];
filled += 1;
}
Err(error) if error.kind() == io::ErrorKind::Interrupted => continue,
Err(error) => return Err(error),
}
}
match reader.read(&mut self.bytes[filled..]) {
Ok(0) => break,
Ok(length) => filled += length,
Err(error) if error.kind() == io::ErrorKind::Interrupted => continue,
Err(error) => return Err(error),
}
}
Ok(self.take_bytes(filled))
}
pub(crate) fn take_bytes(&mut self, length: usize) -> Bytes {
let bytes = std::mem::take(&mut self.bytes);
let reservation = self.reservation.take();
let mut bytes = Bytes::from(bytes.into_vec());
bytes.truncate(length);
Reservation::wrap(bytes, reservation)
}
pub(crate) fn memory_budget(&self) -> Option<Arc<BackgroundMemoryBudget>> {
self.budget.clone()
}
pub(crate) fn as_mut_slice(&mut self) -> &mut [u8] {
&mut self.bytes
}
pub(crate) fn as_slice(&self) -> &[u8] {
&self.bytes
}
}
/// Bounded encoded-frame storage for serializer output and raw read staging.
/// Both paths reserve the actual allocation before initializing its bytes.
pub struct BackgroundEncodeBuffer {
buffer: BackgroundBuffer,
filled: usize,
maximum: usize,
}
impl BackgroundEncodeBuffer {
pub fn new(memory: Arc<BackgroundMemoryBudget>, maximum: usize) -> Self {
Self {
buffer: BackgroundBuffer::new(Some(memory)),
filled: 0,
maximum,
}
}
/// Borrow initialized raw-frame staging; decode only after a successful read.
/// Staging is not encoded output: `into_bytes` remains empty for this path.
pub fn prepare_read(&mut self, length: usize) -> io::Result<&mut [u8]> {
if length > self.maximum {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"background encoded frame exceeds its size limit",
));
}
self.buffer.ensure_len(length)?;
self.filled = 0;
Ok(&mut self.buffer.as_mut_slice()[..length])
}
pub fn into_bytes(mut self) -> Bytes {
self.buffer.take_bytes(self.filled)
}
}
impl io::Write for BackgroundEncodeBuffer {
fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
let required = self.filled.checked_add(bytes.len()).ok_or(io::ErrorKind::OutOfMemory)?;
if required > self.maximum {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"background encoded frame exceeds its size limit",
));
}
if required > self.buffer.as_slice().len() {
let capacity = self
.buffer
.as_slice()
.len()
.saturating_mul(2)
.max(32)
.max(required)
.min(self.maximum);
self.buffer.grow(capacity, self.filled)?;
}
self.buffer.as_mut_slice()[self.filled..required].copy_from_slice(bytes);
self.filled = required;
Ok(bytes.len())
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn encoded_buffer_growth_retains_both_allocations_and_the_last_slice() {
use std::io::Write;
let owner = BackgroundMemoryBudget::new();
owner.set_sampled_limit(95, Duration::from_secs(30));
let mut buffer = BackgroundEncodeBuffer::new(owner.clone(), 64);
buffer.write_all(&[7; 32]).expect("initial allocation fits");
assert_eq!(owner.snapshot().used, 32);
let error = buffer.write_all(&[8]).expect_err("old 32 plus replacement 64 exceed 95");
assert_eq!(error.kind(), io::ErrorKind::WouldBlock);
assert!(is_memory_pressure(&error));
assert_eq!(buffer.filled, 32);
assert_eq!(buffer.buffer.as_slice(), &[7; 32]);
assert_eq!(owner.snapshot().used, 32);
owner.set_sampled_limit(96, Duration::from_secs(30));
buffer.write_all(&[8]).expect("old 32 plus replacement 64 fit 96");
assert_eq!(owner.snapshot().used, 64);
let bytes = buffer.into_bytes();
assert_eq!(&bytes[..32], &[7; 32]);
assert_eq!(bytes[32], 8);
assert_eq!(bytes.len(), 33);
let slice = bytes.slice(32..);
drop(bytes);
assert_eq!(slice.as_ref(), &[8]);
assert_eq!(owner.snapshot().used, 64);
drop(slice);
assert_eq!(owner.snapshot().used, 0);
}
#[test]
fn encoded_buffer_checks_its_limit_before_reserving_storage() {
use std::io::Write;
let owner = BackgroundMemoryBudget::new();
for maximum in [0, 3] {
let mut buffer = BackgroundEncodeBuffer::new(owner.clone(), maximum);
buffer
.write_all(&[])
.expect("empty input needs no storage, including at limit zero");
let error = buffer
.write_all(&[7; 4])
.expect_err("oversized input must fail before memory admission");
assert_eq!(error.kind(), io::ErrorKind::InvalidData);
assert!(buffer.buffer.as_slice().is_empty());
assert!(buffer.into_bytes().is_empty());
assert_eq!(owner.snapshot().used, 0);
}
}
#[test]
fn returned_bytes_keep_the_allocation_charged_through_slices_and_shrink() {
let owner = BackgroundMemoryBudget::new();
owner.set_sampled_limit(8, Duration::from_secs(30));
let bytes = BackgroundBuffer::new(Some(owner.clone()))
.read_to_end(&mut &b"abcdefgh"[..], 8)
.expect("exact allocation fits without an EOF growth allocation");
assert_eq!(bytes.as_ref(), b"abcdefgh");
assert_eq!(owner.snapshot().used, 8);
let slice = bytes.slice(2..6);
owner.set_sampled_limit(0, Duration::ZERO);
drop(bytes);
assert_eq!(slice.as_ref(), b"cdef");
assert_eq!(owner.snapshot().used, 8, "the final slice still owns the complete allocation");
assert!(owner.reserve(1).is_err());
drop(slice);
assert_eq!(owner.snapshot().used, 0);
}
#[test]
fn whole_file_size_hint_never_truncates_growth_or_exposes_unused_bytes() {
for (hint, retained) in [(0, 32), (4, 32), (8, 8), (16, 16)] {
let owner = BackgroundMemoryBudget::new();
owner.set_sampled_limit(36, Duration::from_secs(30));
let bytes = BackgroundBuffer::new(Some(owner.clone()))
.read_to_end(&mut &b"abcdefgh"[..], hint)
.expect("read through EOF independently of the metadata hint");
assert_eq!(bytes.as_ref(), b"abcdefgh", "hint={hint}");
assert_eq!(owner.snapshot().used, retained, "charge allocated capacity, not visible length");
drop(bytes);
assert_eq!(owner.snapshot().used, 0);
}
let owner = BackgroundMemoryBudget::new();
owner.set_sampled_limit(16, Duration::from_secs(30));
let bytes = BackgroundBuffer::new(Some(owner.clone()))
.read_to_end(&mut std::io::empty(), 16)
.expect("file shrank to empty");
assert!(bytes.is_empty());
assert_eq!(owner.snapshot().used, 0, "empty results free unused storage before refund");
}
#[test]
fn preserving_growth_keeps_existing_payload_on_denial_and_reuse() {
let owner = BackgroundMemoryBudget::new();
owner.set_sampled_limit(11, Duration::from_secs(30));
let mut buffer = BackgroundBuffer::new(Some(owner.clone()));
buffer.ensure_len(4).expect("allocate initial payload");
buffer.as_mut_slice().copy_from_slice(b"data");
let error = buffer
.grow(8, 4)
.expect_err("old and replacement allocations need twelve bytes");
assert!(is_memory_pressure(&error));
assert_eq!(buffer.as_slice(), b"data");
assert_eq!(owner.snapshot().used, 4);
owner.set_sampled_limit(12, Duration::from_secs(30));
buffer.grow(8, 4).expect("preserve the payload after admission");
assert_eq!(buffer.as_slice(), b"data\0\0\0\0");
assert_eq!(owner.snapshot().used, 8);
owner.set_sampled_limit(0, Duration::ZERO);
buffer.grow(8, 4).expect("reuse retained storage after expiry");
assert_eq!(buffer.as_slice(), b"data\0\0\0\0");
let bytes = buffer.take_bytes(4);
drop(buffer);
assert_eq!(bytes.as_ref(), b"data");
assert_eq!(owner.snapshot().used, 8);
drop(bytes);
assert_eq!(owner.snapshot().used, 0);
}
#[test]
fn whole_file_growth_denial_and_source_failure_refund_storage() {
let owner = BackgroundMemoryBudget::new();
owner.set_sampled_limit(35, Duration::from_secs(30));
let error = BackgroundBuffer::new(Some(owner.clone()))
.read_to_end(&mut &b"abcdefgh"[..], 4)
.expect_err("old four plus new thirty-two bytes exceed capacity");
assert_eq!(error.kind(), io::ErrorKind::WouldBlock);
assert_eq!(owner.snapshot().used, 0);
struct FailedRead(u8);
impl io::Read for FailedRead {
fn read(&mut self, bytes: &mut [u8]) -> io::Result<usize> {
self.0 += 1;
match self.0 {
1 => Err(io::ErrorKind::Interrupted.into()),
2 => {
bytes[..2].copy_from_slice(b"ab");
Ok(2)
}
_ => Err(io::ErrorKind::PermissionDenied.into()),
}
}
}
let error = BackgroundBuffer::new(Some(owner.clone()))
.read_to_end(&mut FailedRead(0), 8)
.expect_err("failure after partial output is not successful EOF");
assert_eq!(error.kind(), io::ErrorKind::PermissionDenied);
assert_eq!(owner.snapshot().used, 0);
}
#[tokio::test(start_paused = true)]
async fn shrink_and_expiry_keep_retained_storage_charged() {
let owner = BackgroundMemoryBudget::new();
assert!(owner.reserve(1).is_err());
owner.set_sampled_limit(16, Duration::from_secs(3));
let mut buffer = BackgroundBuffer::new(Some(owner.clone()));
buffer.ensure_len(8).expect("initial allocation");
let pointer = buffer.as_slice().as_ptr();
owner.set_sampled_limit(4, Duration::from_secs(3));
buffer.ensure_len(8).expect("existing storage can be reused after shrink");
assert_eq!(buffer.as_slice().as_ptr(), pointer);
assert_eq!(owner.snapshot(), BackgroundMemorySnapshot { limit: 4, used: 8 });
assert_eq!(buffer.ensure_len(9).expect_err("growth must be denied").kind(), io::ErrorKind::WouldBlock);
assert_eq!(owner.snapshot().used, 8, "failed growth preserves the existing reservation");
tokio::time::advance(Duration::from_secs(3)).await;
assert_eq!(owner.snapshot(), BackgroundMemorySnapshot { limit: 0, used: 8 });
assert!(owner.reserve(1).is_err(), "an expired sampler cannot grant memory");
drop(buffer);
assert_eq!(owner.snapshot().used, 0);
}
#[test]
fn growth_accounts_for_both_allocations_and_refunds_failed_attempts() {
let owner = BackgroundMemoryBudget::new();
owner.set_sampled_limit(12, Duration::from_secs(3));
let mut buffer = BackgroundBuffer::new(Some(owner.clone()));
buffer.ensure_len(4).expect("first allocation");
buffer.as_mut_slice().copy_from_slice(b"data");
assert!(buffer.ensure_len(9).is_err(), "four old plus nine new bytes exceed the cap");
assert_eq!(buffer.as_slice(), b"data");
buffer.ensure_len(8).expect("four old plus eight new bytes fit");
assert_eq!(owner.snapshot().used, 8);
assert_eq!(buffer.as_slice(), &[0; 8]);
drop(buffer);
assert_eq!(owner.snapshot().used, 0);
}
#[test]
fn concurrent_reservations_share_one_cap() {
let owner = BackgroundMemoryBudget::new();
owner.set_sampled_limit(1024, Duration::from_secs(30));
let ready = Arc::new(std::sync::Barrier::new(17));
let release = Arc::new(std::sync::Barrier::new(17));
let workers: Vec<_> = (0..16)
.map(|_| {
let owner = owner.clone();
let ready = ready.clone();
let release = release.clone();
std::thread::spawn(move || {
let reservation = owner.reserve(128).ok();
ready.wait();
release.wait();
reservation.is_some()
})
})
.collect();
ready.wait();
let retained = owner.snapshot().used;
release.wait();
let granted = workers
.into_iter()
.map(|worker| worker.join().expect("reservation worker"))
.filter(|granted| *granted)
.count();
assert_eq!(retained, 1024);
assert_eq!(granted, 8);
assert_eq!(owner.snapshot().used, 0);
}
}
@@ -0,0 +1,443 @@
// Copyright 2026 RustFS Team
// Licensed under the Apache License, Version 2.0.
use std::collections::{BTreeMap, BTreeSet};
use std::fs::File;
use std::io::{self, Read};
use std::path::{Path, PathBuf};
const MAX_TOPOLOGY_NODES: usize = 256;
const MAX_TOPOLOGY_DEPTH: usize = 32;
const MAX_ATTRIBUTE_BYTES: u64 = 128;
/// Kernel device numbers can be reused after removal. `diskseq` distinguishes
/// those incarnations; this identity does not authorize storage mutations.
#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
pub(super) struct PhysicalDeviceIdentity {
pub(super) major: u32,
pub(super) minor: u32,
pub(super) generation: u64,
}
fn invalid(message: &'static str) -> io::Error {
io::Error::new(io::ErrorKind::InvalidData, message)
}
pub(super) fn attribute(path: &Path) -> io::Result<String> {
let mut value = String::new();
File::open(path)?.take(MAX_ATTRIBUTE_BYTES + 1).read_to_string(&mut value)?;
if value.len() > MAX_ATTRIBUTE_BYTES as usize {
return Err(invalid("physical device attribute exceeds its bound"));
}
Ok(value.trim().to_owned())
}
fn number(value: &str) -> io::Result<u64> {
if value.is_empty() || !value.bytes().all(|byte| byte.is_ascii_digit()) {
return Err(invalid("invalid physical device number"));
}
value.parse().map_err(|_| invalid("physical device number overflow"))
}
fn identity(path: &Path) -> io::Result<PhysicalDeviceIdentity> {
let generation = number(&attribute(&path.join("diskseq"))?)?;
if generation == 0 {
return Err(invalid("physical device generation is missing"));
}
let device = attribute(&path.join("dev"))?;
let (major, minor) = device
.split_once(':')
.ok_or_else(|| invalid("invalid physical device identity"))?;
let major = u32::try_from(number(major)?).map_err(|_| invalid("physical major number overflow"))?;
let minor = u32::try_from(number(minor)?).map_err(|_| invalid("physical minor number overflow"))?;
if number(&attribute(&path.join("diskseq"))?)? != generation {
return Err(invalid("physical device changed during discovery"));
}
Ok(PhysicalDeviceIdentity {
major,
minor,
generation,
})
}
#[derive(Debug, Eq, PartialEq)]
enum NodeTopology {
Partition {
device: String,
number: u64,
parent: PathBuf,
},
Whole {
identity: PhysicalDeviceIdentity,
children: Vec<PathBuf>,
},
}
fn canonical_beneath(path: &Path, root: &Path) -> io::Result<PathBuf> {
let path = path.canonicalize()?;
if !path.starts_with(root) {
return Err(invalid("physical device topology escapes sysfs"));
}
Ok(path)
}
fn node_topology(path: &Path, root: &Path) -> io::Result<NodeTopology> {
if canonical_beneath(path, root)? != path {
return Err(invalid("physical device path changed during discovery"));
}
match attribute(&path.join("partition")) {
Ok(partition) => {
let number = number(&partition)?;
if number == 0 {
return Err(invalid("physical partition number is missing"));
}
let parent = path.parent().ok_or_else(|| invalid("physical partition has no parent"))?;
return Ok(NodeTopology::Partition {
device: attribute(&path.join("dev"))?,
number,
parent: canonical_beneath(parent, root)?,
});
}
Err(error) if error.kind() == io::ErrorKind::NotFound => {}
Err(error) => return Err(error),
}
let identity = identity(path)?;
let mut children = Vec::new();
for entry in std::fs::read_dir(path.join("slaves"))? {
if children.len() >= MAX_TOPOLOGY_NODES {
return Err(invalid("physical device topology fanout exceeds its bound"));
}
children.push(canonical_beneath(&entry?.path(), root)?);
}
children.sort_unstable();
children.dedup();
Ok(NodeTopology::Whole { identity, children })
}
struct TopologySnapshot {
root: PathBuf,
entry: PathBuf,
first: PathBuf,
nodes: BTreeMap<PathBuf, NodeTopology>,
leaves: BTreeSet<PhysicalDeviceIdentity>,
}
impl TopologySnapshot {
fn discover(sysfs_root: &Path, major: u64, minor: u64) -> io::Result<Self> {
let root = sysfs_root.canonicalize()?;
let entry = root.join(format!("dev/block/{major}:{minor}"));
let first = canonical_beneath(&entry, &root)?;
if attribute(&first.join("dev"))? != format!("{major}:{minor}") {
return Err(invalid("physical device entry does not match its device number"));
}
let mut stack = vec![(first.clone(), 0, false)];
let mut active = BTreeSet::new();
let mut nodes = BTreeMap::new();
let mut leaves = BTreeMap::<PhysicalDeviceIdentity, PathBuf>::new();
while let Some((path, depth, exiting)) = stack.pop() {
if exiting {
active.remove(&path);
continue;
}
if depth > MAX_TOPOLOGY_DEPTH {
return Err(invalid("physical device topology is too deep"));
}
if active.contains(&path) {
return Err(invalid("physical device topology contains a cycle"));
}
if nodes.contains_key(&path) {
continue;
}
if nodes.len() >= MAX_TOPOLOGY_NODES {
return Err(invalid("physical device topology exceeds its node bound"));
}
let topology = node_topology(&path, &root)?;
active.insert(path.clone());
stack.push((path.clone(), depth, true));
match &topology {
NodeTopology::Partition { parent, .. } => stack.push((parent.clone(), depth + 1, false)),
NodeTopology::Whole { identity, children } => {
if stack.len() + children.len() > MAX_TOPOLOGY_NODES * 2 {
return Err(invalid("physical device topology fanout exceeds its bound"));
}
for child in children {
stack.push((child.clone(), depth + 1, false));
}
if children.is_empty()
&& let Some(previous) = leaves.insert(*identity, path.clone())
&& previous != path
{
return Err(invalid("physical device identity is ambiguous"));
}
}
}
nodes.insert(path, topology);
}
if leaves.is_empty() {
return Err(invalid("physical device topology has no leaves"));
}
Ok(Self {
root,
entry,
first,
nodes,
leaves: leaves.into_keys().collect(),
})
}
fn validate(&self) -> io::Result<()> {
if canonical_beneath(&self.entry, &self.root)? != self.first {
return Err(invalid("physical device entry changed during discovery"));
}
for (path, expected) in &self.nodes {
if &node_topology(path, &self.root)? != expected {
return Err(invalid("physical device topology changed during discovery"));
}
}
Ok(())
}
}
/// Resolve a bounded sysfs topology once, before installing an admission owner.
/// Missing topology or generation evidence is unknown, never a name-based owner.
/// Validate intermediate mappings too: stable leaves alone do not prove that a
/// device-mapper table still points to those leaves. This is an observation, not
/// a kernel topology lock; the admission owner must expire and refresh it.
pub(super) fn resolve_physical_devices(sysfs_root: &Path, major: u64, minor: u64) -> io::Result<Vec<PhysicalDeviceIdentity>> {
let snapshot = TopologySnapshot::discover(sysfs_root, major, minor)?;
snapshot.validate()?;
Ok(snapshot.leaves.into_iter().collect())
}
#[cfg(all(test, unix))]
mod tests {
use super::*;
use std::os::unix::fs::symlink;
use std::sync::atomic::{AtomicU64, Ordering};
struct Fixture(PathBuf);
impl Fixture {
fn new() -> Self {
static NEXT: AtomicU64 = AtomicU64::new(0);
let stamp = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("clock")
.as_nanos();
let root = std::env::temp_dir().join(format!(
"physical-identity-{}-{stamp}-{}",
std::process::id(),
NEXT.fetch_add(1, Ordering::Relaxed)
));
std::fs::create_dir(&root).expect("exclusive test root");
std::fs::create_dir_all(root.join("sys/dev/block")).expect("sysfs fixture");
Self(root)
}
fn root(&self) -> PathBuf {
self.0.join("sys")
}
fn device(&self, name: &str, major: u32, minor: u32, generation: u64) -> PathBuf {
let path = self.root().join("devices").join(name);
std::fs::create_dir_all(path.join("slaves")).expect("device fixture");
std::fs::write(path.join("dev"), format!("{major}:{minor}\n")).expect("device number");
std::fs::write(path.join("diskseq"), format!("{generation}\n")).expect("device generation");
symlink(&path, self.root().join(format!("dev/block/{major}:{minor}"))).expect("sysfs link");
path
}
fn link(&self, parent: &Path, child: &Path, name: &str) {
symlink(child, parent.join("slaves").join(name)).expect("slave link");
}
}
impl Drop for Fixture {
fn drop(&mut self) {
std::fs::remove_dir_all(&self.0).expect("remove owned fixture");
}
}
#[test]
fn partitions_share_the_leaf_generation_and_replacement_changes_identity() {
let fixture = Fixture::new();
let leaf = fixture.device("sda", 8, 0, 41);
let first = fixture.device("sda/sda1", 8, 1, 41);
let second = fixture.device("sda/sda2", 8, 2, 41);
std::fs::write(first.join("partition"), "1\n").expect("first partition");
std::fs::write(second.join("partition"), "2\n").expect("second partition");
let old = resolve_physical_devices(&fixture.root(), 8, 1).expect("first leaf");
assert_eq!(old, resolve_physical_devices(&fixture.root(), 8, 2).expect("alias leaf"));
assert_eq!(
old,
vec![PhysicalDeviceIdentity {
major: 8,
minor: 0,
generation: 41
}]
);
std::fs::write(leaf.join("diskseq"), "42\n").expect("replacement generation");
let new = resolve_physical_devices(&fixture.root(), 8, 1).expect("replacement leaf");
assert_ne!(old, new);
assert_eq!(new[0].generation, 42);
}
#[test]
fn overlapping_stacks_deduplicate_shared_physical_leaves() {
let fixture = Fixture::new();
let a = fixture.device("a", 8, 0, 1);
let b = fixture.device("b", 8, 16, 2);
let inner = fixture.device("dm-0", 253, 0, 3);
let outer = fixture.device("dm-1", 253, 1, 4);
fixture.link(&inner, &a, "a");
fixture.link(&inner, &b, "b");
fixture.link(&outer, &inner, "dm-0");
fixture.link(&outer, &a, "a-again");
assert_eq!(
resolve_physical_devices(&fixture.root(), 253, 0).expect("inner"),
resolve_physical_devices(&fixture.root(), 253, 1).expect("outer")
);
assert_eq!(resolve_physical_devices(&fixture.root(), 253, 1).expect("leaves").len(), 2);
}
#[test]
fn missing_or_invalid_generation_never_falls_back_to_a_device_name() {
let fixture = Fixture::new();
let leaf = fixture.device("a", 8, 0, 1);
for value in ["0", "-1", "18446744073709551616", "+1"] {
std::fs::write(leaf.join("diskseq"), value).expect("bad generation");
assert!(resolve_physical_devices(&fixture.root(), 8, 0).is_err(), "accepted {value}");
}
std::fs::remove_file(leaf.join("diskseq")).expect("missing generation");
assert_eq!(
resolve_physical_devices(&fixture.root(), 8, 0)
.expect_err("unknown generation")
.kind(),
io::ErrorKind::NotFound
);
assert!(resolve_physical_devices(&fixture.root(), 99, 99).is_err());
}
#[test]
fn topology_cycles_and_escaping_links_are_rejected() {
let fixture = Fixture::new();
let a = fixture.device("a", 8, 0, 1);
let b = fixture.device("b", 8, 16, 2);
fixture.link(&a, &b, "b");
fixture.link(&b, &a, "a");
assert!(
resolve_physical_devices(&fixture.root(), 8, 0)
.expect_err("cycle")
.to_string()
.contains("cycle")
);
std::fs::remove_file(b.join("slaves/a")).expect("remove cycle edge");
fixture.link(&b, &fixture.0, "outside");
assert!(
resolve_physical_devices(&fixture.root(), 8, 0)
.expect_err("escape")
.to_string()
.contains("escapes")
);
}
#[test]
fn intermediate_mapping_changes_are_rejected_even_when_old_leaves_survive() {
let fixture = Fixture::new();
let a = fixture.device("a", 8, 0, 1);
let b = fixture.device("b", 8, 16, 2);
let mapper = fixture.device("dm-0", 253, 0, 3);
fixture.link(&mapper, &a, "backing");
let snapshot = TopologySnapshot::discover(&fixture.root(), 253, 0).expect("initial mapping");
snapshot.validate().expect("unchanged topology");
std::fs::remove_file(mapper.join("slaves/backing")).expect("replace mapping");
fixture.link(&mapper, &b, "backing");
assert_eq!(identity(&a).expect("old leaf still exists").generation, 1);
assert!(snapshot.validate().is_err(), "stable old leaves cannot certify a changed mapper");
let current = resolve_physical_devices(&fixture.root(), 253, 0).expect("new complete mapping");
assert_eq!(current, vec![identity(&b).expect("new leaf")]);
let snapshot = TopologySnapshot::discover(&fixture.root(), 253, 0).expect("new mapping");
std::fs::write(mapper.join("diskseq"), "4\n").expect("replace mapper incarnation");
assert!(
snapshot.validate().is_err(),
"an intermediate generation change also invalidates discovery"
);
}
#[test]
fn entry_and_partition_changes_are_not_hidden_by_an_unchanged_parent() {
let fixture = Fixture::new();
fixture.device("a", 8, 0, 1);
let partition = fixture.device("a/p1", 8, 1, 1);
std::fs::write(partition.join("partition"), "1\n").expect("partition index");
let snapshot = TopologySnapshot::discover(&fixture.root(), 8, 1).expect("partition topology");
std::fs::write(partition.join("partition"), "2\n").expect("changed partition");
assert!(snapshot.validate().is_err());
let b = fixture.device("b", 8, 16, 2);
let snapshot = TopologySnapshot::discover(&fixture.root(), 8, 0).expect("first entry");
let entry = fixture.root().join("dev/block/8:0");
std::fs::remove_file(&entry).expect("replace entry");
symlink(&b, &entry).expect("changed entry target");
assert!(snapshot.validate().is_err());
assert!(
resolve_physical_devices(&fixture.root(), 8, 0).is_err(),
"entry device numbers must match"
);
}
#[test]
fn missing_slave_directory_is_unknown_not_a_physical_leaf() {
let fixture = Fixture::new();
let mapper = fixture.device("dm-0", 253, 0, 1);
std::fs::remove_dir(mapper.join("slaves")).expect("missing topology evidence");
assert!(resolve_physical_devices(&fixture.root(), 253, 0).is_err());
}
#[test]
fn topology_depth_and_attribute_sizes_are_bounded() {
let fixture = Fixture::new();
let first = fixture.device("d0", 253, 0, 1);
let mut previous = first.clone();
for index in 1..=MAX_TOPOLOGY_DEPTH + 1 {
let child = fixture.device(&format!("d{index}"), 253, u32::try_from(index).expect("bounded index"), 1);
fixture.link(&previous, &child, "next");
previous = child;
}
assert!(
resolve_physical_devices(&fixture.root(), 253, 0)
.expect_err("depth bound")
.to_string()
.contains("deep")
);
std::fs::remove_file(first.join("slaves/next")).expect("make first node a leaf");
std::fs::write(first.join("diskseq"), "1".repeat(129)).expect("oversized attribute");
assert!(
resolve_physical_devices(&fixture.root(), 253, 0)
.expect_err("attribute bound")
.to_string()
.contains("attribute")
);
}
#[test]
fn exact_node_bound_passes_and_one_more_node_is_rejected() {
let fixture = Fixture::new();
let root = fixture.device("root", 253, 0, 1);
for index in 0..MAX_TOPOLOGY_NODES - 1 {
let leaf = fixture.device(&format!("leaf-{index}"), 8, u32::try_from(index).expect("bounded index"), 2);
fixture.link(&root, &leaf, &format!("leaf-{index}"));
}
assert_eq!(
resolve_physical_devices(&fixture.root(), 253, 0)
.expect("exact node bound")
.len(),
MAX_TOPOLOGY_NODES - 1
);
let extra = fixture.device("extra", 9, 0, 3);
fixture.link(&root, &extra, "extra");
assert!(
resolve_physical_devices(&fixture.root(), 253, 0).is_err(),
"reject a graph above the node bound"
);
}
}
@@ -0,0 +1,268 @@
// Copyright 2026 RustFS Team
// Licensed under the Apache License, Version 2.0.
use std::collections::{BTreeMap, BTreeSet};
use std::io;
use std::num::NonZeroU16;
use std::sync::{Arc, Weak};
use std::time::Duration;
use rustfs_concurrency::adaptive::AdaptiveAdmission;
#[cfg(test)]
use rustfs_concurrency::adaptive::{AdaptiveAdmissionSnapshot, AdaptivePermit};
use super::physical_identity::PhysicalDeviceIdentity;
const MAX_PHYSICAL_OWNERS: usize = 256;
/// Resolve aliases outside object work. The sampling owner serializes changes;
/// attached handles and native permits never consult this registry.
pub(super) struct PhysicalOwnerRegistry {
maximum: NonZeroU16,
owners: BTreeMap<PhysicalDeviceIdentity, Weak<AdaptiveAdmission>>,
}
/// One attachment to an existing physical capacity owner, not a reservation.
/// The gate itself owns permit slots, so native permit clones keep a registry
/// entry alive even after the last logical disk attachment disappears.
pub(super) struct PhysicalDeviceAdmission {
#[cfg(test)]
identity: PhysicalDeviceIdentity,
gate: Arc<AdaptiveAdmission>,
}
impl PhysicalOwnerRegistry {
pub(super) fn new(maximum: NonZeroU16) -> Self {
Self {
maximum,
owners: BTreeMap::new(),
}
}
/// Old topology descriptors may still reference an unused gate. Close it
/// before reclaiming its registry slot; only admitted I/O keeps an absent
/// incarnation alive, and its final release still reaches the same owner.
pub(super) fn retire_absent(&mut self, current: &BTreeSet<PhysicalDeviceIdentity>) {
self.owners.retain(|identity, owner| {
let Some(gate) = owner.upgrade() else {
return false;
};
if current.contains(identity) {
return true;
}
gate.set_sampled_limit(0, Duration::ZERO);
gate.snapshot().active != 0
});
}
pub(super) fn bind(&mut self, identities: &[PhysicalDeviceIdentity]) -> io::Result<Vec<PhysicalDeviceAdmission>> {
if identities.is_empty() || identities.len() > MAX_PHYSICAL_OWNERS || identities.iter().any(|id| id.generation == 0) {
return Err(io::Error::new(io::ErrorKind::InvalidData, "invalid physical capacity identities"));
}
let identities = identities.iter().copied().collect::<BTreeSet<_>>();
self.owners.retain(|_, owner| owner.strong_count() != 0);
let new_owners = identities
.iter()
.filter(|identity| !self.owners.contains_key(identity))
.count();
if new_owners > MAX_PHYSICAL_OWNERS - self.owners.len() {
return Err(io::Error::new(io::ErrorKind::WouldBlock, "physical capacity registry is full"));
}
Ok(identities
.into_iter()
.map(|identity| {
let gate = self.owners.get(&identity).and_then(Weak::upgrade).unwrap_or_else(|| {
let gate = AdaptiveAdmission::new(self.maximum, 0);
self.owners.insert(identity, Arc::downgrade(&gate));
gate
});
PhysicalDeviceAdmission {
#[cfg(test)]
identity,
gate,
}
})
.collect())
}
/// Only the sampling owner publishes capacity. New or expired identities
/// remain closed until a fresh observation establishes a bounded limit.
#[cfg(test)]
pub(super) fn publish(&mut self, identity: PhysicalDeviceIdentity, limit: u32, valid_for: Duration) -> bool {
let Some(gate) = self.owners.get(&identity).and_then(Weak::upgrade) else {
return false;
};
gate.set_sampled_limit(limit, valid_for);
true
}
}
impl PhysicalDeviceAdmission {
pub(super) fn gate(&self) -> &Arc<AdaptiveAdmission> {
&self.gate
}
#[cfg(test)]
pub(super) fn identity(&self) -> PhysicalDeviceIdentity {
self.identity
}
#[cfg(test)]
pub(super) fn snapshot(&self) -> AdaptiveAdmissionSnapshot {
self.gate.snapshot()
}
#[cfg(test)]
pub(super) fn try_acquire(&self) -> Option<AdaptivePermit> {
self.gate.try_acquire()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn device(minor: u32, generation: u64) -> PhysicalDeviceIdentity {
PhysicalDeviceIdentity {
major: 259,
minor,
generation,
}
}
fn registry() -> PhysicalOwnerRegistry {
PhysicalOwnerRegistry::new(NonZeroU16::new(2).expect("two units"))
}
#[test]
fn aliases_and_overlapping_stacks_share_one_physical_capacity() {
let mut registry = registry();
let a = device(1, 11);
let b = device(2, 12);
let single = registry.bind(&[a]).expect("single device");
let stack = registry.bind(&[b, a, a]).expect("overlapping stack");
assert_eq!(stack.len(), 2, "duplicates cannot reserve additional units");
assert_eq!(stack[0].identity(), a);
assert!(single[0].try_acquire().is_none(), "unobserved capacity is closed");
assert!(registry.publish(a, 1, Duration::from_secs(60)));
let held = single[0].try_acquire().expect("physical unit");
assert!(stack[0].try_acquire().is_none(), "the alias cannot double the capacity");
assert_eq!(stack[0].snapshot().active, 1);
drop(held);
assert!(stack[0].try_acquire().is_some());
}
#[test]
fn native_tail_keeps_the_owner_alive_after_all_attachments_are_gone() {
let mut registry = registry();
let id = device(1, 11);
let attachment = registry.bind(&[id]).expect("logical disk");
assert!(registry.publish(id, 1, Duration::from_secs(60)));
let waiter = attachment[0].try_acquire().expect("admitted work");
let native = waiter.clone();
drop(waiter);
drop(attachment);
let reconnected = registry.bind(&[id]).expect("reconnect during native I/O");
assert_eq!(reconnected[0].snapshot().active, 1);
assert!(reconnected[0].try_acquire().is_none());
drop(native);
assert!(reconnected[0].try_acquire().is_some(), "only the final native release returns capacity");
}
#[test]
fn replacement_generation_does_not_receive_old_native_refunds() {
let mut registry = registry();
let old_id = device(1, 11);
let new_id = device(1, 12);
let old = registry.bind(&[old_id]).expect("old incarnation");
assert!(registry.publish(old_id, 1, Duration::from_secs(60)));
let native = old[0].try_acquire().expect("old native I/O");
drop(old);
let new = registry.bind(&[new_id]).expect("replacement incarnation");
assert_eq!(new[0].snapshot().active, 0);
assert!(new[0].try_acquire().is_none());
assert!(registry.publish(new_id, 1, Duration::from_secs(60)));
let current = new[0].try_acquire().expect("new work");
drop(native);
assert_eq!(new[0].snapshot().active, 1, "old completion cannot refund the replacement");
assert!(new[0].try_acquire().is_none());
drop(current);
assert_eq!(new[0].snapshot().active, 0);
}
#[test]
fn shrink_and_expiry_preserve_active_units_and_leave_other_disks_available() {
let mut registry = registry();
let a = device(1, 11);
let b = device(2, 12);
let disks = registry.bind(&[a, b]).expect("two physical disks");
assert!(registry.publish(a, 2, Duration::from_secs(60)));
assert!(registry.publish(b, 1, Duration::from_secs(60)));
let first = disks[0].try_acquire().expect("first unit");
let second = disks[0].try_acquire().expect("second unit");
assert!(registry.publish(a, 1, Duration::from_secs(60)));
assert_eq!(disks[0].snapshot().active, 2);
assert!(disks[0].try_acquire().is_none());
assert!(disks[1].try_acquire().is_some(), "a stalled disk cannot consume another disk's capacity");
assert!(registry.publish(a, 2, Duration::ZERO));
assert_eq!(disks[0].snapshot().limit, 0);
assert_eq!(disks[0].snapshot().active, 2);
drop((first, second));
assert!(disks[0].try_acquire().is_none(), "expired observations do not reopen after completion");
}
#[test]
fn registry_bound_preserves_live_native_owners_and_reclaims_only_dead_entries() {
let mut registry = registry();
let ids = (0..MAX_PHYSICAL_OWNERS)
.map(|minor| device(u32::try_from(minor).expect("bounded minor"), 1))
.collect::<Vec<_>>();
let attachments = registry.bind(&ids).expect("exact registry bound");
let extra = device(999, 1);
assert!(matches!(registry.bind(&[extra]), Err(error) if error.kind() == io::ErrorKind::WouldBlock));
assert_eq!(registry.owners.len(), MAX_PHYSICAL_OWNERS, "failed binding must not add partial owners");
assert!(registry.publish(ids[0], 1, Duration::from_secs(60)));
let native = attachments[0].try_acquire().expect("native tail");
drop(attachments);
let replacement = registry.bind(&[extra]).expect("dead owners reclaimed");
assert_eq!(registry.owners.len(), 2, "the native tail still owns its original entry");
let original = registry.bind(&[ids[0]]).expect("original native owner");
assert_eq!(original[0].snapshot().active, 1);
drop((native, original, replacement));
assert!(!registry.publish(extra, 1, Duration::from_secs(60)), "a dead owner cannot be refreshed");
assert!(registry.bind(&[]).is_err());
assert!(registry.bind(&[device(1, 0)]).is_err());
}
#[test]
fn retirement_reclaims_idle_descriptors_but_keeps_active_native_incarnations() {
let mut registry = registry();
let ids = (0..MAX_PHYSICAL_OWNERS)
.map(|minor| device(u32::try_from(minor).expect("bounded minor"), 1))
.collect::<Vec<_>>();
let descriptors = registry.bind(&ids).expect("full registry");
assert!(registry.publish(ids[0], 1, Duration::from_secs(60)));
let native = descriptors[0].try_acquire().expect("active old incarnation");
let mut current = ids[2..].iter().copied().collect::<BTreeSet<_>>();
registry.retire_absent(&current);
assert_eq!(registry.owners.len(), MAX_PHYSICAL_OWNERS - 1);
assert_eq!(descriptors[0].snapshot().active, 1);
assert!(descriptors[0].try_acquire().is_none());
assert!(descriptors[1].try_acquire().is_none());
let first = device(0, 2);
let second = device(1, 2);
let replacement = registry.bind(&[first]).expect("idle descriptor was reclaimed");
current.insert(first);
assert!(matches!(registry.bind(&[second]), Err(error) if error.kind() == io::ErrorKind::WouldBlock));
drop(native);
registry.retire_absent(&current);
let next = registry.bind(&[second]).expect("final native release permits reclamation");
assert_eq!(replacement[0].snapshot().active, 0);
assert_eq!(next[0].snapshot().active, 0);
assert_eq!(descriptors[0].snapshot().active, 0);
assert!(
!registry.publish(ids[0], 1, Duration::from_secs(60)),
"retired identity cannot be refreshed"
);
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,309 @@
// Copyright 2026 RustFS Team
// Licensed under the Apache License, Version 2.0.
use super::{PhysicalDiskBinding, ResourceIoUnit};
use rustfs_concurrency::adaptive::{
NodeAdmission,
permit_set::{NodeResourceCompletion, NodeResourcePermit},
};
use std::cell::{Cell, RefCell};
use std::future::Future;
use std::io;
use std::sync::Arc;
tokio::task_local! {
static REMOTE_RESOURCES: RemoteBackgroundIo;
}
/// Receiver-owned admission for an authenticated background RPC. The owner must
/// be the same node controller used by local Scanner and Heal work.
pub struct RemoteBackgroundIo {
node: Arc<NodeAdmission>,
admitted: RefCell<Option<NodeResourcePermit>>,
progressed: Cell<bool>,
finished: Cell<bool>,
}
impl RemoteBackgroundIo {
pub fn new(node: Arc<NodeAdmission>) -> Self {
Self {
node,
admitted: RefCell::new(None),
progressed: Cell::new(false),
finished: Cell::new(false),
}
}
/// Move the request's reservations into its existing stream producer.
/// Only physical subsets may be cloned into native I/O; never its node unit.
pub fn take_current() -> Option<Self> {
REMOTE_RESOURCES
.try_with(|context| Self {
node: context.node.clone(),
admitted: RefCell::new(context.admitted.take()),
progressed: Cell::new(context.progressed.replace(false)),
finished: Cell::new(false),
})
.ok()
}
/// Defer a response's progress without keeping async admission alive while
/// its body streams. The reader retains its own native resource ownership.
pub fn take_current_completion() -> Option<NodeResourceCompletion> {
REMOTE_RESOURCES
.try_with(|context| {
let admitted = context.admitted.take();
if context.progressed.replace(false) {
admitted.map(NodeResourcePermit::into_completion)
} else {
None
}
})
.ok()
.flatten()
}
/// A failed response or cancelled producer must not credit earlier phases.
pub fn discard_current_progress() {
let _ = REMOTE_RESOURCES.try_with(|context| context.progressed.set(false));
}
}
impl Drop for RemoteBackgroundIo {
fn drop(&mut self) {
if self.finished.get()
&& self.progressed.get()
&& let Some(permit) = self.admitted.get_mut().take()
{
permit.complete();
}
}
}
/// Scope a verified request or its stream producer. An interrupted request
/// returns node capacity while native jobs retain their selected physical units.
pub async fn with_remote_background_io<F: Future>(context: Option<RemoteBackgroundIo>, operation: F) -> F::Output {
match context {
Some(context) => {
REMOTE_RESOURCES
.scope(context, async {
let result = operation.await;
REMOTE_RESOURCES.with(|context| context.finished.set(true));
result
})
.await
}
None => operation.await,
}
}
pub(crate) fn active() -> bool {
REMOTE_RESOURCES.try_with(|_| ()).is_ok()
}
pub(crate) fn resources() -> Option<ResourceIoUnit> {
REMOTE_RESOURCES
.try_with(|context| {
context
.admitted
.borrow()
.as_ref()
.map(|permit| ResourceIoUnit::from(permit.resources().clone()))
})
.ok()
.flatten()
}
pub(crate) fn record_progress() {
let _ = REMOTE_RESOURCES.try_with(|context| context.progressed.set(true));
}
pub(crate) fn admit(binding: &PhysicalDiskBinding) -> io::Result<()> {
REMOTE_RESOURCES
.try_with(|context| {
let deferred = || io::Error::from(io::ErrorKind::WouldBlock);
let plan = binding.plan().ok_or_else(deferred)?;
let mut admitted = context.admitted.borrow_mut();
if let Some(permit) = admitted.as_ref() {
// Reuse the request's receipt for append/read/rename phases at
// capacity one. Additional attachments never wait under locks.
let native = plan.capture_or_acquire_native(permit.resources()).map_err(|_| deferred())?;
drop(native);
if !plan.is_current() {
return Err(deferred());
}
} else {
*admitted = Some(plan.try_acquire(&context.node).ok_or_else(deferred)?);
}
Ok(())
})
.unwrap_or(Ok(()))
}
#[cfg(test)]
mod tests {
use super::*;
use rustfs_concurrency::adaptive::{
AdaptiveAdmission, FeedbackController, FeedbackObservation, ResourcePressure, permit_set::AdaptivePermitSetOwner,
};
use std::num::NonZeroU16;
fn node() -> Arc<NodeAdmission> {
let node = NodeAdmission::new(NonZeroU16::MIN);
node.update(
&mut FeedbackController::new(),
Some(FeedbackObservation {
generation: 1,
max_concurrency: 1,
pressure: ResourcePressure::Unknown,
completed_units: None,
}),
);
node
}
#[tokio::test]
async fn receiver_reuses_one_receipt_across_phases_and_releases_only_native_subsets() {
let node = node();
let a = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let b = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let shared = AdaptivePermitSetOwner::new(NonZeroU16::MIN, vec![a.clone(), b.clone()]).expect("device catalog");
let first = PhysicalDiskBinding::default();
let second = PhysicalDiskBinding::default();
first
.set_test_shared_capacity(shared.clone(), vec![a.clone()])
.expect("device A");
second.set_test_shared_capacity(shared, vec![b.clone()]).expect("device B");
let native = with_remote_background_io(Some(RemoteBackgroundIo::new(node.clone())), async {
assert!(resources().is_none(), "no physical work before target selection");
admit(&first).expect("first phase");
let admitted = resources().expect("admitted resources");
let native = first
.plan()
.expect("current plan")
.capture_native(admitted.as_set().expect("resource set"))
.expect("escaped native read");
admit(&first).expect("second phase reuses capacity while the first reader lives");
admit(&second).expect("additional selected device");
assert_eq!([node.snapshot().active, a.snapshot().active, b.snapshot().active], [1, 1, 1]);
record_progress();
native
})
.await;
assert_eq!([node.snapshot().active, a.snapshot().active, b.snapshot().active], [0, 1, 0]);
assert_eq!([node.completed(), a.completed(), b.completed()], [Some(1), Some(1), Some(1)]);
drop(native);
assert_eq!(a.snapshot().active, 0);
}
#[tokio::test]
async fn receiver_cancellation_returns_node_but_not_escaped_native_capacity() {
let node = node();
let gate = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let binding = PhysicalDiskBinding::default();
binding.set_test_capacity(vec![gate.clone()]).expect("device plan");
let (sender, receiver) = tokio::sync::oneshot::channel();
let task = tokio::spawn(with_remote_background_io(Some(RemoteBackgroundIo::new(node.clone())), async move {
admit(&binding).expect("receiver admission");
let admitted = resources().expect("receiver resources");
let native = binding
.plan()
.expect("plan")
.capture_native(admitted.as_set().expect("set"))
.expect("native subset");
sender.send(native).expect("native owner receiver");
record_progress();
std::future::pending::<()>().await;
}));
let native = receiver.await.expect("native read ownership");
task.abort();
assert!(task.await.expect_err("cancel request").is_cancelled());
assert_eq!([node.snapshot().active, gate.snapshot().active], [0, 1]);
assert_eq!([node.completed(), gate.completed()], [Some(0), Some(0)]);
drop(native);
assert_eq!(gate.snapshot().active, 0);
}
#[tokio::test]
async fn receiver_defers_stream_progress_without_retaining_node_capacity() {
for (progressed, succeeds) in [(false, true), (true, false), (true, true)] {
let node = node();
let gate = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let binding = PhysicalDiskBinding::default();
binding.set_test_capacity(vec![gate.clone()]).expect("device plan");
let completion = with_remote_background_io(Some(RemoteBackgroundIo::new(node.clone())), async {
admit(&binding).expect("receiver admission");
if progressed {
record_progress();
}
let completion = RemoteBackgroundIo::take_current_completion();
assert!(RemoteBackgroundIo::take_current_completion().is_none(), "single progress owner");
assert!(resources().is_none());
completion
})
.await;
assert_eq!([node.snapshot().active, gate.snapshot().active], [0, 0]);
assert_eq!([node.completed(), gate.completed()], [Some(0), Some(0)]);
assert_eq!(completion.is_some(), progressed);
if succeeds && let Some(completion) = completion {
completion.complete();
}
let expected = u64::from(progressed && succeeds);
assert_eq!([node.completed(), gate.completed()], [Some(expected), Some(expected)]);
}
assert!(
RemoteBackgroundIo::take_current_completion().is_none(),
"foreground has no progress owner"
);
}
#[tokio::test]
async fn receiver_failed_completion_does_not_credit_successful_preparation() {
let node = node();
let gate = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let binding = PhysicalDiskBinding::default();
binding.set_test_capacity(vec![gate.clone()]).expect("device plan");
with_remote_background_io(Some(RemoteBackgroundIo::new(node.clone())), async {
admit(&binding).expect("receiver admission");
record_progress();
RemoteBackgroundIo::discard_current_progress();
})
.await;
assert_eq!([node.snapshot().active, gate.snapshot().active], [0, 0]);
assert_eq!([node.completed(), gate.completed()], [Some(0), Some(0)]);
}
#[tokio::test]
async fn receiver_stream_transfer_moves_node_ownership_and_defers_unavailable_devices() {
let node = node();
let gate = AdaptiveAdmission::new(NonZeroU16::MIN, 0);
let binding = PhysicalDiskBinding::default();
binding.set_test_capacity(vec![gate.clone()]).expect("device plan");
let context = with_remote_background_io(Some(RemoteBackgroundIo::new(node.clone())), async {
assert_eq!(admit(&binding).expect_err("closed device").kind(), io::ErrorKind::WouldBlock);
assert_eq!([node.snapshot().active, gate.snapshot().active], [0, 0]);
gate.set_limit(1);
admit(&binding).expect("opened device");
let context = RemoteBackgroundIo::take_current().expect("stream context");
assert!(resources().is_none(), "request cannot keep a second node owner");
context
})
.await;
assert_eq!([node.snapshot().active, gate.snapshot().active], [1, 1]);
assert_eq!(node.completed(), Some(0));
with_remote_background_io(Some(context), async {
admit(&binding).expect("stream reuses the request's receipt");
binding
.set_test_capacity(vec![AdaptiveAdmission::new(NonZeroU16::MIN, 1)])
.expect("replace physical device");
assert_eq!(
admit(&binding).expect_err("old receipt cannot cross topology").kind(),
io::ErrorKind::WouldBlock
);
})
.await;
assert_eq!([node.snapshot().active, gate.snapshot().active], [0, 0]);
assert_eq!(node.completed(), Some(0), "admission alone is not completed work");
}
}
File diff suppressed because it is too large Load Diff
+3 -3
View File
@@ -28,7 +28,7 @@ use super::{
should_remove_staged_meta_before_commit, skip_access_checks, tmp_staging_mutation_path,
};
#[cfg(test)]
use super::{run_inline_before_file_sync_admission, run_owned_file_write_before_open, run_rename_data_after_first_publication};
use super::{run_inline_before_file_sync_admission, run_owned_file_before_open, run_rename_data_after_first_publication};
use crate::crash_inject::{self, CrashPoint};
use crate::disk::{
QUOTA_MUTATION_FENCE_METADATA_SUFFIX, RenameDataResp, STORAGE_FORMAT_FILE, STORAGE_FORMAT_FILE_BACKUP, SnapshotLeaseToken,
@@ -706,7 +706,7 @@ impl LocalDisk {
async move {
os::run_blocking_namespace_operation(owned_mutation_lease, move || {
#[cfg(test)]
run_owned_file_write_before_open(&src_file_path);
run_owned_file_before_open(&src_file_path);
let mut prepared_metadata_source = os::create_prepared_rename_source_with_commit_guard(
&src_file_path,
&dst_file_path,
@@ -867,7 +867,7 @@ impl LocalDisk {
// deterministic rollback backup.
let write_result = os::run_blocking_namespace_operation(commit.mutation_lease.clone(), move || {
#[cfg(test)]
run_owned_file_write_before_open(&backup_path);
run_owned_file_before_open(&backup_path);
backup_path_guard.write_file_for_path_access(
&backup_path,
backup_bytes.as_ref(),
@@ -1245,7 +1245,7 @@ impl LocalDisk {
continue;
};
#[cfg(test)]
run_owned_file_write_before_open(&record_destination.join(name));
run_owned_file_before_open(&record_destination.join(name));
let mut file = std::fs::File::create(staged.path().join(name))?;
std::io::Write::write_all(&mut file, bytes)?;
if sync {
@@ -1713,7 +1713,7 @@ impl LocalDisk {
self.write_all_meta_with_namespace_owner(
volume,
&format!("{object}/{STORAGE_FORMAT_FILE}"),
&target,
target,
true,
owner.clone(),
)
+20 -6
View File
@@ -29,13 +29,19 @@ pub(crate) mod lifecycle_test_io;
pub mod local;
pub mod os;
#[cfg(test)]
pub(crate) use io_admission::DiskIoClass;
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(crate) use io_admission::memory as background_memory;
pub(crate) use io_admission::remote as remote_io_admission;
pub use io_admission::{
BackgroundMemoryBudget, PhysicalDiskBinding, PhysicalDiskPlan, PhysicalDiskSample, PhysicalDiskSampler, PhysicalResourcePlan,
RemoteBackgroundIo, spawn_blocking_with_resource_ownership, with_background_disk_io, with_background_resource_set,
with_background_resource_unit, with_remote_background_io,
};
#[cfg(test)]
pub(crate) use io_admission::{DiskIoClass, test_resource_units};
pub(crate) use io_admission::{
PhysicalHealBinding, acquire_context_disk_io_permit, current_disk_io_context, current_resource_unit, has_resource_io_scope,
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";
@@ -234,6 +240,14 @@ pub enum Disk {
}
impl Disk {
/// Physical ownership belongs to the node hosting the local attachment.
pub fn physical_admission(&self) -> Option<Arc<PhysicalDiskBinding>> {
match self {
Self::Local(disk) => Some(disk.physical_admission()),
Self::Remote(_) => None,
}
}
/// A read-only probe; the physical delete repeats the same leaf checks.
pub async fn is_plain_orphan_data_dir(&self, volume: &str, path: &str) -> Result<bool> {
match self {
+64 -43
View File
@@ -263,7 +263,7 @@ pub(crate) mod fsync_dir_recorder {
}
}
/// Pause a real namespace mutation inside its physical executor.
/// Pause a real native filesystem operation inside its physical executor.
#[cfg(all(any(test, feature = "test-util"), not(windows)))]
pub(crate) mod prepared_publication_test_hooks {
use super::*;
@@ -273,6 +273,7 @@ pub(crate) mod prepared_publication_test_hooks {
PreparedRename,
Rename,
Remove,
NativeReadOpen,
#[cfg(test)]
Rollback,
#[cfg(test)]
@@ -365,6 +366,8 @@ pub enum LocalPublicationStage {
PreparedRename,
Rename,
Remove,
/// Pause an accepted native read/open before its original syscall.
NativeReadOpen,
}
#[cfg(all(feature = "test-util", not(windows)))]
@@ -377,6 +380,7 @@ impl LocalPublicationPause {
LocalPublicationStage::PreparedRename => prepared_publication_test_hooks::Stage::PreparedRename,
LocalPublicationStage::Rename => prepared_publication_test_hooks::Stage::Rename,
LocalPublicationStage::Remove => prepared_publication_test_hooks::Stage::Remove,
LocalPublicationStage::NativeReadOpen => prepared_publication_test_hooks::Stage::NativeReadOpen,
};
let (entered_tx, entered) = oneshot::channel();
let (release, release_rx) = std::sync::mpsc::channel::<()>();
@@ -754,13 +758,13 @@ impl OpenedDstDirFsyncGroup {
struct DstDirFsyncWaiter {
namespace_owner: Option<Arc<dyn Send + Sync>>,
resource_unit: Option<rustfs_concurrency::adaptive::AdaptivePermit>,
resource_unit: Option<super::io_admission::ResourceIoUnit>,
result_tx: oneshot::Sender<SharedDstDirFsyncResult>,
}
struct DstDirFsyncOwner {
_namespace: Option<Arc<dyn Send + Sync>>,
_resource: Option<rustfs_concurrency::adaptive::AdaptivePermit>,
_resource: Option<super::io_admission::ResourceIoUnit>,
}
#[derive(Clone)]
@@ -1202,7 +1206,7 @@ type FileFdatasyncGroupKey = (usize, bool);
struct FileFdatasyncWaiter {
files: Vec<PathBuf>,
resource_unit: Option<rustfs_concurrency::adaptive::AdaptivePermit>,
resource_unit: Option<super::io_admission::ResourceIoUnit>,
enqueued_at: Option<Instant>,
wait_role: &'static str,
result_tx: oneshot::Sender<SharedFileFdatasyncResult>,
@@ -9299,49 +9303,66 @@ mod tests {
#[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;
use crate::disk::io_admission::{test_resource_units, with_optional_resource_unit};
let mut retained = Vec::new();
for lane in [None, Some(DiskBlockingLane::Foreground), Some(DiskBlockingLane::Fsync)] {
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(())
};
match lane {
Some(lane) => {
DiskBlockingQueue::new()
.run_on_lane(lane, "native operation failed", move |_| operation())
.await
for physical_count in [1, 2] {
for lane in [
None,
Some(DiskBlockingLane::Foreground),
Some(DiskBlockingLane::Isolated),
Some(DiskBlockingLane::Fsync),
] {
let (unit, gates) = test_resource_units(physical_count);
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_optional_resource_unit(Some(unit), async move {
let operation = move || {
started_tx.send(()).expect("native operation entered");
let _ = release_rx.blocking_recv();
let _ = finished_tx.send(());
Ok(())
};
match lane {
Some(lane) => {
DiskBlockingQueue::new()
.run_on_lane(lane, "native operation failed", move |_| operation())
.await
}
None => fsync_spawn_blocking(operation).await.expect("native operation joined"),
}
None => 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");
}));
started_rx.await.expect("native operation started");
worker.abort();
assert!(worker.await.expect_err("caller cancelled").is_cancelled());
retained.push(gates.iter().map(|gate| gate.snapshot().active).collect::<Vec<_>>());
// 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 gates.iter().any(|gate| gate.snapshot().active != 0) {
tokio::task::yield_now().await;
}
})
.await
.expect("native completion returns the unit");
}
}
assert_eq!(retained, [1, 1, 1], "fsync and both queue lanes must retain their own producer's unit");
assert_eq!(
retained,
[
vec![1],
vec![1],
vec![1],
vec![1],
vec![1, 1],
vec![1, 1],
vec![1, 1],
vec![1, 1]
],
"native dispatch and all queue lanes must retain their own producer's units"
);
}
#[tokio::test]
+4 -4
View File
@@ -76,14 +76,14 @@ impl DriverReadBudget {
}
}
pub(super) fn start_driver(
pub(super) fn start_driver<O: Send + Unpin + 'static>(
&self,
entries: u32,
shards: usize,
) -> Result<rustfs_uring::UringDriver, rustfs_uring::ProbeFailure> {
) -> Result<rustfs_uring::OwnedUringDriver<O>, rustfs_uring::ProbeFailure> {
match self {
Self::Disabled => rustfs_uring::UringDriver::probe_and_start_sharded(entries, shards),
Self::Limited { pool, per_driver } => rustfs_uring::UringDriver::probe_and_start_with_shared_budget(
Self::Disabled => rustfs_uring::OwnedUringDriver::probe_and_start_sharded(entries, shards),
Self::Limited { pool, per_driver } => rustfs_uring::OwnedUringDriver::probe_and_start_with_shared_budget(
entries,
shards,
rustfs_uring::ReadLimits {
+423 -25
View File
@@ -98,6 +98,9 @@ pin_project! {
// contiguous on-disk `[hash][data]` block so both are pulled in a single
// pass; grown lazily and never shrunk.
buf: Vec<u8>,
// Free scratch storage before refunding its allocation.
buf_memory: Option<crate::disk::background_memory::Reservation>,
memory_budget: Option<std::sync::Arc<crate::disk::BackgroundMemoryBudget>>,
// Reused owned chunk vector for the remote HTTP fast path. Keeping the
// allocation with the reader avoids allocating once per bitrot block.
chunks: Vec<bytes::Bytes>,
@@ -118,6 +121,8 @@ where
hash_algo: algo,
shard_size,
buf: Vec::new(),
buf_memory: None,
memory_budget: crate::disk::current_resource_unit().and_then(|unit| unit.memory_budget()),
chunks: Vec::new(),
skip_verify,
last_verify_duration: Duration::ZERO,
@@ -209,6 +214,7 @@ where
/// On `Ok` the block is `self.buf[..need]`; the caller re-borrows it so the
/// verification can also borrow `self` immutably.
async fn read_scratch_block(&mut self, need: usize, want: usize) -> std::io::Result<()> {
reserve_bitrot_scratch(&mut self.buf, &mut self.buf_memory, self.memory_budget.as_ref(), need)?;
if self.buf.len() < need {
self.buf.resize(need, 0);
}
@@ -229,6 +235,28 @@ where
}
}
// Scratch growth happens only before a complete block or an empty chunk tail.
// Old contents are consumed; keep both allocations charged until replacement.
fn reserve_bitrot_scratch(
buf: &mut Vec<u8>,
reservation: &mut Option<crate::disk::background_memory::Reservation>,
budget: Option<&std::sync::Arc<crate::disk::BackgroundMemoryBudget>>,
capacity: usize,
) -> std::io::Result<()> {
if let Some(budget) = budget
&& buf.capacity() < capacity
{
let memory = budget.reserve(capacity)?;
let mut replacement = Vec::new();
replacement
.try_reserve_exact(capacity)
.map_err(|_| std::io::Error::from(std::io::ErrorKind::OutOfMemory))?;
*buf = replacement;
*reservation = Some(memory);
}
Ok(())
}
/// A truncated shard is `UnexpectedEof`, not a short success (backlog#799 B2).
fn short_shard_read(got: usize, want: usize) -> std::io::Error {
error!(
@@ -339,6 +367,8 @@ where
let inner = &mut self.inner;
let chunks = &mut self.chunks;
let tail_buf = &mut self.buf;
let tail_memory = &mut self.buf_memory;
let memory_budget = self.memory_budget.as_ref();
let mut received = 0usize;
poll_fn(|cx| {
for _ in 0..MAX_CHUNK_POLLS_PER_YIELD {
@@ -381,7 +411,9 @@ where
if chunks.len() == MAX_RETAINED_CHUNKS_PER_BLOCK {
if tail_buf.is_empty() {
tail_buf.reserve_exact(need - (received - chunk.len()));
let tail_capacity = need - (received - chunk.len());
reserve_bitrot_scratch(tail_buf, tail_memory, memory_budget, tail_capacity)?;
tail_buf.reserve_exact(tail_capacity);
}
tail_buf.extend_from_slice(&chunk);
} else {
@@ -657,7 +689,12 @@ pub async fn bitrot_verify<R: AsyncRead + Unpin + Send>(
algo: HashAlgorithm,
mut shard_size: usize,
) -> std::io::Result<()> {
let mut hash_buf = vec![0; algo.size()];
// Every supported digest fits on the stack; only shard payloads need leases.
let mut hash = [0; 64];
let hash_buf = &mut hash[..algo.size()];
let mut background_buffer = crate::disk::current_resource_unit()
.and_then(|unit| unit.memory_budget())
.map(|budget| crate::disk::background_memory::BackgroundBuffer::new(Some(budget)));
let mut left = want_size;
if left != bitrot_shard_file_size(part_size, shard_size, algo.clone()) {
@@ -665,21 +702,30 @@ pub async fn bitrot_verify<R: AsyncRead + Unpin + Send>(
}
while left > 0 {
let n = r.read_exact(&mut hash_buf).await?;
let n = r.read_exact(hash_buf).await?;
left -= n;
if left < shard_size {
shard_size = left;
}
let mut buf = get_ec_buffer(shard_size);
buf.resize(shard_size, 0);
let read = r.read_exact(&mut buf).await?;
let mut pooled = None;
let buf = if let Some(buffer) = background_buffer.as_mut() {
buffer.ensure_len(shard_size)?;
&mut buffer.as_mut_slice()[..shard_size]
} else {
let buffer = pooled.insert(get_ec_buffer(shard_size));
buffer.resize(shard_size, 0);
buffer.as_mut_slice()
};
let read = r.read_exact(buf).await?;
let actual_hash = algo.hash_encode(&buf);
let actual_hash = algo.hash_encode(buf);
let hash_ok = actual_hash.as_ref() == &hash_buf[0..n];
drop(actual_hash); // 释放借用
return_ec_buffer(buf);
if let Some(buf) = pooled {
return_ec_buffer(buf);
}
if !hash_ok {
return Err(std::io::Error::other("bitrot hash mismatch"));
}
@@ -698,15 +744,74 @@ pub async fn bitrot_verify<R: AsyncRead + Unpin + Send>(
/// Custom writer enum that supports inline buffer storage
pub enum CustomWriter {
/// Inline buffer writer - stores data in memory
InlineBuffer(Vec<u8>),
InlineBuffer(InlineBuffer),
/// Disk-based writer using tokio file
Other(Box<dyn AsyncWrite + Unpin + Send + Sync>),
}
/// Inline payload storage retains its background allocation through export.
pub struct InlineBuffer {
foreground: Vec<u8>,
background: Option<crate::disk::background_memory::BackgroundBuffer>,
filled: usize,
}
impl InlineBuffer {
fn new() -> Self {
let background = crate::disk::current_resource_unit()
.and_then(|unit| unit.memory_budget())
.map(|budget| crate::disk::background_memory::BackgroundBuffer::new(Some(budget)));
Self {
foreground: Vec::new(),
background,
filled: 0,
}
}
fn as_slice(&self) -> &[u8] {
match &self.background {
Some(buffer) => &buffer.as_slice()[..self.filled],
None => &self.foreground,
}
}
fn append(&mut self, bufs: &[IoSlice<'_>]) -> std::io::Result<usize> {
let total = bufs
.iter()
.try_fold(0usize, |length, buf| length.checked_add(buf.len()).ok_or(std::io::ErrorKind::OutOfMemory))?;
if let Some(buffer) = &mut self.background {
let length = self.filled.checked_add(total).ok_or(std::io::ErrorKind::OutOfMemory)?;
if length > buffer.as_slice().len() {
let capacity = buffer.as_slice().len().checked_mul(2).unwrap_or(length).max(length);
// Admission includes both allocations until the prefix has been copied.
buffer.grow(capacity, self.filled)?;
}
let mut offset = self.filled;
for buf in bufs {
buffer.as_mut_slice()[offset..offset + buf.len()].copy_from_slice(buf);
offset += buf.len();
}
self.filled = length;
} else {
for buf in bufs {
self.foreground.extend_from_slice(buf);
}
}
Ok(total)
}
fn into_bytes(self) -> bytes::Bytes {
match self.background {
Some(mut buffer) => buffer.take_bytes(self.filled),
None => bytes::Bytes::from(self.foreground),
}
}
}
impl CustomWriter {
/// Create a new inline buffer writer
pub fn new_inline_buffer() -> Self {
Self::InlineBuffer(Vec::new())
Self::InlineBuffer(InlineBuffer::new())
}
/// Create a new disk writer from any AsyncWrite implementation
@@ -720,15 +825,15 @@ impl CustomWriter {
/// Get the inline buffer data if this is an inline buffer writer
pub fn get_inline_data(&self) -> Option<&[u8]> {
match self {
Self::InlineBuffer(data) => Some(data),
Self::InlineBuffer(data) => Some(data.as_slice()),
Self::Other(_) => None,
}
}
/// Extract the inline buffer data, consuming the writer
pub fn into_inline_data(self) -> Option<Vec<u8>> {
pub fn into_inline_data(self) -> Option<bytes::Bytes> {
match self {
Self::InlineBuffer(data) => Some(data),
Self::InlineBuffer(data) => Some(data.into_bytes()),
Self::Other(_) => None,
}
}
@@ -741,10 +846,7 @@ impl AsyncWrite for CustomWriter {
buf: &[u8],
) -> std::task::Poll<std::io::Result<usize>> {
match self.get_mut() {
Self::InlineBuffer(data) => {
data.extend_from_slice(buf);
std::task::Poll::Ready(Ok(buf.len()))
}
Self::InlineBuffer(data) => std::task::Poll::Ready(data.append(&[IoSlice::new(buf)])),
Self::Other(writer) => {
let pinned_writer = std::pin::Pin::new(writer.as_mut());
pinned_writer.poll_write(cx, buf)
@@ -778,13 +880,7 @@ impl AsyncWrite for CustomWriter {
bufs: &[IoSlice<'_>],
) -> std::task::Poll<std::io::Result<usize>> {
match self.get_mut() {
Self::InlineBuffer(data) => {
let total = bufs.iter().map(|buf| buf.len()).sum::<usize>();
for buf in bufs {
data.extend_from_slice(buf);
}
std::task::Poll::Ready(Ok(total))
}
Self::InlineBuffer(data) => std::task::Poll::Ready(data.append(bufs)),
Self::Other(writer) => {
let pinned_writer = std::pin::Pin::new(writer.as_mut());
pinned_writer.poll_write_vectored(cx, bufs)
@@ -859,7 +955,7 @@ impl BitrotWriterWrapper {
}
/// Extract the inline buffer data, consuming the wrapper
pub fn into_inline_data(self) -> Option<Vec<u8>> {
pub fn into_inline_data(self) -> Option<bytes::Bytes> {
match self.writer_type {
WriterType::InlineBuffer => {
let writer = self.bitrot_writer.into_inner();
@@ -1506,6 +1602,184 @@ mod tests {
assert_eq!(data, &out[..]);
}
#[tokio::test]
async fn bitrot_reader_scratch_budget_preserves_growth_and_reuse() {
use crate::disk::{BackgroundMemoryBudget, test_resource_units, with_optional_resource_unit};
let algo = HashAlgorithm::HighwayHash256S;
let mut encoded = Vec::new();
for data in [&b"abcd"[..], &b"abcdefgh"[..], &b"ijklmnop"[..]] {
encoded.extend_from_slice(algo.hash_encode(data).as_ref());
encoded.extend_from_slice(data);
}
let memory = BackgroundMemoryBudget::new();
let (resources, _) = test_resource_units(1);
let mut reader = with_optional_resource_unit(Some(resources.with_memory_budget(memory.clone())), async {
BitrotReader::new(Cursor::new(encoded), 8, algo, false)
})
.await;
let mut first = [0; 4];
assert_eq!(
reader.read(&mut first).await.expect_err("closed scratch budget").kind(),
io::ErrorKind::WouldBlock
);
assert_eq!(reader.inner.position(), 0, "allocation denial must not consume a frame");
memory.set_sampled_limit(36, Duration::from_secs(30));
reader.read(&mut first).await.expect("first scratch allocation fits");
assert_eq!(&first, b"abcd");
assert_eq!(memory.snapshot().used, 36);
memory.set_sampled_limit(75, Duration::from_secs(30));
let mut next = [0; 8];
assert_eq!(
reader
.read(&mut next)
.await
.expect_err("growth requires old36 plus new40")
.kind(),
io::ErrorKind::WouldBlock
);
assert_eq!(reader.inner.position(), 36);
assert_eq!(memory.snapshot().used, 36);
memory.set_sampled_limit(76, Duration::from_secs(30));
reader.read(&mut next).await.expect("growth now fits");
assert_eq!(&next, b"abcdefgh");
assert_eq!(memory.snapshot().used, 40);
memory.set_sampled_limit(0, Duration::ZERO);
reader.read(&mut next).await.expect("reuse after hot shrink");
assert_eq!(&next, b"ijklmnop");
assert_eq!(memory.snapshot().used, 40);
drop(reader);
assert_eq!(memory.snapshot().used, 0);
}
#[tokio::test]
async fn bitrot_chunk_tail_allocation_stays_charged_until_reader_drop() {
use crate::disk::{BackgroundMemoryBudget, test_resource_units, with_optional_resource_unit};
let algo = HashAlgorithm::HighwayHash256S;
let payload = vec![37; 256];
let encoded = encode_one_block(&payload, payload.len(), algo.clone()).await;
let tail_capacity = encoded.len() - MAX_RETAINED_CHUNKS_PER_BLOCK;
for fits in [false, true] {
let memory = BackgroundMemoryBudget::new();
memory.set_sampled_limit(if fits { tail_capacity as u64 } else { 0 }, Duration::from_secs(30));
let (resources, _) = test_resource_units(1);
let mut reader = with_optional_resource_unit(Some(resources.with_memory_budget(memory.clone())), async {
BitrotReader::new(GeneratedChunkSource::new(encoded.clone(), 1), payload.len(), algo.clone(), false)
})
.await;
let mut output = Vec::with_capacity(payload.len());
let result = reader.read_appending(&mut output, payload.len()).await;
if fits {
assert_eq!(result.expect("accounted tail fits"), payload.len());
assert_eq!(output, payload);
assert_eq!(memory.snapshot().used, tail_capacity as u64);
memory.set_sampled_limit(0, Duration::ZERO);
assert_eq!(memory.snapshot().used, tail_capacity as u64);
} else {
assert_eq!(result.expect_err("closed tail budget").kind(), io::ErrorKind::WouldBlock);
assert!(output.is_empty(), "an unverified fragment must not reach output");
assert_eq!(memory.snapshot().used, 0);
}
drop(reader);
assert_eq!(memory.snapshot().used, 0);
}
}
#[tokio::test]
async fn bitrot_verify_background_memory_is_bounded_and_reused() {
use crate::disk::{BackgroundMemoryBudget, test_resource_units, with_optional_resource_unit};
struct ShrinkingReader {
data: Cursor<Vec<u8>>,
memory: Arc<BackgroundMemoryBudget>,
first_block_end: u64,
peak: Arc<std::sync::atomic::AtomicU64>,
}
impl AsyncRead for ShrinkingReader {
fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, out: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
self.peak.fetch_max(self.memory.snapshot().used, Ordering::SeqCst);
let result = Pin::new(&mut self.data).poll_read(cx, out);
if self.data.position() >= self.first_block_end {
self.memory.set_sampled_limit(0, Duration::ZERO);
}
result
}
}
for algo in [
HashAlgorithm::None,
HashAlgorithm::HighwayHash256S,
HashAlgorithm::HighwayHash256SLegacy,
] {
let payload = [37; 17];
let mut writer = BitrotWriter::new(Cursor::new(Vec::new()), 8, algo.clone());
for chunk in payload.chunks(8) {
writer.write(chunk).await.expect("encoded bitrot block");
}
writer.shutdown().await.expect("complete fixture");
let encoded = writer.into_inner().into_inner();
let memory = BackgroundMemoryBudget::new();
let (resources, _) = test_resource_units(1);
let error = with_optional_resource_unit(
Some(resources.with_memory_budget(memory.clone())),
bitrot_verify(Cursor::new(&encoded), encoded.len(), payload.len(), algo.clone(), 8),
)
.await
.expect_err("a closed background budget must reject bitrot payload allocation");
assert_eq!(error.kind(), io::ErrorKind::WouldBlock);
assert_eq!(memory.snapshot().used, 0);
memory.set_sampled_limit(8, Duration::from_secs(30));
let peak = Arc::new(std::sync::atomic::AtomicU64::new(0));
let reader = ShrinkingReader {
data: Cursor::new(encoded.clone()),
memory: memory.clone(),
first_block_end: (algo.size() + 8) as u64,
peak: peak.clone(),
};
let (resources, _) = test_resource_units(1);
with_optional_resource_unit(
Some(resources.with_memory_budget(memory.clone())),
bitrot_verify(reader, encoded.len(), payload.len(), algo, 8),
)
.await
.expect("reuse the same allocation after shrink, including the short tail");
assert_eq!(peak.load(Ordering::SeqCst), 8);
assert_eq!(memory.snapshot().limit, 0);
assert_eq!(
memory.snapshot().used,
0,
"the verifier frees its payload instead of retaining it in an uncharged pool"
);
}
}
#[tokio::test]
async fn bitrot_verify_cancelled_payload_read_refunds_memory() {
use crate::disk::{BackgroundMemoryBudget, test_resource_units, with_optional_resource_unit};
let algo = HashAlgorithm::HighwayHash256S;
let (mut sender, reader) = tokio::io::duplex(64);
sender
.write_all(algo.hash_encode(&[37; 8]).as_ref())
.await
.expect("send only hash");
let memory = BackgroundMemoryBudget::new();
memory.set_sampled_limit(8, Duration::from_secs(30));
let (resources, _) = test_resource_units(1);
let mut verify = Box::pin(with_optional_resource_unit(
Some(resources.with_memory_budget(memory.clone())),
bitrot_verify(reader, algo.size() + 8, 8, algo, 8),
));
assert!(futures::poll!(&mut verify).is_pending(), "the payload read waits after the complete hash");
let retained = memory.snapshot().used;
memory.set_sampled_limit(0, Duration::ZERO);
drop(verify);
drop(sender);
assert_eq!(retained, 8, "pending verification owns its payload allocation");
assert_eq!(memory.snapshot().used, 0);
}
#[tokio::test]
async fn bitrot_verify_accepts_valid_shard_file_and_rejects_size_or_hash_mismatch() {
let data = b"bitrot verify covers every shard";
@@ -2403,4 +2677,128 @@ mod tests {
assert_eq!(output, vec![9u8]);
}
}
mod inline_memory_regressions {
use super::*;
use crate::disk::background_memory::is_memory_pressure;
use crate::disk::{BackgroundMemoryBudget, test_resource_units, with_optional_resource_unit};
use std::sync::Arc;
use std::time::Duration;
use tokio::io::AsyncWriteExt;
async fn writer_for(memory: Arc<BackgroundMemoryBudget>) -> CustomWriter {
let (resources, _) = test_resource_units(1);
with_optional_resource_unit(Some(resources.with_memory_budget(memory)), async { CustomWriter::new_inline_buffer() })
.await
}
#[tokio::test]
async fn inline_writer_rejects_closed_and_expired_owners_before_first_allocation() {
for expired in [false, true] {
let memory = BackgroundMemoryBudget::new();
if expired {
memory.set_sampled_limit(128, Duration::ZERO);
}
let mut writer = writer_for(memory.clone()).await;
let error = writer.write_all(&[7; 64]).await.expect_err("no live grant");
assert_eq!(error.kind(), io::ErrorKind::WouldBlock);
assert!(is_memory_pressure(&error));
assert_eq!(writer.get_inline_data(), Some(&[][..]));
assert_eq!(memory.snapshot().used, 0);
}
}
#[tokio::test]
async fn inline_writer_growth_charges_old_plus_new_and_is_atomic_for_vectored_writes() {
let memory = BackgroundMemoryBudget::new();
memory.set_sampled_limit(128, Duration::from_secs(30));
let mut writer = writer_for(memory.clone()).await;
writer.write_all(&[1; 64]).await.expect("first allocation");
assert_eq!(memory.snapshot().used, 64);
let second = [2; 32];
let third = [3; 32];
let chunks = [IoSlice::new(&second), IoSlice::new(&third)];
let error = writer.write_vectored(&chunks).await.expect_err("64 old + 128 new exceed 128");
assert!(is_memory_pressure(&error));
assert_eq!(writer.get_inline_data(), Some(&[1; 64][..]));
assert_eq!(memory.snapshot().used, 64);
memory.set_sampled_limit(192, Duration::from_secs(30));
assert_eq!(writer.write_vectored(&chunks).await.expect("old and new both fit"), 64);
assert_eq!(memory.snapshot().used, 128);
let data = writer.get_inline_data().expect("inline payload");
assert_eq!(&data[..64], &[1; 64]);
assert_eq!(&data[64..96], &second);
assert_eq!(&data[96..], &third);
drop(writer);
assert_eq!(memory.snapshot().used, 0);
}
#[tokio::test]
async fn inline_writer_expiry_preserves_charged_storage_but_blocks_growth() {
let memory = BackgroundMemoryBudget::new();
memory.set_sampled_limit(192, Duration::from_secs(30));
let mut writer = writer_for(memory.clone()).await;
writer.write_all(&[5; 64]).await.expect("existing storage");
memory.set_sampled_limit(192, Duration::ZERO);
let error = writer
.write_all(&[9])
.await
.expect_err("expired sample cannot authorize replacement");
assert!(is_memory_pressure(&error));
assert_eq!(writer.get_inline_data(), Some(&[5; 64][..]));
assert_eq!(memory.snapshot().used, 64);
writer.write_all(&[]).await.expect("empty write needs no allocation");
writer.flush().await.expect("flush does not revoke storage");
writer.shutdown().await.expect("shutdown does not revoke storage");
let bytes = writer.into_inline_data().expect("export after expiry");
assert_eq!(bytes.as_ref(), &[5; 64]);
assert_eq!(memory.snapshot().used, 64);
drop(bytes);
assert_eq!(memory.snapshot().used, 0);
}
#[tokio::test]
async fn inline_writer_export_keeps_receipt_through_clones_and_the_last_slice() {
let memory = BackgroundMemoryBudget::new();
memory.set_sampled_limit(64, Duration::from_secs(30));
let mut writer = writer_for(memory.clone()).await;
writer.write_all(&[11; 64]).await.expect("allocate payload");
let bytes = writer.into_inline_data().expect("export bytes");
let clone = bytes.clone();
let slice = bytes.slice(3..7);
memory.set_sampled_limit(0, Duration::ZERO);
drop(bytes);
drop(clone);
assert_eq!(slice.as_ref(), &[11; 4]);
assert_eq!(memory.snapshot().used, 64);
drop(slice);
assert_eq!(memory.snapshot().used, 0);
}
#[tokio::test]
async fn inline_writer_keeps_the_constructing_owner_after_leaving_its_scope() {
let first = BackgroundMemoryBudget::new();
first.set_sampled_limit(64, Duration::from_secs(30));
let mut writer = writer_for(first.clone()).await;
let second = BackgroundMemoryBudget::new();
let (resources, _) = test_resource_units(1);
with_optional_resource_unit(Some(resources.with_memory_budget(second.clone())), async {
writer.write_all(&[13; 64]).await.expect("writer still owns the first grant");
})
.await;
assert_eq!(first.snapshot().used, 64);
assert_eq!(second.snapshot().used, 0);
drop(writer);
assert_eq!(first.snapshot().used, 0);
}
#[tokio::test]
async fn foreground_inline_writer_preserves_payload_and_empty_writes() {
let mut writer = CustomWriter::new_inline_buffer();
writer.write_all(&[]).await.expect("empty foreground write");
writer.write_all(b"inline body").await.expect("foreground write");
let bytes = writer.into_inline_data().expect("foreground export");
assert_eq!(bytes.as_ref(), b"inline body");
}
}
}
+208 -33
View File
@@ -16,6 +16,8 @@
//! Supports legacy (reed-solomon-simd) for reading/healing old-version files.
//!
use crate::disk::background_memory::{BackgroundBuffer, Reservation};
use crate::disk::{BackgroundMemoryBudget, current_resource_unit};
use bytes::{Bytes, BytesMut};
use reed_solomon_simd;
use rustfs_erasure_codec::galois_8::ReedSolomon;
@@ -144,12 +146,54 @@ pub fn calc_shard_size_legacy(block_size: usize, data_shards: usize) -> usize {
/// Reed-Solomon encoder for legacy (main branch) format using reed-solomon-simd.
/// Used when decoding/encoding files with uses_legacy_checksum == true.
struct LegacyWorkspace<T> {
codec: T,
// The cached payload is freed before its reservation.
memory: Option<Reservation>,
payload_bound: usize,
}
impl<T> LegacyWorkspace<T> {
fn prepare(
cached: Option<Self>,
payload_bound: usize,
budget: Option<&Arc<BackgroundMemoryBudget>>,
create: impl FnOnce() -> io::Result<T>,
reset: impl FnOnce(&mut T) -> io::Result<()>,
) -> io::Result<Self> {
if let Some(mut cached) = cached {
let reusable = match (&cached.memory, budget) {
(None, None) => true,
(Some(memory), budget) => {
payload_bound <= cached.payload_bound && budget.is_none_or(|budget| memory.belongs_to(budget))
}
// Uncharged foreground caches can have larger retained Vec
// capacities after reset; start a measured background workspace.
(None, Some(_)) => false,
};
if reusable && reset(&mut cached.codec).is_ok() {
return Ok(cached);
}
// A fresh allocation avoids hidden geometric Vec growth. No payload
// references escape codec calls, so dropping this cache is safe.
drop(cached);
}
let memory = budget.map(|budget| budget.reserve(payload_bound)).transpose()?;
let codec = create()?;
Ok(Self {
codec,
memory,
payload_bound,
})
}
}
struct LegacyReedSolomonEncoder {
data_shards: usize,
parity_shards: usize,
cache_workspaces: bool,
encoder_cache: RwLock<Option<reed_solomon_simd::ReedSolomonEncoder>>,
decoder_cache: RwLock<Option<reed_solomon_simd::ReedSolomonDecoder>>,
encoder_cache: RwLock<Option<LegacyWorkspace<reed_solomon_simd::ReedSolomonEncoder>>>,
decoder_cache: RwLock<Option<LegacyWorkspace<reed_solomon_simd::ReedSolomonDecoder>>>,
}
impl LegacyReedSolomonEncoder {
@@ -199,20 +243,29 @@ impl LegacyReedSolomonEncoder {
.write()
.map_err(|_| io::Error::other("Failed to acquire encoder cache lock"))?
.take();
let mut encoder = {
match cached_encoder {
Some(mut cached) => {
if cached.reset(self.data_shards, self.parity_shards, shard_len).is_err() {
reed_solomon_simd::ReedSolomonEncoder::new(self.data_shards, self.parity_shards, shard_len)
.map_err(|e| io::Error::other(format!("Failed to create SIMD encoder: {e:?}")))?
} else {
cached
}
}
None => reed_solomon_simd::ReedSolomonEncoder::new(self.data_shards, self.parity_shards, shard_len)
.map_err(|e| io::Error::other(format!("Failed to create SIMD encoder: {e:?}")))?,
}
};
let budget = current_resource_unit().and_then(|resources| resources.memory_budget());
// reed-solomon-simd 3.1.0 Shards::resize allocates 64-byte units. A
// fresh Vec has at least four units; the decoder work-count bound also
// covers both encoder rate strategies.
let payload_bound = self
.logical_shard_bytes_upper_bound(shard_len)
.ok_or(io::ErrorKind::OutOfMemory)?
.max(256);
let mut workspace = LegacyWorkspace::prepare(
cached_encoder,
payload_bound,
budget.as_ref(),
|| {
reed_solomon_simd::ReedSolomonEncoder::new(self.data_shards, self.parity_shards, shard_len)
.map_err(|e| io::Error::other(format!("Failed to create SIMD encoder: {e:?}")))
},
|encoder| {
encoder
.reset(self.data_shards, self.parity_shards, shard_len)
.map_err(|e| io::Error::other(format!("Failed to reset SIMD encoder: {e:?}")))
},
)?;
let encoder = &mut workspace.codec;
for (i, shard) in shards.iter().enumerate().take(self.data_shards) {
encoder
.add_original_shard(shard)
@@ -233,7 +286,7 @@ impl LegacyReedSolomonEncoder {
.write()
.map_err(|_| io::Error::other("Failed to return encoder to cache"))?;
if cache.is_none() {
*cache = Some(encoder);
*cache = Some(workspace);
}
}
Ok(())
@@ -254,21 +307,26 @@ impl LegacyReedSolomonEncoder {
.write()
.map_err(|_| io::Error::other("Failed to acquire decoder cache lock"))?
.take();
let mut decoder = {
match cached_decoder {
Some(mut cached_decoder) => {
if let Err(e) = cached_decoder.reset(self.data_shards, self.parity_shards, shard_len) {
warn!("Failed to reset SIMD decoder: {:?}, creating new one", e);
reed_solomon_simd::ReedSolomonDecoder::new(self.data_shards, self.parity_shards, shard_len)
.map_err(|e| io::Error::other(format!("Failed to create SIMD decoder: {e:?}")))?
} else {
cached_decoder
}
}
None => reed_solomon_simd::ReedSolomonDecoder::new(self.data_shards, self.parity_shards, shard_len)
.map_err(|e| io::Error::other(format!("Failed to create SIMD decoder: {e:?}")))?,
}
};
let budget = current_resource_unit().and_then(|resources| resources.memory_budget());
let payload_bound = self
.logical_shard_bytes_upper_bound(shard_len)
.ok_or(io::ErrorKind::OutOfMemory)?
.max(256);
let mut workspace = LegacyWorkspace::prepare(
cached_decoder,
payload_bound,
budget.as_ref(),
|| {
reed_solomon_simd::ReedSolomonDecoder::new(self.data_shards, self.parity_shards, shard_len)
.map_err(|e| io::Error::other(format!("Failed to create SIMD decoder: {e:?}")))
},
|decoder| {
decoder
.reset(self.data_shards, self.parity_shards, shard_len)
.map_err(|e| io::Error::other(format!("Failed to reset SIMD decoder: {e:?}")))
},
)?;
let decoder = &mut workspace.codec;
for (i, shard_opt) in shards.iter().enumerate() {
if let Some(shard) = shard_opt {
@@ -308,7 +366,7 @@ impl LegacyReedSolomonEncoder {
.write()
.map_err(|_| io::Error::other("Failed to return decoder to cache"))?;
if cache.is_none() {
*cache = Some(decoder);
*cache = Some(workspace);
}
}
@@ -333,6 +391,13 @@ impl LegacyReedSolomonEncoder {
return Ok(true);
}
let expected_bytes = shards.iter().try_fold(0usize, |total, shard| {
total
.checked_add(shard.len())
.ok_or_else(|| io::Error::from(io::ErrorKind::OutOfMemory))
})?;
let memory = current_resource_unit().and_then(|resources| resources.memory_budget());
let _expected_memory = memory.as_ref().map(|memory| memory.reserve(expected_bytes)).transpose()?;
let mut expected = shards.iter().map(|shard| Some(shard.to_vec())).collect::<Vec<_>>();
self.encode_parity(&mut expected)?;
@@ -431,6 +496,17 @@ impl ReedSolomonEncoder {
}
if let Some(ref rs) = self.encoder {
if let Some(memory) = current_resource_unit().and_then(|resources| resources.memory_budget())
&& let Some(shard_len) = shards.first().map(|shard| shard.len()).filter(|len| *len != 0)
{
let scratch_len = shard_len.checked_mul(self.parity_shards).ok_or(io::ErrorKind::OutOfMemory)?;
let mut scratch = BackgroundBuffer::new(Some(memory));
scratch.ensure_len(scratch_len)?;
let mut buffers: SmallVec<[&mut [u8]; 16]> = scratch.as_mut_slice().chunks_mut(shard_len).collect();
return rs
.verify_with_buffer(shards, &mut buffers)
.map_err(|e| io::Error::other(format!("Reed-Solomon verify failed: {e:?}")));
}
rs.verify(shards)
.map_err(|e| io::Error::other(format!("Reed-Solomon verify failed: {e:?}")))
} else {
@@ -1200,6 +1276,105 @@ impl Erasure {
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn background_legacy_workspaces_keep_their_memory_in_cache() {
use crate::disk::{test_resource_units, with_optional_resource_unit};
use std::time::Duration;
let codec = LegacyReedSolomonEncoder::with_workspace_cache(2, 1, true).expect("isolated cached codec");
let memory = BackgroundMemoryBudget::new();
memory.set_sampled_limit(512, Duration::from_secs(30));
let mut shards = vec![vec![7; 8], vec![9; 8], vec![0; 8]];
let (resources, _) = test_resource_units(1);
with_optional_resource_unit(Some(resources.with_memory_budget(memory.clone())), async {
codec
.encode(shards.iter_mut().map(Vec::as_mut_slice).collect())
.expect("accounted encoder");
})
.await;
assert_eq!(memory.snapshot().used, 256, "fresh SIMD work storage uses at most four 64-byte units");
memory.set_sampled_limit(0, Duration::ZERO);
let (resources, _) = test_resource_units(1);
with_optional_resource_unit(Some(resources.with_memory_budget(memory.clone())), async {
codec
.encode(shards.iter_mut().map(Vec::as_mut_slice).collect())
.expect("cached payload survives shrink");
})
.await;
assert_eq!(memory.snapshot().used, 256);
memory.set_sampled_limit(512, Duration::from_secs(30));
let mut decoded = vec![None, Some(shards[1].clone()), Some(shards[2].clone())];
let (resources, _) = test_resource_units(1);
with_optional_resource_unit(Some(resources.with_memory_budget(memory.clone())), async {
codec.reconstruct_data(&mut decoded).expect("accounted decoder");
})
.await;
assert_eq!(decoded[0].as_deref(), Some(shards[0].as_slice()));
assert_eq!(memory.snapshot().used, 512, "both cached codec payloads remain charged");
drop(codec);
assert_eq!(memory.snapshot().used, 0);
}
#[tokio::test]
async fn background_legacy_cache_cannot_reuse_a_different_memory_owner() {
use crate::disk::{test_resource_units, with_optional_resource_unit};
use std::time::Duration;
let codec = LegacyReedSolomonEncoder::with_workspace_cache(2, 1, true).expect("isolated cached codec");
let first = BackgroundMemoryBudget::new();
first.set_sampled_limit(256, Duration::from_secs(30));
let mut shards = vec![vec![7; 8], vec![9; 8], vec![0; 8]];
let (resources, _) = test_resource_units(1);
with_optional_resource_unit(Some(resources.with_memory_budget(first.clone())), async {
codec
.encode(shards.iter_mut().map(Vec::as_mut_slice).collect())
.expect("first owner");
})
.await;
let second = BackgroundMemoryBudget::new();
let (resources, _) = test_resource_units(1);
let error = with_optional_resource_unit(Some(resources.with_memory_budget(second.clone())), async {
codec.encode(shards.iter_mut().map(Vec::as_mut_slice).collect())
})
.await
.expect_err("a foreign cache must not bypass a closed owner");
assert_eq!(error.kind(), io::ErrorKind::WouldBlock);
assert_eq!(first.snapshot().used, 0);
assert_eq!(second.snapshot().used, 0);
}
#[tokio::test]
async fn background_modern_verification_charges_scratch_and_preserves_invalid_input_errors() {
use crate::disk::{test_resource_units, with_optional_resource_unit};
use std::time::Duration;
let codec = ReedSolomonEncoder::new(2, 1).expect("modern codec");
let mut shards = vec![vec![7; 8], vec![9; 8], vec![0; 8]];
codec
.encode(shards.iter_mut().map(Vec::as_mut_slice).collect())
.expect("valid parity");
let refs: Vec<_> = shards.iter().map(Vec::as_slice).collect();
let memory = BackgroundMemoryBudget::new();
let (resources, _) = test_resource_units(1);
let error =
with_optional_resource_unit(Some(resources.with_memory_budget(memory.clone())), async { codec.verify(&refs) })
.await
.expect_err("closed owner rejects verification scratch");
assert_eq!(error.kind(), io::ErrorKind::WouldBlock);
memory.set_sampled_limit(8, Duration::from_secs(30));
let (resources, _) = test_resource_units(1);
with_optional_resource_unit(Some(resources.with_memory_budget(memory.clone())), async {
assert!(codec.verify(&refs).expect("exact parity scratch fits"));
assert!(
codec.verify(&[&[], &[9; 8], &[0; 8]]).is_err(),
"empty first shard remains an error instead of a panic"
);
})
.await;
assert_eq!(memory.snapshot().used, 0);
}
use proptest::collection::{btree_set, vec};
use proptest::prelude::*;
use std::pin::Pin;
+179 -24
View File
@@ -12,8 +12,10 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use crate::disk::background_memory::{Reservation, is_memory_pressure};
use crate::disk::disk_store::get_object_disk_read_timeout;
use crate::disk::error::{Error, Result};
use crate::disk::{BackgroundMemoryBudget, current_resource_unit};
use crate::erasure::coding::BitrotReader;
use crate::erasure::coding::BitrotWriterWrapper;
use crate::erasure::coding::encode::MultiWriter;
@@ -22,6 +24,7 @@ use futures::StreamExt;
use futures::stream::FuturesUnordered;
use std::io;
use std::io::ErrorKind;
use std::sync::Arc;
use std::time::Duration;
use tokio::io::AsyncRead;
use tracing::{trace, warn};
@@ -30,21 +33,53 @@ const LOG_COMPONENT_ECSTORE: &str = "ecstore";
const LOG_SUBSYSTEM_ERASURE: &str = "erasure";
const EVENT_ERASURE_HEAL_STARTED: &str = "erasure_heal_started";
struct HealShards {
// Drop payloads before refunding their reservations on every error path.
shards: Vec<Option<Vec<u8>>>,
memory: Vec<Option<Reservation>>,
budget: Option<Arc<BackgroundMemoryBudget>>,
}
impl HealShards {
fn reserve_missing(&mut self, shard_size: usize) -> io::Result<()> {
if let Some(budget) = &self.budget {
for (shard, memory) in self.shards.iter().zip(&mut self.memory) {
if shard.is_none() {
*memory = Some(budget.reserve(shard_size)?);
}
}
}
Ok(())
}
fn into_bytes(self) -> Vec<Bytes> {
self.shards
.into_iter()
.zip(self.memory)
.map(|(shard, memory)| Reservation::wrap(Bytes::from(shard.unwrap_or_default()), memory))
.collect()
}
}
async fn read_heal_shards<R>(
readers: &mut [Option<BitrotReader<R>>],
shard_size: usize,
read_timeout: Duration,
) -> (Vec<Option<Vec<u8>>>, Vec<Option<Error>>)
) -> Result<(HealShards, Vec<Option<Error>>)>
where
R: AsyncRead + Unpin + Send + Sync,
{
let num_readers = readers.len();
let mut shards = vec![None; num_readers];
let mut shards = HealShards {
shards: vec![None; num_readers],
memory: std::iter::repeat_with(|| None).take(num_readers).collect(),
budget: current_resource_unit().and_then(|resources| resources.memory_budget()),
};
let mut errs = vec![None; num_readers];
let mut retire_readers = Vec::new();
if shard_size == 0 {
return (shards, errs);
return Ok((shards, errs));
}
{
@@ -55,8 +90,17 @@ where
continue;
};
let budget = shards.budget.clone();
futures.push(Box::pin(async move {
let mut buf = vec![0; shard_size];
let memory = match budget.as_ref().map(|budget| budget.reserve(shard_size)).transpose() {
Ok(memory) => memory,
Err(error) => return (index, Err(Error::from(error)), false),
};
let mut buf = Vec::new();
if buf.try_reserve_exact(shard_size).is_err() {
return (index, Err(Error::from(io::Error::from(ErrorKind::OutOfMemory))), false);
}
buf.resize(shard_size, 0);
let read_result = if read_timeout.is_zero() {
reader.read(&mut buf).await
} else {
@@ -75,7 +119,7 @@ where
match read_result {
Ok(n) => {
buf.truncate(n);
(index, Ok(buf), false)
(index, Ok((buf, memory)), false)
}
Err(err) => {
let should_retire = err.kind() == ErrorKind::TimedOut;
@@ -87,10 +131,14 @@ where
while let Some((index, result, should_retire)) = futures.next().await {
match result {
Ok(shard) => {
shards[index] = Some(shard);
Ok((shard, memory)) => {
shards.shards[index] = Some(shard);
shards.memory[index] = memory;
}
Err(err) => {
if matches!(&err, Error::Io(error) if is_memory_pressure(error)) {
return Err(err);
}
errs[index] = Some(err);
if should_retire {
retire_readers.push(index);
@@ -105,7 +153,7 @@ where
warn!(shard_index = index, "retiring timed-out heal shard reader");
}
(shards, errs)
Ok((shards, errs))
}
fn heal_writer_failure(writers: &mut MultiWriter<'_>, error: io::Error) -> Error {
@@ -164,7 +212,7 @@ impl super::Erasure {
for block_index in start_block..end_block {
let shard_offset = block_index * self.shard_size();
let shard_size = self.shard_size().min(shard_file_size.saturating_sub(shard_offset));
let (mut shards, errs) = read_heal_shards(&mut readers, shard_size, read_timeout).await;
let (mut owned_shards, errs) = read_heal_shards(&mut readers, shard_size, read_timeout).await?;
// Every source shard is already bitrot-verified by its reader, so any
// data_shards survivors are sufficient to reconstruct — requiring more
@@ -183,6 +231,24 @@ impl super::Erasure {
return Err(Error::ErasureReadQuorum);
}
owned_shards.reserve_missing(shard_size)?;
let source_parity_bytes =
owned_shards
.shards
.iter()
.skip(self.data_shards)
.flatten()
.try_fold(0usize, |total, shard| {
total
.checked_add(shard.len())
.ok_or_else(|| io::Error::from(ErrorKind::OutOfMemory))
})?;
let _source_memory = owned_shards
.budget
.as_ref()
.map(|budget| budget.reserve(source_parity_bytes))
.transpose()?;
let shards = &mut owned_shards.shards;
let source_parity = shards
.iter()
.enumerate()
@@ -190,24 +256,23 @@ impl super::Erasure {
.filter_map(|(index, shard)| shard.as_ref().map(|shard| (index, shard.clone())))
.collect::<Vec<_>>();
if self.parity_shards > 0 {
self.decode_data_and_parity(&mut shards)?;
if !source_parity.is_empty() && !self.verify_data_and_parity(&shards)? {
self.decode_data_and_parity(shards)?;
if !source_parity.is_empty() && !self.verify_data_and_parity(shards)? {
return Err(Error::other("can not reconstruct data: inconsistent heal source shards"));
}
for (index, source) in source_parity {
let Some(rebuilt) = shards[index].as_ref() else {
for (index, source) in &source_parity {
let Some(rebuilt) = shards[*index].as_ref() else {
return Err(Error::other("can not reconstruct data: missing rebuilt parity shard"));
};
if rebuilt != &source {
if rebuilt != source {
return Err(Error::other("can not reconstruct data: inconsistent heal source shards"));
}
}
}
let shards = shards
.into_iter()
.map(|s| Bytes::from(s.unwrap_or_default()))
.collect::<Vec<_>>();
drop(source_parity);
drop(_source_memory);
let shards = owned_shards.into_bytes();
if let Err(error) = writers.write(shards).await {
return Err(heal_writer_failure(&mut writers, error));
@@ -306,6 +371,90 @@ mod tests {
BitrotWriterWrapper::new(CustomWriter::new_inline_buffer(), shard_size, HashAlgorithm::None)
}
#[tokio::test]
async fn heal_shards_retain_memory_through_missing_outputs_and_last_slice() {
use crate::disk::{test_resource_units, with_optional_resource_unit};
let memory = BackgroundMemoryBudget::new();
memory.set_sampled_limit(24, Duration::from_secs(30));
let mut readers = vec![
Some(BitrotReader::new(Cursor::new(vec![7; 8]), 8, HashAlgorithm::None, false)),
None,
Some(BitrotReader::new(Cursor::new(vec![9; 8]), 8, HashAlgorithm::None, false)),
];
let (resources, _) = test_resource_units(1);
let (mut shards, _) = with_optional_resource_unit(
Some(resources.with_memory_budget(memory.clone())),
read_heal_shards(&mut readers, 8, Duration::ZERO),
)
.await
.expect("read two accounted shards");
assert_eq!(memory.snapshot().used, 16);
shards
.reserve_missing(8)
.expect("reserve the reconstructed output before allocation");
shards.shards[1] = Some(vec![11; 8]);
assert_eq!(memory.snapshot().used, 24);
let bytes = shards.into_bytes();
let slice = bytes[1].slice(1..3);
memory.set_sampled_limit(0, Duration::ZERO);
drop(bytes);
assert_eq!(slice.as_ref(), &[11; 2]);
assert_eq!(memory.snapshot().used, 8);
drop(slice);
assert_eq!(memory.snapshot().used, 0);
}
#[tokio::test]
async fn heal_memory_denial_is_not_replaced_with_read_quorum() {
use crate::disk::{test_resource_units, with_optional_resource_unit};
let erasure = Erasure::new(2, 1, 16);
let encoded = erasure.encode_data(&[37; 16]).expect("valid source shards");
for limit in [0, 23, 31, 39] {
let memory = BackgroundMemoryBudget::new();
memory.set_sampled_limit(limit, Duration::from_secs(30));
let readers = encoded
.iter()
.map(|shard| Some(BitrotReader::new(Cursor::new(shard.clone()), 8, HashAlgorithm::None, false)))
.collect();
let mut writers = vec![Some(inline_writer(8)), None, None];
let (resources, _) = test_resource_units(1);
let result = with_optional_resource_unit(
Some(resources.with_memory_budget(memory.clone())),
erasure.heal(&mut writers, readers, 16, &[]),
)
.await;
assert!(
matches!(result, Err(Error::Io(error)) if error.kind() == ErrorKind::WouldBlock),
"limit={limit}"
);
assert_eq!(memory.snapshot().used, 0);
}
}
#[tokio::test]
async fn cancelled_heal_shard_reads_release_allocated_payloads() {
use crate::disk::{test_resource_units, with_optional_resource_unit};
let memory = BackgroundMemoryBudget::new();
memory.set_sampled_limit(16, Duration::from_secs(30));
let mut readers = vec![
Some(BitrotReader::new(PendingReader, 8, HashAlgorithm::None, false)),
Some(BitrotReader::new(PendingReader, 8, HashAlgorithm::None, false)),
];
let (resources, _) = test_resource_units(1);
let mut read = Box::pin(with_optional_resource_unit(
Some(resources.with_memory_budget(memory.clone())),
read_heal_shards(&mut readers, 8, Duration::ZERO),
));
assert!(futures::poll!(&mut read).is_pending());
assert_eq!(memory.snapshot().used, 16);
memory.set_sampled_limit(0, Duration::ZERO);
drop(read);
assert_eq!(memory.snapshot().used, 0);
}
#[tokio::test]
async fn read_heal_shards_returns_empty_slots_for_zero_sized_shards() {
let mut readers = vec![
@@ -313,9 +462,11 @@ mod tests {
None,
];
let (shards, errs) = read_heal_shards(&mut readers, 0, Duration::ZERO).await;
let (shards, errs) = read_heal_shards(&mut readers, 0, Duration::ZERO)
.await
.expect("empty read succeeds");
assert_eq!(shards, vec![None, None]);
assert_eq!(shards.shards, vec![None, None]);
assert_eq!(errs, vec![None, None]);
}
@@ -323,9 +474,11 @@ mod tests {
async fn read_heal_shards_records_reader_errors_and_retiring_timeouts() {
let mut readers = vec![Some(BitrotReader::new(FailingReader, 4, HashAlgorithm::None, false))];
let (shards, errs) = read_heal_shards(&mut readers, 4, Duration::ZERO).await;
let (shards, errs) = read_heal_shards(&mut readers, 4, Duration::ZERO)
.await
.expect("ordinary failures remain per-shard");
assert!(shards[0].is_none());
assert!(shards.shards[0].is_none());
assert!(
errs[0]
.as_ref()
@@ -334,9 +487,11 @@ mod tests {
assert!(readers[0].is_some(), "ordinary read errors must not retire readers");
let mut timeout_readers = vec![Some(BitrotReader::new(PendingReader, 4, HashAlgorithm::None, false))];
let (shards, errs) = read_heal_shards(&mut timeout_readers, 4, Duration::from_millis(1)).await;
let (shards, errs) = read_heal_shards(&mut timeout_readers, 4, Duration::from_millis(1))
.await
.expect("timeouts remain per-shard");
assert_eq!(shards, vec![None]);
assert_eq!(shards.shards, vec![None]);
assert!(errs[0].as_ref().is_some_and(|err| err.to_string().contains("timed out")));
assert!(timeout_readers[0].is_none(), "timed-out readers are retired");
}
+122 -4
View File
@@ -1536,7 +1536,7 @@ mod tests {
writer.write(chunk).await.expect("write chunk");
}
let inline_data = Bytes::from(writer.into_inline_data().expect("inline buffer"));
let inline_data = writer.into_inline_data().expect("inline buffer");
let mut reader = create_bitrot_reader_from_bytes(
Some(inline_data),
None,
@@ -1584,7 +1584,7 @@ mod tests {
let inline_data = writer.into_inline_data().expect("inline buffer");
let mut reader = create_deferred_bitrot_reader(
Some(inline_data.into()),
Some(inline_data),
None,
"test-bucket",
"test-path",
@@ -1626,7 +1626,7 @@ mod tests {
for stripe in 0..payload.len() / shard_size {
let (mut reader, handle) = create_deferred_bitrot_reader_with_stripe_handle(
Some(inline_data.clone().into()),
Some(inline_data.clone()),
None,
"test-bucket",
"test-path",
@@ -1671,7 +1671,7 @@ mod tests {
for chunk in payload.chunks(shard_size) {
writer.write(chunk).await.expect("write chunk");
}
let mut inline_data = writer.into_inline_data().expect("inline buffer");
let mut inline_data = writer.into_inline_data().expect("inline buffer").to_vec();
// Corrupt one payload byte of block 2 (blocks are hash + data).
let block_len = algo.size() + shard_size;
@@ -1845,4 +1845,122 @@ mod tests {
assert_eq!(output, data);
}
mod inline_memory_regressions {
use super::*;
use crate::disk::background_memory::Reservation;
use crate::disk::{BackgroundMemoryBudget, test_resource_units, with_optional_resource_unit};
use std::time::Duration;
#[tokio::test]
async fn owned_inline_reader_retains_existing_receipt_without_a_new_grant() {
let memory = BackgroundMemoryBudget::new();
memory.set_sampled_limit(64, Duration::from_secs(30));
let receipt = memory.reserve(64).expect("existing metadata payload grant");
let bytes = Reservation::wrap(Bytes::from(vec![7; 64]), Some(receipt));
memory.set_sampled_limit(0, Duration::ZERO);
let (resources, _) = test_resource_units(1);
let mut reader = with_optional_resource_unit(
Some(resources.with_memory_budget(memory.clone())),
create_bitrot_reader_from_bytes(
Some(bytes.clone()),
None,
"bucket",
"inline",
0,
64,
64,
HashAlgorithm::None,
false,
false,
),
)
.await
.expect("owned data requires no allocation")
.expect("inline reader");
drop(bytes);
assert_eq!(memory.snapshot().used, 64, "the source owner must remain attached to the reader");
let mut output = [0; 64];
assert_eq!(reader.read(&mut output).await.expect("read existing payload"), 64);
assert_eq!(output, [7; 64]);
assert_eq!(memory.snapshot().used, 64);
drop(reader);
assert_eq!(memory.snapshot().used, 0);
}
#[tokio::test]
async fn accounted_inline_writer_round_trips_bitrot_and_preserves_export_lifetime() {
let memory = BackgroundMemoryBudget::new();
memory.set_sampled_limit(192, Duration::from_secs(30));
let (resources, _) = test_resource_units(1);
let bytes = with_optional_resource_unit(Some(resources.with_memory_budget(memory.clone())), async {
let mut writer = create_bitrot_writer(true, None, "bucket", "inline", 64, 64, HashAlgorithm::HighwayHash256S)
.await
.expect("inline writer");
assert_eq!(writer.write(&[19; 64]).await.expect("complete shard"), 64);
writer.shutdown().await.expect("shutdown");
writer.into_inline_data().expect("bitrot bytes")
})
.await;
assert_eq!(bytes.len(), 96);
assert_eq!(memory.snapshot().used, 96, "exported [hash][shard] remains charged");
let (resources, _) = test_resource_units(1);
let mut reader = with_optional_resource_unit(
Some(resources.with_memory_budget(memory.clone())),
create_bitrot_reader_from_bytes(
Some(bytes.clone()),
None,
"bucket",
"inline",
0,
64,
64,
HashAlgorithm::HighwayHash256S,
false,
false,
),
)
.await
.expect("reader")
.expect("inline reader");
let mut output = [0; 64];
assert_eq!(reader.read(&mut output).await.expect("bitrot remains verified"), 64);
assert_eq!(output, [19; 64]);
assert_eq!(memory.snapshot().used, 192, "input owner and verification scratch are both live");
drop(reader);
assert_eq!(memory.snapshot().used, 96);
let slice = bytes.slice(32..36);
drop(bytes);
assert_eq!(memory.snapshot().used, 96);
drop(slice);
assert_eq!(memory.snapshot().used, 0);
}
#[tokio::test]
async fn inline_heal_writer_memory_denial_is_retryable_after_read_and_verify_fit() {
use crate::erasure::coding::Erasure;
let erasure = Erasure::new(2, 1, 16);
let encoded = erasure.encode_data(&[37; 16]).expect("foreground test fixture");
let readers = encoded
.into_iter()
.map(|shard| Some(BitrotReader::new(Cursor::new(shard), 8, HashAlgorithm::None, false)))
.collect();
let memory = BackgroundMemoryBudget::new();
// Read shards: 24 bytes; source parity: 8; verify scratch: 8.
// After verification, input shards still use 24. The target needs a
// new 40-byte [hash][shard] allocation and must fail before copying.
memory.set_sampled_limit(40, Duration::from_secs(30));
let (resources, _) = test_resource_units(1);
let result = with_optional_resource_unit(Some(resources.with_memory_budget(memory.clone())), async {
let writer = create_bitrot_writer(true, None, "bucket", "inline", 8, 8, HashAlgorithm::HighwayHash256S)
.await
.expect("writer setup does not allocate payload");
let mut writers = vec![Some(writer), None, None];
erasure.heal(&mut writers, readers, 16, &[]).await
})
.await;
assert!(matches!(result, Err(DiskError::Io(error)) if error.kind() == io::ErrorKind::WouldBlock));
assert_eq!(memory.snapshot().used, 0);
}
}
}
+3 -1
View File
@@ -55,7 +55,9 @@ mod set_disk;
mod storage_api_contracts;
mod store;
pub use disk::{with_background_disk_io, with_background_resource_unit};
pub use disk::{
spawn_blocking_with_resource_ownership, with_background_disk_io, with_background_resource_set, with_background_resource_unit,
};
pub use store::PoolMetaWriteGateStatus;
// pub mod checksum;
@@ -8716,11 +8716,11 @@ mod tests {
writer.shutdown().await.expect("inline writer should shutdown");
let mut fi = base.clone();
fi.erasure.index = index + 1;
fi.data = Some(Bytes::from(
fi.data = Some(
writer
.into_inline_data()
.expect("inline bitrot writer should retain encoded data"),
));
);
files.push(fi);
}
files
+6 -4
View File
@@ -4466,6 +4466,7 @@ fn resolve_tiered_decommission_write_quorum_result(
pub struct SetDisks {
pub locker_owner: String,
pub disks: Arc<RwLock<Vec<Option<DiskStore>>>>,
pub(crate) physical_heal_admission: Arc<crate::disk::PhysicalHealBinding>,
pub set_endpoints: Vec<Endpoint>,
pub set_drive_count: usize,
pub default_parity_count: usize,
@@ -5427,6 +5428,7 @@ impl SetDisks {
Arc::new(SetDisks {
locker_owner,
disks,
physical_heal_admission: Arc::default(),
set_drive_count,
default_parity_count,
set_index,
@@ -14517,7 +14519,7 @@ mod tests {
file.erasure.index = files.len() + 1;
file.erasure.block_size = erasure.block_size;
file.uses_legacy_checksum = uses_legacy;
file.data = Some(Bytes::from(data));
file.data = Some(data);
files.push(file);
}
@@ -14978,7 +14980,7 @@ mod tests {
}
let encoded = bitrot_writer.into_inline_data().expect("bitrot encoded data should exist");
disk.write_all(bucket, &shard_path, Bytes::from(encoded))
disk.write_all(bucket, &shard_path, encoded)
.await
.expect("encoded shard should be stored");
@@ -15062,7 +15064,7 @@ mod tests {
}
let encoded = bitrot_writer.into_inline_data().expect("bitrot encoded data should exist");
disk.write_all(bucket, &shard_path, Bytes::from(encoded))
disk.write_all(bucket, &shard_path, encoded)
.await
.expect("encoded shard should be stored");
@@ -15175,7 +15177,7 @@ mod tests {
}
let encoded = bitrot_writer.into_inline_data().expect("bitrot encoded data should exist");
disk.write_all(bucket, &shard_path, Bytes::from(encoded))
disk.write_all(bucket, &shard_path, encoded)
.await
.expect("encoded shard should be stored");
}
+11 -12
View File
@@ -13,14 +13,14 @@
// limitations under the License.
use super::super::{
Bytes, CHECK_PART_FILE_CORRUPT, CHECK_PART_SUCCESS, DeleteOptions, DiskError, DiskStore, DriveState, EVENT_SET_DISK_HEAL,
Error, FileInfo, HashAlgorithm, HashMap, HealDriveInfo, HealItemType, HealOpts, HealResultItem, HealScanMode, Infos,
CHECK_PART_FILE_CORRUPT, CHECK_PART_SUCCESS, DeleteOptions, DiskError, DiskStore, DriveState, EVENT_SET_DISK_HEAL, Error,
FileInfo, HashAlgorithm, HashMap, HealDriveInfo, HealItemType, HealOpts, HealResultItem, HealScanMode, Infos,
LOG_SUBSYSTEM_SET_DISK, ObjectInfo, ObjectOptions, ObjectPartInfo, Path, RUSTFS_META_TMP_BUCKET, ReadOptions, Result,
SLASH_SEPARATOR, SetDisks, StorageError, Uuid, coding, count_errs, count_part_not_success, create_bitrot_reader,
create_bitrot_writer, debug, disk, disks_with_all_parts, encode_dir_object, error, file_info_is_valid_for_metadata,
formats_match_reference_slots, get_format_erasure_in_quorum, get_lock_acquire_timeout, has_suffix,
heal_bucket_local_on_disks, is_object_dir_dangling, join_all, load_format_erasure_all, path_join_buf, save_format_file,
should_heal_object_on_disk, stat_all_dirs, to_object_err, warn,
SLASH_SEPARATOR, SetDisks, StorageError, Uuid, coding, count_errs, count_part_not_success, create_bitrot_writer, debug, disk,
disks_with_all_parts, encode_dir_object, error, file_info_is_valid_for_metadata, formats_match_reference_slots,
get_format_erasure_in_quorum, get_lock_acquire_timeout, has_suffix, heal_bucket_local_on_disks, is_object_dir_dangling,
join_all, load_format_erasure_all, path_join_buf, save_format_file, should_heal_object_on_disk, stat_all_dirs, to_object_err,
warn,
};
use crate::bucket::retirement::MarkerRetirementContext;
use crate::diagnostics::object_lock::ObjectLockAttempt;
@@ -44,7 +44,7 @@ struct ExplicitVersionHeal<'a> {
use crate::disk::DataDirDeleteStatus;
use crate::disk::DiskAPI;
use crate::disk::local::{DELETE_DATA_DIR_MARKER_PREFIX, metadata_less_part_file};
use crate::io_support::bitrot::object_mmap_read_enabled;
use crate::io_support::bitrot::{create_bitrot_reader_from_bytes, object_mmap_read_enabled};
use crate::storage_api_contracts::namespace::NamespaceLocking as _;
use rustfs_common::mrf_channel::MrfDeleteMarkerPurge;
use rustfs_common::trace_bus::{TraceEvent, TraceFunc, TraceKind, trace_emit};
@@ -1793,8 +1793,8 @@ impl SetDisks {
checksum_info.algorithm
};
match create_bitrot_reader(
metadata.data.as_deref(),
match create_bitrot_reader_from_bytes(
metadata.data.clone(),
Some(disk),
bucket,
&path_join_buf(&[object, &src_data_dir, &format!("part.{}", part.number)]),
@@ -1891,8 +1891,7 @@ impl SetDisks {
// if let Some(w) = writer.as_any().downcast_ref::<BitrotFileWriter>() {
// parts_metadata[index].data = Some(w.inline_data().to_vec());
// }
parts_metadata[index].data =
Some(writer.into_inline_data().map(Bytes::from).unwrap_or_default());
parts_metadata[index].data = Some(writer.into_inline_data().unwrap_or_default());
}
parts_metadata[index].set_inline_data();
} else {
+1 -1
View File
@@ -4713,7 +4713,7 @@ impl SetDisks {
.ok_or_else(|| Error::other(format!("inline encoder omitted disk shard {i}")))?,
);
} else if let Some(writer) = writers[i].take() {
pfi.data = Some(writer.into_inline_data().map(Bytes::from).unwrap_or_default());
pfi.data = Some(writer.into_inline_data().unwrap_or_default());
}
pfi.set_inline_data();
+125 -12
View File
@@ -599,6 +599,33 @@ impl ECStore {
Ok((write_state, pool_guard, rebalance_guard, pool_meta, rebalance_meta))
}
/// Snapshot only the object sets participating in this request. The sampler
/// resolves their physical leaves before Heal enters namespace/cache locks.
pub fn physical_heal_plan(&self, object: &str, opts: &HealOpts) -> Result<Option<crate::disk::PhysicalResourcePlan>> {
if let Some(index) = opts.pool
&& index >= self.pools.len()
{
return Err(invalid_heal_pool_index(index, self.pools.len()));
}
let object = rustfs_utils::path::encode_dir_object_ref(object);
let mut selected: Option<crate::disk::PhysicalResourcePlan> = None;
for (index, pool) in self.pools.iter().enumerate() {
if opts.pool.is_some_and(|requested| requested != index) {
continue;
}
let set = pool.get_disks_for_heal_object(&object, opts)?;
let Some(plan) = set.physical_heal_admission.plan() else { return Ok(None) };
if let Some(selected) = &mut selected {
if selected.union_with(&plan).is_err() {
return Ok(None);
}
} else {
selected = Some(plan);
}
}
Ok(selected)
}
fn get_pools_for_heal_object(&self, opts: &HealOpts) -> Result<Vec<Arc<Sets>>> {
match opts.pool {
Some(pool_idx) => Ok(vec![
@@ -1554,21 +1581,25 @@ mod tests {
}
}
async fn minimal_heal_pool(pool_idx: usize) -> Arc<Sets> {
let format = FormatV3::new(1, 1);
let endpoint_url = format!("http://127.0.0.1:{}/data", 19000 + pool_idx);
let mut endpoint = Endpoint::try_from(endpoint_url.as_str()).expect("endpoint should parse");
endpoint.set_pool_index(pool_idx);
endpoint.set_set_index(0);
endpoint.set_disk_index(0);
async fn minimal_heal_pool(pool_idx: usize, set_count: usize) -> Arc<Sets> {
let format = FormatV3::new(set_count, 1);
let endpoints = (0..set_count)
.map(|set_idx| {
let endpoint_url = format!("http://127.0.0.1:{}/data", 19000 + pool_idx + set_idx * 2);
let mut endpoint = Endpoint::try_from(endpoint_url.as_str()).expect("endpoint should parse");
endpoint.set_pool_index(pool_idx);
endpoint.set_set_index(set_idx);
endpoint.set_disk_index(0);
endpoint
})
.collect::<Vec<_>>();
Sets::new(
vec![None],
vec![None; set_count],
&PoolEndpoints {
legacy: false,
set_count: 1,
set_count,
drives_per_set: 1,
endpoints: Endpoints::from(vec![endpoint]),
endpoints: Endpoints::from(endpoints),
cmd_line: String::new(),
platform: String::new(),
},
@@ -2033,7 +2064,7 @@ mod tests {
ECStore {
id: Uuid::new_v4(),
disk_map: HashMap::new(),
pools: vec![minimal_heal_pool(0).await, minimal_heal_pool(1).await],
pools: vec![minimal_heal_pool(0, 1).await, minimal_heal_pool(1, 1).await],
peer_sys: S3PeerSys {
clients: Vec::new(),
pools_count: 2,
@@ -2048,6 +2079,88 @@ mod tests {
}
}
#[tokio::test]
async fn physical_heal_plan_uses_the_real_pool_set_and_directory_placement() {
use crate::disk::PhysicalDiskBinding;
use rustfs_concurrency::adaptive::{
AdaptiveAdmission, FeedbackController, FeedbackObservation, NodeAdmission, ResourcePressure,
};
use std::num::NonZeroU16;
let mut store = minimal_heal_store().await;
store.pools[0] = minimal_heal_pool(0, 2).await;
let pool = &store.pools[0];
let a = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let b = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
for (set, gate) in pool.disk_set.iter().zip([&a, &b]) {
let binding = PhysicalDiskBinding::default();
binding.set_test_capacity(vec![gate.clone()]).expect("physical fixture");
set.physical_heal_admission
.set_test_plan(binding.plan().expect("live").resource_plan());
}
let node = NodeAdmission::new(NonZeroU16::MIN);
node.update(
&mut FeedbackController::new(),
Some(FeedbackObservation {
generation: 1,
max_concurrency: 1,
pressure: ResourcePressure::Unknown,
completed_units: None,
}),
);
let directory = (0..100)
.map(|n| format!("dir{n}/"))
.find(|name| {
store.pools[0].get_disks_by_key(name).set_index
!= store.pools[0].get_disks_by_key(&encode_dir_object(name)).set_index
})
.expect("different directory placement");
for object in ["plain-object", directory.as_str()] {
for set in [None, Some(0), Some(1)] {
let opts = HealOpts {
pool: Some(0),
set,
..Default::default()
};
let expected = store.pools[0]
.get_disks_for_heal_object(&encode_dir_object(object), &opts)
.expect("placement")
.set_index;
let permit = store
.physical_heal_plan(object, &opts)
.expect("valid selection")
.expect("known set")
.try_acquire(&node)
.expect("physical admission");
assert_eq!((a.snapshot().active, b.snapshot().active), if expected == 0 { (1, 0) } else { (0, 1) });
drop(permit);
}
}
assert!(
store
.physical_heal_plan("object", &HealOpts::default())
.expect("valid all pools")
.is_none(),
"an undiscovered participating pool cannot be skipped"
);
for opts in [
HealOpts {
pool: Some(2),
..Default::default()
},
HealOpts {
pool: Some(0),
set: Some(2),
..Default::default()
},
] {
assert!(matches!(
store.physical_heal_plan("object", &opts),
Err(StorageError::InvalidArgument(..))
));
}
}
async fn remove_pool_meta_shard(store: &ECStore, pool_idx: usize) -> DiskStore {
let target_set = store.pools[pool_idx].get_disks_by_key(POOL_META_NAME);
let missing_disk = target_set.disks.read().await[0]
+2 -2
View File
@@ -456,8 +456,8 @@ pub(crate) mod utils;
use peer::init_local_peer;
pub use peer::{
BootstrapLocalTarget, all_local_disk, all_local_disk_path, find_local_disk_by_ref, get_disk_infos, init_local_disks,
init_local_disks_with_instance_ctx, init_lock_clients, prewarm_local_disk_id_map,
BootstrapLocalTarget, all_local_disk, all_local_disk_path, find_local_disk_by_ref, find_local_disk_by_ref_for_rpc,
get_disk_infos, init_local_disks, init_local_disks_with_instance_ctx, init_lock_clients, prewarm_local_disk_id_map,
prewarm_local_disk_id_map_with_instance_ctx,
};
+128 -3
View File
@@ -99,6 +99,9 @@ impl ECStore {
fi: FileInfo,
expected: Uuid,
) -> DiskResult<()> {
if crate::disk::remote_io_admission::active() {
local_disk_candidate(&self.ctx, disk_ref).await?;
}
let disk_ref = disk_ref.to_owned();
let path = target.1.to_owned();
self.run_bucket_heal_at_incarnation(target.0, expected, &Default::default(), move |store, volume, _| async move {
@@ -117,6 +120,9 @@ impl ECStore {
opts: DeleteOptions,
expected: Uuid,
) -> DiskResult<()> {
if crate::disk::remote_io_admission::active() {
local_disk_candidate(&self.ctx, disk_ref).await?;
}
let disk_ref = disk_ref.to_owned();
let path = target.1.to_owned();
self.run_bucket_heal_at_incarnation(target.0, expected, &Default::default(), move |store, volume, _| async move {
@@ -144,6 +150,9 @@ impl ECStore {
if expected.is_nil() || opts.undo_write || self.ctx.lock_manager().is_disabled() {
return Err(DiskError::other("invalid incarnation-bound heal deletion"));
}
if crate::disk::remote_io_admission::active() {
local_disk_candidate(&self.ctx, disk_ref).await?;
}
let fence = Arc::new(
self.acquire_bucket_incarnation_fence(target.0, expected)
.await
@@ -281,7 +290,9 @@ impl ECStore {
// The optional ID is a cold lookup to cache only after final admission.
async fn local_disk_candidate(ctx: &Arc<InstanceContext>, disk_ref: &str) -> DiskResult<(DiskStore, Option<Uuid>)> {
let map = ctx.local_disk_map();
if let Some(disk) = map.read().await.get(disk_ref).and_then(Option::as_ref).cloned() {
let path_disk = map.read().await.get(disk_ref).and_then(Option::as_ref).cloned();
if let Some(disk) = path_disk {
admit_rpc_disk(&disk)?;
return Ok((disk, None));
}
let disk_id = Uuid::parse_str(disk_ref).map_err(|_| DiskError::DiskNotFound)?;
@@ -290,11 +301,16 @@ async fn local_disk_candidate(ctx: &Arc<InstanceContext>, disk_ref: &str) -> Dis
let cached_disk = map.read().await.get(&path).and_then(Option::as_ref).cloned();
if let Some(disk) = cached_disk
&& matches!(disk.as_ref(), Disk::Local(_))
&& disk.get_disk_id().await? == Some(disk_id)
{
return Ok((disk, None));
admit_rpc_disk(&disk)?;
if disk.get_disk_id().await? == Some(disk_id) {
return Ok((disk, None));
}
}
}
if crate::disk::remote_io_admission::active() {
return Err(DiskError::Io(std::io::Error::from(std::io::ErrorKind::WouldBlock)));
}
let disks: Vec<_> = map.read().await.values().filter_map(Clone::clone).collect();
// Disk identity may perform format I/O. No registry guard spans this await.
for disk in disks {
@@ -305,6 +321,17 @@ async fn local_disk_candidate(ctx: &Arc<InstanceContext>, disk_ref: &str) -> Dis
Err(DiskError::DiskNotFound)
}
fn admit_rpc_disk(disk: &DiskStore) -> DiskResult<()> {
use crate::disk::remote_io_admission;
if remote_io_admission::active() {
let binding = disk
.physical_admission()
.ok_or_else(|| DiskError::Io(std::io::Error::from(std::io::ErrorKind::WouldBlock)))?;
remote_io_admission::admit(&binding).map_err(DiskError::Io)?;
}
Ok(())
}
async fn admit_local_disk(
ctx: &Arc<InstanceContext>,
disk: &DiskStore,
@@ -456,6 +483,18 @@ pub async fn find_local_disk_by_ref(disk_ref: &str) -> Option<DiskStore> {
None
}
/// Resolve and admit an authenticated RPC's target before its storage locks.
/// Unknown UUID mappings defer background work instead of probing unrelated disks.
pub async fn find_local_disk_by_ref_for_rpc(disk_ref: &str) -> DiskResult<Option<DiskStore>> {
use crate::disk::remote_io_admission;
if !remote_io_admission::active() {
return Ok(find_local_disk_by_ref(disk_ref).await);
}
let (disk, _) = local_disk_candidate(&crate::runtime::global::current_ctx(), disk_ref).await?;
Ok(Some(disk))
}
pub async fn all_local_disk_path() -> Vec<String> {
runtime_sources::local_disk_paths().await
}
@@ -816,6 +855,92 @@ mod tests {
}
}
#[tokio::test]
async fn remote_background_target_lookup_defers_before_io_and_reuses_multi_phase_admission() {
use crate::disk::{RemoteBackgroundIo, with_remote_background_io};
use rustfs_concurrency::adaptive::{
AdaptiveAdmission, FeedbackController, FeedbackObservation, NodeAdmission, ResourcePressure,
};
use std::num::NonZeroU16;
let ctx = Arc::new(InstanceContext::new());
let root = tempfile::tempdir().expect("target directory");
let id = Uuid::new_v4();
let disk = target_disk(&ctx, root.path(), id).await;
disk.make_volume("bucket").await.expect("fixture volume");
let gate = AdaptiveAdmission::new(NonZeroU16::MIN, 0);
disk.physical_admission()
.expect("local binding")
.set_test_capacity(vec![gate.clone()])
.expect("physical device");
let node = NodeAdmission::new(NonZeroU16::MIN);
node.update(
&mut FeedbackController::new(),
Some(FeedbackObservation {
generation: 1,
max_concurrency: 1,
pressure: ResourcePressure::Unknown,
completed_units: None,
}),
);
let path = disk.endpoint().to_string();
with_remote_background_io(Some(RemoteBackgroundIo::new(node.clone())), async {
assert!(
matches!(disk.read_all("bucket", "object").await,
Err(DiskError::Io(error)) if error.kind() == std::io::ErrorKind::WouldBlock),
"no implicit target admission"
);
assert!(
matches!(local_disk_candidate(&ctx, &path).await,
Err(DiskError::Io(error)) if error.kind() == std::io::ErrorKind::WouldBlock),
"closed receiver capacity"
);
gate.set_limit(1);
assert!(
matches!(local_disk_candidate(&ctx, &id.to_string()).await,
Err(DiskError::Io(error)) if error.kind() == std::io::ErrorKind::WouldBlock),
"cold UUID cannot probe unrelated disks"
);
assert_eq!([node.snapshot().active, gate.snapshot().active], [0, 0]);
let (selected, _) = local_disk_candidate(&ctx, &path).await.expect("select known local path");
assert!(Arc::ptr_eq(&selected, &disk));
for body in [b"first phase".as_slice(), b"second phase".as_slice()] {
disk.write_all("bucket", "object", bytes::Bytes::copy_from_slice(body))
.await
.expect("write phase");
let (again, _) = local_disk_candidate(&ctx, &path)
.await
.expect("repeat target resolution at capacity one");
assert_eq!(again.read_all("bucket", "object").await.expect("read phase").as_ref(), body);
assert_eq!([node.snapshot().active, gate.snapshot().active], [1, 1]);
}
})
.await;
assert_eq!([node.snapshot().active, gate.snapshot().active], [0, 0]);
assert_eq!([node.completed(), gate.completed()], [Some(1), Some(1)]);
assert!(record_local_disk_id_if_active(&ctx, &disk, id).await);
let replacement = target_disk(&ctx, root.path(), Uuid::new_v4()).await;
replacement
.physical_admission()
.expect("replacement binding")
.set_test_capacity(vec![gate.clone()])
.expect("replacement device");
with_remote_background_io(Some(RemoteBackgroundIo::new(node.clone())), async {
assert!(
matches!(local_disk_candidate(&ctx, &id.to_string()).await,
Err(DiskError::Io(error)) if error.kind() == std::io::ErrorKind::WouldBlock),
"stale UUID must not resolve replacement data"
);
})
.await;
assert_eq!([node.snapshot().active, gate.snapshot().active], [0, 0]);
assert_eq!(
[node.completed(), gate.completed()],
[Some(1), Some(1)],
"identity checks alone are not work completion"
);
}
#[tokio::test]
async fn target_uuid_cache_cannot_admit_a_different_format_at_the_same_path() {
let ctx = Arc::new(InstanceContext::new());
+2 -1
View File
@@ -84,7 +84,7 @@ rustfs-heal-contracts = { workspace = true }
rustfs-madmin = { workspace = true }
rustfs-utils = { workspace = true }
tokio = { workspace = true, features = ["sync", "io-util", "time", "macros", "fs", "rt-multi-thread"] }
tokio-util = { workspace = true, features = ["io", "compat"] }
tokio-util = { workspace = true, features = ["io", "compat", "rt"] }
tonic = { workspace = true }
tracing = { workspace = true }
serde = { workspace = true, features = ["derive"] }
@@ -100,6 +100,7 @@ crc-fast = { workspace = true }
sha2 = { workspace = true }
[dev-dependencies]
rustfs-ecstore = { workspace = true, features = ["test-util"] }
rustfs-filemeta = { workspace = true }
serde_json = { workspace = true, features = ["raw_value"] }
rustfs-test-utils = { workspace = true }
+33 -4
View File
@@ -185,7 +185,7 @@ pub(crate) enum HealErrorClass {
/// Offline disk, missing volume, unmet quorum, transport failure,
/// slow-down or cancelled RPC: the same request can succeed later.
Transient,
/// Lock contention or a lock/RPC timeout (non-fatal `LockError`).
/// Nonblocking admission/lock contention or a lock/RPC timeout.
LockContention,
/// A storage grace window; `retry_not_before` is the absolute UNIX
/// deadline when storage reported one.
@@ -269,8 +269,7 @@ pub(crate) fn classify_storage_error(error: &EcstoreError) -> HealErrorClass {
}
}
/// Classify a disk error. Nonblocking local CAS and replacement leases report
/// lock contention as `WouldBlock`.
/// Classify a disk error without losing the underlying I/O deferral kind.
pub(crate) fn classify_disk_error(error: &DiskError) -> HealErrorClass {
if error.is_dangling_delete_grace() {
return HealErrorClass::DanglingGrace {
@@ -281,7 +280,6 @@ pub(crate) fn classify_disk_error(error: &DiskError) -> HealErrorClass {
DiskError::FileNotFound | DiskError::FileVersionNotFound => HealErrorClass::Absent,
DiskError::NoHealRequired => HealErrorClass::NoHealRequired,
DiskError::Timeout => HealErrorClass::LockContention,
DiskError::Io(error) if error.kind() == std::io::ErrorKind::WouldBlock => HealErrorClass::LockContention,
DiskError::DiskNotFound
| DiskError::VolumeNotFound
| DiskError::ErasureReadQuorum
@@ -296,6 +294,14 @@ pub(crate) fn classify_disk_error(error: &DiskError) -> HealErrorClass {
}
fn classify_io_error(error: &std::io::Error) -> HealErrorClass {
// Native memory admission and nonblocking CAS keep this kind when their
// errors cross disk, storage and direct I/O wrappers.
if error.kind() == std::io::ErrorKind::WouldBlock {
return HealErrorClass::LockContention;
}
if error.kind() == std::io::ErrorKind::OutOfMemory {
return HealErrorClass::Transient;
}
if DiskError::io_error_is_dangling_delete_grace(error) {
return HealErrorClass::DanglingGrace {
retry_not_before: dangling_deadline(DiskError::io_error_dangling_delete_retry_after(error)),
@@ -826,6 +832,29 @@ mod tests {
assert!(Error::Disk(DiskError::VolumeNotFound).is_recoverable_object_heal());
}
#[test]
fn resource_pressure_stays_retryable_across_error_wrappers() {
for (kind, expected) in [
(std::io::ErrorKind::WouldBlock, super::HealErrorClass::LockContention),
(std::io::ErrorKind::OutOfMemory, super::HealErrorClass::Transient),
] {
let pressure = || DiskError::Io(std::io::Error::new(kind, "background buffer memory budget exhausted"));
for error in [
Error::Disk(pressure()),
Error::Storage(EcstoreError::from(pressure())),
Error::Io(EcstoreError::from(pressure()).into()),
] {
assert_eq!(super::classify_heal_error(&error), expected);
assert!(error.is_recoverable_heal(), "resource pressure must preserve retry responsibility");
assert!(!error.is_object_absent());
}
}
for kind in [std::io::ErrorKind::Other, std::io::ErrorKind::PermissionDenied] {
let error = Error::Io(std::io::Error::new(kind, "background buffer memory budget exhausted"));
assert!(!error.is_recoverable_heal(), "display text must not override the typed error kind");
}
}
#[test]
fn classify_heal_error_table() {
use super::{HealErrorClass as C, classify_heal_error};
+183
View File
@@ -28,6 +28,7 @@ use rustfs_madmin::heal_commands::HealResultItem;
use serde::Serialize;
use std::sync::Arc;
use tokio::sync::{mpsc, oneshot};
use tokio_util::sync::CancellationToken;
use tracing::{debug, error, info};
const LOG_COMPONENT_HEAL: &str = "heal";
@@ -60,6 +61,8 @@ pub struct HealChannelProcessor {
response_sender: mpsc::UnboundedSender<HealChannelResponse>,
/// Response receiver
response_receiver: mpsc::UnboundedReceiver<HealChannelResponse>,
#[cfg(test)]
command_pause: Option<(oneshot::Sender<()>, oneshot::Receiver<()>)>,
}
#[derive(Serialize)]
@@ -161,6 +164,8 @@ impl HealChannelProcessor {
heal_manager,
response_sender: response_tx,
response_receiver: response_rx,
#[cfg(test)]
command_pause: None,
}
}
@@ -215,9 +220,19 @@ impl HealChannelProcessor {
/// Start processing legacy and canonical-receipt heal channel requests.
pub async fn start_with_receipts(
&mut self,
receiver: HealChannelReceiver,
receipt_receiver: HealReceiptReceiver,
) -> Result<()> {
self.start_with_shutdown(receiver, receipt_receiver, CancellationToken::new())
.await
}
pub(crate) async fn start_with_shutdown(
&mut self,
mut receiver: HealChannelReceiver,
mut receipt_receiver: HealReceiptReceiver,
stop: CancellationToken,
) -> Result<()> {
info!(
target: "rustfs::heal::channel",
@@ -230,10 +245,30 @@ impl HealChannelProcessor {
let mut receipt_channel_open = true;
loop {
if stop.is_cancelled() {
receiver.close();
receipt_receiver.close();
while let Some(command) = receiver.recv().await {
Self::reject_stopped_command(command);
}
while let Some(HealReceiptCommand { response_tx, .. }) = receipt_receiver.recv().await {
let _ = response_tx.send(Err(Error::other("Heal channel is stopping").to_string()));
}
self.response_receiver.close();
while self.response_receiver.try_recv().is_ok() {}
break;
}
tokio::select! {
_ = stop.cancelled() => {}
command = receiver.recv() => {
match command {
Some(command) => {
if stop.is_cancelled() {
Self::reject_stopped_command(command);
continue;
}
#[cfg(test)]
self.pause_selected_command().await;
if let Err(e) = self.process_command(command).await {
error!(
target: "rustfs::heal::channel",
@@ -264,6 +299,12 @@ impl HealChannelProcessor {
receipt_channel_open = false;
continue;
};
if stop.is_cancelled() {
let _ = response_tx.send(Err(Error::other("Heal channel is stopping").to_string()));
continue;
}
#[cfg(test)]
self.pause_selected_command().await;
if let Err(e) = self.process_start_request(request, false, true, response_tx).await {
error!(
target: "rustfs::heal::channel",
@@ -305,6 +346,26 @@ impl HealChannelProcessor {
Ok(())
}
fn reject_stopped_command(command: HealChannelCommand) {
let error = Error::other("Heal channel is stopping").to_string();
match command {
HealChannelCommand::Start { response_tx, .. } => {
let _ = response_tx.send(Err(error));
}
HealChannelCommand::Query { response_tx, .. } | HealChannelCommand::Cancel { response_tx, .. } => {
let _ = response_tx.send(Err(error));
}
}
}
#[cfg(test)]
async fn pause_selected_command(&mut self) {
if let Some((started, release)) = self.command_pause.take() {
let _ = started.send(());
let _ = release.await;
}
}
/// Process heal command
async fn process_command(&self, command: HealChannelCommand) -> Result<()> {
match command {
@@ -2459,4 +2520,126 @@ mod tests {
assert_eq!(response.request_id, "missing-token");
assert!(response.error.unwrap_or_default().contains("Heal task not found"));
}
async fn stop_after_selected_start(canonical: bool) {
let manager = create_test_heal_manager();
let mut processor = HealChannelProcessor::new(manager.clone());
let (selected_tx, selected_rx) = oneshot::channel();
let (release_tx, release_rx) = oneshot::channel();
processor.command_pause = Some((selected_tx, release_rx));
let (command_tx, command_rx) = mpsc::unbounded_channel();
let (receipt_tx, receipt_rx) = mpsc::unbounded_channel();
let stop = CancellationToken::new();
let worker_stop = stop.clone();
let worker = tokio::spawn(async move { processor.start_with_shutdown(command_rx, receipt_rx, worker_stop).await });
let request = HealChannelRequest {
id: "active-shutdown-command".to_string(),
bucket: "bucket".to_string(),
object_prefix: Some("active".to_string()),
source: HealRequestSource::Admin,
..Default::default()
};
let mut legacy_response = None;
let mut canonical_response = None;
if canonical {
let (response_tx, response_rx) = oneshot::channel();
receipt_tx.send(HealReceiptCommand { request, response_tx }).unwrap();
canonical_response = Some(response_rx);
} else {
let (response_tx, response_rx) = oneshot::channel();
command_tx.send(HealChannelCommand::Start { request, response_tx }).unwrap();
legacy_response = Some(response_rx);
}
tokio::time::timeout(Duration::from_secs(1), selected_rx)
.await
.unwrap()
.unwrap();
let pending = HealChannelRequest {
id: "queued-shutdown-command".to_string(),
bucket: "bucket".to_string(),
object_prefix: Some("pending".to_string()),
source: HealRequestSource::Admin,
..Default::default()
};
let (start_tx, start_rx) = oneshot::channel();
command_tx
.send(HealChannelCommand::Start {
request: pending.clone(),
response_tx: start_tx,
})
.unwrap();
let (query_tx, query_rx) = oneshot::channel();
command_tx
.send(HealChannelCommand::Query {
heal_path: "bucket/active".to_string(),
client_token: "active-shutdown-command".to_string(),
since_seq: None,
response_tx: query_tx,
})
.unwrap();
let (cancel_tx, cancel_rx) = oneshot::channel();
command_tx
.send(HealChannelCommand::Cancel {
heal_path: "bucket/active".to_string(),
client_token: "active-shutdown-command".to_string(),
response_tx: cancel_tx,
})
.unwrap();
let (receipt_response_tx, receipt_response_rx) = oneshot::channel();
receipt_tx
.send(HealReceiptCommand {
request: pending,
response_tx: receipt_response_tx,
})
.unwrap();
stop.cancel();
tokio::task::yield_now().await;
assert!(!worker.is_finished(), "stop cannot cancel the selected command");
release_tx.send(()).unwrap();
tokio::time::timeout(Duration::from_secs(1), worker)
.await
.unwrap()
.unwrap()
.unwrap();
if let Some(response) = legacy_response {
assert_eq!(response.await.unwrap().unwrap(), HealAdmissionResult::Accepted);
}
if let Some(response) = canonical_response {
let receipt = response.await.unwrap().unwrap();
assert_eq!(receipt.result, HealAdmissionResult::Accepted);
assert_eq!(receipt.task_id, "active-shutdown-command");
}
let expected = Error::other("Heal channel is stopping").to_string();
assert_eq!(start_rx.await.unwrap().unwrap_err(), expected);
assert_eq!(query_rx.await.unwrap().unwrap_err(), expected);
assert_eq!(cancel_rx.await.unwrap().unwrap_err(), expected);
assert_eq!(receipt_response_rx.await.unwrap().unwrap_err(), expected);
assert!(command_tx.is_closed());
assert!(receipt_tx.is_closed());
assert_eq!(
manager
.get_task_status_for_path("bucket/active", "active-shutdown-command")
.await
.unwrap(),
HealTaskStatus::Pending,
"queued Cancel must not execute"
);
assert!(
manager
.get_task_status_for_path("bucket/pending", "queued-shutdown-command")
.await
.is_err(),
"queued starts must not be admitted"
);
manager.stop().await.unwrap();
}
#[tokio::test]
async fn owned_channel_stop_finishes_legacy_command_and_rejects_all_queued_variants() {
stop_after_selected_start(false).await;
}
#[tokio::test]
async fn owned_channel_stop_finishes_receipt_command_and_rejects_all_queued_variants() {
stop_after_selected_start(true).await;
}
}
+30 -5
View File
@@ -125,6 +125,7 @@ pub struct ErasureSetHealer {
replacement_task_id: Option<String>,
replacement_target_identities: Option<Arc<[ReplacementTargetIdentity]>>,
replacement_execution: Option<Arc<super::storage::ReplacementExecution>>,
work_parent: Option<super::manager::HealWork>,
mainline_pacer: Option<Arc<super::pacing::MainlinePacer>>,
}
@@ -351,6 +352,7 @@ impl ErasureSetHealer {
replacement_task_id: None,
replacement_target_identities: None,
replacement_execution: None,
work_parent: None,
mainline_pacer: None,
}
}
@@ -365,6 +367,11 @@ impl ErasureSetHealer {
self
}
pub(super) fn with_work_parent(mut self, parent: Option<super::manager::HealWork>) -> Self {
self.work_parent = parent;
self
}
pub(crate) fn with_replacement_execution(mut self, execution: Option<Arc<super::storage::ReplacementExecution>>) -> Self {
self.replacement_execution = execution;
self
@@ -1694,6 +1701,8 @@ impl ErasureSetHealer {
let admin_identity = self.admin_identity(bucket, &object_name, version_id.as_deref(), incarnation);
let admin_task_id = self.admin_task.as_ref().map(|task| task.id.clone());
let execution = self.replacement_execution.clone();
let replacement_object = execution.is_some();
let work_parent = execution.as_ref().and(self.work_parent.as_ref()).cloned();
let failure_identity = execution
.as_ref()
.map(|_| (dedup_key.clone(), object_name.clone(), version_id.clone()));
@@ -1721,10 +1730,22 @@ impl ErasureSetHealer {
let result = if cancel_token.is_cancelled() {
(0, Err(Error::TaskCancelled))
} else {
match storage
.heal_object(&bucket_name, &object_name, version_id.as_deref(), &heal_opts)
.await
{
let result = if replacement_object {
storage
.heal_object_with_task_cancel(
&bucket_name,
&object_name,
version_id.as_deref(),
&heal_opts,
&cancel_token,
)
.await
} else {
storage
.heal_object(&bucket_name, &object_name, version_id.as_deref(), &heal_opts)
.await
};
match result {
Ok((result, None)) if target_outcomes_complete(&result, &target_endpoints) => {
let object_size = result_object_size_u64(&result);
if !replacement_commit_evidence_required {
@@ -1802,7 +1823,11 @@ impl ErasureSetHealer {
};
page_tasks.push(async move {
if let (Some(execution), Some((key, object, version))) = (execution, failure_identity) {
match execution.run(work).await {
let result = match work_parent {
Some(parent) => execution.run_owned(&parent, work).await,
None => execution.run(work).await,
};
match result {
Ok(result) => result,
Err(error) => (key, object, version, (0, Err(error)), None),
}
+278 -34
View File
@@ -19,10 +19,14 @@ use crate::heal::{
ReplacementPhase, ReplacementRecoveryBlock, ReplacementRecoveryCandidate, ResumeGc, ResumeManager, ResumeState,
ResumeUtils, merge_replacement_recovery_candidate, replacement_retry_is_exhausted_active,
},
storage::HealStorageAPI,
storage::{HealStorageAPI, run_local_metadata},
task::{HealOptions, HealPriority, HealRequest, HealTask, HealTaskStatus, HealType, demote_to_debug_when},
};
use crate::{Error, Result};
use futures::{
FutureExt,
future::{BoxFuture, Shared},
};
use metrics::{counter, gauge, histogram};
use rustfs_concurrency::WorkloadAdmissionSnapshotProvider;
use rustfs_concurrency::workload::{ForegroundPressure, foreground_pressure};
@@ -32,8 +36,10 @@ use rustfs_heal_contracts::heal_channel::{
HealAdmissionDropReason, HealAdmissionReceipt, HealAdmissionResult, HealRequestSource,
};
use rustfs_madmin::heal_commands::HealResultItem;
use std::future::Future;
#[cfg(test)]
use std::sync::LazyLock;
use std::sync::atomic::{AtomicBool, Ordering};
use std::{
collections::{BinaryHeap, HashMap, HashSet},
sync::{Arc, Mutex as StdMutex, MutexGuard as StdMutexGuard},
@@ -44,11 +50,206 @@ use tokio::{
time::{interval, sleep},
};
use tokio_util::sync::CancellationToken;
use tokio_util::task::{TaskTracker, task_tracker::TaskTrackerToken};
use tracing::{debug, error, info, warn};
use uuid::Uuid;
use super::{DiskError, Endpoint, HealDiskExt as _, local_disk_map_read};
// Four long-lived loops have fixed receipts; finished dynamic jobs leave the
// tracker immediately. No manager/cache lock is needed by a drain waiter.
type HealWorkerJoin = Shared<BoxFuture<'static, std::result::Result<(), Arc<tokio::task::JoinError>>>>;
struct HealWorkerRegistration {
open: bool,
stopping: usize,
fixed: [Option<HealWorkerJoin>; 4],
}
#[derive(Clone)]
struct HealWorkers {
registration: Arc<StdMutex<HealWorkerRegistration>>,
tasks: TaskTracker,
startup: TaskTracker,
failed: Arc<AtomicBool>,
}
#[derive(Clone)]
pub(super) struct HealWork {
owner: HealWorkers,
_token: TaskTrackerToken,
}
// Each in-flight stop owns its own intent. Dropping one waiter cannot reopen
// startup while another stop still owns the handoff boundary.
struct HealStopIntent(HealWorkers);
impl Drop for HealStopIntent {
fn drop(&mut self) {
if let Ok(mut registration) = self.0.registration.lock() {
registration.stopping -= 1;
if registration.stopping == 0 && registration.open {
self.0.startup.reopen();
}
}
}
}
struct HealWorkGuard {
work: HealWork,
completed: bool,
}
impl Drop for HealWorkGuard {
fn drop(&mut self) {
if !self.completed {
self.work.owner.failed.store(true, Ordering::Release);
}
}
}
impl HealWork {
fn spawn<F, Fut, T>(&self, work: F) -> tokio::task::JoinHandle<T>
where
F: FnOnce(HealWork) -> Fut + Send + 'static,
Fut: Future<Output = T> + Send + 'static,
T: Send + 'static,
{
// Clone before spawn, while the live parent prevents an empty tracker.
let guard = HealWorkGuard {
work: self.clone(),
completed: false,
};
tokio::spawn(async move {
let mut guard = guard;
let output = work(guard.work.clone()).await;
guard.completed = true;
output
})
}
}
impl HealWorkers {
fn new() -> Self {
Self {
registration: Arc::new(StdMutex::new(HealWorkerRegistration {
open: true,
stopping: 0,
fixed: std::array::from_fn(|_| None),
})),
tasks: TaskTracker::new(),
startup: TaskTracker::new(),
failed: Arc::new(AtomicBool::new(false)),
}
}
fn enter_start(&self) -> Result<(HealWorkGuard, TaskTrackerToken)> {
let registration = self
.registration
.lock()
.map_err(|_| Error::other("heal worker registration poisoned"))?;
if !registration.open || registration.stopping != 0 {
return Err(Error::other("heal manager workers are stopping or terminally closed"));
}
Ok((
HealWorkGuard {
work: HealWork {
owner: self.clone(),
_token: self.tasks.token(),
},
completed: false,
},
self.startup.token(),
))
}
fn begin_stop(&self) -> Result<HealStopIntent> {
let mut registration = self
.registration
.lock()
.map_err(|_| Error::other("heal worker registration poisoned"))?;
registration.stopping += 1;
self.startup.close();
Ok(HealStopIntent(self.clone()))
}
fn spawn_fixed<F, Fut>(&self, slot: usize, work: F) -> Result<()>
where
F: FnOnce(HealWork) -> Fut + Send + 'static,
Fut: Future<Output = ()> + Send + 'static,
{
let mut registration = self
.registration
.lock()
.map_err(|_| Error::other("heal worker registration poisoned"))?;
if !registration.open || registration.fixed[slot].is_some() {
return Err(Error::other("heal manager loop is closed or already registered"));
}
let parent = HealWork {
owner: self.clone(),
_token: self.tasks.token(),
};
let task = parent.spawn(work);
registration.fixed[slot] = Some(async move { task.await.map_err(Arc::new) }.boxed().shared());
Ok(())
}
fn close(&self) -> Result<()> {
let mut registration = self
.registration
.lock()
.map_err(|_| Error::other("heal worker registration poisoned"))?;
registration.open = false;
self.tasks.close();
Ok(())
}
async fn wait(&self) -> Result<()> {
let fixed = {
let registration = self
.registration
.lock()
.map_err(|_| Error::other("heal worker registration poisoned"))?;
if registration.open {
return Err(Error::other("heal worker admission is not closed"));
}
registration.fixed.clone()
};
let mut join_failed = false;
for task in fixed.into_iter().flatten() {
join_failed |= task.await.is_err();
}
self.tasks.wait().await;
if join_failed || self.failed.load(Ordering::Acquire) {
return Err(Error::other("heal owned worker panicked or was unexpectedly cancelled"));
}
Ok(())
}
}
pub(super) fn spawn_heal_work<F, Fut, T>(parent: Option<&HealWork>, work: F) -> tokio::task::JoinHandle<T>
where
F: FnOnce(Option<HealWork>) -> Fut + Send + 'static,
Fut: Future<Output = T> + Send + 'static,
T: Send + 'static,
{
match parent {
Some(parent) => parent.spawn(move |child| work(Some(child))),
None => tokio::spawn(work(None)),
}
}
async fn inspect_resume_disk(
gc: &mut ResumeGc,
disk: &super::DiskStore,
admission: Option<&Arc<rustfs_concurrency::adaptive::NodeAdmission>>,
cancellation: &CancellationToken,
) -> Result<()> {
run_local_metadata(disk, admission, cancellation, gc.inspect_disk(disk))
.await?
.map(|_| ())
}
const KEEP_HEAL_TASK_STATUS_DURATION: Duration = Duration::from_secs(10 * 60);
// Each cache includes alias tokens in its count and byte budget. Eviction
// removes every token sharing a snapshot; neither cache retains repair state.
@@ -994,6 +1195,7 @@ pub struct HealState {
/// Heal manager
pub struct HealManager {
workers: HealWorkers,
/// Heal config
config: Arc<RwLock<HealConfig>>,
/// Heal state
@@ -1168,15 +1370,22 @@ impl HealManager {
)
}
fn schedule_mainline_throttle_recheck(notify: Arc<Notify>, delay: Duration) {
fn schedule_mainline_throttle_recheck(
notify: Arc<Notify>,
delay: Duration,
parent: Option<&HealWork>,
cancel: CancellationToken,
) {
if delay.is_zero() {
notify.notify_one();
return;
}
tokio::spawn(async move {
sleep(delay).await;
notify.notify_one();
spawn_heal_work(parent, move |_parent| async move {
tokio::select! {
_ = cancel.cancelled() => {},
_ = sleep(delay) => notify.notify_one(),
}
});
}
@@ -1571,9 +1780,10 @@ impl HealManager {
/// Start the bounded resume-state inspector. Destructive GC remains
/// disabled until the durable owner/CAS contract from backlog#1927 is
/// available; this task therefore cannot remove an active or stale file.
async fn start_resume_gc(&self) {
async fn start_resume_gc(&self) -> Result<()> {
let cancel = self.cancel_token.clone();
tokio::spawn(async move {
let admission = self.storage.background_node_admission();
self.workers.spawn_fixed(2, move |_parent| async move {
let mut gc_by_disk = HashMap::<String, ResumeGc>::new();
let mut ticker = interval(RESUME_GC_INTERVAL);
loop {
@@ -1589,7 +1799,7 @@ impl HealManager {
let gc = gc_by_disk.entry(disk_key).or_default();
tokio::select! {
_ = cancel.cancelled() => return,
result = gc.inspect_disk(&disk) => {
result = inspect_resume_disk(gc, &disk, admission.as_ref(), &cancel) => {
if let Err(error) = result {
warn!(
target: "rustfs::heal::manager",
@@ -1608,13 +1818,14 @@ impl HealManager {
}
}
}
});
})?;
Ok(())
}
async fn start_root_recovery_terminal_gc(&self) {
async fn start_root_recovery_terminal_gc(&self) -> Result<()> {
let cancel = self.cancel_token.clone();
let root_recovery = self.root_recovery.clone();
tokio::spawn(async move {
self.workers.spawn_fixed(3, move |_parent| async move {
let mut ticker = interval(RESUME_GC_INTERVAL);
loop {
tokio::select! {
@@ -1653,7 +1864,8 @@ impl HealManager {
}
}
}
});
})?;
Ok(())
}
/// Create new HealManager
@@ -1678,6 +1890,7 @@ impl HealManager {
) -> Self {
let config = config.unwrap_or_default();
Self {
workers: HealWorkers::new(),
config: Arc::new(RwLock::new(config)),
state: Arc::new(RwLock::new(HealState::default())),
active_heals: Arc::new(Mutex::new(HashMap::new())),
@@ -1724,8 +1937,10 @@ impl HealManager {
/// Start HealManager
pub async fn start(&self) -> Result<()> {
let (mut starting, _startup_token) = self.workers.enter_start()?;
let mut state = self.state.write().await;
if state.is_running {
starting.completed = true;
warn!(
target: "rustfs::heal::manager",
event = EVENT_HEAL_MANAGER_STATE,
@@ -1778,8 +1993,8 @@ impl HealManager {
self.process_unclean_shutdown().await;
// Inspect resume artifacts in a bounded, fail-closed background task.
self.start_resume_gc().await;
self.start_root_recovery_terminal_gc().await;
self.start_resume_gc().await?;
self.start_root_recovery_terminal_gc().await?;
// start auto disk scanner to heal unformatted disks
if self.config.read().await.enable_auto_heal {
@@ -1803,11 +2018,31 @@ impl HealManager {
state = "running",
"Heal manager started"
);
starting.completed = true;
Ok(())
}
#[cfg(test)]
pub(super) fn test_work_parent(&self) -> HealWork {
assert!(self.workers.registration.lock().unwrap().open, "test root parent precedes stop");
HealWork {
owner: self.workers.clone(),
_token: self.workers.tasks.token(),
}
}
/// Durable stop handoff completed; owned workers may still be joining.
pub(crate) fn stop_handoff_completed(&self) -> bool {
self.cancel_token.is_cancelled()
}
/// Stop HealManager
pub async fn stop(&self) -> Result<()> {
let _stop_intent = self.workers.begin_stop()?;
// TaskTracker registers its notification before checking closed/empty.
// Stop admission first; wait outside admin/active/cache locks so startup
// can finish its existing recovery and register all four loop receipts.
self.workers.startup.wait().await;
let _admin_start_guard = self.admin_start_shutdown.lock().await;
info!(
target: "rustfs::heal::manager",
@@ -1824,29 +2059,32 @@ impl HealManager {
let mut active_heals = self.active_heals.lock().await;
let queue = self.heal_queue.lock().await;
let retrying = self.retrying_heals.lock().await;
for task in active_heals.values() {
if root_recovery::is_admin_heal_recovery(&task.heal_type, task.source) {
if task.get_status().await == HealTaskStatus::Completed {
self.root_recovery.remove(&task.id, &task.heal_type, task.source).await?;
} else {
let mut request = match task.retry_request_with_remaining_timeout().await {
Ok(request) => request,
Err(Error::TaskTimeout) => {
let mut request = task.retry_request();
request.options.timeout = Some(Duration::ZERO);
request
}
Err(error) => return Err(error),
};
request.retry_attempts = task.retry_attempts;
self.root_recovery.persist(&request).await?;
if !self.cancel_token.is_cancelled() {
for task in active_heals.values() {
if root_recovery::is_admin_heal_recovery(&task.heal_type, task.source) {
if task.get_status().await == HealTaskStatus::Completed {
self.root_recovery.remove(&task.id, &task.heal_type, task.source).await?;
} else {
let mut request = match task.retry_request_with_remaining_timeout().await {
Ok(request) => request,
Err(Error::TaskTimeout) => {
let mut request = task.retry_request();
request.options.timeout = Some(Duration::ZERO);
request
}
Err(error) => return Err(error),
};
request.retry_attempts = task.retry_attempts;
self.root_recovery.persist(&request).await?;
}
}
}
for request in queue.requests().chain(retrying.values().map(|retrying| &retrying.request)) {
self.root_recovery.persist(request).await?;
}
self.workers.close()?;
self.cancel_token.cancel();
}
for request in queue.requests().chain(retrying.values().map(|retrying| &retrying.request)) {
self.root_recovery.persist(request).await?;
}
self.cancel_token.cancel();
drop(retrying);
drop(queue);
@@ -1890,6 +2128,12 @@ impl HealManager {
}
crate::set_heal_queue_length(0);
// Completion paths reacquire active/cache locks. Never wait while
// holding admin, active, queue, retry, cache or state guards.
drop(active_heals);
drop(_admin_start_guard);
self.workers.wait().await?;
// update state
let mut state = self.state.write().await;
state.is_running = false;
+49 -5
View File
@@ -36,6 +36,7 @@ impl HealManager {
let displaced_terminals = self.displaced_terminals.clone();
let mrf_repair_notice_targets = self.mrf_repair_notice_targets.clone();
let storage = self.storage.clone();
let admission = storage.background_node_admission();
let replacement_recovery_anchors = self.replacement_recovery_anchors.clone();
let blocks = self.replacement_recovery_blocks.clone();
let cancel_token = self.cancel_token.clone();
@@ -57,7 +58,7 @@ impl HealManager {
"Heal auto disk scanner started"
);
tokio::spawn(async move {
self.workers.spawn_fixed(1, move |_parent| async move {
let mut interval = interval(duration);
loop {
@@ -88,6 +89,10 @@ impl HealManager {
let mut conflicted_recovery_tasks = HashSet::<String>::new();
let mut conflicted_recovery_sets = HashSet::<String>::new();
let mut deferred_recovery_sets = HashSet::<String>::new();
// This tick lacks local evidence for these sets/endpoints;
// later inventory must not turn the deferral into absence.
let mut metadata_deferred_sets = HashSet::<String>::new();
let mut metadata_deferred_endpoints = HashSet::<String>::new();
let mut deferred_replacement_endpoints = HashSet::<String>::new();
let local_disks = {
let local_disk_map = local_disk_map_read().await;
@@ -133,8 +138,32 @@ impl HealManager {
continue;
}
// detect unformatted disk via get_disk_id()
match disk.get_disk_id().await {
// Preserve backend format results separately from admission deferral.
let probe = run_local_metadata(
disk, admission.as_ref(), &cancel_token, disk.get_disk_id(),
).await;
// Native capability projection can defer after the outer plan was acquired.
let resource_deferred = admission.is_some()
&& matches!(&probe, Ok(Err(DiskError::Io(error))) if error.kind() == std::io::ErrorKind::WouldBlock);
let disk_id = match probe {
Ok(result) if !resource_deferred => result,
_ => {
if cancel_token.is_cancelled() {
return;
}
metadata_deferred_endpoints.insert(endpoint.to_string());
if let Some(set) = &set_disk_id {
metadata_deferred_sets.insert(set.clone());
deferred_recovery_sets.insert(set.clone());
}
skipped_invalid_count += 1;
continue;
}
};
if admission.is_some() && cancel_token.is_cancelled() {
return;
}
match disk_id {
Err(DiskError::UnformattedDisk) => {
if !super::super::replacement_readiness::auto_replacement_target_ready(disk, &local_disks)
.await
@@ -203,8 +232,18 @@ impl HealManager {
let mut recovery_disks = BTreeMap::new();
extend_replacement_recovery_disks(&mut recovery_disks, local_disks.iter().cloned());
for set_disk_id in &recovery_sets {
match storage.replacement_intent_disks(set_disk_id).await {
if metadata_deferred_sets.contains(set_disk_id) {
continue;
}
match storage.replacement_intent_disks(set_disk_id, Some(&cancel_token)).await {
Ok(disks) => extend_replacement_recovery_disks(&mut recovery_disks, disks),
Err(Error::TransientSkip { .. }) => {
if cancel_token.is_cancelled() {
return;
}
metadata_deferred_sets.insert(set_disk_id.clone());
deferred_recovery_sets.insert(set_disk_id.clone());
}
Err(error) => {
conflicted_recovery_sets.insert(set_disk_id.clone());
warn!(
@@ -230,6 +269,11 @@ impl HealManager {
let endpoint = disk.endpoint();
let disk_set_disk_id =
crate::heal::utils::format_set_disk_id_from_i32(endpoint.pool_idx, endpoint.set_idx);
if metadata_deferred_endpoints.contains(anchor)
|| disk_set_disk_id.as_ref().is_some_and(|set| metadata_deferred_sets.contains(set))
{
continue;
}
let replacement_task_ids = match ResumeUtils::get_replacement_intent_tasks(disk).await {
Ok(task_ids) => task_ids,
Err(error) => {
@@ -793,7 +837,7 @@ impl HealManager {
}
}
}
});
})?;
Ok(())
}
}
+19 -10
View File
@@ -50,7 +50,7 @@ impl HealManager {
let workload_provider = self.workload_provider.clone();
let disposition_cache = self.disposition_cache.clone();
tokio::spawn(async move {
self.workers.spawn_fixed(0, move |parent| async move {
let mut interval = interval(config.read().await.heal_interval);
loop {
@@ -68,7 +68,7 @@ impl HealManager {
break;
}
_ = notify.notified(), if event_driven_scheduler_enable => {
Self::process_heal_queue(HealQueueContext {
Self::process_owned_heal_queue(HealQueueContext {
heal_queue: &heal_queue,
active_heals: &active_heals,
completed_heals: &completed_heals,
@@ -85,11 +85,11 @@ impl HealManager {
cancel_token: &cancel_token,
workload_provider: &workload_provider,
disposition_cache: &disposition_cache,
})
}, &parent)
.await;
}
_ = interval.tick() => {
Self::process_heal_queue(HealQueueContext {
Self::process_owned_heal_queue(HealQueueContext {
heal_queue: &heal_queue,
active_heals: &active_heals,
completed_heals: &completed_heals,
@@ -106,19 +106,28 @@ impl HealManager {
cancel_token: &cancel_token,
workload_provider: &workload_provider,
disposition_cache: &disposition_cache,
})
}, &parent)
.await;
}
}
}
});
})?;
Ok(())
}
/// Process heal queue
/// Processes multiple tasks per cycle when capacity allows and queue has high-priority items
#[cfg(test)]
pub(super) async fn process_heal_queue(context: HealQueueContext<'_>) {
Self::process_heal_queue_inner(context, None).await;
}
async fn process_owned_heal_queue(context: HealQueueContext<'_>, parent: &HealWork) {
Self::process_heal_queue_inner(context, Some(parent)).await;
}
async fn process_heal_queue_inner(context: HealQueueContext<'_>, parent: Option<&HealWork>) {
let HealQueueContext {
heal_queue,
active_heals,
@@ -242,7 +251,7 @@ impl HealManager {
let config_for_spawn = config.clone();
// start heal task
tokio::spawn(async move {
spawn_heal_work(parent, move |parent| async move {
debug!(
target: "rustfs::heal::manager",
event = EVENT_HEAL_SCHEDULER_STATE,
@@ -255,7 +264,7 @@ impl HealManager {
state = "task_started",
"Heal scheduler task started"
);
let result = task.execute().await;
let result = task.execute_owned(parent.clone()).await;
let retry_request = retry_request_for_result_with_budget(task.as_ref(), &result).await;
record_task_disposition(&disposition_cache_clone, task.as_ref(), &result, retry_request.is_some()).await;
match &result {
@@ -494,7 +503,7 @@ impl HealManager {
let retry_notify = notify_clone.clone();
let retry_manager_cancel_token = manager_cancel_token.clone();
let retry_config = config_for_spawn.clone();
tokio::spawn(async move {
spawn_heal_work(parent.as_ref(), move |_parent| async move {
let mut grace_admission_backoff = Duration::ZERO;
loop {
let now = SystemTime::now()
@@ -723,7 +732,7 @@ impl HealManager {
if delayed_by_mainline_throttle && let Some(pressure) = mainline_pressure {
Self::record_mainline_throttle_delay(pressure, &config);
Self::schedule_mainline_throttle_recheck(notify.clone(), config.mainline_max_sleep);
Self::schedule_mainline_throttle_recheck(notify.clone(), config.mainline_max_sleep, parent, cancel_token.clone());
}
// Log queue status if items remain
+18
View File
@@ -28,6 +28,7 @@ use tempfile::TempDir;
pub(super) mod admin_overlap;
mod admission_timing;
mod owned_workers;
mod root_recovery;
use uuid::Uuid;
mod auto_scan;
@@ -6096,6 +6097,7 @@ fn admin_bucket_discovery_retries_keep_the_existing_budget() {
async fn check_node_admission_task_retry(source: HealRequestSource) {
use crate::heal::outcome::HealFailureClass;
use crate::heal::storage::ECStoreHealStorage;
use crate::heal::storage_api::owner::{PhysicalDiskSample, PhysicalDiskSampler};
use rustfs_concurrency::adaptive::{FeedbackController, FeedbackObservation, NodeAdmission, ResourcePressure};
use rustfs_heal_contracts::heal_channel::HealScanMode;
use std::num::NonZeroU16;
@@ -6108,6 +6110,18 @@ async fn check_node_admission_task_retry(source: HealRequestSource) {
env.make_bucket("admission-retry", false).await;
env.put_object_bytes("admission-retry", "existing", b"repair me".to_vec())
.await;
let mut physical = PhysicalDiskSampler::default();
physical.publish(
PhysicalDiskSample::collect(&env.ecstore).with_test_single_device(),
ResourcePressure::Unknown,
);
assert!(
env.ecstore
.physical_heal_plan("existing", &Default::default())
.expect("fixture placement")
.is_some(),
"physical discovery must be ready before testing node capacity"
);
let gate = NodeAdmission::new(NonZeroU16::new(1).expect("one node unit"));
let storage: Arc<dyn HealStorageAPI> =
Arc::new(ECStoreHealStorage::new(env.ecstore.clone()).with_admission(Some(gate.clone())));
@@ -6148,6 +6162,10 @@ async fn check_node_admission_task_retry(source: HealRequestSource) {
completed_units: None,
}),
);
physical.publish(
PhysicalDiskSample::collect(&env.ecstore).with_test_single_device(),
ResourcePressure::Unknown,
);
let retried = HealTask::from_request(retry, storage);
retried
.execute()
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,366 @@
// Copyright 2026 RustFS Team
// Licensed under the Apache License, Version 2.0.
use super::*;
fn owned_config() -> HealConfig {
HealConfig {
enable_auto_heal: false,
event_driven_scheduler_enable: true,
heal_interval: Duration::from_secs(60),
mainline_throttle_enable: false,
..Default::default()
}
}
async fn until(mut ready: impl AsyncFnMut() -> bool) {
tokio::time::timeout(Duration::from_secs(5), async {
while !ready().await {
tokio::task::yield_now().await;
}
})
.await
.expect("real producer reaches expected boundary");
}
async fn completion_fixture() -> (Arc<HealManager>, Arc<CompletedRetentionHook>, String, String) {
let manager = Arc::new(HealManager::new_without_root_recovery_for_test(
Arc::new(MockStorage),
Some(owned_config()),
));
let bucket = format!("owned-worker-{}", Uuid::new_v4());
let request = HealRequest::object(bucket.clone(), "object".to_owned(), None);
let id = request.id.clone();
let hook = Arc::new(CompletedRetentionHook {
pause_before_publish: true,
..Default::default()
});
{
let mut hooks = COMPLETED_RETENTION_HOOKS.lock().await;
hooks.insert(bucket.clone(), hook.clone());
hooks.insert(id.clone(), hook.clone());
}
manager.submit_heal_request(request).await.expect("queue real object heal");
manager.start().await.expect("start actual scheduler and both GCs");
tokio::time::timeout(Duration::from_secs(5), hook.started.notified())
.await
.expect("actual worker reaches storage");
hook.execute.notify_one();
tokio::time::timeout(Duration::from_secs(5), hook.before_publish.notified())
.await
.expect("actual worker reaches existing publication boundary");
(manager, hook, bucket, id)
}
async fn remove_hook(bucket: &str, id: &str) {
COMPLETED_RETENTION_HOOKS
.lock()
.await
.retain(|key, _| key != bucket && key != id);
}
#[tokio::test]
async fn owned_manager_waits_actual_worker_without_holding_manager_locks() {
let (manager, hook, bucket, id) = completion_fixture().await;
let stopping = manager.clone();
let mut stop = tokio::spawn(async move { stopping.stop().await });
manager.cancel_token.cancelled().await;
assert!(tokio::time::timeout(Duration::from_millis(30), &mut stop).await.is_err());
// The real worker needs these locks after the existing test boundary.
let admin = manager.admin_start_shutdown.try_lock().expect("admin released before join");
let active = manager.active_heals.try_lock().expect("active released before join");
let cache = manager.completed_heals.try_lock().expect("cache released before join");
assert!(manager.workers.tasks.len() >= 1, "actual publication worker is still owned");
drop((admin, active, cache));
hook.publish.notify_one();
tokio::time::timeout(Duration::from_secs(5), hook.handoff.notified())
.await
.expect("real worker reaches existing final handoff boundary");
hook.finish.notify_one();
tokio::time::timeout(Duration::from_secs(5), stop)
.await
.expect("worker drain deadline")
.expect("stop waiter")
.expect("actual worker joined");
assert!(manager.workers.tasks.is_empty());
assert!(
manager
.workers
.registration
.lock()
.unwrap()
.fixed
.iter()
.flatten()
.all(|join| join.peek().is_some())
);
assert!(manager.start().await.is_err(), "terminal owner cannot restart");
manager.stop().await.expect("successful receipt is repeatable");
remove_hook(&bucket, &id).await;
}
#[tokio::test]
async fn owned_manager_cancelled_stop_waiter_keeps_actual_worker_receipt() {
let (manager, hook, bucket, id) = completion_fixture().await;
let stopping = manager.clone();
let stop = tokio::spawn(async move { stopping.stop().await });
manager.cancel_token.cancelled().await;
until(async || manager.admin_start_shutdown.try_lock().is_ok()).await;
stop.abort();
assert!(stop.await.unwrap_err().is_cancelled());
let stopping = manager.clone();
let mut retry = tokio::spawn(async move { stopping.stop().await });
assert!(tokio::time::timeout(Duration::from_millis(30), &mut retry).await.is_err());
hook.publish.notify_one();
tokio::time::timeout(Duration::from_secs(5), hook.handoff.notified())
.await
.expect("real worker reaches existing final handoff boundary");
hook.finish.notify_one();
tokio::time::timeout(Duration::from_secs(5), retry)
.await
.expect("repeat drain deadline")
.expect("retry waiter")
.expect("same actual worker joined");
assert!(manager.workers.tasks.is_empty());
remove_hook(&bucket, &id).await;
}
#[tokio::test]
async fn owned_manager_start_stop_race_counts_independent_stop_intents() {
let manager = Arc::new(HealManager::new_without_root_recovery_for_test(
Arc::new(MockStorage),
Some(owned_config()),
));
let admin = manager.admin_start_shutdown.lock().await;
let starting = manager.clone();
let start = tokio::spawn(async move { starting.start().await });
until(async || manager.state.read().await.is_running).await;
let stopping = manager.clone();
let first = tokio::spawn(async move { stopping.stop().await });
let stopping = manager.clone();
let second = tokio::spawn(async move { stopping.stop().await });
until(async || manager.workers.registration.lock().unwrap().stopping == 2).await;
first.abort();
assert!(first.await.unwrap_err().is_cancelled());
assert_eq!(manager.workers.registration.lock().unwrap().stopping, 1);
assert!(manager.start().await.is_err(), "cancelled waiter must not reopen another waiter's gate");
assert!(
!manager.cancel_token.is_cancelled(),
"handoff waits for startup before signalling cancellation"
);
assert!(!second.is_finished());
drop(admin);
start
.await
.expect("actual startup waiter")
.expect("startup finishes before handoff");
tokio::time::timeout(Duration::from_secs(5), second)
.await
.expect("stop drains all newly registered loops")
.expect("stop waiter")
.expect("joined");
assert!(manager.workers.tasks.is_empty());
assert_eq!(manager.workers.registration.lock().unwrap().fixed.iter().flatten().count(), 3);
}
#[tokio::test]
async fn owned_manager_cancelled_pre_handoff_stop_reopens_start_admission() {
let manager = Arc::new(HealManager::new_without_root_recovery_for_test(
Arc::new(MockStorage),
Some(owned_config()),
));
let active = manager.active_heals.lock().await;
let stopping = manager.clone();
let stop = tokio::spawn(async move { stopping.stop().await });
until(async || manager.workers.registration.lock().unwrap().stopping == 1).await;
until(async || manager.admin_start_shutdown.try_lock().is_err()).await;
stop.abort();
assert!(stop.await.unwrap_err().is_cancelled());
assert_eq!(manager.workers.registration.lock().unwrap().stopping, 0);
assert!(!manager.cancel_token.is_cancelled());
drop(active);
manager
.start()
.await
.expect("pre-handoff cancellation did not terminally close startup");
manager.stop().await.expect("later real startup and stop succeed");
}
#[tokio::test]
async fn owned_manager_failed_durable_handoff_remains_open_and_retryable() {
use super::super::super::RUSTFS_META_BUCKET;
use super::super::root_recovery::RootHealRecovery;
let temp = TempDir::new().expect("owned recovery disk");
let endpoint = Endpoint::try_from(temp.path().to_string_lossy().as_ref()).expect("endpoint");
let disk = new_disk(
&endpoint,
&DiskOption {
cleanup: false,
health_check: false,
},
)
.await
.expect("real local disk");
match disk.make_volume(RUSTFS_META_BUCKET).await {
Ok(()) | Err(DiskError::VolumeExists) => {}
Err(error) => panic!("metadata volume: {error}"),
}
let mut manager = HealManager::new(Arc::new(MockStorage), Some(owned_config()));
manager.root_recovery = Arc::new(RootHealRecovery::with_disks(vec![disk.clone()]));
let mut request = HealRequest::new(HealType::Cluster, HealOptions::default(), HealPriority::High);
request.source = HealRequestSource::Admin;
let id = request.id.clone();
manager.submit_heal_request(request).await.expect("real durable admission");
std::fs::remove_dir_all(temp.path().join(RUSTFS_META_BUCKET)).expect("remove accepted owner volume");
assert!(manager.stop().await.is_err(), "real durable handoff fails");
assert!(!manager.stop_handoff_completed());
assert!(!manager.cancel_token.is_cancelled());
assert!(manager.workers.registration.lock().unwrap().open);
assert_eq!(manager.workers.registration.lock().unwrap().stopping, 0);
assert!(manager.heal_queue.lock().await.requests().any(|request| request.id == id));
disk.make_volume(RUSTFS_META_BUCKET).await.expect("restore original owner");
manager.stop().await.expect("handoff retries on same disk and drains");
assert!(manager.stop_handoff_completed());
assert!(!manager.workers.registration.lock().unwrap().open);
assert_eq!(
manager
.root_recovery
.pending()
.await
.expect("durable successor remains")
.len(),
1
);
}
#[tokio::test]
async fn owned_manager_real_retry_and_throttle_wake_drain() {
let manager = Arc::new(HealManager::new_without_root_recovery_for_test(
Arc::new(MockStorage),
Some(owned_config()),
));
let hook = Arc::new(CompletedRetentionHook::default());
COMPLETED_RETENTION_HOOKS
.lock()
.await
.insert("retry-transition".to_owned(), hook.clone());
manager
.submit_heal_request(HealRequest::object("retry-transition".to_owned(), "object".to_owned(), None))
.await
.expect("queue actual recoverable storage failure");
manager.start().await.expect("actual scheduler");
tokio::time::timeout(Duration::from_secs(5), hook.started.notified())
.await
.expect("real object operation");
hook.execute.notify_one();
until(async || !manager.retrying_heals.lock().await.is_empty()).await;
manager.stop().await.expect("actual delayed retry drained");
assert!(manager.workers.tasks.is_empty());
assert!(manager.retrying_heals.lock().await.is_empty());
COMPLETED_RETENTION_HOOKS.lock().await.remove("retry-transition");
let provider: WorkloadSnapshotProviderRef = Arc::new(FixedWorkloadProvider {
class: WorkloadClass::ForegroundRead,
active: 8,
limit: 10,
state: AdmissionState::Open,
});
let manager = HealManager::new_with_workload_provider_without_root_recovery_for_test(
Arc::new(MockStorage),
Some(HealConfig {
mainline_throttle_enable: true,
mainline_read_utilization_high_percent: 80,
mainline_write_utilization_high_percent: 80,
mainline_max_sleep: Duration::from_secs(60),
..owned_config()
}),
Some(provider),
);
manager.start().await.expect("actual throttled scheduler");
until(async || manager.workers.tasks.len() == 4).await;
manager
.submit_heal_request(bucket_request("owned-throttle", HealPriority::Normal, HealRequestSource::ReadRepair))
.await
.expect("real pressure-delayed request");
until(async || manager.workers.tasks.len() > 4).await;
assert_eq!(manager.get_queue_length().await, 1);
tokio::time::timeout(Duration::from_secs(5), manager.stop())
.await
.expect("wake cancels without waiting 60s")
.expect("wake drained");
assert!(manager.workers.tasks.is_empty());
}
#[tokio::test]
async fn owned_manager_real_scheduler_panic_is_sticky_failure() {
let manager = HealManager::new_without_root_recovery_for_test(
Arc::new(MockStorage),
Some(HealConfig {
heal_interval: Duration::ZERO,
..owned_config()
}),
);
manager.start().await.expect("startup registered actual scheduler");
until(async || manager.workers.failed.load(Ordering::Acquire)).await;
assert!(manager.stop().await.is_err(), "actual interval panic cannot acknowledge joined success");
assert!(manager.stop_handoff_completed());
assert!(manager.stop().await.is_err(), "repeat wait cannot erase failure");
assert!(manager.workers.tasks.is_empty());
}
#[tokio::test]
async fn owned_manager_completed_dynamic_workers_do_not_accumulate_handles() {
let manager = HealManager::new_without_root_recovery_for_test(Arc::new(MockStorage), Some(owned_config()));
manager.start().await.expect("actual producer loops");
for index in 0..64 {
let request = HealRequest::object(format!("owned-completed-{index}"), "object".to_owned(), None);
let id = request.id.clone();
manager
.submit_heal_request(request)
.await
.expect("actual scheduler admission");
until(async || manager.completed_heals.lock().await.contains_key(&id)).await;
}
// Scheduler retains a live parent plus its guard; each GC retains one guard.
// These four fixed-loop tokens remain, irrespective of the 64 finished jobs.
until(async || manager.workers.tasks.len() == 4).await;
assert_eq!(manager.workers.registration.lock().unwrap().fixed.iter().flatten().count(), 3);
manager.stop().await.expect("bounded dynamic bookkeeping drains");
assert!(manager.workers.tasks.is_empty());
}
#[tokio::test]
async fn owned_manager_enabled_auto_scan_has_four_actual_loop_receipts() {
let manager = HealManager::new_without_root_recovery_for_test(
Arc::new(MockStorage),
Some(HealConfig {
enable_auto_heal: true,
..owned_config()
}),
);
manager.start().await.expect("four actual producer loops");
assert_eq!(manager.workers.registration.lock().unwrap().fixed.iter().flatten().count(), 4);
manager.stop().await.expect("four loops joined");
assert!(manager.workers.tasks.is_empty());
}
#[tokio::test]
async fn owned_manager_cancelled_actual_start_is_sticky_failure() {
let manager = Arc::new(HealManager::new_without_root_recovery_for_test(
Arc::new(MockStorage),
Some(owned_config()),
));
let admin = manager.admin_start_shutdown.lock().await;
let starting = manager.clone();
let start = tokio::spawn(async move { starting.start().await });
until(async || manager.state.read().await.is_running).await;
start.abort();
assert!(start.await.unwrap_err().is_cancelled());
drop(admin);
assert!(
manager.stop().await.is_err(),
"cancelled partial startup cannot become successful receipt"
);
assert!(manager.stop().await.is_err(), "failure remains sticky for later waiter");
assert!(manager.workers.tasks.is_empty());
}
@@ -3128,3 +3128,65 @@ async fn root_recovery_conflicting_record_does_not_block_other_replays() {
);
}
}
#[cfg(not(windows))]
#[tokio::test]
async fn owned_terminal_gc_waits_for_actual_native_receipt_read() {
use crate::heal::storage_api::{
EcstoreLocalPublicationPause as LocalPublicationPause, EcstoreLocalPublicationStage as LocalPublicationStage,
};
let (_temp, disk) = recovery_disk().await;
let manager = recovery_manager(vec![disk.clone()]);
let request = root_request();
let completed_at = SystemTime::now()
.checked_sub(KEEP_HEAL_TASK_STATUS_DURATION + Duration::from_secs(1))
.expect("expired fixture time");
let completed = completed_admin_status(&request.heal_type, completed_at);
manager
.publish_admin_terminal(&request.id, &request.heal_type, request.source, &completed)
.await
.expect("persist real expired terminal");
let path = format!("terminal-root-heal-{}.json", request.id);
assert!(
!disk
.read_all(RUSTFS_META_BUCKET, &path)
.await
.expect("nonempty real terminal")
.is_empty()
);
let mut pause =
LocalPublicationPause::install(disk.as_ref(), RUSTFS_META_BUCKET, &path, LocalPublicationStage::NativeReadOpen)
.expect("native bounded terminal-read pause");
manager
.start_root_recovery_terminal_gc()
.await
.expect("actual slot-3 producer");
tokio::time::timeout(Duration::from_secs(5), pause.entered())
.await
.expect("terminal GC native read entry deadline")
.expect("real bounded terminal read entered");
let mut stop = Box::pin(manager.stop());
let first = tokio::time::timeout(Duration::from_millis(30), stop.as_mut()).await;
let held = first.is_err();
let unlocked = manager.admin_start_shutdown.try_lock().is_ok()
&& manager.active_heals.try_lock().is_ok()
&& manager.completed_heals.try_lock().is_ok();
drop(pause);
let result = match first {
Ok(result) => Ok(result),
Err(_) => tokio::time::timeout(Duration::from_secs(5), stop).await,
};
assert!(held, "terminal mutation branch must finish before its receipt");
assert!(unlocked, "join may not hold manager locks");
result.expect("released terminal GC joins").expect("terminal GC stop");
assert!(matches!(disk.read_all(RUSTFS_META_BUCKET, &path).await, Err(DiskError::FileNotFound)));
let registration = manager.workers.registration.lock().unwrap();
assert!(
registration.fixed[3]
.as_ref()
.expect("actual terminal GC receipt")
.peek()
.is_some()
);
}
@@ -240,7 +240,7 @@ impl HealManager {
}
for set_disk_id in &set_disk_ids {
match self.storage.replacement_intent_disks(set_disk_id).await {
match self.storage.replacement_intent_disks(set_disk_id, None).await {
Ok(disks) => recovery_disks.extend(disks),
Err(error) => {
conflicted_replacement_sets.insert(set_disk_id.clone());
+146 -28
View File
@@ -995,6 +995,33 @@ struct MrfRuntime {
}
impl MrfRuntime {
fn admit_best_effort(&mut self, intent: MrfIntent) {
if intent.kind.is_durable() {
if self.admit_partial_write(intent.clone(), None).is_err() {
rustfs_common::mrf_channel::release_mrf_intent(&intent);
counter!("rustfs_heal_mrf_dropped_total", "reason" => "queue_overflow").increment(1);
}
return;
}
if !self.queue.pending_keys.contains(&queue_key(&intent))
&& (self.queue.depth() + self.partial_writes.depth() >= self.queue.capacity
|| self.queue.bytes() + self.partial_writes.bytes() + intent.estimated_bytes() > self.queue.byte_budget)
{
rustfs_common::mrf_channel::release_mrf_intent(&intent);
counter!("rustfs_heal_mrf_dropped_total", "reason" => "queue_overflow").increment(1);
return;
}
match self.queue.try_push_typed(intent.clone()) {
MrfQueuePushResult::Enqueued => {
self.new_since_flush += 1;
self.dirty = true;
}
MrfQueuePushResult::Coalesced | MrfQueuePushResult::Rejected => {
rustfs_common::mrf_channel::release_mrf_intent(&intent);
}
}
}
fn adopt_replayed_partial_writes(&mut self, intents: Vec<(MrfIntent, Option<ResponsibilityCheckpoint>)>) {
// The decoded checkpoint bounds replay; live admission subsequently
// shares these limits with the ordinary pending queue.
@@ -1550,6 +1577,20 @@ fn map_mrf_lifecycle_state_error(error: &'static str) -> MrfLifecycleControlErro
/// ordinary MRF kill-switch.
/// Public for integration tests that drive the real consumer loop.
pub fn spawn_mrf_consumer(manager: Arc<HealManager>) {
let _ = spawn_mrf_consumer_with_stop(manager, None);
}
pub(crate) fn spawn_owned_mrf_consumer(
manager: Arc<HealManager>,
stop: tokio_util::sync::CancellationToken,
) -> Option<tokio::task::JoinHandle<()>> {
spawn_mrf_consumer_with_stop(manager, Some(stop))
}
fn spawn_mrf_consumer_with_stop(
manager: Arc<HealManager>,
stop: Option<tokio_util::sync::CancellationToken>,
) -> Option<tokio::task::JoinHandle<()>> {
let enabled = rustfs_utils::get_env_bool(rustfs_config::ENV_HEAL_MRF_ENABLE, rustfs_config::DEFAULT_HEAL_MRF_ENABLE);
rustfs_common::mrf_channel::set_mrf_delivery_enabled(enabled);
if !enabled {
@@ -1571,13 +1612,21 @@ pub fn spawn_mrf_consumer(manager: Arc<HealManager>) {
error = err,
"MRF channel initialization failed; intents will be dropped at producers"
);
return;
return None;
}
};
tokio::spawn(async move {
run_mrf_consumer(manager, receiver, durable_receiver, lifecycle_receiver).await;
});
let task = match stop {
Some(stop) => tokio::spawn(run_mrf_consumer_with_stop(
manager,
receiver,
durable_receiver,
lifecycle_receiver,
Some(stop),
)),
None => tokio::spawn(run_mrf_consumer(manager, receiver, durable_receiver, lifecycle_receiver)),
};
tracing::info!(target: "rustfs::heal::mrf", "MRF intent consumer started");
Some(task)
}
/// Replay the journal once for integration tests, publishing a durable
@@ -1936,10 +1985,20 @@ async fn replay_into(
/// Replay the journal, then keep draining the channel into the heal manager
/// while persisting the pending snapshot.
async fn run_mrf_consumer(
manager: Arc<HealManager>,
receiver: mpsc::Receiver<MrfIntent>,
durable_receiver: mpsc::Receiver<MrfDurableSubmission>,
lifecycle_receiver: mpsc::Receiver<MrfLifecycleCommand>,
) {
run_mrf_consumer_with_stop(manager, receiver, durable_receiver, lifecycle_receiver, None).await;
}
async fn run_mrf_consumer_with_stop(
manager: Arc<HealManager>,
mut receiver: mpsc::Receiver<MrfIntent>,
mut durable_receiver: mpsc::Receiver<MrfDurableSubmission>,
mut lifecycle_receiver: mpsc::Receiver<MrfLifecycleCommand>,
stop: Option<tokio_util::sync::CancellationToken>,
) {
let config = MrfConsumerConfig::default();
let mut runtime = MrfRuntime {
@@ -1991,9 +2050,25 @@ async fn run_mrf_consumer(
let mut durable_batch = Vec::with_capacity(runtime.config.replay_batch);
loop {
if stop.as_ref().is_some_and(tokio_util::sync::CancellationToken::is_cancelled) {
finalize_mrf_consumer(&mut runtime, &mut receiver, &mut durable_receiver, &mut lifecycle_receiver).await;
return;
}
tokio::select! {
_ = async {
match &stop {
Some(stop) => stop.cancelled().await,
None => std::future::pending::<()>().await,
}
} => {
finalize_mrf_consumer(&mut runtime, &mut receiver, &mut durable_receiver, &mut lifecycle_receiver).await;
return;
}
received = durable_receiver.recv_many(&mut durable_batch, runtime.config.replay_batch) => {
if received == 0 {
if stop.is_some() {
finalize_mrf_consumer(&mut runtime, &mut receiver, &mut durable_receiver, &mut lifecycle_receiver).await;
}
return;
}
let mut responses = Vec::with_capacity(received);
@@ -2013,6 +2088,9 @@ async fn run_mrf_consumer(
}
command = lifecycle_receiver.recv() => {
let Some(command) = command else {
if stop.is_some() {
finalize_mrf_consumer(&mut runtime, &mut receiver, &mut durable_receiver, &mut lifecycle_receiver).await;
}
return;
};
match command {
@@ -2056,6 +2134,10 @@ async fn run_mrf_consumer(
}
received = receiver.recv_many(&mut batch, runtime.config.replay_batch) => {
if received == 0 {
if stop.is_some() {
finalize_mrf_consumer(&mut runtime, &mut receiver, &mut durable_receiver, &mut lifecycle_receiver).await;
return;
}
// Channel closed: flush once more unless the snapshot is
// provably current AND idle (a dirty or pending state
// gets one last persist attempt, matching the shutdown
@@ -2070,30 +2152,7 @@ async fn run_mrf_consumer(
return;
}
for intent in batch.drain(..) {
if intent.kind.is_durable() {
if runtime.admit_partial_write(intent.clone(), None).is_err() {
rustfs_common::mrf_channel::release_mrf_intent(&intent);
counter!("rustfs_heal_mrf_dropped_total", "reason" => "queue_overflow").increment(1);
}
continue;
}
if !runtime.queue.pending_keys.contains(&queue_key(&intent))
&& (runtime.queue.depth() + runtime.partial_writes.depth() >= runtime.queue.capacity
|| runtime.queue.bytes() + runtime.partial_writes.bytes() + intent.estimated_bytes() > runtime.queue.byte_budget)
{
rustfs_common::mrf_channel::release_mrf_intent(&intent);
counter!("rustfs_heal_mrf_dropped_total", "reason" => "queue_overflow").increment(1);
continue;
}
match runtime.queue.try_push_typed(intent.clone()) {
MrfQueuePushResult::Enqueued => {
runtime.new_since_flush += 1;
runtime.dirty = true;
}
MrfQueuePushResult::Coalesced | MrfQueuePushResult::Rejected => {
rustfs_common::mrf_channel::release_mrf_intent(&intent);
}
}
runtime.admit_best_effort(intent);
}
if runtime.dirty && runtime.partial_writes.depth() > 0 {
runtime.flush().await;
@@ -2139,6 +2198,64 @@ async fn run_mrf_consumer(
}
}
// Closing the receivers fences new ingress while preserving all messages that
// were already accepted by these channels. Existing admission limits and replies
// remain unchanged; a failed checkpoint keeps this runtime alive for retry.
async fn finalize_mrf_consumer(
runtime: &mut MrfRuntime,
receiver: &mut mpsc::Receiver<MrfIntent>,
durable_receiver: &mut mpsc::Receiver<MrfDurableSubmission>,
lifecycle_receiver: &mut mpsc::Receiver<MrfLifecycleCommand>,
) {
receiver.close();
durable_receiver.close();
lifecycle_receiver.close();
while let Some(command) = lifecycle_receiver.recv().await {
match command {
MrfLifecycleCommand::List { response, .. } => {
let _ = response.send(Err(MrfLifecycleControlError::Unavailable));
}
MrfLifecycleCommand::Recheck { response, .. }
| MrfLifecycleCommand::Refresh { response, .. }
| MrfLifecycleCommand::AcceptUnverifiedRisk { response, .. } => {
let _ = response.send(Err(MrfLifecycleControlError::Unavailable));
}
}
}
while let Some(intent) = receiver.recv().await {
runtime.admit_best_effort(intent);
}
let mut responses = Vec::new();
while let Some(submission) = durable_receiver.recv().await {
match runtime.admit_partial_write(submission.intent, submission.source_bucket_incarnation_id) {
Ok(()) => responses.push(submission.response),
Err(error) => {
let _ = submission.response.send(Err(error));
}
}
}
let needs_flush = runtime.dirty || runtime.queue.depth() > 0 || runtime.partial_writes.depth() > 0;
let persisted = !needs_flush || runtime.flush().await;
for response in responses {
let _ = response.send(if persisted {
Ok(())
} else {
Err(MrfDurableAdmissionError::Persistence)
});
}
if !persisted {
let mut retry = tokio::time::interval(runtime.config.flush_interval);
retry.tick().await;
loop {
retry.tick().await;
if runtime.flush().await {
break;
}
}
}
runtime.publish_metrics();
}
/// What the periodic tick should do, as a pure function of the runtime state
/// so the decision table is unit-testable.
enum TickAction {
@@ -2185,6 +2302,7 @@ mod tests {
mod empty_events_tests;
mod lifecycle_replay_tests;
mod lifecycle_v2_reader_tests;
mod owned_shutdown_tests;
#[tokio::test(start_paused = true)]
async fn journal_disk_io_bounds_fanout_and_collects_every_result() {
@@ -0,0 +1,256 @@
// Copyright 2026 RustFS Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use super::*;
fn empty_runtime(capacity: usize) -> MrfRuntime {
let config = MrfConsumerConfig::default();
MrfRuntime {
queue: MrfQueue::new(capacity, config.journal_max_bytes),
partial_writes: PartialWrites::default(),
config,
checkpoint_owner: Uuid::new_v4(),
next_checkpoint_sequence: 1,
lifecycle_format_version: MRF_LIFECYCLE_FORMAT_VERSION,
new_since_flush: 0,
dirty: false,
journal_on_disk: false,
retain_replay_journal: false,
replay_readable: true,
durable_replay_anchors: Vec::new(),
replay_cleanup: None,
runtime_checkpoint: None,
backoff_until: None,
}
}
fn partial(object: &str) -> MrfIntent {
let mut item = intent("mrf-owned-shutdown", object, 0);
item.kind = MrfKind::PartialWrite;
item.version_id = None;
item
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn owned_mrf_shutdown_drains_and_persists_before_ack_and_join() {
let _env = rustfs_test_utils::TestECStoreEnv::builder()
.prefix("mrf_owned_shutdown")
.build()
.await;
let mut runtime = empty_runtime(8);
let (incoming, mut receiver) = mpsc::channel(4);
let (durable, mut durable_receiver) = mpsc::channel(4);
let (lifecycle, mut lifecycle_receiver) = mpsc::channel(4);
incoming
.send(intent("mrf-owned-shutdown", "ordinary", 0))
.await
.expect("queued ordinary intent");
let (reply, ack) = tokio::sync::oneshot::channel();
durable
.send(MrfDurableSubmission {
intent: partial("durable"),
source_bucket_incarnation_id: Some(Uuid::new_v4()),
response: reply,
})
.await
.expect("queued durable submission");
let (response, listed) = tokio::sync::oneshot::channel();
lifecycle
.send(MrfLifecycleCommand::List {
after: None,
limit: 1,
response,
})
.await
.expect("queued lifecycle query");
tokio::time::timeout(
StdDuration::from_secs(10),
finalize_mrf_consumer(&mut runtime, &mut receiver, &mut durable_receiver, &mut lifecycle_receiver),
)
.await
.expect("real final checkpoint completes with live sender handles");
assert!(ack.await.expect("explicit durable acknowledgement").is_ok());
assert!(matches!(
listed.await.expect("explicit query rejection"),
Err(MrfLifecycleControlError::Unavailable)
));
assert!(incoming.is_closed() && durable.is_closed() && lifecycle.is_closed());
assert!(!runtime.dirty);
let saved = snapshot::inspect_local_committed_snapshot(runtime.config.journal_max_bytes)
.await
.expect("inspect committed shutdown snapshot")
.expect("shutdown successor exists");
let (items, truncated) = decode_journal(saved.payload());
assert_eq!(truncated, 0);
assert_eq!(items.len(), 2, "both ingress kinds have durable replay storage before successful receipt");
}
// The blocked paths belong to this fresh test environment. Restore them even
// when an assertion fails so its ordinary disk cleanup can finish.
struct BlockedMetadata(Vec<(PathBuf, PathBuf)>);
impl Drop for BlockedMetadata {
fn drop(&mut self) {
for (original, held) in self.0.drain(..) {
let _ = fs::remove_file(&original);
let _ = fs::rename(held, original);
}
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn owned_mrf_failed_flush_retains_runtime_until_disk_recovers() {
let env = rustfs_test_utils::TestECStoreEnv::builder()
.prefix("mrf_owned_retry")
.build()
.await;
let mut runtime = empty_runtime(8);
runtime
.admit_partial_write(partial("previously-committed"), Some(Uuid::new_v4()))
.expect("accept resident responsibility");
assert!(runtime.flush().await, "establish the prior committed responsibility before failure");
let mut blocked = BlockedMetadata(Vec::new());
for root in &env.disk_paths {
let original = root.join(RUSTFS_META_BUCKET);
let held = root.join("mrf-owned-held-metadata");
assert!(!held.exists());
fs::rename(&original, &held).expect("hold only this fixture metadata volume");
blocked.0.push((original.clone(), held));
fs::write(&original, b"block directory traversal").expect("force actual disk writes to fail");
}
let (incoming, mut receiver) = mpsc::channel(4);
let (durable, mut durable_receiver) = mpsc::channel(4);
let (_lifecycle, mut lifecycle_receiver) = mpsc::channel(4);
let (reply, ack) = tokio::sync::oneshot::channel();
durable
.send(MrfDurableSubmission {
intent: partial("retained-after-failed-flush"),
source_bucket_incarnation_id: Some(Uuid::new_v4()),
response: reply,
})
.await
.expect("queue durable responsibility");
let mut shutdown = tokio::spawn(async move {
finalize_mrf_consumer(&mut runtime, &mut receiver, &mut durable_receiver, &mut lifecycle_receiver).await;
runtime
});
assert!(matches!(
tokio::time::timeout(StdDuration::from_secs(10), ack)
.await
.expect("failed checkpoint explicitly responds")
.expect("durable reply is not dropped"),
Err(MrfDurableAdmissionError::Persistence)
));
assert!(
tokio::time::timeout(StdDuration::from_millis(100), &mut shutdown)
.await
.is_err(),
"failed flush is not an exited/joined consumer"
);
assert!(incoming.is_closed() && durable.is_closed());
drop(blocked);
let runtime = tokio::time::timeout(StdDuration::from_secs(10), shutdown)
.await
.expect("same live worker retries after volume recovery")
.expect("real worker joins");
assert!(!runtime.dirty);
assert_eq!(
runtime.partial_writes.depth(),
2,
"prior committed and newly unacknowledged responsibilities both stay owned"
);
let saved = snapshot::inspect_local_committed_snapshot(runtime.config.journal_max_bytes)
.await
.expect("inspect recovered successor")
.expect("recovered successor exists");
assert_eq!(decode_journal(saved.payload()).0.len(), 2);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn owned_mrf_full_admission_replies_full_without_false_ack() {
let _env = rustfs_test_utils::TestECStoreEnv::builder()
.prefix("mrf_owned_full")
.build()
.await;
let mut runtime = empty_runtime(1);
runtime.admit_best_effort(intent("mrf-owned-shutdown", "resident", 0));
let (_incoming, mut receiver) = mpsc::channel(4);
let (durable, mut durable_receiver) = mpsc::channel(4);
let (_lifecycle, mut lifecycle_receiver) = mpsc::channel(4);
let (reply, ack) = tokio::sync::oneshot::channel();
durable
.send(MrfDurableSubmission {
intent: partial("beyond-resident-limit"),
source_bucket_incarnation_id: Some(Uuid::new_v4()),
response: reply,
})
.await
.expect("queue unacknowledged submission");
tokio::time::timeout(
StdDuration::from_secs(10),
finalize_mrf_consumer(&mut runtime, &mut receiver, &mut durable_receiver, &mut lifecycle_receiver),
)
.await
.expect("resident checkpoint completes");
assert!(matches!(
ack.await.expect("Full explicitly returned"),
Err(MrfDurableAdmissionError::Full)
));
assert_eq!(runtime.queue.depth(), 1);
assert_eq!(runtime.partial_writes.depth(), 0);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn owned_mrf_real_consumer_joins_with_live_static_senders() {
const CHILD: &str = "RUSTFS_MRF_OWNED_CONSUMER_TEST_CHILD";
const TEST: &str = "heal::mrf_queue::tests::owned_shutdown_tests::owned_mrf_real_consumer_joins_with_live_static_senders";
if std::env::var_os(CHILD).is_none() {
// Process-global senders cannot be reinitialized by another unit test.
let output = std::process::Command::new(std::env::current_exe().expect("test executable"))
.args(["--exact", TEST, "--nocapture", "--test-threads=1"])
.env(CHILD, "1")
.output()
.expect("isolated MRF consumer test");
let stdout = String::from_utf8_lossy(&output.stdout);
assert!(output.status.success(), "{stdout}{}", String::from_utf8_lossy(&output.stderr));
assert!(stdout.contains("test result: ok. 1 passed; 0 failed; 0 ignored;"));
assert!(stdout.contains("MRF_OWNED_CONSUMER_CHILD_OK"));
return;
}
let env = rustfs_test_utils::TestECStoreEnv::builder()
.prefix("mrf_owned_loop")
.build()
.await;
let manager = Arc::new(HealManager::new_without_root_recovery_for_test(
Arc::new(ECStoreHealStorage::new(env.ecstore.clone())),
None,
));
manager.stop().await.expect("original manager handoff before MRF stop");
let stop = tokio_util::sync::CancellationToken::new();
let task = spawn_owned_mrf_consumer(manager, stop.clone()).expect("fresh process owns actual MRF startup");
tokio::time::timeout(StdDuration::from_secs(10), list_legacy_responsibilities(None, 1))
.await
.expect("actual consumer must process a command before the stop probe")
.expect("empty lifecycle query succeeds");
assert!(!task.is_finished(), "the ready consumer is still running with static senders retained");
stop.cancel();
tokio::time::timeout(StdDuration::from_secs(10), task)
.await
.expect("explicit stop must terminate despite process-global sender lifetime")
.expect("actual consumer joined");
println!("MRF_OWNED_CONSUMER_CHILD_OK");
}
+17 -1
View File
@@ -140,7 +140,23 @@ impl ReplacementExecution {
/// A dropped page waiter must not release the lease of a mutation that is
/// still executing. The owned worker keeps the lease through its I/O.
pub(crate) async fn run<T: Send + 'static>(self: Arc<Self>, work: impl Future<Output = T> + Send + 'static) -> Result<T> {
tokio::spawn(async move {
self.run_with_owner(None, work).await
}
pub(super) async fn run_owned<T: Send + 'static>(
self: Arc<Self>,
parent: &super::manager::HealWork,
work: impl Future<Output = T> + Send + 'static,
) -> Result<T> {
self.run_with_owner(Some(parent), work).await
}
async fn run_with_owner<T: Send + 'static>(
self: Arc<Self>,
parent: Option<&super::manager::HealWork>,
work: impl Future<Output = T> + Send + 'static,
) -> Result<T> {
super::manager::spawn_heal_work(parent, move |_parent| async move {
let _execution = self;
work.await
})
+732 -62
View File
@@ -17,17 +17,22 @@ use async_trait::async_trait;
use base64_simd::URL_SAFE_NO_PAD;
use rustfs_common::mrf_channel::MrfDeleteMarkerPurge;
use rustfs_concurrency::adaptive::NodeAdmission;
use rustfs_concurrency::adaptive::permit_set::NodeResourcePermit;
use rustfs_heal_contracts::heal_channel::{DriveState, HealOpts, HealScanMode};
use rustfs_madmin::heal_commands::HealResultItem;
use serde::{Deserialize, Serialize};
use std::{collections::HashSet, sync::Arc};
use std::{collections::HashSet, future::Future, sync::Arc, time::Duration};
use tokio_util::sync::CancellationToken;
use tracing::{debug, error, warn};
use uuid::Uuid;
use super::outcome::{HealObjectDisposition, HealObjectIdentity, HealObjectKind, HealObjectReceipt};
use super::progress::stable_generation;
pub use super::replacement_execution::ReplacementExecution;
use super::storage_api::owner::{EcstoreHealLifecycleExpiryContext, ecstore_load_admin_data_usage_from_backend_cached};
use super::storage_api::owner::{
EcstoreHealLifecycleExpiryContext, PhysicalResourcePlan, ecstore_load_admin_data_usage_from_backend_cached,
ecstore_with_background_resource_set,
};
use super::storage_api::storage::{
BucketInfo, BucketOperations, DiskSetSelector, EcstoreHealObjectStorageResult, HealOperations as _, ListOperations as _,
ObjectIO as _, ObjectOperations as _, StorageAdminApi,
@@ -423,6 +428,11 @@ pub struct HealListItem {
/// Heal storage layer interface
#[async_trait]
pub trait HealStorageAPI: Send + Sync {
/// Existing process scheduling capability; snapshots alone do not admit I/O.
fn background_node_admission(&self) -> Option<Arc<NodeAdmission>> {
None
}
/// Delete one exact marker from every member of its original erasure set.
async fn purge_delete_marker(
&self,
@@ -566,6 +576,22 @@ pub trait HealStorageAPI: Send + Sync {
opts: &HealOpts,
) -> Result<(HealResultItem, Option<Error>)>;
/// Preserve an issued replacement mutation while yielding its task's node
/// scheduling unit on logical cancellation. New objects remain fenced.
async fn heal_object_with_task_cancel(
&self,
bucket: &str,
object: &str,
version_id: Option<&str>,
opts: &HealOpts,
task_cancel: &CancellationToken,
) -> Result<(HealResultItem, Option<Error>)> {
if task_cancel.is_cancelled() {
return Err(Error::TaskCancelled);
}
self.heal_object(bucket, object, version_id, opts).await
}
async fn heal_object_with_receipt(
&self,
bucket: &str,
@@ -702,7 +728,13 @@ pub trait HealStorageAPI: Send + Sync {
/// Return every formatted disk that may durably host or expose a
/// replacement intent. Implementations should order local disks first so
/// discovery remains deterministic and preserves the local preference.
async fn replacement_intent_disks(&self, _set_disk_id: &str) -> Result<Vec<DiskStore>> {
/// A supplied token uses attached admission for known-local probes; None
/// keeps startup discovery on its original direct path.
async fn replacement_intent_disks(
&self,
_set_disk_id: &str,
_cancellation: Option<&CancellationToken>,
) -> Result<Vec<DiskStore>> {
Ok(Vec::new())
}
@@ -749,32 +781,152 @@ pub trait HealStorageAPI: Send + Sync {
}
}
// One bounded local metadata operation owns only its selected disk. Permanent
// tickers retain the controller capability, never a permit between operations.
pub(super) async fn run_local_metadata<T>(
disk: &super::DiskStore,
admission: Option<&Arc<rustfs_concurrency::adaptive::NodeAdmission>>,
cancellation: &CancellationToken,
operation: impl Future<Output = T>,
) -> Result<T> {
let Some(admission) = admission else {
return Ok(operation.await);
};
let plan = disk
.physical_admission()
.and_then(|binding| binding.plan())
.ok_or_else(|| Error::transient_skip("local metadata physical resource discovery unavailable"))?;
let permit = tokio::select! {
biased;
_ = cancellation.cancelled() => return Err(Error::transient_skip("local metadata admission cancelled")),
result = tokio::time::timeout(Duration::from_millis(250), plan.acquire(admission)) => {
result.ok().flatten()
.ok_or_else(|| Error::transient_skip("local metadata node or physical resource admission deferred"))?
}
};
if !plan.is_current() {
return Err(Error::transient_skip("local metadata physical resource discovery expired"));
}
use super::storage_api::owner::{
ecstore_with_background_disk_io, ecstore_with_background_resource_set, ecstore_with_rpc_background_request,
};
let mut operation = std::pin::pin!(ecstore_with_rpc_background_request(
false,
ecstore_with_background_disk_io(
cancellation.clone(),
ecstore_with_background_resource_set(permit.resources().clone(), operation),
),
));
let mut started = false;
let selected = tokio::select! {
biased;
_ = cancellation.cancelled() => None,
result = std::future::poll_fn(|cx| {
started = true;
std::future::Future::poll(operation.as_mut(), cx)
}) => Some(result),
};
// Cancellation yields only the node unit; the same started operation owns
// its physical/native resources until it finishes. Never credit diagnostics.
drop(permit);
match selected {
Some(result) => Ok(result),
None if started => Ok(operation.await),
None => Err(Error::transient_skip("local metadata cancelled before first poll")),
}
}
// Scheduling precedes the operation's first poll and therefore its namespace
// locks. No-lock callers may already own a namespace lock: never park them.
async fn run_background_object<T>(
admission: Option<&Arc<NodeAdmission>>,
physical_plan: impl FnOnce() -> Result<Option<PhysicalResourcePlan>>,
opts: &HealOpts,
completed: impl FnOnce(&T) -> bool,
operation: impl std::future::Future<Output = Result<T>>,
) -> Result<T> {
run_background_object_inner(admission, physical_plan, opts, None, completed, operation).await
}
async fn run_background_object_inner<T>(
admission: Option<&Arc<NodeAdmission>>,
physical_plan: impl FnOnce() -> Result<Option<PhysicalResourcePlan>>,
opts: &HealOpts,
task_cancel: Option<&CancellationToken>,
completed: impl FnOnce(&T) -> bool,
operation: impl std::future::Future<Output = Result<T>>,
) -> Result<T> {
if task_cancel.is_some_and(CancellationToken::is_cancelled) {
return Err(Error::TaskCancelled);
}
let Some(admission) = admission else {
return operation.await;
};
let permit = if opts.no_lock {
admission.try_acquire()
} else {
tokio::time::timeout(std::time::Duration::from_millis(250), admission.acquire())
.await
.ok()
let plan = physical_plan()?.ok_or_else(|| Error::transient_skip("physical resource discovery unavailable"))?;
let acquire = async {
if opts.no_lock {
plan.try_acquire(admission)
} else {
tokio::time::timeout(std::time::Duration::from_millis(250), plan.acquire(admission))
.await
.ok()
.flatten()
}
};
let permit = match task_cancel {
Some(cancel) => tokio::select! {
biased;
_ = cancel.cancelled() => return Err(Error::TaskCancelled),
permit = acquire => permit,
},
None => acquire.await,
}
.ok_or_else(|| Error::transient_skip("node resource admission deferred"))?;
let result = operation.await;
.ok_or_else(|| Error::transient_skip("node or physical resource admission deferred"))?;
if let Some(cancel) = task_cancel {
return await_admitted_object(permit, opts, cancel, completed, operation).await;
}
let result = ecstore_with_background_resource_set(permit.resources().clone(), operation).await;
if !opts.dry_run && result.as_ref().is_ok_and(completed) {
permit.complete();
}
result
}
// The permit accounts for one logical object, not each shard/native dispatch.
// Physical units live in the pinned operation scope after scheduling is yielded.
async fn await_admitted_object<T>(
permit: NodeResourcePermit,
opts: &HealOpts,
task_cancel: &CancellationToken,
completed: impl FnOnce(&T) -> bool,
operation: impl std::future::Future<Output = Result<T>>,
) -> Result<T> {
let mut operation = std::pin::pin!(ecstore_with_background_resource_set(permit.resources().clone(), operation));
let mut started = false;
let selected = {
tokio::select! {
biased;
_ = task_cancel.cancelled() => None,
result = std::future::poll_fn(|cx| {
started = true;
std::future::Future::poll(operation.as_mut(), cx)
}) => Some(result),
}
};
match selected {
None => {
drop(permit);
if started { operation.await } else { Err(Error::TaskCancelled) }
}
Some(result) => {
if !task_cancel.is_cancelled() && !opts.dry_run && result.as_ref().is_ok_and(completed) {
permit.complete();
}
result
}
}
}
/// ECStore Heal storage layer implementation
pub struct ECStoreHealStorage {
ecstore: Arc<ECStore>,
@@ -796,6 +948,50 @@ impl ECStoreHealStorage {
self
}
async fn heal_object_admitted(
&self,
bucket: &str,
object: &str,
version_id: Option<&str>,
opts: &HealOpts,
task_cancel: Option<&CancellationToken>,
) -> Result<(HealResultItem, Option<Error>)> {
run_background_object_inner(
self.admission.as_ref(),
|| self.ecstore.physical_heal_plan(object, opts).map_err(Error::Storage),
opts,
task_cancel,
|result: &(HealResultItem, Option<Error>)| result.1.is_none(),
async {
debug!(
target: "rustfs::heal::storage",
event = EVENT_HEAL_STORAGE_REPAIR_OP,
component = LOG_COMPONENT_HEAL,
subsystem = LOG_SUBSYSTEM_STORAGE,
operation = "heal_object",
bucket,
object,
version_id = ?version_id,
scan_mode = %opts.scan_mode.as_str(),
dry_run = opts.dry_run,
state = "started",
"Heal storage repair started"
);
let version_id_str = version_id.unwrap_or("");
Self::adapt_object_heal_result(
bucket,
object,
version_id,
opts,
self.ecstore.heal_object(bucket, object, version_id_str, opts).await,
)
},
)
.await
}
fn adapt_object_heal_result(
bucket: &str,
object: &str,
@@ -1016,6 +1212,10 @@ fn is_transient_object_exists_error(err: &StorageError) -> bool {
#[async_trait]
impl HealStorageAPI for ECStoreHealStorage {
fn background_node_admission(&self) -> Option<Arc<NodeAdmission>> {
self.admission.clone()
}
async fn purge_delete_marker(
&self,
bucket: &str,
@@ -1079,6 +1279,7 @@ impl HealStorageAPI for ECStoreHealStorage {
) -> Result<HealStorageObjectResult> {
run_background_object(
self.admission.as_ref(),
|| self.ecstore.physical_heal_plan(object, opts).map_err(Error::Storage),
opts,
|result: &HealStorageObjectResult| result.error.is_none(),
async {
@@ -1105,6 +1306,7 @@ impl HealStorageAPI for ECStoreHealStorage {
) -> Result<HealStorageObjectResult> {
run_background_object(
self.admission.as_ref(),
|| self.ecstore.physical_heal_plan(object, opts).map_err(Error::Storage),
opts,
|result: &HealStorageObjectResult| result.error.is_none(),
async {
@@ -1547,38 +1749,19 @@ impl HealStorageAPI for ECStoreHealStorage {
version_id: Option<&str>,
opts: &HealOpts,
) -> Result<(HealResultItem, Option<Error>)> {
run_background_object(
self.admission.as_ref(),
opts,
|result: &(HealResultItem, Option<Error>)| result.1.is_none(),
async {
debug!(
target: "rustfs::heal::storage",
event = EVENT_HEAL_STORAGE_REPAIR_OP,
component = LOG_COMPONENT_HEAL,
subsystem = LOG_SUBSYSTEM_STORAGE,
operation = "heal_object",
bucket,
object,
version_id = ?version_id,
scan_mode = %opts.scan_mode.as_str(),
dry_run = opts.dry_run,
state = "started",
"Heal storage repair started"
);
self.heal_object_admitted(bucket, object, version_id, opts, None).await
}
let version_id_str = version_id.unwrap_or("");
Self::adapt_object_heal_result(
bucket,
object,
version_id,
opts,
self.ecstore.heal_object(bucket, object, version_id_str, opts).await,
)
},
)
.await
async fn heal_object_with_task_cancel(
&self,
bucket: &str,
object: &str,
version_id: Option<&str>,
opts: &HealOpts,
task_cancel: &CancellationToken,
) -> Result<(HealResultItem, Option<Error>)> {
self.heal_object_admitted(bucket, object, version_id, opts, Some(task_cancel))
.await
}
async fn heal_object_with_receipt(
@@ -1590,6 +1773,7 @@ impl HealStorageAPI for ECStoreHealStorage {
) -> Result<HealStorageObjectResult> {
run_background_object(
self.admission.as_ref(),
|| self.ecstore.physical_heal_plan(object, opts).map_err(Error::Storage),
opts,
|result: &HealStorageObjectResult| result.error.is_none(),
async {
@@ -2045,15 +2229,35 @@ impl HealStorageAPI for ECStoreHealStorage {
})
}
async fn replacement_intent_disks(&self, set_disk_id: &str) -> Result<Vec<DiskStore>> {
async fn replacement_intent_disks(
&self,
set_disk_id: &str,
cancellation: Option<&CancellationToken>,
) -> Result<Vec<DiskStore>> {
let disks = self.resume_disk_inventory(set_disk_id).await?;
let mut formatted = Vec::with_capacity(disks.len());
let admitted_scan = self.admission.as_ref().zip(cancellation);
for prefer_local in [true, false] {
for disk in &disks {
if disk.endpoint().is_local != prefer_local {
continue;
}
match disk.get_disk_id().await {
if admitted_scan.is_some_and(|(_, cancellation)| cancellation.is_cancelled()) {
return Err(Error::transient_skip("replacement inventory admission cancelled"));
}
let disk_id = if prefer_local && let Some((admission, cancellation)) = admitted_scan {
let result = run_local_metadata(disk, Some(admission), cancellation, disk.get_disk_id()).await?;
if cancellation.is_cancelled() {
return Err(Error::transient_skip("replacement inventory admission cancelled"));
}
if matches!(&result, Err(DiskError::Io(error)) if error.kind() == std::io::ErrorKind::WouldBlock) {
return Err(Error::transient_skip("replacement inventory native resources unavailable"));
}
result
} else {
disk.get_disk_id().await
};
match disk_id {
Ok(Some(id)) if !id.is_nil() => formatted.push(disk.clone()),
Ok(_) | Err(DiskError::UnformattedDisk) => {}
Err(error) => {
@@ -2228,6 +2432,79 @@ mod tests {
gate
}
fn physical_plan() -> crate::Result<Option<super::PhysicalResourcePlan>> {
let binding = super::super::storage_api::owner::PhysicalDiskBinding::default();
binding
.set_test_capacity(vec![rustfs_concurrency::adaptive::AdaptiveAdmission::new(NonZeroU16::MIN, 1)])
.expect("synthetic physical capacity");
Ok(Some(binding.plan().expect("live physical binding").resource_plan()))
}
#[tokio::test]
async fn object_admission_releases_node_and_retains_only_a_running_native_disk() {
use super::super::storage_api::owner::{PhysicalDiskBinding, ecstore_spawn_blocking_with_resource_ownership};
let node = admission(true);
let disk = rustfs_concurrency::adaptive::AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let binding = PhysicalDiskBinding::default();
binding.set_test_capacity(vec![disk.clone()]).expect("physical binding");
let plan = binding.plan().expect("physical plan").resource_plan();
let (started_tx, started_rx) = tokio::sync::oneshot::channel();
let (release_tx, release_rx) = std::sync::mpsc::channel();
let handle = tokio::spawn({
let node = node.clone();
async move {
super::run_background_object(Some(&node), || Ok(Some(plan)), &super::HealOpts::default(), |_: &()| true, async {
ecstore_spawn_blocking_with_resource_ownership(move || {
started_tx.send(()).expect("native started");
release_rx.recv().expect("release native");
})
.await
.expect("native job");
Ok(())
})
.await
}
});
started_rx.await.expect("native dispatched");
assert_eq!((node.snapshot().active, disk.snapshot().active), (1, 1));
handle.abort();
assert!(handle.await.expect_err("waiter aborted").is_cancelled());
assert_eq!((node.snapshot().active, disk.snapshot().active), (0, 1));
assert_eq!((node.completed(), disk.completed()), (Some(0), Some(0)));
release_tx.send(()).expect("release native");
tokio::time::timeout(Duration::from_secs(2), async {
while disk.snapshot().active != 0 {
tokio::task::yield_now().await;
}
})
.await
.expect("native releases physical resource");
}
#[tokio::test(start_paused = true)]
async fn object_admission_closed_physical_capacity_defers_before_storage_locks() {
use super::super::storage_api::owner::PhysicalDiskBinding;
let node = admission(true);
let disk = rustfs_concurrency::adaptive::AdaptiveAdmission::new(NonZeroU16::MIN, 0);
let binding = PhysicalDiskBinding::default();
binding.set_test_capacity(vec![disk.clone()]).expect("physical binding");
let plan = binding.plan().expect("physical plan").resource_plan();
for no_lock in [false, true] {
let began = tokio::time::Instant::now();
let opts = super::HealOpts {
no_lock,
..Default::default()
};
let result = super::run_background_object(Some(&node), || Ok(Some(plan.clone())), &opts, |_: &()| true, async {
panic!("closed disk must prevent storage polling");
})
.await;
assert!(matches!(result, Err(super::Error::TransientSkip { .. })));
assert_eq!(began.elapsed(), if no_lock { Duration::ZERO } else { Duration::from_millis(250) });
assert_eq!((node.snapshot().active, disk.snapshot().active), (0, 0));
}
}
#[tokio::test]
async fn object_admission_is_wired_to_all_four_real_storage_entry_points() {
use super::{ECStoreHealStorage, Error, HealOpts, HealStorageAPI};
@@ -2236,7 +2513,7 @@ mod tests {
.base_dir(root.path())
.build()
.await;
let gate = admission(false);
let gate = admission(true);
let storage = ECStoreHealStorage::new(env.ecstore.clone()).with_admission(Some(gate.clone()));
let opts = HealOpts {
no_lock: true,
@@ -2261,7 +2538,7 @@ mod tests {
for result in results {
assert!(
matches!(result, Err(Error::TransientSkip { .. })),
"closed admission precedes even bucket validation: {result:?}"
"unknown physical ownership precedes even bucket validation: {result:?}"
);
}
assert_eq!(gate.completed(), Some(0));
@@ -2285,11 +2562,12 @@ mod tests {
let gate = admission(false);
let polled = std::cell::Cell::new(false);
let began = tokio::time::Instant::now();
let result = super::run_background_object(Some(&gate), &super::HealOpts::default(), |_: &()| true, async {
polled.set(true);
Ok(())
})
.await;
let result =
super::run_background_object(Some(&gate), physical_plan, &super::HealOpts::default(), |_: &()| true, async {
polled.set(true);
Ok(())
})
.await;
assert!(matches!(result, Err(super::Error::TransientSkip { .. })));
assert!(!polled.get(), "no locks or storage mutation before admission");
assert_eq!(began.elapsed(), Duration::from_millis(250));
@@ -2304,9 +2582,10 @@ mod tests {
..Default::default()
};
let began = tokio::time::Instant::now();
let result =
super::run_background_object(Some(&gate), &opts, |_: &()| true, async { panic!("storage must remain unpolled") })
.await;
let result = super::run_background_object(Some(&gate), physical_plan, &opts, |_: &()| true, async {
panic!("storage must remain unpolled")
})
.await;
assert!(matches!(result, Err(super::Error::TransientSkip { .. })));
assert_eq!(began.elapsed(), Duration::ZERO);
}
@@ -2316,15 +2595,17 @@ mod tests {
use std::task::{Context, Poll, Waker};
let gate = admission(true);
let opts = super::HealOpts::default();
let future = super::run_background_object(Some(&gate), &opts, |_: &()| true, std::future::pending());
let future = super::run_background_object(Some(&gate), physical_plan, &opts, |_: &()| true, std::future::pending());
let mut future = Box::pin(future);
assert!(matches!(future.as_mut().poll(&mut Context::from_waker(Waker::noop())), Poll::Pending));
assert_eq!(gate.snapshot().active, 1);
drop(future);
assert_eq!(gate.snapshot().active, 0);
assert_eq!(gate.completed(), Some(0));
let result =
super::run_background_object(Some(&gate), &opts, |_: &()| true, async { Err(super::Error::TaskCancelled) }).await;
let result = super::run_background_object(Some(&gate), physical_plan, &opts, |_: &()| true, async {
Err(super::Error::TaskCancelled)
})
.await;
assert!(matches!(result, Err(super::Error::TaskCancelled)));
assert_eq!(gate.completed(), Some(0));
assert_eq!(gate.snapshot().active, 0);
@@ -2337,7 +2618,7 @@ mod tests {
read_repair: true,
..Default::default()
};
let result = super::run_background_object(Some(&gate), &opts, |_: &()| true, async {
let result = super::run_background_object(Some(&gate), physical_plan, &opts, |_: &()| true, async {
panic!("queued read repair must not bypass the closed node gate");
})
.await;
@@ -2355,7 +2636,7 @@ mod tests {
..Default::default()
};
assert_eq!(
super::run_background_object(Some(&gate), &opts, |ok: &bool| *ok, async { Ok(successful) })
super::run_background_object(Some(&gate), physical_plan, &opts, |ok: &bool| *ok, async { Ok(successful) })
.await
.expect("operation"),
successful
@@ -2365,9 +2646,15 @@ mod tests {
}
gate.update(&mut FeedbackController::new(), None);
assert_eq!(gate.completed(), Some(1));
super::run_background_object(None, &super::HealOpts::default(), |_: &()| true, async { Ok(()) })
.await
.expect("default path stays available");
super::run_background_object(
None,
|| panic!("disabled admission must not read discovery"),
&super::HealOpts::default(),
|_: &()| true,
async { Ok(()) },
)
.await
.expect("default path stays available");
}
#[test]
@@ -2686,4 +2973,387 @@ mod tests {
assert!(matches!(error, crate::Error::ReplacementGenerationConflict { .. }));
assert!(replacement_resume_conflict_for_other_tasks("task", &HashSet::new(), "pool_0_set_0").is_none());
}
// Owner/resource composition only: this does not traverse an erasure page or namespace lock.
#[test]
fn replacement_outside_waiter_drop_keeps_cancelled_native_mutation_and_manager_drain_owned() {
const FLAG: &str = "RUSTFS_TEST_REPLACEMENT_OUTSIDE_WAITER_CHILD";
const NAME: &str =
"heal::storage::tests::replacement_outside_waiter_drop_keeps_cancelled_native_mutation_and_manager_drain_owned";
if std::env::var(FLAG).ok().as_deref() != Some("1") {
let output = std::process::Command::new(std::env::current_exe().unwrap())
.args([NAME, "--exact", "--nocapture", "--test-threads=1"])
.env(FLAG, "1")
.output()
.unwrap();
let stdout = String::from_utf8_lossy(&output.stdout);
assert!(
output.status.success() && stdout.contains("1 passed; 0 failed; 0 ignored;"),
"child failed: {stdout}\n{}",
String::from_utf8_lossy(&output.stderr)
);
println!("OUTSIDE_WAITER_COMPOSITION_CHILD_OK");
return;
}
use super::super::storage_api::owner::{PhysicalDiskBinding, ecstore_spawn_blocking_with_resource_ownership};
use crate::heal::storage_api::{
EcstoreBackgroundEncodeBuffer as BackgroundEncodeBuffer, EcstoreBackgroundMemoryBudget as BackgroundMemoryBudget,
};
use std::io::Write as _;
let runtime = tokio::runtime::Builder::new_multi_thread()
.worker_threads(2)
.enable_all()
.build()
.unwrap();
runtime.block_on(async {
let root = tempfile::tempdir().unwrap();
let env = rustfs_test_utils::TestECStoreEnv::builder()
.base_dir(root.path().join("store"))
.build()
.await;
let manager = Arc::new(crate::HealManager::new_without_root_recovery_for_test(
Arc::new(super::ECStoreHealStorage::new(env.ecstore.clone())),
None,
));
let parent = manager.test_work_parent();
let execution = super::ReplacementExecution::for_test(Vec::new(), Vec::new());
let ownership = Arc::downgrade(&execution);
let node = admission(true);
let disk = rustfs_concurrency::adaptive::AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let binding = PhysicalDiskBinding::default();
binding.set_test_capacity(vec![disk.clone()]).unwrap();
let plan = binding.plan().unwrap().resource_plan();
let memory = BackgroundMemoryBudget::new();
memory.set_sampled_limit(128 * 1024, Duration::from_secs(60));
let mut encoded = BackgroundEncodeBuffer::new(memory.clone(), 64 * 1024);
encoded.write_all(b"same outside waiter mutation").unwrap();
let bytes = encoded.into_bytes();
let charged = memory.snapshot().used;
let path = root.path().join("mutated-object");
let native_path = path.clone();
let task_cancel = tokio_util::sync::CancellationToken::new();
let cancel = task_cancel.clone();
let (entered_tx, entered_rx) = tokio::sync::oneshot::channel();
let (release_tx, release_rx) = std::sync::mpsc::channel();
let (continued_tx, continued_rx) = tokio::sync::oneshot::channel();
let (finish_tx, finish_rx) = tokio::sync::oneshot::channel();
let (result_tx, result_rx) = tokio::sync::oneshot::channel();
let worker_node = node.clone();
let waiter = tokio::spawn(async move {
execution
.run_owned(&parent, async move {
let result = super::run_background_object_inner(
Some(&worker_node),
|| Ok(Some(plan)),
&super::HealOpts::default(),
Some(&cancel),
|_: &bool| true,
async {
ecstore_spawn_blocking_with_resource_ownership(move || -> std::io::Result<()> {
let _ = entered_tx.send(());
let _ = release_rx.recv();
std::fs::write(native_path, &bytes)
})
.await
.map_err(|error| super::Error::other(error.to_string()))??;
let _ = continued_tx.send(());
let _ = finish_rx.await;
Ok(true)
},
)
.await;
let _ = result_tx.send(result.as_ref().is_ok_and(|ok| *ok));
result
})
.await
});
let entered = tokio::time::timeout(Duration::from_secs(5), entered_rx)
.await
.is_ok_and(|result| result.is_ok());
let admitted = (node.snapshot().active, disk.snapshot().active);
waiter.abort();
let waiter_cancelled = waiter.await.is_err_and(|error| error.is_cancelled());
// Observer Drop alone is not logical cancellation of the accepted object.
task_cancel.cancel();
let yielded = tokio::time::timeout(Duration::from_secs(2), async {
while node.snapshot().active != 0 {
tokio::task::yield_now().await;
}
})
.await
.is_ok();
let retained = (disk.snapshot().active, memory.snapshot().used, ownership.upgrade().is_some());
let credit = (node.completed(), disk.completed());
let stopping = manager.clone();
let mut stop = tokio::spawn(async move { stopping.stop().await });
let stop_entered = tokio::time::timeout(Duration::from_secs(5), async {
while !manager.stop_handoff_completed() {
tokio::task::yield_now().await;
}
})
.await
.is_ok();
let first_stop = tokio::time::timeout(Duration::from_millis(30), &mut stop).await;
let stop_pending_while_native_held = first_stop.is_err();
// Always release native work and its continuation before asserting or removing roots.
let _ = release_tx.send(());
let continued = tokio::time::timeout(Duration::from_secs(5), continued_rx)
.await
.is_ok_and(|result| result.is_ok());
let stop_pending_through_continuation = !stop.is_finished();
let _ = finish_tx.send(());
let mutation_completed = tokio::time::timeout(Duration::from_secs(5), result_rx)
.await
.is_ok_and(|result| result.is_ok_and(|completed| completed));
let drained = match first_stop {
Ok(result) => Ok(result),
Err(_) => tokio::time::timeout(Duration::from_secs(5), &mut stop).await,
};
if drained.is_err() {
stop.abort();
let _ = stop.await;
}
if !mutation_completed || drained.is_err() {
// Unknown native retirement must preserve its task-owned filesystem root.
let _ = root.keep();
}
assert!(entered && waiter_cancelled && yielded);
assert_eq!(admitted, (1, 1));
assert!(charged > 0);
assert_eq!(retained, (1, charged, true));
assert_eq!(credit, (Some(0), Some(0)));
assert!(stop_entered && stop_pending_while_native_held && continued && stop_pending_through_continuation);
assert!(mutation_completed, "the original owned operation must complete its actual continuation");
drained.unwrap().unwrap().unwrap();
assert!(ownership.upgrade().is_none());
assert_eq!(std::fs::read(path).unwrap(), b"same outside waiter mutation");
assert_eq!((node.snapshot().active, disk.snapshot().active, memory.snapshot().used), (0, 0, 0));
assert_eq!((node.completed(), disk.completed()), (Some(0), Some(0)));
manager.stop().await.unwrap();
});
drop(runtime);
}
#[tokio::test]
async fn logical_cancel_yields_node_without_dropping_native_buffer_or_result_continuation() {
use super::super::storage_api::owner::{PhysicalDiskBinding, ecstore_spawn_blocking_with_resource_ownership};
use crate::heal::storage_api::{
EcstoreBackgroundEncodeBuffer as BackgroundEncodeBuffer, EcstoreBackgroundMemoryBudget as BackgroundMemoryBudget,
};
use std::io::Write as _;
let node = admission(true);
let disk = rustfs_concurrency::adaptive::AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let binding = PhysicalDiskBinding::default();
binding.set_test_capacity(vec![disk.clone()]).unwrap();
let plan = binding.plan().unwrap().resource_plan();
let memory = BackgroundMemoryBudget::new();
memory.set_sampled_limit(128 * 1024, Duration::from_secs(60));
let mut encoded = BackgroundEncodeBuffer::new(memory.clone(), 64 * 1024);
encoded.write_all(b"same accepted object mutation").unwrap();
let bytes = encoded.into_bytes();
let charged = memory.snapshot().used;
let root = tempfile::tempdir().unwrap();
let path = root.path().join("mutated-object");
let native_path = path.clone();
let task_cancel = tokio_util::sync::CancellationToken::new();
let cancel = task_cancel.clone();
let (entered_tx, entered_rx) = tokio::sync::oneshot::channel();
let (release_tx, release_rx) = std::sync::mpsc::channel();
let (continued_tx, continued_rx) = tokio::sync::oneshot::channel();
let (finish_tx, finish_rx) = tokio::sync::oneshot::channel();
let worker_node = node.clone();
let task = tokio::spawn(async move {
super::run_background_object_inner(
Some(&worker_node),
|| Ok(Some(plan)),
&super::HealOpts::default(),
Some(&cancel),
|_: &bool| true,
async {
ecstore_spawn_blocking_with_resource_ownership(move || -> std::io::Result<()> {
let _ = entered_tx.send(());
let _ = release_rx.recv();
std::fs::write(native_path, &bytes)
})
.await
.map_err(|error| super::Error::other(error.to_string()))??;
let _ = continued_tx.send(());
let _ = finish_rx.await;
Ok(true)
},
)
.await
});
entered_rx.await.unwrap();
let admitted = (node.snapshot().active, disk.snapshot().active);
task_cancel.cancel();
let yielded = tokio::time::timeout(Duration::from_secs(2), async {
while node.snapshot().active != 0 {
tokio::task::yield_now().await;
}
})
.await
.is_ok();
let retained = (disk.snapshot().active, memory.snapshot().used, !task.is_finished());
let credit = (node.completed(), disk.completed());
let new_object_polled = std::sync::atomic::AtomicBool::new(false);
let rejected = super::run_background_object_inner(
Some(&node),
physical_plan,
&super::HealOpts::default(),
Some(&task_cancel),
|_: &()| true,
async {
new_object_polled.store(true, std::sync::atomic::Ordering::Release);
Ok(())
},
)
.await;
// Unconditionally release issued native work and its same-operation
// continuation before inspecting the observations, including red probes.
let _ = release_tx.send(());
let continued = tokio::time::timeout(Duration::from_secs(2), continued_rx)
.await
.is_ok_and(|r| r.is_ok());
let still_owned = !task.is_finished();
let _ = finish_tx.send(());
let result = tokio::time::timeout(Duration::from_secs(5), task).await;
assert_eq!(admitted, (1, 1));
assert!(yielded, "logical cancellation must return only the node scheduling unit");
assert!(charged > 0, "the buffer must have an actual reservation");
assert_eq!(retained, (1, charged, true));
assert_eq!(credit, (Some(0), Some(0)));
assert!(matches!(rejected, Err(super::Error::TaskCancelled)));
assert!(!new_object_polled.load(std::sync::atomic::Ordering::Acquire));
assert!(
continued && still_owned,
"the same issued operation cannot return before its continuation"
);
assert!(result.unwrap().unwrap().unwrap());
assert_eq!(std::fs::read(path).unwrap(), b"same accepted object mutation");
assert_eq!((node.snapshot().active, disk.snapshot().active, memory.snapshot().used), (0, 0, 0));
assert_eq!((node.completed(), disk.completed()), (Some(0), Some(0)));
}
#[tokio::test]
async fn logical_cancel_rejects_entry_and_unpolled_admitted_object() {
let node = admission(true);
let task_cancel = tokio_util::sync::CancellationToken::new();
task_cancel.cancel();
let planned = std::sync::atomic::AtomicBool::new(false);
let polled = std::sync::atomic::AtomicBool::new(false);
let result = super::run_background_object_inner(
Some(&node),
|| {
planned.store(true, std::sync::atomic::Ordering::Release);
physical_plan()
},
&super::HealOpts::default(),
Some(&task_cancel),
|_: &()| true,
async {
polled.store(true, std::sync::atomic::Ordering::Release);
Ok(())
},
)
.await;
assert!(matches!(result, Err(super::Error::TaskCancelled)));
assert!(!planned.load(std::sync::atomic::Ordering::Acquire));
let plan = physical_plan().unwrap().unwrap();
let permit = plan.try_acquire(&node).unwrap();
let result = super::await_admitted_object(permit, &super::HealOpts::default(), &task_cancel, |_: &()| true, async {
polled.store(true, std::sync::atomic::Ordering::Release);
Ok(())
})
.await;
assert!(matches!(result, Err(super::Error::TaskCancelled)));
assert!(
!polled.load(std::sync::atomic::Ordering::Acquire),
"do not start an unpaid object in cancellation's branch"
);
assert_eq!(node.snapshot().active, 0);
assert_eq!(node.completed(), Some(0));
}
#[tokio::test]
async fn logical_cancel_interrupts_admission_without_polling_new_object() {
let node = admission(false);
let task_cancel = tokio_util::sync::CancellationToken::new();
let cancel = task_cancel.clone();
let polled = Arc::new(std::sync::atomic::AtomicBool::new(false));
let observed = polled.clone();
let (planned_tx, planned_rx) = tokio::sync::oneshot::channel();
let worker_node = node.clone();
let task = tokio::spawn(async move {
super::run_background_object_inner(
Some(&worker_node),
|| {
let _ = planned_tx.send(());
physical_plan()
},
&super::HealOpts::default(),
Some(&cancel),
|_: &()| true,
async {
observed.store(true, std::sync::atomic::Ordering::Release);
Ok(())
},
)
.await
});
planned_rx.await.unwrap();
task_cancel.cancel();
let result = tokio::time::timeout(Duration::from_secs(2), task).await.unwrap().unwrap();
assert!(matches!(result, Err(super::Error::TaskCancelled)));
assert!(!polled.load(std::sync::atomic::Ordering::Acquire));
assert_eq!(node.snapshot().active, 0);
assert_eq!(node.completed(), Some(0));
}
#[tokio::test]
async fn logical_cancel_entry_preserves_normal_dry_run_and_legacy_credit() {
for (logical, dry_run) in [(true, false), (true, true), (false, false)] {
let node = admission(true);
let task_cancel = tokio_util::sync::CancellationToken::new();
let opts = super::HealOpts {
dry_run,
..Default::default()
};
let result = super::run_background_object_inner(
Some(&node),
physical_plan,
&opts,
logical.then_some(&task_cancel),
|ok: &bool| *ok,
async { Ok(true) },
)
.await;
assert!(result.unwrap());
assert_eq!(node.snapshot().active, 0);
assert_eq!(node.completed(), Some(u64::from(!dry_run)));
}
}
#[tokio::test]
async fn actual_storage_logical_cancel_entry_refuses_an_already_cancelled_object() {
let env = rustfs_test_utils::TestECStoreEnv::builder()
.prefix("task_cancel_entry")
.build()
.await;
let node = admission(true);
let storage = super::ECStoreHealStorage::new(env.ecstore.clone()).with_admission(Some(node.clone()));
let task_cancel = tokio_util::sync::CancellationToken::new();
task_cancel.cancel();
let result = super::HealStorageAPI::heal_object_with_task_cancel(
&storage,
"bucket",
"object",
None,
&super::HealOpts::default(),
&task_cancel,
)
.await;
assert!(matches!(result, Err(super::Error::TaskCancelled)));
assert_eq!((node.snapshot().active, node.completed()), (0, Some(0)));
}
}
+17
View File
@@ -25,7 +25,15 @@ pub(crate) use rustfs_ecstore::api::disk::{
DiskStore as EcstoreDiskStore, HEALING_MARKER_PATH as ECSTORE_HEALING_MARKER_PATH,
RUSTFS_META_BUCKET as ECSTORE_RUSTFS_META_BUCKET, ReplacementExecutionLease as EcstoreReplacementExecutionLease,
};
#[cfg(test)]
pub(crate) use rustfs_ecstore::api::disk::{
BackgroundEncodeBuffer as EcstoreBackgroundEncodeBuffer, BackgroundMemoryBudget as EcstoreBackgroundMemoryBudget,
};
pub(crate) use rustfs_ecstore::api::disk::{DiskOption as EcstoreDiskOption, new_disk as ecstore_new_disk};
#[cfg(all(test, not(windows)))]
pub(crate) use rustfs_ecstore::api::disk::{
LocalPublicationPause as EcstoreLocalPublicationPause, LocalPublicationStage as EcstoreLocalPublicationStage,
};
pub(crate) use rustfs_ecstore::api::error::{Error as EcstoreErrorType, StorageError as EcstoreStorageError};
#[cfg(test)]
pub(crate) use rustfs_ecstore::api::layout::{
@@ -41,6 +49,15 @@ pub(crate) use rustfs_ecstore::api::storage::{
use rustfs_storage_api as storage_contracts;
pub(crate) mod owner {
pub(crate) use rustfs_ecstore::api::disk::PhysicalResourcePlan;
#[cfg(test)]
pub(crate) use rustfs_ecstore::api::disk::{PhysicalDiskBinding, PhysicalDiskSample, PhysicalDiskSampler};
pub(crate) use rustfs_ecstore::api::rpc::with_rpc_background_request as ecstore_with_rpc_background_request;
#[cfg(test)]
pub(crate) use rustfs_ecstore::spawn_blocking_with_resource_ownership as ecstore_spawn_blocking_with_resource_ownership;
pub(crate) use rustfs_ecstore::with_background_disk_io as ecstore_with_background_disk_io;
pub(crate) use rustfs_ecstore::with_background_resource_set as ecstore_with_background_resource_set;
pub(crate) use super::storage_contracts::{ObjectIO, ObjectOperations};
pub(crate) use super::{
+20 -6
View File
@@ -1027,6 +1027,16 @@ impl HealTask {
#[tracing::instrument(skip(self), fields(task_id = %self.id, heal_type = ?self.heal_type))]
#[hotpath::measure]
pub async fn execute(&self) -> Result<()> {
self.execute_with_owner(None).await
}
#[tracing::instrument(name = "execute", skip(self, parent), fields(task_id = %self.id, heal_type = ?self.heal_type))]
#[hotpath::measure]
pub(super) async fn execute_owned(&self, parent: Option<super::manager::HealWork>) -> Result<()> {
self.execute_with_owner(parent).await
}
async fn execute_with_owner(&self, parent: Option<super::manager::HealWork>) -> Result<()> {
if self.source == HealRequestSource::AutoHeal
&& !self.heal_endpoints.is_empty()
&& matches!(self.heal_type, HealType::ErasureSet { .. })
@@ -1035,14 +1045,17 @@ impl HealTask {
// blocking write. Keep the entire executor, including its leases,
// alive until that write and failure persistence have finished.
let task = self.clone();
return tokio::spawn(async move { task.execute_inner().await })
.await
.map_err(|error| Error::other(format!("replacement executor failed: {error}")))?;
return super::manager::spawn_heal_work(
parent.as_ref(),
move |parent| async move { task.execute_inner(parent).await },
)
.await
.map_err(|error| Error::other(format!("replacement executor failed: {error}")))?;
}
self.execute_inner().await
self.execute_inner(parent).await
}
async fn execute_inner(&self) -> Result<()> {
async fn execute_inner(&self, parent: Option<super::manager::HealWork>) -> Result<()> {
self.outcome.write().await.start();
// update status and timestamps atomically to avoid race conditions
let now = SystemTime::now();
@@ -1111,7 +1124,8 @@ impl HealTask {
purge,
} => self.heal_delete_marker_purge(bucket, object, version_id, purge).await,
HealType::ErasureSet { buckets, set_disk_id } => {
self.heal_erasure_set(buckets.clone(), set_disk_id.clone()).await
self.heal_erasure_set_owned(buckets.clone(), set_disk_id.clone(), parent.as_ref())
.await
}
}
}
@@ -22,7 +22,17 @@ pub(super) async fn load_verified_replacement_resume(disk: &DiskStore, task_id:
}
impl HealTask {
#[cfg(test)]
pub(super) async fn heal_erasure_set(&self, buckets: Vec<String>, set_disk_id: String) -> Result<()> {
self.heal_erasure_set_owned(buckets, set_disk_id, None).await
}
pub(super) async fn heal_erasure_set_owned(
&self,
buckets: Vec<String>,
set_disk_id: String,
parent: Option<&super::super::manager::HealWork>,
) -> Result<()> {
debug!(
target: "rustfs::heal::task",
event = EVENT_HEAL_ERASURE_SET_STAGE,
@@ -532,6 +542,7 @@ impl HealTask {
})
.with_replacement_identity_fence(replacement_target_identities.clone())
.with_replacement_execution(replacement_execution)
.with_work_parent(parent.cloned())
.with_mainline_pacer(self.mainline_pacer.clone())
.with_admin_task((self.source == HealRequestSource::Admin).then(|| self.clone()));
+137
View File
@@ -5147,3 +5147,140 @@ async fn resumable_erasure_set_execution_is_cancelled_while_object_heal_is_pendi
assert!(storage.bucket_heal_opts.lock().unwrap().iter().all(|opts| opts.no_lock));
assert!(storage.object_heal_opts.lock().unwrap().iter().all(|opts| opts.no_lock));
}
#[tokio::test]
async fn replacement_owned_executor_keeps_manager_drain_pending_through_failure_persistence() {
let directory = TempDir::new().expect("survivor root");
let disk = make_resume_disk(&directory).await;
let identity = replacement_identity("replacement-a", "device-a", "mount-a");
let execution = crate::heal::storage::ReplacementExecution::for_test(Vec::new(), vec![identity.clone()]);
let ownership = Arc::downgrade(&execution);
let entered = Arc::new(tokio::sync::Notify::new());
let release = Arc::new(tokio::sync::Notify::new());
let storage = Arc::new(MockStorage {
replacement_target_identities_ready: Mutex::new(true),
replacement_target_identity_sequences: Mutex::new(VecDeque::from([vec![identity.clone()], vec![identity]])),
replacement_execution_fixture: Mutex::new(Some(execution)),
replacement_format_barrier: Some((entered.clone(), release.clone())),
resume_disk: Mutex::new(Some(disk.clone())),
..Default::default()
});
let mut request = directory_backed_replacement_request();
request.heal_endpoints = vec!["replacement-a".to_string()];
let task = Arc::new(HealTask::from_request(request, storage.clone()));
let manager = Arc::new(crate::HealManager::new_without_root_recovery_for_test(storage.clone(), None));
let parent = manager.test_work_parent();
let waiter = tokio::spawn({
let task = task.clone();
async move { task.execute_owned(Some(parent)).await }
});
entered.notified().await;
waiter.abort();
assert!(waiter.await.expect_err("waiter aborted").is_cancelled());
task.cancel_token.cancel();
assert!(ownership.upgrade().is_some(), "issued format I/O still owns execution");
let stopping = manager.clone();
let mut first_stop = tokio::spawn(async move { stopping.stop().await });
let first_pending = tokio::time::timeout(Duration::from_millis(30), &mut first_stop)
.await
.is_err();
first_stop.abort();
let first_cancelled = first_stop.await.is_err_and(|error| error.is_cancelled());
let stopping = manager.clone();
let mut retry = tokio::spawn(async move { stopping.stop().await });
let retry_pending = tokio::time::timeout(Duration::from_millis(30), &mut retry).await.is_err();
// Release actual owned work before the assertions/temporary disk cleanup.
release.notify_one();
let drained = tokio::time::timeout(Duration::from_secs(5), retry).await;
assert!(first_pending && first_cancelled && retry_pending);
drained
.expect("same owned executor drains")
.expect("stop waiter joins")
.expect("normal owned completion");
tokio::time::timeout(Duration::from_secs(5), async {
while ownership.upgrade().is_some() {
tokio::task::yield_now().await;
}
})
.await
.expect("executor drains after storage completes");
let state = ResumeManager::load_replacement_intent(disk, &task.id)
.await
.expect("failure state")
.get_state()
.await;
assert!(
state.error_message.as_deref().is_some_and(|error| error.contains("cancel")),
"{:?}",
state.error_message
);
assert!(!task.replacement_is_running());
assert!(storage.bucket_heal_calls.lock().unwrap().is_empty());
assert!(storage.heal_object_calls.lock().unwrap().is_empty());
}
#[tokio::test]
async fn replacement_owned_mutation_registers_after_handoff_and_survives_cancelled_waiters() {
let manager = Arc::new(crate::HealManager::new_without_root_recovery_for_test(
Arc::new(MockStorage::default()),
None,
));
let parent = manager.test_work_parent();
let stopping = manager.clone();
let mut first = tokio::spawn(async move { stopping.stop().await });
tokio::time::timeout(Duration::from_secs(5), async {
while !manager.stop_handoff_completed() {
tokio::task::yield_now().await;
}
})
.await
.expect("manager closed new admission while the accepted parent remains live");
let execution = crate::heal::storage::ReplacementExecution::for_test(Vec::new(), Vec::new());
let ownership = Arc::downgrade(&execution);
let (entered_tx, entered_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(async move {
execution
.run_owned(&parent, async move {
let _ = entered_tx.send(());
let _ = release_rx.await;
let _ = finished_tx.send(());
})
.await
});
entered_rx.await.expect("accepted descendant starts after tracker closure");
waiter.abort();
let waiter_cancelled = waiter.await.is_err_and(|error| error.is_cancelled());
let first_pending = tokio::time::timeout(Duration::from_millis(30), &mut first).await.is_err();
first.abort();
let first_cancelled = first.await.is_err_and(|error| error.is_cancelled());
let stopping = manager.clone();
let mut retry = tokio::spawn(async move { stopping.stop().await });
let retry_pending = tokio::time::timeout(Duration::from_millis(30), &mut retry).await.is_err();
let retained = ownership.upgrade().is_some();
let _ = release_tx.send(());
let drained = tokio::time::timeout(Duration::from_secs(5), retry).await;
let finished = finished_rx.await.is_ok();
assert!(waiter_cancelled && first_pending && first_cancelled && retry_pending && retained && finished);
drained
.expect("actual mutation finishes")
.expect("repeat waiter joins")
.expect("same manager owner drains");
assert!(ownership.upgrade().is_none());
manager.stop().await.expect("repeat successful receipt remains successful");
}
#[tokio::test]
async fn replacement_owned_mutation_panic_keeps_manager_join_failed() {
let manager = crate::HealManager::new_without_root_recovery_for_test(Arc::new(MockStorage::default()), None);
let parent = manager.test_work_parent();
let execution = crate::heal::storage::ReplacementExecution::for_test(Vec::new(), Vec::new());
let result = execution
.run_owned::<()>(&parent, async { panic!("owned replacement mutation panic") })
.await;
drop(parent);
assert!(result.is_err());
assert!(manager.stop().await.is_err(), "actual child panic must fail the manager owner");
assert!(manager.stop().await.is_err(), "repeat drain cannot erase the child panic");
}
+327 -19
View File
@@ -18,6 +18,10 @@ mod error;
pub mod heal;
pub use error::{Error, Result};
use futures::{
FutureExt,
future::{BoxFuture, Shared},
};
pub use heal::{
HealAdmissionTelemetry, HealManager, HealOperationsSnapshot, HealOptions, HealPriority, HealPriorityCounts, HealRequest,
HealSourceCounts, HealType,
@@ -62,8 +66,8 @@ pub fn create_ahm_services_cancel_token() -> CancellationToken {
/// Shutdown all heal services gracefully
pub async fn shutdown_ahm_services() -> Result<()> {
if let Some(manager) = get_heal_manager() {
manager.stop().await?;
if let Some(runtime) = GLOBAL_HEAL_RUNTIME.get() {
return runtime.shutdown(GLOBAL_AHM_SERVICES_CANCEL_TOKEN.get()).await;
}
if let Some(cancel_token) = GLOBAL_AHM_SERVICES_CANCEL_TOKEN.get() {
cancel_token.cancel();
@@ -71,9 +75,90 @@ pub async fn shutdown_ahm_services() -> Result<()> {
Ok(())
}
// The runtime retains the shared future that owns the actual consumer handle.
// Each waiter only clones the completion receipt; dropping one cannot detach
// the worker, and completed success/failure remains available to later waiters.
struct ConsumerJoin {
completion: Option<Shared<BoxFuture<'static, std::result::Result<(), Arc<Error>>>>>,
}
impl ConsumerJoin {
fn new(task: Option<tokio::task::JoinHandle<()>>) -> Self {
Self {
completion: task.map(|task| {
async move {
task.await
.map_err(|error| Arc::new(Error::other(format!("Heal consumer failed to join: {error}"))))
}
.boxed()
.shared()
}),
}
}
fn from_result(task: tokio::task::JoinHandle<Result<()>>) -> Self {
Self {
completion: Some(
async move {
task.await
.map_err(|error| Arc::new(Error::other(format!("Heal consumer failed to join: {error}"))))?
.map_err(Arc::new)
}
.boxed()
.shared(),
),
}
}
async fn wait(&self) -> Result<()> {
let Some(completion) = &self.completion else {
return Err(Error::other("Heal consumer did not start with an owned handle"));
};
completion
.clone()
.await
.map_err(|error| Error::other(format!("Heal consumer failed: {error}")))
}
}
struct HealRuntime {
manager: Arc<HealManager>,
channel_processor: Arc<Mutex<HealChannelProcessor>>,
channel_stop: CancellationToken,
channel_task: ConsumerJoin,
mrf_stop: CancellationToken,
mrf_task: ConsumerJoin,
}
impl HealRuntime {
async fn shutdown(&self, ahm_stop: Option<&CancellationToken>) -> Result<()> {
let manager_result = self.manager.stop().await;
self.finish_shutdown(manager_result, ahm_stop).await
}
async fn finish_shutdown(&self, manager_result: Result<()>, ahm_stop: Option<&CancellationToken>) -> Result<()> {
if !self.manager.stop_handoff_completed() {
return manager_result.and_then(|()| Err(Error::other("Heal manager stop returned without durable handoff")));
}
let mut first_error = manager_result.err();
self.channel_stop.cancel();
if let Some(stop) = ahm_stop {
stop.cancel();
}
if let Err(error) = self.channel_task.wait().await {
first_error.get_or_insert(error);
}
// The active channel command must finish its final publication/reply
// before the MRF consumer closes its receivers.
self.mrf_stop.cancel();
if let Err(error) = self.mrf_task.wait().await {
first_error.get_or_insert(error);
}
match first_error {
Some(error) => Err(error),
None => Ok(()),
}
}
}
/// Process-wide heal runtime. Async single-flight initialization prevents two
@@ -267,10 +352,6 @@ pub async fn init_heal_manager_with_workload_provider(
return Err(err);
}
// Start the MRF intent consumer (error-path repair intents + durable
// journal replay) now that the manager can accept submissions.
heal::mrf_queue::spawn_mrf_consumer(heal_manager.clone());
#[cfg(test)]
test_hook_after_manager_start().await;
@@ -293,16 +374,19 @@ pub async fn init_heal_manager_with_workload_provider(
}
};
let channel_processor = Arc::new(Mutex::new(HealChannelProcessor::new(heal_manager.clone())));
GLOBAL_HEAL_RUNTIME
.set(HealRuntime {
manager: heal_manager.clone(),
channel_processor: channel_processor.clone(),
})
.map_err(|_| Error::Config("Heal manager already initialized".to_string()))?;
tokio::spawn(async move {
let mut processor = channel_processor.lock().await;
if let Err(e) = processor.start_with_receipts(receiver, receipt_receiver).await {
// Channel initialization can fail and be retried. Do not claim the
// process-wide MRF channels before this recoverable boundary succeeds.
let mrf_stop = CancellationToken::new();
let mrf_task = ConsumerJoin::new(heal::mrf_queue::spawn_owned_mrf_consumer(heal_manager.clone(), mrf_stop.clone()));
let channel_stop = CancellationToken::new();
let processor_stop = channel_stop.clone();
let processor = channel_processor.clone();
let channel_task = ConsumerJoin::from_result(tokio::spawn(async move {
let mut processor = processor.lock().await;
let result = processor
.start_with_shutdown(receiver, receipt_receiver, processor_stop)
.await;
if let Err(e) = &result {
error!(
target: "rustfs::heal",
event = EVENT_HEAL_RUNTIME_STATE,
@@ -313,7 +397,18 @@ pub async fn init_heal_manager_with_workload_provider(
"Heal runtime channel processor failed"
);
}
});
result
}));
GLOBAL_HEAL_RUNTIME
.set(HealRuntime {
manager: heal_manager.clone(),
channel_processor,
channel_stop,
channel_task,
mrf_stop,
mrf_task,
})
.map_err(|_| Error::Config("Heal manager already initialized".to_string()))?;
info!(
target: "rustfs::heal",
@@ -491,8 +586,8 @@ pub(crate) fn set_heal_queue_length(count: usize) {
#[cfg(test)]
mod tests {
use super::{
Error, HEAL_RUNTIME_INIT_TEST_HOOK, HealRuntimeInitTestHook, get_heal_channel_processor, get_heal_manager,
heal::DiskStore, heal::manager::HealConfig, heal::storage::HealListItem, heal::storage::HealObjectInfo,
ConsumerJoin, Error, HEAL_RUNTIME_INIT_TEST_HOOK, HealRuntime, HealRuntimeInitTestHook, get_heal_channel_processor,
get_heal_manager, heal::DiskStore, heal::manager::HealConfig, heal::storage::HealListItem, heal::storage::HealObjectInfo,
heal::storage::HealStorageAPI, init_heal_manager, run_owned_initialization,
};
use crate::heal::storage_api::status::BucketInfo;
@@ -566,6 +661,51 @@ mod tests {
}
}
#[tokio::test]
async fn mrf_owned_join_survives_cancelled_shutdown_waiter() {
let release = Arc::new(Notify::new());
let completed = Arc::new(AtomicBool::new(false));
let worker_release = release.clone();
let worker_completed = completed.clone();
let owner = Arc::new(ConsumerJoin::new(Some(tokio::spawn(async move {
worker_release.notified().await;
worker_completed.store(true, Ordering::Release);
}))));
let waiting_owner = owner.clone();
let (started, entered) = oneshot::channel();
let waiter = tokio::spawn(async move {
let _ = started.send(());
waiting_owner.wait().await
});
entered.await.expect("shutdown waiter polls the shared actual handle");
waiter.abort();
let _ = waiter.await;
assert!(!completed.load(Ordering::Acquire));
assert!(owner.completion.as_ref().expect("owned handle").peek().is_none());
let second_owner = owner.clone();
let second = tokio::spawn(async move { second_owner.wait().await });
release.notify_one();
tokio::time::timeout(std::time::Duration::from_secs(1), owner.wait())
.await
.expect("owned MRF worker completes after release")
.expect("successful real join");
second
.await
.expect("concurrent waiter completes")
.expect("same successful receipt");
assert!(completed.load(Ordering::Acquire));
owner.wait().await.expect("repeat join receipt remains successful");
}
#[tokio::test]
async fn mrf_owned_join_failure_stays_failed() {
let owner = ConsumerJoin::new(Some(tokio::spawn(async { panic!("MRF owned-join test failure") })));
assert!(owner.wait().await.is_err());
assert!(owner.wait().await.is_err(), "a failed join must not become a successful receipt");
let unavailable = ConsumerJoin::new(None);
assert!(unavailable.wait().await.is_err(), "a missing owned handle is not a successful receipt");
}
#[tokio::test]
async fn owned_initialization_survives_waiter_cancellation() {
let completed = Arc::new(AtomicBool::new(false));
@@ -644,5 +784,173 @@ mod tests {
assert_eq!(hook.stops.load(Ordering::SeqCst), 1);
manager.stop().await.expect("manager should stop cleanly");
super::shutdown_ahm_services()
.await
.expect("retried initialization retains an actual MRF join receipt");
let runtime = super::GLOBAL_HEAL_RUNTIME.get().expect("published complete runtime");
assert!(
runtime
.mrf_task
.completion
.as_ref()
.expect("actual owned handle")
.peek()
.is_some_and(|result| result.is_ok())
);
}
#[tokio::test]
async fn channel_owned_join_inner_failure_stays_failed() {
let owner = ConsumerJoin::from_result(tokio::spawn(async { Err(Error::Config("channel loop failure".into())) }));
for _ in 0..2 {
assert!(owner.wait().await.unwrap_err().to_string().contains("channel loop failure"));
}
}
#[tokio::test]
async fn shutdown_pre_handoff_error_keeps_all_consumer_tokens_open() {
let manager = Arc::new(crate::HealManager::new_without_root_recovery_for_test(Arc::new(MockStorage), None));
let channel_stop = tokio_util::sync::CancellationToken::new();
let mrf_stop = tokio_util::sync::CancellationToken::new();
let ahm_stop = tokio_util::sync::CancellationToken::new();
let channel = channel_stop.clone();
let mrf = mrf_stop.clone();
let runtime = HealRuntime {
channel_processor: Arc::new(tokio::sync::Mutex::new(crate::HealChannelProcessor::new(manager.clone()))),
manager,
channel_stop,
channel_task: ConsumerJoin::from_result(tokio::spawn(async move {
channel.cancelled().await;
Ok(())
})),
mrf_stop,
mrf_task: ConsumerJoin::new(Some(tokio::spawn(async move { mrf.cancelled().await }))),
};
// Inject only the orchestration phase result; actual persistence failure
// and the witness ordering are exercised by the manager recovery test.
let error = runtime
.finish_shutdown(Err(Error::Config("handoff failed".into())), Some(&ahm_stop))
.await
.unwrap_err();
assert!(matches!(error, Error::Config(ref message) if message == "handoff failed"));
assert!(!runtime.manager.stop_handoff_completed());
assert!(!runtime.channel_stop.is_cancelled());
assert!(!runtime.mrf_stop.is_cancelled());
assert!(!ahm_stop.is_cancelled());
runtime.shutdown(Some(&ahm_stop)).await.unwrap();
assert!(runtime.manager.stop_handoff_completed());
}
#[tokio::test]
async fn shutdown_failed_manager_waits_channel_before_mrf_and_survives_cancelled_waiter() {
let manager = Arc::new(crate::HealManager::new_without_root_recovery_for_test(
Arc::new(MockStorage),
Some(HealConfig {
heal_interval: std::time::Duration::ZERO,
enable_auto_heal: false,
event_driven_scheduler_enable: false,
..Default::default()
}),
));
manager.start().await.unwrap();
let manager_error = manager
.stop()
.await
.expect_err("actual scheduler interval panic is a sticky join failure")
.to_string();
assert!(manager.stop_handoff_completed());
let channel_stop = tokio_util::sync::CancellationToken::new();
let mrf_stop = tokio_util::sync::CancellationToken::new();
let ahm_stop = tokio_util::sync::CancellationToken::new();
let channel = channel_stop.clone();
let mrf = mrf_stop.clone();
let channel_mrf = mrf_stop.clone();
let completed = Arc::new(AtomicBool::new(false));
let mrf_completed = completed.clone();
let (selected_tx, selected_rx) = oneshot::channel();
let (release_tx, release_rx) = oneshot::channel();
let runtime = Arc::new(HealRuntime {
channel_processor: Arc::new(tokio::sync::Mutex::new(crate::HealChannelProcessor::new(manager.clone()))),
manager,
channel_stop,
channel_task: ConsumerJoin::from_result(tokio::spawn(async move {
channel.cancelled().await;
assert!(!channel_mrf.is_cancelled());
selected_tx.send(()).unwrap();
release_rx.await.unwrap();
assert!(!channel_mrf.is_cancelled(), "MRF stays open through final channel publication");
Ok(())
})),
mrf_stop,
mrf_task: ConsumerJoin::new(Some(tokio::spawn(async move {
mrf.cancelled().await;
mrf_completed.store(true, Ordering::Release);
}))),
});
let stopping = runtime.clone();
let ahm = ahm_stop.clone();
let waiter = tokio::spawn(async move { stopping.shutdown(Some(&ahm)).await });
tokio::time::timeout(std::time::Duration::from_secs(1), selected_rx)
.await
.unwrap()
.unwrap();
assert!(ahm_stop.is_cancelled());
assert!(!runtime.mrf_stop.is_cancelled());
waiter.abort();
assert!(waiter.await.unwrap_err().is_cancelled());
let stopping = runtime.clone();
let retry = tokio::spawn(async move { stopping.shutdown(None).await });
tokio::task::yield_now().await;
assert!(!retry.is_finished(), "cancelled waiter cannot detach the held channel worker");
assert!(!runtime.mrf_stop.is_cancelled());
release_tx.send(()).unwrap();
let result = tokio::time::timeout(std::time::Duration::from_secs(1), retry)
.await
.unwrap()
.unwrap();
assert_eq!(result.unwrap_err().to_string(), manager_error);
assert!(completed.load(Ordering::Acquire));
assert_eq!(runtime.shutdown(None).await.unwrap_err().to_string(), manager_error);
}
#[tokio::test]
async fn shutdown_channel_failure_still_awaits_actual_mrf_failure() {
let manager = Arc::new(crate::HealManager::new_without_root_recovery_for_test(Arc::new(MockStorage), None));
let channel_stop = tokio_util::sync::CancellationToken::new();
let mrf_stop = tokio_util::sync::CancellationToken::new();
let channel = channel_stop.clone();
let mrf = mrf_stop.clone();
let (entered_tx, entered_rx) = oneshot::channel();
let (release_tx, release_rx) = oneshot::channel();
let runtime = Arc::new(HealRuntime {
channel_processor: Arc::new(tokio::sync::Mutex::new(crate::HealChannelProcessor::new(manager.clone()))),
manager,
channel_stop,
channel_task: ConsumerJoin::from_result(tokio::spawn(async move {
channel.cancelled().await;
Err(Error::Config("channel inner failure".into()))
})),
mrf_stop,
mrf_task: ConsumerJoin::new(Some(tokio::spawn(async move {
mrf.cancelled().await;
entered_tx.send(()).unwrap();
release_rx.await.unwrap();
panic!("actual MRF join failure");
}))),
});
let stopping = runtime.clone();
let waiter = tokio::spawn(async move { stopping.shutdown(None).await });
tokio::time::timeout(std::time::Duration::from_secs(1), entered_rx)
.await
.unwrap()
.unwrap();
assert!(!waiter.is_finished(), "channel error cannot skip the actual MRF join");
release_tx.send(()).unwrap();
let result = tokio::time::timeout(std::time::Duration::from_secs(1), waiter)
.await
.unwrap()
.unwrap();
assert!(result.unwrap_err().to_string().contains("channel inner failure"));
assert!(runtime.channel_task.wait().await.is_err());
assert!(runtime.mrf_task.wait().await.is_err());
}
}
+637 -10
View File
@@ -30,8 +30,8 @@ use std::io::{self, Error};
use std::net::IpAddr;
use std::ops::Not as _;
use std::pin::Pin;
use std::sync::LazyLock;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, LazyLock};
use std::task::{Context, Poll};
use std::time::{Duration, Instant};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
@@ -1409,6 +1409,24 @@ impl Stream for ReceiverStream {
}
}
/// Owns a caller's memory reservation until the associated allocation is freed.
pub type HttpWriterMemoryReservation = Box<dyn Send + Sync + 'static>;
/// Reserves a full buffer capacity before the writer allocates or copies data.
pub type HttpWriterMemoryReserver = Arc<dyn Fn(usize) -> io::Result<HttpWriterMemoryReservation> + Send + Sync + 'static>;
struct ReservedHttpBytes {
// Release the allocation before its reservation, including its last slice.
bytes: Bytes,
_reservation: Arc<HttpWriterMemoryReservation>,
}
impl AsRef<[u8]> for ReservedHttpBytes {
fn as_ref(&self) -> &[u8] {
&self.bytes
}
}
// Not pin-projected: every field is `Unpin` and the `AsyncWrite` impl accesses
// them through `get_mut()`, so a plain struct lets us add a manual `Drop`
// (pin-project forbids one) to abort the background HTTP task — see below.
@@ -1421,6 +1439,8 @@ pub struct HttpWriter {
sender: PollSender<Option<Bytes>>,
handle: tokio::task::JoinHandle<std::io::Result<()>>,
pending_chunk: BytesMut,
pending_reservation: Option<Arc<HttpWriterMemoryReservation>>,
memory_reserver: Option<HttpWriterMemoryReserver>,
finish: bool,
track_internode_metrics: bool,
internode_operation: Option<&'static str>,
@@ -1432,6 +1452,32 @@ const HTTP_WRITER_BUFFER_SIZE: usize = 1024 * 1024;
impl HttpWriter {
/// Create a new HttpWriter for the given URL. The HTTP request is performed in the background.
pub async fn new(url: String, method: Method, headers: HeaderMap) -> io::Result<Self> {
Self::with_memory_reserver(url, method, headers, None).await
}
/// Keep each pending and queued allocation charged through its last byte slice.
pub async fn new_with_memory_reserver(
url: String,
method: Method,
headers: HeaderMap,
memory_reserver: HttpWriterMemoryReserver,
) -> io::Result<Self> {
Self::with_memory_reserver(url, method, headers, Some(memory_reserver)).await
}
async fn with_memory_reserver(
url: String,
method: Method,
headers: HeaderMap,
memory_reserver: Option<HttpWriterMemoryReserver>,
) -> io::Result<Self> {
// A rejected reservation must not leave an HTTP task or payload allocation.
let pending_reservation = memory_reserver
.as_ref()
.map(|reserve| reserve(HTTP_WRITER_BUFFER_SIZE))
.transpose()?
.map(Arc::new);
let pending_chunk = BytesMut::with_capacity(HTTP_WRITER_BUFFER_SIZE);
// http_log!("[HttpWriter::new] url: {url}, method: {method:?}, headers: {headers:?}");
let url_clone = url.clone();
let method_clone = method.clone();
@@ -1523,7 +1569,9 @@ impl HttpWriter {
start_tx: Some(start_tx),
sender: PollSender::new(sender),
handle,
pending_chunk: BytesMut::with_capacity(HTTP_WRITER_BUFFER_SIZE),
pending_chunk,
pending_reservation,
memory_reserver,
finish: false,
track_internode_metrics,
internode_operation,
@@ -1718,6 +1766,50 @@ fn send_error_to_io<T>(err: tokio_util::sync::PollSendError<T>, context: &str) -
}
impl HttpWriter {
fn ensure_pending_capacity(&mut self) -> io::Result<()> {
if let Some(reserve) = &self.memory_reserver {
if self.pending_reservation.is_some() {
if self.pending_chunk.capacity() != 0 || self.pending_chunk.try_reclaim(HTTP_WRITER_BUFFER_SIZE) {
return Ok(());
}
// Outstanding output still owns this allocation. Release the
// empty tail before its guard, then reserve separate storage.
self.pending_chunk = BytesMut::new();
self.pending_reservation = None;
}
let reservation = reserve(HTTP_WRITER_BUFFER_SIZE)?;
self.pending_chunk = BytesMut::with_capacity(HTTP_WRITER_BUFFER_SIZE);
self.pending_reservation = Some(Arc::new(reservation));
}
Ok(())
}
fn take_pending_chunk(&mut self) -> io::Result<Bytes> {
if self.memory_reserver.is_none() {
return Ok(self.pending_chunk.split().freeze());
}
let reservation = self
.pending_reservation
.as_ref()
.cloned()
.ok_or_else(|| io::Error::other("HTTP writer buffer has no memory reservation"))?;
let bytes = self.pending_chunk.split().freeze();
// The tail keeps this allocation charged until it can be reclaimed
// without allocating, or is dropped before requesting new storage.
Ok(Bytes::from_owner(ReservedHttpBytes {
bytes,
_reservation: reservation,
}))
}
fn pending_capacity(&self) -> usize {
if self.memory_reserver.is_some() {
self.pending_chunk.capacity()
} else {
HTTP_WRITER_BUFFER_SIZE
}
}
fn start_request(&mut self) {
if let Some(start_tx) = self.start_tx.take() {
let _ = start_tx.send(());
@@ -1738,7 +1830,7 @@ impl HttpWriter {
match self.sender.poll_reserve(cx) {
Poll::Ready(Ok(())) => {
let chunk = self.pending_chunk.split().freeze();
let chunk = self.take_pending_chunk()?;
self.sender
.send_item(Some(chunk))
.map_err(|e| send_error_to_io(e, "HttpWriter send error"))?;
@@ -1763,7 +1855,7 @@ impl AsyncWrite for HttpWriter {
return Poll::Ready(Err(err));
}
if this.pending_chunk.len() >= HTTP_WRITER_BUFFER_SIZE {
if !this.pending_chunk.is_empty() && this.pending_chunk.len() >= this.pending_capacity() {
match this.poll_send_pending_chunk(cx) {
Poll::Ready(Ok(())) => {}
Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
@@ -1774,11 +1866,22 @@ impl AsyncWrite for HttpWriter {
if buf.len() >= HTTP_WRITER_BUFFER_SIZE && this.pending_chunk.is_empty() {
match this.sender.poll_reserve(cx) {
Poll::Ready(Ok(())) => {
let (chunk, written) = if this.memory_reserver.is_some() {
if let Err(error) = this.ensure_pending_capacity() {
this.sender.abort_send();
return Poll::Ready(Err(error));
}
let written = buf.len().min(this.pending_capacity());
this.pending_chunk.extend_from_slice(&buf[..written]);
(this.take_pending_chunk()?, written)
} else {
(Bytes::copy_from_slice(&buf[..HTTP_WRITER_BUFFER_SIZE]), HTTP_WRITER_BUFFER_SIZE)
};
this.sender
.send_item(Some(Bytes::copy_from_slice(&buf[..HTTP_WRITER_BUFFER_SIZE])))
.send_item(Some(chunk))
.map_err(|e| send_error_to_io(e, "HttpWriter send error"))?;
this.start_request();
return Poll::Ready(Ok(HTTP_WRITER_BUFFER_SIZE));
return Poll::Ready(Ok(written));
}
Poll::Ready(Err(err)) => return Poll::Ready(Err(poll_send_error_to_io(err, "HttpWriter send error"))),
Poll::Pending => return Poll::Pending,
@@ -1786,7 +1889,10 @@ impl AsyncWrite for HttpWriter {
}
// A slot bounds memory only when every queued and pending chunk is bounded.
let written = buf.len().min(HTTP_WRITER_BUFFER_SIZE - this.pending_chunk.len());
if !buf.is_empty() {
this.ensure_pending_capacity()?;
}
let written = buf.len().min(this.pending_capacity() - this.pending_chunk.len());
this.pending_chunk.extend_from_slice(&buf[..written]);
if written > 0 {
this.start_request();
@@ -1813,7 +1919,7 @@ impl AsyncWrite for HttpWriter {
return Poll::Ready(Err(err));
}
if this.pending_chunk.len() >= HTTP_WRITER_BUFFER_SIZE {
if !this.pending_chunk.is_empty() && this.pending_chunk.len() >= this.pending_capacity() {
match this.poll_send_pending_chunk(cx) {
Poll::Ready(Ok(())) => {}
Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
@@ -1821,10 +1927,16 @@ impl AsyncWrite for HttpWriter {
}
}
let available = HTTP_WRITER_BUFFER_SIZE - this.pending_chunk.len();
if this.memory_reserver.is_some() && bufs.iter().any(|buf| !buf.is_empty()) {
this.ensure_pending_capacity()?;
}
let available = this.pending_capacity() - this.pending_chunk.len();
let mut written = 0;
for buf in bufs {
let len = buf.len().min(available - written);
if len > 0 {
this.ensure_pending_capacity()?;
}
this.pending_chunk.extend_from_slice(&buf[..len]);
written += len;
if written == available {
@@ -1884,6 +1996,11 @@ impl AsyncWrite for HttpWriter {
// );
self.finish = true;
if self.memory_reserver.is_some() {
// EOF has been queued, so this writer no longer needs its tail.
self.pending_chunk = BytesMut::new();
self.pending_reservation = None;
}
}
// Wait for the HTTP request to complete
use futures::FutureExt;
@@ -2943,6 +3060,17 @@ mod tests {
}
fn paused_http_writer() -> (HttpWriter, tokio::sync::mpsc::Receiver<Option<Bytes>>) {
paused_http_writer_with_memory(None).expect("unrestricted writer")
}
fn paused_http_writer_with_memory(
memory_reserver: Option<HttpWriterMemoryReserver>,
) -> io::Result<(HttpWriter, tokio::sync::mpsc::Receiver<Option<Bytes>>)> {
let pending_reservation = memory_reserver
.as_ref()
.map(|reserve| reserve(HTTP_WRITER_BUFFER_SIZE))
.transpose()?
.map(Arc::new);
let (sender, receiver) = tokio::sync::mpsc::channel(HTTP_WRITER_CHANNEL_CAPACITY);
let (_, err_rx) = tokio::sync::oneshot::channel();
let writer = HttpWriter {
@@ -2954,11 +3082,510 @@ mod tests {
sender: PollSender::new(sender),
handle: tokio::spawn(std::future::pending()),
pending_chunk: BytesMut::with_capacity(HTTP_WRITER_BUFFER_SIZE),
pending_reservation,
memory_reserver,
finish: false,
track_internode_metrics: false,
internode_operation: None,
};
(writer, receiver)
Ok((writer, receiver))
}
struct HttpWriterTestMemory {
limit: AtomicUsize,
used: AtomicUsize,
calls: AtomicUsize,
}
struct HttpWriterTestReservation {
memory: Arc<HttpWriterTestMemory>,
capacity: usize,
}
impl Drop for HttpWriterTestReservation {
fn drop(&mut self) {
self.memory.used.fetch_sub(self.capacity, Ordering::SeqCst);
}
}
impl HttpWriterTestMemory {
fn new(limit: usize) -> Arc<Self> {
Arc::new(Self {
limit: AtomicUsize::new(limit),
used: AtomicUsize::new(0),
calls: AtomicUsize::new(0),
})
}
fn reserver(self: &Arc<Self>) -> HttpWriterMemoryReserver {
let memory = self.clone();
Arc::new(move |capacity| {
assert_eq!(capacity, HTTP_WRITER_BUFFER_SIZE, "reserve full allocation capacity");
memory.calls.fetch_add(1, Ordering::SeqCst);
memory
.used
.try_update(Ordering::SeqCst, Ordering::SeqCst, |used| {
used.checked_add(capacity)
.filter(|next| *next <= memory.limit.load(Ordering::SeqCst))
})
.map_err(|_| io::Error::from(io::ErrorKind::WouldBlock))?;
Ok(Box::new(HttpWriterTestReservation {
memory: memory.clone(),
capacity,
}) as HttpWriterMemoryReservation)
})
}
fn used(&self) -> usize {
self.used.load(Ordering::SeqCst)
}
}
#[test]
fn http_writer_memory_denial_precedes_runtime_or_allocation() {
use std::future::Future;
for limit in [0, HTTP_WRITER_BUFFER_SIZE - 1] {
let memory = HttpWriterTestMemory::new(limit);
let mut create = Box::pin(HttpWriter::new_with_memory_reserver(
"http://127.0.0.1/unpolled".to_owned(),
Method::PUT,
HeaderMap::new(),
memory.reserver(),
));
let mut cx = Context::from_waker(std::task::Waker::noop());
// Spawning before admission would panic without a Tokio runtime.
assert!(
matches!(create.as_mut().poll(&mut cx), Poll::Ready(Err(error)) if error.kind() == io::ErrorKind::WouldBlock)
);
assert_eq!(memory.calls.load(Ordering::SeqCst), 1);
assert_eq!(memory.used(), 0);
}
}
#[tokio::test]
async fn http_writer_memory_tracks_pending_and_last_output_slice() {
let memory = HttpWriterTestMemory::new(2 * HTTP_WRITER_BUFFER_SIZE);
let (mut writer, mut receiver) = paused_http_writer_with_memory(Some(memory.reserver())).expect("reserved writer");
assert_eq!(writer.pending_chunk.capacity(), HTTP_WRITER_BUFFER_SIZE);
assert_eq!(memory.used(), HTTP_WRITER_BUFFER_SIZE);
memory.limit.store(0, Ordering::SeqCst);
writer.write_all(b"first").await.expect("pending payload");
writer.flush().await.expect("emit partial allocation");
assert_eq!(writer.pending_chunk.capacity(), HTTP_WRITER_BUFFER_SIZE - 5);
assert_eq!(memory.used(), HTTP_WRITER_BUFFER_SIZE, "the empty tail still owns this allocation");
let frame = receiver.recv().await.expect("queued frame").expect("payload");
let clone = frame.clone();
let last = frame.slice(4..);
memory.limit.store(2 * HTTP_WRITER_BUFFER_SIZE, Ordering::SeqCst);
writer
.write_all(b"next")
.await
.expect("reuse the allocation's remaining tail");
assert_eq!(memory.used(), HTTP_WRITER_BUFFER_SIZE);
drop(frame);
drop(clone);
drop(writer);
assert_eq!(last.as_ref(), b"t");
assert_eq!(memory.used(), HTTP_WRITER_BUFFER_SIZE);
drop(last);
assert_eq!(memory.used(), 0);
}
#[tokio::test]
async fn http_writer_memory_shrink_rejects_new_scalar_and_vectored_storage() {
for mode in ["scalar", "fast", "vectored"] {
let memory = HttpWriterTestMemory::new(2 * HTTP_WRITER_BUFFER_SIZE);
let (mut writer, mut receiver) = paused_http_writer_with_memory(Some(memory.reserver())).expect("reserved writer");
writer
.write_all(&vec![0x5a; HTTP_WRITER_BUFFER_SIZE])
.await
.expect("fill the first allocation");
writer
.flush()
.await
.expect("first allocation transferred without a remaining tail");
let retained = receiver.recv().await.expect("first frame").expect("payload");
memory.limit.store(0, Ordering::SeqCst);
let error = match mode {
"vectored" => writer.write_vectored(&[IoSlice::new(b"lost"), IoSlice::new(b"bytes")]).await,
"fast" => writer.write(&vec![0xff; HTTP_WRITER_BUFFER_SIZE]).await,
_ => writer.write(b"lost bytes").await,
}
.expect_err("new storage must be rejected after shrink");
assert_eq!(error.kind(), io::ErrorKind::WouldBlock);
assert_eq!(writer.pending_chunk.capacity(), 0);
assert!(receiver.try_recv().is_err(), "rejected bytes must not be queued");
assert_eq!(memory.used(), HTTP_WRITER_BUFFER_SIZE);
let calls = memory.calls.load(Ordering::SeqCst);
assert_eq!(writer.write(&[]).await.expect("empty write"), 0);
assert_eq!(memory.calls.load(Ordering::SeqCst), calls);
memory.limit.store(2 * HTTP_WRITER_BUFFER_SIZE, Ordering::SeqCst);
writer.write_all(b"accepted").await.expect("restored capacity");
writer.flush().await.expect("emit accepted bytes");
let accepted = receiver.recv().await.expect("accepted frame").expect("payload");
assert_eq!(accepted.as_ref(), b"accepted");
drop(writer);
drop(retained);
drop(accepted);
assert_eq!(memory.used(), 0);
}
}
#[tokio::test]
async fn http_writer_memory_full_queue_keeps_backpressure_and_capacity() {
for vectored in [false, true] {
let memory = HttpWriterTestMemory::new((HTTP_WRITER_CHANNEL_CAPACITY + 1) * HTTP_WRITER_BUFFER_SIZE);
let (mut writer, mut receiver) = paused_http_writer_with_memory(Some(memory.reserver())).expect("reserved writer");
let payload = vec![0xa5; HTTP_WRITER_BUFFER_SIZE];
for _ in 0..HTTP_WRITER_CHANNEL_CAPACITY {
let written = if vectored {
writer
.write_vectored(&[IoSlice::new(&payload[..32]), IoSlice::new(&payload[32..])])
.await
} else {
writer.write(&payload).await
}
.expect("free queue slot");
assert_eq!(written, HTTP_WRITER_BUFFER_SIZE);
writer.flush().await.expect("queue the complete allocation");
}
assert_eq!(memory.used(), HTTP_WRITER_CHANNEL_CAPACITY * HTTP_WRITER_BUFFER_SIZE);
assert_eq!(writer.pending_chunk.capacity(), 0);
if vectored {
let written = writer
.write_vectored(&[IoSlice::new(&payload[..32]), IoSlice::new(&payload[32..])])
.await
.expect("one bounded pending allocation beside the full queue");
assert_eq!(written, HTTP_WRITER_BUFFER_SIZE);
assert_eq!(memory.used(), (HTTP_WRITER_CHANNEL_CAPACITY + 1) * HTTP_WRITER_BUFFER_SIZE);
}
let calls = memory.calls.load(Ordering::SeqCst);
{
let write = writer.write(&payload);
tokio::pin!(write);
assert!(
futures::poll!(&mut write).is_pending(),
"a full queue must not accept another large write"
);
}
assert_eq!(
memory.calls.load(Ordering::SeqCst),
calls,
"backpressure must not allocate a fast-path copy"
);
if !vectored {
drop(receiver.recv().await.expect("free a queue slot"));
assert_eq!(writer.write(&payload).await.expect("resumed write"), HTTP_WRITER_BUFFER_SIZE);
}
drop(writer);
assert_eq!(memory.used(), HTTP_WRITER_CHANNEL_CAPACITY * HTTP_WRITER_BUFFER_SIZE);
while let Ok(frame) = receiver.try_recv() {
let frame = frame.expect("payload");
assert_eq!(frame.as_ref(), payload.as_slice());
}
assert_eq!(memory.used(), 0);
}
}
#[test]
fn http_writer_memory_cancellation_and_runtime_shutdown_release_queued_storage() {
for cancel_writer in [false, true] {
let memory = HttpWriterTestMemory::new(2 * HTTP_WRITER_BUFFER_SIZE);
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("writer runtime");
let writer = runtime.block_on(async {
let (mut writer, receiver) = paused_http_writer_with_memory(Some(memory.reserver())).expect("reserved writer");
writer.handle.abort();
writer.handle = tokio::spawn(async move {
std::future::pending::<()>().await;
drop(receiver);
Ok(())
});
writer.write_all(b"queued").await.expect("pending bytes");
writer.flush().await.expect("queued allocation");
writer
});
assert_eq!(memory.used(), HTTP_WRITER_BUFFER_SIZE);
if cancel_writer {
drop(writer);
runtime.block_on(tokio::task::yield_now());
assert_eq!(memory.used(), 0, "cancelled task must release its queue");
drop(runtime);
} else {
drop(runtime);
assert_eq!(memory.used(), HTTP_WRITER_BUFFER_SIZE, "the writer still owns the shared tail");
drop(writer);
assert_eq!(memory.used(), 0, "the final pending owner releases its allocation");
}
}
}
#[tokio::test]
async fn http_writer_memory_short_flushes_share_the_pending_allocation() {
let memory = HttpWriterTestMemory::new(2 * HTTP_WRITER_BUFFER_SIZE);
let (mut writer, mut receiver) =
paused_http_writer_with_memory(Some(memory.reserver())).expect("reserve the shared short-frame allocation");
let base = writer.pending_chunk.as_ptr();
for index in 0..16 {
writer
.write_all(&[0x37; 4096])
.await
.expect("write a short frame into the existing tail");
writer
.flush()
.await
.expect("flush a short frame without replacing its allocation");
let frame = receiver
.recv()
.await
.expect("receive the short frame")
.expect("short frame payload");
assert_eq!(frame.as_ptr(), base.wrapping_add(index * 4096), "split the original payload allocation");
assert_eq!(frame.as_ref(), &[0x37; 4096]);
drop(frame);
assert_eq!(memory.used(), HTTP_WRITER_BUFFER_SIZE, "the pending tail remains a real owner");
}
assert_eq!(memory.calls.load(Ordering::SeqCst), 1);
assert_eq!(writer.pending_chunk.capacity(), HTTP_WRITER_BUFFER_SIZE - 16 * 4096);
drop(writer);
assert_eq!(memory.used(), 0);
}
#[tokio::test]
async fn http_writer_memory_consumed_full_frames_reclaim_without_a_new_reservation() {
let memory = HttpWriterTestMemory::new(HTTP_WRITER_BUFFER_SIZE);
let (mut writer, mut receiver) =
paused_http_writer_with_memory(Some(memory.reserver())).expect("reserve the reclaimable allocation");
let base = writer.pending_chunk.as_ptr();
let payload = vec![0x37; HTTP_WRITER_BUFFER_SIZE];
memory.limit.store(0, Ordering::SeqCst);
for _ in 0..16 {
assert_eq!(
writer
.write(&payload)
.await
.expect("reuse charged storage after the cap closes"),
payload.len()
);
writer.flush().await.expect("flush a complete reusable allocation");
let frame = receiver
.recv()
.await
.expect("receive the complete frame")
.expect("complete frame payload");
assert_eq!(frame.as_ptr(), base, "try_reclaim must reuse the allocation in place");
assert_eq!(frame.as_ref(), payload.as_slice());
drop(frame);
assert_eq!(memory.used(), HTTP_WRITER_BUFFER_SIZE);
}
assert_eq!(memory.calls.load(Ordering::SeqCst), 1, "reusing old storage is allowed after shrink");
drop(writer);
assert_eq!(memory.used(), 0);
}
#[tokio::test]
async fn http_writer_memory_last_slice_prevents_reclaim_and_keeps_both_allocations_charged() {
let memory = HttpWriterTestMemory::new(2 * HTTP_WRITER_BUFFER_SIZE);
let (mut writer, mut receiver) =
paused_http_writer_with_memory(Some(memory.reserver())).expect("reserve the first retained allocation");
writer
.write_all(&vec![0x5a; HTTP_WRITER_BUFFER_SIZE])
.await
.expect("queue the complete first allocation");
let first = receiver
.recv()
.await
.expect("receive the retained first frame")
.expect("first frame payload");
let last = first.slice(0..1);
drop(first);
memory.limit.store(0, Ordering::SeqCst);
let error = writer
.write(b"n")
.await
.expect_err("a live slice prevents reclaim under the closed cap");
assert_eq!(error.kind(), io::ErrorKind::WouldBlock);
assert_eq!(writer.pending_chunk.capacity(), 0);
assert!(writer.pending_reservation.is_none(), "failed reclaim drops the empty view and its guard");
assert_eq!(memory.used(), HTTP_WRITER_BUFFER_SIZE);
memory.limit.store(2 * HTTP_WRITER_BUFFER_SIZE, Ordering::SeqCst);
writer
.write_all(b"new")
.await
.expect("reserve a separate allocation beside the live slice");
writer.flush().await.expect("queue bytes from the second allocation");
let second = receiver
.recv()
.await
.expect("receive the second frame")
.expect("second frame payload");
assert_ne!(second.as_ptr(), last.as_ptr(), "a live slice prevents storage reuse");
assert_eq!(memory.used(), 2 * HTTP_WRITER_BUFFER_SIZE);
drop(last);
assert_eq!(memory.used(), HTTP_WRITER_BUFFER_SIZE);
drop(second);
assert_eq!(
memory.used(),
HTTP_WRITER_BUFFER_SIZE,
"new pending tail still owns the second allocation"
);
drop(writer);
assert_eq!(memory.used(), 0);
}
#[tokio::test]
async fn http_writer_memory_half_full_tail_bounds_scalar_fast_and_vectored_writes() {
for mode in ["scalar", "fast", "vectored"] {
let memory = HttpWriterTestMemory::new(2 * HTTP_WRITER_BUFFER_SIZE);
let (mut writer, mut receiver) =
paused_http_writer_with_memory(Some(memory.reserver())).expect("reserve the half-filled allocation");
let prefix = vec![0x11; HTTP_WRITER_BUFFER_SIZE / 2];
writer.write_all(&prefix).await.expect("fill half of the original allocation");
writer.flush().await.expect("split off the first half and retain the tail");
let first = receiver
.recv()
.await
.expect("receive the first half")
.expect("first half payload");
assert_eq!(writer.pending_chunk.capacity(), HTTP_WRITER_BUFFER_SIZE / 2);
memory.limit.store(0, Ordering::SeqCst);
let payload = vec![0x5a; HTTP_WRITER_BUFFER_SIZE];
let written = match mode {
"fast" => writer.write(&payload).await,
"vectored" => {
writer
.write_vectored(&[IoSlice::new(&payload[..32]), IoSlice::new(&payload[32..])])
.await
}
_ => writer.write(&payload[..payload.len() - 1]).await,
}
.expect("the half-capacity tail remains usable after the budget shrinks");
assert_eq!(written, HTTP_WRITER_BUFFER_SIZE / 2, "bound {mode} writes by actual tail capacity");
writer
.flush()
.await
.expect("flush the full remaining half without growing storage");
let tail = receiver
.recv()
.await
.expect("receive the remaining half")
.expect("tail payload");
assert_eq!(tail.as_ref(), &payload[..written]);
assert_eq!(tail.as_ptr(), first.as_ptr().wrapping_add(prefix.len()));
assert_eq!(memory.calls.load(Ordering::SeqCst), 1);
assert_eq!(memory.used(), HTTP_WRITER_BUFFER_SIZE);
let error = writer
.write(b"new storage")
.await
.expect_err("exhausted shared storage needs a new reservation");
assert_eq!(error.kind(), io::ErrorKind::WouldBlock);
assert!(receiver.try_recv().is_err(), "rejected bytes must not be queued");
drop(writer);
drop(first);
assert_eq!(memory.used(), HTTP_WRITER_BUFFER_SIZE);
drop(tail);
assert_eq!(memory.used(), 0);
}
}
#[tokio::test]
async fn http_writer_memory_short_queue_keeps_backpressure_and_one_allocation() {
let memory = HttpWriterTestMemory::new(HTTP_WRITER_BUFFER_SIZE);
let (mut writer, mut receiver) =
paused_http_writer_with_memory(Some(memory.reserver())).expect("reserve the shared queued allocation");
for _ in 0..HTTP_WRITER_CHANNEL_CAPACITY {
writer.write_all(&[0x37; 4096]).await.expect("write one short queued frame");
writer.flush().await.expect("occupy one available queue slot");
}
assert_eq!(memory.calls.load(Ordering::SeqCst), 1);
assert_eq!(memory.used(), HTTP_WRITER_BUFFER_SIZE);
let large = vec![0x37; HTTP_WRITER_BUFFER_SIZE];
{
let blocked = writer.write(&large);
tokio::pin!(blocked);
assert!(futures::poll!(&mut blocked).is_pending());
}
assert_eq!(memory.calls.load(Ordering::SeqCst), 1, "full queue cannot allocate a fast path");
writer
.write_all(b"pending")
.await
.expect("the charged tail accepts one pending short write");
{
let blocked = writer.flush();
tokio::pin!(blocked);
assert!(futures::poll!(&mut blocked).is_pending());
}
drop(writer);
for remaining in (0..HTTP_WRITER_CHANNEL_CAPACITY).rev() {
let frame = receiver
.recv()
.await
.expect("drain a queued short frame")
.expect("queued short payload");
assert_eq!(frame.as_ref(), &[0x37; 4096]);
drop(frame);
assert_eq!(memory.used(), if remaining == 0 { 0 } else { HTTP_WRITER_BUFFER_SIZE });
}
assert!(receiver.recv().await.is_none());
}
#[tokio::test]
async fn http_writer_memory_clean_shutdown_releases_only_its_unused_tail() {
let memory = HttpWriterTestMemory::new(HTTP_WRITER_BUFFER_SIZE);
let (mut writer, mut receiver) =
paused_http_writer_with_memory(Some(memory.reserver())).expect("reserve the shutdown body allocation");
writer.handle.abort();
writer.handle = tokio::spawn(async { Ok(()) });
writer.write_all(b"body").await.expect("buffer the final body bytes");
writer
.shutdown()
.await
.expect("send the body and EOF, then join the completed request");
assert_eq!(writer.pending_chunk.capacity(), 0);
assert!(writer.pending_reservation.is_none());
assert_eq!(memory.used(), HTTP_WRITER_BUFFER_SIZE, "queued body still owns this allocation");
let frame = receiver
.recv()
.await
.expect("receive the final body frame")
.expect("final body payload");
let last = frame.slice(1..2);
drop(frame);
assert!(receiver.recv().await.expect("receive the queued EOF").is_none());
drop(writer);
assert_eq!(memory.used(), HTTP_WRITER_BUFFER_SIZE);
drop(last);
assert_eq!(memory.used(), 0);
}
#[tokio::test]
async fn http_writer_memory_round_trip_preserves_vectored_body_and_shutdown() {
let state = TestState::default();
let (url, handle) = start_test_server(state.clone()).await.expect("memory writer test server");
let memory = HttpWriterTestMemory::new((HTTP_WRITER_CHANNEL_CAPACITY + 1) * HTTP_WRITER_BUFFER_SIZE);
let header = [0x11; 32];
let payload = vec![0xa5; 3 * HTTP_WRITER_BUFFER_SIZE + 17];
let mut writer = HttpWriter::new_with_memory_reserver(url, Method::PUT, HeaderMap::new(), memory.reserver())
.await
.expect("reserved HTTP writer");
let mut slices = [IoSlice::new(&[]), IoSlice::new(&header), IoSlice::new(&payload)];
let mut remaining = slices.as_mut_slice();
while remaining.iter().any(|slice| !slice.is_empty()) {
let written = writer.write_vectored(remaining).await.expect("reserved vectored write");
assert!(written > 0 && written <= HTTP_WRITER_BUFFER_SIZE);
IoSlice::advance_slices(&mut remaining, written);
}
writer.flush().await.expect("flush final reserved allocation");
writer.shutdown().await.expect("complete reserved HTTP request");
let bodies = state.put_bodies.lock().await;
assert_eq!(bodies.len(), 1);
assert_eq!(&bodies[0][..header.len()], &header);
assert_eq!(&bodies[0][header.len()..], payload.as_slice());
assert_eq!(memory.used(), 0, "completed transport retains no payload reservation");
handle.abort();
}
#[tokio::test]
+70 -55
View File
@@ -31,7 +31,8 @@ pub use rustfs_data_usage::{
DataUsageEntry, DataUsageHash, DataUsageHashMap, DataUsageInfo, DataUsageSegmentInvalidationProof, DataUsageSnapshotSetState,
LEGACY_DATA_USAGE_OBJECT_NAME, PrefixUsageEntry, PrefixUsageQuery, PrefixUsageSummary, ReplTargetSizeSummary,
SizeReconciliationEntry, SizeReconciliationScope, SizeSummary, TierAccountingProof, TierStats, UNKNOWN_TIER,
UNKNOWN_TIER_DIAGNOSTIC_BYTE_CAP, UNKNOWN_TIER_DIAGNOSTIC_ENTRY_CAP, UnknownTierStats, hash_path, prefix_usage_in_cache,
UNKNOWN_TIER_DIAGNOSTIC_BYTE_CAP, UNKNOWN_TIER_DIAGNOSTIC_ENTRY_CAP, UnknownTierStats, clone_cache_value, hash_path,
prefix_usage_in_cache,
};
use rustfs_heal_contracts::heal_channel::HealScanMode;
use rustfs_utils::path::{SLASH_SEPARATOR, path_join_buf};
@@ -793,7 +794,7 @@ impl DataUsageCache {
let known_tiers = tier_names.iter().map(String::as_str).collect::<HashSet<_>>();
for entry in self.cache.values_mut() {
let Some(tiers) = entry.all_tier_stats.as_mut() else { continue };
let existing_unknown = tiers.tiers.get(UNKNOWN_TIER).cloned().unwrap_or_default();
let existing_unknown = tiers.tiers.get(UNKNOWN_TIER).map(clone_cache_value).unwrap_or_default();
let companion_present = entry.unknown_tier_stats.as_ref().is_some_and(|stats| !stats.is_empty());
let migrate_existing_unknown = !companion_present;
let mut retired = TierStats::default();
@@ -970,7 +971,7 @@ impl DataUsageCache {
|| self.validated_raw_enumeration_cursor().is_some()
|| self.validated_raw_enumeration_page_index().is_some());
let handoff_frontier = cross_epoch_checkpoint
.then(|| self.validated_scan_frontier().map(str::to_owned))
.then(|| self.validated_scan_frontier().map(clone_cache_value))
.flatten();
let reusable = identity.is_valid()
&& name != DATA_USAGE_ROOT
@@ -1256,18 +1257,18 @@ impl DataUsageCache {
let frontier = frontier.filter(|path| path.len() <= 16 * 1024 && self.find(path).is_some());
self.info.scan_coverage_receipt = match frontier {
Some(through) => Some(DataUsageScanCoverageReceipt {
through: through.to_owned(),
through: clone_cache_value(through),
digest: self.coverage_prefix_digest(through)?,
}),
None => None,
};
self.info.scan_resume_after = frontier.map(str::to_owned);
self.info.scan_resume_after = frontier.map(clone_cache_value);
let reason = self
.info
.scan_checkpoint
.as_ref()
.map_or(DataUsageScanCheckpointReason::Unknown, |checkpoint| checkpoint.reason);
self.info.scan_checkpoint = frontier.map(|through| DataUsageScanCheckpoint::new(through.to_owned(), reason));
self.info.scan_checkpoint = frontier.map(|through| DataUsageScanCheckpoint::new(clone_cache_value(through), reason));
Ok(())
}
@@ -1303,26 +1304,32 @@ impl DataUsageCache {
pub fn replace(&mut self, path: &str, parent: &str, e: DataUsageEntry) {
let hash = hash_path(path);
self.cache.insert(hash.key(), e);
self.cache.insert(clone_cache_value(&hash.0), e);
if !parent.is_empty() {
let parent_hash = hash_path(parent);
self.cache.entry(parent_hash.key()).or_default().add_child(&hash);
self.cache
.entry(clone_cache_value(&parent_hash.0))
.or_default()
.add_child(&hash);
}
}
pub fn replace_hashed(&mut self, hash: &DataUsageHash, parent: &Option<DataUsageHash>, e: &DataUsageEntry) {
self.cache.insert(hash.key(), e.clone());
self.cache.insert(clone_cache_value(&hash.0), clone_cache_value(e));
if let Some(parent) = parent {
self.cache.entry(parent.key()).or_default().add_child(hash);
self.cache.entry(clone_cache_value(&parent.0)).or_default().add_child(hash);
}
}
pub fn find(&self, path: &str) -> Option<&DataUsageEntry> {
self.cache.get(&hash_path(path).key())
self.cache.get(&clone_cache_value(&hash_path(path).0))
}
pub fn find_children_copy(&self, h: &DataUsageHash) -> DataUsageHashMap {
self.cache.get(&h.0).map(|entry| entry.children.clone()).unwrap_or_default()
self.cache
.get(&h.0)
.map(|entry| clone_cache_value(&entry.children))
.unwrap_or_default()
}
pub fn flatten(&self, root: &DataUsageEntry) -> DataUsageEntry {
@@ -1346,7 +1353,7 @@ impl DataUsageCache {
}
let (entry, visited) = self.checked_flatten_inner(path)?;
let root_parent_only = {
let path_key = hash_path(path).key();
let path_key = clone_cache_value(&hash_path(path).0);
self.cache
.get(DATA_USAGE_ROOT)
.is_some_and(|root| root_is_parent_only(root, &path_key))
@@ -1379,7 +1386,7 @@ impl DataUsageCache {
}
fn checked_flatten_inner(&self, path: &str) -> Option<(DataUsageEntry, usize)> {
let root_key = hash_path(path).key();
let root_key = clone_cache_value(&hash_path(path).0);
let (root_key, root) = self.cache.get_key_value(&root_key)?;
if root.compacted && !root.children.is_empty() {
return None;
@@ -1411,18 +1418,18 @@ impl DataUsageCache {
}
fn flatten_with_guard(&self, root: &DataUsageEntry, visited: &mut HashSet<String>, depth: usize) -> DataUsageEntry {
let mut root = root.clone();
let mut root = clone_cache_value(root);
if depth >= MAX_DATA_USAGE_CACHE_DEPTH {
root.children.clear();
return root;
}
for id in root.children.clone().iter() {
if !visited.insert(id.clone()) {
for id in clone_cache_value(&root.children).iter() {
if !visited.insert(clone_cache_value(id)) {
continue;
}
if let Some(e) = self.cache.get(id) {
let mut e = e.clone();
let mut e = clone_cache_value(e);
if !e.children.is_empty() {
e = self.flatten_with_guard(&e, visited, depth + 1);
}
@@ -1446,22 +1453,28 @@ impl DataUsageCache {
visited: &mut HashSet<String>,
depth: usize,
) {
if !visited.insert(hash.key()) {
if !visited.insert(clone_cache_value(&hash.0)) {
return;
}
if let Some(e) = src.cache.get(&hash.string()) {
self.cache.insert(hash.key(), e.clone());
if let Some(e) = src.cache.get(&clone_cache_value(&hash.0)) {
self.cache.insert(clone_cache_value(&hash.0), clone_cache_value(e));
if depth < MAX_DATA_USAGE_CACHE_DEPTH {
for ch in e.children.iter() {
if *ch == hash.key() {
if *ch == clone_cache_value(&hash.0) {
continue;
}
self.copy_with_children_guard(src, &DataUsageHash(ch.to_string()), &Some(hash.clone()), visited, depth + 1);
self.copy_with_children_guard(
src,
&DataUsageHash(clone_cache_value(ch)),
&Some(clone_cache_value(hash)),
visited,
depth + 1,
);
}
}
if let Some(parent) = parent {
self.cache.entry(parent.key()).or_default().add_child(hash);
self.cache.entry(clone_cache_value(&parent.0)).or_default().add_child(hash);
}
}
}
@@ -1472,17 +1485,17 @@ impl DataUsageCache {
}
fn delete_recursive_guard(&mut self, hash: &DataUsageHash, visited: &mut HashSet<String>, depth: usize) {
if !visited.insert(hash.key()) {
if !visited.insert(clone_cache_value(&hash.0)) {
return;
}
let mut need_remove = Vec::new();
if let Some(v) = self.cache.get(&hash.string()) {
if let Some(v) = self.cache.get(&clone_cache_value(&hash.0)) {
for child in v.children.iter() {
need_remove.push(child.clone());
need_remove.push(clone_cache_value(child));
}
}
self.cache.remove(&hash.string());
self.cache.remove(&clone_cache_value(&hash.0));
if depth >= MAX_DATA_USAGE_CACHE_DEPTH {
return;
}
@@ -1495,10 +1508,10 @@ impl DataUsageCache {
match self.find(path) {
Some(root) => {
if root.children.is_empty() {
return Some(root.clone());
return Some(clone_cache_value(root));
}
let mut visited = HashSet::new();
visited.insert(hash_path(path).key());
visited.insert(clone_cache_value(&hash_path(path).0));
let mut flat = self.flatten_with_guard(root, &mut visited, 0);
if flat.replication_stats.as_ref().is_some_and(|stats| stats.is_empty()) {
flat.replication_stats = None;
@@ -1510,7 +1523,7 @@ impl DataUsageCache {
}
pub fn search_parent(&self, hash: &DataUsageHash) -> Option<DataUsageHash> {
let want = hash.key();
let want = clone_cache_value(&hash.0);
if let Some(last_index) = want.rfind('/')
&& let Some(v) = self.find(&want[0..last_index])
&& v.children.contains(&want)
@@ -1520,22 +1533,22 @@ impl DataUsageCache {
for (k, v) in self.cache.iter() {
if v.children.contains(&want) {
return Some(DataUsageHash(k.clone()));
return Some(DataUsageHash(clone_cache_value(k)));
}
}
None
}
pub fn is_compacted(&self, hash: &DataUsageHash) -> bool {
self.cache.get(&hash.key()).is_some_and(|due| due.compacted)
self.cache.get(&clone_cache_value(&hash.0)).is_some_and(|due| due.compacted)
}
pub fn force_compact(&mut self, limit: usize) {
if self.cache.len() < limit {
return;
}
let top = hash_path(&self.info.name).key();
let Some(top_e) = self.find(&top).cloned() else {
let top = clone_cache_value(&hash_path(&self.info.name).0);
let Some(top_e) = self.find(&top).map(clone_cache_value) else {
return;
};
@@ -1553,7 +1566,7 @@ impl DataUsageCache {
}
pub fn reduce_children_of(&mut self, path: &DataUsageHash, limit: usize, compact_self: bool) {
let Some(e) = self.cache.get(&path.key()).cloned() else {
let Some(e) = self.cache.get(&clone_cache_value(&path.0)).map(clone_cache_value) else {
return;
};
@@ -1562,14 +1575,14 @@ impl DataUsageCache {
}
if e.children.len() > limit && compact_self {
let mut flat = self.size_recursive(&path.key()).unwrap_or_default();
let mut flat = self.size_recursive(&clone_cache_value(&path.0)).unwrap_or_default();
flat.compacted = true;
self.delete_recursive(path);
self.replace_hashed(path, &None, &flat);
return;
}
let total = self.total_children_rec(&path.key());
let total = self.total_children_rec(&clone_cache_value(&path.0));
if total < limit {
return;
}
@@ -1582,12 +1595,12 @@ impl DataUsageCache {
let mut candidate_index = 0;
while remove > 0 && candidate_index < candidates.len() {
let e = &candidates[candidate_index];
let candidate = e.path.clone();
let candidate = clone_cache_value(&e.path);
if candidate == *path && !compact_self {
break;
}
let removing = self.total_children_rec(&candidate.key());
let mut flat = match self.size_recursive(&candidate.key()) {
let removing = self.total_children_rec(&clone_cache_value(&candidate.0));
let mut flat = match self.size_recursive(&clone_cache_value(&candidate.0)) {
Some(flat) => flat,
None => {
candidate_index += 1;
@@ -1606,7 +1619,7 @@ impl DataUsageCache {
pub fn total_children_rec(&self, path: &str) -> usize {
let mut visited = HashSet::new();
visited.insert(hash_path(path).key());
visited.insert(clone_cache_value(&hash_path(path).0));
self.total_children_rec_guard(path, &mut visited, 0)
}
@@ -1620,7 +1633,7 @@ impl DataUsageCache {
let mut n = 0;
for ch in root.children.iter() {
if visited.insert(ch.clone()) {
if visited.insert(clone_cache_value(ch)) {
n += 1 + self.total_children_rec_guard(ch, visited, depth + 1);
}
}
@@ -1632,7 +1645,7 @@ impl DataUsageCache {
if o.root().is_none() {
return;
}
*self = o.clone();
*self = clone_cache_value(o);
return;
};
@@ -1645,7 +1658,8 @@ impl DataUsageCache {
}
existing_root.merge(&other_root);
self.cache.insert(hash_path(&self.info.name).key(), existing_root);
self.cache
.insert(clone_cache_value(&hash_path(&self.info.name).0), existing_root);
let root_hash = self.root_hash();
for key in other_root.children.iter() {
@@ -1656,7 +1670,7 @@ impl DataUsageCache {
if let Some(existing) = self.cache.get_mut(key) {
existing.merge(&flat);
} else {
self.replace_hashed(&DataUsageHash(key.clone()), &Some(root_hash.clone()), &flat);
self.replace_hashed(&DataUsageHash(clone_cache_value(key)), &Some(clone_cache_value(&root_hash)), &flat);
}
}
}
@@ -1666,7 +1680,7 @@ impl DataUsageCache {
}
pub fn root(&self) -> Option<DataUsageEntry> {
self.find(&self.info.name).cloned()
self.find(&self.info.name).map(clone_cache_value)
}
/// Convert cache to DataUsageInfo for a specific path
@@ -1700,7 +1714,7 @@ impl DataUsageCache {
for (arn, stat) in rs.targets.iter() {
bui.replication_info.insert(
arn.clone(),
clone_cache_value(arn),
BucketTargetUsageInfo {
replication_pending_size: stat.pending_size,
replicated_size: stat.replicated_size,
@@ -1713,7 +1727,7 @@ impl DataUsageCache {
);
}
}
buckets_usage.insert(bucket_name.clone(), bui);
buckets_usage.insert(clone_cache_value(bucket_name), bui);
}
DataUsageInfo {
@@ -1752,7 +1766,7 @@ struct Inner {
fn add(data_usage_cache: &DataUsageCache, path: &DataUsageHash, candidates: &mut Vec<Inner>) -> usize {
let mut visited = HashSet::new();
visited.insert(path.key());
visited.insert(clone_cache_value(&path.0));
add_with_guard(data_usage_cache, path, candidates, &mut visited, 0)
}
@@ -1763,15 +1777,16 @@ fn add_with_guard(
visited: &mut HashSet<String>,
depth: usize,
) -> usize {
let e = match data_usage_cache.cache.get(&path.key()) {
let e = match data_usage_cache.cache.get(&clone_cache_value(&path.0)) {
Some(e) => e,
None => return 0,
};
let mut objects = e.objects;
if depth < MAX_DATA_USAGE_CACHE_DEPTH {
for ch in e.children.iter() {
if visited.insert(ch.clone()) {
objects += add_with_guard(data_usage_cache, &DataUsageHash(ch.clone()), candidates, visited, depth + 1);
if visited.insert(clone_cache_value(ch)) {
objects +=
add_with_guard(data_usage_cache, &DataUsageHash(clone_cache_value(ch)), candidates, visited, depth + 1);
}
}
}
@@ -1781,7 +1796,7 @@ fn add_with_guard(
if !e.children.is_empty() {
candidates.push(Inner {
objects,
path: path.clone(),
path: clone_cache_value(path),
});
}
objects
@@ -1797,7 +1812,7 @@ fn mark_with_depth(duc: &DataUsageCache, entry: &DataUsageEntry, found: &mut Has
}
for k in entry.children.iter() {
if !found.insert(k.to_string()) {
if !found.insert(clone_cache_value(k)) {
continue;
}
if let Some(ch) = duc.cache.get(k) {
+4
View File
@@ -48,6 +48,10 @@ pub enum ScannerError {
#[error("Remote namespace scanner disk is already active")]
RemoteDiskBusy,
/// Node resource pressure deferred work before its first storage poll.
#[error("Scanner node resource admission deferred")]
ResourceDeferred,
/// Partial data usage cache produced before the scanner stopped.
#[error("Scanner stopped with partial data usage cache")]
PartialCache(Box<DataUsageCache>),
+3 -2
View File
@@ -85,7 +85,8 @@ pub use prefix_usage::{BucketPrefixUsageResponse, bucket_prefix_usage, invalidat
pub use remote_scanner::{
NS_SCANNER_MAX_REQUEST_BODY_SIZE, RemoteScannerAdmission, RemoteScannerRequest, admit_remote_scanner_request,
claim_remote_scanner_request, decode_remote_scanner_request, preflight_remote_scanner_request,
remote_scanner_request_matches_envelope, serve_remote_scanner_request, validate_remote_scanner_request_fence,
remote_scanner_request_matches_envelope, serve_remote_scanner_request, serve_remote_scanner_request_with_memory,
validate_remote_scanner_request_fence,
};
pub use runtime_config::{apply_scanner_runtime_config, scanner_runtime_config_status, validate_scanner_runtime_config};
pub use rustfs_scanner_metrics::last_minute;
@@ -123,7 +124,7 @@ pub use storage_api::scan::{
SCANNER_DIRTY_USAGE_SNAPSHOT_PROTOCOL_VERSION, SCANNER_DIRTY_USAGE_SNAPSHOT_RPC_MAX_MESSAGE_SIZE,
SCANNER_SCOPED_DIRTY_USAGE_ACK_MAX_ENTRIES,
};
pub use workload_admission::set_scanner_workload_admission_snapshot_provider;
pub use workload_admission::{set_scanner_node_admission, set_scanner_workload_admission_snapshot_provider};
static SCANNER_ACTIVE_WORK_UNITS: AtomicU64 = AtomicU64::new(0);
static SCANNER_RUNTIME_INSTANCES: AtomicU64 = AtomicU64::new(0);
+17 -12
View File
@@ -12,6 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use crate::data_usage_define::clone_cache_value;
use serde::{Deserialize, Serialize};
use sha2::{Digest as _, Sha256};
use thiserror::Error;
@@ -226,7 +227,7 @@ impl RawEnumerationPageIndex {
return Err(RawEnumerationPageIndexError::IdentityMismatch);
}
let mut indexed_entries = committed_entries.clone();
let mut indexed_entries = clone_cache_value(&committed_entries);
if let Some(building) = &inner.building {
building.validate(
u64::try_from(inner.pages.len()).unwrap_or(u64::MAX),
@@ -236,7 +237,7 @@ impl RawEnumerationPageIndex {
if source_complete && !entries_contain_all(&entries, &building.entries) {
return Err(RawEnumerationPageIndexError::IdentityMismatch);
}
indexed_entries.extend(building.entries.iter().cloned());
indexed_entries.extend(building.entries.iter().map(clone_cache_value));
}
let indexed_entries = normalize_owner_entries(indexed_entries)?;
if source_complete && entry_sets_match(&entries, &indexed_entries) && inner.building.is_none() {
@@ -266,7 +267,7 @@ impl RawEnumerationPageIndex {
.iter()
.filter(|entry| indexed_entries.binary_search(entry).is_err())
.take(max_new_entries.min(remaining_page_slots))
.cloned()
.map(clone_cache_value)
.collect::<Vec<_>>();
let append_count = append_entries.len();
if append_count == 0 {
@@ -282,7 +283,7 @@ impl RawEnumerationPageIndex {
building.entries.dedup();
let source_entries_indexed = source_complete && {
let mut indexed_after_append = indexed_entries;
indexed_after_append.extend(building.entries.iter().cloned());
indexed_after_append.extend(building.entries.iter().map(clone_cache_value));
let indexed_after_append = normalize_owner_entries(indexed_after_append)?;
entry_sets_match(&indexed_after_append, &entries)
};
@@ -311,7 +312,11 @@ impl RawEnumerationPageIndex {
}
let entries_start = inner.validated_committed_entries()?.len();
inner.commit_building_page(entries_start)?;
inner.pages.last().cloned().ok_or(RawEnumerationPageIndexError::EmptyCommit)
inner
.pages
.last()
.map(clone_cache_value)
.ok_or(RawEnumerationPageIndexError::EmptyCommit)
}
pub fn page(&self, page_index: u64, expected_digest: [u8; 32]) -> Result<&RawEnumerationPage, RawEnumerationPageIndexError> {
@@ -377,7 +382,7 @@ impl RawEnumerationPageIndexInner {
if let Some(building) = &self.building {
RawEnumerationPageOwnerStatus::Building {
generation: self.generation,
parent: self.parent.clone(),
parent: clone_cache_value(&self.parent),
page_index: building.page_index,
indexed_entries,
buffered_entries: building.entries.len(),
@@ -385,7 +390,7 @@ impl RawEnumerationPageIndexInner {
} else {
RawEnumerationPageOwnerStatus::Ready {
generation: self.generation,
parent: self.parent.clone(),
parent: clone_cache_value(&self.parent),
committed_pages: self.pages.len(),
indexed_entries,
complete: self.complete,
@@ -404,12 +409,12 @@ impl RawEnumerationPageIndexInner {
if page.terminal && page_index + 1 != self.pages.len() {
return Err(RawEnumerationPageIndexError::CorruptIndex);
}
entries.extend(page.entries.iter().cloned());
entries.extend(page.entries.iter().map(clone_cache_value));
}
if self.complete && self.pages.last().is_some_and(|page| !page.terminal) {
return Err(RawEnumerationPageIndexError::CorruptIndex);
}
let mut unique_entries = entries.clone();
let mut unique_entries = clone_cache_value(&entries);
unique_entries.sort();
unique_entries.dedup();
if unique_entries.len() != entries.len() {
@@ -426,7 +431,7 @@ impl RawEnumerationPageIndexInner {
u64::try_from(entries.len()).unwrap_or(u64::MAX),
self.page_entry_limit,
)?;
entries.extend(building.entries.iter().cloned());
entries.extend(building.entries.iter().map(clone_cache_value));
}
Ok(entries)
}
@@ -484,7 +489,7 @@ impl RawEnumerationPage {
let digest = raw_page_digest(parent, &building);
Self {
version: RAW_PAGE_INDEX_VERSION,
parent: parent.to_string(),
parent: clone_cache_value(parent),
page_index: building.page_index,
entries_start: building.entries_start,
entries: building.entries,
@@ -497,7 +502,7 @@ impl RawEnumerationPage {
let builder = RawEnumerationPageBuilder {
page_index: self.page_index,
entries_start: self.entries_start,
entries: self.entries.clone(),
entries: clone_cache_value(&self.entries),
terminal: self.terminal,
};
if self.version != RAW_PAGE_INDEX_VERSION
+3 -2
View File
@@ -16,6 +16,7 @@ use super::{
RawEnumerationPageBuilder, RawEnumerationPageIndex, RawEnumerationPageIndexError, RawEnumerationPageIndexInner,
RawEnumerationPageIndexState, owner_entry_is_valid,
};
use crate::data_usage_define::clone_cache_value;
use std::collections::{BTreeSet, HashSet};
/// Exclusive, validated owner of an index during one directory enumeration.
@@ -111,7 +112,7 @@ impl RawEnumerationPageWriter {
entries: Vec::new(),
terminal: false,
});
self.indexed.insert(entry.clone());
self.indexed.insert(clone_cache_value(&entry));
// The building page is normalized (strictly increasing) on every
// path that reaches here, so a positional insert keeps it so.
let position = building.entries.binary_search(&entry).unwrap_or_else(|position| position);
@@ -137,7 +138,7 @@ impl RawEnumerationPageWriter {
}
pub(crate) fn checkpoint(&self) -> Result<RawEnumerationPageIndex, RawEnumerationPageIndexError> {
let mut inner = self.inner.clone();
let mut inner = clone_cache_value(&self.inner);
if inner.building.is_some() {
inner.commit_building_page(self.committed_entries)?;
}
+3 -1
View File
@@ -55,7 +55,9 @@ mod stream;
#[cfg(test)]
pub(crate) use stream::checkpoint_fixture_partial_return;
pub use stream::{RemoteScannerAdmission, RemoteScannerRequest, serve_remote_scanner_request};
pub use stream::{
RemoteScannerAdmission, RemoteScannerRequest, serve_remote_scanner_request, serve_remote_scanner_request_with_memory,
};
pub(crate) use stream::{RemoteScannerOutcome, RemoteScannerScanSpec, scan_remote_bucket};
use stream::{RemoteScannerReplayCache, RemoteScannerRequestWire, RemoteScannerValidatedCycle};
+344 -29
View File
@@ -21,7 +21,7 @@ use crate::scanner_io::{
cache_root_entry_info, current_cache_root_or_prepare_with_generation, log_saved_partial_bucket_cache,
persist_scanner_checkpoint, scanner_set_disk_inventory,
};
use crate::storage_api::owner::NS_SCANNER_PROTOCOL_VERSION;
use crate::storage_api::owner::{EcstoreBackgroundEncodeBuffer, EcstoreBackgroundMemoryBudget, NS_SCANNER_PROTOCOL_VERSION};
use crate::{
DATA_USAGE_CACHE_NAME, DataUsageCache, DataUsageCachePrepareOutcome, DataUsageCacheSource, DataUsageEntryInfo,
DataUsageScanPlanDigest, Disk, ScannerError, StorageError, resolve_scanner_object_store_handle,
@@ -531,10 +531,23 @@ fn update_frame_mac(mac: &mut HmacSha256, request_id: Uuid, sequence: u64, paylo
}
pub async fn serve_remote_scanner_request<W>(
disk: Arc<Disk>,
request: RemoteScannerRequest,
writer: W,
disconnect: CancellationToken,
) -> Result<(), ScannerError>
where
W: AsyncWrite + Unpin + Send,
{
serve_remote_scanner_request_with_memory(disk, request, writer, disconnect, None).await
}
pub async fn serve_remote_scanner_request_with_memory<W>(
disk: Arc<Disk>,
request: RemoteScannerRequest,
mut writer: W,
disconnect: CancellationToken,
memory: Option<Arc<EcstoreBackgroundMemoryBudget>>,
) -> Result<(), ScannerError>
where
W: AsyncWrite + Unpin + Send,
@@ -577,7 +590,7 @@ where
Ok(result) => result,
Err(err) => RemoteScannerFrameResult::Error(err.into_frame()),
};
write_frame_bounded(
write_frame_bounded_with_memory(
&mut writer,
&authenticator,
&mut sequence,
@@ -588,6 +601,7 @@ where
),
&disconnect,
NS_SCANNER_WRITE_TIMEOUT,
memory.as_ref(),
)
.await?;
shutdown_writer_bounded(&mut writer, &disconnect, NS_SCANNER_WRITE_TIMEOUT).await?;
@@ -605,7 +619,7 @@ where
}
_ = phase_changed.notified(), if !persistence_announced => {
persistence_announced = true;
if let Err(err) = write_frame_bounded(
if let Err(err) = write_frame_bounded_with_memory(
&mut writer,
&authenticator,
&mut sequence,
@@ -616,6 +630,7 @@ where
),
&disconnect,
NS_SCANNER_WRITE_TIMEOUT,
memory.as_ref(),
)
.await
{
@@ -630,7 +645,7 @@ where
}
}
_ = heartbeat.tick() => {
if let Err(err) = write_frame_bounded(
if let Err(err) = write_frame_bounded_with_memory(
&mut writer,
&authenticator,
&mut sequence,
@@ -641,6 +656,7 @@ where
),
&disconnect,
NS_SCANNER_WRITE_TIMEOUT,
memory.as_ref(),
)
.await
{
@@ -705,7 +721,7 @@ async fn scan_and_persist_local_bucket(
RemoteScannerServerError::worker("remote namespace scanner cache publication is blocked by data movement")
})?;
let cache_name = path_join_buf(&[&bucket, DATA_USAGE_CACHE_NAME]);
let guard = acquire_scanner_cache_locks(set.as_ref(), &cache_name, source)
let mut guard = acquire_scanner_cache_locks(set.as_ref(), &cache_name, source)
.await
.map_err(|err| {
if err.is_contention() {
@@ -791,19 +807,23 @@ async fn scan_and_persist_local_bucket(
let set_disks = scanner_set_disk_inventory(set.as_ref()).await;
let (checkpoint_tx, mut checkpoint_rx) = mpsc::channel::<DataUsageCache>(1);
let (node_tx, mut node_rx) = crate::workload_admission::scanner_node_request_channel();
let scan_ctx = ctx.child_token();
let scan = ScannerIODisk::nsscanner_disk(
disk.clone(),
scan_ctx.clone(),
budget,
set_disks,
cache,
None,
ScannerDiskScanOptions {
scan_mode,
prefix_scan_scope: None,
checkpoint_tx: Some(checkpoint_tx),
},
let scan = crate::workload_admission::with_scanner_node_requests(
node_tx,
ScannerIODisk::nsscanner_disk(
disk.clone(),
scan_ctx.clone(),
budget,
set_disks,
cache,
None,
ScannerDiskScanOptions {
scan_mode,
prefix_scan_scope: None,
checkpoint_tx: Some(checkpoint_tx),
},
),
);
tokio::pin!(scan);
let fence_watch = watch_remote_scanner_request_fence(next_cycle, leader_epoch, store.clone(), NS_SCANNER_FENCE_POLL_INTERVAL);
@@ -813,6 +833,27 @@ async fn scan_and_persist_local_bucket(
let mut checkpoint_channel_closed = false;
let outcome = loop {
tokio::select! {
request = async {
match node_rx.as_mut() {
Some(receiver) => receiver.recv().await,
None => std::future::pending().await,
}
} => {
let Some(request) = request else {
node_rx = None;
continue;
};
if let Err(error) = wait_for_remote_node_resumption(
guard.resume_node_operation(set.clone(), &cache_name, source, &revisions, request),
fence_watch.as_mut(),
).await {
scan_ctx.cancel();
let _ = tokio::time::timeout(NS_SCANNER_LOCK_LOSS_SHUTDOWN_TIMEOUT, scan.as_mut()).await;
return Err(RemoteScannerServerError::retry_bucket(format!(
"remote namespace scanner could not resume node admission: {error}",
)));
}
}
result = &mut scan => {
match result {
Ok(outcome) => break outcome,
@@ -1121,6 +1162,22 @@ pub(crate) async fn scan_remote_bucket(
result
}
/// Resource waits must keep checking the persisted leader fence.
async fn wait_for_remote_node_resumption<F, G>(resume: F, fence_watch: Pin<&mut G>) -> Result<(), StorageError>
where
F: Future<Output = Result<(), StorageError>>,
G: Future<Output = Result<(), ScannerError>>,
{
tokio::select! {
biased;
result = fence_watch => Err(match result {
Ok(()) => StorageError::other("remote namespace scanner fence watcher stopped unexpectedly"),
Err(error) => StorageError::other(error),
}),
result = resume => result,
}
}
fn finish_remote_scanner_stream(
result: RemoteScannerStreamResult<RemoteScannerOutcome>,
budget: &ScannerCycleBudget,
@@ -1167,7 +1224,7 @@ pub(crate) async fn checkpoint_fixture_partial_return(progress: (u64, u64), entr
)
.await
.expect("checkpoint partial frame must encode");
let frame = read_frame(&mut std::io::Cursor::new(bytes.as_slice()), &reader_auth, &mut 0)
let frame = read_frame(&mut std::io::Cursor::new(bytes.as_slice()), &reader_auth, &mut 0, None)
.await
.expect("checkpoint progress frame must authenticate");
assert_eq!(frame.progress.entries_visited, entries_visited);
@@ -1229,6 +1286,9 @@ async fn consume_remote_scanner_stream_until<R>(
where
R: AsyncRead + Unpin,
{
// Namespace consumption is background work even in a plain spawned worker.
// Capture only its installed local memory policy, never scheduler receipts.
let memory = EcstoreBackgroundMemoryBudget::current();
let mut expected_sequence = 0_u64;
let mut last_progress = RemoteScannerProgress::default();
let mut last_phase = RemoteScannerPhase::Scanning;
@@ -1254,7 +1314,7 @@ where
}
result = tokio::time::timeout_at(
read_deadline,
read_frame(&mut reader, &authenticator, &mut expected_sequence),
read_frame(&mut reader, &authenticator, &mut expected_sequence, memory.as_ref()),
) => match result {
Ok(Ok(frame)) => frame,
Ok(Err(err)) => {
@@ -1452,6 +1512,7 @@ fn writer_timeout_error(operation: &str, timeout_duration: Duration) -> ScannerE
))
}
#[cfg(test)]
async fn write_frame_bounded<W>(
writer: &mut W,
authenticator: &FrameAuthenticator,
@@ -1460,13 +1521,33 @@ async fn write_frame_bounded<W>(
disconnect: &CancellationToken,
timeout_duration: Duration,
) -> Result<(), ScannerError>
where
W: AsyncWrite + Unpin,
{
write_frame_bounded_with_memory(writer, authenticator, sequence, frame, disconnect, timeout_duration, None).await
}
async fn write_frame_bounded_with_memory<W>(
writer: &mut W,
authenticator: &FrameAuthenticator,
sequence: &mut u64,
frame: &RemoteScannerFrame,
disconnect: &CancellationToken,
timeout_duration: Duration,
memory: Option<&Arc<EcstoreBackgroundMemoryBudget>>,
) -> Result<(), ScannerError>
where
W: AsyncWrite + Unpin,
{
tokio::select! {
biased;
_ = disconnect.cancelled() => Err(disconnected_writer_error()),
result = tokio::time::timeout(timeout_duration, write_frame(writer, authenticator, sequence, frame)) => {
result = tokio::time::timeout(timeout_duration, async {
match memory {
Some(memory) => write_background_frame(writer, authenticator, sequence, frame, memory).await,
None => write_frame(writer, authenticator, sequence, frame).await,
}
}) => {
result.map_err(|_| writer_timeout_error("write", timeout_duration))?
}
}
@@ -1546,10 +1627,215 @@ where
Ok(())
}
#[derive(Serialize)]
struct BorrowedRemoteScannerFrameEnvelope<'a> {
version: u16,
sequence: u64,
payload: &'a [u8],
mac: &'a [u8],
}
fn encode_background<T: Serialize>(value: &T, memory: &Arc<EcstoreBackgroundMemoryBudget>) -> Result<bytes::Bytes, ScannerError> {
let mut buffer = EcstoreBackgroundEncodeBuffer::new(memory.clone(), NS_SCANNER_MAX_FRAME_SIZE);
// Preserve the serializer's original I/O error and its memory-denial cause.
rmp_serde::encode::write_named(&mut buffer, value).map_err(|error| match error {
rmp_serde::encode::Error::InvalidValueWrite(error) => ScannerError::Io(error.into()),
error => ScannerError::Io(IoError::other(error)),
})?;
Ok(buffer.into_bytes())
}
async fn write_background_frame<W>(
writer: &mut W,
authenticator: &FrameAuthenticator,
sequence: &mut u64,
frame: &RemoteScannerFrame,
memory: &Arc<EcstoreBackgroundMemoryBudget>,
) -> Result<(), ScannerError>
where
W: AsyncWrite + Unpin,
{
let payload = encode_background(frame, memory)?;
let mut mac = HmacSha256::new_from_slice(authenticator.secret.as_bytes())
.map_err(|_| ScannerError::Other("invalid remote namespace scanner authentication key".to_string()))?;
update_frame_mac(&mut mac, authenticator.request_id, *sequence, &payload);
let signature = mac.finalize().into_bytes();
let envelope = BorrowedRemoteScannerFrameEnvelope {
version: NS_SCANNER_PROTOCOL_VERSION,
sequence: *sequence,
payload: &payload,
mac: &signature,
};
let encoded = encode_background(&envelope, memory)?;
if encoded.is_empty() {
return Err(ScannerError::Other("remote namespace scanner encoded frame is empty".to_string()));
}
let len = u32::try_from(encoded.len())
.map_err(|_| ScannerError::Other("remote namespace scanner frame is too large".to_string()))?;
writer.write_all(&len.to_be_bytes()).await?;
writer.write_all(&encoded).await?;
drop(encoded);
drop(payload);
*sequence = sequence
.checked_add(1)
.ok_or_else(|| ScannerError::Other("remote namespace scanner frame sequence overflow".to_string()))?;
Ok(())
}
// These owners account only envelope payload and MAC storage, not the decoded model.
struct BackgroundRemoteScannerEnvelope {
version: u16,
sequence: u64,
payload: bytes::Bytes,
mac: bytes::Bytes,
}
#[derive(Deserialize)]
#[serde(field_identifier, rename_all = "lowercase")]
enum RemoteScannerEnvelopeField {
Version,
Sequence,
Payload,
Mac,
}
struct BackgroundEnvelopeSeed<'a> {
memory: &'a Arc<EcstoreBackgroundMemoryBudget>,
maximum: usize,
failure: &'a mut Option<IoError>,
}
impl BackgroundEnvelopeSeed<'_> {
fn bytes(&mut self) -> BackgroundEnvelopeBytesSeed<'_> {
BackgroundEnvelopeBytesSeed {
buffer: EcstoreBackgroundEncodeBuffer::new(self.memory.clone(), self.maximum),
failure: self.failure,
}
}
}
impl<'de> serde::de::DeserializeSeed<'de> for BackgroundEnvelopeSeed<'_> {
type Value = BackgroundRemoteScannerEnvelope;
fn deserialize<D: serde::Deserializer<'de>>(self, deserializer: D) -> Result<Self::Value, D::Error> {
deserializer.deserialize_struct("RemoteScannerFrameEnvelope", &["version", "sequence", "payload", "mac"], self)
}
}
impl<'de> serde::de::Visitor<'de> for BackgroundEnvelopeSeed<'_> {
type Value = BackgroundRemoteScannerEnvelope;
fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str("remote namespace scanner frame envelope")
}
fn visit_seq<A: serde::de::SeqAccess<'de>>(mut self, mut seq: A) -> Result<Self::Value, A::Error> {
use serde::de::Error;
let version = seq.next_element()?.ok_or_else(|| A::Error::invalid_length(0, &self))?;
let sequence = seq.next_element()?.ok_or_else(|| A::Error::invalid_length(1, &self))?;
let payload = seq
.next_element_seed(self.bytes())?
.ok_or_else(|| A::Error::invalid_length(2, &self))?;
let mac = seq
.next_element_seed(self.bytes())?
.ok_or_else(|| A::Error::invalid_length(3, &self))?;
Ok(BackgroundRemoteScannerEnvelope {
version,
sequence,
payload,
mac,
})
}
fn visit_map<A: serde::de::MapAccess<'de>>(mut self, mut map: A) -> Result<Self::Value, A::Error> {
use serde::de::Error;
let (mut version, mut sequence, mut payload, mut mac) = (None, None, None, None);
while let Some(field) = map.next_key()? {
match field {
RemoteScannerEnvelopeField::Version => {
if version.is_some() {
return Err(A::Error::duplicate_field("version"));
}
version = Some(map.next_value()?);
}
RemoteScannerEnvelopeField::Sequence => {
if sequence.is_some() {
return Err(A::Error::duplicate_field("sequence"));
}
sequence = Some(map.next_value()?);
}
RemoteScannerEnvelopeField::Payload => {
if payload.is_some() {
return Err(A::Error::duplicate_field("payload"));
}
payload = Some(map.next_value_seed(self.bytes())?);
}
RemoteScannerEnvelopeField::Mac => {
if mac.is_some() {
return Err(A::Error::duplicate_field("mac"));
}
mac = Some(map.next_value_seed(self.bytes())?);
}
}
}
Ok(BackgroundRemoteScannerEnvelope {
version: version.ok_or_else(|| A::Error::missing_field("version"))?,
sequence: sequence.ok_or_else(|| A::Error::missing_field("sequence"))?,
payload: payload.ok_or_else(|| A::Error::missing_field("payload"))?,
mac: mac.ok_or_else(|| A::Error::missing_field("mac"))?,
})
}
}
struct BackgroundEnvelopeBytesSeed<'a> {
buffer: EcstoreBackgroundEncodeBuffer,
failure: &'a mut Option<IoError>,
}
impl<'de> serde::de::DeserializeSeed<'de> for BackgroundEnvelopeBytesSeed<'_> {
type Value = bytes::Bytes;
fn deserialize<D: serde::Deserializer<'de>>(self, deserializer: D) -> Result<Self::Value, D::Error> {
// deserialize_seq matches Vec<u8>, including MessagePack bin input.
deserializer.deserialize_seq(self)
}
}
impl<'de> serde::de::Visitor<'de> for BackgroundEnvelopeBytesSeed<'_> {
type Value = bytes::Bytes;
fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str("a sequence of bytes")
}
fn visit_seq<A: serde::de::SeqAccess<'de>>(mut self, mut seq: A) -> Result<Self::Value, A::Error> {
use serde::de::Error;
// Never allocate from an untrusted sequence size_hint.
while let Some(byte) = seq.next_element::<u8>()? {
if let Err(error) = std::io::Write::write_all(&mut self.buffer, &[byte]) {
*self.failure = Some(error);
return Err(A::Error::custom("background envelope byte storage unavailable"));
}
}
Ok(self.buffer.into_bytes())
}
}
fn decode_background_envelope(
encoded: &[u8],
memory: &Arc<EcstoreBackgroundMemoryBudget>,
) -> std::io::Result<BackgroundRemoteScannerEnvelope> {
use serde::de::DeserializeSeed;
let mut failure = None;
let seed = BackgroundEnvelopeSeed {
memory,
maximum: encoded.len(),
failure: &mut failure,
};
seed.deserialize(&mut rmp_serde::Deserializer::from_read_ref(encoded))
.map_err(|_| {
failure.unwrap_or_else(|| IoError::new(ErrorKind::InvalidData, "invalid remote namespace scanner frame envelope"))
})
}
async fn read_frame<R>(
reader: &mut R,
authenticator: &FrameAuthenticator,
expected_sequence: &mut u64,
memory: Option<&Arc<EcstoreBackgroundMemoryBudget>>,
) -> std::io::Result<RemoteScannerFrame>
where
R: AsyncRead + Unpin,
@@ -1571,21 +1857,50 @@ where
));
}
let mut encoded = vec![0_u8; len];
reader.read_exact(&mut encoded).await?;
let envelope: RemoteScannerFrameEnvelope = rmp_serde::from_slice(&encoded)
.map_err(|_| IoError::new(ErrorKind::InvalidData, "invalid remote namespace scanner frame envelope"))?;
if envelope.version != NS_SCANNER_PROTOCOL_VERSION {
let mut background = memory.map(|memory| EcstoreBackgroundEncodeBuffer::new(memory.clone(), len));
let mut foreground;
let encoded = match background.as_mut() {
Some(buffer) => buffer.prepare_read(len)?,
None => {
foreground = vec![0_u8; len];
foreground.as_mut_slice()
}
};
reader.read_exact(encoded).await?;
let background_envelope;
let foreground_envelope;
let (version, sequence, payload, mac) = match memory {
Some(memory) => {
background_envelope = decode_background_envelope(encoded, memory)?;
(
background_envelope.version,
background_envelope.sequence,
background_envelope.payload.as_ref(),
background_envelope.mac.as_ref(),
)
}
None => {
foreground_envelope = rmp_serde::from_slice::<RemoteScannerFrameEnvelope>(encoded)
.map_err(|_| IoError::new(ErrorKind::InvalidData, "invalid remote namespace scanner frame envelope"))?;
(
foreground_envelope.version,
foreground_envelope.sequence,
foreground_envelope.payload.as_slice(),
foreground_envelope.mac.as_slice(),
)
}
};
if version != NS_SCANNER_PROTOCOL_VERSION {
return Err(IoError::new(
ErrorKind::InvalidData,
format!("unsupported remote namespace scanner frame version: {}", envelope.version),
format!("unsupported remote namespace scanner frame version: {}", version),
));
}
if envelope.sequence != *expected_sequence {
if sequence != *expected_sequence {
return Err(IoError::new(ErrorKind::InvalidData, "remote namespace scanner frame sequence is invalid"));
}
authenticator.verify(envelope.sequence, &envelope.payload, &envelope.mac)?;
let frame: RemoteScannerFrame = rmp_serde::from_slice(&envelope.payload)
authenticator.verify(sequence, payload, mac)?;
let frame: RemoteScannerFrame = rmp_serde::from_slice(payload)
.map_err(|_| IoError::new(ErrorKind::InvalidData, "invalid remote namespace scanner frame payload"))?;
*expected_sequence = expected_sequence
.checked_add(1)
@@ -452,6 +452,102 @@ fn request_envelope_must_match_every_pre_body_replay_field() {
));
}
#[tokio::test]
async fn background_frame_matches_foreground_wire_mac_and_sequence() {
let authenticator = FrameAuthenticator::for_test(Uuid::from_u128(7));
let frames = [
RemoteScannerFrame::progress(RemoteScannerProgress::default()),
RemoteScannerFrame::terminal(RemoteScannerProgress::default(), eof_test_terminal(0)),
RemoteScannerFrame::terminal(
RemoteScannerProgress::default(),
RemoteScannerFrameResult::Error(RemoteScannerErrorFrame {
scope: RemoteScannerErrorScope::Worker,
message: "目录/🙂/\0".to_string(),
}),
),
];
for frame in frames {
for initial_sequence in [0, 128] {
let owner = EcstoreBackgroundMemoryBudget::new();
owner.set_sampled_limit(64 * 1024, Duration::from_secs(30));
let (mut foreground, mut default, mut background) = (Vec::new(), Vec::new(), Vec::new());
let (mut foreground_sequence, mut default_sequence, mut background_sequence) =
(initial_sequence, initial_sequence, initial_sequence);
write_frame(&mut foreground, &authenticator, &mut foreground_sequence, &frame)
.await
.expect("foreground frame should encode");
write_frame_bounded(
&mut default,
&authenticator,
&mut default_sequence,
&frame,
&CancellationToken::new(),
Duration::from_secs(1),
)
.await
.expect("default writer should use the foreground encoding");
write_background_frame(&mut background, &authenticator, &mut background_sequence, &frame, &owner)
.await
.expect("background frame should encode");
assert_eq!(foreground, default);
assert_eq!(foreground, background);
assert_eq!(foreground_sequence, initial_sequence + 1);
assert_eq!(default_sequence, foreground_sequence);
assert_eq!(background_sequence, foreground_sequence);
assert_eq!(owner.snapshot().used, 0);
let mut expected_sequence = initial_sequence;
let decoded = read_frame(&mut Cursor::new(background), &authenticator, &mut expected_sequence, None)
.await
.expect("the original decoder must verify the background MAC and sequence");
assert_eq!(expected_sequence, foreground_sequence);
assert_eq!(
rmp_serde::to_vec_named(&decoded).expect("decoded payload"),
rmp_serde::to_vec_named(&frame).expect("original payload")
);
}
}
}
#[tokio::test]
async fn background_frame_payload_and_envelope_denial_preserve_io_cause_without_output() {
let authenticator = FrameAuthenticator::for_test(Uuid::from_u128(7));
let frame = RemoteScannerFrame::progress(RemoteScannerProgress::default());
let probe_owner = EcstoreBackgroundMemoryBudget::new();
probe_owner.set_sampled_limit(64 * 1024, Duration::from_secs(30));
let payload = encode_background(&frame, &probe_owner).expect("measure retained payload capacity");
let payload_capacity = probe_owner.snapshot().used;
assert!(payload_capacity > 0);
drop(payload);
assert_eq!(probe_owner.snapshot().used, 0);
for limit in [0, payload_capacity * 2] {
let owner = EcstoreBackgroundMemoryBudget::new();
owner.set_sampled_limit(limit, Duration::from_secs(30));
if limit != 0 {
let payload = encode_background(&frame, &owner).expect("this budget admits the entire inner encoding");
assert_eq!(owner.snapshot().used, payload_capacity);
drop(payload);
assert_eq!(owner.snapshot().used, 0);
}
let mut output = Vec::new();
let mut sequence = 17;
let error = write_background_frame(&mut output, &authenticator, &mut sequence, &frame, &owner)
.await
.expect_err("deny the payload at zero or the envelope after an admitted payload");
let ScannerError::Io(error) = error else {
panic!("serializer memory denial must remain an I/O error");
};
assert_eq!(error.kind(), ErrorKind::WouldBlock);
assert!(error.get_ref().is_some(), "the serializer must preserve the original memory-denial cause");
assert!(output.is_empty());
assert_eq!(sequence, 17);
assert_eq!(owner.snapshot().used, 0);
}
}
#[tokio::test]
async fn bounded_frame_write_releases_a_stalled_response() {
let request_id = Uuid::new_v4();
@@ -906,6 +1002,58 @@ async fn fence_watcher_stops_work_after_persisted_leader_changes() {
assert!(err.to_string().contains("requested_epoch=9"));
}
#[tokio::test]
async fn node_admission_wait_does_not_hide_replacement_leader() {
use rustfs_concurrency::adaptive::{FeedbackController, FeedbackObservation, NodeAdmission, ResourcePressure};
use std::num::NonZeroU16;
let store = Arc::new(MutableCycleStateStore {
state: RwLock::new(crate::scanner::encode_scanner_cycle_fence_for_test(7, 9)),
});
let cache = Mutex::new(None);
let refresh = AsyncMutex::new(());
let watcher = watch_remote_scanner_request_fence_with_cache(
7,
9,
store.clone(),
Duration::from_millis(5),
&cache,
&refresh,
Duration::from_millis(5),
);
tokio::pin!(watcher);
let gate = NodeAdmission::new(NonZeroU16::new(1).expect("one shared node unit"));
let storage_polled = AtomicBool::new(false);
{
let resume = async {
let _permit = gate.acquire().await;
storage_polled.store(true, Ordering::Relaxed);
Ok(())
};
let waiting = wait_for_remote_node_resumption(resume, watcher.as_mut());
tokio::pin!(waiting);
assert!(futures::poll!(&mut waiting).is_pending(), "no resource sample grants a node unit");
*store.state.write().expect("persist replacement leader") = crate::scanner::encode_scanner_cycle_fence_for_test(7, 10);
let error = tokio::time::timeout(Duration::from_millis(100), waiting)
.await
.expect("resource pause must not hide a changed leader")
.expect_err("old leader cannot resume storage work");
assert!(error.to_string().contains("requested_epoch=9"));
}
assert!(!storage_polled.load(Ordering::Relaxed));
assert_eq!(gate.snapshot().active, 0);
gate.update(
&mut FeedbackController::new(),
Some(FeedbackObservation {
generation: 1,
max_concurrency: 1,
pressure: ResourcePressure::Unknown,
completed_units: None,
}),
);
assert!(gate.try_acquire().is_some(), "failed fencing drops the old node queue ticket");
}
#[tokio::test]
async fn concurrent_cycle_cache_misses_share_one_backend_read() {
let request = RemoteScannerRequest(test_request(Uuid::new_v4()));
@@ -1904,7 +2052,7 @@ async fn wrong_frame_version_and_sequence_are_rejected() {
input.set_position(0);
let mut expected_sequence = 0;
let error = read_frame(&mut input, &auth, &mut expected_sequence)
let error = read_frame(&mut input, &auth, &mut expected_sequence, None)
.await
.expect_err("invalid frame must fail");
assert!(error.to_string().contains(expected_error));
@@ -1996,7 +2144,321 @@ fn oversized_frame_is_rejected_before_allocation() {
let runtime = tokio::runtime::Runtime::new().expect("runtime should start");
let mut sequence = 0;
let error = runtime
.block_on(read_frame(&mut input, &auth, &mut sequence))
.block_on(read_frame(&mut input, &auth, &mut sequence, None))
.expect_err("oversized frame must fail");
assert_eq!(error.kind(), ErrorKind::InvalidData);
}
#[tokio::test]
async fn raw_frame_staging_denial_consumes_only_the_length_prefix() {
let auth = FrameAuthenticator::for_test(Uuid::new_v4());
let frame = RemoteScannerFrame::progress(RemoteScannerProgress::default());
let mut wire = Vec::new();
write_frame(&mut wire, &auth, &mut 0, &frame).await.expect("wire fixture");
let owner = EcstoreBackgroundMemoryBudget::new();
let mut input = Cursor::new(wire.clone());
let mut sequence = 0;
let error = read_frame(&mut input, &auth, &mut sequence, Some(&owner))
.await
.expect_err("closed staging budget");
assert_eq!(error.kind(), ErrorKind::WouldBlock);
assert_eq!(input.position(), 4, "do not read a frame body after staging denial");
assert_eq!(sequence, 0);
assert_eq!(owner.snapshot().used, 0);
let length = wire.len() - 4;
owner.set_sampled_limit(u64::try_from(length).expect("frame length"), Duration::from_secs(30));
let decoded = read_frame(&mut input, &auth, &mut sequence, Some(&owner)).await;
assert!(decoded.is_err(), "denial is terminal; do not resume from the consumed prefix");
let mut input = Cursor::new(wire.clone());
let error = read_frame(&mut input, &auth, &mut sequence, Some(&owner))
.await
.expect_err("raw-only capacity cannot fund an envelope-owned payload");
assert_eq!(error.kind(), ErrorKind::WouldBlock);
assert!(error.get_ref().is_some(), "preserve the original memory-denial cause");
assert_eq!(sequence, 0);
assert_eq!(owner.snapshot().used, 0);
// Each allocation is bounded by length: raw, the retained other field,
// and both old and new field growth allocations can coexist.
owner.set_sampled_limit(u64::try_from(length * 4).expect("growth peak bound"), Duration::from_secs(30));
let mut input = Cursor::new(wire);
read_frame(&mut input, &auth, &mut sequence, Some(&owner))
.await
.expect("a fresh reader with all retained storage and growth funded decodes");
assert_eq!(sequence, 1);
assert_eq!(owner.snapshot().used, 0);
}
#[tokio::test]
async fn raw_frame_staging_failure_keeps_existing_validation_errors() {
let auth = FrameAuthenticator::for_test(Uuid::new_v4());
let mut wire = Vec::new();
write_frame(&mut wire, &auth, &mut 0, &RemoteScannerFrame::progress(RemoteScannerProgress::default()))
.await
.expect("wire fixture");
let owner = EcstoreBackgroundMemoryBudget::new();
owner.set_sampled_limit(4 * 1024 * 1024, Duration::from_secs(30));
let mut truncated = wire.clone();
truncated.pop();
for bytes in [
vec![0, 0, 0],
vec![0, 0, 0, 0],
(2_097_153_u32).to_be_bytes().to_vec(),
truncated,
] {
let original = read_frame(&mut Cursor::new(bytes.clone()), &auth, &mut 0, None)
.await
.expect_err("invalid original frame");
let mut sequence = 0;
let error = read_frame(&mut Cursor::new(bytes), &auth, &mut sequence, Some(&owner))
.await
.expect_err("invalid admitted frame");
assert_eq!(error.kind(), original.kind());
assert_eq!(error.to_string(), original.to_string());
assert_eq!(sequence, 0);
assert_eq!(owner.snapshot().used, 0);
}
}
struct ShrinkDecodeBudgetAfterFrame {
input: Cursor<Vec<u8>>,
first_frame_end: u64,
memory: Arc<EcstoreBackgroundMemoryBudget>,
}
impl AsyncRead for ShrinkDecodeBudgetAfterFrame {
fn poll_read(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
buf: &mut tokio::io::ReadBuf<'_>,
) -> std::task::Poll<std::io::Result<()>> {
if self.input.position() >= self.first_frame_end {
self.memory.set_sampled_limit(0, Duration::from_secs(30));
}
std::pin::Pin::new(&mut self.input).poll_read(cx, buf)
}
}
#[tokio::test]
async fn namespace_decode_denial_keeps_previously_reported_progress() {
const CHILD: &str = "RUSTFS_NAMESPACE_DECODE_PROGRESS_CHILD";
const TEST: &str = "remote_scanner::stream::tests::namespace_decode_denial_keeps_previously_reported_progress";
if std::env::var_os(CHILD).is_none() {
let output = std::process::Command::new(std::env::current_exe().expect("test executable"))
.args(["--exact", TEST, "--nocapture", "--test-threads=1"])
.env(CHILD, "1")
.output()
.expect("isolated installed-policy consumer test");
assert!(
output.status.success(),
"{}{}",
String::from_utf8_lossy(&output.stdout),
String::from_utf8_lossy(&output.stderr)
);
assert!(String::from_utf8_lossy(&output.stdout).contains("NAMESPACE_DECODE_PROGRESS_OK"));
return;
}
let owner = EcstoreBackgroundMemoryBudget::new();
assert!(owner.install());
owner.set_sampled_limit(4 * 1024 * 1024, Duration::from_secs(30));
let auth = FrameAuthenticator::for_test(Uuid::new_v4());
let mut wire = Vec::new();
let mut sequence = 0;
write_frame(
&mut wire,
&auth,
&mut sequence,
&RemoteScannerFrame::progress(RemoteScannerProgress {
objects_scanned: 3,
directories_started: 1,
entries_visited: 3,
}),
)
.await
.expect("reported progress frame");
let first_frame_end = u64::try_from(wire.len()).expect("frame position");
write_frame(
&mut wire,
&auth,
&mut sequence,
&RemoteScannerFrame::terminal(
RemoteScannerProgress {
objects_scanned: 3,
directories_started: 1,
entries_visited: 3,
},
RemoteScannerFrameResult::NamespaceNotFound,
),
)
.await
.expect("terminal frame that must not be accepted after denial");
let mut input = ShrinkDecodeBudgetAfterFrame {
input: Cursor::new(wire),
first_frame_end,
memory: owner.clone(),
};
let budget = ScannerCycleBudget::new_with_progress_tracking(
&CancellationToken::new(),
ScannerCycleBudgetConfig {
max_objects: Some(100),
..ScannerCycleBudgetConfig::default()
},
);
let result = consume_remote_scanner_stream_until(
&mut input,
budget.token(),
budget.clone(),
RemoteScannerResponseExpectation {
bucket: "bucket",
source: TEST_SOURCE,
next_cycle: 7,
scan_plan_digest: TEST_PLAN_DIGEST,
tier_registry_generation: 0,
},
auth,
Instant::now() + Duration::from_secs(30),
)
.await;
let failure = finish_remote_scanner_stream(result, &budget).expect_err("staging denial must not become completion");
assert!(failure.retire_worker());
assert_eq!(budget.progress(), (3, 1));
assert!(budget.token().is_cancelled());
assert_eq!(input.input.position(), first_frame_end + 4);
assert_eq!(owner.snapshot().used, 0);
println!("NAMESPACE_DECODE_PROGRESS_OK");
}
#[test]
fn background_envelope_does_not_reserve_an_unread_sequence_hint() {
let owner = EcstoreBackgroundMemoryBudget::new();
// Compact envelope: [version, sequence, payload, mac]. Neither forged
// byte-array hint contains a byte, so a closed budget must not be queried.
for encoded in [
&[0x94, 3, 0, 0xdd, 0xff, 0xff, 0xff, 0xff][..],
&[0x94, 3, 0, 0x90, 0xdd, 0xff, 0xff, 0xff, 0xff][..],
] {
let error = decode_background_envelope(encoded, &owner)
.err()
.expect("truncated byte array");
assert_eq!(error.kind(), ErrorKind::InvalidData);
assert_eq!(owner.snapshot().used, 0);
}
let error = decode_background_envelope(&[0x94, 3, 0, 0x91, 7, 0x90], &owner)
.err()
.expect("the first actual payload byte needs admission");
assert_eq!(error.kind(), ErrorKind::WouldBlock);
assert!(error.get_ref().is_some(), "preserve the original I/O cause");
assert_eq!(owner.snapshot().used, 0);
}
#[tokio::test]
async fn background_envelope_mac_denial_follows_same_budget_payload_success() {
let auth = FrameAuthenticator::for_test(Uuid::from_u128(7));
let frame = RemoteScannerFrame::progress(RemoteScannerProgress::default());
let payload = rmp_serde::to_vec(&frame).expect("compact frame");
assert!(!payload.is_empty() && payload.len() <= 32, "single 32-byte payload allocation");
let envelope = RemoteScannerFrameEnvelope {
version: NS_SCANNER_PROTOCOL_VERSION,
sequence: 0,
mac: auth.sign(0, &payload).expect("MAC"),
payload,
};
let encoded = rmp_serde::to_vec_named(&envelope).expect("envelope");
let owner = EcstoreBackgroundMemoryBudget::new();
let limit = u64::try_from(encoded.len() + 32).expect("raw plus payload capacity");
owner.set_sampled_limit(limit, Duration::from_secs(30));
let mut raw = EcstoreBackgroundEncodeBuffer::new(owner.clone(), encoded.len());
raw.prepare_read(encoded.len())
.expect("raw admission")
.copy_from_slice(&encoded);
let payload_wire = rmp_serde::to_vec(&envelope.payload).expect("payload byte-array wire");
let mut failure = None;
let payload_seed = BackgroundEnvelopeBytesSeed {
buffer: EcstoreBackgroundEncodeBuffer::new(owner.clone(), encoded.len()),
failure: &mut failure,
};
let payload =
serde::de::DeserializeSeed::deserialize(payload_seed, &mut rmp_serde::Deserializer::from_read_ref(&payload_wire))
.expect("the identical budget admits raw staging and the complete payload");
assert_eq!(payload.as_ref(), envelope.payload.as_slice());
assert_eq!(owner.snapshot().used, limit);
drop(payload);
drop(raw);
assert_eq!(owner.snapshot().used, 0);
let mut wire = u32::try_from(encoded.len()).expect("frame length").to_be_bytes().to_vec();
wire.extend_from_slice(&encoded);
let mut input = Cursor::new(wire);
let mut sequence = 0;
let error = read_frame(&mut input, &auth, &mut sequence, Some(&owner))
.await
.expect_err("MAC allocation must fail after the admitted payload");
assert_eq!(error.kind(), ErrorKind::WouldBlock);
assert!(error.get_ref().is_some(), "preserve the original I/O cause");
assert_eq!(input.position(), u64::try_from(encoded.len() + 4).expect("consumed frame"));
assert_eq!(sequence, 0);
assert_eq!(owner.snapshot().used, 0);
}
#[tokio::test]
async fn background_envelope_accepts_the_full_wire_limit_and_matches_none() {
let auth = FrameAuthenticator::for_test(Uuid::from_u128(7));
let encode = |length| {
let frame = RemoteScannerFrame {
progress: RemoteScannerProgress {
objects_scanned: 3,
directories_started: 2,
entries_visited: 4,
},
phase: RemoteScannerPhase::Persisting,
result: RemoteScannerFrameResult::Error(RemoteScannerErrorFrame {
scope: RemoteScannerErrorScope::Bucket,
message: "x".repeat(length),
}),
};
let payload = rmp_serde::to_vec_named(&frame).expect("frame");
rmp_serde::to_vec_named(&RemoteScannerFrameEnvelope {
version: NS_SCANNER_PROTOCOL_VERSION,
sequence: 0,
mac: auth.sign(0, &payload).expect("MAC"),
payload,
})
.expect("envelope")
};
let (mut low, mut high) = (0, NS_SCANNER_MAX_FRAME_SIZE);
while low < high {
let length = low + (high - low).div_ceil(2);
if encode(length).len() <= NS_SCANNER_MAX_FRAME_SIZE {
low = length;
} else {
high = length - 1;
}
}
// MAC integers have variable MessagePack widths; do not assume monotonicity
// within this small neighborhood when finding the exact valid boundary.
let encoded = std::iter::once(low)
.chain(low.saturating_sub(64)..=low.saturating_add(64).min(NS_SCANNER_MAX_FRAME_SIZE))
.map(encode)
.find(|encoded| encoded.len() == NS_SCANNER_MAX_FRAME_SIZE)
.expect("valid authenticated frame exactly at the existing wire limit");
let mut wire = u32::try_from(encoded.len()).expect("wire length").to_be_bytes().to_vec();
wire.extend_from_slice(&encoded);
let owner = EcstoreBackgroundMemoryBudget::new();
// Each individual allocation is bounded by encoded.len(): raw staging,
// the retained other field, and both old and new field growth allocations.
owner.set_sampled_limit(u64::try_from(encoded.len() * 4).expect("growth peak bound"), Duration::from_secs(30));
let (mut foreground_sequence, mut background_sequence) = (0, 0);
let foreground = read_frame(&mut Cursor::new(&wire), &auth, &mut foreground_sequence, None)
.await
.expect("legacy None decoding");
let background = read_frame(&mut Cursor::new(&wire), &auth, &mut background_sequence, Some(&owner))
.await
.expect("budgeted full-limit decoding");
assert_eq!(
rmp_serde::to_vec_named(&background).expect("background frame"),
rmp_serde::to_vec_named(&foreground).expect("foreground frame")
);
assert_eq!((foreground_sequence, background_sequence), (1, 1));
assert_eq!(owner.snapshot().used, 0);
}
+16 -7
View File
@@ -1072,15 +1072,17 @@ async fn run_scanner_usage_recovery_intents_for_startup(
Ok(attempted)
}
/// Start normal scanning when enabled, or one bounded recovery attempt.
/// The disabled branch returns a finite task for the startup owner to join;
/// it never enables ordinary namespace scanning while it replays durable reset
/// intents and cleanup markers.
/// Return the normal scanner supervisor or one finite disabled recovery task.
/// The disabled branch never enables ordinary namespace scanning while it
/// replays durable reset intents and cleanup markers.
pub async fn init_scanner_with_recovery(
ctx: CancellationToken,
storeapi: Arc<ECStore>,
enabled: bool,
) -> Option<tokio::task::JoinHandle<()>> {
#[cfg(feature = "hotpath-alloc")]
drop(hotpath::MeasurementGuardSync::new("cache_model.owned_copy", false, false));
if enabled {
if let Err(error) = run_scanner_usage_recovery_intents_for_startup(ctx.clone(), storeapi.clone()).await {
warn!(
@@ -1093,8 +1095,7 @@ pub async fn init_scanner_with_recovery(
"Scanner recovery intent startup discovery failed"
);
}
init_data_scanner(ctx, storeapi).await;
return None;
return Some(init_data_scanner_task_with_storage(ctx, storeapi).await);
}
Some(tokio::spawn(async move {
if let Err(error) = run_scanner_usage_recovery_intents_for_startup(ctx.clone(), storeapi.clone()).await {
@@ -1123,6 +1124,14 @@ pub async fn init_scanner_with_recovery(
}
async fn init_data_scanner_with_storage<S>(ctx: CancellationToken, storeapi: Arc<S>)
where
S: ScannerStorage,
{
// Preserve the existing detached-start contract for legacy callers.
drop(init_data_scanner_task_with_storage(ctx, storeapi).await);
}
async fn init_data_scanner_task_with_storage<S>(ctx: CancellationToken, storeapi: Arc<S>) -> tokio::task::JoinHandle<()>
where
S: ScannerStorage,
{
@@ -1230,7 +1239,7 @@ where
break;
}
}
});
})
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
@@ -239,6 +239,45 @@ async fn issue_8297_supervisor_recovers_paused_follower_without_restart() {
crate::runtime_config::refresh_scanner_runtime_config_for_tests();
}
#[tokio::test]
#[serial]
async fn enabled_scanner_startup_returns_owned_supervisor_handle() {
temp_env::async_with_vars([(ENV_SCANNER_CYCLE, Some("1")), (ENV_SCANNER_START_DELAY_SECS, Some("0"))], async {
crate::runtime_config::refresh_scanner_runtime_config_for_tests();
crate::scanner_io::clear_dirty_usage_buckets_for_tests();
global_metrics().set_cycle(None).await;
let (_dir, store) = setup_scanner_cycle_store().await;
assert!(!crate::scanner_runtime_initialized(), "test must own the scanner runtime");
let ctx = CancellationToken::new();
let _cancel_on_drop = ctx.clone().drop_guard();
let startup_probe = ScannerStartupObservedProbe::install();
let mut supervisor = init_scanner_with_recovery(ctx.clone(), store, true)
.await
.expect("enabled startup must return its actual supervisor handle");
startup_probe.wait().await;
assert!(crate::scanner_runtime_initialized(), "the observed supervisor owns its runtime guard");
assert!(!supervisor.is_finished(), "the observed supervisor remains in flight");
ctx.cancel();
assert!(
futures::poll!(&mut supervisor).is_pending(),
"cancellation cannot join the supervisor before its observed work is released"
);
startup_probe.resume();
drop(startup_probe);
tokio::time::timeout(Duration::from_secs(15), supervisor)
.await
.expect("cancelled supervisor must join within the existing test bound")
.expect("the exact enabled supervisor must exit successfully");
assert!(
!crate::scanner_runtime_initialized(),
"joining releases the actual supervisor runtime guard"
);
})
.await;
crate::runtime_config::refresh_scanner_runtime_config_for_tests();
}
async fn run_disabled_startup(ctx: CancellationToken, store: Arc<ECStore>) {
let initialized_before = crate::scanner_runtime_initialized();
let cleanup = init_scanner_with_recovery(ctx, store, false).await;
+143 -77
View File
@@ -12,6 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use crate::data_usage_define::clone_cache_value;
use std::collections::{HashMap, HashSet};
use std::fs::FileType;
use std::io::ErrorKind;
@@ -42,6 +43,7 @@ use crate::storage_api::owner::{
};
#[cfg(test)]
use crate::storage_api::owner::{EcstoreExpirationStatus as ExpirationStatus, EcstoreLifecycleRule as LifecycleRule};
use crate::workload_admission::{run_scanner_metadata_operation, run_scanner_node_operation};
use metrics::{counter, describe_counter};
use rustfs_common::trace_bus::{TraceEvent, TraceFunc, TraceKind, trace_emit, trace_subscriber_count};
use rustfs_filemeta::{
@@ -507,7 +509,7 @@ fn failed_scan_paths(cache: &DataUsageCache) -> Vec<String> {
.cache
.iter()
.filter(|(_, entry)| entry.failed_objects > 0)
.map(|(path, _)| path.clone())
.map(|(path, _)| clone_cache_value(path))
.collect::<Vec<_>>();
paths.sort_unstable();
paths
@@ -518,12 +520,12 @@ fn data_usage_entry_has_progress(entry: &DataUsageEntry) -> bool {
}
fn set_scan_checkpoint(cache: &mut DataUsageCache, reason: DataUsageScanCheckpointReason) {
let resume_after = cache.info.scan_resume_after.clone().or_else(|| {
let resume_after = clone_cache_value(&cache.info.scan_resume_after).or_else(|| {
cache
.info
.scan_checkpoint
.as_ref()
.map(|checkpoint| checkpoint.resume_after.clone())
.map(|checkpoint| clone_cache_value(&checkpoint.resume_after))
});
if let Some(resume_after) = resume_after {
@@ -769,6 +771,7 @@ pub struct FolderScanner {
resolved_failed_object_paths: HashSet<String>,
sleeper: DynamicSleeper,
node_admission: Option<Arc<rustfs_concurrency::adaptive::NodeAdmission>>,
// should_heal: Arc<dyn Fn() -> bool + Send + Sync>,
disks: Vec<Arc<Disk>>,
disks_quorum: usize,
@@ -807,6 +810,8 @@ pub struct FolderScanner {
pending_size_reconciliation_truncated: bool,
#[cfg(test)]
list_path_raw_options_observer: Option<mpsc::UnboundedSender<ListPathRawTimeoutSnapshot>>,
#[cfg(test)]
defer_abandoned_entries_until_finished: bool,
}
struct RawEnumerationProgress {
@@ -849,8 +854,8 @@ impl RawEnumerationProgress {
return None;
}
Some(DataUsageRawEnumerationCursor::new(
self.parent.clone(),
self.last_entry.clone(),
clone_cache_value(&self.parent),
clone_cache_value(&self.last_entry),
self.entries_seen,
self.digest.clone().finalize().into(),
))
@@ -924,7 +929,7 @@ fn prune_size_reconciliation(info: &mut DataUsageCacheInfo, now: u64) {
.min_by(|(left_key, left), (right_key, right)| {
left.first_seen.cmp(&right.first_seen).then_with(|| left_key.cmp(right_key))
})
.map(|(key, _)| key.clone());
.map(|(key, _)| clone_cache_value(key));
let Some(oldest) = oldest else {
break;
};
@@ -941,7 +946,7 @@ async fn wait_for_scanner_io<T>(ctx: &CancellationToken, wait: impl std::future:
}
struct ScannerDirectory {
reader: tokio::fs::ReadDir,
reader: native_io::ReadDir,
batch_entries: usize,
#[cfg(test)]
ordered: Option<tests::erasure_probe::OrderedEntries>,
@@ -949,7 +954,7 @@ struct ScannerDirectory {
impl ScannerDirectory {
async fn open(path: &str) -> std::io::Result<Self> {
let reader = tokio::fs::read_dir(path).await?;
let reader = native_io::ReadDir::open(path).await?;
#[cfg(test)]
let (reader, ordered) = {
let mut reader = reader;
@@ -964,7 +969,7 @@ impl ScannerDirectory {
})
}
async fn next_entry(&mut self) -> std::io::Result<Option<tokio::fs::DirEntry>> {
async fn next_entry(&mut self) -> std::io::Result<Option<native_io::DirEntry>> {
#[cfg(test)]
if let Some(entries) = &mut self.ordered {
return entries.pop_front().transpose();
@@ -977,12 +982,13 @@ async fn next_scanner_dir_entry(
ctx: &CancellationToken,
sleeper: &DynamicSleeper,
entries: &mut ScannerDirectory,
) -> Result<Option<tokio::fs::DirEntry>, ScannerError> {
admission: Option<&Arc<rustfs_concurrency::adaptive::NodeAdmission>>,
) -> Result<Option<native_io::DirEntry>, ScannerError> {
if entries.batch_entries == SCANNER_DIRECTORY_BATCH_ENTRIES {
wait_for_scanner_io(ctx, tokio::task::yield_now()).await?;
entries.batch_entries = 0;
}
let entry = wait_for_scanner_io(ctx, entries.next_entry()).await??;
let entry = wait_for_scanner_io(ctx, run_scanner_node_operation(admission, ctx, entries.next_entry())).await???;
if entry.is_some() {
// Tokio prefetches bounded chunks. Pacing consumption also paces the
// next refill, without charging an extra entry for EOF.
@@ -1019,7 +1025,7 @@ impl FolderScanner {
|| folder.parent.is_some()
|| folder.name != self.old_cache.info.name
|| into.compacted
|| !abandoned_children.contains(&child_hash.key())
|| !abandoned_children.contains(&clone_cache_value(&child_hash.0))
{
return false;
}
@@ -1135,7 +1141,7 @@ impl FolderScanner {
return;
}
let mut entries: Vec<(String, u64)> = failed.iter().map(|(k, v)| (k.clone(), *v)).collect();
let mut entries: Vec<(String, u64)> = failed.iter().map(|(k, v)| (clone_cache_value(k), *v)).collect();
entries.sort_by(|(k1, ts1), (k2, ts2)| ts1.cmp(ts2).then_with(|| k1.cmp(k2)));
let remove_count = failed.len().saturating_sub(max_entries);
@@ -1150,7 +1156,7 @@ impl FolderScanner {
fn apply_size_reconciliation(&mut self, summary: &SizeSummary) {
let now = Self::now_secs();
self.pending_size_reconciliation_keys
.extend(summary.size_reconciliation.iter().map(|entry| entry.key.clone()));
.extend(summary.size_reconciliation.iter().map(|entry| clone_cache_value(&entry.key)));
self.pending_size_reconciliation_scopes.extend(
summary
.reconciliation_scopes
@@ -1162,10 +1168,10 @@ impl FolderScanner {
for info in [&mut self.new_cache.info, &mut self.update_cache.info] {
for incoming in &summary.size_reconciliation {
if let Some(existing) = info.size_reconciliation.get_mut(&incoming.key) {
existing.reason = incoming.reason.clone();
existing.reason = clone_cache_value(&incoming.reason);
existing.physical_size = incoming.physical_size;
existing.generation = incoming.generation.clone();
existing.version_id = incoming.version_id.clone();
existing.generation = clone_cache_value(&incoming.generation);
existing.version_id = clone_cache_value(&incoming.version_id);
existing.attempts = existing.attempts.saturating_add(1);
continue;
}
@@ -1174,10 +1180,10 @@ impl FolderScanner {
continue;
}
let mut entry = incoming.clone();
let mut entry = clone_cache_value(incoming);
entry.first_seen = now;
entry.attempts = 1;
info.size_reconciliation.insert(entry.key.clone(), entry);
info.size_reconciliation.insert(clone_cache_value(&entry.key), entry);
}
}
}
@@ -1199,15 +1205,19 @@ impl FolderScanner {
}
fn record_scan_resume_hint(&mut self, folder: &str) {
self.new_cache.info.scan_resume_after = Some(folder.to_string());
self.update_cache.info.scan_resume_after = Some(folder.to_string());
let checkpoint = DataUsageScanCheckpoint::new(folder.to_string(), DataUsageScanCheckpointReason::Unknown);
self.record_scan_resume_hint_with(folder, str::to_owned);
}
fn record_scan_resume_hint_with(&mut self, folder: &str, copy: impl Fn(&str) -> String) {
self.new_cache.info.scan_resume_after = Some(copy(folder));
self.update_cache.info.scan_resume_after = Some(copy(folder));
let checkpoint = DataUsageScanCheckpoint::new(copy(folder), DataUsageScanCheckpointReason::Unknown);
global_metrics().record_scanner_checkpoint_set(
checkpoint.version,
checkpoint.resume_after.clone(),
checkpoint.reason.as_str(),
);
self.new_cache.info.scan_checkpoint = Some(checkpoint.clone());
self.new_cache.info.scan_checkpoint = Some(clone_cache_value(&checkpoint));
self.update_cache.info.scan_checkpoint = Some(checkpoint);
}
@@ -1263,7 +1273,7 @@ impl FolderScanner {
.as_deref()
.is_some_and(|existing| matches!(folder_resume_match(folder, existing), Some(FolderResumeMatch::Descendant)));
if !keep_existing {
self.record_scan_resume_hint(folder);
self.record_scan_resume_hint_with(folder, clone_cache_value::<str>);
}
}
@@ -1278,7 +1288,7 @@ impl FolderScanner {
{
self.raw_enumeration_progress.truncate(position + 1);
} else {
let page_index = self.raw_enumeration_page_index_for_parent(parent).cloned();
let page_index = self.raw_enumeration_page_index_for_parent(parent).map(clone_cache_value);
self.raw_enumeration_progress
.push(RawEnumerationProgress::new(parent, page_index));
}
@@ -1434,11 +1444,11 @@ impl FolderScanner {
}) {
return None;
}
let mut snapshot = self.new_cache.clone();
let mut snapshot = clone_cache_value(&self.new_cache);
let mut parent = None;
for ancestor in &self.checkpoint_ancestors[..self.checkpoint_depth] {
snapshot.replace_hashed(&ancestor.hash, &parent, &ancestor.entry);
parent = Some(ancestor.hash.clone());
parent = Some(clone_cache_value(&ancestor.hash));
}
snapshot.replace_hashed(hash, &folder.parent, into);
snapshot.checked_flatten_complete_scope(&snapshot.info.name)?;
@@ -1469,11 +1479,11 @@ impl FolderScanner {
return;
}
let Some(old_entry) = self.old_cache.cache.get(&parent_hash.key()) else {
let Some(old_entry) = self.old_cache.cache.get(&clone_cache_value(&parent_hash.0)) else {
return;
};
let old_children = old_entry.children.iter().cloned().collect::<Vec<_>>();
let old_children = old_entry.children.iter().map(clone_cache_value).collect::<Vec<_>>();
for child in old_children {
if entry.children.contains(&child) {
continue;
@@ -1482,9 +1492,9 @@ impl FolderScanner {
continue;
}
let child_hash = DataUsageHash(child.clone());
let child_hash = DataUsageHash(clone_cache_value(&child));
self.new_cache
.copy_with_children(&self.old_cache, &child_hash, &Some(parent_hash.clone()));
.copy_with_children(&self.old_cache, &child_hash, &Some(clone_cache_value(parent_hash)));
entry.children.insert(child);
}
}
@@ -1506,9 +1516,9 @@ impl FolderScanner {
self.coverage_frontier = None;
self.coverage_gap = true;
}
let mut child_entry = child_entry.clone();
let mut child_entry = clone_cache_value(child_entry);
self.carry_forward_old_children(child_hash, &mut child_entry);
self.record_scan_resume_hint_if_not_ancestor(&child_hash.key());
self.record_scan_resume_hint_if_not_ancestor(&clone_cache_value(&child_hash.0));
parent_entry.add_child(child_hash);
self.new_cache.replace_hashed(child_hash, parent, &child_entry);
self.update_cache.delete_recursive(child_hash);
@@ -1861,7 +1871,12 @@ impl FolderScanner {
// directory that vanished during iteration instead of failing the
// whole parent subtree. The bucket root (no parent) still propagates
// NotFound so namespace removal is detected by the caller.
let mut dir_reader = match wait_for_scanner_io(&ctx, ScannerDirectory::open(&dir_path)).await? {
let mut dir_reader = match wait_for_scanner_io(
&ctx,
run_scanner_node_operation(self.node_admission.as_ref(), &ctx, ScannerDirectory::open(&dir_path)),
)
.await??
{
Ok(dir_reader) => Some(dir_reader),
Err(e)
if folder.parent.is_some()
@@ -1899,7 +1914,7 @@ impl FolderScanner {
// Hold at most one metadata entry until a candidate reaches
// EOF: a later part must not leave object accounting or ILM work.
let (entry, replayed_metadata) = if let Some(reader) = dir_reader.as_mut() {
let entry = match next_scanner_dir_entry(&ctx, &self.sleeper, reader).await {
let entry = match next_scanner_dir_entry(&ctx, &self.sleeper, reader, self.node_admission.as_ref()).await {
Ok(Some(entry)) => entry,
Ok(None) => {
raw_enumeration_complete = true;
@@ -2010,7 +2025,12 @@ impl FolderScanner {
// Ignore entries that disappeared during traversal or hit symlink
// loops, but propagate other walk errors.
let mut entry_type = match wait_for_scanner_io(&ctx, entry.file_type()).await? {
let mut entry_type = match wait_for_scanner_io(
&ctx,
run_scanner_node_operation(self.node_admission.as_ref(), &ctx, entry.file_type()),
)
.await??
{
Ok(entry_type) => entry_type,
Err(e) if e.kind() == ErrorKind::NotFound => {
debug!(
@@ -2061,7 +2081,12 @@ impl FolderScanner {
let mut entry_metadata_len = None;
if entry_type.is_symlink() {
wait_for_scanner_io(&ctx, self.sleeper.admit_disk_io(1, 0)).await?;
let metadata = match wait_for_scanner_io(&ctx, tokio::fs::metadata(&file_path)).await? {
let metadata = match wait_for_scanner_io(
&ctx,
run_scanner_node_operation(self.node_admission.as_ref(), &ctx, native_io::metadata(&file_path)),
)
.await??
{
Ok(metadata) => metadata,
Err(e) if e.kind() == ErrorKind::NotFound => {
debug!(
@@ -2134,7 +2159,7 @@ impl FolderScanner {
continue;
}
let exists = self.old_cache.cache.contains_key(&h.key());
let exists = self.old_cache.cache.contains_key(&clone_cache_value(&h.0));
let this = CachedFolder {
name: entry_name.clone(),
@@ -2143,15 +2168,15 @@ impl FolderScanner {
};
if exists && self.should_reuse_clean_root_child(&folder, into, &this, &h, &abandoned_children) {
abandoned_children.remove(&h.key());
abandoned_children.remove(&clone_cache_value(&h.0));
self.reuse_clean_root_child(&this, &h, into);
} else if exists {
abandoned_children.remove(&h.key());
abandoned_children.remove(&clone_cache_value(&h.0));
existing_folders.push(this);
self.update_cache
.copy_with_children(&self.old_cache, &h, &Some(this_hash.clone()));
.copy_with_children(&self.old_cache, &h, &Some(clone_cache_value(&this_hash)));
} else {
abandoned_children.remove(&h.key());
abandoned_children.remove(&clone_cache_value(&h.0));
new_folders.push(this);
}
continue;
@@ -2203,21 +2228,26 @@ impl FolderScanner {
len
} else {
wait_for_scanner_io(&ctx, self.sleeper.admit_disk_io(1, 0)).await?;
wait_for_scanner_io(&ctx, tokio::fs::metadata(&item.path))
.await?
.map(|metadata| metadata.len())
.unwrap_or(0)
wait_for_scanner_io(
&ctx,
run_scanner_node_operation(self.node_admission.as_ref(), &ctx, native_io::metadata(&item.path)),
)
.await??
.map(|metadata| metadata.len())
.unwrap_or(0)
}
} else {
0
};
wait_for_scanner_io(&ctx, self.sleeper.admit_disk_io(1, budgeted_bytes)).await?;
self.object_metadata_visits.increment(1);
let sz = match self
.local_disk
.get_size_with_tier_names(item.clone(), &self.tier_registry.names)
.await
let sz = match run_scanner_metadata_operation(self.node_admission.as_ref(), &ctx, async {
self.object_metadata_visits.increment(1);
self.local_disk
.get_size_with_tier_names(item.clone(), &self.tier_registry.names)
.await
})
.await?
{
Ok(sz) => sz,
Err(e) => {
@@ -2481,11 +2511,11 @@ impl FolderScanner {
// Add new folders to the update tree so totals update for these.
if !into.compacted {
let mut found_any = false;
let mut parent = this_hash.clone();
let mut parent = clone_cache_value(&this_hash);
let update_cache_name_hash = hash_path(&self.update_cache.info.name);
while parent != update_cache_name_hash {
let parent_key = parent.key();
let parent_key = clone_cache_value(&parent.0);
let e = self.update_cache.find(&parent_key);
if e.is_none_or(|v| v.compacted) {
found_any = true;
@@ -2500,8 +2530,11 @@ impl FolderScanner {
}
if !found_any {
// Add non-compacted empty entry.
self.update_cache
.replace_hashed(&h, &Some(this_hash.clone()), &DataUsageEntry::default());
self.update_cache.replace_hashed(
&h,
&Some(clone_cache_value(&this_hash)),
&DataUsageEntry::default(),
);
}
}
}
@@ -2548,7 +2581,7 @@ impl FolderScanner {
{
Ok(outcome) => outcome,
Err(e) => {
if ctx.is_cancelled() {
if ctx.is_cancelled() || matches!(&e, ScannerError::ResourceDeferred) {
self.preserve_partial_child_progress(&folder_item.parent, &h, into, &dst)
.await;
return Err(e);
@@ -2575,13 +2608,13 @@ impl FolderScanner {
if outcome == FolderScanOutcome::Vanished {
self.new_cache.delete_recursive(&h);
self.update_cache.delete_recursive(&h);
into.children.remove(&h.key());
into.children.remove(&clone_cache_value(&h.0));
continue;
}
if outcome == FolderScanOutcome::ErasureData {
self.new_cache.delete_recursive(&h);
self.update_cache.delete_recursive(&h);
into.children.remove(&h.key());
into.children.remove(&clone_cache_value(&h.0));
// Earlier namespace siblings may have completed, but none
// can certify coverage of this unresolved parent object.
self.coverage_gap |= forward_sweep;
@@ -2611,11 +2644,14 @@ impl FolderScanner {
if queued_folder.source == FolderScanSource::New
&& !into.compacted
&& self.update_cache.find(&this_hash.key()).is_some_and(|v| !v.compacted)
&& self
.update_cache
.find(&clone_cache_value(&this_hash.0))
.is_some_and(|v| !v.compacted)
{
self.update_cache.delete_recursive(&h);
self.update_cache
.copy_with_children(&self.new_cache, &h, &Some(this_hash.clone()));
.copy_with_children(&self.new_cache, &h, &Some(clone_cache_value(&this_hash)));
}
}
@@ -2795,19 +2831,23 @@ impl FolderScanner {
let mut seen_truncated_objects: HashSet<String> = HashSet::new();
loop {
if agreed_closed && partial_closed && finished_closed {
if child_ctx.is_cancelled() || (agreed_closed && partial_closed && finished_closed) {
break;
}
#[cfg(test)]
let defer_entries = self.defer_abandoned_entries_until_finished && !finished_closed;
#[cfg(not(test))]
let defer_entries = false;
select! {
entry_name = agreed_rx.recv(), if !agreed_closed => {
entry_name = agreed_rx.recv(), if !agreed_closed && !defer_entries => {
let Some(entry_name) = entry_name else {
agreed_closed = true;
continue;
};
(self.update_current_path)(&entry_name).await;
}
entries = partial_rx.recv(), if !partial_closed => {
entries = partial_rx.recv(), if !partial_closed && !defer_entries => {
let Some(entries) = entries else {
partial_closed = true;
continue;
@@ -2933,6 +2973,7 @@ impl FolderScanner {
finished_closed = true;
continue;
};
finished_closed = true;
if disk_errors_are_only_missing_paths(&errs) {
debug!(
target: "rustfs::scanner::folder",
@@ -2954,7 +2995,9 @@ impl FolderScanner {
"Scanner list_path_raw finished with disk errors"
);
}
child_ctx.cancel();
// The producer awaited callback sends, not their
// consumption. Drain the remaining candidate and
// agreed frames before closing this child walk.
}
_ = child_ctx.cancelled() => {
break;
@@ -2972,7 +3015,10 @@ impl FolderScanner {
if into.compacted {
// In compacted mode child totals are accumulated directly into the parent entry.
let fut = Box::pin(self.scan_folder(ctx.clone(), folder_item.clone(), into));
fut.await.map_err(|e| ScannerError::Other(e.to_string()))?;
fut.await.map_err(|error| match error {
ScannerError::ResourceDeferred => ScannerError::ResourceDeferred,
other => ScannerError::Other(other.to_string()),
})?;
self.send_update_for_entry(&this_hash, &folder.parent, into).await;
tokio::task::yield_now().await;
} else {
@@ -2991,7 +3037,7 @@ impl FolderScanner {
)
.await
{
if ctx.is_cancelled() {
if ctx.is_cancelled() || matches!(&e, ScannerError::ResourceDeferred) {
self.preserve_partial_child_progress(&folder_item.parent, &h, into, &dst)
.await;
return Err(e);
@@ -3033,7 +3079,7 @@ impl FolderScanner {
if !into.compacted
&& self.old_cache.info.scan_progress.is_none()
&& self.new_cache.info.name != folder.name
&& let Some(mut flat) = self.new_cache.size_recursive(&this_hash.key())
&& let Some(mut flat) = self.new_cache.size_recursive(&clone_cache_value(&this_hash.0))
{
flat.compacted = true;
let mut should_compact = false;
@@ -3070,9 +3116,9 @@ impl FolderScanner {
done_compact();
}
if self.update_cache.cache.contains_key(&this_hash.key()) && !was_compacted {
if self.update_cache.cache.contains_key(&clone_cache_value(&this_hash.0)) && !was_compacted {
// Replace if existed before.
if let Some(flat) = self.new_cache.size_recursive(&this_hash.key()) {
if let Some(flat) = self.new_cache.size_recursive(&clone_cache_value(&this_hash.0)) {
self.update_cache.delete_recursive(&this_hash);
self.update_cache.replace_hashed(&this_hash, &folder.parent, &flat);
}
@@ -3156,11 +3202,11 @@ pub(crate) async fn scan_data_folder_scoped(
cache.fold_retired_tiers(&tier_registry.names);
cache.info.tier_registry_generation = Some(tier_registry.generation);
let resume_frontier = cache.validated_scan_frontier().map(str::to_owned);
let resume_frontier = cache.validated_scan_frontier().map(clone_cache_value);
if resume_frontier.is_some() {
global_metrics().record_scanner_checkpoint_used();
}
let bucket_name = cache.info.name.clone();
let bucket_name = clone_cache_value(&cache.info.name);
let object_metadata_visits = metrics::counter!(
"rustfs_scanner_object_metadata_visits_total",
"bucket" => bucket_name.clone(),
@@ -3170,11 +3216,11 @@ pub(crate) async fn scan_data_folder_scoped(
let resume_position_on_entry = cache.has_scan_position();
let had_scan_checkpoint = cache.info.scan_checkpoint.is_some();
let new_cache = DataUsageCache {
info: cache.info.clone(),
info: clone_cache_value(&cache.info),
..Default::default()
};
let update_cache = DataUsageCache {
info: cache.info.clone(),
info: clone_cache_value(&cache.info),
..Default::default()
};
let resume_failed_paths = if resume_frontier.is_some() {
@@ -3198,6 +3244,7 @@ pub(crate) async fn scan_data_folder_scoped(
failed_object_paths_seen: HashSet::new(),
resolved_failed_object_paths: HashSet::new(),
sleeper,
node_admission: crate::workload_admission::scanner_node_admission(),
disks,
disks_quorum,
updates,
@@ -3216,7 +3263,7 @@ pub(crate) async fn scan_data_folder_scoped(
local_disk,
tier_registry,
pending_heals_changed: false,
coverage_frontier: resume_frontier.clone(),
coverage_frontier: clone_cache_value(&resume_frontier),
resume_frontier,
resume_failed_paths,
coverage_gap: false,
@@ -3230,6 +3277,8 @@ pub(crate) async fn scan_data_folder_scoped(
pending_size_reconciliation_truncated: false,
#[cfg(test)]
list_path_raw_options_observer: None,
#[cfg(test)]
defer_abandoned_entries_until_finished: false,
};
let now = FolderScanner::now_secs();
@@ -3252,7 +3301,10 @@ pub(crate) async fn scan_data_folder_scoped(
};
// Scan the folder
match scanner.scan_folder(ctx.clone(), folder, &mut root).await {
let physical = scanner.local_disk.physical_admission();
match crate::workload_admission::with_scanner_physical_disk(physical, scanner.scan_folder(ctx.clone(), folder, &mut root))
.await
{
Ok(()) => {
// Get the new cache and finalize it
let coverage_gap = scanner.coverage_gap;
@@ -3320,7 +3372,7 @@ pub(crate) async fn scan_data_folder_scoped(
}
}
Err(e) => {
if ctx.is_cancelled() {
if ctx.is_cancelled() || matches!(&e, ScannerError::ResourceDeferred) {
let root_hash = hash_path(&bucket_name);
let root_has_progress = data_usage_root_has_progress(&root);
let pending_heals_changed = scanner.pending_heals_changed;
@@ -3331,7 +3383,7 @@ pub(crate) async fn scan_data_folder_scoped(
if root_has_progress {
scanner.carry_forward_old_children(&root_hash, &mut root);
}
let coverage_frontier = scanner.coverage_frontier.clone();
let coverage_frontier = clone_cache_value(&scanner.coverage_frontier);
let new_cache = scanner.as_mut_new_cache();
if root_has_progress {
new_cache.replace_hashed(&root_hash, &None, &root);
@@ -3395,6 +3447,19 @@ pub(crate) async fn scan_data_folder_scoped(
close_disk_guard.close().await;
return Err(ScannerError::NamespaceNotFoundCache(Box::new(partial_cache)));
}
if matches!(&e, ScannerError::ResourceDeferred) {
// Admission pressure is not disk failure or an empty census.
// Keep the previous tree when no new useful progress exists.
let mut partial_cache = scanner.old_cache;
if partial_cache.cache.is_empty() {
partial_cache
.cache
.insert(hash_path(&bucket_name).0, DataUsageEntry::default());
}
partial_cache.info.snapshot_complete = false;
close_disk_guard.close().await;
return Err(ScannerError::PartialCache(Box::new(partial_cache)));
}
close_disk_guard.close().await;
// No useful information, return original cache
Err(e)
@@ -3404,6 +3469,7 @@ pub(crate) async fn scan_data_folder_scoped(
mod item_actions;
mod ledger;
mod native_io;
use item_actions::*;
pub use item_actions::{GetSizeFn, ScannerItem};
+16 -11
View File
@@ -58,7 +58,7 @@ impl FolderScanner {
self.pending_heal_batch_dirty = true;
return;
}
self.update_cache.info.pending_heals = self.new_cache.info.pending_heals.clone();
self.update_cache.info.pending_heals = clone_cache_value(&self.new_cache.info.pending_heals);
#[cfg(test)]
{
self.pending_heal_sync_count += 1;
@@ -246,7 +246,7 @@ impl FolderScanner {
return Ok(());
}
let bucket = scanner.new_cache.info.name.clone();
let bucket = clone_cache_value(&scanner.new_cache.info.name);
// Legacy notices cannot bind a verified disposition to the current
// incarnation, kind, set scope and responsibility generation.
let repaired = rustfs_common::mrf_channel::take_mrf_repaired_events_for(&bucket);
@@ -295,9 +295,9 @@ impl FolderScanner {
scanner
.send_required_scanner_heal_request(
pending.kind,
pending.bucket.clone(),
pending.object.clone(),
pending.version_id.clone(),
clone_cache_value(&pending.bucket),
clone_cache_value(&pending.object),
clone_cache_value(&pending.version_id),
request,
)
.await?;
@@ -371,7 +371,7 @@ pub(super) fn pending_scanner_heal_retry_candidates_at(
.then_with(|| pending_scanner_heal_identity(a).cmp(&pending_scanner_heal_identity(b)))
});
entries.truncate(MAX_PENDING_SCANNER_HEAL_RETRIES_PER_BUCKET);
entries.into_iter().cloned().collect()
entries.into_iter().map(clone_cache_value).collect()
}
pub(super) fn build_pending_scanner_heal_request(entry: &PendingScannerHeal) -> Option<HealChannelRequest> {
@@ -383,15 +383,20 @@ pub(super) fn build_pending_scanner_heal_request(entry: &PendingScannerHeal) ->
HealChannelPriority::Low
};
match entry.kind {
PendingScannerHealKind::Bucket => Some(build_bucket_heal_request(entry.bucket.clone(), priority)),
PendingScannerHealKind::Bucket => Some(build_bucket_heal_request(clone_cache_value(&entry.bucket), priority)),
PendingScannerHealKind::Object => entry.object.as_ref().map(|object| {
if entry.version_id.is_none() {
build_non_destructive_object_heal_request(entry.bucket.clone(), object.clone(), entry.scan_mode, priority)
build_non_destructive_object_heal_request(
clone_cache_value(&entry.bucket),
clone_cache_value(object),
entry.scan_mode,
priority,
)
} else {
build_object_heal_request(
entry.bucket.clone(),
object.clone(),
entry.version_id.clone(),
clone_cache_value(&entry.bucket),
clone_cache_value(object),
clone_cache_value(&entry.version_id),
entry.scan_mode,
priority,
)
@@ -0,0 +1,339 @@
// Copyright 2026 RustFS Team
// Licensed under the Apache License, Version 2.0.
//! Scanner filesystem jobs retain resource ownership inside the native worker.
use std::collections::VecDeque;
use std::ffi::OsString;
use std::fs::{FileType, Metadata};
use std::future::{Future, poll_fn};
use std::io;
use std::path::PathBuf;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll, ready};
use crate::storage_api::ecstore_spawn_blocking_with_resource_ownership as spawn_native;
// Keep Tokio's existing bounded prefetch and one native job per refill.
const DIRECTORY_CHUNK_SIZE: usize = 32;
pub(super) struct DirEntry {
entry: Arc<std::fs::DirEntry>,
file_type: Option<FileType>,
}
impl DirEntry {
pub(super) fn path(&self) -> PathBuf {
self.entry.path()
}
pub(super) fn file_name(&self) -> OsString {
self.entry.file_name()
}
pub(super) async fn file_type(&self) -> io::Result<FileType> {
if let Some(file_type) = self.file_type {
return Ok(file_type);
}
let entry = self.entry.clone();
spawn_native(move || entry.file_type()).await.map_err(io::Error::other)?
}
}
struct DirectoryBuffer {
reader: std::fs::ReadDir,
entries: VecDeque<io::Result<DirEntry>>,
more: bool,
}
impl DirectoryBuffer {
fn refill(&mut self) {
for _ in 0..DIRECTORY_CHUNK_SIZE {
let Some(entry) = self.reader.next() else {
self.more = false;
break;
};
let success = entry.is_ok();
self.entries.push_back(entry.map(|entry| DirEntry {
file_type: entry.file_type().ok(),
entry: Arc::new(entry),
}));
if !success {
break;
}
}
}
}
enum DirectoryState {
Idle(Option<DirectoryBuffer>),
Pending(tokio::task::JoinHandle<DirectoryBuffer>),
Failed,
}
pub(super) struct ReadDir {
state: DirectoryState,
}
impl ReadDir {
pub(super) async fn open(path: &str) -> io::Result<Self> {
let path = PathBuf::from(path);
let buffer = spawn_native(move || {
let mut buffer = DirectoryBuffer {
reader: std::fs::read_dir(path)?,
entries: VecDeque::with_capacity(DIRECTORY_CHUNK_SIZE),
more: true,
};
buffer.refill();
Ok::<_, io::Error>(buffer)
})
.await
.map_err(io::Error::other)??;
Ok(Self {
state: DirectoryState::Idle(Some(buffer)),
})
}
pub(super) async fn next_entry(&mut self) -> io::Result<Option<DirEntry>> {
// Keep a pending refill on the reader, not in the borrowed next-entry
// future. Cancelling that future must not skip its prefetched entries.
poll_fn(|cx| self.poll_next_entry(cx)).await
}
fn poll_next_entry(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<Option<DirEntry>>> {
loop {
match &mut self.state {
DirectoryState::Idle(buffer) => {
let Some(ready) = buffer.as_mut() else {
return Poll::Ready(Err(io::Error::other("scanner directory buffer is unavailable")));
};
if let Some(entry) = ready.entries.pop_front() {
return Poll::Ready(entry.map(Some));
}
if !ready.more {
return Poll::Ready(Ok(None));
}
let Some(mut buffer) = buffer.take() else {
return Poll::Ready(Err(io::Error::other("scanner directory buffer is unavailable")));
};
self.state = DirectoryState::Pending(spawn_native(move || {
buffer.refill();
buffer
}));
}
DirectoryState::Pending(worker) => match ready!(Pin::new(worker).poll(cx)) {
Ok(buffer) => self.state = DirectoryState::Idle(Some(buffer)),
Err(error) => {
self.state = DirectoryState::Failed;
return Poll::Ready(Err(io::Error::other(error)));
}
},
DirectoryState::Failed => {
return Poll::Ready(Err(io::Error::other("scanner directory worker failed")));
}
}
}
}
}
pub(super) async fn metadata(path: &str) -> io::Result<Metadata> {
let path = PathBuf::from(path);
spawn_native(move || std::fs::metadata(path))
.await
.map_err(io::Error::other)?
}
#[cfg(test)]
mod tests {
use super::*;
use rustfs_concurrency::adaptive::{
AdaptiveAdmission, AdaptivePermit,
permit_set::{AdaptivePermitSet, AdaptivePermitSetOwner},
};
use std::num::NonZeroU16;
enum Resources {
Single(AdaptivePermit),
Layered(AdaptivePermitSet),
}
fn runtime() -> tokio::runtime::Runtime {
tokio::runtime::Builder::new_current_thread()
.enable_all()
.max_blocking_threads(1)
.build()
.expect("single native worker runtime")
}
fn directory(count: usize) -> tempfile::TempDir {
let root = tempfile::tempdir().expect("owned directory");
for index in 0..count {
std::fs::write(root.path().join(format!("object-{index}")), b"data").expect("test file");
}
root
}
async fn occupy_native_worker() -> (std::sync::mpsc::Sender<()>, tokio::task::JoinHandle<()>) {
let (entered, started) = tokio::sync::oneshot::channel();
let (release, resume) = std::sync::mpsc::channel();
let worker = tokio::task::spawn_blocking(move || {
entered.send(()).expect("worker-start observer");
let _ = resume.recv();
});
started.await.expect("native worker occupied");
(release, worker)
}
#[test]
fn scanner_native_jobs_keep_single_and_layered_units_after_waiter_cancellation() {
runtime().block_on(async {
let root = directory(DIRECTORY_CHUNK_SIZE + 3);
let path = root.path().to_str().expect("test path");
for count in [1, 2] {
for operation in ["open", "refill", "metadata", "file_type"] {
let gates: Vec<_> = (0..count).map(|_| AdaptiveAdmission::new(NonZeroU16::MIN, 1)).collect();
let mut units: Vec<_> = gates.iter().map(|gate| gate.try_acquire().expect("physical unit")).collect();
let owner = AdaptivePermitSetOwner::new(NonZeroU16::MIN, gates.clone()).expect("physical subset");
let resources = if count == 1 {
Resources::Single(units.pop().expect("single physical unit"))
} else {
Resources::Layered(owner.capture(&units.iter().collect::<Vec<_>>()).expect("native ownership"))
};
drop(units);
let mut reader = ReadDir::open(path).await.expect("prepare reader");
let mut entry = reader.next_entry().await.expect("first entry").expect("file");
entry.file_type = None;
for _ in 1..DIRECTORY_CHUNK_SIZE {
reader.next_entry().await.expect("consume prefetched entry").expect("file");
}
let (release, blocker) = occupy_native_worker().await;
let work = async {
match operation {
"open" => ReadDir::open(path).await.map(|_| ()),
"refill" => reader.next_entry().await.map(|_| ()),
"metadata" => metadata(entry.path().to_str().expect("entry path")).await.map(|_| ()),
"file_type" => entry.file_type().await.map(|_| ()),
_ => unreachable!("fixed test operations"),
}
};
let mut waiter = Box::pin(async {
match resources {
Resources::Single(unit) => {
crate::storage_api::ecstore_with_background_resource_unit(unit, work).await
}
Resources::Layered(units) => {
crate::storage_api::ecstore_with_background_resource_set(units, work).await
}
}
});
assert!(futures::poll!(&mut waiter).is_pending(), "native job must be queued: {operation}");
drop(waiter);
drop(reader);
let held_after_cancel = gates.iter().all(|gate| gate.snapshot().active == 1);
gates[0].set_limit(0);
release.send(()).expect("release native worker");
blocker.await.expect("blocker finished");
// The sentinel runs after the queued filesystem job on the
// sole blocking worker, even after its waiter was dropped.
tokio::task::spawn_blocking(|| ()).await.expect("native queue drained");
assert!(
held_after_cancel,
"cancelled {operation} refunded {count} physical units before native completion"
);
assert!(gates.iter().all(|gate| gate.snapshot().active == 0), "native {operation} leaked capacity");
assert!(gates[0].try_acquire().is_none(), "native completion cannot undo a capacity reduction");
}
}
});
}
#[test]
fn cancelled_refill_resumes_without_skipping_or_duplicating_entries() {
runtime().block_on(async {
let count = DIRECTORY_CHUNK_SIZE * 3 + 7;
let root = directory(count);
let mut reader = ReadDir::open(root.path().to_str().expect("root path")).await.expect("open");
let mut names = std::collections::BTreeSet::new();
for _ in 0..DIRECTORY_CHUNK_SIZE {
let entry = reader.next_entry().await.expect("entry").expect("file");
assert!(names.insert(entry.file_name()));
}
let (release, blocker) = occupy_native_worker().await;
let mut next = Box::pin(reader.next_entry());
assert!(futures::poll!(&mut next).is_pending());
drop(next);
release.send(()).expect("release refill");
blocker.await.expect("blocker ended");
while let Some(entry) = reader.next_entry().await.expect("resume refill") {
assert!(names.insert(entry.file_name()), "duplicate entry after cancelled refill");
}
assert_eq!(names.len(), count, "cancelling the borrowed future must retain the pending chunk");
assert!(reader.next_entry().await.expect("repeat EOF").is_none());
});
}
#[test]
fn directory_prefetch_preserves_empty_and_exact_chunk_boundaries() {
runtime().block_on(async {
for count in [0, 1, DIRECTORY_CHUNK_SIZE, DIRECTORY_CHUNK_SIZE + 1] {
let root = directory(count);
let mut reader = ReadDir::open(root.path().to_str().expect("root path")).await.expect("open");
let mut names = std::collections::BTreeSet::new();
while let Some(entry) = reader.next_entry().await.expect("read entry") {
assert!(names.insert(entry.file_name()), "prefetch duplicated an entry");
}
assert_eq!(names.len(), count);
assert!(reader.next_entry().await.expect("repeat EOF").is_none());
}
});
}
#[test]
fn failed_refill_is_an_error_after_partial_enumeration() {
runtime().block_on(async {
let root = directory(DIRECTORY_CHUNK_SIZE + 1);
let mut reader = ReadDir::open(root.path().to_str().expect("root path")).await.expect("open");
for _ in 0..DIRECTORY_CHUNK_SIZE {
assert!(reader.next_entry().await.expect("prefetched entry").is_some());
}
let (release, blocker) = occupy_native_worker().await;
let mut next = Box::pin(reader.next_entry());
assert!(futures::poll!(&mut next).is_pending());
drop(next);
let DirectoryState::Pending(worker) = &reader.state else {
panic!("refill must be queued");
};
worker.abort();
release.send(()).expect("release native worker");
blocker.await.expect("blocker finished");
for _ in 0..2 {
assert!(
matches!(reader.next_entry().await, Err(error) if error.kind() == io::ErrorKind::Other),
"a failed refill cannot become successful EOF"
);
}
});
}
#[test]
fn scanner_native_errors_and_symlink_metadata_keep_filesystem_semantics() {
runtime().block_on(async {
let root = directory(1);
let path = root.path().join("missing");
let missing = path.to_str().expect("missing path");
assert!(matches!(ReadDir::open(missing).await, Err(error) if error.kind() == io::ErrorKind::NotFound));
assert!(matches!(metadata(missing).await, Err(error) if error.kind() == io::ErrorKind::NotFound));
let file = root.path().join("object-0");
assert!(ReadDir::open(file.to_str().expect("file path")).await.is_err());
#[cfg(unix)]
{
let link = root.path().join("link");
std::os::unix::fs::symlink(&file, &link).expect("file symlink");
let info = metadata(link.to_str().expect("link path")).await.expect("follow symlink");
assert!(info.is_file());
assert_eq!(info.len(), 4);
}
});
}
}
@@ -30,7 +30,9 @@ use std::sync::Mutex;
pub(super) mod checkpoint_fixture;
pub(super) mod enumeration_restart;
mod incremental_enumeration;
mod listing_completion;
mod metadata_visits;
mod node_admission;
/// Reset the process-global alert cooldown map; test-only.
fn reset_alert_cooldowns() {
@@ -364,6 +366,7 @@ async fn build_test_scanner() -> (FolderScanner, std::path::PathBuf) {
failed_object_paths_seen: HashSet::new(),
resolved_failed_object_paths: HashSet::new(),
sleeper: SCANNER_SLEEPER.clone(),
node_admission: None,
disks: Vec::new(),
disks_quorum: 0,
updates: None,
@@ -398,6 +401,7 @@ async fn build_test_scanner() -> (FolderScanner, std::path::PathBuf) {
pending_size_reconciliation_scopes: HashSet::new(),
pending_size_reconciliation_truncated: false,
list_path_raw_options_observer: None,
defer_abandoned_entries_until_finished: false,
};
(scanner, temp_dir)
@@ -12,7 +12,7 @@ enum DirectoryOrder {
FailOpen(std::io::ErrorKind),
}
pub(in crate::scanner_folder) type OrderedEntries = VecDeque<std::io::Result<tokio::fs::DirEntry>>;
pub(in crate::scanner_folder) type OrderedEntries = VecDeque<std::io::Result<native_io::DirEntry>>;
static ORDERED_DIRECTORIES: LazyLock<Mutex<HashMap<String, DirectoryOrder>>> = LazyLock::new(|| Mutex::new(HashMap::new()));
struct OrderedDirectory(String);
@@ -52,7 +52,7 @@ impl Drop for OrderedDirectory {
pub(in crate::scanner_folder) async fn ordered_entries(
path: &str,
reader: &mut tokio::fs::ReadDir,
reader: &mut native_io::ReadDir,
) -> std::io::Result<Option<OrderedEntries>> {
let order = ORDERED_DIRECTORIES
.lock()
@@ -0,0 +1,104 @@
// Copyright 2026 RustFS Team
// Licensed under the Apache License, Version 2.0.
use super::*;
#[tokio::test]
#[serial]
async fn abandoned_listing_drains_candidate_buffered_before_finished() {
let (mut scanner, root) = build_test_scanner().await;
let _guard = TestGuard {
temp_dir: Some(root.clone()),
};
let bucket = "bucket";
let object = "abandoned/object";
tokio::fs::create_dir_all(root.join(bucket)).await.expect("local bucket");
let mut disks = vec![scanner.local_disk.clone()];
let mut original_metadata = Vec::new();
for (index, version) in [(2, 0x1935), (3, 0x1935), (4, 0x1934)] {
let disk_root = root.join(format!("disk-{index}"));
let metadata = metadata_for_object_version(bucket, object, Some(Uuid::from_u128(version)));
write_test_object_metadata_bytes(&disk_root, bucket, object, &metadata).await;
original_metadata.push((disk_root.join(bucket).join(object).join(STORAGE_FORMAT_FILE), metadata));
let endpoint = Endpoint::try_from(disk_root.to_string_lossy().as_ref()).expect("listing fixture endpoint");
disks.push(
new_disk(
&endpoint,
&DiskOption {
cleanup: false,
health_check: false,
},
)
.await
.expect("listing fixture disk"),
);
}
for cache in [&mut scanner.old_cache, &mut scanner.new_cache, &mut scanner.update_cache] {
cache.info.name = bucket.to_owned();
}
scanner.old_cache.replace(
"bucket/abandoned/object",
bucket,
DataUsageEntry {
objects: 1,
..Default::default()
},
);
scanner.heal_object_select = 1;
scanner.disks = disks;
scanner.disks_quorum = 2;
scanner.sleeper = DynamicSleeper::new(rustfs_config::ScannerSpeed::Fastest);
// Force the valid select ordering that receives terminal errors before the
// final candidate. One candidate frame fits in the bounded partial channel.
scanner.defer_abandoned_entries_until_finished = true;
let (_commands, mut heal_rx) = rustfs_heal_contracts::heal_channel::init_heal_channels().expect("isolated heal channel");
let (captured_tx, mut captured_rx) = mpsc::unbounded_channel();
let responder = tokio::spawn(async move {
while let Some(command) = heal_rx.recv().await {
let rustfs_heal_contracts::heal_channel::HealReceiptCommand { request, response_tx } = command;
captured_tx
.send((request.object_prefix.clone(), request.object_version_id.clone()))
.expect("candidate capture remains open");
let _ = response_tx.send(Ok(rustfs_heal_contracts::heal_channel::HealAdmissionReceipt {
result: HealAdmissionResult::Accepted,
task_id: request.id,
}));
}
});
let result = tokio::time::timeout(
Duration::from_secs(2),
scanner.scan_folder(
CancellationToken::new(),
CachedFolder {
name: bucket.to_owned(),
parent: None,
object_heal_prob_div: 1,
},
&mut DataUsageEntry::default(),
),
)
.await;
responder.abort();
result
.expect("listing completion must close every callback channel")
.expect("complete folder scan");
let mut candidates = HashSet::new();
while let Ok((object, version)) = captured_rx.try_recv() {
if let (Some(object), Some(version)) = (object, version) {
candidates.insert((object, version));
}
}
let expected = [0x1934, 0x1935]
.into_iter()
// Raw base-object listings include the trailing separator. Preserve
// the producer's candidate identity when draining its callback frame.
.map(|version| (format!("{object}/"), Uuid::from_u128(version).to_string()))
.collect();
assert_eq!(
candidates, expected,
"terminal disk errors must not discard buffered exact-version candidates"
);
for (path, metadata) in original_metadata {
assert_eq!(tokio::fs::read(path).await.expect("retained metadata"), metadata);
}
}
@@ -0,0 +1,710 @@
// Copyright 2026 RustFS Team
// Licensed under the Apache License, Version 2.0.
use super::*;
use rustfs_concurrency::adaptive::{FeedbackController, FeedbackObservation, NodeAdmission, ResourcePressure};
use std::num::NonZeroU16;
use std::sync::atomic::AtomicU64;
fn owner(open: bool) -> Arc<NodeAdmission> {
let owner = NodeAdmission::new(NonZeroU16::new(1).expect("one node unit"));
if open {
owner.update(
&mut FeedbackController::new(),
Some(FeedbackObservation {
generation: 1,
max_concurrency: 1,
pressure: ResourcePressure::Unknown,
completed_units: None,
}),
);
}
owner
}
#[tokio::test]
async fn scanner_physical_native_tail_releases_node_and_allows_unrelated_disk() {
use crate::storage_api::owner::{PhysicalDiskBinding, ecstore_spawn_blocking_with_resource_ownership};
use crate::workload_admission::with_scanner_physical_disk;
use rustfs_concurrency::adaptive::AdaptiveAdmission;
use rustfs_concurrency::adaptive::permit_set::AdaptivePermitSetOwner;
let node = owner(true);
let disk = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let other_disk = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let shared = AdaptivePermitSetOwner::new(NonZeroU16::MIN, vec![disk.clone(), other_disk.clone()])
.expect("one shared carrier for both devices");
let binding = Arc::new(PhysicalDiskBinding::default());
binding
.set_test_shared_capacity(shared.clone(), vec![disk.clone()])
.expect("physical fixture");
let (entered, started) = tokio::sync::oneshot::channel();
let (finish, blocked) = std::sync::mpsc::channel();
let cancel = CancellationToken::new();
let mut operation = Box::pin(with_scanner_physical_disk(
Some(binding),
run_scanner_node_operation(Some(&node), &cancel, async {
ecstore_spawn_blocking_with_resource_ownership(move || {
let _ = entered.send(());
blocked.recv_timeout(Duration::from_secs(2)).expect("native worker released");
})
.await
.map_err(std::io::Error::other)
}),
));
assert!(futures::poll!(&mut operation).is_pending());
tokio::time::timeout(Duration::from_secs(1), started)
.await
.expect("native worker starts")
.expect("start signal");
drop(operation);
assert_eq!(node.snapshot().active, 0, "cancelled waiter returns scheduling capacity");
assert_eq!(node.completed(), Some(0));
assert_eq!(disk.snapshot().active, 1, "native tail still owns its physical capacity");
let other = Arc::new(PhysicalDiskBinding::default());
other
.set_test_shared_capacity(shared, vec![other_disk])
.expect("independent disk using the same carrier pool");
with_scanner_physical_disk(
Some(other),
run_scanner_node_operation(Some(&node), &cancel, async { Ok::<_, std::io::Error>(()) }),
)
.await
.expect("independent disk admission")
.expect("independent operation");
finish.send(()).expect("finish native work");
tokio::time::timeout(Duration::from_secs(1), async {
while disk.snapshot().active != 0 {
tokio::task::yield_now().await;
}
})
.await
.expect("native completion returns capacity");
assert_eq!(disk.completed(), Some(0), "cancelled native work is not successful scanner progress");
}
#[tokio::test]
async fn scanner_physical_unknown_binding_never_falls_back_to_node_only() {
let node = owner(true);
let polled = std::cell::Cell::new(false);
let cancel = CancellationToken::new();
let result = crate::workload_admission::with_scanner_physical_disk(
Some(Arc::new(crate::storage_api::owner::PhysicalDiskBinding::default())),
run_scanner_node_operation(Some(&node), &cancel, async {
polled.set(true);
Ok::<_, std::io::Error>(())
}),
)
.await;
assert!(matches!(result, Err(ScannerError::ResourceDeferred)));
assert!(!polled.get());
assert_eq!(node.snapshot().active, 0);
}
#[tokio::test]
async fn scanner_physical_metadata_scope_retains_only_the_local_read() {
use crate::storage_api::owner::{PhysicalDiskBinding, ecstore_spawn_blocking_with_resource_ownership};
use crate::workload_admission::{run_scanner_metadata_operation, with_scanner_metadata_io, with_scanner_physical_disk};
use rustfs_concurrency::adaptive::AdaptiveAdmission;
for local_read in [true, false] {
let node = owner(true);
let disk = AdaptiveAdmission::new(NonZeroU16::MIN, 1);
let binding = Arc::new(PhysicalDiskBinding::default());
binding.set_test_capacity(vec![disk.clone()]).expect("physical fixture");
let (entered, started) = tokio::sync::oneshot::channel();
let (finish, blocked) = std::sync::mpsc::channel();
let cancel = CancellationToken::new();
let mut operation = Box::pin(with_scanner_physical_disk(
Some(binding),
run_scanner_metadata_operation(Some(&node), &cancel, async {
let io = async {
ecstore_spawn_blocking_with_resource_ownership(move || {
let _ = entered.send(());
blocked.recv_timeout(Duration::from_secs(2)).expect("finish native job");
})
.await
.map_err(std::io::Error::other)
};
if local_read {
with_scanner_metadata_io(io).await
} else {
io.await
}
}),
));
assert!(futures::poll!(&mut operation).is_pending());
tokio::time::timeout(Duration::from_secs(1), started)
.await
.expect("native start")
.expect("start signal");
drop(operation);
assert_eq!(node.snapshot().active, 0);
assert_eq!(
disk.snapshot().active,
u32::from(local_read),
"later lifecycle work cannot retain the metadata disk's unit"
);
finish.send(()).expect("release native job");
tokio::time::timeout(Duration::from_secs(1), async {
while disk.snapshot().active != 0 {
tokio::task::yield_now().await;
}
})
.await
.expect("physical ownership eventually returned");
}
}
#[tokio::test]
#[serial]
async fn scanner_node_admission_blocks_real_walk_before_open_or_metadata() {
let (mut scanner, root) = build_test_scanner().await;
let _guard = TestGuard {
temp_dir: Some(root.clone()),
};
write_test_object_metadata(&root, "bucket", "object").await;
let visits = super::metadata_visits::observe(&mut scanner);
let gate = owner(false);
scanner.node_admission = Some(gate.clone());
scanner.sleeper = DynamicSleeper::new(rustfs_config::ScannerSpeed::Fastest);
for name in ["bucket", "does-not-exist"] {
let result = scanner
.scan_folder(
CancellationToken::new(),
CachedFolder {
name: name.to_string(),
parent: None,
object_heal_prob_div: 1,
},
&mut DataUsageEntry::default(),
)
.await;
assert!(
matches!(result, Err(ScannerError::ResourceDeferred)),
"node gate must precede both real reads and missing-path lookup: {result:?}"
);
}
assert_eq!(visits.load(Ordering::Relaxed), 0);
assert_eq!(gate.snapshot().active, 0);
assert_eq!(gate.completed(), Some(0));
}
#[tokio::test]
#[serial]
async fn scanner_node_admission_shares_capacity_with_heal_owner() {
let (mut scanner, root) = build_test_scanner().await;
let _guard = TestGuard {
temp_dir: Some(root.clone()),
};
write_test_object_metadata(&root, "bucket", "object").await;
let gate = owner(true);
let heal_unit = gate.acquire().await;
scanner.node_admission = Some(gate.clone());
scanner.sleeper = DynamicSleeper::new(rustfs_config::ScannerSpeed::Fastest);
let folder = CachedFolder {
name: "bucket".to_string(),
parent: None,
object_heal_prob_div: 1,
};
let result = scanner
.scan_folder(CancellationToken::new(), folder.clone(), &mut DataUsageEntry::default())
.await;
assert!(
matches!(result, Err(ScannerError::ResourceDeferred)),
"scanner must not create a separate node budget: {result:?}"
);
assert_eq!(gate.snapshot().active, 1);
heal_unit.complete();
let mut entry = DataUsageEntry::default();
scanner
.scan_folder(CancellationToken::new(), folder, &mut entry)
.await
.expect("released heal capacity admits the real scan");
assert_eq!(
scanner
.new_cache
.size_recursive("bucket")
.expect("complete bucket tree")
.objects,
1
);
assert_eq!(gate.snapshot().active, 0);
assert!(
gate.completed().is_some_and(|count| count > 1),
"successful scanner work must reach shared feedback"
);
}
#[tokio::test(start_paused = true)]
async fn scanner_node_admission_precancellation_never_polls_storage() {
let gate = owner(false);
let cancel = CancellationToken::new();
let polled = std::cell::Cell::new(false);
cancel.cancel();
let operation = run_scanner_node_operation(Some(&gate), &cancel, async {
polled.set(true);
Ok::<_, std::io::Error>(())
});
assert!(matches!(operation.await, Err(ScannerError::Other(message)) if message == "Operation cancelled"));
assert!(!polled.get());
assert_eq!(gate.completed(), Some(0));
assert_eq!(gate.snapshot().active, 0);
}
#[tokio::test(start_paused = true)]
async fn scanner_node_admission_shrink_keeps_running_unit_and_preserves_errors() {
let gate = owner(true);
let cancel = CancellationToken::new();
let mut running = Box::pin(run_scanner_node_operation(
Some(&gate),
&cancel,
std::future::pending::<Result<(), std::io::Error>>(),
));
assert!(futures::poll!(&mut running).is_pending());
assert_eq!(gate.snapshot().active, 1);
gate.update(&mut FeedbackController::new(), None);
let polled = std::cell::Cell::new(false);
let deferred = run_scanner_node_operation(Some(&gate), &cancel, async {
polled.set(true);
Ok::<_, std::io::Error>(())
})
.await;
assert!(matches!(deferred, Err(ScannerError::ResourceDeferred)));
assert!(!polled.get());
assert_eq!(gate.snapshot().active, 1, "shrink cannot revoke admitted work");
drop(running);
assert_eq!(gate.snapshot().active, 0);
assert_eq!(gate.completed(), Some(0), "abandoned work is not progress");
gate.update(
&mut FeedbackController::new(),
Some(FeedbackObservation {
generation: 2,
max_concurrency: 1,
pressure: ResourcePressure::Unknown,
completed_units: None,
}),
);
let error = run_scanner_node_operation(Some(&gate), &cancel, async {
Err::<(), _>(std::io::Error::from(std::io::ErrorKind::PermissionDenied))
})
.await
.expect("resource admission succeeds")
.expect_err("storage error survives");
assert_eq!(error.kind(), std::io::ErrorKind::PermissionDenied);
assert_eq!(gate.completed(), Some(0));
run_scanner_node_operation(Some(&gate), &cancel, async { Ok::<_, std::io::Error>(()) })
.await
.expect("admitted")
.expect("successful work");
assert_eq!(gate.completed(), Some(1));
assert_eq!(gate.snapshot().active, 0);
}
struct ShrinkOnMetadata {
owner: Arc<NodeAdmission>,
visited: AtomicU64,
}
impl metrics::CounterFn for ShrinkOnMetadata {
fn increment(&self, value: u64) {
if self.visited.fetch_add(value, Ordering::SeqCst) == 0 {
self.owner.update(&mut FeedbackController::new(), None);
assert_eq!(self.owner.snapshot().active, 1, "shrink occurs during the admitted object");
}
}
fn absolute(&self, _: u64) {
panic!("visit counter must be cumulative");
}
}
struct ShrinkRecorder(Arc<ShrinkOnMetadata>);
impl metrics::Recorder for ShrinkRecorder {
fn describe_counter(&self, _: metrics::KeyName, _: Option<metrics::Unit>, _: metrics::SharedString) {}
fn describe_gauge(&self, _: metrics::KeyName, _: Option<metrics::Unit>, _: metrics::SharedString) {}
fn describe_histogram(&self, _: metrics::KeyName, _: Option<metrics::Unit>, _: metrics::SharedString) {}
fn register_counter(&self, key: &metrics::Key, _: &metrics::Metadata<'_>) -> metrics::Counter {
if key.name() == "rustfs_scanner_object_metadata_visits_total" {
metrics::Counter::from_arc(self.0.clone())
} else {
metrics::Counter::noop()
}
}
fn register_gauge(&self, _: &metrics::Key, _: &metrics::Metadata<'_>) -> metrics::Gauge {
metrics::Gauge::noop()
}
fn register_histogram(&self, _: &metrics::Key, _: &metrics::Metadata<'_>) -> metrics::Histogram {
metrics::Histogram::noop()
}
}
#[tokio::test]
#[serial]
async fn scanner_node_admission_mid_walk_shrink_preserves_partial_and_resumes_all_objects() {
let (scanner, root) = build_test_scanner().await;
let _guard = TestGuard {
temp_dir: Some(root.clone()),
};
for path in ["prefix/a/first", "prefix/b/second"] {
write_test_object_metadata(&root, "bucket", path).await;
}
let gate = owner(true);
let recorder = ShrinkRecorder(Arc::new(ShrinkOnMetadata {
owner: gate.clone(),
visited: AtomicU64::new(0),
}));
let scan = |cache| {
let cancel = CancellationToken::new();
scan_data_folder(
cancel.clone(),
ScannerCycleBudget::new(&cancel, Default::default()),
vec![scanner.local_disk.clone()],
scanner.local_disk.clone(),
cache,
None,
HealScanMode::Normal,
DynamicSleeper::new(rustfs_config::ScannerSpeed::Fastest),
)
};
let cache = DataUsageCache {
info: crate::DataUsageCacheInfo {
name: "bucket".to_string(),
skip_healing: true,
..Default::default()
},
..Default::default()
};
let baseline = scan(cache.clone()).await.expect("independent full walk");
scanner
.local_disk
.physical_admission()
.expect("local physical binding")
.set_test_capacity(vec![rustfs_concurrency::adaptive::AdaptiveAdmission::new(
std::num::NonZeroU16::MIN,
1,
)])
.expect("synthetic physical device");
let partial = crate::workload_admission::TEST_SCANNER_NODE_ADMISSION
.scope(gate.clone(), async {
let mut scanning = std::pin::pin!(scan(cache));
std::future::poll_fn(|cx| metrics::with_local_recorder(&recorder, || scanning.as_mut().poll(cx))).await
})
.await;
let Err(ScannerError::PartialCache(partial)) = partial else {
panic!("deferral below root cannot certify completion: {partial:?}");
};
assert!(!partial.info.snapshot_complete);
assert_eq!(
partial
.size_recursive("bucket")
.expect("partial keeps accounted work")
.objects,
1
);
assert_eq!(recorder.0.visited.load(Ordering::SeqCst), 1);
assert_eq!(gate.snapshot().active, 0);
gate.update(
&mut FeedbackController::new(),
Some(FeedbackObservation {
generation: 2,
max_concurrency: 1,
pressure: ResourcePressure::Unknown,
completed_units: None,
}),
);
let resumed = crate::workload_admission::TEST_SCANNER_NODE_ADMISSION
.scope(gate.clone(), scan(*partial))
.await
.expect("resume completes with restored node capacity");
assert!(resumed.info.snapshot_complete);
let expected = baseline.size_recursive("bucket").expect("baseline bucket");
let actual = resumed.size_recursive("bucket").expect("resumed bucket");
assert_eq!(
(actual.objects, actual.size, actual.versions),
(expected.objects, expected.size, expected.versions)
);
assert_eq!(actual.objects, 2);
assert_eq!(gate.snapshot().active, 0);
}
#[tokio::test(start_paused = true)]
async fn scanner_node_admission_does_not_park_while_cache_lease_may_be_held() {
let gate = owner(false);
let cancel = CancellationToken::new();
let polled = std::cell::Cell::new(false);
let mut operation = Box::pin(run_scanner_node_operation(Some(&gate), &cancel, async {
polled.set(true);
Ok::<_, std::io::Error>(())
}));
assert!(matches!(
futures::poll!(&mut operation),
std::task::Poll::Ready(Err(ScannerError::ResourceDeferred))
));
assert!(!polled.get());
assert_eq!(gate.snapshot().active, 0);
}
#[tokio::test]
#[serial]
async fn scanner_node_admission_early_deferral_keeps_old_tree_incomplete() {
check_early_admission_deferral(false, 1).await;
}
#[tokio::test]
#[serial]
async fn scanner_physical_early_deferral_keeps_real_walk_incomplete() {
check_early_admission_deferral(true, 0).await;
}
async fn check_early_admission_deferral(node_open: bool, physical_limit: u32) {
let (scanner, root) = build_test_scanner().await;
let _guard = TestGuard { temp_dir: Some(root) };
let mut cache = DataUsageCache {
info: crate::DataUsageCacheInfo {
name: "bucket".to_string(),
skip_healing: true,
snapshot_complete: true,
..Default::default()
},
..Default::default()
};
cache.replace(
"bucket",
"",
DataUsageEntry {
objects: 7,
size: 19,
..Default::default()
},
);
scanner
.local_disk
.physical_admission()
.expect("local binding")
.set_test_capacity(vec![rustfs_concurrency::adaptive::AdaptiveAdmission::new(
NonZeroU16::MIN,
physical_limit,
)])
.expect("independent physical capacity");
let gate = owner(node_open);
let cancel = CancellationToken::new();
let result = crate::workload_admission::TEST_SCANNER_NODE_ADMISSION
.scope(
gate.clone(),
scan_data_folder(
cancel.clone(),
ScannerCycleBudget::new(&cancel, Default::default()),
vec![scanner.local_disk.clone()],
scanner.local_disk,
cache,
None,
HealScanMode::Normal,
DynamicSleeper::new(rustfs_config::ScannerSpeed::Fastest),
),
)
.await;
let Err(ScannerError::PartialCache(partial)) = result else {
panic!("resource pressure is a partial census, not a disk error: {result:?}");
};
assert!(!partial.info.snapshot_complete);
let totals = partial.size_recursive("bucket").expect("old tree retained");
assert_eq!((totals.objects, totals.size), (7, 19));
assert_eq!(gate.completed(), Some(0));
}
#[tokio::test]
#[serial]
async fn scanner_node_admission_compacted_heal_child_preserves_deferral() {
let (mut scanner, root) = build_test_scanner().await;
let _guard = TestGuard {
temp_dir: Some(root.clone()),
};
let bucket = "bucket";
let object = "prefix/abandoned";
for index in 0..DATA_SCANNER_COMPACT_AT_FOLDERS {
tokio::fs::create_dir_all(root.join(bucket).join("prefix").join(format!("child-{index:04}")))
.await
.expect("create real compaction threshold");
}
let mut disks = vec![scanner.local_disk.clone()];
for (index, version) in [(2, 0x1935), (3, 0x1935), (4, 0x1934)] {
let disk_root = root.join(format!("disk-{index}"));
// Two partial frames exceed the coordinator channel's capacity, so
// discovery must consume the first before the finished frame can race it.
for child in ["first", "second"] {
let child_object = format!("{object}/{child}");
let metadata = metadata_for_object_version(bucket, &child_object, Some(Uuid::from_u128(version)));
write_test_object_metadata_bytes(&disk_root, bucket, &child_object, &metadata).await;
}
let endpoint = Endpoint::try_from(disk_root.to_string_lossy().as_ref()).expect("remote fixture endpoint");
disks.push(
new_disk(
&endpoint,
&DiskOption {
cleanup: false,
health_check: false,
},
)
.await
.expect("remote fixture disk"),
);
}
scanner.old_cache.info.name = bucket.to_string();
scanner.new_cache.info.name = bucket.to_string();
scanner.update_cache.info.name = bucket.to_string();
scanner.old_cache.replace(
&format!("{bucket}/{object}"),
"bucket/prefix",
DataUsageEntry {
objects: 1,
..Default::default()
},
);
scanner.heal_object_select = 1;
scanner.disks = disks;
scanner.disks_quorum = 2;
scanner.sleeper = DynamicSleeper::new(rustfs_config::ScannerSpeed::Fastest);
let gate = owner(true);
scanner.node_admission = Some(gate.clone());
let shrunk = Arc::new(AtomicBool::new(false));
let response_gate = gate.clone();
let response_shrunk = shrunk.clone();
let (_commands, mut heal_rx) = rustfs_heal_contracts::heal_channel::init_heal_channels().expect("isolated heal channel");
let responder = tokio::spawn(async move {
while let Some(command) = heal_rx.recv().await {
let rustfs_heal_contracts::heal_channel::HealReceiptCommand { request, response_tx } = command;
if request.object_version_id.is_some() {
response_shrunk.store(true, Ordering::SeqCst);
response_gate.update(&mut FeedbackController::new(), None);
}
let _ = response_tx.send(Ok(rustfs_heal_contracts::heal_channel::HealAdmissionReceipt {
result: HealAdmissionResult::Accepted,
task_id: request.id,
}));
}
});
let mut into = DataUsageEntry::default();
let folder = CachedFolder {
name: "bucket/prefix".to_string(),
parent: Some(hash_path(bucket)),
object_heal_prob_div: 1,
};
let result = {
let mut walk = std::pin::pin!(scanner.scan_folder(CancellationToken::new(), folder, &mut into));
let mut controller = FeedbackController::new();
let mut generation = 0;
std::future::poll_fn(|cx| {
// Keep telemetry fresh until the real heal-discovery receipt shrinks it.
if !shrunk.load(Ordering::SeqCst) {
generation += 1;
gate.update(
&mut controller,
Some(FeedbackObservation {
generation,
max_concurrency: 1,
pressure: ResourcePressure::Unknown,
completed_units: None,
}),
);
}
walk.as_mut().poll(cx)
})
.await
};
responder.abort();
assert!(into.compacted, "real directory count must enter the compacted recursion branch");
assert!(shrunk.load(Ordering::SeqCst), "remote orphan discovery must submit an exact-version heal");
assert!(
matches!(result, Err(ScannerError::ResourceDeferred)),
"compacted recursion must retain typed resource pressure for partial-cache handling: {result:?}"
);
assert_eq!(gate.snapshot().active, 0);
}
#[tokio::test]
#[serial]
async fn scanner_node_admission_progresses_during_continuous_heal_queue() {
let (mut scanner, root) = build_test_scanner().await;
let _guard = TestGuard {
temp_dir: Some(root.clone()),
};
write_test_object_metadata(&root, "bucket", "object").await;
for index in 0..65 {
tokio::fs::write(root.join("bucket/object").join(format!("part.{index}.integrity.orphan")), b"")
.await
.expect("sidecars require more than one directory refill");
}
scanner.is_erasure_mode = true;
scanner.budget = ScannerCycleBudget::new_with_progress_tracking(&CancellationToken::new(), Default::default());
let visits = super::metadata_visits::observe(&mut scanner);
let gate = owner(true);
scanner.node_admission = Some(gate.clone());
scanner.sleeper = DynamicSleeper::new(rustfs_config::ScannerSpeed::Fastest);
let heal_stop = CancellationToken::new();
let heal_stop_worker = heal_stop.clone();
let heal_gate = gate.clone();
let heal_completions = Arc::new(AtomicU64::new(0));
let completed = heal_completions.clone();
let (started, ready) = tokio::sync::oneshot::channel();
let heals = tokio::spawn(async move {
let mut held = heal_gate.acquire().await;
started.send(()).expect("scanner waits for the first heal");
loop {
let next = heal_gate.acquire();
tokio::pin!(next);
assert!(futures::poll!(&mut next).is_pending());
held.complete();
completed.fetch_add(1, Ordering::Relaxed);
held = tokio::select! {
biased;
_ = heal_stop_worker.cancelled() => return,
permit = &mut next => permit,
};
tokio::task::yield_now().await;
}
});
ready.await.expect("heal starts before scanner");
let (node_tx, mut node_rx) = mpsc::channel::<crate::workload_admission::ScannerNodeRequest>(1);
let folder = CachedFolder {
name: "bucket".to_string(),
parent: None,
object_heal_prob_div: 1,
};
let mut handoffs = 0;
let result = tokio::time::timeout(Duration::from_secs(2), async {
let mut totals = DataUsageEntry::default();
let scan = crate::workload_admission::with_scanner_node_requests(
Some(node_tx),
scanner.scan_folder(CancellationToken::new(), folder, &mut totals),
);
tokio::pin!(scan);
loop {
tokio::select! {
result = &mut scan => break result,
Some(request) = node_rx.recv() => {
handoffs += 1;
// This direct folder fixture owns no cache lock. Production
// coordinators release and revalidate their lease here.
let permit = request.owner.acquire().await;
assert_eq!(gate.snapshot().active, 1);
request.respond.send(Ok(crate::workload_admission::ScannerNodePermit::Node(permit)))
.unwrap_or_else(|_| panic!("real walk still waits"));
}
}
}
})
.await;
heal_stop.cancel();
heals.await.expect("continuous heal worker exits");
result
.expect("scanner must finish while Heal keeps queueing")
.expect("real directory scan");
assert_eq!(scanner.new_cache.size_recursive("bucket").expect("complete tree").objects, 1);
assert!(handoffs > 1, "the walk must yield through multiple shared-node turns");
assert_eq!(scanner.budget.entries_visited(), 67, "resumption must not replay directory entries");
assert_eq!(
visits.load(Ordering::Relaxed),
1,
"classification must account for the object exactly once"
);
assert!(heal_completions.load(Ordering::Relaxed) > 0, "Heal must also advance during the walk");
assert_eq!(gate.snapshot().active, 0);
}
+7 -3
View File
@@ -12,6 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use crate::data_usage_define::clone_cache_value;
use crate::data_usage_define::{DATA_USAGE_CACHE_KEY_FORMAT, DataUsageCacheRevisions};
use crate::scanner_budget::ScannerCycleBudget;
use crate::scanner_folder::{ScannerBucketPrefixScanScope, ScannerItem, scan_data_folder_scoped};
@@ -864,7 +865,7 @@ fn scanner_completed_set_segment_invalidation_proof(
proof: &Option<crate::DataUsageSegmentInvalidationProof>,
cold_zero_walk_reuse_candidate: bool,
) -> Option<crate::DataUsageSegmentInvalidationProof> {
proof.clone().map(|mut proof| {
clone_cache_value(proof).map(|mut proof| {
proof.cold_zero_walk_oracle = cold_zero_walk_reuse_candidate;
proof
})
@@ -1377,7 +1378,7 @@ pub(crate) fn cache_root_entry_info(cache: &DataUsageCache) -> std::result::Resu
.ok_or_else(|| ScannerError::Other(format!("scanner cache root is missing or corrupt: {}", cache.info.name)))?;
Ok(DataUsageEntryInfo {
name: cache.info.name.clone(),
name: clone_cache_value(&cache.info.name),
parent: DATA_USAGE_ROOT.to_string(),
entry,
bucket_incarnation: cache
@@ -1416,7 +1417,10 @@ fn apply_bucket_result_to_cache(
}
match result.bucket_incarnation {
Some(incarnation) if !incarnation.is_nil() => {
cache.info.scan_bucket_incarnations.insert(result.name.clone(), incarnation);
cache
.info
.scan_bucket_incarnations
.insert(clone_cache_value(&result.name), incarnation);
}
_ => {
cache.info.scan_bucket_incarnations.remove(&result.name);
+72 -6
View File
@@ -11,8 +11,9 @@
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
/// scanner cache locks and the cache snapshot persist/publish path.
use super::*;
/// scanner cache locks and the cache snapshot persist/publish path.
use crate::data_usage_define::clone_cache_value;
use crate::{ScannerConfigObjectDelete, ScannerObjectIO};
use rustfs_scanner_metrics::usage_candidate::{
ScannerUsageCandidateKind as CandidateKind, ScannerUsageCandidateOutcome as CandidateOutcome,
@@ -172,12 +173,68 @@ fn record_checkpoint_save_failure(cache_name: &str, checkpoint: &DataUsageCache,
#[derive(Debug)]
pub(crate) struct ScannerCacheLockGuards {
scoped: NamespaceLockGuard,
scoped: Option<NamespaceLockGuard>,
}
impl ScannerCacheLockGuards {
pub(crate) fn is_lock_lost(&self) -> bool {
self.scoped.is_lock_lost()
self.scoped.as_ref().is_none_or(NamespaceLockGuard::is_lock_lost)
}
/// Suspend the same traversal while yielding its cache lease. Lock order on
/// resumption is physical/node admission -> cache lease -> revision reads.
pub(crate) async fn resume_node_operation(
&mut self,
store: Arc<SetDisks>,
cache_name: &str,
source: DataUsageCacheSource,
revisions: &DataUsageCacheRevisions,
request: crate::workload_admission::ScannerNodeRequest,
) -> Result<()> {
if self.is_lock_lost() {
return Err(Error::other("scanner cache lock was lost before node admission"));
}
if request.respond.is_closed() {
return Ok(());
}
let crate::workload_admission::ScannerNodeRequest {
owner,
physical,
respond,
cancel,
} = request;
drop(self.scoped.take());
let resume = async {
use crate::workload_admission::ScannerNodePermit;
let permit = match physical {
Some(plan) => plan
.acquire(&owner)
.await
.map(|permit| ScannerNodePermit::Physical { permit, plan })
.ok_or(crate::ScannerError::ResourceDeferred),
None => Ok(ScannerNodePermit::Node(owner.acquire().await)),
};
*self = acquire_scanner_cache_locks(store.as_ref(), cache_name, source)
.await
.map_err(Error::other)?;
let current = crate::storage_api::ecstore_with_background_disk_io(
cancel.clone(),
DataUsageCache::read_revisions(store, cache_name),
)
.await?;
if self.is_lock_lost() || &current != revisions {
return Err(Error::other("scanner cache changed while waiting for node admission"));
}
// If the walk was dropped, send returns ownership and drops the
// permit without recording completed work.
let _ = respond.send(permit);
Ok(())
};
tokio::select! {
biased;
_ = cancel.cancelled() => Err(Error::other("scanner node admission cancelled")),
result = resume => result,
}
}
}
@@ -215,6 +272,15 @@ impl std::fmt::Display for ScannerCacheLockError {
}
}
impl std::error::Error for ScannerCacheLockError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Create { source, .. } => Some(source),
Self::Acquire { source, .. } => Some(source),
}
}
}
pub(crate) async fn acquire_scanner_cache_locks(
store: &SetDisks,
cache_name: &str,
@@ -237,7 +303,7 @@ pub(crate) async fn acquire_scanner_cache_locks(
source,
})?;
Ok(ScannerCacheLockGuards { scoped })
Ok(ScannerCacheLockGuards { scoped: Some(scoped) })
}
pub(super) async fn await_scanner_disk_shutdown<F>(scan: Pin<&mut F>)
@@ -758,7 +824,7 @@ pub(super) fn completed_usage_candidate(
scan_plan_digest: Some(result.info.scan_plan_digest.ok_or(CandidateOutcome::SetState)?.0),
complete: true,
tombstone: false,
segment_invalidation_proof: result.info.segment_invalidation_proof.clone(),
segment_invalidation_proof: clone_cache_value(&result.info.segment_invalidation_proof),
})
})
.collect::<std::result::Result<Vec<_>, CandidateOutcome>>()?;
@@ -954,7 +1020,7 @@ pub(super) fn observational_data_usage_info(
selected.info.scan_plan_digest.map(|digest| digest.0),
selected.info.last_update,
true,
selected.info.segment_invalidation_proof.clone(),
clone_cache_value(&selected.info.segment_invalidation_proof),
)
} else {
(
+33 -13
View File
@@ -11,8 +11,9 @@
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
/// ScannerIOCache implementation for SetDisks: bucket ordering, worker fan-out, merge, and publish.
use super::*;
/// ScannerIOCache implementation for SetDisks: bucket ordering, worker fan-out, merge, and publish.
use crate::data_usage_define::clone_cache_value;
pub(super) async fn collect_completed_bucket_results(
ctx: CancellationToken,
@@ -176,7 +177,7 @@ pub(super) fn prepare_scoped_set_scan(
lkg_last_update: old_cache.info.last_update,
lkg_leader_epoch: Some(old_cache.info.leader_epoch),
lkg_scan_plan_digest: old_cache.info.scan_plan_digest,
scan_bucket_incarnations: old_cache.info.scan_bucket_incarnations.clone(),
scan_bucket_incarnations: clone_cache_value(&old_cache.info.scan_bucket_incarnations),
..Default::default()
},
cache: HashMap::new(),
@@ -369,8 +370,8 @@ impl ScannerIOCache for SetDisks {
scan_plan_digest: Some(scan_plan_digest),
scan_coverage_digest: Some(bucket_coverage_digest),
cache_key_format: DATA_USAGE_CACHE_KEY_FORMAT,
segment_invalidation_proof: completed_segment_invalidation_proof.clone(),
scan_bucket_incarnations: current_bucket_incarnations.clone().unwrap_or_default(),
segment_invalidation_proof: clone_cache_value(&completed_segment_invalidation_proof),
scan_bucket_incarnations: clone_cache_value(&current_bucket_incarnations).unwrap_or_default(),
..Default::default()
},
cache: HashMap::new(),
@@ -425,8 +426,8 @@ impl ScannerIOCache for SetDisks {
"Scanner set state found no online disks"
);
reset_disk_bucket_scan_gauges(&pool_label, &set_label);
let lkg = old_cache.info.snapshot_complete.then(|| old_cache.clone());
let mut incomplete_scope = lkg.clone().unwrap_or_default();
let lkg = old_cache.info.snapshot_complete.then(|| clone_cache_value(&old_cache));
let mut incomplete_scope = clone_cache_value(&lkg).unwrap_or_default();
incomplete_scope.info.name = DATA_USAGE_ROOT.to_string();
incomplete_scope.info.next_cycle = want_cycle;
incomplete_scope.info.last_update = None;
@@ -525,8 +526,8 @@ impl ScannerIOCache for SetDisks {
"Scanner set state found no usable namespace scanner disks"
);
reset_disk_bucket_scan_gauges(&pool_label, &set_label);
let lkg = old_cache.info.snapshot_complete.then(|| old_cache.clone());
let mut incomplete_scope = lkg.clone().unwrap_or_default();
let lkg = old_cache.info.snapshot_complete.then(|| clone_cache_value(&old_cache));
let mut incomplete_scope = clone_cache_value(&lkg).unwrap_or_default();
incomplete_scope.info.name = DATA_USAGE_ROOT.to_string();
incomplete_scope.info.next_cycle = want_cycle;
incomplete_scope.info.last_update = None;
@@ -702,7 +703,7 @@ impl ScannerIOCache for SetDisks {
lkg_leader_epoch: old_cache.info.lkg_leader_epoch,
lkg_scan_plan_digest: old_cache.info.lkg_scan_plan_digest,
segment_invalidation_proof: None,
scan_bucket_incarnations: current_bucket_incarnations.clone().unwrap_or_default(),
scan_bucket_incarnations: clone_cache_value(&current_bucket_incarnations).unwrap_or_default(),
..Default::default()
},
cache: HashMap::new(),
@@ -728,7 +729,7 @@ impl ScannerIOCache for SetDisks {
for bucket in scan_order.iter() {
if let Some(c) = old_cache.find(&bucket.name) {
cache.replace(&bucket.name, DATA_USAGE_ROOT, c.clone());
cache.replace(&bucket.name, DATA_USAGE_ROOT, clone_cache_value(c));
}
if let Err(e) = bucket_tx.send(bucket.clone()).await {
@@ -1064,7 +1065,7 @@ impl ScannerIOCache for SetDisks {
// Lock order: scanner leader fence -> set-scoped per-bucket cache lock ->
// cache object read/write.
let cache_guard = match acquire_scanner_cache_locks(store_clone_clone.as_ref(), &cache_name, source).await {
let mut cache_guard = match acquire_scanner_cache_locks(store_clone_clone.as_ref(), &cache_name, source).await {
Ok(guard) => guard,
Err(e) => {
if e.is_contention() {
@@ -1294,9 +1295,10 @@ impl ScannerIOCache for SetDisks {
);
let (checkpoint_tx, mut checkpoint_rx) = mpsc::channel::<DataUsageCache>(1);
let (node_tx, mut node_rx) = crate::workload_admission::scanner_node_request_channel();
let scan_ctx = ctx_clone.child_token();
let scan = disk.clone().nsscanner_disk(
let scan = crate::workload_admission::with_scanner_node_requests(node_tx, disk.clone().nsscanner_disk(
scan_ctx.clone(),
budget_clone.clone(),
set_disk_inventory_clone.as_ref().clone(),
@@ -1307,7 +1309,7 @@ impl ScannerIOCache for SetDisks {
prefix_scan_scope,
checkpoint_tx: Some(checkpoint_tx),
},
);
));
tokio::pin!(scan);
let mut lock_watch = tokio::time::interval(SCANNER_CACHE_LOCK_POLL_INTERVAL);
lock_watch.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
@@ -1315,6 +1317,24 @@ impl ScannerIOCache for SetDisks {
let scan_result = loop {
tokio::select! {
result = &mut scan => break result,
request = async {
match node_rx.as_mut() {
Some(receiver) => receiver.recv().await,
None => std::future::pending().await,
}
} => {
let Some(request) = request else {
node_rx = None;
continue;
};
if let Err(error) = cache_guard.resume_node_operation(
store_clone_clone.clone(), &cache_name, source, &revisions, request,
).await {
scan_ctx.cancel();
await_scanner_disk_shutdown(scan.as_mut()).await;
break Err(error);
}
}
_ = lock_watch.tick() => {
if cache_guard.is_lock_lost() {
scan_ctx.cancel();
+19 -15
View File
@@ -11,9 +11,10 @@
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
/// ScannerIODisk implementation for Disk: get_size and the per-disk bucket scan.
use super::*;
use crate::UNKNOWN_TIER;
/// ScannerIODisk implementation for Disk: get_size and the per-disk bucket scan.
use crate::data_usage_define::clone_cache_value;
///
/// Seed [`SizeSummary::tier_stats`] from the cached tier-name list.
@@ -52,19 +53,22 @@ impl ScannerIODisk for Disk {
}
let metadata_object_path = item.object_path();
let data = match self.read_metadata(&item.bucket, &metadata_object_path).await {
Ok(data) => data,
Err(e) if DiskError::is_err_object_not_found(&e) || DiskError::is_err_version_not_found(&e) => {
return Err(StorageError::other(SCANNER_SKIP_FILE_ERROR.to_string()));
}
Err(e) => {
return Err(scanner_metadata_transient_error(
format!("failed to read metadata: {e}"),
&item.bucket,
&metadata_object_path,
));
}
};
let data =
match crate::workload_admission::with_scanner_metadata_io(self.read_metadata(&item.bucket, &metadata_object_path))
.await
{
Ok(data) => data,
Err(e) if DiskError::is_err_object_not_found(&e) || DiskError::is_err_version_not_found(&e) => {
return Err(StorageError::other(SCANNER_SKIP_FILE_ERROR.to_string()));
}
Err(e) => {
return Err(scanner_metadata_transient_error(
format!("failed to read metadata: {e}"),
&item.bucket,
&metadata_object_path,
));
}
};
item.transform_meta_dir();
let object_path = item.object_path();
@@ -163,7 +167,7 @@ impl ScannerIODisk for Disk {
} = options;
let done_drive = Metrics::time(Metric::ScanBucketDrive);
let drive_start = std::time::Instant::now();
let bucket = cache.info.name.clone();
let bucket = clone_cache_value(&cache.info.name);
let disk_path = self.path().to_string_lossy().to_string();
let source = match scan_mode {
HealScanMode::Deep => rustfs_scanner_metrics::metrics::ScannerWorkSource::Bitrot,
+1
View File
@@ -47,6 +47,7 @@ use uuid::Uuid;
mod checkpoint_coverage_identity;
mod collector_lifecycle;
mod listing_window_scope;
mod node_admission;
mod partial_bucket_carry;
mod per_bucket_certification;
mod scoped_entry_fallback;
@@ -0,0 +1,240 @@
// Copyright 2026 RustFS Team
// Licensed under the Apache License, Version 2.0.
use super::*;
use crate::workload_admission::ScannerNodeRequest;
use rustfs_concurrency::adaptive::{FeedbackController, FeedbackObservation, NodeAdmission, ResourcePressure};
use std::num::NonZeroU16;
use tokio::sync::oneshot;
#[derive(Clone, Copy, PartialEq, Eq)]
enum ResumeCase {
Unchanged,
Physical,
Updated,
Cancelled,
}
async fn check_cache_lease_resume(case: ResumeCase) {
let (_directory, store) = setup_local_scanner_store(1, 4).await;
let set = store.pools[0].disk_set[0].clone();
let source = DataUsageCacheSource::new(0, 0);
let cache_name = "node-resume/.usage-cache.bin";
let cache = DataUsageCache::default();
cache.save(set.clone(), cache_name).await.expect("seed cache");
let revisions = DataUsageCache::read_revisions(set.clone(), cache_name)
.await
.expect("initial revisions");
let gate = NodeAdmission::new(NonZeroU16::new(1).expect("one shared unit"));
gate.update(
&mut FeedbackController::new(),
Some(FeedbackObservation {
generation: 1,
max_concurrency: 1,
pressure: ResourcePressure::Unknown,
completed_units: None,
}),
);
let mut held = Some(gate.acquire().await);
let mut guard = acquire_scanner_cache_locks(set.as_ref(), cache_name, source)
.await
.expect("initial cache lease");
let ctx = CancellationToken::new();
let (respond, response) = oneshot::channel();
let physical = (case == ResumeCase::Physical).then(|| {
let binding = crate::storage_api::owner::PhysicalDiskBinding::default();
binding
.set_test_capacity(vec![rustfs_concurrency::adaptive::AdaptiveAdmission::new(NonZeroU16::MIN, 1)])
.expect("available physical device");
binding.plan().expect("fresh topology")
});
let request = ScannerNodeRequest {
owner: gate.clone(),
physical,
respond,
cancel: ctx.clone(),
};
let mut expected = revisions.clone();
let result = {
let resume = guard.resume_node_operation(set.clone(), cache_name, source, &revisions, request);
tokio::pin!(resume);
assert!(futures::poll!(&mut resume).is_pending(), "shared node is occupied");
let other = tokio::time::timeout(Duration::from_secs(1), acquire_scanner_cache_locks(set.as_ref(), cache_name, source))
.await
.expect("cache lock must be released before joining the node queue")
.expect("another cache worker can acquire the lease");
if case == ResumeCase::Updated {
let mut newer = cache;
newer.info.next_cycle = 1;
newer
.save(set.clone(), cache_name)
.await
.expect("another worker updates the cache");
expected = DataUsageCache::read_revisions(set.clone(), cache_name)
.await
.expect("new cache revisions");
assert_ne!(expected, revisions);
}
drop(other);
if case == ResumeCase::Physical {
tokio::time::pause();
tokio::time::advance(Duration::from_millis(350)).await;
assert!(futures::poll!(&mut resume).is_pending());
held.take().expect("held node unit").complete();
assert!(gate.try_acquire().is_none(), "fresh work cannot bypass the waiting physical scanner");
tokio::time::resume();
}
if case == ResumeCase::Cancelled {
ctx.cancel();
} else if case != ResumeCase::Physical {
held.take().expect("held node unit").complete();
}
tokio::time::timeout(Duration::from_secs(2), resume)
.await
.expect("node resumption must finish or observe cancellation")
};
if matches!(case, ResumeCase::Unchanged | ResumeCase::Physical) {
result.expect("unchanged cache resumes the same walk");
let permit = response
.await
.expect("coordinator replies")
.expect("coordinator returns the node permit");
assert_eq!(gate.snapshot().active, 1);
assert_eq!(gate.completed(), Some(1), "admitting scanner work must not count it as completed");
assert!(!guard.is_lock_lost());
let lock = set
.new_ns_lock(RUSTFS_META_BUCKET, &scanner_cache_lock_resource(cache_name, source))
.await
.expect("competing cache lock");
assert!(
lock.get_write_lock_quiet(Duration::from_millis(20)).await.is_err(),
"cache lease must be restored before the operation resumes"
);
permit.complete();
} else {
result.expect_err("cancelled or stale work cannot resume");
assert!(response.await.is_err(), "failed resumption must not authorize storage work");
if case == ResumeCase::Cancelled {
assert!(guard.is_lock_lost());
assert_eq!(gate.snapshot().active, 1, "the running heal keeps its unit");
drop(held.take());
assert_eq!(gate.completed(), Some(0));
} else {
assert_eq!(gate.completed(), Some(1));
}
}
drop(guard);
assert_eq!(gate.snapshot().active, 0);
assert!(gate.try_acquire().is_some(), "dropped requests must leave no queue or capacity leak");
assert_eq!(
DataUsageCache::read_revisions(set.clone(), cache_name)
.await
.expect("final revisions"),
expected
);
acquire_scanner_cache_locks(set.as_ref(), cache_name, source)
.await
.expect("final cache lease is available");
}
#[tokio::test]
#[serial]
async fn scanner_node_wait_yields_cache_lease_before_fifo_admission() {
check_cache_lease_resume(ResumeCase::Unchanged).await;
}
#[tokio::test]
#[serial]
async fn scanner_physical_node_wait_keeps_fifo_past_one_observation_interval() {
check_cache_lease_resume(ResumeCase::Physical).await;
}
#[tokio::test]
#[serial]
async fn scanner_node_wait_rejects_cache_changed_while_unlocked() {
check_cache_lease_resume(ResumeCase::Updated).await;
}
#[tokio::test]
#[serial]
async fn scanner_node_wait_cancellation_releases_queue_without_stealing_running_unit() {
check_cache_lease_resume(ResumeCase::Cancelled).await;
}
#[tokio::test]
#[serial]
async fn scanner_physical_expiry_yields_cache_then_resumes_with_partial_deferral() {
use crate::storage_api::owner::PhysicalDiskBinding;
use rustfs_concurrency::adaptive::AdaptiveAdmission;
let (_directory, store) = setup_local_scanner_store(1, 4).await;
let set = store.pools[0].disk_set[0].clone();
let source = DataUsageCacheSource::new(0, 0);
let cache_name = "physical-resume/.usage-cache.bin";
DataUsageCache::default()
.save(set.clone(), cache_name)
.await
.expect("seed cache");
let revisions = DataUsageCache::read_revisions(set.clone(), cache_name)
.await
.expect("revisions");
let node = NodeAdmission::new(NonZeroU16::MIN);
node.update(
&mut FeedbackController::new(),
Some(FeedbackObservation {
generation: 1,
max_concurrency: 1,
pressure: ResourcePressure::Unknown,
completed_units: None,
}),
);
let disk = AdaptiveAdmission::new(NonZeroU16::MIN, 0);
let binding = PhysicalDiskBinding::default();
binding
.set_test_capacity(vec![disk.clone()])
.expect("closed physical fixture");
let mut guard = acquire_scanner_cache_locks(set.as_ref(), cache_name, source)
.await
.expect("cache lease");
let (respond, response) = oneshot::channel();
let request = ScannerNodeRequest {
owner: node.clone(),
physical: binding.plan(),
respond,
cancel: CancellationToken::new(),
};
{
let resume = guard.resume_node_operation(set.clone(), cache_name, source, &revisions, request);
tokio::pin!(resume);
assert!(futures::poll!(&mut resume).is_pending());
assert_eq!(node.snapshot().active, 0, "a closed disk does not reserve node capacity");
let competing =
tokio::time::timeout(Duration::from_millis(100), acquire_scanner_cache_locks(set.as_ref(), cache_name, source))
.await
.expect("cache released during physical wait")
.expect("competing lease");
drop(competing);
binding
.set_test_capacity(vec![disk.clone()])
.expect("refresh invalidates the queued topology");
tokio::time::timeout(Duration::from_secs(2), resume)
.await
.expect("bounded wait")
.expect("coordinator resumes traversal");
}
assert!(
matches!(response.await.expect("coordinator replies"), Err(crate::ScannerError::ResourceDeferred)),
"the walk receives a deferral so it can return its partial cache"
);
assert!(!guard.is_lock_lost(), "cache lease restored before returning partial state");
assert_eq!(node.snapshot().active, 0);
assert_eq!(disk.snapshot().active, 0);
disk.set_sampled_limit(1, Duration::from_secs(1));
assert!(disk.try_acquire().is_some(), "expired topology releases its physical FIFO turn");
assert_eq!(
DataUsageCache::read_revisions(set, cache_name)
.await
.expect("final revisions"),
revisions
);
}
+11
View File
@@ -21,7 +21,11 @@ use rustfs_lock::NamespaceLockWrapper;
use serde::{Deserialize, Serialize};
use tokio::sync::Notify;
pub(crate) use rustfs_ecstore::spawn_blocking_with_resource_ownership as ecstore_spawn_blocking_with_resource_ownership;
pub(crate) use rustfs_ecstore::with_background_disk_io as ecstore_with_background_disk_io;
pub(crate) use rustfs_ecstore::with_background_resource_set as ecstore_with_background_resource_set;
#[cfg(test)]
pub(crate) use rustfs_ecstore::with_background_resource_unit as ecstore_with_background_resource_unit;
pub(crate) use s3s::dto::{
BucketLifecycleConfiguration as EcstoreBucketLifecycleConfiguration, LifecycleRuleFilter as EcstoreLifecycleRuleFilter,
@@ -164,6 +168,9 @@ pub(crate) mod owner {
#[cfg(test)]
pub(crate) use super::ecstore_set_object_store_resolver;
#[cfg(test)]
pub(crate) use super::ecstore_spawn_blocking_with_resource_ownership;
pub(crate) use super::ecstore_with_background_resource_set;
#[cfg(test)]
pub(crate) use super::storage_contracts::DiskSetSelector;
pub(crate) use super::{
MAX_SCANNER_PAUSE_BACKLOG_BYTES, ScannerPauseBacklogRetirementError, ScannerPauseBacklogRetirementPlan,
@@ -171,6 +178,10 @@ pub(crate) mod owner {
};
#[cfg(test)]
pub(crate) use super::{NativeScannerPauseBacklogWriteFault, SourceCleanupDeleteBarrier};
pub(crate) use rustfs_ecstore::api::disk::{
BackgroundEncodeBuffer as EcstoreBackgroundEncodeBuffer, BackgroundMemoryBudget as EcstoreBackgroundMemoryBudget,
PhysicalDiskBinding, PhysicalDiskPlan,
};
#[cfg(test)]
pub(crate) use rustfs_ecstore::api::object::{WriteCommitGuard, WriteCompletion};
+189 -1
View File
@@ -12,9 +12,197 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::sync::{Arc, LazyLock, RwLock};
use std::sync::{Arc, LazyLock, OnceLock, RwLock};
use rustfs_concurrency::adaptive::permit_set::{AdaptivePermitSet, NodeResourcePermit};
use rustfs_concurrency::adaptive::{NodeAdmission, NodePermit};
use rustfs_concurrency::{AdmissionState, WorkloadAdmissionRegistrySnapshot, WorkloadAdmissionSnapshotProvider, WorkloadClass};
use tokio::sync::{mpsc, oneshot};
use tokio_util::sync::CancellationToken;
use crate::ScannerError;
use crate::storage_api::owner::{PhysicalDiskBinding, PhysicalDiskPlan, ecstore_with_background_resource_set};
static SCANNER_NODE_ADMISSION: OnceLock<Arc<NodeAdmission>> = OnceLock::new();
tokio::task_local! {
static SCANNER_NODE_REQUESTS: ScannerNodeRequestSender;
static SCANNER_PHYSICAL_DISK: Option<Arc<PhysicalDiskBinding>>;
static SCANNER_METADATA_UNITS: AdaptivePermitSet;
}
pub(crate) struct ScannerNodeRequest {
pub(crate) owner: Arc<NodeAdmission>,
pub(crate) physical: Option<Arc<PhysicalDiskPlan>>,
pub(crate) respond: oneshot::Sender<Result<ScannerNodePermit, ScannerError>>,
pub(crate) cancel: CancellationToken,
}
pub(crate) enum ScannerNodePermit {
Node(NodePermit),
Physical {
permit: NodeResourcePermit,
plan: Arc<PhysicalDiskPlan>,
},
}
impl ScannerNodePermit {
pub(crate) fn complete(self) {
match self {
Self::Node(permit) => permit.complete(),
Self::Physical { permit, .. } => permit.complete(),
}
}
}
pub(crate) async fn with_scanner_physical_disk<F: std::future::Future>(
binding: Option<Arc<PhysicalDiskBinding>>,
operation: F,
) -> F::Output {
SCANNER_PHYSICAL_DISK.scope(binding, operation).await
}
pub(crate) type ScannerNodeRequestSender = mpsc::Sender<ScannerNodeRequest>;
pub(crate) type ScannerNodeRequestReceiver = mpsc::Receiver<ScannerNodeRequest>;
/// One suspended operation per walk; disabled admission allocates no channel.
pub(crate) fn scanner_node_request_channel() -> (Option<ScannerNodeRequestSender>, Option<ScannerNodeRequestReceiver>) {
if scanner_node_admission().is_none() {
return (None, None);
}
let (sender, receiver) = mpsc::channel(1);
(Some(sender), Some(receiver))
}
pub(crate) async fn with_scanner_node_requests<F: std::future::Future>(
sender: Option<ScannerNodeRequestSender>,
operation: F,
) -> F::Output {
match sender {
Some(sender) => SCANNER_NODE_REQUESTS.scope(sender, operation).await,
None => operation.await,
}
}
#[cfg(test)]
tokio::task_local! {
pub(crate) static TEST_SCANNER_NODE_ADMISSION: Arc<NodeAdmission>;
}
/// Install the heal manager's node owner before starting scanner workers.
/// Refuse replacement: existing scans must not outlive an independently renewed owner.
pub fn set_scanner_node_admission(owner: Arc<NodeAdmission>) -> bool {
SCANNER_NODE_ADMISSION.set(owner).is_ok()
}
pub(crate) fn scanner_node_admission() -> Option<Arc<NodeAdmission>> {
#[cfg(test)]
if let Ok(owner) = TEST_SCANNER_NODE_ADMISSION.try_with(Arc::clone) {
return Some(owner);
}
SCANNER_NODE_ADMISSION.get().cloned()
}
/// Try the fast path first. A scoped coordinator releases the cache lease before
/// joining the node queue and validates the cache before returning a permit.
/// Uncoordinated callers defer instead of parking under an unknown lock.
pub(crate) async fn run_scanner_node_operation<T, E>(
owner: Option<&Arc<NodeAdmission>>,
cancel: &CancellationToken,
operation: impl std::future::Future<Output = Result<T, E>>,
) -> Result<Result<T, E>, ScannerError> {
run_scanner_operation(owner, cancel, operation, false).await
}
/// Metadata parsing may subsequently dispatch lifecycle work on other disks.
/// Only its local metadata read inherits the scanner's physical resource set.
pub(crate) async fn run_scanner_metadata_operation<T, E>(
owner: Option<&Arc<NodeAdmission>>,
cancel: &CancellationToken,
operation: impl std::future::Future<Output = Result<T, E>>,
) -> Result<Result<T, E>, ScannerError> {
run_scanner_operation(owner, cancel, operation, true).await
}
pub(crate) async fn with_scanner_metadata_io<F: std::future::Future>(operation: F) -> F::Output {
match SCANNER_METADATA_UNITS.try_with(Clone::clone) {
Ok(units) => ecstore_with_background_resource_set(units, operation).await,
Err(_) => operation.await,
}
}
async fn run_scanner_operation<T, E>(
owner: Option<&Arc<NodeAdmission>>,
cancel: &CancellationToken,
operation: impl std::future::Future<Output = Result<T, E>>,
metadata: bool,
) -> Result<Result<T, E>, ScannerError> {
let Some(owner) = owner else {
return Ok(operation.await);
};
if cancel.is_cancelled() {
return Err(ScannerError::Other("Operation cancelled".to_string()));
}
let physical = SCANNER_PHYSICAL_DISK.try_with(|binding| binding.as_ref().and_then(|binding| binding.plan()));
let physical = match physical {
Ok(Some(plan)) => Some(plan),
Ok(None) => return Err(ScannerError::ResourceDeferred),
Err(_) => None,
};
let available = match &physical {
Some(plan) => plan.try_acquire(owner).map(|permit| ScannerNodePermit::Physical {
permit,
plan: plan.clone(),
}),
None => owner.try_acquire().map(ScannerNodePermit::Node),
};
let permit = match available {
Some(permit) => permit,
None => {
let sender = SCANNER_NODE_REQUESTS
.try_with(Clone::clone)
.map_err(|_| ScannerError::ResourceDeferred)?;
let acquire = async {
let (respond, response) = oneshot::channel();
sender
.send(ScannerNodeRequest {
owner: owner.clone(),
physical,
respond,
cancel: cancel.clone(),
})
.await
.map_err(|_| ScannerError::ResourceDeferred)?;
response.await.map_err(|_| ScannerError::ResourceDeferred)?
};
tokio::select! {
biased;
_ = cancel.cancelled() => return Err(ScannerError::Other("Operation cancelled".to_string())),
permit = acquire => permit?,
}
}
};
let result = match &permit {
ScannerNodePermit::Node(_) => operation.await,
ScannerNodePermit::Physical { permit, plan } => {
if !plan.is_current() {
return Err(ScannerError::ResourceDeferred);
}
let native = plan
.capture_native(permit.resources())
.map_err(|_| ScannerError::ResourceDeferred)?;
if metadata {
SCANNER_METADATA_UNITS.scope(native, operation).await
} else {
ecstore_with_background_resource_set(native, operation).await
}
}
};
if result.is_ok() {
permit.complete();
}
Ok(result)
}
type WorkloadSnapshotProviderRef = Arc<dyn WorkloadAdmissionSnapshotProvider + Send + Sync>;
+6 -3
View File
@@ -33,6 +33,9 @@ unknown-registry = "deny"
unknown-git = "deny"
allow-registry = ["https://github.com/rust-lang/crates.io-index"]
allow-git = [
# Ordered drain patch fixed while hotpath-rs/hotpath#759 is reviewed.
# owner: overtrue review: 2026-10
"https://github.com/rustfs/hotpath.git",
# Official s3s repository pinned for Iceberg SigV4 path encoding support.
# owner: GatewayJ review: 2026-10
"https://github.com/s3s-project/s3s",
@@ -43,9 +46,9 @@ allow-git = [
# owner: rustfs-maintainers review: 2026-10
"https://github.com/ramosbugs/openidconnect-rs",
"https://github.com/ramosbugs/oauth2-rs",
# RustFS io_uring backend is now consumed from crates.io at 0.3.0.
# owner: rustfs-maintainers review: 2026-10
# "https://github.com/rustfs/uring",
# Completion ownership patch fixed to rustfs/uring#20.
# owner: overtrue review: 2026-10
"https://github.com/rustfs/uring.git",
]
[bans]
+18 -1
View File
@@ -218,7 +218,7 @@ retain the failure and reconcile the service before continuing.
The `measure` response echoes the observed `evidence`, `fixed`, `build`,
`data_dir`, and `background`, plus `sample_count` (1..3600), `elapsed_seconds`,
and `metrics`. All metrics must be finite nonnegative numbers: `p99_ms`,
and `metrics`. All metrics must be finite nonnegative numbers: `p95_ms`, `p99_ms`,
`throughput_ops`, `rss_bytes`, `cpu_seconds`, `iops`, `rpc_count`,
`cache_clone_bytes`, `encode_bytes`, `save_bytes`, `oldest_age_seconds`,
`walk_objects`, `cold_walk_objects`, `healed_objects`, `errors`, and `requests`.
@@ -235,6 +235,23 @@ in the cell artifact directory; aggregate values alone do not establish their
measurement provenance. Missing production instrumentation is a pending gate,
not permission to report a fabricated zero.
The response must include `client_request_samples` with a cell-relative `path`,
the file's `sha256`, a `clock_id`, and integer `window_start_ns` and
`window_end_ns` from the client's monotonic clock. Each NDJSON row records a
unique `request_id`, that same `clock_id`, integer `start_ns` and `end_ns`, and
an `outcome` of `success`, `error`, `timeout`, or `cancelled`. Record each real
request from dispatch through completion, including failures and requests that
cross a boundary; do not discard samples to make a window pass. Validation
rejects crossing requests, mixed clocks, failed requests, and fewer than 100
samples. It recomputes nearest-rank p95/p99 from integer nanoseconds and checks
that request count, errors, throughput and elapsed time match the response.
The live telemetry collector starts its HTTP captures concurrently and records
their actual transport intervals. These receipts explicitly report
`synchronized_source_window=false`: concurrent collection does not establish
that the servers measured the same source window. The adapter must separately
prove alignment with the client window before claiming same-window acceptance.
For usage-cache saves, `encode_bytes` uses the window delta of
`rustfs_scanner_cache_serialized_bytes_total`. `save_bytes` uses the sum of
`rustfs_scanner_cache_save_payload_bytes_total` deltas for `cache="main"` and
+41 -5
View File
@@ -486,19 +486,55 @@ The OTLP exporter records these node-local histograms for completed manager admi
The two new histograms aggregate all observed admission sources and results without per-request labels. An ordinary uninterrupted admission partitions the phase into wait and hold. Admin forceStart releases these guards while cancelling overlaps, then reacquires them: the cancellation gap is included only in the original total phase, and the second acquisition contributes to wait. Cancellation subroutine work is outside these outer admission guard intervals. Recording happens after all three guards are released. They measure manager admission ownership, not namespace/object-lock waiting during repair. The status field `maxLockPhaseMicros` remains a manager-lifetime maximum, not a p95 sample. Keep same-window histogram counts and bucket resolution when deriving percentiles; missing observations do not mean zero latency. These metrics do not establish performance acceptance by themselves.
## Experimental heal node admission
## Cache copy allocation profiling
`RUSTFS_HEAL_ADAPTIVE_ADMISSION=true` enables Linux cgroup-v2 feedback for the four object-heal entry points owned by the heal manager. It is off by default while full SSD performance acceptance is pending. Unsupported platforms reject an enabled setting. This does not yet control scanner walks, bucket/format repair or physical devices, and it is not a replacement for existing worker, disk or byte budgets.
With `hotpath-alloc`, an installed Hotpath counting allocator, an active profiler guard, and the Hotpath Prometheus exporter, `hotpath_function_alloc_bytes_total{function="cache_model.owned_copy"}` counts requested allocation bytes while copying existing cache-model values. Use `HOTPATH_ALLOC_CUMULATIVE=0` for exclusive allocation accounting. The shared data-usage model and Scanner use the same fixed label; it includes their projections outside background scanning. Fresh listing/decode ownership, collection growth around individual copies, serialization, and new aggregate outputs are outside this scope. Moving or sharing an owner does not count as a payload copy.
The node samples its effective CPU quota/cpuset, ancestor memory headroom, CPU usage/throttling and CPU/memory/I/O pressure once per second. Sampling uses one blocking worker at a time. Resource reads retain their real start/end times: reads spanning over 20 ms or older than one second after all inputs return cannot renew admission. Counter rates use both read brackets conservatively. The concurrency cap is at most the rounded-up effective CPU count and 256 scheduling units; existing operator limits still apply independently. At most 1% memory headroom blocks new work. This is a pressure safeguard, not a memory reservation or a claim that concurrent repair buffers fit in the remaining bytes.
Scanner startup records an empty scope to establish a genuine zero observation. Export is asynchronous: a missing series is not zero, and separate scrapes do not establish a shared source window across metrics or nodes. This counter measures cumulative allocation requests, not retained bytes, peak memory, RSS, or copied logical string lengths. Compare builds with identical profiling features and record profiling overhead separately from default-build service performance.
## Experimental shared background admission
`RUSTFS_HEAL_ADAPTIVE_ADMISSION=true` enables Linux cgroup-v2 feedback shared by the heal manager’s four object-heal entry points and local Scanner namespace reads and object accounting. Scanner directory and local metadata reads acquire the physical leaves backing their local disk; object Heal admits the selected erasure sets across its requested pools. It is off by default while full SSD performance acceptance is pending. Unsupported platforms reject an enabled setting. Physical admission does not yet cover cache publication, lifecycle actions, the abandoned-child streaming list coordinator or independent bucket/format repair. Existing worker, disk and byte budgets still apply.
The node samples its effective CPU quota/cpuset, ancestor memory headroom, CPU usage/throttling and CPU/memory/I/O pressure once per second. Sampling uses one blocking worker at a time. Resource reads retain their real start/end times: reads spanning over 20 ms or older than one second after all inputs return cannot renew admission. Counter rates use both read brackets conservatively. The concurrency cap is at most the rounded-up effective CPU count and 256 scheduling units; existing operator limits still apply independently. At most 1% memory headroom blocks new work. This scheduling safeguard alone does not reserve memory for repair buffers.
The same sampler also controls the reusable staging storage of admitted buffered file readers and writers, whole-file results, and positioned-read results. Its process-wide cap is the smaller of 10% of the effective memory limit and the remaining headroom above 20% of that limit. Missing or expired samples block new allocations. A buffer reserves its payload size before allocation; growth counts both the old and replacement allocations until the old storage is freed. Shrinking the cap does not revoke an existing allocation, and reuse needs no new reservation. A denied allocation returns `WouldBlock`; reusable file streams accept no write bytes and preserve their read position. Cancellation and runtime shutdown keep storage charged until the actual native job releases it. Returned `Bytes` retain the allocation's charge through the last clone or slice. Standard-backend direct reads charge both the aligned bounce allocation and logical output while they coexist. Native io_uring reads retain their full allocation, including alignment padding, through the returned result; chunked reads also charge the assembled output while a native chunk is pending. Memory-budget denial does not switch either read path to a different backend.
Direct writers reserve their full aligned staging allocation when opened and retain it through native flushes, cancellation and runtime shutdown until the buffer is freed. An existing writer can reuse its allocation after the cap shrinks; a denied open does not disable direct I/O or fall back to buffered writes.
Admitted bitrot verification and reader scratch, Heal read and reconstruction shard buffers, legacy codec workspaces, and modern parity-verification scratch also reserve their payload storage before allocation. Cached legacy workspaces retain their charge while cached and can only reuse it for the same memory owner. Heal shard output retains its charge through the last `Bytes` clone or slice. HTTP read-file, walk-directory and namespace-scan response streams charge staging storage separately from the source reader. Short frames share the unused tail of that allocation; exhausted storage is reclaimed without allocating only after its last output reference drops. A retained output chunk keeps the full backing allocation charged and can prevent a replacement allocation after a cap reduction. EOF, read failure and dropping the stream release its own view, while output clones or slices keep their shared reservation.
Inline Heal readers share existing metadata bytes. Inline writers charge replacement storage while the old allocation remains live, and their returned `Bytes` retain the charge. Admitted remote Heal HTTP writers reserve each full 1 MiB staging allocation before copying input. Short flushes share that allocation with its unused tail; fully consumed storage is reused only after its last output reference drops. Queued chunks and the pending tail retain one shared reservation through the HTTP body and its final clone or slice, independently of the sender's disk capacity. A denied allocation accepts no new input bytes; existing storage remains usable after a cap reduction.
Admitted walk-directory and namespace-scan pipes reserve a fixed 1 MiB and 64 KiB ring, respectively, before their first nonempty write. Read and write ends retain the same allocation until both release it; consuming bytes does not refund storage still owned by either end. The ring never grows when writes wrap around. This also means a one-byte response reserves its full configured pipe capacity. Pipe admission failure terminates the response body, including legacy directory streams.
An authenticated namespace-scan POST also uses the receiver's enabled local admission policy when its sender omitted the background marker. Local admission preserves the authenticated wire class for nested RPC signatures. An admission-disabled receiver keeps its existing unmarked behavior; explicitly marked requests still require an enabled receiver.
Namespace frame payload and envelope encoding use the memory owner captured by the pipe before preflight scheduling capacity is dropped. Each replacement reserves old and new storage together. Both encoded allocations remain charged while a frame write is pending, and cancellation, timeout or write failure releases them. Encoding denial happens before writing the frame; it cannot fall back to unreserved encoding or become successful empty EOF. The existing MessagePack envelope, MAC and 2 MiB wire limit are preserved. Without a captured memory owner, encoding follows the existing path.
This cap does not yet cover other encode/decode paths, other RPC serializers and transport copies, allocator overhead or total process RSS. Already resident format-cache bytes are shared without a new payload allocation. The `rustfs_background_buffer_memory_limit_bytes` and `rustfs_background_buffer_memory_retained_bytes` gauges expose the sampled cap and retained payload. Like the node gauges, they can remain stale when the sampler stalls; the in-process admission deadline still expires after three seconds.
Each complete sample leases admission for three seconds. Missing resource caps block new work; missing or low-confidence feedback retains at most one maintenance unit. CPU use above 90%, throttling or pressure above 1% of elapsed time, memory headroom at most 10%, foreground errors/cancellations, a foreground queue or a comparable foreground p99 regression reduce concurrency immediately. Recovery requires CPU use at most 70%, throttling/stalls at most 0.1% of elapsed time, at least 20% memory headroom, fresh foreground evidence and actual successful work. The existing controller probes one additional unit only after eight consecutive healthy seconds and cooldown. Limits affect new admission and do not revoke active work.
The Scanner captures the same node owner once per bucket walk. Directory open/refill, entry type and size lookup, and object accounting try admission before their first storage poll; each scheduling unit ends when its operation future finishes. Physical leaves take FIFO turns in a stable order before node admission, and partial grants are returned before any resource wait. On contention, local and remote coordinators suspend the same traversal, release the bucket cache lease, and keep their physical/node FIFO position while topology remains current. The existing cycle budget and cancellation bound this wait. Topology validity is rechecked every 250 ms; expired or replaced plans yield a typed deferral. The coordinators reacquire the cache lease and verify both cache revisions before resuming with a permit or a typed deferral that preserves partial progress. A changed cache rejects the suspended result. Cancellation also drops the queued request.
The same sampling worker resolves local attachments through Linux sysfs, partitions and layered-device leaves. Device numbers plus `diskseq` distinguish replacements. Aliases share capacity; an attachment has its own three-second topology lease, so a healthy alias cannot renew a stale attachment. Missing identity or disk statistics prevents new physical admission. Rotational or unknown media stay at one unit; non-rotational media are bounded by a positive `nr_requests` and 256. Busy time above 90% or weighted queue time above one window cuts capacity; recovery also requires busy/queue time at most 50%, no outstanding I/O, valid foreground evidence and successful Scanner progress. Counter resets and gaps cannot justify growth. The node and physical controllers take turns probing, allowing at most one upward probe per sampling window.
Native directory work and the local metadata read retain only that disk's physical leaves after async cancellation. They do not retain the node scheduling unit. Later lifecycle/configuration work in object accounting does not inherit those physical units. Physical owners and carrier slots are preallocated during attachment refresh; object work does not discover devices or allocate a new resource group.
Resource-controlled peer requests carry `x-rustfs-rpc-background-admission: v1`, bound to their HTTP or gRPC signature. Stream retries preserve that classification without retaining the sender's disk quota. The receiver uses its existing shared node owner and admits the selected local physical device before storage locks or I/O. Repeated phases reuse the request's reservation, including when capacity is one. Missing capacity, an unknown UUID mapping, or an unavailable physical plan defers the request instead of probing other disks or reporting object absence. Cancellation returns node capacity; native work retains only its selected physical resources until it finishes.
An HTTP file response releases node admission after opening its reader, but records progress only after producing the declared byte count or reaching an unbounded EOF. Body errors, early EOF and cancellation do not credit successful preparation. The deferred progress record retains counter identities without keeping node slots or physical units.
Enable this experimental setting only after all peers support the admission marker. Older peers reject marked signatures, and a new peer with admission disabled rejects marked work; neither path silently runs it outside resource control. Requests without the marker retain their existing signature and foreground behavior. The receiver reports deferred admission as HTTP 503 or gRPC `RESOURCE_EXHAUSTED`.
Calls without a coordinating cache lease retain immediate resource deferral and useful partial-cache handling. Existing remote single-session ownership and local-disk admission shares remain in force. Disabled admission creates no request channel. Enabled walks use one bounded channel, with a response allocation and cache revision reads only on contention; the uncontended per-entry path adds no global lookup or task. The streaming abandoned-child coordinator requires a separate ownership boundary: retaining a node unit while it waits for a bounded-channel consumer that needs the same unit would deadlock at capacity one.
Foreground comparison uses the existing complete response-body windows and shared admission snapshots. Each active class needs at least 100 completions and a finite p99 interval; this permits the percentile rank but is not a statistical confidence guarantee. Comparisons require the same active classes and rates within 80–125% of a retained reference. A changed workload resets recovery. The reference is an observed workload window, not an independent background-off baseline. A p99 interval can justify stability only when its upper bound is within 5% of the reference lower bound; a lower bound above 105% of the reference upper bound is pressure. Overlapping/uncertain comparisons cannot justify growth. Missing queues, stale or repeated cuts, uncertain cut boundaries and low counts are not idle evidence.
Admission waits at most 250 ms before returning the existing typed transient deferral, before polling storage or taking namespace locks. No-lock callers never wait. Queued read repair uses the same node gate; direct foreground ECStore calls remain outside this heal-manager adapter. Cancellation releases the non-cloneable node scheduling unit. Native work retains separate physical-I/O units, and unsuccessful/cancelled/dry-run work cannot increment the node's completed-work counter. This scheduling counter is not an authoritative heal receipt and cannot discharge repair responsibility.
Heal admission waits at most 250 ms before returning the existing typed transient deferral, before polling storage or taking namespace locks. No-lock callers never wait. Queued read repair uses the same node gate; direct foreground ECStore calls remain outside this heal-manager adapter. Cancellation releases the non-cloneable node scheduling unit. Native work retains separate physical-I/O units, and unsuccessful/cancelled/dry-run work cannot increment the node's completed-work counter. This scheduling counter is not an authoritative heal receipt and cannot discharge repair responsibility.
The `rustfs_background_node_admission_limit` and `rustfs_background_node_admission_active` gauges describe the last published sample, not a live reservation proof. A blocked sampler can leave these exports stale even though the admission lease has expired. Full disk-owner feedback, shared scanner scheduling, escaping-buffer byte ownership and the matched-release SSD matrix/hot-shrink/72-hour soak remain required before enabling this by default or claiming optimal performance.
The `rustfs_background_node_admission_limit` and `rustfs_background_node_admission_active` gauges describe the last published sample, not a live reservation proof. A blocked sampler can leave these exports stale even though the admission lease has expired. Full-service validation of multi-disk Heal and remote RPC ownership, coverage of the remaining escaping buffers, and the matched-release SSD matrix/hot-shrink/72-hour soak remain required before enabling this by default or claiming optimal performance.
## Heal runtime controls
+1 -1
View File
@@ -301,7 +301,7 @@ moka = { workspace = true, features = ["sync"] }
xxhash-rust = { workspace = true, features = ["xxh3"] }
aws-sdk-s3 = { workspace = true, default-features = false, features = ["sigv4a", "default-https-client", "rt-tokio"] }
tokio-stream.workspace = true
tokio-util = { workspace = true, features = ["io", "compat", "time"] }
tokio-util = { workspace = true, features = ["io", "compat", "time", "rt"] }
tonic = { workspace = true, features = ["gzip", "deflate"] }
tower = { workspace = true, features = ["timeout"] }
tower-http = { workspace = true, features = ["trace", "compression-full", "cors", "catch-panic", "timeout", "limit", "request-id", "add-extension"] }
+54 -12
View File
@@ -1,8 +1,7 @@
// Copyright 2026 RustFS Team
// Licensed under the Apache License, Version 2.0.
//! Opt-in node scheduling feedback. This is not a physical-disk reservation or
//! a byte budget: native work and escaping buffers require their own owners.
//! Opt-in node scheduling, physical-device feedback and native stream-buffer caps.
use rustfs_concurrency::adaptive::{FeedbackController, FeedbackObservation, NodeAdmission, ResourcePressure};
use rustfs_concurrency::foreground::ForegroundWindow;
@@ -13,6 +12,7 @@ use std::time::{Duration, Instant};
use tokio_util::sync::CancellationToken;
use crate::cgroup_resources::live::ResourceSample;
use crate::storage_api::startup::background::{BackgroundMemoryBudget, ECStore, PhysicalDiskSample, PhysicalDiskSampler};
const SAMPLE_INTERVAL: Duration = Duration::from_secs(1);
pub(crate) const MAXIMUM: NonZeroU16 = NonZeroU16::new(256).expect("fixed nonzero scheduling bound");
@@ -237,46 +237,72 @@ fn resource_pressure(current: TimedResources, old: Option<TimedResources>) -> Re
}
}
struct Lifetime(Arc<NodeAdmission>);
// Match the resource policy's low-pressure headroom, and cap this staging
// storage at ten percent of the effective memory limit. Headroom is already net
// of retained allocations; using it directly conservatively includes other work.
fn memory_capacity(sample: ResourceSample) -> u64 {
(sample.memory_limit / 10).min(sample.memory_headroom.saturating_sub(sample.memory_limit / 5))
}
struct Lifetime(Arc<NodeAdmission>, Arc<BackgroundMemoryBudget>);
impl Drop for Lifetime {
fn drop(&mut self) {
self.1.set_sampled_limit(0, Duration::ZERO);
self.0.update(&mut FeedbackController::new(), None);
}
}
pub(crate) fn start(
owner: Arc<NodeAdmission>,
memory: Arc<BackgroundMemoryBudget>,
store: Arc<ECStore>,
cancel: CancellationToken,
provider: Arc<dyn WorkloadAdmissionSnapshotProvider + Send + Sync>,
) {
let lifetime = Lifetime(owner);
let lifetime = Lifetime(owner, memory);
tokio::spawn(async move {
let lifetime = lifetime;
let mut policy = Policy::default();
let mut controller = FeedbackController::new();
let mut physical = PhysicalDiskSampler::default();
let mut tick = tokio::time::interval(SAMPLE_INTERVAL);
tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
loop {
tokio::select! { biased; _ = cancel.cancelled() => break, _ = tick.tick() => {} }
// One blocking sample at a time. A stuck reader cannot create an
// unbounded worker queue; the previous admission lease expires.
let sample = tokio::select! {
let disks = PhysicalDiskSample::collect(&store);
let samples = tokio::select! {
biased;
_ = cancel.cancelled() => break,
result = tokio::task::spawn_blocking(|| {
result = tokio::task::spawn_blocking(move || {
let began = Instant::now();
let sample = crate::cgroup_resources::live::sample()?;
Ok::<_, std::io::Error>(TimedResources { sample, began, ended: Instant::now() })
}) => result.ok().and_then(Result::ok),
let sample = crate::cgroup_resources::live::sample().ok()
.map(|sample| TimedResources { sample, began, ended: Instant::now() });
(sample, disks.read())
}) => result.ok(),
};
let foreground = crate::foreground_observation::latest_window();
let registry = provider.workload_admission_snapshot();
// Recheck age after every input read, including a delayed provider.
let observation = policy.observe(Instant::now(), sample, foreground, &registry);
let decision = lifetime.0.update(&mut controller, observation);
let observation =
policy.observe(Instant::now(), samples.as_ref().and_then(|(sample, _)| *sample), foreground, &registry);
let memory_limit = observation
.and_then(|_| samples.as_ref().and_then(|(sample, _)| *sample))
.map_or(0, |sample| memory_capacity(sample.sample));
lifetime.1.set_sampled_limit(memory_limit, Duration::from_secs(3));
let decision = lifetime
.0
.update_with_probe(&mut controller, observation, physical.node_may_probe());
if let Some((_, disks)) = samples {
physical.publish(disks, observation.map_or(ResourcePressure::Unknown, |sample| sample.pressure));
}
if rustfs_io_metrics::metrics_enabled() {
metrics::gauge!("rustfs_background_node_admission_limit").set(f64::from(decision.limit));
metrics::gauge!("rustfs_background_node_admission_active").set(f64::from(lifetime.0.snapshot().active));
let memory = lifetime.1.snapshot();
metrics::gauge!("rustfs_background_buffer_memory_limit_bytes").set(memory.limit as f64);
metrics::gauge!("rustfs_background_buffer_memory_retained_bytes").set(memory.used as f64);
}
}
});
@@ -302,6 +328,19 @@ mod tests {
}
}
#[test]
fn buffer_capacity_respects_effective_limit_and_headroom() {
let mut sample = resources();
for (headroom, expected) in [(0, 0), (199, 0), (200, 0), (201, 1), (250, 50), (300, 100), (1000, 100)] {
sample.memory_headroom = headroom;
assert_eq!(memory_capacity(sample), expected, "headroom={headroom}");
}
sample.memory_limit = 500;
assert_eq!(memory_capacity(sample), 50, "a hot shrink changes the absolute cap");
sample.memory_limit = 0;
assert_eq!(memory_capacity(sample), 0, "missing effective capacity cannot grant memory");
}
fn timed(at: Instant, sample: ResourceSample) -> TimedResources {
TimedResources {
sample,
@@ -656,7 +695,10 @@ mod tests {
}),
);
let active = owner.acquire().await;
drop(Lifetime(owner.clone()));
let memory = BackgroundMemoryBudget::new();
memory.set_sampled_limit(1024, Duration::from_secs(3));
drop(Lifetime(owner.clone(), memory.clone()));
assert_eq!(memory.snapshot().limit, 0, "sampler shutdown closes new buffer admission");
assert_eq!(owner.snapshot().limit, 0);
assert_eq!(owner.snapshot().active, 1);
assert!(owner.try_acquire().is_none());
+1 -3
View File
@@ -71,9 +71,7 @@ pub fn run() -> ExitCode {
};
}
rustfs::startup_entrypoint::run_process();
ExitCode::SUCCESS
rustfs::startup_entrypoint::run_process()
}
fn run_offline(arguments: &[String]) -> Result<(), String> {
+2 -2
View File
@@ -61,10 +61,10 @@ static GLOBAL: hotpath::CountingAllocator<std::alloc::System> = hotpath::Countin
#[global_allocator]
static GLOBAL: rustfs_mimalloc::MiMalloc = rustfs_mimalloc::MiMalloc;
fn main() {
fn main() -> std::process::ExitCode {
let _hotpath_guard = hotpath::HotpathGuardBuilder::new("main").build();
rustfs::startup_entrypoint::run_process();
rustfs::startup_entrypoint::run_process()
}
#[cfg(all(test, feature = "hotpath", feature = "hotpath-alloc", not(target_os = "windows")))]
+1411 -101
View File
File diff suppressed because it is too large Load Diff
+295 -6
View File
@@ -36,11 +36,17 @@ mod http3;
use tracing::warn;
const EVENT_SERVER_SHUTDOWN_STATE: &str = "server_shutdown_state";
const LOG_COMPONENT_SERVER: &str = "server";
const LOG_SUBSYSTEM_TRANSPORT: &str = "transport";
// Items used by main.rs (binary crate) and/or embedded.rs — must be fully pub.
pub(crate) use audit::apply_audit_module_switch_for_context;
pub use audit::{is_audit_module_enabled, refresh_audit_module_enabled, start_audit_system, stop_audit_system};
pub use event::{init_event_notifier, is_notify_module_enabled, refresh_notify_module_enabled, shutdown_event_notifier};
pub use http::start_http_server;
#[cfg(test)]
pub(crate) use http::test_retire_http12_scope;
pub use prefix::LOGO;
pub use runtime::build_tokio_runtime;
pub use service_state::SHUTDOWN_TIMEOUT;
@@ -102,16 +108,154 @@ pub(crate) fn strip_valid_port_suffix(host: &str) -> &str {
}
}
/// Observations produced by the actual accepted HTTP/1-2 task owner.
#[derive(Debug, Clone)]
pub(crate) struct Http12ShutdownOutcome {
deadline_missed: bool,
first_child_error: Option<std::sync::Arc<tokio::task::JoinError>>,
child_failures: u64,
unverified_h3: bool,
unverified_tls_reloader: bool,
tls_reloader: Option<tls_material::TlsReloadRetirement>,
}
impl Http12ShutdownOutcome {
#[cfg(test)]
pub(crate) fn with_test_tls_receipt(mut self, receipt: tls_material::TlsReloadRetirement) -> Self {
self.unverified_tls_reloader = true;
self.tls_reloader = Some(receipt);
self
}
pub(crate) fn deadline_missed(&self) -> bool {
self.deadline_missed
}
pub(crate) fn verified_scope_complete(&self) -> bool {
!self.unverified_h3
&& (!self.unverified_tls_reloader || self.tls_reloader.as_ref().is_some_and(|receipt| receipt.group_joined()))
}
fn tls_reload_error(&self) -> Option<&std::sync::Arc<tokio::task::JoinError>> {
self.tls_reloader.as_ref().and_then(|receipt| receipt.join_error())
}
pub(crate) fn late_failure(&self) -> Option<ShutdownWaitError> {
self.tls_reload_error().map(|_| ShutdownWaitError::Http12(self.clone()))
}
fn failed(&self) -> bool {
self.deadline_missed || self.child_failures != 0 || self.tls_reload_error().is_some()
}
}
#[derive(Debug, Clone)]
pub(crate) enum ShutdownReceipt {
OuterTaskJoined,
Http12ChildrenRetired(Http12ShutdownOutcome),
}
/// A failed or unknown observation retains the actual task error, when known.
#[derive(Debug, Clone)]
pub(crate) enum ShutdownWaitError {
Join(std::sync::Arc<tokio::task::JoinError>),
MissingTask,
Http12(Http12ShutdownOutcome),
}
impl ShutdownWaitError {
fn reason(&self) -> &'static str {
match self {
Self::Join(error) if error.is_cancelled() => "join_cancelled",
Self::Join(_) => "join_panicked",
Self::MissingTask => "join_unconfirmed",
Self::Http12(outcome) if outcome.deadline_missed => "http12_deadline_missed",
Self::Http12(outcome) if outcome.tls_reload_error().is_some_and(|error| error.is_cancelled()) => {
"tls_reload_join_cancelled"
}
Self::Http12(outcome) if outcome.tls_reload_error().is_some() => "tls_reload_join_panicked",
Self::Http12(_) => "http12_child_join_failed",
}
}
pub(crate) fn http12_outcome(&self) -> Option<&Http12ShutdownOutcome> {
match self {
Self::Http12(outcome) => Some(outcome),
_ => None,
}
}
}
impl std::fmt::Display for ShutdownWaitError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "server task shutdown {}", self.reason())?;
if let Self::Http12(outcome) = self {
write!(f, " ({} child join failures)", outcome.child_failures)?;
}
Ok(())
}
}
impl std::error::Error for ShutdownWaitError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Join(error) => Some(error.as_ref()),
Self::Http12(outcome) => outcome
.first_child_error
.as_ref()
.or_else(|| outcome.tls_reload_error())
.map(|error| error.as_ref() as _),
Self::MissingTask => None,
}
}
}
enum ShutdownTask {
Outer(tokio::task::JoinHandle<()>),
Http12(tokio::task::JoinHandle<Http12ShutdownOutcome>),
}
impl ShutdownTask {
#[cfg(test)]
fn is_finished(&self) -> bool {
match self {
Self::Outer(task) => task.is_finished(),
Self::Http12(task) => task.is_finished(),
}
}
#[cfg(test)]
fn id(&self) -> tokio::task::Id {
match self {
Self::Outer(task) => task.id(),
Self::Http12(task) => task.id(),
}
}
}
pub struct ShutdownHandle {
shutdown_tx: Option<tokio::sync::broadcast::Sender<()>>,
task_handle: Option<tokio::task::JoinHandle<()>>,
task_handle: Option<ShutdownTask>,
failure: Option<ShutdownWaitError>,
}
impl ShutdownHandle {
pub fn new(shutdown_tx: tokio::sync::broadcast::Sender<()>, task_handle: tokio::task::JoinHandle<()>) -> Self {
Self {
shutdown_tx: Some(shutdown_tx),
task_handle: Some(task_handle),
task_handle: Some(ShutdownTask::Outer(task_handle)),
failure: None,
}
}
pub(crate) fn new_http12(
shutdown_tx: tokio::sync::broadcast::Sender<()>,
task_handle: tokio::task::JoinHandle<Http12ShutdownOutcome>,
) -> Self {
Self {
shutdown_tx: Some(shutdown_tx),
task_handle: Some(ShutdownTask::Http12(task_handle)),
failure: None,
}
}
@@ -126,11 +270,41 @@ impl ShutdownHandle {
self.wait().await;
}
/// A cancelled borrowed observer leaves the same original task in this handle.
/// A legacy outer join never certifies HTTP children. Failed observations are sticky.
pub(crate) async fn wait_checked(&mut self) -> Result<ShutdownReceipt, ShutdownWaitError> {
if let Some(error) = &self.failure {
return Err(error.clone());
}
let Some(task) = self.task_handle.as_mut() else {
return Err(ShutdownWaitError::MissingTask);
};
let result = match task {
ShutdownTask::Outer(task) => task.await.map(|()| ShutdownReceipt::OuterTaskJoined),
ShutdownTask::Http12(task) => task.await.map(ShutdownReceipt::Http12ChildrenRetired),
}
.map_err(|error| ShutdownWaitError::Join(std::sync::Arc::new(error)))
.and_then(|receipt| match receipt {
ShutdownReceipt::Http12ChildrenRetired(outcome) if outcome.failed() => Err(ShutdownWaitError::Http12(outcome)),
receipt => Ok(receipt),
});
let _ = self.task_handle.take();
if let Err(error) = &result {
self.failure = Some(error.clone());
}
result
}
pub async fn wait(mut self) {
if let Some(task_handle) = self.task_handle.take()
&& let Err(err) = task_handle.await
{
warn!(?err, "Server task join failed during shutdown");
if let Err(err) = self.wait_checked().await {
warn!(
event = EVENT_SERVER_SHUTDOWN_STATE,
component = LOG_COMPONENT_SERVER,
subsystem = LOG_SUBSYSTEM_TRANSPORT,
state = "join_failed",
reason = err.reason(),
"Server task join failed during shutdown"
);
}
}
}
@@ -142,3 +316,118 @@ impl Drop for ShutdownHandle {
}
}
}
#[cfg(test)]
mod shutdown_handle_tests {
use super::{ShutdownHandle, ShutdownWaitError};
use std::{future::Future, task::Poll, time::Duration};
use tokio::sync::{broadcast, oneshot};
#[tokio::test]
async fn checked_wait_cancelled_observer_keeps_original_held_task() {
let (tx, mut rx) = broadcast::channel(1);
let (entered_tx, entered_rx) = oneshot::channel();
let (release_tx, release_rx) = oneshot::channel();
let (done_tx, done_rx) = oneshot::channel();
let task = tokio::spawn(async move {
let _ = rx.recv().await;
let _ = entered_tx.send(());
let _ = release_rx.await;
let _ = done_tx.send(());
});
let id = task.id();
let mut handle = ShutdownHandle::new(tx, task);
handle.signal();
let entered = tokio::time::timeout(Duration::from_secs(5), entered_rx)
.await
.is_ok_and(|r| r.is_ok());
let mut observer = Box::pin(handle.wait_checked());
let pending = std::future::poll_fn(|cx| Poll::Ready(observer.as_mut().poll(cx).is_pending())).await;
drop(observer);
let retained = handle.task_handle.as_ref().is_some_and(|task| task.id() == id);
let _ = release_tx.send(());
let completed = tokio::time::timeout(Duration::from_secs(5), done_rx)
.await
.is_ok_and(|r| r.is_ok());
let joined = handle.wait_checked().await;
assert!(entered && pending && retained && completed);
joined.expect("same original task joined after observer cancellation");
assert!(handle.task_handle.is_none());
}
#[tokio::test]
async fn checked_wait_preserves_actual_panic_and_cancel_sources() {
let (tx, _rx) = broadcast::channel(1);
let task = tokio::spawn(async { panic!("expected checked server task panic") });
let panic_id = task.id();
let mut handle = ShutdownHandle::new(tx, task);
let panic = handle.wait_checked().await.expect_err("actual panic must survive");
assert!(panic.to_string().contains("join_panicked"));
assert!(
!panic.to_string().contains("expected checked server task panic"),
"Display must not emit the panic payload"
);
let source = std::error::Error::source(&panic).expect("actual JoinError source");
let source = source.downcast_ref::<tokio::task::JoinError>().expect("typed JoinError");
assert!(source.is_panic());
assert_eq!(source.id(), panic_id);
let (tx, _rx) = broadcast::channel(1);
let (entered_tx, entered_rx) = oneshot::channel();
let task = tokio::spawn(async move {
let _ = entered_tx.send(());
std::future::pending::<()>().await;
});
let cancel_id = task.id();
let abort = task.abort_handle();
let mut handle = ShutdownHandle::new(tx, task);
let entered = tokio::time::timeout(Duration::from_secs(5), entered_rx)
.await
.is_ok_and(|r| r.is_ok());
abort.abort();
let cancelled = handle.wait_checked().await.expect_err("actual cancellation must survive");
let source = std::error::Error::source(&cancelled).expect("actual JoinError source");
let source = source.downcast_ref::<tokio::task::JoinError>().expect("typed JoinError");
assert!(entered && source.is_cancelled());
assert_eq!(source.id(), cancel_id);
}
#[tokio::test]
async fn checked_wait_missing_or_consumed_receipt_stays_unknown() {
let mut missing = ShutdownHandle {
shutdown_tx: None,
task_handle: None,
failure: None,
};
assert!(matches!(missing.wait_checked().await, Err(ShutdownWaitError::MissingTask)));
let (tx, _rx) = broadcast::channel(1);
let mut completed = ShutdownHandle::new(tx, tokio::spawn(async {}));
completed.wait_checked().await.expect("actual task joined once");
assert!(matches!(completed.wait_checked().await, Err(ShutdownWaitError::MissingTask)));
}
#[tokio::test]
async fn checked_outer_join_does_not_certify_accepted_child_retirement() {
let (tx, _rx) = broadcast::channel(1);
let (child_tx, child_rx) = oneshot::channel();
let (entered_tx, entered_rx) = oneshot::channel();
let (release_tx, release_rx) = oneshot::channel();
let outer = tokio::spawn(async move {
let child = tokio::spawn(async move {
let _ = entered_tx.send(());
let _ = release_rx.await;
});
let _ = child_tx.send(child);
});
let mut handle = ShutdownHandle::new(tx, outer);
let child = child_rx.await.expect("actual accepted child owner");
let entered = tokio::time::timeout(Duration::from_secs(5), entered_rx)
.await
.is_ok_and(|r| r.is_ok());
let outer_joined = handle.wait_checked().await.is_ok();
let child_still_held = !child.is_finished();
let _ = release_tx.send(());
child.await.expect("accepted child joined only after its own release");
assert!(entered && outer_joined && child_still_held);
}
}
+577 -25
View File
@@ -21,6 +21,10 @@
//! TLS acceptor directly from the pre-loaded server material — no double reads.
use crate::startup_runtime_sources;
use futures::{
FutureExt,
future::{BoxFuture, Shared},
};
use rustfs_common::MtlsIdentityPem;
use rustfs_config::{
DEFAULT_SERVER_MTLS_ENABLE, DEFAULT_TLS_KEYLOG, DEFAULT_TLS_RELOAD_ENABLE, DEFAULT_TLS_RELOAD_INTERVAL,
@@ -42,6 +46,7 @@ use std::sync::RwLock;
use std::sync::{Arc, LazyLock, Mutex, Weak};
use std::time::Duration;
use tokio_rustls::TlsAcceptor;
use tokio_util::task::TaskTracker;
use tracing::{debug, info, warn};
const LOG_COMPONENT_TLS: &str = "tls";
@@ -71,20 +76,47 @@ const SYSTEM_CA_PATHS: &[&str] = &[
/// Returns the fully enriched snapshot ready for both outbound publishing and
/// acceptor construction.
pub async fn load_tls_material(tls_path: &str) -> Result<rustfs_tls_runtime::TlsMaterialSnapshot, TlsMaterialError> {
load_tls_material_observed(
tls_path,
#[cfg(test)]
None,
)
.await
}
async fn load_tls_material_observed(
tls_path: &str,
#[cfg(test)] read_control: Option<Arc<ReloadReadTestControl>>,
) -> Result<rustfs_tls_runtime::TlsMaterialSnapshot, TlsMaterialError> {
if tls_path.is_empty() {
return Err(TlsMaterialError::Io("TLS path is empty".into()));
}
let tls_source = TlsSource::from_directory(tls_path);
let mut snapshot = tokio::task::spawn_blocking(move || {
let pinned_source = tls_source.snapshot().map_err(map_runtime_tls_error)?;
if let Some(snapshot) =
super::tls_publish::configured_server_snapshot(&tls_source.base_dir.to_string_lossy(), &pinned_source)
.map_err(|_| TlsMaterialError::Io("invalid certificate publication source".into()))?
{
return Ok(snapshot);
#[cfg(test)]
if let Some(control) = &read_control {
control.entered.cancel();
let _ = control.release.lock().unwrap_or_else(|e| e.into_inner()).recv();
}
rustfs_tls_runtime::TlsMaterialSnapshot::load_from_snapshot(&pinned_source).map_err(map_runtime_tls_error)
let result = (|| {
let pinned_source = tls_source.snapshot().map_err(map_runtime_tls_error)?;
if let Some(snapshot) =
super::tls_publish::configured_server_snapshot(&tls_source.base_dir.to_string_lossy(), &pinned_source)
.map_err(|_| TlsMaterialError::Io("invalid certificate publication source".into()))?
{
return Ok(snapshot);
}
rustfs_tls_runtime::TlsMaterialSnapshot::load_from_snapshot(&pinned_source).map_err(map_runtime_tls_error)
})();
#[cfg(test)]
if let Some(control) = &read_control {
control
.result
.store(if result.is_ok() { 1 } else { 2 }, std::sync::atomic::Ordering::Release);
control.finished.cancel();
}
result
})
.await
.map_err(|_| TlsMaterialError::Io("TLS material loader failed".into()))??;
@@ -571,6 +603,13 @@ fn tls_key_log() -> bool {
pub enum TlsMaterialError {
Io(String),
Parse(String),
ReloadStopping,
ReloadTaskFailed(Arc<tokio::task::JoinError>),
#[cfg(feature = "http3")]
ReloadStartupCleanup {
reload: Box<TlsMaterialError>,
http3_join_error: Arc<tokio::task::JoinError>,
},
}
impl std::fmt::Display for TlsMaterialError {
@@ -578,11 +617,34 @@ impl std::fmt::Display for TlsMaterialError {
match self {
TlsMaterialError::Io(msg) => write!(f, "TLS material I/O error: {msg}"),
TlsMaterialError::Parse(msg) => write!(f, "TLS material parse error: {msg}"),
TlsMaterialError::ReloadStopping => write!(f, "TLS reload group is stopping"),
TlsMaterialError::ReloadTaskFailed(error) => {
write!(f, "TLS reload group {}", if error.is_cancelled() { "cancelled" } else { "panicked" })
}
#[cfg(feature = "http3")]
TlsMaterialError::ReloadStartupCleanup { http3_join_error, .. } => write!(
f,
"TLS reload startup failed; HTTP/3 outer join {}",
if http3_join_error.is_cancelled() {
"cancelled"
} else {
"panicked"
}
),
}
}
}
impl std::error::Error for TlsMaterialError {}
impl std::error::Error for TlsMaterialError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::ReloadTaskFailed(error) => Some(error.as_ref()),
#[cfg(feature = "http3")]
Self::ReloadStartupCleanup { reload, .. } => Some(reload.as_ref()),
_ => None,
}
}
}
// ── TLS Handshake Error Classification ──
@@ -723,11 +785,14 @@ async fn reload_acceptors(
#[derive(Clone)]
struct ReloadTarget {
holder: Arc<TlsAcceptorHolder>,
rounds: TaskTracker,
#[cfg(feature = "http3")]
http3_endpoint: Option<quinn::Endpoint>,
}
struct ReloadState {
closing: bool,
round_tokens: Option<Vec<tokio_util::task::task_tracker::TaskTrackerToken>>,
targets: HashMap<u64, ReloadTarget>,
accepted: Arc<rustfs_tls_runtime::TlsMaterialSnapshot>,
tcp: Arc<TlsAcceptor>,
@@ -741,6 +806,15 @@ struct ReloadGroup {
state: Mutex<ReloadState>,
changed: tokio::sync::Notify,
next_id: std::sync::atomic::AtomicU64,
driver: std::sync::OnceLock<ReloadJoin>,
#[cfg(test)]
read_control: Mutex<Option<Arc<ReloadReadTestControl>>>,
#[cfg(test)]
panic_before_read: std::sync::atomic::AtomicBool,
#[cfg(test)]
failure_entered: tokio_util::sync::CancellationToken,
#[cfg(test)]
driver_abort: std::sync::OnceLock<tokio::task::AbortHandle>,
#[cfg(test)]
pass_completed: tokio::sync::Notify,
}
@@ -774,6 +848,14 @@ pub(crate) fn notify_server_publication(tls_path: &str) {
}
}
#[cfg(test)]
struct ReloadReadTestControl {
entered: tokio_util::sync::CancellationToken,
finished: tokio_util::sync::CancellationToken,
release: Mutex<std::sync::mpsc::Receiver<()>>,
result: std::sync::atomic::AtomicU8,
}
struct ReloadRegistration {
group: Arc<ReloadGroup>,
id: u64,
@@ -781,16 +863,123 @@ struct ReloadRegistration {
impl Drop for ReloadRegistration {
fn drop(&mut self) {
self.group
.state
.lock()
.unwrap_or_else(|e| e.into_inner())
.targets
.remove(&self.id);
{
let _registry = RELOAD_GROUPS.lock().unwrap_or_else(|e| e.into_inner());
let mut state = self.group.state.lock().unwrap_or_else(|e| e.into_inner());
if let Some(target) = state.targets.remove(&self.id) {
// Round tokens are captured under this same lock before a read starts.
target.rounds.close();
}
if state.targets.is_empty() {
state.closing = true;
}
}
self.group.changed.notify_one();
}
}
struct ReloadRound {
group: Arc<ReloadGroup>,
}
impl Drop for ReloadRound {
fn drop(&mut self) {
let tokens = self.group.state.lock().unwrap_or_else(|e| e.into_inner()).round_tokens.take();
drop(tokens);
}
}
type ReloadJoin = Shared<BoxFuture<'static, Result<(), Arc<tokio::task::JoinError>>>>;
/// A stopped registration has drained every round captured before its stop fence.
/// The shared result records the original group task's eventual join, including failures.
#[derive(Clone)]
pub(crate) struct TlsReloadRetirement {
joined: ReloadJoin,
}
impl std::fmt::Debug for TlsReloadRetirement {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("TlsReloadRetirement")
.field("group_joined", &self.group_joined())
.finish()
}
}
impl TlsReloadRetirement {
pub(crate) fn group_joined(&self) -> bool {
matches!(self.joined.peek(), Some(Ok(())))
}
pub(crate) fn join_error(&self) -> Option<&Arc<tokio::task::JoinError>> {
self.joined.peek().and_then(|result| result.as_ref().err())
}
}
#[cfg(test)]
pub(crate) fn test_retirement_from_task(task: tokio::task::JoinHandle<()>) -> TlsReloadRetirement {
TlsReloadRetirement {
joined: task.map(|result| result.map_err(Arc::new)).boxed().shared(),
}
}
#[cfg(test)]
impl TlsReloadRetirement {
pub(crate) async fn test_observe_join(&self) {
let _ = self.joined.clone().await;
}
}
pub(crate) struct TlsReloadTask {
registration: Mutex<Option<ReloadRegistration>>,
group: Arc<ReloadGroup>,
rounds: TaskTracker,
joined: ReloadJoin,
}
impl TlsReloadTask {
pub(crate) fn request_stop(&self) {
// Handle registration -> registry -> group state; no inverse acquisition or await.
self.registration.lock().unwrap_or_else(|e| e.into_inner()).take();
}
pub(crate) async fn wait_stopped(&self) -> TlsReloadRetirement {
self.request_stop();
self.rounds.wait().await;
let _ = self.joined.clone().now_or_never();
let empty = self.group.state.lock().unwrap_or_else(|e| e.into_inner()).targets.is_empty();
if empty {
// The same-lock last-stop fence marks this generation closing.
// New listeners cannot revive it; they receive ReloadStopping until it exits.
let _ = self.joined.clone().await;
}
TlsReloadRetirement {
joined: self.joined.clone(),
}
}
#[cfg(test)]
pub(super) fn abort(&self) {
self.request_stop();
}
}
impl Drop for TlsReloadTask {
fn drop(&mut self) {
self.request_stop();
}
}
// Existing module tests request stop with abort(), then await the same real group.
#[cfg(test)]
impl std::future::Future for TlsReloadTask {
type Output = Result<(), Arc<tokio::task::JoinError>>;
fn poll(self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> std::task::Poll<Self::Output> {
self.get_mut().joined.poll_unpin(cx)
}
}
fn server_names(server: &Option<RuntimeServerTlsMaterial>) -> Result<BTreeMap<String, BTreeSet<String>>, TlsMaterialError> {
let chains: Vec<_> = match server {
Some(RuntimeServerTlsMaterial::SingleCert { certs, .. }) => vec![("default", certs)],
@@ -868,7 +1057,7 @@ pub(crate) fn spawn_reload_loop(
initial: rustfs_tls_runtime::TlsMaterialSnapshot,
prepared: PreparedServerTls,
#[cfg(feature = "http3")] http3_endpoint: Option<quinn::Endpoint>,
) -> Result<Option<tokio::task::JoinHandle<()>>, TlsMaterialError> {
) -> Result<Option<TlsReloadTask>, TlsMaterialError> {
if !get_env_bool(ENV_TLS_RELOAD_ENABLE, DEFAULT_TLS_RELOAD_ENABLE) {
return Ok(None);
}
@@ -877,11 +1066,24 @@ pub(crate) fn spawn_reload_loop(
let mut registry = RELOAD_GROUPS.lock().unwrap_or_else(|e| e.into_inner());
registry.retain(|_, group| group.strong_count() != 0);
let existing = registry.get(&tls_path).and_then(Weak::upgrade);
if let Some(group) = &existing {
if group.state.lock().unwrap_or_else(|e| e.into_inner()).closing {
return Err(TlsMaterialError::ReloadStopping);
}
if let Some(result) = group.driver.get().and_then(|joined| joined.clone().now_or_never()) {
return Err(match result {
Ok(()) => TlsMaterialError::ReloadStopping,
Err(error) => TlsMaterialError::ReloadTaskFailed(error),
});
}
}
let first = existing.is_none();
let group = existing.unwrap_or_else(|| {
Arc::new(ReloadGroup {
required_startup_files,
state: Mutex::new(ReloadState {
closing: false,
round_tokens: None,
targets: HashMap::new(),
accepted: initial.clone(),
tcp: holder.get(),
@@ -891,11 +1093,21 @@ pub(crate) fn spawn_reload_loop(
}),
changed: tokio::sync::Notify::new(),
next_id: std::sync::atomic::AtomicU64::new(1),
driver: std::sync::OnceLock::new(),
#[cfg(test)]
read_control: Mutex::new(None),
#[cfg(test)]
panic_before_read: std::sync::atomic::AtomicBool::new(false),
#[cfg(test)]
failure_entered: tokio_util::sync::CancellationToken::new(),
#[cfg(test)]
driver_abort: std::sync::OnceLock::new(),
#[cfg(test)]
pass_completed: tokio::sync::Notify::new(),
})
});
let id = group.next_id.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let rounds = TaskTracker::new();
{
let mut state = group.state.lock().unwrap_or_else(|e| e.into_inner());
#[cfg(feature = "http3")]
@@ -908,26 +1120,45 @@ pub(crate) fn spawn_reload_loop(
endpoint.set_server_config(state.quic.clone());
}
holder.swap(&TlsAcceptorHolder::new(state.tcp.clone()));
if let Some(tokens) = &mut state.round_tokens {
// A listener joining mid-round can receive this publication too.
tokens.push(rounds.token());
}
state.targets.insert(
id,
ReloadTarget {
holder,
rounds: rounds.clone(),
#[cfg(feature = "http3")]
http3_endpoint,
},
);
}
registry.insert(tls_path.clone(), Arc::downgrade(&group));
drop(registry);
if first {
tokio::spawn(run_reload_group(tls_path, group.clone()));
let driver = tokio::spawn(run_reload_group(tls_path, group.clone()));
#[cfg(test)]
let _ = group.driver_abort.set(driver.abort_handle());
let joined = driver.map(|result| result.map_err(Arc::new)).boxed().shared();
let _ = group.driver.set(joined);
}
let joined = group
.driver
.get()
.expect("registered group retains its original driver")
.clone();
drop(registry);
group.changed.notify_one();
let registration = ReloadRegistration { group, id };
Ok(Some(tokio::spawn(async move {
let _registration = registration;
std::future::pending::<()>().await;
})))
let registration = ReloadRegistration {
group: group.clone(),
id,
};
Ok(Some(TlsReloadTask {
registration: Mutex::new(Some(registration)),
group,
rounds,
joined,
}))
}
async fn run_reload_group(tls_path: String, group: Arc<ReloadGroup>) {
@@ -941,19 +1172,39 @@ async fn run_reload_group(tls_path: String, group: Arc<ReloadGroup>) {
_ = group.changed.notified() => {},
}
let (accepted, empty) = {
let state = group.state.lock().unwrap_or_else(|e| e.into_inner());
let mut state = group.state.lock().unwrap_or_else(|e| e.into_inner());
state.round_tokens = Some(state.targets.values().map(|target| target.rounds.token()).collect());
(state.accepted.clone(), state.targets.is_empty())
};
let round = ReloadRound { group: group.clone() };
if empty {
let mut registry = RELOAD_GROUPS.lock().unwrap_or_else(|e| e.into_inner());
if group.state.lock().unwrap_or_else(|e| e.into_inner()).targets.is_empty() {
registry.remove(&tls_path);
if registry
.get(&tls_path)
.and_then(Weak::upgrade)
.is_some_and(|current| Arc::ptr_eq(&current, &group))
{
registry.remove(&tls_path);
}
return;
}
continue;
}
#[cfg(test)]
if group.panic_before_read.swap(false, std::sync::atomic::Ordering::AcqRel) {
group.failure_entered.cancel();
panic!("test TLS group driver failure");
}
let result = async {
let candidate = load_tls_material(&tls_path).await?;
#[cfg(test)]
let read_control = group.read_control.lock().unwrap_or_else(|e| e.into_inner()).take();
let candidate = load_tls_material_observed(
&tls_path,
#[cfg(test)]
read_control,
)
.await?;
validate_candidate_replacement(&accepted, &candidate)?;
validate_required_material_files(&group.required_startup_files, &candidate)?;
// Listener initialization can load a newer bundle than startup published.
@@ -1031,6 +1282,8 @@ async fn run_reload_group(tls_path: String, group: Arc<ReloadGroup>) {
result = "rejected", error = %error, "TLS material update rejected");
}
}
// Complete reads, parsing, publication and result processing before retiring this round.
drop(round);
#[cfg(test)]
group.pass_completed.notify_one();
}
@@ -1060,6 +1313,305 @@ mod tests {
fs::write(dir.join(rustfs_config::RUSTFS_TLS_KEY), signing_key.serialize_pem()).unwrap();
}
fn register_test_reload(path: &str, initial: rustfs_tls_runtime::TlsMaterialSnapshot) -> TlsReloadTask {
let prepared = PreparedServerTls::from_snapshot(&initial)
.expect("prepare fixture")
.expect("server material");
let holder = prepared.tcp_acceptor().expect("fixture acceptor");
spawn_reload_loop(
path.into(),
holder,
initial,
prepared,
#[cfg(feature = "http3")]
None,
)
.expect("register actual reloader")
.expect("reload enabled")
}
async fn assert_owned_reload_stop(failed_read: bool, cancel_observer: bool, sharing: &str) {
ensure_rustls_crypto_provider();
temp_env::async_with_vars(
[
(ENV_TLS_RELOAD_ENABLE, Some("1")),
(ENV_SERVER_MTLS_ENABLE, Some("0")),
(ENV_TRUST_SYSTEM_CA, Some("0")),
(ENV_TRUST_LEAF_CERT_AS_CA, Some("0")),
],
async {
let dir = TempDir::new().expect("owned TLS input");
write_test_cert_pair(dir.path(), "owned-stop.example");
let path = dir.path().to_str().expect("fixture path");
let initial = load_tls_material(path).await.expect("initial real material");
// Load the second owned input before any reload task or native pause exists.
let mut registration_snapshot = Some(load_tls_material(path).await.expect("second real registration snapshot"));
let task = register_test_reload(path, initial);
let existing_peer = (sharing == "existing")
.then(|| register_test_reload(path, registration_snapshot.take().expect("preloaded existing registration")));
let group = task.group.clone();
let (release_tx, release_rx) = std::sync::mpsc::channel();
let control = Arc::new(ReloadReadTestControl {
entered: tokio_util::sync::CancellationToken::new(),
finished: tokio_util::sync::CancellationToken::new(),
release: Mutex::new(release_rx),
result: std::sync::atomic::AtomicU8::new(0),
});
*group.read_control.lock().unwrap_or_else(|e| e.into_inner()) = Some(control.clone());
let entered = tokio::time::timeout(Duration::from_secs(5), control.entered.cancelled())
.await
.is_ok();
let mid = (sharing == "mid-round")
.then(|| register_test_reload(path, registration_snapshot.take().expect("preloaded mid-round registration")));
let input_changed =
!failed_read || fs::write(dir.path().join(RUSTFS_TLS_CERT), b"invalid certificate fixture").is_ok();
let first = mid.as_ref().unwrap_or(&task);
first.request_stop();
let mut observer = Box::pin(first.wait_stopped());
let pending = std::future::poll_fn(|cx| {
std::task::Poll::Ready(std::future::Future::poll(observer.as_mut(), cx).is_pending())
})
.await;
let native_still_running = !control.finished.is_cancelled();
let closing_rejection = if sharing.is_empty() {
let initial = registration_snapshot.take().expect("preloaded closing registration");
let prepared = PreparedServerTls::from_snapshot(&initial)
.expect("same prepared input")
.expect("server");
let holder = prepared.tcp_acceptor().expect("same acceptor");
Some(spawn_reload_loop(
path.into(),
holder,
initial,
prepared,
#[cfg(feature = "http3")]
None,
))
} else {
None
};
// Disconnecting this sender also releases native work during assertion unwinding.
let first_receipt = if cancel_observer {
drop(observer);
drop(release_tx);
tokio::time::timeout(Duration::from_secs(5), first.wait_stopped()).await
} else {
drop(release_tx);
tokio::time::timeout(Duration::from_secs(5), observer).await
};
let first_was_group_joined = first_receipt.as_ref().is_ok_and(|receipt| receipt.group_joined());
let closing_was_rejected =
matches!(&closing_rejection, Some(Err(TlsMaterialError::ReloadStopping))) || !sharing.is_empty();
if let Some(Ok(Some(unexpected))) = closing_rejection {
unexpected.request_stop();
let _ = tokio::time::timeout(Duration::from_secs(5), unexpected.wait_stopped()).await;
}
let remaining = if sharing == "mid-round" {
Some(&task)
} else {
existing_peer.as_ref()
};
let later_receipt = if let Some(remaining) = remaining {
remaining.request_stop();
Some(tokio::time::timeout(Duration::from_secs(5), remaining.wait_stopped()).await)
} else {
None
};
// Cleanup observes the same actual driver even when a bounded test observation failed.
let final_join = group.driver.get().expect("same original driver").clone().await;
// The original selected native closure is joined by the original group, not a noop barrier.
let finished = control.finished.is_cancelled();
let result = control.result.load(std::sync::atomic::Ordering::Acquire);
let receipt = first_receipt.expect("original stopped registration drains");
if let Some(receipt) = later_receipt {
assert!(receipt.expect("last group drains").group_joined());
}
assert!(final_join.is_ok(), "original group joined after all registrations stopped");
assert!(entered, "selected original native read entered");
assert!(pending && native_still_running, "stop cannot discard an entered native read");
assert!(finished);
assert!(input_changed, "actual parser failure input was installed");
assert!(closing_was_rejected, "a closing old group explicitly rejects new registrations");
assert_eq!(result, if failed_read { 2 } else { 1 }, "actual native result retained through stop");
assert!(receipt.group_joined(), "earlier receipt observes the same eventual real group join");
assert_eq!(
first_was_group_joined,
sharing.is_empty(),
"an independent live registration remains active"
);
if cancel_observer {
assert!(
group
.driver
.get()
.expect("same original driver")
.peek()
.is_some_and(Result::is_ok)
);
}
assert!(!RELOAD_GROUPS.lock().unwrap_or_else(|e| e.into_inner()).contains_key(path));
},
)
.await;
}
#[tokio::test]
#[serial_test::serial]
async fn owned_reload_stop_waits_for_original_native_success() {
assert_owned_reload_stop(false, false, "").await;
}
#[tokio::test]
#[serial_test::serial]
async fn owned_reload_stop_waits_for_original_native_error() {
assert_owned_reload_stop(true, false, "").await;
}
#[tokio::test]
#[serial_test::serial]
async fn owned_reload_cancelled_observer_retains_original_native_owner() {
assert_owned_reload_stop(false, true, "").await;
}
#[tokio::test]
#[serial_test::serial]
async fn owned_reload_mid_round_listener_waits_for_its_publication_round() {
assert_owned_reload_stop(false, false, "mid-round").await;
}
#[tokio::test]
#[serial_test::serial]
async fn owned_reload_existing_shared_listener_waits_for_its_captured_round() {
assert_owned_reload_stop(false, false, "existing").await;
}
#[tokio::test]
#[serial_test::serial]
async fn owned_reload_late_ready_keeps_original_http_deadline_failed() {
ensure_rustls_crypto_provider();
temp_env::async_with_vars(
[
(ENV_TLS_RELOAD_ENABLE, Some("1")),
(ENV_SERVER_MTLS_ENABLE, Some("0")),
(ENV_TRUST_SYSTEM_CA, Some("0")),
(ENV_TRUST_LEAF_CERT_AS_CA, Some("0")),
],
async {
let dir = TempDir::new().expect("owned late observation input");
write_test_cert_pair(dir.path(), "late-stop.example");
let path = dir.path().to_str().expect("fixture path");
let initial = load_tls_material(path).await.expect("initial real material");
let task = register_test_reload(path, initial);
let group = task.group.clone();
let (release_tx, release_rx) = std::sync::mpsc::channel();
let control = Arc::new(ReloadReadTestControl {
entered: tokio_util::sync::CancellationToken::new(),
finished: tokio_util::sync::CancellationToken::new(),
release: Mutex::new(release_rx),
result: std::sync::atomic::AtomicU8::new(0),
});
*group.read_control.lock().unwrap_or_else(|e| e.into_inner()) = Some(control.clone());
let entered = tokio::time::timeout(Duration::from_secs(5), control.entered.cancelled())
.await
.is_ok();
task.request_stop();
tokio::time::pause();
let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
let mut outcome = super::super::http::test_retire_http12_scope(std::future::ready(()), deadline).await;
let pending;
let joined;
{
let mut observer = Box::pin(super::super::http::retire_tls_reloader(&task, &mut outcome, deadline));
pending = std::future::poll_fn(|cx| {
std::task::Poll::Ready(std::future::Future::poll(observer.as_mut(), cx).is_pending())
})
.await;
tokio::time::advance(Duration::from_secs(10)).await;
drop(release_tx);
joined = group.driver.get().expect("same actual driver").clone().await;
// The owned result is already ready before the original observer is polled again.
observer.await;
}
tokio::time::resume();
assert!(entered && pending);
assert!(joined.is_ok() && control.finished.is_cancelled());
assert_eq!(control.result.load(std::sync::atomic::Ordering::Acquire), 1);
assert!(
outcome.deadline_missed(),
"late Ready cannot erase the original absolute 10-second observation failure"
);
assert!(outcome.verified_scope_complete(), "cleanup genuinely joined the TLS group");
},
)
.await;
}
async fn assert_reload_driver_failure(cancel: bool) {
ensure_rustls_crypto_provider();
temp_env::async_with_vars(
[
(ENV_TLS_RELOAD_ENABLE, Some("1")),
(ENV_SERVER_MTLS_ENABLE, Some("0")),
(ENV_TRUST_SYSTEM_CA, Some("0")),
(ENV_TRUST_LEAF_CERT_AS_CA, Some("0")),
],
async {
let dir = TempDir::new().expect("owned driver fixture");
write_test_cert_pair(dir.path(), "driver-stop.example");
let path = dir.path().to_str().expect("path");
let initial = load_tls_material(path).await.expect("initial material");
let task = register_test_reload(path, initial);
let group = task.group.clone();
let entered = if cancel {
// This task has not been polled; no native read is detached by this error injection.
group.driver_abort.get().expect("actual original task").abort();
true
} else {
group.panic_before_read.store(true, std::sync::atomic::Ordering::Release);
tokio::time::timeout(Duration::from_secs(5), group.failure_entered.cancelled())
.await
.is_ok()
};
task.request_stop();
let receipt = tokio::time::timeout(Duration::from_secs(5), task.wait_stopped())
.await
.expect("actual failed driver joins");
let error = receipt.join_error().expect("actual typed JoinError retained");
assert!(entered, "actual driver reached the failure boundary");
assert_eq!(error.is_cancelled(), cancel);
assert_eq!(error.is_panic(), !cancel);
assert!(!receipt.group_joined());
let typed = TlsMaterialError::ReloadTaskFailed(error.clone());
assert_eq!(
typed.to_string(),
if cancel {
"TLS reload group cancelled"
} else {
"TLS reload group panicked"
}
);
let source = std::error::Error::source(&typed).expect("actual join source");
assert!(std::ptr::eq(
source.downcast_ref::<tokio::task::JoinError>().expect("typed source"),
error.as_ref()
));
},
)
.await;
}
#[tokio::test]
#[serial_test::serial]
async fn owned_reload_original_driver_panic_remains_failure() {
assert_reload_driver_failure(false).await;
}
#[tokio::test]
#[serial_test::serial]
async fn owned_reload_original_driver_cancel_remains_failure() {
assert_reload_driver_failure(true).await;
}
#[tokio::test]
#[serial_test::serial]
async fn prepared_generation_shares_crypto_and_serves_sni_with_mtls() {

Some files were not shown because too many files have changed in this diff Show More