support cancel coroutine/task (#399)

Author: loongs-zhang

CHANGELOG.md

@@ -1,3 +1,7 @@
+### 0.7.x
+
+- [x] support cancel coroutine/task
+
 ### 0.6.x
 
 - [x] support custom task and coroutine priority.

README.md

@@ -34,7 +34,6 @@ English | [中文](README_ZH.md)
 
 - [ ] add
   performance [benchmark](https://github.com/TechEmpower/FrameworkBenchmarks/wiki/Project-Information-Framework-Tests-Overview);
-- [ ] cancel coroutine/task;
 - [ ] add metrics;
 - [ ] add synchronization toolkit;
 - [ ] support and compatibility for AF_XDP socket;

README_ZH.md

@@ -30,7 +30,6 @@
 
 - [ ] 
   增加性能[基准测试](https://github.com/TechEmpower/FrameworkBenchmarks/wiki/Project-Information-Framework-Tests-Overview);
-- [ ] 取消协程/任务;
 - [ ] 增加性能指标监控;
 - [ ] 增加并发工具包;
 - [ ] 支持AF_XDP套接字;

core/src/co_pool/creator.rs

@@ -29,7 +29,7 @@ impl Listener<(), Option<usize>> for CoroutineCreator {
                         .store(pool.get_running_size().saturating_sub(1), Ordering::Release);
                 }
             }
-            CoroutineState::Error(_) => {
+            CoroutineState::Cancelled | CoroutineState::Error(_) => {
                 if let Some(pool) = CoroutinePool::current() {
                     //worker协程异常退出，需要先回收再创建
                     pool.running

core/src/co_pool/mod.rs

@@ -6,8 +6,9 @@ use crate::common::ordered_work_steal::{OrderedLocalQueue, OrderedWorkStealQueue
 use crate::common::{get_timeout_time, now, CondvarBlocker};
 use crate::coroutine::suspender::Suspender;
 use crate::scheduler::{SchedulableCoroutine, Scheduler};
-use crate::{error, impl_current_for, impl_display_by_debug, impl_for_named, trace};
+use crate::{error, impl_current_for, impl_display_by_debug, impl_for_named, trace, warn};
 use dashmap::{DashMap, DashSet};
+use once_cell::sync::Lazy;
 use std::cell::Cell;
 use std::ffi::c_longlong;
 use std::io::{Error, ErrorKind};
@@ -25,6 +26,11 @@ mod state;
 /// Creator for coroutine pool.
 mod creator;
 
+/// `task_name` -> `co_name`
+static RUNNING_TASKS: Lazy<DashMap<&str, &str>> = Lazy::new(DashMap::new);
+
+static CANCEL_TASKS: Lazy<DashSet<&str>> = Lazy::new(DashSet::new);
+
 /// The coroutine pool impls.
 #[repr(C)]
 #[derive(Debug)]
@@ -383,7 +389,17 @@ impl<'p> CoroutinePool<'p> {
 
     fn try_run(&self) -> Option<()> {
         self.task_queue.pop().map(|task| {
+            let tname = task.get_name().to_string().leak();
+            if CANCEL_TASKS.contains(tname) {
+                _ = CANCEL_TASKS.remove(tname);
+                warn!("Cancel task:{} successfully !", tname);
+                return;
+            }
+            if let Some(co) = SchedulableCoroutine::current() {
+                _ = RUNNING_TASKS.insert(tname, co.name());
+            }
             let (task_name, result) = task.run();
+            _ = RUNNING_TASKS.remove(tname);
             let n = task_name.clone().leak();
             if self.no_waits.contains(n) {
                 _ = self.no_waits.remove(n);
@@ -406,6 +422,44 @@ impl<'p> CoroutinePool<'p> {
         }
     }
 
+    /// Try to cancel a task.
+    pub fn try_cancel_task(task_name: &str) {
+        // 检查正在运行的任务是否是要取消的任务
+        if let Some(info) = RUNNING_TASKS.get(task_name) {
+            let co_name = *info;
+            // todo windows support
+            #[allow(unused_variables)]
+            if let Some(pthread) = Scheduler::get_scheduling_thread(co_name) {
+                // 发送SIGVTALRM信号，在运行时取消任务
+                #[cfg(unix)]
+                if nix::sys::pthread::pthread_kill(pthread, nix::sys::signal::Signal::SIGVTALRM)
+                    .is_ok()
+                {
+                    warn!(
+                        "Attempt to cancel task:{} running on coroutine:{} by thread:{}, cancelling...",
+                        task_name, co_name, pthread
+                    );
+                } else {
+                    error!(
+                        "Attempt to cancel task:{} running on coroutine:{} by thread:{} failed !",
+                        task_name, co_name, pthread
+                    );
+                }
+            } else {
+                // 添加到待取消队列
+                Scheduler::try_cancel_coroutine(co_name);
+                warn!(
+                    "Attempt to cancel task:{} running on coroutine:{}, cancelling...",
+                    task_name, co_name
+                );
+            }
+        } else {
+            // 添加到待取消队列
+            _ = CANCEL_TASKS.insert(Box::leak(Box::from(task_name)));
+            warn!("Attempt to cancel task:{}, cancelling...", task_name);
+        }
+    }
+
     /// Schedule the tasks.
     ///
     /// Allow multiple threads to concurrently submit task to the pool,

core/src/co_pool/task.rs

@@ -33,6 +33,12 @@ impl<'t> Task<'t> {
         }
     }
 
+    /// get the task name.
+    #[must_use]
+    pub fn get_name(&self) -> &str {
+        &self.name
+    }
+
     /// execute the task
     ///
     /// # Errors

core/src/common/constants.rs

@@ -195,9 +195,11 @@ pub enum CoroutineState<Y, R> {
     Suspend(Y, u64),
     ///The coroutine enters the syscall.
     Syscall(Y, SyscallName, SyscallState),
+    /// The coroutine cancelled.
+    Cancelled,
     /// The coroutine completed with a return value.
     Complete(R),
-    /// The coroutine completed with a error message.
+    /// The coroutine completed with an error message.
     Error(&'static str),
 }
 

core/src/coroutine/korosensei.rs

@@ -445,6 +445,10 @@ where
                 let current = self.state();
                 match current {
                     CoroutineState::Running => {
+                        if Suspender::<Yield, Param>::is_cancel() {
+                            self.cancel()?;
+                            return Ok(CoroutineState::Cancelled);
+                        }
                         let timestamp = Suspender::<Yield, Param>::timestamp();
                         self.suspend(y, timestamp)?;
                         Ok(CoroutineState::Suspend(y, timestamp))

core/src/coroutine/listener.rs

@@ -27,6 +27,9 @@ pub trait Listener<Yield, Return>: Debug {
     /// callback when the coroutine enters syscall.
     fn on_syscall(&self, local: &CoroutineLocal, old_state: CoroutineState<Yield, Return>) {}
 
+    /// Callback when the coroutine is cancelled.
+    fn on_cancel(&self, local: &CoroutineLocal, old_state: CoroutineState<Yield, Return>) {}
+
     /// Callback when the coroutine is completed.
     fn on_complete(
         &self,
@@ -91,6 +94,11 @@ where
         old_state: CoroutineState<Yield, Return>
     ), "on_syscall");
 
+    broadcast!(on_cancel(
+        local: &CoroutineLocal,
+        old_state: CoroutineState<Yield, Return>
+    ), "on_cancel");
+
     broadcast!(on_complete(
         local: &CoroutineLocal,
         old_state: CoroutineState<Yield, Return>,

core/src/coroutine/mod.rs

@@ -135,6 +135,35 @@ impl<'c, Param, Yield, Return> Coroutine<'c, Param, Yield, Return> {
             callback,
         )
     }
+
+    /// handle SIGVTALRM
+    #[cfg(unix)]
+    fn setup_sigvtalrm_handler() {
+        use nix::sys::signal::{sigaction, SaFlags, SigAction, SigHandler, SigSet, Signal};
+        use std::sync::atomic::{AtomicBool, Ordering};
+        static CANCEL_HANDLER_INITED: AtomicBool = AtomicBool::new(false);
+        if CANCEL_HANDLER_INITED
+            .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
+            .is_ok()
+        {
+            extern "C" fn sigvtalrm_handler<Param, Yield>(_: libc::c_int) {
+                if let Some(suspender) = suspender::Suspender::<Param, Yield>::current() {
+                    suspender.cancel();
+                }
+            }
+            // install SIGVTALRM signal handler
+            let mut set = SigSet::empty();
+            set.add(Signal::SIGVTALRM);
+            let sa = SigAction::new(
+                SigHandler::Handler(sigvtalrm_handler::<Param, Yield>),
+                SaFlags::SA_RESTART,
+                set,
+            );
+            unsafe {
+                _ = sigaction(Signal::SIGVTALRM, &sa).expect("install SIGVTALRM handler failed !");
+            }
+        }
+    }
 }
 
 impl<Yield, Return> Coroutine<'_, (), Yield, Return>
@@ -170,6 +199,8 @@ where
         }
         Self::init_current(self);
         self.running()?;
+        #[cfg(unix)]
+        Self::setup_sigvtalrm_handler();
         let r = self.raw_resume(arg);
         Self::clean_current();
         r