Add spans to points along the message send/recv path (#1833)

Summary:
Pull Request resolved: #1833

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

Reviewed By: vidhyav

Differential Revision: D86802704

fbshipit-source-id: bf0708ef5021e38f74658071a36a08960e8658eb

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/client.rs

@@ -17,6 +17,7 @@ use std::ops::DerefMut;
 use backoff::ExponentialBackoffBuilder;
 use backoff::backoff::Backoff;
 use enum_as_inner::EnumAsInner;
+use hyperactor_telemetry::skip_record;
 use tokio::io::AsyncWriteExt;
 use tokio::io::ReadHalf;
 use tokio::io::WriteHalf;
@@ -789,8 +790,6 @@ where
         .as_ref()
         .map(|acked_seq| AckedSeqValue(acked_seq.clone()));
 
-    use hyperactor_telemetry::skip_record;
-
     tracing::span!(
         Level::ERROR,
         "net i/o loop",
@@ -926,7 +925,7 @@ where
             tokio::select! {
                 biased;
 
-                ack_result = reader.next() => {
+                ack_result = reader.next().instrument(tracing::span!(Level::ERROR, "read ack", skip_record)) => {
                     match ack_result {
                         Ok(Some(buffer)) => {
                             match deserialize_response(buffer) {
@@ -1018,7 +1017,7 @@ where
 
                 // We have to be careful to manage outgoing write states, so that we never write
                 // partial frames in the presence cancellation.
-                send_result = write_state.send() => {
+                send_result = write_state.send().instrument(tracing::span!(Level::ERROR, "write bytes", skip_record)) => {
                     match send_result {
                         Ok(()) => {
                             let mut message = outbox.pop_front().expect("outbox should not be empty");