From abbe89e52520e693581754a770168bc68ce279fa Mon Sep 17 00:00:00 2001 From: am <6e30cd56c30e030cd31bb0939b94a7c257c9a09d5ba2d92cf2735da45629f248@buzz.block.builderlab.xyz> Date: Wed, 23 Sep 2026 17:53:03 -0700 Subject: [PATCH] fix: preserve live fan-out capacity during websocket thread windows Reserve outbound queue headroom for live subscriptions and clear consecutive backpressure on successful finite enqueues. Exercise interleaving and deadline behavior with the production send paths. Co-authored-by: Kalvin Chau Signed-off-by: Kalvin Chau --- .../bridge/thread_window/postgres_tests.rs | 2 +- crates/buzz-relay/src/handlers/req.rs | 143 ++++++++++++++++-- 2 files changed, 135 insertions(+), 10 deletions(-) diff --git a/crates/buzz-relay/src/api/bridge/thread_window/postgres_tests.rs b/crates/buzz-relay/src/api/bridge/thread_window/postgres_tests.rs index c2e767521e..fe820ec391 100644 --- a/crates/buzz-relay/src/api/bridge/thread_window/postgres_tests.rs +++ b/crates/buzz-relay/src/api/bridge/thread_window/postgres_tests.rs @@ -554,7 +554,7 @@ async fn websocket_batch_enforces_token_scope_and_rejects_queue_truncation() { let drained: Vec<_> = std::iter::from_fn(|| slow_rx.try_recv().ok()) .map(|m| serde_json::from_str::(m.to_text().unwrap()).unwrap()) .collect(); - assert_eq!(drained.len(), 4); + assert_eq!(drained.len(), 3); assert!( drained.iter().all(|frame| frame[0] == "EVENT"), "{drained:?}" diff --git a/crates/buzz-relay/src/handlers/req.rs b/crates/buzz-relay/src/handlers/req.rs index b08d4ef48b..cadaf44ea9 100644 --- a/crates/buzz-relay/src/handlers/req.rs +++ b/crates/buzz-relay/src/handlers/req.rs @@ -589,20 +589,46 @@ pub(crate) async fn serve_thread_windows<'a>( Ok(true) } -// Finite pages must not fail merely because the producer briefly outruns the -// socket writer. The same deadline bounds both query and queue delivery. +// Leave space for live fan-out, which uses try_send rather than a queued +// reservation. Waiting on send() would claim every freed slot before fan-out +// can use it; a bounded retry only takes a slot when headroom remains. async fn send_thread_window_frame( conn: &ConnectionState, frame: String, deadline: tokio::time::Instant, ) -> bool { - tokio::select! { - biased; - _ = conn.cancel.cancelled() => false, - result = tokio::time::timeout_at( - deadline, - conn.send_tx.send(axum::extract::ws::Message::Text(frame.into())), - ) => matches!(result, Ok(Ok(()))), + let reserve = (conn.send_tx.max_capacity() / 4) + .max(1) + .min(conn.send_tx.max_capacity().saturating_sub(1)); + let mut frame = axum::extract::ws::Message::Text(frame.into()); + loop { + if conn.cancel.is_cancelled() || tokio::time::Instant::now() >= deadline { + return false; + } + if conn.send_tx.is_closed() { + return false; + } + if conn.send_tx.capacity() > reserve { + match conn.send_tx.try_send(frame) { + Ok(()) => { + conn.backpressure_count + .store(0, std::sync::atomic::Ordering::Relaxed); + return true; + } + Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => return false, + Err(tokio::sync::mpsc::error::TrySendError::Full(message)) => { + frame = message; + } + } + } + // No reservation: a live producer can take the reserved slots before + // our next attempt. Polling is bounded by the common request deadline. + tokio::select! { + biased; + _ = conn.cancel.cancelled() => return false, + _ = tokio::time::sleep_until(deadline) => return false, + _ = tokio::time::sleep(std::time::Duration::from_millis(1)) => {} + } } } @@ -2680,3 +2706,102 @@ mod tests { assert!(!result_gated_count_safe_for_pushdown(&f, &owner)); } } + +#[cfg(test)] +mod thread_window_queue_tests { + use super::*; + use std::sync::atomic::Ordering; + use tokio::time::{timeout, Duration, Instant}; + + #[tokio::test] + async fn finite_frames_leave_live_capacity_and_reset_stale_backpressure() { + let (conn, mut rx) = crate::connection::tests::test_conn_with_auth( + crate::connection::tests::authenticated_state(), + ); + let manager = crate::state::ConnectionManager::new(); + manager.register( + conn.conn_id, + conn.send_tx.clone(), + conn.ctrl_tx.clone(), + None, + conn.cancel.clone(), + conn.tenant.community(), + conn.backpressure_count.clone(), + conn.subscriptions.clone(), + conn.grace_limit, + ); + conn.backpressure_count.store(2, Ordering::Relaxed); + let deadline = Instant::now() + Duration::from_secs(2); + for n in 0..3 { + assert!(send_thread_window_frame(&conn, format!("finite-{n}"), deadline).await); + } + assert_eq!(conn.backpressure_count.load(Ordering::Relaxed), 0); + assert_eq!( + conn.send_tx.capacity(), + 1, + "one slot reserved for live frames" + ); + assert!(manager.send_to(conn.conn_id, "live-0".into())); + assert!(!manager.send_to(conn.conn_id, "transient-full".into())); + assert_eq!(conn.backpressure_count.load(Ordering::Relaxed), 1); + assert!(!conn.cancel.is_cancelled()); + // A queued finite frame must not steal capacity before the live + // producer can use it when the reader frees a slot. + let pending = tokio::spawn({ + let conn = conn.clone(); + async move { send_thread_window_frame(&conn, "finite-3".into(), deadline).await } + }); + tokio::task::yield_now().await; + assert!(!pending.is_finished()); + assert_eq!(rx.recv().await.unwrap().to_text().unwrap(), "finite-0"); + assert!(manager.send_to(conn.conn_id, "live-1".into())); + assert_eq!(rx.recv().await.unwrap().to_text().unwrap(), "finite-1"); + assert_eq!(rx.recv().await.unwrap().to_text().unwrap(), "finite-2"); + assert!(timeout(Duration::from_secs(1), pending) + .await + .unwrap() + .unwrap()); + let remaining: Vec<_> = std::iter::from_fn(|| rx.try_recv().ok()) + .map(|msg| msg.to_text().unwrap().to_owned()) + .collect(); + assert_eq!(remaining, ["live-0", "live-1", "finite-3"]); + for n in 4..7 { + assert!(send_thread_window_frame(&conn, format!("finite-{n}"), deadline).await); + } + for n in 4..7 { + assert_eq!( + rx.recv().await.unwrap().to_text().unwrap(), + format!("finite-{n}") + ); + } + assert!(!conn.cancel.is_cancelled()); + } + + #[tokio::test] + async fn finite_frame_stalled_reader_expires_without_filling_reserved_slot() { + let (conn, mut rx) = crate::connection::tests::test_conn_with_auth( + crate::connection::tests::authenticated_state(), + ); + for n in 0..3 { + assert!(conn.send(format!("padding-{n}"))); + } + assert!( + !send_thread_window_frame( + &conn, + "finite".into(), + Instant::now() + Duration::from_millis(20), + ) + .await + ); + assert_eq!(rx.try_recv().unwrap().to_text().unwrap(), "padding-0"); + conn.cancel.cancel(); + assert!( + !send_thread_window_frame( + &conn, + "cancelled".into(), + Instant::now() + Duration::from_secs(1), + ) + .await + ); + } +}