feat: Added ListQueueSentinelBroker.

Author: DABND19

docker-compose.yml

@@ -58,3 +58,24 @@ services:
       REDIS_CLUSTER_CREATOR: "yes"
     ports:
       - 7001:6379
+
+  redis-master:
+    image: bitnami/redis:6.2.5
+    environment:
+      ALLOW_EMPTY_PASSWORD: "yes"
+    healthcheck:
+      test: ["CMD", "redis-cli", "ping"]
+      interval: 5s
+      timeout: 5s
+      retries: 3
+      start_period: 10s
+
+  redis-sentinel:
+    image: bitnami/redis-sentinel:latest
+    depends_on:
+      - redis-master
+    environment:
+      ALLOW_EMPTY_PASSWORD: "yes"
+      REDIS_MASTER_HOST: "redis-master"
+    ports:
+      - 7002:26379

taskiq_redis/__init__.py

@@ -5,6 +5,7 @@
 )
 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.schedule_source import (
     RedisClusterScheduleSource,
     RedisScheduleSource,
@@ -16,6 +17,7 @@
     "ListQueueBroker",
     "PubSubBroker",
     "ListQueueClusterBroker",
+    "ListQueueSentinelBroker",
     "RedisScheduleSource",
     "RedisClusterScheduleSource",
 ]

taskiq_redis/redis_sentinel_broker.py

@@ -0,0 +1,98 @@
+import sys
+from contextlib import asynccontextmanager
+from logging import getLogger
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncGenerator,
+    AsyncIterator,
+    Callable,
+    List,
+    Optional,
+    Tuple,
+    TypeVar,
+)
+
+from redis.asyncio import Redis, Sentinel
+from taskiq import AsyncResultBackend, BrokerMessage
+from taskiq.abc.broker import AsyncBroker
+
+if sys.version_info >= (3, 10):
+    from typing import TypeAlias
+else:
+    from typing_extensions import TypeAlias
+
+if TYPE_CHECKING:
+    _Redis: TypeAlias = Redis[bytes]
+else:
+    _Redis: TypeAlias = Redis
+
+_T = TypeVar("_T")
+
+logger = getLogger("taskiq.redis_sentinel_broker")
+
+
+class BaseSentinelBroker(AsyncBroker):
+    """Base broker that works with Sentinel."""
+
+    def __init__(
+        self,
+        sentinels: List[Tuple[str, int]],
+        master_name: str,
+        result_backend: Optional[AsyncResultBackend[_T]] = None,
+        task_id_generator: Optional[Callable[[], str]] = None,
+        queue_name: str = "taskiq",
+        min_other_sentinels: int = 0,
+        sentinel_kwargs: Optional[Any] = None,
+        **connection_kwargs: Any,
+    ) -> None:
+        super().__init__(
+            result_backend=result_backend,
+            task_id_generator=task_id_generator,
+        )
+
+        self.sentinel = Sentinel(
+            sentinels=sentinels,
+            min_other_sentinels=min_other_sentinels,
+            sentinel_kwargs=sentinel_kwargs,
+            **connection_kwargs,
+        )
+        self.master_name = master_name
+        self.queue_name = queue_name
+
+    @asynccontextmanager
+    async def _acquire_master_conn(self) -> AsyncIterator[_Redis]:
+        async with self.sentinel.master_for(self.master_name) as redis_conn:
+            yield redis_conn
+
+
+class ListQueueSentinelBroker(BaseSentinelBroker):
+    """Broker that works with Sentinel and distributes tasks between workers."""
+
+    async def kick(self, message: BrokerMessage) -> None:
+        """
+        Put a message in a list.
+
+        This method appends a message to the list of all messages.
+
+        :param message: message to append.
+        """
+        queue_name = message.labels.get("queue_name") or self.queue_name
+        async with self._acquire_master_conn() as redis_conn:
+            await redis_conn.lpush(queue_name, message.message)
+
+    async def listen(self) -> AsyncGenerator[bytes, None]:
+        """
+        Listen redis queue for new messages.
+
+        This function listens to the queue
+        and yields new messages if they have BrokerMessage type.
+
+        :yields: broker messages.
+        """
+        redis_brpop_data_position = 1
+        async with self._acquire_master_conn() as redis_conn:
+            while True:
+                yield (await redis_conn.brpop(self.queue_name))[
+                    redis_brpop_data_position
+                ]

tests/conftest.py

@@ -1,4 +1,5 @@
 import os
+from typing import List, Tuple
 
 import pytest
 
@@ -40,3 +41,33 @@ def redis_cluster_url() -> str:
     :return: URL string.
     """
     return os.environ.get("TEST_REDIS_CLUSTER_URL", "redis://localhost:7001")
+
+
+@pytest.fixture
+def redis_sentinels() -> List[Tuple[str, int]]:
+    """
+    List of redis sentinel hosts.
+
+    It tries to get it from environ,
+    and return default one if the variable is
+    not set.
+
+    :return: list of host and port pairs.
+    """
+    sentinels = os.environ.get("TEST_REDIS_SENTINELS", "localhost:7002")
+    host, _, port = sentinels.partition(":")
+    return [(host, int(port))]
+
+
+@pytest.fixture
+def redis_sentinel_master_name() -> str:
+    """
+    Redis sentinel master name.
+
+    It tries to get it from environ,
+    and return default one if the variable is
+    not set.
+
+    :return: redis sentinel master name string.
+    """
+    return os.environ.get("TEST_REDIS_SENTINEL_MASTER_NAME", "mymaster")

tests/test_broker.py

@@ -1,11 +1,16 @@
 import asyncio
 import uuid
-from typing import Union
+from typing import List, Tuple, Union
 
 import pytest
 from taskiq import AckableMessage, AsyncBroker, BrokerMessage
 
-from taskiq_redis import ListQueueBroker, ListQueueClusterBroker, PubSubBroker
+from taskiq_redis import (
+    ListQueueBroker,
+    ListQueueClusterBroker,
+    ListQueueSentinelBroker,
+    PubSubBroker,
+)
 
 
 def test_no_url_should_raise_typeerror() -> None:
@@ -169,3 +174,32 @@ async def test_list_queue_cluster_broker(
     assert worker_task.result() == valid_broker_message.message
     worker_task.cancel()
     await broker.shutdown()
+
+
+@pytest.mark.anyio
+async def test_list_queue_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 ListQueueSentinelBroker.
+
+    We create two workers that listen and send a message to them.
+    Expect only one worker to receive the same message we sent.
+    """
+    broker = ListQueueSentinelBroker(
+        sentinels=redis_sentinels,
+        master_name=redis_sentinel_master_name,
+        queue_name=uuid.uuid4().hex,
+    )
+    worker_task = asyncio.create_task(get_message(broker))
+    await asyncio.sleep(0.3)
+
+    await broker.kick(valid_broker_message)
+    await asyncio.sleep(0.3)
+
+    assert worker_task.done()
+    assert worker_task.result() == valid_broker_message.message
+    worker_task.cancel()
+    await broker.shutdown()