Files
dify/api/libs/broadcast_channel/redis/sharded_channel.py
QuantumGhost 03e3acfc71 feat(api): Human Input Node (backend part) (#31646)
The backend part of the human in the loop (HITL) feature and relevant architecture / workflow engine changes.

Signed-off-by: yihong0618 <zouzou0208@gmail.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: -LAN- <laipz8200@outlook.com>
Co-authored-by: 盐粒 Yanli <yanli@dify.ai>
Co-authored-by: CrabSAMA <40541269+CrabSAMA@users.noreply.github.com>
Co-authored-by: Stephen Zhou <38493346+hyoban@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: yihong <zouzou0208@gmail.com>
Co-authored-by: Joel <iamjoel007@gmail.com>
2026-01-30 10:18:49 +08:00

75 lines
2.3 KiB
Python

from __future__ import annotations
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
# NOTE(QuantumGhost): this is an issue in
# upstream code. If Sharded PubSub is used with Cluster, the
# `ClusterPubSub.get_sharded_message` will return `None` regardless of
# message['type'].
#
# Since we have already filtered at the caller's site, we can safely set
# `ignore_subscribe_messages=False`.
return self._pubsub.get_sharded_message(ignore_subscribe_messages=False, timeout=0.1) # type: ignore[attr-defined]
def _get_message_type(self) -> str:
return "smessage"