hive_router/background_tasks/
mod.rs1use 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}