feat: add lock heartbeat to prevent mid-execution expiry (#64)

README.md

@@ -21,7 +21,7 @@ Redis-backed deduplication middleware for Taskiq that prevents duplicate tasks f
 ## Installation
 
 ```bash
-uv add taskiq-deduplication
+uv add "taskiq-deduplication"
 ```
 
 ## Quick Start
@@ -53,6 +53,7 @@ except DuplicateTaskError:
 - **Sender-side deduplication** — rejects duplicate tasks at dispatch time via a Redis lock, before they reach the broker.
 - **Atomic lock release** — lock is released on completion or error via a Lua check-and-delete; only the owning task can release its lock.
 - **Configurable TTL** — set a global default or override per task with the `deduplication_ttl` label.
+- **Lock heartbeat** — a background task re-extends the lock TTL while the task runs, so long-running tasks keep their lock instead of expiring mid-execution and admitting a duplicate.
 - **Explicit lock key** — pin any task to a fixed Redis key with `deduplication_key`, bypassing fingerprint computation entirely.
 - **Partial fingerprint** — deduplicate on a subset of kwargs with `deduplication_key_fields`, ignoring irrelevant arguments.
 - **Per-task opt-out** — disable deduplication for individual tasks with the `deduplication` label.

docs/index.md

@@ -53,6 +53,7 @@ except DuplicateTaskError:
 - **Sender-side deduplication** — rejects duplicate tasks at dispatch time via a Redis lock, before they reach the broker.
 - **Atomic lock release** — lock is released on completion or error via a Lua check-and-delete; only the owning task can release its lock.
 - **Configurable TTL** — set a global default or override per task with the `deduplication_ttl` label.
+- **Lock heartbeat** — a background task re-extends the lock TTL while the task runs, so long-running tasks keep their lock instead of expiring mid-execution and admitting a duplicate.
 - **Explicit lock key** — pin any task to a fixed Redis key with `deduplication_key`, bypassing fingerprint computation entirely.
 - **Partial fingerprint** — deduplicate on a subset of kwargs with `deduplication_key_fields`, ignoring irrelevant arguments.
 - **Per-task opt-out** — disable deduplication for individual tasks with the `deduplication` label.

docs/usage.md

@@ -23,6 +23,8 @@ broker = ListQueueBroker("redis://localhost:6379").with_middlewares(
 | `key_prefix` | `str` | `"taskiq:deduplication"` | Prefix for all Redis lock keys. |
 | `startup_retries` | `int` | `3` | Number of connection attempts during broker startup. |
 | `startup_retry_delay` | `float` | `1.0` | Base delay in seconds between retries (exponential backoff: delay × 2^n). |
+| `heartbeat` | `bool` | `True` | Whether to periodically re-extend the lock TTL while the task runs (see [Long-running tasks](#long-running-tasks-and-the-heartbeat)). |
+| `heartbeat_interval` | `float \| None` | `None` | Seconds between heartbeat refreshes. When `None`, defaults to a third of the task's TTL (1s floor). |
 
 ```python
 broker = ListQueueBroker("redis://localhost:6379").with_middlewares(
@@ -37,6 +39,31 @@ broker = ListQueueBroker("redis://localhost:6379").with_middlewares(
 )
 ```
 
+## Long-running tasks and the heartbeat
+
+The lock is created with a TTL so a crashed worker cannot leak it forever. Without
+any refresh, a task that runs longer than its TTL would let the lock expire
+**mid-execution**, allowing a duplicate to be dispatched.
+
+To prevent this, the middleware starts a background **heartbeat** in `pre_execute`
+that re-extends the lock TTL while the task runs (atomically, only if the lock is
+still owned by the running task). It is cancelled when the task completes or fails.
+This means you do **not** need to size `default_ttl` to your slowest task — the TTL
+only needs to outlive a single heartbeat interval; it acts purely as a safety net
+for worker crashes.
+
+```python
+RedisDeduplicationMiddleware(
+    redis_url="redis://localhost:6379",
+    default_ttl=60,        # safety-net TTL; refreshed every ~20s while running
+    heartbeat_interval=20, # optional; defaults to default_ttl / 3
+)
+```
+
+If you disable the heartbeat (`heartbeat=False`), the invariant **TTL must exceed
+the slowest task** applies: set `default_ttl` (or the per-task `deduplication_ttl`
+label) above your worst-case task duration, or duplicates may slip through.
+
 ## Startup resilience
 
 On startup the middleware verifies the Redis connection with a `PING`. If Redis is

src/taskiq_deduplication/middleware.py

@@ -2,7 +2,8 @@ import asyncio
 import hashlib
 import json
 import logging
-from typing import Any, Awaitable, cast
+from collections.abc import Awaitable
+from typing import Any, cast
 
 from pydantic import RedisDsn
 from redis.asyncio import Redis
@@ -10,8 +11,10 @@ from taskiq import TaskiqMessage, TaskiqResult
 from taskiq.abc.middleware import TaskiqMiddleware
 
 from .utils import (
+    REFRESH_LUA_SCRIPT,
     RELEASE_LUA_SCRIPT,
     check_and_delete,
+    check_and_refresh,
     parse_bool_label,
     parse_int_label,
     parse_list_label,
@@ -45,6 +48,10 @@ class RedisDeduplicationMiddleware(TaskiqMiddleware):
         default_deduplication: Whether deduplication is enabled by default.
         default_ttl: Default lock TTL in seconds.
         key_prefix: Prefix for all Redis lock keys.
+        heartbeat: Whether to periodically re-extend the lock TTL during task
+            execution so long-running tasks keep their lock.
+        heartbeat_interval: Seconds between heartbeat refreshes. When ``None`` it
+            defaults to a third of the task's TTL (with a 1s floor).
     """
 
     def __init__(
@@ -55,6 +62,8 @@ class RedisDeduplicationMiddleware(TaskiqMiddleware):
         key_prefix: str = "taskiq:deduplication",
         startup_retries: int = 3,
         startup_retry_delay: float = 1.0,
+        heartbeat: bool = True,
+        heartbeat_interval: float | None = None,
     ) -> None:
         self.redis_url = redis_url
         self.default_deduplication = default_deduplication
@@ -62,8 +71,12 @@ class RedisDeduplicationMiddleware(TaskiqMiddleware):
         self.key_prefix = key_prefix
         self.startup_retries = startup_retries
         self.startup_retry_delay = startup_retry_delay
+        self.heartbeat = heartbeat
+        self.heartbeat_interval = heartbeat_interval
         self._redis: Redis | None = None
         self._release_script: Any = None
+        self._refresh_script: Any = None
+        self._heartbeats: dict[str, asyncio.Task[None]] = {}
 
     async def startup(self) -> None:
         last_error: BaseException | None = None
@@ -73,6 +86,7 @@ class RedisDeduplicationMiddleware(TaskiqMiddleware):
                 await cast(Awaitable[bool], client.ping())
                 self._redis = client
                 self._release_script = self._redis.register_script(RELEASE_LUA_SCRIPT)
+                self._refresh_script = self._redis.register_script(REFRESH_LUA_SCRIPT)
                 return
             except Exception as exc:
                 await client.aclose()
@@ -96,6 +110,8 @@ class RedisDeduplicationMiddleware(TaskiqMiddleware):
         ) from last_error
 
     async def shutdown(self) -> None:
+        for task_id in list(self._heartbeats):
+            await self._cancel_heartbeat(task_id)
         if self._redis is not None:
             await self._redis.aclose()
 
@@ -147,6 +163,15 @@ class RedisDeduplicationMiddleware(TaskiqMiddleware):
         else:
             logger.debug("Skipped release of lock %s: not owned by this task", key)
 
+    async def _refresh_if_owned(self, key: str, task_id: str, ttl: int) -> bool:
+        if self._redis is None:
+            raise RuntimeError(
+                "RedisDeduplicationMiddleware.startup() was never called."
+            )
+        if self._refresh_script is None:
+            self._refresh_script = self._redis.register_script(REFRESH_LUA_SCRIPT)
+        return await check_and_refresh(self._refresh_script, key, task_id, ttl)
+
     @staticmethod
     def _get_cached_key(message: TaskiqMessage) -> str | None:
         return message.labels.get(_CACHED_KEY_LABEL)
@@ -189,7 +214,60 @@ class RedisDeduplicationMiddleware(TaskiqMiddleware):
         logger.debug("Lock %s acquired for task %s", key, message.task_name)
         return message
 
+    def _get_heartbeat_interval(self, ttl: int) -> float:
+        if self.heartbeat_interval is not None:
+            return self.heartbeat_interval
+        return max(ttl / 3, 1.0)
+
+    async def _heartbeat_loop(
+        self, key: str, task_id: str, ttl: int, interval: float
+    ) -> None:
+        try:
+            while True:
+                await asyncio.sleep(interval)
+                try:
+                    refreshed = await self._refresh_if_owned(key, task_id, ttl)
+                except Exception as exc:
+                    logger.warning("Failed to refresh lock %s: %s", key, exc)
+                    continue
+                if refreshed:
+                    logger.debug("Refreshed lock %s (ttl=%ds)", key, ttl)
+                else:
+                    logger.warning(
+                        "Lock %s no longer owned by task %s; stopping heartbeat.",
+                        key,
+                        task_id,
+                    )
+                    return
+        except asyncio.CancelledError:
+            pass
+
+    async def pre_execute(self, message: TaskiqMessage) -> TaskiqMessage:
+        if not self.heartbeat:
+            return message
+        # The cached key is set by pre_send() only when deduplication is enabled.
+        key = self._get_cached_key(message)
+        if key is None:
+            return message
+        ttl = self._get_ttl(message.labels)
+        interval = self._get_heartbeat_interval(ttl)
+        self._heartbeats[message.task_id] = asyncio.create_task(
+            self._heartbeat_loop(key, message.task_id, ttl, interval)
+        )
+        return message
+
+    async def _cancel_heartbeat(self, task_id: str) -> None:
+        task = self._heartbeats.pop(task_id, None)
+        if task is None:
+            return
+        task.cancel()
+        try:
+            await task
+        except asyncio.CancelledError:
+            pass
+
     async def _release_lock(self, message: TaskiqMessage) -> None:
+        await self._cancel_heartbeat(message.task_id)
         # The cached key is set by pre_send() only when deduplication is enabled.
         key = self._get_cached_key(message)
         if key is None:

src/taskiq_deduplication/utils.py

@@ -12,6 +12,14 @@ else
 end
 """
 
+REFRESH_LUA_SCRIPT = """
+if redis.call('get', KEYS[1]) == ARGV[1] then
+    return redis.call('expire', KEYS[1], ARGV[2])
+else
+    return 0
+end
+"""
+
 
 async def check_and_delete(script: Any, key: str, owner: str) -> bool:
     """Delete *key* only if its value equals *owner*.
@@ -28,6 +36,23 @@ async def check_and_delete(script: Any, key: str, owner: str) -> bool:
     return bool(released)
 
 
+async def check_and_refresh(script: Any, key: str, owner: str, ttl: int) -> bool:
+    """Extend *key*'s TTL to *ttl* only if its value equals *owner*.
+
+    Args:
+        script: Pre-registered Lua script object (from ``Redis.register_script``).
+        key: Lock key to refresh.
+        owner: Expected value of the key (task_id).
+        ttl: New TTL in seconds.
+
+    Returns:
+        True if the TTL was extended, False if the key is missing or owned by
+        another task.
+    """
+    refreshed: int = await script(keys=[key], args=[owner, ttl])
+    return bool(refreshed)
+
+
 def parse_bool_label(value: Any, default: bool, label_name: str = "") -> bool:
     if isinstance(value, bool):
         return value

tests/test_integration.py

@@ -7,7 +7,7 @@ In CI a Redis service is started before the test step.
 import pytest
 
 from taskiq_deduplication import DuplicateTaskError, RedisDeduplicationMiddleware
-from taskiq_deduplication.utils import RELEASE_LUA_SCRIPT
+from taskiq_deduplication.utils import REFRESH_LUA_SCRIPT, RELEASE_LUA_SCRIPT
 
 
 @pytest.fixture
@@ -15,6 +15,7 @@ def mw(real_redis):
     middleware = RedisDeduplicationMiddleware(redis_url="redis://localhost:6379/15")
     middleware._redis = real_redis
     middleware._release_script = real_redis.register_script(RELEASE_LUA_SCRIPT)
+    middleware._refresh_script = real_redis.register_script(REFRESH_LUA_SCRIPT)
     return middleware
 
 
@@ -68,6 +69,31 @@ async def test_ttl_is_applied(mw, real_redis, make_message):
     assert 0 < ttl <= mw.default_ttl
 
 
+@pytest.mark.integration
+async def test_heartbeat_keeps_long_running_lock_alive(
+    mw, real_redis, make_message, make_result
+):
+    import asyncio
+
+    from taskiq_deduplication.middleware import DEDUP_TTL_LABEL
+
+    # 1s TTL with a sub-second heartbeat: without refresh the lock would expire.
+    mw.heartbeat_interval = 0.2
+    msg = make_message(labels={DEDUP_TTL_LABEL: 1})
+    await mw.pre_send(msg)
+    key = mw._build_deduplication_key(msg)
+    await mw.pre_execute(msg)
+    try:
+        # outlive the original TTL; the heartbeat should keep the lock present
+        await asyncio.sleep(1.5)
+        assert await real_redis.exists(key)
+        with pytest.raises(DuplicateTaskError):
+            await mw.pre_send(make_message(labels={DEDUP_TTL_LABEL: 1}))
+    finally:
+        await mw.post_execute(msg, make_result())
+    assert not await real_redis.exists(key)
+
+
 @pytest.mark.integration
 async def test_explicit_key_end_to_end(mw, real_redis, make_message, make_result):
     from taskiq_deduplication.middleware import DEDUP_EXPLICIT_KEY_LABEL

tests/test_middleware.py

@@ -704,6 +704,131 @@ class TestTTLExpiry:
         await middleware.pre_send(make_message())
 
 
+class TestHeartbeat:
+    async def test_pre_execute_starts_heartbeat(self, middleware, make_message):
+        msg = make_message()
+        await middleware.pre_send(msg)
+        await middleware.pre_execute(msg)
+        assert msg.task_id in middleware._heartbeats
+        await middleware._cancel_heartbeat(msg.task_id)
+
+    async def test_pre_execute_noop_when_heartbeat_disabled(
+        self, fake_redis, make_message
+    ):
+        mw = RedisDeduplicationMiddleware(
+            redis_url="redis://localhost", heartbeat=False
+        )
+        mw._redis = fake_redis
+        msg = make_message()
+        await mw.pre_send(msg)
+        await mw.pre_execute(msg)
+        assert msg.task_id not in mw._heartbeats
+
+    async def test_pre_execute_noop_without_cached_key(self, middleware, make_message):
+        # deduplication disabled -> pre_send never caches a key
+        msg = make_message(labels={DEDUP_LABEL: False})
+        await middleware.pre_send(msg)
+        await middleware.pre_execute(msg)
+        assert msg.task_id not in middleware._heartbeats
+
+    async def test_heartbeat_refreshes_ttl(self, middleware, fake_redis, make_message):
+        import asyncio
+
+        # short ttl, tiny heartbeat interval so the lock would expire without refresh
+        middleware.heartbeat_interval = 0.05
+        msg = make_message(labels={DEDUP_TTL_LABEL: 1})
+        await middleware.pre_send(msg)
+        key = middleware._build_deduplication_key(msg)
+        await middleware.pre_execute(msg)
+        try:
+            # let several heartbeats elapse — longer than the original 1s ttl
+            await asyncio.sleep(0.3)
+            assert await fake_redis.exists(key)
+            ttl = await fake_redis.ttl(key)
+            assert 0 < ttl <= 1
+        finally:
+            await middleware._cancel_heartbeat(msg.task_id)
+
+    async def test_release_lock_cancels_heartbeat(
+        self, middleware, fake_redis, make_message, make_result
+    ):
+        msg = make_message()
+        await middleware.pre_send(msg)
+        await middleware.pre_execute(msg)
+        assert msg.task_id in middleware._heartbeats
+        await middleware.post_execute(msg, make_result())
+        assert msg.task_id not in middleware._heartbeats
+        key = middleware._build_deduplication_key(msg)
+        assert not await fake_redis.exists(key)
+
+    async def test_heartbeat_stops_when_lock_lost(
+        self, middleware, fake_redis, make_message
+    ):
+        import asyncio
+
+        middleware.heartbeat_interval = 0.05
+        msg = make_message(labels={DEDUP_TTL_LABEL: 1})
+        await middleware.pre_send(msg)
+        key = middleware._build_deduplication_key(msg)
+        await middleware.pre_execute(msg)
+        # another task steals the key
+        await fake_redis.set(key, "other-task", ex=10)
+        await asyncio.sleep(0.15)
+        task = middleware._heartbeats.get(msg.task_id)
+        # heartbeat loop should have returned on its own
+        assert task is None or task.done()
+        await middleware._cancel_heartbeat(msg.task_id)
+
+    async def test_shutdown_cancels_heartbeats(self, fake_redis, make_message):
+        mw = RedisDeduplicationMiddleware(redis_url="redis://localhost")
+        mw._redis = fake_redis
+        msg = make_message()
+        await mw.pre_send(msg)
+        await mw.pre_execute(msg)
+        assert msg.task_id in mw._heartbeats
+        await mw.shutdown()
+        assert not mw._heartbeats
+
+    async def test_default_heartbeat_interval_is_third_of_ttl(self, middleware):
+        assert middleware._get_heartbeat_interval(300) == 100.0
+        assert middleware._get_heartbeat_interval(1) == 1.0
+
+    async def test_explicit_heartbeat_interval_overrides(self, fake_redis):
+        mw = RedisDeduplicationMiddleware(
+            redis_url="redis://localhost", heartbeat_interval=5.0
+        )
+        mw._redis = fake_redis
+        assert mw._get_heartbeat_interval(300) == 5.0
+
+    async def test_refresh_if_owned_raises_without_redis(self, middleware):
+        middleware._redis = None
+        with pytest.raises(RuntimeError, match="startup"):
+            await middleware._refresh_if_owned("some-key", "some-task", 60)
+
+    async def test_heartbeat_continues_after_refresh_error(
+        self, middleware, make_message, caplog
+    ):
+        import asyncio
+        import logging
+
+        middleware.heartbeat_interval = 0.02
+        # first refresh raises, subsequent ones succeed; the loop must survive
+        middleware._refresh_if_owned = AsyncMock(
+            side_effect=[ConnectionError("boom"), True, True, True]
+        )
+        msg = make_message()
+        await middleware.pre_send(msg)
+        with caplog.at_level(logging.WARNING, logger="taskiq_deduplication.middleware"):
+            await middleware.pre_execute(msg)
+            await asyncio.sleep(0.1)
+            task = middleware._heartbeats.get(msg.task_id)
+            # loop swallowed the error and kept running
+            assert task is not None and not task.done()
+            await middleware._cancel_heartbeat(msg.task_id)
+        assert any("Failed to refresh lock" in r.message for r in caplog.records)
+        assert middleware._refresh_if_owned.call_count >= 2
+
+
 class TestExplicitKeyEdgeCases:
     def test_empty_string_key_produces_prefix_only_key(self, middleware, make_message):
         m = make_message(labels={DEDUP_EXPLICIT_KEY_LABEL: ""})