Add spans to points along the message send/recv path

Summary: Our observability only shows us when a message is sent out and when it is delivered to its final destination. We want to add spans to some of the important stops along the way

Differential Revision: D86802704

Author: thomasywang

hyperactor/src/channel.rs

@@ -119,11 +119,13 @@ pub trait Tx<M: RemoteMessage>: std::fmt::Debug {
     /// message is either delivered, or we eventually discover that
     /// the channel has failed and it will be sent back on `return_channel`.
     #[allow(clippy::result_large_err)] // TODO: Consider reducing the size of `SendError`.
+    #[hyperactor::instrument_infallible]
     fn try_post(&self, message: M, return_channel: oneshot::Sender<SendError<M>>) {
         self.do_post(message, Some(return_channel));
     }
 
     /// Enqueue a message to be sent on the channel.
+    #[hyperactor::instrument_infallible]
     fn post(&self, message: M) {
         self.do_post(message, None);
     }
@@ -803,6 +805,7 @@ enum ChannelRxKind<M: RemoteMessage> {
 
 #[async_trait]
 impl<M: RemoteMessage> Rx<M> for ChannelRx<M> {
+    #[hyperactor::instrument]
     async fn recv(&mut self) -> Result<M, ChannelError> {
         match &mut self.inner {
             ChannelRxKind::Local(rx) => rx.recv().await,

hyperactor/src/channel/net.rs

@@ -86,6 +86,7 @@ use tokio::time::Duration;
 use tokio::time::Instant;
 use tokio_util::net::Listener;
 use tokio_util::sync::CancellationToken;
+use tracing::Instrument;
 
 use super::*;
 use crate::RemoteMessage;
@@ -652,7 +653,11 @@ impl<M: RemoteMessage> NetTx<M> {
                         unacked,
                     }),
                     conn,
-                ) if outbox.is_empty() && unacked.is_empty() => match receiver.recv().await {
+                ) if outbox.is_empty() && unacked.is_empty() => match receiver
+                    .recv()
+                    .instrument(tracing::debug_span!("receiver.recv"))
+                    .await
+                {
                     Some(msg) => match outbox.push_back(msg) {
                         Ok(()) => {
                             let running = State::Running(Deliveries { outbox, unacked });
@@ -828,7 +833,7 @@ impl<M: RemoteMessage> NetTx<M> {
                         // UnboundedReceiver::recv() is cancel safe.
                         // Only checking mspc channel when outbox is empty. In this way, we prioritize
                         // sending messages already in outbox.
-                        work_result = receiver.recv(), if outbox.is_empty() => {
+                        work_result = receiver.recv().instrument(tracing::debug_span!("receiver.recv")), if outbox.is_empty() => {
                             match work_result {
                                 Some(msg) => {
                                     match outbox.push_back(msg) {

hyperactor/src/channel/net/framed.rs

@@ -66,6 +66,7 @@ impl<R: AsyncRead + Unpin> FrameReader<R> {
     ///   `max_frame_length`. **This error is fatal:** once returned,
     ///   the `FrameReader` must be dropped; the underlying connection
     ///   is no longer valid.
+    #[tracing::instrument(skip_all)]
     pub async fn next(&mut self) -> io::Result<Option<Bytes>> {
         loop {
             match &mut self.state {
@@ -218,6 +219,7 @@ impl<W: AsyncWrite + Unpin, B: Buf> FrameWrite<W, B> {
     /// returned futures at any time. Upon completion, the frame is guaranteed to be
     /// written, unless an error was encountered, in which case the underlying stream
     /// is in an undefined state.
+    #[tracing::instrument(skip_all)]
     pub async fn send(&mut self) -> io::Result<()> {
         loop {
             if self.len_buf.has_remaining() {

hyperactor/src/mailbox.rs

@@ -1161,6 +1161,7 @@ impl MailboxClient {
 }
 
 impl MailboxSender for MailboxClient {
+    #[hyperactor::instrument_infallible]
     fn post_unchecked(
         &self,
         envelope: MessageEnvelope,

hyperactor_mesh/src/actor_mesh.rs

@@ -77,6 +77,7 @@ declare_attrs! {
 /// Common implementation for `ActorMesh`s and `ActorMeshRef`s to cast
 /// an `M`-typed message
 #[allow(clippy::result_large_err)] // TODO: Consider reducing the size of `CastError`.
+#[hyperactor::instrument]
 pub(crate) fn actor_mesh_cast<A, M>(
     cx: &impl context::Actor,
     actor_mesh_id: ActorMeshId,

hyperactor_mesh/src/comm.rs

@@ -345,6 +345,7 @@ impl Handler<CommActorMode> for CommActor {
 // TODO(T218630526): reliable casting for mutable topology
 #[async_trait]
 impl Handler<CastMessage> for CommActor {
+    #[hyperactor::instrument]
     async fn handle(&mut self, cx: &Context<Self>, cast_message: CastMessage) -> Result<()> {
         // Always forward the message to the root rank of the slice, casting starts from there.
         let slice = cast_message.dest.slice.clone();
@@ -375,6 +376,7 @@ impl Handler<CastMessage> for CommActor {
 
 #[async_trait]
 impl Handler<ForwardMessage> for CommActor {
+    #[hyperactor::instrument]
     async fn handle(&mut self, cx: &Context<Self>, fwd_message: ForwardMessage) -> Result<()> {
         let ForwardMessage {
             sender,

monarch_hyperactor/src/actor.rs

@@ -712,6 +712,7 @@ impl PanicFlag {
 
 #[async_trait]
 impl Handler<PythonMessage> for PythonActor {
+    #[hyperactor::instrument]
     async fn handle(
         &mut self,
         cx: &Context<PythonActor>,

monarch_hyperactor/src/actor_mesh.rs

@@ -137,6 +137,7 @@ pub(crate) fn to_hy_sel(selection: &str) -> PyResult<Selection> {
 
 #[pymethods]
 impl PythonActorMesh {
+    #[hyperactor::instrument]
     fn cast(
         &self,
         message: &PythonMessage,