use std::sync::Arc;
use std::time::Duration;
use crate::error::{Error, Result};
use crate::handler::IHostedService;
pub use tokio_util::sync::CancellationToken;
pub const STOP_GRACE: Duration = Duration::from_secs(5);
#[async_trait::async_trait]
pub trait BackgroundService: Send + Sync + 'static {
async fn execute(&self, ct: CancellationToken) -> Result<()>;
}
pub async fn delay(ct: &CancellationToken, duration: Duration) -> bool {
tokio::select! {
_ = ct.cancelled() => false,
_ = tokio::time::sleep(duration) => true,
}
}
#[async_trait::async_trait]
pub trait PeriodicBackgroundService: Send + Sync + 'static {
fn interval(&self) -> Duration;
async fn tick(&self) -> Result<()>;
fn stop_on_error(&self) -> bool {
false
}
}
#[async_trait::async_trait]
impl<T: PeriodicBackgroundService + ?Sized> BackgroundService for T {
async fn execute(&self, ct: CancellationToken) -> Result<()> {
let interval = self.interval();
if interval.is_zero() {
return Err(Error::Internal(
"periodic background service interval must be greater than zero".to_string(),
));
}
while delay(&ct, interval).await {
if let Err(err) = self.tick().await {
tracing::error!(
error = %err,
interval_secs = interval.as_secs_f64(),
"periodic background service tick failed"
);
if self.stop_on_error() {
return Err(err);
}
}
}
Ok(())
}
}
pub struct BackgroundServiceHost {
service: Arc<dyn BackgroundService>,
run: tokio::sync::Mutex<Option<Running>>,
}
struct Running {
stop: CancellationToken,
handle: tokio::task::JoinHandle<Result<()>>,
}
impl BackgroundServiceHost {
pub fn new(service: Arc<dyn BackgroundService>) -> Self {
Self {
service,
run: tokio::sync::Mutex::new(None),
}
}
pub(crate) async fn stop_with_grace(&self, grace: Duration) -> Result<()> {
let Some(running) = self.run.lock().await.take() else {
return Ok(());
};
running.stop.cancel();
let mut handle = running.handle;
match tokio::time::timeout(grace, &mut handle).await {
Ok(Ok(result)) => result,
Ok(Err(join_error)) if join_error.is_cancelled() => Ok(()),
Ok(Err(join_error)) => Err(Error::Internal(format!(
"background service panicked: {join_error}"
))),
Err(_) => {
handle.abort();
Err(Error::Internal(format!(
"background service did not observe cancellation within {}s and was aborted",
grace.as_secs()
)))
}
}
}
}
#[async_trait::async_trait]
impl IHostedService for BackgroundServiceHost {
async fn start(&self) -> Result<()> {
let mut slot = self.run.lock().await;
if slot.is_some() {
return Err(Error::State(
"background service is already running".to_string(),
));
}
let stop = CancellationToken::new();
let service = Arc::clone(&self.service);
let token = stop.clone();
let handle = tokio::spawn(async move {
let result = service.execute(token).await;
if let Err(err) = &result {
tracing::error!(error = %err, "background service exited with an error");
}
result
});
*slot = Some(Running { stop, handle });
Ok(())
}
async fn stop(&self) -> Result<()> {
self.stop_with_grace(STOP_GRACE).await
}
}
pub(crate) struct OrderedHostedService<T: ?Sized> {
order: i32,
inner: Arc<T>,
}
impl<T: IHostedService + ?Sized> OrderedHostedService<T> {
pub(crate) fn new(order: i32, inner: Arc<T>) -> Self {
Self { order, inner }
}
}
#[async_trait::async_trait]
impl<T: IHostedService + ?Sized> IHostedService for OrderedHostedService<T> {
async fn start(&self) -> Result<()> {
self.inner.start().await
}
async fn stop(&self) -> Result<()> {
self.inner.stop().await
}
fn order(&self) -> i32 {
self.order
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use std::time::Instant;
struct Idle;
#[async_trait::async_trait]
impl BackgroundService for Idle {
async fn execute(&self, _ct: CancellationToken) -> Result<()> {
Ok(())
}
}
#[tokio::test]
async fn delay_is_interrupted_by_cancellation() {
let ct = CancellationToken::new();
let cancel = ct.clone();
tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(50)).await;
cancel.cancel();
});
let started = Instant::now();
let slept = delay(&ct, Duration::from_secs(60)).await;
assert!(!slept, "a cancelled delay must report cancellation");
assert!(
started.elapsed() < Duration::from_millis(100),
"cancellation must interrupt the sleep within 100ms, took {:?}",
started.elapsed()
);
}
#[tokio::test]
async fn delay_waits_the_full_duration_without_cancellation() {
let ct = CancellationToken::new();
assert!(delay(&ct, Duration::from_millis(10)).await);
}
struct Counting(Arc<AtomicBool>);
#[async_trait::async_trait]
impl BackgroundService for Counting {
async fn execute(&self, ct: CancellationToken) -> Result<()> {
while delay(&ct, Duration::from_secs(60)).await {
self.0.store(true, Ordering::SeqCst);
}
Ok(())
}
}
#[tokio::test]
async fn stop_cancels_the_worker_and_joins_it() {
let stopped = Arc::new(AtomicBool::new(false));
let host = BackgroundServiceHost::new(Arc::new(Counting(Arc::clone(&stopped))));
host.start().await.expect("start must succeed");
tokio::time::sleep(Duration::from_millis(20)).await;
let started = Instant::now();
host.stop().await.expect("stop must succeed");
assert!(
started.elapsed() < Duration::from_millis(500),
"stop must interrupt the worker's sleep"
);
assert!(!stopped.load(Ordering::SeqCst), "the loop never ticked");
host.stop().await.expect("second stop must be a no-op");
}
struct Failing;
#[async_trait::async_trait]
impl BackgroundService for Failing {
async fn execute(&self, _ct: CancellationToken) -> Result<()> {
Err(Error::Internal("worker exploded".to_string()))
}
}
#[tokio::test]
async fn worker_failure_is_reported_as_a_typed_error_from_stop() {
let host = BackgroundServiceHost::new(Arc::new(Failing));
host.start().await.expect("start must succeed");
tokio::time::sleep(Duration::from_millis(20)).await;
let err = host
.stop()
.await
.expect_err("failure must not be swallowed");
assert!(err.to_string().contains("worker exploded"), "{err}");
}
struct Panicking;
#[async_trait::async_trait]
impl BackgroundService for Panicking {
async fn execute(&self, _ct: CancellationToken) -> Result<()> {
panic!("boom");
}
}
#[tokio::test]
async fn worker_panic_becomes_a_typed_error_instead_of_propagating() {
let host = BackgroundServiceHost::new(Arc::new(Panicking));
host.start().await.expect("start must succeed");
tokio::time::sleep(Duration::from_millis(20)).await;
let err = host.stop().await.expect_err("panic must be a typed error");
assert!(
err.to_string().contains("background service panicked"),
"{err}"
);
}
struct Stubborn;
#[async_trait::async_trait]
impl BackgroundService for Stubborn {
async fn execute(&self, _ct: CancellationToken) -> Result<()> {
tokio::time::sleep(Duration::from_secs(600)).await;
Ok(())
}
}
#[tokio::test]
async fn a_worker_that_ignores_cancellation_is_aborted() {
let host = BackgroundServiceHost::new(Arc::new(Stubborn));
host.start().await.expect("start must succeed");
let result = tokio::time::timeout(Duration::from_secs(10), host.stop()).await;
let err = result
.expect("stop must not hang forever")
.expect_err("aborted worker is reported as an error");
assert!(err.to_string().contains("aborted"), "{err}");
}
#[tokio::test]
async fn idle_worker_stops_cleanly() {
let host = BackgroundServiceHost::new(Arc::new(Idle));
host.start().await.expect("start must succeed");
tokio::time::sleep(Duration::from_millis(20)).await;
host.stop().await.expect("a finished worker stops cleanly");
}
struct Ticker {
ticks: Arc<AtomicU32>,
}
#[async_trait::async_trait]
impl PeriodicBackgroundService for Ticker {
fn interval(&self) -> Duration {
Duration::from_millis(10)
}
async fn tick(&self) -> Result<()> {
self.ticks.fetch_add(1, Ordering::SeqCst);
Ok(())
}
}
#[tokio::test]
async fn periodic_worker_ticks_until_stopped() {
let ticks = Arc::new(AtomicU32::new(0));
let host = BackgroundServiceHost::new(Arc::new(Ticker {
ticks: Arc::clone(&ticks),
}));
host.start().await.expect("start must succeed");
let deadline = Instant::now() + Duration::from_secs(2);
while ticks.load(Ordering::SeqCst) < 2 && Instant::now() < deadline {
tokio::time::sleep(Duration::from_millis(5)).await;
}
assert!(
ticks.load(Ordering::SeqCst) >= 2,
"the periodic worker must tick, saw {}",
ticks.load(Ordering::SeqCst)
);
let started = Instant::now();
host.stop().await.expect("stop must succeed");
assert!(
started.elapsed() < Duration::from_millis(500),
"stop must not wait for the next interval"
);
let after_stop = ticks.load(Ordering::SeqCst);
tokio::time::sleep(Duration::from_millis(50)).await;
assert_eq!(ticks.load(Ordering::SeqCst), after_stop);
}
struct LongInterval(Arc<AtomicU32>);
#[async_trait::async_trait]
impl PeriodicBackgroundService for LongInterval {
fn interval(&self) -> Duration {
Duration::from_secs(60)
}
async fn tick(&self) -> Result<()> {
self.0.fetch_add(1, Ordering::SeqCst);
Ok(())
}
}
#[tokio::test]
async fn periodic_worker_shutdown_interrupts_the_interval_wait() {
let ticks = Arc::new(AtomicU32::new(0));
let host = BackgroundServiceHost::new(Arc::new(LongInterval(Arc::clone(&ticks))));
host.start().await.expect("start must succeed");
tokio::time::sleep(Duration::from_millis(30)).await;
assert_eq!(
ticks.load(Ordering::SeqCst),
0,
"first tick is one interval away"
);
let started = Instant::now();
host.stop().await.expect("stop must succeed");
assert!(
started.elapsed() < Duration::from_millis(500),
"the 60s interval must not delay shutdown, took {:?}",
started.elapsed()
);
assert_eq!(ticks.load(Ordering::SeqCst), 0);
}
struct FailingTicker {
ticks: Arc<AtomicU32>,
fatal: bool,
}
#[async_trait::async_trait]
impl PeriodicBackgroundService for FailingTicker {
fn interval(&self) -> Duration {
Duration::from_millis(5)
}
async fn tick(&self) -> Result<()> {
self.ticks.fetch_add(1, Ordering::SeqCst);
Err(Error::Internal("tick failed".to_string()))
}
fn stop_on_error(&self) -> bool {
self.fatal
}
}
#[tokio::test]
async fn a_failed_tick_is_logged_and_the_loop_continues() {
let ticks = Arc::new(AtomicU32::new(0));
let host = BackgroundServiceHost::new(Arc::new(FailingTicker {
ticks: Arc::clone(&ticks),
fatal: false,
}));
host.start().await.expect("start must succeed");
let deadline = Instant::now() + Duration::from_secs(2);
while ticks.load(Ordering::SeqCst) < 3 && Instant::now() < deadline {
tokio::time::sleep(Duration::from_millis(5)).await;
}
assert!(
ticks.load(Ordering::SeqCst) >= 3,
"a transient failure must not stop a poller"
);
host.stop().await.expect("stop must succeed");
}
#[tokio::test]
async fn stop_on_error_surfaces_the_typed_tick_error() {
let ticks = Arc::new(AtomicU32::new(0));
let host = BackgroundServiceHost::new(Arc::new(FailingTicker { ticks, fatal: true }));
host.start().await.expect("start must succeed");
tokio::time::sleep(Duration::from_millis(50)).await;
let err = host
.stop()
.await
.expect_err("a fatal tick error must be reported");
assert!(err.to_string().contains("tick failed"), "{err}");
}
struct ZeroInterval;
#[async_trait::async_trait]
impl PeriodicBackgroundService for ZeroInterval {
fn interval(&self) -> Duration {
Duration::ZERO
}
async fn tick(&self) -> Result<()> {
Ok(())
}
}
#[tokio::test]
async fn zero_interval_is_a_typed_configuration_error() {
let host = BackgroundServiceHost::new(Arc::new(ZeroInterval));
host.start().await.expect("start must succeed");
tokio::time::sleep(Duration::from_millis(20)).await;
let err = host.stop().await.expect_err("zero interval must fail");
assert!(err.to_string().contains("greater than zero"), "{err}");
}
}