fix Redis stream reclaim edge cases

Author: vvanglro

README.md

@@ -113,8 +113,6 @@ For a single-stream broker, `xread_count` also prevents one worker from hoarding
 
 When a listener is closed with messages already fetched from redis but not yet yielded to taskiq, those buffered entries are claimed to an internal `abandoned` consumer and stamped as very idle. The next reclaim sweep can recover them immediately instead of waiting for `idle_timeout`.
 
-If the consumer group is removed out of band, the broker recreates it on the next `XREADGROUP` so an in-use queue heals instead of spinning on errors.
-
 The legacy `unacknowledged_lock_timeout` parameter is deprecated and ignored; reclaim correctness now relies on the `XCLAIM` min-idle-time check instead of a broker-side redis lock.
 
 ## RedisAsyncResultBackend configuration

taskiq_redis/redis_broker.py

@@ -13,7 +13,13 @@
     cast,
 )
 
-from redis.asyncio import BlockingConnectionPool, Connection, Redis, ResponseError
+from redis.asyncio import (
+    BlockingConnectionPool,
+    Connection,
+    Redis,
+    RedisError,
+    ResponseError,
+)
 from taskiq import AckableMessage
 from taskiq.abc.broker import AsyncBroker
 from taskiq.abc.result_backend import AsyncResultBackend
@@ -327,16 +333,15 @@ def _ack_generator(
         async def _ack() -> None:
             nonlocal acked
             async with Redis(connection_pool=self.connection_pool) as redis_conn:
-                try:
-                    await redis_conn.xack(
-                        queue_name,
-                        self.consumer_group_name,
-                        id,
-                    )
-                finally:
-                    if not acked and on_ack is not None:
-                        acked = True
-                        on_ack()
+                await redis_conn.xack(
+                    queue_name,
+                    self.consumer_group_name,
+                    id,
+                )
+            if not acked:
+                acked = True
+                if on_ack is not None:
+                    on_ack()
 
         return _ack
 
@@ -380,45 +385,39 @@ async def _claim_timed_out_messages(
         stream: str,
         count: int,
         protected_message_ids: set[tuple[str, str]],
-    ) -> list[tuple[str, dict[bytes, bytes]]]:
+        pending_start: str,
+    ) -> tuple[list[tuple[str, dict[bytes, bytes]]], str]:
         """Claim pending messages that exceeded their reclaim deadline.
 
         protected_message_ids contains entries delivered by this listener
         instance and not acked yet. We skip those to avoid re-delivering our own
         in-flight work, while still allowing a restarted worker with the same
         Redis consumer name to recover messages left by its predecessor.
         """
-        try:
-            pending = await redis_conn.xpending_range(
-                stream,
-                self.consumer_group_name,
-                min="-",
-                max="+",
-                count=self.unacknowledged_batch_size,
-                idle=0,
-            )
-        except ResponseError as exc:
-            if "NOGROUP" not in str(exc):
-                raise
-            logger.info("Consumer group missing for %s, recreating", stream)
-            await self._declare_consumer_group()
-            return []
+        pending = await redis_conn.xpending_range(
+            stream,
+            self.consumer_group_name,
+            min=pending_start,
+            max="+",
+            count=self.unacknowledged_batch_size,
+            idle=0,
+        )
 
         claimed: list[tuple[str, dict[bytes, bytes]]] = []
+        last_checked_id: str | None = None
         for pending_message in pending:
             if len(claimed) >= count:
                 break
 
