diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d5e3aef..e6d783d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -27,6 +27,7 @@ jobs: - "3.14" django-version: - "6.1a1" + runs-on: ${{ matrix.os }} services: redis: image: redis @@ -35,7 +36,6 @@ jobs: options: --entrypoint redis-server env: REDIS_URL: redis:///0 - runs-on: ${{ matrix.os }} steps: - uses: actions/checkout@v7 - uses: astral-sh/setup-uv@v7 @@ -45,26 +45,3 @@ jobs: - uses: codecov/codecov-action@v7 with: token: ${{ secrets.CODECOV_TOKEN }} - pytest-windows-macos: - name: Pytest - permissions: - contents: read - strategy: - matrix: - os: - - "windows-latest" - - "macos-latest" - python-version: - - "3.13" - django-version: - - "6.1a1" - runs-on: ${{ matrix.os }} - steps: - - uses: actions/checkout@v7 - - uses: astral-sh/setup-uv@v7 - with: - python-version: ${{ matrix.python-version }} - - uses: codecov/codecov-action@v7 - with: - token: ${{ secrets.CODECOV_TOKEN }} - - run: uv run --with django~=${{ matrix.django-version }} pytest -m "not benchmark" diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index d942904..da729f9 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -9,8 +9,8 @@ curl -sSL https://raw.githubusercontent.com/codingjoe/naming-things/refs/heads/m ## Design Principles -- Consistency – We never lose data, even if someone unplugs the power or network. - Durability – We recover from any failures, even poorly written tasks. +- Consistency – We never lose data, even if someone unplugs the power or network. - Utilization – We keep the CPU saturated with tasks, not with idle time or waiting for locks. ## Testing diff --git a/README.md b/README.md index d995df4..b0055dd 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ - Django Grinder: A queue agnostic worker for Django's task framework. + Threadmill: A queue agnostic worker for Django's task framework.
Documentation | @@ -34,20 +34,31 @@ ## Setup -You need to have [Django's Task framework][django-tasks] setup properly. +You need to have [Django's Task framework][django-tasks] set up properly. ```console -uv add threadmill +uv add threadmill[redis] ``` -Add `threadmill` to your `INSTALLED_APPS` in `settings.py`: +Add `threadmill` to your `INSTALLED_APPS` in `settings.py` +and configure the task backend: ```python # settings.py +import os + INSTALLED_APPS = [ "threadmill", # ... ] + +TASKS = { + "default": { + "BACKEND": "threadmill.backends.redis.RedisTaskBackend", + "REDIS_URL": os.getenv("REDIS_URL", "redis://localhost:6379/0"), + }, + # ... +} ``` Finally, you launch the worker pool: @@ -58,9 +69,11 @@ uv run manage.py threadmill ## Usage +### Workers + The workers are inspired by Gunicorn, and the CLI is very similar. -### Utilization +#### Utilization Depending on your workload, you can tweak the number of processes and threads. Processes allow for parallel compute (no GIL) while threads are great for low-memory concurrent IO. @@ -69,7 +82,7 @@ Processes allow for parallel compute (no GIL) while threads are great for low-me uv run manage.py threadmill --processes 4 --threads 2 ``` -### Health +#### Health If your tasks leak memory, you can recycle (restart) the workers after a certain number of tasks have been processed: @@ -81,61 +94,33 @@ This will restart the workers after 1000 tasks have been processed, with a rando Should a worker crash or be killed, the pool will automatically restart it. -### Shutdown +#### Shutdown A graceful shutdown is possible with the `SIGTERM` or a keyboard interrupt. -All workers will finish the tasks they acquired and publish them. +All workers will finish the tasks they acquired and acknowledge them. You can use `--exit-empty` to exit immediately after all tasks have been processed, which might be useful for draining a one-off queue. -### Task Backlog - -You can prefetch tasks from a queue to avoid IO latency bottlenecks. -However, this will increase the memory usage of the worker pool. - -```console -uv run manage.py threadmill --prefetch 100 -``` - -### Task Timeouts - -> [!WARNING] -> Work in progress, this feature is not yet stable. +### Redis Backend Options -Task timeouts are important to ensure the long-term health of your pool. -However, they need to be aligned with your queueing system's timeout settings. -The message queue needs to requeue a task that hasn't been acknowledged within the timeout. +The `RedisTaskBackend` accepts the following options under `OPTIONS` in your +`TASKS` configuration: -## Integration +| Option | Default | Description | +| ----------------- | ---------------------- | ------------------------------------------------------------ | +| `lease_ttl` | `timedelta(hours=1)` | Max processing time before a started task is marked FAILED. | +| `result_ttl` | `timedelta(days=1)` | How long task results are retained before automatic removal. | +| `broker_interval` | `timedelta(seconds=1)` | Interval between background broker maintenance passes. | +| `batch_size` | `100` | Max tasks to move or requeue per broker pass. | -> [!NOTE] -> This section is for people who want to integrate Threadmill into their queueing system. +A task that is started but never acknowledged (lease expired) is marked FAILED +with an `AcknowledgementTimeout` error. Set `lease_ttl` comfortably above your +worst-case task runtime. -Threadmill is designed to be durable and requires a queueing system to support late acknowledgement. +All keys for one backend alias share a Redis Cluster hash tag (`{alias}`), so +every multi-key operation — including the cross-queue acquire — runs on a single +shard. Scale horizontally by running additional backend aliases, not by relying +on cross-slot operations. -To use Threadmill, your backend will need to inherit from `threadmill.backends.AcknowledgeableTaskBackend` and implement the following methods: - -```python -class AcknowledgeableTaskBackend(BaseTaskBackend, ABC): - """Provide an interface for tasks queues to be processed by the executor.""" - - def acquire( - self, *queue_names: str, timeout: datetime.timedelta | None = None - ) -> TaskResult: - """ - Return and lock the next task to be processed without removing it from the queue. - - Args: - queue_names: The names of the queues to acquire tasks from. - timeout: The maximum time to wait for a task. If None, wait indefinitely. - - Raises: - TimeoutError: If no task is available within the specified timeout. - """ - raise NotImplementedError - - def acknowledge(self, task_result: TaskResult) -> None: - """Remove the task from the queue and publish the result.""" - raise NotImplementedError -``` +[django-tasks]: https://docs.djangoproject.com/en/stable/topics/tasks/ diff --git a/pyproject.toml b/pyproject.toml index 54f137d..3746b1b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,9 @@ classifiers = [ requires-python = ">=3.12" dependencies = ["django>=6.1a1"] +[project.optional-dependencies] +redis = ["redis>=5.0"] + [project.urls] # https://packaging.python.org/en/latest/specifications/well-known-project-urls/#well-known-labels Homepage = "https://github.com/codingjoe/threadmill" @@ -56,9 +59,9 @@ minversion = "6.0" 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" +asyncio_mode = "auto" markers = [ "benchmark: mark benchmark tests.", - "integration: mark integration tests.", ] [tool.coverage.run] @@ -91,6 +94,7 @@ combine-as-imports = true split-on-trailing-comma = true section-order = ["future", "standard-library", "third-party", "first-party", "local-folder"] force-wrap-aliases = true +known-first-party = ["threadmill", "tests"] [tool.ruff.lint.pydocstyle] convention = "pep257" @@ -105,4 +109,5 @@ test = [ "pytest-asyncio", "pytest-cov", "pytest-django", + "redis>=5.0", ] diff --git a/tests/backends/__init__.py b/tests/backends/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/backends/test_base.py b/tests/backends/test_base.py new file mode 100644 index 0000000..2ee24a0 --- /dev/null +++ b/tests/backends/test_base.py @@ -0,0 +1,116 @@ +from __future__ import annotations + +import datetime +import time +import uuid + +import pytest +from django.tasks import TaskResult, TaskResultStatus +from django.tasks.base import TaskError +from django.utils import timezone + +from threadmill.backends.base import Broker, ThreadmillTaskBackend +from threadmill.exceptions import AcknowledgementTimeout + + +class BackendDouble(ThreadmillTaskBackend): + def enqueue(self, task, args, kwargs): + return TaskResult( + task=task, + id=str(uuid.uuid4()), + status=TaskResultStatus.READY, + enqueued_at=timezone.now(), + started_at=None, + finished_at=None, + last_attempted_at=None, + backend=self.alias, + errors=[], + worker_ids=[], + args=args, + kwargs=kwargs, + ) + + +class TestAcknowledgeableTaskBackend: + def test_acquire__raise_not_implemented_error(self) -> None: + """Raise NotImplementedError for backend acquire API.""" + with pytest.raises(NotImplementedError): + BackendDouble(alias="default", params={}).acquire( + timeout=datetime.timedelta(seconds=1) + ) + + def test_acknowledge__raise_not_implemented_error(self) -> None: + """Raise NotImplementedError for backend acknowledge API.""" + with pytest.raises(NotImplementedError): + BackendDouble(alias="default", params={}).acknowledge(task_result=None) + + def test_peek__raise_not_implemented_error(self) -> None: + """Raise NotImplementedError for backend peek_results API.""" + with pytest.raises(NotImplementedError): + list(BackendDouble(alias="default", params={}).peek("default")) + + +class TestAcknowledgementTimeout: + """Tests for the AcknowledgementTimeout exception.""" + + def test_exception_can_be_instantiated(self) -> None: + """AcknowledgementTimeout can be instantiated.""" + exc = AcknowledgementTimeout() + assert isinstance(exc, Exception) + + def test_exception_can_be_used_in_task_error(self) -> None: + """AcknowledgementTimeout can be used as a TaskError's exception_class_path.""" + error = TaskError( + exception_class_path="threadmill.exceptions.AcknowledgementTimeout", + traceback="Task processing lease expired.", + ) + assert ( + error.exception_class_path == "threadmill.exceptions.AcknowledgementTimeout" + ) + + +class FakeBroker(Broker): + """Broker that records main() calls for testing.""" + + def __init__( + self, *, interval: datetime.timedelta = datetime.timedelta(seconds=0.01) + ) -> None: + super().__init__(interval=interval) + self.maintain_calls: list[float] = [] + + def main(self) -> None: + self.maintain_calls.append(time.monotonic()) + + +class TestBroker: + def test_main__is_noop(self) -> None: + """Base Broker.main() is a no-op.""" + Broker(interval=datetime.timedelta(seconds=1)).main() + + def test_run__calls_maintain_then_exits_on_shutdown(self) -> None: + """run() loops calling main() and exits after shutdown().""" + broker = FakeBroker(interval=datetime.timedelta(seconds=0.01)) + broker.start() + time.sleep(0.05) + broker.shutdown() + broker.join(timeout=1) + assert not broker.is_alive() + assert len(broker.maintain_calls) >= 1 + + def test_interval_is_honored(self) -> None: + """Broker waits at least interval between main() calls.""" + broker = FakeBroker(interval=datetime.timedelta(seconds=0.1)) + broker.start() + time.sleep(0.25) + broker.shutdown() + broker.join(timeout=1) + assert len(broker.maintain_calls) >= 2 + for i in range(1, len(broker.maintain_calls)): + assert broker.maintain_calls[i] - broker.maintain_calls[i - 1] >= 0.09 + + def test_shutdown__sets_event(self) -> None: + """shutdown() sets the shutdown_requested event.""" + broker = Broker(interval=datetime.timedelta(seconds=1)) + assert not broker.shutdown_requested.is_set() + broker.shutdown() + assert broker.shutdown_requested.is_set() diff --git a/tests/backends/test_redis.py b/tests/backends/test_redis.py new file mode 100644 index 0000000..296d769 --- /dev/null +++ b/tests/backends/test_redis.py @@ -0,0 +1,230 @@ +from __future__ import annotations + +import dataclasses +import datetime +import logging +import time +from dataclasses import replace +from unittest.mock import patch + +from django.tasks import default_task_backend +from django.tasks.base import TaskResultStatus +from django.utils import timezone + +from tests.testapp.tasks import compute_workload, echo +from threadmill.backends.redis import RedisBroker, RedisTaskBackend # noqa: E402 + + +class TestRedisBroker: + def test_mover__moves_deferred_task_to_ready(self): + """Mover promotes due deferred tasks to the ready queue.""" + deferred_task = replace( + compute_workload, + run_after=timezone.now() - datetime.timedelta(seconds=10), + ) + task_result = default_task_backend.enqueue(deferred_task, args=[]) + broker = RedisBroker(default_task_backend) + broker.main() + acquired = default_task_backend.acquire(timeout=datetime.timedelta(seconds=1)) + assert acquired is not None + assert acquired.id == task_result.id + + def test_error_path__maintain_continues_after_exception(self, caplog): + """main() logs and continues when _move_queue raises.""" + broker = RedisBroker(default_task_backend) + with caplog.at_level(logging.ERROR): + with patch.object( + broker, + "_move_queue", + side_effect=RuntimeError("boom"), + ): + broker.main() + assert "Mover error for queue" in caplog.text + + +class TestRedisTaskBackend: + """Tests for the RedisTaskBackend update and lease functionality.""" + + def test_acquire__moves_to_running_set(self): + """acquire() moves task directly to running set with worker info.""" + backend = RedisTaskBackend( + "acquire_running_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(echo, args=[42]) + acquired = backend.acquire( + timeout=datetime.timedelta(seconds=1), worker="worker-1" + ) + assert acquired is not None + assert acquired.id == task_result.id + + # Verify task is in running set, not in any processing set + running_key = backend.RUNNING_KEY.format( + prefix=backend.key_prefix, queue_name="default" + ) + assert backend.client.zscore(running_key, task_result.id) is not None + + # Verify task data was updated with worker info + task_key = backend.TASK_KEY.format( + prefix=backend.key_prefix, task_id=task_result.id + ) + stored_data = backend.client.hget(task_key, "data") + deserialized = backend.deserialize_task_result(stored_data) + assert deserialized.status == TaskResultStatus.RUNNING + assert deserialized.worker_ids == ["worker-1"] + assert deserialized.last_attempted_at is not None + finally: + backend.close() + + def test_acquire__sets_last_attempted_at(self): + """acquire() sets last_attempted_at and worker_ids in the stored task data.""" + backend = RedisTaskBackend( + "last_attempted_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(echo, args=[42]) + acquired = backend.acquire( + timeout=datetime.timedelta(seconds=1), worker="test-worker" + ) + assert acquired is not None + assert acquired.last_attempted_at is not None + assert acquired.worker_ids == ["test-worker"] + + # Verify it's persisted in Redis + task_key = backend.TASK_KEY.format( + prefix=backend.key_prefix, task_id=task_result.id + ) + stored_data = backend.client.hget(task_key, "data") + deserialized = backend.deserialize_task_result(stored_data) + assert deserialized.last_attempted_at is not None + assert deserialized.worker_ids == ["test-worker"] + finally: + backend.close() + + def test_running_reaper__fails_expired_tasks(self): + """Running reaper creates FAILED results for tasks with expired lease.""" + backend = RedisTaskBackend( + "running_reaper_test", + { + "QUEUES": ["default"], + "REDIS_URL": "redis://localhost:6379/0", + "OPTIONS": { + "lease_ttl": datetime.timedelta(seconds=1), + "result_ttl": datetime.timedelta(seconds=60), + }, + }, + ) + try: + task_result = backend.enqueue(echo, args=[42]) + acquired = backend.acquire( + timeout=datetime.timedelta(seconds=1), worker="reaper-test" + ) + assert acquired is not None + + # Wait for lease to expire + time.sleep(1.1) + + # Run the broker + broker = RedisBroker(backend) + broker.main() + + # Verify the task result exists and is FAILED + result = backend.get_result(task_result.id) + assert result.status == TaskResultStatus.FAILED + assert len(result.errors) == 1 + assert "AcknowledgementTimeout" in result.errors[0].exception_class_path + finally: + backend.close() + + def test_stale_acknowledge__is_noop(self): + """acknowledge() is a no-op when the task is no longer in the running set.""" + backend = RedisTaskBackend( + "stale_ack_test", + { + "QUEUES": ["default"], + "REDIS_URL": "redis://localhost:6379/0", + "OPTIONS": { + "lease_ttl": datetime.timedelta(seconds=1), + "result_ttl": datetime.timedelta(seconds=60), + }, + }, + ) + try: + task_result = backend.enqueue(echo, args=[42]) + acquired = backend.acquire( + timeout=datetime.timedelta(seconds=1), worker="stale-ack-test" + ) + assert acquired is not None + + # Wait for lease to expire + time.sleep(1.1) + + # Run the broker to reap the running set + broker = RedisBroker(backend) + broker.main() + + # Try to acknowledge the task (should be a no-op since it was reaped) + finished = dataclasses.replace( + acquired, + status=TaskResultStatus.SUCCESSFUL, + finished_at=timezone.now(), + ) + # This should not raise + backend.acknowledge(finished) + + # The result should still be the FAILED one from the reaper + result = backend.get_result(task_result.id) + assert result.status == TaskResultStatus.FAILED + finally: + backend.close() + + def test_lease_ttl_defaults(self): + """lease_ttl defaults to 1h.""" + backend = RedisTaskBackend( + "default_lease_test", + { + "QUEUES": ["default"], + "REDIS_URL": "redis://localhost:6379/0", + "OPTIONS": { + "result_ttl": datetime.timedelta(seconds=60), + }, + }, + ) + try: + assert backend.lease_ttl == datetime.timedelta(hours=1) + finally: + backend.close() + + def test_explicit_lease_ttl(self): + """lease_ttl is used when set.""" + backend = RedisTaskBackend( + "explicit_lease_test", + { + "QUEUES": ["default"], + "REDIS_URL": "redis://localhost:6379/0", + "OPTIONS": { + "lease_ttl": datetime.timedelta(seconds=120), + "result_ttl": datetime.timedelta(seconds=60), + }, + }, + ) + try: + assert backend.lease_ttl == datetime.timedelta(seconds=120) + finally: + backend.close() diff --git a/tests/conftest.py b/tests/conftest.py index e69de29..c6a49eb 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -0,0 +1,17 @@ +import pytest +from django.tasks import default_task_backend + + +def _flush_keys(backend) -> None: + """Delete all threadmill-prefixed Redis keys.""" + keys = backend.client.keys("threadmill:*") + if keys: + backend.client.delete(*keys) + + +@pytest.fixture(autouse=True) +def flush_default_backend(): + """Flush all threadmill keys before and after each test.""" + _flush_keys(default_task_backend) + yield + _flush_keys(default_task_backend) diff --git a/tests/test_backends.py b/tests/test_backends.py deleted file mode 100644 index 2c9ef4f..0000000 --- a/tests/test_backends.py +++ /dev/null @@ -1,25 +0,0 @@ -from __future__ import annotations - -import datetime - -import pytest -from threadmill.backends import AcknowledgeableTaskBackend - - -class BackendDouble(AcknowledgeableTaskBackend): - def enqueue(self, task): - return task - - -class TestAcknowledgeableTaskBackend: - def test_acquire__raise_not_implemented_error(self) -> None: - """Raise NotImplementedError for backend acquire API.""" - with pytest.raises(NotImplementedError): - BackendDouble(alias="default", params={}).acquire( - timeout=datetime.timedelta(seconds=1) - ) - - def test_acknowledge__raise_not_implemented_error(self) -> None: - """Raise NotImplementedError for backend acknowledge API.""" - with pytest.raises(NotImplementedError): - BackendDouble(alias="default", params={}).acknowledge(task_result=None) diff --git a/tests/test_command.py b/tests/test_command.py index 0947aa9..c7ca7ec 100644 --- a/tests/test_command.py +++ b/tests/test_command.py @@ -5,7 +5,9 @@ import pytest from django.core.management import call_command -from django.tasks import default_task_backend +from django.tasks import Task, default_task_backend + +from tests.testapp.tasks import compute_workload, io_workload, memory_workload from threadmill.management.commands import threadmill @@ -29,15 +31,17 @@ def test_add_arguments__register_all_worker_options(self): assert parsed_arguments.threads == 1 assert parsed_arguments.max_tasks == 0 assert parsed_arguments.max_tasks_jitter == 0 - assert parsed_arguments.task_timeout == 3600.0 @pytest.mark.benchmark def test_call_command__benchmark_compute( self, benchmark, ): - """Benchmark command execution for one CPU intense task solved 100 times.""" - default_task_backend.reset() + """Benchmark command execution for compute tasks.""" + backend = default_task_backend + for _ in range(100): + backend.enqueue(compute_workload, args=[]) + benchmark.pedantic( lambda: call_command( "threadmill", @@ -50,15 +54,16 @@ def test_call_command__benchmark_compute( warmup_rounds=0, ) - assert default_task_backend.solved_task_count == 100 - @pytest.mark.benchmark def test_call_command__benchmark_io( self, benchmark, ): - """Benchmark command execution for one CPU intense task solved 100 times.""" - default_task_backend.reset() + """Benchmark command execution for IO tasks.""" + backend = default_task_backend + for _ in range(100): + backend.enqueue(io_workload, args=[]) + benchmark.pedantic( lambda: call_command( "threadmill", @@ -72,15 +77,17 @@ def test_call_command__benchmark_io( warmup_rounds=0, ) - assert default_task_backend.solved_task_count == 100 - @pytest.mark.benchmark def test_call_command__benchmark_compute_and_io( self, benchmark, ): - """Benchmark command execution for one CPU intense task solved 100 times.""" - default_task_backend.reset() + """Benchmark command execution for compute and IO tasks.""" + backend = default_task_backend + for _ in range(100): + backend.enqueue(compute_workload, args=[]) + backend.enqueue(io_workload, args=[]) + benchmark.pedantic( lambda: call_command( "threadmill", @@ -94,15 +101,16 @@ def test_call_command__benchmark_compute_and_io( warmup_rounds=0, ) - assert default_task_backend.solved_task_count == 200 - @pytest.mark.benchmark def test_call_command__benchmark_memory_leak_recovery( self, benchmark, ): - """Benchmark command execution for one CPU intense task solved 100 times.""" - default_task_backend.reset(1000) + """Benchmark command execution for memory leak recovery.""" + backend = default_task_backend + for _ in range(1000): + backend.enqueue(memory_workload, args=[]) + benchmark.pedantic( lambda: call_command( "threadmill", @@ -117,9 +125,13 @@ def test_call_command__benchmark_memory_leak_recovery( ) @pytest.mark.benchmark - def test_call_command__benchmark_random_crash(self, benchmark): - """Benchmark command execution for one CPU intense task solved 100 times.""" - default_task_backend.reset() + 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") + for _ in range(100): + backend.enqueue(task, args=[]) + benchmark.pedantic( lambda: call_command( "threadmill", @@ -130,6 +142,3 @@ def test_call_command__benchmark_random_crash(self, benchmark): iterations=1, warmup_rounds=0, ) - assert default_task_backend.issued_task_count == 100, ( - "All tasks should be issued." - ) diff --git a/tests/test_executor.py b/tests/test_executor.py new file mode 100644 index 0000000..520d060 --- /dev/null +++ b/tests/test_executor.py @@ -0,0 +1,316 @@ +from __future__ import annotations + +import dataclasses +import multiprocessing +import sys +import threading +import time +import uuid + +from django.tasks import Task, TaskResult, TaskResultStatus, default_task_backend +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 + + +def _add(x, y): + return x + y + + +ADD_TASK = Task(func=_add, queue_name="default") + + +def _context_captor(context): + """Task function that captures the context it receives.""" + _context_captor.captured = context + return 42 + + +CONTEXT_TASK = Task(func=_context_captor, queue_name="default", takes_context=True) + + +async def _async_task(): + return 99 + + +ASYNC_TASK = Task(func=_async_task, queue_name="default") + + +def _task_result(task, *args, **kwargs) -> TaskResult: + """Build a READY `TaskResult` without touching Redis.""" + return TaskResult( + task=task, + id=str(uuid.uuid4()), + status=TaskResultStatus.READY, + enqueued_at=timezone.now(), + started_at=None, + finished_at=None, + last_attempted_at=None, + args=list(args), + kwargs=dict(kwargs), + backend="default", + errors=[], + worker_ids=[], + ) + + +def _make_worker(*, max_tasks: int | None = None) -> WorkerProcess: + """Build an unstarted `WorkerProcess`.""" + return WorkerProcess( + thread_count=1, + max_tasks=max_tasks, + backend_alias="default", + queues=("default",), + ) + + +class TestTaskExecutor: + """Tests for the TaskExecutor dataclass and its methods.""" + + def test_post_init__sets_process_count_from_workers(self): + """__post_init__ uses explicit workers value.""" + executor = TaskExecutor( + backend=default_task_backend, workers=3, queues=("default",) + ) + assert executor.process_count == 3 + assert executor.thread_count == 1 + + def test_post_init__defaults_process_count_to_cpu_minus_one(self): + """__post_init__ defaults to cpu_count - 1 when workers is None.""" + executor = TaskExecutor(backend=default_task_backend, queues=("default",)) + expected = max(multiprocessing.cpu_count() - 1, 1) + assert executor.process_count == expected + + def test_post_init__thread_count_at_least_one(self): + """__post_init__ ensures thread_count is at least 1.""" + executor = TaskExecutor( + backend=default_task_backend, threads=0, queues=("default",) + ) + assert executor.thread_count == 1 + + def test_get_maximum_tasks_per_child__returns_none_when_max_tasks_is_zero(self): + """get_maximum_tasks_per_child returns None when max_tasks is 0.""" + executor = TaskExecutor( + backend=default_task_backend, max_tasks=0, queues=("default",) + ) + assert executor.get_maximum_tasks_per_child() is None + + def test_get_maximum_tasks_per_child__returns_value_when_set(self): + """get_maximum_tasks_per_child returns max_tasks // thread_count with jitter.""" + executor = TaskExecutor( + backend=default_task_backend, + max_tasks=100, + max_tasks_jitter=0, + threads=4, + queues=("default",), + ) + assert executor.get_maximum_tasks_per_child() == 25 # 100 // 4 + + def test_get_maximum_tasks_per_child__applies_jitter(self): + """get_maximum_tasks_per_child adds random jitter to max_tasks.""" + executor = TaskExecutor( + backend=default_task_backend, + max_tasks=100, + max_tasks_jitter=10, + threads=1, + queues=("default",), + ) + result = executor.get_maximum_tasks_per_child() + assert 100 <= result <= 110 # (100 + randint(0, 10)) // 1 + + def test_create_worker_process__starts_worker(self): + """create_worker_process creates and starts a WorkerProcess.""" + executor = TaskExecutor(backend=default_task_backend, queues=("default",)) + worker = executor.create_worker_process() + assert worker.is_alive() + worker.shutdown() + + def test_run__processes_enqueued_tasks_end_to_end(self): + """run() acquires, executes, and acknowledges tasks through to Redis.""" + count = 3 + enqueued = [default_task_backend.enqueue(echo, args=[i]) for i in range(count)] + executor = TaskExecutor( + backend=default_task_backend, + workers=1, + threads=2, + queues=("default",), + ) + + run_thread = threading.Thread(target=executor.run, daemon=True) + run_thread.start() + time.sleep(2) + executor.shutdown() + run_thread.join(timeout=5) + assert not run_thread.is_alive() + + results = list( + default_task_backend.peek( + "default", status=TaskResultStatus.SUCCESSFUL, count=count + ) + ) + assert {r.id for r in results} == {r.id for r in enqueued} + assert all(r.status == TaskResultStatus.SUCCESSFUL for r in results) + + def test_worker_acquires_updates_and_acknowledges(self): + """Worker acquires, executes, and acknowledges via its own backend.""" + enqueued = default_task_backend.enqueue(echo, args=[42]) + + 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() + + persisted = default_task_backend.get_result(enqueued.id) + assert persisted.status == TaskResultStatus.SUCCESSFUL + + def test_shutdown__stops_publishing(self): + """Shutdown stops publishing.""" + executor = TaskExecutor(backend=default_task_backend, queues=("default",)) + executor.shutdown() + assert not executor.is_publishing + + def test_shutdown__shuts_down_broker(self): + """Shutdown calls broker.shutdown when a broker is set.""" + executor = TaskExecutor(backend=default_task_backend, queues=("default",)) + executor.broker = Broker(default_task_backend) + executor.shutdown() + assert executor.broker.shutdown_requested.is_set() + + def test_shutdown__shuts_down_worker_processes(self): + """Shutdown calls shutdown on all worker processes.""" + executor = TaskExecutor(backend=default_task_backend, queues=("default",)) + worker = executor.create_worker_process() + executor.worker_processes = [worker] + executor.shutdown() + assert not worker.is_alive() + + def test_maintain_worker_pool__restarts_dead_workers(self): + """maintain_worker_pool replaces dead workers with new ones.""" + executor = TaskExecutor( + backend=default_task_backend, workers=1, threads=1, queues=("default",) + ) + worker = executor.create_worker_process() + executor.worker_processes = [worker] + worker.shutdown() + assert not worker.is_alive() + + maintain_thread = threading.Thread( + target=executor.maintain_worker_pool, daemon=True + ) + maintain_thread.start() + time.sleep(0.1) + executor.is_publishing = False + maintain_thread.join(timeout=2) + + assert executor.worker_processes[0] is not worker + assert executor.worker_processes[0].is_alive() + executor.worker_processes[0].shutdown() + + +class TestWorkerProcess: + """Tests for the WorkerProcess class.""" + + def test_record_task__increments_count(self): + """record_task increments task_count.""" + worker = _make_worker(max_tasks=5) + worker.lock = threading.Lock() + worker.expired = threading.Event() + worker.record_task() + assert worker.task_count == 1 + + def test_record_task__sets_expired_when_max_reached(self): + """record_task sets expired event when max_tasks is reached.""" + worker = _make_worker(max_tasks=1) + worker.lock = threading.Lock() + worker.expired = threading.Event() + worker.record_task() + assert worker.expired.is_set() + + def test_record_task__noop_when_max_tasks_is_none(self): + """record_task is a no-op when max_tasks is None.""" + worker = _make_worker(max_tasks=None) + worker.lock = threading.Lock() + worker.expired = threading.Event() + worker.record_task() + assert worker.task_count == 0 + assert not worker.expired.is_set() + + def test_record_task__noop_before_run_sets_lock_and_expired(self): + """record_task is a safe no-op before run() initializes lock/expired.""" + worker = _make_worker(max_tasks=5) + worker.record_task() + assert worker.task_count == 0 + + def test_shutdown_requested__is_settable(self): + """shutdown_requested event can be set on an unstarted worker.""" + worker = _make_worker() + worker.shutdown_requested.set() + assert worker.shutdown_requested.is_set() + + +class TestWorkerThread: + """Tests for the WorkerThread class.""" + + def test_execute_task_result__successful_execution(self): + """execute_task_result runs a task and returns SUCCESSFUL result.""" + result = WorkerThread( + worker=_make_worker(), index=0, backend=default_task_backend + ).execute_task_result(_task_result(echo, 42)) + assert result.status == TaskResultStatus.SUCCESSFUL + assert result.started_at is not None + assert result.finished_at is not None + assert result._return_value == 42 + + def test_execute_task_result__failed_execution(self): + """execute_task_result returns FAILED result when task raises.""" + result = WorkerThread( + worker=_make_worker(), index=0, backend=default_task_backend + ).execute_task_result(_task_result(boom)) + assert result.status == TaskResultStatus.FAILED + assert len(result.errors) == 1 + assert "ValueError" in result.errors[0].exception_class_path + + def test_execute_task_result__preserves_worker_ids(self): + """execute_task_result preserves worker_ids set by acquire.""" + thread = WorkerThread( + worker=_make_worker(), index=0, backend=default_task_backend + ) + task_result = _task_result(echo, 1) + task_result = dataclasses.replace(task_result, worker_ids=["pre-set-worker"]) + result = thread.execute_task_result(task_result) + assert result.worker_ids == ["pre-set-worker"] + + def test_call_task__calls_function_with_args(self): + """call_task invokes the task function with args and kwargs.""" + result = WorkerThread.call_task(_task_result(ADD_TASK, 1, y=2)) + assert result == 3 + + def test_call_task__passes_context_when_takes_context(self): + """call_task passes TaskContext when task.takes_context is True.""" + _context_captor.captured = None + result = WorkerThread.call_task(_task_result(CONTEXT_TASK)) + assert result == 42 + assert _context_captor.captured is not None + assert _context_captor.captured.task_result.task is CONTEXT_TASK + + def test_call_task__runs_async_function(self): + """call_task runs async task functions with asyncio.run.""" + result = WorkerThread.call_task(_task_result(ASYNC_TASK)) + assert result == 99 + + def test_create_task_error__builds_task_error(self): + """create_task_error builds a TaskError with exception info.""" + try: + raise RuntimeError("test error") + except RuntimeError: + error = WorkerThread.create_task_error(sys.exc_info()[1]) + assert "RuntimeError" in error.exception_class_path + assert "test error" in error.traceback diff --git a/tests/testapp/backends.py b/tests/testapp/backends.py deleted file mode 100644 index f237f90..0000000 --- a/tests/testapp/backends.py +++ /dev/null @@ -1,111 +0,0 @@ -from queue import Empty - -from django.tasks import TaskResult, TaskResultStatus -from django.utils import timezone -from django.utils.module_loading import import_string -from threadmill.backends import AcknowledgeableTaskBackend - - -class GeneratingTaskBackend(AcknowledgeableTaskBackend): - solved_task_count = 0 - issued_task_count = 0 - target_task_count = 100 - supports_async_task = True - - def __init__(self, alias, params): - super().__init__(alias=alias, params=params) - self._queues = None - - def reset(self, task_count=100): - GeneratingTaskBackend.solved_task_count = 0 - GeneratingTaskBackend.issued_task_count = 0 - self._queues = { - "default": [ - TaskResult( - task=import_string("tests.testapp.tasks.random_crash"), - enqueued_at=timezone.now(), - status=TaskResultStatus.READY, - id=f"default-{i + 1}", - args=[], - kwargs={}, - worker_ids=[], - started_at=None, - finished_at=None, - errors=[], - backend=self.alias, - last_attempted_at=None, - ) - for i in range(task_count) - ], - "compute": [ - TaskResult( - task=import_string("tests.testapp.tasks.compute_workload"), - enqueued_at=timezone.now(), - status=TaskResultStatus.READY, - id=f"compute-{i + 1}", - args=[], - kwargs={}, - worker_ids=[], - started_at=None, - finished_at=None, - errors=[], - backend=self.alias, - last_attempted_at=None, - ) - for i in range(task_count) - ], - "io": [ - TaskResult( - task=import_string("tests.testapp.tasks.io_workload"), - enqueued_at=timezone.now(), - status=TaskResultStatus.READY, - id=f"io-{i + 1}", - args=[], - kwargs={}, - worker_ids=[], - started_at=None, - finished_at=None, - errors=[], - backend=self.alias, - last_attempted_at=None, - ) - for i in range(task_count) - ], - "memory": [ - TaskResult( - task=import_string("tests.testapp.tasks.memory_workload"), - enqueued_at=timezone.now(), - status=TaskResultStatus.READY, - id=f"memory-{i + 1}", - args=[], - kwargs={}, - worker_ids=[], - started_at=None, - finished_at=None, - errors=[], - backend=self.alias, - last_attempted_at=None, - ) - for i in range(task_count) - ], - } - - def enqueue(self, task): - return task - - def acquire(self, *queue_names, timeout=None): - if self._queues is None: - self.reset() - GeneratingTaskBackend.issued_task_count += 1 - queues = [self._queues[queue_name] for queue_name in queue_names] - try: - # pop from the longest queue first to simulate a more realistic scenario - for queue in sorted(queues, key=len, reverse=True): - return queue.pop(0) - except IndexError as e: - GeneratingTaskBackend.issued_task_count -= 1 - raise Empty("No more tasks to solve.") from e - raise Empty("No more tasks to solve.") - - def acknowledge(self, task_result: TaskResult) -> None: - GeneratingTaskBackend.solved_task_count += 1 diff --git a/tests/testapp/settings.py b/tests/testapp/settings.py index 4884ec4..91740fa 100644 --- a/tests/testapp/settings.py +++ b/tests/testapp/settings.py @@ -10,6 +10,7 @@ https://docs.djangoproject.com/en/4.2/ref/settings/ """ +import os from pathlib import Path from django.tasks import DEFAULT_TASK_QUEUE_NAME @@ -53,7 +54,7 @@ "django.middleware.clickjacking.XFrameOptionsMiddleware", ] -ROOT_URLCONF = "testapp.urls" +ROOT_URLCONF = "tests.testapp.urls" TEMPLATES = [ { @@ -86,8 +87,12 @@ TASKS = { "default": { - "BACKEND": "tests.testapp.backends.GeneratingTaskBackend", + "BACKEND": "threadmill.backends.redis.RedisTaskBackend", "QUEUES": [DEFAULT_TASK_QUEUE_NAME, "compute", "io", "memory"], + "REDIS_URL": os.environ.get("REDIS_URL", "redis://localhost:6379/0"), + }, + "immediate": { + "BACKEND": "django.tasks.backends.immediate.ImmediateTaskBackend", }, } diff --git a/tests/testapp/tasks.py b/tests/testapp/tasks.py index 87d912c..f13db68 100644 --- a/tests/testapp/tasks.py +++ b/tests/testapp/tasks.py @@ -8,6 +8,18 @@ logger = logging.getLogger(__name__) +@task() +def echo(value): + """Return the given value (fast, deterministic, for tests).""" + return value + + +@task() +def boom(): + """Raise ValueError (deterministic failure, for tests).""" + raise ValueError("boom") + + @task(queue_name="compute") def compute_workload(): """Calculate the first 1000 prime numbers.""" diff --git a/threadmill/__init__.py b/threadmill/__init__.py index 5f32987..5d97ec7 100644 --- a/threadmill/__init__.py +++ b/threadmill/__init__.py @@ -1,4 +1,4 @@ -"""A queue agnostic worker for Django's task framework.""" +"""The most reliable backend for Django's task framework.""" from . import _version diff --git a/threadmill/backends.py b/threadmill/backends.py deleted file mode 100644 index cb0ab7b..0000000 --- a/threadmill/backends.py +++ /dev/null @@ -1,34 +0,0 @@ -from __future__ import annotations - -import datetime -from abc import ABC - -from django.tasks import TaskResult -from django.tasks.backends.base import BaseTaskBackend - - -class AcknowledgeableTaskBackend(BaseTaskBackend, ABC): - """Provide an interface for tasks queues to be processed by the executor.""" - - supports_async_task = True - supports_get_result = True - - def acquire( - self, *queue_names: str, timeout: datetime.timedelta | None = None - ) -> TaskResult: - """ - Return and lock the next task to be processed without removing it from the queue. - - Args: - queue_names: The names of the queues to acquire tasks from. - timeout: The maximum time to wait for a task. If None, wait indefinitely. - - Raises: - TimeoutError: If no task is available within the specified timeout. - queue.Empty: If no task is available and timeout is None. - """ - raise NotImplementedError - - def acknowledge(self, task_result: TaskResult) -> None: - """Remove the task from the queue and publish the result.""" - raise NotImplementedError diff --git a/threadmill/backends/__init__.py b/threadmill/backends/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/threadmill/backends/base.py b/threadmill/backends/base.py new file mode 100644 index 0000000..799988d --- /dev/null +++ b/threadmill/backends/base.py @@ -0,0 +1,145 @@ +from __future__ import annotations + +import collections.abc +import dataclasses +import datetime +import json +import threading +from abc import ABC + +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.utils.module_loading import import_string + + +class Broker(threading.Thread): + """Backend maintenance thread launched by the task executor.""" + + def __init__( + self, + backend: ThreadmillTaskBackend | None = None, + *, + interval: datetime.timedelta = datetime.timedelta(seconds=1), + ) -> None: + super().__init__(daemon=True) + self.backend = backend + self.interval = interval + self.shutdown_requested = threading.Event() + + def main(self) -> None: + """Perform one maintenance pass.""" + + def run(self) -> None: + while not self.shutdown_requested.wait(self.interval.total_seconds()): + self.main() + + def shutdown(self) -> None: + """Request graceful shutdown.""" + self.shutdown_requested.set() + + +def _parse_datetime(value: object) -> object: + """Parse an ISO datetime string, returning the value unchanged if not parseable.""" + if isinstance(value, str): + try: + return datetime.datetime.fromisoformat(value) + except ValueError: + return value + return value + + +class TaskResultEncoder(DjangoJSONEncoder): + """JSON encoder for TaskResult and TaskError objects.""" + + def default(self, o): + if isinstance(o, (TaskResult, TaskError)): + return { + field.name: getattr(o, field.name) + for field in dataclasses.fields(type(o)) + } + if isinstance(o, Task): + return { + 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} + return super().default(o) + + +class ThreadmillTaskBackend(BaseTaskBackend, ABC): + """Interface for task queues to be processed by the executor.""" + + supports_async_task = True + supports_get_result = True + broker_class: type[Broker] | None = None + + result_ttl: datetime.timedelta | None = None + + @staticmethod + def serialize_task_result(task_result: TaskResult) -> str: + return json.dumps(task_result, cls=TaskResultEncoder) + + @classmethod + 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 + }, + ) + d["status"] = TaskResultStatus(d["status"]) + d["errors"] = [TaskError(**error) for error in d["errors"]] + return_value = d.pop("_return_value", None) + for key, value in d.items(): + d[key] = _parse_datetime(value) + result = TaskResult(**d) + object.__setattr__(result, "_return_value", return_value) + return result + return d + + return json.loads(data, object_hook=_object_hook) + + def acquire( + self, + *queue_names: str, + timeout: datetime.timedelta | None = None, + worker: str = "", + ) -> TaskResult: + """ + Return and lock the next task to be processed without removing it from the queue. + + Args: + queue_names: The names of the queues to acquire tasks from. + timeout: The maximum time to wait for a task. If None, wait indefinitely. + worker: The name of the worker thread acquiring the task. + + Raises: + TimeoutError: If no task is available within the specified timeout. + queue.Empty: If no task is available and timeout is None. + """ + raise NotImplementedError + + def acknowledge(self, task_result: TaskResult) -> None: + """Remove the task from the queue and publish the result.""" + raise NotImplementedError + + def peek( + self, + queue_name: str = DEFAULT_TASK_QUEUE_NAME, + *, + status: TaskResultStatus | None = None, + count: int = 1, + ) -> collections.abc.Generator[TaskResult, None, None]: + """Yield acknowledged results from a queue, optionally filtered by status.""" + raise NotImplementedError diff --git a/threadmill/backends/lua/acknowledge.lua b/threadmill/backends/lua/acknowledge.lua new file mode 100644 index 0000000..27a6d40 --- /dev/null +++ b/threadmill/backends/lua/acknowledge.lua @@ -0,0 +1,24 @@ +-- Finalize a completed task: remove it from the running set, persist the +-- result with a TTL, delete the task data hash, and add the result to the +-- results history set. Also evict results whose finish score falls outside the +-- retention window. +-- +-- KEYS[1] -- running set (ZSET) +-- KEYS[2] -- result key (STRING, stores serialized TaskResult) +-- KEYS[3] -- task data key (HASH, deleted after acknowledge) +-- KEYS[4] -- results history set (ZSET, ordered by finish time) +-- ARGV[1] -- task ID +-- ARGV[2] -- serialized TaskResult JSON +-- ARGV[3] -- result TTL in seconds +-- ARGV[4] -- finish timestamp in milliseconds (score for the results set) +-- Returns: 1 on success, 0 if task was not in the running set + +local removed = redis.call('ZREM', KEYS[1], ARGV[1]) +if removed == 0 then + return 0 -- Task already reaped, skip +end +redis.call('SET', KEYS[2], ARGV[2], 'EX', ARGV[3]) +redis.call('DEL', KEYS[3]) +redis.call('ZADD', KEYS[4], ARGV[4], ARGV[1]) +redis.call('ZREMRANGEBYSCORE', KEYS[4], 0, tonumber(ARGV[4]) - tonumber(ARGV[3]) * 1000) +return 1 diff --git a/threadmill/backends/lua/acquire.lua b/threadmill/backends/lua/acquire.lua new file mode 100644 index 0000000..5b8523f --- /dev/null +++ b/threadmill/backends/lua/acquire.lua @@ -0,0 +1,44 @@ +-- Atomically pop the lowest-scored task from any of the given priority queues, +-- update its JSON data with worker info, and move it directly to the running +-- set. Iterates queues in key order and returns the first available task. +-- +-- KEYS[1..N] -- interleaved running keys and queue keys, one pair per queue: +-- KEYS[1] = running set, KEYS[2] = queue set, KEYS[3] = running, +-- KEYS[4] = queue, etc. +-- ARGV[1] -- current time in milliseconds +-- ARGV[2] -- current time as ISO-8601 string +-- ARGV[3] -- task key prefix (e.g. "threadmill:default:task:") +-- ARGV[4] -- number of queue pairs (N/2) +-- ARGV[5] -- worker name +-- ARGV[6] -- lease TTL in milliseconds +-- Returns: updated serialized data on success, nil if all queues are empty. + +local num_queues = tonumber(ARGV[4]) +local lease_ttl_ms = tonumber(ARGV[6]) +for i = 1, num_queues do + local result = redis.call('ZPOPMIN', KEYS[i * 2]) + if #result > 0 then + local task_id = result[1] + local data = redis.call('HGET', ARGV[3] .. task_id, 'data') + if data then + local ok, parsed = pcall(cjson.decode, data) + if ok then + parsed.status = 'RUNNING' + parsed.last_attempted_at = ARGV[2] + if not parsed.started_at then + parsed.started_at = ARGV[2] + end + if not parsed.worker_ids then + parsed.worker_ids = {} + end + table.insert(parsed.worker_ids, ARGV[5]) + local updated_data = cjson.encode(parsed) + local deadline = tonumber(ARGV[1]) + lease_ttl_ms + redis.call('ZADD', KEYS[i * 2 - 1], deadline, task_id) + redis.call('HSET', ARGV[3] .. task_id, 'data', updated_data) + return updated_data + end + end + end +end +return nil diff --git a/threadmill/backends/lua/mover.lua b/threadmill/backends/lua/mover.lua new file mode 100644 index 0000000..b1134f2 --- /dev/null +++ b/threadmill/backends/lua/mover.lua @@ -0,0 +1,19 @@ +-- Move tasks whose scheduled time has passed from the deferred set to the active +-- priority queue. Only processes up to a batch limit per call. +-- +-- KEYS[1] -- deferred set (ZSET, scored by run_after timestamp) +-- KEYS[2] -- active priority queue (ZSET, scored by priority+time) +-- ARGV[1] -- current time in milliseconds (all scores <= this are due) +-- ARGV[2] -- task key prefix (e.g. "threadmill:default:task:") +-- ARGV[3] -- maximum number of tasks to move per call (batch size) +-- Returns: number of tasks moved + +local due = redis.call('ZRANGEBYSCORE', KEYS[1], 0, ARGV[1], 'LIMIT', 0, tonumber(ARGV[3])) +for _, task_id in ipairs(due) do + redis.call('ZREM', KEYS[1], task_id) + local score = redis.call('HGET', ARGV[2] .. task_id, 'score') + if score then + redis.call('ZADD', KEYS[2], score, task_id) + end +end +return #due diff --git a/threadmill/backends/lua/reaper.lua b/threadmill/backends/lua/reaper.lua new file mode 100644 index 0000000..278744a --- /dev/null +++ b/threadmill/backends/lua/reaper.lua @@ -0,0 +1,36 @@ +-- Fail tasks whose processing lease has expired. +-- +-- KEYS[1] -- running set (ZSET) +-- KEYS[2] -- results history set (ZSET) +-- ARGV[1] -- current time in milliseconds (for score comparison) +-- ARGV[2] -- task key prefix (e.g. "threadmill:default:task:") +-- ARGV[3] -- result key prefix (e.g. "threadmill:default:result:") +-- ARGV[4] -- batch size +-- ARGV[5] -- result TTL in seconds +-- ARGV[6] -- finished_at as ISO format string +-- Returns: number of tasks failed + +local stale = redis.call('ZRANGEBYSCORE', KEYS[1], 0, ARGV[1], 'LIMIT', 0, tonumber(ARGV[4])) +for _, task_id in ipairs(stale) do + local data = redis.call('HGET', ARGV[2] .. task_id, 'data') + if data then + local ok, parsed = pcall(cjson.decode, data) + if ok then + parsed.status = 'FAILED' + parsed.finished_at = ARGV[6] + if not parsed.errors then + parsed.errors = {} + end + table.insert(parsed.errors, { + exception_class_path = 'threadmill.exceptions.AcknowledgementTimeout', + traceback = 'Task processing lease expired.' + }) + local failed_data = cjson.encode(parsed) + redis.call('ZREM', KEYS[1], task_id) + redis.call('SET', ARGV[3] .. task_id, failed_data, 'EX', ARGV[5]) + redis.call('DEL', ARGV[2] .. task_id) + redis.call('ZADD', KEYS[2], tonumber(ARGV[1]), task_id) + end + end +end +return #stale diff --git a/threadmill/backends/redis.py b/threadmill/backends/redis.py new file mode 100644 index 0000000..d8b46fe --- /dev/null +++ b/threadmill/backends/redis.py @@ -0,0 +1,324 @@ +"""Redis-backed durable priority queue backend for Django's task framework.""" + +from __future__ import annotations + +import datetime +import logging +import queue +import time +import uuid +from collections.abc import Generator, Sequence +from pathlib import Path + +import redis +from django.tasks import DEFAULT_TASK_QUEUE_NAME, TaskResult, TaskResultStatus +from django.tasks.exceptions import TaskResultDoesNotExist +from django.tasks.signals import task_enqueued +from django.utils import timezone + +from threadmill.backends.base import Broker, ThreadmillTaskBackend + +logger = logging.getLogger(__name__) + +_LUA_DIR = Path(__file__).resolve().parent / "lua" + + +def _load_lua(name: str) -> str: + """Load a Lua script from the lua directory.""" + return (_LUA_DIR / f"{name}.lua").read_text() + + +class RedisBroker(Broker): + """Background maintenance broker for the Redis backend.""" + + backend: RedisTaskBackend + + MOVER_SCRIPT = _load_lua("mover") + """Move tasks whose scheduled time has passed from the deferred to the active queue.""" + REAPER_SCRIPT = _load_lua("reaper") + """Fail tasks whose processing lease has expired from the running set.""" + + def __init__(self, backend: RedisTaskBackend) -> None: + interval = backend.options.get("broker_interval", datetime.timedelta(seconds=1)) + super().__init__(backend, interval=interval) + self._mover_script = self.backend.client.register_script(self.MOVER_SCRIPT) + self._reaper_script = self.backend.client.register_script(self.REAPER_SCRIPT) + + def _move_queue(self, queue_name: str) -> None: + """Move due deferred tasks from a single deferred set.""" + deferred_key = self.backend.DEFERRED_KEY.format( + prefix=self.backend.key_prefix, queue_name=queue_name + ) + queue_key = self.backend.QUEUE_KEY.format( + prefix=self.backend.key_prefix, queue_name=queue_name + ) + self._mover_script( + keys=[deferred_key, queue_key], + args=[ + str(time.time() * 1000), + self.backend.key_prefix + ":task:", + str(self.backend.batch_size), + ], + ) + + def _reap_running_queue(self, queue_name: str) -> None: + """Fail tasks whose processing lease has expired from the running set.""" + now = timezone.now() + now_ms = now.timestamp() * 1000 + finished_at_iso = now.isoformat() + running_key = self.backend.RUNNING_KEY.format( + prefix=self.backend.key_prefix, queue_name=queue_name + ) + results_key = self.backend.RESULTS_KEY.format( + prefix=self.backend.key_prefix, queue_name=queue_name + ) + self._reaper_script( + keys=[running_key, results_key], + args=[ + str(now_ms), + self.backend.key_prefix + ":task:", + self.backend.key_prefix + ":result:", + str(self.backend.batch_size), + str(int(self.backend.result_ttl.total_seconds())), + finished_at_iso, + ], + ) + + def main(self) -> None: + """Run mover and running reaper passes for all queues.""" + for queue_name in self.backend.queues: + try: + self._move_queue(queue_name) + except Exception: # noqa: BLE001 + logger.exception("Mover error for queue %r", queue_name) + + try: + self._reap_running_queue(queue_name) + except Exception: # noqa: BLE001 + logger.exception("Running reaper error for queue %r", queue_name) + + +class RedisTaskBackend(ThreadmillTaskBackend): + """Redis-backed durable priority queue backend. + + Uses sorted sets for priority ordering, a running set for in-flight + tracking, and a deferred set for scheduled tasks. All multi-step operations + are atomic via Lua scripts. + """ + + supports_async_task = True + supports_get_result = True + supports_priority = True + supports_defer = True + + broker_class = RedisBroker + + QUEUE_KEY = "{prefix}:queue:{queue_name}" + RUNNING_KEY = "{prefix}:running:{queue_name}" + DEFERRED_KEY = "{prefix}:deferred:{queue_name}" + TASK_KEY = "{prefix}:task:{task_id}" + RESULT_KEY = "{prefix}:result:{result_id}" + RESULTS_KEY = "{prefix}:results:{queue_name}" + + ACQUIRE_SCRIPT = _load_lua("acquire") + """Pop the next task from a priority queue and move it directly to the running set.""" + ACKNOWLEDGE_SCRIPT = _load_lua("acknowledge") + """Remove from running, persist the result, and clean up.""" + + def __init__(self, alias: str, params: dict) -> None: + super().__init__(alias=alias, params=params) + + try: + redis_url = params["REDIS_URL"] + except KeyError as e: + raise ValueError( + f"REDIS_URL must be specified in your settings for the {type(self).__name__}." + ) from e + self.client = redis.from_url(redis_url) + self.key_prefix = f"threadmill:{{{alias}}}" + self.lease_ttl = self.options.get("lease_ttl", datetime.timedelta(hours=1)) + self.result_ttl = self.options.get("result_ttl", datetime.timedelta(days=1)) + self.batch_size = self.options.get("batch_size", 100) + self._acquire_script = self.client.register_script(self.ACQUIRE_SCRIPT) + self._acknowledge_script = self.client.register_script(self.ACKNOWLEDGE_SCRIPT) + + def _compute_score(self, priority: int, enqueued_at: datetime.datetime) -> float: + """Compute a ZSET score for priority-ordered FIFO queueing. + + Higher priority (more positive) tasks are popped first. Within the same + priority, earlier enqueued tasks are popped first. + """ + enqueued_at_ms = enqueued_at.timestamp() * 1e3 + return -priority * 1e13 + enqueued_at_ms + + def enqueue( + self, + task, + args: Sequence | None = None, + kwargs: dict | None = None, + ) -> TaskResult: + """Enqueue a task for execution. + + If the task has a run_after datetime, it is stored in the deferred set + instead of the active priority queue. + """ + self.validate_task(task) + + enqueued_at = timezone.now() + task_result = TaskResult( + task=task, + id=str(uuid.uuid4()), + status=TaskResultStatus.READY, + enqueued_at=enqueued_at, + started_at=None, + finished_at=None, + last_attempted_at=None, + args=list(args or []), + kwargs=dict(kwargs or {}), + backend=self.alias, + errors=[], + worker_ids=[], + ) + + score = self._compute_score(task.priority, enqueued_at) + serialized = self.serialize_task_result(task_result) + task_key = self.TASK_KEY.format(prefix=self.key_prefix, task_id=task_result.id) + task_data_ttl = int( + self.lease_ttl.total_seconds() * 3 + self.result_ttl.total_seconds() + ) + + pipe = self.client.pipeline() + pipe.hset( + task_key, + mapping={ + "data": serialized, + "score": str(score), + "queue_name": task.queue_name, + }, + ) + pipe.expire(task_key, task_data_ttl) + + if task.run_after is not None: + deferred_key = self.DEFERRED_KEY.format( + prefix=self.key_prefix, queue_name=task.queue_name + ) + run_after_ms = task.run_after.timestamp() * 1000 + pipe.zadd(deferred_key, {task_result.id: run_after_ms}) + else: + queue_key = self.QUEUE_KEY.format( + prefix=self.key_prefix, queue_name=task.queue_name + ) + pipe.zadd(queue_key, {task_result.id: score}) + + pipe.execute() + + task_enqueued.send(self.__class__, task_result=task_result) + return task_result + + def acquire( + self, + *queue_names: str, + timeout: datetime.timedelta | None = None, + worker: str = "", + ) -> TaskResult: + queue_names = queue_names or tuple(self.queues) + deadline = time.monotonic() + timeout.total_seconds() if timeout else None + keys = [ + key + for queue_name in queue_names + for key in ( + self.RUNNING_KEY.format(prefix=self.key_prefix, queue_name=queue_name), + self.QUEUE_KEY.format(prefix=self.key_prefix, queue_name=queue_name), + ) + ] + + while True: + now = timezone.now() + now_ms = now.timestamp() * 1000 + now_iso = now.isoformat() + + if data := self._acquire_script( + keys=keys, + args=[ + str(now_ms), + now_iso, + self.key_prefix + ":task:", + str(len(queue_names)), + worker, + str(int(self.lease_ttl.total_seconds() * 1000)), + ], + ): + return self.deserialize_task_result(data) + + try: + if deadline - time.monotonic() <= 0: + raise TimeoutError( + "No task available within the specified timeout." + ) + except TypeError: + raise queue.Empty("No task available.") + else: + time.sleep(0.01) + + def acknowledge(self, task_result: TaskResult) -> 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 + ) + result_key = self.RESULT_KEY.format( + prefix=self.key_prefix, result_id=task_result.id + ) + task_key = self.TASK_KEY.format(prefix=self.key_prefix, task_id=task_result.id) + results_key = self.RESULTS_KEY.format( + prefix=self.key_prefix, queue_name=task_result.task.queue_name + ) + finished_at = task_result.finished_at or timezone.now() + finish_score = finished_at.timestamp() * 1000 + + self._acknowledge_script( + keys=[running_key, result_key, task_key, results_key], + args=[ + task_result.id, + serialized, + str(int(self.result_ttl.total_seconds())), + str(finish_score), + ], + ) + + def peek( + self, + queue_name: str = DEFAULT_TASK_QUEUE_NAME, + *, + status: TaskResultStatus | None = None, + count: int = 1, + ) -> Generator[TaskResult]: + result_ids = [ + rid.decode() if isinstance(rid, bytes) else rid + for rid in self.client.zrange( + self.RESULTS_KEY.format(prefix=self.key_prefix, queue_name=queue_name), + 0, + count - 1, + ) + ] + pipe = self.client.pipeline() + for result_id in result_ids: + pipe.get( + self.RESULT_KEY.format(prefix=self.key_prefix, result_id=result_id) + ) + for _result_id, data in zip(result_ids, pipe.execute(), strict=False): + if data: + data_str = data.decode() if isinstance(data, bytes) else data + task_result = self.deserialize_task_result(data_str) + if status is None or task_result.status == status: + yield task_result + + def get_result(self, result_id: str) -> TaskResult: + if data := self.client.get( + self.RESULT_KEY.format(prefix=self.key_prefix, result_id=result_id) + ): + return self.deserialize_task_result(data) + raise TaskResultDoesNotExist(f"Task result {result_id!r} does not exist.") + + def close(self) -> None: + """Close the Redis connection.""" + self.client.close() diff --git a/threadmill/exceptions.py b/threadmill/exceptions.py new file mode 100644 index 0000000..acc4c54 --- /dev/null +++ b/threadmill/exceptions.py @@ -0,0 +1,7 @@ +"""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 a0a3510..3531e22 100644 --- a/threadmill/executor.py +++ b/threadmill/executor.py @@ -13,21 +13,18 @@ import time import typing from concurrent.futures import ThreadPoolExecutor -from contextlib import suppress from inspect import iscoroutinefunction -from multiprocessing.queues import JoinableQueue from queue import Empty from traceback import format_exception -from django.tasks import TaskResult +from django.tasks import TaskResult, task_backends from django.tasks.base import TaskContext, TaskError, TaskResultStatus from django.tasks.signals import task_finished, task_started from django.utils import timezone from django.utils.json import normalize_json if typing.TYPE_CHECKING: - from .backends import AcknowledgeableTaskBackend - + from .backends.base import Broker, ThreadmillTaskBackend logger = multiprocessing.get_logger() formatter = logging.Formatter( @@ -41,15 +38,13 @@ @dataclasses.dataclass(kw_only=True, slots=True) class TaskExecutor: - """Consume tasks from shared joinable queues with process and thread pools.""" + """Tasks consumed from shared joinable queues via process and thread pools.""" - backend: AcknowledgeableTaskBackend + backend: ThreadmillTaskBackend workers: int | None = None threads: int = 1 max_tasks: int = 0 max_tasks_jitter: int = 0 - task_timeout: datetime.timedelta = datetime.timedelta(hours=1) - is_acquiring: bool = dataclasses.field(default=True, init=False) is_publishing: bool = dataclasses.field(default=True, init=False) worker_processes: list[WorkerProcess] = dataclasses.field( default_factory=list, init=False @@ -57,22 +52,13 @@ class TaskExecutor: process_count: int = dataclasses.field(init=False) thread_count: int = dataclasses.field(init=False) queues: tuple[str] - shared_task_queue: multiprocessing.JoinableQueue[TaskResult] = dataclasses.field( - init=False - ) - processed_task_queue: multiprocessing.JoinableQueue[TaskResult] = dataclasses.field( - init=False - ) + broker: Broker | None = dataclasses.field(default=None, init=False) exit_empty: bool = False def __post_init__(self) -> None: """Initialize derived orchestration fields and queues.""" self.process_count = self.workers or max(multiprocessing.cpu_count() - 1, 1) self.thread_count = max(self.threads, 1) - self.shared_task_queue = multiprocessing.JoinableQueue( - maxsize=self.process_count * self.thread_count, - ) - self.processed_task_queue = multiprocessing.JoinableQueue() def get_maximum_tasks_per_child(self) -> int | None: """Return worker recycling limit based on config and thread count.""" @@ -84,11 +70,11 @@ def get_maximum_tasks_per_child(self) -> int | None: def create_worker_process(self) -> WorkerProcess: """Create and start a new worker process.""" worker = WorkerProcess( - self.shared_task_queue, - self.processed_task_queue, self.thread_count, - self.task_timeout, self.get_maximum_tasks_per_child(), + self.backend.alias, + self.queues, + self.exit_empty, ) worker.start() return worker @@ -99,48 +85,22 @@ def run(self) -> None: self.create_worker_process() for _ in range(self.process_count) ] threads = [ - threading.Thread(target=self.acknowledge_tasks, daemon=True), threading.Thread(target=self.maintain_worker_pool, daemon=True), - threading.Thread(target=self.acquire_tasks, daemon=True), ] + if self.backend.broker_class: + self.broker = self.backend.broker_class(self.backend) + threads.append(self.broker) + for thread in threads: thread.start() for thread in threads: thread.join() - def acquire_tasks(self) -> None: - """Buffer tasks in shared task queue.""" - while self.is_acquiring: - try: - work = self.backend.acquire(*self.queues) - except Empty: - if self.exit_empty: - logger.info("No more tasks to solve. Shutting down.") - self.shutdown() - return - time.sleep(0.01) - else: - self.shared_task_queue.put(work) - - def acknowledge_tasks(self) -> None: - """Acknowledge processed tasks and publish updated results in main process.""" - while self.is_publishing: - try: - task = self.processed_task_queue.get_nowait() - except Empty: - time.sleep(0.01) - else: - self.backend.acknowledge(task) - self.processed_task_queue.task_done() - def shutdown(self) -> None: """Stop queue consumption and terminate all worker processes.""" logger.info("Shutting down task executor") - self.is_acquiring = False - with suppress(ValueError): - self.shared_task_queue.join() - with suppress(ValueError): - self.processed_task_queue.join() + if self.broker is not None: + self.broker.shutdown() with ThreadPoolExecutor(max_workers=self.process_count) as executor: executor.map(lambda worker: worker.shutdown(), self.worker_processes) self.is_publishing = False @@ -148,11 +108,18 @@ def shutdown(self) -> None: def maintain_worker_pool(self) -> None: """Restart worker processes that have exited.""" while self.is_publishing: + all_dead = True for index, worker in enumerate(self.worker_processes): if worker.is_alive(): + all_dead = False continue worker.join(timeout=0) + if self.exit_empty: + continue self.worker_processes[index] = self.create_worker_process() + if all_dead and self.exit_empty: + self.shutdown() + return time.sleep(1) @@ -161,20 +128,20 @@ class WorkerProcess(multiprocessing.Process): def __init__( self, - task_queue: JoinableQueue[TaskResult], - processed_task_queue: JoinableQueue[TaskResult], thread_count: int, - task_timeout: datetime.timedelta, max_tasks: int | None = None, + backend_alias: str = "", + queues: tuple[str, ...] = (), + exit_empty: bool = False, ) -> None: """Create process with dedicated thread pool for task execution.""" self.shutdown_requested = multiprocessing.Event() super().__init__(daemon=True) - self.task_queue = task_queue - self.processed_task_queue = processed_task_queue self.thread_count = thread_count - self.task_timeout = task_timeout self.max_tasks = max_tasks + self.backend_alias = backend_alias + self.queues = queues + self.exit_empty = exit_empty self.task_count = 0 self.lock: threading.Lock | None = None self.expired: threading.Event | None = None @@ -184,13 +151,17 @@ def run(self) -> None: logger.info("Starting worker process %s", self.name) self.lock = threading.Lock() self.expired = threading.Event() + backend = task_backends[self.backend_alias] consumer_threads = [ - WorkerThread(worker=self, index=index) for index in range(self.thread_count) + WorkerThread(worker=self, index=index, backend=backend) + for index in range(self.thread_count) ] for consumer_thread in consumer_threads: consumer_thread.start() for consumer_thread in consumer_threads: - consumer_thread.join(self.task_timeout.total_seconds()) + consumer_thread.join( + backend.result_ttl.total_seconds() if backend.result_ttl else None + ) def record_task(self) -> None: """Record one processed task and stop when max_tasks is reached.""" @@ -218,28 +189,31 @@ def __init__( *, worker: WorkerProcess, index: int, + backend: ThreadmillTaskBackend, ) -> None: """Create worker thread bound to process worker state.""" super().__init__(name=f"{socket.gethostname()}:{worker.pid}-{index}") self.worker = worker + self.backend = backend def run(self) -> None: """Start consuming tasks for this thread.""" while self.worker.expired is None or not self.worker.expired.is_set(): try: - task_result = self.worker.task_queue.get(timeout=1.0) - except Empty: - if self.worker.shutdown_requested.is_set(): + task_result = self.backend.acquire( + *self.worker.queues, + timeout=datetime.timedelta(seconds=1), + worker=self.name, + ) + except (Empty, TimeoutError): + if self.worker.shutdown_requested.is_set() or self.worker.exit_empty: return continue + try: - self.worker.processed_task_queue.put( - self.execute_task_result( - task_result, - ) - ) + result = self.execute_task_result(task_result) + self.backend.acknowledge(result) finally: - self.worker.task_queue.task_done() self.worker.record_task() def execute_task_result(self, task_result: TaskResult) -> TaskResult: @@ -249,9 +223,8 @@ def execute_task_result(self, task_result: TaskResult) -> TaskResult: task_result = dataclasses.replace( task_result, status=TaskResultStatus.RUNNING, - started_at=started_at, + started_at=task_result.started_at or started_at, last_attempted_at=started_at, - worker_ids=[*task_result.worker_ids, self.name], ) task_started.send(TaskExecutor, task_result=task_result) @@ -262,22 +235,20 @@ def execute_task_result(self, task_result: TaskResult) -> TaskResult: task_result, status=TaskResultStatus.FAILED, errors=[*task_result.errors, WorkerThread.create_task_error(exception)], + finished_at=timezone.now(), ) logger.exception("Task failed %r", task_result.id) else: task_result = dataclasses.replace( task_result, status=TaskResultStatus.SUCCESSFUL, + finished_at=timezone.now(), ) object.__setattr__( task_result, "_return_value", normalize_json(return_value) ) logger.info("Task successful %r", task_result.id) finally: - task_result = dataclasses.replace( - task_result, - finished_at=timezone.now(), - ) task_finished.send(TaskExecutor, task_result=task_result) return task_result diff --git a/threadmill/management/commands/threadmill.py b/threadmill/management/commands/threadmill.py index 0fa7249..1cd8da8 100644 --- a/threadmill/management/commands/threadmill.py +++ b/threadmill/management/commands/threadmill.py @@ -1,4 +1,3 @@ -import datetime import signal import sys @@ -66,12 +65,6 @@ def add_arguments(self, parser): default=0, help="Maximum random jitter to add to the max-tasks value by randint(0, max_tasks_jitter).", ) - parser.add_argument( - "--task-timeout", - type=float, - default=3600.0, - help="Kill hung tasks after timeout seconds. Defaults to one hour.", - ) parser.add_argument( "--task-backlog-size", type=int, @@ -94,7 +87,6 @@ def handle( threads, max_tasks, max_tasks_jitter, - task_timeout, exit_empty, **options, ): @@ -120,7 +112,6 @@ def handle( threads=threads, max_tasks=max_tasks, max_tasks_jitter=max_tasks_jitter, - task_timeout=datetime.timedelta(seconds=task_timeout), exit_empty=exit_empty, queues=queues, )