From 581942567b5427260308f89d623766233179720f Mon Sep 17 00:00:00 2001
From: sycatle <sycatle@users.noreply.github.com>
Date: Mon, 20 Jul 2026 18:46:34 +0200
Subject: [PATCH] fix(cdp): drain live connections before persisting cookies on
 shutdown

Connection threads are detached and never joined, so `save_cookies()` ran on
the main thread while they were still serving: a connection writing a
Set-Cookie after that point lost it, and the process then exited and killed the
thread mid-flight. Before the per-connection move this could not happen -- the
single processor saved on its own way out, ordered against all connection work
on one LocalSet. Draining on the live-connection counter restores that ordering.

Scope, honestly: this bounds the loss window rather than closing it. The drain
does not finish early today because a connection's frame reader is parked in a
socket read and never observes `notify_waiters`, so a client that holds a
connection open costs the full 3s deadline and then gets the warning. Making
readers wake on shutdown is a separate change and not #430's business.

Measured: 2ms to shut down with no connections (unchanged), 3005ms with one
held open, against the 10s `docker stop` allows before SIGKILL. Cookies are
persisted in both cases.

Refs #430, #333.
---
 crates/obscura-cdp/src/server.rs | 30 +++++++++++++++++++++++++++++-
 1 file changed, 29 insertions(+), 1 deletion(-)

diff --git a/crates/obscura-cdp/src/server.rs b/crates/obscura-cdp/src/server.rs
index e52b356..48102dc 100644
--- a/crates/obscura-cdp/src/server.rs
+++ b/crates/obscura-cdp/src/server.rs
@@ -40,6 +40,10 @@ const MAX_PENDING_WS_HANDOFFS: usize = 128;
 // which is there at zero connections. Override with `--max-connections`.
 pub const DEFAULT_MAX_CONNECTIONS: usize = 128;
 
+// How long shutdown waits for connection threads to finish before persisting
+// the cookie jar. Well under the 10s `docker stop` gives us before SIGKILL.
+const SHUTDOWN_DRAIN_MS: u64 = 3_000;
+
 // Sent to a client that arrives while the server is at `max_connections`, in
 // place of dropping the socket unexplained. The client sees a refusal it can
 // retry rather than a bare connection reset.
@@ -343,7 +347,31 @@ pub async fn start_with_serve_options_and_limit(
         );
     }
 
-    // Server is shutting down: persist the shared cookie jar once.
+    // Server is shutting down. Connection threads are detached, so saving the
+    // jar right here would race them: a connection still writing a Set-Cookie
+    // loses it, and the process then exits and kills the thread mid-flight.
+    // Before the per-connection move, the single processor saved on its own way
+    // out, ordered against all connection work on one LocalSet -- draining here
+    // is what restores that ordering. `notify_waiters` above has already woken
+    // every processor, so this is bounded in practice; the deadline only covers
+    // a connection wedged in V8, where its own command watchdog is the backstop.
+    let drain_deadline =
+        tokio::time::Instant::now() + tokio::time::Duration::from_millis(SHUTDOWN_DRAIN_MS);
+    loop {
+        let live = live_connections.load(Ordering::Acquire);
+        if live == 0 {
+            break;
+        }
+        if tokio::time::Instant::now() >= drain_deadline {
+            warn!(
+                "shutting down with {} connection(s) still live after {}ms; \
+                 cookies they write from here are lost",
+                live, SHUTDOWN_DRAIN_MS
+            );
+            break;
+        }
+        tokio::time::sleep(tokio::time::Duration::from_millis(20)).await;
+    }
     shared_ctx.save_cookies();
     Ok(())
 }