feat: Added PubSubSentinelBroker.

Author: DABND19

taskiq_redis/__init__.py

@@ -6,7 +6,10 @@
 )
 from taskiq_redis.redis_broker import ListQueueBroker, PubSubBroker
 from taskiq_redis.redis_cluster_broker import ListQueueClusterBroker
-from taskiq_redis.redis_sentinel_broker import ListQueueSentinelBroker
+from taskiq_redis.redis_sentinel_broker import (
+    ListQueueSentinelBroker,
+    PubSubSentinelBroker,
+)
 from taskiq_redis.schedule_source import (
     RedisClusterScheduleSource,
     RedisScheduleSource,
@@ -21,6 +24,7 @@
     "PubSubBroker",
     "ListQueueClusterBroker",
     "ListQueueSentinelBroker",
+    "PubSubSentinelBroker",
     "RedisScheduleSource",
     "RedisClusterScheduleSource",
     "RedisSentinelScheduleSource",

taskiq_redis/redis_sentinel_broker.py

@@ -66,6 +66,40 @@ async def _acquire_master_conn(self) -> AsyncIterator[_Redis]:
             yield redis_conn
 
 
+class PubSubSentinelBroker(BaseSentinelBroker):
+    """Broker that works with Redis and broadcasts tasks to all workers."""
+
+    async def kick(self, message: BrokerMessage) -> None:
+        """
+        Publish message over PUBSUB channel.
+
+        :param message: message to send.
+        """
+        queue_name = message.labels.get("queue_name") or self.queue_name
+        async with self._acquire_master_conn() as redis_conn:
+            await redis_conn.publish(queue_name, message.message)
+
+    async def listen(self) -> AsyncGenerator[bytes, None]:
+        """
+        Listen redis queue for new messages.
+
+        This function listens to the pubsub channel
+        and yields all messages with proper types.
+
+        :yields: broker messages.
+        """
+        async with self._acquire_master_conn() as redis_conn:
+            redis_pubsub_channel = redis_conn.pubsub()
+            await redis_pubsub_channel.subscribe(self.queue_name)
+            async for message in redis_pubsub_channel.listen():
+                if not message:
+                    continue
+                if message["type"] != "message":
+                    logger.debug("Received non-message from redis: %s", message)
+                    continue
+                yield message["data"]
+
+
 class ListQueueSentinelBroker(BaseSentinelBroker):
     """Broker that works with Sentinel and distributes tasks between workers."""
 

tests/test_broker.py

@@ -10,6 +10,7 @@
     ListQueueClusterBroker,
     ListQueueSentinelBroker,
     PubSubBroker,
+    PubSubSentinelBroker,
 )
 
 
@@ -176,6 +177,37 @@ async def test_list_queue_cluster_broker(
     await broker.shutdown()
 
 
+@pytest.mark.anyio
+async def test_pub_sub_sentinel_broker(
+    valid_broker_message: BrokerMessage,
+    redis_sentinels: List[Tuple[str, int]],
+    redis_sentinel_master_name: str,
+) -> None:
+    """
+    Test that messages are published and read correctly by PubSubBroker.
+
+    We create two workers that listen and send a message to them.
+    Expect both workers to receive the same message we sent.
+    """
+    broker = PubSubSentinelBroker(
+        sentinels=redis_sentinels,
+        master_name=redis_sentinel_master_name,
+        queue_name=uuid.uuid4().hex,
+    )
+    worker1_task = asyncio.create_task(get_message(broker))
+    worker2_task = asyncio.create_task(get_message(broker))
+    await asyncio.sleep(0.3)
+
+    await broker.kick(valid_broker_message)
+    await asyncio.sleep(0.3)
+
+    message1 = worker1_task.result()
+    message2 = worker2_task.result()
+    assert message1 == valid_broker_message.message
+    assert message1 == message2
+    await broker.shutdown()
+
+
 @pytest.mark.anyio
 async def test_list_queue_sentinel_broker(
     valid_broker_message: BrokerMessage,