dify

dify/api/libs/broadcast_channel/redis/sharded_channel.py

@@ -0,0 +1,65 @@
+from libs.broadcast_channel.channel import Producer, Subscriber, Subscription
+from redis import Redis
+
+from ._subscription import RedisSubscriptionBase
+
+
+class ShardedRedisBroadcastChannel:
+    """
+    Redis 7.0+ Sharded Pub/Sub based broadcast channel implementation.
+
+    Provides "at most once" delivery semantics using SPUBLISH/SSUBSCRIBE commands,
+    distributing channels across Redis cluster nodes for better scalability.
+    """
+
+    def __init__(
+        self,
+        redis_client: Redis,
+    ):
+        self._client = redis_client
+
+    def topic(self, topic: str) -> "ShardedTopic":
+        return ShardedTopic(self._client, topic)
+
+
+class ShardedTopic:
+    def __init__(self, redis_client: Redis, topic: str):
+        self._client = redis_client
+        self._topic = topic
+
+    def as_producer(self) -> Producer:
+        return self
+
+    def publish(self, payload: bytes) -> None:
+        self._client.spublish(self._topic, payload)  # type: ignore[attr-defined]
+
+    def as_subscriber(self) -> Subscriber:
+        return self
+
+    def subscribe(self) -> Subscription:
+        return _RedisShardedSubscription(
+            pubsub=self._client.pubsub(),
+            topic=self._topic,
+        )
+
+
+class _RedisShardedSubscription(RedisSubscriptionBase):
+    """Redis 7.0+ sharded pub/sub subscription implementation."""
+
+    def _get_subscription_type(self) -> str:
+        return "sharded"
+
+    def _subscribe(self) -> None:
+        assert self._pubsub is not None
+        self._pubsub.ssubscribe(self._topic)  # type: ignore[attr-defined]
+
+    def _unsubscribe(self) -> None:
+        assert self._pubsub is not None
+        self._pubsub.sunsubscribe(self._topic)  # type: ignore[attr-defined]
+
+    def _get_message(self) -> dict | None:
+        assert self._pubsub is not None
+        return self._pubsub.get_sharded_message(ignore_subscribe_messages=True, timeout=0.1)  # type: ignore[attr-defined]
+
+    def _get_message_type(self) -> str:
+        return "smessage"