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

This commit is contained in:
cxymds
2026-08-07 19:02:42 +08:00
committed by GitHub
parent bd15dd5784
commit 2a44985037
11 changed files with 293 additions and 17 deletions
Generated
+1
View File
@@ -9784,6 +9784,7 @@ dependencies = [
"temp-env",
"tempfile",
"tokio",
"tokio-util",
"tracing",
"uuid",
"walkdir",
+1
View File
@@ -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 }
+176 -8
View File
@@ -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() {
+26 -1
View File
@@ -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() {
+9 -1
View File
@@ -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<CapacityBackgroundTasks> {
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)> {
+2 -1
View File
@@ -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,
};
+13 -3
View File
@@ -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<uuid::Uuid>) {
}
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<CapacityBackgroundTasks> {
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)> {
+3
View File
@@ -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,
+4
View File
@@ -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<ServiceStateManager>,
pub(crate) s3_shutdown_tx: Option<ShutdownHandle>,
pub(crate) console_shutdown_tx: Option<ShutdownHandle>,
pub(crate) capacity_tasks: Option<CapacityBackgroundTasks>,
pub(crate) service_runtime: StartupServiceRuntime,
pub(crate) store: Arc<ECStore>,
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;
-2
View File
@@ -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<GlobalReadiness>,
server_ctx: Arc<ServerContextSlot>,
) -> Result<StartupHttpServers> {
init_capacity_management().await;
let state_manager = Arc::new(ServiceStateManager::new());
state_manager.update(ServiceState::Starting);
+58 -1
View File
@@ -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<ShutdownHandle>,
console_shutdown_handle: Option<ShutdownHandle>,
optional_runtimes: OptionalRuntimeServices,
capacity_tasks: Option<CapacityBackgroundTasks>,
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);