redis_broker.py

import uuid
from collections.abc import AsyncGenerator, Awaitable, Callable
from logging import getLogger
from typing import (
    TYPE_CHECKING,
    Any,
    TypeAlias,
    TypeVar,
)

from redis.asyncio import BlockingConnectionPool, Connection, Redis, ResponseError
from taskiq import AckableMessage
from taskiq.abc.broker import AsyncBroker
from taskiq.abc.result_backend import AsyncResultBackend
from taskiq.message import BrokerMessage

_T = TypeVar("_T")

logger = getLogger("taskiq.redis_broker")


if TYPE_CHECKING:
    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool[Connection]  # type: ignore
else:
    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool


class BaseRedisBroker(AsyncBroker):
    """Base broker that works with Redis."""

    def __init__(
        self,
        url: str,
        task_id_generator: Callable[[], str] | None = None,
        result_backend: AsyncResultBackend[_T] | None = None,
        queue_name: str = "taskiq",
        max_connection_pool_size: int | None = None,
        **connection_kwargs: Any,
    ) -> None:
        """
        Constructs a new broker.

        :param url: url to redis.
        :param task_id_generator: custom task_id generator.
        :param result_backend: custom result backend.
        :param queue_name: name for a list in redis.
        :param max_connection_pool_size: maximum number of connections in pool.
            Each worker opens its own connection. Therefore this value has to be
            at least number of workers + 1.
        :param connection_kwargs: additional arguments for redis BlockingConnectionPool.
        """
        super().__init__(
            result_backend=result_backend,
            task_id_generator=task_id_generator,
        )

        self.connection_pool: _BlockingConnectionPool = BlockingConnectionPool.from_url(
            url=url,
            max_connections=max_connection_pool_size,
            **connection_kwargs,
        )
        self.queue_name = queue_name

    async def shutdown(self) -> None:
        """Closes redis connection pool."""
        await super().shutdown()
        await self.connection_pool.disconnect()


class PubSubBroker(BaseRedisBroker):
    """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 Redis(connection_pool=self.connection_pool) 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 Redis(connection_pool=self.connection_pool) 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 ListQueueBroker(BaseRedisBroker):
    """Broker that works with Redis 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 Redis(connection_pool=self.connection_pool) as redis_conn:
            await redis_conn.lpush(queue_name, message.message)  # type: ignore

    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
        while True:
            try:
                async with Redis(connection_pool=self.connection_pool) as redis_conn:
                    yield (await redis_conn.brpop(self.queue_name))[  # type: ignore
                        redis_brpop_data_position
                    ]
            except ConnectionError as exc:
                logger.warning("Redis connection error: %s", exc)
                continue


