fix(flybus): stop connection teardown being starved by the frame it waits for

The write gate was one mutex the writer held across every synchronous transport
poll, and close_conn took that same mutex before reclaiming a connection's
owners. Under a transport that accepts a byte per poll the writer re-acquired it
hundreds of times while teardown queued behind it, so teardown could be starved
for a whole frame and the delivery completed just before its owner was
reclaimed. bus-v1 section 9 puts release and route-health control out of reach of
telemetry starvation; tests/sol_rereview_regressions.rs's poll-gate regression
failed 6 release runs out of 6 alone, and about one debug run in five.

The gate now separates "teardown has begun" from "a poll is in progress": a
mutex plus a condvar. begin_close sets closing once without waiting, then waits
for the poll already in progress and records whether it left a frame half
written. The writer brackets each poll with enter_poll (refused once closing)
and leave_poll(bytes) instead of holding a lock across it, so teardown's window
is one poll rather than a frame, and no byte follows it. close_conn no longer
holds the gate across reclamation. Every item is pub(crate): no public
signature changed.

The regression test no longer synchronises with a sleep. The shutdown thread
announces itself on a channel, and what follows is structural: teardown cannot
pass the gate until the held poll returns, which only release() allows. Its
delivery is 50 KB now, which a writer resuming a byte per poll cannot finish
inside the window, so the test observes the ordering instead of a coin toss and
asserts the two legal shapes of the stream: cut short with nothing appended, or
never begun with exactly the closing notices following.

services/flysim/crates/flybus/src/router/mod.rs

