diff --git a/crates/buzz-relay/src/connection.rs b/crates/buzz-relay/src/connection.rs index 5b9951c9db..ac794ba895 100644 --- a/crates/buzz-relay/src/connection.rs +++ b/crates/buzz-relay/src/connection.rs @@ -79,6 +79,25 @@ impl PendingSubscription { } } + /// 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; @@ -182,7 +201,7 @@ impl ConnectionState { } /// Forget a completed REQ only if it still owns this subscription ID. - async fn finish_pending_subscription( + pub(crate) async fn finish_pending_subscription( &self, sub_id: &str, completed: &Arc, diff --git a/crates/buzz-relay/src/handlers/req.rs b/crates/buzz-relay/src/handlers/req.rs index 501fd0057f..5dc56637a7 100644 --- a/crates/buzz-relay/src/handlers/req.rs +++ b/crates/buzz-relay/src/handlers/req.rs @@ -497,10 +497,10 @@ async fn register_subscription_if_current( pending: &Arc, ) -> bool { let pending_subscriptions = conn.pending_subscriptions.lock().await; - let still_current = pending_subscriptions + let permitted_to_commit = pending_subscriptions .get(sub_id) - .is_some_and(|current| Arc::ptr_eq(current, pending)); - if conn.cancel.is_cancelled() || pending.is_cancelled() || !still_current { + .is_some_and(|current| PendingSubscription::permits_commit(current, pending)); + if conn.cancel.is_cancelled() || !permitted_to_commit { return false; } @@ -1659,6 +1659,64 @@ mod tests { 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, + Some(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 + ); + } + #[test] fn disconnect_during_pre_registration_await_drains_req_before_cleanup() { let recorder = metrics_util::debugging::DebuggingRecorder::new();