tighten the usage of bind_to (meta-pytorch#1437)

Summary:
Pull Request resolved: meta-pytorch#1437

`PortHandle.bind_to` is not meant to be a public method. This diff makes it `pub(crate)`, rename it to `bind_to_actor_port` and remove its `port_index` parameter.

https://www.internalfb.com/code/fbsource/[5e07de58da822924da9c2108a46ba4cd173a75c8]/fbcode/monarch/hyperactor/src/mailbox.rs?lines=1600%2C1613

However, this method is already used outside the `hyperactor` crate in a couple of places. This diff hides those usage behind a new method `Instance::<()>::bind_actor_port`.

This diff also removes `Ports::bind_to` and fold its logic to where it is currently used.

https://www.internalfb.com/code/fbsource/[5e07de58da822924da9c2108a46ba4cd173a75c8]/fbcode/monarch/hyperactor/src/proc.rs?lines=1862%2C1958

Reviewed By: mariusae

Differential Revision: D83931993

fbshipit-source-id: 35dce72c235ec54b883df4c12c44d6141446db2c

Author: pzhan9

controller/src/lib.rs

@@ -634,7 +634,6 @@ mod tests {
     use hyperactor::clock::Clock;
     use hyperactor::clock::RealClock;
     use hyperactor::context::Mailbox as _;
-    use hyperactor::data::Named;
     use hyperactor::id;
     use hyperactor::mailbox::BoxedMailboxSender;
     use hyperactor::mailbox::DialMailboxRouter;
@@ -1129,8 +1128,7 @@ mod tests {
 
         // Build a supervisor.
         let sup_mail = system.attach().await.unwrap();
-        let (sup_tx, _sup_rx) = sup_mail.open_port::<ProcSupervisionMessage>();
-        sup_tx.bind_to(ProcSupervisionMessage::port());
+        let (_sup_tx, _sup_rx) = sup_mail.bind_actor_port::<ProcSupervisionMessage>();
         let sup_ref = ActorRef::<ProcSupervisor>::attest(sup_mail.self_id().clone());
 
         // Construct a system sender.
@@ -1360,8 +1358,7 @@ mod tests {
 
         // Build a supervisor.
         let sup_mail = system.attach().await.unwrap();
-        let (sup_tx, _sup_rx) = sup_mail.open_port::<ProcSupervisionMessage>();
-        sup_tx.bind_to(ProcSupervisionMessage::port());
+        let (_sup_tx, _sup_rx) = sup_mail.bind_actor_port::<ProcSupervisionMessage>();
         let sup_ref = ActorRef::<ProcSupervisor>::attest(sup_mail.self_id().clone());
 
         // Construct a system sender.
@@ -1665,9 +1662,8 @@ mod tests {
             .await
             .unwrap();
 
-        let (client_supervision_tx, mut client_supervision_rx) =
-            client_mailbox.open_port::<ClientMessage>();
-        client_supervision_tx.bind_to(ClientMessage::port());
+        let (_client_supervision_tx, mut client_supervision_rx) =
+            client_mailbox.bind_actor_port::<ClientMessage>();
 
         // mock a proc actor that doesn't update supervision state
         let (
@@ -1726,9 +1722,8 @@ mod tests {
         // Client actor.
         let mut system = System::new(server_handle.local_addr().clone());
         let client_mailbox = system.attach().await.unwrap();
-        let (client_supervision_tx, mut client_supervision_rx) =
-            client_mailbox.open_port::<ClientMessage>();
-        client_supervision_tx.bind_to(ClientMessage::port());
+        let (_client_supervision_tx, mut client_supervision_rx) =
+            client_mailbox.bind_actor_port::<ClientMessage>();
 
         // Bootstrap the controller
         let controller_id = id!(controller[0].root);
@@ -1865,9 +1860,8 @@ mod tests {
         // Client actor.
         let mut system = System::new(server_handle.local_addr().clone());
         let client_mailbox = system.attach().await.unwrap();
-        let (client_supervision_tx, mut client_supervision_rx) =
+        let (_client_supervision_tx, mut client_supervision_rx) =
             client_mailbox.open_port::<ClientMessage>();
-        client_supervision_tx.bind_to(ClientMessage::port());
 
         // Bootstrap the controller
         let controller_id = id!(controller[0].root);

hyperactor/src/clock.rs

@@ -22,7 +22,6 @@ use serde::Serialize;
 
 use crate::Mailbox;
 use crate::channel::ChannelAddr;
-use crate::data::Named;
 use crate::id;
 use crate::mailbox::DeliveryError;
 use crate::mailbox::MailboxSender;
@@ -261,9 +260,8 @@ impl SimClock {
         static SIMCLOCK_MAILBOX: OnceLock<Mailbox> = OnceLock::new();
         SIMCLOCK_MAILBOX.get_or_init(|| {
             let mailbox = Mailbox::new_detached(id!(proc[0].proc).clone());
-            let (undeliverable_messages, mut rx) =
-                mailbox.open_port::<Undeliverable<MessageEnvelope>>();
-            undeliverable_messages.bind_to(Undeliverable::<MessageEnvelope>::port());
+            let (_undeliverable_messages, mut rx) =
+                mailbox.bind_actor_port::<Undeliverable<MessageEnvelope>>();
             tokio::spawn(async move {
                 while let Ok(Undeliverable(mut envelope)) = rx.recv().await {
                     envelope.set_error(DeliveryError::BrokenLink(

hyperactor/src/mailbox.rs

@@ -1244,6 +1244,18 @@ impl Mailbox {
         )
     }
 
+    /// Bind this message's actor port to this actor's mailbox. This method is
+    /// normally used:
+    ///   1. when we need to intercept a message sent to a handler, and re-route
+    ///      that message to the returned receiver;
+    ///   2. mock this message's handler when it is not implemented for this actor
+    ///      type, with the returned receiver.
+    pub(crate) fn bind_actor_port<M: RemoteMessage>(&self) -> (PortHandle<M>, PortReceiver<M>) {
+        let (handle, receiver) = self.open_port();
+        handle.bind_actor_port();
+        (handle, receiver)
+    }
+
     /// Open a new port with an accumulator with default reduce options.
     /// See [`open_accum_port_opts`] for more details.
     pub fn open_accum_port<A>(&self, accum: A) -> (PortHandle<A::Update>, PortReceiver<A::State>)
@@ -1389,13 +1401,14 @@ impl Mailbox {
         PortRef::attest(port_id)
     }
 
-    fn bind_to<M: RemoteMessage>(&self, handle: &PortHandle<M>, port_index: u64) {
+    fn bind_to_actor_port<M: RemoteMessage>(&self, handle: &PortHandle<M>) {
         assert_eq!(
             handle.mailbox.actor_id(),
             self.actor_id(),
             "port does not belong to mailbox"
         );
 
+        let port_index = M::port();
         let port_id = self.actor_id().port_id(port_index);
         match self.inner.ports.entry(port_index) {
             Entry::Vacant(entry) => {
@@ -1625,10 +1638,13 @@ impl<M: RemoteMessage> PortHandle<M> {
         )
     }
 
-    /// Bind to a specific port index. This is used by [`actor::Binder`] implementations to
-    /// bind actor refs. This is not intended for general use.
-    pub fn bind_to(&self, port_index: u64) {
-        self.mailbox.bind_to(self, port_index);
+    /// Bind to this message's actor port. This method will panic if the handle
+    /// is already bound.
+    ///
+    /// This is used by [`actor::Binder`] implementations to bind actor refs.
+    /// This is not intended for general use.
+    pub(crate) fn bind_actor_port(&self) {
+        self.mailbox.bind_to_actor_port(self);
     }
 }
 
@@ -3648,8 +3664,7 @@ mod tests {
             actor_id.clone(),
             BoxedMailboxSender::new(AsyncLoopForwarder),
         );
-        let (ret_port, mut ret_rx) = mailbox.open_port::<Undeliverable<MessageEnvelope>>();
-        ret_port.bind_to(Undeliverable::<MessageEnvelope>::port());
+        let (ret_port, mut ret_rx) = mailbox.bind_actor_port::<Undeliverable<MessageEnvelope>>();
 
         // Create a destination not owned by this mailbox to force
         // forwarding.
@@ -3696,9 +3711,8 @@ mod tests {
             actor_id.clone(),
             BoxedMailboxSender::new(PanickingMailboxSender),
         );
-        let (undeliverable_tx, mut undeliverable_rx) =
-            mailbox.open_port::<Undeliverable<MessageEnvelope>>();
-        undeliverable_tx.bind_to(Undeliverable::<MessageEnvelope>::port());
+        let (_undeliverable_tx, mut undeliverable_rx) =
+            mailbox.bind_actor_port::<Undeliverable<MessageEnvelope>>();
 
         // Open a local user u64 port.
         let (user_port, mut user_rx) = mailbox.open_port::<u64>();

hyperactor/src/message.rs

@@ -249,7 +249,7 @@ impl<M: Bind> IndexedErasedUnbound<M> {
                 Ok(())
             }
         });
-        port_handle.bind_to(IndexedErasedUnbound::<M>::port());
+        port_handle.bind_actor_port();
         Ok(())
     }
 }

hyperactor/src/proc.rs

@@ -477,8 +477,7 @@ impl Proc {
         R: Referable + RemoteHandles<M>,
     {
         let (instance, _handle) = self.instance(name)?;
-        let (handle, rx) = instance.open_port::<M>();
-        handle.bind_to(M::port());
+        let (_handle, rx) = instance.bind_actor_port::<M>();
         let actor_ref = ActorRef::attest(instance.self_id().clone());
         Ok((instance, actor_ref, rx))
     }
@@ -1509,6 +1508,17 @@ impl<A: Actor> context::Actor for &Context<'_, A> {
     }
 }
 
+impl Instance<()> {
+    /// See [Mailbox::bind_actor_port] for details.
+    pub fn bind_actor_port<M: RemoteMessage>(&self) -> (PortHandle<M>, PortReceiver<M>) {
+        assert!(
+            self.actor_task_handle().is_none(),
+            "can only bind actor port on instance with no running actor task"
+        );
+        self.mailbox.bind_actor_port()
+    }
+}
+
 #[derive(Debug)]
 enum ActorType {
     Named(&'static TypeInfo),
@@ -1947,24 +1957,15 @@ impl<A: Actor> Ports<A> {
         }
     }
 
-    /// Bind the given message type to its default port.
+    /// Bind the given message type to its actor port.
     pub fn bind<M: RemoteMessage>(&self)
     where
         A: Handler<M>,
     {
-        self.bind_to::<M>(M::port());
-    }
-
-    /// Bind the given message type to the provided port.
-    /// Ports cannot be rebound to different message types;
-    /// and attempting to do so will result in a panic.
-    pub fn bind_to<M: RemoteMessage>(&self, port_index: u64)
-    where
-        A: Handler<M>,
-    {
+        let port_index = M::port();
         match self.bound.entry(port_index) {
             Entry::Vacant(entry) => {
-                self.get::<M>().bind_to(port_index);
+                self.get::<M>().bind_actor_port();
                 entry.insert(M::typename());
             }
             Entry::Occupied(entry) => {

hyperactor_mesh/src/proc_mesh.rs

@@ -22,7 +22,6 @@ use hyperactor::ActorHandle;
 use hyperactor::ActorId;
 use hyperactor::ActorRef;
 use hyperactor::Instance;
-use hyperactor::Named;
 use hyperactor::RemoteMessage;
 use hyperactor::WorldId;
 use hyperactor::actor::ActorStatus;
@@ -188,9 +187,8 @@ pub fn global_root_client() -> &'static Instance<()> {
         // The hook logs each undeliverable, along with whether a sink
         // was present at the time of receipt, which helps diagnose
         // lost or misrouted events.
-        let (undeliverable_tx, undeliverable_rx) =
-            client.open_port::<Undeliverable<MessageEnvelope>>();
-        undeliverable_tx.bind_to(Undeliverable::<MessageEnvelope>::port());
+        let (_undeliverable_tx, undeliverable_rx) =
+            client.bind_actor_port::<Undeliverable<MessageEnvelope>>();
         hyperactor::mailbox::supervise_undeliverable_messages_with(
             undeliverable_rx,
             crate::proc_mesh::get_global_supervision_sink,
@@ -346,9 +344,8 @@ impl ProcMesh {
         // `global_root_client()`.
         let (client, _handle) = client_proc.instance("client")?;
         // Bind an undeliverable message port in the client.
-        let (undeliverable_messages, client_undeliverable_receiver) =
-            client.open_port::<Undeliverable<MessageEnvelope>>();
-        undeliverable_messages.bind_to(Undeliverable::<MessageEnvelope>::port());
+        let (_undeliverable_messages, client_undeliverable_receiver) =
+            client.bind_actor_port::<Undeliverable<MessageEnvelope>>();
         hyperactor::mailbox::supervise_undeliverable_messages(
             supervision_port.clone(),
             client_undeliverable_receiver,

hyperactor_multiprocess/src/proc_actor.rs

@@ -878,7 +878,6 @@ mod tests {
     use hyperactor::channel::ChannelTransport;
     use hyperactor::clock::Clock;
     use hyperactor::clock::RealClock;
-    use hyperactor::data::Named;
     use hyperactor::forward;
     use hyperactor::id;
     use hyperactor::reference::ActorRef;
@@ -1205,9 +1204,8 @@ mod tests {
         // A test supervisor.
         let mut system = System::new(server_handle.local_addr().clone());
         let supervisor = system.attach().await.unwrap();
-        let (supervisor_supervision_tx, mut supervisor_supervision_receiver) =
-            supervisor.open_port::<ProcSupervisionMessage>();
-        supervisor_supervision_tx.bind_to(ProcSupervisionMessage::port());
+        let (_supervisor_supervision_tx, mut supervisor_supervision_receiver) =
+            supervisor.bind_actor_port::<ProcSupervisionMessage>();
         let supervisor_actor_ref: ActorRef<ProcSupervisor> =
             ActorRef::attest(supervisor.self_id().clone());
 
@@ -1388,8 +1386,7 @@ mod tests {
 
         // Build a supervisor.
         let supervisor = system.attach().await.unwrap();
-        let (sup_tx, _sup_rx) = supervisor.open_port::<ProcSupervisionMessage>();
-        sup_tx.bind_to(ProcSupervisionMessage::port());
+        let (_sup_tx, _sup_rx) = supervisor.bind_actor_port::<ProcSupervisionMessage>();
         let sup_ref = ActorRef::<ProcSupervisor>::attest(supervisor.self_id().clone());
 
         // Construct a system sender.
@@ -1512,8 +1509,7 @@ mod tests {
 
         // Build a supervisor.
         let supervisor = system.attach().await.unwrap();
-        let (sup_tx, _sup_rx) = supervisor.open_port::<ProcSupervisionMessage>();
-        sup_tx.bind_to(ProcSupervisionMessage::port());
+        let (_sup_tx, _sup_rx) = supervisor.bind_actor_port::<ProcSupervisionMessage>();
         let sup_ref = ActorRef::<ProcSupervisor>::attest(supervisor.self_id().clone());
 
         // Construct a system sender.

hyperactor_multiprocess/src/system_actor.rs

@@ -2242,8 +2242,7 @@ mod tests {
 
         // Build a supervisor.
         let supervisor = system.attach().await.unwrap();
-        let (sup_tx, _sup_rx) = supervisor.open_port::<ProcSupervisionMessage>();
-        sup_tx.bind_to(ProcSupervisionMessage::port());
+        let (_sup_tx, _sup_rx) = supervisor.bind_actor_port::<ProcSupervisionMessage>();
         let sup_ref = ActorRef::<ProcSupervisor>::attest(supervisor.self_id().clone());
 
         // Construct a system sender.

monarch_hyperactor/src/proc.rs

@@ -25,7 +25,6 @@ use std::time::SystemTime;
 
 use anyhow::Result;
 use hyperactor::ActorRef;
-use hyperactor::Named;
 use hyperactor::RemoteMessage;
 use hyperactor::actor::Signal;
 use hyperactor::channel;
@@ -453,11 +452,9 @@ impl<M: RemoteMessage> InstanceWrapper<M> {
 
     fn new_with_instance_and_clock(instance: Instance<()>, clock: ClockKind) -> Result<Self> {
         // TEMPORARY: remove after using fixed message ports.
-        let (message_port, message_receiver) = instance.open_port::<M>();
-        message_port.bind_to(M::port());
+        let (_message_port, message_receiver) = instance.bind_actor_port::<M>();
 
-        let (signal_port, signal_receiver) = instance.open_port::<Signal>();
-        signal_port.bind_to(<Signal as Named>::port());
+        let (signal_port, signal_receiver) = instance.bind_actor_port::<Signal>();
 
         let (controller_error_sender, controller_error_receiver) = watch::channel("".to_string());
         let actor_id = instance.self_id().clone();

monarch_tensor_worker/src/lib.rs

@@ -1174,7 +1174,6 @@ mod tests {
 
     use anyhow::Result;
     use hyperactor::Instance;
-    use hyperactor::Named;
     use hyperactor::WorldId;
     use hyperactor::actor::ActorStatus;
     use hyperactor::channel::ChannelAddr;
@@ -2348,8 +2347,7 @@ mod tests {
 
         // Create a fake controller for the workers to talk to.
         let client = System::new(system_addr.clone()).attach().await?;
-        let (handle, mut controller_rx) = client.open_port::<ControllerMessage>();
-        handle.bind_to(ControllerMessage::port());
+        let (_handle, mut controller_rx) = client.bind_actor_port::<ControllerMessage>();
         let controller_ref: ActorRef<ControllerActor> = ActorRef::attest(client.self_id().clone());
 
         // Create the worker world