gsadmin/rustfs
1
0
Fork 0
mirror of https://github.com/rustfs/rustfs.git synced 2026-08-10 15:16:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(capacity): harden scope registry, scan symlink guard, and test temp dir cleanup (#2432)

Browse Source 
Co-authored-by: heihutu <heihutu@gmail.com>
Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: houseme <4829346+houseme@users.noreply.github.com>
This commit is contained in:
houseme
2026-04-08 20:58:17 +08:00
committed by GitHub
parent d4ea14c2ba
commit 064e21062d
32 changed files with 2751 additions and 1217 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rustfs/Cargo.toml
+2 -2
View File
@@ -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
rustfs/src/app/admin_usecase.rs
+3 -695
View File
@@ -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<T> = Result<T, ApiError>;
@@ -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<CapacityScanResult, Box<dyn std::error::Error + Send + Sync>> {
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<PathBuf>,
/// 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<CapacityScanResult, std::io::Error> {
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<Arc<AppContext>>,
@@ -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());
}
}
rustfs/src/app/bucket_usecase.rs
+3 -6
View File
@@ -1494,10 +1494,7 @@ impl DefaultBucketUsecase {
&& let Some(store) = new_object_layer_fn()
{
let bucket_name = bucket.clone();
let request_context = req
.extensions
.get::<crate::storage::request_context::RequestContext>()
.cloned();
let request_context = req.extensions.get::<request_context::RequestContext>().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();
rustfs/src/app/capacity_dirty_scope_test.rs
+224
View File
@@ -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<PathBuf>, Arc<ECStore>, 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<PathBuf>, Arc<ECStore>) {
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<PathBuf> {
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<u8> = (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:?}"
);
}
rustfs/src/app/lifecycle_transition_api_test.rs
+48
View File
@@ -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);
}
rustfs/src/app/mod.rs
+2
View File
@@ -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;
rustfs/src/app/multipart_usecase.rs
+9 -4
View File
@@ -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,
rustfs/src/app/object_usecase.rs
+14 -13
View File
@@ -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::<crate::storage::request_context::RequestContext>()
.get::<request_context::RequestContext>()
.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::<S3Result<SelectObjectContentEvent>>(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;
rustfs/src/app/object_usecase/app_adapters.rs
+2 -2
View File
@@ -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<dyn NotifyInterface>,
request_context: Option<crate::storage::request_context::RequestContext>,
request_context: Option<request_context::RequestContext>,
bucket: String,
req_params: HashMap<String, String>,
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;
});
}
rustfs/src/app/object_usecase/put_object_extract.rs
+5 -5
View File
@@ -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::<crate::storage::request_context::RequestContext>()
.cloned();
let tracing_context = request_context.extensions.get::<request_context::RequestContext>().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));
rustfs/src/app/object_usecase/put_object_flow.rs
+9 -8
View File
@@ -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<bytes::Bytes> {
fn resolved_checksum_bytes(checksums: &PutObjectChecksums) -> Option<Bytes> {
[
(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<std::io::Result<()>> {
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<S, B, E>(body: S, size: i64, pool: std::sync::Arc<BytesPool>) -> S3Result<PooledBuffer>
async fn read_small_put_body_eager<S, B, E>(body: S, size: i64, pool: Arc<BytesPool>) -> S3Result<PooledBuffer>
where
S: Stream<Item = Result<B, E>>,
B: Buf,
@@ -243,7 +243,7 @@ where
async fn build_small_put_eager_hash_stage<S, B, E>(
body: S,
size: i64,
pool: std::sync::Arc<BytesPool>,
pool: Arc<BytesPool>,
hash_values: PutObjectLegacyHashValues,
headers: &HeaderMap,
trailing_headers: Option<s3s::TrailingHeaders>,
@@ -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;
rustfs/src/app/object_usecase/zero_copy_tests.rs
+5 -5
View File
@@ -85,7 +85,7 @@ async fn setup_direct_chunk_test_env() -> (Vec<PathBuf>, Arc<ECStore>) {
.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<PathBuf>, Arc<ECStore>
.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<PathBuf> {
fn find_part_file(root: &Path, part_name: &str) -> Option<PathBuf> {
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<PathBuf> {
None
}
fn find_part_files(root: &std::path::Path, part_name: &str, out: &mut Vec<PathBuf>) {
fn find_part_files(root: &Path, part_name: &str, out: &mut Vec<PathBuf>) {
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<PathBu
}
}
async fn remove_part_files(root: &std::path::Path, part_name: &str) -> Vec<(PathBuf, Vec<u8>)> {
async fn remove_part_files(root: &Path, part_name: &str) -> Vec<(PathBuf, Vec<u8>)> {
let mut paths = Vec::new();
find_part_files(root, part_name, &mut paths);
rustfs/src/capacity/capacity_integration.rs
+6 -51
View File
@@ -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<rustfs_madmin::Disk> = 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() {
rustfs/src/capacity/capacity_manager.rs
-945
View File
@@ -1,945 +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.
//! Hybrid Capacity Manager for efficient capacity statistics
use crate::app::admin_usecase::calculate_data_dir_used_capacity;
use futures::FutureExt;
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,
DEFAULT_FAST_UPDATE_THRESHOLD_SECS, DEFAULT_MAX_FILES_THRESHOLD, DEFAULT_SAMPLE_RATE, DEFAULT_SCHEDULED_UPDATE_INTERVAL_SECS,
DEFAULT_STAT_TIMEOUT_SECS, DEFAULT_WRITE_FREQUENCY_THRESHOLD, DEFAULT_WRITE_TRIGGER_DELAY_SECS,
ENV_CAPACITY_ENABLE_DYNAMIC_TIMEOUT, ENV_CAPACITY_FAST_UPDATE_THRESHOLD, ENV_CAPACITY_FOLLOW_SYMLINKS,
ENV_CAPACITY_MAX_FILES_THRESHOLD, ENV_CAPACITY_MAX_SYMLINK_DEPTH, ENV_CAPACITY_MAX_TIMEOUT, ENV_CAPACITY_MIN_TIMEOUT,
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_utils::{get_env_bool, get_env_u64, get_env_usize};
use std::future::Future;
use std::panic::AssertUnwindSafe;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::{Mutex, RwLock, watch};
use tracing::{debug, info, warn};
// ============================================================================
// Configuration Functions
// ============================================================================
/// Cached capacity configuration to avoid repeated environment variable reads
#[derive(Clone, Debug)]
struct CachedCapacityConfig {
/// Scheduled update interval
scheduled_update_interval: Duration,
/// Write trigger delay
write_trigger_delay: Duration,
/// Write frequency threshold
write_frequency_threshold: usize,
/// Fast update threshold
fast_update_threshold: Duration,
/// Max files threshold for sampling
max_files_threshold: usize,
/// Stat timeout
stat_timeout: Duration,
/// Sample rate
sample_rate: usize,
/// Follow symlinks flag
follow_symlinks: bool,
/// Max symlink depth
max_symlink_depth: u8,
/// Enable dynamic timeout flag
enable_dynamic_timeout: bool,
/// Min timeout
min_timeout: Duration,
/// Max timeout
max_timeout: Duration,
/// Stall timeout
stall_timeout: Duration,
}
impl CachedCapacityConfig {
/// Build configuration from environment variables
fn from_env() -> Self {
Self {
scheduled_update_interval: Duration::from_secs(get_env_u64(
ENV_CAPACITY_SCHEDULED_INTERVAL,
DEFAULT_SCHEDULED_UPDATE_INTERVAL_SECS,
)),
write_trigger_delay: Duration::from_secs(get_env_u64(
ENV_CAPACITY_WRITE_TRIGGER_DELAY,
DEFAULT_WRITE_TRIGGER_DELAY_SECS,
)),
write_frequency_threshold: get_env_usize(ENV_CAPACITY_WRITE_FREQUENCY_THRESHOLD, DEFAULT_WRITE_FREQUENCY_THRESHOLD),
fast_update_threshold: Duration::from_secs(get_env_u64(
ENV_CAPACITY_FAST_UPDATE_THRESHOLD,
DEFAULT_FAST_UPDATE_THRESHOLD_SECS,
)),
max_files_threshold: get_env_usize(ENV_CAPACITY_MAX_FILES_THRESHOLD, DEFAULT_MAX_FILES_THRESHOLD),
stat_timeout: Duration::from_secs(get_env_u64(ENV_CAPACITY_STAT_TIMEOUT, DEFAULT_STAT_TIMEOUT_SECS)),
sample_rate: get_env_usize(ENV_CAPACITY_SAMPLE_RATE, DEFAULT_SAMPLE_RATE),
follow_symlinks: get_env_bool(ENV_CAPACITY_FOLLOW_SYMLINKS, DEFAULT_CAPACITY_FOLLOW_SYMLINKS),
max_symlink_depth: get_env_u64(ENV_CAPACITY_MAX_SYMLINK_DEPTH, DEFAULT_CAPACITY_MAX_SYMLINK_DEPTH as u64) as u8,
enable_dynamic_timeout: get_env_bool(ENV_CAPACITY_ENABLE_DYNAMIC_TIMEOUT, DEFAULT_CAPACITY_ENABLE_DYNAMIC_TIMEOUT),
min_timeout: Duration::from_secs(get_env_u64(ENV_CAPACITY_MIN_TIMEOUT, DEFAULT_CAPACITY_MIN_TIMEOUT_SECS)),
max_timeout: Duration::from_secs(get_env_u64(ENV_CAPACITY_MAX_TIMEOUT, DEFAULT_CAPACITY_MAX_TIMEOUT_SECS)),
stall_timeout: Duration::from_secs(get_env_u64(ENV_CAPACITY_STALL_TIMEOUT, DEFAULT_CAPACITY_STALL_TIMEOUT_SECS)),
}
}
}
/// Get cached capacity configuration (reads environment variables once)
#[cfg(not(test))]
fn get_cached_config() -> &'static CachedCapacityConfig {
static CONFIG: std::sync::OnceLock<CachedCapacityConfig> = std::sync::OnceLock::new();
CONFIG.get_or_init(CachedCapacityConfig::from_env)
}
#[cfg(test)]
fn get_cached_config() -> CachedCapacityConfig {
// Don't cache in tests to allow temp_env::with_var to work
CachedCapacityConfig::from_env()
}
/// Get scheduled update interval from environment or default
#[cfg(not(test))]
pub fn get_scheduled_update_interval() -> Duration {
get_cached_config().scheduled_update_interval
}
/// Get scheduled update interval from environment or default (test mode)
#[cfg(test)]
pub fn get_scheduled_update_interval() -> Duration {
get_cached_config().scheduled_update_interval
}
/// Get write trigger delay from environment or default
#[cfg(not(test))]
pub fn get_write_trigger_delay() -> Duration {
get_cached_config().write_trigger_delay
}
/// Get write trigger delay from environment or default (test mode)
#[cfg(test)]
pub fn get_write_trigger_delay() -> Duration {
get_cached_config().write_trigger_delay
}
/// Get write frequency threshold from environment or default
#[cfg(not(test))]
pub fn get_write_frequency_threshold() -> usize {
get_cached_config().write_frequency_threshold
}
/// Get write frequency threshold from environment or default (test mode)
#[cfg(test)]
pub fn get_write_frequency_threshold() -> usize {
get_cached_config().write_frequency_threshold
}
/// Get fast update threshold from environment or default
#[cfg(not(test))]
pub fn get_fast_update_threshold() -> Duration {
get_cached_config().fast_update_threshold
}
/// Get fast update threshold from environment or default (test mode)
#[cfg(test)]
pub fn get_fast_update_threshold() -> Duration {
get_cached_config().fast_update_threshold
}
/// Get max files threshold from environment or default
#[cfg(not(test))]
pub fn get_max_files_threshold() -> usize {
get_cached_config().max_files_threshold
}
/// Get max files threshold from environment or default (test mode)
#[cfg(test)]
pub fn get_max_files_threshold() -> usize {
get_cached_config().max_files_threshold
}
/// Get stat timeout from environment or default
#[cfg(not(test))]
pub fn get_stat_timeout() -> Duration {
get_cached_config().stat_timeout
}
/// Get stat timeout from environment or default (test mode)
#[cfg(test)]
pub fn get_stat_timeout() -> Duration {
get_cached_config().stat_timeout
}
/// Get sample rate from environment or default
#[cfg(not(test))]
pub fn get_sample_rate() -> usize {
get_cached_config().sample_rate
}
/// Get sample rate from environment or default (test mode)
#[cfg(test)]
pub fn get_sample_rate() -> usize {
get_cached_config().sample_rate
}
/// Get follow symlinks flag from environment or default
#[cfg(not(test))]
pub fn get_follow_symlinks() -> bool {
get_cached_config().follow_symlinks
}
/// Get follow symlinks flag from environment or default (test mode)
#[cfg(test)]
pub fn get_follow_symlinks() -> bool {
get_cached_config().follow_symlinks
}
/// Get max symlink depth from environment or default
#[cfg(not(test))]
pub fn get_max_symlink_depth() -> u8 {
get_cached_config().max_symlink_depth
}
/// Get max symlink depth from environment or default (test mode)
#[cfg(test)]
pub fn get_max_symlink_depth() -> u8 {
get_cached_config().max_symlink_depth
}
/// Get enable dynamic timeout flag from environment or default
#[cfg(not(test))]
pub fn get_enable_dynamic_timeout() -> bool {
get_cached_config().enable_dynamic_timeout
}
/// Get enable dynamic timeout flag from environment or default (test mode)
#[cfg(test)]
pub fn get_enable_dynamic_timeout() -> bool {
get_cached_config().enable_dynamic_timeout
}
/// Get min timeout from environment or default
#[cfg(not(test))]
pub fn get_min_timeout() -> Duration {
get_cached_config().min_timeout
}
/// Get min timeout from environment or default (test mode)
#[cfg(test)]
pub fn get_min_timeout() -> Duration {
get_cached_config().min_timeout
}
/// Get max timeout from environment or default
#[cfg(not(test))]
pub fn get_max_timeout() -> Duration {
get_cached_config().max_timeout
}
/// Get max timeout from environment or default (test mode)
#[cfg(test)]
pub fn get_max_timeout() -> Duration {
get_cached_config().max_timeout
}
/// Get stall timeout from environment or default
#[cfg(not(test))]
pub fn get_stall_timeout() -> Duration {
get_cached_config().stall_timeout
}
/// Get stall timeout from environment or default (test mode)
#[cfg(test)]
pub fn get_stall_timeout() -> Duration {
get_cached_config().stall_timeout
}
// ============================================================================
// Data Structures
// ============================================================================
/// Cached capacity data
#[derive(Clone, Debug)]
pub struct CachedCapacity {
/// Total used capacity in bytes
pub total_used: u64,
/// Last update time
pub last_update: Instant,
/// File count (optional)
pub file_count: usize,
/// Whether it's an estimated value
pub is_estimated: bool,
/// Data source
pub source: DataSource,
}
/// Structured capacity update payload.
#[derive(Clone, Debug)]
pub struct CapacityUpdate {
/// Total used capacity in bytes.
pub total_used: u64,
/// Number of files observed during scan.
pub file_count: usize,
/// Whether the value is estimated instead of exact.
pub is_estimated: bool,
}
impl CapacityUpdate {
/// Create an exact capacity update.
pub fn exact(total_used: u64, file_count: usize) -> Self {
Self {
total_used,
file_count,
is_estimated: false,
}
}
/// Create an estimated capacity update.
pub fn estimated(total_used: u64, file_count: usize) -> Self {
Self {
total_used,
file_count,
is_estimated: true,
}
}
/// Create a fallback capacity update.
pub fn fallback(total_used: u64) -> Self {
Self {
total_used,
file_count: 0,
is_estimated: true,
}
}
}
#[derive(Clone, Debug, PartialEq, Copy, Eq)]
pub enum DataSource {
/// Real-time statistics
RealTime,
/// Scheduled update
Scheduled,
/// Write triggered
WriteTriggered,
/// Fallback value
#[allow(dead_code)]
Fallback,
}
impl DataSource {
fn as_metric_label(self) -> &'static str {
match self {
Self::RealTime => "realtime",
Self::Scheduled => "scheduled",
Self::WriteTriggered => "write_triggered",
Self::Fallback => "fallback",
}
}
}
/// Write record for tracking write operations
#[derive(Debug)]
pub struct WriteRecord {
/// Last write time
pub last_write_time: Instant,
/// Write count
pub write_count: usize,
/// Write time window (for frequency calculation)
pub write_window: Vec<Instant>,
}
/// Hybrid strategy configuration
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub struct HybridStrategyConfig {
/// Scheduled update interval
pub scheduled_update_interval: Duration,
/// Write trigger delay
pub write_trigger_delay: Duration,
/// Write frequency threshold (writes/minute)
pub write_frequency_threshold: usize,
/// Fast update threshold
pub fast_update_threshold: Duration,
/// Enable smart update
pub enable_smart_update: bool,
/// Enable write trigger
pub enable_write_trigger: bool,
}
impl Default for HybridStrategyConfig {
fn default() -> Self {
Self {
scheduled_update_interval: get_scheduled_update_interval(),
write_trigger_delay: get_write_trigger_delay(),
write_frequency_threshold: get_write_frequency_threshold(),
fast_update_threshold: get_fast_update_threshold(),
enable_smart_update: true,
enable_write_trigger: true,
}
}
}
impl HybridStrategyConfig {
/// Create config from environment variables
pub fn from_env() -> Self {
Self::default()
}
}
// ============================================================================
// Hybrid Capacity Manager
// ============================================================================
struct RefreshState {
running: bool,
/// Sender for the current refresh cycle. Joiners subscribe to this before releasing the
/// mutex so they cannot miss the completion notification. A new channel is created at the
/// start of every refresh cycle so stale subscribers from previous cycles are not confused
/// by results that were already published.
result_tx: watch::Sender<Option<Result<CapacityUpdate, String>>>,
}
impl Default for RefreshState {
fn default() -> Self {
let (tx, _) = watch::channel(None);
Self {
running: false,
result_tx: tx,
}
}
}
/// Hybrid capacity manager
pub struct HybridCapacityManager {
/// Capacity cache
cache: Arc<RwLock<Option<CachedCapacity>>>,
/// Write record
write_record: Arc<RwLock<WriteRecord>>,
/// Configuration
config: HybridStrategyConfig,
/// Shared singleflight refresh state
refresh_state: Arc<Mutex<RefreshState>>,
}
impl HybridCapacityManager {
fn max_stale_age(&self) -> Duration {
self.config
.scheduled_update_interval
.max(self.config.fast_update_threshold.checked_mul(3).unwrap_or(Duration::MAX))
}
/// Create a new hybrid capacity manager
pub fn new(config: HybridStrategyConfig) -> Self {
Self {
cache: Arc::new(RwLock::new(None)),
write_record: Arc::new(RwLock::new(WriteRecord {
last_write_time: Instant::now(),
write_count: 0,
write_window: Vec::new(),
})),
config,
refresh_state: Arc::new(Mutex::new(RefreshState::default())),
}
}
/// Create with default config from environment
pub fn from_env() -> Self {
Self::new(HybridStrategyConfig::from_env())
}
/// Get capacity (core method)
pub async fn get_capacity(&self) -> Option<CachedCapacity> {
let cache = self.cache.read().await;
cache.clone()
}
/// Update capacity
pub async fn update_capacity(&self, update: CapacityUpdate, source: DataSource) {
let start = Instant::now();
let mut cache = self.cache.write().await;
*cache = Some(CachedCapacity {
total_used: update.total_used,
last_update: Instant::now(),
file_count: update.file_count,
is_estimated: update.is_estimated,
source,
});
debug!(
"Capacity updated: {} bytes, files={}, estimated={}, source: {:?}",
update.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 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);
}
record_capacity_write_operation(record.write_window.len());
debug!(
"Write operation recorded: total writes = {}, recent writes = {}",
record.write_count,
record.write_window.len()
);
}
/// Check if fast update is needed
pub async fn needs_fast_update(&self) -> bool {
if !self.config.enable_smart_update {
return false;
}
let cache = self.cache.read().await;
if let Some(cached) = cache.as_ref() {
let cache_age = cached.last_update.elapsed();
// Cache is fresh, no need to update
if cache_age < self.config.fast_update_threshold {
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;
}
// 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);
return true;
}
}
false
}
/// Get cache age
#[allow(dead_code)]
pub async fn get_cache_age(&self) -> Option<Duration> {
let cache = self.cache.read().await;
cache.as_ref().map(|c| c.last_update.elapsed())
}
/// Get write frequency (writes/minute)
#[allow(dead_code)]
pub async fn get_write_frequency(&self) -> usize {
let record = self.write_record.read().await;
record.write_window.len()
}
/// Run a singleflight refresh. Callers either join an existing in-flight refresh or become the leader.
///
/// Joiners subscribe to the watch channel *before* releasing the mutex, which guarantees
/// they cannot miss the completion notification even if the leader finishes very quickly.
pub async fn refresh_or_join<F, Fut>(&self, source: DataSource, refresh_fn: F) -> Result<CapacityUpdate, String>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<CapacityUpdate, String>>,
{
let maybe_rx = {
let mut state = self.refresh_state.lock().await;
if state.running {
// Subscribe while holding the lock so the send that completes the current
// refresh cycle cannot happen before we are subscribed.
Some(state.result_tx.subscribe())
} else {
// Become the leader. Create a fresh channel so that joiners from a previous
// cycle cannot observe the result that was published for the new cycle.
let (tx, _) = watch::channel(None);
state.result_tx = tx;
state.running = true;
None
}
};
if let Some(mut result_rx) = maybe_rx {
// Wait until the leader publishes Some(result). Because we subscribed before
// releasing the mutex, we cannot miss the notification.
if result_rx.wait_for(|v| v.is_some()).await.is_err() {
// The leader's sender was dropped (e.g. due to a panic) without publishing
// a result. Surface a clear error rather than silently returning the default.
return Err("capacity refresh leader exited without publishing a result".to_string());
}
return result_rx
.borrow()
.as_ref()
.cloned()
.unwrap_or_else(|| Err("capacity refresh completed without a result".to_string()));
}
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())
});
if let Ok(update) = &result {
self.update_capacity(update.clone(), source).await;
}
{
let mut state = self.refresh_state.lock().await;
state.running = false;
let _ = state.result_tx.send(Some(result.clone()));
}
result
}
/// Start a background refresh if one is not already in flight.
pub async fn spawn_refresh_if_needed<F, Fut>(self: Arc<Self>, source: DataSource, refresh_fn: F) -> bool
where
F: FnOnce() -> Fut + Send + 'static,
Fut: Future<Output = Result<CapacityUpdate, String>> + Send + 'static,
{
let should_spawn = {
let mut state = self.refresh_state.lock().await;
if state.running {
false
} else {
let (tx, _) = watch::channel(None);
state.result_tx = tx;
state.running = true;
true
}
};
if !should_spawn {
return false;
}
tokio::spawn(async move {
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())
});
if let Ok(update) = &result {
self.update_capacity(update.clone(), source).await;
}
let mut state = self.refresh_state.lock().await;
state.running = false;
let _ = state.result_tx.send(Some(result));
});
true
}
/// Get config
pub fn get_config(&self) -> &HybridStrategyConfig {
&self.config
}
/// Check if the cache is too stale to keep serving without a foreground refresh.
pub fn should_block_on_refresh(&self, cache_age: Duration) -> bool {
cache_age >= self.max_stale_age()
}
/// Return whether a refresh is currently in flight.
#[cfg(test)]
pub async fn refresh_in_progress(&self) -> bool {
self.refresh_state.lock().await.running
}
}
/// Global capacity manager instance
static GLOBAL_CAPACITY_MANAGER: std::sync::OnceLock<Arc<HybridCapacityManager>> = std::sync::OnceLock::new();
/// Get or initialize the global capacity manager
pub fn get_capacity_manager() -> Arc<HybridCapacityManager> {
GLOBAL_CAPACITY_MANAGER
.get_or_init(|| Arc::new(HybridCapacityManager::from_env()))
.clone()
}
/// Create an isolated capacity manager instance for testing
///
/// This factory function allows tests to create independent instances
/// without affecting the global singleton, avoiding test pollution.
///
/// # Example
/// ```no_run
/// 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<HybridCapacityManager> {
Arc::new(HybridCapacityManager::new(config))
}
/// Start background update task
pub async fn start_background_task(disks: Vec<rustfs_madmin::Disk>) {
let manager = get_capacity_manager();
let mut interval = manager.get_config().scheduled_update_interval;
// Prevent panic in tokio::time::interval when misconfigured to 0
if interval.is_zero() {
warn!("RUSTFS_CAPACITY_SCHEDULED_INTERVAL is configured as 0; clamping to 1s to avoid panic");
interval = Duration::from_secs(1);
}
tokio::spawn(async move {
let mut timer = tokio::time::interval(interval);
loop {
timer.tick().await;
info!("Starting scheduled capacity update");
let start = Instant::now();
let manager = manager.clone();
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())
})
.await;
if started {
debug!("Scheduled capacity refresh started in {:?}", start.elapsed());
} else {
debug!("Scheduled capacity refresh skipped because another refresh is already in progress");
}
}
});
}
// ============================================================================
// Tests
// ============================================================================
#[cfg(test)]
mod tests {
use super::*;
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;
#[test]
#[serial]
fn test_get_scheduled_update_interval() {
let interval = get_scheduled_update_interval();
assert_eq!(interval, Duration::from_secs(120));
}
#[test]
#[serial]
fn test_get_write_trigger_delay() {
let delay = get_write_trigger_delay();
assert_eq!(delay, Duration::from_secs(5));
}
#[test]
#[serial]
fn test_get_write_frequency_threshold() {
let threshold = get_write_frequency_threshold();
assert_eq!(threshold, 5);
}
#[test]
#[serial]
fn test_get_fast_update_threshold() {
let threshold = get_fast_update_threshold();
assert_eq!(threshold, Duration::from_secs(30));
}
#[test]
#[serial]
fn test_get_max_files_threshold() {
let threshold = get_max_files_threshold();
assert_eq!(threshold, 200_000);
}
#[test]
#[serial]
fn test_get_stat_timeout() {
let timeout = get_stat_timeout();
assert_eq!(timeout, Duration::from_secs(3));
}
#[test]
#[serial]
fn test_get_sample_rate() {
let rate = get_sample_rate();
assert_eq!(rate, 200);
}
#[test]
#[serial]
fn test_env_var_override_scheduled_interval() {
temp_env::with_var(ENV_CAPACITY_SCHEDULED_INTERVAL, Some("600"), || {
let interval = get_scheduled_update_interval();
assert_eq!(interval, Duration::from_secs(600));
});
}
#[test]
#[serial]
fn test_env_var_override_write_trigger_delay() {
temp_env::with_var(ENV_CAPACITY_WRITE_TRIGGER_DELAY, Some("20"), || {
let delay = get_write_trigger_delay();
assert_eq!(delay, Duration::from_secs(20));
});
}
#[test]
#[serial]
fn test_env_var_override_write_frequency_threshold() {
temp_env::with_var(ENV_CAPACITY_WRITE_FREQUENCY_THRESHOLD, Some("20"), || {
let threshold = get_write_frequency_threshold();
assert_eq!(threshold, 20);
});
}
#[test]
#[serial]
fn test_env_var_override_fast_update_threshold() {
temp_env::with_var(ENV_CAPACITY_FAST_UPDATE_THRESHOLD, Some("120"), || {
let threshold = get_fast_update_threshold();
assert_eq!(threshold, Duration::from_secs(120));
});
}
#[test]
#[serial]
fn test_env_var_override_max_files_threshold() {
temp_env::with_var(ENV_CAPACITY_MAX_FILES_THRESHOLD, Some("2000000"), || {
let threshold = get_max_files_threshold();
assert_eq!(threshold, 2_000_000);
});
}
#[test]
#[serial]
fn test_env_var_override_stat_timeout() {
temp_env::with_var(ENV_CAPACITY_STAT_TIMEOUT, Some("10"), || {
let timeout = get_stat_timeout();
assert_eq!(timeout, Duration::from_secs(10));
});
}
#[test]
#[serial]
fn test_env_var_override_sample_rate() {
temp_env::with_var(ENV_CAPACITY_SAMPLE_RATE, Some("200"), || {
let rate = get_sample_rate();
assert_eq!(rate, 200);
});
}
#[tokio::test]
#[serial]
async fn test_capacity_manager_creation() {
let config = HybridStrategyConfig::default();
let manager = HybridCapacityManager::new(config);
assert!(manager.get_capacity().await.is_none());
}
#[tokio::test]
#[serial]
async fn test_update_capacity() {
let manager = HybridCapacityManager::from_env();
manager
.update_capacity(CapacityUpdate::exact(1000, 0), DataSource::RealTime)
.await;
let cached = manager.get_capacity().await;
assert!(cached.is_some());
assert_eq!(cached.unwrap().total_used, 1000);
}
#[tokio::test]
#[serial]
async fn test_record_write_operation() {
let manager = HybridCapacityManager::from_env();
manager.record_write_operation().await;
let frequency = manager.get_write_frequency().await;
assert_eq!(frequency, 1);
}
#[tokio::test]
#[serial]
async fn test_needs_fast_update() {
let manager = HybridCapacityManager::from_env();
// No cache, should not need update
assert!(!manager.needs_fast_update().await);
// Update cache
manager
.update_capacity(CapacityUpdate::exact(1000, 0), DataSource::RealTime)
.await;
// Fresh cache, should not need update
assert!(!manager.needs_fast_update().await);
}
#[tokio::test]
#[serial]
async fn test_config_from_env() {
let config = HybridStrategyConfig::from_env();
// Check default values
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]
#[serial]
async fn test_config_from_env_with_override() {
temp_env::with_var(ENV_CAPACITY_SCHEDULED_INTERVAL, Some("600"), || {
let config = HybridStrategyConfig::from_env();
assert_eq!(config.scheduled_update_interval, Duration::from_secs(600));
});
}
}
rustfs/src/capacity/capacity_manager_test.rs
-267
View File
@@ -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);
}
}
rustfs/src/capacity/mod.rs
+6 -17
View File
@@ -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,
};
rustfs/src/capacity/service.rs
+238
View File
@@ -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<String>, drive_path: impl Into<String>) -> CapacityDiskRef {
CapacityDiskRef {
endpoint: endpoint.into(),
drive_path: drive_path.into(),
}
}
fn capacity_disk_refs(disks: &[rustfs_madmin::Disk]) -> Vec<CapacityDiskRef> {
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<CapacityDiskRef>,
allow_dirty_subset: bool,
) -> Result<capacity_manager::CapacityUpdate, String> {
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<capacity_manager::HybridCapacityManager>,
source: capacity_manager::DataSource,
disks: &[rustfs_madmin::Disk],
allow_dirty_subset: bool,
) -> Result<capacity_manager::CapacityUpdate, String> {
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<capacity_manager::HybridCapacityManager>,
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<uuid::Uuid>) {
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",
}
}
rustfs/src/capacity/write_trigger_test.rs
-87
View File
@@ -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;
}
}