-            message_id = pending_message["message_id"]
-            if isinstance(message_id, bytes):
-                message_id = message_id.decode()
+            message_id = self._to_str(pending_message["message_id"])
+            last_checked_id = message_id
             if self._message_key(stream, message_id) in protected_message_ids:
                 continue
 
             message = await self._get_pending_message(
                 redis_conn,
                 stream,
-                cast(str, message_id),
+                message_id,
             )
             if message is None:
                 await redis_conn.xclaim(
@@ -448,26 +447,37 @@ async def _claim_timed_out_messages(
                 message_ids=[message_id],
             )
             claimed.extend(cast("list[tuple[str, dict[bytes, bytes]]]", result))
-        return claimed
+        next_pending_start = "-"
+        if (
+            len(pending) == self.unacknowledged_batch_size
+            and last_checked_id is not None
+        ):
+            next_pending_start = f"({last_checked_id}"
+        return claimed, next_pending_start
 
     async def _claim_available_timed_out_messages(
         self,
         redis_conn: Redis,
         count: int,
         protected_message_ids: set[tuple[str, str]],
+        pending_starts: dict[str, str],
     ) -> list[tuple[str, dict[bytes, bytes], str]]:
         """Claim overdue messages across all configured streams."""
         claimed: list[tuple[str, dict[bytes, bytes], str]] = []
         for stream in [self.queue_name, *self.additional_streams.keys()]:
             remaining_count = count - len(claimed)
             if remaining_count <= 0:
                 break
-            for msg_id, msg in await self._claim_timed_out_messages(
-                redis_conn,
-                stream,
-                remaining_count,
-                protected_message_ids,
-            ):
+            claimed_messages, pending_starts[stream] = (
+                await self._claim_timed_out_messages(
+                    redis_conn,
+                    stream,
+                    remaining_count,
+                    protected_message_ids,
+                    pending_starts.get(stream, "-"),
+                )
+            )
+            for msg_id, msg in claimed_messages:
                 claimed.append((msg_id, msg, stream))
         return claimed
 
@@ -476,25 +486,18 @@ async def _read_new_messages(
         redis_conn: Redis,
         count: int | None,
     ) -> Any:
-        """Read newly added stream messages, recreating a missing group once."""
-        try:
-            return 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=count,
-            )
-        except ResponseError as exc:
-            if "NOGROUP" not in str(exc):
-                raise
-            logger.info("Consumer group missing, recreating")
-            await self._declare_consumer_group()
-            return []
+        """Read newly added stream messages."""
+        return 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=count,
+        )
 
     async def _abandon_buffered_messages(
         self,
@@ -517,13 +520,12 @@ async def _abandon_buffered_messages(
                     idle=ABANDONED_IDLE_MS,
                     justid=True,
                 )
-            except ResponseError as exc:
-                if "NOGROUP" not in str(exc):
-                    logger.warning(
-                        "Failed to abandon messages in stream %s",
-                        stream,
-                        exc_info=True,
-                    )
+            except RedisError:
+                logger.warning(
+                    "Failed to abandon messages in stream %s",
+                    stream,
+                    exc_info=True,
+                )
 
     def _build_ackable_message(
         self,
@@ -578,12 +580,14 @@ async def _build_reclaimed_ackable_messages(
         redis_conn: Redis,
         count: int,
         delivered: set[tuple[str, str]],
+        pending_starts: dict[str, str],
     ) -> list[tuple[str, dict[bytes, bytes], str]]:
         """Claim overdue messages for taskiq's receiver."""
         return await self._claim_available_timed_out_messages(
             redis_conn,
             count or self.unacknowledged_batch_size,
             delivered,
+            pending_starts,
         )
 
     async def _build_due_reclaimed_messages(
@@ -592,6 +596,7 @@ async def _build_due_reclaimed_messages(
         count: int,
         delivered: set[tuple[str, str]],
         last_reclaim: float,
+        pending_starts: dict[str, str],
     ) -> tuple[float, list[tuple[str, dict[bytes, bytes], str]]]:
         """Return overdue pending messages only when the reclaim interval elapsed."""
         if not self._should_reclaim(last_reclaim):
@@ -600,6 +605,7 @@ async def _build_due_reclaimed_messages(
             redis_conn,
             count,
             delivered,
+            pending_starts,
         )
 
     async def _build_new_ackable_messages(
@@ -637,6 +643,7 @@ async def listen(self) -> AsyncGenerator[AckableMessage, None]:
         buffered: list[tuple[str, dict[bytes, bytes], str]] = []
         slot_freed = asyncio.Event()
         last_reclaim = 0.0
+        pending_starts: dict[str, str] = {}
 
         def on_ack(message_key: tuple[str, str]) -> None:
             nonlocal unacked
@@ -665,6 +672,7 @@ def _on_ack() -> None:
                         count or self.unacknowledged_batch_size,
                         delivered,
                         last_reclaim,
+                        pending_starts,
                     )
                     unacked += len(buffered)
                     async for message in self._yield_buffered_messages(