Skip to main content

kerniq_runtime/
tasks.rs

1//! Background task management.
2
3use std::future::Future;
4use std::time::Duration;
5
6use tokio_util::task::TaskTracker;
7
8use crate::Shutdown;
9
10/// Spawns and tracks background tasks tied to the server lifecycle.
11///
12/// Tasks receive a [`Shutdown`] handle and are expected to return soon after
13/// it triggers. During graceful shutdown the server waits (up to the
14/// shutdown timeout) for every tracked task to finish.
15///
16/// This is intentionally small: it is not a job queue, it has no retries and
17/// no persistence.
18///
19/// ```
20/// use std::time::Duration;
21/// use kerniq_runtime::{BackgroundTasks, Shutdown};
22///
23/// # async fn demo() {
24/// let shutdown = Shutdown::new();
25/// let tasks = BackgroundTasks::new(shutdown.clone());
26/// tasks.spawn("ticker", |shutdown| async move {
27///     let mut interval = tokio::time::interval(Duration::from_millis(10));
28///     loop {
29///         tokio::select! {
30///             _ = interval.tick() => { /* periodic work */ }
31///             () = shutdown.wait() => break,
32///         }
33///     }
34/// });
35/// shutdown.trigger();
36/// assert!(tasks.shutdown(Duration::from_secs(1)).await);
37/// # }
38/// ```
39#[derive(Debug, Clone)]
40pub struct BackgroundTasks {
41    tracker: TaskTracker,
42    shutdown: Shutdown,
43}
44
45impl BackgroundTasks {
46    /// Creates a task set bound to `shutdown`.
47    pub fn new(shutdown: Shutdown) -> Self {
48        Self {
49            tracker: TaskTracker::new(),
50            shutdown,
51        }
52    }
53
54    /// Spawns a named task on the current Tokio runtime.
55    ///
56    /// # Panics
57    ///
58    /// Panics if called outside a Tokio runtime.
59    pub fn spawn<F, Fut>(&self, name: &'static str, task: F)
60    where
61        F: FnOnce(Shutdown) -> Fut,
62        Fut: Future<Output = ()> + Send + 'static,
63    {
64        let future = task(self.shutdown.clone());
65        tracing::debug!(task = name, "background task started");
66        self.tracker.spawn(async move {
67            future.await;
68            tracing::debug!(task = name, "background task finished");
69        });
70    }
71
72    /// Number of tasks still running.
73    pub fn len(&self) -> usize {
74        self.tracker.len()
75    }
76
77    /// Whether no tasks are running.
78    pub fn is_empty(&self) -> bool {
79        self.tracker.is_empty()
80    }
81
82    /// Triggers shutdown and waits up to `timeout` for all tasks to finish.
83    /// Returns `false` if the timeout elapsed first.
84    pub async fn shutdown(&self, timeout: Duration) -> bool {
85        self.shutdown.trigger();
86        self.tracker.close();
87        match tokio::time::timeout(timeout, self.tracker.wait()).await {
88            Ok(()) => true,
89            Err(_) => {
90                tracing::warn!(
91                    remaining = self.tracker.len(),
92                    "background tasks did not finish before the shutdown timeout"
93                );
94                false
95            }
96        }
97    }
98}