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

Author: ElFantasma

concurrency/src/tasks/actor.rs

@@ -10,10 +10,39 @@ use spawned_rt::{
     tasks::{self as rt, mpsc, oneshot, timeout, CancellationToken, JoinHandle},
     threads,
 };
-use std::{fmt::Debug, future::Future, panic::AssertUnwindSafe, time::Duration};
+use std::{
+    fmt::Debug,
+    future::Future,
+    panic::AssertUnwindSafe,
+    sync::{Arc, Mutex},
+    time::Duration,
+};
 
 const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(5);
 
+/// Wrapper for different JoinHandle types based on backend.
+#[derive(Debug)]
+enum ActorJoinHandle {
+    /// Tokio task JoinHandle (for Async and Blocking backends)
+    Task(JoinHandle<()>),
+    /// OS thread JoinHandle (for Thread backend)
+    Thread(threads::JoinHandle<()>),
+}
+
+impl ActorJoinHandle {
+    /// Waits for the actor to finish.
+    async fn join(self) {
+        match self {
+            ActorJoinHandle::Task(h) => {
+                let _ = h.await;
+            }
+            ActorJoinHandle::Thread(h) => {
+                let _ = h.join();
+            }
+        }
+    }
+}
+
 /// Execution backend for Actor.
 ///
 /// Determines how the Actor's async loop is executed. Choose based on
@@ -106,13 +135,16 @@ pub struct ActorRef<A: Actor + 'static> {
     pub tx: mpsc::Sender<ActorInMsg<A>>,
     /// Cancellation token to stop the Actor
     cancellation_token: CancellationToken,
+    /// JoinHandle for waiting on actor completion
+    join_handle: Arc<Mutex<Option<ActorJoinHandle>>>,
 }
 
 impl<A: Actor> Clone for ActorRef<A> {
     fn clone(&self) -> Self {
         Self {
             tx: self.tx.clone(),
             cancellation_token: self.cancellation_token.clone(),
+            join_handle: self.join_handle.clone(),
         }
     }
 }
