fix(#1462): buffer transient wakes so they survive when no recv is parked

This commit is contained in:
damocles 2026-06-08 20:40:40 +02:00 committed by mara
commit 8b66bf78df

View file

@ -1,7 +1,7 @@
//! Sqlite-backed message broker. Survives `hive-c0re` restart, and taps every
//! send/recv onto a broadcast channel so the dashboard can stream it.
use std::collections::{HashMap, HashSet};
use std::collections::{HashMap, HashSet, VecDeque};
use std::path::Path;
use std::sync::Mutex;
use std::time::{SystemTime, UNIX_EPOCH};
@ -147,6 +147,13 @@ struct RecipientInflight {
requeued_ids: HashSet<i64>,
}
/// Hard cap on buffered transient pings per recipient. Pings are
/// ephemeral wakes (bash completions, matrix events); an agent drains
/// them on its next `recv`, so the buffer is normally near-empty. The
/// cap only bounds memory if a recipient stops recv'ing entirely (e.g.
/// a stopped container) — past it the oldest ping is dropped.
const MAX_PENDING_PINGS: usize = 256;
pub struct Broker {
conn: Mutex<Connection>,
events: broadcast::Sender<MessageEvent>,
@ -155,6 +162,15 @@ pub struct Broker {
/// boot, which rebuilds the `requeued_ids` set from the DB and
/// clears any stale `unacked_ids`).
inflight: Mutex<HashMap<String, RecipientInflight>>,
/// Per-recipient buffer of transient wake signals (`ping`) that
/// haven't been consumed by a `recv` yet. The broadcast channel only
/// reaches a `recv_blocking_batch` that is *currently parked*; a ping
/// fired while the harness is mid-turn (no live subscriber) would
/// otherwise be lost. Buffering here makes a transient wake reliable
/// — the next `recv` drains it — without sqlite persistence or
/// redelivery-on-restart (the buffer is in-memory, dropped on
/// hive-c0re restart, which is fine: a wake that old is stale).
pending_pings: Mutex<HashMap<String, VecDeque<(String, String)>>>,
}
impl Broker {
@ -173,6 +189,7 @@ impl Broker {
conn: Mutex::new(conn),
events,
inflight: Mutex::new(HashMap::new()),
pending_pings: Mutex::new(HashMap::new()),
})
}
@ -201,11 +218,25 @@ impl Broker {
}
/// Deliver a transient wake signal to `to` without writing to sqlite.
/// The signal wakes a long-polling `recv_blocking_batch` for the target
/// agent but is not persisted, not redelivered on restart, and not shown
/// in message history. Use for ephemeral notifications (bash task
/// completions) where persistence would cause duplicate delivery.
/// The signal wakes the target agent's next `recv_blocking_batch` but is
/// not persisted, not redelivered on restart, and not shown in message
/// history. Use for ephemeral notifications (bash task completions,
/// matrix events) where persistence would cause duplicate delivery.
///
/// The ping is buffered in `pending_pings` *before* the broadcast so a
/// `recv` that wasn't parked at fire time still drains it on its next
/// call — the broadcast alone reaches only a currently-parked receiver.
pub fn ping(&self, to: &str, from: &str, body: &str) {
{
let mut pending = self.pending_pings.lock().unwrap();
let queue = pending.entry(to.to_owned()).or_default();
queue.push_back((from.to_owned(), body.to_owned()));
// Bound memory if the recipient never recv's (stopped container):
// drop the oldest so the most recent wakes survive.
while queue.len() > MAX_PENDING_PINGS {
queue.pop_front();
}
}
let _ = self.events.send(MessageEvent::Ping {
to: to.to_owned(),
from: from.to_owned(),
@ -213,6 +244,50 @@ impl Broker {
});
}
/// Drain up to `max` buffered transient pings for `recipient`, mapping
/// each to a `Delivery` with the sentinel `id = 0` (never pushed to
/// `unacked_ids`, so `ack_turn` ignores it). FIFO; leaves any remainder
/// buffered for the next call.
fn drain_pings(&self, recipient: &str, max: usize) -> Vec<Delivery> {
if max == 0 {
return Vec::new();
}
let mut pending = self.pending_pings.lock().unwrap();
let Some(queue) = pending.get_mut(recipient) else {
return Vec::new();
};
let take = queue.len().min(max);
let drained: Vec<Delivery> = queue
.drain(..take)
.map(|(from, body)| Delivery {
id: 0,
redelivered: false,
message: Message {
from,
to: recipient.to_owned(),
body,
in_reply_to: None,
},
})
.collect();
if queue.is_empty() {
pending.remove(recipient);
}
drained
}
/// Collect a batch for `recipient`: buffered transient pings first
/// (they're wakes — surface them ahead of queued mail), then sqlite
/// messages up to the remaining budget. Shared by the immediate check
/// and the post-wake path of `recv_blocking_batch`.
fn collect_batch(&self, recipient: &str, max: usize) -> Result<Vec<Delivery>> {
let mut batch = self.drain_pings(recipient, max);
if batch.len() < max {
batch.extend(self.recv_batch(recipient, max - batch.len())?);
}
Ok(batch)
}
/// Latest `limit` messages addressed to `recipient`, newest-first.
/// Includes delivered + undelivered alike — used for the operator
/// inbox view on the dashboard. Caller decides what to show.
@ -440,7 +515,9 @@ impl Broker {
return Ok(Vec::new());
}
let mut rx = self.subscribe();
let batch = self.recv_batch(recipient, max)?;
// Immediate check: buffered pings (incl. any fired while no recv was
// parked) + queued sqlite mail.
let batch = self.collect_batch(recipient, max)?;
if !batch.is_empty() {
return Ok(batch);
}
@ -452,37 +529,21 @@ impl Broker {
};
match tokio::time::timeout(remaining, rx.recv()).await {
Err(_) => return Ok(Vec::new()),
// Channel lagged or closed — fall back to a single direct
// pop (in case we missed our notification while behind).
Ok(Err(_)) => return self.recv_batch(recipient, max),
Ok(Ok(MessageEvent::Sent { to, .. })) if to == recipient => {
let batch = self.recv_batch(recipient, max)?;
// Channel lagged or closed — fall back to a direct collect
// (in case we missed our notification while behind).
Ok(Err(_)) => return self.collect_batch(recipient, max),
// A relevant event landed (real message or transient ping).
// Re-collect from the buffers; the ping payload on the event
// is ignored — `drain_pings` is the source of truth, so a
// ping can't be double-delivered.
Ok(Ok(MessageEvent::Sent { to, .. } | MessageEvent::Ping { to, .. }))
if to == recipient =>
{
let batch = self.collect_batch(recipient, max)?;
if !batch.is_empty() {
return Ok(batch);
}
// Lost a race (concurrent recv elsewhere). Keep waiting.
}
// Transient ping — not sqlite-backed. Return it directly as
// a Delivery with id=0 (sentinel: never pushed to unacked_ids
// so ack_turn silently ignores it).
Ok(Ok(MessageEvent::Ping { to, from, body })) if to == recipient => {
// Also drain any real sqlite messages that may have landed
// concurrently; prepend the ping so the agent sees both.
let mut batch = self.recv_batch(recipient, max.saturating_sub(1))?;
batch.insert(
0,
Delivery {
id: 0,
redelivered: false,
message: Message {
from,
to,
body,
in_reply_to: None,
},
},
);
return Ok(batch);
// Lost a race (concurrent recv drained it). Keep waiting.
}
Ok(Ok(_)) => {}
}
@ -1291,6 +1352,47 @@ mod tests {
assert_eq!(broker.ack_turn("b").unwrap(), 5);
}
/// The #1462 fix: a transient `ping` fired while no `recv` is parked
/// must NOT be lost — it's buffered and drained by the next collect.
#[test]
fn transient_ping_buffered_when_no_receiver_parked() {
let h = open_broker();
let broker = &h.broker;
// Nobody is parked on recv when these fire (the old broadcast-only
// path dropped them here).
broker.ping("b", "matrix", "wake1");
broker.ping("b", "bash-task", "wake2");
let batch = broker.collect_batch("b", 10).unwrap();
let got: Vec<_> = batch
.iter()
.map(|d| (d.id, d.message.from.as_str(), d.message.body.as_str()))
.collect();
assert_eq!(
got,
vec![(0, "matrix", "wake1"), (0, "bash-task", "wake2")],
"buffered pings drain FIFO as id=0 deliveries"
);
// Drained — the buffer is now empty.
assert!(broker.collect_batch("b", 10).unwrap().is_empty());
// Pings are sentinel id=0, so ack_turn has nothing to close out.
assert_eq!(broker.ack_turn("b").unwrap(), 0);
}
/// Buffered pings surface ahead of queued sqlite mail in one batch.
#[test]
fn pending_pings_surface_before_queued_mail() {
let h = open_broker();
let broker = &h.broker;
broker.send(&msg("a", "b", "mail")).unwrap();
broker.ping("b", "matrix", "wake");
let batch = broker.collect_batch("b", 10).unwrap();
assert_eq!(batch.len(), 2);
assert_eq!(batch[0].id, 0, "ping first");
assert_eq!(batch[0].message.body, "wake");
assert!(batch[1].id > 0, "sqlite mail after");
assert_eq!(batch[1].message.body, "mail");
}
/// `recv_batch` with no pending traffic returns an empty vec
/// (the "(empty)" path), not an error.
#[test]