use futures::future::BoxFuture;
use futures::FutureExt;
use std::sync::Arc;
use std::time::Duration;
use tokio::time::Instant;
use tokio_util::sync::CancellationToken;
use tracing::warn;
pub(crate) struct CleanupPass {
pub name: &'static str,
pub interval: Duration,
pub run: Box<dyn Fn() -> BoxFuture<'static, ()> + Send + Sync>,
}
pub(crate) fn spawn_cleanup_scheduler(
passes: Vec<CleanupPass>,
cleanup_lock: Arc<tokio::sync::Mutex<()>>,
cancel: CancellationToken,
) -> tokio::task::JoinHandle<()> {
tokio::spawn(async move {
let mut due: Vec<Instant> = vec![Instant::now(); passes.len()];
loop {
let Some(next) = due.iter().min().copied() else {
return;
};
tokio::select! {
_ = cancel.cancelled() => return,
_ = tokio::time::sleep_until(next) => {}
}
let guard = tokio::select! {
_ = cancel.cancelled() => return,
g = cleanup_lock.lock() => g,
};
let now = Instant::now();
for (i, pass) in passes.iter().enumerate() {
if due[i] > now {
continue;
}
due[i] = now + pass.interval;
if std::panic::AssertUnwindSafe((pass.run)())
.catch_unwind()
.await
.is_err()
{
warn!(
pass = pass.name,
"cleanup pass panicked — continuing with remaining passes"
);
}
}
drop(guard);
}
})
}
#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
use super::*;
type Recorder = Arc<parking_lot::Mutex<Vec<&'static str>>>;
fn recorder() -> Recorder {
Arc::new(parking_lot::Mutex::new(Vec::new()))
}
fn recording_pass(name: &'static str, interval: Duration, rec: &Recorder) -> CleanupPass {
let rec = rec.clone();
CleanupPass {
name,
interval,
run: Box::new(move || {
let rec = rec.clone();
async move {
rec.lock().push(name);
}
.boxed()
}),
}
}
fn panicking_pass(name: &'static str, interval: Duration, rec: &Recorder) -> CleanupPass {
let rec = rec.clone();
CleanupPass {
name,
interval,
run: Box::new(move || {
let rec = rec.clone();
async move {
rec.lock().push(name);
panic!("pass exploded");
}
.boxed()
}),
}
}
async fn wait_for(rec: &Recorder, count: usize) -> Vec<&'static str> {
let deadline = Instant::now() + Duration::from_secs(10);
loop {
let seen = rec.lock().clone();
if seen.len() >= count {
return seen;
}
assert!(Instant::now() < deadline, "only recorded {seen:?}");
tokio::time::sleep(Duration::from_millis(10)).await;
}
}
#[tokio::test]
async fn cycle_runs_passes_in_order_at_boot() {
let rec = recorder();
let interval = Duration::from_secs(10);
let cancel = CancellationToken::new();
let start = Instant::now();
let handle = spawn_cleanup_scheduler(
vec![
recording_pass("retention", interval, &rec),
recording_pass("gc", interval, &rec),
],
Arc::new(tokio::sync::Mutex::new(())),
cancel.clone(),
);
let seen = wait_for(&rec, 2).await;
assert_eq!(&seen[..2], &["retention", "gc"]);
assert!(
start.elapsed() < Duration::from_secs(2),
"boot cycle waited for the interval"
);
cancel.cancel();
handle.await.unwrap();
}
#[tokio::test]
async fn pass_panic_does_not_prevent_later_passes() {
let rec = recorder();
let interval = Duration::from_millis(100);
let cancel = CancellationToken::new();
let handle = spawn_cleanup_scheduler(
vec![
panicking_pass("retention", interval, &rec),
recording_pass("gc", interval, &rec),
],
Arc::new(tokio::sync::Mutex::new(())),
cancel.clone(),
);
let seen = wait_for(&rec, 4).await;
assert_eq!(&seen[..4], &["retention", "gc", "retention", "gc"]);
cancel.cancel();
handle.await.unwrap();
}
#[tokio::test]
async fn cycle_waits_for_cleanup_lock_instead_of_skipping() {
let rec = recorder();
let interval = Duration::from_millis(100);
let cleanup_lock = Arc::new(tokio::sync::Mutex::new(()));
let held = cleanup_lock.clone().lock_owned().await;
let cancel = CancellationToken::new();
let handle = spawn_cleanup_scheduler(
vec![
recording_pass("retention", interval, &rec),
recording_pass("gc", interval, &rec),
],
cleanup_lock,
cancel.clone(),
);
tokio::time::sleep(Duration::from_millis(300)).await;
assert!(rec.lock().is_empty(), "ran while the lock was held");
drop(held);
let seen = wait_for(&rec, 2).await;
assert_eq!(&seen[..2], &["retention", "gc"]);
cancel.cancel();
handle.await.unwrap();
}
#[tokio::test]
async fn cancel_while_parked_on_lock_stops_scheduler() {
let rec = recorder();
let cleanup_lock = Arc::new(tokio::sync::Mutex::new(()));
let held = cleanup_lock.clone().lock_owned().await;
let cancel = CancellationToken::new();
let handle = spawn_cleanup_scheduler(
vec![recording_pass(
"retention",
Duration::from_millis(100),
&rec,
)],
cleanup_lock,
cancel.clone(),
);
tokio::time::sleep(Duration::from_millis(300)).await;
cancel.cancel();
tokio::time::timeout(Duration::from_secs(5), handle)
.await
.expect("scheduler did not stop when cancelled while parked on the lock")
.unwrap();
assert!(rec.lock().is_empty());
drop(held);
}
}