Introduce SharedSubscribe for state management.

Author: ioquatix

lib/async/redis/cluster_client.rb

@@ -5,6 +5,7 @@
 # Copyright, 2025, by Travis Bell.
 
 require_relative "client"
+require_relative "context/shard_subscribe"
 require "io/stream"
 
 module Async
@@ -138,8 +139,7 @@ def client_for(slot, role = :master)
 			end
 
 			# Get any available client from the cluster.
-			# This is useful for operations that don't require slot-specific routing,
-			# such as global pub/sub operations, INFO commands, or other cluster-wide operations.
+			# This is useful for operations that don't require slot-specific routing, such as global pub/sub operations, INFO commands, or other cluster-wide operations.
 			# @parameter role [Symbol] The role of node to get (:master or :slave).
 			# @returns [Client] A Redis client for any available node.
 			def any_client(role = :master)
@@ -160,15 +160,6 @@ def any_client(role = :master)
 				client_for(0, role)
 			end
 
-			# Execute a Redis command on any available cluster node.
-			# This is useful for commands that don't require slot-specific routing.
-			# @parameter command [String] The Redis command to execute.
-			# @parameter arguments [Array] The command arguments.
-			# @returns [Object] The result of the Redis command.
-			def call(command, *arguments)
-				any_client.call(command, *arguments)
-			end
-			
 			protected
 
 			def reload_cluster!(endpoints = @endpoints)
@@ -339,23 +330,24 @@ def psubscribe(*patterns)
 			end
 
 			# Subscribe to one or more sharded channels for pub/sub messaging in cluster environment (Redis 7.0+).
-			# The subscription will be created on the node responsible for the channel's hash slot.
+			# The subscription will be created on the appropriate nodes responsible for each channel's hash slot.
 			# @parameter channels [Array(String)] The sharded channels to subscribe to.
 			# @yields {|context| ...} If a block is given, it will be executed within the subscription context.
-			# 	@parameter context [Context::Subscribe] The subscription context.
+			# 	@parameter context [Context::ShardSubscribe] The shard subscription context.
 			# @returns [Object] The result of the block if block given.
-			# @returns [Context::Subscribe] The subscription context if no block given.
+			# @returns [Context::ShardSubscribe] The shard subscription context if no block given.
 			def ssubscribe(*channels)
-				# For sharded subscriptions, route to appropriate node based on channel hash
-				slot = channels.any? ? slot_for(channels.first) : 0
-				client = client_for(slot)
+				context = Context::ShardSubscribe.new(self)
+				context.subscribe(channels) if channels.any?
 
-				client.ssubscribe(*channels) do |context|
-					if block_given?
+				if block_given?
+					begin
 						yield context
-					else
-						return context
+					ensure
+						context.close
 					end
+				else
+					return context
 				end
 			end
 		end

lib/async/redis/context/shard_subscribe.rb

@@ -0,0 +1,122 @@
+# frozen_string_literal: true
+
+# Released under the MIT License.
+# Copyright, 2025, by Samuel Williams.
+
+require_relative "subscribe"
+
+module Async
+	module Redis
+		module Context
+			# Context for managing sharded subscriptions across multiple Redis cluster nodes.
+			# This class handles the complexity of subscribing to channels that may be distributed
+			# across different shards in a Redis cluster.
+			class ShardSubscribe
+				# Initialize a new shard subscription context.
+				# @parameter cluster_client [ClusterClient] The cluster client to use.
+				def initialize(cluster_client)
+					@cluster_client = cluster_client
+					@subscriptions = {}
+					@channels = []
+				end
+				
+				# Close all shard subscriptions.
+				def close
+					@subscriptions.each_value(&:close)
+					@subscriptions.clear
+				end
+				
+				# Listen for the next message from any subscribed shard.
+				# This uses a simple round-robin approach to check each shard.
+				# @returns [Array] The next message response, or nil if all connections closed.
+				def listen
+					return nil if @subscriptions.empty?
+					
+					# Simple round-robin checking of subscriptions
+					@subscriptions.each_value do |subscription|
+						# Non-blocking check for messages
+						begin
+							if response = subscription.listen
+								return response
+							end
+						rescue => error
+							# Handle connection errors gracefully
+							Console.warn(self, "Error reading from shard subscription: #{error}")
+						end
+					end
+					
+					# If no immediate messages, do a blocking wait on the first subscription
+					if first_subscription = @subscriptions.values.first
+						first_subscription.listen
+					end
+				end
+				
+				# Iterate over all messages from all subscribed shards.
+				# @yields {|response| ...} Block called for each message.
+				# 	@parameter response [Array] The message response.
+				def each
+					return to_enum unless block_given?
+					
+					while response = self.listen
+						yield response
+					end
+				end
+				
+				# Subscribe to additional sharded channels.
+				# @parameter channels [Array(String)] The channels to subscribe to.
+				def subscribe(channels)
+					slots = @cluster_client.slots_for(channels)
+					
+					slots.each do |slot, channels_for_slot|
+						if subscription = @subscriptions[slot]
+							# Add to existing subscription for this shard
+							subscription.ssubscribe(channels_for_slot)
+						else
+							# Create new subscription for this shard
+							client = @cluster_client.client_for(slot)
+							@subscriptions[slot] = client.ssubscribe(*channels_for_slot)
+						end
+					end
+					
+					@channels.concat(channels)
+				end
+				
+				# Unsubscribe from sharded channels.
+				# @parameter channels [Array(String)] The channels to unsubscribe from.
+				def unsubscribe(channels)
+					slots = @cluster_client.slots_for(channels)
+					
+					slots.each do |slot, channels_for_slot|
+						if subscription = @subscriptions[slot]
+							subscription.sunsubscribe(channels_for_slot)
+							
+							# Remove channels from our tracking
+							@channels -= channels_for_slot
+							
+							# Check if this shard still has channels
+							remaining_channels_for_slot = @channels.select { |ch| @cluster_client.slot_for(ch) == slot }
+							
+							# If no channels left for this shard, close and remove it
+							if remaining_channels_for_slot.empty?
+								subscription.close
+								@subscriptions.delete(slot)
+							end
+						end
+					end
+				end
+				
+				# Get the list of currently subscribed channels.
+				# @returns [Array(String)] The list of subscribed channels.
+				def channels
+					@channels.dup
+				end
+				
+				# Get the number of active shard subscriptions.
+				# @returns [Integer] The number of shard connections.
+				def shard_count
+					@subscriptions.size
+				end
+			end
+		end
+	end
+end