diff --git a/Cargo.lock b/Cargo.lock index 62a41583b..1331b0d99 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7712,6 +7712,7 @@ dependencies = [ "rustfs-madmin", "rustfs-metrics", "rustfs-notify", + "rustfs-object-capacity", "rustfs-object-io", "rustfs-obs", "rustfs-policy", @@ -7757,7 +7758,6 @@ dependencies = [ "url", "urlencoding", "uuid", - "walkdir", "zip", ] @@ -8229,6 +8229,25 @@ dependencies = [ "wildmatch", ] +[[package]] +name = "rustfs-object-capacity" +version = "0.0.5" +dependencies = [ + "criterion", + "futures", + "rustfs-common", + "rustfs-config", + "rustfs-io-metrics", + "rustfs-utils", + "serial_test", + "temp-env", + "tempfile", + "tokio", + "tracing", + "uuid", + "walkdir", +] + [[package]] name = "rustfs-object-io" version = "0.0.5" @@ -9726,6 +9745,7 @@ version = "0.3.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "96374855068f47402c3121c6eed88d29cb1de8f3ab27090e273e420bdabcf050" dependencies = [ + "futures", "parking_lot 0.12.5", ] diff --git a/Cargo.toml b/Cargo.toml index 0b886abbd..febbf9eec 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,6 +35,7 @@ members = [ "crates/metrics", # Metrics collection and reporting "crates/notify", # Notification system for events "crates/obs", # Observability utilities + "crates/object-capacity", # Capacity scan and refresh core "crates/policy", # Policy management "crates/protocols", # Protocol implementations (FTPS, SFTP, etc.) "crates/protos", # Protocol buffer definitions @@ -99,6 +100,7 @@ rustfs-notify = { path = "crates/notify", version = "0.0.5" } rustfs-io-metrics = { path = "crates/io-metrics", version = "0.0.5" } rustfs-io-core = { path = "crates/io-core", version = "0.0.5" } rustfs-object-io = { path = "crates/object-io", version = "0.0.5" } +rustfs-object-capacity = { path = "crates/object-capacity", version = "0.0.5" } rustfs-obs = { path = "crates/obs", version = "0.0.5" } rustfs-policy = { path = "crates/policy", version = "0.0.5" } rustfs-protos = { path = "crates/protos", version = "0.0.5" } diff --git a/crates/common/src/capacity_scope.rs b/crates/common/src/capacity_scope.rs new file mode 100644 index 000000000..9f4ea49a2 --- /dev/null +++ b/crates/common/src/capacity_scope.rs @@ -0,0 +1,247 @@ +// Copyright 2024 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 std::collections::{HashMap, HashSet}; +use std::sync::{Mutex, OnceLock}; +use std::time::{Duration, Instant}; +use uuid::Uuid; + +const CAPACITY_SCOPE_REGISTRY_SOFT_LIMIT: usize = 2_048; +const CAPACITY_SCOPE_REGISTRY_HARD_LIMIT: usize = 4_096; +const CAPACITY_SCOPE_TTL: Duration = Duration::from_secs(300); + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct CapacityScopeDisk { + pub endpoint: String, + pub drive_path: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub struct CapacityScope { + pub disks: Vec, +} + +#[derive(Debug, Clone)] +struct CapacityScopeEntry { + scope: CapacityScope, + recorded_at: Instant, +} + +fn capacity_scope_registry() -> &'static Mutex> { + static REGISTRY: OnceLock>> = OnceLock::new(); + REGISTRY.get_or_init(|| Mutex::new(HashMap::new())) +} + +fn global_dirty_scope_registry() -> &'static Mutex> { + static REGISTRY: OnceLock>> = OnceLock::new(); + REGISTRY.get_or_init(|| Mutex::new(HashSet::new())) +} + +fn prune_expired_entries(entries: &mut HashMap, now: Instant) { + entries.retain(|_, entry| now.duration_since(entry.recorded_at) <= CAPACITY_SCOPE_TTL); +} + +fn enforce_hard_limit(entries: &mut HashMap, max_len: usize) { + if entries.len() < max_len { + return; + } + + let evict_count = entries.len() - max_len + 1; + let mut eviction_order: Vec<_> = entries.iter().map(|(token, entry)| (*token, entry.recorded_at)).collect(); + eviction_order.sort_unstable_by_key(|(_, recorded_at)| *recorded_at); + + for (token, _) in eviction_order.into_iter().take(evict_count) { + entries.remove(&token); + } +} + +fn merge_capacity_scopes(existing: &mut CapacityScope, incoming: CapacityScope) { + let mut seen: HashSet = existing.disks.iter().cloned().collect(); + for disk in incoming.disks { + if seen.insert(disk.clone()) { + existing.disks.push(disk); + } + } +} + +pub fn record_capacity_scope(token: Uuid, scope: CapacityScope) { + let now = Instant::now(); + let mut entries = capacity_scope_registry().lock().unwrap_or_else(|p| p.into_inner()); + if !entries.contains_key(&token) && entries.len() >= CAPACITY_SCOPE_REGISTRY_SOFT_LIMIT { + prune_expired_entries(&mut entries, now); + enforce_hard_limit(&mut entries, CAPACITY_SCOPE_REGISTRY_HARD_LIMIT); + } + if let Some(entry) = entries.get_mut(&token) { + merge_capacity_scopes(&mut entry.scope, scope); + entry.recorded_at = now; + } else { + entries.insert(token, CapacityScopeEntry { scope, recorded_at: now }); + } +} + +pub fn take_capacity_scope(token: Uuid) -> Option { + let now = Instant::now(); + let mut entries = capacity_scope_registry().lock().unwrap_or_else(|p| p.into_inner()); + let entry = entries.remove(&token)?; + if now.duration_since(entry.recorded_at) > CAPACITY_SCOPE_TTL { + return None; + } + Some(entry.scope) +} + +pub fn record_global_dirty_scope(scope: CapacityScope) { + if scope.disks.is_empty() { + return; + } + + let mut dirty_scopes = global_dirty_scope_registry().lock().unwrap_or_else(|p| p.into_inner()); + dirty_scopes.extend(scope.disks); +} + +pub fn drain_global_dirty_scopes() -> Vec { + let mut dirty_scopes = global_dirty_scope_registry().lock().unwrap_or_else(|p| p.into_inner()); + dirty_scopes.drain().collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Mutex; + + fn test_lock() -> &'static Mutex<()> { + static LOCK: OnceLock> = OnceLock::new(); + LOCK.get_or_init(|| Mutex::new(())) + } + + fn clear_capacity_scope_registry_for_test() { + capacity_scope_registry() + .lock() + .expect("capacity scope registry poisoned") + .clear(); + global_dirty_scope_registry() + .lock() + .expect("global dirty scope registry poisoned") + .clear(); + } + + #[test] + fn record_and_take_capacity_scope_round_trips() { + let _guard = test_lock().lock().expect("test lock poisoned"); + clear_capacity_scope_registry_for_test(); + let token = Uuid::new_v4(); + let scope = CapacityScope { + disks: vec![CapacityScopeDisk { + endpoint: "node-a".to_string(), + drive_path: "/tmp/disk-a".to_string(), + }], + }; + + record_capacity_scope(token, scope.clone()); + + assert_eq!(take_capacity_scope(token), Some(scope)); + assert_eq!(take_capacity_scope(token), None); + clear_capacity_scope_registry_for_test(); + } + + #[test] + fn record_capacity_scope_merges_disks_for_same_token() { + let _guard = test_lock().lock().expect("test lock poisoned"); + clear_capacity_scope_registry_for_test(); + let token = Uuid::new_v4(); + record_capacity_scope( + token, + CapacityScope { + disks: vec![CapacityScopeDisk { + endpoint: "node-a".to_string(), + drive_path: "/tmp/disk-a".to_string(), + }], + }, + ); + record_capacity_scope( + token, + CapacityScope { + disks: vec![ + CapacityScopeDisk { + endpoint: "node-b".to_string(), + drive_path: "/tmp/disk-b".to_string(), + }, + CapacityScopeDisk { + endpoint: "node-a".to_string(), + drive_path: "/tmp/disk-a".to_string(), + }, + ], + }, + ); + + let scope = take_capacity_scope(token).expect("scope should exist"); + assert_eq!(scope.disks.len(), 2); + assert!(scope.disks.iter().any(|disk| disk.endpoint == "node-a")); + assert!(scope.disks.iter().any(|disk| disk.endpoint == "node-b")); + clear_capacity_scope_registry_for_test(); + } + + #[test] + fn record_capacity_scope_enforces_hard_limit() { + let _guard = test_lock().lock().expect("test lock poisoned"); + clear_capacity_scope_registry_for_test(); + + for _ in 0..(CAPACITY_SCOPE_REGISTRY_HARD_LIMIT + 32) { + record_capacity_scope( + Uuid::new_v4(), + CapacityScope { + disks: vec![CapacityScopeDisk { + endpoint: "node-a".to_string(), + drive_path: "/tmp/disk-a".to_string(), + }], + }, + ); + } + + let entries = capacity_scope_registry().lock().expect("capacity scope registry poisoned"); + assert!(entries.len() <= CAPACITY_SCOPE_REGISTRY_HARD_LIMIT); + drop(entries); + clear_capacity_scope_registry_for_test(); + } + + #[test] + fn record_and_drain_global_dirty_scope_round_trips() { + let _guard = test_lock().lock().expect("test lock poisoned"); + clear_capacity_scope_registry_for_test(); + record_global_dirty_scope(CapacityScope { + disks: vec![CapacityScopeDisk { + endpoint: "node-a".to_string(), + drive_path: "/tmp/disk-a".to_string(), + }], + }); + record_global_dirty_scope(CapacityScope { + disks: vec![ + CapacityScopeDisk { + endpoint: "node-b".to_string(), + drive_path: "/tmp/disk-b".to_string(), + }, + CapacityScopeDisk { + endpoint: "node-a".to_string(), + drive_path: "/tmp/disk-a".to_string(), + }, + ], + }); + + let drained = drain_global_dirty_scopes(); + assert_eq!(drained.len(), 2); + assert!(drained.iter().any(|disk| disk.endpoint == "node-a")); + assert!(drained.iter().any(|disk| disk.endpoint == "node-b")); + assert!(drain_global_dirty_scopes().is_empty()); + clear_capacity_scope_registry_for_test(); + } +} diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs index 3c532fcf0..df06c5966 100644 --- a/crates/common/src/lib.rs +++ b/crates/common/src/lib.rs @@ -13,6 +13,7 @@ // limitations under the License. pub mod bucket_stats; +pub mod capacity_scope; // pub mod error; pub mod data_usage; pub mod globals; diff --git a/crates/ecstore/src/set_disk.rs b/crates/ecstore/src/set_disk.rs index 29877ced6..82034c058 100644 --- a/crates/ecstore/src/set_disk.rs +++ b/crates/ecstore/src/set_disk.rs @@ -67,6 +67,7 @@ use http::HeaderMap; use md5::{Digest as Md5Digest, Md5}; use rand::{Rng, seq::SliceRandom}; use regex::Regex; +use rustfs_common::capacity_scope::{CapacityScope, CapacityScopeDisk, record_capacity_scope, record_global_dirty_scope}; use rustfs_common::heal_channel::{DriveState, HealChannelPriority, HealItemType, HealOpts, HealScanMode, send_heal_disk}; use rustfs_config::MI_B; use rustfs_filemeta::{ @@ -133,6 +134,36 @@ fn env_non_negative_usize(name: &str) -> Option { rustfs_utils::get_env_opt_usize(name) } +fn capacity_scope_from_disks(disks: &[Option]) -> CapacityScope { + let mut unique = HashSet::with_capacity(disks.len()); + let mut scoped_disks = Vec::with_capacity(disks.len()); + + for disk in disks.iter().flatten() { + let scope_disk = CapacityScopeDisk { + endpoint: disk.endpoint().to_string(), + drive_path: disk.to_string(), + }; + if unique.insert(scope_disk.clone()) { + scoped_disks.push(scope_disk); + } + } + + CapacityScope { disks: scoped_disks } +} + +fn record_capacity_scope_if_needed(scope_token: Option, disks: &[Option]) { + let scope = capacity_scope_from_disks(disks); + if scope.disks.is_empty() { + return; + } + + record_global_dirty_scope(scope.clone()); + + if let Some(token) = scope_token { + record_capacity_scope(token, scope); + } +} + fn resolved_put_inline_buffer_enabled(object_size: i64, inline_by_topology: bool) -> bool { if !inline_by_topology || object_size < 0 { return false; @@ -1087,6 +1118,8 @@ impl ObjectIO for SetDisks { } } + record_capacity_scope_if_needed(opts.capacity_scope_token, &online_disks); + fi.replication_state_internal = Some(opts.put_replication_state()); fi.is_latest = true; @@ -1681,6 +1714,8 @@ impl ObjectOperations for SetDisks { } } + record_capacity_scope_if_needed(opts.capacity_scope_token, &disks); + // TODO: add_partial if dist_erasure { @@ -1829,6 +1864,10 @@ impl ObjectOperations for SetDisks { .await .map_err(|e| to_object_err(e, vec![bucket, object]))?; + if let Ok(disks) = self.get_disks(0, 0).await { + record_capacity_scope_if_needed(opts.capacity_scope_token, &disks); + } + let mut oi = ObjectInfo::from_file_info(&fi, bucket, object, opts.versioned || opts.version_suspended); oi.replication_decision = goi.replication_decision; return Ok(oi); @@ -1855,6 +1894,10 @@ impl ObjectOperations for SetDisks { .await .map_err(|e| to_object_err(e, vec![bucket, object]))?; + if let Ok(disks) = self.get_disks(0, 0).await { + record_capacity_scope_if_needed(opts.capacity_scope_token, &disks); + } + let mut obj_info = ObjectInfo::from_file_info(&dfi, bucket, object, opts.versioned || opts.version_suspended); obj_info.size = goi.size; Ok(obj_info) @@ -2157,6 +2200,8 @@ impl ObjectOperations for SetDisks { error = ?err, "transition completed on remote tier but source cleanup failed; skipping external lifecycle transition notification" ); + } else { + record_capacity_scope_if_needed(opts.capacity_scope_token, &disks); } for disk in disks.iter() { @@ -3515,6 +3560,8 @@ impl MultipartOperations for SetDisks { } } + record_capacity_scope_if_needed(opts.capacity_scope_token, &online_disks); + fi.is_latest = true; Ok(ObjectInfo::from_file_info(&fi, bucket, object, opts.versioned || opts.version_suspended)) diff --git a/crates/ecstore/src/set_disk/heal.rs b/crates/ecstore/src/set_disk/heal.rs index 41137bfc7..ac5762e6b 100644 --- a/crates/ecstore/src/set_disk/heal.rs +++ b/crates/ecstore/src/set_disk/heal.rs @@ -558,6 +558,8 @@ impl SetDisks { } } + record_capacity_scope_if_needed(None, &out_dated_disks); + Ok((result, None)) } Err(err) => Ok((result, Some(err))), diff --git a/crates/ecstore/src/store_api/types.rs b/crates/ecstore/src/store_api/types.rs index efbcf396d..51bd69d61 100644 --- a/crates/ecstore/src/store_api/types.rs +++ b/crates/ecstore/src/store_api/types.rs @@ -73,6 +73,7 @@ pub struct ObjectOptions { pub resolved_checksum: Option, pub want_checksum: Option, pub skip_verify_bitrot: bool, + pub capacity_scope_token: Option, } impl ObjectOptions { diff --git a/crates/io-metrics/src/capacity_metrics.rs b/crates/io-metrics/src/capacity_metrics.rs index a032727ee..be659c1dc 100644 --- a/crates/io-metrics/src/capacity_metrics.rs +++ b/crates/io-metrics/src/capacity_metrics.rs @@ -29,6 +29,12 @@ pub fn record_capacity_cache_miss() { counter!("rustfs.capacity.cache.misses").increment(1); } +/// Record how capacity cache was served to the caller. +#[inline(always)] +pub fn record_capacity_cache_served(state: &'static str) { + counter!("rustfs.capacity.cache.served.total", "state" => state).increment(1); +} + /// Record current capacity gauge. #[inline(always)] pub fn record_capacity_current_bytes(used_bytes: u64) { @@ -55,6 +61,44 @@ pub fn record_capacity_update_failed(source: &'static str) { counter!("rustfs.capacity.update.failures", "source" => source).increment(1); } +/// Record a capacity refresh request. +#[inline(always)] +pub fn record_capacity_refresh_request(mode: &'static str, source: &'static str) { + counter!("rustfs.capacity.refresh.requests.total", "mode" => mode, "source" => source).increment(1); +} + +/// Record a refresh joiner waiting for an inflight refresh. +#[inline(always)] +pub fn record_capacity_refresh_joiner(source: &'static str) { + counter!("rustfs.capacity.refresh.joiners.total", "source" => source).increment(1); +} + +/// Record the number of inflight capacity refreshes. +#[inline(always)] +pub fn record_capacity_refresh_inflight(count: usize) { + gauge!("rustfs.capacity.refresh.inflight").set(count as f64); +} + +/// Record the final result of a capacity refresh. +#[inline(always)] +pub fn record_capacity_refresh_result(source: &'static str, result: &'static str, duration: Duration) { + counter!("rustfs.capacity.refresh.result.total", "source" => source, "result" => result).increment(1); + histogram!("rustfs.capacity.refresh.duration.seconds", "source" => source, "result" => result).record(duration.as_secs_f64()); +} + +/// Record the refresh scope selected for a capacity refresh. +#[inline(always)] +pub fn record_capacity_refresh_scope(scope: &'static str, disk_count: usize) { + counter!("rustfs.capacity.refresh.scope.total", "scope" => scope).increment(1); + histogram!("rustfs.capacity.refresh.scope.disks", "scope" => scope).record(disk_count as f64); +} + +/// Record the current number of dirty disks tracked by capacity management. +#[inline(always)] +pub fn record_capacity_dirty_disk_count(count: usize) { + gauge!("rustfs.capacity.dirty.disks").set(count as f64); +} + /// Record capacity write activity. #[inline(always)] pub fn record_capacity_write_operation(write_frequency: usize) { @@ -98,3 +142,33 @@ pub fn record_capacity_scan_sampling(sampled_count: usize, estimated: bool) { ) .increment(1); } + +/// Record the scan mode used for a capacity result. +#[inline(always)] +pub fn record_capacity_scan_mode(mode: &'static str) { + counter!("rustfs.capacity.scan.mode.total", "mode" => mode).increment(1); +} + +/// Record per-disk capacity scan statistics. +#[inline(always)] +pub fn record_capacity_scan_disk( + disk: &str, + duration: Duration, + file_count: usize, + sampled_count: usize, + estimated: bool, + partial_errors: bool, +) { + histogram!("rustfs.capacity.scan.disk.duration.seconds", "disk" => disk.to_owned()).record(duration.as_secs_f64()); + histogram!("rustfs.capacity.scan.disk.files", "disk" => disk.to_owned()).record(file_count as f64); + histogram!("rustfs.capacity.scan.disk.sampled", "disk" => disk.to_owned()).record(sampled_count as f64); + counter!( + "rustfs.capacity.scan.disk.estimated.total", + "disk" => disk.to_owned(), + "estimated" => if estimated { "true" } else { "false" } + ) + .increment(1); + if partial_errors { + counter!("rustfs.capacity.scan.disk.partial_errors.total", "disk" => disk.to_owned()).increment(1); + } +} diff --git a/crates/io-metrics/src/lib.rs b/crates/io-metrics/src/lib.rs index 04257ff4d..c75a8041f 100644 --- a/crates/io-metrics/src/lib.rs +++ b/crates/io-metrics/src/lib.rs @@ -74,9 +74,12 @@ pub use adaptive_ttl::{ // Capacity metrics exports pub use capacity_metrics::{ - record_capacity_cache_hit, record_capacity_cache_miss, record_capacity_current_bytes, record_capacity_dynamic_timeout, - record_capacity_scan_sampling, record_capacity_stall_detected, record_capacity_symlink, record_capacity_timeout_fallback, - record_capacity_update_completed, record_capacity_update_failed, record_capacity_write_operation, + record_capacity_cache_hit, record_capacity_cache_miss, record_capacity_cache_served, record_capacity_current_bytes, + record_capacity_dirty_disk_count, record_capacity_dynamic_timeout, record_capacity_refresh_inflight, + record_capacity_refresh_joiner, record_capacity_refresh_request, record_capacity_refresh_result, + record_capacity_refresh_scope, record_capacity_scan_disk, record_capacity_scan_mode, record_capacity_scan_sampling, + record_capacity_stall_detected, record_capacity_symlink, record_capacity_timeout_fallback, record_capacity_update_completed, + record_capacity_update_failed, record_capacity_write_operation, }; // I/O metrics exports diff --git a/crates/object-capacity/Cargo.toml b/crates/object-capacity/Cargo.toml new file mode 100644 index 000000000..08cdda9de --- /dev/null +++ b/crates/object-capacity/Cargo.toml @@ -0,0 +1,53 @@ +# Copyright 2024 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. + +[package] +name = "rustfs-object-capacity" +version.workspace = true +edition.workspace = true +license.workspace = true +repository.workspace = true +rust-version.workspace = true +homepage.workspace = true +description = "Capacity scan and refresh core for RustFS." +keywords = ["capacity", "storage", "rustfs", "scan", "metrics"] +categories = ["filesystem", "development-tools"] + +[lib] +doctest = false + +[[bench]] +name = "capacity_scan" +harness = false + +[lints] +workspace = true + +[dependencies] +rustfs-common = { workspace = true } +rustfs-config = { workspace = true, features = ["constants"] } +rustfs-io-metrics = { workspace = true } +rustfs-utils = { workspace = true } +futures = { workspace = true } +tokio = { workspace = true, features = ["sync", "time"] } +tracing = { workspace = true } +uuid = { workspace = true } +walkdir = { workspace = true } + +[dev-dependencies] +criterion = { workspace = true } +serial_test = { workspace = true } +temp-env = { workspace = true, features = ["async_closure"] } +tempfile = { workspace = true } +tokio = { workspace = true, features = ["test-util"] } diff --git a/crates/object-capacity/benches/capacity_scan.rs b/crates/object-capacity/benches/capacity_scan.rs new file mode 100644 index 000000000..ca44ab2c1 --- /dev/null +++ b/crates/object-capacity/benches/capacity_scan.rs @@ -0,0 +1,136 @@ +use criterion::{Criterion, criterion_group, criterion_main}; +use rustfs_object_capacity::{CapacityDiskRef, scan_used_capacity_disks}; +use std::fs; +use std::hint::black_box; +use std::path::{Path, PathBuf}; +use std::time::Duration; +use tempfile::TempDir; + +const EXACT_FILE_SIZE: usize = 4 * 1024; +const SAMPLED_FILE_SIZE: usize = 1; +const DEFAULT_SAMPLE_TRIGGER_FILE_COUNT: usize = 202_048; + +#[derive(Clone, Copy)] +struct DiskSpec { + file_count: usize, + file_size: usize, +} + +struct CapacityScanFixture { + _dirs: Vec, + disks: Vec, +} + +impl CapacityScanFixture { + fn new(specs: &[DiskSpec]) -> Self { + let mut dirs = Vec::with_capacity(specs.len()); + let mut disks = Vec::with_capacity(specs.len()); + + for (idx, spec) in specs.iter().enumerate() { + let dir = TempDir::new().expect("create temp dir"); + populate_files(dir.path(), spec.file_count, spec.file_size).expect("populate files"); + disks.push(CapacityDiskRef { + endpoint: format!("bench-disk-{idx}"), + drive_path: dir.path().to_string_lossy().into_owned(), + }); + dirs.push(dir); + } + + Self { _dirs: dirs, disks } + } +} + +fn populate_files(root: &Path, file_count: usize, file_size: usize) -> std::io::Result<()> { + let payload = vec![b'x'; file_size]; + let shard_count = (file_count / 512).clamp(1, 256); + + for shard_idx in 0..shard_count { + fs::create_dir_all(root.join(format!("bucket-{shard_idx:03}")))?; + } + + for file_idx in 0..file_count { + let subdir = root.join(format!("bucket-{:03}", file_idx % shard_count)); + let file_path: PathBuf = subdir.join(format!("object-{file_idx:08}.bin")); + fs::write(file_path, &payload)?; + } + + Ok(()) +} + +fn bench_capacity_scan(c: &mut Criterion) { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("create runtime"); + + let exact_fixture = CapacityScanFixture::new(&[DiskSpec { + file_count: 10_000, + file_size: EXACT_FILE_SIZE, + }]); + + let sampled_fixture = CapacityScanFixture::new(&[DiskSpec { + file_count: DEFAULT_SAMPLE_TRIGGER_FILE_COUNT, + file_size: SAMPLED_FILE_SIZE, + }]); + + let multi_disk_fixture = CapacityScanFixture::new(&[ + DiskSpec { + file_count: 4_000, + file_size: 1024, + }, + DiskSpec { + file_count: 6_000, + file_size: 2048, + }, + DiskSpec { + file_count: 8_000, + file_size: 4096, + }, + DiskSpec { + file_count: 10_000, + file_size: 1024, + }, + ]); + + let mut exact_group = c.benchmark_group("capacity_scan_exact"); + exact_group.sample_size(10); + exact_group.measurement_time(Duration::from_secs(10)); + exact_group.bench_function("single_disk_10k_4k", |b| { + b.iter(|| { + let summary = runtime + .block_on(scan_used_capacity_disks(black_box(&exact_fixture.disks))) + .expect("exact scan"); + black_box(summary); + }); + }); + exact_group.finish(); + + let mut sampled_group = c.benchmark_group("capacity_scan_sampled"); + sampled_group.sample_size(10); + sampled_group.measurement_time(Duration::from_secs(10)); + sampled_group.bench_function("single_disk_202k_1b", |b| { + b.iter(|| { + let summary = runtime + .block_on(scan_used_capacity_disks(black_box(&sampled_fixture.disks))) + .expect("sampled scan"); + black_box(summary); + }); + }); + sampled_group.finish(); + + let mut multi_disk_group = c.benchmark_group("capacity_scan_multi_disk"); + multi_disk_group.sample_size(10); + multi_disk_group.measurement_time(Duration::from_secs(10)); + multi_disk_group.bench_function("four_disks_mixed_exact", |b| { + b.iter(|| { + let summary = runtime + .block_on(scan_used_capacity_disks(black_box(&multi_disk_fixture.disks))) + .expect("multi-disk scan"); + black_box(summary); + }); + }); + multi_disk_group.finish(); +} + +criterion_group!(benches, bench_capacity_scan); +criterion_main!(benches); diff --git a/rustfs/src/capacity/capacity_manager.rs b/crates/object-capacity/src/capacity_manager.rs similarity index 56% rename from rustfs/src/capacity/capacity_manager.rs rename to crates/object-capacity/src/capacity_manager.rs index d8586770f..32cfc3c85 100644 --- a/rustfs/src/capacity/capacity_manager.rs +++ b/crates/object-capacity/src/capacity_manager.rs @@ -14,8 +14,10 @@ //! Hybrid Capacity Manager for efficient capacity statistics -use crate::app::admin_usecase::calculate_data_dir_used_capacity; +use super::scan::refresh_capacity_with_scope; +use super::types::CapacityDiskRef; use futures::FutureExt; +use rustfs_common::capacity_scope::{CapacityScope, CapacityScopeDisk, drain_global_dirty_scopes, take_capacity_scope}; use rustfs_config::{ DEFAULT_CAPACITY_ENABLE_DYNAMIC_TIMEOUT, DEFAULT_CAPACITY_FOLLOW_SYMLINKS, DEFAULT_CAPACITY_MAX_SYMLINK_DEPTH, DEFAULT_CAPACITY_MAX_TIMEOUT_SECS, DEFAULT_CAPACITY_MIN_TIMEOUT_SECS, DEFAULT_CAPACITY_STALL_TIMEOUT_SECS, @@ -26,12 +28,17 @@ use rustfs_config::{ ENV_CAPACITY_SAMPLE_RATE, ENV_CAPACITY_SCHEDULED_INTERVAL, ENV_CAPACITY_STALL_TIMEOUT, ENV_CAPACITY_STAT_TIMEOUT, ENV_CAPACITY_WRITE_FREQUENCY_THRESHOLD, ENV_CAPACITY_WRITE_TRIGGER_DELAY, }; -use rustfs_io_metrics::{record_capacity_current_bytes, record_capacity_update_completed, record_capacity_write_operation}; +use rustfs_io_metrics::capacity_metrics::{ + record_capacity_current_bytes, record_capacity_dirty_disk_count, record_capacity_refresh_inflight, + record_capacity_refresh_joiner, record_capacity_refresh_result, record_capacity_update_completed, + record_capacity_update_failed, record_capacity_write_operation, +}; use rustfs_utils::{get_env_bool, get_env_u64, get_env_usize}; +use std::collections::{HashMap, HashSet}; use std::future::Future; use std::panic::AssertUnwindSafe; use std::sync::Arc; -use std::time::{Duration, Instant}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use tokio::sync::{Mutex, RwLock, watch}; use tracing::{debug, info, warn}; @@ -297,6 +304,14 @@ pub struct CapacityUpdate { pub file_count: usize, /// Whether the value is estimated instead of exact. pub is_estimated: bool, + /// Per-disk breakdown captured from a successful refresh. + pub per_disk: Vec, + /// Expected disk count for a complete disk cache. + pub expected_disk_count: Option, + /// Whether this update should replace the current disk cache. + pub replaces_disk_cache: bool, + /// Dirty disks that can be cleared after the update is committed. + pub clear_dirty_disks: Vec, } impl CapacityUpdate { @@ -306,6 +321,10 @@ impl CapacityUpdate { total_used, file_count, is_estimated: false, + per_disk: Vec::new(), + expected_disk_count: None, + replaces_disk_cache: false, + clear_dirty_disks: Vec::new(), } } @@ -315,6 +334,10 @@ impl CapacityUpdate { total_used, file_count, is_estimated: true, + per_disk: Vec::new(), + expected_disk_count: None, + replaces_disk_cache: false, + clear_dirty_disks: Vec::new(), } } @@ -324,10 +347,27 @@ impl CapacityUpdate { total_used, file_count: 0, is_estimated: true, + per_disk: Vec::new(), + expected_disk_count: None, + replaces_disk_cache: false, + clear_dirty_disks: Vec::new(), } } } +#[derive(Clone, Debug)] +pub struct DiskCapacityUpdate { + pub disk: CapacityScopeDisk, + pub used_bytes: u64, + pub file_count: usize, + pub is_estimated: bool, +} + +#[derive(Clone, Debug)] +struct CachedDiskCapacity { + used_bytes: u64, +} + #[derive(Clone, Debug, PartialEq, Copy, Eq)] pub enum DataSource { /// Real-time statistics @@ -342,7 +382,7 @@ pub enum DataSource { } impl DataSource { - fn as_metric_label(self) -> &'static str { + pub fn as_metric_label(self) -> &'static str { match self { Self::RealTime => "realtime", Self::Scheduled => "scheduled", @@ -352,15 +392,60 @@ impl DataSource { } } +const WRITE_WINDOW_SECS: u64 = 60; +const WRITE_WINDOW_BUCKETS: usize = WRITE_WINDOW_SECS as usize; + +#[derive(Clone, Copy, Debug, Default)] +struct WriteBucket { + second: u64, + count: usize, +} + /// Write record for tracking write operations #[derive(Debug)] pub struct WriteRecord { /// Last write time - pub last_write_time: Instant, + pub last_write_time: Option, /// Write count pub write_count: usize, - /// Write time window (for frequency calculation) - pub write_window: Vec, + /// Fixed-size time buckets for the recent write window. + write_buckets: [WriteBucket; WRITE_WINDOW_BUCKETS], +} + +impl WriteRecord { + fn current_unix_second() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or(Duration::ZERO) + .as_secs() + } + + fn recent_write_count(&self, now_second: u64) -> usize { + self.write_buckets + .iter() + .filter(|bucket| bucket.count > 0 && now_second.saturating_sub(bucket.second) < WRITE_WINDOW_SECS) + .map(|bucket| bucket.count) + .sum() + } + + fn record_write(&mut self, now: Instant) -> usize { + let now_second = Self::current_unix_second(); + let bucket_idx = (now_second % WRITE_WINDOW_BUCKETS as u64) as usize; + let bucket = &mut self.write_buckets[bucket_idx]; + + if bucket.second != now_second { + *bucket = WriteBucket { + second: now_second, + count: 0, + }; + } + + bucket.count = bucket.count.saturating_add(1); + self.last_write_time = Some(now); + self.write_count = self.write_count.saturating_add(1); + + self.recent_write_count(now_second) + } } /// Hybrid strategy configuration @@ -430,6 +515,12 @@ pub struct HybridCapacityManager { cache: Arc>>, /// Write record write_record: Arc>, + /// Dirty disks recorded from write-side scope propagation. + dirty_disks: Arc>>, + /// Per-disk cache populated after a successful full refresh and updated by dirty subset refreshes. + disk_cache: Arc>>, + /// Whether the per-disk cache currently covers all known disks. + disk_cache_complete: Arc>, /// Configuration config: HybridStrategyConfig, /// Shared singleflight refresh state @@ -437,6 +528,17 @@ pub struct HybridCapacityManager { } impl HybridCapacityManager { + async fn sync_global_dirty_scopes(&self) { + let scopes = drain_global_dirty_scopes(); + if scopes.is_empty() { + return; + } + + let mut dirty_disks = self.dirty_disks.write().await; + dirty_disks.extend(scopes); + record_capacity_dirty_disk_count(dirty_disks.len()); + } + fn max_stale_age(&self) -> Duration { self.config .scheduled_update_interval @@ -448,10 +550,13 @@ impl HybridCapacityManager { Self { cache: Arc::new(RwLock::new(None)), write_record: Arc::new(RwLock::new(WriteRecord { - last_write_time: Instant::now(), + last_write_time: None, write_count: 0, - write_window: Vec::new(), + write_buckets: [WriteBucket::default(); WRITE_WINDOW_BUCKETS], })), + dirty_disks: Arc::new(RwLock::new(HashSet::new())), + disk_cache: Arc::new(RwLock::new(HashMap::new())), + disk_cache_complete: Arc::new(RwLock::new(false)), config, refresh_state: Arc::new(Mutex::new(RefreshState::default())), } @@ -471,49 +576,97 @@ impl HybridCapacityManager { /// Update capacity pub async fn update_capacity(&self, update: CapacityUpdate, source: DataSource) { let start = Instant::now(); + let mut total_used = update.total_used; + + if !update.per_disk.is_empty() { + let mut disk_cache = self.disk_cache.write().await; + let mut disk_cache_complete = self.disk_cache_complete.write().await; + + if update.replaces_disk_cache && update.expected_disk_count == Some(update.per_disk.len()) { + disk_cache.clear(); + for entry in &update.per_disk { + disk_cache.insert( + entry.disk.clone(), + CachedDiskCapacity { + used_bytes: entry.used_bytes, + }, + ); + } + *disk_cache_complete = true; + total_used = disk_cache.values().map(|entry| entry.used_bytes).sum(); + } else if *disk_cache_complete { + for entry in &update.per_disk { + disk_cache.insert( + entry.disk.clone(), + CachedDiskCapacity { + used_bytes: entry.used_bytes, + }, + ); + } + total_used = disk_cache.values().map(|entry| entry.used_bytes).sum(); + } + } + let mut cache = self.cache.write().await; *cache = Some(CachedCapacity { - total_used: update.total_used, + total_used, last_update: Instant::now(), file_count: update.file_count, is_estimated: update.is_estimated, source, }); + if !update.clear_dirty_disks.is_empty() { + let mut dirty_disks = self.dirty_disks.write().await; + for disk in &update.clear_dirty_disks { + dirty_disks.remove(disk); + } + record_capacity_dirty_disk_count(dirty_disks.len()); + } + debug!( "Capacity updated: {} bytes, files={}, estimated={}, source: {:?}", - update.total_used, update.file_count, update.is_estimated, source + total_used, update.file_count, update.is_estimated, source ); - record_capacity_current_bytes(update.total_used); - record_capacity_update_completed(source.as_metric_label(), start.elapsed(), update.total_used, update.is_estimated); + record_capacity_current_bytes(total_used); + record_capacity_update_completed(source.as_metric_label(), start.elapsed(), total_used, update.is_estimated); } /// Record write operation pub async fn record_write_operation(&self) { let mut record = self.write_record.write().await; - record.last_write_time = Instant::now(); - record.write_count += 1; - - // Maintain write time window (keep last 1 minute) - // Cap the window size to prevent unbounded memory growth at high write rates - const MAX_WRITE_WINDOW_SIZE: usize = 10000; let now = Instant::now(); - record - .write_window - .retain(|&t| now.duration_since(t) < Duration::from_secs(60)); - // Only push if under the cap to prevent unbounded growth - if record.write_window.len() < MAX_WRITE_WINDOW_SIZE { - record.write_window.push(now); - } + let recent_write_count = record.record_write(now); - record_capacity_write_operation(record.write_window.len()); + record_capacity_write_operation(recent_write_count); debug!( "Write operation recorded: total writes = {}, recent writes = {}", - record.write_count, - record.write_window.len() + record.write_count, recent_write_count ); } + /// Record write scope propagated from the storage layer. + pub async fn mark_dirty_scope(&self, scope: &CapacityScope) { + if scope.disks.is_empty() { + return; + } + + let mut dirty_disks = self.dirty_disks.write().await; + dirty_disks.extend(scope.disks.iter().cloned()); + record_capacity_dirty_disk_count(dirty_disks.len()); + } + + /// Record a write operation and consume any propagated disk scope bound to the token. + pub async fn record_write_operation_with_scope_token(&self, scope_token: Option) { + if let Some(token) = scope_token + && let Some(scope) = take_capacity_scope(token) + { + self.mark_dirty_scope(&scope).await; + } + + self.record_write_operation().await; + } + /// Check if fast update is needed pub async fn needs_fast_update(&self) -> bool { if !self.config.enable_smart_update { @@ -529,19 +682,31 @@ impl HybridCapacityManager { return false; } - let write_record = self.write_record.read().await; - let time_since_write = write_record.last_write_time.elapsed(); - - // Recent write, trigger fast update - if time_since_write < self.config.fast_update_threshold { - debug!("Recent write detected ({:?} ago), needs fast update", time_since_write); - return true; + if !self.config.enable_write_trigger { + return false; } - // High write frequency, trigger update - let write_frequency = write_record.write_window.len(); - if write_frequency > self.config.write_frequency_threshold { - debug!("High write frequency detected ({} writes/min), needs fast update", write_frequency); + let write_record = self.write_record.read().await; + let write_frequency = write_record.recent_write_count(WriteRecord::current_unix_second()); + if write_frequency <= self.config.write_frequency_threshold { + return false; + } + + if let Some(last_write_time) = write_record.last_write_time { + let time_since_write = last_write_time.elapsed(); + + if time_since_write < self.config.write_trigger_delay { + debug!( + "Write-triggered refresh still debounced ({:?} ago, trigger_delay={:?}, writes/min={})", + time_since_write, self.config.write_trigger_delay, write_frequency + ); + return false; + } + + debug!( + "Write-triggered refresh eligible after debounce ({:?} ago, trigger_delay={:?}, writes/min={})", + time_since_write, self.config.write_trigger_delay, write_frequency + ); return true; } } @@ -560,7 +725,19 @@ impl HybridCapacityManager { #[allow(dead_code)] pub async fn get_write_frequency(&self) -> usize { let record = self.write_record.read().await; - record.write_window.len() + record.recent_write_count(WriteRecord::current_unix_second()) + } + + /// Snapshot the currently dirty disks recorded from write-side scope propagation. + pub async fn get_dirty_disks(&self) -> Vec { + self.sync_global_dirty_scopes().await; + let dirty_disks = self.dirty_disks.read().await; + dirty_disks.iter().cloned().collect() + } + + /// Returns true if the manager has a complete per-disk cache and can safely refresh only dirty disks. + pub async fn can_refresh_dirty_subset(&self) -> bool { + *self.disk_cache_complete.read().await } /// Run a singleflight refresh. Callers either join an existing in-flight refresh or become the leader. @@ -577,6 +754,7 @@ impl HybridCapacityManager { if state.running { // Subscribe while holding the lock so the send that completes the current // refresh cycle cannot happen before we are subscribed. + record_capacity_refresh_joiner(source.as_metric_label()); Some(state.result_tx.subscribe()) } else { // Become the leader. Create a fresh channel so that joiners from a previous @@ -584,6 +762,7 @@ impl HybridCapacityManager { let (tx, _) = watch::channel(None); state.result_tx = tx; state.running = true; + record_capacity_refresh_inflight(1); None } }; @@ -603,6 +782,7 @@ impl HybridCapacityManager { .unwrap_or_else(|| Err("capacity refresh completed without a result".to_string())); } + let refresh_start = Instant::now(); let result = AssertUnwindSafe(refresh_fn()).catch_unwind().await.unwrap_or_else(|err| { warn!(error = ?err, "capacity refresh function panicked"); Err("capacity refresh panicked".to_string()) @@ -610,10 +790,20 @@ impl HybridCapacityManager { if let Ok(update) = &result { self.update_capacity(update.clone(), source).await; } + let refresh_duration = refresh_start.elapsed(); + if result.is_err() { + record_capacity_update_failed(source.as_metric_label()); + } + record_capacity_refresh_result( + source.as_metric_label(), + if result.is_ok() { "success" } else { "error" }, + refresh_duration, + ); { let mut state = self.refresh_state.lock().await; state.running = false; + record_capacity_refresh_inflight(0); let _ = state.result_tx.send(Some(result.clone())); } @@ -634,6 +824,7 @@ impl HybridCapacityManager { let (tx, _) = watch::channel(None); state.result_tx = tx; state.running = true; + record_capacity_refresh_inflight(1); true } }; @@ -643,6 +834,7 @@ impl HybridCapacityManager { } tokio::spawn(async move { + let refresh_start = Instant::now(); let result = AssertUnwindSafe(refresh_fn()).catch_unwind().await.unwrap_or_else(|err| { warn!(error = ?err, "capacity refresh function panicked"); Err("capacity refresh panicked".to_string()) @@ -650,9 +842,19 @@ impl HybridCapacityManager { if let Ok(update) = &result { self.update_capacity(update.clone(), source).await; } + let refresh_duration = refresh_start.elapsed(); + if result.is_err() { + record_capacity_update_failed(source.as_metric_label()); + } + record_capacity_refresh_result( + source.as_metric_label(), + if result.is_ok() { "success" } else { "error" }, + refresh_duration, + ); let mut state = self.refresh_state.lock().await; state.running = false; + record_capacity_refresh_inflight(0); let _ = state.result_tx.send(Some(result)); }); @@ -670,7 +872,6 @@ impl HybridCapacityManager { } /// Return whether a refresh is currently in flight. - #[cfg(test)] pub async fn refresh_in_progress(&self) -> bool { self.refresh_state.lock().await.running } @@ -692,20 +893,19 @@ pub fn get_capacity_manager() -> Arc { /// without affecting the global singleton, avoiding test pollution. /// /// # Example -/// ```no_run +/// ```ignore /// let manager = create_isolated_manager(HybridStrategyConfig::default()); /// manager /// .update_capacity(CapacityUpdate::exact(1000, 0), DataSource::RealTime) /// .await; /// ``` -#[cfg(test)] #[allow(dead_code)] pub fn create_isolated_manager(config: HybridStrategyConfig) -> Arc { Arc::new(HybridCapacityManager::new(config)) } /// Start background update task -pub async fn start_background_task(disks: Vec) { +pub async fn start_background_task(disks: Vec) { let manager = get_capacity_manager(); let mut interval = manager.get_config().scheduled_update_interval; @@ -727,12 +927,10 @@ pub async fn start_background_task(disks: Vec) { let disks = disks.clone(); let started = manager .clone() - .spawn_refresh_if_needed(DataSource::Scheduled, move || async move { - calculate_data_dir_used_capacity(&disks) - .await - .map(|scan| scan.to_capacity_update()) - .map_err(|e| e.to_string()) - }) + .spawn_refresh_if_needed( + DataSource::Scheduled, + move || async move { refresh_capacity_with_scope(disks, false).await }, + ) .await; if started { @@ -751,11 +949,14 @@ pub async fn start_background_task(disks: Vec) { #[cfg(test)] mod tests { use super::*; + use rustfs_common::capacity_scope::{CapacityScope, CapacityScopeDisk, record_capacity_scope, record_global_dirty_scope}; use rustfs_config::{ ENV_CAPACITY_FAST_UPDATE_THRESHOLD, ENV_CAPACITY_MAX_FILES_THRESHOLD, ENV_CAPACITY_SAMPLE_RATE, ENV_CAPACITY_STAT_TIMEOUT, ENV_CAPACITY_WRITE_FREQUENCY_THRESHOLD, ENV_CAPACITY_WRITE_TRIGGER_DELAY, }; use serial_test::serial; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; #[test] #[serial] @@ -892,6 +1093,22 @@ mod tests { assert_eq!(cached.unwrap().total_used, 1000); } + #[tokio::test] + #[serial] + async fn test_update_capacity_preserves_retrieval_metadata() { + let manager = HybridCapacityManager::from_env(); + + manager + .update_capacity(CapacityUpdate::exact(1000, 10), DataSource::RealTime) + .await; + + let cached = manager.get_capacity().await.unwrap(); + assert_eq!(cached.total_used, 1000); + assert_eq!(cached.file_count, 10); + assert_eq!(cached.source, DataSource::RealTime); + assert!(!cached.is_estimated); + } + #[tokio::test] #[serial] async fn test_record_write_operation() { @@ -903,6 +1120,18 @@ mod tests { assert_eq!(frequency, 1); } + #[tokio::test] + #[serial] + async fn test_write_frequency_window() { + let manager = HybridCapacityManager::from_env(); + + for _ in 0..20 { + manager.record_write_operation().await; + } + + assert_eq!(manager.get_write_frequency().await, 20); + } + #[tokio::test] #[serial] async fn test_needs_fast_update() { @@ -920,6 +1149,326 @@ mod tests { assert!(!manager.needs_fast_update().await); } + #[tokio::test] + #[serial] + async fn test_cache_age_tracking() { + let manager = HybridCapacityManager::from_env(); + + assert!(manager.get_cache_age().await.is_none()); + + manager + .update_capacity(CapacityUpdate::exact(1000, 1), DataSource::RealTime) + .await; + + let age = manager.get_cache_age().await.unwrap(); + assert!(age < Duration::from_secs(1)); + + tokio::time::sleep(Duration::from_millis(100)).await; + + let age = manager.get_cache_age().await.unwrap(); + assert!(age >= Duration::from_millis(100)); + } + + #[tokio::test] + #[serial] + async fn test_data_source_tracking() { + let manager = HybridCapacityManager::from_env(); + + for source in [ + DataSource::RealTime, + DataSource::Scheduled, + DataSource::WriteTriggered, + DataSource::Fallback, + ] { + manager.update_capacity(CapacityUpdate::exact(1000, 1), source).await; + assert_eq!(manager.get_capacity().await.unwrap().source, source); + } + } + + #[tokio::test] + #[serial] + async fn test_needs_fast_update_waits_for_write_trigger_delay() { + let manager = create_isolated_manager(HybridStrategyConfig { + scheduled_update_interval: Duration::from_secs(60), + write_trigger_delay: Duration::from_millis(50), + write_frequency_threshold: 1, + fast_update_threshold: Duration::from_millis(10), + enable_smart_update: true, + enable_write_trigger: true, + }); + + manager + .update_capacity(CapacityUpdate::exact(1000, 0), DataSource::RealTime) + .await; + tokio::time::sleep(Duration::from_millis(15)).await; + + manager.record_write_operation().await; + manager.record_write_operation().await; + tokio::time::sleep(Duration::from_millis(5)).await; + + assert!( + !manager.needs_fast_update().await, + "write-triggered refresh should wait for debounce delay after a qualifying burst" + ); + + tokio::time::sleep(Duration::from_millis(60)).await; + assert!(manager.needs_fast_update().await); + } + + #[tokio::test] + #[serial] + async fn test_needs_fast_update_respects_enable_write_trigger() { + let manager = create_isolated_manager(HybridStrategyConfig { + scheduled_update_interval: Duration::from_secs(60), + write_trigger_delay: Duration::from_secs(60), + write_frequency_threshold: 1, + fast_update_threshold: Duration::from_millis(10), + enable_smart_update: true, + enable_write_trigger: false, + }); + + manager + .update_capacity(CapacityUpdate::exact(1000, 0), DataSource::RealTime) + .await; + tokio::time::sleep(Duration::from_millis(15)).await; + + manager.record_write_operation().await; + manager.record_write_operation().await; + + assert!( + !manager.needs_fast_update().await, + "write-triggered refresh should be disabled when enable_write_trigger is false" + ); + } + + #[tokio::test] + #[serial] + async fn test_concurrent_access() { + let manager = Arc::new(HybridCapacityManager::from_env()); + let mut handles = Vec::new(); + + for i in 0..10 { + let mgr = manager.clone(); + handles.push(tokio::spawn(async move { + mgr.update_capacity(CapacityUpdate::exact(i as u64 * 100, i), DataSource::RealTime) + .await; + mgr.record_write_operation().await; + })); + } + + for handle in handles { + handle.await.unwrap(); + } + + assert!(manager.get_capacity().await.is_some()); + assert_eq!(manager.get_write_frequency().await, 10); + } + + #[tokio::test] + #[serial] + async fn test_performance_overhead() { + let manager = Arc::new(HybridCapacityManager::from_env()); + let start = Instant::now(); + + for i in 0..1000 { + manager + .update_capacity(CapacityUpdate::exact(i as u64, i), DataSource::RealTime) + .await; + manager.record_write_operation().await; + let _ = manager.get_capacity().await; + } + + assert!(start.elapsed() < Duration::from_secs(1)); + } + + #[tokio::test] + #[serial] + async fn test_refresh_or_join_singleflight() { + let manager = Arc::new(HybridCapacityManager::from_env()); + let calls = Arc::new(AtomicUsize::new(0)); + + let mgr1 = manager.clone(); + let calls1 = calls.clone(); + let first = tokio::spawn(async move { + mgr1.refresh_or_join(DataSource::Scheduled, move || async move { + calls1.fetch_add(1, Ordering::SeqCst); + tokio::time::sleep(Duration::from_millis(50)).await; + Ok(CapacityUpdate::exact(2048, 8)) + }) + .await + }); + + tokio::time::sleep(Duration::from_millis(10)).await; + + let mgr2 = manager.clone(); + let calls2 = calls.clone(); + let second = tokio::spawn(async move { + mgr2.refresh_or_join(DataSource::WriteTriggered, move || async move { + calls2.fetch_add(1, Ordering::SeqCst); + Ok(CapacityUpdate::exact(4096, 16)) + }) + .await + }); + + let first = first.await.unwrap().unwrap(); + let second = second.await.unwrap().unwrap(); + + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!(first.total_used, 2048); + assert_eq!(second.total_used, 2048); + let cached = manager.get_capacity().await.unwrap(); + assert_eq!(cached.total_used, 2048); + assert_eq!(cached.file_count, 8); + } + + #[tokio::test] + #[serial] + async fn test_spawn_refresh_if_needed_deduplicates_background_refresh() { + let manager = Arc::new(HybridCapacityManager::from_env()); + let calls = Arc::new(AtomicUsize::new(0)); + + let first_manager = manager.clone(); + let first_calls = calls.clone(); + let started = first_manager + .clone() + .spawn_refresh_if_needed(DataSource::Scheduled, move || async move { + first_calls.fetch_add(1, Ordering::SeqCst); + tokio::time::sleep(Duration::from_millis(50)).await; + Ok(CapacityUpdate::estimated(8192, 32)) + }) + .await; + assert!(started); + + let second_manager = manager.clone(); + let second_calls = calls.clone(); + let started = second_manager + .clone() + .spawn_refresh_if_needed(DataSource::Scheduled, move || async move { + second_calls.fetch_add(1, Ordering::SeqCst); + Ok(CapacityUpdate::exact(1, 1)) + }) + .await; + assert!(!started); + + tokio::time::sleep(Duration::from_millis(100)).await; + + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert!(!manager.refresh_in_progress().await); + let cached = manager.get_capacity().await.unwrap(); + assert_eq!(cached.total_used, 8192); + assert!(cached.is_estimated); + } + + #[tokio::test] + #[serial] + async fn test_record_write_operation_with_scope_token_marks_dirty_disks() { + let manager = create_isolated_manager(HybridStrategyConfig::default()); + let token = uuid::Uuid::new_v4(); + record_capacity_scope( + token, + CapacityScope { + disks: vec![CapacityScopeDisk { + endpoint: "node-a".to_string(), + drive_path: "/tmp/disk-a".to_string(), + }], + }, + ); + + manager.record_write_operation_with_scope_token(Some(token)).await; + + let dirty_disks = manager.get_dirty_disks().await; + assert_eq!(dirty_disks.len(), 1); + assert_eq!(dirty_disks[0].endpoint, "node-a"); + assert_eq!(dirty_disks[0].drive_path, "/tmp/disk-a"); + assert_eq!(manager.get_write_frequency().await, 1); + } + + #[tokio::test] + #[serial] + async fn test_get_dirty_disks_drains_global_dirty_scope_registry() { + let manager = create_isolated_manager(HybridStrategyConfig::default()); + record_global_dirty_scope(CapacityScope { + disks: vec![CapacityScopeDisk { + endpoint: "node-bg".to_string(), + drive_path: "/tmp/disk-bg".to_string(), + }], + }); + + let dirty_disks = manager.get_dirty_disks().await; + assert_eq!(dirty_disks.len(), 1); + assert_eq!(dirty_disks[0].endpoint, "node-bg"); + assert_eq!(dirty_disks[0].drive_path, "/tmp/disk-bg"); + + let second_read = manager.get_dirty_disks().await; + assert_eq!(second_read.len(), 1); + } + + #[tokio::test] + #[serial] + async fn test_update_capacity_recomputes_total_from_disk_cache_for_subset_refresh() { + let manager = create_isolated_manager(HybridStrategyConfig::default()); + + manager + .update_capacity( + CapacityUpdate { + total_used: 300, + file_count: 3, + is_estimated: false, + per_disk: vec![ + DiskCapacityUpdate { + disk: CapacityScopeDisk { + endpoint: "node-a".to_string(), + drive_path: "/tmp/disk-a".to_string(), + }, + used_bytes: 100, + file_count: 1, + is_estimated: false, + }, + DiskCapacityUpdate { + disk: CapacityScopeDisk { + endpoint: "node-b".to_string(), + drive_path: "/tmp/disk-b".to_string(), + }, + used_bytes: 200, + file_count: 2, + is_estimated: false, + }, + ], + expected_disk_count: Some(2), + replaces_disk_cache: true, + clear_dirty_disks: Vec::new(), + }, + DataSource::RealTime, + ) + .await; + + manager + .update_capacity( + CapacityUpdate { + total_used: 150, + file_count: 1, + is_estimated: false, + per_disk: vec![DiskCapacityUpdate { + disk: CapacityScopeDisk { + endpoint: "node-a".to_string(), + drive_path: "/tmp/disk-a".to_string(), + }, + used_bytes: 150, + file_count: 1, + is_estimated: false, + }], + expected_disk_count: Some(1), + replaces_disk_cache: false, + clear_dirty_disks: Vec::new(), + }, + DataSource::WriteTriggered, + ) + .await; + + let cached = manager.get_capacity().await.unwrap(); + assert_eq!(cached.total_used, 350); + } + #[tokio::test] #[serial] async fn test_config_from_env() { diff --git a/crates/object-capacity/src/lib.rs b/crates/object-capacity/src/lib.rs new file mode 100644 index 000000000..c5730cdbf --- /dev/null +++ b/crates/object-capacity/src/lib.rs @@ -0,0 +1,20 @@ +// Copyright 2024 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. + +pub mod capacity_manager; +pub mod scan; +pub mod types; + +pub use scan::scan_used_capacity_disks; +pub use types::{CapacityDiskRef, CapacityScanSummary}; diff --git a/crates/object-capacity/src/scan.rs b/crates/object-capacity/src/scan.rs new file mode 100644 index 000000000..9fab908ab --- /dev/null +++ b/crates/object-capacity/src/scan.rs @@ -0,0 +1,903 @@ +// Copyright 2024 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::capacity_manager::{ + CapacityUpdate, DiskCapacityUpdate, HybridCapacityManager, get_enable_dynamic_timeout, get_follow_symlinks, + get_max_files_threshold, get_max_symlink_depth, get_max_timeout, get_min_timeout, get_sample_rate, get_stall_timeout, + get_stat_timeout, +}; +use super::types::{CapacityDiskRef, CapacityScanResult, CapacityScanSummary}; +use futures::{StreamExt, stream}; +use rustfs_common::capacity_scope::CapacityScopeDisk; +use rustfs_io_metrics::capacity_metrics::{ + record_capacity_dynamic_timeout, record_capacity_scan_disk, record_capacity_scan_mode, record_capacity_scan_sampling, + record_capacity_stall_detected, record_capacity_symlink, record_capacity_timeout_fallback, +}; +use std::collections::HashSet; +use std::path::{Path, PathBuf}; +use std::time::{Duration, Instant}; +use tracing::{debug, info, warn}; +use walkdir::WalkDir; + +const MAX_CAPACITY_SCAN_CONCURRENCY: usize = 4; +const CAPACITY_PROGRESS_CHECK_STRIDE: usize = 512; + +#[derive(Debug)] +struct DiskScanOutcome { + disk_label: String, + drive_path: String, + duration: Duration, + result: Result, +} + +#[derive(Debug, Clone)] +struct DiskCapacityScanResult { + disk: CapacityScopeDisk, + scan: CapacityScanResult, +} + +#[derive(Debug, Clone)] +struct CapacityScanReport { + summary: CapacityScanResult, + per_disk: Vec, +} + +impl CapacityScanReport { + fn into_capacity_update(self, expected_disk_count: usize, replaces_disk_cache: bool) -> CapacityUpdate { + let mut update = if self.summary.is_estimated { + CapacityUpdate::estimated(self.summary.used_bytes, self.summary.file_count) + } else { + CapacityUpdate::exact(self.summary.used_bytes, self.summary.file_count) + }; + + if !self.summary.had_partial_errors && self.per_disk.len() == expected_disk_count { + update.per_disk = self + .per_disk + .into_iter() + .map(|entry| DiskCapacityUpdate { + disk: entry.disk, + used_bytes: entry.scan.used_bytes, + file_count: entry.scan.file_count, + is_estimated: entry.scan.is_estimated, + }) + .collect(); + update.expected_disk_count = Some(expected_disk_count); + update.replaces_disk_cache = replaces_disk_cache; + update.clear_dirty_disks = update.per_disk.iter().map(|entry| entry.disk.clone()).collect(); + } + + update + } +} + +fn disk_metric_label(disk: &CapacityDiskRef) -> String { + let mount_name = Path::new(&disk.drive_path) + .file_name() + .and_then(|value| value.to_str()) + .filter(|value| !value.is_empty()) + .unwrap_or(disk.drive_path.as_str()); + format!("{}:{mount_name}", disk.endpoint) +} + +fn disk_scope_key(disk: &CapacityDiskRef) -> CapacityScopeDisk { + CapacityScopeDisk { + endpoint: disk.endpoint.clone(), + drive_path: disk.drive_path.clone(), + } +} + +async fn scan_disk_used_capacity(disk: CapacityDiskRef) -> DiskScanOutcome { + let disk_label = disk_metric_label(&disk); + let drive_path = disk.drive_path.clone(); + let start = Instant::now(); + let result = get_dir_size_async(Path::new(&drive_path)).await; + + DiskScanOutcome { + disk_label, + drive_path, + duration: start.elapsed(), + result, + } +} + +async fn calculate_data_dir_used_capacity_report( + disks: &[CapacityDiskRef], +) -> Result> { + let start = Instant::now(); + let mut total_used = 0u64; + let mut total_files = 0usize; + let mut total_sampled = 0usize; + let mut has_failure = false; + let mut has_success = false; + let mut is_estimated = false; + let mut per_disk = Vec::with_capacity(disks.len()); + + let concurrency_limit = disks.len().clamp(1, MAX_CAPACITY_SCAN_CONCURRENCY); + let mut scans = stream::iter(disks.iter().cloned().map(scan_disk_used_capacity)).buffer_unordered(concurrency_limit); + + while let Some(outcome) = scans.next().await { + match outcome.result { + Ok(scan) => { + record_capacity_scan_disk( + outcome.disk_label.as_str(), + outcome.duration, + scan.file_count, + scan.sampled_count, + scan.is_estimated, + scan.had_partial_errors, + ); + debug!( + "Data directory {} size: {} bytes, files={}, sampled={}, estimated={}, duration={:?}", + outcome.drive_path, scan.used_bytes, scan.file_count, scan.sampled_count, scan.is_estimated, outcome.duration + ); + total_used += scan.used_bytes; + total_files += scan.file_count; + total_sampled += scan.sampled_count; + is_estimated |= scan.is_estimated; + has_failure |= scan.had_partial_errors; + has_success = true; + if let Some(disk) = disks + .iter() + .find(|disk| disk.drive_path == outcome.drive_path && disk_metric_label(disk) == outcome.disk_label) + { + per_disk.push(DiskCapacityScanResult { + disk: disk_scope_key(disk), + scan, + }); + } + } + Err(e) => { + record_capacity_scan_disk(outcome.disk_label.as_str(), outcome.duration, 0, 0, false, true); + warn!("Failed to get size for directory {}: {:?}", outcome.drive_path, e); + has_failure = true; + } + } + } + + if !has_success { + return Err("All directories failed to calculate size".into()); + } + + if has_failure { + warn!("Some directories failed to calculate size, result may be incomplete"); + } + + let mut summary = CapacityScanResult { + used_bytes: total_used, + file_count: total_files, + sampled_count: total_sampled, + is_estimated, + scan_duration: start.elapsed(), + had_partial_errors: false, + }; + + if has_failure { + summary = summary.with_partial_errors(); + } + + Ok(CapacityScanReport { summary, per_disk }) +} + +/// Calculate actual used capacity of all data directories. +pub(crate) async fn calculate_data_dir_used_capacity( + disks: &[CapacityDiskRef], +) -> Result> { + Ok(calculate_data_dir_used_capacity_report(disks).await?.summary) +} + +pub async fn select_capacity_refresh_disks( + capacity_manager: &HybridCapacityManager, + disks: &[CapacityDiskRef], +) -> (Vec, bool) { + if !capacity_manager.can_refresh_dirty_subset().await { + return (disks.to_vec(), false); + } + + let dirty_disks = capacity_manager.get_dirty_disks().await; + if dirty_disks.is_empty() { + return (disks.to_vec(), false); + } + + let dirty_set: HashSet = dirty_disks.into_iter().collect(); + let selected: Vec<_> = disks + .iter() + .filter(|disk| dirty_set.contains(&disk_scope_key(disk))) + .cloned() + .collect(); + + if selected.is_empty() || selected.len() >= disks.len() { + (disks.to_vec(), false) + } else { + (selected, true) + } +} + +pub async fn refresh_capacity_with_scope(disks: Vec, dirty_subset: bool) -> Result { + let report = calculate_data_dir_used_capacity_report(&disks) + .await + .map_err(|e| e.to_string())?; + + if dirty_subset && report.summary.had_partial_errors { + return Err("dirty subset refresh had partial errors".to_string()); + } + + Ok(report.into_capacity_update(disks.len(), !dirty_subset)) +} + +/// Scan the provided local disk roots and return a summarized used-capacity result. +/// +/// This is primarily intended for benchmarks and operational tooling that need to exercise +/// the same scan path as admin capacity queries without going through the full admin stack. +pub async fn scan_used_capacity_disks( + disks: &[CapacityDiskRef], +) -> Result> { + Ok(calculate_data_dir_used_capacity(disks).await?.into()) +} + +/// Tracker for symlink resolution with circular reference detection. +struct SymlinkTracker { + visited: HashSet, + symlink_count: usize, + symlink_size: u64, + max_depth: u8, +} + +impl SymlinkTracker { + fn new(max_depth: u8) -> Self { + Self { + visited: HashSet::new(), + symlink_count: 0, + symlink_size: 0, + max_depth, + } + } + + fn should_follow(&self, path: &Path, depth: u8) -> bool { + if depth >= self.max_depth { + debug!("Symlink depth limit reached: {} >= {}, not following {:?}", depth, self.max_depth, path); + return false; + } + + if self.visited.contains(path) { + warn!("Circular symlink reference detected: {:?}, skipping", path); + return false; + } + + true + } + + fn record_symlink(&mut self, path: PathBuf, size: u64) { + if self.visited.insert(path) { + self.symlink_count += 1; + self.symlink_size += size; + record_capacity_symlink(size); + } + } + + fn get_stats(&self) -> (usize, u64) { + (self.symlink_count, self.symlink_size) + } +} + +/// Monitor for directory traversal progress with timeout and stall detection. +struct ProgressMonitor { + start_time: Instant, + last_check: Instant, + last_checkpoint_files: usize, + timeout: Duration, + min_timeout: Duration, + max_timeout: Duration, + stall_timeout: Duration, + enable_dynamic_timeout: bool, + used_dynamic_timeout: bool, +} + +impl ProgressMonitor { + fn new( + base_timeout: Duration, + min_timeout: Duration, + max_timeout: Duration, + stall_timeout: Duration, + enable_dynamic: bool, + ) -> Self { + Self { + start_time: Instant::now(), + last_check: Instant::now(), + last_checkpoint_files: 0, + timeout: base_timeout, + min_timeout, + max_timeout, + stall_timeout, + enable_dynamic_timeout: enable_dynamic, + used_dynamic_timeout: false, + } + } + + fn calculate_dynamic_timeout(&mut self, file_count: usize, avg_file_size: u64) -> Duration { + if !self.enable_dynamic_timeout { + return self.timeout; + } + + self.used_dynamic_timeout = true; + + let file_factor = (file_count as f64).sqrt() * 0.01; + let size_factor = if avg_file_size > 0 { + (avg_file_size as f64).log(10.0) * 0.05 + } else { + 0.0 + }; + + let multiplier = 1.0 + file_factor + size_factor; + let adjusted_timeout = self.timeout.mul_f64(multiplier.min(5.0)); + let clamped_timeout = adjusted_timeout.max(self.min_timeout).min(self.max_timeout); + + debug!( + "Dynamic timeout calculation: files={}, avg_size={}, multiplier={:.2}, base_timeout={:?}, adjusted_timeout={:?}, clamped_timeout={:?}", + file_count, avg_file_size, multiplier, self.timeout, adjusted_timeout, clamped_timeout + ); + + clamped_timeout + } + + fn update_and_check_timeout(&mut self, files_processed: usize, avg_file_size: u64) -> Result<(), std::io::Error> { + let elapsed = self.start_time.elapsed(); + let dynamic_timeout = if self.enable_dynamic_timeout { + self.calculate_dynamic_timeout(files_processed, avg_file_size) + } else { + self.timeout + }; + + if elapsed >= dynamic_timeout { + warn!( + "Directory size calculation timeout after {} files, elapsed: {:?}, timeout: {:?}", + files_processed, elapsed, dynamic_timeout + ); + + if self.enable_dynamic_timeout { + record_capacity_dynamic_timeout(dynamic_timeout); + } + + return Err(std::io::Error::new( + std::io::ErrorKind::TimedOut, + format!("Timeout after {} files", files_processed), + )); + } + + let now = Instant::now(); + if now.duration_since(self.last_check) >= self.stall_timeout { + let files_per_checkpoint = files_processed.saturating_sub(self.last_checkpoint_files); + + if files_per_checkpoint == 0 && files_processed > 0 { + warn!( + "No progress detected for {:?}, possible stall at {} files", + self.stall_timeout, files_processed + ); + + record_capacity_stall_detected(); + + return Err(std::io::Error::new( + std::io::ErrorKind::TimedOut, + format!("Stall detected at {} files", files_processed), + )); + } + + self.last_check = now; + self.last_checkpoint_files = files_processed; + } + + Ok(()) + } + + fn record_timeout_fallback(&self) { + record_capacity_timeout_fallback(); + } +} + +async fn get_dir_size_async(path: &Path) -> Result { + let path = path.to_path_buf(); + + let max_files_threshold = get_max_files_threshold(); + let base_timeout = get_stat_timeout(); + let min_timeout = get_min_timeout(); + let max_timeout = get_max_timeout(); + let stall_timeout = get_stall_timeout(); + let sample_rate = get_sample_rate(); + let enable_dynamic_timeout = get_enable_dynamic_timeout(); + let follow_symlinks = get_follow_symlinks(); + let max_symlink_depth = get_max_symlink_depth(); + + let effective_sample_rate = if sample_rate == 0 { + warn!("Invalid sampling configuration: sample_rate=0. Clamping to 1 to avoid panic."); + 1 + } else { + sample_rate + }; + + tokio::task::spawn_blocking(move || { + if !path.exists() { + return Err(std::io::Error::new( + std::io::ErrorKind::NotFound, + format!("Directory not found: {:?}", path), + )); + } + + let start_time = Instant::now(); + let mut exact_prefix_bytes = 0u64; + let mut overflow_sampled_bytes = 0u64; + let mut file_count = 0usize; + let mut sampled_count = 0usize; + let mut had_partial_errors = false; + let mut last_progress_check_files = 0usize; + + let mut symlink_tracker = SymlinkTracker::new(max_symlink_depth); + let mut progress_monitor = + ProgressMonitor::new(base_timeout, min_timeout, max_timeout, stall_timeout, enable_dynamic_timeout); + + let walker = WalkDir::new(&path) + .follow_links(follow_symlinks) + .follow_root_links(follow_symlinks) + .into_iter(); + + for entry_result in walker { + let entry = match entry_result { + Ok(entry) => entry, + Err(err) => { + warn!("Failed to traverse directory entry under {:?}: {}", path, err); + had_partial_errors = true; + continue; + } + }; + + if follow_symlinks + && entry.path_is_symlink() + && let Ok(target) = std::fs::read_link(entry.path()) + && symlink_tracker.should_follow(&target, entry.depth().min(u8::MAX as usize) as u8) + { + symlink_tracker.record_symlink(target, 0); + } + + let file_type = entry.file_type(); + if file_type.is_dir() { + continue; + } + + if file_type.is_symlink() || !file_type.is_file() { + continue; + } + + let metadata = match entry.metadata() { + Ok(meta) => meta, + Err(err) => { + warn!("Failed to get metadata for {:?}: {}", entry.path(), err); + had_partial_errors = true; + continue; + } + }; + + file_count += 1; + let exact_count = file_count.min(max_files_threshold); + let avg_size = if exact_count > 0 { + exact_prefix_bytes / exact_count as u64 + } else { + 0 + }; + + let should_check_progress = + file_count == 1 || file_count.saturating_sub(last_progress_check_files) >= CAPACITY_PROGRESS_CHECK_STRIDE; + + if should_check_progress && let Err(e) = progress_monitor.update_and_check_timeout(file_count, avg_size) { + if sampled_count > 0 { + let overflow_count = file_count.saturating_sub(max_files_threshold); + let estimated_overflow = overflow_sampled_bytes.saturating_mul(overflow_count as u64) / sampled_count as u64; + let estimated_total = exact_prefix_bytes.saturating_add(estimated_overflow); + info!( + "Timeout/stall at {} files, using sampled estimate: exact_prefix={} overflow_estimate={} sampled={}", + file_count, exact_prefix_bytes, estimated_overflow, sampled_count + ); + progress_monitor.record_timeout_fallback(); + record_capacity_scan_sampling(sampled_count, true); + record_capacity_scan_mode("timeout_fallback"); + return Ok(CapacityScanResult { + used_bytes: estimated_total, + file_count, + sampled_count, + is_estimated: true, + scan_duration: start_time.elapsed(), + had_partial_errors, + }); + } + return Err(e); + } + if should_check_progress { + last_progress_check_files = file_count; + } + + if file_count <= max_files_threshold { + exact_prefix_bytes += metadata.len(); + } else { + let overflow_index = file_count - max_files_threshold; + if overflow_index.is_multiple_of(effective_sample_rate) { + overflow_sampled_bytes += metadata.len(); + sampled_count += 1; + } + + if file_count.is_multiple_of(100_000) { + debug!( + "Processed {} files, exact_prefix_bytes={}, sampled_overflow={} files/{} bytes", + file_count, exact_prefix_bytes, sampled_count, overflow_sampled_bytes + ); + } + } + } + + if file_count > last_progress_check_files { + let exact_count = file_count.min(max_files_threshold); + let avg_size = if exact_count > 0 { + exact_prefix_bytes / exact_count as u64 + } else { + 0 + }; + + if let Err(e) = progress_monitor.update_and_check_timeout(file_count, avg_size) { + if sampled_count > 0 { + let overflow_count = file_count.saturating_sub(max_files_threshold); + let estimated_overflow = overflow_sampled_bytes.saturating_mul(overflow_count as u64) / sampled_count as u64; + let estimated_total = exact_prefix_bytes.saturating_add(estimated_overflow); + info!( + "Timeout/stall at {} files during final check, using sampled estimate: exact_prefix={} overflow_estimate={} sampled={}", + file_count, exact_prefix_bytes, estimated_overflow, sampled_count + ); + progress_monitor.record_timeout_fallback(); + record_capacity_scan_sampling(sampled_count, true); + record_capacity_scan_mode("timeout_fallback"); + return Ok(CapacityScanResult { + used_bytes: estimated_total, + file_count, + sampled_count, + is_estimated: true, + scan_duration: start_time.elapsed(), + had_partial_errors, + }); + } + return Err(e); + } + } + + let (symlink_count, symlink_size) = symlink_tracker.get_stats(); + if symlink_count > 0 { + info!( + "Symlink tracking: {} symlinks processed, total tracked size: {} bytes", + symlink_count, symlink_size + ); + } + + if file_count > max_files_threshold && sampled_count > 0 { + let overflow_count = file_count - max_files_threshold; + let estimated_overflow = overflow_sampled_bytes.saturating_mul(overflow_count as u64) / sampled_count as u64; + let estimated_size = exact_prefix_bytes.saturating_add(estimated_overflow); + info!( + "Large directory detected: {} files, estimated size: {} bytes (exact prefix: {}, sampled overflow {}/{})", + file_count, estimated_size, exact_prefix_bytes, sampled_count, overflow_count + ); + record_capacity_scan_sampling(sampled_count, true); + record_capacity_scan_mode("estimated"); + Ok(CapacityScanResult { + used_bytes: estimated_size, + file_count, + sampled_count, + is_estimated: true, + scan_duration: start_time.elapsed(), + had_partial_errors, + }) + } else if file_count > max_files_threshold { + let overflow_count = file_count - max_files_threshold; + let exact_prefix_count = file_count.min(max_files_threshold) as u64; + let avg_prefix_size = if exact_prefix_count > 0 { + exact_prefix_bytes / exact_prefix_count + } else { + 0 + }; + let estimated_overflow = avg_prefix_size.saturating_mul(overflow_count as u64); + let estimated_size = exact_prefix_bytes.saturating_add(estimated_overflow); + info!( + "Large directory detected: {} files, estimated size: {} bytes (no overflow samples, used prefix average {} bytes/file)", + file_count, estimated_size, avg_prefix_size + ); + record_capacity_scan_sampling(0, true); + record_capacity_scan_mode("estimated"); + Ok(CapacityScanResult { + used_bytes: estimated_size, + file_count, + sampled_count: 0, + is_estimated: true, + scan_duration: start_time.elapsed(), + had_partial_errors, + }) + } else { + record_capacity_scan_sampling(0, false); + debug!( + "Directory size calculation completed: {} files, {} bytes, took {:?}", + file_count, + exact_prefix_bytes, + start_time.elapsed() + ); + record_capacity_scan_mode("exact"); + Ok(CapacityScanResult { + used_bytes: exact_prefix_bytes, + file_count, + sampled_count, + is_estimated: false, + scan_duration: start_time.elapsed(), + had_partial_errors, + }) + } + }) + .await + .map_err(std::io::Error::other)? +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::capacity_manager::{DataSource, HybridStrategyConfig, create_isolated_manager}; + use rustfs_common::capacity_scope::{CapacityScope, CapacityScopeDisk}; + use rustfs_config::ENV_CAPACITY_FOLLOW_SYMLINKS; + use serial_test::serial; + + #[tokio::test] + async fn test_get_dir_size_async_empty_directory() { + use tempfile::TempDir; + + let temp_dir = TempDir::new().unwrap(); + let size = get_dir_size_async(temp_dir.path()).await.unwrap(); + assert_eq!(size.used_bytes, 0); + assert_eq!(size.file_count, 0); + } + + #[tokio::test] + async fn test_get_dir_size_async_single_file() { + use std::fs::File; + use std::io::Write; + use tempfile::TempDir; + + let temp_dir = TempDir::new().unwrap(); + let file_path = temp_dir.path().join("test.txt"); + let mut file = File::create(&file_path).unwrap(); + file.write_all(b"Hello, World!").unwrap(); + + let size = get_dir_size_async(temp_dir.path()).await.unwrap(); + assert_eq!(size.used_bytes, 13); + assert_eq!(size.file_count, 1); + } + + #[tokio::test] + async fn test_get_dir_size_async_multiple_files() { + use std::fs::File; + use std::io::Write; + use tempfile::TempDir; + + let temp_dir = TempDir::new().unwrap(); + + for i in 0..10 { + let file_path = temp_dir.path().join(format!("file_{}.txt", i)); + let mut file = File::create(&file_path).unwrap(); + file.write_all(b"test").unwrap(); + } + + let size = get_dir_size_async(temp_dir.path()).await.unwrap(); + assert_eq!(size.used_bytes, 40); + assert_eq!(size.file_count, 10); + } + + #[tokio::test] + async fn test_get_dir_size_async_nested_directories() { + use std::fs::File; + use std::io::Write; + use tempfile::TempDir; + + let temp_dir = TempDir::new().unwrap(); + let subdir = temp_dir.path().join("subdir"); + std::fs::create_dir(&subdir).unwrap(); + + let file1 = temp_dir.path().join("file1.txt"); + let mut f1 = File::create(&file1).unwrap(); + f1.write_all(b"content1").unwrap(); + + let file2 = subdir.join("file2.txt"); + let mut f2 = File::create(&file2).unwrap(); + f2.write_all(b"content2").unwrap(); + + let size = get_dir_size_async(temp_dir.path()).await.unwrap(); + assert_eq!(size.used_bytes, 16); + assert_eq!(size.file_count, 2); + } + + #[tokio::test] + #[serial] + async fn test_get_dir_size_async_nonexistent_directory() { + let result = get_dir_size_async(Path::new("/nonexistent/path")).await; + assert!(result.is_err()); + } + + #[tokio::test] + async fn test_calculate_data_dir_used_capacity_returns_partial_success() { + use std::fs::File; + use std::io::Write; + use tempfile::TempDir; + + let temp_dir = TempDir::new().unwrap(); + let file_path = temp_dir.path().join("test.txt"); + let mut file = File::create(&file_path).unwrap(); + file.write_all(b"Hello, World!").unwrap(); + + let disks = vec![ + CapacityDiskRef { + endpoint: "disk-1".to_string(), + drive_path: temp_dir.path().to_string_lossy().into_owned(), + }, + CapacityDiskRef { + endpoint: "disk-2".to_string(), + drive_path: "/nonexistent/path".to_string(), + }, + ]; + + let result = calculate_data_dir_used_capacity(&disks).await.unwrap(); + assert_eq!(result.used_bytes, 13); + assert_eq!(result.file_count, 1); + assert!(result.had_partial_errors); + } + + #[tokio::test] + async fn test_select_capacity_refresh_disks_returns_full_when_disk_cache_incomplete() { + let manager = create_isolated_manager(HybridStrategyConfig::default()); + manager + .mark_dirty_scope(&CapacityScope { + disks: vec![CapacityScopeDisk { + endpoint: "disk-1".to_string(), + drive_path: "/tmp/disk-1".to_string(), + }], + }) + .await; + + let disks = vec![ + CapacityDiskRef { + endpoint: "disk-1".to_string(), + drive_path: "/tmp/disk-1".to_string(), + }, + CapacityDiskRef { + endpoint: "disk-2".to_string(), + drive_path: "/tmp/disk-2".to_string(), + }, + ]; + + let (selected, dirty_subset) = select_capacity_refresh_disks(manager.as_ref(), &disks).await; + assert!(!dirty_subset); + assert_eq!(selected.len(), 2); + } + + #[tokio::test] + async fn test_select_capacity_refresh_disks_returns_dirty_subset_when_cache_complete() { + let manager = create_isolated_manager(HybridStrategyConfig::default()); + manager + .update_capacity( + CapacityUpdate { + total_used: 300, + file_count: 3, + is_estimated: false, + per_disk: vec![ + DiskCapacityUpdate { + disk: CapacityScopeDisk { + endpoint: "disk-1".to_string(), + drive_path: "/tmp/disk-1".to_string(), + }, + used_bytes: 100, + file_count: 1, + is_estimated: false, + }, + DiskCapacityUpdate { + disk: CapacityScopeDisk { + endpoint: "disk-2".to_string(), + drive_path: "/tmp/disk-2".to_string(), + }, + used_bytes: 200, + file_count: 2, + is_estimated: false, + }, + ], + expected_disk_count: Some(2), + replaces_disk_cache: true, + clear_dirty_disks: Vec::new(), + }, + DataSource::RealTime, + ) + .await; + manager + .mark_dirty_scope(&CapacityScope { + disks: vec![CapacityScopeDisk { + endpoint: "disk-2".to_string(), + drive_path: "/tmp/disk-2".to_string(), + }], + }) + .await; + + let disks = vec![ + CapacityDiskRef { + endpoint: "disk-1".to_string(), + drive_path: "/tmp/disk-1".to_string(), + }, + CapacityDiskRef { + endpoint: "disk-2".to_string(), + drive_path: "/tmp/disk-2".to_string(), + }, + ]; + + let (selected, dirty_subset) = select_capacity_refresh_disks(manager.as_ref(), &disks).await; + assert!(dirty_subset); + assert_eq!(selected.len(), 1); + assert_eq!(selected[0].endpoint, "disk-2"); + assert_eq!(selected[0].drive_path, "/tmp/disk-2"); + } + + #[cfg(unix)] + #[tokio::test] + #[serial] + async fn test_get_dir_size_async_ignores_symlink_targets_when_follow_disabled() { + use std::fs::File; + use std::io::Write; + use std::os::unix::fs::symlink; + use tempfile::TempDir; + + let scan_dir = TempDir::new().unwrap(); + let target_dir = TempDir::new().unwrap(); + let target_path = target_dir.path().join("external.txt"); + let mut file = File::create(&target_path).unwrap(); + file.write_all(b"external-bytes").unwrap(); + symlink(&target_path, scan_dir.path().join("external-link")).unwrap(); + + let size = temp_env::async_with_vars([(ENV_CAPACITY_FOLLOW_SYMLINKS, Some("false"))], async { + get_dir_size_async(scan_dir.path()).await + }) + .await + .unwrap(); + + assert_eq!(size.used_bytes, 0); + assert_eq!(size.file_count, 0); + } + + #[cfg(unix)] + #[tokio::test] + #[serial] + async fn test_get_dir_size_async_counts_symlink_targets_when_follow_enabled() { + use std::fs::File; + use std::io::Write; + use std::os::unix::fs::symlink; + use tempfile::TempDir; + + let scan_dir = TempDir::new().unwrap(); + let target_dir = TempDir::new().unwrap(); + let target_path = target_dir.path().join("external.txt"); + let mut file = File::create(&target_path).unwrap(); + file.write_all(b"external-bytes").unwrap(); + symlink(&target_path, scan_dir.path().join("external-link")).unwrap(); + + let size = temp_env::async_with_vars([(ENV_CAPACITY_FOLLOW_SYMLINKS, Some("true"))], async { + get_dir_size_async(scan_dir.path()).await + }) + .await + .unwrap(); + + assert_eq!(size.used_bytes, "external-bytes".len() as u64); + assert_eq!(size.file_count, 1); + } +} diff --git a/crates/object-capacity/src/types.rs b/crates/object-capacity/src/types.rs new file mode 100644 index 000000000..6f867ac2b --- /dev/null +++ b/crates/object-capacity/src/types.rs @@ -0,0 +1,62 @@ +// Copyright 2024 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 std::time::Duration; + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct CapacityDiskRef { + pub endpoint: String, + pub drive_path: String, +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub(crate) struct CapacityScanResult { + pub used_bytes: u64, + pub file_count: usize, + pub sampled_count: usize, + pub is_estimated: bool, + pub scan_duration: Duration, + pub had_partial_errors: bool, +} + +impl CapacityScanResult { + pub(crate) fn with_partial_errors(mut self) -> Self { + self.had_partial_errors = true; + self + } +} + +/// Public summary type for external tooling such as benches. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct CapacityScanSummary { + pub used_bytes: u64, + pub file_count: usize, + pub sampled_count: usize, + pub is_estimated: bool, + pub had_partial_errors: bool, + pub scan_duration: Duration, +} + +impl From for CapacityScanSummary { + fn from(scan: CapacityScanResult) -> Self { + Self { + used_bytes: scan.used_bytes, + file_count: scan.file_count, + sampled_count: scan.sampled_count, + is_estimated: scan.is_estimated, + had_partial_errors: scan.had_partial_errors, + scan_duration: scan.scan_duration, + } + } +} diff --git a/rustfs/Cargo.toml b/rustfs/Cargo.toml index 3b314d08e..2242d4598 100644 --- a/rustfs/Cargo.toml +++ b/rustfs/Cargo.toml @@ -91,6 +91,7 @@ rustfs-zip = { workspace = true } rustfs-io-core = { workspace = true } rustfs-io-metrics = { workspace = true } rustfs-object-io = { workspace = true } +rustfs-object-capacity = { workspace = true } rustfs-concurrency = { workspace = true } rustfs-scanner = { workspace = true } tempfile = { workspace = true } @@ -119,7 +120,6 @@ tower-http = { workspace = true, features = ["trace", "compression-full", "cors" # Serialization and Data Formats bytes = { workspace = true } flatbuffers.workspace = true -walkdir = { workspace = true } rmp-serde.workspace = true rustfs-signer.workspace = true serde.workspace = true @@ -197,7 +197,7 @@ tempfile = { workspace = true } aws-config = { workspace = true } anyhow = { workspace = true } tokio = { workspace = true, features = ["test-util"] } -temp-env = { workspace = true } +temp-env = { workspace = true, features = ["async_closure"] } [build-dependencies] http.workspace = true diff --git a/rustfs/src/app/admin_usecase.rs b/rustfs/src/app/admin_usecase.rs index 8b4d5117f..616b21d2c 100644 --- a/rustfs/src/app/admin_usecase.rs +++ b/rustfs/src/app/admin_usecase.rs @@ -15,10 +15,7 @@ //! Admin application use-case contracts. use crate::app::context::{AppContext, get_global_app_context}; -use crate::capacity::capacity_manager::{ - CapacityUpdate, DataSource, get_capacity_manager, get_enable_dynamic_timeout, get_follow_symlinks, get_max_files_threshold, - get_max_symlink_depth, get_max_timeout, get_min_timeout, get_sample_rate, get_stall_timeout, get_stat_timeout, -}; +use crate::capacity::resolve_admin_used_capacity; use crate::error::ApiError; use rustfs_common::data_usage::DataUsageInfo; use rustfs_ecstore::admin_server_info::get_server_info; @@ -27,18 +24,10 @@ use rustfs_ecstore::endpoints::EndpointServerPools; use rustfs_ecstore::new_object_layer_fn; use rustfs_ecstore::pools::{PoolStatus, get_total_usable_capacity, get_total_usable_capacity_free}; use rustfs_ecstore::store_api::StorageAPI; -use rustfs_io_metrics::{ - record_capacity_dynamic_timeout, record_capacity_scan_sampling, record_capacity_stall_detected, record_capacity_symlink, - record_capacity_timeout_fallback, -}; use rustfs_madmin::{InfoMessage, StorageInfo}; use s3s::S3ErrorCode; -use std::collections::HashSet; -use std::path::{Path, PathBuf}; use std::sync::Arc; -use std::time::{Duration, Instant}; use tracing::{debug, error, info, warn}; -use walkdir::WalkDir; pub type AdminUsecaseResult = Result; @@ -47,31 +36,6 @@ pub struct QueryServerInfoRequest { pub include_pools: bool, } -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub(crate) struct CapacityScanResult { - pub used_bytes: u64, - pub file_count: usize, - pub sampled_count: usize, - pub is_estimated: bool, - pub scan_duration: Duration, - pub had_partial_errors: bool, -} - -impl CapacityScanResult { - fn with_partial_errors(mut self) -> Self { - self.had_partial_errors = true; - self - } - - pub(crate) fn to_capacity_update(self) -> CapacityUpdate { - if self.is_estimated { - CapacityUpdate::estimated(self.used_bytes, self.file_count) - } else { - CapacityUpdate::exact(self.used_bytes, self.file_count) - } - } -} - pub struct QueryServerInfoResponse { pub info: InfoMessage, } @@ -94,475 +58,6 @@ pub struct QueryPoolStatusRequest { pub by_id: bool, } -/// Calculate actual used capacity of all data directories -pub(crate) async fn calculate_data_dir_used_capacity( - disks: &[rustfs_madmin::Disk], -) -> Result> { - let start = Instant::now(); - let mut total_used = 0u64; - let mut total_files = 0usize; - let mut total_sampled = 0usize; - let mut has_failure = false; - let mut has_success = false; - let mut is_estimated = false; - - for disk in disks { - let path = Path::new(&disk.drive_path); - - if !path.exists() { - warn!("Data directory does not exist: {}", disk.drive_path); - has_failure = true; - continue; - } - - match get_dir_size_async(path).await { - Ok(scan) => { - debug!( - "Data directory {} size: {} bytes, files={}, sampled={}, estimated={}", - disk.drive_path, scan.used_bytes, scan.file_count, scan.sampled_count, scan.is_estimated - ); - total_used += scan.used_bytes; - total_files += scan.file_count; - total_sampled += scan.sampled_count; - is_estimated |= scan.is_estimated; - has_failure |= scan.had_partial_errors; - has_success = true; - } - Err(e) => { - warn!("Failed to get size for directory {}: {:?}", disk.drive_path, e); - has_failure = true; - } - } - } - - if !has_success { - return Err("All directories failed to calculate size".into()); - } - - if has_failure { - warn!("Some directories failed to calculate size, result may be incomplete"); - } - - let mut result = CapacityScanResult { - used_bytes: total_used, - file_count: total_files, - sampled_count: total_sampled, - is_estimated, - scan_duration: start.elapsed(), - had_partial_errors: false, - }; - - if has_failure { - result = result.with_partial_errors(); - } - - Ok(result) -} - -// ============================================================================ -// Symlink Tracker for Circular Reference Detection -// ============================================================================ - -/// Tracker for symlink resolution with circular reference detection -struct SymlinkTracker { - /// Set of visited symlink paths to detect circular references - visited: HashSet, - /// Count of symlinks encountered - symlink_count: usize, - /// Total size of symlink targets - symlink_size: u64, - /// Maximum symlink depth to follow - max_depth: u8, -} - -impl SymlinkTracker { - /// Create a new symlink tracker - fn new(max_depth: u8) -> Self { - Self { - visited: HashSet::new(), - symlink_count: 0, - symlink_size: 0, - max_depth, - } - } - - /// Check if we should follow a symlink at the given depth - fn should_follow(&self, path: &Path, depth: u8) -> bool { - if depth >= self.max_depth { - debug!("Symlink depth limit reached: {} >= {}, not following {:?}", depth, self.max_depth, path); - return false; - } - - if self.visited.contains(path) { - warn!("Circular symlink reference detected: {:?}, skipping", path); - return false; - } - - true - } - - /// Record a visited symlink path and update metrics - fn record_symlink(&mut self, path: PathBuf, size: u64) { - self.visited.insert(path); - self.symlink_count += 1; - self.symlink_size += size; - record_capacity_symlink(size); - } - - /// Get symlink statistics - fn get_stats(&self) -> (usize, u64) { - (self.symlink_count, self.symlink_size) - } -} - -// ============================================================================ -// Progress Monitor for Timeout and Stall Detection -// ============================================================================ - -/// Monitor for directory traversal progress with timeout and stall detection -struct ProgressMonitor { - /// Start time of the operation - start_time: Instant, - /// Last check time for stall detection - last_check: Instant, - /// Number of files processed at last checkpoint - last_checkpoint_files: usize, - /// Base timeout for this operation - timeout: Duration, - /// Minimum allowed timeout - min_timeout: Duration, - /// Maximum allowed timeout - max_timeout: Duration, - /// Stall detection timeout - stall_timeout: Duration, - /// Enable dynamic timeout calculation - enable_dynamic_timeout: bool, - /// Track if dynamic timeout was used - used_dynamic_timeout: bool, -} - -impl ProgressMonitor { - /// Create a new progress monitor - fn new( - base_timeout: Duration, - min_timeout: Duration, - max_timeout: Duration, - stall_timeout: Duration, - enable_dynamic: bool, - ) -> Self { - Self { - start_time: Instant::now(), - last_check: Instant::now(), - last_checkpoint_files: 0, - timeout: base_timeout, - min_timeout, - max_timeout, - stall_timeout, - enable_dynamic_timeout: enable_dynamic, - used_dynamic_timeout: false, - } - } - - /// Calculate dynamic timeout based on directory characteristics - fn calculate_dynamic_timeout(&mut self, file_count: usize, avg_file_size: u64) -> Duration { - if !self.enable_dynamic_timeout { - return self.timeout; - } - - // Mark that we're using dynamic timeout - self.used_dynamic_timeout = true; - - // Calculate multipliers based on directory characteristics - let file_factor = (file_count as f64).sqrt() * 0.01; // File count influence - let size_factor = if avg_file_size > 0 { - (avg_file_size as f64).log(10.0) * 0.05 // File size influence - } else { - 0.0 - }; - - let multiplier = 1.0 + file_factor + size_factor; - let adjusted_timeout = self.timeout.mul_f64(multiplier.min(5.0)); // Max 5x multiplier - - // Clamp to min/max bounds - let clamped_timeout = adjusted_timeout.max(self.min_timeout).min(self.max_timeout); - - debug!( - "Dynamic timeout calculation: files={}, avg_size={}, multiplier={:.2}, base_timeout={:?}, adjusted_timeout={:?}, clamped_timeout={:?}", - file_count, avg_file_size, multiplier, self.timeout, adjusted_timeout, clamped_timeout - ); - - clamped_timeout - } - - /// Update and check for timeout or stall - fn update_and_check_timeout(&mut self, files_processed: usize, avg_file_size: u64) -> Result<(), std::io::Error> { - let elapsed = self.start_time.elapsed(); - - // Calculate dynamic timeout based on current state - let dynamic_timeout = if self.enable_dynamic_timeout { - self.calculate_dynamic_timeout(files_processed, avg_file_size) - } else { - self.timeout - }; - - // Check for hard timeout - if elapsed >= dynamic_timeout { - warn!( - "Directory size calculation timeout after {} files, elapsed: {:?}, timeout: {:?}", - files_processed, elapsed, dynamic_timeout - ); - - if self.enable_dynamic_timeout { - record_capacity_dynamic_timeout(dynamic_timeout); - } - - return Err(std::io::Error::new( - std::io::ErrorKind::TimedOut, - format!("Timeout after {} files", files_processed), - )); - } - - // Check for stall (no progress) - let now = Instant::now(); - if now.duration_since(self.last_check) >= self.stall_timeout { - let files_per_checkpoint = files_processed.saturating_sub(self.last_checkpoint_files); - - if files_per_checkpoint == 0 && files_processed > 0 { - // No progress for stall_timeout duration - warn!( - "No progress detected for {:?}, possible stall at {} files", - self.stall_timeout, files_processed - ); - - record_capacity_stall_detected(); - - return Err(std::io::Error::new( - std::io::ErrorKind::TimedOut, - format!("Stall detected at {} files", files_processed), - )); - } - - self.last_check = now; - self.last_checkpoint_files = files_processed; - } - - Ok(()) - } - - /// Record timeout fallback to sampling - fn record_timeout_fallback(&self) { - record_capacity_timeout_fallback(); - } -} - -/// Asynchronously get directory size with enhanced symlink handling and dynamic timeout -async fn get_dir_size_async(path: &Path) -> Result { - let path = path.to_path_buf(); - - let max_files_threshold = get_max_files_threshold(); - let base_timeout = get_stat_timeout(); - let min_timeout = get_min_timeout(); - let max_timeout = get_max_timeout(); - let stall_timeout = get_stall_timeout(); - let sample_rate = get_sample_rate(); - let enable_dynamic_timeout = get_enable_dynamic_timeout(); - let follow_symlinks = get_follow_symlinks(); - let max_symlink_depth = get_max_symlink_depth(); - - let effective_sample_rate = if sample_rate == 0 { - warn!("Invalid sampling configuration: sample_rate=0. Clamping to 1 to avoid panic."); - 1 - } else { - sample_rate - }; - - if !path.exists() { - return Err(std::io::Error::new( - std::io::ErrorKind::NotFound, - format!("Directory not found: {:?}", path), - )); - } - - tokio::task::spawn_blocking(move || { - let start_time = Instant::now(); - let mut exact_prefix_bytes = 0u64; - let mut overflow_sampled_bytes = 0u64; - let mut file_count = 0usize; - let mut sampled_count = 0usize; - let mut had_partial_errors = false; - - let mut symlink_tracker = if follow_symlinks { - Some(SymlinkTracker::new(max_symlink_depth)) - } else { - None - }; - - let mut progress_monitor = - ProgressMonitor::new(base_timeout, min_timeout, max_timeout, stall_timeout, enable_dynamic_timeout); - - let mut walker_builder = WalkDir::new(&path); - if !follow_symlinks { - walker_builder = walker_builder.follow_links(false); - } - let walker = walker_builder.into_iter(); - - for entry_result in walker { - let entry = match entry_result { - Ok(entry) => entry, - Err(err) => { - warn!("Failed to traverse directory entry under {:?}: {}", path, err); - had_partial_errors = true; - continue; - } - }; - - let metadata = match entry.metadata() { - Ok(meta) => meta, - Err(err) => { - warn!("Failed to get metadata for {:?}: {}", entry.path(), err); - had_partial_errors = true; - continue; - } - }; - - if metadata.is_symlink() { - if let Some(ref mut tracker) = symlink_tracker - && let Ok(target) = std::fs::read_link(entry.path()) - && tracker.should_follow(&target, 0) - { - tracker.record_symlink(target, metadata.len()); - } - continue; - } - - if !metadata.is_file() { - continue; - } - - file_count += 1; - let exact_count = file_count.min(max_files_threshold); - let avg_size = if exact_count > 0 { - exact_prefix_bytes / exact_count as u64 - } else { - 0 - }; - - if let Err(e) = progress_monitor.update_and_check_timeout(file_count, avg_size) { - if sampled_count > 0 { - let overflow_count = file_count.saturating_sub(max_files_threshold); - let estimated_overflow = overflow_sampled_bytes.saturating_mul(overflow_count as u64) / sampled_count as u64; - let estimated_total = exact_prefix_bytes.saturating_add(estimated_overflow); - info!( - "Timeout/stall at {} files, using sampled estimate: exact_prefix={} overflow_estimate={} sampled={}", - file_count, exact_prefix_bytes, estimated_overflow, sampled_count - ); - progress_monitor.record_timeout_fallback(); - record_capacity_scan_sampling(sampled_count, true); - return Ok(CapacityScanResult { - used_bytes: estimated_total, - file_count, - sampled_count, - is_estimated: true, - scan_duration: start_time.elapsed(), - had_partial_errors, - }); - } - return Err(e); - } - - if file_count <= max_files_threshold { - exact_prefix_bytes += metadata.len(); - } else { - let overflow_index = file_count - max_files_threshold; - if overflow_index.is_multiple_of(effective_sample_rate) { - overflow_sampled_bytes += metadata.len(); - sampled_count += 1; - } - - if file_count.is_multiple_of(100_000) { - debug!( - "Processed {} files, exact_prefix_bytes={}, sampled_overflow={} files/{} bytes", - file_count, exact_prefix_bytes, sampled_count, overflow_sampled_bytes - ); - } - } - } - - if let Some(tracker) = symlink_tracker { - let (count, size) = tracker.get_stats(); - if count > 0 { - info!("Symlink tracking: {} symlinks processed, total target size: {} bytes", count, size); - } - } - - if file_count > max_files_threshold && sampled_count > 0 { - let overflow_count = file_count - max_files_threshold; - let estimated_overflow = overflow_sampled_bytes.saturating_mul(overflow_count as u64) / sampled_count as u64; - let estimated_size = exact_prefix_bytes.saturating_add(estimated_overflow); - info!( - "Large directory detected: {} files, estimated size: {} bytes (exact prefix: {}, sampled overflow {}/{})", - file_count, estimated_size, exact_prefix_bytes, sampled_count, overflow_count - ); - record_capacity_scan_sampling(sampled_count, true); - Ok(CapacityScanResult { - used_bytes: estimated_size, - file_count, - sampled_count, - is_estimated: true, - scan_duration: start_time.elapsed(), - had_partial_errors, - }) - } else if file_count > max_files_threshold { - // sampled_count == 0: too few overflow files to reach the sample rate threshold. - // Fall back to estimating the overflow using the average file size from the exact - // prefix so that overflow files are not silently dropped from the total. - let overflow_count = file_count - max_files_threshold; - // Use the actual number of files counted in the exact prefix, not the threshold - // value, to avoid a divide-by-zero or incorrect average when fewer files were - // processed than max_files_threshold. - let exact_prefix_count = file_count.min(max_files_threshold) as u64; - let avg_prefix_size = if exact_prefix_count > 0 { - exact_prefix_bytes / exact_prefix_count - } else { - 0 - }; - let estimated_overflow = avg_prefix_size.saturating_mul(overflow_count as u64); - let estimated_size = exact_prefix_bytes.saturating_add(estimated_overflow); - info!( - "Large directory detected: {} files, estimated size: {} bytes (no overflow samples, used prefix average {} bytes/file)", - file_count, estimated_size, avg_prefix_size - ); - record_capacity_scan_sampling(0, true); - Ok(CapacityScanResult { - used_bytes: estimated_size, - file_count, - sampled_count: 0, - is_estimated: true, - scan_duration: start_time.elapsed(), - had_partial_errors, - }) - } else { - record_capacity_scan_sampling(0, false); - debug!( - "Directory size calculation completed: {} files, {} bytes, took {:?}", - file_count, - exact_prefix_bytes, - start_time.elapsed() - ); - Ok(CapacityScanResult { - used_bytes: exact_prefix_bytes, - file_count, - sampled_count, - is_estimated: false, - scan_duration: start_time.elapsed(), - had_partial_errors, - }) - } - }) - .await - .map_err(std::io::Error::other)? -} - #[derive(Clone, Default)] pub struct DefaultAdminUsecase { context: Option>, @@ -688,115 +183,8 @@ impl DefaultAdminUsecase { info.total_free_capacity = free_u64; } - // Use hybrid strategy for capacity calculation - let capacity_manager = get_capacity_manager(); - - // Check if we have a valid cache - if let Some(cached) = capacity_manager.get_capacity().await { - let cache_age = cached.last_update.elapsed(); - let fast_update_threshold = capacity_manager.get_config().fast_update_threshold; - - // If cache is fresh (< fast_update_threshold), use it directly - if cache_age < fast_update_threshold { - info.total_used_capacity = cached.total_used; - debug!( - "Using cached capacity: {} bytes (age: {:?}, source: {:?}, files={}, estimated={})", - cached.total_used, cache_age, cached.source, cached.file_count, cached.is_estimated - ); - } else { - // Cache is stale, check if we need fast update - let needs_update = capacity_manager.needs_fast_update().await; - let should_block = capacity_manager.should_block_on_refresh(cache_age); - - if needs_update && should_block { - let start = Instant::now(); - match capacity_manager - .refresh_or_join(DataSource::WriteTriggered, || async { - calculate_data_dir_used_capacity(&storage_info.disks) - .await - .map(|scan| scan.to_capacity_update()) - .map_err(|e| e.to_string()) - }) - .await - { - Ok(update) => { - info.total_used_capacity = update.total_used; - - let elapsed = start.elapsed(); - debug!( - "Foreground capacity refresh completed in {:?} (files={}, estimated={})", - elapsed, update.file_count, update.is_estimated - ); - } - Err(e) => { - warn!("Foreground capacity refresh failed: {}, using cached value", e); - info.total_used_capacity = cached.total_used; - } - } - } else { - info.total_used_capacity = cached.total_used; - debug!( - "Using stale cached capacity: {} bytes (age: {:?}, source: {:?}, files={}, estimated={}, needs_update={}, blocking={})", - cached.total_used, - cache_age, - cached.source, - cached.file_count, - cached.is_estimated, - needs_update, - should_block - ); - - let disks = storage_info.disks.clone(); - let manager = capacity_manager.clone(); - if manager - .clone() - .spawn_refresh_if_needed(DataSource::Scheduled, move || async move { - calculate_data_dir_used_capacity(&disks) - .await - .map(|scan| scan.to_capacity_update()) - .map_err(|e| e.to_string()) - }) - .await - { - debug!("Background capacity update started"); - } else { - debug!("Background update already in progress, skipping spawn"); - } - } - } - } else { - // No cache, perform initial calculation - let start = Instant::now(); - match capacity_manager - .refresh_or_join(DataSource::RealTime, || async { - calculate_data_dir_used_capacity(&storage_info.disks) - .await - .map(|scan| scan.to_capacity_update()) - .map_err(|e| e.to_string()) - }) - .await - { - Ok(update) => { - info.total_used_capacity = update.total_used; - - let elapsed = start.elapsed(); - info!( - "Initial capacity calculation completed: {} bytes in {:?} (files={}, estimated={})", - update.total_used, elapsed, update.file_count, update.is_estimated - ); - } - Err(e) => { - warn!( - "Failed to calculate data directory used capacity: {}, falling back to disk used capacity", - e - ); - info.total_used_capacity = info.total_capacity.saturating_sub(info.total_free_capacity); - capacity_manager - .update_capacity(CapacityUpdate::fallback(info.total_used_capacity), DataSource::Fallback) - .await; - } - } - } + info.total_used_capacity = + resolve_admin_used_capacity(&storage_info.disks, info.total_capacity.saturating_sub(info.total_free_capacity)).await; debug!( "Capacity statistics: total={:.2} TiB, free={:.2} TiB, used={:.2} TiB", info.total_capacity as f64 / (1024.0_f64.powi(4)), @@ -885,7 +273,6 @@ impl DefaultAdminUsecase { #[cfg(test)] mod tests { use super::*; - use serial_test::serial; #[tokio::test] async fn execute_query_storage_info_returns_internal_error_when_store_uninitialized() { @@ -911,83 +298,4 @@ mod tests { let _ = readiness.storage_ready; let _ = readiness.iam_ready; } - - // Tests for directory size calculation functions - #[tokio::test] - async fn test_get_dir_size_async_empty_directory() { - use tempfile::TempDir; - - let temp_dir = TempDir::new().unwrap(); - let size = get_dir_size_async(temp_dir.path()).await.unwrap(); - assert_eq!(size.used_bytes, 0); - assert_eq!(size.file_count, 0); - } - - #[tokio::test] - async fn test_get_dir_size_async_single_file() { - use std::fs::File; - use std::io::Write; - use tempfile::TempDir; - - let temp_dir = TempDir::new().unwrap(); - let file_path = temp_dir.path().join("test.txt"); - let mut file = File::create(&file_path).unwrap(); - file.write_all(b"Hello, World!").unwrap(); - - let size = get_dir_size_async(temp_dir.path()).await.unwrap(); - assert_eq!(size.used_bytes, 13); - assert_eq!(size.file_count, 1); - } - - #[tokio::test] - async fn test_get_dir_size_async_multiple_files() { - use std::fs::File; - use std::io::Write; - use tempfile::TempDir; - - let temp_dir = TempDir::new().unwrap(); - - // Create multiple files - for i in 0..10 { - let file_path = temp_dir.path().join(format!("file_{}.txt", i)); - let mut file = File::create(&file_path).unwrap(); - file.write_all(b"test").unwrap(); - } - - let size = get_dir_size_async(temp_dir.path()).await.unwrap(); - assert_eq!(size.used_bytes, 40); // 10 files * 4 bytes - assert_eq!(size.file_count, 10); - } - - #[tokio::test] - async fn test_get_dir_size_async_nested_directories() { - use std::fs::File; - use std::io::Write; - use tempfile::TempDir; - - let temp_dir = TempDir::new().unwrap(); - - // Create nested directories and files - let subdir = temp_dir.path().join("subdir"); - std::fs::create_dir(&subdir).unwrap(); - - let file1 = temp_dir.path().join("file1.txt"); - let mut f1 = File::create(&file1).unwrap(); - f1.write_all(b"content1").unwrap(); - - let file2 = subdir.join("file2.txt"); - let mut f2 = File::create(&file2).unwrap(); - f2.write_all(b"content2").unwrap(); - - let size = get_dir_size_async(temp_dir.path()).await.unwrap(); - assert_eq!(size.used_bytes, 16); // "content1" (8) + "content2" (8) - assert_eq!(size.file_count, 2); - } - - #[tokio::test] - #[serial] - async fn test_get_dir_size_async_nonexistent_directory() { - let result = get_dir_size_async(Path::new("/nonexistent/path")).await; - assert!(result.is_err()); - } } diff --git a/rustfs/src/app/bucket_usecase.rs b/rustfs/src/app/bucket_usecase.rs index 073753add..d08eae838 100644 --- a/rustfs/src/app/bucket_usecase.rs +++ b/rustfs/src/app/bucket_usecase.rs @@ -1494,10 +1494,7 @@ impl DefaultBucketUsecase { && let Some(store) = new_object_layer_fn() { let bucket_name = bucket.clone(); - let request_context = req - .extensions - .get::() - .cloned(); + let request_context = req.extensions.get::().cloned(); spawn_background_with_context(request_context, async move { if let Err(err) = enqueue_transition_for_existing_objects(store, &bucket_name).await { warn!(bucket = %bucket_name, error = ?err, "failed to enqueue transition for existing objects"); @@ -1883,7 +1880,7 @@ impl DefaultBucketUsecase { let store = get_validated_store(&bucket).await?; - let incl_deleted = rustfs_utils::http::get_header(&req.headers, rustfs_utils::http::SUFFIX_INCLUDE_DELETED) + let incl_deleted = get_header(&req.headers, rustfs_utils::http::SUFFIX_INCLUDE_DELETED) .map(|v| v.as_ref() == "true") .unwrap_or_default(); @@ -1963,7 +1960,7 @@ impl DefaultBucketUsecase { .transpose()?; let store = get_validated_store(&bucket).await?; - let incl_deleted = rustfs_utils::http::get_header(&req.headers, rustfs_utils::http::SUFFIX_INCLUDE_DELETED) + let incl_deleted = get_header(&req.headers, rustfs_utils::http::SUFFIX_INCLUDE_DELETED) .map(|value| value.as_ref() == "true") .unwrap_or_default(); diff --git a/rustfs/src/app/capacity_dirty_scope_test.rs b/rustfs/src/app/capacity_dirty_scope_test.rs new file mode 100644 index 000000000..ad751340e --- /dev/null +++ b/rustfs/src/app/capacity_dirty_scope_test.rs @@ -0,0 +1,224 @@ +// Copyright 2024 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 rustfs_common::heal_channel::{HealOpts, HealScanMode}; +use rustfs_ecstore::{ + bucket::metadata_sys, + disk::endpoint::Endpoint, + endpoints::{EndpointServerPools, Endpoints, PoolEndpoints}, + store::ECStore, + store_api::{ + BucketOperations, BucketOptions, ChunkNativePutData, HealOperations, MakeBucketOptions, ObjectIO, ObjectOptions, + }, +}; +use rustfs_object_capacity::capacity_manager::{HybridStrategyConfig, create_isolated_manager}; +use serial_test::serial; +use std::{ + collections::HashSet, + fs as stdfs, + path::Path, + path::PathBuf, + sync::{Arc, Once, OnceLock}, +}; +use tempfile::TempDir; +use tokio::fs; +use tokio_util::sync::CancellationToken; +use uuid::Uuid; + +static CAPACITY_DIRTY_SCOPE_ENV: OnceLock<(Vec, Arc, TempDir)> = OnceLock::new(); +static CAPACITY_DIRTY_SCOPE_INIT: Once = Once::new(); + +fn init_capacity_dirty_scope_tracing() { + CAPACITY_DIRTY_SCOPE_INIT.call_once(|| {}); +} + +async fn setup_capacity_dirty_scope_env() -> (Vec, Arc) { + init_capacity_dirty_scope_tracing(); + + if let Some((paths, store, _)) = CAPACITY_DIRTY_SCOPE_ENV.get() { + return (paths.clone(), store.clone()); + } + + let temp_dir = TempDir::new().expect("create temp dir for capacity dirty scope test"); + let temp_path = temp_dir.path().to_path_buf(); + + let disk_paths = vec![ + temp_path.join("disk1"), + temp_path.join("disk2"), + temp_path.join("disk3"), + temp_path.join("disk4"), + ]; + for disk_path in &disk_paths { + fs::create_dir_all(disk_path).await.unwrap(); + } + + let mut endpoints = Vec::new(); + for (i, disk_path) in disk_paths.iter().enumerate() { + let mut endpoint = Endpoint::try_from(disk_path.to_str().unwrap()).unwrap(); + endpoint.set_pool_index(0); + endpoint.set_set_index(0); + endpoint.set_disk_index(i); + endpoints.push(endpoint); + } + + let pool_endpoints = PoolEndpoints { + legacy: false, + set_count: 1, + drives_per_set: 4, + endpoints: Endpoints::from(endpoints), + cmd_line: "capacity-dirty-scope-test".to_string(), + platform: format!("OS: {} | Arch: {}", std::env::consts::OS, std::env::consts::ARCH), + }; + + let endpoint_pools = EndpointServerPools(vec![pool_endpoints]); + rustfs_ecstore::store::init_local_disks(endpoint_pools.clone()).await.unwrap(); + + let server_addr: std::net::SocketAddr = "127.0.0.1:0".parse().unwrap(); + let ecstore = ECStore::new(server_addr, endpoint_pools, CancellationToken::new()) + .await + .unwrap(); + + let buckets_list = ecstore + .list_bucket(&BucketOptions { + no_metadata: true, + ..Default::default() + }) + .await + .unwrap(); + let buckets = buckets_list.into_iter().map(|v| v.name).collect(); + metadata_sys::init_bucket_metadata_sys(ecstore.clone(), buckets).await; + + let _ = CAPACITY_DIRTY_SCOPE_ENV.set((disk_paths.clone(), ecstore.clone(), temp_dir)); + (disk_paths, ecstore) +} + +fn find_part_file(root: &Path, part_name: &str) -> Option { + let entries = stdfs::read_dir(root).ok()?; + for entry in entries.flatten() { + let path = entry.path(); + if path.is_dir() { + if let Some(found) = find_part_file(&path, part_name) { + return Some(found); + } + continue; + } + + if path.file_name().and_then(|name| name.to_str()) == Some(part_name) { + return Some(path); + } + } + + None +} + +#[tokio::test] +#[serial] +async fn data_movement_put_object_marks_dirty_disks_for_capacity_manager() { + let (disk_paths, ecstore) = setup_capacity_dirty_scope_env().await; + let bucket_name = format!("dirty-scope-{}", Uuid::new_v4()); + + ecstore + .make_bucket(&bucket_name, &MakeBucketOptions::default()) + .await + .expect("create test bucket"); + + let manager = create_isolated_manager(HybridStrategyConfig::default()); + let _ = manager.get_dirty_disks().await; + + let payload = b"data-movement-dirty-scope".to_vec(); + let mut reader = ChunkNativePutData::from_vec(payload); + let opts = ObjectOptions { + data_movement: true, + src_pool_idx: 0, + ..Default::default() + }; + + ecstore + .put_object(&bucket_name, "object.bin", &mut reader, &opts) + .await + .expect("data movement put_object should succeed"); + + let dirty_disks = manager.get_dirty_disks().await; + assert_eq!(dirty_disks.len(), disk_paths.len()); + + let actual_paths: HashSet<_> = dirty_disks + .into_iter() + .map(|disk| stdfs::canonicalize(&disk.drive_path).unwrap().to_string_lossy().into_owned()) + .collect(); + let expected_paths: HashSet<_> = disk_paths + .iter() + .map(|path| stdfs::canonicalize(path).unwrap().to_string_lossy().into_owned()) + .collect(); + assert_eq!(actual_paths, expected_paths); +} + +#[tokio::test] +#[serial] +async fn heal_object_marks_missing_shard_disk_dirty_for_capacity_manager() { + let (disk_paths, ecstore) = setup_capacity_dirty_scope_env().await; + let bucket_name = format!("dirty-heal-{}", Uuid::new_v4()); + + ecstore + .make_bucket(&bucket_name, &MakeBucketOptions::default()) + .await + .expect("create test bucket"); + + let manager = create_isolated_manager(HybridStrategyConfig::default()); + let _ = manager.get_dirty_disks().await; + + let payload_len = 3 * 1024 * 1024 + 137; + let payload: Vec = (0..payload_len).map(|idx| (idx % 251) as u8).collect(); + let mut reader = ChunkNativePutData::from_vec(payload); + let object_name = "test/heal.bin"; + let put_info = ecstore + .put_object(&bucket_name, object_name, &mut reader, &ObjectOptions::default()) + .await + .expect("put object for heal test"); + assert!(put_info.data_blocks > 1, "expected multi-shard object for heal test"); + + let _ = manager.get_dirty_disks().await; + + let object_root = disk_paths[0].join(&bucket_name).join("test").join("heal.bin"); + let missing_part = find_part_file(&object_root, "part.1").expect("part file on first disk"); + fs::remove_file(&missing_part).await.expect("remove shard to force heal"); + + let heal_opts = HealOpts { + recursive: false, + dry_run: false, + remove: false, + recreate: true, + scan_mode: HealScanMode::Deep, + update_parity: true, + no_lock: false, + pool: None, + set: None, + }; + + let (_result, error) = ecstore + .heal_object(&bucket_name, object_name, "", &heal_opts) + .await + .expect("heal_object call should succeed"); + + let dirty_disks = manager.get_dirty_disks().await; + let actual_paths: HashSet<_> = dirty_disks + .into_iter() + .map(|disk| stdfs::canonicalize(&disk.drive_path).unwrap().to_string_lossy().into_owned()) + .collect(); + let expected_missing_disk = stdfs::canonicalize(&disk_paths[0]).unwrap().to_string_lossy().into_owned(); + + assert!( + error.is_none() || actual_paths.contains(&expected_missing_disk), + "heal returned {error:?} and did not mark the repaired shard disk dirty: {actual_paths:?}" + ); +} diff --git a/rustfs/src/app/lifecycle_transition_api_test.rs b/rustfs/src/app/lifecycle_transition_api_test.rs index 8ed844525..03193697c 100644 --- a/rustfs/src/app/lifecycle_transition_api_test.rs +++ b/rustfs/src/app/lifecycle_transition_api_test.rs @@ -35,12 +35,14 @@ use rustfs_ecstore::{ warm_backend::{WarmBackend, WarmBackendGetOpts}, }, }; +use rustfs_object_capacity::capacity_manager::{HybridStrategyConfig, create_isolated_manager}; use rustfs_utils::http::{SUFFIX_FORCE_DELETE, insert_header}; use s3s::{S3Request, dto::*}; use serial_test::serial; use std::{ collections::HashMap, convert::Infallible, + fs as stdfs, io::Cursor, path::PathBuf, sync::{Arc, Once, OnceLock}, @@ -539,3 +541,49 @@ async fn delete_transitioned_object_removes_remote_tier_copy_via_usecase() { "transitioned object should be removed from remote tier after delete usecase" ); } + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[serial] +#[ignore = "requires isolated global object layer state"] +async fn lifecycle_transition_marks_dirty_disks_for_capacity_manager() { + let (disk_paths, ecstore) = setup_test_env().await; + let manager = create_isolated_manager(HybridStrategyConfig::default()); + let _ = manager.get_dirty_disks().await; + + let tier_name = format!("COLDTIER{}", &Uuid::new_v4().simple().to_string()[..8]).to_uppercase(); + let _backend = register_mock_tier(&tier_name).await; + + let bucket = format!("test-capacity-transition-{}", &Uuid::new_v4().simple().to_string()[..8]); + let object = "test/object.txt"; + let payload = b"transition should mark dirty scope"; + + create_test_bucket(&ecstore, bucket.as_str()).await; + set_bucket_lifecycle_transition_with_tier(bucket.as_str(), &tier_name) + .await + .expect("Failed to set lifecycle configuration"); + let _ = upload_test_object(&ecstore, bucket.as_str(), object, payload).await; + + rustfs_ecstore::bucket::lifecycle::bucket_lifecycle_ops::enqueue_transition_for_existing_objects( + ecstore.clone(), + bucket.as_str(), + ) + .await + .expect("Failed to enqueue transitioned object"); + + let _ = wait_for_transition(&ecstore, bucket.as_str(), object, TRANSITION_WAIT_TIMEOUT) + .await + .expect("object should transition before dirty scope assertion"); + + let dirty_disks = manager.get_dirty_disks().await; + assert_eq!(dirty_disks.len(), disk_paths.len()); + + let actual_paths: std::collections::HashSet<_> = dirty_disks + .into_iter() + .map(|disk| stdfs::canonicalize(&disk.drive_path).unwrap().to_string_lossy().into_owned()) + .collect(); + let expected_paths: std::collections::HashSet<_> = disk_paths + .iter() + .map(|path| stdfs::canonicalize(path).unwrap().to_string_lossy().into_owned()) + .collect(); + assert_eq!(actual_paths, expected_paths); +} diff --git a/rustfs/src/app/mod.rs b/rustfs/src/app/mod.rs index ab8391102..3084c1090 100644 --- a/rustfs/src/app/mod.rs +++ b/rustfs/src/app/mod.rs @@ -21,5 +21,7 @@ pub mod context; pub mod multipart_usecase; pub mod object_usecase; +#[cfg(test)] +mod capacity_dirty_scope_test; #[cfg(test)] mod lifecycle_transition_api_test; diff --git a/rustfs/src/app/multipart_usecase.rs b/rustfs/src/app/multipart_usecase.rs index 16fc92990..8daab532c 100644 --- a/rustfs/src/app/multipart_usecase.rs +++ b/rustfs/src/app/multipart_usecase.rs @@ -16,6 +16,7 @@ use crate::app::context::{AppContext, get_global_app_context}; use crate::app::object_usecase::{build_put_like_object_lock_metadata, validate_existing_object_lock_for_write}; +use crate::capacity::record_capacity_write; use crate::error::ApiError; use crate::storage::access::has_bypass_governance_header; use crate::storage::entity; @@ -66,6 +67,7 @@ use tokio::sync::RwLock; use tokio_util::io::StreamReader; use tracing::{info, instrument, warn}; use urlencoding::encode; +use uuid::Uuid; async fn maybe_enqueue_transition_immediate(obj_info: &rustfs_ecstore::store_api::ObjectInfo, src: LcEventSrc) { enqueue_transition_immediate(obj_info, src).await; @@ -286,7 +288,9 @@ impl DefaultMultipartUsecase { let Some(multipart_upload) = multipart_upload else { return Err(s3_error!(InvalidPart)) }; - let opts = &get_complete_multipart_upload_opts(&req.headers).map_err(ApiError::from)?; + let mut opts = get_complete_multipart_upload_opts(&req.headers).map_err(ApiError::from)?; + let capacity_scope_token = Uuid::new_v4(); + opts.capacity_scope_token = Some(capacity_scope_token); let uploaded_parts_vec = multipart_upload .parts @@ -360,9 +364,10 @@ impl DefaultMultipartUsecase { let obj_info = store .clone() - .complete_multipart_upload(&bucket, &key, &upload_id, uploaded_parts, opts) + .complete_multipart_upload(&bucket, &key, &upload_id, uploaded_parts, &opts) .await .map_err(ApiError::from)?; + record_capacity_write(Some(capacity_scope_token)).await; // check quota after completing multipart upload if let Some(metadata_sys) = self.bucket_metadata_sys() { @@ -1056,7 +1061,7 @@ impl DefaultMultipartUsecase { let mut src_opts = copy_src_opts(&src_bucket, &src_key, &req.headers).map_err(ApiError::from)?; src_opts.version_id = src_version_id.clone(); - let h = http::HeaderMap::new(); + let h = HeaderMap::new(); let get_opts = ObjectOptions { version_id: src_opts.version_id.clone(), versioned: src_opts.versioned, @@ -1108,7 +1113,7 @@ impl DefaultMultipartUsecase { (0, src_info.size) }; - let h = http::HeaderMap::new(); + let h = HeaderMap::new(); let get_opts = ObjectOptions { version_id: src_opts.version_id.clone(), versioned: src_opts.versioned, diff --git a/rustfs/src/app/object_usecase.rs b/rustfs/src/app/object_usecase.rs index 3df36f1c4..71669260b 100644 --- a/rustfs/src/app/object_usecase.rs +++ b/rustfs/src/app/object_usecase.rs @@ -27,14 +27,14 @@ use self::get_object_flow::{GetObjectBootstrap, GetObjectFlowRuntime}; use self::types::*; use crate::app::context::{AppContext, default_notify_interface, get_global_app_context}; -use crate::capacity::capacity_manager::get_capacity_manager; +use crate::capacity::record_capacity_write; use crate::config::RustFSBufferConfig; use crate::error::ApiError; use crate::storage::access::{PostObjectRequestMarker, authorize_request, has_bypass_governance_header, req_info_mut}; use crate::storage::concurrency::{GetObjectGuard, get_concurrency_manager}; use crate::storage::ecfs::*; use crate::storage::head_prefix::{head_prefix_not_found_message, probe_prefix_has_children}; -use crate::storage::helper::OperationHelper; +use crate::storage::helper::{OperationHelper, spawn_background}; use crate::storage::options::{ copy_dst_opts, copy_src_opts, del_opts, extract_metadata, extract_metadata_from_mime_with_object_name, filter_object_metadata, get_content_sha256_with_query, get_opts, normalize_content_encoding_for_storage, put_opts, @@ -1020,9 +1020,9 @@ impl DefaultObjectUsecase { let request_id = req .extensions - .get::() + .get::() .map(|ctx| ctx.request_id.clone()) - .unwrap_or_else(|| crate::storage::request_context::RequestContext::fallback().request_id); + .unwrap_or_else(|| request_context::RequestContext::fallback().request_id); let bootstrap = init_get_object_bootstrap(&req.input.bucket, &req.input.key, &request_id)?; let request_context = prepare_get_object_request_context(&req).await?; let base_buffer_size = self.base_buffer_size(); @@ -1848,6 +1848,7 @@ impl DefaultObjectUsecase { let version_cfg = BucketVersioningSys::get(&bucket).await.unwrap_or_default(); let bypass_governance = has_bypass_governance_header(&req.headers); + let capacity_scope_token = Uuid::new_v4(); #[derive(Default, Clone)] struct DeleteResult { @@ -1970,6 +1971,7 @@ impl DefaultObjectUsecase { object_to_delete.clone(), ObjectOptions { version_suspended: version_cfg.suspended(), + capacity_scope_token: Some(capacity_scope_token), ..Default::default() }, ) @@ -2076,7 +2078,7 @@ impl DefaultObjectUsecase { .as_ref() .map(|context| context.notify()) .unwrap_or_else(default_notify_interface); - crate::storage::helper::spawn_background(async move { + spawn_background(async move { for res in delete_results { if let Some(dobj) = res.delete_object { let event_name = if dobj.delete_marker { @@ -2108,8 +2110,7 @@ impl DefaultObjectUsecase { let result = Ok(S3Response::new(output)); let _ = helper.complete(&result); // Record write operation for capacity management (inline to avoid per-request tokio::spawn overhead) - let manager = get_capacity_manager(); - manager.record_write_operation().await; + record_capacity_write(Some(capacity_scope_token)).await; result } @@ -2144,6 +2145,8 @@ impl DefaultObjectUsecase { let mut opts: ObjectOptions = del_opts(&bucket, &key, version_id, &req.headers, metadata) .await .map_err(ApiError::from)?; + let capacity_scope_token = Uuid::new_v4(); + opts.capacity_scope_token = Some(capacity_scope_token); let force_delete = opts.delete_prefix; let lock_cfg = BucketObjectLockSys::get(&bucket).await; @@ -2255,8 +2258,7 @@ impl DefaultObjectUsecase { }) .version_id(String::new()); let result = Ok(S3Response::with_status(DeleteObjectOutput::default(), StatusCode::NO_CONTENT)); - let manager = get_capacity_manager(); - manager.record_write_operation().await; + record_capacity_write(Some(capacity_scope_token)).await; let _ = helper.complete(&result); return result; } @@ -2304,8 +2306,7 @@ impl DefaultObjectUsecase { let result = Ok(S3Response::new(output)); // Record write operation for capacity management (inline to avoid per-request tokio::spawn overhead) - let manager = get_capacity_manager(); - manager.record_write_operation().await; + record_capacity_write(Some(capacity_scope_token)).await; let _ = helper.complete(&result); result } @@ -2860,7 +2861,7 @@ impl DefaultObjectUsecase { let rreq_clone = rreq.clone(); let version_id_clone = version_id.clone(); - crate::storage::request_context::spawn_traced(async move { + request_context::spawn_traced(async move { let opts = ObjectOptions { transition: TransitionOptions { restore_request: rreq_clone, @@ -2953,7 +2954,7 @@ impl DefaultObjectUsecase { let (tx, rx) = mpsc::channel::>(2); let stream = ReceiverStream::new(rx); - crate::storage::request_context::spawn_traced(async move { + request_context::spawn_traced(async move { let _ = tx .send(Ok(SelectObjectContentEvent::Cont(ContinuationEvent::default()))) .await; diff --git a/rustfs/src/app/object_usecase/app_adapters.rs b/rustfs/src/app/object_usecase/app_adapters.rs index a8156c486..268c10eb5 100644 --- a/rustfs/src/app/object_usecase/app_adapters.rs +++ b/rustfs/src/app/object_usecase/app_adapters.rs @@ -400,7 +400,7 @@ pub(super) fn complete_put_response(helper: OperationHelper, output: PutObjectOu #[allow(clippy::too_many_arguments)] pub(super) fn spawn_put_extract_notification( notify: Arc, - request_context: Option, + request_context: Option, bucket: String, req_params: HashMap, version_id: String, @@ -422,7 +422,7 @@ pub(super) fn spawn_put_extract_notification( user_agent, }; - crate::storage::helper::spawn_background_with_context(request_context, async move { + helper::spawn_background_with_context(request_context, async move { notify.notify(event_args).await; }); } diff --git a/rustfs/src/app/object_usecase/put_object_extract.rs b/rustfs/src/app/object_usecase/put_object_extract.rs index 8f89d3e80..697f03880 100644 --- a/rustfs/src/app/object_usecase/put_object_extract.rs +++ b/rustfs/src/app/object_usecase/put_object_extract.rs @@ -167,10 +167,7 @@ impl DefaultObjectUsecase { let host = get_request_host(&request_context.headers); let port = get_request_port(&request_context.headers); let user_agent = get_request_user_agent(&request_context.headers); - let tracing_context = request_context - .extensions - .get::() - .cloned(); + let tracing_context = request_context.extensions.get::().cloned(); while let Some(entry) = entries.next().await { let mut f = match entry { @@ -299,7 +296,9 @@ impl DefaultObjectUsecase { opts.user_defined.extend(encryption_metadata); } opts.user_defined.extend(metadata); - let mut reader = rustfs_ecstore::store_api::ChunkNativePutData::new(hrd); + let capacity_scope_token = Uuid::new_v4(); + opts.capacity_scope_token = Some(capacity_scope_token); + let mut reader = ChunkNativePutData::new(hrd); let obj_info = match store.put_object(&bucket, &fpath, &mut reader, &opts).await { Ok(info) => info, @@ -311,6 +310,7 @@ impl DefaultObjectUsecase { return Err(ApiError::from(e).into()); } }; + record_capacity_write(Some(capacity_scope_token)).await; let e_tag = obj_info.etag.clone().map(|etag| to_s3s_etag(&etag)); diff --git a/rustfs/src/app/object_usecase/put_object_flow.rs b/rustfs/src/app/object_usecase/put_object_flow.rs index d7cf9b436..9ae0e5a8d 100644 --- a/rustfs/src/app/object_usecase/put_object_flow.rs +++ b/rustfs/src/app/object_usecase/put_object_flow.rs @@ -32,7 +32,7 @@ const SLOW_PUT_PHASE_DEBUG_THRESHOLD_MS: u64 = 100; const SLOW_PUT_PHASE_WARN_THRESHOLD_MS: u64 = 1_000; const SLOW_PUT_PHASE_ERROR_THRESHOLD_MS: u64 = 5_000; -fn resolved_checksum_bytes(checksums: &PutObjectChecksums) -> Option { +fn resolved_checksum_bytes(checksums: &PutObjectChecksums) -> Option { [ (rustfs_rio::ChecksumType::CRC32, checksums.crc32.as_deref()), (rustfs_rio::ChecksumType::CRC32C, checksums.crc32c.as_deref()), @@ -117,7 +117,7 @@ fn log_put_flow_phase( bucket: &str, key: &str, phase: &str, - elapsed: std::time::Duration, + elapsed: Duration, object_size: i64, small_eager: bool, reduced_copy: bool, @@ -159,11 +159,11 @@ impl PooledBufferReader { } } -impl tokio::io::AsyncRead for PooledBufferReader { +impl AsyncRead for PooledBufferReader { fn poll_read( mut self: std::pin::Pin<&mut Self>, _cx: &mut std::task::Context<'_>, - buf: &mut tokio::io::ReadBuf<'_>, + buf: &mut ReadBuf<'_>, ) -> std::task::Poll> { let remaining = &self.buffer[self.position..]; if remaining.is_empty() { @@ -199,7 +199,7 @@ impl HashReaderDetector for PooledBufferReader {} impl TryGetIndex for PooledBufferReader {} -async fn read_small_put_body_eager(body: S, size: i64, pool: std::sync::Arc) -> S3Result +async fn read_small_put_body_eager(body: S, size: i64, pool: Arc) -> S3Result where S: Stream>, B: Buf, @@ -243,7 +243,7 @@ where async fn build_small_put_eager_hash_stage( body: S, size: i64, - pool: std::sync::Arc, + pool: Arc, hash_values: PutObjectLegacyHashValues, headers: &HeaderMap, trailing_headers: Option, @@ -587,6 +587,8 @@ impl DefaultObjectUsecase { let mt2 = metadata.clone(); opts.user_defined.extend(metadata); + let capacity_scope_token = Uuid::new_v4(); + opts.capacity_scope_token = Some(capacity_scope_token); let repoptions = get_must_replicate_options(&mt2, "".to_string(), ReplicationStatusType::Empty, ReplicationType::Object, opts.clone()); @@ -689,8 +691,7 @@ impl DefaultObjectUsecase { ..Default::default() }; - let manager = get_capacity_manager(); - manager.record_write_operation().await; + record_capacity_write(Some(capacity_scope_token)).await; { let duration_ms = start_time.elapsed().as_millis() as f64; diff --git a/rustfs/src/app/object_usecase/zero_copy_tests.rs b/rustfs/src/app/object_usecase/zero_copy_tests.rs index 5d9dc4f99..3aa253bdf 100644 --- a/rustfs/src/app/object_usecase/zero_copy_tests.rs +++ b/rustfs/src/app/object_usecase/zero_copy_tests.rs @@ -85,7 +85,7 @@ async fn setup_direct_chunk_test_env() -> (Vec, Arc) { .unwrap(); let buckets_list = ecstore - .list_bucket(&rustfs_ecstore::store_api::BucketOptions { + .list_bucket(&BucketOptions { no_metadata: true, ..Default::default() }) @@ -151,7 +151,7 @@ async fn setup_direct_chunk_multi_disk_test_env() -> (Vec, Arc .unwrap(); let buckets_list = ecstore - .list_bucket(&rustfs_ecstore::store_api::BucketOptions { + .list_bucket(&BucketOptions { no_metadata: true, ..Default::default() }) @@ -212,7 +212,7 @@ async fn create_direct_chunk_test_multipart_object( parts } -fn find_part_file(root: &std::path::Path, part_name: &str) -> Option { +fn find_part_file(root: &Path, part_name: &str) -> Option { let entries = std::fs::read_dir(root).ok()?; for entry in entries.flatten() { let path = entry.path(); @@ -231,7 +231,7 @@ fn find_part_file(root: &std::path::Path, part_name: &str) -> Option { None } -fn find_part_files(root: &std::path::Path, part_name: &str, out: &mut Vec) { +fn find_part_files(root: &Path, part_name: &str, out: &mut Vec) { let Ok(entries) = std::fs::read_dir(root) else { return; }; @@ -248,7 +248,7 @@ fn find_part_files(root: &std::path::Path, part_name: &str, out: &mut Vec Vec<(PathBuf, Vec)> { +async fn remove_part_files(root: &Path, part_name: &str) -> Vec<(PathBuf, Vec)> { let mut paths = Vec::new(); find_part_files(root, part_name, &mut paths); diff --git a/rustfs/src/capacity/capacity_integration.rs b/rustfs/src/capacity/capacity_integration.rs index 82dcffb29..e161736b7 100644 --- a/rustfs/src/capacity/capacity_integration.rs +++ b/rustfs/src/capacity/capacity_integration.rs @@ -14,71 +14,26 @@ //! Capacity management integration for application startup -use crate::capacity::capacity_manager::{DataSource, get_capacity_manager, start_background_task}; -use rustfs_ecstore::disk::DiskAPI; -use rustfs_io_metrics::{record_capacity_cache_hit, record_capacity_cache_miss}; -use tracing::{info, warn}; +use crate::capacity::{get_cached_capacity_with_metrics, init_capacity_management_for_local_disks}; /// Initialize capacity management system /// This should be called during application startup after local disks are initialized pub async fn init_capacity_management() { - info!("Initializing capacity management system..."); - - // Get all local disks - let disks = rustfs_ecstore::store::all_local_disk().await; - - if disks.is_empty() { - warn!("No local disks found, capacity management will not run"); - return; - } - - info!("Found {} local disk(s)", disks.len()); - - // Convert DiskStore to Disk (for compatibility with capacity_manager) - let disk_refs: Vec = disks - .iter() - .map(|ds| rustfs_madmin::Disk { - endpoint: ds.endpoint().to_string(), - drive_path: ds.to_string(), - root_disk: true, - ..Default::default() - }) - .collect(); - - // Start background update task - info!("Starting background capacity update task..."); - start_background_task(disk_refs).await; - - info!("Capacity management system initialized successfully"); + init_capacity_management_for_local_disks().await; } /// Get capacity statistics with metrics #[allow(dead_code)] pub async fn get_capacity_with_metrics() -> Option<(u64, String)> { - let manager = get_capacity_manager(); - - // Check cache - if let Some(cached) = manager.get_capacity().await { - record_capacity_cache_hit(); - - let source = match cached.source { - DataSource::RealTime => "real-time", - DataSource::Scheduled => "scheduled", - DataSource::WriteTriggered => "write-triggered", - DataSource::Fallback => "fallback", - }; - - return Some((cached.total_used, source.to_string())); - } - - record_capacity_cache_miss(); - None + get_cached_capacity_with_metrics() + .await + .map(|(capacity, source)| (capacity, source.to_string())) } #[cfg(test)] mod tests { use super::*; - use crate::capacity::capacity_manager::{CapacityUpdate, DataSource, get_capacity_manager}; + use rustfs_object_capacity::capacity_manager::{CapacityUpdate, DataSource, get_capacity_manager}; #[tokio::test] async fn test_get_capacity_with_metrics() { diff --git a/rustfs/src/capacity/capacity_manager_test.rs b/rustfs/src/capacity/capacity_manager_test.rs deleted file mode 100644 index 33158b1f9..000000000 --- a/rustfs/src/capacity/capacity_manager_test.rs +++ /dev/null @@ -1,267 +0,0 @@ -// Copyright 2024 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. - -//! Comprehensive tests for Hybrid Capacity Manager - -#[cfg(test)] -mod tests { - use crate::capacity::capacity_manager::{CapacityUpdate, DataSource, HybridCapacityManager, HybridStrategyConfig}; - use serial_test::serial; - use std::sync::Arc; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::time::Duration; - use tokio::time::sleep; - - #[tokio::test] - #[serial] - async fn test_capacity_manager_initialization() { - let manager = HybridCapacityManager::from_env(); - assert!(manager.get_capacity().await.is_none()); - } - - #[tokio::test] - async fn test_capacity_update_and_retrieval() { - let manager = HybridCapacityManager::from_env(); - - assert!(manager.get_capacity().await.is_none()); - - manager - .update_capacity(CapacityUpdate::exact(1000, 10), DataSource::RealTime) - .await; - - let cached = manager.get_capacity().await; - assert!(cached.is_some()); - let cached = cached.unwrap(); - assert_eq!(cached.total_used, 1000); - assert_eq!(cached.file_count, 10); - assert_eq!(cached.source, DataSource::RealTime); - assert!(!cached.is_estimated); - } - - #[tokio::test] - async fn test_write_operation_recording() { - let manager = HybridCapacityManager::from_env(); - - manager.record_write_operation().await; - manager.record_write_operation().await; - manager.record_write_operation().await; - - let frequency = manager.get_write_frequency().await; - assert_eq!(frequency, 3); - } - - #[tokio::test] - async fn test_fast_update_detection() { - let manager = HybridCapacityManager::from_env(); - - assert!(!manager.needs_fast_update().await); - - manager - .update_capacity(CapacityUpdate::exact(1000, 1), DataSource::RealTime) - .await; - - assert!(!manager.needs_fast_update().await); - - manager.record_write_operation().await; - sleep(Duration::from_millis(100)).await; - - let _needs_update = manager.needs_fast_update().await; - } - - #[tokio::test] - async fn test_cache_age_tracking() { - let manager = HybridCapacityManager::from_env(); - - assert!(manager.get_cache_age().await.is_none()); - - manager - .update_capacity(CapacityUpdate::exact(1000, 1), DataSource::RealTime) - .await; - - let age = manager.get_cache_age().await; - assert!(age.is_some()); - let age = age.unwrap(); - assert!(age < Duration::from_secs(1)); - - sleep(Duration::from_millis(100)).await; - - let age = manager.get_cache_age().await.unwrap(); - assert!(age >= Duration::from_millis(100)); - } - - #[tokio::test] - async fn test_data_source_tracking() { - let manager = HybridCapacityManager::from_env(); - - let sources = vec![ - DataSource::RealTime, - DataSource::Scheduled, - DataSource::WriteTriggered, - DataSource::Fallback, - ]; - - for source in sources { - manager.update_capacity(CapacityUpdate::exact(1000, 1), source).await; - let cached = manager.get_capacity().await.unwrap(); - assert_eq!(cached.source, source); - } - } - - #[tokio::test] - async fn test_config_from_env() { - let config = HybridStrategyConfig::from_env(); - - assert_eq!(config.scheduled_update_interval, Duration::from_secs(120)); - assert_eq!(config.write_trigger_delay, Duration::from_secs(5)); - assert_eq!(config.write_frequency_threshold, 5); - assert_eq!(config.fast_update_threshold, Duration::from_secs(30)); - assert!(config.enable_smart_update); - assert!(config.enable_write_trigger); - } - - #[tokio::test] - async fn test_write_frequency_window() { - let manager = HybridCapacityManager::from_env(); - - for _ in 0..20 { - manager.record_write_operation().await; - } - - let frequency = manager.get_write_frequency().await; - assert_eq!(frequency, 20); - } - - #[tokio::test] - #[serial] - async fn test_concurrent_access() { - let manager = Arc::new(HybridCapacityManager::from_env()); - let mut handles = vec![]; - - for i in 0..10 { - let mgr = manager.clone(); - let handle = tokio::spawn(async move { - mgr.update_capacity(CapacityUpdate::exact(i as u64 * 100, i), DataSource::RealTime) - .await; - mgr.record_write_operation().await; - }); - handles.push(handle); - } - - for handle in handles { - handle.await.unwrap(); - } - - let cached = manager.get_capacity().await; - assert!(cached.is_some()); - - let frequency = manager.get_write_frequency().await; - assert_eq!(frequency, 10); - } - - #[tokio::test] - #[serial] - async fn test_performance_overhead() { - let manager = Arc::new(HybridCapacityManager::from_env()); - let start = std::time::Instant::now(); - - for i in 0..1000 { - manager - .update_capacity(CapacityUpdate::exact(i as u64, i), DataSource::RealTime) - .await; - manager.record_write_operation().await; - let _ = manager.get_capacity().await; - } - - let elapsed = start.elapsed(); - assert!(elapsed < Duration::from_secs(1)); - - println!("1000 operations completed in {:?}", elapsed); - } - - #[tokio::test] - async fn test_refresh_or_join_singleflight() { - let manager = Arc::new(HybridCapacityManager::from_env()); - let calls = Arc::new(AtomicUsize::new(0)); - - let mgr1 = manager.clone(); - let calls1 = calls.clone(); - let first = tokio::spawn(async move { - mgr1.refresh_or_join(DataSource::Scheduled, move || async move { - calls1.fetch_add(1, Ordering::SeqCst); - sleep(Duration::from_millis(50)).await; - Ok(CapacityUpdate::exact(2048, 8)) - }) - .await - }); - - sleep(Duration::from_millis(10)).await; - - let mgr2 = manager.clone(); - let calls2 = calls.clone(); - let second = tokio::spawn(async move { - mgr2.refresh_or_join(DataSource::WriteTriggered, move || async move { - calls2.fetch_add(1, Ordering::SeqCst); - Ok(CapacityUpdate::exact(4096, 16)) - }) - .await - }); - - let first = first.await.unwrap().unwrap(); - let second = second.await.unwrap().unwrap(); - - assert_eq!(calls.load(Ordering::SeqCst), 1); - assert_eq!(first.total_used, 2048); - assert_eq!(second.total_used, 2048); - let cached = manager.get_capacity().await.unwrap(); - assert_eq!(cached.total_used, 2048); - assert_eq!(cached.file_count, 8); - } - - #[tokio::test] - async fn test_spawn_refresh_if_needed_deduplicates_background_refresh() { - let manager = Arc::new(HybridCapacityManager::from_env()); - let calls = Arc::new(AtomicUsize::new(0)); - - let first_manager = manager.clone(); - let first_calls = calls.clone(); - let started = first_manager - .clone() - .spawn_refresh_if_needed(DataSource::Scheduled, move || async move { - first_calls.fetch_add(1, Ordering::SeqCst); - sleep(Duration::from_millis(50)).await; - Ok(CapacityUpdate::estimated(8192, 32)) - }) - .await; - assert!(started); - - let second_manager = manager.clone(); - let second_calls = calls.clone(); - let started = second_manager - .clone() - .spawn_refresh_if_needed(DataSource::Scheduled, move || async move { - second_calls.fetch_add(1, Ordering::SeqCst); - Ok(CapacityUpdate::exact(1, 1)) - }) - .await; - assert!(!started); - - sleep(Duration::from_millis(100)).await; - - assert_eq!(calls.load(Ordering::SeqCst), 1); - assert!(!manager.refresh_in_progress().await); - let cached = manager.get_capacity().await.unwrap(); - assert_eq!(cached.total_used, 8192); - assert!(cached.is_estimated); - } -} diff --git a/rustfs/src/capacity/mod.rs b/rustfs/src/capacity/mod.rs index 10e0c7749..1bb9ae2aa 100644 --- a/rustfs/src/capacity/mod.rs +++ b/rustfs/src/capacity/mod.rs @@ -48,22 +48,11 @@ //! Capacity metrics flow through the existing observability pipeline via the `metrics` //! crate and `rustfs-io-metrics`; this module does not expose a Prometheus HTTP endpoint. //! -//! ## Testing -//! -//! For isolated tests, use `create_isolated_manager()` to create independent -//! instances instead of the global singleton: -//! -//! ```ignore -//! use crate::capacity::create_isolated_manager; -//! -//! let manager = create_isolated_manager(HybridStrategyConfig::default()); -//! // Test without affecting global state -//! ``` -//! pub mod capacity_integration; -pub mod capacity_manager; -#[cfg(test)] -mod capacity_manager_test; -#[cfg(test)] -mod write_trigger_test; +pub mod service; + +pub use service::{ + capacity_disk_ref, get_cached_capacity_with_metrics, init_capacity_management_for_local_disks, record_capacity_write, + refresh_or_join_admin_disks, resolve_admin_used_capacity, spawn_refresh_if_needed_admin_disks, +}; diff --git a/rustfs/src/capacity/service.rs b/rustfs/src/capacity/service.rs new file mode 100644 index 000000000..7d64d3173 --- /dev/null +++ b/rustfs/src/capacity/service.rs @@ -0,0 +1,238 @@ +// Copyright 2024 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 rustfs_ecstore::disk::DiskAPI; +use rustfs_io_metrics::capacity_metrics::{ + record_capacity_cache_hit, record_capacity_cache_miss, record_capacity_cache_served, record_capacity_refresh_request, + record_capacity_scan_mode, +}; +use rustfs_object_capacity::{CapacityDiskRef, capacity_manager, scan}; +use std::sync::Arc; +use std::time::Instant; +use tracing::{debug, info, warn}; + +pub fn capacity_disk_ref(endpoint: impl Into, drive_path: impl Into) -> CapacityDiskRef { + CapacityDiskRef { + endpoint: endpoint.into(), + drive_path: drive_path.into(), + } +} + +fn capacity_disk_refs(disks: &[rustfs_madmin::Disk]) -> Vec { + disks + .iter() + .map(|disk| capacity_disk_ref(disk.endpoint.clone(), disk.drive_path.clone())) + .collect() +} + +async fn refresh_admin_disks_with_subset_fallback( + capacity_manager: &capacity_manager::HybridCapacityManager, + all_disks: Vec, + allow_dirty_subset: bool, +) -> Result { + let (refresh_disks, dirty_subset) = if allow_dirty_subset { + scan::select_capacity_refresh_disks(capacity_manager, &all_disks).await + } else { + (all_disks.clone(), false) + }; + + match scan::refresh_capacity_with_scope(refresh_disks.clone(), dirty_subset).await { + Ok(update) => Ok(update), + Err(err) if dirty_subset => { + warn!("Dirty-subset capacity refresh failed: {}. Retrying full-disk refresh for recovery", err); + scan::refresh_capacity_with_scope(all_disks, false).await + } + Err(err) => Err(err), + } +} + +pub async fn refresh_or_join_admin_disks( + capacity_manager: Arc, + source: capacity_manager::DataSource, + disks: &[rustfs_madmin::Disk], + allow_dirty_subset: bool, +) -> Result { + let all_disks = capacity_disk_refs(disks); + let refresh_manager = capacity_manager.clone(); + + capacity_manager + .refresh_or_join(source, move || { + let capacity_manager = refresh_manager.clone(); + let all_disks = all_disks.clone(); + async move { + refresh_admin_disks_with_subset_fallback(capacity_manager.as_ref(), all_disks, allow_dirty_subset).await + } + }) + .await +} + +pub async fn spawn_refresh_if_needed_admin_disks( + capacity_manager: Arc, + source: capacity_manager::DataSource, + disks: &[rustfs_madmin::Disk], + allow_dirty_subset: bool, +) -> bool { + let all_disks = capacity_disk_refs(disks); + let refresh_manager = capacity_manager.clone(); + + capacity_manager + .spawn_refresh_if_needed(source, move || async move { + refresh_admin_disks_with_subset_fallback(refresh_manager.as_ref(), all_disks, allow_dirty_subset).await + }) + .await +} + +pub async fn record_capacity_write(scope_token: Option) { + capacity_manager::get_capacity_manager() + .record_write_operation_with_scope_token(scope_token) + .await; +} + +pub async fn resolve_admin_used_capacity(disks: &[rustfs_madmin::Disk], fallback_used_capacity: u64) -> u64 { + let capacity_manager = capacity_manager::get_capacity_manager(); + + if let Some(cached) = capacity_manager.get_capacity().await { + record_capacity_cache_hit(); + let cache_age = cached.last_update.elapsed(); + let fast_update_threshold = capacity_manager.get_config().fast_update_threshold; + + if cache_age < fast_update_threshold { + record_capacity_cache_served("fresh"); + debug!( + "Using cached capacity: {} bytes (age: {:?}, source: {:?}, files={}, estimated={})", + cached.total_used, cache_age, cached.source, cached.file_count, cached.is_estimated + ); + return cached.total_used; + } + + let needs_update = capacity_manager.needs_fast_update().await; + let should_block = capacity_manager.should_block_on_refresh(cache_age); + + if needs_update && should_block { + let start = Instant::now(); + record_capacity_refresh_request("blocking", capacity_manager::DataSource::WriteTriggered.as_metric_label()); + return match refresh_or_join_admin_disks( + capacity_manager.clone(), + capacity_manager::DataSource::WriteTriggered, + disks, + true, + ) + .await + { + Ok(update) => { + let elapsed = start.elapsed(); + debug!( + "Foreground capacity refresh completed in {:?} (files={}, estimated={})", + elapsed, update.file_count, update.is_estimated + ); + update.total_used + } + Err(err) => { + warn!("Foreground capacity refresh failed: {}, using cached value", err); + record_capacity_cache_served("stale"); + cached.total_used + } + }; + } + + record_capacity_cache_served("stale"); + debug!( + "Using stale cached capacity: {} bytes (age: {:?}, source: {:?}, files={}, estimated={}, needs_update={}, blocking={})", + cached.total_used, cache_age, cached.source, cached.file_count, cached.is_estimated, needs_update, should_block + ); + + record_capacity_refresh_request("background", capacity_manager::DataSource::Scheduled.as_metric_label()); + if spawn_refresh_if_needed_admin_disks(capacity_manager.clone(), capacity_manager::DataSource::Scheduled, disks, true) + .await + { + debug!("Background capacity update started"); + } else { + debug!("Background update already in progress, skipping spawn"); + } + + return cached.total_used; + } + + let start = Instant::now(); + record_capacity_cache_miss(); + record_capacity_refresh_request("initial", capacity_manager::DataSource::RealTime.as_metric_label()); + match refresh_or_join_admin_disks(capacity_manager.clone(), capacity_manager::DataSource::RealTime, disks, false).await { + Ok(update) => { + let elapsed = start.elapsed(); + info!( + "Initial capacity calculation completed: {} bytes in {:?} (files={}, estimated={})", + update.total_used, elapsed, update.file_count, update.is_estimated + ); + update.total_used + } + Err(err) => { + warn!( + "Failed to calculate data directory used capacity: {}, falling back to disk used capacity", + err + ); + record_capacity_cache_served("fallback"); + record_capacity_scan_mode("fallback"); + capacity_manager + .update_capacity( + capacity_manager::CapacityUpdate::fallback(fallback_used_capacity), + capacity_manager::DataSource::Fallback, + ) + .await; + fallback_used_capacity + } + } +} + +pub async fn init_capacity_management_for_local_disks() { + info!("Initializing capacity management system..."); + + let disks = rustfs_ecstore::store::all_local_disk().await; + if disks.is_empty() { + warn!("No local disks found, capacity management will not run"); + return; + } + + info!("Found {} local disk(s)", disks.len()); + + let disk_refs = disks + .iter() + .map(|ds| capacity_disk_ref(ds.endpoint().to_string(), ds.to_string())) + .collect(); + + info!("Starting background capacity update task..."); + capacity_manager::start_background_task(disk_refs).await; + + info!("Capacity management system initialized successfully"); +} + +pub async fn get_cached_capacity_with_metrics() -> Option<(u64, &'static str)> { + let manager = capacity_manager::get_capacity_manager(); + + if let Some(cached) = manager.get_capacity().await { + record_capacity_cache_hit(); + return Some((cached.total_used, capacity_source_label(cached.source))); + } + + record_capacity_cache_miss(); + None +} + +fn capacity_source_label(source: capacity_manager::DataSource) -> &'static str { + match source { + capacity_manager::DataSource::RealTime => "real-time", + capacity_manager::DataSource::Scheduled => "scheduled", + capacity_manager::DataSource::WriteTriggered => "write-triggered", + capacity_manager::DataSource::Fallback => "fallback", + } +} diff --git a/rustfs/src/capacity/write_trigger_test.rs b/rustfs/src/capacity/write_trigger_test.rs deleted file mode 100644 index 1a606e005..000000000 --- a/rustfs/src/capacity/write_trigger_test.rs +++ /dev/null @@ -1,87 +0,0 @@ -// Copyright 2024 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. - -//! Write trigger integration tests - -#[cfg(test)] -mod tests { - use crate::capacity::capacity_manager::{CapacityUpdate, DataSource, HybridCapacityManager}; - use serial_test::serial; - use std::time::Duration; - - #[tokio::test] - #[serial] - async fn test_write_trigger_integration() { - let manager = HybridCapacityManager::from_env(); - - manager.record_write_operation().await; - manager.record_write_operation().await; - manager.record_write_operation().await; - - let frequency = manager.get_write_frequency().await; - assert_eq!(frequency, 3); - } - - #[tokio::test] - #[serial] - async fn test_write_trigger_with_capacity_update() { - let manager = HybridCapacityManager::from_env(); - - manager - .update_capacity(CapacityUpdate::exact(1000, 4), DataSource::WriteTriggered) - .await; - - let cached = manager.get_capacity().await; - assert!(cached.is_some()); - let cached = cached.unwrap(); - assert_eq!(cached.total_used, 1000); - assert_eq!(cached.file_count, 4); - assert_eq!(cached.source, DataSource::WriteTriggered); - } - - #[tokio::test] - async fn test_write_frequency_tracking() { - let manager = HybridCapacityManager::from_env(); - - assert_eq!(manager.get_write_frequency().await, 0); - - for _ in 0..5 { - manager.record_write_operation().await; - } - - assert_eq!(manager.get_write_frequency().await, 5); - - tokio::time::sleep(Duration::from_millis(10)).await; - assert_eq!(manager.get_write_frequency().await, 5); - } - - #[tokio::test] - async fn test_needs_fast_update() { - let manager = HybridCapacityManager::from_env(); - - assert!(!manager.needs_fast_update().await); - - manager - .update_capacity(CapacityUpdate::exact(1000, 1), DataSource::Scheduled) - .await; - - assert!(!manager.needs_fast_update().await); - - manager.record_write_operation().await; - - let needs_update = manager.needs_fast_update().await; - #[allow(clippy::overly_complex_bool_expr)] - let _ = needs_update || !needs_update; - } -}