Add more logs for server (#492)

Summary:
Pull Request resolved: #492

helpful to debug

Reviewed By: shayne-fletcher

Differential Revision: D78116522

fbshipit-source-id: 9889b1c4d01cf3bae95b40f9cb1a4b5e75521b40

Author: pzhan9

hyperactor/src/channel/net.rs

@@ -751,7 +751,8 @@ impl<M: RemoteMessage> Rx<M> for NetRx<M> {
 
 impl<M: RemoteMessage> Drop for NetRx<M> {
     fn drop(&mut self) {
-        self.2.stop();
+        self.2
+            .stop(&format!("NetRx dropped; channel address: {}", self.1));
     }
 }
 
@@ -780,12 +781,13 @@ impl ServerHandle {
     /// incoming connection requests, and drain pending operations
     /// on active connections. After draining is completed, the
     /// connections are closed.
-    pub fn stop(&self) {
+    fn stop(&self, reason: &str) {
+        tracing::info!("stopping server: {}; reason: {}", self, reason);
         self.cancel_token.cancel();
     }
 
     /// Reference to the channel that was used to start the server.
-    pub fn local_channel_addr(&self) -> &ChannelAddr {
+    fn local_channel_addr(&self) -> &ChannelAddr {
         &self.channel_addr
     }
 }
@@ -1145,7 +1147,7 @@ where
         let _ = tracing::info_span!("channel_listen_accept_loop");
         tokio::select! {
             result = listener.accept() => {
-                tracing::debug!("listener accepted a new connection.");
+                tracing::debug!("listener accepted a new connection to {}", listener_channel_addr);
                 match result {
                     Ok((stream, addr)) => {
                         let tx = tx.clone();
@@ -1180,13 +1182,13 @@ where
                     });
                     }
                     Err(err) => {
-                        tracing::info!("serve {}: accept error: {}", listener_channel_addr,err)
+                        tracing::info!("serve {}: accept error: {}", listener_channel_addr, err)
                     }
                 }
             }
 
             _ = parent_cancel_token.cancelled() => {
-                tracing::info!("received parent token cancellation");
+                tracing::info!("serve {}: received parent token cancellation", listener_channel_addr);
                 break Ok(());
             }
 
@@ -1962,7 +1964,7 @@ mod tests {
         RealClock.sleep(Duration::from_secs(5)).await;
 
         // Stop the server.
-        rx1.2.stop();
+        rx1.2.stop("from testing");
         assert_matches!(rx1.recv().await.unwrap_err(), ChannelError::Closed);
 
         // Send the message is allowed even when the server is down.

hyperactor/src/mailbox.rs

@@ -809,7 +809,8 @@ impl MailboxServerHandle {
     /// join the handle by awaiting the [`MailboxServerHandle`] future.
     ///
     /// Stop should be called at most once.
-    pub fn stop(&self) {
+    pub fn stop(&self, reason: &str) {
+        tracing::info!("stopping mailbox server; reason: {}", reason);
         self.stopped_tx.send(true).expect("stop called twice");
     }
 }
@@ -2554,7 +2555,7 @@ mod tests {
             .serialize_and_send_once(port, 123u64, monitored_return_handle())
             .unwrap();
         assert_eq!(receiver.recv().await.unwrap(), 123u64);
-        serve_handle.stop();
+        serve_handle.stop("fromt test");
         serve_handle.await.unwrap().unwrap();
     }
 
@@ -2595,7 +2596,7 @@ mod tests {
             .serialize_and_send_once(port, msg, monitored_return_handle())
             .unwrap();
         assert_eq!(receiver.recv().await.unwrap(), msg);
-        serve_handle.stop();
+        serve_handle.stop("from test");
         serve_handle.await.unwrap().unwrap();
     }
 

hyperactor_mesh/src/alloc/local.rs

@@ -220,7 +220,7 @@ impl Alloc for LocalAlloc {
                     };
 
                     // Stop serving the mailbox.
-                    proc_to_stop.handle.stop();
+                    proc_to_stop.handle.stop("received Action::Stop");
 
                     if let Err(err) = proc_to_stop
                         .proc

hyperactor_mesh/src/alloc/remoteprocess.rs

@@ -371,7 +371,7 @@ impl RemoteProcessAllocator {
         )
         .await;
 
-        mailbox_handle.stop();
+        mailbox_handle.stop("alloc stopped");
         if let Err(e) = mailbox_handle.await {
             tracing::error!("failed to join forwarder: {}", e);
         }

hyperactor_multiprocess/src/proc_actor.rs

@@ -488,7 +488,7 @@ impl ProcActor {
     /// Shutdown the mailbox server to free up rx and its cooresponding listen address.
     /// Because in the next bootstrap attempt, the same listen address will be used.
     async fn failed_proc_bootstrap_cleanup(mailbox_handle: MailboxServerHandle) {
-        mailbox_handle.stop();
+        mailbox_handle.stop("failed proc bootstrap cleanup");
         if let Err(shutdown_err) = mailbox_handle.await {
             // Ignore the shutdown error and populate the original error.
             tracing::error!(

hyperactor_multiprocess/src/system.rs

@@ -145,7 +145,7 @@ impl ServerHandle {
     pub async fn stop(&self) -> Result<(), ActorError> {
         // TODO: this needn't be async
         self.actor_handle.drain_and_stop()?;
-        self.mailbox_handle.stop();
+        self.mailbox_handle.stop("system server stopped");
         Ok(())
     }
 

hyperactor_multiprocess/src/system_actor.rs

@@ -2722,7 +2722,7 @@ mod tests {
 
         // Now kill pong's mailbox server making message delivery
         // between procs impossible.
-        proc_actor_1.mailbox.stop();
+        proc_actor_1.mailbox.stop("from testing");
         proc_actor_1.mailbox.await.unwrap().unwrap();
 
         // That in itself shouldn't be a problem. Check the world