Skip to main content

hive_router/background_tasks/
mod.rs

1use async_trait::async_trait;
2use ntex::rt::Arbiter;
3use std::future::Future;
4use tokio_util::sync::CancellationToken;
5use tracing::info;
6
7#[async_trait]
8pub trait BackgroundTask: Send + Sync {
9    fn id(&self) -> &str;
10    async fn run(&self, token: CancellationToken);
11}
12
13pub struct BackgroundTasksManager {
14    cancellation_token: CancellationToken,
15    arbiter: Arbiter,
16}
17
18impl Default for BackgroundTasksManager {
19    fn default() -> Self {
20        Self::new()
21    }
22}
23
24impl BackgroundTasksManager {
25    pub fn new() -> Self {
26        Self {
27            cancellation_token: CancellationToken::new(),
28            arbiter: Arbiter::new(),
29        }
30    }
31
32    pub fn register_task<T>(&mut self, task: T)
33    where
34        T: BackgroundTask + 'static,
35    {
36        info!("registering background task: {}", task.id());
37        let child_token = self.cancellation_token.clone();
38
39        self.arbiter.spawn(async move {
40            task.run(child_token).await;
41        });
42    }
43
44    pub fn register_handle<F>(&mut self, f: F)
45    where
46        F: Future<Output = ()> + Send + 'static,
47    {
48        self.arbiter.spawn(f);
49    }
50
51    pub fn shutdown(&mut self) {
52        info!("shutdown triggered, stopping all background tasks...");
53
54        self.cancellation_token.cancel();
55        self.arbiter.stop();
56
57        info!("all background tasks have been shut down gracefully.");
58    }
59}