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/.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/README.md b/README.md index 4a9430e..87054fe 100644 --- a/README.md +++ b/README.md @@ -130,6 +130,75 @@ 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. + +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 + +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): ... +``` + ## Sponsors [![Sponsors](https://django.the-box.sh/sponsors/codingjoe/threadmill.svg)](https://github.com/sponsors/codingjoe) 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/backends/test_base.py b/tests/backends/test_base.py index 9f6b23d..76c659d 100644 --- a/tests/backends/test_base.py +++ b/tests/backends/test_base.py @@ -1,15 +1,33 @@ -from __future__ import annotations - import datetime +import json +import pickle import time 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, + _parse_datetime, +) from threadmill.exceptions import AcknowledgementTimeout @@ -31,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): @@ -44,6 +64,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): @@ -58,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.""" @@ -123,3 +304,42 @@ def test_shutdown__sets_event(self) -> None: assert not broker.shutdown_requested.is_set() broker.shutdown() assert broker.shutdown_requested.is_set() + + +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 DeconstructibleRetry: + """A deconstructible retry callable for testing.""" + + def __init__(self, *, seconds: int = 1) -> None: + self.seconds = seconds + + def __call__(self, context) -> datetime.timedelta: + return datetime.timedelta(seconds=self.seconds) + + def deconstruct(self): + return ( + "tests.backends.test_base.DeconstructibleRetry", + (), + {"seconds": self.seconds}, + ) + + 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 cbb24c9..a4adff3 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 @@ -11,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, @@ -585,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]) @@ -687,3 +681,154 @@ def test_peek__empty_history_returns_nothing(self): ) assert successful == [] assert failed == [] + + 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_command.py b/tests/test_command.py index c4f755f..d808520 100644 --- a/tests/test_command.py +++ b/tests/test_command.py @@ -1,12 +1,10 @@ -from __future__ import annotations - import argparse import signal from unittest.mock import patch 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 @@ -133,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_executor.py b/tests/test_executor.py index 19e1c09..456e376 100644 --- a/tests/test_executor.py +++ b/tests/test_executor.py @@ -1,6 +1,6 @@ -from __future__ import annotations - import dataclasses +import datetime +import logging import multiprocessing import sys import threading @@ -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") @@ -309,3 +316,144 @@ 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 + + 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) + + 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/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/tests/test_retry.py b/tests/test_retry.py new file mode 100644 index 0000000..ada3aff --- /dev/null +++ b/tests/test_retry.py @@ -0,0 +1,150 @@ +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: + """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=60), + factor=2.0, + max_retries=5, + expected_exceptions=(ValueError,), + ) + for attempt in range(5): + context = _context(attempt=attempt, exception_class=ValueError) + delay = backoff(context) + expected_delay = backoff.base_delay * (backoff.factor**attempt) + assert delay == expected_delay + + 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 + + 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__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), + factor=2.0, + max_retries=5, + expected_exceptions=(ValueError,), + ) + class_path, args, kwargs = backoff.deconstruct() + assert class_path == "threadmill.retry.ExponentialBackoff" + assert args == () + assert kwargs == { + "base_delay": 1.0, + "max_delay": 10.0, + "factor": 2.0, + "max_retries": 5, + "expected_exceptions": ("builtins.ValueError",), + } + + 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), + factor=2.0, + max_retries=5, + expected_exceptions=(ValueError,), + ) + _, _, 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/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/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..a29f246 100644 --- a/threadmill/backends/base.py +++ b/threadmill/backends/base.py @@ -1,39 +1,63 @@ -from __future__ import annotations - import collections.abc import dataclasses import datetime import json import threading +import typing 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 -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): - # 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 +@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)) - return (self.__class__._reconstruct, (kwargs,)) + @task(retry=retry_callback) + def my_task(): + ... + + """ + + retry: collections.abc.Callable[[TaskContext], datetime.timedelta | None] | None = ( + None + ) + + @property + 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): + try: + path, pos_args, kw_args = retry_path + except ValueError: + kwargs["retry"] = import_string(retry_path) + else: + kwargs["retry"] = import_string(path)(*pos_args, **kw_args) + 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) @@ -116,19 +140,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 +169,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 +183,15 @@ 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) or hasattr(task.retry, "deconstruct") + ): + raise InvalidTask( + "Task's retry function must be defined at a module level or be a deconstructible callable." + ) + def acquire( self, *queue_names: str, @@ -192,15 +216,19 @@ 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, *, 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. + 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..48e9311 100644 --- a/threadmill/backends/redis.py +++ b/threadmill/backends/redis.py @@ -1,7 +1,6 @@ """Redis-backed durable priority queue backend for Django's task framework.""" -from __future__ import annotations - +import dataclasses import datetime import logging import queue @@ -315,6 +314,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/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 3531e22..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,17 +203,32 @@ 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 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) 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 diff --git a/threadmill/retry.py b/threadmill/retry.py new file mode 100644 index 0000000..25e720e --- /dev/null +++ b/threadmill/retry.py @@ -0,0 +1,71 @@ +"""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: + """Delay retries exponentially, capped at ``max_delay`` and limited to ``max_retries`` attempts. + + Only retry exceptions listed in ``expected_exceptions``. + + 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 __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 + ): + return min(self.base_delay * (self.factor**context.attempt), self.max_delay) + + def deconstruct(self): + return ( + f"{self.__class__.__module__}.{self.__class__.__qualname__}", + (), + { + "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 + ), + }, + )