shared/extensions/
background_tasks.rs

1use crate::State;
2use futures_util::FutureExt;
3use std::{
4    borrow::Cow,
5    collections::{HashMap, HashSet},
6    panic::AssertUnwindSafe,
7    sync::Arc,
8};
9use tokio::sync::{OwnedRwLockReadGuard, RwLock};
10
11pub struct BackgroundTask {
12    pub name: &'static str,
13    pub last_execution: std::time::Instant,
14    pub last_error: Option<anyhow::Error>,
15
16    pub task: tokio::task::JoinHandle<()>,
17}
18
19pub struct BackgroundTaskBuilder {
20    state: State,
21    tasks: Arc<RwLock<HashMap<&'static str, BackgroundTask>>>,
22}
23
24impl BackgroundTaskBuilder {
25    pub fn new(state: State) -> Self {
26        Self {
27            state,
28            tasks: Arc::new(RwLock::new(HashMap::new())),
29        }
30    }
31
32    pub(crate) fn share(&self) -> Self {
33        Self {
34            state: self.state.clone(),
35            tasks: Arc::clone(&self.tasks),
36        }
37    }
38
39    pub(crate) async fn task_names(&self) -> HashSet<&'static str> {
40        self.tasks.read().await.keys().copied().collect()
41    }
42
43    /// Aborts and removes every task not in `keep`.
44    pub(crate) async fn rollback(&self, keep: &HashSet<&'static str>) {
45        self.tasks.write().await.retain(|name, task| {
46            if keep.contains(name) {
47                return true;
48            }
49
50            task.task.abort();
51            false
52        });
53    }
54
55    /// Adds a background task that will be executed periodically, depending on your loop function implementation.
56    /// This will only run on primary instances, so be aware of that when implementing your task.
57    pub async fn add_task<
58        F: Fn(State) -> Fut + Send + Sync + 'static,
59        Fut: Future<Output = Result<(), anyhow::Error>> + Send + 'static,
60    >(
61        &self,
62        name: &'static str,
63        loop_fn: F,
64    ) {
65        if !self.state.env.app_primary {
66            return;
67        }
68
69        let state = self.state.clone();
70        let tasks = Arc::clone(&self.tasks);
71
72        let previous = self.tasks.write().await.insert(
73            name,
74            BackgroundTask {
75                name,
76                last_execution: std::time::Instant::now(),
77                last_error: None,
78                task: tokio::spawn(async move {
79                    loop {
80                        if let Some(task) = tasks.write().await.get_mut(name) {
81                            task.last_execution = std::time::Instant::now();
82                        }
83
84                        tracing::debug!(name, "running background task function");
85                        let result = AssertUnwindSafe(loop_fn(state.clone()))
86                            .catch_unwind()
87                            .await;
88
89                        let result = match result {
90                            Ok(result) => result,
91                            Err(err) => {
92                                let err_msg: Cow<'_, str> =
93                                    if let Some(s) = err.downcast_ref::<&str>() {
94                                        (*s).into()
95                                    } else if let Some(s) = err.downcast_ref::<String>() {
96                                        s.clone().into()
97                                    } else {
98                                        "Unknown panic".into()
99                                    };
100
101                                tracing::error!(name, "background task panicked: {}", err_msg);
102                                sentry::capture_message(
103                                    &format!("Background task '{}' panicked: {}", name, err_msg),
104                                    sentry::Level::Error,
105                                );
106
107                                if let Some(task) = tasks.write().await.get_mut(name) {
108                                    task.last_error = Some(anyhow::anyhow!(err_msg));
109                                }
110
111                                return;
112                            }
113                        };
114
115                        if let Err(err) = &result {
116                            tracing::error!(name, "a background task error occurred: {:?}", err);
117                            sentry_anyhow::capture_anyhow(err);
118                        }
119
120                        if let Some(task) = tasks.write().await.get_mut(name) {
121                            task.last_error = result.err();
122                        }
123                    }
124                }),
125            },
126        );
127
128        if let Some(previous) = previous {
129            tracing::warn!(
130                name,
131                "background task registered twice, aborting the previous loop"
132            );
133            previous.task.abort();
134        }
135    }
136
137    /// Adds a background task that will be executed periodically, depending on the cron you provide, with the UTC timezone.
138    /// This will only run on primary instances, so be aware of that when implementing your task.
139    pub async fn add_cron_task<
140        F: Fn(State) -> Fut + Send + Sync + 'static,
141        Fut: Future<Output = Result<(), anyhow::Error>> + Send + 'static,
142    >(
143        &self,
144        name: &'static str,
145        cron: croner::Cron,
146        r#fn: F,
147    ) {
148        if !self.state.env.app_primary {
149            return;
150        }
151
152        let state = self.state.clone();
153        let tasks = Arc::clone(&self.tasks);
154
155        let previous = self.tasks.write().await.insert(
156            name,
157            BackgroundTask {
158                name,
159                last_execution: std::time::Instant::now(),
160                last_error: None,
161                task: tokio::spawn(async move {
162                    let schedule_iter = cron.iter_after(chrono::Utc::now());
163
164                    for target_datetime in schedule_iter {
165                        let target_timestamp = target_datetime.timestamp();
166                        let now_timestamp = chrono::Utc::now().timestamp();
167                        let sleep_duration = target_timestamp - now_timestamp;
168                        if sleep_duration <= 0 {
169                            tokio::time::sleep(std::time::Duration::from_secs(1)).await;
170                            continue;
171                        }
172
173                        tokio::time::sleep(std::time::Duration::from_secs(sleep_duration as u64))
174                            .await;
175
176                        if let Some(task) = tasks.write().await.get_mut(name) {
177                            task.last_execution = std::time::Instant::now();
178                        }
179
180                        tracing::debug!(name, "running background task function");
181                        let result = AssertUnwindSafe(r#fn(state.clone())).catch_unwind().await;
182
183                        let result = match result {
184                            Ok(result) => result,
185                            Err(err) => {
186                                let err_msg: Cow<'_, str> =
187                                    if let Some(s) = err.downcast_ref::<&str>() {
188                                        (*s).into()
189                                    } else if let Some(s) = err.downcast_ref::<String>() {
190                                        s.clone().into()
191                                    } else {
192                                        "Unknown panic".into()
193                                    };
194
195                                tracing::error!(name, "background task panicked: {}", err_msg);
196                                sentry::capture_message(
197                                    &format!("Background task '{}' panicked: {}", name, err_msg),
198                                    sentry::Level::Error,
199                                );
200
201                                if let Some(task) = tasks.write().await.get_mut(name) {
202                                    task.last_error = Some(anyhow::anyhow!(err_msg));
203                                }
204
205                                return;
206                            }
207                        };
208
209                        if let Err(err) = &result {
210                            tracing::error!(name, "a background task error occurred: {:?}", err);
211                            sentry_anyhow::capture_anyhow(err);
212                        }
213
214                        if let Some(task) = tasks.write().await.get_mut(name) {
215                            task.last_error = result.err();
216                        }
217                    }
218                }),
219            },
220        );
221
222        if let Some(previous) = previous {
223            tracing::warn!(
224                name,
225                "background task registered twice, aborting the previous loop"
226            );
227            previous.task.abort();
228        }
229    }
230}
231
232#[derive(Default)]
233pub struct BackgroundTaskManager {
234    builder: RwLock<Option<BackgroundTaskBuilder>>,
235}
236
237impl BackgroundTaskManager {
238    pub async fn merge_builder(&self, builder: BackgroundTaskBuilder) {
239        self.builder.write().await.replace(builder);
240    }
241
242    pub async fn get_tasks(&self) -> OwnedRwLockReadGuard<HashMap<&'static str, BackgroundTask>> {
243        let inner = self.builder.read().await;
244
245        inner.as_ref().unwrap().tasks.clone().read_owned().await
246    }
247}
shared/extensions/background_tasks.rs

shared/extensions/
background_tasks.rs
