use std::sync::Arc;
use std::time::Duration;
use tokio_util::sync::CancellationToken;
pub(crate) async fn start_auto_cleanup(
rate_limit_tracker: Arc<crate::proxy::rate_limit::RateLimitTracker>,
auto_cleanup_handle: &Arc<tokio::sync::Mutex<Option<tokio::task::JoinHandle<()>>>>,
cancel_token: &CancellationToken,
) {
let tracker = rate_limit_tracker;
let cancel = cancel_token.child_token();
let handle = tokio::spawn(async move {
let mut interval = tokio::time::interval(Duration::from_secs(15));
loop {
tokio::select! {
_ = cancel.cancelled() => {
tracing::info!("Auto-cleanup task received cancel signal");
break;
}
_ = interval.tick() => {
let cleaned = tracker.cleanup_expired();
if cleaned > 0 {
tracing::info!(
"Auto-cleanup: Removed {} expired rate limit record(s)",
cleaned
);
}
}
}
}
});
let mut guard = auto_cleanup_handle.lock().await;
if let Some(old) = guard.take() {
old.abort();
tracing::warn!("Aborted previous auto-cleanup task");
}
*guard = Some(handle);
tracing::info!("Rate limit auto-cleanup task started (interval: 15s)");
}
pub(crate) async fn graceful_shutdown(
cancel_token: &CancellationToken,
auto_cleanup_handle: &Arc<tokio::sync::Mutex<Option<tokio::task::JoinHandle<()>>>>,
timeout: Duration,
) {
tracing::info!("Initiating graceful shutdown of background tasks...");
cancel_token.cancel();
match tokio::time::timeout(timeout, abort_background_tasks(auto_cleanup_handle)).await {
Ok(_) => tracing::info!("All background tasks cleaned up gracefully"),
Err(_) => tracing::warn!(
"Graceful cleanup timed out after {:?}, tasks were force-aborted",
timeout
),
}
}
pub(crate) async fn abort_background_tasks(
auto_cleanup_handle: &Arc<tokio::sync::Mutex<Option<tokio::task::JoinHandle<()>>>>,
) {
abort_task(auto_cleanup_handle, "Auto-cleanup task").await;
}
async fn abort_task(
handle: &Arc<tokio::sync::Mutex<Option<tokio::task::JoinHandle<()>>>>,
task_name: &str,
) {
let Some(handle) = handle.lock().await.take() else {
return;
};
handle.abort();
match handle.await {
Ok(()) => tracing::debug!("{} completed", task_name),
Err(e) if e.is_cancelled() => tracing::info!("{} aborted", task_name),
Err(e) => tracing::warn!("{} error: {}", task_name, e),
}
}