fix(server): graceful-shutdown connection drain + guard-first parked teardown (#438 F1/F2)

CHANGELOG.md

@@ -25,6 +25,36 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   `--experimental-per-shard-rewrite` is deprecated (warns, no-op).
 
 ### Fixed
+- **Graceful shutdown now drains connections (#438 F1, conn#8).** On
+  SIGTERM/SIGINT each shard previously tore its runtime down the moment its
+  event loop observed the cancellation, dropping every pending connection
+  task mid-poll: in-flight replies were truncated and the blocking/subscriber
+  shutdown arms (`-ERR server shutting down`) never ran — measured 19/50
+  BLPOP-blocked clients losing their shutdown reply on Linux io_uring
+  (macOS passed only by scheduler luck). Shards now run a bounded drain
+  before persistence teardown: parked stage-1/2 reads are woken through
+  their cancellers (re-fired per tick to close the re-park race), token arms
+  cover the tracking/subscriber/blocking parks, and the loop waits for the
+  shard's live connection tasks to exit through the normal flush+FIN
+  epilogue, up to a 5 s ceiling so a wedged peer cannot hold up shutdown.
+  Writing the drain's red test surfaced a second, tokio-only leak in the
+  same class: connections accepted by the CENTRAL listener were io-bound to
+  the MAIN runtime's driver (tokio io resources bind at creation), so their
+  reads/writes died with main's runtime no matter what the shard drained —
+  with SO_REUSEPORT splitting accepts roughly evenly, about half of all
+  tokio connections failed their final writes with "A Tokio 1.x context was
+  found, but it is being shutdown". Forwarded streams are now re-registered
+  with the owning shard's runtime driver at spawn (the monoio path already
+  did this by forwarding std streams).
+- **Parked-connection teardown can no longer close a reused fd out from
+  under a kill scan (#438 F2, conn#9 / sec L1).** `kill_clients`'
+  fd-liveness invariant (registry deregister strictly before fd close) held
+  in handler tasks by local-before-parameter drop order, but a task-parked
+  connection's watcher co-owned guard and stream as future upvars, whose
+  drop order is merely capture order — and the F1 drain makes dropping that
+  future a routine path. Both are now wrapped in a `ParkedSession` whose
+  hand-written `Drop` deregisters before closing, with the invariant
+  restated at the kill site.
 - **Remotely-triggerable shard-thread crash on pipelined batch tails
   (#438 follow-on).** On any `--shards >= 2` deployment, one pipelined write
   shaped `[<remote-key cmd>…, SUBSCRIBE|BLPOP|…]` aborted the whole server:

src/client_registry.rs

@@ -410,14 +410,23 @@ pub fn client_info(id: u64) -> Option<String> {
 /// self-kill, then closes).
 pub fn kill_clients(filter: &KillFilter, self_id: Option<u64>) -> u64 {
     // Lock ordering / fd-liveness invariant, per stripe (c10k W5): the raw-fd
-    // `shutdown` is race-free because a connection's `RegistryGuard` (a local)
-    // drops — calling `deregister`, which needs its OWN stripe's WRITE lock —
-    // strictly before its `stream` (a parameter) drops and closes the fd. So
-    // while we hold a stripe's READ lock and an entry is present in it,
-    // `deregister` for that entry is blocked, its `stream` has not dropped,
-    // and `kill_fd` is still an open socket. The invariant only ever involves
-    // one entry and its own stripe, so striping preserves it; entries in
-    // other stripes are simply not visited while their lock is free.
+    // `shutdown` is race-free because a connection's `RegistryGuard` drops —
+    // calling `deregister`, which needs its OWN stripe's WRITE lock —
+    // strictly before its `stream` drops and closes the fd. So while we hold
+    // a stripe's READ lock and an entry is present in it, `deregister` for
+    // that entry is blocked, its `stream` has not dropped, and `kill_fd` is
+    // still an open socket. The invariant only ever involves one entry and
+    // its own stripe, so striping preserves it; entries in other stripes are
+    // simply not visited while their lock is free.
+    //
+    // WHO upholds guard-before-stream (F2, #438): in a handler task it falls
+    // out of drop order — the guard is a local, the stream a parameter, and
+    // locals drop first. A task-parked connection has NO handler task; its
+    // watcher co-owns both as future upvars (drop order = capture order, no
+    // language guarantee that helps), so they are wrapped in
+    // `conn_accept::ParkedSession`, whose hand-written Drop deregisters
+    // before closing. Any future owner of a {guard, stream} pair must
+    // preserve this order or wrap in ParkedSession.
     let mut count = 0u64;
     let kill_entry = |entry: &ClientEntry, count: &mut u64| {
         entry.live.kill_flag.store(true, Ordering::Relaxed);

src/server/conn/handler_monoio/idle_park.rs

@@ -164,6 +164,32 @@ pub(crate) fn sweep(now_ms: u64) -> usize {
     })
 }
 
+/// F1 (#438): cancel EVERY currently-parked stage-1/2 read, regardless of
+/// age. Called from the shard event loop's shutdown drain — the parked
+/// reads are plain awaits (not selects), so the shutdown token alone cannot
+/// wake them; the canceller is the only same-thread wake the design has.
+/// The woken handler sees the sweep-cancel error, checks the shutdown token
+/// and exits through the normal flush+FIN epilogue instead of re-parking.
+///
+/// Re-fired every drain tick: a connection that was mid-batch at the first
+/// call parks again only if it raced the token check, and the next tick
+/// catches it. Idempotent on unparked slots.
+pub(crate) fn cancel_all_parked() -> usize {
+    REGISTRY.with(|r| {
+        let mut cancelled = 0usize;
+        for slot in r.borrow().values() {
+            if slot.parked_since_ms.get() != 0 {
+                let old = slot.canceller.replace(Canceller::new());
+                let fresh = old.cancel();
+                slot.canceller.replace(fresh);
+                slot.parked_since_ms.set(0);
+                cancelled += 1;
+            }
+        }
+        cancelled
+    })
+}
+
 /// Release the parked working set. The rent buffer is dropped outright (the
 /// pre-park sizing reallocates the probe size next iteration); the scratch
 /// buffers are only released when empty — a non-empty `read_buf` holds a
@@ -355,6 +381,21 @@ mod idle_park_tests {
         assert_eq!(sweep(300_000 + IDLE_DOWNSHIFT_MS), 1);
     }
 
+    #[test]
+    fn cancel_all_parked_ignores_age_and_unparked() {
+        let reg_young = register(90_010);
+        let reg_old = register(90_011);
+        let reg_idle = register(90_012);
+        reg_young.slot.mark_parked(10_000); // just parked — sweep would skip
+        reg_old.slot.mark_parked_stage2(1); // ancient stage-2 park
+        // reg_idle: not parked at all.
+        assert_eq!(cancel_all_parked(), 2, "both parked slots, any age/stage");
+        assert_eq!(reg_young.slot.parked_since_ms.get(), 0);
+        assert_eq!(reg_old.slot.parked_since_ms.get(), 0);
+        assert_eq!(cancel_all_parked(), 0, "idempotent once unparked");
+        drop(reg_idle);
+    }
+
     #[test]
     fn registration_drop_deregisters() {
         {

src/server/conn/handler_monoio/mod.rs

@@ -847,6 +847,11 @@ pub(crate) async fn handle_connection_sharded_monoio<
             let track_buf = std::mem::take(&mut tmp_buf);
             let mut push_frame: Option<Frame> = None;
             monoio::select! {
+                // F1 (#438): tracking conns park in this select (not the
+                // cancel-registered reads below), so the shutdown drain
+                // reaches them via the token. Losing the read future drops
+                // `track_buf` — acceptable, the connection is exiting.
+                _ = shutdown.cancelled() => { break; }
                 read_result = stream.read(track_buf) => {
                     let (result, returned_buf) = read_result;
                     tmp_buf = returned_buf;
@@ -937,6 +942,11 @@ pub(crate) async fn handle_connection_sharded_monoio<
                     // ECONNRESET, …): terminate. Parking instead would spin
                     // park→wake→park forever — the dead fd stays readable.
                     Err(ref e) if !idle_park::is_sweep_cancel(e) => break,
+                    // F1 (#438): cancelled by the shutdown drain, not the
+                    // stage-2 sweep — exit through the flush+FIN epilogue
+                    // instead of task-parking into a watcher that would just
+                    // be dropped.
+                    Err(_) if shutdown.is_cancelled() => break,
                     Err(_) => {
                         // Cancelled by the stage-2 sweep: exit the task.
                         // read_buf holds at most MAX_PARKED_REMAINDER bytes
@@ -997,6 +1007,10 @@ pub(crate) async fn handle_connection_sharded_monoio<
                 // task-park path parks WITHOUT reading, so a mistaken
                 // downshift here would feed the park→wake→park spin.)
                 Err(ref e) if !idle_park::is_sweep_cancel(e) => break,
+                // F1 (#438): cancelled by the shutdown drain, not the idle
+                // sweep — exit through the flush+FIN epilogue instead of
+                // re-parking a read nothing will ever complete.
+                Err(_) if shutdown.is_cancelled() => break,
                 Err(_) => {
                     // Cancelled by the idle sweep: shed the working set,
                     // re-park small.