Assign seq for PortHandle/PortRef

Differential Revision: D83591305

Author: pzhan9

hyperactor/src/actor.rs

@@ -716,9 +716,12 @@ pub trait RemoteHandles<M: RemoteMessage>: Referable {}
 
 #[cfg(test)]
 mod tests {
+    use std::collections::HashMap;
     use std::sync::Mutex;
     use std::time::Duration;
 
+    use timed_test::async_timed_test;
+    use tokio::sync::mpsc;
     use tokio::time::timeout;
 
     use super::*;
@@ -728,6 +731,13 @@ mod tests {
     use crate::PortRef;
     use crate::checkpoint::CheckpointError;
     use crate::checkpoint::Checkpointable;
+    use crate::config;
+    use crate::id;
+    use crate::mailbox::BoxedMailboxSender;
+    use crate::mailbox::MailboxSender;
+    use crate::mailbox::monitored_return_handle;
+    use crate::proc::SEQ_INFO;
+    use crate::proc::SeqInfo;
     use crate::test_utils::pingpong::PingPongActor;
     use crate::test_utils::pingpong::PingPongActorParams;
     use crate::test_utils::pingpong::PingPongMessage;
@@ -1072,4 +1082,210 @@ mod tests {
         handle.drain_and_stop().unwrap();
         handle.await;
     }
+
+    // Returning the sequence number assigned to the message.
+    #[derive(Debug)]
+    #[hyperactor::export(handlers = [String])]
+    struct GetSeqActor(PortRef<(String, SeqInfo)>);
+
+    #[async_trait]
+    impl Actor for GetSeqActor {
+        type Params = PortRef<(String, SeqInfo)>;
+
+        async fn new(params: PortRef<(String, SeqInfo)>) -> Result<Self, anyhow::Error> {
+            Ok(Self(params))
+        }
+    }
+
+    #[async_trait]
+    impl Handler<String> for GetSeqActor {
+        async fn handle(
+            &mut self,
+            cx: &Context<Self>,
+            message: String,
+        ) -> Result<(), anyhow::Error> {
+            let Self(port) = self;
+            let seq_info = cx.headers().get(SEQ_INFO).unwrap();
+            port.send(cx, (message, seq_info.clone()))?;
+            Ok(())
+        }
+    }
+
+    #[async_timed_test(timeout_secs = 30)]
+    async fn test_sequencing_actor_handle_basic() {
+        let proc = Proc::local();
+        let (client, _) = proc.instance("client").unwrap();
+        let (tx, mut rx) = client.open_port();
+
+        let actor_handle = proc
+            .spawn::<GetSeqActor>("get_seq", tx.bind())
+            .await
+            .unwrap();
+        let actor_ref: ActorRef<GetSeqActor> = actor_handle.bind();
+
+        let session_id = client.sequencer().session_id();
+        let mut expected_seq = 0;
+        // Interleave messages sent through the handle and the reference.
+        for _ in 0..10 {
+            actor_handle.send(&client, "".to_string()).unwrap();
+            expected_seq += 1;
+            assert_eq!(
+                rx.recv().await.unwrap().1,
+                SeqInfo {
+                    session_id,
+                    seq: expected_seq,
+                }
+            );
+
+            for _ in 0..2 {
+                actor_ref.port().send(&client, "".to_string()).unwrap();
+                expected_seq += 1;
+                assert_eq!(
+                    rx.recv().await.unwrap().1,
+                    SeqInfo {
+                        session_id,
+                        seq: expected_seq,
+                    }
+                );
+            }
+        }
+    }
+
+    // Adding a delay before sending the destination proc. Useful for tests
+    // requiring latency injection.
+    #[derive(Debug)]
+    struct DelayedMailboxSender {
+        relay_tx: mpsc::UnboundedSender<MessageEnvelope>,
+    }
+
+    impl DelayedMailboxSender {
+        // Use a random latency between 0 and 1 second if the plan is empty.
+        fn boxed(dest_proc: Proc, latency_plan: HashMap<u64, Duration>) -> BoxedMailboxSender {
+            let (relay_tx, mut relay_rx) = mpsc::unbounded_channel();
+            tokio::spawn(async move {
+                let mut count = 0;
+                while let Some(envelope) = relay_rx.recv().await {
+                    count += 1;
+
+                    let latency = if latency_plan.is_empty() {
+                        Duration::from_millis(1000)
+                    } else {
+                        latency_plan.get(&count).unwrap().clone()
+                    };
+
+                    let dest_proc_clone = dest_proc.clone();
+                    tokio::spawn(async move {
+                        // Need Clock::sleep is an async function.
+                        #[allow(clippy::disallowed_methods)]
+                        tokio::time::sleep(latency).await;
+                        dest_proc_clone.post(envelope, monitored_return_handle());
+                    });
+                }
+            });
+
+            BoxedMailboxSender::new(Self { relay_tx })
+        }
+    }
+
+    impl MailboxSender for DelayedMailboxSender {
+        fn post_unchecked(
+            &self,
+            envelope: MessageEnvelope,
+            _return_handle: PortHandle<Undeliverable<MessageEnvelope>>,
+        ) {
+            self.relay_tx.send(envelope).unwrap();
+        }
+    }
+
+    async fn assert_out_of_order_delivery(
+        expected: Vec<(String, u64)>,
+        latency_plan: HashMap<u64, Duration>,
+    ) {
+        let local_proc: Proc = Proc::local();
+        let (client, _) = local_proc.instance("local").unwrap();
+        let (tx, mut rx) = client.open_port();
+
+        let handle = local_proc
+            .spawn::<GetSeqActor>("get_seq", tx.bind())
+            .await
+            .unwrap();
+
+        let actor_ref: ActorRef<GetSeqActor> = handle.bind();
+
+        let remote_proc = Proc::new(
+            id!(remote[0]),
+            DelayedMailboxSender::boxed(local_proc.clone(), latency_plan),
+        );
+        let (remote_client, _) = remote_proc.instance("remote").unwrap();
+        // Send the messages out in the order of their expected sequence numbers.
+        let mut messages = expected.clone();
+        messages.sort_by_key(|v| v.1);
+        for (message, _seq) in messages {
+            actor_ref.send(&remote_client, message).unwrap();
+        }
+        let session_id = remote_client.sequencer().session_id();
+        for expect in expected {
+            let expected = (
+                expect.0,
+                SeqInfo {
+                    session_id,
+                    seq: expect.1,
+                },
+            );
+            assert_eq!(rx.recv().await.unwrap(), expected);
+        }
+
+        handle.drain_and_stop().unwrap();
+        handle.await;
+    }
+
+    // Send several messages, use DelayedMailboxSender and the latency plan to
+    // ensure these messages will arrive at handler's workq in a determinstic
+    // out-of-order way. Then verify the actor handler will still process these
+    // messages based on their sending order if reordering buffer is enabled.
+    #[async_timed_test(timeout_secs = 30)]
+    async fn test_sequencing_actor_ref_out_of_order_deterministic() {
+        let config = config::global::lock();
+
+        let latency_plan = maplit::hashmap! {
+            1 => Duration::from_millis(1000),
+            2 => Duration::from_millis(0),
+        };
+
+        // By disabling the actor side re-ordering buffer, the mssages will
+        // be processed in the same order as they sent out.
+        let _guard = config.override_key(config::ENABLE_CLIENT_SEQ_ASSIGNMENT, false);
+        assert_out_of_order_delivery(
+            vec![("second".to_string(), 2), ("first".to_string(), 1)],
+            latency_plan.clone(),
+        )
+        .await;
+
+        // By enabling the actor side re-ordering buffer, the mssages will
+        // be re-ordered before being processed.
+        let _guard = config.override_key(config::ENABLE_CLIENT_SEQ_ASSIGNMENT, true);
+        assert_out_of_order_delivery(
+            vec![("first".to_string(), 1), ("second".to_string(), 2)],
+            latency_plan.clone(),
+        )
+        .await;
+    }
+
+    // Send a large nubmer of messages, use DelayedMailboxSender to ensure these
+    // messages will arrive at handler's workq in a random order. Then verify the
+    // actor handler will still process these messages based on their sending
+    // order with reordering buffer enabled.
+    #[async_timed_test(timeout_secs = 30)]
+    async fn test_sequencing_actor_ref_out_of_order_random() {
+        let config = config::global::lock();
+
+        // By enabling the actor side re-ordering buffer, the mssages will
+        // be re-ordered before being processed.
+        let _guard = config.override_key(config::ENABLE_CLIENT_SEQ_ASSIGNMENT, true);
+        let expected = (1..10000)
+            .map(|i| (format!("msg{i}"), i))
+            .collect::<Vec<_>>();
+
+        assert_out_of_order_delivery(expected, HashMap::new()).await;
+    }
 }

hyperactor/src/data.rs

@@ -23,6 +23,9 @@ use serde::de::DeserializeOwned;
 use crate as hyperactor;
 use crate::config;
 
+/// Actor handler port should have its most significant bit set to 1.
+pub(crate) static ACTOR_PORT_BIT: u64 = 1 << 63;
+
 /// A [`Named`] type is a type that has a globally unique name.
 pub trait Named: Sized + 'static {
     /// The globally unique type name for the type.
@@ -46,7 +49,7 @@ pub trait Named: Sized + 'static {
     /// The globally unique port for this type. Typed ports are in the range
     /// of 1<<63..1<<64-1.
     fn port() -> u64 {
-        Self::typehash() | (1 << 63)
+        Self::typehash() | ACTOR_PORT_BIT
     }
 
     /// If the named type is an enum, this returns the name of the arm

hyperactor/src/mailbox.rs

@@ -121,6 +121,8 @@ use crate::context;
 use crate::data::Serialized;
 use crate::id;
 use crate::metrics;
+use crate::proc::SEQ_INFO;
+use crate::proc::SeqInfo;
 use crate::reference::ActorId;
 use crate::reference::PortId;
 use crate::reference::Reference;
@@ -1575,20 +1577,35 @@ impl<M: Message> PortHandle<M> {
         }
     }
 
