Add stop behaviour to GenServer

concurrency/src/tasks/gen_server.rs

@@ -1,28 +1,34 @@
 //! GenServer trait and structs to create an abstraction similar to Erlang gen_server.
 //! See examples/name_server for a usage example.
 use futures::future::FutureExt as _;
-use spawned_rt::tasks::{self as rt, mpsc, oneshot};
+use spawned_rt::tasks::{self as rt, mpsc, oneshot, CancellationToken};
 use std::{fmt::Debug, future::Future, panic::AssertUnwindSafe};
 
 use crate::error::GenServerError;
 
-#[derive(Debug)]
 pub struct GenServerHandle<G: GenServer + 'static> {
     pub tx: mpsc::Sender<GenServerInMsg<G>>,
+    /// Cancellation token to stop the GenServer
+    cancellation_token: CancellationToken,
 }
 
 impl<G: GenServer> Clone for GenServerHandle<G> {
     fn clone(&self) -> Self {
         Self {
             tx: self.tx.clone(),
+            cancellation_token: self.cancellation_token.clone(),
         }
     }
 }
 
 impl<G: GenServer> GenServerHandle<G> {
     pub(crate) fn new(initial_state: G::State) -> Self {
         let (tx, mut rx) = mpsc::channel::<GenServerInMsg<G>>();
-        let handle = GenServerHandle { tx };
+        let cancellation_token = CancellationToken::new();
+        let handle = GenServerHandle {
+            tx,
+            cancellation_token,
+        };
         let mut gen_server: G = GenServer::new();
         let handle_clone = handle.clone();
         // Ignore the JoinHandle for now. Maybe we'll use it in the future
@@ -40,7 +46,11 @@ impl<G: GenServer> GenServerHandle<G> {
 
     pub(crate) fn new_blocking(initial_state: G::State) -> Self {
         let (tx, mut rx) = mpsc::channel::<GenServerInMsg<G>>();
-        let handle = GenServerHandle { tx };
+        let cancellation_token = CancellationToken::new();
+        let handle = GenServerHandle {
+            tx,
+            cancellation_token,
+        };
         let mut gen_server: G = GenServer::new();
         let handle_clone = handle.clone();
         // Ignore the JoinHandle for now. Maybe we'll use it in the future
@@ -79,6 +89,14 @@ impl<G: GenServer> GenServerHandle<G> {
             .send(GenServerInMsg::Cast { message })
             .map_err(|_error| GenServerError::Server)
     }
+
+    pub fn cancellation_token(&self) -> CancellationToken {
+        self.cancellation_token.clone()
+    }
+
+    pub fn teardown(&mut self) {
+        self.cancellation_token.cancel();
+    }
 }
 
 pub enum GenServerInMsg<G: GenServer> {

concurrency/src/tasks/time.rs

@@ -22,9 +22,16 @@ where
 {
     let cancellation_token = CancellationToken::new();
     let cloned_token = cancellation_token.clone();
+    let gen_server_cancellation_token = handle.cancellation_token();
     let join_handle = rt::spawn(async move {
-        let _ = select(
+        // Timer action is ignored if it was either cancelled or the associated GenServer is no longer running.
+        let cancel_conditions = select(
             Box::pin(cloned_token.cancelled()),
+            Box::pin(gen_server_cancellation_token.cancelled()),
+        );
+
+        let _ = select(
+            cancel_conditions,
             Box::pin(async {
                 rt::sleep(period).await;
                 let _ = handle.cast(message.clone()).await;
@@ -49,10 +56,17 @@ where
 {
     let cancellation_token = CancellationToken::new();
     let cloned_token = cancellation_token.clone();
+    let gen_server_cancellation_token = handle.cancellation_token();
     let join_handle = rt::spawn(async move {
         loop {
-            let result = select(
+            // Timer action is ignored if it was either cancelled or the associated GenServer is no longer running.
+            let cancel_conditions = select(
                 Box::pin(cloned_token.cancelled()),
+                Box::pin(gen_server_cancellation_token.cancelled()),
+            );
+
+            let result = select(
+                Box::pin(cancel_conditions),
                 Box::pin(async {
                     rt::sleep(period).await;
                     let _ = handle.cast(message.clone()).await;

concurrency/src/tasks/timer_tests.rs

@@ -246,3 +246,47 @@ pub fn test_send_after_and_cancellation() {
         assert_eq!(DelayedOutMessage::Count(1), count2);
     });
 }
+
+#[test]
+pub fn test_send_after_gen_server_teardown() {
+    let runtime = rt::Runtime::new().unwrap();
+    runtime.block_on(async move {
+        // Start a Delayed
+        let mut repeater = Delayed::start(DelayedState { count: 0 });
+
+        // Set a just once timed message
+        let _ = send_after(
+            Duration::from_millis(100),
+            repeater.clone(),
+            DelayedCastMessage::Inc,
+        );
+
+        // Wait for 200 milliseconds
+        rt::sleep(Duration::from_millis(200)).await;
+
+        // Check count
+        let count = Delayed::get_count(&mut repeater).await.unwrap();
+
+        // Only one message (no repetition)
+        assert_eq!(DelayedOutMessage::Count(1), count);
+
+        // New timer
+        let _ = send_after(
+            Duration::from_millis(100),
+            repeater.clone(),
+            DelayedCastMessage::Inc,
+        );
+
+        // Cancel the new timer before timeout
+        repeater.teardown();
+
+        // Wait another 200 milliseconds
+        rt::sleep(Duration::from_millis(200)).await;
+
+        // Check count again
+        let count2 = Delayed::get_count(&mut repeater).await.unwrap();
+
+        // As timer was cancelled, count should remain at 1
+        assert_eq!(DelayedOutMessage::Count(1), count2);
+    });
+}