channel: net: framed: new fn FrameWrite::write_frame (#925)

Summary:
Pull Request resolved: #925

replace ad-hoc `send_frame` helper with a new `FrameWrite::write_frame` utility method (per https://www.internalfb.com/diff/D80477756?dst_version_fbid=1325594515669182&transaction_fbid=942604394698850).
update net.rs tests to call the method directly (returning the writer). add a smoke test in framed.rs to verify round-trip behavior.

Reviewed By: mariusae

Differential Revision: D80541567

fbshipit-source-id: c573f01a259a705271f32d8dac8d774d24b7d432

Author: shayne-fletcher

hyperactor/src/channel/net.rs

@@ -2577,15 +2577,6 @@ mod tests {
         (join_handle, reader, writer, rx, cancel_token)
     }
 
-    async fn send_frame<W>(writer: W, bytes: bytes::Bytes) -> W
-    where
-        W: AsyncWrite + Unpin,
-    {
-        let mut fw = FrameWrite::new(writer, bytes);
-        fw.send().await.unwrap();
-        fw.complete()
-    }
-
     async fn write_stream<M, W>(
         mut writer: W,
         session_id: u64,
@@ -2866,7 +2857,9 @@ mod tests {
             .await;
 
             for i in 0u64..5u64 {
-                writer = send_frame(writer, serialize_ack(i)).await;
+                writer = FrameWrite::write_frame(writer, serialize_ack(i))
+                    .await
+                    .unwrap();
             }
             // Wait for the acks to be processed by NetTx.
             RealClock.sleep(Duration::from_secs(3)).await;
@@ -2927,7 +2920,9 @@ mod tests {
                 // In the last iteration, ack part of the messages. This should
                 // prune them from future resent.
                 if i == n - 1 {
-                    writer = send_frame(writer, serialize_ack(1)).await;
+                    writer = FrameWrite::write_frame(writer, serialize_ack(1))
+                        .await
+                        .unwrap();
                     // Wait for the acks to be processed by NetTx.
                     RealClock.sleep(Duration::from_secs(3)).await;
                 }
@@ -2983,9 +2978,15 @@ mod tests {
                 if i == n - 1 {
                     // Intentionally ack 1 again to verify it is okay to ack
                     // messages that was already acked.
-                    writer = send_frame(writer, serialize_ack(1)).await;
-                    writer = send_frame(writer, serialize_ack(2)).await;
-                    writer = send_frame(writer, serialize_ack(3)).await;
+                    writer = FrameWrite::write_frame(writer, serialize_ack(1))
+                        .await
+                        .unwrap();
+                    writer = FrameWrite::write_frame(writer, serialize_ack(2))
+                        .await
+                        .unwrap();
+                    writer = FrameWrite::write_frame(writer, serialize_ack(3))
+                        .await
+                        .unwrap();
                     // Wait for the acks to be processed by NetTx.
                     RealClock.sleep(Duration::from_secs(3)).await;
                 }
@@ -3017,7 +3018,9 @@ mod tests {
 
                 // In the last iteration, ack part of the messages from the 2nd send.
                 if i == n - 1 {
-                    writer = send_frame(writer, serialize_ack(7)).await;
+                    writer = FrameWrite::write_frame(writer, serialize_ack(7))
+                        .await
+                        .unwrap();
                     // Wait for the acks to be processed by NetTx.
                     RealClock.sleep(Duration::from_secs(3)).await;
                 }
@@ -3061,7 +3064,9 @@ mod tests {
         let (mut reader, mut writer) = take_receiver(&receiver_storage).await;
         verify_stream(&mut reader, &[(0u64, 100u64)], None, line!()).await;
         // ack it
-        writer = send_frame(writer, serialize_ack(0)).await;
+        writer = FrameWrite::write_frame(writer, serialize_ack(0))
+            .await
+            .unwrap();
         // confirm Tx received ack
         //
         // Using `is_err` to confirm the message is delivered/acked is confusing,
@@ -3072,7 +3077,9 @@ mod tests {
         // Although Tx did not actually send seq=1, we still ack it from Rx to
         // pretend Tx already sent it, just it did not know it was sent
         // successfully.
-        let _ = send_frame(writer, serialize_ack(1)).await;
+        let _ = FrameWrite::write_frame(writer, serialize_ack(1))
+            .await
+            .unwrap();
 
         let (return_channel_tx, return_channel_rx) = oneshot::channel();
         net_tx.try_post(101, return_channel_tx).unwrap();
@@ -3104,7 +3111,9 @@ mod tests {
         // Confirm message is sent to rx.
         verify_stream(&mut reader, &[(0u64, 100u64)], None, line!()).await;
         // ack it
-        let _ = send_frame(writer, serialize_ack(0)).await;
+        let _ = FrameWrite::write_frame(writer, serialize_ack(0))
+            .await
+            .unwrap();
         RealClock.sleep(Duration::from_secs(3)).await;
         // Channel should be still alive because ack was sent.
         assert!(!tx_status.has_changed().unwrap());

hyperactor/src/channel/net/framed.rs

@@ -170,6 +170,41 @@ impl<W: AsyncWrite + Unpin> FrameWrite<W> {
         let Self { writer, .. } = self;
         writer
     }
+
+    /// Writes a single frame into the underlying writer and returns
+    /// it.
+    ///
+    /// This is a convenience for the common pattern:
+    /// `FrameWrite::new(writer, bytes).send().await?.complete()`.
+    ///
+    /// Frame writes are atomic: either the entire frame is sent, or
+    /// an error is returned. No partial frames are observed by the
+    /// receiver.
+    ///
+    /// # Arguments
+    ///
+    /// * `writer` — the `AsyncWrite` sink to write into.
+    /// * `bytes` — the serialized frame body to send.
+    ///
+    /// # Returns
+    ///
+    /// On success, returns the underlying writer so the caller can
+    /// continue using it for further frames. On error, returns the
+    /// I/O error from the underlying write.
+    ///
+    /// # Examples
+    ///
+    /// ```ignore
+    /// use bytes::Bytes;
+    ///
+    /// // `writer` is any AsyncWrite + Unpin (e.g. a tokio `WriteHalf`)
+    /// let writer = FrameWrite::write_frame(writer, Bytes::from_static(b"hello")).await?;
+    /// ```
+    pub async fn write_frame(writer: W, bytes: Bytes) -> std::io::Result<W> {
+        let mut fw = FrameWrite::new(writer, bytes);
+        fw.send().await?;
+        Ok(fw.complete())
+    }
 }
 
 #[cfg(test)]
@@ -208,4 +243,26 @@ mod tests {
     }
 
     // todo: test cancellation, frame size
+
+    #[tokio::test]
+    async fn test_write_frame_smoke() {
+        let (a, b) = tokio::io::duplex(4096);
+        let (r, _w_unused) = tokio::io::split(a);
+        let (_r_unused, w) = tokio::io::split(b);
+
+        let mut reader = FrameReader::new(r, 1024);
+
+        let w = FrameWrite::write_frame(w, Bytes::from_static(b"hello"))
+            .await
+            .unwrap();
+        let _ = FrameWrite::write_frame(w, Bytes::from_static(b"world"))
+            .await
+            .unwrap();
+
+        let f1 = reader.next().await.unwrap().unwrap();
+        let f2 = reader.next().await.unwrap().unwrap();
+
+        assert_eq!(f1.as_ref(), b"hello");
+        assert_eq!(f2.as_ref(), b"world");
+    }
 }