Close channel when NetRx is being stopped (meta-pytorch#2010)

Summary:

When `NetRx` is stopped, we should let its `NetTx` half know as well, so it will exit too and stop reconnecting. Otherwise, those reconnections would all fail, and lead to a log spew.

Reviewed By: shayne-fletcher

Differential Revision: D87926632

Author: pzhan9

hyperactor/src/channel/net.rs

@@ -108,8 +108,10 @@ enum Frame<M> {
 #[derive(Debug, Serialize, Deserialize, EnumAsInner)]
 enum NetRxResponse {
     Ack(u64),
-    /// This channel is closed with the given reason. NetTx should stop reconnecting.
+    /// This session is rejected with the given reason. NetTx should stop reconnecting.
     Reject(String),
+    /// This channel is closed.
+    Closed,
 }
 
 fn serialize_response(response: NetRxResponse) -> Result<Bytes, bincode::Error> {
@@ -1625,6 +1627,9 @@ mod tests {
             handle.await.unwrap().unwrap();
             // mpsc is closed too and there should be no unread message left.
             assert!(rx.recv().await.is_none());
+            // should send NetRxResponse::Closed before stopping server.
+            let bytes = reader.next().await.unwrap().unwrap();
+            assert!(deserialize_response(bytes).unwrap().is_closed());
             // No more acks from server.
             assert!(reader.next().await.unwrap().is_none());
         };
@@ -1659,6 +1664,9 @@ mod tests {
         handle.await.unwrap().unwrap();
         // mpsc is closed too and there should be no unread message left.
         assert!(rx.recv().await.is_none());
+        // should send NetRxResponse::Closed before stopping server.
+        let bytes = reader.next().await.unwrap().unwrap();
+        assert!(deserialize_response(bytes).unwrap().is_closed());
         // No more acks from server.
         assert!(reader.next().await.unwrap().is_none());
     }
@@ -2398,4 +2406,41 @@ mod tests {
         let bytes = reader.next().await.unwrap().unwrap();
         assert!(deserialize_response(bytes).unwrap().is_reject());
     }
+
+    #[async_timed_test(timeout_secs = 60)]
+    // TODO: OSS: called `Result::unwrap()` on an `Err` value: Listen(Tcp([::1]:0), Os { code: 99, kind: AddrNotAvailable, message: "Cannot assign requested address" })
+    #[cfg_attr(not(fbcode_build), ignore)]
+    async fn test_stop_net_tx_after_stopping_net_rx() {
+        hyperactor_telemetry::initialize_logging_for_test();
+
+        let config = config::global::lock();
+        let _guard =
+            config.override_key(config::MESSAGE_DELIVERY_TIMEOUT, Duration::from_secs(300));
+        let (addr, mut rx) = tcp::serve::<u64>("[::1]:0".parse().unwrap()).unwrap();
+        let socket_addr = match addr {
+            ChannelAddr::Tcp(a) => a,
+            _ => panic!("unexpected channel type"),
+        };
+        let tx = tcp::dial::<u64>(socket_addr);
+        // NetTx will not establish a connection until it sends the 1st message.
+        // Without a live connection, NetTx cannot received the Closed message
+        // from NetRx. Therefore, we need to send a message to establish the
+        //connection.
+        tx.send(100).await.unwrap();
+        assert_eq!(rx.recv().await.unwrap(), 100);
+        // Drop rx will close the NetRx server.
+        rx.2.stop("testing");
+        assert!(rx.recv().await.is_err());
+
+        // NetTx will only read from the stream when it needs to send a message
+        // or wait for an ack. Therefore we need to send a message to trigger that.
+        tx.post(101);
+        let mut watcher = tx.status().clone();
+        // When NetRx exits, it should notify NetTx to exit as well.
+        let _ = watcher.wait_for(|val| *val == TxStatus::Closed).await;
+        // wait_for could return Err due to race between when watch's sender was
+        // dropped and when wait_for was called. So we still need to do an
+        // equality check.
+        assert_eq!(*watcher.borrow(), TxStatus::Closed);
+    }
 }

hyperactor/src/channel/net/client.rs

@@ -940,7 +940,19 @@ where
                                                     );
                                             (State::Closing {
                                                 deliveries: Deliveries{outbox, unacked},
-                                                reason: format!("{log_id}: {error_msg}"),
+                                                reason: error_msg,
+                                            }, Conn::reconnect_with_default())
+                                        }
+                                        NetRxResponse::Closed => {
+                                            let msg = "server closed the channel".to_string();
+                                            tracing::info!(
+                                                        dest = %link.dest(),
+                                                        session_id = session_id,
+                                                        "{}", msg
+                                                    );
+                                            (State::Closing {
+                                                deliveries: Deliveries{outbox, unacked},
+                                                reason: msg,
                                             }, Conn::reconnect_with_default())
                                         }
                                     }

