From d8835733edf751f96240cf332db8a9c786932546 Mon Sep 17 00:00:00 2001 From: d3vyce <44915747+d3vyce@users.noreply.github.com> Date: Wed, 29 Jul 2026 10:27:25 +0200 Subject: [PATCH] fix: bound the lock leak when broker.kick() fails (#83) --- docs/usage.md | 6 +++ src/taskiq_deduplication/middleware.py | 28 ++++++++++- tests/test_integration.py | 4 +- tests/test_middleware.py | 66 +++++++++++++++++++++++++- 4 files changed, 100 insertions(+), 4 deletions(-) diff --git a/docs/usage.md b/docs/usage.md index 4393cef..f582be8 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -88,6 +88,12 @@ fingerprint. Any subsequent dispatch with the same fingerprint raises `DuplicateTaskError` while the lock is held. The lock is released automatically when the task completes or fails. +The lock is acquired in two phases. It is first taken with a short grace TTL of 10 +seconds, then extended to its full TTL once the broker has accepted the message. +taskiq fires no middleware hook when the send itself fails, so this bounds the +damage: if the broker is unreachable, the lock of a task that was never queued +expires within seconds instead of blocking its fingerprint for the full TTL. + ## Handling duplicates When a duplicate is detected, the middleware logs a warning and raises diff --git a/src/taskiq_deduplication/middleware.py b/src/taskiq_deduplication/middleware.py index bee30af..8a90a3d 100644 --- a/src/taskiq_deduplication/middleware.py +++ b/src/taskiq_deduplication/middleware.py @@ -25,6 +25,7 @@ DEDUP_LABEL = "deduplication" DEDUP_TTL_LABEL = "deduplication_ttl" DEDUP_KEY_FIELDS_LABEL = "deduplication_key_fields" DEDUP_EXPLICIT_KEY_LABEL = "deduplication_key" +SEND_GRACE_TTL = 10 _CACHED_KEY_LABEL = "__taskiq_dedup_cached_key" @@ -238,7 +239,9 @@ class RedisDeduplicationMiddleware(TaskiqMiddleware): ttl = self._get_ttl(message.labels) logger.debug("Acquiring lock %s for task %s", key, message.task_name) - acquired = await self._redis.set(key, message.task_id, ex=ttl, nx=True) + acquired = await self._redis.set( + key, message.task_id, ex=min(ttl, SEND_GRACE_TTL), nx=True + ) if not acquired: holder_task_id = await self._redis.get(key) if isinstance(holder_task_id, bytes): @@ -258,6 +261,29 @@ class RedisDeduplicationMiddleware(TaskiqMiddleware): logger.debug("Lock %s acquired for task %s", key, message.task_name) return message + async def post_send(self, message: TaskiqMessage) -> None: + # The cached key is set by pre_send() only when deduplication is enabled. + key = self._get_cached_key(message) + if key is None: + return + ttl = self._get_ttl(message.labels) + if ttl <= SEND_GRACE_TTL: + return + try: + # Returns False when a fast worker already ran and released the lock. + extended = await self._refresh_if_owned(key, message.task_id, ttl) + except Exception as exc: + # The task is already queued; never fail the send. The lock just keeps + # its grace TTL. + logger.warning("Failed to extend lock %s after send: %s", key, exc) + return + logger.debug( + "Lock %s %s to the full TTL (%ds) after send", + key, + "extended" if extended else "not extended", + ttl, + ) + def _get_heartbeat_interval(self, ttl: int) -> float: if self.heartbeat_interval is not None: return self.heartbeat_interval diff --git a/tests/test_integration.py b/tests/test_integration.py index 9fa3e79..4192267 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -7,6 +7,7 @@ In CI a Redis service is started before the test step. import pytest from taskiq_deduplication import DuplicateTaskError, RedisDeduplicationMiddleware +from taskiq_deduplication.middleware import SEND_GRACE_TTL from taskiq_deduplication.utils import REFRESH_LUA_SCRIPT, RELEASE_LUA_SCRIPT @@ -64,9 +65,10 @@ async def test_lua_only_owner_can_release(mw, real_redis, make_message): async def test_ttl_is_applied(mw, real_redis, make_message): msg = make_message() await mw.pre_send(msg) + await mw.post_send(msg) key = mw._build_deduplication_key(msg) ttl = await real_redis.ttl(key) - assert 0 < ttl <= mw.default_ttl + assert SEND_GRACE_TTL < ttl <= mw.default_ttl @pytest.mark.integration diff --git a/tests/test_middleware.py b/tests/test_middleware.py index 69f90cf..7789cbb 100644 --- a/tests/test_middleware.py +++ b/tests/test_middleware.py @@ -9,6 +9,7 @@ from taskiq_deduplication.middleware import ( DEDUP_KEY_FIELDS_LABEL, DEDUP_LABEL, DEDUP_TTL_LABEL, + SEND_GRACE_TTL, ) @@ -214,9 +215,10 @@ class TestPreSend: async def test_ttl_applied(self, middleware, fake_redis, make_message): msg = make_message(labels={DEDUP_TTL_LABEL: 42}) await middleware.pre_send(msg) + await middleware.post_send(msg) key = middleware._build_deduplication_key(msg) ttl = await fake_redis.ttl(key) - assert 0 < ttl <= 42 + assert SEND_GRACE_TTL < ttl <= 42 async def test_different_kwargs_both_pass(self, middleware, make_message): await middleware.pre_send(make_message(kwargs={"x": 1})) @@ -235,9 +237,69 @@ class TestPreSend: mw._redis = fake_redis msg = make_message() await mw.pre_send(msg) + await mw.post_send(msg) key = mw._build_deduplication_key(msg) ttl = await fake_redis.ttl(key) - assert 0 < ttl <= 77 + assert SEND_GRACE_TTL < ttl <= 77 + + +class TestPostSend: + async def test_pre_send_only_uses_grace_ttl( + self, middleware, fake_redis, make_message + ): + msg = make_message(labels={DEDUP_TTL_LABEL: 300}) + await middleware.pre_send(msg) + key = middleware._build_deduplication_key(msg) + assert 0 < await fake_redis.ttl(key) <= SEND_GRACE_TTL + + async def test_post_send_extends_to_full_ttl( + self, middleware, fake_redis, make_message + ): + msg = make_message(labels={DEDUP_TTL_LABEL: 300}) + await middleware.pre_send(msg) + await middleware.post_send(msg) + key = middleware._build_deduplication_key(msg) + assert await fake_redis.ttl(key) > SEND_GRACE_TTL + + async def test_post_send_does_not_extend_short_ttl( + self, middleware, fake_redis, make_message + ): + msg = make_message(labels={DEDUP_TTL_LABEL: 5}) + await middleware.pre_send(msg) + await middleware.post_send(msg) + key = middleware._build_deduplication_key(msg) + assert 0 < await fake_redis.ttl(key) <= 5 + + async def test_post_send_does_not_resurrect_released_lock( + self, middleware, fake_redis, make_message, make_result + ): + msg = make_message(labels={DEDUP_TTL_LABEL: 300}) + await middleware.pre_send(msg) + await middleware.post_execute(msg, make_result()) + await middleware.post_send(msg) + key = middleware._build_deduplication_key(msg) + assert not await fake_redis.exists(key) + + async def test_post_send_survives_redis_error( + self, middleware, fake_redis, make_message + ): + msg = make_message(labels={DEDUP_TTL_LABEL: 300}) + await middleware.pre_send(msg) + middleware._refresh_script = AsyncMock(side_effect=ConnectionError("boom")) + await middleware.post_send(msg) # should not raise + key = middleware._build_deduplication_key(msg) + assert 0 < await fake_redis.ttl(key) <= SEND_GRACE_TTL + + async def test_post_send_deduplication_disabled_noop( + self, fake_redis, make_message + ): + mw = RedisDeduplicationMiddleware( + redis_url="redis://localhost", default_deduplication=False + ) + mw._redis = fake_redis + msg = make_message() + await mw.pre_send(msg) + await mw.post_send(msg) # should not raise class TestPostExecute: