fix(capacity): stop background schedulers on shutdown (#5797)

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 }

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<CapacityUpdate, String>) -> bool
     matches!(result, Ok(update) if !update.timed_out && !update.degraded)
 }
 
-async fn run_scheduled_refresh_loop<F, Fut>(refresh_interval: Duration, mut refresh: F)
+async fn run_scheduled_refresh_loop<F, Fut>(refresh_interval: Duration, shutdown: CancellationToken, mut refresh: F)
 where
     F: FnMut() -> Fut,
     Fut: Future<Output = bool>,
 {
     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<CapacityDiskRef>) {
+    start_background_tasks(disks).await.detach();
+}
+
+/// Start capacity refresh and metrics schedulers with lifecycle ownership.
+pub async fn start_background_tasks(disks: Vec<CapacityDiskRef>) -> 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<CapacityDiskRef>) {
     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<CapacityDiskRef>) {
         .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() {

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<AtomicBool>);
+
+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<CapacityScanResult, std::io::Error> {
     let path = path.to_path_buf();
     let limits = ScanLimits::from_env();
@@ -655,11 +669,12 @@ async fn get_dir_size_async(path: &Path) -> Result<CapacityScanResult, std::io::
     let cancelled = Arc::new(AtomicBool::new(false));
     let scan_cancelled = cancelled.clone();
     let scan = tokio::task::spawn_blocking(move || scan_dir_blocking(&path, &limits, &scan_cancelled));
+    let cancel_on_drop = ScanCancellationGuard(cancelled);
 
     match tokio::time::timeout(budget, scan).await {
         Ok(join_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() {