hirofa_utils/
task_manager.rs

1use futures::Future;
2use log::trace;
3use tokio::runtime::{Handle, Runtime};
4use tokio::task::JoinError;
5
6pub struct TaskManager {
7    handle: Handle,
8    _runtime: Option<Runtime>,
9}
10
11impl TaskManager {
12    pub fn new(thread_count: usize) -> Self {
13        // start threads
14
15        let runtime = tokio::runtime::Builder::new_multi_thread()
16            .enable_all()
17            .max_blocking_threads(thread_count)
18            .build()
19            .expect("tokio rt failed");
20
21        let handle = runtime.handle().clone();
22        TaskManager {
23            handle,
24            _runtime: Some(runtime),
25        }
26    }
27
28    pub fn from_handle(handle: Handle) -> Self {
29        TaskManager {
30            handle,
31            _runtime: None,
32        }
33    }
34
35    pub fn add_task<T: FnOnce() + Send + 'static>(&self, task: T) {
36        trace!("adding a task");
37        self.handle.spawn_blocking(task);
38    }
39
40    /// start an async task
41    /// # Example
42    /// ```rust
43    /// use hirofa_utils::task_manager::TaskManager;
44    /// let tm = TaskManager::new(2);
45    /// let task = async {
46    ///     println!("foo");
47    /// };
48    /// tm.add_task_async(task);
49    /// ```
50    pub fn add_task_async<R: Send + 'static, T: Future<Output = R> + Send + 'static>(
51        &self,
52        task: T,
53    ) -> impl Future<Output = Result<R, JoinError>> {
54        self.handle.spawn(task)
55    }
56
57    #[allow(dead_code)]
58    pub fn run_task_blocking<R: Send + 'static, T: FnOnce() -> R + Send + 'static>(
59        &self,
60        task: T,
61    ) -> R {
62        trace!("adding a sync task from thread {}", thread_id::get());
63        // check if the current thread is not a worker thread, because that would be bad
64        let join_handle = self.handle.spawn_blocking(task);
65        self.handle.block_on(join_handle).expect("task failed")
66    }
67}
68
69#[cfg(test)]
70mod tests {
71    use crate::task_manager::TaskManager;
72    use log::trace;
73    use std::thread;
74    use std::time::Duration;
75
76    #[test]
77    fn test() {
78        trace!("testing");
79
80        let tm = TaskManager::new(1);
81        for _x in 0..5 {
82            tm.add_task(|| {
83                thread::sleep(Duration::from_secs(1));
84            })
85        }
86
87        let s = tm.run_task_blocking(|| {
88            thread::sleep(Duration::from_secs(1));
89            "res"
90        });
91
92        assert_eq!(s, "res");
93
94        for _x in 0..10 {
95            let s = tm.run_task_blocking(|| "res");
96
97            assert_eq!(s, "res");
98        }
99    }
100
101    #[test]
102    fn test_from_handle() {
103        let rt = tokio::runtime::Builder::new_multi_thread()
104            .enable_all()
105            .build()
106            .unwrap();
107        let handle = rt.handle().clone();
108        let tm = TaskManager::from_handle(handle);
109
110        let s = tm.run_task_blocking(|| "res");
111        assert_eq!(s, "res");
112    }
113}
hirofa_utils/task_manager.rs

hirofa_utils/
task_manager.rs