hyperactor/src/channel/net/server.rs

@@ -89,7 +89,11 @@ impl<S: AsyncRead + AsyncWrite + Send + 'static + Unpin> ServerConn<S> {
     ) -> (Next, Result<(), anyhow::Error>) {
         #[derive(Debug)]
         enum RejectConn {
-            Yes(String),
+            /// Reject the connection due to the given error.
+            EncounterError(String),
+            /// The server is being closed.
+            ServerClosing,
+            /// Do not reject the connection.
             No,
         }
 
@@ -170,7 +174,7 @@ impl<S: AsyncRead + AsyncWrite + Send + 'static + Unpin> ServerConn<S> {
                 // Have a tick to abort select! call to make sure the ack for the last message can get the chance
                 // to be sent as a result of time interval being reached.
                 _ = RealClock.sleep_until(last_ack_time + ack_time_interval), if next.ack < next.seq => {},
-                _ = cancel_token.cancelled() => break (next, Ok(()), RejectConn::No),
+                _ = cancel_token.cancelled() => break (next, Ok(()), RejectConn::ServerClosing),
                 bytes_result = self.reader.next() => {
                     rcv_raw_frame_count += 1;
                     // First handle transport-level I/O errors, and EOFs.
@@ -231,7 +235,7 @@ impl<S: AsyncRead + AsyncWrite + Send + 'static + Unpin> ServerConn<S> {
                             break (
                                 next,
                                 Err(anyhow::anyhow!("{log_id}: unexpected init frame")),
-                                RejectConn::Yes("expect Frame::Message; got Frame::Int".to_string()),
+                                RejectConn::EncounterError("expect Frame::Message; got Frame::Int".to_string()),
                             )
                         },
                         // Ignore retransmits.
@@ -260,7 +264,7 @@ impl<S: AsyncRead + AsyncWrite + Send + 'static + Unpin> ServerConn<S> {
                                 break (
                                     next,
                                     Err(anyhow::anyhow!(format!("{log_id}: {error_msg}"))),
-                                    RejectConn::Yes(error_msg),
+                                    RejectConn::EncounterError(error_msg),
                                 )
                             }
                             match self.send_with_buffer_metric(session_id, &tx, message).await {
@@ -391,12 +395,20 @@ impl<S: AsyncRead + AsyncWrite + Send + 'static + Unpin> ServerConn<S> {
         }
 
         if self.write_state.is_idle()
-            && let RejectConn::Yes(reason) = reject_conn
+            && matches!(
+                reject_conn,
+                RejectConn::EncounterError(_) | RejectConn::ServerClosing
+            )
         {
             let Ok(writer) = replace(&mut self.write_state, WriteState::Broken).into_idle() else {
                 panic!("illegal state");
             };
-            if let Ok(data) = serialize_response(NetRxResponse::Reject(reason)) {
+            let rsp = match reject_conn {
+                RejectConn::EncounterError(reason) => NetRxResponse::Reject(reason),
+                RejectConn::ServerClosing => NetRxResponse::Closed,
+                RejectConn::No => panic!("illegal state"),
+            };
+            if let Ok(data) = serialize_response(rsp) {
                 match FrameWrite::new(
                     writer,
                     data,