Skip to content

Commit 0fcfca8

Browse files
committed
Fixed race-condition for streambroker.
Signed-off-by: Pavel Kirilin <s3riussan@gmail.com>
1 parent da1271d commit 0fcfca8

1 file changed

Lines changed: 25 additions & 21 deletions

File tree

taskiq_redis/redis_broker.py

Lines changed: 25 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import uuid
22
from collections.abc import AsyncGenerator, Awaitable, Callable
3+
from datetime import timedelta
34
from logging import getLogger
45
from typing import (
56
TYPE_CHECKING,
@@ -295,28 +296,31 @@ async def listen(self) -> AsyncGenerator[AckableMessage, None]:
295296
)
296297
logger.debug("Starting fetching unacknowledged messages")
297298
for stream in [self.queue_name, *self.additional_streams.keys()]:
298-
lock = redis_conn.lock(
299+
pipe = redis_conn.pipeline()
300+
lock = pipe.lock(
299301
f"autoclaim:{self.consumer_group_name}:{stream}",
300302
timeout=self.unacknowledged_lock_timeout,
301303
)
302-
if await lock.locked():
303-
continue
304-
async with lock:
305-
pending = await redis_conn.xautoclaim(
306-
name=stream,
307-
groupname=self.consumer_group_name,
308-
consumername=self.consumer_name,
309-
min_idle_time=self.idle_timeout,
310-
count=self.unacknowledged_batch_size,
311-
)
312-
logger.debug(
313-
"Found %d pending messages in stream %s",
314-
len(pending[1]),
315-
stream,
304+
await lock.acquire()
305+
await pipe.xautoclaim(
306+
name=stream,
307+
groupname=self.consumer_group_name,
308+
consumername=self.consumer_name,
309+
min_idle_time=self.idle_timeout,
310+
count=self.unacknowledged_batch_size,
311+
)
312+
await lock.release()
313+
results = await pipe.execute()
314+
pending = results[1]
315+
316+
logger.debug(
317+
"Found %d pending messages in stream %s",
318+
len(pending[1]),
319+
stream,
320+
)
321+
for msg_id, msg in pending[1]:
322+
logger.debug("Received message: %s", msg)
323+
yield AckableMessage(
324+
data=msg[b"data"],
325+
ack=self._ack_generator(id=msg_id, queue_name=stream),
316326
)
317-
for msg_id, msg in pending[1]:
318-
logger.debug("Received message: %s", msg)
319-
yield AckableMessage(
320-
data=msg[b"data"],
321-
ack=self._ack_generator(id=msg_id, queue_name=stream),
322-
)

0 commit comments

Comments
 (0)