feat: add maxlen to stream sentinel and cluster brokers

Author: spikeninja

taskiq_redis/redis_cluster_broker.py

@@ -92,6 +92,7 @@ def __init__(
         consumer_id: str = "$",
         mkstream: bool = True,
         xread_block: int = 10000,
+        maxlen: Optional[int] = None,
         additional_streams: Optional[Dict[str, str]] = None,
         **connection_kwargs: Any,
     ) -> None:
@@ -111,6 +112,8 @@ def __init__(
         :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 additional_streams: additional streams to read from.
             Each key is a stream name, value is a consumer id.
         """
@@ -125,6 +128,7 @@ def __init__(
         self.consumer_id = consumer_id
         self.mkstream = mkstream
         self.block = xread_block
+        self.maxlen = maxlen
         self.additional_streams = additional_streams or {}
 
     async def _declare_consumer_group(self) -> None:
@@ -154,7 +158,11 @@ async def kick(self, message: BrokerMessage) -> None:
 
         :param message: message to append.
         """
-        await self.redis.xadd(self.queue_name, {b"data": message.message})
+        await self.redis.xadd(
+            self.queue_name,
+            {b"data": message.message},
+            maxlen=self.maxlen,
+        )
 
     def _ack_generator(self, id: str) -> Callable[[], Awaitable[None]]:
         async def _ack() -> None:

taskiq_redis/redis_sentinel_broker.py

@@ -157,6 +157,7 @@ def __init__(
         consumer_id: str = "$",
         mkstream: bool = True,
         xread_block: int = 10000,
+        maxlen: Optional[int] = None,
         additional_streams: Optional[Dict[str, str]] = None,
         **connection_kwargs: Any,
     ) -> None:
@@ -176,6 +177,8 @@ def __init__(
         :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 additional_streams: additional streams to read from.
             Each key is a stream name, value is a consumer id.
         """
@@ -193,6 +196,7 @@ def __init__(
         self.consumer_id = consumer_id
         self.mkstream = mkstream
         self.block = xread_block
+        self.maxlen = maxlen
         self.additional_streams = additional_streams or {}
 
     async def _declare_consumer_group(self) -> None:
@@ -223,7 +227,11 @@ async def kick(self, message: BrokerMessage) -> None:
         :param message: message to append.
         """
         async with self._acquire_master_conn() as redis_conn:
-            await redis_conn.xadd(self.queue_name, {b"data": message.message})
+            await redis_conn.xadd(
+                self.queue_name,
+                {b"data": message.message},
+                maxlen=self.maxlen,
+            )
 
     def _ack_generator(self, id: str) -> Callable[[], Awaitable[None]]:
         async def _ack() -> None: