diff --git a/services/flysim/crates/flybus/src/router/mod.rs b/services/flysim/crates/flybus/src/router/mod.rs index 2429a4b..f97d4e4 100644 --- a/services/flysim/crates/flybus/src/router/mod.rs +++ b/services/flysim/crates/flybus/src/router/mod.rs @@ -364,75 +364,58 @@ async fn write_selected( 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( 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 => {} diff --git a/services/flysim/crates/flybus/src/router/state.rs b/services/flysim/crates/flybus/src/router/state.rs index 4dc52e6..33380d0 100644 --- a/services/flysim/crates/flybus/src/router/state.rs +++ b/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, - /// Serializes synchronous transport polls with connection teardown. It is never held - /// across an await. - pub write_gate: std::sync::Mutex, + /// 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>>, } +/// 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, + /// 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(); } diff --git a/services/flysim/crates/flybus/tests/sol_rereview_regressions.rs b/services/flysim/crates/flybus/tests/sol_rereview_regressions.rs index 38042c9..244b981 100644 --- a/services/flysim/crates/flybus/tests/sol_rereview_regressions.rs +++ b/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)]