Cleanup/refactor connection helper library (#412)

Summary:
Pull Request resolved: #412

This diff does some refactor of the connection helper:
1) Add a `ActorConnection` wrapper which to model the connection which, similar
    to `TcpStream` can be split into owned read and write halves.
2) Make the `AsyncRead`/`AsyncWrite`s that represent the connection have a
    `.peer()` method to return the `ActorId` of the other side.

Reviewed By: shayne-fletcher, mariusae

Differential Revision: D77596371

fbshipit-source-id: 96aa5d5ea271d0dc14d7b1f52f6eb818e06b6459

Author: andrewjcg

hyperactor_mesh/Cargo.toml

@@ -40,6 +40,7 @@ libc = "0.2.139"
 mockall = "0.13.1"
 ndslice = { version = "0.0.0", path = "../ndslice" }
 nix = { version = "0.29.0", features = ["dir", "event", "hostname", "inotify", "ioctl", "mman", "mount", "net", "poll", "ptrace", "reboot", "resource", "sched", "signal", "term", "time", "user", "zerocopy"] }
+pin-project = "0.4.30"
 preempt_rwlock = { version = "0.0.0", path = "../preempt_rwlock" }
 rand = { version = "0.8", features = ["small_rng"] }
 serde = { version = "1.0.185", features = ["derive", "rc"] }

hyperactor_mesh/src/code_sync/rsync.rs

@@ -18,6 +18,9 @@ use anyhow::Result;
 use anyhow::bail;
 use anyhow::ensure;
 use async_trait::async_trait;
+use futures::StreamExt;
+use futures::TryStreamExt;
+use futures::stream;
 use futures::try_join;
 use hyperactor::Actor;
 use hyperactor::Handler;
@@ -40,7 +43,6 @@ use crate::actor_mesh::ActorMesh;
 use crate::code_sync::WorkspaceLocation;
 use crate::connect::Connect;
 use crate::connect::accept;
-use crate::connect::connect_mesh;
 
 pub async fn do_rsync(addr: &SocketAddr, workspace: &Path) -> Result<()> {
     let output = Command::new("rsync")
@@ -190,36 +192,45 @@ impl Handler<Connect> for RsyncActor {
     ) -> Result<(), anyhow::Error> {
         let (mut local, mut stream) = try_join!(
             async { Ok(TcpStream::connect(self.daemon.addr()).await?) },
-            async {
-                let (rd, wr) = accept(cx, message).await?;
-                anyhow::Ok(tokio::io::join(rd, wr))
-            }
+            accept(cx, cx.self_id().clone(), message),
         )?;
         tokio::io::copy_bidirectional(&mut local, &mut stream).await?;
         Ok(())
     }
 }
 
-pub async fn rsync_mesh<M>(actor_mesh: M, workspace: PathBuf) -> Result<()>
+pub async fn rsync_mesh<M>(actor_mesh: &M, workspace: PathBuf) -> Result<()>
 where
     M: ActorMesh<Actor = RsyncActor>,
 {
-    connect_mesh(actor_mesh, async move |rd, wr| {
-        let workspace = workspace.clone();
-        let listener = TcpListener::bind(("::1", 0)).await?;
-        let addr = listener.local_addr()?;
-        let mut local = tokio::io::join(rd, wr);
-        try_join!(
-            async move { do_rsync(&addr, &workspace).await },
-            async move {
-                let (mut stream, _) = listener.accept().await?;
-                tokio::io::copy_bidirectional(&mut stream, &mut local).await?;
-                anyhow::Ok(())
-            },
-        )?;
-        anyhow::Ok(())
-    })
-    .await
+    // We avoid casting here as we need point-to-point connections to each individual actor.
+    stream::iter(actor_mesh.iter_actor_refs())
+        // Connect to all actors in the mesh.
+        .map(|actor| async move {
+            let mailbox = actor_mesh.proc_mesh().client();
+            let (connect, completer) = Connect::allocate(mailbox.actor_id().clone(), mailbox);
+            actor.send(mailbox, connect)?;
+            completer.complete().await
+        })
+        // Max the connections run in parallel.
+        .buffer_unordered(usize::MAX)
+        // Initiate the rsync, in parallel.
+        .try_for_each_concurrent(None, |mut local| async {
+            let workspace = workspace.clone();
+            let listener = TcpListener::bind(("::1", 0)).await?;
+            let addr = listener.local_addr()?;
+            try_join!(
+                async move { do_rsync(&addr, &workspace).await },
+                async move {
+                    let (mut stream, _) = listener.accept().await?;
+                    tokio::io::copy_bidirectional(&mut stream, &mut local).await?;
+                    anyhow::Ok(())
+                },
+            )?;
+            anyhow::Ok(())
+        })
+        .await?;
+    Ok(())
 }
 
 #[cfg(test)]