diff --git a/src/ws-ckpt/src/crates/daemon/src/listener.rs b/src/ws-ckpt/src/crates/daemon/src/listener.rs index 4d7ea239c9..8f4c984fff 100644 --- a/src/ws-ckpt/src/crates/daemon/src/listener.rs +++ b/src/ws-ckpt/src/crates/daemon/src/listener.rs @@ -62,6 +62,14 @@ pub async fn run_listener( break; } } + + // Reap finished connections. JoinSet holds each task's handle and output + // slot until it is joined, so without this the set grows by one entry per + // CLI invocation and is only released at shutdown. try_join_next never + // awaits, so reaping cannot delay the next accept; a join_next branch in + // the select! above would instead need a second mutable borrow of + // join_set while the accept arm still spawns into it. + while join_set.try_join_next().is_some() {} } // 7. Wait for in-flight tasks to complete (with timeout) @@ -136,3 +144,29 @@ async fn handle_connection( Ok(()) } + +#[cfg(test)] +mod tests { + use tokio::task::JoinSet; + + /// Pins the semantics the accept loop relies on: completed tasks stay in the + /// set until reaped, and `try_join_next` drains them without awaiting. If the + /// reap call in the accept loop is dropped, the daemon leaks one entry per + /// connection, which is what caused ~8 KB of growth per CLI invocation. + #[tokio::test] + async fn try_join_next_reaps_completed_tasks() { + let mut set = JoinSet::new(); + for _ in 0..16 { + set.spawn(async {}); + } + while !set.is_empty() { + let before = set.len(); + if set.try_join_next().is_none() { + tokio::task::yield_now().await; + continue; + } + assert_eq!(set.len(), before - 1, "reaping must release one entry"); + } + assert!(set.try_join_next().is_none(), "empty set must not block"); + } +}