Merge branch 'stack/subscription-routing' into stack/connect-transport-rewrite

Author: sanity

crates/core/src/operations/put.rs

@@ -581,6 +581,11 @@ impl Operation for PutOp {
                     return_msg = None;
                 }
                 PutMsg::SuccessfulPut { id, .. } => {
+                    tracing::debug!(
+                        tx = %id,
+                        current_state = ?self.state,
+                        "PutOp::process_message: handling SuccessfulPut"
+                    );
                     match self.state {
                         Some(PutState::AwaitingResponse {
                             key,
@@ -665,13 +670,24 @@ impl Operation for PutOp {
 
                             // Forward success message upstream if needed
                             if let Some(upstream) = upstream {
+                                tracing::trace!(
+                                    tx = %id,
+                                    %key,
+                                    upstream = %upstream.peer,
+                                    "PutOp::process_message: Forwarding SuccessfulPut upstream"
+                                );
                                 return_msg = Some(PutMsg::SuccessfulPut {
                                     id: *id,
                                     target: upstream,
                                     key,
                                     sender: op_manager.ring.connection_manager.own_location(),
                                 });
                             } else {
+                                tracing::trace!(
+                                    tx = %id,
+                                    %key,
+                                    "PutOp::process_message: SuccessfulPut originated locally; no upstream"
+                                );
                                 return_msg = None;
                             }
                         }
@@ -894,6 +910,14 @@ async fn try_to_broadcast(
         _ => false,
     };
 
+    let preserved_upstream = match &state {
+        Some(PutState::AwaitingResponse {
+            upstream: Some(existing),
+            ..
+        }) => Some(existing.clone()),
+        _ => None,
+    };
+
     match state {
         // Handle initiating node that's also the target (single node or targeting self)
         Some(PutState::AwaitingResponse {
@@ -932,9 +956,12 @@ async fn try_to_broadcast(
                     key
                 );
                 // means the whole tx finished so can return early
+                let upstream_for_completion = preserved_upstream
+                    .clone()
+                    .or_else(|| Some(upstream.clone()));
                 new_state = Some(PutState::AwaitingResponse {
                     key,
-                    upstream: Some(upstream),
+                    upstream: upstream_for_completion,
                     contract: contract.clone(), // No longer optional
                     state: new_value.clone(),
                     subscribe,

crates/core/src/operations/update.rs

@@ -786,10 +786,29 @@ impl OpManager {
             .subscribers_of(key)
             .map(|subs| {
                 let self_peer = self.ring.connection_manager.get_peer_key();
+                let allow_self = self_peer.as_ref().map(|me| me == sender).unwrap_or(false);
                 subs.value()
                     .iter()
-                    .filter(|pk| &pk.peer != sender)
-                    .filter(|pk| self_peer.as_ref().map(|me| &pk.peer != me).unwrap_or(true))
+                    .filter(|pk| {
+                        // Allow the sender to remain in the broadcast list when we're the sender,
+                        // so local auto-subscribe via GET/PUT still receives notifications.
+                        if &pk.peer == sender {
+                            allow_self
+                        } else {
+                            true
+                        }
+                    })
+                    .filter(|pk| {
+                        if let Some(self_peer) = &self_peer {
+                            if &pk.peer == self_peer {
+                                allow_self
+                            } else {
+                                true
+                            }
+                        } else {
+                            true
+                        }
+                    })
                     .cloned()
                     .collect::<Vec<_>>()
             })