fix(relay): cancel pending subscriptions on close

crates/buzz-pubsub/src/lib.rs

@@ -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,10 +207,22 @@ impl PubSubManager {
         };
 
         if should_subscribe {
-            let _ = self
+            // 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
-                .send(subscriber::SubscriptionCommand::Subscribe(topic_key))
-                .await;
+                .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; reconciliation scheduled"
+                    );
+                }
+            }
         }
     }
 
@@ -377,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};
@@ -589,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();

crates/buzz-pubsub/src/subscriber.rs

@@ -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()));