jerry/buzz
1
0
Fork 0
mirror of https://github.com/block/buzz.git synced 2026-10-27 12:40:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(acp): pace targeted overflow recovery on consumer capacity (#7325)

Browse Source 
🤖
## Summary

When a Buzz agent falls behind on incoming messages, its connection can
make the backlog worse while trying to recover. The connection buffers
messages from the relay server until the agent is ready to process them;
if that buffer overflows, recovery previously requested history for
**every subscribed channel** and paused socket reads while sending those
requests. That adds traffic to an already overloaded connection. This
change requests history only for affected subscriptions, once the code
consuming those messages has room, with at least five seconds between
attempts.

The recovery path now:

- Combines repeated losses into one pending recovery per affected
subscription, keeping the oldest dropped timestamp so replay starts
early enough.
- Waits until at least half the consumer queue is free and the relay's
existing rate-limit delay has expired. The queue wakes recovery when
space becomes available; recovery does not periodically sample capacity
or hold queue space away from live messages.
- Attempts one subscription at a time, choosing the least recently
attempted so a busy channel cannot crowd out other channels or
membership notifications. The five-second delay starts when an attempt
finishes, including a failed write; failed writes leave recovery
pending.

Recovery is paced by available capacity, not by how often messages are
lost. This is not a larger buffer or a cutoff that abandons recovery.
Subscription identifiers, message filters, replay timestamp overlap and
duplicate filtering are unchanged; no downstream agent changes are
required.

This targets a reproducible overload **amplifier**, not every cause of
overload or every catch-up limitation. The initial live overload's cause
has not been established. Recovery remains best effort: a successful
request write is not proof of delivery, and existing history/retention
limits, bounded duplicate tracking and replay limitations still apply.
There is no exactly-once or complete catch-up guarantee. A stalled write
can still pause socket reads for the existing ten-second timeout; the
pacing bound does not cover initial subscriptions, reconnects or other
retry paths.

### Related issue

Closest related: #5014 (channel re-subscription); also #6661 (membership
reconciliation) and #6090 (relay backpressure gap signaling). This
addresses local overflow recovery scheduling, not those separate
mechanisms.

### Testing

Recorded offline comparisons against the previous behavior, with the
final implementation at `8000636f3073167c5a5107bb179c7d91160f1729`:

| Same fixture: 18 subscriptions, three overload rounds | Before | After
|
| --- | --- | --- |
| Recovery history requests | 108 | 3 |
| Ping-response delay | About 4.6 seconds | Below the measurement's 1 ms
resolution |

A separate bounded-history fixture delivered all 320 events plus
subsequent live traffic in **both** versions. Regression coverage
exercises the real socket-handling task, including intermittent consumer
capacity, fairness, failed writes and cancellation of capacity waits
before live delivery. These are synthetic results, not production
throughput measurements or evidence of a deployed cure.

