From 4d9b1ef7674aa2879d694659ad47ad74eb8616fd Mon Sep 17 00:00:00 2001 From: Johannes Maron Date: Sun, 5 Jul 2026 21:52:04 +0200 Subject: [PATCH 1/6] Resolve #22 -- Add retry callback support --- README.md | 35 ++++++++ tests/backends/test_base.py | 144 +++++++++++++++++++++++++++++++ tests/backends/test_redis.py | 158 ++++++++++++++++++++++++++++++++++ tests/test_executor.py | 159 +++++++++++++++++++++++++++++++++++ tests/testapp/tasks.py | 48 +++++++++++ threadmill/backends/base.py | 82 +++++++++++++----- threadmill/backends/redis.py | 29 +++++++ threadmill/executor.py | 17 +++- 8 files changed, 651 insertions(+), 21 deletions(-) diff --git a/README.md b/README.md index 4a9430e..5a28276 100644 --- a/README.md +++ b/README.md @@ -130,6 +130,41 @@ every multi-key operation — including the cross-queue acquire — runs on a si shard. Scale horizontally by running additional backend aliases, not by relying on cross-slot operations. +### Retrying failed tasks + +Pass a `retry` callback to `@task()` to retry failed tasks with a delay. +The callback receives a `TaskContext` — use `context.attempt` for the current +attempt count and `context.task_result.errors[-1]` for the latest error. +Return a `timedelta` to schedule the next attempt, or `None` to stop retrying. + +```python +import datetime + +from django.tasks import task +from django.tasks.base import TaskContext +from requests import HTTPError + + +def retry_on_rate_limit(context: TaskContext) -> datetime.timedelta | None: + """Retry HTTP 429 responses with exponential backoff, up to 5 attempts.""" + if context.attempt >= 5: + return None + error = context.task_result.errors[-1] + if not issubclass(error.exception_class, HTTPError): + return None + return min( + datetime.timedelta(seconds=2**context.attempt), + datetime.timedelta(seconds=60), + ) + + +@task(retry=retry_on_rate_limit) +def fetch_github_api(url: str): ... +``` + +The worker re-queues the failed task, preserving its ID and error history; +the broker promotes it back to the ready queue once the delay elapses. + ## Sponsors [![Sponsors](https://django.the-box.sh/sponsors/codingjoe/threadmill.svg)](https://github.com/sponsors/codingjoe) diff --git a/tests/backends/test_base.py b/tests/backends/test_base.py index 9f6b23d..36ba540 100644 --- a/tests/backends/test_base.py +++ b/tests/backends/test_base.py @@ -44,6 +44,14 @@ def test_acknowledge__raise_not_implemented_error(self) -> None: with pytest.raises(NotImplementedError): BackendDouble(alias="default", params={}).acknowledge(task_result=None) + def test_requeue__raise_not_implemented_error(self) -> None: + """Raise NotImplementedError for backend requeue API.""" + with pytest.raises(NotImplementedError): + BackendDouble(alias="default", params={}).requeue( + task_result=None, + run_after=timezone.now(), + ) + def test_peek__raise_not_implemented_error(self) -> None: """Raise NotImplementedError for backend peek_results API.""" with pytest.raises(NotImplementedError): @@ -123,3 +131,139 @@ def test_shutdown__sets_event(self) -> None: assert not broker.shutdown_requested.is_set() broker.shutdown() assert broker.shutdown_requested.is_set() + + +from django.tasks import TaskResultStatus as TRS # noqa: E402 +from django.tasks.exceptions import InvalidTask # noqa: E402 +from django.utils import timezone as tz # noqa: E402 + +from tests.testapp.tasks import ( # noqa: E402 + boom_with_retry, + echo, + retry_always, +) +from threadmill.backends.base import RetryTask # noqa: E402 + + +def _task_result_with_retry(task, *, errors=None, worker_ids=None) -> TaskResult: + """Build a FAILED TaskResult with errors for retry testing.""" + now = tz.now() + return TaskResult( + task=task, + id=str(uuid.uuid4()), + status=TRS.FAILED, + enqueued_at=now, + started_at=now, + finished_at=now, + last_attempted_at=now, + backend="default", + errors=errors or [], + worker_ids=worker_ids or [], + args=[], + kwargs={}, + ) + + +class TestRetryTask: + """Tests for the RetryTask class.""" + + def test_retry_path__none_when_no_retry(self) -> None: + """Return None when retry callback is not set.""" + task = RetryTask(func=echo.func) + assert task.retry_path is None + + def test_retry_path__returns_dotted_path(self) -> None: + """Return the importable dotted path of the retry callback.""" + task = RetryTask(func=echo.func, retry=retry_always) + assert task.retry_path == "tests.testapp.tasks.retry_always" + + def test_validate_task__accepts_module_level_function(self) -> None: + """validate_task accepts a module-level retry callback.""" + backend = BackendDouble(alias="default", params={}) + backend.validate_task(boom_with_retry) + + def test_validate_task__rejects_lambda(self) -> None: + """Reject a non-module-level retry callback during task creation.""" + with pytest.raises(InvalidTask): + RetryTask( + func=echo.func, + retry=lambda ctx: None, + ) + + +class TestRetryTaskSerialization: + """Tests for RetryTask JSON and pickle serialization.""" + + def test_serialize__encodes_retry_as_path(self) -> None: + """TaskResultEncoder serializes retry as a dotted path.""" + result = _task_result_with_retry(boom_with_retry) + data = ThreadmillTaskBackend.serialize_task_result(result) + assert "tests.testapp.tasks.retry_always" in data + + def test_serialize__omits_retry_when_none(self) -> None: + """TaskResultEncoder omits retry when it is None.""" + task = RetryTask(func=echo.func) + result = TaskResult( + task=task, + id=str(uuid.uuid4()), + status=TRS.READY, + enqueued_at=tz.now(), + started_at=None, + finished_at=None, + last_attempted_at=None, + backend="default", + errors=[], + worker_ids=[], + args=[], + kwargs={}, + ) + data = ThreadmillTaskBackend.serialize_task_result(result) + import json + + decoded = json.loads(data) + assert decoded["task"]["retry"] is None + + def test_deserialize__roundtrips_retry_callback(self) -> None: + """deserialize_task_result restores the retry callable from path.""" + result = _task_result_with_retry(boom_with_retry) + data = ThreadmillTaskBackend.serialize_task_result(result) + restored = ThreadmillTaskBackend.deserialize_task_result(data) + assert restored.task.retry is retry_always + assert restored.task.retry_path == "tests.testapp.tasks.retry_always" + + def test_deserialize__roundtrips_without_retry(self) -> None: + """deserialize_task_result works when retry is None.""" + result = _task_result_with_retry(echo) + data = ThreadmillTaskBackend.serialize_task_result(result) + restored = ThreadmillTaskBackend.deserialize_task_result(data) + assert restored.task.retry is None + + def test_pickle__roundtrips_retry_callback(self) -> None: + """Pickle preserves the retry callable by importable path.""" + import pickle + + result = _task_result_with_retry(boom_with_retry) + restored = pickle.loads(pickle.dumps(result)) # noqa: S301 + assert restored.task.retry is retry_always + + def test_pickle__roundtrips_without_retry(self) -> None: + """Pickle works when retry is None.""" + import pickle + + result = _task_result_with_retry(echo) + restored = pickle.loads(pickle.dumps(result)) # noqa: S301 + assert restored.task.retry is None + + def test_deserialize__preserves_errors(self) -> None: + """deserialize_task_result preserves accumulated errors.""" + from django.tasks.base import TaskError + + error = TaskError( + exception_class_path="ValueError", + traceback="traceback", + ) + result = _task_result_with_retry(boom_with_retry, errors=[error]) + data = ThreadmillTaskBackend.serialize_task_result(result) + restored = ThreadmillTaskBackend.deserialize_task_result(data) + assert len(restored.errors) == 1 + assert restored.errors[0].exception_class_path == "ValueError" diff --git a/tests/backends/test_redis.py b/tests/backends/test_redis.py index cbb24c9..026f4b6 100644 --- a/tests/backends/test_redis.py +++ b/tests/backends/test_redis.py @@ -687,3 +687,161 @@ def test_peek__empty_history_returns_nothing(self): ) assert successful == [] assert failed == [] + + +from tests.testapp.tasks import boom_with_retry # noqa: E402 + + +class TestRedisRequeue: + """Tests for RedisTaskBackend.requeue.""" + + def test_requeue__moves_from_running_to_deferred(self) -> None: + """requeue() removes from running set and adds to deferred set.""" + backend = RedisTaskBackend( + "requeue_test", + { + "QUEUES": ["default"], + "REDIS_URL": "redis://localhost:6379/0", + "OPTIONS": { + "lease_ttl": datetime.timedelta(hours=1), + "result_ttl": datetime.timedelta(seconds=60), + }, + }, + ) + try: + task_result = backend.enqueue(boom_with_retry, args=[]) + acquired = backend.acquire( + timeout=datetime.timedelta(seconds=1), worker="requeue-test" + ) + assert acquired is not None + + # Simulate a failed execution + from django.tasks.base import TaskError + + failed = dataclasses.replace( + acquired, + status=TaskResultStatus.FAILED, + finished_at=timezone.now(), + errors=[ + TaskError( + exception_class_path="ValueError", + traceback="ValueError: boom", + ) + ], + ) + + run_after = timezone.now() + datetime.timedelta(seconds=10) + backend.requeue(failed, run_after) + + running_key = backend.RUNNING_KEY.format( + prefix=backend.key_prefix, queue_name="default" + ) + deferred_key = backend.DEFERRED_KEY.format( + prefix=backend.key_prefix, queue_name="default" + ) + assert backend.client.zscore(running_key, task_result.id) is None + assert backend.client.zscore(deferred_key, task_result.id) is not None + finally: + backend.close() + + def test_requeue__preserves_id_and_errors(self) -> None: + """requeue() preserves the task ID and accumulated errors.""" + backend = RedisTaskBackend( + "requeue_preserve_test", + { + "QUEUES": ["default"], + "REDIS_URL": "redis://localhost:6379/0", + "OPTIONS": { + "lease_ttl": datetime.timedelta(hours=1), + "result_ttl": datetime.timedelta(seconds=60), + }, + }, + ) + try: + task_result = backend.enqueue(boom_with_retry, args=[]) + acquired = backend.acquire( + timeout=datetime.timedelta(seconds=1), worker="preserve-test" + ) + assert acquired is not None + + from django.tasks.base import TaskError + + error = TaskError( + exception_class_path="ValueError", + traceback="ValueError: boom", + ) + failed = dataclasses.replace( + acquired, + status=TaskResultStatus.FAILED, + finished_at=timezone.now(), + errors=[error], + ) + + run_after = timezone.now() + datetime.timedelta(seconds=10) + backend.requeue(failed, run_after) + + # Verify the stored data preserves ID and errors + task_key = backend.TASK_KEY.format( + prefix=backend.key_prefix, task_id=task_result.id + ) + stored_data = backend.client.hget(task_key, "data") + restored = backend.deserialize_task_result(stored_data) + assert restored.id == task_result.id + assert len(restored.errors) == 1 + assert restored.errors[0].exception_class_path == "ValueError" + assert restored.status == TaskResultStatus.READY + assert restored.started_at is None + assert restored.finished_at is None + finally: + backend.close() + + def test_requeue__task_is_re_acquirable_after_delay(self) -> None: + """Requeued task can be acquired after run_after has elapsed.""" + backend = RedisTaskBackend( + "requeue_acquire_test", + { + "QUEUES": ["default"], + "REDIS_URL": "redis://localhost:6379/0", + "OPTIONS": { + "lease_ttl": datetime.timedelta(hours=1), + "result_ttl": datetime.timedelta(seconds=60), + }, + }, + ) + try: + task_result = backend.enqueue(boom_with_retry, args=[]) + acquired = backend.acquire( + timeout=datetime.timedelta(seconds=1), worker="requeue-acq-test" + ) + assert acquired is not None + + from django.tasks.base import TaskError + + failed = dataclasses.replace( + acquired, + status=TaskResultStatus.FAILED, + finished_at=timezone.now(), + errors=[ + TaskError( + exception_class_path="ValueError", + traceback="ValueError: boom", + ) + ], + ) + + # Requeue with a past run_after so it's immediately due + run_after = timezone.now() - datetime.timedelta(seconds=1) + backend.requeue(failed, run_after) + + # Run the broker to move the deferred task to the ready queue + broker = RedisBroker(backend) + broker.main() + + # The task should be acquirable again + re_acquired = backend.acquire( + timeout=datetime.timedelta(seconds=1), worker="requeue-acq-test-2" + ) + assert re_acquired is not None + assert re_acquired.id == task_result.id + finally: + backend.close() diff --git a/tests/test_executor.py b/tests/test_executor.py index 19e1c09..8958e43 100644 --- a/tests/test_executor.py +++ b/tests/test_executor.py @@ -1,6 +1,8 @@ from __future__ import annotations import dataclasses +import datetime +import logging import multiprocessing import sys import threading @@ -309,3 +311,160 @@ def test_create_task_error__builds_task_error(self): error = WorkerThread.create_task_error(sys.exc_info()[1]) assert "RuntimeError" in error.exception_class_path assert "test error" in error.traceback + + +from tests.testapp.tasks import ( # noqa: E402 + boom_no_retry, + boom_retry_raises, + boom_retry_thrice, + boom_with_retry, +) + + +class TestRetryDelay: + """Tests for WorkerThread.retry_delay.""" + + def test_retry_delay__none_when_task_has_no_retry(self) -> None: + """Return None when the task has no retry callback.""" + result = _task_result(boom) + assert WorkerThread.retry_delay(result) is None + + def test_retry_delay__returns_timedelta_from_callback(self) -> None: + """Return the timedelta from the retry callback.""" + result = _task_result(boom_with_retry) + result = dataclasses.replace( + result, + status=TaskResultStatus.FAILED, + errors=[WorkerThread.create_task_error(ValueError("boom"))], + ) + delay = WorkerThread.retry_delay(result) + assert delay == datetime.timedelta(seconds=1) + + def test_retry_delay__none_when_callback_returns_none(self) -> None: + """Return None when the retry callback returns None.""" + result = _task_result(boom_no_retry) + result = dataclasses.replace( + result, + status=TaskResultStatus.FAILED, + errors=[WorkerThread.create_task_error(ValueError("boom"))], + ) + assert WorkerThread.retry_delay(result) is None + + def test_retry_delay__none_when_callback_raises(self, caplog) -> None: + """Return None and log when the retry callback raises an exception.""" + mp_logger = multiprocessing.get_logger() + mp_logger.addHandler(caplog.handler) + mp_logger.setLevel(logging.ERROR) + result = _task_result(boom_retry_raises) + result = dataclasses.replace( + result, + status=TaskResultStatus.FAILED, + errors=[WorkerThread.create_task_error(ValueError("boom"))], + ) + try: + assert WorkerThread.retry_delay(result) is None + finally: + mp_logger.removeHandler(caplog.handler) + assert "Retry callback failed" in caplog.text + + def test_retry_delay__passes_task_context(self) -> None: + """Pass TaskContext with task_result to the retry callback.""" + result = _task_result(boom_retry_thrice, worker_ids=["w1"]) + result = dataclasses.replace( + result, + status=TaskResultStatus.FAILED, + errors=[WorkerThread.create_task_error(ValueError("boom"))], + ) + delay = WorkerThread.retry_delay(result) + assert delay == datetime.timedelta(seconds=1) + + def test_retry_delay__callback_receives_errors(self) -> None: + """Retry callback can access the latest error via context.task_result.errors.""" + result = _task_result(boom_with_retry) + error = WorkerThread.create_task_error(ValueError("boom")) + result = dataclasses.replace( + result, + status=TaskResultStatus.FAILED, + errors=[error], + ) + WorkerThread.retry_delay(result) + # The retry_always callback doesn't inspect errors, but the context + # must contain them. Verify via retry_thrice which checks attempt count. + result_with_workers = dataclasses.replace(result, worker_ids=["w1", "w2"]) + delay = WorkerThread.retry_delay(result_with_workers) + assert delay == datetime.timedelta(seconds=1) + + +class TestWorkerThreadRetry: + """Tests for retry integration in WorkerThread.run().""" + + def test_run__requeues_failed_task_with_retry(self) -> None: + """run() requeues a FAILED task when retry_delay returns a timedelta.""" + enqueued = default_task_backend.enqueue(boom_with_retry, args=[]) + worker = _make_worker(max_tasks=1) + worker.lock = threading.Lock() + worker.expired = threading.Event() + + thread = WorkerThread( + worker=worker, + index=0, + backend=default_task_backend, + ) + thread.run() + + # The task should have been requeued to the deferred set, not acknowledged + from threadmill.backends.redis import RedisTaskBackend + + deferred_key = RedisTaskBackend.DEFERRED_KEY.format( + prefix=default_task_backend.key_prefix, queue_name="default" + ) + assert default_task_backend.client.zscore(deferred_key, enqueued.id) is not None + + def test_run__acknowledges_failed_task_without_retry(self) -> None: + """run() acknowledges a FAILED task when retry_delay returns None.""" + enqueued = default_task_backend.enqueue(boom_no_retry, args=[]) + worker = _make_worker(max_tasks=1) + worker.lock = threading.Lock() + worker.expired = threading.Event() + + thread = WorkerThread( + worker=worker, + index=0, + backend=default_task_backend, + ) + thread.run() + + # The task should be acknowledged (FAILED result, not in deferred) + from threadmill.backends.redis import RedisTaskBackend + + deferred_key = RedisTaskBackend.DEFERRED_KEY.format( + prefix=default_task_backend.key_prefix, queue_name="default" + ) + assert default_task_backend.client.zscore(deferred_key, enqueued.id) is None + + result = default_task_backend.get_result(enqueued.id) + assert result.status == TaskResultStatus.FAILED + + def test_run__acknowledges_failed_task_when_callback_raises(self) -> None: + """run() acknowledges a FAILED task when the retry callback raises.""" + enqueued = default_task_backend.enqueue(boom_retry_raises, args=[]) + worker = _make_worker(max_tasks=1) + worker.lock = threading.Lock() + worker.expired = threading.Event() + + thread = WorkerThread( + worker=worker, + index=0, + backend=default_task_backend, + ) + thread.run() + + from threadmill.backends.redis import RedisTaskBackend + + deferred_key = RedisTaskBackend.DEFERRED_KEY.format( + prefix=default_task_backend.key_prefix, queue_name="default" + ) + assert default_task_backend.client.zscore(deferred_key, enqueued.id) is None + + result = default_task_backend.get_result(enqueued.id) + assert result.status == TaskResultStatus.FAILED diff --git a/tests/testapp/tasks.py b/tests/testapp/tasks.py index f13db68..0be41ac 100644 --- a/tests/testapp/tasks.py +++ b/tests/testapp/tasks.py @@ -1,9 +1,11 @@ import asyncio +import datetime import logging import random import uuid from django.tasks import task +from django.tasks.base import TaskContext logger = logging.getLogger(__name__) @@ -65,3 +67,49 @@ def random_crash(): """Raise a random exception.""" if random.random() < 0.75: # noqa: S311 exit(1) + + +def retry_always(context: TaskContext) -> datetime.timedelta: + """Retry with a fixed 1-second delay.""" + return datetime.timedelta(seconds=1) + + +def retry_never(context: TaskContext) -> datetime.timedelta | None: + """Never retry — always return None.""" + return None + + +def retry_thrice(context: TaskContext) -> datetime.timedelta | None: + """Retry up to 3 attempts, then stop.""" + if context.attempt >= 3: + return None + return datetime.timedelta(seconds=1) + + +def retry_raise(context: TaskContext) -> datetime.timedelta | None: + """Raise an exception to test retry callback error handling.""" + raise RuntimeError("retry callback crashed") + + +@task(retry=retry_always) +def boom_with_retry(): + """Raise ValueError, but always schedule a retry.""" + raise ValueError("boom") + + +@task(retry=retry_never) +def boom_no_retry(): + """Raise ValueError, retry callback returns None.""" + raise ValueError("boom") + + +@task(retry=retry_thrice) +def boom_retry_thrice(): + """Raise ValueError, retry up to 3 attempts.""" + raise ValueError("boom") + + +@task(retry=retry_raise) +def boom_retry_raises(): + """Raise ValueError, retry callback itself raises.""" + raise ValueError("boom") diff --git a/threadmill/backends/base.py b/threadmill/backends/base.py index 14e9195..f80ae89 100644 --- a/threadmill/backends/base.py +++ b/threadmill/backends/base.py @@ -11,7 +11,9 @@ from django.core.serializers.json import DjangoJSONEncoder from django.tasks import DEFAULT_TASK_QUEUE_NAME, Task, TaskResult, TaskResultStatus from django.tasks.backends.base import BaseTaskBackend -from django.tasks.base import TaskError +from django.tasks.base import TaskContext, TaskError +from django.tasks.exceptions import InvalidTask +from django.utils.inspect import is_module_level_function from django.utils.module_loading import import_string if django.VERSION == (6, 0): @@ -36,6 +38,44 @@ def __reduce__(self): return (self.__class__._reconstruct, (kwargs,)) +@dataclasses.dataclass(frozen=True, slots=True, kw_only=True) +class RetryTask(Task): + """Task with an optional retry callback for backoff scheduling. + + Usage: + + def retry_callback(context: TaskContext) -> datetime.timedelta | None: + if context.attempt < 3: + return datetime.timedelta(seconds=60 * (2 ** context.attempt)) + + @task(retry=retry_callback) + def my_task(): + ... + + """ + + retry: collections.abc.Callable[[TaskContext], datetime.timedelta | None] | None = ( + None + ) + + @property + def retry_path(self) -> str | None: + """Importable dotted path of the retry callback, or None.""" + if self.retry: + return f"{self.retry.__module__}.{self.retry.__qualname__}" + + @classmethod + def _reconstruct(cls, kwargs): + if retry_path := kwargs.pop("retry", None): + kwargs["retry"] = import_string(retry_path) + return super()._reconstruct(kwargs) + + def __reduce__(self): + reconstructor, (kwargs,) = super().__reduce__() + kwargs["retry"] = self.retry_path + return (reconstructor, (kwargs,)) + + @dataclasses.dataclass(kw_only=True, slots=True) class QueueCounts: """Point-in-time cardinality of each queue segment.""" @@ -116,19 +156,20 @@ def default(self, o): field.name: getattr(o, field.name) for field in dataclasses.fields(type(o)) } - if isinstance(o, Task): - return { + if isinstance(o, RetryTask): + data = { field.name: getattr(o, field.name) - for field in dataclasses.fields(Task) - if field.name != "func" # Exclude the function object itself - } | {"func": o.module_path} + for field in dataclasses.fields(RetryTask) + if field.name not in {"func", "retry"} + } | {"func": o.module_path, "retry": o.retry_path} + return data return super().default(o) class ThreadmillTaskBackend(BaseTaskBackend, ABC): """Interface for task queues to be processed by the executor.""" - task_class = Task # can be removed in the future when Django 6.0 support is dropped + task_class = RetryTask supports_async_task = True supports_get_result = True broker_class: type[Broker] | None = None @@ -144,18 +185,8 @@ def deserialize_task_result(cls, data: str) -> TaskResult: def _object_hook(d: dict) -> dict | TaskResult: if "task" in d and isinstance(d["task"], dict) and "func" in d["task"]: task_data = d["task"] - func = import_string(task_data["func"]) - if isinstance(func, cls.task_class): - func = func.func - d["task"] = cls.task_class( - func=func, - **{ - field.name: _parse_datetime(task_data[field.name]) - for field in dataclasses.fields(cls.task_class) - if field.name not in {"func", "takes_context"} - and field.name in task_data - }, - ) + task_data["run_after"] = _parse_datetime(task_data["run_after"]) + d["task"] = cls.task_class._reconstruct(task_data) d["status"] = TaskResultStatus(d["status"]) d["errors"] = [TaskError(**error) for error in d["errors"]] return_value = d.pop("_return_value", None) @@ -168,6 +199,13 @@ def _object_hook(d: dict) -> dict | TaskResult: return json.loads(data, object_hook=_object_hook) + def validate_task(self, task: RetryTask) -> None: + super().validate_task(task) + if task.retry is not None and not is_module_level_function(task.retry): + raise InvalidTask( + "Task's retry function must be defined at a module level." + ) + def acquire( self, *queue_names: str, @@ -192,6 +230,10 @@ def acknowledge(self, task_result: TaskResult) -> None: """Remove the task from the queue and publish the result.""" raise NotImplementedError + def requeue(self, task_result: TaskResult, run_after: datetime.datetime) -> None: + """Re-queue a failed task result for a retry attempt after `run_after`.""" + raise NotImplementedError + def peek( self, queue_name: str = DEFAULT_TASK_QUEUE_NAME, @@ -200,7 +242,7 @@ def peek( count: int = 1, ) -> collections.abc.Generator[TaskResult, None, None]: """ - Yield up to ``count`` tasks from a queue in the given status segment. + Yield up to `count` tasks from a queue in the given status segment. Args: queue_name: The name of the queue to peek into. diff --git a/threadmill/backends/redis.py b/threadmill/backends/redis.py index 23fb0c9..20c55f1 100644 --- a/threadmill/backends/redis.py +++ b/threadmill/backends/redis.py @@ -2,6 +2,7 @@ from __future__ import annotations +import dataclasses import datetime import logging import queue @@ -315,6 +316,34 @@ def acknowledge(self, task_result: TaskResult) -> None: ], ) + def requeue(self, task_result: TaskResult, run_after: datetime.datetime) -> None: + task_result = dataclasses.replace( + task_result, + status=TaskResultStatus.READY, + started_at=None, + finished_at=None, + ) + serialized = self.serialize_task_result(task_result) + running_key = self.RUNNING_KEY.format( + prefix=self.key_prefix, queue_name=task_result.task.queue_name + ) + deferred_key = self.DEFERRED_KEY.format( + prefix=self.key_prefix, queue_name=task_result.task.queue_name + ) + task_key = self.TASK_KEY.format(prefix=self.key_prefix, task_id=task_result.id) + score = self._compute_score(task_result.task.priority, task_result.enqueued_at) + run_after_ms = run_after.timestamp() * 1000 + task_data_ttl = int( + self.lease_ttl.total_seconds() * 3 + self.result_ttl.total_seconds() + ) + + pipe = self.client.pipeline() + pipe.zrem(running_key, task_result.id) + pipe.hset(task_key, mapping={"data": serialized, "score": str(score)}) + pipe.expire(task_key, task_data_ttl) + pipe.zadd(deferred_key, {task_result.id: run_after_ms}) + pipe.execute() + def peek( self, queue_name: str = DEFAULT_TASK_QUEUE_NAME, diff --git a/threadmill/executor.py b/threadmill/executor.py index 3531e22..adb4d5a 100644 --- a/threadmill/executor.py +++ b/threadmill/executor.py @@ -212,10 +212,25 @@ def run(self) -> None: try: result = self.execute_task_result(task_result) - self.backend.acknowledge(result) + if ( + result.status is TaskResultStatus.FAILED + and (delay := self.retry_delay(result)) is not None + ): + self.backend.requeue(result, timezone.now() + delay) + else: + self.backend.acknowledge(result) finally: self.worker.record_task() + @staticmethod + def retry_delay(task_result: TaskResult) -> datetime.timedelta | None: + """Return the retry delay for a failed task, or None to stop retrying.""" + if task_result.task.retry: + try: + return task_result.task.retry(TaskContext(task_result=task_result)) + except Exception: + logger.exception("Retry callback failed for task %r", task_result.id) + def execute_task_result(self, task_result: TaskResult) -> TaskResult: """Execute task from task result and update result lifecycle state.""" logger.info("Executing task %r", task_result.id) From bc66db8a76fc2dea2fa8be52654ae95d4d246341 Mon Sep 17 00:00:00 2001 From: Johannes Maron Date: Mon, 6 Jul 2026 09:46:37 +0200 Subject: [PATCH 2/6] Drop old Python and Django versions --- .github/workflows/ci.yml | 3 --- pyproject.toml | 7 ++----- tests/testapp/settings.py | 4 ++-- threadmill/backends/base.py | 24 +----------------------- 4 files changed, 5 insertions(+), 33 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c5caa81..83093ec 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -22,11 +22,8 @@ jobs: os: - "ubuntu-latest" python-version: - - "3.12" - - "3.13" - "3.14" django-version: - - "6.0.0" - "6.1a1" runs-on: ${{ matrix.os }} services: diff --git a/pyproject.toml b/pyproject.toml index fee79cb..136db4b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,14 +26,11 @@ classifiers = [ "Programming Language :: Python", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.12", - "Programming Language :: Python :: 3.13", "Programming Language :: Python :: 3.14", "Framework :: Django", - "Framework :: Django :: 6.0", "Framework :: Django :: 6.1", ] -requires-python = ">=3.12" +requires-python = ">=3.14" dependencies = ["django>=6.0"] [project.optional-dependencies] @@ -59,7 +56,7 @@ name = "threadmill" write_to = "threadmill/_version.py" [tool.pytest.ini_options] -minversion = "6.0" +minversion = "6.1a1" addopts = "--cov --cov-report=xml --cov-report=term --tb=short -rxs --benchmark-autosave --benchmark-group-by=fullname --benchmark-min-rounds=10" testpaths = ["tests"] DJANGO_SETTINGS_MODULE = "tests.testapp.settings" diff --git a/tests/testapp/settings.py b/tests/testapp/settings.py index 2b421ef..1d4e2fd 100644 --- a/tests/testapp/settings.py +++ b/tests/testapp/settings.py @@ -14,7 +14,7 @@ import os from pathlib import Path -from django.tasks import DEFAULT_TASK_QUEUE_NAME +from django.tasks import DEFAULT_TASK_BACKEND_ALIAS, DEFAULT_TASK_QUEUE_NAME # Build paths inside the project like this: BASE_DIR / 'subdir'. BASE_DIR = Path(__file__).resolve().parent.parent @@ -87,7 +87,7 @@ } TASKS = { - "default": { + DEFAULT_TASK_BACKEND_ALIAS: { "BACKEND": "threadmill.backends.redis.RedisTaskBackend", "QUEUES": [DEFAULT_TASK_QUEUE_NAME, "compute", "io", "memory"], "REDIS_URL": os.environ.get("REDIS_URL", "redis://localhost:6379/0"), diff --git a/threadmill/backends/base.py b/threadmill/backends/base.py index f80ae89..b61b3ef 100644 --- a/threadmill/backends/base.py +++ b/threadmill/backends/base.py @@ -7,7 +7,6 @@ import threading from abc import ABC -import django from django.core.serializers.json import DjangoJSONEncoder from django.tasks import DEFAULT_TASK_QUEUE_NAME, Task, TaskResult, TaskResultStatus from django.tasks.backends.base import BaseTaskBackend @@ -16,27 +15,6 @@ from django.utils.inspect import is_module_level_function from django.utils.module_loading import import_string -if django.VERSION == (6, 0): - # https://github.com/django/django/commit/8c8b833d32c02d3ae6f43b04bb1e45968796b402 - @dataclasses.dataclass(frozen=True, slots=True, kw_only=True) - class Task(Task): - @classmethod - def _reconstruct(cls, kwargs): - func_path = kwargs["func"] - try: - func = import_string(func_path) - kwargs["func"] = func.func - except (ImportError, AttributeError) as e: - msg = f"Expected {func_path!r} to point to a Task instance." - raise ValueError(msg) from e - return cls(**kwargs) - - def __reduce__(self): - kwargs = {f.name: getattr(self, f.name) for f in dataclasses.fields(self)} - kwargs["func"] = self.module_path - - return (self.__class__._reconstruct, (kwargs,)) - @dataclasses.dataclass(frozen=True, slots=True, kw_only=True) class RetryTask(Task): @@ -240,7 +218,7 @@ def peek( *, status: TaskResultStatus, count: int = 1, - ) -> collections.abc.Generator[TaskResult, None, None]: + ) -> collections.abc.Generator[TaskResult]: """ Yield up to `count` tasks from a queue in the given status segment. From 33576a7f21cfd8607223253efa1f6c78f0e71927 Mon Sep 17 00:00:00 2001 From: Johannes Maron Date: Mon, 6 Jul 2026 09:55:03 +0200 Subject: [PATCH 3/6] Update pyupgrade --- .pre-commit-config.yaml | 1 + tests/backends/test_base.py | 2 -- tests/backends/test_redis.py | 2 -- tests/test_command.py | 2 -- tests/test_executor.py | 2 -- tests/test_inspector.py | 2 -- threadmill/backends/base.py | 2 -- threadmill/backends/redis.py | 2 -- threadmill/exceptions.py | 2 -- threadmill/executor.py | 4 +--- threadmill/inspector/app.py | 2 -- 11 files changed, 2 insertions(+), 21 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index e764ff4..914248c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -20,6 +20,7 @@ repos: rev: v3.21.2 hooks: - id: pyupgrade + args: [--py314-plus] - repo: https://github.com/adamchainz/django-upgrade rev: 1.31.1 hooks: diff --git a/tests/backends/test_base.py b/tests/backends/test_base.py index 36ba540..6cc55d8 100644 --- a/tests/backends/test_base.py +++ b/tests/backends/test_base.py @@ -1,5 +1,3 @@ -from __future__ import annotations - import datetime import time import uuid diff --git a/tests/backends/test_redis.py b/tests/backends/test_redis.py index 026f4b6..42430c7 100644 --- a/tests/backends/test_redis.py +++ b/tests/backends/test_redis.py @@ -1,5 +1,3 @@ -from __future__ import annotations - import dataclasses import datetime import logging diff --git a/tests/test_command.py b/tests/test_command.py index c4f755f..99ee66a 100644 --- a/tests/test_command.py +++ b/tests/test_command.py @@ -1,5 +1,3 @@ -from __future__ import annotations - import argparse import signal from unittest.mock import patch diff --git a/tests/test_executor.py b/tests/test_executor.py index 8958e43..bcd8661 100644 --- a/tests/test_executor.py +++ b/tests/test_executor.py @@ -1,5 +1,3 @@ -from __future__ import annotations - import dataclasses import datetime import logging diff --git a/tests/test_inspector.py b/tests/test_inspector.py index 11aaa62..6da6ad2 100644 --- a/tests/test_inspector.py +++ b/tests/test_inspector.py @@ -1,5 +1,3 @@ -from __future__ import annotations - import dataclasses import datetime diff --git a/threadmill/backends/base.py b/threadmill/backends/base.py index b61b3ef..6bc94c8 100644 --- a/threadmill/backends/base.py +++ b/threadmill/backends/base.py @@ -1,5 +1,3 @@ -from __future__ import annotations - import collections.abc import dataclasses import datetime diff --git a/threadmill/backends/redis.py b/threadmill/backends/redis.py index 20c55f1..48e9311 100644 --- a/threadmill/backends/redis.py +++ b/threadmill/backends/redis.py @@ -1,7 +1,5 @@ """Redis-backed durable priority queue backend for Django's task framework.""" -from __future__ import annotations - import dataclasses import datetime import logging diff --git a/threadmill/exceptions.py b/threadmill/exceptions.py index acc4c54..e7280d3 100644 --- a/threadmill/exceptions.py +++ b/threadmill/exceptions.py @@ -1,7 +1,5 @@ """Custom exceptions for the threadmill task framework.""" -from __future__ import annotations - class AcknowledgementTimeout(Exception): """Raised when a task's lease has expired before it could be acknowledged.""" diff --git a/threadmill/executor.py b/threadmill/executor.py index adb4d5a..cfd4d64 100644 --- a/threadmill/executor.py +++ b/threadmill/executor.py @@ -1,7 +1,5 @@ """Task worker executor implementation.""" -from __future__ import annotations - import asyncio import dataclasses import datetime @@ -205,7 +203,7 @@ def run(self) -> None: timeout=datetime.timedelta(seconds=1), worker=self.name, ) - except (Empty, TimeoutError): + except Empty, TimeoutError: if self.worker.shutdown_requested.is_set() or self.worker.exit_empty: return continue diff --git a/threadmill/inspector/app.py b/threadmill/inspector/app.py index eb6193c..9cc2d5f 100644 --- a/threadmill/inspector/app.py +++ b/threadmill/inspector/app.py @@ -1,7 +1,5 @@ """Textual app for the inspector TUI.""" -from __future__ import annotations - import datetime import logging import math From 1fea7084a77e08a6e6b5a5926f5dfbe3f97dca74 Mon Sep 17 00:00:00 2001 From: Johannes Maron Date: Mon, 6 Jul 2026 19:14:09 +0200 Subject: [PATCH 4/6] Reorganize imports --- tests/backends/test_base.py | 31 ++++++++++++++++--------------- tests/test_executor.py | 21 ++++++++++----------- 2 files changed, 26 insertions(+), 26 deletions(-) diff --git a/tests/backends/test_base.py b/tests/backends/test_base.py index 6cc55d8..4077998 100644 --- a/tests/backends/test_base.py +++ b/tests/backends/test_base.py @@ -3,11 +3,24 @@ import uuid import pytest -from django.tasks import TaskResult, TaskResultStatus +from django.tasks import ( + TaskResult, + TaskResultStatus, + TaskResultStatus as TRS, +) from django.tasks.base import TaskError -from django.utils import timezone +from django.tasks.exceptions import InvalidTask +from django.utils import ( + timezone, + timezone as tz, +) -from threadmill.backends.base import Broker, ThreadmillTaskBackend +from tests.testapp.tasks import ( + boom_with_retry, + echo, + retry_always, +) +from threadmill.backends.base import Broker, RetryTask, ThreadmillTaskBackend from threadmill.exceptions import AcknowledgementTimeout @@ -131,18 +144,6 @@ def test_shutdown__sets_event(self) -> None: assert broker.shutdown_requested.is_set() -from django.tasks import TaskResultStatus as TRS # noqa: E402 -from django.tasks.exceptions import InvalidTask # noqa: E402 -from django.utils import timezone as tz # noqa: E402 - -from tests.testapp.tasks import ( # noqa: E402 - boom_with_retry, - echo, - retry_always, -) -from threadmill.backends.base import RetryTask # noqa: E402 - - def _task_result_with_retry(task, *, errors=None, worker_ids=None) -> TaskResult: """Build a FAILED TaskResult with errors for retry testing.""" now = tz.now() diff --git a/tests/test_executor.py b/tests/test_executor.py index bcd8661..643f721 100644 --- a/tests/test_executor.py +++ b/tests/test_executor.py @@ -16,9 +16,16 @@ ) from django.utils import timezone -from tests.testapp.tasks import boom, echo -from threadmill.backends.base import Broker # noqa: E402 -from threadmill.executor import TaskExecutor, WorkerProcess, WorkerThread # noqa: E402 +from tests.testapp.tasks import ( + boom, + boom_no_retry, + boom_retry_raises, + boom_retry_thrice, + boom_with_retry, + echo, +) +from threadmill.backends.base import Broker +from threadmill.executor import TaskExecutor, WorkerProcess, WorkerThread @task(queue_name="default") @@ -311,14 +318,6 @@ def test_create_task_error__builds_task_error(self): assert "test error" in error.traceback -from tests.testapp.tasks import ( # noqa: E402 - boom_no_retry, - boom_retry_raises, - boom_retry_thrice, - boom_with_retry, -) - - class TestRetryDelay: """Tests for WorkerThread.retry_delay.""" From a739d24dec9a632e27610ba844ffa4e23856acc6 Mon Sep 17 00:00:00 2001 From: Johannes Maron Date: Mon, 6 Jul 2026 20:18:34 +0200 Subject: [PATCH 5/6] Add callable --- tests/test_command.py | 6 ++- tests/test_retry.py | 81 +++++++++++++++++++++++++++++++++++++ threadmill/backends/base.py | 18 +++++++-- threadmill/retry.py | 46 +++++++++++++++++++++ 4 files changed, 145 insertions(+), 6 deletions(-) create mode 100644 tests/test_retry.py create mode 100644 threadmill/retry.py diff --git a/tests/test_command.py b/tests/test_command.py index 99ee66a..d808520 100644 --- a/tests/test_command.py +++ b/tests/test_command.py @@ -4,7 +4,7 @@ import pytest from django.core.management import CommandError, call_command -from django.tasks import Task, default_task_backend +from django.tasks import default_task_backend from tests.testapp.tasks import compute_workload, io_workload, memory_workload from threadmill.management.commands import threadmill @@ -131,7 +131,9 @@ def test_call_command__benchmark_memory_leak_recovery( def test_call_command__benchmark_default_queue(self, benchmark): """Benchmark command execution for default queue tasks.""" backend = default_task_backend - task = Task(func=compute_workload.func, queue_name="default") + task = default_task_backend.task_class( + func=compute_workload.func, queue_name="default" + ) for _ in range(100): backend.enqueue(task, args=[]) diff --git a/tests/test_retry.py b/tests/test_retry.py new file mode 100644 index 0000000..3fc45f6 --- /dev/null +++ b/tests/test_retry.py @@ -0,0 +1,81 @@ +import datetime + +from threadmill import retry + + +class TestExponentialBackoff: + def test_call__exponential_backoff(self): + """Test the ExponentialBackoff retry strategy.""" + backoff = retry.ExponentialBackoff( + base_delay=datetime.timedelta(seconds=1), + max_delay=datetime.timedelta(seconds=10), + factor=2.0, + max_retries=5, + expected_exceptions=(ValueError,), + ) + + class DummyContext: + def __init__(self, attempt, exception_class): + self.attempt = attempt + self.task_result = type( + "TaskResult", + (), + { + "errors": [ + type("Error", (), {"exception_class": exception_class})() + ] + }, + )() + + # Test that it returns the correct delay for each attempt + for attempt in range(5): + context = DummyContext(attempt, ValueError) + delay = backoff(context) + expected_delay = min( + backoff.base_delay * (backoff.factor**attempt), backoff.max_delay + ) + assert delay == expected_delay + + # Test that it returns None after max_retries + context = DummyContext(5, ValueError) + assert backoff(context) is None + + # Test that it returns None for unexpected exceptions + context = DummyContext(0, KeyError) + assert backoff(context) is None + + def test_deconstruct(self): + """Test the deconstruct method of ExponentialBackoff.""" + backoff = retry.ExponentialBackoff( + base_delay=datetime.timedelta(seconds=1), + max_delay=datetime.timedelta(seconds=10), + factor=2.0, + max_retries=5, + expected_exceptions=(ValueError,), + ) + class_path, args, kwargs = backoff.deconstruct() + assert ( + class_path + == f"{backoff.__class__.__module__}.{backoff.__class__.__qualname__}" + ) + assert args == () + assert kwargs == { + "base_delay": datetime.timedelta(seconds=1), + "max_delay": datetime.timedelta(seconds=10), + "factor": 2.0, + "max_retries": 5, + "expected_exceptions": (ValueError,), + } + + def test_deconstruct__reconstruct(self): + """Test that deconstruct and reconstruct work together.""" + backoff = retry.ExponentialBackoff( + base_delay=datetime.timedelta(seconds=1), + max_delay=datetime.timedelta(seconds=10), + factor=2.0, + max_retries=5, + expected_exceptions=(ValueError,), + ) + class_path, args, kwargs = backoff.deconstruct() + reconstructed_backoff = retry.ExponentialBackoff(**kwargs) + assert backoff == reconstructed_backoff diff --git a/threadmill/backends/base.py b/threadmill/backends/base.py index 6bc94c8..54f61b3 100644 --- a/threadmill/backends/base.py +++ b/threadmill/backends/base.py @@ -3,6 +3,7 @@ import datetime import json import threading +import typing from abc import ABC from django.core.serializers.json import DjangoJSONEncoder @@ -35,15 +36,22 @@ def my_task(): ) @property - def retry_path(self) -> str | None: + def retry_path(self) -> tuple[str, tuple, dict[str, typing.Any]] | str | None: """Importable dotted path of the retry callback, or None.""" if self.retry: + if hasattr(self.retry, "deconstruct"): + return self.retry.deconstruct() return f"{self.retry.__module__}.{self.retry.__qualname__}" @classmethod def _reconstruct(cls, kwargs): if retry_path := kwargs.pop("retry", None): - kwargs["retry"] = import_string(retry_path) + try: + retry_path, args, kwargs = retry_path + except ValueError: + kwargs["retry"] = import_string(retry_path) + else: + kwargs["retry"] = import_string(retry_path)(*args, **kwargs) return super()._reconstruct(kwargs) def __reduce__(self): @@ -177,9 +185,11 @@ def _object_hook(d: dict) -> dict | TaskResult: def validate_task(self, task: RetryTask) -> None: super().validate_task(task) - if task.retry is not None and not is_module_level_function(task.retry): + if task.retry is not None and not ( + is_module_level_function(task.retry) or hasattr(task.retry, "deconstruct") + ): raise InvalidTask( - "Task's retry function must be defined at a module level." + "Task's retry function must be defined at a module level or be a deconstructible callable." ) def acquire( diff --git a/threadmill/retry.py b/threadmill/retry.py new file mode 100644 index 0000000..41b30dd --- /dev/null +++ b/threadmill/retry.py @@ -0,0 +1,46 @@ +"""Reusable callables for retrying tasks.""" + +import dataclasses +import datetime +import typing + +if typing.TYPE_CHECKING: + from django.tasks import TaskContext + + +@dataclasses.dataclass(frozen=True, slots=True, kw_only=True) +class ExponentialBackoff: + """ + Exponential backoff retry strategy. + + Usage: + + @task(retry=ExponentialBackoff(base_delay=datetime.timedelta(seconds=1), max_delay=datetime.timedelta(minutes=5), factor=2.0, max_retries=5)) + def my_task(): + ... + """ + + base_delay: datetime.timedelta = dataclasses.field( + default=datetime.timedelta(seconds=1), doc="Base delay in seconds" + ) + max_delay: datetime.timedelta = dataclasses.field( + default=datetime.timedelta(minutes=60), doc="Maximum delay in seconds" + ) + factor: float = dataclasses.field(default=2.0, doc="Exponential factor for backoff") + max_retries: int = dataclasses.field(default=5, doc="Maximum number of retries") + expected_exceptions: tuple[type[Exception], ...] = dataclasses.field( + default=(Exception,), doc="Tuple of exception classes that trigger a retry" + ) + + def __call__(self, context: TaskContext) -> datetime.timedelta | None: + if context.attempt < self.max_retries and issubclass( + context.task_result.errors[-1].exception_class, self.expected_exceptions + ): + return min(self.base_delay * (self.factor**context.attempt), self.max_delay) + + def deconstruct(self): + return ( + f"{self.__class__.__module__}.{self.__class__.__qualname__}", + (), + dataclasses.asdict(self), + ) From 407e32f1015c54cb9e68f4a2ab30e3e6aa7987fc Mon Sep 17 00:00:00 2001 From: Johannes Maron Date: Mon, 6 Jul 2026 21:20:48 +0200 Subject: [PATCH 6/6] Add backoff callable --- README.md | 40 ++++- tests/backends/test_base.py | 277 ++++++++++++++++++++++------------- tests/backends/test_redis.py | 13 +- tests/test_executor.py | 8 - tests/test_retry.py | 149 ++++++++++++++----- threadmill/backends/base.py | 4 +- threadmill/retry.py | 31 +++- 7 files changed, 354 insertions(+), 168 deletions(-) diff --git a/README.md b/README.md index 5a28276..87054fe 100644 --- a/README.md +++ b/README.md @@ -137,6 +137,43 @@ The callback receives a `TaskContext` — use `context.attempt` for the current attempt count and `context.task_result.errors[-1]` for the latest error. Return a `timedelta` to schedule the next attempt, or `None` to stop retrying. +The worker re-queues the failed task, preserving its ID and error history; +the broker promotes it back to the ready queue once the delay elapses. + +#### Built-in `ExponentialBackoff` + +`threadmill.retry.ExponentialBackoff` provides a serializable exponential +backoff strategy out of the box. It caps the delay at `max_delay`, stops +after `max_retries` attempts, and only retries exceptions listed in +`expected_exceptions`. + +```python +import datetime + +from django.tasks import task +from requests import HTTPError + +from threadmill.retry import ExponentialBackoff + + +@task( + retry=ExponentialBackoff( + base_delay=datetime.timedelta(seconds=1), + max_delay=datetime.timedelta(minutes=5), + factor=2.0, + max_retries=5, + expected_exceptions=(HTTPError,), + ) +) +def fetch_github_api(url: str): ... +``` + +#### Custom retry callbacks + +For cases that need logic beyond what `ExponentialBackoff` supports, +write a callable that accepts a `TaskContext` and returns a `timedelta` +or `None`. Use `TaskError.exception_class` to filter by exception type: + ```python import datetime @@ -162,9 +199,6 @@ def retry_on_rate_limit(context: TaskContext) -> datetime.timedelta | None: def fetch_github_api(url: str): ... ``` -The worker re-queues the failed task, preserving its ID and error history; -the broker promotes it back to the ready queue once the delay elapses. - ## Sponsors [![Sponsors](https://django.the-box.sh/sponsors/codingjoe/threadmill.svg)](https://github.com/sponsors/codingjoe) diff --git a/tests/backends/test_base.py b/tests/backends/test_base.py index 4077998..76c659d 100644 --- a/tests/backends/test_base.py +++ b/tests/backends/test_base.py @@ -1,4 +1,6 @@ import datetime +import json +import pickle import time import uuid @@ -20,7 +22,12 @@ echo, retry_always, ) -from threadmill.backends.base import Broker, RetryTask, ThreadmillTaskBackend +from threadmill.backends.base import ( + Broker, + RetryTask, + ThreadmillTaskBackend, + _parse_datetime, +) from threadmill.exceptions import AcknowledgementTimeout @@ -42,7 +49,9 @@ def enqueue(self, task, args, kwargs): ) -class TestAcknowledgeableTaskBackend: +class TestThreadmillTaskBackend: + """Tests for the ThreadmillTaskBackend class.""" + def test_acquire__raise_not_implemented_error(self) -> None: """Raise NotImplementedError for backend acquire API.""" with pytest.raises(NotImplementedError): @@ -77,6 +86,159 @@ def test_telemetry__raise_not_implemented_error(self) -> None: with pytest.raises(NotImplementedError): BackendDouble(alias="default", params={}).telemetry() + def test_validate_task__accepts_module_level_function(self) -> None: + """validate_task accepts a module-level retry callback.""" + backend = BackendDouble(alias="default", params={}) + backend.validate_task(boom_with_retry) + + def test_validate_task__accepts_deconstructible_callable(self) -> None: + """validate_task accepts a retry callback with a deconstruct method.""" + backend = BackendDouble(alias="default", params={}) + task = RetryTask(func=echo.func, retry=DeconstructibleRetry()) + backend.validate_task(task) + + def test_validate_task__accepts_none_retry(self) -> None: + """validate_task accepts a task with retry=None.""" + backend = BackendDouble(alias="default", params={}) + task = RetryTask(func=echo.func) + backend.validate_task(task) + + def test_validate_task__rejects_lambda(self) -> None: + """Reject a non-module-level retry callback during task creation.""" + with pytest.raises(InvalidTask): + RetryTask( + func=echo.func, + retry=lambda ctx: None, + ) + + def test_serialize_task_result__encodes_retry_as_path(self) -> None: + """TaskResultEncoder serializes retry as a dotted path.""" + result = _task_result_with_retry(boom_with_retry) + data = ThreadmillTaskBackend.serialize_task_result(result) + assert "tests.testapp.tasks.retry_always" in data + + def test_serialize_task_result__omits_retry_when_none(self) -> None: + """TaskResultEncoder omits retry when it is None.""" + task = RetryTask(func=echo.func) + result = TaskResult( + task=task, + id=str(uuid.uuid4()), + status=TRS.READY, + enqueued_at=tz.now(), + started_at=None, + finished_at=None, + last_attempted_at=None, + backend="default", + errors=[], + worker_ids=[], + args=[], + kwargs={}, + ) + data = ThreadmillTaskBackend.serialize_task_result(result) + decoded = json.loads(data) + assert decoded["task"]["retry"] is None + + def test_serialize_task_result__raises_type_error_for_unknown_object( + self, + ) -> None: + """TaskResultEncoder raises TypeError for objects it cannot serialize.""" + with pytest.raises(TypeError): + ThreadmillTaskBackend.serialize_task_result(object()) + + def test_deserialize_task_result__roundtrips_retry_callback(self) -> None: + """deserialize_task_result restores the retry callable from path.""" + result = _task_result_with_retry(boom_with_retry) + data = ThreadmillTaskBackend.serialize_task_result(result) + restored = ThreadmillTaskBackend.deserialize_task_result(data) + assert restored.task.retry is retry_always + assert restored.task.retry_path == "tests.testapp.tasks.retry_always" + + def test_deserialize_task_result__roundtrips_without_retry(self) -> None: + """deserialize_task_result works when retry is None.""" + result = _task_result_with_retry(echo) + data = ThreadmillTaskBackend.serialize_task_result(result) + restored = ThreadmillTaskBackend.deserialize_task_result(data) + assert restored.task.retry is None + + def test_deserialize_task_result__roundtrips_deconstructible_retry(self) -> None: + """deserialize_task_result restores a deconstructible retry callable.""" + task = RetryTask(func=echo.func, retry=DeconstructibleRetry(seconds=5)) + result = _task_result_with_retry(task) + data = ThreadmillTaskBackend.serialize_task_result(result) + restored = ThreadmillTaskBackend.deserialize_task_result(data) + assert isinstance(restored.task.retry, DeconstructibleRetry) + assert restored.task.retry.seconds == 5 + + def test_deserialize_task_result__preserves_errors(self) -> None: + """deserialize_task_result preserves accumulated errors.""" + error = TaskError( + exception_class_path="ValueError", + traceback="traceback", + ) + result = _task_result_with_retry(boom_with_retry, errors=[error]) + data = ThreadmillTaskBackend.serialize_task_result(result) + restored = ThreadmillTaskBackend.deserialize_task_result(data) + assert len(restored.errors) == 1 + assert restored.errors[0].exception_class_path == "ValueError" + + def test_deserialize_task_result__preserves_return_value(self) -> None: + """deserialize_task_result preserves the _return_value attribute.""" + result = _task_result_with_retry(echo) + object.__setattr__(result, "_return_value", 42) + data = ThreadmillTaskBackend.serialize_task_result(result) + restored = ThreadmillTaskBackend.deserialize_task_result(data) + assert restored._return_value == 42 + + +class TestRetryTask: + """Tests for the RetryTask class.""" + + def test_retry_path__none_when_no_retry(self) -> None: + """Return None when retry callback is not set.""" + task = RetryTask(func=echo.func) + assert task.retry_path is None + + def test_retry_path__returns_dotted_path(self) -> None: + """Return the importable dotted path of the retry callback.""" + task = RetryTask(func=echo.func, retry=retry_always) + assert task.retry_path == "tests.testapp.tasks.retry_always" + + def test_retry_path__returns_deconstruct_tuple(self) -> None: + """Return the deconstruct() tuple for a deconstructible callable.""" + retry = DeconstructibleRetry(seconds=5) + task = RetryTask(func=echo.func, retry=retry) + assert task.retry_path == retry.deconstruct() + + def test_pickle__roundtrips_retry_callback(self) -> None: + """Pickle preserves the retry callable by importable path.""" + result = _task_result_with_retry(boom_with_retry) + restored = pickle.loads(pickle.dumps(result)) # noqa: S301 + assert restored.task.retry is retry_always + + def test_pickle__roundtrips_without_retry(self) -> None: + """Pickle works when retry is None.""" + result = _task_result_with_retry(echo) + restored = pickle.loads(pickle.dumps(result)) # noqa: S301 + assert restored.task.retry is None + + +class TestParseDateTime: + """Tests for the _parse_datetime helper.""" + + def test_parse_datetime__parses_iso_string(self) -> None: + """Parse an ISO datetime string into a datetime object.""" + now = timezone.now() + assert _parse_datetime(now.isoformat()) == now + + def test_parse_datetime__returns_non_string_unchanged(self) -> None: + """Non-string values pass through unchanged.""" + assert _parse_datetime(None) is None + assert _parse_datetime(42) == 42 + + def test_parse_datetime__returns_invalid_string_unchanged(self) -> None: + """Invalid datetime strings pass through unchanged.""" + assert _parse_datetime("not-a-date") == "not-a-date" + class TestAcknowledgementTimeout: """Tests for the AcknowledgementTimeout exception.""" @@ -163,106 +325,21 @@ def _task_result_with_retry(task, *, errors=None, worker_ids=None) -> TaskResult ) -class TestRetryTask: - """Tests for the RetryTask class.""" +class DeconstructibleRetry: + """A deconstructible retry callable for testing.""" - def test_retry_path__none_when_no_retry(self) -> None: - """Return None when retry callback is not set.""" - task = RetryTask(func=echo.func) - assert task.retry_path is None + def __init__(self, *, seconds: int = 1) -> None: + self.seconds = seconds - def test_retry_path__returns_dotted_path(self) -> None: - """Return the importable dotted path of the retry callback.""" - task = RetryTask(func=echo.func, retry=retry_always) - assert task.retry_path == "tests.testapp.tasks.retry_always" + def __call__(self, context) -> datetime.timedelta: + return datetime.timedelta(seconds=self.seconds) - def test_validate_task__accepts_module_level_function(self) -> None: - """validate_task accepts a module-level retry callback.""" - backend = BackendDouble(alias="default", params={}) - backend.validate_task(boom_with_retry) - - def test_validate_task__rejects_lambda(self) -> None: - """Reject a non-module-level retry callback during task creation.""" - with pytest.raises(InvalidTask): - RetryTask( - func=echo.func, - retry=lambda ctx: None, - ) - - -class TestRetryTaskSerialization: - """Tests for RetryTask JSON and pickle serialization.""" - - def test_serialize__encodes_retry_as_path(self) -> None: - """TaskResultEncoder serializes retry as a dotted path.""" - result = _task_result_with_retry(boom_with_retry) - data = ThreadmillTaskBackend.serialize_task_result(result) - assert "tests.testapp.tasks.retry_always" in data - - def test_serialize__omits_retry_when_none(self) -> None: - """TaskResultEncoder omits retry when it is None.""" - task = RetryTask(func=echo.func) - result = TaskResult( - task=task, - id=str(uuid.uuid4()), - status=TRS.READY, - enqueued_at=tz.now(), - started_at=None, - finished_at=None, - last_attempted_at=None, - backend="default", - errors=[], - worker_ids=[], - args=[], - kwargs={}, + def deconstruct(self): + return ( + "tests.backends.test_base.DeconstructibleRetry", + (), + {"seconds": self.seconds}, ) - data = ThreadmillTaskBackend.serialize_task_result(result) - import json - - decoded = json.loads(data) - assert decoded["task"]["retry"] is None - def test_deserialize__roundtrips_retry_callback(self) -> None: - """deserialize_task_result restores the retry callable from path.""" - result = _task_result_with_retry(boom_with_retry) - data = ThreadmillTaskBackend.serialize_task_result(result) - restored = ThreadmillTaskBackend.deserialize_task_result(data) - assert restored.task.retry is retry_always - assert restored.task.retry_path == "tests.testapp.tasks.retry_always" - - def test_deserialize__roundtrips_without_retry(self) -> None: - """deserialize_task_result works when retry is None.""" - result = _task_result_with_retry(echo) - data = ThreadmillTaskBackend.serialize_task_result(result) - restored = ThreadmillTaskBackend.deserialize_task_result(data) - assert restored.task.retry is None - - def test_pickle__roundtrips_retry_callback(self) -> None: - """Pickle preserves the retry callable by importable path.""" - import pickle - - result = _task_result_with_retry(boom_with_retry) - restored = pickle.loads(pickle.dumps(result)) # noqa: S301 - assert restored.task.retry is retry_always - - def test_pickle__roundtrips_without_retry(self) -> None: - """Pickle works when retry is None.""" - import pickle - - result = _task_result_with_retry(echo) - restored = pickle.loads(pickle.dumps(result)) # noqa: S301 - assert restored.task.retry is None - - def test_deserialize__preserves_errors(self) -> None: - """deserialize_task_result preserves accumulated errors.""" - from django.tasks.base import TaskError - - error = TaskError( - exception_class_path="ValueError", - traceback="traceback", - ) - result = _task_result_with_retry(boom_with_retry, errors=[error]) - data = ThreadmillTaskBackend.serialize_task_result(result) - restored = ThreadmillTaskBackend.deserialize_task_result(data) - assert len(restored.errors) == 1 - assert restored.errors[0].exception_class_path == "ValueError" + def __eq__(self, other): + return isinstance(other, DeconstructibleRetry) and self.seconds == other.seconds diff --git a/tests/backends/test_redis.py b/tests/backends/test_redis.py index 42430c7..a4adff3 100644 --- a/tests/backends/test_redis.py +++ b/tests/backends/test_redis.py @@ -9,7 +9,7 @@ from django.tasks.base import TaskResultStatus from django.utils import timezone -from tests.testapp.tasks import boom, compute_workload, echo +from tests.testapp.tasks import boom, boom_with_retry, compute_workload, echo from threadmill.backends.base import ( BackendTelemetry, QueueCounts, @@ -583,10 +583,6 @@ def test_telemetry__egress_window_follows_display_interval(self): finally: backend.close() - -class TestRedisTaskBackendPeek: - """Tests for the RedisTaskBackend peek API.""" - def _acknowledge(self, status: TaskResultStatus) -> str: """Enqueue, acquire, and acknowledge a task with the given status.""" task_result = default_task_backend.enqueue(echo, args=[1]) @@ -686,13 +682,6 @@ def test_peek__empty_history_returns_nothing(self): assert successful == [] assert failed == [] - -from tests.testapp.tasks import boom_with_retry # noqa: E402 - - -class TestRedisRequeue: - """Tests for RedisTaskBackend.requeue.""" - def test_requeue__moves_from_running_to_deferred(self) -> None: """requeue() removes from running set and adds to deferred set.""" backend = RedisTaskBackend( diff --git a/tests/test_executor.py b/tests/test_executor.py index 643f721..456e376 100644 --- a/tests/test_executor.py +++ b/tests/test_executor.py @@ -317,10 +317,6 @@ def test_create_task_error__builds_task_error(self): assert "RuntimeError" in error.exception_class_path assert "test error" in error.traceback - -class TestRetryDelay: - """Tests for WorkerThread.retry_delay.""" - def test_retry_delay__none_when_task_has_no_retry(self) -> None: """Return None when the task has no retry callback.""" result = _task_result(boom) @@ -391,10 +387,6 @@ def test_retry_delay__callback_receives_errors(self) -> None: delay = WorkerThread.retry_delay(result_with_workers) assert delay == datetime.timedelta(seconds=1) - -class TestWorkerThreadRetry: - """Tests for retry integration in WorkerThread.run().""" - def test_run__requeues_failed_task_with_retry(self) -> None: """run() requeues a FAILED task when retry_delay returns a timedelta.""" enqueued = default_task_backend.enqueue(boom_with_retry, args=[]) diff --git a/tests/test_retry.py b/tests/test_retry.py index 3fc45f6..ada3aff 100644 --- a/tests/test_retry.py +++ b/tests/test_retry.py @@ -1,51 +1,87 @@ import datetime +import uuid + +from django.tasks import TaskResult, TaskResultStatus +from django.tasks.base import TaskContext, TaskError +from django.utils import timezone from threadmill import retry +def _context(*, attempt: int, exception_class: type[Exception]) -> TaskContext: + """Build a TaskContext with a single error of the given exception class.""" + result = TaskResult( + task=print, + id=str(uuid.uuid4()), + status=TaskResultStatus.FAILED, + enqueued_at=timezone.now(), + started_at=timezone.now(), + finished_at=timezone.now(), + last_attempted_at=timezone.now(), + backend="default", + errors=[ + TaskError( + exception_class_path=f"{exception_class.__module__}.{exception_class.__qualname__}", + traceback="", + ) + ], + worker_ids=[f"w{i}" for i in range(attempt)], + args=[], + kwargs={}, + ) + return TaskContext(task_result=result) + + class TestExponentialBackoff: - def test_call__exponential_backoff(self): - """Test the ExponentialBackoff retry strategy.""" + """Tests for the ExponentialBackoff retry strategy.""" + + def test_call__returns_exponential_delay(self) -> None: + """Return exponentially increasing delay for each retry attempt.""" backoff = retry.ExponentialBackoff( base_delay=datetime.timedelta(seconds=1), - max_delay=datetime.timedelta(seconds=10), + max_delay=datetime.timedelta(seconds=60), factor=2.0, max_retries=5, expected_exceptions=(ValueError,), ) - - class DummyContext: - def __init__(self, attempt, exception_class): - self.attempt = attempt - self.task_result = type( - "TaskResult", - (), - { - "errors": [ - type("Error", (), {"exception_class": exception_class})() - ] - }, - )() - - # Test that it returns the correct delay for each attempt for attempt in range(5): - context = DummyContext(attempt, ValueError) + context = _context(attempt=attempt, exception_class=ValueError) delay = backoff(context) - expected_delay = min( - backoff.base_delay * (backoff.factor**attempt), backoff.max_delay - ) + expected_delay = backoff.base_delay * (backoff.factor**attempt) assert delay == expected_delay - # Test that it returns None after max_retries - context = DummyContext(5, ValueError) + def test_call__caps_delay_at_max_delay(self) -> None: + """Return max_delay when the exponential delay exceeds it.""" + backoff = retry.ExponentialBackoff( + base_delay=datetime.timedelta(seconds=1), + max_delay=datetime.timedelta(seconds=10), + factor=2.0, + max_retries=10, + expected_exceptions=(ValueError,), + ) + context = _context(attempt=8, exception_class=ValueError) + delay = backoff(context) + assert delay == backoff.max_delay + + def test_call__returns_none_after_max_retries(self) -> None: + """Return None when the attempt count reaches max_retries.""" + backoff = retry.ExponentialBackoff( + max_retries=5, + expected_exceptions=(ValueError,), + ) + context = _context(attempt=5, exception_class=ValueError) assert backoff(context) is None - # Test that it returns None for unexpected exceptions - context = DummyContext(0, KeyError) + def test_call__returns_none_for_unexpected_exception(self) -> None: + """Return None when the error is not in expected_exceptions.""" + backoff = retry.ExponentialBackoff( + expected_exceptions=(ValueError,), + ) + context = _context(attempt=0, exception_class=KeyError) assert backoff(context) is None - def test_deconstruct(self): - """Test the deconstruct method of ExponentialBackoff.""" + def test_deconstruct__returns_path_args_kwargs(self) -> None: + """Deconstruct returns the importable path, empty args, and field kwargs.""" backoff = retry.ExponentialBackoff( base_delay=datetime.timedelta(seconds=1), max_delay=datetime.timedelta(seconds=10), @@ -54,21 +90,18 @@ def test_deconstruct(self): expected_exceptions=(ValueError,), ) class_path, args, kwargs = backoff.deconstruct() - assert ( - class_path - == f"{backoff.__class__.__module__}.{backoff.__class__.__qualname__}" - ) + assert class_path == "threadmill.retry.ExponentialBackoff" assert args == () assert kwargs == { - "base_delay": datetime.timedelta(seconds=1), - "max_delay": datetime.timedelta(seconds=10), + "base_delay": 1.0, + "max_delay": 10.0, "factor": 2.0, "max_retries": 5, - "expected_exceptions": (ValueError,), + "expected_exceptions": ("builtins.ValueError",), } - def test_deconstruct__reconstruct(self): - """Test that deconstruct and reconstruct work together.""" + def test_deconstruct__reconstructs_roundtrip(self) -> None: + """Reconstructing from deconstruct output yields an equal object.""" backoff = retry.ExponentialBackoff( base_delay=datetime.timedelta(seconds=1), max_delay=datetime.timedelta(seconds=10), @@ -76,6 +109,42 @@ def test_deconstruct__reconstruct(self): max_retries=5, expected_exceptions=(ValueError,), ) - class_path, args, kwargs = backoff.deconstruct() - reconstructed_backoff = retry.ExponentialBackoff(**kwargs) - assert backoff == reconstructed_backoff + _, _, kwargs = backoff.deconstruct() + reconstructed = retry.ExponentialBackoff(**kwargs) + assert backoff == reconstructed + + def test_serialize__roundtrips_through_task_result(self) -> None: + """ExponentialBackoff survives JSON serialization as a retry callback.""" + import uuid + + from django.tasks import TaskResult, TaskResultStatus + from django.utils import timezone + + from threadmill.backends.base import RetryTask, ThreadmillTaskBackend + + backoff = retry.ExponentialBackoff( + base_delay=datetime.timedelta(seconds=2), + max_delay=datetime.timedelta(minutes=5), + max_retries=3, + ) + from tests.testapp.tasks import echo + + task = RetryTask(func=echo.func, retry=backoff) + result = TaskResult( + task=task, + id=str(uuid.uuid4()), + status=TaskResultStatus.FAILED, + enqueued_at=timezone.now(), + started_at=timezone.now(), + finished_at=timezone.now(), + last_attempted_at=timezone.now(), + backend="default", + errors=[], + worker_ids=[], + args=[], + kwargs={}, + ) + data = ThreadmillTaskBackend.serialize_task_result(result) + restored = ThreadmillTaskBackend.deserialize_task_result(data) + assert isinstance(restored.task.retry, retry.ExponentialBackoff) + assert restored.task.retry == backoff diff --git a/threadmill/backends/base.py b/threadmill/backends/base.py index 54f61b3..a29f246 100644 --- a/threadmill/backends/base.py +++ b/threadmill/backends/base.py @@ -47,11 +47,11 @@ def retry_path(self) -> tuple[str, tuple, dict[str, typing.Any]] | str | None: def _reconstruct(cls, kwargs): if retry_path := kwargs.pop("retry", None): try: - retry_path, args, kwargs = retry_path + path, pos_args, kw_args = retry_path except ValueError: kwargs["retry"] = import_string(retry_path) else: - kwargs["retry"] = import_string(retry_path)(*args, **kwargs) + kwargs["retry"] = import_string(path)(*pos_args, **kw_args) return super()._reconstruct(kwargs) def __reduce__(self): diff --git a/threadmill/retry.py b/threadmill/retry.py index 41b30dd..25e720e 100644 --- a/threadmill/retry.py +++ b/threadmill/retry.py @@ -10,8 +10,9 @@ @dataclasses.dataclass(frozen=True, slots=True, kw_only=True) class ExponentialBackoff: - """ - Exponential backoff retry strategy. + """Delay retries exponentially, capped at ``max_delay`` and limited to ``max_retries`` attempts. + + Only retry exceptions listed in ``expected_exceptions``. Usage: @@ -32,6 +33,21 @@ def my_task(): default=(Exception,), doc="Tuple of exception classes that trigger a retry" ) + def __post_init__(self) -> None: + """Resolve deserialized fields back to their original types.""" + for field_name in ("base_delay", "max_delay"): + value = getattr(self, field_name) + if isinstance(value, (int, float)): + object.__setattr__(self, field_name, datetime.timedelta(seconds=value)) + resolved = [] + for exc in self.expected_exceptions: + if isinstance(exc, str): + from django.utils.module_loading import import_string + + exc = import_string(exc) + resolved.append(exc) + object.__setattr__(self, "expected_exceptions", tuple(resolved)) + def __call__(self, context: TaskContext) -> datetime.timedelta | None: if context.attempt < self.max_retries and issubclass( context.task_result.errors[-1].exception_class, self.expected_exceptions @@ -42,5 +58,14 @@ def deconstruct(self): return ( f"{self.__class__.__module__}.{self.__class__.__qualname__}", (), - dataclasses.asdict(self), + { + "base_delay": self.base_delay.total_seconds(), + "max_delay": self.max_delay.total_seconds(), + "factor": self.factor, + "max_retries": self.max_retries, + "expected_exceptions": tuple( + f"{exc.__module__}.{exc.__qualname__}" + for exc in self.expected_exceptions + ), + }, )