fix(pubsub): reconcile dropped subscription hints

Signed-off-by: npub13n66s06epmqf2kc3v373ez8hj65cuzyvxzjf93vwpervxqn2u7jq2qd9je <8cf5a83f590ec0955b11647d1c88f796a98e088c30a492c58e0e46c3026ae7a4@buzz.block.builderlab.xyz>
Signed-off-by: Larry <8cf5a83f590ec0955b11647d1c88f796a98e088c30a492c58e0e46c3026ae7a4@buzz.block.builderlab.xyz>
Signed-off-by: loganj <loganj@squareup.com>
This commit is contained in:
npub13n66s06epmqf2kc3v373ez8hj65cuzyvxzjf93vwpervxqn2u7jq2qd9je
2026-07-28 10:44:23 -04:00
committed by loganj
parent 58a871affb
commit 7a6de053c8
2 changed files with 106 additions and 8 deletions

View File

@@ -49,7 +49,7 @@ use std::time::Duration;
use buzz_core::TenantContext;
use nostr::PublicKey;
use tokio::sync::{broadcast, mpsc, Mutex};
use tokio::sync::{broadcast, mpsc, Mutex, Notify};
use crate::cache_invalidation::{
cache_invalidation_channel, CacheInvalidation, ScopedCacheInvalidation,
@@ -96,6 +96,9 @@ impl PubSubConfig {
}
}
/// Capacity for live subscription hints. Desired state is reconciled separately.
const SUBSCRIPTION_COMMAND_CAPACITY: usize = 4096;
/// Central pub/sub manager for a Buzz relay instance.
pub struct PubSubManager {
pool: deadpool_redis::Pool,
@@ -107,6 +110,8 @@ pub struct PubSubManager {
desired_topics: subscriber::DesiredTopics,
subscription_tx: mpsc::Sender<subscriber::SubscriptionCommand>,
subscription_rx: Mutex<Option<mpsc::Receiver<subscriber::SubscriptionCommand>>>,
/// Coalesced wake-up for reconciling desired topics with the live Redis connection.
subscription_changed: Arc<Notify>,
broadcast_tx: broadcast::Sender<ChannelEvent>,
cache_invalidation_tx: broadcast::Sender<ScopedCacheInvalidation>,
conn_control_tx: broadcast::Sender<ScopedConnControl>,
@@ -126,7 +131,7 @@ impl PubSubManager {
let (broadcast_tx, _) = broadcast::channel(4096);
let (cache_invalidation_tx, _) = broadcast::channel(4096);
let (conn_control_tx, _) = broadcast::channel(4096);
let (subscription_tx, subscription_rx) = mpsc::channel(4096);
let (subscription_tx, subscription_rx) = mpsc::channel(SUBSCRIPTION_COMMAND_CAPACITY);
Ok(Self {
pool,
@@ -135,6 +140,7 @@ impl PubSubManager {
desired_topics: Arc::new(Mutex::new(HashMap::new())),
subscription_tx,
subscription_rx: Mutex::new(Some(subscription_rx)),
subscription_changed: Arc::new(Notify::new()),
broadcast_tx,
cache_invalidation_tx,
conn_control_tx,
@@ -156,6 +162,7 @@ impl PubSubManager {
self.broadcast_tx.clone(),
self.desired_topics.clone(),
subscription_rx,
self.subscription_changed.clone(),
)
.await;
}
@@ -200,17 +207,20 @@ impl PubSubManager {
};
if should_subscribe {
// `desired_topics` is the source of truth. A full command queue is
// only a stale live-connection hint: reconnect snapshots desired
// topics, while blocking here can deadlock REQ/CLOSE lifecycle
// commits during a Redis outage.
// This wake-up is lossless and coalesced: if the bounded command
// queue is full, the subscriber still reconciles the complete
// desired-topic snapshot without waiting for a Redis reconnect.
self.subscription_changed.notify_one();
match self
.subscription_tx
.try_send(subscriber::SubscriptionCommand::Subscribe(topic_key))
{
Ok(()) | Err(mpsc::error::TrySendError::Closed(_)) => {}
Err(mpsc::error::TrySendError::Full(_)) => {
tracing::warn!(?topic_key, "pubsub subscription command queue full");
tracing::warn!(
?topic_key,
"pubsub subscription command queue full; reconciliation scheduled"
);
}
}
}
@@ -386,6 +396,8 @@ pub(crate) mod test_util {
#[cfg(test)]
mod tests {
use std::collections::HashSet;
use super::*;
use crate::test_util::make_test_pool;
use buzz_core::{CommunityId, TenantContext};
@@ -598,6 +610,38 @@ mod tests {
assert_eq!(manager.topic_refcount(&ctx_b, topic).await, 0);
}
#[tokio::test]
async fn queue_full_retain_schedules_live_reconciliation() {
let manager = make_manager().await;
let filler_ctx = ctx(0xaaaa, "a.example");
for id in 0..SUBSCRIPTION_COMMAND_CAPACITY {
manager
.retain_topic(
&filler_ctx,
EventTopic::Channel(Uuid::from_u128(id as u128)),
)
.await;
}
let dropped_ctx = ctx(0xbbbb, "b.example");
let dropped_topic = EventTopic::Channel(Uuid::from_u128(0xcccc));
manager.retain_topic(&dropped_ctx, dropped_topic).await;
assert_eq!(manager.topic_refcount(&dropped_ctx, dropped_topic).await, 1);
tokio::time::timeout(
Duration::from_millis(10),
manager.subscription_changed.notified(),
)
.await
.expect("queue-full retain must leave a reconciliation wake-up");
let topic_key = EventTopicKey::from_context(&dropped_ctx, dropped_topic);
let missing =
subscriber::missing_desired_topics(&manager.desired_topics, &HashSet::new()).await;
assert!(missing.contains(&topic_key));
}
#[tokio::test]
async fn retain_release_refcounts_and_debounces_last_release() {
let pool = make_test_pool();

View File

@@ -6,7 +6,7 @@ use std::time::Duration;
use futures_util::StreamExt;
use nostr::JsonUtil;
use tokio::sync::{broadcast, mpsc, Mutex};
use tokio::sync::{broadcast, mpsc, Mutex, Notify};
use crate::topic::EventTopicKey;
use crate::ChannelEvent;
@@ -39,6 +39,7 @@ pub(crate) async fn run_subscriber(
broadcast_tx: broadcast::Sender<ChannelEvent>,
desired_topics: DesiredTopics,
mut subscription_rx: mpsc::Receiver<SubscriptionCommand>,
subscription_changed: Arc<Notify>,
) {
let mut backoff_secs = BACKOFF_INITIAL_SECS;
@@ -48,6 +49,7 @@ pub(crate) async fn run_subscriber(
&broadcast_tx,
desired_topics.clone(),
&mut subscription_rx,
subscription_changed.clone(),
)
.await
{
@@ -77,6 +79,7 @@ async fn connect_and_subscribe(
broadcast_tx: &broadcast::Sender<ChannelEvent>,
desired_topics: DesiredTopics,
subscription_rx: &mut mpsc::Receiver<SubscriptionCommand>,
subscription_changed: Arc<Notify>,
) -> Result<(), redis::RedisError> {
let client = redis::Client::open(redis_url)?;
let conn = client.get_async_pubsub().await?;
@@ -104,6 +107,16 @@ async fn connect_and_subscribe(
loop {
tokio::select! {
_ = subscription_changed.notified() => {
// Retains publish a coalesced wake-up in addition to their
// best-effort queue hint. Re-read the source of truth so a
// queue-full hint cannot leave this live connection stale.
for topic in missing_desired_topics(&desired_topics, &active_topics).await {
let channel = topic.redis_channel();
sink.subscribe(&channel).await?;
active_topics.insert(channel);
}
}
Some(command) = subscription_rx.recv() => {
match command {
SubscriptionCommand::Subscribe(topic) => {
@@ -171,6 +184,20 @@ async fn connect_and_subscribe(
}
}
pub(crate) async fn missing_desired_topics(
desired_topics: &DesiredTopics,
active_topics: &HashSet<String>,
) -> Vec<EventTopicKey> {
desired_topics
.lock()
.await
.iter()
.filter_map(|(topic, count)| {
(*count > 0 && !active_topics.contains(&topic.redis_channel())).then_some(*topic)
})
.collect()
}
async fn desired_refcount(desired_topics: &DesiredTopics, topic: EventTopicKey) -> usize {
desired_topics
.lock()
@@ -191,6 +218,33 @@ mod tests {
EventTopicKey::from_context(&ctx, crate::EventTopic::Global)
}
#[tokio::test]
async fn reconciliation_finds_desired_topic_when_bounded_queue_is_full() {
let desired = Arc::new(Mutex::new(HashMap::new()));
let (tx, _rx) = mpsc::channel(1);
let queued_topic = topic(1);
let dropped_topic = topic(2);
tx.try_send(SubscriptionCommand::Subscribe(queued_topic))
.expect("fill command queue");
desired.lock().await.insert(dropped_topic, 1);
let subscription_changed = Notify::new();
subscription_changed.notify_one();
assert!(matches!(
tx.try_send(SubscriptionCommand::Subscribe(dropped_topic)),
Err(mpsc::error::TrySendError::Full(_))
));
tokio::time::timeout(Duration::from_millis(10), subscription_changed.notified())
.await
.expect("queue-full retain must leave a reconciliation wake-up");
let missing = missing_desired_topics(&desired, &HashSet::new()).await;
assert_eq!(missing, vec![dropped_topic]);
let active = HashSet::from([dropped_topic.redis_channel()]);
assert!(missing_desired_topics(&desired, &active).await.is_empty());
}
#[tokio::test]
async fn desired_refcount_returns_zero_for_absent_topic() {
let desired = Arc::new(Mutex::new(HashMap::new()));