Plumbing: Use context::Actor for Rust PortRef and Port (#1370)

Summary:
Pull Request resolved: #1370

This diff is part of the effect to adding sequencing logic to sender actor. See D83371710 for details.

This diff updates the `send` methods in Rust `PortRef` and `Port` to use `context::Actor`.

Reviewed By: pablorfb-meta, mariusae

Differential Revision: D83507394

fbshipit-source-id: fae0a83192628a33092526d06ef27b7e6daf9308

Author: pzhan9

hyperactor/src/mailbox.rs

@@ -1164,27 +1164,27 @@ impl MailboxSender for MailboxClient {
 }
 
 /// Wrapper to turn `PortRef` into a `Sink`.
-pub struct PortSink<C: context::Mailbox, M: RemoteMessage> {
-    caps: C,
+pub struct PortSink<C: context::Actor, M: RemoteMessage> {
+    cx: C,
     port: PortRef<M>,
 }
 
-impl<C: context::Mailbox, M: RemoteMessage> PortSink<C, M> {
+impl<C: context::Actor, M: RemoteMessage> PortSink<C, M> {
     /// Create new PortSink
-    pub fn new(caps: C, port: PortRef<M>) -> Self {
-        Self { caps, port }
+    pub fn new(cx: C, port: PortRef<M>) -> Self {
+        Self { cx, port }
     }
 }
 
-impl<C: context::Mailbox, M: RemoteMessage> Sink<M> for PortSink<C, M> {
+impl<C: context::Actor, M: RemoteMessage> Sink<M> for PortSink<C, M> {
     type Error = MailboxSenderError;
 
     fn poll_ready(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         Poll::Ready(Ok(()))
     }
 
     fn start_send(self: Pin<&mut Self>, item: M) -> Result<(), Self::Error> {
-        self.port.send(&self.caps, item)
+        self.port.send(&self.cx, item)
     }
 
     fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {

hyperactor/src/reference.rs

@@ -72,9 +72,6 @@ use parse::ParseError;
 use parse::Token;
 use parse::parse;
 
-use crate::proc::SEQ_INFO;
-use crate::proc::SeqInfo;
-
 /// A universal reference to hierarchical identifiers in Hyperactor.
 ///
 /// References implement a concrete syntax which can be parsed via
@@ -723,7 +720,7 @@ impl<A: RemoteActor> ActorRef<A> {
     /// Send an [`M`]-typed message to the referenced actor.
     pub fn send<M: RemoteMessage>(
         &self,
-        cx: &impl context::Mailbox,
+        cx: &impl context::Actor,
         message: M,
     ) -> Result<(), MailboxSenderError>
     where
@@ -736,7 +733,7 @@ impl<A: RemoteActor> ActorRef<A> {
     /// headers.
     pub fn send_with_headers<M: RemoteMessage>(
         &self,
-        cx: &impl context::Mailbox,
+        cx: &impl context::Actor,
         headers: Attrs,
         message: M,
     ) -> Result<(), MailboxSenderError>
@@ -883,7 +880,7 @@ impl PortId {
     /// Send a serialized message to this port, provided a sending capability,
     /// such as [`crate::actor::Instance`]. It is the sender's responsibility
     /// to ensure that the provided message is well-typed.
-    pub fn send(&self, cx: &impl context::Mailbox, serialized: Serialized) {
+    pub fn send(&self, cx: &impl context::Actor, serialized: Serialized) {
         let mut headers = Attrs::new();
         crate::mailbox::headers::set_send_timestamp(&mut headers);
         cx.post(self.clone(), headers, serialized);
@@ -894,7 +891,7 @@ impl PortId {
     /// It is the sender's responsibility to ensure that the provided message is well-typed.
     pub fn send_with_headers(
         &self,
-        cx: &impl context::Mailbox,
+        cx: &impl context::Actor,
         serialized: Serialized,
         mut headers: Attrs,
     ) {
@@ -1014,7 +1011,7 @@ impl<M: RemoteMessage> PortRef<M> {
 
     /// Send a message to this port, provided a sending capability, such as
     /// [`crate::actor::Instance`].
-    pub fn send(&self, cx: &impl context::Mailbox, message: M) -> Result<(), MailboxSenderError> {
+    pub fn send(&self, cx: &impl context::Actor, message: M) -> Result<(), MailboxSenderError> {
         self.send_with_headers(cx, Attrs::new(), message)
     }
 
@@ -1023,7 +1020,7 @@ impl<M: RemoteMessage> PortRef<M> {
     /// headers.
     pub fn send_with_headers(
         &self,
-        cx: &impl context::Mailbox,
+        cx: &impl context::Actor,
         headers: Attrs,
         message: M,
     ) -> Result<(), MailboxSenderError> {
@@ -1041,7 +1038,7 @@ impl<M: RemoteMessage> PortRef<M> {
     /// [`crate::actor::Instance`].
     pub fn send_serialized(
         &self,
-        cx: &impl context::Mailbox,
+        cx: &impl context::Actor,
         mut headers: Attrs,
         message: Serialized,
     ) {
@@ -1050,7 +1047,7 @@ impl<M: RemoteMessage> PortRef<M> {
     }
 
     /// Convert this port into a sink that can be used to send messages using the given capability.
-    pub fn into_sink<C: context::Mailbox>(self, cx: C) -> PortSink<C, M> {
+    pub fn into_sink<C: context::Actor>(self, cx: C) -> PortSink<C, M> {
         PortSink::new(cx, self)
     }
 }
@@ -1138,15 +1135,15 @@ impl<M: RemoteMessage> OncePortRef<M> {
 
     /// Send a message to this port, provided a sending capability, such as
     /// [`crate::actor::Instance`].
-    pub fn send(self, cx: &impl context::Mailbox, message: M) -> Result<(), MailboxSenderError> {
+    pub fn send(self, cx: &impl context::Actor, message: M) -> Result<(), MailboxSenderError> {
         self.send_with_headers(cx, Attrs::new(), message)
     }
 
     /// Send a message to this port, provided a sending capability, such as
     /// [`crate::actor::Instance`]. Additional context can be provided in the form of headers.
     pub fn send_with_headers(
         self,
-        cx: &impl context::Mailbox,
+        cx: &impl context::Actor,
         mut headers: Attrs,
         message: M,
     ) -> Result<(), MailboxSenderError> {

hyperactor_macros/tests/export.rs

@@ -111,7 +111,7 @@ mod tests {
     #[async_timed_test(timeout_secs = 30)]
     async fn test_binds() {
         let proc = Proc::local();
-        let client = proc.attach("client").unwrap();
+        let (client, _) = proc.instance("client").unwrap();
         let (tx, mut rx) = client.open_port();
         let params = TestActorParams {
             forward_port: tx.bind(),
@@ -192,7 +192,7 @@ mod tests {
     #[async_timed_test(timeout_secs = 30)]
     async fn test_ref_alias() {
         let proc = Proc::local();
-        let client = proc.attach("client").unwrap();
+        let (client, _) = proc.instance("client").unwrap();
         let (tx, mut rx) = client.open_port();
         let params = TestActorParams {
             forward_port: tx.bind(),

monarch_hyperactor/src/mailbox.rs

@@ -48,6 +48,7 @@ use serde::Serialize;
 use crate::actor::PythonMessage;
 use crate::actor::PythonMessageKind;
 use crate::context::PyInstance;
+use crate::instance_dispatch;
 use crate::proc::PyActorId;
 use crate::pytokio::PyPythonTask;
 use crate::pytokio::PythonTask;
@@ -321,9 +322,11 @@ impl PythonPortRef {
     }
 
     fn send(&self, instance: &PyInstance, message: PythonMessage) -> PyResult<()> {
-        self.inner
-            .send(&instance._mailbox().inner, message)
-            .map_err(|err| PyErr::new::<PyEOFError, _>(format!("Port closed: {}", err)))?;
+        instance_dispatch!(instance, |cx_instance| {
+            self.inner
+                .send(cx_instance, message)
+                .map_err(|err| PyErr::new::<PyEOFError, _>(format!("Port closed: {}", err)))?;
+        });
         Ok(())
     }
 
@@ -537,9 +540,11 @@ impl PythonOncePortRef {
             return Err(PyErr::new::<PyValueError, _>("OncePortRef is already used"));
         };
 
-        port_ref
-            .send(&instance._mailbox().inner, message)
-            .map_err(|err| PyErr::new::<PyEOFError, _>(format!("Port closed: {}", err)))?;
+        instance_dispatch!(instance, |cx_instance| {
+            port_ref
+                .send(cx_instance, message)
+                .map_err(|err| PyErr::new::<PyEOFError, _>(format!("Port closed: {}", err)))?;
+        });
         Ok(())
     }