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        let mut tasks = self.tasks.lock().unwrap();
52        tasks.retain(|h| !h.is_finished());
53        tasks.push(handle);
54    }
55
56    /// Waits for all managed tasks to complete.
57    ///
58    /// This method will block until all tasks that were spawned through this
59    /// manager have finished executing. Tasks are joined in reverse order
60    /// (most recently spawned first).
61    pub async fn join_all(&self) {
62        let handles = {
63            let mut tasks = self.tasks.lock().unwrap();
64            std::mem::take(&mut *tasks)
65        };
66
67        for handle in handles {
68            let _ = handle.await;
69        }
70    }
71
72    /// Aborts all managed tasks.
73    ///
74    /// This method immediately cancels all tasks that were spawned through this
75    /// manager. The tasks will be terminated without waiting for them to complete.
76    pub async fn abort_all(&self) {
77        let mut tasks = self.tasks.lock().unwrap();
78        for handle in tasks.drain(..) {
79            handle.abort();
80        }
81    }
82}
stratum_apps/task_manager.rs

stratum_apps/
task_manager.rs
