Add DirectSend variant to SeqInfo

Differential Revision: D83839619

Author: pzhan9

hyperactor/src/actor.rs

@@ -1143,31 +1143,45 @@ mod tests {
             .spawn::<GetSeqActor>("get_seq", tx.bind())
             .await
             .unwrap();
+
+        // Verify that unbound handle can send message.
+        actor_handle.send(&client, "unbound".to_string()).unwrap();
+        assert_eq!(
+            rx.recv().await.unwrap(),
+            ("unbound".to_string(), SeqInfo::Unordered)
+        );
+
         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();
+        for m in 0..10 {
+            actor_handle.send(&client, format!("{m}")).unwrap();
             expected_seq += 1;
             assert_eq!(
-                rx.recv().await.unwrap().1,
-                SeqInfo {
-                    session_id,
-                    seq: expected_seq,
-                }
+                rx.recv().await.unwrap(),
+                (
+                    format!("{m}"),
+                    SeqInfo::Session {
+                        session_id,
+                        seq: expected_seq,
+                    }
+                )
             );
 
-            for _ in 0..2 {
-                actor_ref.port().send(&client, "".to_string()).unwrap();
+            for n in 0..2 {
+                actor_ref.port().send(&client, format!("{m}-{n}")).unwrap();
                 expected_seq += 1;
                 assert_eq!(
-                    rx.recv().await.unwrap().1,
-                    SeqInfo {
-                        session_id,
-                        seq: expected_seq,
-                    }
+                    rx.recv().await.unwrap(),
+                    (
+                        format!("{m}-{n}"),
+                        SeqInfo::Session {
+                            session_id,
+                            seq: expected_seq,
+                        }
+                    )
                 );
             }
         }
@@ -1200,7 +1214,10 @@ mod tests {
         let session_id = client.sequencer().session_id();
         assert_eq!(
             rx.recv().await.unwrap(),
-            ("finally".to_string(), SeqInfo { session_id, seq: 1 })
+            (
+                "finally".to_string(),
+                SeqInfo::Session { session_id, seq: 1 }
+            )
         );
     }
 
@@ -1280,7 +1297,7 @@ mod tests {
         for expect in expected {
             let expected = (
                 expect.0,
-                SeqInfo {
+                SeqInfo::Session {
                     session_id,
                     seq: expect.1,
                 },

hyperactor/src/mailbox.rs

@@ -66,8 +66,6 @@
 //! implementation to avoid a serialization roundtrip when passing
 //! messages locally.
 
-#![allow(dead_code)] // Allow until this is used outside of tests.
-
 use std::any::Any;
 use std::collections::BTreeMap;
 use std::collections::BTreeSet;
@@ -946,43 +944,6 @@ impl Future for MailboxServerHandle {
     }
 }
 
-// A `MailboxServer` (such as a router) can receive a message
-// that couldn't reach its destination. We can use the fact that
-// servers are `MailboxSender`s to attempt to forward them back to
-// their senders.
-fn server_return_handle<T: MailboxServer>(server: T) -> PortHandle<Undeliverable<MessageEnvelope>> {
-    let (return_handle, mut rx) = undeliverable::new_undeliverable_port();
-
-    tokio::task::spawn(async move {
-        while let Ok(Undeliverable(mut envelope)) = rx.recv().await {
-            if let Ok(Undeliverable(e)) = envelope.deserialized::<Undeliverable<MessageEnvelope>>()
-            {
-                // A non-returnable undeliverable.
-                UndeliverableMailboxSender.post(e, monitored_return_handle());
-                continue;
-            }
-            envelope.set_error(DeliveryError::BrokenLink(
-                "message was undeliverable".to_owned(),
-            ));
-            server.post(
-                MessageEnvelope::new(
-                    envelope.sender().clone(),
-                    PortRef::<Undeliverable<MessageEnvelope>>::attest_message_port(
-                        envelope.sender(),
-                    )
-                    .port_id()
-                    .clone(),
-                    Serialized::serialize(&Undeliverable(envelope)).unwrap(),
-                    Attrs::new(),
-                ),
-                monitored_return_handle(),
-            );
-        }
-    });
-
-    return_handle
-}
-
 /// Serve a port on the provided [`channel::Rx`]. This dispatches all
 /// channel messages directly to the port.
 pub trait MailboxServer: MailboxSender + Clone + Sized + 'static {
@@ -1011,6 +972,9 @@ pub trait MailboxServer: MailboxSender + Clone + Sized + 'static {
                 envelope.set_error(DeliveryError::BrokenLink(
                     "message was undeliverable".to_owned(),
                 ));
+                let mut headers = Attrs::new();
+                // Ordering is not required when returning Undeliverable.
+                headers.set(SEQ_INFO, SeqInfo::Unordered);
                 server.post(
                     MessageEnvelope::new(
                         envelope.sender().clone(),
@@ -1020,7 +984,7 @@ pub trait MailboxServer: MailboxSender + Clone + Sized + 'static {
                         .port_id()
                         .clone(),
                         Serialized::serialize(&Undeliverable(envelope)).unwrap(),
-                        Attrs::new(),
+                        headers,
                     ),
                     monitored_return_handle(),
                 );
@@ -1589,19 +1553,30 @@ impl<M: Message> PortHandle<M> {
         let mut headers = Attrs::new();
 
         crate::mailbox::headers::set_send_timestamp(&mut headers);
-        // 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);
+
+        match self.bound.get() {
+            Some(bound_port) => {
+                // Message sent from handle is delivered immediately. It could
+                // race with messages from refs. So we need to assign seq to
+                // preserve the ordering.
+                if bound_port.is_actor_port() {
+                    let sequencer = cx.instance().sequencer();
+                    let seq = sequencer.assign_seq(self.mailbox.actor_id());
+                    let seq_info = SeqInfo::Session {
+                        session_id: sequencer.session_id(),
+                        seq,
+                    };
+                    headers.set(SEQ_INFO, seq_info);
+                }
+            }
+            None => {
+                // we do not have info to know whether this handle is used for
+                // enqueue port or not. Since enqueue port requires the SEQ_INFO
+                // header, we set it in for all messages sent from unbound handles.
+                headers.set(SEQ_INFO, SeqInfo::Unordered);
+            }
         }
+
         // 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
@@ -1619,6 +1594,7 @@ 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);
+        headers.set(SEQ_INFO, SeqInfo::Unordered);
         self.sender.send(headers, message).map_err(|err| {
             MailboxSenderError::new_unbound::<M>(
                 self.mailbox.actor_id().clone(),

hyperactor/src/ordering.rs

@@ -164,7 +164,6 @@ impl<T> OrderedSender<T> {
     }
 
     pub(crate) fn direct_send(&self, msg: T) -> Result<(), SendError<T>> {
-        assert!(!self.enable_buffering);
         self.tx.send(msg)
     }
 }

hyperactor/src/proc.rs

@@ -13,6 +13,7 @@
 
 use std::any::Any;
 use std::any::TypeId;
+use std::any::type_name;
 use std::collections::HashMap;
 use std::fmt;
 use std::future::Future;
@@ -1847,30 +1848,42 @@ pub struct Ports<A: Actor> {
 
 /// A message's sequencer number infomation.
 #[derive(Debug, Serialize, Deserialize, Clone, Named, AttrValue, PartialEq)]
-pub struct SeqInfo {
-    /// Message's session ID
-    pub session_id: Uuid,
-    /// Message's sequence number in the given session.
-    pub seq: u64,
+pub enum SeqInfo {
+    /// Messages with the same session ID should be delivered in order.
+    Session {
+        /// Message's session ID
+        session_id: Uuid,
+        /// Message's sequence number in the given session.
+        seq: u64,
+    },
+    /// This message does not require ordering and thus have no sequence number.
+    Unordered,
 }
 
 impl fmt::Display for SeqInfo {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "{}:{}", self.session_id, self.seq)
+        match self {
+            Self::Unordered => write!(f, "unordered"),
+            Self::Session { session_id, seq } => write!(f, "{}:{}", session_id, seq),
+        }
     }
 }
 
 impl std::str::FromStr for SeqInfo {
     type Err = anyhow::Error;
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
+        if s == "unordered" {
+            return Ok(SeqInfo::Unordered);
+        }
+
         let parts: Vec<_> = s.split(':').collect();
         if parts.len() != 2 {
             return Err(anyhow::anyhow!("invalid SeqInfo: {}", s));
         }
         let session_id: Uuid = parts[0].parse()?;
         let seq: u64 = parts[1].parse()?;
-        Ok(SeqInfo { session_id, seq })
+        Ok(SeqInfo::Session { session_id, seq })
     }
 }
 
@@ -1926,18 +1939,37 @@ impl<A: Actor> Ports<A> {
                         hyperactor_telemetry::kv_pairs!("actor_id" => actor_id.clone()),
                     );
                     if workq.enable_buffering {
-                        let SeqInfo { session_id, seq } =
-                            seq_info.expect("SEQ_INFO must be set when buffering is enabled");
-
-                        // TODO: return the message contained in the error instead of dropping them when converting
-                        // to anyhow::Error. In that way, the message can be picked up by mailbox and returned to sender.
-                        workq.send(session_id, seq, work).map_err(|e| match e {
-                            OrderedSenderError::InvalidZeroSeq(_) => {
-                                anyhow::anyhow!("seq must be greater than 0")
+                        match seq_info {
+                            Some(SeqInfo::Session { session_id, seq }) => {
+                                // TODO: return the message contained in the error instead of dropping them when converting
+                                // to anyhow::Error. In that way, the message can be picked up by mailbox and returned to sender.
+                                workq.send(session_id, seq, work).map_err(|e| match e {
+                                    OrderedSenderError::InvalidZeroSeq(_) => {
+                                        let error_msg = format!(
+                                             "in enqueue func for {}, got seq 0 for message type {}",
+                                            actor_id,
+                                            std::any::type_name::<M>(),
+                                        );
+                                        tracing::error!(error_msg);
+                                        anyhow::anyhow!(error_msg)
+                                    }
+                                    OrderedSenderError::SendError(e) => anyhow::Error::from(e),
+                                    OrderedSenderError::FlushError(e) => e,
+                                })
                             }
-                            OrderedSenderError::SendError(e) => anyhow::Error::from(e),
-                            OrderedSenderError::FlushError(e) => e,
-                        })
+                            Some(SeqInfo::Unordered) => {
+                                workq.direct_send(work).map_err(anyhow::Error::from)
+                            }
+                            None => {
+                                let error_msg = format!(
+                                    "in enqueue func for {}, buffering is enabled, but SEQ_INFO is not set for message type {}",
+                                    actor_id,
+                                    std::any::type_name::<M>(),
+                                    );
+                                tracing::error!(error_msg);
+                                anyhow::bail!(error_msg);
+                            }
+                        }
                     } else {
                         workq.direct_send(work).map_err(anyhow::Error::from)
                     }

hyperactor/src/reference.rs

@@ -918,7 +918,7 @@ impl PortId {
             // without worrying about rollback.
             let sequencer = cx.instance().sequencer();
             let seq = sequencer.assign_seq(self.actor_id());
-            let seq_info = SeqInfo {
+            let seq_info = SeqInfo::Session {
                 session_id: sequencer.session_id(),
                 seq,
             };
@@ -1563,10 +1563,13 @@ mod tests {
         let port_ref = PortRef::attest(port_id.clone());
 
         port_handle.send(&client, ()).unwrap();
-        let SeqInfo {
+        let SeqInfo::Session {
             session_id,
             mut seq,
-        } = rx.try_recv().unwrap().unwrap();
+        } = rx.try_recv().unwrap().unwrap()
+        else {
+            panic!("expected session info");
+        };
         assert_eq!(session_id, client.sequencer().session_id());
         assert_eq!(seq, 1);
 
@@ -1576,10 +1579,13 @@ mod tests {
             seq: &mut u64,
         ) {
             *seq += 1;
-            let SeqInfo {
+            let SeqInfo::Session {
                 session_id: rcved_session_id,
                 seq: rcved_seq,
-            } = rx.try_recv().unwrap().unwrap();
+            } = rx.try_recv().unwrap().unwrap()
+            else {
+                panic!("expected session info");
+            };
             assert_eq!(rcved_session_id, session_id);
             assert_eq!(rcved_seq, *seq);
         }
@@ -1617,8 +1623,9 @@ mod tests {
             Ok(())
         });
         port_handle.send(&client, ()).unwrap();
-        // No seq will be assigned for unbound port handle.
-        assert!(rx.try_recv().unwrap().is_none());
+        // Unordered be set for unbound port handle since handler's ordered
+        // channel is expecting the SEQ_INFO header to be set.
+        assert_eq!(rx.try_recv().unwrap().unwrap(), SeqInfo::Unordered);
 
         // Bind to the allocated port.
         port_handle.bind();

hyperactor_mesh/src/comm.rs

@@ -521,7 +521,7 @@ impl Handler<ForwardMessageV1> for CommActor {
                 &mut headers,
                 cast_point,
                 message.cast_headers.sender.clone(),
-                Some(SeqInfo {
+                Some(SeqInfo::Session {
                     session_id: message.cast_headers.session_id,
                     seq,
                 }),

hyperactor_mesh/src/v1/testactor.rs

@@ -281,7 +281,7 @@ pub async fn assert_casting_correctness(
             .iter()
             .zip(
                 seqs.into_iter()
-                    .map(|seq| Some(SeqInfo { session_id, seq })),
+                    .map(|seq| Some(SeqInfo::Session { session_id, seq })),
             )
             .collect(),
     };