fix: pool shutdown hang + per-connection leak from the peer-state callback (#4)

PR #3 added an `on_peer_connection_state_change` callback that captures a
clone of the per-session `SupervisorEvent` sender and a strong `Arc` back to
the `RTCPeerConnection`. Two regressions followed:

1. Shutdown hang. The pool's per-worker forwarder ran
   `while let Some(event) = worker_events_rx.recv().await`, which only ends
   when every sender clone drops. The clone captured by the WebRTC callback
   outlives the session, so the channel never closed — the pool future never
   completed after cancellation and `SharingHandle::stop` hung forever.
   (Surfaced by spark-sharing's `starts_emits_and_stops_the_unprivileged_pool`
   once its pin advanced onto a main containing #3; unbounded-rs's own tests
   missed it because they drive a mock `run_session` that never clones the
   sender into a real WebRTC callback.)

2. Per-connection memory leak. `state_connection = connection.clone()` inside
   the callback forms an `Arc<RTCPeerConnection>` <-> handler reference cycle,
   so the connection (and its captured sender) is never freed even after
   `close()` — one leaked connection per session over a long-running pool.

Fixes:

- supervisor.rs: the worker forwards inline via `tokio::select!` and stops when
  `supervise_peer_proxy` completes, draining the queue once (the final
  `Stopped` is enqueued before it returns, so nothing is lost). Shutdown no
  longer depends on a leaked sender clone dropping.
- peer_proxy.rs: the callback holds `Weak<RTCPeerConnection>` and upgrades it
  on the Connected transition, breaking the cycle so the connection and its
  captured sender are freed when the session ends.

Adds `pool_shuts_down_when_a_session_registers_a_state_callback`, which drives
the real peer proxy through the pool and asserts shutdown completes after
cancellation — it hangs (5s timeout) without either fix.

Co-authored-by: Adam Fisk <a@lantern.io>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Author: 3 people

src/peer_proxy.rs

@@ -188,7 +188,11 @@ pub async fn run_peer_proxy_until_cancelled(
     // so a session that fails straight to Failed/Closed emits nothing.
     let connect_emitted = Arc::new(AtomicBool::new(false));
 
-    let state_connection = connection.clone();
+    // Hold a `Weak` here, not a strong `Arc`: a strong clone captured by the
+    // state-change callback (which the connection itself owns) forms a reference
+    // cycle that leaks the whole `RTCPeerConnection` — and keeps the `events`
+    // sender clone below alive — for every session, hanging pool shutdown.
+    let state_connection = Arc::downgrade(&connection);
     let state_events = events.clone();
     let state_session_id = session_id.clone();
     let state_disconnect_emitted = disconnect_emitted.clone();
@@ -220,7 +224,12 @@ pub async fn run_peer_proxy_until_cancelled(
             match state {
                 RTCPeerConnectionState::Connected => {
                     connect_emitted.store(true, Ordering::SeqCst);
-                    let selected_pair = selected_candidate_pair(&connection).await;
+                    // The connection is normally still alive here; read the selected
+                    // pair when it is, otherwise fall back to `None` (remote unknown).
+                    let selected_pair = match connection.upgrade() {
+                        Some(connection) => selected_candidate_pair(&connection).await,
+                        None => None,
+                    };
                     let _ = events.send(peer_connected_event(session_id, selected_pair.as_ref()));
                 }
                 // Emit a disconnect only for a session that reached Connected, and

src/supervisor.rs

@@ -104,15 +104,35 @@ pub async fn supervise_peer_proxy_pool(
         let events = events.clone();
         async move {
             let (worker_events_tx, mut worker_events_rx) = mpsc::unbounded_channel();
-            let forwarder = tokio::spawn(async move {
-                while let Some(event) = worker_events_rx.recv().await {
-                    if let Some(events) = &events {
-                        let _ = events.send(PoolEvent { slot, event });
+            let mut supervise = std::pin::pin!(supervise_peer_proxy(
+                config,
+                cancellation,
+                Some(worker_events_tx)
+            ));
+            // Forward inline and terminate when the supervised session finishes —
+            // not when the channel closes. A session's WebRTC state-change callback
+            // can outlive the session while still holding a clone of the event
+            // sender; waiting for that clone to drop would hang shutdown. The final
+            // `Stopped` event is enqueued before `supervise_peer_proxy` returns, so
+            // draining once it completes forwards every event without loss.
+            let summary = loop {
+                tokio::select! {
+                    biased;
+                    summary = &mut supervise => {
+                        while let Ok(event) = worker_events_rx.try_recv() {
+                            if let Some(events) = &events {
+                                let _ = events.send(PoolEvent { slot, event });
+                            }
+                        }
+                        break summary;
+                    }
+                    Some(event) = worker_events_rx.recv() => {
+                        if let Some(events) = &events {
+                            let _ = events.send(PoolEvent { slot, event });
+                        }
                     }
                 }
-            });
-            let summary = supervise_peer_proxy(config, cancellation, Some(worker_events_tx)).await;
-            let _ = forwarder.await;
+            };
             (slot, summary)
         }
     });
@@ -411,6 +431,55 @@ mod tests {
         assert!(second_retry.unwrap() <= Duration::from_millis(12));
     }
 
+    #[tokio::test]
+    async fn pool_shuts_down_when_a_session_registers_a_state_callback() {
+        // Regression: `run_peer_proxy_until_cancelled` registers an
+        // `on_peer_connection_state_change` callback that captures a clone of the
+        // per-session `SupervisorEvent` sender (and, before the fix, a strong
+        // `Arc` back to the connection). WebRTC keeps that callback alive past the
+        // session, so the sender clone outlived the session and kept the worker's
+        // forwarder channel open — the pool future never completed after
+        // cancellation and `SharingHandle::stop` hung. The pool must shut down
+        // regardless of a leaked event-sender clone.
+        let cancellation = CancellationToken::new();
+        let (events_tx, mut events_rx) = mpsc::unbounded_channel::<PoolEvent>();
+        let mut config = test_config();
+        // Park in backoff after the first failure so cancellation, not the retry
+        // cadence, drives shutdown.
+        config.initial_backoff = Duration::from_secs(30);
+        config.max_backoff = Duration::from_secs(30);
+
+        let pool = tokio::spawn(supervise_peer_proxy_pool(
+            config,
+            1,
+            cancellation.clone(),
+            Some(events_tx),
+        ));
+
+        // Wait until the real peer proxy has attempted and failed — this proves it
+        // built the RTCPeerConnection and registered the state-change callback.
+        let mut saw_failure = false;
+        while let Some(PoolEvent { event, .. }) =
+            tokio::time::timeout(Duration::from_secs(5), events_rx.recv())
+                .await
+                .expect("expected peer-proxy events before timeout")
+        {
+            if matches!(event, SupervisorEvent::AttemptFailed { .. }) {
+                saw_failure = true;
+                break;
+            }
+        }
+        assert!(saw_failure, "peer proxy should have attempted and failed");
+
+        cancellation.cancel();
+        let summary = tokio::time::timeout(Duration::from_secs(5), pool)
+            .await
+            .expect("pool did not shut down within 5s of cancellation")
+            .expect("pool task panicked");
+        assert_eq!(summary.workers.len(), 1);
+        assert!(summary.attempts() >= 1);
+    }
+
     #[test]
     fn pool_summary_aggregates_worker_counts() {
         let summary = SupervisorPoolSummary {