@@ -121,9 +153,11 @@ impl<A: Actor> ActorRef<A> {
     fn new(actor: A) -> Self {
         let (tx, mut rx) = mpsc::channel::<ActorInMsg<A>>();
         let cancellation_token = CancellationToken::new();
+        let join_handle = Arc::new(Mutex::new(None));
         let handle = ActorRef {
             tx,
             cancellation_token,
+            join_handle: join_handle.clone(),
         };
         let handle_clone = handle.clone();
         let inner_future = async move {
@@ -136,47 +170,62 @@ impl<A: Actor> ActorRef<A> {
         // Optionally warn if the Actor future blocks for too much time
         let inner_future = warn_on_block::WarnOnBlocking::new(inner_future);
 
-        // Ignore the JoinHandle for now. Maybe we'll use it in the future
-        let _join_handle = rt::spawn(inner_future);
+        let task_handle = rt::spawn(inner_future);
+        let mut guard = join_handle
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner());
+        *guard = Some(ActorJoinHandle::Task(task_handle));
 
         handle_clone
     }
 
     fn new_blocking(actor: A) -> Self {
         let (tx, mut rx) = mpsc::channel::<ActorInMsg<A>>();
         let cancellation_token = CancellationToken::new();
+        let join_handle = Arc::new(Mutex::new(None));
         let handle = ActorRef {
             tx,
             cancellation_token,
+            join_handle: join_handle.clone(),
         };
         let handle_clone = handle.clone();
-        // Ignore the JoinHandle for now. Maybe we'll use it in the future
-        let _join_handle = rt::spawn_blocking(|| {
+        let task_handle = rt::spawn_blocking(|| {
             rt::block_on(async move {
                 if let Err(error) = actor.run(&handle, &mut rx).await {
                     tracing::trace!(%error, "Actor crashed")
                 };
             })
         });
+        let mut guard = join_handle
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner());
+        *guard = Some(ActorJoinHandle::Task(task_handle));
+
         handle_clone
     }
 
     fn new_on_thread(actor: A) -> Self {
         let (tx, mut rx) = mpsc::channel::<ActorInMsg<A>>();
         let cancellation_token = CancellationToken::new();
+        let join_handle = Arc::new(Mutex::new(None));
         let handle = ActorRef {
             tx,
             cancellation_token,
+            join_handle: join_handle.clone(),
         };
         let handle_clone = handle.clone();
-        // Ignore the JoinHandle for now. Maybe we'll use it in the future
-        let _join_handle = threads::spawn(|| {
+        let thread_handle = threads::spawn(|| {
             threads::block_on(async move {
                 if let Err(error) = actor.run(&handle, &mut rx).await {
                     tracing::trace!(%error, "Actor crashed")
                 };
             })
         });
+        let mut guard = join_handle
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner());
+        *guard = Some(ActorJoinHandle::Thread(thread_handle));
+
         handle_clone
     }
 
@@ -220,9 +269,19 @@ impl<A: Actor> ActorRef<A> {
     /// Waits for the actor to stop.
     ///
     /// This method returns a future that completes when the actor has finished
-    /// processing and exited its main loop.
+    /// processing and exited its main loop. Can only be called once; subsequent
+    /// calls return immediately.
     pub async fn join(&self) {
-        self.cancellation_token.cancelled().await
+        let handle = {
+            let mut guard = self
+                .join_handle
+                .lock()
+                .unwrap_or_else(|poisoned| poisoned.into_inner());
+            guard.take()
+        };
+        if let Some(h) = handle {
+            h.join().await;
+        }
     }
 }
 

concurrency/src/threads/actor.rs

@@ -1,11 +1,12 @@
 //! 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,
+    self as rt, mpsc, oneshot, oneshot::RecvTimeoutError, CancellationToken, JoinHandle,
 };
 use std::{
     fmt::Debug,
     panic::{catch_unwind, AssertUnwindSafe},
+    sync::{Arc, Mutex},
     time::Duration,
 };
 
@@ -17,13 +18,15 @@ const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(5);
 pub struct ActorRef<A: Actor + 'static> {
     pub tx: mpsc::Sender<ActorInMsg<A>>,
     cancellation_token: CancellationToken,
+    join_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
 }
 
 impl<A: Actor> Clone for ActorRef<A> {
     fn clone(&self) -> Self {
         Self {
             tx: self.tx.clone(),
             cancellation_token: self.cancellation_token.clone(),
+            join_handle: self.join_handle.clone(),
         }
     }
 }
@@ -32,17 +35,22 @@ impl<A: Actor> ActorRef<A> {
     pub(crate) fn new(actor: A) -> Self {
         let (tx, mut rx) = mpsc::channel::<ActorInMsg<A>>();
         let cancellation_token = CancellationToken::new();
+        let join_handle = Arc::new(Mutex::new(None));
         let handle = ActorRef {
             tx,
             cancellation_token,
+            join_handle: join_handle.clone(),
         };
         let handle_clone = handle.clone();
-        // Ignore the JoinHandle for now. Maybe we'll use it in the future
-        let _join_handle = rt::spawn(move || {
+        let thread_handle = rt::spawn(move || {
             if actor.run(&handle, &mut rx).is_err() {
                 tracing::trace!("Actor crashed")
             };
         });
+        let mut guard = join_handle
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner());
+        *guard = Some(thread_handle);
         handle_clone
     }
 
@@ -84,11 +92,16 @@ impl<A: Actor> ActorRef<A> {
     /// Blocks until the actor has stopped.
     ///
     /// This method blocks the current thread until the actor has finished
-    /// processing and exited its main loop.
+    /// processing and exited its main loop. Can only be called once; subsequent
+    /// calls return immediately.
     pub fn join(&self) {
-        let mut token = self.cancellation_token.clone();
-        while !token.is_cancelled() {
-            std::thread::sleep(std::time::Duration::from_millis(10));
+        // Recover from poisoned lock if another thread panicked while holding it
+        let mut guard = self
+            .join_handle
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner());
+        if let Some(h) = guard.take() {
+            let _ = h.join();
         }
     }
 }