Create separate events for sending & receiving messages (#250)

Summary:
Pull Request resolved: #250

We have been treating the entire message send & recv as a single event but this doesn't really make sense since each of these should have their own duration (sending -> in flight -> receiving), and they happen on different processes (This will matter in the next diff when we render our events onto Perfetto).

Differential Revision: D76433477

Author: Thomas Wang

hyperactor/src/channel/sim.rs

@@ -28,7 +28,6 @@ use crate::mailbox::MessageEnvelope;
 use crate::simnet;
 use crate::simnet::Dispatcher;
 use crate::simnet::Event;
-use crate::simnet::ScheduledEvent;
 use crate::simnet::SimNetConfig;
 use crate::simnet::SimNetEdge;
 use crate::simnet::SimNetError;
@@ -152,15 +151,16 @@ impl fmt::Display for SimAddr {
 
 /// Message Event that can be passed around in the simnet.
 #[derive(Debug)]
-pub(crate) struct MessageDeliveryEvent {
+pub(crate) struct MessageSendEvent {
     src_addr: Option<AddressProxyPair>,
     dest_addr: AddressProxyPair,
     data: Serialized,
     duration_ms: u64,
+    inflight_time_ms: u64,
 }
 
-impl MessageDeliveryEvent {
-    /// Creates a new MessageDeliveryEvent.
+impl MessageSendEvent {
+    /// Creates a new MessageSendEvent.
     pub fn new(
         src_addr: Option<AddressProxyPair>,
         dest_addr: AddressProxyPair,
@@ -170,22 +170,32 @@ impl MessageDeliveryEvent {
             src_addr,
             dest_addr,
             data,
-            duration_ms: 100,
+            duration_ms: 1,
+            inflight_time_ms: 1,
         }
     }
 }
 
 #[async_trait]
-impl Event for MessageDeliveryEvent {
+impl Event for MessageSendEvent {
     async fn handle(&self) -> Result<(), SimNetError> {
-        // Send the message to the correct receiver.
-        SENDER
-            .send(
-                self.src_addr.clone(),
-                self.dest_addr.clone(),
-                self.data.clone(),
-            )
-            .await?;
+        let inflight_time_ms = self.inflight_time_ms;
+        let event = Box::new(MessageRecvEvent::new(
+            self.src_addr.clone(),
+            self.dest_addr.clone(),
+            self.data.clone(),
+        ));
+
+        tokio::task::spawn(async move {
+            SimClock
+                .sleep(tokio::time::Duration::from_millis(inflight_time_ms))
+                .await;
+
+            if let Ok(handle) = simnet_handle() {
+                let _ = handle.send_event(event);
+            }
+        });
+
         Ok(())
     }
 
@@ -195,7 +205,7 @@ impl Event for MessageDeliveryEvent {
 
     fn summary(&self) -> String {
         format!(
-            "Sending message from {} to {}",
+            "{} sending message to to {}",
             self.src_addr
                 .as_ref()
                 .map_or("unknown".to_string(), |addr| addr.address.to_string()),
@@ -209,7 +219,7 @@ impl Event for MessageDeliveryEvent {
                 src: src_addr.address.clone(),
                 dst: self.dest_addr.address.clone(),
             };
-            self.duration_ms = topology
+            self.inflight_time_ms = topology
                 .lock()
                 .await
                 .topology
@@ -219,6 +229,60 @@ impl Event for MessageDeliveryEvent {
     }
 }
 
+/// Message Recv Event that can be passed around in the simnet.
+#[derive(Debug)]
+pub(crate) struct MessageRecvEvent {
+    src_addr: Option<AddressProxyPair>,
+    dest_addr: AddressProxyPair,
+    data: Serialized,
+    duration_ms: u64,
+}
+
+impl MessageRecvEvent {
+    /// Creates a new MessageRecvEvent.
+    pub fn new(
+        src_addr: Option<AddressProxyPair>,
+        dest_addr: AddressProxyPair,
+        data: Serialized,
+    ) -> Self {
+        Self {
+            src_addr,
+            dest_addr,
+            data,
+            duration_ms: 1,
+        }
+    }
+}
+
+#[async_trait]
+impl Event for MessageRecvEvent {
+    async fn handle(&self) -> Result<(), SimNetError> {
+        // Send the message to the correct receiver.
+        SENDER
+            .send(
+                self.src_addr.clone(),
+                self.dest_addr.clone(),
+                self.data.clone(),
+            )
+            .await?;
+        Ok(())
+    }
+
+    fn duration_ms(&self) -> u64 {
+        self.duration_ms
+    }
+
+    fn summary(&self) -> String {
+        format!(
+            "{} received message from {}",
+            self.dest_addr.address.clone(),
+            self.src_addr
+                .as_ref()
+                .map_or("unknown".to_string(), |addr| addr.address.to_string()),
+        )
+    }
+}
+
 /// Bind a channel address to the simnet. It will register the address as a node in simnet,
 /// and configure default latencies between this node and all other existing nodes.
 pub async fn bind(addr: ChannelAddr) -> anyhow::Result<(), SimNetError> {
@@ -372,16 +436,24 @@ impl<M: RemoteMessage> Tx<M> for SimTx<M> {
         };
         match simnet_handle() {
             Ok(handle) => match &self.src_addr {
-                Some(src_addr) if src_addr.proxy != *handle.proxy_addr() => handle
-                    .send_scheduled_event(ScheduledEvent {
-                        event: Box::new(MessageDeliveryEvent::new(
-                            self.src_addr.clone(),
-                            self.dst_addr.clone(),
-                            data,
-                        )),
-                        time: SimClock.millis_since_start(RealClock.now()),
-                    }),
-                _ => handle.send_event(Box::new(MessageDeliveryEvent::new(
+                Some(src_addr) if src_addr.proxy != *handle.proxy_addr() => {
+                    let event = Box::new(MessageRecvEvent::new(
+                        self.src_addr.clone(),
+                        self.dst_addr.clone(),
+                        data,
+                    ));
+                    let recv_time = RealClock.now();
+
+                    tokio::task::spawn(async move {
+                        SimClock.sleep_until(recv_time).await;
+
+                        if let Ok(handle) = simnet_handle() {
+                            let _ = handle.send_event(event);
+                        }
+                    });
+                    Ok(())
+                }
+                _ => handle.send_event(Box::new(MessageSendEvent::new(
                     self.src_addr.clone(),
                     self.dst_addr.clone(),
                     data,
@@ -602,8 +674,15 @@ mod tests {
             // Messages have not been receive since 10 seconds have not elapsed
             assert!(rx.rx.try_recv().is_err());
         }
-        // Advance "real" time by 100 seconds
-        tokio::time::advance(tokio::time::Duration::from_secs(100)).await;
+        tokio::time::advance(
+            // Advance "real" time by 1 ms for send time
+            tokio::time::Duration::from_millis(1)
+            // Advance "real" time by 100 seconds for inflight time
+            + tokio::time::Duration::from_secs(100)
+            // Advance "real" time by 1 ms for recv time
+            + tokio::time::Duration::from_millis(1),
+        )
+        .await;
         {
             // Allow some time for simnet to run
             tokio::task::yield_now().await;
@@ -622,16 +701,6 @@ mod tests {
             1000,
         )
         .unwrap();
-        let controller_to_dst = SimAddr::new_with_src(
-            AddressProxyPair {
-                address: "unix!@controller".parse::<ChannelAddr>().unwrap(),
-                proxy: proxy_addr.clone(),
-            },
-            "unix!@dst".parse::<ChannelAddr>().unwrap(),
-            proxy_addr.clone(),
-        )
-        .unwrap();
-        let controller_tx = sim::dial::<()>(controller_to_dst.clone()).unwrap();
 
         let client_to_dst = SimAddr::new_with_src(
             AddressProxyPair {
@@ -644,35 +713,18 @@ mod tests {
         .unwrap();
         let client_tx = sim::dial::<()>(client_to_dst).unwrap();
 
-        // 1 second of latency
-        let simnet_config_yaml = r#"
-        edges:
-        - src: unix!@controller
-          dst: unix!@dst
-          metadata:
-            latency: 1
-        "#;
-        update_config(NetworkConfig::from_yaml(simnet_config_yaml).unwrap())
-            .await
-            .unwrap();
-
         assert_eq!(SimClock.millis_since_start(RealClock.now()), 0);
         // Fast forward real time to 5 seconds
         tokio::time::advance(tokio::time::Duration::from_secs(5)).await;
         {
             // Send client message
             client_tx.try_post((), oneshot::channel().0).unwrap();
-            // Send system message
-            controller_tx.try_post((), oneshot::channel().0).unwrap();
+            tokio::time::advance(tokio::time::Duration::from_millis(1)).await;
             // Allow some time for simnet to run
-            RealClock.sleep(tokio::time::Duration::from_secs(1)).await;
+            tokio::task::yield_now().await;
         }
         let recs = simnet::simnet_handle().unwrap().close().await.unwrap();
-        assert_eq!(recs.len(), 2);
-        let end_times = recs.iter().map(|rec| rec.end_at).collect::<Vec<_>>();
         // client message was delivered at "real" time = 5 seconds
-        assert!(end_times.contains(&5000));
-        // system message was delivered at simulated time = 1 second
-        assert!(end_times.contains(&1000));
+        assert_eq!(recs.first().map(|rec| rec.end_at).unwrap(), 5000);
     }
 }