diff --git a/README.md b/README.md index 7212e99..5e345fd 100644 --- a/README.md +++ b/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. diff --git a/docs/index.md b/docs/index.md index 406b8fb..5e345fd 100644 --- a/docs/index.md +++ b/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. diff --git a/docs/usage.md b/docs/usage.md index 268586a..c02fd6a 100644 --- a/docs/usage.md +++ b/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 diff --git a/src/taskiq_deduplication/middleware.py b/src/taskiq_deduplication/middleware.py index 57e410a..4dd72a2 100644 --- a/src/taskiq_deduplication/middleware.py +++ b/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: diff --git a/src/taskiq_deduplication/utils.py b/src/taskiq_deduplication/utils.py index a366712..1fdeeb3 100644 --- a/src/taskiq_deduplication/utils.py +++ b/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 diff --git a/tests/test_integration.py b/tests/test_integration.py index cad1b98..9fa3e79 100644 --- a/tests/test_integration.py +++ b/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 diff --git a/tests/test_middleware.py b/tests/test_middleware.py index f746544..4060843 100644 --- a/tests/test_middleware.py +++ b/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: ""})