diode_base/
daemon.rs

1use std::any::TypeId;
2use std::collections::HashMap;
3use std::sync::Arc;
4
5use async_trait::async_trait;
6use diode::{App, AppBuilder, StdError};
7use tokio::task::JoinSet;
8
9pub use tokio_util::sync::CancellationToken;
10
11use crate::defer;
12
13#[derive(Default)]
14struct DaemonRegistry {
15    daemons: HashMap<TypeId, Arc<dyn DynDaemon>>,
16}
17
18impl DaemonRegistry {
19    pub fn add_daemon<T>(&mut self, daemon: Arc<T>)
20    where
21        T: Daemon + 'static,
22    {
23        let type_id = TypeId::of::<T>();
24        self.daemons.insert(type_id, daemon);
25    }
26
27    pub fn has_daemon<T>(&self) -> bool
28    where
29        T: Daemon + 'static,
30    {
31        let type_id = TypeId::of::<T>();
32        self.daemons.contains_key(&type_id)
33    }
34
35    pub async fn run_daemons(
36        &self,
37        app: Arc<App>,
38        shutdown: CancellationToken,
39    ) -> Result<(), StdError> {
40        let span = tracing::info_span!("daemons");
41        let mut futures = JoinSet::new();
42        tracing::info!(parent: &span, "Daemons starting");
43        for daemon in self.daemons.values() {
44            let shutdown = shutdown.child_token();
45            let app = app.clone();
46            let daemon = daemon.clone();
47            futures.spawn(async move { daemon.run(&app, shutdown).await });
48        }
49        tracing::info!(parent: &span, "Daemons running");
50        defer! {
51            tracing::info!(parent: &span, "Daemons stopped");
52        };
53        let first_result = futures.join_next().await;
54        shutdown.cancel();
55        if let Some(result) = first_result {
56            result.map_err(Box::new)??;
57            while let Some(result) = futures.join_next().await {
58                result.map_err(Box::new)??;
59            }
60        }
61        Ok(())
62    }
63}
64
65pub trait Daemon: Send + Sync {
66    fn run(
67        &self,
68        app: &App,
69        shutdown: CancellationToken,
70    ) -> impl Future<Output = Result<(), StdError>> + Send {
71        let _ = app;
72        async move {
73            shutdown.cancelled_owned().await;
74            Ok(())
75        }
76    }
77}
78
79#[async_trait]
80trait DynDaemon: Send + Sync {
81    async fn run(&self, app: &App, shutdown: CancellationToken) -> Result<(), StdError>;
82}
83
84#[async_trait]
85impl<T> DynDaemon for T
86where
87    T: Daemon,
88{
89    async fn run(&self, app: &App, shutdown: CancellationToken) -> Result<(), StdError> {
90        self.run(app, shutdown).await
91    }
92}
93
94pub trait RunDaemonsExt {
95    fn run_daemons(
96        self,
97        shutdown: CancellationToken,
98    ) -> impl Future<Output = Result<(), StdError>> + Send;
99}
100
101impl RunDaemonsExt for App {
102    async fn run_daemons(self, shutdown: CancellationToken) -> Result<(), StdError> {
103        Arc::new(self).run_daemons(shutdown).await
104    }
105}
106
107impl RunDaemonsExt for Arc<App> {
108    async fn run_daemons(self, shutdown: CancellationToken) -> Result<(), StdError> {
109        match self.get_component_ref::<DaemonRegistry>() {
110            Some(v) => v.run_daemons(self.clone(), shutdown).await,
111            None => Ok(()),
112        }
113    }
114}
115
116pub trait AddDaemonExt {
117    fn add_daemon<T>(&mut self, daemon: impl Into<Arc<T>>) -> &mut Self
118    where
119        T: Daemon + 'static;
120
121    fn has_daemon<T>(&self) -> bool
122    where
123        T: Daemon + 'static;
124}
125
126impl AddDaemonExt for AppBuilder {
127    fn add_daemon<T>(&mut self, daemon: impl Into<Arc<T>>) -> &mut Self
128    where
129        T: Daemon + 'static,
130    {
131        if !self.has_component::<DaemonRegistry>() {
132            self.add_component(DaemonRegistry::default());
133        }
134        self.get_component_mut::<DaemonRegistry>()
135            .unwrap()
136            .add_daemon(daemon.into());
137        self
138    }
139
140    fn has_daemon<T>(&self) -> bool
141    where
142        T: Daemon + 'static,
143    {
144        self.get_component_ref::<DaemonRegistry>()
145            .is_some_and(|v| v.has_daemon::<T>())
146    }
147}
diode_base/daemon.rs

diode_base/
daemon.rs