The full local `RUST_TEST_THREADS=4 just ci` run passed on September 4,
2026. Earlier unsuccessful local runs remain part of the validation
history. The [recorded validation evidence and separate desktop
follow-up](https://github.com/block/buzz/pull/7325#issuecomment-5540592398)
preserve the original desktop mock-history scroll failure, its passing
rerun and the remaining investigation. That desktop path does not run
the agent connection code; neither this repair nor the passing rerun
fixes the observed scroll problem.

---------

Signed-off-by: Logan Johnson <loganj@squareup.com>
This commit is contained in:
Logan Johnson
2026-09-08 17:51:47 -04:00
committed by GitHub
parent 218633b8fd
commit c045321a7f
5 changed files with 1052 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
crates/buzz-acp/README.md
View File

@@ -270,6 +270,10 @@ Forum event kinds:
4. **Prompting** — When events are pending and no prompt is in flight for that channel, drains all queued events for the oldest channel into a single batched prompt via ACP `session/prompt`.
5. **Agent response** — The agent processes the prompt and uses the Buzz CLI (`send_message`, `get_messages`, etc.) to interact with Buzz.
6. **Recovery** — If the agent crashes, the harness respawns it. If the relay disconnects, the harness reconnects with a `since` filter to avoid missing events.
If the inbound queue overflows, the harness attempts replay for affected
subscriptions when capacity and relay quota permit, with at least five seconds
between attempts. Recovery depends on available relay history and the consumer
making progress; complete delivery is not guaranteed.
Each channel has at most one prompt in flight. Multiple channels can be processed concurrently when agents > 1.

112
crates/buzz-acp/src/relay.rs
View File

@@ -1160,10 +1160,8 @@ struct BgState {
/// On reconnect resubscribe, `since` = min(last_seen, channel_dropped_since).
/// Cleared per-channel after a successful resubscribe.
channel_dropped_since: HashMap<Uuid, u64>,
/// Set by the backpressure handler when the event channel is full.
/// The main loop checks this flag and triggers a proactive resubscribe
/// (without waiting for a disconnect) so dropped events are replayed.
proactive_resubscribe_needed: bool,
/// Rate/fairness bookkeeping only; replay cursors retain baseline semantics.
recovery: recovery::RecoverySchedule,
/// Unix timestamp captured just before the relay connection was established.
/// Used as the floor `since` for membership notification replay so events
/// predating this session are never re-delivered.
@@ -1237,7 +1235,7 @@ impl BgState {
membership_sub_active: false,
observer_control_sub_active: false,
channel_dropped_since: HashMap::new(),
proactive_resubscribe_needed: false,
recovery: recovery::RecoverySchedule::default(),
startup_watermark: None,
subscribe_since: HashMap::new(),
rate_limit_gate: None,
@@ -1298,6 +1296,9 @@ impl BgState {
/// Prevents stale replay on re-subscribe and avoids unbounded state growth
/// for channels that are removed and never re-added.
fn clear_channel_state(&mut self, channel_id: &Uuid) {
self.recovery
.last_attempt
.remove(&channel_sub_id(*channel_id));
self.last_seen.remove(channel_id);
self.subscribe_since.remove(channel_id);
self.channel_dropped_since.remove(channel_id);
@@ -1826,82 +1827,6 @@ async fn run_background_task(
let mut drain_pacing_next: Option<tokio::time::Instant> = None;
loop {
if state.proactive_resubscribe_needed {
state.proactive_resubscribe_needed = false;
info!("proactive resubscribe triggered by backpressure event loss");
// Proactive resubscribe runs on the EXISTING socket — do NOT clear the
// rate-limit gate or pending queues.
match resubscribe_after_reconnect(
&mut ws,
&mut cmd_rx,
&mut state,
&agent_pubkey_hex,
false, // existing socket — preserve gate state
)
.await
{
ResubscribeResult::Ok => {}
ResubscribeResult::Shutdown => return,
ResubscribeResult::RetryConnection => {
warn!("proactive resubscribe had failures — triggering reconnect");
let _ = event_tx.try_send(None);
match try_autonomous_reconnect(
&mut ws,
&mut cmd_rx,
&mut state,
&keys,
&relay_url,
&agent_pubkey_hex,
&event_tx,
&observer_control_tx,
auth_tag.as_ref(),
)
.await
{
ReconnectOutcome::Ok => {
if matches!(
drain_post_reconnect(
&mut ws,
&mut cmd_rx,
&mut state,
&agent_pubkey_hex
)
.await,
ReconnectOutcome::Shutdown
) {
return;
}
}
ReconnectOutcome::Shutdown => return,
ReconnectOutcome::Failed => {
if matches!(
wait_for_reconnect(
&mut ws,
&mut cmd_rx,
&mut state,
&keys,
&relay_url,
&agent_pubkey_hex,
&event_tx,
&observer_control_tx,
true,
auth_tag.as_ref(),
)
.await,
ReconnectOutcome::Shutdown
) {
return;
}
}
}
ping_sent = false;
last_pong = Instant::now();
connected_since = Instant::now();
stable_logged = false;
}
}
}
// Drain pending subs, one REQ per pacing tick within the relay's
// admission window.
let drain_window_open = drain_pacing_next.is_none_or(|t| tokio::time::Instant::now() >= t);
@@ -1982,7 +1907,11 @@ async fn run_background_task(
}
}
let recovery_at = recovery::ready_at(&mut state);
tokio::select! {
_ = recovery::ready(&event_tx, recovery_at) => {
recovery::recover_one(&mut ws, &mut state, &event_tx, &agent_pubkey_hex).await;
}
raw = ws.next() => {
// Determine if the socket is lost.
let socket_lost = match raw {
@@ -2358,12 +2287,10 @@ async fn handle_ws_message(
// replay starts early enough to re-deliver it.
state.membership_dropped_since =
Some(state.membership_dropped_since.map_or(ts, |d| d.min(ts)));
// Proactively trigger resubscribe without waiting for a disconnect.
state.proactive_resubscribe_needed = true;
warn!(
channel_id = %channel_uuid,
ts,
"membership notification dropped (backpressure) — proactive resubscribe queued"
"membership notification dropped (backpressure) — targeted recovery pending"
);
}
Err(mpsc::error::TrySendError::Closed(_)) => return false,
@@ -2400,12 +2327,10 @@ async fn handle_ws_message(
.entry(channel_id)
.and_modify(|d| *d = (*d).min(ts))
.or_insert(ts);
// Proactively trigger resubscribe without waiting for a disconnect.
state.proactive_resubscribe_needed = true;
warn!(
channel_id = %channel_id,
ts,
"event channel full — dropping event for channel {channel_id} — proactive resubscribe queued"
"event channel full — dropping event for channel {channel_id} — targeted recovery pending"
);
}
Err(mpsc::error::TrySendError::Closed(_)) => {
@@ -4248,6 +4173,11 @@ async fn wait_for_any_ok(
}
}
mod recovery;
#[cfg(test)]
mod recovery_tests;
#[cfg(test)]
mod tests {
use super::*;
@@ -4775,7 +4705,7 @@ mod tests {
.expect("signing should succeed")
}
async fn test_ws_pair() -> (WsStream, WebSocketStream<tokio::net::TcpStream>) {
pub(super) async fn test_ws_pair() -> (WsStream, WebSocketStream<tokio::net::TcpStream>) {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind test websocket");
@@ -4792,7 +4722,7 @@ mod tests {
(client, server.await.expect("join test websocket server"))
}
async fn next_test_frame(
pub(super) async fn next_test_frame(
server: &mut WebSocketStream<tokio::net::TcpStream>,
) -> serde_json::Value {
let message = timeout(Duration::from_secs(1), server.next())
@@ -5035,14 +4965,14 @@ mod tests {
));
}
fn test_channel_filter() -> ChannelFilter {
pub(super) fn test_channel_filter() -> ChannelFilter {
ChannelFilter {
kinds: Some(vec![9]),
require_mention: false,
}
}
fn seed_test_subscription(state: &mut BgState, channel_id: Uuid) {
pub(super) fn seed_test_subscription(state: &mut BgState, channel_id: Uuid) {
apply_command_to_state(
state,
RelayCommand::Subscribe {

132
crates/buzz-acp/src/relay/recovery.rs Normal file
View File

@@ -0,0 +1,132 @@
//! Overflow recovery is an attempted replay, not an EOSE/consumer receipt.
//! Keep the existing IDs and cursor retirement rules; bound when work is sent.
use super::*;
pub(super) const RECOVERY_INTERVAL: Duration = Duration::from_secs(5);
#[derive(Default)]
pub(super) struct RecoverySchedule {
next_attempt: Option<tokio::time::Instant>,
pub(super) last_attempt: HashMap<String, tokio::time::Instant>,
}
/// Attempt at most one affected subscription, with space for replay to arrive.
/// Failed writes retain the loss cursor and are paced too. No EOSE is interpreted
/// as completion: overlapping requests keep their existing stable wire IDs.
pub(super) async fn recover_one(
ws: &mut WsStream,
state: &mut BgState,
event_tx: &mpsc::Sender<Option<BuzzEvent>>,
agent_pubkey_hex: &str,
) {
let now = tokio::time::Instant::now();
if event_tx.is_closed()
|| event_tx.capacity() < event_tx.max_capacity().div_ceil(2)
|| state.recovery.next_attempt.is_some_and(|next| now < next)
|| state.check_rate_gate().is_some()
{
return;
}
let channel = next_channel(state);
let Some(channel) = channel else { return };
let sub = channel.map_or_else(|| MEMBERSHIP_NOTIF_SUB_ID.to_owned(), channel_sub_id);
state.recovery.last_attempt.insert(sub.clone(), now);
info!(subscription = sub, "attempting targeted overflow replay");
if let Some(ch) = channel {
if let Some(filter) = state.active_filters.get(&ch).cloned() {
let since = state.channel_since(&ch);
if send_subscribe(ws, state, ch, agent_pubkey_hex, since, &filter).await {
// Baseline retirement point: REQ write, NOT proven delivery.
// New overflow after this attempt creates another pending cursor.
state.channel_dropped_since.remove(&ch);
}
}
} else {
let since = match (state.membership_dropped_since, state.membership_last_seen) {
(Some(d), Some(l)) => Some(d.min(l)),
(Some(d), None) => Some(d),
(None, Some(l)) => Some(l),
(None, None) => state.startup_watermark,
};
if send_membership_subscribe(ws, agent_pubkey_hex, since).await {
state.membership_dropped_since = None;
}
}
// Pace from the end of a potentially backpressured write. No catch-up burst.
// The existing bounded write timeout and read/ping owner detect socket loss.
state.recovery.next_attempt = Some(tokio::time::Instant::now() + RECOVERY_INTERVAL);
}
/// No timer or capacity waiter when another authority owns all pending loss.
/// Only actual attempts advance the cooldown; closed gates wake at expiry.
pub(super) fn ready_at(state: &mut BgState) -> Option<tokio::time::Instant> {
next_channel(state)?;
Some(
state
.recovery
.next_attempt
.unwrap_or_else(tokio::time::Instant::now)
.max(
state
.check_rate_gate()
.unwrap_or_else(tokio::time::Instant::now),
),
)
}
/// Select-local readiness, not a send or a reservation carried across reads.
/// The socket task is the sole producer. `select!` drops this future (including
/// partial permits) BEFORE handling another frame/command, so live try_send
/// never competes with a recovery reservation. Receives only add capacity.
/// Keep this future inside select!: awaiting it alone would block the reader;
/// persisting it across iterations would steal capacity from live delivery.
pub(super) async fn ready(
event_tx: &mpsc::Sender<Option<BuzzEvent>>,
at: Option<tokio::time::Instant>,
) {
if let Some(at) = at {
if tokio::time::Instant::now() < at {
tokio::time::sleep_until(at).await;
}
// Use the channel's own race-free capacity wake, not periodic samples.
// Return ALL permits before recover_one rechecks capacity and intent.
if let Ok(permits) = event_tx
.reserve_many(event_tx.max_capacity().div_ceil(2))
.await
{
drop(permits);
return;
}
}
// No loss or a closed receiver: no immediate-ready/error wake loop.
std::future::pending::<()>().await;
}
fn next_channel(state: &BgState) -> Option<Option<Uuid>> {
// One record per active intent, not per loss or per request generation.
// Least recently attempted prevents a repeatedly overflowing channel from
// starving other channels or membership. Missing filters fail closed.
state
.channel_dropped_since
.keys()
.filter(|ch| {
state.active_subscriptions.contains_key(ch)
&& state.active_filters.contains_key(ch)
&& !state.rate_limited_pending.contains_key(ch)
&& !state.resubscribe_retry.contains(ch)
})
.copied()
.map(Some)
.chain(
(state.membership_sub_active
&& state.membership_dropped_since.is_some()
&& !state.membership_resub_needed)
.then_some(None),
)
.min_by_key(|ch| {
let sub = ch.map_or_else(|| MEMBERSHIP_NOTIF_SUB_ID.to_owned(), channel_sub_id);
(state.recovery.last_attempt.get(&sub).copied(), sub)
})
}

388
crates/buzz-acp/src/relay/recovery_tests.rs Normal file
View File

@@ -0,0 +1,388 @@
//! Bounded synthetic WebSocket fixtures; no relay service, proxy or real agent.
use super::tests::{next_test_frame, seed_test_subscription, test_channel_filter, test_ws_pair};
use super::*;
fn fixture_event(channel: Uuid, n: u64, kind: u16) -> Event {
let keys =
Keys::parse("0000000000000000000000000000000000000000000000000000000000000001").unwrap();
EventBuilder::new(Kind::Custom(kind), format!("synthetic-{n}"))
.tags([Tag::parse(["h", &channel.to_string()]).unwrap()])
.custom_created_at(nostr::Timestamp::from(1_000 + n))
.sign_with_keys(&keys)
.unwrap()
}
async fn dispatch(
client: &mut WsStream,
state: &mut BgState,
tx: &mpsc::Sender<Option<BuzzEvent>>,
frame: Value,
) {
let (control_tx, _control_rx) = mpsc::channel(1);
assert!(
handle_ws_message(
Message::Text(frame.to_string().into()),
client,
tx,
&control_tx,
state,
&Keys::generate(),
"ws://127.0.0.1:1",
"synthetic-agent",
None,
)
.await
);
}
#[tokio::test]
async fn repeated_overflow_recovers_only_affected_channel_after_capacity() {
let (mut client, mut server) = test_ws_pair().await;
let mut state = BgState::new();
let channels: Vec<_> = (0..18).map(|_| Uuid::new_v4()).collect();
for ch in &channels {
seed_test_subscription(&mut state, *ch);
}
let ch = channels[0];
let sub = channel_sub_id(ch);
let (tx, mut rx) = mpsc::channel(256);
// Relay history is newest-first; the oldest dropped event must survive
// a watermark already advanced by much newer successfully-enqueued events.
let events: Vec<_> = (0..320).rev().map(|n| fixture_event(ch, n, 9)).collect();
for event in &events {
dispatch(&mut client, &mut state, &tx, json!(["EVENT", sub, event])).await;
recovery::recover_one(&mut client, &mut state, &tx, "synthetic-agent").await;
}
assert_eq!(rx.len(), 256);
assert_eq!(state.channel_dropped_since[&ch], 1_000);
assert!(timeout(Duration::from_millis(30), server.next())
.await
.is_err());
for event in &events[..256] {
assert_eq!(rx.recv().await.unwrap().unwrap().event.id, event.id);
}
recovery::recover_one(&mut client, &mut state, &tx, "synthetic-agent").await;
let req = next_test_frame(&mut server).await;
assert_eq!(req[0], "REQ");
assert_eq!(req[1], sub);
assert_eq!(req[2]["#h"], json!([ch.to_string()]));
assert_eq!(req[2]["kinds"], json!([9]));
assert_eq!(req[2]["since"], 995);
// Concurrent timer ticks / duplicate arrivals cannot replace the replay.
for _ in 0..40 {
recovery::recover_one(&mut client, &mut state, &tx, "synthetic-agent").await;
}
assert!(timeout(Duration::from_millis(30), server.next())
.await
.is_err());
for event in &events {
dispatch(&mut client, &mut state, &tx, json!(["EVENT", sub, event])).await;
}
dispatch(&mut client, &mut state, &tx, json!(["EOSE", sub])).await;
assert_eq!(rx.len(), 64, "delivered IDs must remain deduplicated");
for event in &events[256..] {
assert_eq!(rx.recv().await.unwrap().unwrap().event.id, event.id);
}
assert_eq!(state.channel_since(&ch), Some(1_319));
recovery::recover_one(&mut client, &mut state, &tx, "synthetic-agent").await;
assert!(timeout(Duration::from_millis(30), server.next())
.await
.is_err());
let live = fixture_event(ch, 400, 9);
dispatch(&mut client, &mut state, &tx, json!(["EVENT", sub, live])).await;
assert_eq!(rx.recv().await.unwrap().unwrap().event.id, live.id);
println!("18 subscriptions; 320 newest-first arrivals; 64 losses coalesced; 0 REQ while full; 1 targeted REQ; 320 unique deliveries + live");
}
#[tokio::test]
async fn socket_owner_services_ping_shutdown_and_coalesces_overflow_ticks() {
let (client, mut server) = test_ws_pair().await;
let (tx, mut rx) = mpsc::channel(1);
let (control_tx, _control_rx) = mpsc::channel(1);
let (cmd_tx, cmd_rx) = mpsc::channel(64);
let task = tokio::spawn(run_background_task(
client,
VecDeque::new(),
tx,
control_tx,
cmd_rx,
Keys::generate(),
"ws://127.0.0.1:1".into(),
"synthetic-agent".into(),
None,
));
let channels: Vec<_> = (0..18).map(|_| Uuid::new_v4()).collect();
for ch in &channels {
cmd_tx
.send(RelayCommand::Subscribe {
channel_id: *ch,
filter: test_channel_filter(),
replay_since: Some(1_000),
})
.await
.unwrap();
}
let mut subscriptions = 0;
while subscriptions < 18 {
let frame = timeout(Duration::from_secs(2), server.next())
.await
.unwrap()
.unwrap()
.unwrap();
match frame {
Message::Text(text) => {
let req: Value = serde_json::from_str(&text).unwrap();
assert_eq!(req[0], "REQ");
server
.send(Message::Text(json!(["EOSE", req[1]]).to_string().into()))
.await
.unwrap();
subscriptions += 1;
}
Message::Ping(payload) => server.send(Message::Pong(payload)).await.unwrap(),
other => panic!("unexpected {other:?}"),
}
}
let sub = channel_sub_id(channels[0]);
for n in 0..40 {
server
.send(Message::Text(
json!(["EVENT", sub, fixture_event(channels[0], n, 9)])
.to_string()
.into(),
))
.await
.unwrap();
}
server.send(Message::Ping(vec![42].into())).await.unwrap();
timeout(Duration::from_secs(2), async {
loop {
match server.next().await.unwrap().unwrap() {
Message::Ping(payload) => server.send(Message::Pong(payload)).await.unwrap(),
Message::Pong(payload) => {
assert_eq!(payload.as_ref(), &[42]);
break;
}
other => panic!("no immediate all-channel recovery before ping: {other:?}"),
}
}
})
.await
.unwrap();
// Wait across a recovery tick with the consumer still full.
assert!(socket_frame(&mut server, Duration::from_millis(5_100))
.await
.is_none());
rx.recv().await.unwrap().unwrap();
let frame = timeout(Duration::from_secs(6), server.next())
.await
.unwrap()
.unwrap()
.unwrap();
let req: Value = serde_json::from_str(frame.to_text().unwrap()).unwrap();
assert_eq!(req[1], sub);
assert_eq!(req[2]["since"], 996);
assert!(timeout(Duration::from_millis(100), server.next())
.await
.is_err());
// Sustained lag: each replay is followed by another burst, without EOSE.
// Recovery must stay paced, not permanently stall or sweep healthy channels.
for round in 1..=3 {
let started = tokio::time::Instant::now();
for n in round * 40..(round + 1) * 40 {
server
.send(Message::Text(
json!(["EVENT", sub, fixture_event(channels[0], n, 9)])
.to_string()
.into(),
))
.await
.unwrap();
}
server.send(Message::Ping(vec![43].into())).await.unwrap();
let pong = timeout(Duration::from_secs(1), server.next())
.await
.unwrap()
.unwrap()
.unwrap();
assert!(
matches!(pong, Message::Pong(_)),
"recovery preempted ping: {pong:?}"
);
rx.recv().await.unwrap().unwrap();
let frame = timeout(Duration::from_secs(6), server.next())
.await
.unwrap()
.unwrap()
.unwrap();
let req: Value = serde_json::from_str(frame.to_text().unwrap()).unwrap();
assert_eq!(req[1], sub, "healthy channels must not be swept");
assert!(
started.elapsed() >= Duration::from_secs(4),
"unpaced repeat"
);
}
cmd_tx.send(RelayCommand::Shutdown).await.unwrap();
timeout(Duration::from_secs(1), task)
.await
.unwrap()
.unwrap();
println!("socket-owner seam: 18 live REQs, 39 coalesced losses, PONG while full, zero recovery across full-capacity timer tick, four paced targeted REQs over sustained lag without EOSE, responsive shutdown");
}
#[tokio::test]
async fn recovery_is_fair_and_paced_even_with_new_loss_and_stale_eose() {
let (mut client, mut server) = test_ws_pair().await;
let mut state = BgState::new();
let channels = [Uuid::new_v4(), Uuid::new_v4()];
for ch in channels {
seed_test_subscription(&mut state, ch);
state.channel_dropped_since.insert(ch, 600);
}
state.membership_sub_active = true;
state.membership_dropped_since = Some(500);
let (tx, _rx) = mpsc::channel(1);
let mut visited = HashSet::new();
for round in 0..9 {
timeout(
Duration::from_millis(100),
recovery::ready(&tx, recovery::ready_at(&mut state)),
)
.await
.unwrap();
recovery::recover_one(&mut client, &mut state, &tx, "agent").await;
let req = next_test_frame(&mut server).await;
let sub = req[1].as_str().unwrap();
if round < 3 {
assert!(visited.insert(sub.to_owned()), "starved intent");
}
if sub == MEMBERSHIP_NOTIF_SUB_ID {
assert_eq!(req[2]["since"], 495);
state.membership_dropped_since = Some(500);
} else {
let ch = channel_id_from_sub_id(sub).unwrap();
assert_eq!(req[2]["since"], 595);
state.channel_dropped_since.insert(ch, 600);
}
// Neither stale nor current EOSE creates completion state or erases loss.
dispatch(&mut client, &mut state, &tx, json!(["EOSE", sub])).await;
for _ in 0..30 {
recovery::recover_one(&mut client, &mut state, &tx, "agent").await;
}
assert!(timeout(Duration::from_millis(1), server.next())
.await
.is_err());
advance_clock(recovery::RECOVERY_INTERVAL).await;
}
for ch in channels {
state.active_subscriptions.remove(&ch);
state.clear_channel_state(&ch);
assert!(!state
.recovery
.last_attempt
.contains_key(&channel_sub_id(ch)));
}
assert_eq!(state.recovery.last_attempt.len(), 1);
}
#[tokio::test]
async fn gate_headroom_failed_writes_and_reconnect_preserve_pending_attempts() {
let (mut client, mut server) = test_ws_pair().await;
let mut state = BgState::new();
let ch = Uuid::new_v4();
seed_test_subscription(&mut state, ch);
state.channel_dropped_since.insert(ch, 700);
let (tx, mut rx) = mpsc::channel(4);
for _ in 0..3 {
tx.try_send(None).unwrap();
}
recovery::recover_one(&mut client, &mut state, &tx, "agent").await;
assert!(state.recovery.last_attempt.is_empty());
rx.recv().await;
state.rate_limit_gate = Some(tokio::time::Instant::now() + Duration::from_secs(10));
recovery::recover_one(&mut client, &mut state, &tx, "agent").await;
assert!(state.recovery.last_attempt.is_empty());
advance_clock(Duration::from_secs(10)).await;
// Close locally, so the actual production writer fails deterministically.
client.close(None).await.unwrap();
server.next().await;
recovery::recover_one(&mut client, &mut state, &tx, "agent").await;
assert_eq!(state.channel_dropped_since[&ch], 700);
let attempted = state.recovery.last_attempt.clone();
for _ in 0..30 {
recovery::recover_one(&mut client, &mut state, &tx, "agent").await;
}
assert_eq!(state.recovery.last_attempt, attempted);
advance_clock(recovery::RECOVERY_INTERVAL).await;
recovery::recover_one(&mut client, &mut state, &tx, "agent").await;
assert_ne!(
state.recovery.last_attempt, attempted,
"failed write must be retried"
);
assert_eq!(state.channel_dropped_since[&ch], 700);
let (mut client, mut server) = test_ws_pair().await;
let (_cmd_tx, mut cmd_rx) = mpsc::channel(1);
assert!(matches!(
resubscribe_after_reconnect(&mut client, &mut cmd_rx, &mut state, "agent", true,).await,
ResubscribeResult::Ok
));
let req = next_test_frame(&mut server).await;
assert_eq!(req[1], channel_sub_id(ch));
assert_eq!(req[2]["since"], 695);
assert!(!state.channel_dropped_since.contains_key(&ch));
// This is deliberately the baseline write-retirement contract, not a receipt.
}
async fn advance_clock(duration: Duration) {
tokio::time::pause();
tokio::time::advance(duration).await;
tokio::time::resume();
}
#[tokio::test]
async fn blocked_recovery_write_is_bounded_and_retains_loss() {
let (mut client, _stalled_server) = test_ws_pair().await;
let mut state = BgState::new();
let ch = Uuid::new_v4();
seed_test_subscription(&mut state, ch);
// Bounded 16MB JSON request exceeds loopback TCP buffering. The server does
// not read it. This tests the real production write/timeout, not a mock sink.
state.active_filters.get_mut(&ch).unwrap().kinds = Some(vec![9; 8_000_000]);
state.channel_dropped_since.insert(ch, 700);
let (tx, _rx) = mpsc::channel(1);
let started = tokio::time::Instant::now();
timeout(
Duration::from_secs(15),
recovery::recover_one(&mut client, &mut state, &tx, "agent"),
)
.await
.unwrap();
assert!(started.elapsed() >= Duration::from_secs(WS_SEND_TIMEOUT_SECS));
assert_eq!(state.channel_dropped_since[&ch], 700);
let attempted = state.recovery.last_attempt.clone();
recovery::recover_one(&mut client, &mut state, &tx, "agent").await;
assert_eq!(state.recovery.last_attempt, attempted);
}
// Keep the fixture responsive to independent client keepalives while checking
// recovery traffic. Wall-clock scheduling may deliver the initial ping late.
async fn socket_frame(
server: &mut WebSocketStream<tokio::net::TcpStream>,
duration: Duration,
) -> Option<Message> {
let deadline = tokio::time::Instant::now() + duration;
loop {
match tokio::time::timeout_at(deadline, server.next()).await {
Err(_) => return None,
Ok(Some(Ok(Message::Ping(payload)))) => {
server.send(Message::Pong(payload)).await.unwrap();
}
Ok(Some(Ok(frame))) => return Some(frame),
other => panic!("unexpected socket state: {other:?}"),
}
}
}
#[path = "recovery_wake_tests.rs"]
mod wake;

507
crates/buzz-acp/src/relay/recovery_wake_tests.rs Normal file
View File

@@ -0,0 +1,507 @@
//! Capacity-wake boundaries and the independently reproduced R1 schedule.
use super::*;
// Reviewer-authored, exact-source comparison of recurring headroom at the socket owner.
// A small periodically refilled queue is empty for most of each five-second period.
async fn review_count_until(
server: &mut WebSocketStream<tokio::net::TcpStream>,
deadline: tokio::time::Instant,
) -> usize {
let mut count = 0;
while tokio::time::Instant::now() < deadline {
let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
match socket_frame(server, remaining).await {
None => break,
Some(Message::Text(text)) => {
let frame: Value = serde_json::from_str(&text).unwrap();
assert_eq!(frame[0], "REQ");
count += 1;
}
other => panic!("unexpected {other:?}"),
}
}
count
}
async fn review_barrier(server: &mut WebSocketStream<tokio::net::TcpStream>) -> usize {
server.send(Message::Ping(vec![77].into())).await.unwrap();
let mut count = 0;
loop {
match socket_frame(server, Duration::from_secs(2)).await.unwrap() {
Message::Pong(payload) => {
assert_eq!(payload.as_ref(), &[77]);
return count;
}
Message::Text(text) => {
let frame: Value = serde_json::from_str(&text).unwrap();
assert_eq!(frame[0], "REQ");
count += 1;
}
other => panic!("unexpected {other:?}"),
}
}
}
#[tokio::test]
async fn review_recurring_headroom_between_ticks_gets_an_attempt() {
let (client, mut server) = test_ws_pair().await;
let (tx, mut rx) = mpsc::channel(1);
let (control_tx, _control_rx) = mpsc::channel(1);
let (cmd_tx, cmd_rx) = mpsc::channel(8);
let start = tokio::time::Instant::now();
let task = tokio::spawn(run_background_task(
client,
VecDeque::new(),
tx,
control_tx,
cmd_rx,
Keys::generate(),
"ws://127.0.0.1:1".into(),
"agent".into(),
None,
));
let ch = Uuid::new_v4();
let sub = channel_sub_id(ch);
cmd_tx
.send(RelayCommand::Subscribe {
channel_id: ch,
filter: test_channel_filter(),
replay_since: Some(1000),
})
.await
.unwrap();
let frame = socket_frame(&mut server, Duration::from_secs(2))
.await
.unwrap();
let req: Value = serde_json::from_str(frame.to_text().unwrap()).unwrap();
assert_eq!(req[1], sub);
let lost = fixture_event(ch, 1, 9);
for event in [fixture_event(ch, 0, 9), lost.clone()] {
server
.send(Message::Text(
json!(["EVENT", sub, event]).to_string().into(),
))
.await
.unwrap();
}
let mut attempts = review_barrier(&mut server).await;
assert_eq!(rx.len(), 1);
rx.recv().await.unwrap().unwrap();
for round in 1..=4 {
// Headroom until 1s before the tick; then only one live arrival, no new
// overflow. The queue stays full across the tick and drains 0.5s later.
attempts += review_count_until(
&mut server,
start + Duration::from_millis(round * 5000 - 1000),
)
.await;
assert_eq!(rx.len(), 0);
server
.send(Message::Text(
json!(["EVENT", sub, fixture_event(ch, 10 + round, 9)])
.to_string()
.into(),
))
.await
.unwrap();
attempts += review_barrier(&mut server).await;
assert_eq!(rx.len(), 1);
attempts += review_count_until(
&mut server,
start + Duration::from_millis(round * 5000 + 500),
)
.await;
rx.recv().await.unwrap().unwrap();
}
println!("REVIEW periodic consumer: 4 full-at-tick windows, empty >=3.5s each period, recovery_requests={attempts}");
assert_eq!(
attempts, 1,
"recurring headroom must not strand the first attempt"
);
// Return the missing event using the actual requested stable subscription.
server
.send(Message::Text(
json!(["EVENT", sub, lost]).to_string().into(),
))
.await
.unwrap();
assert_eq!(review_barrier(&mut server).await, 0);
assert_eq!(rx.recv().await.unwrap().unwrap().event.id, lost.id);
let after = review_count_until(&mut server, start + Duration::from_millis(25_500)).await;
assert_eq!(
after, 0,
"no timer churn or extra requests after successful write"
);
println!(
"REVIEW continuous-headroom control: additional_requests={after}; missing event delivered"
);
cmd_tx.send(RelayCommand::Shutdown).await.unwrap();
timeout(Duration::from_secs(2), task)
.await
.unwrap()
.unwrap();
}
/// Same real socket owner as production, with no control of its internal state.
struct Owner {
server: WebSocketStream<tokio::net::TcpStream>,
rx: mpsc::Receiver<Option<BuzzEvent>>,
cmd: mpsc::Sender<RelayCommand>,
task: tokio::task::JoinHandle<()>,
ch: Uuid,
}
impl Owner {
async fn new(capacity: usize) -> Self {
let (client, server) = test_ws_pair().await;
let (tx, rx) = mpsc::channel(capacity);
let (control_tx, _control_rx) = mpsc::channel(1);
let (cmd, cmd_rx) = mpsc::channel(8);
let task = tokio::spawn(run_background_task(
client,
VecDeque::new(),
tx,
control_tx,
cmd_rx,
Keys::generate(),
"ws://127.0.0.1:1".into(),
"agent".into(),
None,
));
let mut owner = Self {
server,
rx,
cmd,
task,
ch: Uuid::new_v4(),
};
owner.subscribe().await;
owner
}
async fn subscribe(&mut self) {
self.cmd
.send(RelayCommand::Subscribe {
channel_id: self.ch,
filter: test_channel_filter(),
replay_since: Some(1000),
})
.await
.unwrap();
let req = self.request(Duration::from_secs(2)).await;
assert_eq!(req[1], channel_sub_id(self.ch));
}
async fn event(&mut self, n: u64) {
self.server
.send(Message::Text(
json!([
"EVENT",
channel_sub_id(self.ch),
fixture_event(self.ch, n, 9)
])
.to_string()
.into(),
))
.await
.unwrap();
}
async fn request(&mut self, duration: Duration) -> Value {
let frame = socket_frame(&mut self.server, duration)
.await
.expect("recovery not woken");
let req: Value = serde_json::from_str(frame.to_text().unwrap()).unwrap();
assert_eq!(req[0], "REQ");
req
}
async fn shutdown(self) {
self.cmd.send(RelayCommand::Shutdown).await.unwrap();
timeout(Duration::from_secs(1), self.task)
.await
.unwrap()
.unwrap();
}
}
#[tokio::test]
async fn capacity_flapping_cannot_storm_or_delay_an_allowed_attempt() {
let mut owner = Owner::new(1).await;
owner.event(0).await;
owner.event(1).await;
assert_eq!(review_barrier(&mut owner.server).await, 0);
owner.rx.recv().await.unwrap();
let first = owner.request(Duration::from_millis(500)).await;
let first_at = tokio::time::Instant::now();
assert_eq!(first[2]["since"], 996);
// New loss, then rapid full/empty transitions during the attempt cooldown.
// No socket activity or capacity transition may reset or bypass that bound.
for n in 1..=20 {
owner.event(n * 2).await;
owner.event(n * 2 + 1).await;
assert_eq!(review_barrier(&mut owner.server).await, 0);
owner.rx.recv().await.unwrap();
assert!(socket_frame(&mut owner.server, Duration::from_millis(10))
.await
.is_none());
}
let req = owner.request(Duration::from_secs(6)).await;
assert_eq!(req[1], channel_sub_id(owner.ch));
// Arrival timestamps approximate send completion; leave tolerance for TCP.
assert!(first_at.elapsed() >= Duration::from_millis(4_900));
assert!(first_at.elapsed() < Duration::from_secs(6));
assert_eq!(req[2]["since"], 998);
assert!(socket_frame(&mut owner.server, Duration::from_millis(100))
.await
.is_none());
owner.shutdown().await;
}
#[tokio::test]
async fn partial_capacity_wait_does_not_steal_live_slots_and_cancels_on_unsubscribe() {
let mut owner = Owner::new(5).await; // odd capacity: threshold rounds UP to 3
for n in 0..6 {
owner.event(n).await;
}
assert_eq!(review_barrier(&mut owner.server).await, 0);
owner.rx.recv().await.unwrap(); // partial reservation, insufficient for replay
assert!(socket_frame(&mut owner.server, Duration::from_millis(30))
.await
.is_none());
owner.event(6).await;
assert_eq!(review_barrier(&mut owner.server).await, 0);
assert_eq!(
owner.rx.len(),
5,
"select must release partial permits BEFORE try_send"
);
for _ in 0..2 {
owner.rx.recv().await.unwrap();
}
assert!(socket_frame(&mut owner.server, Duration::from_millis(30))
.await
.is_none());
owner.rx.recv().await.unwrap(); // exactly three slots free, prompt capacity wake
let req = owner.request(Duration::from_millis(500)).await;
assert_eq!(req[2]["since"], 1000);
assert_eq!(
owner
.rx
.recv()
.await
.unwrap()
.unwrap()
.event
.created_at
.as_secs(),
1004
);
assert_eq!(
owner
.rx
.recv()
.await
.unwrap()
.unwrap()
.event
.created_at
.as_secs(),
1006
);
// Wait out cooldown then create another pending loss and partial reservation.
tokio::time::sleep(recovery::RECOVERY_INTERVAL).await;
for n in 7..13 {
owner.event(n).await;
}
assert_eq!(review_barrier(&mut owner.server).await, 0);
owner.rx.recv().await.unwrap();
assert!(socket_frame(&mut owner.server, Duration::from_millis(30))
.await
.is_none());
owner
.cmd
.send(RelayCommand::Unsubscribe {
channel_id: owner.ch,
})
.await
.unwrap();
let close = socket_frame(&mut owner.server, Duration::from_secs(1))
.await
.unwrap();
let close: Value = serde_json::from_str(close.to_text().unwrap()).unwrap();
assert_eq!(close[0], "CLOSE");
while owner.rx.try_recv().is_ok() {}
assert!(socket_frame(&mut owner.server, Duration::from_millis(100))
.await
.is_none());
owner.subscribe().await;
assert!(socket_frame(&mut owner.server, Duration::from_millis(100))
.await
.is_none());
// A re-added intent can record and recover fresh loss immediately.
for n in 20..26 {
owner.event(n).await;
}
assert_eq!(review_barrier(&mut owner.server).await, 0);
while owner.rx.try_recv().is_ok() {}
let req = owner.request(Duration::from_millis(500)).await;
assert_eq!(req[2]["since"], 1020);
owner.shutdown().await;
}
#[tokio::test]
async fn shutdown_and_transport_loss_cancel_a_capacity_wait() {
let mut owner = Owner::new(1).await;
owner.event(0).await;
owner.event(1).await;
assert_eq!(review_barrier(&mut owner.server).await, 0);
// Full queue cannot block commands or processing an actual socket close.
owner.server.close(None).await.unwrap();
owner.shutdown().await;
let mut owner = Owner::new(1).await;
owner.event(0).await;
owner.event(1).await;
assert_eq!(review_barrier(&mut owner.server).await, 0);
owner.shutdown().await;
}
#[tokio::test]
async fn readiness_gate_ownership_and_attempt_deadlines_are_not_polling_ticks() {
let (mut client, mut server) = test_ws_pair().await;
let mut state = BgState::new();
let ch = Uuid::new_v4();
seed_test_subscription(&mut state, ch);
let (tx, mut rx) = mpsc::channel(4);
assert!(recovery::ready_at(&mut state).is_none());
state.channel_dropped_since.insert(ch, 700);
state
.rate_limited_pending
.insert(ch, tokio::time::Instant::now());
assert!(recovery::ready_at(&mut state).is_none());
state.rate_limited_pending.clear();
state.resubscribe_retry.insert(ch);
assert!(recovery::ready_at(&mut state).is_none());
state.resubscribe_retry.clear();
state.rate_limit_gate = Some(tokio::time::Instant::now() + Duration::from_millis(80));
let at = recovery::ready_at(&mut state);
assert_eq!(at, state.rate_limit_gate);
assert!(timeout(Duration::from_millis(20), recovery::ready(&tx, at))
.await
.is_err());
timeout(Duration::from_millis(200), recovery::ready(&tx, at))
.await
.unwrap();
recovery::recover_one(&mut client, &mut state, &tx, "agent").await;
next_test_frame(&mut server).await;
state.channel_dropped_since.insert(ch, 800);
let at = recovery::ready_at(&mut state).unwrap();
let remaining = at.saturating_duration_since(tokio::time::Instant::now());
assert!(remaining > Duration::from_millis(4900));
// Neither new loss nor an intervening gate shorter than cooldown delays it.
state.rate_limit_gate = Some(tokio::time::Instant::now() + Duration::from_millis(10));
assert_eq!(recovery::ready_at(&mut state), Some(at));
state.rate_limit_gate = Some(at + Duration::from_secs(1));
assert_eq!(recovery::ready_at(&mut state), state.rate_limit_gate);
advance_clock(Duration::from_secs(6)).await;
// Cancellation releases partial permits. No hidden reservation survives it.
for _ in 0..3 {
tx.try_send(None).unwrap();
}
assert!(timeout(
Duration::from_millis(20),
recovery::ready(&tx, recovery::ready_at(&mut state))
)
.await
.is_err());
assert_eq!(tx.capacity(), 1);
rx.recv().await.unwrap();
timeout(
Duration::from_millis(100),
recovery::ready(&tx, recovery::ready_at(&mut state)),
)
.await
.unwrap();
assert_eq!(tx.capacity(), 2);
drop(rx);
assert!(timeout(
Duration::from_millis(20),
recovery::ready(&tx, recovery::ready_at(&mut state))
)
.await
.is_err());
}
#[tokio::test]
async fn readiness_uses_channel_wakes_without_idle_churn_or_lost_capacity() {
use std::future::Future;
use std::sync::{
atomic::{AtomicUsize, Ordering},
Arc,
};
use std::task::{Context, Wake, Waker};
#[derive(Default)]
struct Wakes(AtomicUsize);
impl Wake for Wakes {
fn wake(self: Arc<Self>) {
self.0.fetch_add(1, Ordering::SeqCst);
}
fn wake_by_ref(self: &Arc<Self>) {
self.0.fetch_add(1, Ordering::SeqCst);
}
}
let wakes = Arc::new(Wakes::default());
let waker = Waker::from(wakes.clone());
let mut cx = Context::from_waker(&waker);
let (tx, mut rx) = mpsc::channel(4);
for _ in 0..4 {
tx.try_send(None).unwrap();
}
let mut idle = Box::pin(recovery::ready(&tx, None));
assert!(idle.as_mut().poll(&mut cx).is_pending());
rx.recv().await.unwrap();
rx.recv().await.unwrap();
assert_eq!(
wakes.0.load(Ordering::SeqCst),
0,
"no loss: no capacity subscription"
);
drop(idle);
// Capacity freed before registration cannot be lost; ready on first poll.
let now = Some(tokio::time::Instant::now());
let mut ready = Box::pin(recovery::ready(&tx, now));
assert!(ready.as_mut().poll(&mut cx).is_ready());
assert_eq!(tx.capacity(), 2, "successful readiness returns all permits");
drop(ready);
for _ in 0..2 {
tx.try_send(None).unwrap();
}
let mut ready = Box::pin(recovery::ready(&tx, now));
assert!(ready.as_mut().poll(&mut cx).is_pending());
assert_eq!(wakes.0.load(Ordering::SeqCst), 0);
rx.recv().await.unwrap();
assert_eq!(
wakes.0.load(Ordering::SeqCst),
0,
"below threshold: do not wake"
);
rx.recv().await.unwrap();
assert_eq!(
wakes.0.load(Ordering::SeqCst),
1,
"threshold: channel wakes its waiter"
);
assert!(ready.as_mut().poll(&mut cx).is_ready());
assert_eq!(tx.capacity(), 2);
drop(ready);
drop(rx);
let mut closed = Box::pin(recovery::ready(&tx, now));
let before = wakes.0.load(Ordering::SeqCst);
assert!(closed.as_mut().poll(&mut cx).is_pending());
assert_eq!(
wakes.0.load(Ordering::SeqCst),
before,
"closed: no self-wake loop"
);
}