Merge pull request taskiq-python#55 from DABND19/feature/redis-sentinel

Author: s3rius

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

@@ -2,20 +2,30 @@
 from taskiq_redis.redis_backend import (
  RedisAsyncClusterResultBackend,
  RedisAsyncResultBackend,
+ RedisAsyncSentinelResultBackend,
 )
 from taskiq_redis.redis_broker import ListQueueBroker, PubSubBroker
 from taskiq_redis.redis_cluster_broker import ListQueueClusterBroker
+from taskiq_redis.redis_sentinel_broker import (
+ ListQueueSentinelBroker,
+ PubSubSentinelBroker,
+)
 from taskiq_redis.schedule_source import (
  RedisClusterScheduleSource,
  RedisScheduleSource,
+ RedisSentinelScheduleSource,
 )
 
 __all__ = [
  "RedisAsyncClusterResultBackend",
  "RedisAsyncResultBackend",
+ "RedisAsyncSentinelResultBackend",
  "ListQueueBroker",
  "PubSubBroker",
  "ListQueueClusterBroker",
+ "ListQueueSentinelBroker",
+ "PubSubSentinelBroker",
  "RedisScheduleSource",
  "RedisClusterScheduleSource",
+ "RedisSentinelScheduleSource",
 ]

taskiq_redis/redis_backend.py

@@ -1,16 +1,40 @@
 import pickle
-from typing import Any, Dict, Optional, TypeVar, Union
+import sys
+from contextlib import asynccontextmanager
+from typing import (
+ TYPE_CHECKING,
+ Any,
+ AsyncIterator,
+ Dict,
+ List,
+ Optional,
+ Tuple,
+ TypeVar,
+ Union,
+)
 
-from redis.asyncio import BlockingConnectionPool, Redis
+from redis.asyncio import BlockingConnectionPool, Redis, Sentinel
 from redis.asyncio.cluster import RedisCluster
 from taskiq import AsyncResultBackend
 from taskiq.abc.result_backend import TaskiqResult
+from taskiq.abc.serializer import TaskiqSerializer
 
 from taskiq_redis.exceptions import (
  DuplicateExpireTimeSelectedError,
  ExpireTimeMustBeMoreThanZeroError,
  ResultIsMissingError,
 )
+from taskiq_redis.serializer import PickleSerializer
+
+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
 
 _ReturnType = TypeVar("_ReturnType")
 
@@ -267,3 +291,142 @@ async def get_result(
  taskiq_result.log = None
 
  return taskiq_result
+
+
+class RedisAsyncSentinelResultBackend(AsyncResultBackend[_ReturnType]):
+ """Async result based on redis sentinel."""
+
+ def __init__(
+ self,
+ sentinels: List[Tuple[str, int]],
+ master_name: str,
+ keep_results: bool = True,
+ result_ex_time: Optional[int] = None,
+ result_px_time: Optional[int] = None,
+ min_other_sentinels: int = 0,
+ sentinel_kwargs: Optional[Any] = None,
+ serializer: Optional[TaskiqSerializer] = None,
+ **connection_kwargs: Any,
+ ) -> None:
+ """
+ Constructs a new result backend.
+
+ :param sentinels: list of sentinel host and ports pairs.
+ :param master_name: sentinel master name.
+ :param keep_results: flag to not remove results from Redis after reading.
+ :param result_ex_time: expire time in seconds for result.
+ :param result_px_time: expire time in milliseconds for result.
+ :param max_connection_pool_size: maximum number of connections in pool.
+ :param connection_kwargs: additional arguments for redis BlockingConnectionPool.
+
+ :raises DuplicateExpireTimeSelectedError: if result_ex_time
+ and result_px_time are selected.
+ :raises ExpireTimeMustBeMoreThanZeroError: if result_ex_time
+ and result_px_time are equal zero.
+ """
+ self.sentinel = Sentinel(
+ sentinels=sentinels,
+ min_other_sentinels=min_other_sentinels,
+ sentinel_kwargs=sentinel_kwargs,
+ **connection_kwargs,
+ )
+ self.master_name = master_name
+ if serializer is None:
+ serializer = PickleSerializer()
+ self.serializer = serializer
+ self.keep_results = keep_results
+ self.result_ex_time = result_ex_time
+ self.result_px_time = result_px_time
+
+ unavailable_conditions = any(
+ (
+ self.result_ex_time is not None and self.result_ex_time <= 0,
+ self.result_px_time is not None and self.result_px_time <= 0,
+ ),
+ )
+ if unavailable_conditions:
+ raise ExpireTimeMustBeMoreThanZeroError(
+ "You must select one expire time param and it must be more than zero.",
+ )
+
+ if self.result_ex_time and self.result_px_time:
+ raise DuplicateExpireTimeSelectedError(
+ "Choose either result_ex_time or result_px_time.",
+ )
+
+ @asynccontextmanager
+ async def _acquire_master_conn(self) -> AsyncIterator[_Redis]:
+ async with self.sentinel.master_for(self.master_name) as redis_conn:
+ yield redis_conn
+
+ async def set_result(
+ self,
+ task_id: str,
+ result: TaskiqResult[_ReturnType],
+ ) -> None:
+ """
+ Sets task result in redis.
+
+ Dumps TaskiqResult instance into the bytes and writes
+ it to redis.
+
+ :param task_id: ID of the task.
+ :param result: TaskiqResult instance.
+ """
+ redis_set_params: Dict[str, Union[str, bytes, int]] = {
+ "name": task_id,
+ "value": self.serializer.dumpb(result),
+ }
+ if self.result_ex_time:
+ redis_set_params["ex"] = self.result_ex_time
+ elif self.result_px_time:
+ redis_set_params["px"] = self.result_px_time
+
+ async with self._acquire_master_conn() as redis:
+ await redis.set(**redis_set_params) # type: ignore
+
+ async def is_result_ready(self, task_id: str) -> bool:
+ """
+ Returns whether the result is ready.
+
+ :param task_id: ID of the task.
+
+ :returns: True if the result is ready else False.
+ """
+ async with self._acquire_master_conn() as redis:
+ return bool(await redis.exists(task_id))
+
+ async def get_result(
+ self,
+ task_id: str,
+ with_logs: bool = False,
+ ) -> TaskiqResult[_ReturnType]:
+ """
+ Gets result from the task.
+
+ :param task_id: task's id.
+ :param with_logs: if True it will download task's logs.
+ :raises ResultIsMissingError: if there is no result when trying to get it.
+ :return: task's return value.
+ """
+ async with self._acquire_master_conn() as redis:
+ if self.keep_results:
+ result_value = await redis.get(
+ name=task_id,
+ )
+ else:
+ result_value = await redis.getdel(
+ name=task_id,
+ )
+
+ if result_value is None:
+ raise ResultIsMissingError
+
+ taskiq_result: TaskiqResult[_ReturnType] = pickle.loads( # noqa: S301
+ result_value,
+ )
+
+ if not with_logs:
+ taskiq_result.log = None
+
+ return taskiq_result

taskiq_redis/redis_sentinel_broker.py

@@ -0,0 +1,132 @@
+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 PubSubSentinelBroker(BaseSentinelBroker):
+ """Broker that works with Sentinel 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."""
+
+ 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
+ ]