diff --git a/Cargo.lock b/Cargo.lock index 92a3ec9a9..5577cff08 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9784,6 +9784,7 @@ dependencies = [ "temp-env", "tempfile", "tokio", + "tokio-util", "tracing", "uuid", "walkdir", diff --git a/crates/object-capacity/Cargo.toml b/crates/object-capacity/Cargo.toml index fd17815f4..a1eea7740 100644 --- a/crates/object-capacity/Cargo.toml +++ b/crates/object-capacity/Cargo.toml @@ -66,6 +66,7 @@ rustfs-io-metrics = { workspace = true } rustfs-utils = { workspace = true, features = ["os"] } futures = { workspace = true } tokio = { workspace = true, features = ["sync", "time", "fs", "rt-multi-thread"] } +tokio-util = { workspace = true } tracing = { workspace = true } uuid = { workspace = true, features = ["v4", "fast-rng", "macro-diagnostics"] } walkdir = { workspace = true } diff --git a/crates/object-capacity/src/capacity_manager.rs b/crates/object-capacity/src/capacity_manager.rs index b7a7a1386..231d3c6a0 100644 --- a/crates/object-capacity/src/capacity_manager.rs +++ b/crates/object-capacity/src/capacity_manager.rs @@ -40,6 +40,8 @@ use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, Instant}; use tokio::sync::{Mutex, RwLock, watch}; +use tokio::task::{JoinError, JoinSet}; +use tokio_util::sync::CancellationToken; use tracing::{debug, info, warn}; const LOG_COMPONENT_CAPACITY: &str = "capacity"; @@ -1359,19 +1361,63 @@ fn scheduled_refresh_was_clean(result: &Result) -> bool matches!(result, Ok(update) if !update.timed_out && !update.degraded) } -async fn run_scheduled_refresh_loop(refresh_interval: Duration, mut refresh: F) +async fn run_scheduled_refresh_loop(refresh_interval: Duration, shutdown: CancellationToken, mut refresh: F) where F: FnMut() -> Fut, Fut: Future, { let mut backoff = ScheduledRefreshBackoff::new(refresh_interval); loop { - tokio::time::sleep(backoff.delay()).await; + tokio::select! { + biased; + _ = shutdown.cancelled() => break, + _ = tokio::time::sleep(backoff.delay()) => {} + } + // Widen the cross-thread cancellation window deterministically in tests. + #[cfg(test)] + tokio::task::yield_now().await; + if shutdown.is_cancelled() { + break; + } backoff.record_result(refresh().await); } } +/// Owned capacity scheduler tasks for one server runtime. +#[must_use = "capacity background tasks stop when their lifecycle handle is dropped"] +pub struct CapacityBackgroundTasks { + shutdown: CancellationToken, + tasks: JoinSet<()>, +} + +impl CapacityBackgroundTasks { + /// Leave the schedulers running without lifecycle ownership. + pub fn detach(mut self) { + self.tasks.detach_all(); + } + + /// Cancel the schedulers and wait for them and any active scheduled refresh. + pub async fn shutdown(mut self) -> Result<(), JoinError> { + self.shutdown.cancel(); + let mut join_error = None; + while let Some(result) = self.tasks.join_next().await { + if let Err(err) = result + && join_error.is_none() + { + join_error = Some(err); + } + } + join_error.map_or(Ok(()), Err) + } +} + +/// Start capacity refresh and metrics schedulers without lifecycle ownership. pub async fn start_background_task(disks: Vec) { + start_background_tasks(disks).await.detach(); +} + +/// Start capacity refresh and metrics schedulers with lifecycle ownership. +pub async fn start_background_tasks(disks: Vec) -> CapacityBackgroundTasks { let manager = get_capacity_manager(); let manager_for_refresh = manager.clone(); let manager_for_metrics = manager.clone(); @@ -1381,8 +1427,13 @@ pub async fn start_background_task(disks: Vec) { refresh_interval = clamp_background_interval(refresh_interval, ENV_CAPACITY_SCHEDULED_INTERVAL); metrics_interval = clamp_background_interval(metrics_interval, ENV_CAPACITY_METRICS_INTERVAL); - tokio::spawn(async move { - run_scheduled_refresh_loop(refresh_interval, move || { + let shutdown = CancellationToken::new(); + let refresh_shutdown = shutdown.clone(); + let metrics_shutdown = shutdown.clone(); + let mut tasks = JoinSet::new(); + + tasks.spawn(async move { + run_scheduled_refresh_loop(refresh_interval, refresh_shutdown, move || { let start = Instant::now(); let manager = manager_for_refresh.clone(); let disks = disks.clone(); @@ -1410,13 +1461,19 @@ pub async fn start_background_task(disks: Vec) { .await; }); - tokio::spawn(async move { + tasks.spawn(async move { let mut timer = tokio::time::interval_at(tokio::time::Instant::now() + metrics_interval, metrics_interval); loop { - timer.tick().await; + tokio::select! { + biased; + _ = metrics_shutdown.cancelled() => break, + _ = timer.tick() => {} + } manager_for_metrics.log_runtime_summary().await; } }); + + CapacityBackgroundTasks { shutdown, tasks } } // ============================================================================ @@ -1484,7 +1541,8 @@ mod tests { let results = Arc::new(StdMutex::new(VecDeque::from([false, false, true, true]))); let task_calls = calls.clone(); let task_results = results.clone(); - let task = tokio::spawn(run_scheduled_refresh_loop(Duration::from_secs(10), move || { + let shutdown = CancellationToken::new(); + let task = tokio::spawn(run_scheduled_refresh_loop(Duration::from_secs(10), shutdown.clone(), move || { let task_calls = task_calls.clone(); let task_results = task_results.clone(); async move { @@ -1500,23 +1558,77 @@ mod tests { tokio::task::yield_now().await; tokio::time::advance(Duration::from_secs(10)).await; tokio::task::yield_now().await; + tokio::task::yield_now().await; assert_eq!(calls.load(Ordering::SeqCst), 1); tokio::time::advance(Duration::from_secs(19)).await; tokio::task::yield_now().await; + tokio::task::yield_now().await; assert_eq!(calls.load(Ordering::SeqCst), 1); tokio::time::advance(Duration::from_secs(1)).await; tokio::task::yield_now().await; + tokio::task::yield_now().await; assert_eq!(calls.load(Ordering::SeqCst), 2); tokio::time::advance(Duration::from_secs(40)).await; tokio::task::yield_now().await; + tokio::task::yield_now().await; assert_eq!(calls.load(Ordering::SeqCst), 3); tokio::time::advance(Duration::from_secs(10)).await; tokio::task::yield_now().await; + tokio::task::yield_now().await; assert_eq!(calls.load(Ordering::SeqCst), 4); - task.abort(); + shutdown.cancel(); + task.await.expect("scheduled refresh loop should stop cleanly"); + } + + #[tokio::test(start_paused = true)] + async fn scheduled_refresh_shutdown_waits_for_active_refresh() { + let shutdown = CancellationToken::new(); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let mut started_tx = Some(started_tx); + let mut release_rx = Some(release_rx); + let task_shutdown = shutdown.clone(); + let task = tokio::spawn(run_scheduled_refresh_loop(Duration::from_secs(1), task_shutdown, move || { + let started_tx = started_tx.take().expect("test runs one refresh"); + let release_rx = release_rx.take().expect("test runs one refresh"); + async move { + started_tx.send(()).expect("test should observe refresh start"); + release_rx.await.expect("test should release active refresh"); + true + } + })); + + tokio::task::yield_now().await; + tokio::time::advance(Duration::from_secs(1)).await; + started_rx.await.expect("scheduled refresh should start"); + shutdown.cancel(); + assert!(!task.is_finished()); + + release_tx.send(()).expect("active refresh should still be running"); + task.await.expect("scheduled refresh loop should stop cleanly"); + } + + #[tokio::test(start_paused = true)] + async fn scheduled_refresh_shutdown_at_due_boundary_skips_refresh() { + let shutdown = CancellationToken::new(); + let calls = Arc::new(AtomicUsize::new(0)); + let task_calls = calls.clone(); + let task_shutdown = shutdown.clone(); + let task = tokio::spawn(run_scheduled_refresh_loop(Duration::from_secs(1), task_shutdown, move || { + task_calls.fetch_add(1, Ordering::SeqCst); + async { true } + })); + + tokio::task::yield_now().await; + tokio::time::advance(Duration::from_secs(1)).await; + tokio::task::yield_now().await; + shutdown.cancel(); + task.await.expect("scheduled refresh loop should stop cleanly"); + + assert_eq!(calls.load(Ordering::SeqCst), 0); } type ConfigGetterCase = (&'static str, fn() -> u64, u64, &'static str, u64); @@ -1660,6 +1772,62 @@ mod tests { let _ = tokio::time::Instant::now() + MAX_BACKGROUND_INTERVAL; } + #[tokio::test] + async fn background_tasks_shutdown_cancels_and_joins_schedulers() { + let tasks = start_background_tasks(Vec::new()).await; + + tokio::time::timeout(Duration::from_secs(1), tasks.shutdown()) + .await + .expect("capacity background tasks should stop promptly") + .expect("capacity background tasks should join cleanly"); + } + + #[tokio::test] + async fn background_tasks_shutdown_joins_remaining_tasks_after_failure() { + let shutdown = CancellationToken::new(); + let mut tasks = JoinSet::new(); + let failed_task = tasks.spawn(async { panic!("expected scheduler failure") }); + while !failed_task.is_finished() { + tokio::task::yield_now().await; + } + + let task_shutdown = shutdown.clone(); + let (ready_tx, ready_rx) = tokio::sync::oneshot::channel(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + tasks.spawn(async move { + task_shutdown.cancelled().await; + ready_tx.send(()).expect("test should observe scheduler cancellation"); + release_rx.await.expect("test should release scheduler shutdown"); + }); + let release_task = tokio::spawn(async move { + ready_rx.await.expect("scheduler should observe cancellation"); + release_tx.send(()).expect("scheduler should still be joinable"); + }); + + let background_tasks = CapacityBackgroundTasks { shutdown, tasks }; + assert!(background_tasks.shutdown().await.is_err()); + release_task.await.expect("scheduler release task should join"); + } + + #[tokio::test] + async fn background_tasks_detach_keeps_schedulers_running() { + let shutdown = CancellationToken::new(); + let mut tasks = JoinSet::new(); + let (ready_tx, ready_rx) = tokio::sync::oneshot::channel(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let (completed_tx, completed_rx) = tokio::sync::oneshot::channel(); + tasks.spawn(async move { + ready_tx.send(()).expect("test should observe scheduler start"); + release_rx.await.expect("test should release detached scheduler"); + completed_tx.send(()).expect("test should observe scheduler completion"); + }); + + CapacityBackgroundTasks { shutdown, tasks }.detach(); + ready_rx.await.expect("detached scheduler should start"); + release_tx.send(()).expect("detached scheduler should still be running"); + completed_rx.await.expect("detached scheduler should complete"); + } + #[test] #[serial] fn test_recent_write_count_ignores_future_buckets() { diff --git a/crates/object-capacity/src/scan.rs b/crates/object-capacity/src/scan.rs index 267b27ad2..bf974cda6 100644 --- a/crates/object-capacity/src/scan.rs +++ b/crates/object-capacity/src/scan.rs @@ -643,6 +643,20 @@ fn outer_scan_budget(limits: &ScanLimits) -> Duration { limits.max_timeout.saturating_mul(2).max(Duration::from_secs(5)) } +struct ScanCancellationGuard(Arc); + +impl ScanCancellationGuard { + fn cancel(&self) { + self.0.store(true, Ordering::Relaxed); + } +} + +impl Drop for ScanCancellationGuard { + fn drop(&mut self) { + self.cancel(); + } +} + async fn get_dir_size_async(path: &Path) -> Result { let path = path.to_path_buf(); let limits = ScanLimits::from_env(); @@ -655,11 +669,12 @@ async fn get_dir_size_async(path: &Path) -> Result join_result.map_err(std::io::Error::other)?, Err(_) => { - cancelled.store(true, Ordering::Relaxed); + cancel_on_drop.cancel(); warn!( event = EVENT_CAPACITY_SCAN_HARD_TIMEOUT, component = LOG_COMPONENT_CAPACITY, @@ -1456,6 +1471,16 @@ mod tests { assert_eq!(outer_scan_budget(&tight_limits(Duration::ZERO)), Duration::from_secs(5)); } + #[test] + fn test_scan_cancellation_guard_sets_flag_on_drop() { + let cancelled = Arc::new(AtomicBool::new(false)); + let guard = ScanCancellationGuard(cancelled.clone()); + + drop(guard); + + assert!(cancelled.load(Ordering::Relaxed)); + } + #[cfg(unix)] #[test] fn test_scan_dir_blocking_resolves_symlink_root() { diff --git a/rustfs/src/capacity/capacity_integration.rs b/rustfs/src/capacity/capacity_integration.rs index e161736b7..9389049bf 100644 --- a/rustfs/src/capacity/capacity_integration.rs +++ b/rustfs/src/capacity/capacity_integration.rs @@ -14,7 +14,10 @@ //! Capacity management integration for application startup -use crate::capacity::{get_cached_capacity_with_metrics, init_capacity_management_for_local_disks}; +use crate::capacity::{ + get_cached_capacity_with_metrics, init_capacity_management_for_local_disks, init_capacity_management_for_local_disks_managed, +}; +use rustfs_object_capacity::capacity_manager::CapacityBackgroundTasks; /// Initialize capacity management system /// This should be called during application startup after local disks are initialized @@ -22,6 +25,11 @@ pub async fn init_capacity_management() { init_capacity_management_for_local_disks().await; } +/// Initialize capacity management with lifecycle ownership. +pub async fn init_capacity_management_managed() -> Option { + init_capacity_management_for_local_disks_managed().await +} + /// Get capacity statistics with metrics #[allow(dead_code)] pub async fn get_capacity_with_metrics() -> Option<(u64, String)> { diff --git a/rustfs/src/capacity/mod.rs b/rustfs/src/capacity/mod.rs index 400322f69..de485895d 100644 --- a/rustfs/src/capacity/mod.rs +++ b/rustfs/src/capacity/mod.rs @@ -53,5 +53,6 @@ pub mod capacity_integration; pub mod service; pub use service::{ - capacity_disk_ref, get_cached_capacity_with_metrics, init_capacity_management_for_local_disks, record_capacity_write, + capacity_disk_ref, get_cached_capacity_with_metrics, init_capacity_management_for_local_disks, + init_capacity_management_for_local_disks_managed, record_capacity_write, }; diff --git a/rustfs/src/capacity/service.rs b/rustfs/src/capacity/service.rs index 91d21ac7a..aa81faa4f 100644 --- a/rustfs/src/capacity/service.rs +++ b/rustfs/src/capacity/service.rs @@ -14,7 +14,10 @@ use crate::storage_api::capacity::service::{all_local_disk, disk_drive_path, disk_endpoint}; use rustfs_io_metrics::capacity_metrics::{record_capacity_cache_hit, record_capacity_cache_miss}; -use rustfs_object_capacity::{CapacityDiskRef, capacity_manager}; +use rustfs_object_capacity::{ + CapacityDiskRef, + capacity_manager::{self, CapacityBackgroundTasks}, +}; use tracing::{info, warn}; const LOG_COMPONENT_CAPACITY: &str = "capacity"; @@ -34,6 +37,12 @@ pub async fn record_capacity_write(scope_token: Option) { } pub async fn init_capacity_management_for_local_disks() { + if let Some(tasks) = init_capacity_management_for_local_disks_managed().await { + tasks.detach(); + } +} + +pub async fn init_capacity_management_for_local_disks_managed() -> Option { info!( component = LOG_COMPONENT_CAPACITY, subsystem = LOG_SUBSYSTEM_CAPACITY, @@ -52,7 +61,7 @@ pub async fn init_capacity_management_for_local_disks() { reason = "no_local_disks", "Capacity manager state changed" ); - return; + return None; } info!( @@ -75,7 +84,7 @@ pub async fn init_capacity_management_for_local_disks() { state = "starting_background_task", "Capacity manager state changed" ); - capacity_manager::start_background_task(disk_refs).await; + let tasks = capacity_manager::start_background_tasks(disk_refs).await; info!( component = LOG_COMPONENT_CAPACITY, @@ -84,6 +93,7 @@ pub async fn init_capacity_management_for_local_disks() { state = "initialized", "Capacity manager state changed" ); + Some(tasks) } pub async fn get_cached_capacity_with_metrics() -> Option<(u64, &'static str)> { diff --git a/rustfs/src/startup_entrypoint.rs b/rustfs/src/startup_entrypoint.rs index 26b5c73a9..a4b7d0c3d 100644 --- a/rustfs/src/startup_entrypoint.rs +++ b/rustfs/src/startup_entrypoint.rs @@ -144,6 +144,8 @@ async fn run(config: Config) -> Result<()> { shutdown_token: ctx, } = init_startup_storage_runtime(server_addr, &endpoint_pools, readiness.clone(), instance_ctx).await?; + let capacity_tasks = crate::capacity::capacity_integration::init_capacity_management_managed().await; + let service_runtime = init_startup_runtime_services( &config, endpoint_pools, @@ -160,6 +162,7 @@ async fn run(config: Config) -> Result<()> { state_manager, s3_shutdown_tx, console_shutdown_tx, + capacity_tasks, service_runtime, store, shutdown_token: ctx, diff --git a/rustfs/src/startup_lifecycle.rs b/rustfs/src/startup_lifecycle.rs index 47abbf248..c4d09872b 100644 --- a/rustfs/src/startup_lifecycle.rs +++ b/rustfs/src/startup_lifecycle.rs @@ -21,6 +21,7 @@ use crate::{ startup_shutdown::run_startup_shutdown_sequence, }; use rustfs_common::GlobalReadiness; +use rustfs_object_capacity::capacity_manager::CapacityBackgroundTasks; use rustfs_scanner::init_data_scanner; use std::{ io::{Error, Result}, @@ -106,6 +107,7 @@ pub(crate) struct StartupRuntimeLifecycle { pub(crate) state_manager: Arc, pub(crate) s3_shutdown_tx: Option, pub(crate) console_shutdown_tx: Option, + pub(crate) capacity_tasks: Option, pub(crate) service_runtime: StartupServiceRuntime, pub(crate) store: Arc, pub(crate) shutdown_token: CancellationToken, @@ -118,6 +120,7 @@ pub(crate) async fn run_startup_runtime_lifecycle(lifecycle: StartupRuntimeLifec state_manager, s3_shutdown_tx, console_shutdown_tx, + capacity_tasks, service_runtime, store, shutdown_token, @@ -155,6 +158,7 @@ pub(crate) async fn run_startup_runtime_lifecycle(lifecycle: StartupRuntimeLifec s3_shutdown_tx, console_shutdown_tx, optional_runtimes, + capacity_tasks, shutdown_token, ) .await; diff --git a/rustfs/src/startup_server.rs b/rustfs/src/startup_server.rs index 1daa15d6c..3f637f235 100644 --- a/rustfs/src/startup_server.rs +++ b/rustfs/src/startup_server.rs @@ -13,7 +13,6 @@ // limitations under the License. use crate::{ - capacity::capacity_integration::init_capacity_management, config::Config, server::{ServiceState, ServiceStateManager, ShutdownHandle, start_http_server}, startup_runtime_sources, @@ -259,7 +258,6 @@ pub(crate) async fn init_startup_http_servers( readiness: Arc, server_ctx: Arc, ) -> Result { - init_capacity_management().await; let state_manager = Arc::new(ServiceStateManager::new()); state_manager.update(ServiceState::Starting); diff --git a/rustfs/src/startup_shutdown.rs b/rustfs/src/startup_shutdown.rs index 2830502e5..70b1a0701 100644 --- a/rustfs/src/startup_shutdown.rs +++ b/rustfs/src/startup_shutdown.rs @@ -14,7 +14,10 @@ use crate::storage_api::startup::shutdown::shutdown_background_services; use crate::{ - server::{ServiceState, ServiceStateManager, ShutdownHandle, ShutdownSignal, shutdown_event_notifier, stop_audit_system}, + server::{ + SHUTDOWN_TIMEOUT, ServiceState, ServiceStateManager, ShutdownHandle, ShutdownSignal, shutdown_event_notifier, + stop_audit_system, + }, startup_optional_runtime_sidecars::{ OptionalRuntimeServices, prepare_optional_runtime_shutdowns, shutdown_optional_runtime_services, }, @@ -22,6 +25,7 @@ use crate::{ }; use rustfs_heal::shutdown_ahm_services; use rustfs_notify::NotificationLifecycleTransition; +use rustfs_object_capacity::capacity_manager::CapacityBackgroundTasks; use rustfs_utils::get_env_bool_with_aliases; use std::future::Future; use std::path::PathBuf; @@ -42,6 +46,7 @@ const EVENT_EMBEDDED_SERVER_STATE: &str = "embedded_server_state"; const EVENT_EMBEDDED_SHUTDOWN_CLEANUP_FAILED: &str = "embedded_shutdown_cleanup_failed"; const EVENT_SHUTDOWN_SIGNAL_RECEIVED: &str = "shutdown_signal_received"; const EVENT_BACKGROUND_SERVICE_SHUTDOWN: &str = "background_service_shutdown"; +const BACKGROUND_SERVICE_CAPACITY: &str = "capacity"; const EVENT_EVENT_NOTIFIER_SHUTDOWN: &str = "event_notifier_shutdown"; const EVENT_PROFILING_SHUTDOWN: &str = "profiling_shutdown"; const EVENT_SERVER_SHUTDOWN_STATE: &str = "server_shutdown_state"; @@ -200,6 +205,7 @@ pub(crate) async fn run_startup_shutdown_sequence( s3_shutdown_handle: Option, console_shutdown_handle: Option, optional_runtimes: OptionalRuntimeServices, + capacity_tasks: Option, ctx: CancellationToken, ) { ctx.cancel(); @@ -220,6 +226,57 @@ pub(crate) async fn run_startup_shutdown_sequence( ); state_manager.update(ServiceState::Stopping); + if let Some(handle) = &s3_shutdown_handle { + handle.signal(); + } + if let Some(handle) = &console_shutdown_handle { + handle.signal(); + } + + if let Some(capacity_tasks) = capacity_tasks { + info!( + target: "rustfs::main::handle_shutdown", + event = EVENT_BACKGROUND_SERVICE_SHUTDOWN, + component = LOG_COMPONENT_MAIN, + subsystem = LOG_SUBSYSTEM_STARTUP, + service = BACKGROUND_SERVICE_CAPACITY, + state = "stopping", + "Background service shutdown started" + ); + match tokio::time::timeout(SHUTDOWN_TIMEOUT, capacity_tasks.shutdown()).await { + Ok(Ok(())) => info!( + target: "rustfs::main::handle_shutdown", + event = EVENT_BACKGROUND_SERVICE_SHUTDOWN, + component = LOG_COMPONENT_MAIN, + subsystem = LOG_SUBSYSTEM_STARTUP, + service = BACKGROUND_SERVICE_CAPACITY, + state = "stopped", + "Background service shutdown completed" + ), + Ok(Err(err)) => error!( + target: "rustfs::main::handle_shutdown", + event = EVENT_BACKGROUND_SERVICE_SHUTDOWN, + component = LOG_COMPONENT_MAIN, + subsystem = LOG_SUBSYSTEM_STARTUP, + service = BACKGROUND_SERVICE_CAPACITY, + state = "stop_failed", + reason = join_failure_reason(&err), + "Background service shutdown failed" + ), + Err(_) => error!( + target: "rustfs::main::handle_shutdown", + event = EVENT_BACKGROUND_SERVICE_SHUTDOWN, + component = LOG_COMPONENT_MAIN, + subsystem = LOG_SUBSYSTEM_STARTUP, + service = BACKGROUND_SERVICE_CAPACITY, + state = "stop_failed", + reason = "timeout", + timeout_secs = SHUTDOWN_TIMEOUT.as_secs(), + "Background service shutdown timed out" + ), + } + } + let enable_scanner = get_env_bool_with_aliases(ENV_SCANNER_ENABLED, &[ENV_SCANNER_ENABLED_DEPRECATED], true); let enable_heal = get_env_bool_with_aliases(ENV_HEAL_ENABLED, &[ENV_HEAL_ENABLED_DEPRECATED], true);