Plumbing: Add cx: context::Actor to PortHandle (#1350)

Summary: Pull Request resolved: #1350

Differential Revision: D83385432

Author: pzhan9

hyperactor/src/actor.rs

@@ -624,14 +624,13 @@ impl<A: Actor> ActorHandle<A> {
     /// are always queued in process, and do not require serialization.
     pub fn send<M: Message>(
         &self,
-        // TODO(pzhang): use this parameter to generate sequence number.
-        _cx: &impl context::Actor,
+        cx: &impl context::Actor,
         message: M,
     ) -> Result<(), MailboxSenderError>
     where
         A: Handler<M>,
     {
-        self.ports.get().send(message)
+        self.ports.get().send(cx, message)
     }
 
     /// Return a port for the provided message type handled by the actor.

hyperactor/src/mailbox.rs

@@ -17,14 +17,17 @@
 //!
 //! ```
 //! # use hyperactor::mailbox::Mailbox;
+//! # use hyperactor::Proc;
 //! # use hyperactor::reference::{ActorId, ProcId, WorldId};
 //! # tokio_test::block_on(async {
+//! # let proc = Proc::local();
+//! # let (client, _) = proc.instance("client").unwrap();
 //! # let proc_id = ProcId::Ranked(WorldId("world".to_string()), 0);
 //! # let actor_id = ActorId(proc_id, "actor".to_string(), 0);
 //! let mbox = Mailbox::new_detached(actor_id);
 //! let (port, mut receiver) = mbox.open_port::<u64>();
 //!
-//! port.send(123).unwrap();
+//! port.send(&client, 123).unwrap();
 //! assert_eq!(receiver.recv().await.unwrap(), 123u64);
 //! # })
 //! ```
@@ -1088,7 +1091,9 @@ impl MailboxClient {
             tokio::spawn(async move {
                 let result = return_receiver.await;
                 if let Ok(message) = result {
-                    let _ = return_handle_0.send(Undeliverable(message));
+                    // When returning messages, we do not care whether the messages are delivered
+                    // out of order.
+                    let _ = return_handle_0.anon_send(Undeliverable(message));
                 } else {
                     // Sender dropped, this task can end.
                 }
@@ -1578,9 +1583,24 @@ impl<M: Message> PortHandle<M> {
     }
 
     /// Send a message to this port.
-    pub fn send(&self, message: M) -> Result<(), MailboxSenderError> {
+    pub fn send(&self, _cx: &impl context::Actor, message: M) -> Result<(), MailboxSenderError> {
         let mut headers = Attrs::new();
 
+        crate::mailbox::headers::set_send_timestamp(&mut headers);
+        // TODO(pzhang) Use cx to add SEQ_INFO header.
+
+        self.sender.send(headers, message).map_err(|err| {
+            MailboxSenderError::new_unbound::<M>(
+                self.mailbox.actor_id().clone(),
+                MailboxSenderErrorKind::Other(err),
+            )
+        })
+    }
+
+    /// Send a message to this port without a known sender. This method should
+    /// only be used if you do not care about out-of-ordering delivery.
+    pub fn anon_send(&self, message: M) -> Result<(), MailboxSenderError> {
+        let mut headers = Attrs::new();
         crate::mailbox::headers::set_send_timestamp(&mut headers);
 
         self.sender.send(headers, message).map_err(|err| {
@@ -2610,6 +2630,7 @@ mod tests {
     use crate::channel::sim::SimAddr;
     use crate::clock::Clock;
     use crate::clock::RealClock;
+    use crate::context::Mailbox as _;
     use crate::data::Serialized;
     use crate::id;
     use crate::proc::Proc;
@@ -2653,32 +2674,33 @@ mod tests {
 
     #[tokio::test]
     async fn test_mailbox_accum() {
-        let mbox = Mailbox::new_detached(id!(test[0].test));
-        let (port, mut receiver) = mbox.open_accum_port(accum::max::<i64>());
+        let proc = Proc::local();
+        let (client, _) = proc.instance("client").unwrap();
+        let (port, mut receiver) = client.mailbox().open_accum_port(accum::max::<i64>());
 
         for i in -3..4 {
-            port.send(i).unwrap();
+            port.send(&client, i).unwrap();
             let received: accum::Max<i64> = receiver.recv().await.unwrap();
             let msg = received.get();
             assert_eq!(msg, &i);
         }
         // Send a smaller or same value. Should still receive the previous max.
         for i in -3..4 {
-            port.send(i).unwrap();
+            port.send(&client, i).unwrap();
             assert_eq!(receiver.recv().await.unwrap().get(), &3);
         }
         // send a larger value. Should receive the new max.
-        port.send(4).unwrap();
+        port.send(&client, 4).unwrap();
         assert_eq!(receiver.recv().await.unwrap().get(), &4);
 
         // Send multiple updates. Should only receive the final change.
         for i in 5..10 {
-            port.send(i).unwrap();
+            port.send(&client, i).unwrap();
         }
         assert_eq!(receiver.recv().await.unwrap().get(), &9);
-        port.send(1).unwrap();
-        port.send(3).unwrap();
-        port.send(2).unwrap();
+        port.send(&client, 1).unwrap();
+        port.send(&client, 3).unwrap();
+        port.send(&client, 2).unwrap();
         assert_eq!(receiver.recv().await.unwrap().get(), &9);
     }
 
@@ -2706,9 +2728,10 @@ mod tests {
     #[tokio::test]
     #[ignore] // error behavior changed, but we will bring it back
     async fn test_mailbox_once() {
-        let mbox = Mailbox::new_detached(id!(test[0].test));
+        let proc = Proc::local();
+        let (client, _) = proc.instance("client").unwrap();
 
-        let (port, receiver) = mbox.open_once_port::<u64>();
+        let (port, receiver) = client.open_once_port::<u64>();
 
         // let port_id = port.port_id().clone();
 
@@ -2977,19 +3000,20 @@ mod tests {
 
     #[tokio::test]
     async fn test_enqueue_port() {
-        let mbox = Mailbox::new_detached(id!(test[0].test));
+        let proc = Proc::local();
+        let (client, _) = proc.instance("client").unwrap();
 
         let count = Arc::new(AtomicUsize::new(0));
         let count_clone = count.clone();
-        let port = mbox.open_enqueue_port(move |_, n| {
+        let port = client.mailbox().open_enqueue_port(move |_, n| {
             count_clone.fetch_add(n, Ordering::SeqCst);
             Ok(())
         });
 
-        port.send(10).unwrap();
-        port.send(5).unwrap();
-        port.send(1).unwrap();
-        port.send(0).unwrap();
+        port.send(&client, 10).unwrap();
+        port.send(&client, 5).unwrap();
+        port.send(&client, 1).unwrap();
+        port.send(&client, 0).unwrap();
 
         assert_eq!(count.load(Ordering::SeqCst), 16);
     }
@@ -3054,6 +3078,7 @@ mod tests {
         let proc_id = id!(quux[0]);
         let mut proc = Proc::new(proc_id.clone(), proc_forwarder);
         ProcSupervisionCoordinator::set(&proc).await.unwrap();
+        let (client, _) = proc.instance("client").unwrap();
 
         let foo = proc.spawn::<Foo>("foo", ()).await.unwrap();
         let return_handle = foo.port::<Undeliverable<MessageEnvelope>>();
@@ -3063,7 +3088,7 @@ mod tests {
             Serialized::serialize(&1u64).unwrap(),
             Attrs::new(),
         );
-        return_handle.send(Undeliverable(message)).unwrap();
+        return_handle.send(&client, Undeliverable(message)).unwrap();
 
         RealClock
             .sleep(tokio::time::Duration::from_millis(100))
@@ -3095,7 +3120,9 @@ mod tests {
             Serialized::serialize(&1u64).unwrap(),
             Attrs::new(),
         );
-        return_handle.send(Undeliverable(envelope.clone())).unwrap();
+        return_handle
+            .anon_send(Undeliverable(envelope.clone()))
+            .unwrap();
         // Check we receive the undelivered message.
         assert!(
             RealClock
@@ -3747,12 +3774,13 @@ mod tests {
 
     #[tokio::test]
     async fn test_port_contramap() {
-        let mbox = Mailbox::new_detached(id!(test[0].test));
-        let (handle, mut rx) = mbox.open_port();
+        let proc = Proc::local();
+        let (client, _) = proc.instance("client").unwrap();
+        let (handle, mut rx) = client.open_port();
 
         handle
             .contramap(|m| (1, m))
-            .send("hello".to_string())
+            .send(&client, "hello".to_string())
             .unwrap();
         assert_eq!(rx.recv().await.unwrap(), (1, "hello".to_string()));
     }

hyperactor/src/mailbox/undeliverable.rs

@@ -17,6 +17,7 @@ use crate::ActorId;
 use crate::Message;
 use crate::Named;
 use crate::PortId;
+use crate::Proc;
 use crate::actor::ActorStatus;
 use crate::id;
 use crate::mailbox::DeliveryError;
@@ -94,7 +95,7 @@ pub(crate) fn return_undeliverable(
     envelope: MessageEnvelope,
 ) {
     let envelope_copy = envelope.clone();
-    if (return_handle.send(Undeliverable(envelope))).is_err() {
+    if (return_handle.anon_send(Undeliverable(envelope))).is_err() {
         UndeliverableMailboxSender.post(envelope_copy, /*unsued*/ return_handle)
     }
 }
@@ -157,6 +158,9 @@ pub fn supervise_undeliverable_messages_with<R, F>(
     F: Fn(&MessageEnvelope) + Send + Sync + 'static,
 {
     crate::init::get_runtime().spawn(async move {
+        // Create a local client for this task.
+        let proc = Proc::local();
+        let (client, _) = proc.instance("undeliverable_supervisor").unwrap();
         while let Ok(Undeliverable(mut env)) = rx.recv().await {
             // Let caller log/trace before we mutate.
             on_undeliverable(&env);
@@ -173,12 +177,15 @@ pub fn supervise_undeliverable_messages_with<R, F>(
                     let actor_id = env.dest().actor_id().clone();
                     let headers = env.headers().clone();
 
-                    if let Err(e) = sink.send(ActorSupervisionEvent::new(
-                        actor_id,
-                        ActorStatus::Failed(format!("message not delivered: {}", env)),
-                        Some(headers),
-                        None,
-                    )) {
+                    if let Err(e) = sink.send(
+                        &client,
+                        ActorSupervisionEvent::new(
+                            actor_id,
+                            ActorStatus::Failed(format!("message not delivered: {}", env)),
+                            Some(headers),
+                            None,
+                        ),
+                    ) {
                         tracing::warn!(
                             %e,
                             actor=%env.dest().actor_id(),

hyperactor/src/proc.rs

@@ -389,9 +389,9 @@ impl Proc {
             .map_err(|existing| anyhow::anyhow!("coordinator port is already set to {existing}"))
     }
 
-    fn handle_supervision_event(&self, event: ActorSupervisionEvent) {
+    fn handle_supervision_event(&self, cx: &impl context::Actor, event: ActorSupervisionEvent) {
         let result = match self.state().supervision_coordinator_port.get() {
-            Some(port) => port.send(event).map_err(anyhow::Error::from),
+            Some(port) => port.send(cx, event).map_err(anyhow::Error::from),
             None => Err(anyhow::anyhow!(
                 "coordinator port is not set for proc {}",
                 self.proc_id()
@@ -1044,7 +1044,9 @@ impl<A: Actor> Instance<A> {
         let clock = self.proc.state().clock.clone();
         tokio::spawn(async move {
             clock.non_advancing_sleep(delay).await;
-            if let Err(e) = port.send(message) {
+            // There is only one message from this context, so there is no need
+            // to worry out-of-order delivery.
+            if let Err(e) = port.anon_send(message) {
                 // TODO: this is a fire-n-forget thread. We need to
                 // handle errors in a better way.
                 tracing::info!("{}: error sending delayed message: {}", self_id, e);
@@ -1113,7 +1115,7 @@ impl<A: Actor> Instance<A> {
         if let Some(parent) = self.cell.maybe_unlink_parent() {
             if let Some(event) = event {
                 // Parent exists, failure should be propagated to the parent.
-                parent.send_supervision_event_or_crash(event);
+                parent.send_supervision_event_or_crash(&self, event);
             }
             // TODO: we should get rid of this signal, and use *only* supervision events for
             // the purpose of conveying lifecycle changes
@@ -1132,7 +1134,7 @@ impl<A: Actor> Instance<A> {
             // Note that orphaned actor is unexpected and would only happen if
             // there is a bug.
             if let Some(event) = event {
-                self.proc.handle_supervision_event(event);
+                self.proc.handle_supervision_event(&self, event);
             }
         }
         self.change_status(actor_status);
@@ -1667,7 +1669,9 @@ impl InstanceCell {
     #[allow(clippy::result_large_err)] // TODO: Consider reducing the size of `ActorError`.
     pub fn signal(&self, signal: Signal) -> Result<(), ActorError> {
         if let Some((signal_port, _)) = &self.inner.actor_loop {
-            signal_port.send(signal).map_err(ActorError::from)
+            // The owner of InstanceCell would not use PortRef to send signal.
+            // So we do not need to worry about out-of-order delivery here.
+            signal_port.anon_send(signal).map_err(ActorError::from)
         } else {
             tracing::warn!(
                 "{}: attempted to send signal {} to detached actor",
@@ -1686,10 +1690,14 @@ impl InstanceCell {
     /// Note that "let it crash" is the default behavior when a supervision event
     /// cannot be delivered upstream. It is the upstream's responsibility to
     /// detect and handle crashes.
-    pub fn send_supervision_event_or_crash(&self, event: ActorSupervisionEvent) {
+    pub fn send_supervision_event_or_crash(
+        &self,
+        cx: &impl context::Actor,
+        event: ActorSupervisionEvent,
+    ) {
         match &self.inner.actor_loop {
             Some((_, supervision_port)) => {
-                if let Err(err) = supervision_port.send(event) {
+                if let Err(err) = supervision_port.send(cx, event) {
                     tracing::error!(
                         "{}: failed to send supervision event to actor: {:?}. Crash the process.",
                         self.actor_id(),
@@ -2662,7 +2670,7 @@ mod tests {
         let (tx, rx) = client.open_once_port();
         handle.send(&client, tx).unwrap();
         let usize_handle = rx.recv().await.unwrap();
-        usize_handle.send(123).unwrap();
+        usize_handle.send(&client, 123).unwrap();
 
         handle.drain_and_stop().unwrap();
         handle.await;

hyperactor_multiprocess/src/system_actor.rs

@@ -1652,11 +1652,14 @@ impl Handler<MaintainWorldHealth> for SystemActor {
 
                 // The proc has expired heartbeating and it manages the lifecycle of system, schedule system stop
                 let (tx, _) = cx.open_once_port::<()>();
-                cx.port().send(SystemMessage::Stop {
-                    worlds: None,
-                    proc_timeout: Duration::from_secs(5),
-                    reply_port: tx.bind(),
-                })?;
+                cx.port().send(
+                    &cx,
+                    SystemMessage::Stop {
+                        worlds: None,
+                        proc_timeout: Duration::from_secs(5),
+                        reply_port: tx.bind(),
+                    },
+                )?;
             }
 
             if world.state.status == WorldStatus::Live {

monarch_hyperactor/src/mailbox.rs

@@ -249,11 +249,12 @@ pub(super) struct PythonPortHandle {
 
 #[pymethods]
 impl PythonPortHandle {
-    // TODO(pzhang) Use instance after its required by PortHandle.
-    fn send(&self, _instance: &PyInstance, message: PythonMessage) -> PyResult<()> {
-        self.inner
-            .send(message)
-            .map_err(|err| PyErr::new::<PyEOFError, _>(format!("Port closed: {}", err)))?;
+    fn send(&self, instance: &PyInstance, message: PythonMessage) -> PyResult<()> {
+        instance_dispatch!(instance, |cx_instance| {
+            self.inner
+                .send(cx_instance, message)
+                .map_err(|err| PyErr::new::<PyEOFError, _>(format!("Port closed: {}", err)))?;
+        });
         Ok(())
     }
 

monarch_hyperactor/src/proc.rs

@@ -680,7 +680,7 @@ async fn check_actor_supervision_state(
             // TODO: should allow for multiple attempts
             tracing::error!("system actor is not alive, aborting!");
             // Send a signal to the client to abort.
-            signal_port.send(Signal::Stop).unwrap();
+            signal_port.send(&instance, Signal::Stop).unwrap();
         }
     }
     Ok(())