stratum_apps/
task_manager.rs

1use std::sync::Mutex as StdMutex;
2use tokio::task::JoinHandle;
3
4/// Manages a collection of spawned tokio tasks.
5///
6/// This struct provides a centralized way to spawn, track, and manage the lifecycle
7/// of async tasks in the translator. It maintains a list of join handles that can
8/// be used to wait for all tasks to complete or abort them during shutdown.
9pub struct TaskManager {
10    tasks: StdMutex<Vec<JoinHandle<()>>>,
11}
12
13impl Default for TaskManager {
14    fn default() -> Self {
15        Self::new()
16    }
17}
18
19impl TaskManager {
20    /// Creates a new TaskManager instance.
21    ///
22    /// Initializes an empty task manager ready to spawn and track tasks.
23    pub fn new() -> Self {
24        Self {
25            tasks: StdMutex::new(Vec::new()),
26        }
27    }
28
29    /// Spawns a new async task and adds it to the managed collection.
30    ///
31    /// The task will be tracked by this manager and can be waited for or aborted
32    /// using the other methods.
33    ///
34    /// # Arguments
35    /// * `fut` - The future to spawn as a task
36    #[track_caller]
37    pub fn spawn<F>(&self, fut: F)
38    where
39        F: std::future::Future<Output = ()> + Send + 'static,
40    {
41        use tracing::Instrument;
42        let location = std::panic::Location::caller();
43        let span = tracing::trace_span!(
44            "task",
45            file = location.file(),
46            line = location.line(),
47            column = location.column(),
48        );
49
50        let handle = tokio::spawn(fut.instrument(span));
51        self.tasks.lock().unwrap().push(handle);
52    }
53
54    /// Waits for all managed tasks to complete.
55    ///
56    /// This method will block until all tasks that were spawned through this
57    /// manager have finished executing. Tasks are joined in reverse order
58    /// (most recently spawned first).
59    pub async fn join_all(&self) {
60        let handles = {
61            let mut tasks = self.tasks.lock().unwrap();
62            std::mem::take(&mut *tasks)
63        };
64
65        for handle in handles {
66            let _ = handle.await;
67        }
68    }
69
70    /// Aborts all managed tasks.
71    ///
72    /// This method immediately cancels all tasks that were spawned through this
73    /// manager. The tasks will be terminated without waiting for them to complete.
74    pub async fn abort_all(&self) {
75        let mut tasks = self.tasks.lock().unwrap();
76        for handle in tasks.drain(..) {
77            handle.abort();
78        }
79    }
80}
stratum_apps/task_manager.rs

stratum_apps/
task_manager.rs
