From 838dfe3dcc50a64dab95b770361809b0c8408392 Mon Sep 17 00:00:00 2001 From: Lawrence Chen <54008264+lawrencecchen@users.noreply.github.com> Date: Thu, 27 Aug 2026 08:30:08 -0700 Subject: [PATCH 1/4] cmux-tui: abort detached local forward task on drop --- cmux-tui/crates/cmux-remote/src/bridge.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/cmux-tui/crates/cmux-remote/src/bridge.rs b/cmux-tui/crates/cmux-remote/src/bridge.rs index 46ac0e66db8d..7897df090182 100644 --- a/cmux-tui/crates/cmux-remote/src/bridge.rs +++ b/cmux-tui/crates/cmux-remote/src/bridge.rs @@ -115,6 +115,12 @@ impl Drop for LocalPortForward { if let Some(shutdown) = self.shutdown.take() { let _ = shutdown.send(()); } + // A dropped JoinHandle detaches the accept loop. Request cancellation + // here so the listener and in-flight tunnel tasks cannot run forever + // after the forward's owner is gone. + if let Some(task) = self.task.take() { + task.abort(); + } } } From a72c750c591cf6c6c6e741c789a74e5d32cb6a9b Mon Sep 17 00:00:00 2001 From: Lawrence Chen <54008264+lawrencecchen@users.noreply.github.com> Date: Thu, 27 Aug 2026 09:06:16 -0700 Subject: [PATCH 2/4] cmux-tui: join local forward tunnel tasks --- cmux-tui/crates/cmux-remote/src/bridge.rs | 119 ++++++++++++++++++---- 1 file changed, 99 insertions(+), 20 deletions(-) diff --git a/cmux-tui/crates/cmux-remote/src/bridge.rs b/cmux-tui/crates/cmux-remote/src/bridge.rs index 7897df090182..5249edb67a5b 100644 --- a/cmux-tui/crates/cmux-remote/src/bridge.rs +++ b/cmux-tui/crates/cmux-remote/src/bridge.rs @@ -6,7 +6,7 @@ use std::sync::Arc; use bytes::Bytes; use cmux_remote_protocol::{MUX_INPUT_V1_FEATURE, RouteId, Service, ServiceControl}; use tokio::io::{AsyncRead, AsyncWrite}; -use tokio::sync::{Semaphore, oneshot}; +use tokio::sync::{Mutex, Semaphore, oneshot, watch}; use crate::mux_codec::{MAX_MUX_DOWNLOAD_LINE_BYTES, MAX_MUX_LINE_BYTES, mux_line_payload_len}; use crate::service::{ServiceError, ServiceMultiplexer, ServiceStream}; @@ -14,10 +14,39 @@ use crate::services::ServicesError; pub const DEFAULT_MAX_FORWARD_CONNECTIONS: usize = 128; +struct ForwardConnections { + tasks: Mutex>, +} + +impl ForwardConnections { + fn new() -> Arc { + Arc::new(Self { tasks: Mutex::new(tokio::task::JoinSet::new()) }) + } + + async fn reap_finished(&self) { + let mut tasks = self.tasks.lock().await; + while tasks.try_join_next().is_some() {} + } + + async fn shutdown(&self) { + let mut tasks = self.tasks.lock().await; + tasks.abort_all(); + while tasks.join_next().await.is_some() {} + } + + fn abort_all(&self) { + if let Ok(tasks) = self.tasks.try_lock() { + tasks.abort_all(); + } + } +} + pub struct LocalPortForward { local_addr: SocketAddr, shutdown: Option>, task: Option>, + connections: Arc, + cancellation: watch::Sender, } impl LocalPortForward { @@ -45,10 +74,13 @@ impl LocalPortForward { let local_addr = listener.local_addr()?; let permits = Arc::new(Semaphore::new(maximum_connections)); let (shutdown_tx, mut shutdown_rx) = oneshot::channel(); + let connections = ForwardConnections::new(); + let task_connections = connections.clone(); + let (cancellation, _) = watch::channel(false); + let task_cancellation = cancellation.clone(); let task = tokio::spawn(async move { - let mut connections = tokio::task::JoinSet::new(); loop { - while connections.try_join_next().is_some() {} + task_connections.reap_finished().await; let permit = tokio::select! { _ = &mut shutdown_rx => break, permit = permits.clone().acquire_owned() => { @@ -65,24 +97,37 @@ impl LocalPortForward { continue; } let multiplexer = multiplexer.clone(); - connections.spawn(async move { - let _permit = permit; - let mut metadata = BTreeMap::new(); - metadata.insert("route".into(), route.0.to_string()); - let Ok(stream) = multiplexer.open(Service::TcpTunnel, metadata).await else { - return; + let mut cancellation = task_cancellation.subscribe(); + task_connections.tasks.lock().await.spawn(async move { + let handler = async move { + let _permit = permit; + let mut metadata = BTreeMap::new(); + metadata.insert("route".into(), route.0.to_string()); + let Ok(stream) = multiplexer.open(Service::TcpTunnel, metadata).await + else { + return; + }; + if await_opened(&stream).await.is_err() { + return; + } + let (reader, writer) = socket.into_split(); + let _ = pump_client(Arc::new(stream), reader, writer).await; }; - if await_opened(&stream).await.is_err() { - return; + tokio::select! { + _ = cancellation.changed() => {} + _ = handler => {} } - let (reader, writer) = socket.into_split(); - let _ = pump_client(Arc::new(stream), reader, writer).await; }); } - connections.abort_all(); - while connections.join_next().await.is_some() {} + task_connections.shutdown().await; }); - Ok(Self { local_addr, shutdown: Some(shutdown_tx), task: Some(task) }) + Ok(Self { + local_addr, + shutdown: Some(shutdown_tx), + task: Some(task), + connections, + cancellation, + }) } pub fn local_addr(&self) -> SocketAddr { @@ -97,6 +142,7 @@ impl LocalPortForward { } pub async fn shutdown(mut self) { + self.cancellation.send_replace(true); if let Some(shutdown) = self.shutdown.take() { let _ = shutdown.send(()); } @@ -112,12 +158,14 @@ fn configure_forward_socket(socket: &tokio::net::TcpStream) -> std::io::Result<( impl Drop for LocalPortForward { fn drop(&mut self) { + self.cancellation.send_replace(true); + self.connections.abort_all(); if let Some(shutdown) = self.shutdown.take() { let _ = shutdown.send(()); } - // A dropped JoinHandle detaches the accept loop. Request cancellation - // here so the listener and in-flight tunnel tasks cannot run forever - // after the forward's owner is gone. + // A dropped JoinHandle detaches the accept loop. Abort it after the + // shared cancellation signal and child-task abort request so Drop does + // not leave either owner or tunnel handlers running indefinitely. if let Some(task) = self.task.take() { task.abort(); } @@ -462,7 +510,7 @@ impl From for BridgeError { #[cfg(test)] mod tests { - use std::sync::atomic::{AtomicU64, Ordering}; + use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use async_trait::async_trait; use cmux_remote_protocol::{FrameFlags, Lane}; @@ -956,6 +1004,37 @@ mod tests { forward.shutdown().await; } + #[tokio::test] + async fn emergency_forward_cleanup_aborts_active_tunnel_handlers() { + struct DropFlag(Arc); + + impl Drop for DropFlag { + fn drop(&mut self) { + self.0.store(true, Ordering::Release); + } + } + + let connections = ForwardConnections::new(); + let dropped = Arc::new(AtomicBool::new(false)); + connections.tasks.lock().await.spawn({ + let dropped = dropped.clone(); + async move { + let _flag = DropFlag(dropped); + std::future::pending::<()>().await; + } + }); + + connections.abort_all(); + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while !dropped.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + }) + .await + .expect("emergency cleanup did not abort the active tunnel handler"); + connections.shutdown().await; + } + #[tokio::test] async fn forward_rejects_zero_connection_limit() { let (client_endpoint, _) = endpoint_pair(); From 859f94144e9122c1c5eacfea87affd6f915afca4 Mon Sep 17 00:00:00 2001 From: Lawrence Chen <54008264+lawrencecchen@users.noreply.github.com> Date: Thu, 27 Aug 2026 09:23:28 -0700 Subject: [PATCH 3/4] cmux-tui: fix mutable forward task guard --- cmux-tui/crates/cmux-remote/src/bridge.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cmux-tui/crates/cmux-remote/src/bridge.rs b/cmux-tui/crates/cmux-remote/src/bridge.rs index 5249edb67a5b..400764751b68 100644 --- a/cmux-tui/crates/cmux-remote/src/bridge.rs +++ b/cmux-tui/crates/cmux-remote/src/bridge.rs @@ -35,7 +35,7 @@ impl ForwardConnections { } fn abort_all(&self) { - if let Ok(tasks) = self.tasks.try_lock() { + if let Ok(mut tasks) = self.tasks.try_lock() { tasks.abort_all(); } } From 1e0c3eefaf43e733c967131199361d587f56a34b Mon Sep 17 00:00:00 2001 From: Lawrence Chen <54008264+lawrencecchen@users.noreply.github.com> Date: Thu, 27 Aug 2026 09:47:43 -0700 Subject: [PATCH 4/4] test(remote): synchronize tunnel cleanup assertion --- cmux-tui/crates/cmux-remote/src/bridge.rs | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/cmux-tui/crates/cmux-remote/src/bridge.rs b/cmux-tui/crates/cmux-remote/src/bridge.rs index 400764751b68..3d904fef730d 100644 --- a/cmux-tui/crates/cmux-remote/src/bridge.rs +++ b/cmux-tui/crates/cmux-remote/src/bridge.rs @@ -515,7 +515,7 @@ mod tests { use async_trait::async_trait; use cmux_remote_protocol::{FrameFlags, Lane}; use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader, split}; - use tokio::sync::{Mutex, mpsc, watch}; + use tokio::sync::{Mutex, mpsc, oneshot, watch}; use super::*; use crate::service::{EndpointRole, SessionEndpoint}; @@ -1016,23 +1016,20 @@ mod tests { let connections = ForwardConnections::new(); let dropped = Arc::new(AtomicBool::new(false)); + let (started_tx, started_rx) = oneshot::channel(); connections.tasks.lock().await.spawn({ let dropped = dropped.clone(); async move { let _flag = DropFlag(dropped); + let _ = started_tx.send(()); std::future::pending::<()>().await; } }); + started_rx.await.unwrap(); connections.abort_all(); - tokio::time::timeout(std::time::Duration::from_secs(1), async { - while !dropped.load(Ordering::Acquire) { - tokio::task::yield_now().await; - } - }) - .await - .expect("emergency cleanup did not abort the active tunnel handler"); connections.shutdown().await; + assert!(dropped.load(Ordering::Acquire)); } #[tokio::test]