From 0438bef278b989e215764144f2cff84746b41bc2 Mon Sep 17 00:00:00 2001
From: npub13n66s06epmqf2kc3v373ez8hj65cuzyvxzjf93vwpervxqn2u7jq2qd9je
 <8cf5a83f590ec0955b11647d1c88f796a98e088c30a492c58e0e46c3026ae7a4@buzz.block.builderlab.xyz>
Date: Tue, 28 Jul 2026 10:55:06 -0400
Subject: [PATCH] fix(relay): skip completed request predecessors

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>
---
 crates/buzz-relay/src/connection.rs | 57 ++++++++++++++++++++++++++---
 1 file changed, 52 insertions(+), 5 deletions(-)

diff --git a/crates/buzz-relay/src/connection.rs b/crates/buzz-relay/src/connection.rs
index 9b53e82f2b..5b9951c9db 100644
--- a/crates/buzz-relay/src/connection.rs
+++ b/crates/buzz-relay/src/connection.rs
@@ -44,6 +44,7 @@ pub(crate) struct RestartClose {
 pub(crate) struct PendingSubscription {
     cancel: CancellationToken,
     committed: AtomicBool,
+    completed: AtomicBool,
     predecessor: Option<Arc<PendingSubscription>>,
 }
 
@@ -52,6 +53,7 @@ impl PendingSubscription {
         Self {
             cancel: CancellationToken::new(),
             committed: AtomicBool::new(false),
+            completed: AtomicBool::new(false),
             predecessor,
         }
     }
@@ -78,11 +80,18 @@ impl PendingSubscription {
     }
 
     fn live_predecessor(&self) -> Option<Arc<PendingSubscription>> {
-        (!self.committed.load(Ordering::Acquire))
-            .then_some(self.predecessor.as_ref())
-            .flatten()
-            .filter(|predecessor| !predecessor.is_cancelled())
-            .cloned()
+        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 {
@@ -178,6 +187,7 @@ impl ConnectionState {
         sub_id: &str,
         completed: &Arc<PendingSubscription>,
     ) {
+        completed.completed.store(true, Ordering::Release);
         let mut pending = self.pending_subscriptions.lock().await;
         if pending
             .get(sub_id)
@@ -926,6 +936,43 @@ mod tests {
         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 committed_same_id_replacement_cancels_the_previous_lease() {
         let original = Arc::new(PendingSubscription::new(None));