Add asserts for actor status transition (#1825)

Summary:

Adding asserts so we can be more certain on the the actor status transition.

Reviewed By: mariusae

Differential Revision: D86772033

Author: pzhan9

hyperactor/src/actor.rs

@@ -20,6 +20,7 @@ use std::sync::Arc;
 use std::time::SystemTime;
 
 use async_trait::async_trait;
+use enum_as_inner::EnumAsInner;
 use futures::FutureExt;
 use futures::future::BoxFuture;
 use serde::Deserialize;
@@ -472,7 +473,16 @@ impl fmt::Display for Signal {
 }
 
 /// The runtime status of an actor.
-#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone, Named)]
+#[derive(
+    Debug,
+    Serialize,
+    Deserialize,
+    PartialEq,
+    Eq,
+    Clone,
+    Named,
+    EnumAsInner
+)]
 pub enum ActorStatus {
     /// The actor status is unknown.
     Unknown,
@@ -505,12 +515,7 @@ pub enum ActorStatus {
 impl ActorStatus {
     /// Tells whether the status is a terminal state.
     pub fn is_terminal(&self) -> bool {
-        matches!(self, Self::Stopped | Self::Failed(_))
-    }
-
-    /// Tells whether the status represents a failure.
-    pub fn is_failed(&self) -> bool {
-        matches!(self, Self::Failed(_))
+        self.is_stopped() || self.is_failed()
     }
 
     /// Create a generic failure status with the provided error message.

hyperactor/src/proc.rs

@@ -1028,6 +1028,14 @@ impl<A: Actor> Instance<A> {
         self.status_tx.send_replace(new.clone());
     }
 
+    fn is_terminal(&self) -> bool {
+        self.status_tx.borrow().is_terminal()
+    }
+
+    fn is_stopped(&self) -> bool {
+        self.status_tx.borrow().is_stopped()
+    }
+
     /// This instance's actor ID.
     pub fn self_id(&self) -> &ActorId {
         self.mailbox.actor_id()
@@ -1119,23 +1127,35 @@ impl<A: Actor> Instance<A> {
             .run_actor_tree(&mut actor, actor_loop_receivers, &mut work_rx)
             .await;
 
-        let (actor_status, event) = match result {
-            Ok(_) => (ActorStatus::Stopped, None),
-            Err(ActorError {
-                kind: box ActorErrorKind::UnhandledSupervisionEvent(box event),
-                ..
-            }) => (event.actor_status.clone(), Some(event)),
+        let event = match result {
+            Ok(_) => {
+                // actor should have been stopped by run_actor_tree
+                assert!(self.is_stopped());
+                None
+            }
             Err(err) => {
-                let error_kind = ActorErrorKind::Generic(err.kind.to_string());
-                (
-                    ActorStatus::Failed(error_kind.clone()),
-                    Some(ActorSupervisionEvent::new(
-                        self.cell.actor_id().clone(),
-                        ActorStatus::Failed(error_kind),
-                        None,
-                        None,
-                    )),
-                )
+                assert!(!self.is_terminal());
+                match *err.kind {
+                    ActorErrorKind::UnhandledSupervisionEvent(box event) => {
+                        // Currently only terminated actors are allowed to raise supervision events.
+                        // If we want to change that in the future, we need to modify the exit
+                        // status here too, because we use event's actor_status as this actor's
+                        // terminal status.
+                        assert!(event.actor_status.is_terminal());
+                        self.change_status(event.actor_status.clone());
+                        Some(event)
+                    }
+                    _ => {
+                        let error_kind = ActorErrorKind::Generic(err.kind.to_string());
+                        self.change_status(ActorStatus::Failed(error_kind.clone()));
+                        Some(ActorSupervisionEvent::new(
+                            self.cell.actor_id().clone(),
+                            ActorStatus::Failed(error_kind),
+                            None,
+                            None,
+                        ))
+                    }
+                }
             }
         };
 
@@ -1164,7 +1184,6 @@ impl<A: Actor> Instance<A> {
                 self.proc.handle_supervision_event(event);
             }
         }
-        self.change_status(actor_status);
     }
 
     /// Runs the actor, and manages its supervision tree. When the function returns,
@@ -1207,10 +1226,16 @@ impl<A: Actor> Instance<A> {
             }
         };
 
-        if let Err(ref err) = result {
-            tracing::error!("{}: actor failure: {}", self.self_id(), err);
+        match &result {
+            Ok(_) => assert!(self.is_stopped()),
+            Err(err) => {
+                tracing::error!("{}: actor failure: {}", self.self_id(), err);
+                assert!(!self.is_terminal());
+                // Send Stopping instead of Failed, because we still need to
+                // unlink child actors.
+                self.change_status(ActorStatus::Stopping);
+            }
         }
-        self.change_status(ActorStatus::Stopping);
 
         // After this point, we know we won't spawn any more children,
         // so we can safely read the current child keys.