mirror of
https://github.com/block/buzz.git
synced 2026-10-31 06:21:34 +00:00
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:
committed by
loganj
parent
58a871affb
commit
7a6de053c8
@@ -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();
|
||||
|
||||
@@ -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()));
|
||||
|
||||
Reference in New Issue
Block a user