@@ -364,75 +364,58 @@ async fn write_selected<W: AsyncWrite + Unpin>(
     let mut frame = Vec::with_capacity(bytes.len() + 4);
     frame.extend_from_slice(&(bytes.len() as u32).to_le_bytes());
     frame.extend_from_slice(bytes);
-    {
-        let mut gate = signals.write_gate.lock().unwrap_or_else(|e| e.into_inner());
-        if !gate.begin_frame(frame.len()) {
-            return SelectedWrite::Closing {
-                partial: gate.cut_partial(),
-            };
-        }
+    let closing = || SelectedWrite::Closing {
+        partial: signals.write_gate.cut_partial(),
+    };
+    let interrupted = || io::Error::new(io::ErrorKind::Interrupted, "connection closing");
+    if !signals.write_gate.begin_frame(frame.len()) {
+        return closing();
     }
 
     let mut written = 0;
     while written < frame.len() {
         let polled = std::future::poll_fn(|cx| {
-            let mut gate = signals.write_gate.lock().unwrap_or_else(|e| e.into_inner());
-            if gate.closing() {
-                return std::task::Poll::Ready(Err(io::Error::new(
-                    io::ErrorKind::Interrupted,
-                    "connection closing",
-                )));
+            if !signals.write_gate.enter_poll() {
+                return std::task::Poll::Ready(Err(interrupted()));
             }
-            match std::pin::Pin::new(&mut *wr).poll_write(cx, &frame[written..]) {
+            let polled = std::pin::Pin::new(&mut *wr).poll_write(cx, &frame[written..]);
+            let wrote = match polled {
+                std::task::Poll::Ready(Ok(n)) => n,
+                _ => 0,
+            };
+            signals.write_gate.leave_poll(wrote);
+            match polled {
                 std::task::Poll::Ready(Ok(0)) => std::task::Poll::Ready(Err(io::Error::new(
                     io::ErrorKind::WriteZero,
                     "failed to write router frame",
                 ))),
-                std::task::Poll::Ready(Ok(n)) => {
-                    gate.wrote(n);
-                    std::task::Poll::Ready(Ok(n))
-                }
                 other => other,
             }
         })
         .await;
         match polled {
             Ok(n) => written += n,
-            Err(e) if e.kind() == io::ErrorKind::Interrupted => {
-                let gate = signals.write_gate.lock().unwrap_or_else(|e| e.into_inner());
-                return SelectedWrite::Closing {
-                    partial: gate.cut_partial(),
-                };
-            }
+            Err(e) if e.kind() == io::ErrorKind::Interrupted => return closing(),
             Err(_) => return SelectedWrite::Failed,
         }
     }
 
     let flushed = std::future::poll_fn(|cx| {
-        let gate = signals.write_gate.lock().unwrap_or_else(|e| e.into_inner());
-        if gate.closing() {
-            return std::task::Poll::Ready(Err(io::Error::new(
-                io::ErrorKind::Interrupted,
-                "connection closing",
-            )));
+        if !signals.write_gate.enter_poll() {
+            return std::task::Poll::Ready(Err(interrupted()));
         }
-        std::pin::Pin::new(&mut *wr).poll_flush(cx)
+        let polled = std::pin::Pin::new(&mut *wr).poll_flush(cx);
+        signals.write_gate.leave_poll(0);
+        polled
     })
     .await;
     if let Err(e) = flushed {
         if e.kind() == io::ErrorKind::Interrupted {
-            let gate = signals.write_gate.lock().unwrap_or_else(|e| e.into_inner());
-            return SelectedWrite::Closing {
-                partial: gate.cut_partial(),
-            };
+            return closing();
         }
         return SelectedWrite::Failed;
     }
-    signals
-        .write_gate
-        .lock()
-        .unwrap_or_else(|e| e.into_inner())
-        .finish_frame();
+    signals.write_gate.finish_frame();
     SelectedWrite::Complete
 }
 
@@ -452,13 +435,9 @@ async fn write_loop<W: AsyncWrite + Unpin>(
                 tokio::pin!(write);
                 let selected = tokio::select! {
                     result = &mut write => result,
-                    _ = stopped(&mut shutdown) => {
-                        let gate = signals
-                            .write_gate
-                            .lock()
-                            .unwrap_or_else(|e| e.into_inner());
-                        SelectedWrite::Closing { partial: gate.cut_partial() }
-                    }
+                    _ = stopped(&mut shutdown) => SelectedWrite::Closing {
+                        partial: signals.write_gate.cut_partial(),
+                    },
                 };
                 match selected {
                     SelectedWrite::Complete => {}

services/flysim/crates/flybus/src/router/state.rs

@@ -42,56 +42,98 @@ pub(crate) struct ConnSignals {
     pub wake: Notify,
     /// Flips to true when the connection is closed; both tasks watch it.
     pub shutdown: watch::Sender<bool>,
-    /// Serializes synchronous transport polls with connection teardown. It is never held
-    /// across an await.
-    pub write_gate: std::sync::Mutex<WriteGate>,
+    /// Orders connection teardown against the synchronous transport polls of the writer.
+    pub write_gate: WriteGate,
     /// The last frames to write before closing: a refusal or `connection.closing` notice,
     /// preceded on router shutdown by `subscription.closed` notices.
     pub final_frames: std::sync::Mutex<Vec<Vec<u8>>>,
 }
 
+/// Orders teardown against the writer's synchronous transport polls.
+///
+/// Teardown marks the stream closing *before* it waits for a poll already in progress, so at
+/// most that one poll can still write and every later one is refused, whichever task reaches
+/// the lock first. The earlier design held one mutex across each poll instead, which a writer
+/// sending a frame a byte per poll re-acquired hundreds of times while teardown waited for it:
+/// teardown could be starved for a whole frame and the frame completed just before its
+/// delivery owner was reclaimed. The lock is held only for these bookkeeping steps, never
+/// across an await.
 #[derive(Default)]
 pub(crate) struct WriteGate {
+    state: std::sync::Mutex<GateState>,
+    /// Signalled when a poll leaves the transport.
+    idle: std::sync::Condvar,
+}
+
+#[derive(Default)]
+struct GateState {
     closing: bool,
+    /// A writer is inside a synchronous transport poll right now.
+    polling: bool,
     frame_len: usize,
     written: usize,
     cut_partial: bool,
 }
 
 impl WriteGate {
-    pub(crate) fn begin_frame(&mut self, len: usize) -> bool {
-        if self.closing {
+    fn lock(&self) -> std::sync::MutexGuard<'_, GateState> {
+        self.state.lock().unwrap_or_else(|e| e.into_inner())
+    }
+
+    /// Starts one frame. False once teardown has begun.
+    pub(crate) fn begin_frame(&self, len: usize) -> bool {
+        let mut g = self.lock();
+        if g.closing {
             return false;
         }
-        self.frame_len = len;
-        self.written = 0;
+        g.frame_len = len;
+        g.written = 0;
         true
     }
 
-    pub(crate) fn wrote(&mut self, len: usize) {
-        self.written += len;
-        debug_assert!(self.written <= self.frame_len);
+    /// Claims the transport for one synchronous poll. False once teardown has begun.
+    pub(crate) fn enter_poll(&self) -> bool {
+        let mut g = self.lock();
+        if g.closing {
+            return false;
+        }
+        debug_assert!(!g.polling, "one writer task polls one connection");
+        g.polling = true;
+        true
     }
 
-    pub(crate) fn finish_frame(&mut self) {
-        if !self.closing {
-            debug_assert_eq!(self.written, self.frame_len);
-            self.frame_len = 0;
-            self.written = 0;
+    /// Releases the transport, accounting for what that poll wrote.
+    pub(crate) fn leave_poll(&self, wrote: usize) {
+        let mut g = self.lock();
+        g.polling = false;
+        g.written += wrote;
+        debug_assert!(g.written <= g.frame_len);
+        drop(g);
+        self.idle.notify_all();
+    }
+
+    pub(crate) fn finish_frame(&self) {
+        let mut g = self.lock();
+        if !g.closing {
+            debug_assert_eq!(g.written, g.frame_len);
+            g.frame_len = 0;
+            g.written = 0;
         }
     }
 
-    fn begin_close(&mut self) {
-        self.closing = true;
-        self.cut_partial = self.written > 0 && self.written < self.frame_len;
-    }
-
-    pub(crate) fn closing(&self) -> bool {
-        self.closing
+    /// Refuses every later poll, then waits for one already in progress and records whether it
+    /// left a frame half written. The caller may reclaim owners once this returns.
+    fn begin_close(&self) {
+        let mut g = self.lock();
+        g.closing = true;
+        while g.polling {
+            g = self.idle.wait(g).unwrap_or_else(|e| e.into_inner());
+        }
+        g.cut_partial = g.written > 0 && g.written < g.frame_len;
     }
 
     pub(crate) fn cut_partial(&self) -> bool {
-        self.cut_partial
+        self.lock().cut_partial
     }
 }
 
@@ -100,7 +142,7 @@ impl ConnSignals {
         ConnSignals {
             wake: Notify::new(),
             shutdown: watch::Sender::new(false),
-            write_gate: std::sync::Mutex::new(WriteGate::default()),
+            write_gate: WriteGate::default(),
             final_frames: std::sync::Mutex::new(Vec::new()),
         }
     }
@@ -705,15 +747,14 @@ impl State {
             .lock()
             .unwrap_or_else(|e| e.into_inner())
             .extend(final_frames);
-        // A transport poll that is already in progress finishes before this lock is acquired.
-        // Once acquired, teardown marks the stream closing before reclaiming any owner, and no
-        // later normal-frame poll is allowed through.
-        let mut write_gate = signals.write_gate.lock().unwrap_or_else(|e| e.into_inner());
-        write_gate.begin_close();
+        // Teardown refuses every later transport poll first, then waits for a poll already in
+        // progress, and only then reclaims what the connection owned. So a delivery frame is
+        // either complete before its owner is reclaimed, or left truncated with nothing more
+        // appended to the stream; the writer can never finish it afterwards.
+        signals.write_gate.begin_close();
         self.disconnect(c);
         signals.shutdown.send_replace(true);
         signals.wake.notify_one();
-        drop(write_gate);
         self.flush_notices();
     }
 

services/flysim/crates/flybus/tests/sol_rereview_regressions.rs

@@ -734,8 +734,17 @@ async fn teardown_waits_for_an_active_transport_poll_before_reclaiming() {
         .store_dir()
         .join("sealed")
         .join(&artifact.reference().artifact_id);
+    // A deliberately large delivery: the transport under the gate writes one byte per poll,
+    // so a writer that resumes cannot possibly finish this frame inside the window between
+    // releasing the held poll and teardown marking the stream closing. Without that, a short
+    // frame sometimes completes first, which is teardown's other legal arm and would make the
+    // assertions below a coin toss rather than a test of the ordering.
     publisher
-        .publish("t.poll-gate", obj(json!({})), &[("data", &artifact)])
+        .publish(
+            "t.poll-gate",
+            obj(json!({"blob": "p".repeat(50_000)})),
+            &[("data", &artifact)],
+        )
         .await
         .unwrap();
     drop(artifact);
@@ -745,19 +754,27 @@ async fn teardown_waits_for_an_active_transport_poll_before_reclaiming() {
     let done = Arc::new(AtomicBool::new(false));
     let shutdown_done = done.clone();
     let shutdown_router = router.clone();
+    // The thread announces itself before calling shutdown, so the assertions below need no
+    // sleep: teardown cannot get past the write gate until the held poll returns, which only
+    // `hold.release()` allows.
+    let (started_tx, started_rx) = std::sync::mpsc::channel();
     let shutdown = std::thread::spawn(move || {
+        started_tx.send(()).expect("the test is waiting");
         shutdown_router.shutdown();
         shutdown_done.store(true, Ordering::SeqCst);
     });
-    tokio::time::sleep(Duration::from_millis(50)).await;
-    assert!(
-        !done.load(Ordering::SeqCst),
-        "teardown completed while poll_write was active"
-    );
-    assert!(
-        sealed_path.exists(),
-        "artifact was reclaimed while poll_write was active"
-    );
+    started_rx.recv().expect("shutdown thread started");
+    for _ in 0..64 {
+        assert!(
+            !done.load(Ordering::SeqCst),
+            "teardown completed while poll_write was active"
+        );
+        assert!(
+            sealed_path.exists(),
+            "artifact was reclaimed while poll_write was active"
+        );
+        tokio::task::yield_now().await;
+    }
 
     hold.release();
     shutdown.join().unwrap();
@@ -765,12 +782,23 @@ async fn teardown_waits_for_an_active_transport_poll_before_reclaiming() {
     assert_eq!(router.stats().owners, 0);
     assert_eq!(router.stats().artifacts, 0);
     assert!(!sealed_path.exists());
+    let mut ops = Vec::new();
     while let Some(envelope) = within("poll-gate close", raw.recv()).await {
         assert_ne!(
             envelope.op, "topic.message",
             "delivery completed after teardown reclaimed its owner"
         );
+        ops.push(envelope.op);
     }
+    // The delivery never completes, so only teardown's two shapes are legal: the frame was
+    // cut short and nothing whatever follows it, or it never began and the stream is still
+    // frame aligned, in which case the closing notices are all that follow.
+    assert!(
+        ops.is_empty()
+            || ops == ["subscription.closed".to_owned(), "connection.closing".to_owned()],
+        "a cut stream carries nothing more and an aligned one exactly the closing notices: \
+         {ops:?}"
+    );
 }
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]