taskflow_rs/executor/
task_executor.rs

1use std::collections::HashSet;
2use std::sync::Arc;
3use tokio::sync::{Notify, RwLock};
4
5#[derive(Clone)]
6pub struct TaskExecutor {
7    running_tasks: Arc<RwLock<HashSet<String>>>,
8    shutdown_signal: Arc<Notify>,
9}
10
11impl TaskExecutor {
12    pub fn new() -> Self {
13        Self {
14            running_tasks: Arc::new(RwLock::new(HashSet::new())),
15            shutdown_signal: Arc::new(Notify::new()),
16        }
17    }
18
19    pub async fn add_running_task(&self, task_id: String) {
20        let mut running = self.running_tasks.write().await;
21        running.insert(task_id);
22    }
23
24    pub async fn remove_running_task(&self, task_id: &str) {
25        let mut running = self.running_tasks.write().await;
26        running.remove(task_id);
27    }
28
29    pub async fn running_count(&self) -> usize {
30        let running = self.running_tasks.read().await;
31        running.len()
32    }
33
34    pub async fn is_empty(&self) -> bool {
35        let running = self.running_tasks.read().await;
36        running.is_empty()
37    }
38
39    pub fn shutdown(&self) {
40        self.shutdown_signal.notify_waiters();
41    }
42
43    pub fn wait_for_shutdown(&self) -> &Notify {
44        &self.shutdown_signal
45    }
46
47    pub async fn wait_for_running_tasks(&self) {
48        while !self.is_empty().await {
49            tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
50        }
51    }
52}
taskflow_rs/executor/task_executor.rs

taskflow_rs/executor/
task_executor.rs
