kellnr_common/
background_tasks.rs

1use std::time::Duration;
2
3use tokio::task::JoinHandle;
4use tokio::time::Instant;
5use tokio_util::sync::CancellationToken;
6use tokio_util::task::TaskTracker;
7use tracing::{error, trace, warn};
8
9/// Pause before a task runs again after a round panicked, unless the task
10/// runs periodically, then the next round follows the interval.
11const PANIC_PAUSE: Duration = Duration::from_secs(1);
12
13/// When the first round of a periodic task runs.
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum FirstRound {
16    Immediately,
17    AfterInterval,
18}
19
20/// Background tasks that run in rounds until shutdown, or once.
21///
22/// Every round runs as its own tokio task. A round that panics is logged and
23/// the task continues. On shutdown no new round starts, a running round or
24/// one-off task is not interrupted.
25#[derive(Debug, Clone, Default)]
26pub struct BackgroundTasks {
27    shutdown: CancellationToken,
28    tracker: TaskTracker,
29}
30
31impl BackgroundTasks {
32    pub fn new() -> Self {
33        Self::default()
34    }
35
36    /// Cancelled on shutdown. Rounds that may run for a long time, e.g. docs
37    /// builds, use it to stop early.
38    pub fn shutdown_token(&self) -> CancellationToken {
39        self.shutdown.clone()
40    }
41
42    pub fn request_shutdown(&self) {
43        self.shutdown.cancel();
44    }
45
46    /// Run `round` every `interval`. A round that takes longer than the
47    /// interval delays the next one.
48    pub fn spawn_periodic<F, Fut>(
49        &self,
50        name: &'static str,
51        interval: Duration,
52        first_round: FirstRound,
53        mut round: F,
54    ) where
55        F: FnMut() -> Fut + Send + 'static,
56        Fut: Future<Output = ()> + Send + 'static,
57    {
58        let first_delay = match first_round {
59            FirstRound::Immediately => Duration::ZERO,
60            FirstRound::AfterInterval => interval,
61        };
62        let timed_round = move || {
63            let round = round();
64            async move {
65                let start = Instant::now();
66                round.await;
67                interval.saturating_sub(start.elapsed())
68            }
69        };
70        self.spawn_rounds(name, first_delay, interval, timed_round);
71    }
72
73    /// Run `round` repeatedly, starting immediately. Each round returns the
74    /// delay until the next one.
75    pub fn spawn_loop<F, Fut>(&self, name: &'static str, round: F)
76    where
77        F: FnMut() -> Fut + Send + 'static,
78        Fut: Future<Output = Duration> + Send + 'static,
79    {
80        self.spawn_rounds(name, Duration::ZERO, PANIC_PAUSE, round);
81    }
82
83    /// Run `task` once. Shutdown waits for it like for a running round.
84    pub fn spawn_tracked<Fut>(&self, task: Fut) -> JoinHandle<Fut::Output>
85    where
86        Fut: Future + Send + 'static,
87        Fut::Output: Send + 'static,
88    {
89        self.tracker.spawn(task)
90    }
91
92    /// Run `task` once and log it if it panics. Shutdown waits for it like
93    /// for a running round.
94    pub fn spawn_once<Fut>(&self, name: &'static str, task: Fut) -> JoinHandle<()>
95    where
96        Fut: Future<Output = ()> + Send + 'static,
97    {
98        self.spawn_tracked(async move {
99            if let Err(e) = tokio::spawn(task).await {
100                error!("Background task {name} stopped: {e}");
101            }
102        })
103    }
104
105    fn spawn_rounds<F, Fut>(
106        &self,
107        name: &'static str,
108        first_delay: Duration,
109        delay_after_panic: Duration,
110        mut round: F,
111    ) where
112        F: FnMut() -> Fut + Send + 'static,
113        Fut: Future<Output = Duration> + Send + 'static,
114    {
115        let shutdown = self.shutdown.clone();
116        self.tracker.spawn(async move {
117            let mut delay = first_delay;
118            loop {
119                tokio::select! {
120                    biased;
121                    () = shutdown.cancelled() => break,
122                    () = tokio::time::sleep(delay) => {}
123                }
124                delay = match tokio::spawn(round()).await {
125                    Ok(next_delay) => next_delay,
126                    Err(e) if e.is_panic() => {
127                        error!("Background task {name} panicked, continuing: {e}");
128                        delay_after_panic
129                    }
130                    Err(_) => break,
131                };
132            }
133            trace!("Background task {name} stopped");
134        });
135    }
136
137    /// Wait until shutdown is requested, then up to `timeout` for the running
138    /// rounds to finish.
139    pub async fn wait_for_shutdown(&self, timeout: Duration) {
140        self.shutdown.cancelled().await;
141        self.tracker.close();
142        if tokio::time::timeout(timeout, self.tracker.wait())
143            .await
144            .is_err()
145        {
146            warn!(
147                "Background tasks did not finish within {}s, shutting down anyway",
148                timeout.as_secs()
149            );
150        }
151    }
152}
153
154#[cfg(test)]
155mod tests {
156    use std::sync::Arc;
157    use std::sync::atomic::{AtomicUsize, Ordering};
158
159    use super::*;
160
161    const INTERVAL: Duration = Duration::from_secs(10);
162
163    fn counting_round(
164        rounds: &Arc<AtomicUsize>,
165    ) -> impl FnMut() -> std::future::Ready<()> + Send + 'static {
166        let rounds = rounds.clone();
167        move || {
168            rounds.fetch_add(1, Ordering::SeqCst);
169            std::future::ready(())
170        }
171    }
172
173    async fn sleep_secs(secs: u64) {
174        tokio::time::sleep(Duration::from_secs(secs)).await;
175    }
176
177    #[tokio::test(start_paused = true)]
178    async fn periodic_task_runs_once_per_interval() {
179        let tasks = BackgroundTasks::new();
180        let rounds = Arc::new(AtomicUsize::new(0));
181
182        tasks.spawn_periodic(
183            "test",
184            INTERVAL,
185            FirstRound::AfterInterval,
186            counting_round(&rounds),
187        );
188
189        sleep_secs(5).await;
190        assert_eq!(0, rounds.load(Ordering::SeqCst));
191        sleep_secs(30).await;
192        assert_eq!(3, rounds.load(Ordering::SeqCst));
193    }
194
195    #[tokio::test(start_paused = true)]
196    async fn periodic_task_can_start_immediately() {
197        let tasks = BackgroundTasks::new();
198        let rounds = Arc::new(AtomicUsize::new(0));
199
200        tasks.spawn_periodic(
201            "test",
202            INTERVAL,
203            FirstRound::Immediately,
204            counting_round(&rounds),
205        );
206
207        sleep_secs(1).await;
208        assert_eq!(1, rounds.load(Ordering::SeqCst));
209        sleep_secs(10).await;
210        assert_eq!(2, rounds.load(Ordering::SeqCst));
211    }
212
213    #[tokio::test(start_paused = true)]
214    async fn periodic_task_continues_after_a_panicking_round() {
215        let tasks = BackgroundTasks::new();
216        let rounds = Arc::new(AtomicUsize::new(0));
217        let counter = rounds.clone();
218
219        tasks.spawn_periodic("test", INTERVAL, FirstRound::Immediately, move || {
220            let round = counter.fetch_add(1, Ordering::SeqCst);
221            async move {
222                assert!(round != 0, "first round panics");
223            }
224        });
225
226        sleep_secs(25).await;
227        assert_eq!(3, rounds.load(Ordering::SeqCst));
228    }
229
230    #[tokio::test(start_paused = true)]
231    async fn loop_task_waits_the_delay_returned_by_each_round() {
232        let tasks = BackgroundTasks::new();
233        let rounds = Arc::new(AtomicUsize::new(0));
234        let counter = rounds.clone();
235
236        tasks.spawn_loop("test", move || {
237            let round = counter.fetch_add(1, Ordering::SeqCst);
238            async move { Duration::from_secs(if round == 0 { 100 } else { 10 }) }
239        });
240
241        sleep_secs(99).await;
242        assert_eq!(1, rounds.load(Ordering::SeqCst));
243        sleep_secs(20).await;
244        assert_eq!(3, rounds.load(Ordering::SeqCst));
245    }
246
247    #[tokio::test(start_paused = true)]
248    async fn loop_task_pauses_after_a_panicking_round() {
249        let tasks = BackgroundTasks::new();
250        let rounds = Arc::new(AtomicUsize::new(0));
251        let counter = rounds.clone();
252
253        tasks.spawn_loop("test", move || {
254            let round = counter.fetch_add(1, Ordering::SeqCst);
255            async move {
256                assert!(round != 0, "first round panics");
257                Duration::from_secs(100)
258            }
259        });
260
261        sleep_secs(2).await;
262        assert_eq!(2, rounds.load(Ordering::SeqCst));
263    }
264
265    #[tokio::test(start_paused = true)]
266    async fn shutdown_lets_the_running_round_finish_and_stops_the_task() {
267        let tasks = BackgroundTasks::new();
268        let started = Arc::new(AtomicUsize::new(0));
269        let finished = Arc::new(AtomicUsize::new(0));
270        let (started_in_round, finished_in_round) = (started.clone(), finished.clone());
271
272        tasks.spawn_periodic("test", INTERVAL, FirstRound::Immediately, move || {
273            started_in_round.fetch_add(1, Ordering::SeqCst);
274            let finished = finished_in_round.clone();
275            async move {
276                sleep_secs(5).await;
277                finished.fetch_add(1, Ordering::SeqCst);
278            }
279        });
280
281        sleep_secs(1).await;
282        tasks.request_shutdown();
283        tasks.wait_for_shutdown(Duration::from_secs(30)).await;
284
285        assert_eq!(1, finished.load(Ordering::SeqCst));
286        sleep_secs(60).await;
287        assert_eq!(1, started.load(Ordering::SeqCst));
288    }
289
290    #[tokio::test(start_paused = true)]
291    async fn shutdown_waits_for_a_task_spawned_once() {
292        let tasks = BackgroundTasks::new();
293        let finished = Arc::new(AtomicUsize::new(0));
294        let finished_in_task = finished.clone();
295
296        tasks.spawn_once("test", async move {
297            sleep_secs(5).await;
298            finished_in_task.fetch_add(1, Ordering::SeqCst);
299        });
300
301        tasks.request_shutdown();
302        tasks.wait_for_shutdown(Duration::from_secs(30)).await;
303
304        assert_eq!(1, finished.load(Ordering::SeqCst));
305    }
306
307    #[tokio::test(start_paused = true)]
308    async fn task_spawned_once_logs_a_panic_instead_of_passing_it_on() {
309        let tasks = BackgroundTasks::new();
310
311        tasks
312            .spawn_once("test", async { panic!("task panics") })
313            .await
314            .unwrap();
315    }
316
317    #[tokio::test(start_paused = true)]
318    async fn wait_for_shutdown_gives_up_after_the_timeout() {
319        let tasks = BackgroundTasks::new();
320        tasks.spawn_loop("test", || async {
321            sleep_secs(3600).await;
322            Duration::ZERO
323        });
324
325        sleep_secs(1).await;
326        tasks.request_shutdown();
327        let start = Instant::now();
328        tasks.wait_for_shutdown(Duration::from_secs(30)).await;
329
330        assert_eq!(Duration::from_secs(30), start.elapsed());
331    }
332}
kellnr_common/background_tasks.rs

kellnr_common/
background_tasks.rs
