fix: bound the lock leak when broker.kick() fails (#83)

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

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

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

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: