v2

vnt-core/src/utils/task_control.rs

@@ -0,0 +1,255 @@
+use parking_lot::Mutex;
+use std::collections::HashMap;
+use std::future::Future;
+use std::sync::{Arc, Weak};
+use tokio::sync::Notify;
+use tokio::task::{Id, JoinHandle};
+
+struct TaskGroupState {
+    stopped: bool,
+    tasks: HashMap<Id, JoinHandle<()>>,
+}
+
+struct TaskGroupInner {
+    state: Mutex<TaskGroupState>,
+    all_stopped_notify: Notify,
+}
+
+impl TaskGroupInner {
+    fn new() -> Self {
+        Self {
+            state: Mutex::new(TaskGroupState {
+                stopped: false,
+                tasks: HashMap::new(),
+            }),
+            all_stopped_notify: Notify::new(),
+        }
+    }
+
+    fn spawn<F>(self: &Arc<Self>, f: F) -> Option<Id>
+    where
+        F: Future + Send + 'static,
+        F::Output: Send + 'static,
+    {
+        let mut state = self.state.lock();
+        if state.stopped {
+            return None;
+        }
+
+        let guard = TaskGuard {
+            inner: Arc::downgrade(self),
+        };
+
+        let handle = tokio::spawn(async move {
+            let _guard = guard;
+            f.await;
+        });
+
+        let task_id = handle.id();
+        state.tasks.insert(task_id, handle);
+        Some(task_id)
+    }
+
+    fn stop(&self) {
+        let mut state = self.state.lock();
+        state.stopped = true;
+        for (_, handle) in state.tasks.drain() {
+            handle.abort();
+        }
+    }
+
+    fn is_stopped(&self) -> bool {
+        self.state.lock().stopped
+    }
+
+    fn remove_task(&self, task_id: Id) {
+        let all_stopped = {
+            let mut state = self.state.lock();
+            state.tasks.remove(&task_id);
+            if state.tasks.is_empty() {
+                state.stopped = true;
+                true
+            } else {
+                false
+            }
+        };
+        if all_stopped {
+            self.all_stopped_notify.notify_waiters();
+        }
+    }
+
+    async fn abort_task(&self, task_id: Id) {
+        let handle = self.state.lock().tasks.remove(&task_id);
+        if let Some(handle) = handle {
+            handle.abort();
+            _ = handle.await;
+        }
+    }
+
+    async fn join_all(&self) {
+        let tasks = std::mem::take(&mut self.state.lock().tasks);
+        for (_, h) in tasks {
+            let _ = h.await;
+        }
+    }
+
+    fn all_tasks_stopped(&self) -> bool {
+        let state = self.state.lock();
+        state.stopped && state.tasks.is_empty()
+    }
+}
+
+impl Drop for TaskGroupInner {
+    fn drop(&mut self) {
+        self.stop();
+    }
+}
+
+struct TaskGuard {
+    inner: Weak<TaskGroupInner>,
+}
+
+impl Drop for TaskGuard {
+    fn drop(&mut self) {
+        if let Some(inner) = self.inner.upgrade() {
+            let task_id = tokio::task::id();
+            inner.remove_task(task_id);
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct TaskGroup {
+    inner: Arc<TaskGroupInner>,
+}
+
+impl TaskGroup {
+    fn new() -> Self {
+        Self {
+            inner: Arc::new(TaskGroupInner::new()),
+        }
+    }
+
+    pub fn stop(&self) {
+        self.inner.stop();
+    }
+
+    pub fn is_stopped(&self) -> bool {
+        self.inner.is_stopped()
+    }
+
+    pub fn spawn<F>(&self, f: F) -> SubTask
+    where
+        F: Future + Send + 'static,
+        F::Output: Send + 'static,
+    {
+        match self.inner.spawn(f) {
+            Some(task_id) => SubTask::new(task_id, Arc::downgrade(&self.inner)),
+            None => SubTask::empty(),
+        }
+    }
+
+    pub async fn join_all(&self) {
+        self.inner.join_all().await;
+    }
+
+    pub async fn wait_all_stopped(&self) {
+        loop {
+            if self.inner.all_tasks_stopped() {
+                return;
+            }
+            self.inner.all_stopped_notify.notified().await;
+        }
+    }
+}
+
+pub struct SubTask {
+    task_id: Option<Id>,
+    inner: Weak<TaskGroupInner>,
+}
+
+impl SubTask {
+    fn new(task_id: Id, inner: Weak<TaskGroupInner>) -> Self {
+        Self {
+            task_id: Some(task_id),
+            inner,
+        }
+    }
+
+    fn empty() -> Self {
+        Self {
+            task_id: None,
+            inner: Weak::new(),
+        }
+    }
+
+    pub async fn stop(&self) {
+        if let Some(task_id) = self.task_id
+            && let Some(inner) = self.inner.upgrade()
+        {
+            inner.abort_task(task_id).await;
+        }
+    }
+
+    pub fn is_running(&self) -> bool {
+        if let Some(task_id) = self.task_id
+            && let Some(inner) = self.inner.upgrade()
+        {
+            return inner.state.lock().tasks.contains_key(&task_id);
+        }
+        false
+    }
+
+    pub fn id(&self) -> Option<Id> {
+        self.task_id
+    }
+}
+#[derive(Clone, Default)]
+pub struct TaskGroupManager {
+    task_group: Arc<Mutex<Option<TaskGroup>>>,
+}
+
+impl TaskGroupManager {
+    pub fn new() -> Self {
+        TaskGroupManager::default()
+    }
+
+    pub fn is_running(&self) -> bool {
+        self.task_group.lock().is_some()
+    }
+
+    pub fn is_stopped(&self) -> bool {
+        self.task_group.lock().is_none()
+    }
+
+    pub fn create_task(&self) -> anyhow::Result<(TaskGroup, TaskGroupGuard)> {
+        let mut guard = self.task_group.lock();
+        if guard.is_some() {
+            anyhow::bail!("运行中")
+        }
+
+        let task_group = TaskGroup::new();
+        guard.replace(task_group.clone());
+        let stop_guard = TaskGroupGuard {
+            task_group: self.task_group.clone(),
+        };
+        Ok((task_group, stop_guard))
+    }
+
+    pub fn stop(&self) {
+        let option = self.task_group.lock();
+        if let Some(task_group) = option.as_ref() {
+            task_group.stop();
+        }
+    }
+}
+pub struct TaskGroupGuard {
+    task_group: Arc<Mutex<Option<TaskGroup>>>,
+}
+impl Drop for TaskGroupGuard {
+    fn drop(&mut self) {
+        if let Some(task_group) = self.task_group.lock().take() {
+            task_group.stop();
+        }
+    }
+}