diff --git a/crates/tx5-connection/src/framed.rs b/crates/tx5-connection/src/framed.rs index db6175c..ebccaf8 100644 --- a/crates/tx5-connection/src/framed.rs +++ b/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(); diff --git a/crates/tx5-connection/src/hub.rs b/crates/tx5-connection/src/hub.rs index 38ec486..ede2645 100644 --- a/crates/tx5-connection/src/hub.rs +++ b/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(); diff --git a/crates/tx5/src/ep.rs b/crates/tx5/src/ep.rs index 389c69e..f34fcc5 100644 --- a/crates/tx5/src/ep.rs +++ b/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 { + // 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; } }) diff --git a/crates/tx5/src/peer.rs b/crates/tx5/src/peer.rs index c5c5b2f..68c4733 100644 --- a/crates/tx5/src/peer.rs +++ b/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 { - self.ready.wait_for_ready().await + pub async fn wait_for_ready( + &self, + timeout: std::time::Duration, + ) -> Option { + 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, }; diff --git a/crates/tx5/src/peer/maybe_ready.rs b/crates/tx5/src/peer/maybe_ready.rs index b4dd2e0..5eed332 100644 --- a/crates/tx5/src/peer/maybe_ready.rs +++ b/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 { + /// 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 { 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(); } });