diff --git a/crates/buzz-pubsub/src/lib.rs b/crates/buzz-pubsub/src/lib.rs index 4f1690beefb..a439de4e6e5 100644 --- a/crates/buzz-pubsub/src/lib.rs +++ b/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, subscription_rx: Mutex>>, + /// Coalesced wake-up for reconciling desired topics with the live Redis connection. + subscription_changed: Arc, broadcast_tx: broadcast::Sender, cache_invalidation_tx: broadcast::Sender, conn_control_tx: broadcast::Sender, @@ -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(); diff --git a/crates/buzz-pubsub/src/subscriber.rs b/crates/buzz-pubsub/src/subscriber.rs index 88826ed99be..14c49d261c9 100644 --- a/crates/buzz-pubsub/src/subscriber.rs +++ b/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, desired_topics: DesiredTopics, mut subscription_rx: mpsc::Receiver, + subscription_changed: Arc, ) { 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, desired_topics: DesiredTopics, subscription_rx: &mut mpsc::Receiver, + subscription_changed: Arc, ) -> 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, +) -> Vec { + 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())); diff --git a/crates/buzz-relay/src/connection.rs b/crates/buzz-relay/src/connection.rs index 5fcfe70b91c..45ab439c874 100644 --- a/crates/buzz-relay/src/connection.rs +++ b/crates/buzz-relay/src/connection.rs @@ -2,13 +2,14 @@ use std::collections::HashMap; use std::net::SocketAddr; -use std::sync::atomic::{AtomicU8, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU8, Ordering}; use std::sync::Arc; use std::time::Duration; use axum::extract::ws::{Message as WsMessage, WebSocket}; use futures_util::{Sink, SinkExt, StreamExt}; use tokio::sync::{mpsc, watch, Mutex, RwLock}; +use tokio::task::JoinSet; use tokio_util::sync::CancellationToken; use tracing::Instrument as _; use tracing::{debug, info, trace, warn}; @@ -36,6 +37,94 @@ pub(crate) struct RestartClose { pub(crate) flushed: tokio::sync::oneshot::Sender, } +/// A cancellable lease for one in-flight REQ. The lease is installed before +/// the handler task starts so CLOSE can invalidate a subscription that has not +/// reached registration yet. +#[derive(Debug)] +pub(crate) struct PendingSubscription { + cancel: CancellationToken, + committed: AtomicBool, + completed: AtomicBool, + predecessor: Option>, +} + +impl PendingSubscription { + fn new(predecessor: Option>) -> Self { + Self { + cancel: CancellationToken::new(), + committed: AtomicBool::new(false), + completed: AtomicBool::new(false), + predecessor, + } + } + + pub(crate) fn cancel(&self) { + self.cancel.cancel(); + } + + /// Cancel this request and every older same-ID request it superseded. + pub(crate) fn cancel_lineage(&self) { + let mut current = Some(self); + while let Some(request) = current { + request.cancel(); + current = request.predecessor.as_deref(); + } + } + + /// Supersede older same-ID requests after this replacement commits. + pub(crate) fn commit(&self) { + self.committed.store(true, Ordering::Release); + if let Some(predecessor) = self.predecessor.as_ref() { + predecessor.cancel_lineage(); + } + } + + /// Return whether `candidate` is allowed to commit while this request owns + /// the subscription ID. A newer same-ID REQ may take map ownership while + /// an older request is still doing validation or historical setup. Until + /// the newer request commits, any uncancelled request in its predecessor + /// lineage may still establish the subscription it started. + pub(crate) fn permits_commit( + current: &Arc, + candidate: &Arc, + ) -> bool { + let mut request = Some(Arc::clone(current)); + while let Some(request_in_lineage) = request { + if Arc::ptr_eq(&request_in_lineage, candidate) { + return !candidate.is_cancelled(); + } + request = request_in_lineage.predecessor.clone(); + } + false + } + + fn live_predecessor(&self) -> Option> { + if self.committed.load(Ordering::Acquire) { + return None; + } + + let mut current = self.predecessor.clone(); + while let Some(predecessor) = current { + if !predecessor.is_cancelled() && !predecessor.completed.load(Ordering::Acquire) { + return Some(predecessor); + } + current = predecessor.predecessor.clone(); + } + None + } + + pub(crate) fn is_cancelled(&self) -> bool { + self.cancel.is_cancelled() + } + + pub(crate) async fn cancelled(&self) { + self.cancel.cancelled().await; + } +} + +/// In-flight REQs keyed by client-supplied subscription ID. +pub(crate) type PendingSubscriptions = Arc>>>; + /// Maximum outbound data frames buffered into the websocket sink before one flush. const MAX_WS_SEND_BATCH: usize = 64; @@ -70,6 +159,8 @@ pub struct ConnectionState { pub auth_state: RwLock, /// Active subscriptions keyed by subscription ID. pub subscriptions: ConnectionSubscriptions, + /// REQs that have started but may not yet be registered for fan-out. + pub(crate) pending_subscriptions: PendingSubscriptions, /// Sender for outbound data messages (EVENT, NOTICE, OK, etc.). pub send_tx: mpsc::Sender, /// Sender for outbound control frames (Pong, Close). @@ -87,6 +178,86 @@ pub struct ConnectionState { } impl ConnectionState { + /// Install a pending REQ before its handler task starts. A repeated REQ + /// with the same subscription ID becomes the new owner, but the previous + /// request stays alive until the replacement commits successfully. + pub(crate) async fn begin_pending_subscription( + &self, + sub_id: &str, + ) -> Arc { + let mut pending_subscriptions = self.pending_subscriptions.lock().await; + let predecessor = pending_subscriptions.get(sub_id).cloned(); + let pending = Arc::new(PendingSubscription::new(predecessor)); + pending_subscriptions.insert(sub_id.to_owned(), Arc::clone(&pending)); + pending + } + + /// Cancel every in-flight REQ before connection registry cleanup begins. + async fn cancel_all_pending_subscriptions(&self) { + let mut pending = self.pending_subscriptions.lock().await; + for (_, request) in pending.drain() { + request.cancel_lineage(); + } + } + + /// Forget a completed REQ only if it still owns this subscription ID. + pub(crate) async fn finish_pending_subscription( + &self, + sub_id: &str, + completed: &Arc, + ) { + completed.completed.store(true, Ordering::Release); + let mut pending = self.pending_subscriptions.lock().await; + if pending + .get(sub_id) + .is_some_and(|current| Arc::ptr_eq(current, completed)) + { + if let Some(predecessor) = completed.live_predecessor() { + pending.insert(sub_id.to_owned(), predecessor); + } else { + pending.remove(sub_id); + } + } + } + + /// Queue output owned by one REQ only while its lease remains live. + /// + /// The lifecycle lock covers the complete batch so a paired NOTICE/CLOSED + /// response cannot straddle a same-ID replacement or CLOSE cutover. The + /// same fence also covers historical EVENT/EOSE output. + pub(crate) async fn send_req_messages_if_permitted( + &self, + sub_id: &str, + pending: &Arc, + messages: I, + ) -> bool + where + I: IntoIterator, + { + let pending_subscriptions = self.pending_subscriptions.lock().await; + let permitted = pending_subscriptions + .get(sub_id) + .is_some_and(|current| PendingSubscription::permits_commit(current, pending)); + if self.cancel.is_cancelled() || !permitted { + return false; + } + + // Keep the lifecycle lock through every queue insertion. Replacement + // commit and CLOSE take the same lock, so neither can publish its + // cutover and then be followed by output from this lease. + messages.into_iter().all(|msg| self.send(msg)) + } + + pub(crate) async fn send_historical_if_permitted( + &self, + sub_id: &str, + pending: &Arc, + msg: String, + ) -> bool { + self.send_req_messages_if_permitted(sub_id, pending, [msg]) + .await + } + /// Sends a data message to this connection's outbound channel. /// /// On a full buffer, increments the backpressure counter. The first @@ -178,6 +349,7 @@ async fn handle_active_connection( let backpressure_count = Arc::new(AtomicU8::new(0)); let subscriptions = Arc::new(Mutex::new(HashMap::new())); + let pending_subscriptions = Arc::new(Mutex::new(HashMap::new())); let conn = Arc::new(ConnectionState { conn_id, @@ -187,6 +359,7 @@ async fn handle_active_connection( challenge: challenge.clone(), }), subscriptions: Arc::clone(&subscriptions), + pending_subscriptions, send_tx: tx.clone(), ctrl_tx: ctrl_tx.clone(), cancel: cancel.clone(), @@ -470,6 +643,8 @@ async fn recv_loop( missed_pongs: Arc, cancel: CancellationToken, ) { + let mut req_tasks = JoinSet::new(); + loop { tokio::select! { msg = ws_recv.next() => { @@ -491,7 +666,12 @@ async fn recv_loop( break; } trace!(len = text.len(), "frame received"); - handle_text_message(text.to_string(), Arc::clone(&conn), Arc::clone(&state)).await; + handle_text_message( + text.to_string(), + Arc::clone(&conn), + Arc::clone(&state), + &mut req_tasks, + ).await; } Some(Ok(WsMessage::Binary(bytes))) => { let max_frame_bytes = state.config.max_frame_bytes; @@ -513,7 +693,12 @@ async fn recv_loop( // (notably certain Nostr libraries) send text payloads in binary frames. // NIP-01 is text-only, but accepting binary is a common relay extension. if let Ok(text) = String::from_utf8(bytes.to_vec()) { - handle_text_message(text, Arc::clone(&conn), Arc::clone(&state)).await; + handle_text_message( + text, + Arc::clone(&conn), + Arc::clone(&state), + &mut req_tasks, + ).await; } } Some(Ok(WsMessage::Pong(_))) => { @@ -539,12 +724,31 @@ async fn recv_loop( } } } + completed = req_tasks.join_next(), if !req_tasks.is_empty() => { + if let Some(Err(error)) = completed { + debug!(conn_id = %conn.conn_id, %error, "REQ handler task failed"); + } + } _ = cancel.cancelled() => break, } } + + shutdown_pending_requests(&conn, &mut req_tasks).await; } -async fn handle_text_message(text: String, conn: Arc, state: Arc) { +pub(crate) async fn shutdown_pending_requests(conn: &ConnectionState, req_tasks: &mut JoinSet<()>) { + conn.cancel.cancel(); + conn.cancel_all_pending_subscriptions().await; + req_tasks.abort_all(); + while req_tasks.join_next().await.is_some() {} +} + +async fn handle_text_message( + text: String, + conn: Arc, + state: Arc, + req_tasks: &mut JoinSet<()>, +) { let msg = match ClientMessage::parse(&text) { Ok(m) => m, Err(e) => { @@ -606,10 +810,22 @@ async fn handle_text_message(text: String, conn: Arc, state: Ar return; } }; + let pending = conn.begin_pending_subscription(&sub_id).await; let span = tracing::info_span!("ws.req", conn_id = %conn.conn_id, sub_id = %sub_id); - tokio::spawn( + req_tasks.spawn( async move { - handlers::req::handle_req(sub_id, filters, conn, state).await; + tokio::select! { + biased; + _ = pending.cancelled() => {} + _ = handlers::req::handle_req( + sub_id.clone(), + filters, + Arc::clone(&conn), + state, + Arc::clone(&pending), + ) => {} + } + conn.finish_pending_subscription(&sub_id, &pending).await; drop(permit); } .instrument(span), @@ -741,6 +957,291 @@ mod tests { use super::*; use std::sync::{Arc, Mutex}; + #[tokio::test] + async fn failed_same_id_replacement_restores_the_previous_lease() { + let conn = ConnectionState { + conn_id: Uuid::new_v4(), + tenant: TenantContext::resolved( + buzz_core::CommunityId::from_uuid(Uuid::new_v4()), + "lease.test", + ), + remote_addr: "127.0.0.1:1234".parse().expect("socket address"), + auth_state: RwLock::new(AuthState::Failed), + subscriptions: Arc::new(tokio::sync::Mutex::new(HashMap::new())), + pending_subscriptions: Arc::new(tokio::sync::Mutex::new(HashMap::new())), + send_tx: mpsc::channel(1).0, + ctrl_tx: mpsc::channel(1).0, + cancel: CancellationToken::new(), + backpressure_count: Arc::new(AtomicU8::new(0)), + grace_limit: 3, + }; + + let original = conn.begin_pending_subscription("same-id").await; + let replacement = conn.begin_pending_subscription("same-id").await; + assert!(!original.is_cancelled()); + + conn.finish_pending_subscription("same-id", &replacement) + .await; + let current = conn + .pending_subscriptions + .lock() + .await + .get("same-id") + .cloned() + .expect("original lease restored"); + assert!(Arc::ptr_eq(¤t, &original)); + assert!(!original.is_cancelled()); + } + + #[tokio::test] + async fn failed_replacement_does_not_restore_a_completed_predecessor() { + let conn = ConnectionState { + conn_id: Uuid::new_v4(), + tenant: TenantContext::resolved( + buzz_core::CommunityId::from_uuid(Uuid::new_v4()), + "lease.test", + ), + remote_addr: "127.0.0.1:1234".parse().expect("socket address"), + auth_state: RwLock::new(AuthState::Failed), + subscriptions: Arc::new(tokio::sync::Mutex::new(HashMap::new())), + pending_subscriptions: Arc::new(tokio::sync::Mutex::new(HashMap::new())), + send_tx: mpsc::channel(1).0, + ctrl_tx: mpsc::channel(1).0, + cancel: CancellationToken::new(), + backpressure_count: Arc::new(AtomicU8::new(0)), + grace_limit: 3, + }; + + let original = conn.begin_pending_subscription("same-id").await; + let replacement = conn.begin_pending_subscription("same-id").await; + + // The predecessor finishes while the replacement still owns the map. + conn.finish_pending_subscription("same-id", &original).await; + assert!(conn + .pending_subscriptions + .lock() + .await + .contains_key("same-id")); + + // If the replacement also fails, the completed predecessor must not be + // restored as an owner whose completion callback will never run again. + conn.finish_pending_subscription("same-id", &replacement) + .await; + assert!(conn.pending_subscriptions.lock().await.is_empty()); + } + + #[tokio::test] + async fn replacement_commit_fences_predecessor_historical_output() { + let (send_tx, mut send_rx) = mpsc::channel(8); + let conn = Arc::new(ConnectionState { + conn_id: Uuid::new_v4(), + tenant: TenantContext::resolved( + buzz_core::CommunityId::from_uuid(Uuid::new_v4()), + "lease.test", + ), + remote_addr: "127.0.0.1:1234".parse().expect("socket address"), + auth_state: RwLock::new(AuthState::Failed), + subscriptions: Arc::new(tokio::sync::Mutex::new(HashMap::new())), + pending_subscriptions: Arc::new(tokio::sync::Mutex::new(HashMap::new())), + send_tx, + ctrl_tx: mpsc::channel(1).0, + cancel: CancellationToken::new(), + backpressure_count: Arc::new(AtomicU8::new(0)), + grace_limit: 3, + }); + let predecessor = conn.begin_pending_subscription("same-id").await; + let replacement = conn.begin_pending_subscription("same-id").await; + let lifecycle = conn.pending_subscriptions.lock().await; + + let old_event = tokio::spawn({ + let conn = Arc::clone(&conn); + let predecessor = Arc::clone(&predecessor); + async move { + conn.send_historical_if_permitted("same-id", &predecessor, "old-event".into()) + .await + } + }); + let old_eose = tokio::spawn({ + let conn = Arc::clone(&conn); + let predecessor = Arc::clone(&predecessor); + async move { + conn.send_historical_if_permitted("same-id", &predecessor, "old-eose".into()) + .await + } + }); + tokio::task::yield_now().await; + + replacement.commit(); + assert!(conn.send("new-owner".into())); + drop(lifecycle); + + assert!(!old_event.await.expect("old event task")); + assert!(!old_eose.await.expect("old EOSE task")); + assert_eq!( + send_rx.recv().await, + Some(WsMessage::Text("new-owner".into())) + ); + assert!(send_rx.try_recv().is_err()); + } + + #[tokio::test] + async fn close_fences_predecessor_historical_output_before_closed() { + let (send_tx, mut send_rx) = mpsc::channel(8); + let conn = Arc::new(ConnectionState { + conn_id: Uuid::new_v4(), + tenant: TenantContext::resolved( + buzz_core::CommunityId::from_uuid(Uuid::new_v4()), + "lease.test", + ), + remote_addr: "127.0.0.1:1234".parse().expect("socket address"), + auth_state: RwLock::new(AuthState::Failed), + subscriptions: Arc::new(tokio::sync::Mutex::new(HashMap::new())), + pending_subscriptions: Arc::new(tokio::sync::Mutex::new(HashMap::new())), + send_tx, + ctrl_tx: mpsc::channel(1).0, + cancel: CancellationToken::new(), + backpressure_count: Arc::new(AtomicU8::new(0)), + grace_limit: 3, + }); + let pending = conn.begin_pending_subscription("same-id").await; + let mut lifecycle = conn.pending_subscriptions.lock().await; + + let old_event = tokio::spawn({ + let conn = Arc::clone(&conn); + let pending = Arc::clone(&pending); + async move { + conn.send_historical_if_permitted("same-id", &pending, "old-event".into()) + .await + } + }); + tokio::task::yield_now().await; + + lifecycle + .remove("same-id") + .expect("pending owner") + .cancel_lineage(); + assert!(conn.send("closed".into())); + drop(lifecycle); + + assert!(!old_event.await.expect("old event task")); + assert_eq!(send_rx.recv().await, Some(WsMessage::Text("closed".into()))); + assert!(send_rx.try_recv().is_err()); + } + + #[tokio::test] + async fn stale_req_rejection_is_fenced_after_same_id_replacement_commits() { + let (send_tx, mut send_rx) = mpsc::channel(8); + let conn = Arc::new(ConnectionState { + conn_id: Uuid::new_v4(), + tenant: TenantContext::resolved( + buzz_core::CommunityId::from_uuid(Uuid::new_v4()), + "lease.test", + ), + remote_addr: "127.0.0.1:1234".parse().expect("socket address"), + auth_state: RwLock::new(AuthState::Failed), + subscriptions: Arc::new(tokio::sync::Mutex::new(HashMap::new())), + pending_subscriptions: Arc::new(tokio::sync::Mutex::new(HashMap::new())), + send_tx, + ctrl_tx: mpsc::channel(1).0, + cancel: CancellationToken::new(), + backpressure_count: Arc::new(AtomicU8::new(0)), + grace_limit: 3, + }); + let predecessor = conn.begin_pending_subscription("same-id").await; + let replacement = conn.begin_pending_subscription("same-id").await; + let lifecycle = conn.pending_subscriptions.lock().await; + + // Model A resuming from an access-resolution error while B is at its + // registration cutover. A must wait for the lifecycle lock rather than + // queue its terminal response directly. + let stale_rejection = tokio::spawn({ + let conn = Arc::clone(&conn); + let predecessor = Arc::clone(&predecessor); + async move { + conn.send_req_messages_if_permitted( + "same-id", + &predecessor, + ["old-notice".into(), "old-closed".into()], + ) + .await + } + }); + tokio::task::yield_now().await; + + replacement.commit(); + assert!(conn.send("new-owner".into())); + drop(lifecycle); + + assert!(!stale_rejection.await.expect("stale rejection task")); + assert_eq!( + send_rx.recv().await, + Some(WsMessage::Text("new-owner".into())) + ); + assert!(send_rx.try_recv().is_err()); + } + + #[tokio::test] + async fn stale_req_rejection_is_fenced_after_close_acknowledgement() { + let (send_tx, mut send_rx) = mpsc::channel(8); + let conn = Arc::new(ConnectionState { + conn_id: Uuid::new_v4(), + tenant: TenantContext::resolved( + buzz_core::CommunityId::from_uuid(Uuid::new_v4()), + "lease.test", + ), + remote_addr: "127.0.0.1:1234".parse().expect("socket address"), + auth_state: RwLock::new(AuthState::Failed), + subscriptions: Arc::new(tokio::sync::Mutex::new(HashMap::new())), + pending_subscriptions: Arc::new(tokio::sync::Mutex::new(HashMap::new())), + send_tx, + ctrl_tx: mpsc::channel(1).0, + cancel: CancellationToken::new(), + backpressure_count: Arc::new(AtomicU8::new(0)), + grace_limit: 3, + }); + let pending = conn.begin_pending_subscription("same-id").await; + let mut lifecycle = conn.pending_subscriptions.lock().await; + + let stale_rejection = tokio::spawn({ + let conn = Arc::clone(&conn); + let pending = Arc::clone(&pending); + async move { + conn.send_req_messages_if_permitted( + "same-id", + &pending, + ["old-notice".into(), "old-closed".into()], + ) + .await + } + }); + tokio::task::yield_now().await; + + lifecycle + .remove("same-id") + .expect("pending owner") + .cancel_lineage(); + assert!(conn.send("close-ack".into())); + drop(lifecycle); + + assert!(!stale_rejection.await.expect("stale rejection task")); + assert_eq!( + send_rx.recv().await, + Some(WsMessage::Text("close-ack".into())) + ); + assert!(send_rx.try_recv().is_err()); + } + + #[tokio::test] + async fn committed_same_id_replacement_cancels_the_previous_lease() { + let original = Arc::new(PendingSubscription::new(None)); + let replacement = PendingSubscription::new(Some(Arc::clone(&original))); + + replacement.commit(); + + assert!(original.is_cancelled()); + assert!(!replacement.is_cancelled()); + } + #[derive(Debug, Default)] struct MockSinkState { messages: Vec, diff --git a/crates/buzz-relay/src/handlers/close.rs b/crates/buzz-relay/src/handlers/close.rs index 86f3d0da79b..0e61e006e6e 100644 --- a/crates/buzz-relay/src/handlers/close.rs +++ b/crates/buzz-relay/src/handlers/close.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use tracing::debug; use crate::connection::ConnectionState; +use crate::handlers::req::SubscriptionTopics; use crate::protocol::RelayMessage; use crate::state::AppState; @@ -10,26 +11,41 @@ use crate::state::AppState; pub async fn handle_close(sub_id: String, conn: Arc, state: Arc) { let conn_id = conn.conn_id; - conn.subscriptions.lock().await.remove(&sub_id); + remove_subscription(&sub_id, &conn, &state.sub_registry, state.pubsub.as_ref()).await; - // Deregister from the fan-out index before sending CLOSED so no new - // messages are routed to this sub after the client's CLOSE is acknowledged. - if let Some(removed) = state.sub_registry.remove_subscription(conn_id, &sub_id) { + conn.send(RelayMessage::closed(&sub_id, "")); + + debug!(conn_id = %conn_id, sub_id = %sub_id, "Subscription closed"); +} + +/// Cancel an in-flight REQ or remove its committed subscription. The pending +/// lock is held through topic release to serialize this cleanup with the REQ's +/// registration/topic-retain commit. +pub(crate) async fn remove_subscription( + sub_id: &str, + conn: &ConnectionState, + registry: &crate::subscription::SubscriptionRegistry, + pubsub: &dyn SubscriptionTopics, +) { + let mut pending_subscriptions = conn.pending_subscriptions.lock().await; + if let Some(pending) = pending_subscriptions.remove(sub_id) { + pending.cancel_lineage(); + } + + conn.subscriptions.lock().await.remove(sub_id); + + if let Some(removed) = registry.remove_subscription(conn.conn_id, sub_id) { if removed.scope.is_global() { - state - .pubsub + pubsub .release_topic(&conn.tenant, buzz_pubsub::EventTopic::Global) .await; } for &channel_id in removed.scope.channel_ids() { - state - .pubsub + pubsub .release_topic(&conn.tenant, buzz_pubsub::EventTopic::Channel(channel_id)) .await; } } - conn.send(RelayMessage::closed(&sub_id, "")); - - debug!(conn_id = %conn_id, sub_id = %sub_id, "Subscription closed"); + drop(pending_subscriptions); } diff --git a/crates/buzz-relay/src/handlers/event.rs b/crates/buzz-relay/src/handlers/event.rs index ccba40f3282..cf30e58a029 100644 --- a/crates/buzz-relay/src/handlers/event.rs +++ b/crates/buzz-relay/src/handlers/event.rs @@ -1405,6 +1405,7 @@ mod tests { }, )), subscriptions: Arc::new(Mutex::new(HashMap::new())), + pending_subscriptions: Arc::new(Mutex::new(HashMap::new())), send_tx, ctrl_tx, cancel: CancellationToken::new(), diff --git a/crates/buzz-relay/src/handlers/req.rs b/crates/buzz-relay/src/handlers/req.rs index 250fb4f9b92..ec5538dae6b 100644 --- a/crates/buzz-relay/src/handlers/req.rs +++ b/crates/buzz-relay/src/handlers/req.rs @@ -18,10 +18,31 @@ use nostr::Filter; use buzz_auth::Scope; -use crate::connection::{AuthState, ConnectionState}; +use crate::connection::{AuthState, ConnectionState, PendingSubscription}; use crate::protocol::RelayMessage; use crate::state::AppState; +/// Topic refcount operations used by subscription lifecycle commits. +/// +/// Keeping this boundary explicit lets race tests pause inside retain/release +/// rather than only before the commit begins. +#[async_trait::async_trait] +pub(crate) trait SubscriptionTopics: Send + Sync { + async fn retain_topic(&self, tenant: &TenantContext, topic: EventTopic); + async fn release_topic(&self, tenant: &TenantContext, topic: EventTopic); +} + +#[async_trait::async_trait] +impl SubscriptionTopics for buzz_pubsub::PubSubManager { + async fn retain_topic(&self, tenant: &TenantContext, topic: EventTopic) { + self.retain_topic(tenant, topic).await; + } + + async fn release_topic(&self, tenant: &TenantContext, topic: EventTopic) { + self.release_topic(tenant, topic).await; + } +} + const MAX_SUBSCRIPTIONS: usize = 1024; /// Maximum `query_events` calls in flight per multi-filter REQ / bridge query. @@ -47,23 +68,40 @@ pub(crate) const MAX_EXPLICIT_CHANNEL_VALUES: usize = 128; // the range fails the build. const _: () = assert!(FILTER_QUERY_CONCURRENCY >= 2 && FILTER_QUERY_CONCURRENCY <= 8); +async fn resolve_accessible_channel_ids( + state: &AppState, + conn: &ConnectionState, + pubkey_bytes: &[u8], +) -> Result, buzz_db::DbError> { + #[cfg(test)] + pause_access_resolution_for_test(conn.conn_id).await?; + + state + .get_accessible_channel_ids_cached(conn.tenant.community(), pubkey_bytes) + .await +} + /// Handle a REQ message: register the subscription, deliver historical events, then send EOSE. -pub async fn handle_req( +pub(crate) async fn handle_req( sub_id: String, filters: Vec, conn: Arc, state: Arc, + pending: Arc, ) { let (conn_id, pubkey_bytes, token_channel_ids) = { let auth = conn.auth_state.read().await; match &*auth { AuthState::Authenticated(ctx) => { if !ctx.scopes.is_empty() && !ctx.scopes.contains(&Scope::MessagesRead) { - conn.send(RelayMessage::notice("restricted: insufficient scope")); - conn.send(RelayMessage::closed( + send_req_rejection( + &conn, &sub_id, + &pending, + Some("restricted: insufficient scope"), "restricted: insufficient scope", - )); + ) + .await; return; } @@ -71,23 +109,28 @@ pub async fn handle_req( let subs = conn.subscriptions.lock().await; if !subs.contains_key(&sub_id) && subs.len() >= MAX_SUBSCRIPTIONS { - conn.send(RelayMessage::closed( + send_req_rejection( + &conn, &sub_id, + &pending, + None, "error: too many subscriptions", - )); + ) + .await; return; } (conn.conn_id, pk_bytes, ctx.channel_ids.clone()) } _ => { - conn.send(RelayMessage::notice( - "auth-required: authenticate before subscribing", - )); - conn.send(RelayMessage::closed( + send_req_rejection( + &conn, &sub_id, + &pending, + Some("auth-required: authenticate before subscribing"), "auth-required: not authenticated", - )); + ) + .await; return; } } @@ -97,10 +140,14 @@ pub async fn handle_req( let requested_channel_ids = match extract_channel_ids_from_filters_limited(&filters) { Ok(ids) => ids, Err(()) => { - conn.send(RelayMessage::closed( + send_req_rejection( + &conn, &sub_id, + &pending, + None, "restricted: too many explicit channels", - )); + ) + .await; return; } }; @@ -110,14 +157,11 @@ pub async fn handle_req( .increment(1); Vec::new() } else { - match state - .get_accessible_channel_ids_cached(conn.tenant.community(), &pubkey_bytes) - .await - { + match resolve_accessible_channel_ids(&state, &conn, &pubkey_bytes).await { Ok(ids) => ids, Err(e) => { warn!(conn_id = %conn_id, "Failed to get accessible channels: {e}"); - conn.send(RelayMessage::closed(&sub_id, "error: database error")); + send_req_rejection(&conn, &sub_id, &pending, None, "error: database error").await; return; } } @@ -170,7 +214,8 @@ pub async fn handle_req( } Err(e) => { warn!(conn_id = %conn_id, "Channel membership confirmation failed: {e}"); - conn.send(RelayMessage::closed(&sub_id, "error: database error")); + send_req_rejection(&conn, &sub_id, &pending, None, "error: database error") + .await; return; } } @@ -201,10 +246,14 @@ pub async fn handle_req( .as_ref() .is_some_and(|authorized| authorized.is_empty()) { - conn.send(RelayMessage::closed( + send_req_rejection( + &conn, &sub_id, + &pending, + None, "restricted: not a channel member", - )); + ) + .await; return; } @@ -219,24 +268,36 @@ pub async fn handle_req( if channel_id.is_none() { let authed_pubkey_hex = hex::encode(&pubkey_bytes); if !p_gated_filters_authorized(&filters, &authed_pubkey_hex) { - conn.send(RelayMessage::closed( + send_req_rejection( + &conn, &sub_id, + &pending, + None, "restricted: p-gated events require #p matching your pubkey", - )); + ) + .await; return; } if !engram_filters_authorized(&filters, &authed_pubkey_hex) { - conn.send(RelayMessage::closed( + send_req_rejection( + &conn, &sub_id, + &pending, + None, "restricted: agent-engram reads require authors=[self] or #p=[self]", - )); + ) + .await; return; } if !author_only_filters_authorized(&filters, &authed_pubkey_hex) { - conn.send(RelayMessage::closed( + send_req_rejection( + &conn, &sub_id, + &pending, + None, "restricted: author-only kinds require authors=[self]", - )); + ) + .await; return; } } @@ -248,10 +309,14 @@ pub async fn handle_req( let has_search = filters.iter().any(|f| f.search.is_some()); if has_search { if filters.iter().any(|f| f.search.is_none()) { - conn.send(RelayMessage::closed( + send_req_rejection( + &conn, &sub_id, + &pending, + None, "error: mixed search and non-search filters not supported", - )); + ) + .await; return; } handle_search_req( @@ -264,48 +329,28 @@ pub async fn handle_req( &conn, &state, trace_state.as_ref(), + &pending, ) .await; return; } + let subscription_scope = authorized_requested_channels + .clone() + .map(crate::subscription::SubscriptionScope::Channels) + .unwrap_or(crate::subscription::SubscriptionScope::Global); + if !register_subscription_if_current( + &sub_id, + &filters, + &subscription_scope, + &conn, + &state.sub_registry, + state.pubsub.as_ref(), + &pending, + ) + .await { - let mut subs = conn.subscriptions.lock().await; - subs.insert(sub_id.clone(), filters.clone()); - } - - let replaced = if let Some(channel_ids) = authorized_requested_channels.as_ref() { - state.sub_registry.register_channels_scoped( - conn.tenant.community(), - conn_id, - sub_id.clone(), - filters.clone(), - channel_ids.clone(), - ) - } else { - state.sub_registry.register_scoped( - conn.tenant.community(), - conn_id, - sub_id.clone(), - filters.clone(), - None, - ) - }; - if let Some(replaced) = replaced { - release_subscription_topics(&state, &conn.tenant, &replaced.scope).await; - } - if let Some(channel_ids) = authorized_requested_channels.as_ref() { - for &channel_id in channel_ids { - state - .pubsub - .retain_topic(&conn.tenant, EventTopic::Channel(channel_id)) - .await; - } - } else { - state - .pubsub - .retain_topic(&conn.tenant, EventTopic::Global) - .await; + return; } debug!(conn_id = %conn_id, sub_id = %sub_id, "Subscription registered"); @@ -382,7 +427,9 @@ pub async fn handle_req( Ok(evs) => evs, Err(e) => { warn!(conn_id = %conn_id, sub_id = %sub_id, "Historical query failed: {e}"); - conn.send(RelayMessage::eose(&sub_id)); + let _ = conn + .send_historical_if_permitted(&sub_id, &pending, RelayMessage::eose(&sub_id)) + .await; return; } }; @@ -459,7 +506,10 @@ pub async fn handle_req( } let msg = RelayMessage::event(&sub_id, &stored.event); - if !conn.send(msg) { + if !conn + .send_historical_if_permitted(&sub_id, &pending, msg) + .await + { return; } total_sent += 1; @@ -469,7 +519,9 @@ pub async fn handle_req( } } - conn.send(RelayMessage::eose(&sub_id)); + let _ = conn + .send_historical_if_permitted(&sub_id, &pending, RelayMessage::eose(&sub_id)) + .await; debug!( conn_id = %conn_id, @@ -485,6 +537,92 @@ pub async fn handle_req( /// pages to the request. const SEARCH_PAGE_SIZE: u32 = 100; +/// Send a terminal response owned by this REQ while its lifecycle lease is current. +/// A paired NOTICE and CLOSED are queued as one fenced batch. +async fn send_req_rejection( + conn: &ConnectionState, + sub_id: &str, + pending: &Arc, + notice: Option<&str>, + reason: &str, +) -> bool { + let messages = notice + .into_iter() + .map(RelayMessage::notice) + .chain(std::iter::once(RelayMessage::closed(sub_id, reason))); + conn.send_req_messages_if_permitted(sub_id, pending, messages) + .await +} + +/// Atomically commit a REQ into the per-connection map, fan-out registry, and +/// Redis topic refcount only while it still owns the pending lease. CLOSE uses +/// the same pending-subscription lock, so it either invalidates this lease +/// first or observes and removes the completed registration afterward. +async fn register_subscription_if_current( + sub_id: &str, + filters: &[Filter], + scope: &crate::subscription::SubscriptionScope, + conn: &ConnectionState, + registry: &crate::subscription::SubscriptionRegistry, + pubsub: &dyn SubscriptionTopics, + pending: &Arc, +) -> bool { + let pending_subscriptions = conn.pending_subscriptions.lock().await; + let permitted_to_commit = pending_subscriptions + .get(sub_id) + .is_some_and(|current| PendingSubscription::permits_commit(current, pending)); + if conn.cancel.is_cancelled() || !permitted_to_commit { + return false; + } + + conn.subscriptions + .lock() + .await + .insert(sub_id.to_owned(), filters.to_vec()); + + let replaced = match scope { + crate::subscription::SubscriptionScope::Channels(channel_ids) => registry + .register_channels_scoped( + conn.tenant.community(), + conn.conn_id, + sub_id.to_owned(), + filters.to_vec(), + channel_ids.clone(), + ), + crate::subscription::SubscriptionScope::Global => registry.register_scoped( + conn.tenant.community(), + conn.conn_id, + sub_id.to_owned(), + filters.to_vec(), + None, + ), + }; + if let Some(replaced) = replaced { + release_subscription_topics(pubsub, &conn.tenant, &replaced.scope).await; + } + match scope { + crate::subscription::SubscriptionScope::Channels(channel_ids) => { + for &channel_id in channel_ids { + pubsub + .retain_topic(&conn.tenant, EventTopic::Channel(channel_id)) + .await; + } + } + crate::subscription::SubscriptionScope::Global => { + pubsub.retain_topic(&conn.tenant, EventTopic::Global).await; + } + } + + // Only a successful replacement commit supersedes historical delivery from + // the previous same-ID REQ. Validation/search failures leave it untouched. + pending.commit(); + + drop(pending_subscriptions); + true +} + +/// Handle a NIP-50 search REQ: query Postgres FTS, fetch full events, deliver results, EOSE. +/// Search subscriptions are one-shot — no persistent subscription is registered. /// Maximum FTS pages to fetch per filter (prevents unbounded loops). /// /// Derived from the advertised page ceiling rather than fixed: the scan @@ -591,6 +729,7 @@ async fn handle_search_req( conn: &ConnectionState, state: &AppState, trace_state: Option<&crate::conformance::AbstractState>, + pending: &Arc, ) { // The community-wide channel scope (no #h tag on the filter). `None` means // "no accessible channels and no global access" → EOSE, exactly as the @@ -599,7 +738,9 @@ async fn handle_search_req( match build_search_channel_scope_filter(accessible_channels, include_global) { Some(scope) => scope, None => { - conn.send(RelayMessage::eose(sub_id)); + let _ = conn + .send_historical_if_permitted(sub_id, pending, RelayMessage::eose(sub_id)) + .await; return; } }; @@ -789,7 +930,14 @@ async fn handle_search_req( if !seen_ids.insert(stored.event.id) { continue; } - if !conn.send(RelayMessage::event(sub_id, &stored.event)) { + if !conn + .send_historical_if_permitted( + sub_id, + pending, + RelayMessage::event(sub_id, &stored.event), + ) + .await + { return; } emitted += 1; @@ -802,7 +950,9 @@ async fn handle_search_req( } } - conn.send(RelayMessage::eose(sub_id)); + let _ = conn + .send_historical_if_permitted(sub_id, pending, RelayMessage::eose(sub_id)) + .await; } /// Convert a single NIP-01 filter into an [`EventQuery`] for the database. @@ -1134,16 +1284,15 @@ pub(crate) fn extract_channel_ids_from_filters(filters: &[Filter]) -> Option>>, + resume: Arc, + fail: bool, +} + +#[cfg(test)] +fn access_resolution_test_hooks( +) -> &'static std::sync::Mutex>> +{ + static HOOKS: std::sync::OnceLock< + std::sync::Mutex>>, + > = std::sync::OnceLock::new(); + HOOKS.get_or_init(|| std::sync::Mutex::new(std::collections::HashMap::new())) +} + +#[cfg(test)] +async fn pause_access_resolution_for_test(conn_id: uuid::Uuid) -> Result<(), buzz_db::DbError> { + let hook = access_resolution_test_hooks() + .lock() + .expect("access-resolution hook mutex") + .get(&conn_id) + .cloned(); + let Some(hook) = hook else { + return Ok(()); + }; + + if let Some(started) = hook + .started + .lock() + .expect("access-resolution started mutex") + .take() + { + let _ = started.send(()); + } + hook.resume.notified().await; + access_resolution_test_hooks() + .lock() + .expect("access-resolution hook mutex") + .remove(&conn_id); + + if hook.fail { + Err(buzz_db::DbError::InvalidData( + "injected access-resolution failure".to_owned(), + )) + } else { + Ok(()) + } +} + #[cfg(test)] mod tests { use super::*; - use nostr::{Alphabet, Filter, SingleLetterTag}; + use std::collections::HashMap; + use std::sync::atomic::{AtomicBool, AtomicU8, Ordering}; + + use axum::extract::ws::Message as WsMessage; + use buzz_core::{CommunityId, StoredEvent}; + use buzz_pubsub::{EventTopic, PubSubManager}; + use chrono::Utc; + use nostr::{Alphabet, EventBuilder, Filter, Keys, Kind, SingleLetterTag}; + use tokio::sync::{mpsc, Barrier, Mutex, Notify, RwLock}; + use tokio::task::JoinSet; + use tokio_util::sync::CancellationToken; + + fn test_connection_with_receiver( + tenant: TenantContext, + ) -> (Arc, mpsc::Receiver) { + let (send_tx, send_rx) = mpsc::channel::(8); + let (ctrl_tx, _ctrl_rx) = mpsc::channel::(8); + let conn = Arc::new(ConnectionState { + conn_id: uuid::Uuid::new_v4(), + tenant, + remote_addr: "127.0.0.1:1234".parse().expect("socket address"), + auth_state: RwLock::new(AuthState::Failed), + subscriptions: Arc::new(Mutex::new(HashMap::new())), + pending_subscriptions: Arc::new(Mutex::new(HashMap::new())), + send_tx, + ctrl_tx, + cancel: CancellationToken::new(), + backpressure_count: Arc::new(AtomicU8::new(0)), + grace_limit: 3, + }); + (conn, send_rx) + } + + fn test_connection(tenant: TenantContext) -> Arc { + test_connection_with_receiver(tenant).0 + } + + async fn test_pubsub() -> Arc { + let config = deadpool_redis::Config::from_url("redis://127.0.0.1:6379"); + let pool = config + .create_pool(Some(deadpool_redis::Runtime::Tokio1)) + .expect("Redis test pool"); + Arc::new( + PubSubManager::new("redis://127.0.0.1:6379", pool) + .await + .expect("pubsub manager"), + ) + } + + fn matching_event(channel_id: uuid::Uuid) -> StoredEvent { + let event = EventBuilder::new(Kind::TextNote, "race test") + .tags([]) + .sign_with_keys(&Keys::generate()) + .expect("sign event"); + StoredEvent::with_received_at(event, Utc::now(), Some(channel_id), true) + } + + fn subscription_gauge_value(snapshotter: &metrics_util::debugging::Snapshotter) -> f64 { + snapshotter + .snapshot() + .into_vec() + .into_iter() + .find(|(key, ..)| key.key().name() == "buzz_subscriptions_active") + .map(|(_, _, _, value)| { + let metrics_util::debugging::DebugValue::Gauge(value) = value else { + panic!("buzz_subscriptions_active must be a gauge"); + }; + value.into_inner() + }) + .unwrap_or(0.0) + } + + struct RetainBarrierTopics { + inner: Arc, + entered: Arc, + resume: Arc, + pause_next_retain: AtomicBool, + } + + impl RetainBarrierTopics { + fn new(inner: Arc) -> Self { + Self { + inner, + entered: Arc::new(Barrier::new(2)), + resume: Arc::new(Notify::new()), + pause_next_retain: AtomicBool::new(true), + } + } + } + + #[async_trait::async_trait] + impl SubscriptionTopics for RetainBarrierTopics { + async fn retain_topic(&self, tenant: &TenantContext, topic: EventTopic) { + if self.pause_next_retain.swap(false, Ordering::SeqCst) { + self.entered.wait().await; + self.resume.notified().await; + } + self.inner.retain_topic(tenant, topic).await; + } + + async fn release_topic(&self, tenant: &TenantContext, topic: EventTopic) { + self.inner.release_topic(tenant, topic).await; + } + } + + #[test] + fn close_during_topic_retain_serializes_after_registration_commit() { + let recorder = metrics_util::debugging::DebuggingRecorder::new(); + let snapshotter = recorder.snapshotter(); + metrics::with_local_recorder(&recorder, || { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("test runtime"); + let local = tokio::task::LocalSet::new(); + local.block_on(&runtime, async { + metrics::gauge!("buzz_subscriptions_active").increment(0.0); + + let community = CommunityId::from_uuid(uuid::Uuid::new_v4()); + let conn = test_connection(TenantContext::resolved(community, "race.test")); + let registry = Arc::new(crate::subscription::SubscriptionRegistry::new()); + let inner_pubsub = test_pubsub().await; + let pubsub = Arc::new(RetainBarrierTopics::new(Arc::clone(&inner_pubsub))); + let channel_id = uuid::Uuid::new_v4(); + let sub_id = "close-race".to_owned(); + let filters = vec![Filter::new().kind(Kind::TextNote)]; + let pending = conn.begin_pending_subscription(&sub_id).await; + + let req_task = tokio::task::spawn_local({ + let conn = Arc::clone(&conn); + let registry = Arc::clone(®istry); + let pubsub = Arc::clone(&pubsub); + let pending = Arc::clone(&pending); + let sub_id = sub_id.clone(); + let filters = filters.clone(); + async move { + register_subscription_if_current( + &sub_id, + &filters, + &crate::subscription::SubscriptionScope::Channels(vec![channel_id]), + &conn, + ®istry, + pubsub.as_ref(), + &pending, + ) + .await + } + }); + + pubsub.entered.wait().await; + assert_eq!(registry.total_subscriptions(), 1); + assert_eq!( + inner_pubsub + .topic_refcount(&conn.tenant, EventTopic::Channel(channel_id)) + .await, + 0 + ); + + let close_finished = Arc::new(AtomicBool::new(false)); + let close_task = tokio::task::spawn_local({ + let close_finished = Arc::clone(&close_finished); + let conn = Arc::clone(&conn); + let registry = Arc::clone(®istry); + let pubsub = Arc::clone(&pubsub); + let sub_id = sub_id.clone(); + async move { + crate::handlers::close::remove_subscription( + &sub_id, + &conn, + ®istry, + pubsub.as_ref(), + ) + .await; + close_finished.store(true, Ordering::SeqCst); + } + }); + tokio::task::yield_now().await; + assert!(!close_finished.load(Ordering::SeqCst)); + + pubsub.resume.notify_one(); + assert!(req_task.await.expect("REQ task")); + close_task.await.expect("CLOSE task"); + + assert!(conn.subscriptions.lock().await.is_empty()); + assert!(conn.pending_subscriptions.lock().await.is_empty()); + assert_eq!(registry.total_subscriptions(), 0); + assert!(registry + .fan_out_scoped(community, &matching_event(channel_id)) + .is_empty()); + assert_eq!( + inner_pubsub + .topic_refcount(&conn.tenant, EventTopic::Channel(channel_id)) + .await, + 0 + ); + }); + }); + assert_eq!(subscription_gauge_value(&snapshotter), 0.0); + } + + #[tokio::test] + async fn predecessor_commits_while_newer_same_id_req_is_pending() { + let community = CommunityId::from_uuid(uuid::Uuid::new_v4()); + let conn = test_connection(TenantContext::resolved(community, "race.test")); + let registry = crate::subscription::SubscriptionRegistry::new(); + let pubsub = test_pubsub().await; + let channel_id = uuid::Uuid::new_v4(); + let sub_id = "replacement-pending"; + let filters = vec![Filter::new().kind(Kind::TextNote)]; + + let predecessor = conn.begin_pending_subscription(sub_id).await; + let replacement = conn.begin_pending_subscription(sub_id).await; + + assert!( + register_subscription_if_current( + sub_id, + &filters, + &crate::subscription::SubscriptionScope::Channels(vec![channel_id]), + &conn, + ®istry, + pubsub.as_ref(), + &predecessor, + ) + .await + ); + assert!(!predecessor.is_cancelled()); + assert_eq!(registry.total_subscriptions(), 1); + assert_eq!( + pubsub + .topic_refcount(&conn.tenant, EventTopic::Channel(channel_id)) + .await, + 1 + ); + + // The predecessor's handler completes after establishing its persistent + // subscription. If the newer request then fails, pending ownership is + // cleared without disturbing that already-registered subscription. + conn.finish_pending_subscription(sub_id, &predecessor).await; + conn.finish_pending_subscription(sub_id, &replacement).await; + + assert!(conn.pending_subscriptions.lock().await.is_empty()); + assert_eq!(registry.total_subscriptions(), 1); + assert_eq!( + conn.subscriptions + .lock() + .await + .get(sub_id) + .map(Vec::as_slice), + Some(filters.as_slice()) + ); + assert_eq!( + pubsub + .topic_refcount(&conn.tenant, EventTopic::Channel(channel_id)) + .await, + 1 + ); + } + + #[tokio::test] + async fn stale_handle_req_rejection_after_same_id_replacement_is_suppressed() { + let community = CommunityId::from_uuid(uuid::Uuid::new_v4()); + let (conn, mut send_rx) = + test_connection_with_receiver(TenantContext::resolved(community, "race.test")); + let keys = Keys::generate(); + *conn.auth_state.write().await = AuthState::Authenticated(buzz_auth::AuthContext { + pubkey: keys.public_key(), + scopes: vec![], + channel_ids: None, + auth_method: buzz_auth::AuthMethod::Nip42, + agent_owner_pubkey: None, + }); + + let mut config = crate::config::Config::from_env().expect("default config loads"); + config.redis_url = "redis://127.0.0.1:6379".to_owned(); + let pool = sqlx::PgPool::connect_lazy(&config.database_url).expect("lazy pg pool"); + let db = buzz_db::Db::from_pool(pool.clone()); + let redis_pool = deadpool_redis::Config::from_url(&config.redis_url) + .create_pool(Some(deadpool_redis::Runtime::Tokio1)) + .expect("redis pool"); + let pubsub = test_pubsub().await; + let auth = buzz_auth::AuthService::new(config.auth.clone()); + let search = buzz_search::SearchService::new(pool); + let workflow_engine = Arc::new(buzz_workflow::WorkflowEngine::new( + db.clone(), + buzz_workflow::WorkflowConfig::default(), + )); + let media_storage = buzz_media::MediaStorage::new(&config.media).expect("media storage"); + let (state, _audit_shutdown) = AppState::new( + config, + db, + redis_pool, + None, + Arc::clone(&pubsub), + auth, + search, + workflow_engine, + Keys::generate(), + media_storage, + ); + let state = Arc::new(state); + + let sub_id = "handler-replacement-race".to_owned(); + let stale_filters = vec![Filter::new().kind(Kind::TextNote)]; + let survivor_filters = vec![Filter::new().kind(Kind::Metadata)]; + let stale = conn.begin_pending_subscription(&sub_id).await; + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let resume = Arc::new(Notify::new()); + access_resolution_test_hooks() + .lock() + .expect("access-resolution hook mutex") + .insert( + conn.conn_id, + Arc::new(AccessResolutionTestHook { + started: std::sync::Mutex::new(Some(started_tx)), + resume: Arc::clone(&resume), + fail: true, + }), + ); + + let stale_task = tokio::spawn(handle_req( + sub_id.clone(), + stale_filters, + Arc::clone(&conn), + Arc::clone(&state), + Arc::clone(&stale), + )); + started_rx + .await + .expect("stale REQ reached access resolution"); + + let survivor = conn.begin_pending_subscription(&sub_id).await; + assert!( + register_subscription_if_current( + &sub_id, + &survivor_filters, + &crate::subscription::SubscriptionScope::Global, + &conn, + &state.sub_registry, + state.pubsub.as_ref(), + &survivor, + ) + .await + ); + resume.notify_one(); + stale_task.await.expect("stale handle_req task"); + + assert!( + send_rx.try_recv().is_err(), + "stale handler must not queue NOTICE or CLOSED after replacement" + ); + assert_eq!( + conn.subscriptions.lock().await.get(&sub_id), + Some(&survivor_filters), + "same-ID survivor must remain registered" + ); + assert_eq!(state.sub_registry.total_subscriptions(), 1); + } + + #[test] + fn disconnect_during_pre_registration_await_drains_req_before_cleanup() { + let recorder = metrics_util::debugging::DebuggingRecorder::new(); + let snapshotter = recorder.snapshotter(); + metrics::with_local_recorder(&recorder, || { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("test runtime"); + runtime.block_on(async { + metrics::gauge!("buzz_subscriptions_active").increment(0.0); + + let community = CommunityId::from_uuid(uuid::Uuid::new_v4()); + let conn = test_connection(TenantContext::resolved(community, "race.test")); + let registry = Arc::new(crate::subscription::SubscriptionRegistry::new()); + let pubsub = test_pubsub().await; + let channel_id = uuid::Uuid::new_v4(); + let sub_id = "disconnect-race".to_owned(); + let filters = vec![Filter::new().kind(Kind::TextNote)]; + let pending = conn.begin_pending_subscription(&sub_id).await; + let gate = Arc::new(Notify::new()); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let mut req_tasks = JoinSet::new(); + + req_tasks.spawn({ + let conn = Arc::clone(&conn); + let registry = Arc::clone(®istry); + let pubsub = Arc::clone(&pubsub); + let pending = Arc::clone(&pending); + let gate = Arc::clone(&gate); + let sub_id = sub_id.clone(); + let filters = filters.clone(); + async move { + started_tx.send(()).expect("signal pre-registration await"); + gate.notified().await; + let _ = register_subscription_if_current( + &sub_id, + &filters, + &crate::subscription::SubscriptionScope::Channels(vec![channel_id]), + &conn, + ®istry, + pubsub.as_ref(), + &pending, + ) + .await; + } + }); + + started_rx + .await + .expect("REQ reached pre-registration await"); + crate::connection::shutdown_pending_requests(&conn, &mut req_tasks).await; + for removed in registry.remove_connection(conn.conn_id) { + release_subscription_topics(pubsub.as_ref(), &conn.tenant, &removed.scope) + .await; + } + gate.notify_one(); + tokio::task::yield_now().await; + + assert!(conn.cancel.is_cancelled()); + assert!(req_tasks.is_empty()); + assert!(conn.subscriptions.lock().await.is_empty()); + assert!(conn.pending_subscriptions.lock().await.is_empty()); + assert_eq!(registry.total_subscriptions(), 0); + assert!(registry + .fan_out_scoped(community, &matching_event(channel_id)) + .is_empty()); + assert_eq!( + pubsub + .topic_refcount(&conn.tenant, EventTopic::Channel(channel_id)) + .await, + 0 + ); + }); + }); + assert_eq!(subscription_gauge_value(&snapshotter), 0.0); + } #[test] fn global_queries_push_access_scope_before_limit() { diff --git a/crates/buzz-relay/src/state.rs b/crates/buzz-relay/src/state.rs index 2f544e188c0..5f96b5b9a0e 100644 --- a/crates/buzz-relay/src/state.rs +++ b/crates/buzz-relay/src/state.rs @@ -1502,6 +1502,7 @@ mod tests { remote_addr: "127.0.0.1:1234".parse().unwrap(), auth_state: RwLock::new(AuthState::Failed), subscriptions: Arc::new(Mutex::new(HashMap::new())), + pending_subscriptions: Arc::new(Mutex::new(HashMap::new())), send_tx: tx.clone(), ctrl_tx, cancel: cancel.clone(),