mailbox: undeliverable: custom_monitored_return_handle (#754)

Summary:
Pull Request resolved: #754

new function `custom_monitored_return_handle` to replace the few remaining `monitored_return_handle` in order that we get better insight into where they still exist.

Reviewed By: mariusae

Differential Revision: D79593009

fbshipit-source-id: eca4b267333e0f1201508135a9710441ba10beda

Author: shayne-fletcher

hyperactor/src/mailbox.rs

@@ -127,6 +127,7 @@ mod undeliverable;
 /// For [`Undeliverable`], a message type for delivery failures.
 pub use undeliverable::Undeliverable;
 pub use undeliverable::UndeliverableMessageError;
+pub use undeliverable::custom_monitored_return_handle;
 pub use undeliverable::monitored_return_handle; // TODO: Audit
 pub use undeliverable::supervise_undeliverable_messages;
 /// For [`MailboxAdminMessage`], a message type for mailbox administration.

hyperactor/src/mailbox/undeliverable.rs

@@ -70,6 +70,24 @@ pub fn monitored_return_handle() -> PortHandle<Undeliverable<MessageEnvelope>> {
     return_handle.clone()
 }
 
+/// Now that monitored return handles are rare, it's becoming helpful
+/// to get insights into where they are getting used (so that they can
+/// be eliminated and replaced with something better).
+#[track_caller]
+pub fn custom_monitored_return_handle(caller: &str) -> PortHandle<Undeliverable<MessageEnvelope>> {
+    let caller = caller.to_owned();
+    let (return_handle, mut rx) = new_undeliverable_port();
+    tokio::task::spawn(async move {
+        while let Ok(Undeliverable(mut envelope)) = rx.recv().await {
+            envelope.try_set_error(DeliveryError::BrokenLink(
+                "message returned to undeliverable port".to_string(),
+            ));
+            tracing::error!("{caller} took back an undeliverable message: {}", envelope);
+        }
+    });
+    return_handle
+}
+
 /// Returns a message envelope to its original sender.
 pub(crate) fn return_undeliverable(
     return_handle: PortHandle<Undeliverable<MessageEnvelope>>,

hyperactor_mesh/src/proc_mesh.rs

@@ -193,7 +193,7 @@ impl ProcMesh {
         }
         router
             .clone()
-            .serve(router_rx, mailbox::monitored_return_handle());
+            .serve(router_rx, mailbox::custom_monitored_return_handle("router"));
 
         // Set up a client proc for the mesh itself, so that we can attach ourselves
         // to it, and communicate with the agents. We wire it into the same router as
@@ -207,9 +207,10 @@ impl ProcMesh {
             client_proc_id.clone(),
             BoxedMailboxSender::new(router.clone()),
         );
-        client_proc
-            .clone()
-            .serve(client_rx, mailbox::monitored_return_handle());
+        client_proc.clone().serve(
+            client_rx,
+            mailbox::custom_monitored_return_handle("client proc"),
+        );
         router.bind(client_proc_id.clone().into(), client_proc_addr.clone());
 
         // Bind this router to the global router, to enable cross-mesh routing.