diode_base/
daemon.rs

1use std::marker::PhantomData;
2use std::sync::Arc;
3
4use async_trait::async_trait;
5use diode::{
6    AddServiceExt as _, App, AppBuilder, AppContext, Dependencies, Plugin, Service,
7    ServiceDependencyExt as _, StdError,
8};
9use tokio::task::JoinSet;
10
11pub use tokio_util::sync::CancellationToken;
12
13use crate::defer;
14
15#[derive(Default)]
16struct DaemonRegistry {
17    daemons: Vec<Arc<dyn DynDaemon>>,
18}
19
20impl DaemonRegistry {
21    pub fn add_daemon<T>(&mut self, daemon: Arc<T>)
22    where
23        T: Daemon + 'static,
24    {
25        self.daemons.push(daemon);
26    }
27
28    pub async fn run_daemons(
29        &self,
30        app: Arc<App>,
31        shutdown: CancellationToken,
32    ) -> Result<(), StdError> {
33        let span = tracing::info_span!("daemons");
34        let mut futures = JoinSet::new();
35        tracing::info!(parent: &span, "Daemons starting");
36        for daemon in self.daemons.iter() {
37            let shutdown = shutdown.child_token();
38            let app = app.clone();
39            let daemon = daemon.clone();
40            futures.spawn(async move { daemon.run(&app, shutdown).await });
41        }
42        tracing::info!(parent: &span, "Daemons running");
43        defer! {
44            tracing::info!(parent: &span, "Daemons stopped");
45        };
46        let first_result = futures.join_next().await;
47        shutdown.cancel();
48        if let Some(result) = first_result {
49            result.map_err(Box::new)??;
50            while let Some(result) = futures.join_next().await {
51                result.map_err(Box::new)??;
52            }
53        }
54        Ok(())
55    }
56}
57
58pub trait Daemon: Send + Sync {
59    fn run(
60        &self,
61        app: &App,
62        shutdown: CancellationToken,
63    ) -> impl Future<Output = Result<(), StdError>> + Send {
64        let _ = app;
65        async move {
66            shutdown.cancelled_owned().await;
67            Ok(())
68        }
69    }
70}
71
72#[async_trait]
73trait DynDaemon: Send + Sync {
74    async fn run(&self, app: &App, shutdown: CancellationToken) -> Result<(), StdError>;
75}
76
77#[async_trait]
78impl<T> DynDaemon for T
79where
80    T: Daemon,
81{
82    async fn run(&self, app: &App, shutdown: CancellationToken) -> Result<(), StdError> {
83        self.run(app, shutdown).await
84    }
85}
86
87pub trait RunDaemonsExt {
88    fn run_daemons(
89        self,
90        shutdown: CancellationToken,
91    ) -> impl Future<Output = Result<(), StdError>> + Send;
92}
93
94impl RunDaemonsExt for App {
95    async fn run_daemons(self, shutdown: CancellationToken) -> Result<(), StdError> {
96        Arc::new(self).run_daemons(shutdown).await
97    }
98}
99
100impl RunDaemonsExt for Arc<App> {
101    async fn run_daemons(self, shutdown: CancellationToken) -> Result<(), StdError> {
102        match self.get_component_ref::<DaemonRegistry>() {
103            Some(v) => v.run_daemons(self.clone(), shutdown).await,
104            None => Ok(()),
105        }
106    }
107}
108
109pub trait AddDaemonExt {
110    fn add_daemon<T>(&self, daemon: impl Into<Arc<T>>)
111    where
112        T: Daemon + 'static;
113}
114
115impl AddDaemonExt for AppContext {
116    fn add_daemon<T>(&self, daemon: impl Into<Arc<T>>)
117    where
118        T: Daemon + 'static,
119    {
120        if !self.has_component::<DaemonRegistry>() {
121            self.add_component(DaemonRegistry::default());
122        }
123        self.get_component_mut::<DaemonRegistry>()
124            .unwrap()
125            .add_daemon(daemon.into());
126    }
127}
128
129struct DaemonServiceProvider<T>(PhantomData<T>);
130
131impl<T> Plugin for DaemonServiceProvider<T>
132where
133    T: Service<Handle = Arc<T>> + Daemon + 'static,
134{
135    async fn build(&self, ctx: &AppContext) -> Result<(), StdError> {
136        let handle = ctx.get_component::<T::Handle>().unwrap();
137        ctx.add_daemon::<T>(handle);
138        Ok(())
139    }
140
141    fn dependencies(&self) -> Dependencies {
142        T::dependencies().service::<T>()
143    }
144}
145
146pub trait AddDaemonServiceExt {
147    fn add_daemon_service<T>(&mut self) -> &mut Self
148    where
149        T: Service<Handle = Arc<T>> + Daemon + 'static;
150
151    fn has_daemon_service<T>(&self) -> bool
152    where
153        T: Service<Handle = Arc<T>> + Daemon + 'static;
154}
155
156impl AddDaemonServiceExt for AppBuilder {
157    fn add_daemon_service<T>(&mut self) -> &mut Self
158    where
159        T: Service<Handle = Arc<T>> + Daemon + 'static,
160    {
161        if !self.has_service::<T>() {
162            self.add_service::<T>();
163        }
164        self.add_plugin(DaemonServiceProvider::<T>(PhantomData));
165        self
166    }
167
168    fn has_daemon_service<T>(&self) -> bool
169    where
170        T: Service<Handle = Arc<T>> + Daemon + 'static,
171    {
172        self.has_plugin::<DaemonServiceProvider<T>>()
173    }
174}
diode_base/daemon.rs

diode_base/
daemon.rs
