Shutdown NetTx stream when it exits (meta-pytorch#1993)

Summary:
Pull Request resolved: meta-pytorch#1993

We currently do not explicitly close `NetTx`'s write stream when it exits. As a result, we will see the following log on the `NetRx` side when using `MetaTls`:

```
[-]I1124 14:09:27.934423 629996 fbcode/monarch/hyperactor/src/channel/net/server.rs:592] error processing peer connection, source:tcp:[2401:db00:eef0:1120:3520:0:6c09:3cba]:57588, dest:metatls:devvm16922.vll0.facebook.com:35865, error:session metatls:devvm16922.vll0.facebook.com:35865.18243036951658569703<-tcp:[2401:db00:eef0:1120:3520:0:6c09:3cba]:57588: reading into Frame with M = u64

Caused by:
    peer closed connection without sending TLS close_notify: https://docs.rs/rustls/latest/rustls/manual/_03_howto/index.html#unexpected-eof
```

This diff fixes that.

Reviewed By: mariusae

Differential Revision: D87823699

fbshipit-source-id: a84c57d5a2b0ebc384931154908ffd9c03eb1851

Author: pzhan9

hyperactor/src/channel/net.rs

@@ -1086,11 +1086,12 @@ mod tests {
         assert!(rx.await.is_err());
     }
 
-    #[tracing_test::traced_test]
-    #[tokio::test]
+    #[async_timed_test(timeout_secs = 60)]
     // TODO: OSS: failed to retrieve ipv6 address
     #[cfg_attr(not(fbcode_build), ignore)]
     async fn test_meta_tls_basic() {
+        hyperactor_telemetry::initialize_logging_for_test();
+
         let addr = ChannelAddr::any(ChannelTransport::MetaTls(TlsMode::IpV6));
         let meta_addr = match addr {
             ChannelAddr::MetaTls(meta_addr) => meta_addr,

hyperactor/src/channel/net/client.rs

@@ -645,29 +645,39 @@ async fn run<M: RemoteMessage>(
         );
     }
 
-    if let Conn::Connected {
-        write_state: WriteState::Writing(mut frame_writer, ()),
-        ..
-    } = conn
-    {
-        if let Err(err) = frame_writer.send().await {
-            tracing::info!(
-                parent: &span,
-                dest = %dest,
-                error = %err,
-                session_id = session_id,
-                "write error during cleanup"
-            );
-        } else if let Err(err) = frame_writer.complete().flush().await {
-            tracing::info!(
-                parent: &span,
-                dest = %dest,
-                error = %err,
-                session_id = session_id,
-                "flush error during cleanup"
-            );
+    match conn {
+        Conn::Connected { write_state, .. } => {
+            let write_half = match write_state {
+                WriteState::Writing(mut frame_writer, ()) => {
+                    if let Err(err) = frame_writer.send().await {
+                        tracing::info!(
+                            parent: &span,
+                            dest = %dest,
+                            error = %err,
+                            session_id = session_id,
+                            "write error during cleanup"
+                        );
+                    }
+                    Some(frame_writer.complete())
+                }
+                WriteState::Idle(writer) => Some(writer),
+                WriteState::Broken => None,
+            };
+
+            if let Some(mut w) = write_half {
+                if let Err(err) = w.shutdown().await {
+                    tracing::info!(
+                        parent: &span,
+                        dest = %dest,
+                        error = %err,
+                        session_id = session_id,
+                        "failed to shutdown NetTx write stream during cleanup"
+                    );
+                }
+            }
         }
-    }
+        Conn::Disconnected(_) => (),
+    };
 
     tracing::info!(
         parent: &span,