class RedisStreamBroker(BaseRedisBroker):
    """
    Redis broker that uses streams for task distribution.

    You can read more about streams here:
    https://redis.io/docs/latest/develop/data-types/streams

    This broker supports acknowledgment of messages.
    """

    def __init__(
        self,
        url: str,
        queue_name: str = "taskiq",
        max_connection_pool_size: int | None = None,
        consumer_group_name: str = "taskiq",
        consumer_name: str | None = None,
        consumer_id: str = "$",
        mkstream: bool = True,
        xread_block: int = 2000,
        maxlen: int | None = None,
        approximate: bool = True,
        idle_timeout: int = 600000,  # 10 minutes
        unacknowledged_batch_size: int = 100,
        unacknowledged_lock_timeout: float | None = None,
        xread_count: int | None = 100,
        additional_streams: dict[str, str | int] | None = None,
        **connection_kwargs: Any,
    ) -> None:
        """
        Constructs a new broker that uses streams.

        :param url: url to redis.
        :param queue_name: name for a key with stream in redis.
        :param max_connection_pool_size: maximum number of connections in pool.
            Each worker opens its own connection. Therefore this value has to be
            at least number of workers + 1.
        :param consumer_group_name: name for a consumer group.
            Redis will keep track of acked messages for this group.
        :param consumer_name: name for a consumer. By default it is a random uuid.
        :param consumer_id: id for a consumer. ID of a message to start reading from.
            $ means start from the latest message.
        :param mkstream: create stream if it does not exist.
        :param xread_block: block time in ms for xreadgroup.
            Better to set it to a bigger value, to avoid unnecessary calls.
        :param maxlen: sets the maximum length of the stream
            trims (the old values of) the stream each time a new element is added
        :param approximate: decides wether to trim the stream immediately (False) or
            later on (True)
        :param xread_count: number of messages to fetch from the stream at once.
        :param additional_streams: additional streams to read from.
            Each key is a stream name, value is a consumer id.
        :param unacknowledged_batch_size: number of unacknowledged messages to fetch.
        :param unacknowledged_lock_timeout: time in seconds before auto-releasing
            the lock. Useful when the worker crashes or gets killed.
            If not set, the lock can remain locked indefinitely.
        """
        super().__init__(
            url,
            task_id_generator=None,
            result_backend=None,
            queue_name=queue_name,
            max_connection_pool_size=max_connection_pool_size,
            **connection_kwargs,
        )
        self.consumer_group_name = consumer_group_name
        self.consumer_name = consumer_name or str(uuid.uuid4())
        self.consumer_id = consumer_id
        self.mkstream = mkstream
        self.block = xread_block
        self.maxlen = maxlen
        self.approximate = approximate
        self.additional_streams = additional_streams or {}
        self.idle_timeout = idle_timeout
        self.unacknowledged_batch_size = unacknowledged_batch_size
        self.unacknowledged_lock_timeout = unacknowledged_lock_timeout
        self.count = xread_count

    async def _declare_consumer_group(self) -> None:
        """
        Declare consumber group.

        Required for proper work of the broker.
        """
        streams = {self.queue_name, *self.additional_streams.keys()}
        async with Redis(connection_pool=self.connection_pool) as redis_conn:
            for stream_name in streams:
                try:
                    await redis_conn.xgroup_create(
                        stream_name,
                        self.consumer_group_name,
                        id=self.consumer_id,
                        mkstream=self.mkstream,
                    )
                except ResponseError as err:
                    logger.debug(err)

    async def startup(self) -> None:
        """Declare consumer group on startup."""
        await super().startup()
        await self._declare_consumer_group()

    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 Redis(connection_pool=self.connection_pool) as redis_conn:
            await redis_conn.xadd(
                queue_name,
                {b"data": message.message},
                maxlen=self.maxlen,
                approximate=self.approximate,
            )

    def _ack_generator(self, id: str, queue_name: str) -> Callable[[], Awaitable[None]]:
        async def _ack() -> None:
            async with Redis(connection_pool=self.connection_pool) as redis_conn:
                await redis_conn.xack(
                    queue_name,
                    self.consumer_group_name,
                    id,
                )

        return _ack

    async def listen(self) -> AsyncGenerator[AckableMessage, None]:
        """Listen to incoming messages."""
        async with Redis(connection_pool=self.connection_pool) as redis_conn:
            while True:
                logger.debug("Starting fetching new messages")
                fetched = await redis_conn.xreadgroup(
                    self.consumer_group_name,
                    self.consumer_name,
                    {
                        self.queue_name: ">",
                        **self.additional_streams,  # type: ignore[dict-item]
                    },
                    block=self.block,
                    noack=False,
                    count=self.count,
                )
                for stream, msg_list in fetched:
                    for msg_id, msg in msg_list:
                        logger.debug("Received message: %s", msg)
                        yield AckableMessage(
                            data=msg[b"data"],
                            ack=self._ack_generator(id=msg_id, queue_name=stream),
                        )
                logger.debug("Starting fetching unacknowledged messages")
                for stream in [self.queue_name, *self.additional_streams.keys()]:
                    lock = redis_conn.lock(
                        f"autoclaim:{self.consumer_group_name}:{stream}",
                        timeout=self.unacknowledged_lock_timeout,
                    )
                    if await lock.locked():
                        continue
                    async with lock:
                        pending = await redis_conn.xautoclaim(
                            name=stream,
                            groupname=self.consumer_group_name,
                            consumername=self.consumer_name,
                            min_idle_time=self.idle_timeout,
                            count=self.unacknowledged_batch_size,
                        )
                        logger.debug(
                            "Found %d pending messages in stream %s",
                            len(pending[1]),
                            stream,
                        )
                        for msg_id, msg in pending[1]:
                            logger.debug("Received message: %s", msg)
                            yield AckableMessage(
                                data=msg[b"data"],
                                ack=self._ack_generator(id=msg_id, queue_name=stream),
                            )