mirror of
https://github.com/d3vyce/taskiq-deduplication.git
synced 2026-08-04 19:14:07 +00:00
feat: enrich DuplicateTaskError with structured attributes (#66)
This commit is contained in:
@@ -103,6 +103,25 @@ except DuplicateTaskError:
|
||||
pass # task is already queued or running
|
||||
```
|
||||
|
||||
`DuplicateTaskError` carries structured attributes describing the collision:
|
||||
|
||||
```python
|
||||
try:
|
||||
await my_task.kiq(user_id=42)
|
||||
except DuplicateTaskError as err:
|
||||
logger.info(
|
||||
"Skipped %s; already held by %s (key=%s)",
|
||||
err.task_name,
|
||||
err.holder_task_id,
|
||||
err.key,
|
||||
)
|
||||
```
|
||||
|
||||
- `task_name` — name of the task that was rejected.
|
||||
- `key` — Redis lock key whose owner caused the rejection.
|
||||
- `holder_task_id` — `task_id` of the task currently holding the lock, or `None`
|
||||
if it could not be retrieved.
|
||||
|
||||
## Per-task label overrides
|
||||
|
||||
Labels can be set at the task level (applied to every call) or at call time.
|
||||
|
||||
@@ -31,7 +31,29 @@ _CACHED_KEY_LABEL = "__taskiq_dedup_cached_key"
|
||||
|
||||
|
||||
class DuplicateTaskError(Exception):
|
||||
"""Raised when a task with identical name and kwargs is already queued or running."""
|
||||
"""Raised when a task with identical name and kwargs is already queued or running.
|
||||
|
||||
Attributes:
|
||||
task_name: Name of the task that was rejected.
|
||||
key: Redis lock key whose owner caused the rejection.
|
||||
holder_task_id: ``task_id`` of the task currently holding the lock, or
|
||||
``None`` if it could not be retrieved (e.g. the lock was released
|
||||
between the failed acquisition and the lookup).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
task_name: str,
|
||||
key: str,
|
||||
holder_task_id: str | None = None,
|
||||
) -> None:
|
||||
self.task_name = task_name
|
||||
self.key = key
|
||||
self.holder_task_id = holder_task_id
|
||||
super().__init__(
|
||||
f"Task {task_name!r} with the same arguments is already queued or "
|
||||
f"running (key={key!r}, holder_task_id={holder_task_id!r})."
|
||||
)
|
||||
|
||||
|
||||
class RedisDeduplicationMiddleware(TaskiqMiddleware):
|
||||
@@ -202,13 +224,19 @@ class RedisDeduplicationMiddleware(TaskiqMiddleware):
|
||||
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)
|
||||
if not acquired:
|
||||
holder_task_id = await self._redis.get(key)
|
||||
if isinstance(holder_task_id, bytes):
|
||||
holder_task_id = holder_task_id.decode()
|
||||
logger.warning(
|
||||
"Duplicate task %s dropped (key=%s).",
|
||||
"Duplicate task %s dropped (key=%s, holder_task_id=%s).",
|
||||
message.task_name,
|
||||
key,
|
||||
holder_task_id,
|
||||
)
|
||||
raise DuplicateTaskError(
|
||||
f"Task {message.task_name!r} with the same arguments is already queued or running."
|
||||
task_name=message.task_name,
|
||||
key=key,
|
||||
holder_task_id=holder_task_id,
|
||||
)
|
||||
|
||||
logger.debug("Lock %s acquired for task %s", key, message.task_name)
|
||||
|
||||
@@ -146,6 +146,34 @@ class TestPreSend:
|
||||
with pytest.raises(DuplicateTaskError):
|
||||
await middleware.pre_send(make_message())
|
||||
|
||||
async def test_duplicate_error_carries_structured_attributes(
|
||||
self, middleware, make_message
|
||||
):
|
||||
holder = make_message(task_id="holder-task")
|
||||
await middleware.pre_send(holder)
|
||||
key = middleware._build_deduplication_key(holder)
|
||||
with pytest.raises(DuplicateTaskError) as exc_info:
|
||||
await middleware.pre_send(make_message(task_id="loser-task"))
|
||||
err = exc_info.value
|
||||
assert err.task_name == "my_task"
|
||||
assert err.key == key
|
||||
assert err.holder_task_id == "holder-task"
|
||||
assert key in str(err)
|
||||
|
||||
async def test_duplicate_error_holder_none_when_lock_released_in_race(
|
||||
self, make_message
|
||||
):
|
||||
# The lock is released between the failed SET NX and the GET lookup, so
|
||||
# GET returns None and holder_task_id is left unset.
|
||||
redis = AsyncMock()
|
||||
redis.set.return_value = False
|
||||
redis.get.return_value = None
|
||||
mw = RedisDeduplicationMiddleware(redis_url="redis://localhost")
|
||||
mw._redis = redis
|
||||
with pytest.raises(DuplicateTaskError) as exc_info:
|
||||
await mw.pre_send(make_message())
|
||||
assert exc_info.value.holder_task_id is None
|
||||
|
||||
async def test_deduplication_disabled_label(self, middleware, make_message):
|
||||
msg1 = make_message(labels={DEDUP_LABEL: False})
|
||||
msg2 = make_message(labels={DEDUP_LABEL: False})
|
||||
|
||||
Reference in New Issue
Block a user