fix: honor read_from_replicas for sharded pubsub subscriptions in ClusterPubSub

redis/asyncio/cluster.py

@@ -2205,7 +2205,7 @@ def get_node_from_slot(
                 )
                 return self.slots_cache[slot][node_idx]
             return self.slots_cache[slot][0]
-        except (IndexError, TypeError):
+        except (KeyError, IndexError, TypeError):
             raise SlotNotCoveredError(
                 f'Slot "{slot}" not covered by the cluster. '
                 f'"require_full_coverage={self.require_full_coverage}"'
@@ -3961,6 +3961,19 @@ async def get_sharded_message(
                 return None
         return message
 
+    def _eligible_subscription_nodes(self, slot: int) -> Dict[str, "ClusterNode"]:
+        """
+        Nodes a shard subscription may live on for ``slot``, keyed by node
+        name. Mirrors regular pubsub node selection: replicas are eligible
+        only when replica reads are enabled.
+        """
+        nodes = self.cluster.nodes_manager.slots_cache.get(slot) or []
+        if not (
+            self.cluster.read_from_replicas or self.cluster.load_balancing_strategy
+        ):
+            nodes = nodes[:1]
+        return {node.name: node for node in nodes}
+
     async def ssubscribe(
         self, *args: ChannelT | Subscription, **kwargs: PubSubHandler
     ) -> None:
@@ -3978,25 +3991,38 @@ async def ssubscribe(
         # re-acquire this lock (asyncio.Lock is non-reentrant).
         async with self._shard_state_lock:
             for s_channel, handler in s_channels.items():
-                node = self.cluster.get_node_from_key(s_channel)
+                slot = self.cluster.keyslot(s_channel)
+                node = self.cluster.nodes_manager.get_node_from_slot(
+                    slot,
+                    self.cluster.read_from_replicas,
+                    self.cluster.load_balancing_strategy,
+                )
                 if not node:
                     continue
                 # Lazy re-route: if this channel is already tracked against a
-                # different node (e.g. after a slot migration), migrate it now
-                # so the caller's intent is applied on the current owner.
+                # different node, migrate it now so the caller's intent is
+                # applied on the current owner — unless the tracked node
+                # still serves the slot, in which case leave it alone.
                 normalized_key = next(iter(self._normalize_keys({s_channel: None})))
                 old_name = self._shard_channel_to_node.get(normalized_key)
                 if old_name and old_name != node.name:
-                    # Match PubSub.ssubscribe() dict.update() semantics: the
-                    # caller's newly supplied handler (including None) always
-                    # overrides any previously registered handler.
-                    await self._migrate_shard_channel(
-                        normalized_key,
-                        handler,
-                        old_name,
-                        node,
-                    )
-                    continue
+                    tracked_node = self._eligible_subscription_nodes(slot).get(old_name)
+                    if tracked_node is not None:
+                        # The tracked node still serves the channel's slot
+                        # (the balancer merely picked a sibling); keep the
+                        # subscription where it is to avoid churn.
+                        node = tracked_node
+                    else:
+                        # Match PubSub.ssubscribe() dict.update() semantics: the
+                        # caller's newly supplied handler (including None) always
+                        # overrides any previously registered handler.
+                        await self._migrate_shard_channel(
+                            normalized_key,
+                            handler,
+                            old_name,
+                            node,
+                        )
+                        continue
                 pubsub = self._get_node_pubsub(node)
                 if handler:
                     await pubsub.ssubscribe(Subscription(s_channel, handler))
@@ -4032,10 +4058,26 @@ async def sunsubscribe(self, *args: Any) -> None:
                 if name and name in self.node_pubsub_mapping:
                     pubsub = self.node_pubsub_mapping[name]
                 else:
-                    node = self.cluster.get_node_from_key(s_channel)
-                    if not node or node.name not in self.node_pubsub_mapping:
+                    # No reverse-index entry: find the node whose pubsub
+                    # actually holds the subscription (with replica reads
+                    # enabled that is not necessarily the primary).
+                    slot = self.cluster.keyslot(s_channel)
+                    eligible = self._eligible_subscription_nodes(slot)
+                    if not eligible:
+                        raise SlotNotCoveredError(
+                            f'Slot "{slot}" is not covered by the cluster.'
+                        )
+                    pubsub = None
+                    for name in eligible:
+                        candidate = self.node_pubsub_mapping.get(name)
+                        if (
+                            candidate is not None
+                            and normalized_key in candidate.shard_channels
+                        ):
+                            pubsub = candidate
+                            break
+                    if pubsub is None:
                         continue
-                    pubsub = self.node_pubsub_mapping[node.name]
                 await pubsub.sunsubscribe(s_channel)
                 self.pending_unsubscribe_shard_channels.update(
                     pubsub.pending_unsubscribe_shard_channels
@@ -4044,18 +4086,28 @@ async def sunsubscribe(self, *args: Any) -> None:
     async def reinitialize_shard_subscriptions(self) -> None:
         """
         Reconcile per-node shard subscriptions against the cluster's current
-        slot ownership map. For each tracked shard channel whose owning node
-        has changed (e.g. after CLUSTER SETSLOT / failover), sunsubscribe on
-        the old node's pubsub and ssubscribe on the new owner's pubsub,
-        preserving any registered handler.
+        slot ownership map. For each tracked shard channel whose node no
+        longer serves the channel's slot (e.g. after CLUSTER SETSLOT /
+        failover), sunsubscribe on the old node's pubsub and ssubscribe on a
+        node currently serving the slot, preserving any registered handler.
         """
         uncovered: list = []
         made_progress = False
         first_migrate_error: Optional[BaseException] = None
         async with self._shard_state_lock:
             for channel, handler in list(self.shard_channels.items()):
+                slot = self.cluster.keyslot(channel)
+                old_name = self._shard_channel_to_node.get(channel)
+                if old_name and old_name in self._eligible_subscription_nodes(slot):
+                    # The tracked node still serves the channel's slot under
+                    # the current routing configuration; nothing to do.
+                    continue
                 try:
-                    new_node = self.cluster.get_node_from_key(channel)
+                    new_node = self.cluster.nodes_manager.get_node_from_slot(
+                        slot,
+                        self.cluster.read_from_replicas,
+                        self.cluster.load_balancing_strategy,
+                    )
                 except SlotNotCoveredError:
                     # Slot is transiently uncovered (mid-migration / partial
                     # topology refresh). Defer this channel so coverable
@@ -4065,9 +4117,6 @@ async def reinitialize_shard_subscriptions(self) -> None:
                     # slots-cache change notification.
                     uncovered.append(channel)
                     continue
-                old_name = self._shard_channel_to_node.get(channel)
-                if old_name == new_node.name:
-                    continue
                 try:
                     await self._migrate_shard_channel(
                         channel, handler, old_name, new_node

redis/cluster.py

@@ -3199,6 +3199,19 @@ def get_sharded_message(
                 return None
         return message
 
+    def _eligible_subscription_nodes(self, slot):
+        """
+        Nodes a shard subscription may live on for ``slot``, keyed by node
+        name. Mirrors regular pubsub node selection: replicas are eligible
+        only when replica reads are enabled.
+        """
+        nodes = self.cluster.nodes_manager.slots_cache.get(slot) or []
+        if not (
+            self.cluster.read_from_replicas or self.cluster.load_balancing_strategy
+        ):
+            nodes = nodes[:1]
+        return {node.name: node for node in nodes}
+
     def ssubscribe(
         self, *args: ChannelT | Subscription, **kwargs: PubSubHandler
     ) -> None:
@@ -3215,25 +3228,38 @@ def ssubscribe(
         # not mutated concurrently.
         with self._shard_state_lock:
             for s_channel, handler in s_channels.items():
-                node = self.cluster.get_node_from_key(s_channel)
+                slot = self.cluster.keyslot(s_channel)
+                node = self.cluster.nodes_manager.get_node_from_slot(
+                    slot,
+                    self.cluster.read_from_replicas,
+                    self.cluster.load_balancing_strategy,
+                )
                 if not node:
                     continue
                 # Lazy re-route: if this channel is already tracked against a
-                # different node (e.g. after a slot migration), migrate it now
-                # so the caller's intent is applied on the current owner.
+                # different node, migrate it now so the caller's intent is
+                # applied on the current owner — unless the tracked node
+                # still serves the slot, in which case leave it alone.
                 normalized_key = next(iter(self._normalize_keys({s_channel: None})))
                 old_name = self._shard_channel_to_node.get(normalized_key)
                 if old_name and old_name != node.name:
-                    # Match PubSub.ssubscribe() dict.update() semantics: the
-                    # caller's newly supplied handler (including None) always
-                    # overrides any previously registered handler.
-                    self._migrate_shard_channel(
-                        normalized_key,
-                        handler,
-                        old_name,
-                        node,
-                    )
-                    continue
+                    tracked_node = self._eligible_subscription_nodes(slot).get(old_name)
+                    if tracked_node is not None:
+                        # The tracked node still serves the channel's slot
+                        # (the balancer merely picked a sibling); keep the
+                        # subscription where it is to avoid churn.
+                        node = tracked_node
+                    else:
+                        # Match PubSub.ssubscribe() dict.update() semantics: the
+                        # caller's newly supplied handler (including None) always
+                        # overrides any previously registered handler.
+                        self._migrate_shard_channel(
+                            normalized_key,
+                            handler,
+                            old_name,
+                            node,
+                        )
+                        continue
                 pubsub = self._get_node_pubsub(node)
                 if handler:
                     pubsub.ssubscribe(Subscription(s_channel, handler))
@@ -3267,10 +3293,26 @@ def sunsubscribe(self, *args):
                 if name and name in self.node_pubsub_mapping:
                     p = self.node_pubsub_mapping[name]
                 else:
-                    node = self.cluster.get_node_from_key(s_channel)
-                    if not node or node.name not in self.node_pubsub_mapping:
+                    # No reverse-index entry: find the node whose pubsub
+                    # actually holds the subscription (with replica reads
+                    # enabled that is not necessarily the primary).
+                    slot = self.cluster.keyslot(s_channel)
+                    eligible = self._eligible_subscription_nodes(slot)
+                    if not eligible:
+                        raise SlotNotCoveredError(
+                            f'Slot "{slot}" is not covered by the cluster.'
+                        )
+                    p = None
+                    for name in eligible:
+                        candidate = self.node_pubsub_mapping.get(name)
+                        if (
+                            candidate is not None
+                            and normalized_key in candidate.shard_channels
+                        ):
+                            p = candidate
+                            break
+                    if p is None:
                         continue
-                    p = self.node_pubsub_mapping[node.name]
                 p.sunsubscribe(s_channel)
                 self.pending_unsubscribe_shard_channels.update(
                     p.pending_unsubscribe_shard_channels
@@ -3279,18 +3321,28 @@ def sunsubscribe(self, *args):
     def reinitialize_shard_subscriptions(self):
         """
         Reconcile per-node shard subscriptions against the cluster's current
-        slot ownership map. For each tracked shard channel whose owning node
-        has changed (e.g. after CLUSTER SETSLOT / failover), sunsubscribe on
-        the old node's pubsub and ssubscribe on the new owner's pubsub,
-        preserving any registered handler.
+        slot ownership map. For each tracked shard channel whose node no
+        longer serves the channel's slot (e.g. after CLUSTER SETSLOT /
+        failover), sunsubscribe on the old node's pubsub and ssubscribe on a
+        node currently serving the slot, preserving any registered handler.
         """
         uncovered: list = []
         made_progress = False
         first_migrate_error: Optional[BaseException] = None
         with self._shard_state_lock:
             for channel, handler in list(self.shard_channels.items()):
+                slot = self.cluster.keyslot(channel)
+                old_name = self._shard_channel_to_node.get(channel)
+                if old_name and old_name in self._eligible_subscription_nodes(slot):
+                    # The tracked node still serves the channel's slot under
+                    # the current routing configuration; nothing to do.
+                    continue
                 try:
-                    new_node = self.cluster.get_node_from_key(channel)
+                    new_node = self.cluster.nodes_manager.get_node_from_slot(
+                        slot,
+                        self.cluster.read_from_replicas,
+                        self.cluster.load_balancing_strategy,
+                    )
                 except SlotNotCoveredError:
                     # Slot is transiently uncovered (mid-migration / partial
                     # topology refresh). Defer this channel so coverable
@@ -3300,9 +3352,6 @@ def reinitialize_shard_subscriptions(self):
                     # slots-cache change notification.
                     uncovered.append(channel)
                     continue
-                old_name = self._shard_channel_to_node.get(channel)
-                if old_name == new_node.name:
-                    continue
                 try:
                     self._migrate_shard_channel(channel, handler, old_name, new_node)
                     made_progress = True