[runtime] Fix Network Flake (#3494)

Author: patrick-ogrady

runtime/src/network/mod.rs

@@ -17,8 +17,10 @@ stability_scope!(ALPHA, cfg(all(not(target_arch = "wasm32"), feature = "iouring-
 #[cfg(test)]
 mod tests {
     use crate::{IoBuf, IoBufs, Listener, Sink, Stream};
+    use commonware_utils::sync::Barrier;
     use futures::join;
-    use std::net::SocketAddr;
+    use std::{net::SocketAddr, sync::Arc};
+    use tokio::task::JoinSet;
 
     const CLIENT_SEND_DATA: &[u8] = b"client_send_data";
     const SERVER_SEND_DATA: &[u8] = b"server_send_data";
@@ -47,12 +49,11 @@ mod tests {
         // Get the local address of the listener
         let listener_addr = listener.local_addr().expect("Failed to get local address");
 
-        let runtime = tokio::runtime::Handle::current();
-
-        // Spawn server
-        let server = runtime.spawn(async move {
+        // Spawn server. Returning the socket halves keeps them alive until both
+        // join handles are awaited below.
+        let server = tokio::spawn(async move {
+            // Server accepts a client, verifies the payload, and sends a reply.
             let (_, mut sink, mut stream) = listener.accept().await.expect("Failed to accept");
-
             let received = stream
                 .recv(CLIENT_SEND_DATA.len())
                 .await
@@ -61,10 +62,14 @@ mod tests {
             sink.send(IoBuf::from(SERVER_SEND_DATA))
                 .await
                 .expect("Failed to send");
+            (sink, stream)
         });
 
-        // Spawn client, connect to server, send and receive data over connection
-        let client = runtime.spawn(async move {
+        // Spawn client, connect to server, send and receive data over connection.
+        // Returning the socket halves keeps them alive until both join handles
+        // are awaited below.
+        let client = tokio::spawn(async move {
+            // Client connects to the server, sends a payload, and reads the reply.
             // Connect to the server
             let (mut sink, mut stream) = network
                 .dial(listener_addr)
@@ -74,18 +79,18 @@ mod tests {
             sink.send(IoBuf::from(CLIENT_SEND_DATA))
                 .await
                 .expect("Failed to send data");
-
             let received = stream
                 .recv(SERVER_SEND_DATA.len())
                 .await
                 .expect("Failed to receive data");
             assert_eq!(received.coalesce(), SERVER_SEND_DATA);
+            (sink, stream)
         });
 
         // Wait for both tasks to complete
         let (server_result, client_result) = join!(server, client);
-        assert!(server_result.is_ok());
-        assert!(client_result.is_ok());
+        server_result.expect("Server task failed");
+        client_result.expect("Client task failed");
     }
 
     // Test sending a multi-buffer payload.
@@ -99,8 +104,6 @@ mod tests {
         // Get the local address of the listener
         let listener_addr = listener.local_addr().expect("Failed to get local address");
 
-        let runtime = tokio::runtime::Handle::current();
-
         // Build one logical message from multiple chunks so this test exercises
         // the `IoBufs` send path (instead of the single-buffer fast path).
         let message = IoBufs::from(vec![
@@ -112,36 +115,40 @@ mod tests {
 
         // Spawn a server and read exactly the logical message size. The receive
         // side should observe the same byte stream regardless of send chunking.
-        let server = runtime.spawn(async move {
-            let (_, _sink, mut stream) = listener.accept().await.expect("Failed to accept");
+        let server = tokio::spawn(async move {
+            // Server receives the vectored payload as one logical byte stream.
+            let (_, sink, mut stream) = listener.accept().await.expect("Failed to accept");
             let received = stream
                 .recv(expected.len())
                 .await
                 .expect("Failed to receive");
             assert_eq!(received.coalesce(), expected.as_ref());
+            (sink, stream)
         });
 
         // Spawn client
-        let client = runtime.spawn(async move {
+        let client = tokio::spawn(async move {
+            // Client connects and sends the pre-built vectored message.
             // Connect to the server
-            let (mut sink, _stream) = network
+            let (mut sink, stream) = network
                 .dial(listener_addr)
                 .await
                 .expect("Failed to dial server");
 
             // Send the pre-built vectored message.
             sink.send(message).await.expect("Failed to send data");
+            (sink, stream)
         });
 
         // Wait for both tasks to complete
         let (server_result, client_result) = join!(server, client);
-        assert!(server_result.is_ok());
-        assert!(client_result.is_ok());
+        server_result.expect("Server task failed");
+        client_result.expect("Client task failed");
     }
 
     // Test handling multiple clients
     async fn test_network_multiple_clients<N: crate::Network>(network: N) {
-        let runtime = tokio::runtime::Handle::current();
+        const NUM_CLIENTS: usize = 3;
 
         // Start a server
         let mut listener = network
@@ -150,27 +157,42 @@ mod tests {
             .expect("Failed to bind");
         let listener_addr = listener.local_addr().expect("Failed to get local address");
 
+        // Keep all sockets alive until every participant finishes.
+        let barrier = Arc::new(Barrier::new(NUM_CLIENTS * 2));
+
         // Server task
-        let server = runtime.spawn(async move {
+        let server_barrier = barrier.clone();
+        let server = tokio::spawn(async move {
             // Handle multiple clients
-            for _ in 0..3 {
+            let mut set = JoinSet::new();
+            for _ in 0..NUM_CLIENTS {
                 let (_, mut sink, mut stream) = listener.accept().await.expect("Failed to accept");
-
-                let received = stream
-                    .recv(CLIENT_SEND_DATA.len())
-                    .await
-                    .expect("Failed to receive");
-                assert_eq!(received.coalesce(), CLIENT_SEND_DATA);
-
-                sink.send(IoBuf::from(SERVER_SEND_DATA))
-                    .await
-                    .expect("Failed to send");
+                let barrier = server_barrier.clone();
+                set.spawn(async move {
+                    let received = stream
+                        .recv(CLIENT_SEND_DATA.len())
+                        .await
+                        .expect("Failed to receive");
+                    assert_eq!(received.coalesce(), CLIENT_SEND_DATA);
+                    sink.send(IoBuf::from(SERVER_SEND_DATA))
+                        .await
+                        .expect("Failed to send");
+
+                    // Hold the connection open until every peer has finished.
+                    barrier.wait().await;
+                });
+            }
+            while let Some(result) = set.join_next().await {
+                result.expect("Server connection task failed");
             }
         });
 
         // Start multiple clients
-        let client = runtime.spawn(async move {
-            for _ in 0..3 {
+        let mut set = JoinSet::new();
+        for _ in 0..NUM_CLIENTS {
+            let network = network.clone();
+            let barrier = barrier.clone();
+            set.spawn(async move {
                 // Connect to the server
                 let (mut sink, mut stream) = network
                     .dial(listener_addr)
@@ -187,14 +209,20 @@ mod tests {
                     .recv(SERVER_SEND_DATA.len())
                     .await
                     .expect("Failed to receive data");
+
                 // Verify the received data
                 assert_eq!(received.coalesce(), SERVER_SEND_DATA);
-            }
-        });
 
-        // Wait for server and all clients
+                // Hold the connection open until every peer has finished.
+                barrier.wait().await;
+            });
+        }
+
+        // Wait for all servers and clients to complete.
+        while let Some(result) = set.join_next().await {
+            result.expect("Client task failed");
+        }
         server.await.expect("Server task failed");
-        client.await.expect("Client task failed");
     }
 
     // Test large data transfer
@@ -209,8 +237,9 @@ mod tests {
             .expect("Failed to bind");
         let listener_addr = listener.local_addr().expect("Failed to get local address");
 
-        let runtime = tokio::runtime::Handle::current();
-        let server = runtime.spawn(async move {
+        // Spawn server. Returning the socket halves keeps them alive until both
+        // join handles are awaited below.
+        let server = tokio::spawn(async move {
             let (_, mut sink, mut stream) = listener.accept().await.expect("Failed to accept");
 
             // Receive and echo large data in chunks
@@ -221,10 +250,12 @@ mod tests {
                     .expect("Failed to receive chunk");
                 sink.send(received).await.expect("Failed to send chunk");
             }
+            (sink, stream)
         });
 
-        // Client task
-        let client = runtime.spawn(async move {
+        // Client task. Returning the socket halves keeps them alive until both
+        // join handles are awaited below.
+        let client = tokio::spawn(async move {
             // Connect to the server
             let (mut sink, mut stream) = network
                 .dial(listener_addr)
@@ -245,11 +276,13 @@ mod tests {
                     .expect("Failed to receive chunk");
                 assert_eq!(received.coalesce(), &pattern[..]);
             }
+            (sink, stream)
         });
 
         // Wait for both tasks to complete
-        server.await.expect("Server task failed");
-        client.await.expect("Client task failed");
+        let (server_result, client_result) = join!(server, client);
+        server_result.expect("Server task failed");
+        client_result.expect("Client task failed");
     }
 
     // Tests dialing and binding errors
@@ -281,17 +314,17 @@ mod tests {
             .expect("Failed to bind");
         let listener_addr = listener.local_addr().expect("Failed to get local address");
 
-        let runtime = tokio::runtime::Handle::current();
-
         // Server sends data
-        let server = runtime.spawn(async move {
-            let (_, mut sink, _) = listener.accept().await.expect("Failed to accept");
+        let server = tokio::spawn(async move {
+            let (_, mut sink, stream) = listener.accept().await.expect("Failed to accept");
             sink.send(IoBuf::from(DATA)).await.expect("Failed to send");
+            (sink, stream)
         });
 
         // Client receives and tests peek
-        let client = runtime.spawn(async move {
-            let (_, mut stream) = network
+        let client = tokio::spawn(async move {
+            // Connect to the server
+            let (sink, mut stream) = network
                 .dial(listener_addr)
                 .await
                 .expect("Failed to dial server");
@@ -323,11 +356,13 @@ mod tests {
             // After consuming all data, peek should return empty
             let final_peek = stream.peek(100);
             assert!(final_peek.is_empty());
+            (sink, stream)
         });
 
+        // Wait for both tasks to complete
         let (server_result, client_result) = join!(server, client);
-        assert!(server_result.is_ok());
-        assert!(client_result.is_ok());
+        server_result.expect("Server task failed");
+        client_result.expect("Client task failed");
     }
 
     /// Network stress tests
@@ -352,36 +387,55 @@ mod tests {
             .unwrap();
         let addr = listener.local_addr().unwrap();
 
+        // Keep every connection alive until both the client and server halves finish.
+        let barrier = Arc::new(Barrier::new(NUM_CLIENTS * 2));
+
         // Spawn a server task that echoes messages from many clients.
+        let server_barrier = barrier.clone();
         let server = tokio::spawn(async move {
+            let mut set = JoinSet::new();
             for _ in 0..NUM_CLIENTS {
                 let (_, mut sink, mut stream) = listener.accept().await.unwrap();
-                tokio::spawn(async move {
+                let barrier = server_barrier.clone();
+                set.spawn(async move {
+                    // Echo every message back to the connected client.
                     for _ in 0..NUM_MESSAGES {
                         let received = stream.recv(MESSAGE_SIZE).await.unwrap();
                         sink.send(received).await.unwrap();
                     }
+
+                    // Hold the connection open until every peer has finished.
+                    barrier.wait().await;
                 });
             }
+            while let Some(result) = set.join_next().await {
+                result.unwrap();
+            }
         });
 
         // Spawn all clients.
-        let mut clients = Vec::new();
+        let mut set = JoinSet::new();
         for _ in 0..NUM_CLIENTS {
             let network = network.clone();
-            clients.push(tokio::spawn(async move {
+            let barrier = barrier.clone();
+            set.spawn(async move {
+                // Dial the server and repeatedly verify the echoed payload.
                 let (mut sink, mut stream) = network.dial(addr).await.unwrap();
                 let payload = vec![42u8; MESSAGE_SIZE];
                 for _ in 0..NUM_MESSAGES {
                     sink.send(payload.clone()).await.unwrap();
                     let received = stream.recv(MESSAGE_SIZE).await.unwrap();
                     assert_eq!(received.coalesce(), &payload[..]);
                 }
-            }));
+
+                // Hold the connection open until every peer has finished.
+                barrier.wait().await;
+            });
         }
 
-        for client in clients {
-            client.await.unwrap();
+        // Wait for all servers and clients to complete.
+        while let Some(result) = set.join_next().await {
+            result.unwrap();
         }
         server.await.unwrap();
     }

utils/src/sync/mod.rs

@@ -26,7 +26,7 @@ pub use parking_lot::{
     Condvar, Mutex, MutexGuard, Once, RwLock, RwLockReadGuard, RwLockWriteGuard,
 };
 pub use tokio::sync::{
-    Mutex as AsyncMutex, MutexGuard as AsyncMutexGuard, RwLock as AsyncRwLock,
+    Barrier, Mutex as AsyncMutex, MutexGuard as AsyncMutexGuard, RwLock as AsyncRwLock,
     RwLockReadGuard as AsyncRwLockReadGuard, RwLockWriteGuard as AsyncRwLockWriteGuard,
 };