-    fn location(&self) -> PortLocation {
+    pub(crate) fn location(&self) -> PortLocation {
         match self.bound.get() {
             Some(port_id) => PortLocation::Bound(port_id.clone()),
             None => PortLocation::new_unbound::<M>(self.mailbox.actor_id().clone()),
         }
     }
 
     /// Send a message to this port.
-    pub fn send(&self, _cx: &impl context::Actor, 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.
-
+        // Message sent from handle is delivered immediately. It could race with
+        // messages from refs. So we need to assign seq if the handle is bound.
+        if let Some(bound_port) = self.bound.get()
+            && bound_port.is_actor_port()
+        {
+            let sequencer = cx.instance().sequencer();
+            let seq = sequencer.assign_seq(self.mailbox.actor_id());
+            let seq_info = SeqInfo {
+                session_id: sequencer.session_id(),
+                seq,
+            };
+            headers.set(SEQ_INFO, seq_info);
+        }
+        // Encountering error means the port is closed. So we do not need to
+        // rollback the seq, because no message can be delivered to it, and
+        // subsequently do not need to worry about out-of-sequence for messages
+        // after this seq.
         self.sender.send(headers, message).map_err(|err| {
             MailboxSenderError::new_unbound::<M>(
                 self.mailbox.actor_id().clone(),
@@ -1602,7 +1619,6 @@ impl<M: Message> PortHandle<M> {
     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| {
             MailboxSenderError::new_unbound::<M>(
                 self.mailbox.actor_id().clone(),

hyperactor/src/proc.rs

@@ -1846,7 +1846,7 @@ pub struct Ports<A: Actor> {
 }
 
 /// A message's sequencer number infomation.
-#[derive(Serialize, Deserialize, Clone, Named, AttrValue)]
+#[derive(Debug, Serialize, Deserialize, Clone, Named, AttrValue, PartialEq)]
 pub struct SeqInfo {
     /// Message's session ID
     pub session_id: Uuid,