fix: timers wake immediately on actor shutdown in threads module (#142)

* feat: add unified ctrl_c() signal handling to spawned-rt

* Formatted several files

* fix: use JoinHandle for Actor::join() instead of polling

* fix: timers wake immediately on actor shutdown in threads module

* Removed message about reverting state that no longer holds

* fix: allow multiple actors to subscribe to ctrl_c signal

* test: update signal_test to use two actors

* fix: use spawn_blocking for Thread backend join and fix flaky stream test

* fix: handle poisoned mutex in ctrl_c subscriber registration

* fix: race condition in CancellationToken::on_cancel

* Fixed typo

Author: ElFantasma

concurrency/src/tasks/actor.rs

@@ -456,10 +456,10 @@ where
     U: Future + Send + 'static,
     <U as Future>::Output: Send,
 {
-    let cancelation_token = handle.cancellation_token();
+    let cancellation_token = handle.cancellation_token();
     let mut handle_clone = handle.clone();
     let join_handle = rt::spawn(async move {
-        let is_cancelled = pin!(cancelation_token.cancelled());
+        let is_cancelled = pin!(cancellation_token.cancelled());
         let signal = pin!(future);
         match future::select(is_cancelled, signal).await {
             future::Either::Left(_) => tracing::debug!("Actor stopped"),

concurrency/src/tasks/stream.rs

@@ -13,10 +13,10 @@ where
     T: Actor,
     S: Send + Stream<Item = T::Message> + 'static,
 {
-    let cancelation_token = handle.cancellation_token();
+    let cancellation_token = handle.cancellation_token();
     let join_handle = spawned_rt::tasks::spawn(async move {
         let mut pinned_stream = core::pin::pin!(stream);
-        let is_cancelled = core::pin::pin!(cancelation_token.cancelled());
+        let is_cancelled = core::pin::pin!(cancellation_token.cancelled());
         let listener_loop = core::pin::pin!(async {
             loop {
                 match pinned_stream.next().await {

concurrency/src/threads/actor.rs

@@ -1,6 +1,8 @@
 //! Actor trait and structs to create an abstraction similar to Erlang gen_server.
 //! See examples/name_server for a usage example.
-use spawned_rt::threads::{self as rt, mpsc, oneshot, oneshot::RecvTimeoutError, CancellationToken};
+use spawned_rt::threads::{
+    self as rt, mpsc, oneshot, oneshot::RecvTimeoutError, CancellationToken,
+};
 use std::{
     fmt::Debug,
     panic::{catch_unwind, AssertUnwindSafe},
@@ -164,7 +166,7 @@ pub trait Actor: Send + Sized {
         handle: &ActorRef<Self>,
         rx: &mut mpsc::Receiver<ActorInMsg<Self>>,
     ) -> Result<(), ActorError> {
-        let mut cancellation_token = handle.cancellation_token.clone();
+        let cancellation_token = handle.cancellation_token.clone();
 
         let res = match self.init(handle) {
             Ok(InitResult::Success(new_state)) => {
@@ -235,7 +237,7 @@ pub trait Actor: Send + Sized {
                         }
                     },
                     Err(error) => {
-                        tracing::trace!("Error in callback, reverting state - Error: '{error:?}'");
+                        tracing::error!("Error in callback: '{error:?}'");
                         (true, Err(ActorError::Callback))
                     }
                 };
@@ -256,7 +258,7 @@ pub trait Actor: Send + Sized {
                         }
                     },
                     Err(error) => {
-                        tracing::trace!("Error in callback, reverting state - Error: '{error:?}'");
+                        tracing::error!("Error in callback: '{error:?}'");
                         true
                     }
                 }
@@ -301,7 +303,7 @@ where
     T: Actor,
     F: FnOnce() + Send + 'static,
 {
-    let mut cancellation_token = handle.cancellation_token();
+    let cancellation_token = handle.cancellation_token();
     let mut handle_clone = handle.clone();
     rt::spawn(move || {
         f();

concurrency/src/threads/stream.rs

@@ -12,7 +12,7 @@ where
     <I as IntoIterator>::IntoIter: std::marker::Send + 'static,
 {
     let mut iter = stream.into_iter();
-    let mut cancelation_token = handle.cancellation_token();
+    let cancellation_token = handle.cancellation_token();
     let join_handle = spawned_rt::threads::spawn(move || loop {
         match iter.next() {
             Some(msg) => match handle.send(msg) {
@@ -27,7 +27,7 @@ where
                 break;
             }
         }
-        if cancelation_token.is_cancelled() {
+        if cancellation_token.is_cancelled() {
             tracing::trace!("Actor stopped");
             break;
         }

concurrency/src/threads/time.rs

@@ -1,3 +1,4 @@
+use std::sync::mpsc::{self, RecvTimeoutError};
 use std::time::Duration;
 
 use spawned_rt::threads::{self as rt, CancellationToken, JoinHandle};
@@ -9,29 +10,58 @@ pub struct TimerHandle {
     pub cancellation_token: CancellationToken,
 }
 
-// Sends a message after a given period to the specified Actor. The task terminates
-// once the send has completed
+/// Sends a message after a given period to the specified Actor.
+///
+/// The timer respects both its own cancellation token and the Actor's
+/// cancellation token. If either is cancelled, the timer wakes up immediately
+/// and exits without sending the message.
 pub fn send_after<T>(period: Duration, mut handle: ActorRef<T>, message: T::Message) -> TimerHandle
 where
     T: Actor + 'static,
 {
     let cancellation_token = CancellationToken::new();
-    let mut cloned_token = cancellation_token.clone();
-    let mut actor_cancellation_token = handle.cancellation_token();
+    let timer_token = cancellation_token.clone();
+    let actor_token = handle.cancellation_token();
+
+    // Channel to wake the timer thread on cancellation
+    let (wake_tx, wake_rx) = mpsc::channel::<()>();
+
+    // Register wake-up on timer cancellation
+    let wake_tx1 = wake_tx.clone();
+    timer_token.on_cancel(Box::new(move || {
+        let _ = wake_tx1.send(());
+    }));
+
+    // Register wake-up on actor cancellation
+    actor_token.on_cancel(Box::new(move || {
+        let _ = wake_tx.send(());
+    }));
+
     let join_handle = rt::spawn(move || {
-        rt::sleep(period);
-        // Only send if neither the timer nor the actor has been cancelled
-        if !cloned_token.is_cancelled() && !actor_cancellation_token.is_cancelled() {
-            let _ = handle.send(message);
-        };
+        match wake_rx.recv_timeout(period) {
+            Err(RecvTimeoutError::Timeout) => {
+                // Timer expired - send if still valid
+                if !timer_token.is_cancelled() && !actor_token.is_cancelled() {
+                    let _ = handle.send(message);
+                }
+            }
+            Ok(()) | Err(RecvTimeoutError::Disconnected) => {
+                // Woken early by cancellation - exit without sending
+            }
+        }
     });
+
     TimerHandle {
         join_handle,
         cancellation_token,
     }
 }
 
-// Sends a message to the specified Actor repeatedly after `Time` milliseconds.
+/// Sends a message to the specified Actor repeatedly at the given interval.
+///
+/// The timer respects both its own cancellation token and the Actor's
+/// cancellation token. If either is cancelled, the timer wakes up immediately
+/// and exits.
 pub fn send_interval<T>(
     period: Duration,
     mut handle: ActorRef<T>,
@@ -41,17 +71,34 @@ where
     T: Actor + 'static,
 {
     let cancellation_token = CancellationToken::new();
-    let mut cloned_token = cancellation_token.clone();
-    let mut actor_cancellation_token = handle.cancellation_token();
-    let join_handle = rt::spawn(move || loop {
-        rt::sleep(period);
-        // Stop if either the timer or the actor has been cancelled
-        if cloned_token.is_cancelled() || actor_cancellation_token.is_cancelled() {
-            break;
-        } else {
+    let timer_token = cancellation_token.clone();
+    let actor_token = handle.cancellation_token();
+
+    // Channel to wake the timer thread on cancellation
+    let (wake_tx, wake_rx) = mpsc::channel::<()>();
+
+    // Register wake-up on timer cancellation
+    let wake_tx1 = wake_tx.clone();
+    timer_token.on_cancel(Box::new(move || {
+        let _ = wake_tx1.send(());
+    }));
+
+    // Register wake-up on actor cancellation
+    actor_token.on_cancel(Box::new(move || {
+        let _ = wake_tx.send(());
+    }));
+
+    let join_handle = rt::spawn(move || {
+        while let Err(RecvTimeoutError::Timeout) = wake_rx.recv_timeout(period) {
+            // Timer expired - send if still valid
+            if timer_token.is_cancelled() || actor_token.is_cancelled() {
+                break;
+            }
             let _ = handle.send(message.clone());
-        };
+        }
+        // If we exit the loop via Ok(()) or Disconnected, cancellation occurred
     });
+
     TimerHandle {
         join_handle,
         cancellation_token,

concurrency/src/threads/timer_tests.rs

@@ -86,7 +86,7 @@ impl Actor for Repeater {
                 self.count += 1;
             }
             RepeaterCastMessage::StopTimer => {
-                if let Some(mut ct) = self.cancellation_token.clone() {
+                if let Some(ct) = self.cancellation_token.clone() {
                     ct.cancel()
                 };
             }
@@ -132,6 +132,7 @@ enum DelayedCastMessage {
 #[derive(Clone)]
 enum DelayedCallMessage {
     GetCount,
+    Stop,
 }
 
 #[derive(PartialEq, Debug)]
@@ -156,6 +157,10 @@ impl Delayed {
     pub fn get_count(server: &mut DelayedHandle) -> Result<DelayedOutMessage, ()> {
         server.request(DelayedCallMessage::GetCount).map_err(|_| ())
     }
+
+    pub fn stop(server: &mut DelayedHandle) -> Result<DelayedOutMessage, ()> {
+        server.request(DelayedCallMessage::Stop).map_err(|_| ())
+    }
 }
 
 impl Actor for Delayed {
@@ -166,11 +171,17 @@ impl Actor for Delayed {
 
     fn handle_request(
         &mut self,
-        _message: Self::Request,
+        message: Self::Request,
         _handle: &DelayedHandle,
     ) -> RequestResponse<Self> {
-        let count = self.count;
-        RequestResponse::Reply(DelayedOutMessage::Count(count))
+        match message {
+            DelayedCallMessage::GetCount => {
+                RequestResponse::Reply(DelayedOutMessage::Count(self.count))
+            }
+            DelayedCallMessage::Stop => {
+                RequestResponse::Stop(DelayedOutMessage::Count(self.count))
+            }
+        }
     }
 
     fn handle_message(
@@ -209,7 +220,7 @@ pub fn test_send_after_and_cancellation() {
     assert_eq!(DelayedOutMessage::Count(1), count);
 
     // New timer
-    let mut timer = send_after(
+    let timer = send_after(
         Duration::from_millis(100),
         repeater.clone(),
         DelayedCastMessage::Inc,
@@ -227,3 +238,41 @@ pub fn test_send_after_and_cancellation() {
     // As timer was cancelled, count should remain at 1
     assert_eq!(DelayedOutMessage::Count(1), count2);
 }
+
+#[test]
+pub fn test_send_after_actor_shutdown() {
+    // Start a Delayed
+    let mut actor = Delayed::new(0).start();
+
+    // Set a just once timed message
+    let _ = send_after(
+        Duration::from_millis(100),
+        actor.clone(),
+        DelayedCastMessage::Inc,
+    );
+
+    // Wait for 200 milliseconds
+    rt::sleep(Duration::from_millis(200));
+
+    // Check count
+    let count = Delayed::get_count(&mut actor).unwrap();
+
+    // Only one message (no repetition)
+    assert_eq!(DelayedOutMessage::Count(1), count);
+
+    // New timer with long delay
+    let _ = send_after(
+        Duration::from_millis(100),
+        actor.clone(),
+        DelayedCastMessage::Inc,
+    );
+
+    // Stop the Actor before timeout - this should wake up the timer immediately
+    let count2 = Delayed::stop(&mut actor).unwrap();
+
+    // Wait another 200 milliseconds
+    rt::sleep(Duration::from_millis(200));
+
+    // As actor was stopped, count should remain at 1 (timer didn't fire)
+    assert_eq!(DelayedOutMessage::Count(1), count2);
+}

rt/src/threads/mod.rs

@@ -37,24 +37,69 @@ where
     spawn(f)
 }
 
-#[derive(Clone, Debug, Default)]
+type CancelCallback = Box<dyn FnOnce() + Send>;
+
+/// A token that can be used to signal cancellation.
+///
+/// Supports registering callbacks via `on_cancel()` that fire when
+/// the token is cancelled, enabling efficient waiting patterns.
+#[derive(Clone, Default)]
 pub struct CancellationToken {
     is_cancelled: Arc<AtomicBool>,
+    callbacks: Arc<Mutex<Vec<CancelCallback>>>,
+}
+
+impl std::fmt::Debug for CancellationToken {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("CancellationToken")
+            .field("is_cancelled", &self.is_cancelled())
+            .finish()
+    }
 }
 
 impl CancellationToken {
     pub fn new() -> Self {
         CancellationToken {
             is_cancelled: Arc::new(false.into()),
+            callbacks: Arc::new(Mutex::new(Vec::new())),
         }
     }
 
-    pub fn is_cancelled(&mut self) -> bool {
-        self.is_cancelled.fetch_and(false, Ordering::SeqCst)
+    pub fn is_cancelled(&self) -> bool {
+        self.is_cancelled.load(Ordering::SeqCst)
     }
 
-    pub fn cancel(&mut self) {
-        self.is_cancelled.fetch_or(true, Ordering::SeqCst);
+    pub fn cancel(&self) {
+        self.is_cancelled.store(true, Ordering::SeqCst);
+        // Fire all registered callbacks
+        let callbacks: Vec<_> = self
+            .callbacks
+            .lock()
+            .unwrap_or_else(|e| e.into_inner())
+            .drain(..)
+            .collect();
+        for cb in callbacks {
+            cb();
+        }
+    }
+
+    /// Register a callback to be invoked when this token is cancelled.
+    /// If already cancelled, the callback fires immediately.
+    ///
+    /// This method is thread-safe: the callback is guaranteed to fire exactly
+    /// once, either immediately (if already cancelled) or when `cancel()` is called.
+    pub fn on_cancel(&self, callback: CancelCallback) {
+        // Hold the lock while checking is_cancelled to avoid a race with cancel().
+        // cancel() sets the flag BEFORE acquiring the lock, so if we see
+        // is_cancelled=false while holding the lock, cancel() hasn't drained
+        // callbacks yet and will drain ours after we release the lock.
+        let mut callbacks = self.callbacks.lock().unwrap_or_else(|e| e.into_inner());
+        if self.is_cancelled() {
+            drop(callbacks);
+            callback();
+        } else {
+            callbacks.push(callback);
+        }
     }
 }