WIP: Harden Connection Management to Prevent Zombie Connections and Indefinite Blocking

crates/tx5-connection/src/framed.rs

@@ -61,8 +61,8 @@ impl FramedConn {
         let (a, b, c, d) = crate::proto::PROTO_VER_2.encode()?;
         conn.send(vec![a, b, c, d]).await?;
 
-        let (cmd_send, mut cmd_recv) = tokio::sync::mpsc::channel(32);
-        let (msg_send, msg_recv) = tokio::sync::mpsc::channel(32);
+        let (cmd_send, mut cmd_recv) = tokio::sync::mpsc::channel(256);
+        let (msg_send, msg_recv) = tokio::sync::mpsc::channel(256);
 
         // set up the receive to just feed straight into the cmd task
         let cmd_send2 = cmd_send.clone();

crates/tx5-connection/src/hub.rs

@@ -139,7 +139,7 @@ impl Hub {
 
         let mut task_list = Vec::new();
 
-        let (hub_cmd_send, mut cmd_recv) = tokio::sync::mpsc::channel(32);
+        let (hub_cmd_send, mut cmd_recv) = tokio::sync::mpsc::channel(256);
 
         // forward received messages to the cmd task
         let hub_cmd_send2 = hub_cmd_send.clone();

crates/tx5/src/ep.rs

@@ -139,6 +139,12 @@ impl EpInner {
 
     /// Get an existing peer connection or create a new outgoing one.
     pub fn connect_peer(&mut self, peer_url: PeerUrl) -> Arc<Peer> {
+        // Check for stale failed peers before attempting reuse
+        if let Some(existing_peer) = self.peer_map.get(&peer_url) {
+            if existing_peer.is_failed() {
+                self.peer_map.remove(&peer_url);
+            }
+        }
         self.peer_map
             .entry(peer_url.clone())
             .or_insert_with(|| {
@@ -199,7 +205,9 @@ impl Endpoint {
             config.incoming_message_bytes_max as usize,
         ));
 
-        let (evt_send, evt_recv) = tokio::sync::mpsc::channel(32);
+        let (evt_send, evt_recv) = tokio::sync::mpsc::channel(
+            config.internal_event_channel_size as usize,
+        );
 
         (
             Self {
@@ -250,9 +258,11 @@ impl Endpoint {
         }
         tokio::time::timeout(self.config.timeout, async {
             let peer = self.inner.lock().unwrap().connect_peer(peer_url.clone());
-            if let Some(conn) = peer.wait_for_ready().await {
+            if let Some(conn) = peer.wait_for_ready(self.config.timeout).await {
                 conn.send(data).await
             } else {
+                // Remove the failed peer from peer_map to prevent reuse of zombie connections
+                self.inner.lock().unwrap().drop_peer_url(&peer_url);
                 // Check to see if a new incoming connection was established
                 // in the mean time. If so, we can send over that.
                 let maybe_peer = self.inner.lock().unwrap().peer_if_ready(
@@ -286,7 +296,7 @@ impl Endpoint {
                 let data = data.to_vec();
                 tokio::task::spawn(async move {
                     let _ = tokio::time::timeout(timeout, async {
-                        if let Some(conn) = peer.wait_for_ready().await {
+                        if let Some(conn) = peer.wait_for_ready(timeout).await {
                             let _ = conn.send(data).await;
                         }
                     })

crates/tx5/src/peer.rs

@@ -117,11 +117,19 @@ impl Peer {
         self.ready.query_ready(|c| c.clone())
     }
 
+    /// Check if this peer connection has failed
+    pub(crate) fn is_failed(&self) -> bool {
+        self.ready.is_failed()
+    }
+
     /// This future resolves when the connection is ready to use or has failed to connect.
     ///
     /// If the connection is not usable, it will return `None`.
-    pub async fn wait_for_ready(&self) -> Option<DynBackCon> {
-        self.ready.wait_for_ready().await
+    pub async fn wait_for_ready(
+        &self,
+        timeout: std::time::Duration,
+    ) -> Option<DynBackCon> {
+        self.ready.wait_for_ready(timeout).await
     }
 }
 
@@ -241,14 +249,22 @@ async fn task(
 ) {
     // establish our cleanup drop guard
     let drop_guard = DropPeer {
-        ep,
+        ep: ep.clone(),
         peer_url: peer_url.clone(),
         evt_send: evt_send.clone(),
         ready: ready.clone(),
     };
 
     let mut wc = match conn {
-        None => return,
+        None => {
+            // Connection failed during negotiation - clean up immediately
+            // to prevent zombie connection reuse
+            if let Some(ep_inner) = ep.upgrade() {
+                ep_inner.lock().unwrap().drop_peer_url(&peer_url);
+            }
+            ready.set_failed();
+            return;
+        }
         Some(wc) => wc,
     };
 

crates/tx5/src/peer/maybe_ready.rs

@@ -73,7 +73,16 @@ impl MaybeReady {
         *lock = MaybeReadyState::Failed;
     }
 
-    pub(super) async fn wait_for_ready(&self) -> Option<DynBackCon> {
+    /// Check if the connection is in a failed state
+    pub(super) fn is_failed(&self) -> bool {
+        let lock = self.0.lock().expect("poisoned");
+        matches!(*lock, MaybeReadyState::Failed)
+    }
+
+    pub(super) async fn wait_for_ready(
+        &self,
+        timeout: std::time::Duration,
+    ) -> Option<DynBackCon> {
         let wait = {
             let lock = self.0.lock().expect("poisoned");
             match &*lock {
@@ -83,15 +92,29 @@ impl MaybeReady {
             }
         };
 
-        let _ = wait.acquire().await;
-
-        let lock = self.0.lock().expect("poisoned");
-        match &*lock {
-            MaybeReadyState::Ready(back) => Some(back.clone()),
-            MaybeReadyState::Failed => None,
-            MaybeReadyState::Wait(_) => {
-                // Unexpected state, this struct is supposed to prevent this being possible.
-                tracing::warn!("Waited for ready but still in a wait state");
+        // Add timeout wrapper to prevent indefinite blocking
+        let result = tokio::time::timeout(timeout, async {
+            let _ = wait.acquire().await;
+        })
+        .await;
+
+        match result {
+            Ok(_) => {
+                // Semaphore acquired, check final state
+                let lock = self.0.lock().expect("poisoned");
+                match &*lock {
+                    MaybeReadyState::Ready(back) => Some(back.clone()),
+                    MaybeReadyState::Failed => None,
+                    MaybeReadyState::Wait(_) => {
+                        // Unexpected state, this struct is supposed to prevent this being possible.
+                        tracing::warn!("Waited for ready but still in a wait state");
+                        None
+                    }
+                }
+            }
+            Err(_) => {
+                // Timeout - connection failed to become ready in time
+                tracing::debug!("wait_for_ready timed out after {:?}", timeout);
                 None
             }
         }
@@ -136,8 +159,13 @@ mod tests {
         tokio::spawn({
             let maybe_ready = maybe_ready.clone();
             async move {
-                tx.send(maybe_ready.wait_for_ready().await.is_some())
-                    .unwrap();
+                tx.send(
+                    maybe_ready
+                        .wait_for_ready(std::time::Duration::from_secs(60))
+                        .await
+                        .is_some(),
+                )
+                .unwrap();
             }
         });
 
@@ -178,8 +206,13 @@ mod tests {
         tokio::spawn({
             let maybe_ready = maybe_ready.clone();
             async move {
-                tx.send(maybe_ready.wait_for_ready().await.is_none())
-                    .unwrap();
+                tx.send(
+                    maybe_ready
+                        .wait_for_ready(std::time::Duration::from_secs(60))
+                        .await
+                        .is_none(),
+                )
+                .unwrap();
             }
         });