diff --git a/README.md b/README.md index 87054fe..47f6f50 100644 --- a/README.md +++ b/README.md @@ -114,12 +114,14 @@ uv run manage.py threadmill inspector The `RedisTaskBackend` accepts the following options under `OPTIONS` in your `TASKS` configuration: -| 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. | +| 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. | +| `poll_interval` | `timedelta(seconds=0.01)` | Base wait between idle acquire attempts, doubled after each empty poll. | +| `poll_max_interval` | `timedelta(seconds=1)` | Max wait between idle acquire attempts. | A task that is started but never acknowledged (lease expired) is marked FAILED with an `AcknowledgementTimeout` error. Set `lease_ttl` comfortably above your diff --git a/tests/backends/test_redis.py b/tests/backends/test_redis.py index 76c2fbf..587e408 100644 --- a/tests/backends/test_redis.py +++ b/tests/backends/test_redis.py @@ -1,10 +1,14 @@ +import collections.abc import dataclasses import datetime import logging +import queue import time +import typing from dataclasses import replace from unittest.mock import patch +import pytest from django.tasks import default_task_backend from django.tasks.base import TaskResultStatus from django.utils import timezone @@ -38,6 +42,39 @@ def _stats(**overrides: int | datetime.timedelta) -> QueueStats: return QueueStats(counts=counts, rates=rates) +class CountingAcquireScript: + """Delegate to the real acquire script while recording each invocation.""" + + def __init__(self, script: collections.abc.Callable[..., typing.Any]) -> None: + self.script: collections.abc.Callable[..., typing.Any] = script + self.calls: list[float] = [] + + def __call__(self, **kwargs: typing.Any) -> typing.Any: + self.calls.append(time.monotonic()) + return self.script(**kwargs) + + +def _make_backend(alias: str, **options: datetime.timedelta) -> RedisTaskBackend: + """Build a single-queue backend with per-test options.""" + return RedisTaskBackend( + alias, + { + "QUEUES": ["default"], + "REDIS_URL": "redis://localhost:6379/0", + "OPTIONS": { + "lease_ttl": datetime.timedelta(hours=1), + "result_ttl": datetime.timedelta(seconds=60), + **options, + }, + }, + ) + + +def _measure_wait_deltas(calls: list[float]) -> list[float]: + """Return the seconds elapsed between consecutive recorded script calls.""" + return [calls[index + 1] - calls[index] for index in range(len(calls) - 1)] + + class TestRedisBroker: def test_mover__moves_deferred_task_to_ready(self): """Mover promotes due deferred tasks to the ready queue.""" @@ -973,3 +1010,97 @@ def test_purge__empty_queue_is_noop(self) -> None: ) finally: backend.close() + + def test_acquire__backs_off_when_idle(self): + """Idle waits grow so a one-second acquire polls far less than every 10ms.""" + backend = _make_backend("acquire_backoff_test") + backend._acquire_script = CountingAcquireScript(backend._acquire_script) + try: + started_at = time.monotonic() + with pytest.raises(TimeoutError): + backend.acquire(timeout=datetime.timedelta(seconds=1)) + elapsed_secs = time.monotonic() - started_at + poll_count = len(backend._acquire_script.calls) + assert elapsed_secs >= 0.99 + assert 1 < poll_count <= 20 + finally: + backend.close() + + def test_acquire__doubles_wait_up_to_poll_max_interval(self): + """Consecutive empty polls double their wait, capped at poll_max_interval.""" + poll_interval_secs = 0.05 + poll_max_secs = 0.2 + backend = _make_backend( + "acquire_double_test", + poll_interval=datetime.timedelta(seconds=poll_interval_secs), + poll_max_interval=datetime.timedelta(seconds=poll_max_secs), + ) + backend._acquire_script = CountingAcquireScript(backend._acquire_script) + try: + with pytest.raises(TimeoutError): + backend.acquire(timeout=datetime.timedelta(seconds=1.5)) + deltas = _measure_wait_deltas(backend._acquire_script.calls) + assert len(deltas) >= 4 + assert deltas[0] >= poll_interval_secs * 0.8 + assert deltas[1] >= poll_interval_secs * 1.6 + assert deltas[2] >= poll_max_secs * 0.75 + assert max(deltas) <= poll_max_secs + 0.15 + finally: + backend.close() + + def test_acquire__resets_wait_after_success(self): + """A successful acquire restarts the next idle sequence at poll_interval.""" + poll_interval_secs = 0.05 + backend = _make_backend( + "acquire_reset_test", + poll_interval=datetime.timedelta(seconds=poll_interval_secs), + ) + script = CountingAcquireScript(backend._acquire_script) + backend._acquire_script = script + try: + with pytest.raises(TimeoutError): + backend.acquire(timeout=datetime.timedelta(seconds=1)) + buildup_end = len(script.calls) + buildup_deltas = _measure_wait_deltas(script.calls) + # Slow runners inflate deltas, so only bind timing with slack. + if len(buildup_deltas) > 2: + assert buildup_deltas[2] >= 0.15 + + backend.enqueue(echo, args=[1]) + acquired = backend.acquire(timeout=datetime.timedelta(seconds=1)) + assert acquired is not None + assert len(script.calls) == buildup_end + 1 + + with pytest.raises(TimeoutError): + backend.acquire(timeout=datetime.timedelta(seconds=0.5)) + reset_deltas = _measure_wait_deltas(script.calls[buildup_end + 1 :]) + # Slow runners may only fit the first wait into the budget. + assert reset_deltas[0] < max(buildup_deltas) + if len(reset_deltas) > 1: + assert reset_deltas[1] >= poll_interval_secs * 1.6 + finally: + backend.close() + + def test_acquire__raise_timeout_error_at_deadline(self): + """Acquire raises TimeoutError when the deadline passes without a task.""" + backend = _make_backend("acquire_timeout_test") + try: + started_at = time.monotonic() + with pytest.raises(TimeoutError): + backend.acquire(timeout=datetime.timedelta(seconds=0.2)) + elapsed_secs = time.monotonic() - started_at + assert elapsed_secs >= 0.2 + assert elapsed_secs < 2 + finally: + backend.close() + + def test_acquire__raise_queue_empty_when_timeout_is_none(self): + """Acquire without a timeout raises queue.Empty after a single attempt.""" + backend = _make_backend("acquire_empty_test") + backend._acquire_script = CountingAcquireScript(backend._acquire_script) + try: + with pytest.raises(queue.Empty): + backend.acquire() + assert len(backend._acquire_script.calls) == 1 + finally: + backend.close() diff --git a/tests/test_command.py b/tests/test_command.py index 08fbed1..aeb93c2 100644 --- a/tests/test_command.py +++ b/tests/test_command.py @@ -1,4 +1,5 @@ import argparse +import datetime import logging import re import signal @@ -40,8 +41,20 @@ 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.poll_interval == 0.01 + assert parsed_arguments.poll_max_interval == 1 assert parsed_arguments.log_format is None + def test_add_arguments__parse_poll_intervals_as_floats(self): + """Parse poll options as floats.""" + parser = argparse.ArgumentParser() + threadmill.WorkerCommand().add_arguments(parser) + parsed_arguments = parser.parse_args( + ["--poll-interval", "0.05", "--poll-max-interval", "0.2"] + ) + assert parsed_arguments.poll_interval == 0.05 + assert parsed_arguments.poll_max_interval == 0.2 + def test_call_command__log_format(self): """Run the worker with the given log format string.""" call_command( @@ -100,6 +113,28 @@ def test_call_command__log_format__empty_string(self): ) assert handler.formatter.format(record) == "Hello world" + def test_call_command__poll_intervals(self): + """Convert poll options to timedeltas for the task executor.""" + with patch.object(threadmill.TaskExecutor, "run", autospec=True) as run: + call_command( + "threadmill", + "worker", + verbosity=0, + poll_interval=0.05, + poll_max_interval=0.2, + ) + executor = run.call_args.args[0] + assert executor.poll_interval == datetime.timedelta(seconds=0.05) + assert executor.poll_max_interval == datetime.timedelta(seconds=0.2) + + def test_call_command__poll_intervals__default_to_backend_defaults(self): + """Pass the backend default poll options to the task executor by default.""" + with patch.object(threadmill.TaskExecutor, "run", autospec=True) as run: + call_command("threadmill", "worker", verbosity=0) + executor = run.call_args.args[0] + assert executor.poll_interval == datetime.timedelta(seconds=0.01) + assert executor.poll_max_interval == datetime.timedelta(seconds=1) + @pytest.mark.benchmark def test_call_command__benchmark_compute( self, diff --git a/tests/test_executor.py b/tests/test_executor.py index 298a4df..c6bb440 100644 --- a/tests/test_executor.py +++ b/tests/test_executor.py @@ -70,13 +70,20 @@ def _task_result(task, *args, **kwargs) -> TaskResult: ) -def _make_worker(*, max_tasks: int | None = None) -> WorkerProcess: +def _make_worker( + *, + max_tasks: int | None = None, + poll_interval: datetime.timedelta | None = None, + poll_max_interval: datetime.timedelta | None = None, +) -> WorkerProcess: """Build an unstarted `WorkerProcess`.""" return WorkerProcess( thread_count=1, max_tasks=max_tasks, backend_alias="default", queues=("default",), + poll_interval=poll_interval, + poll_max_interval=poll_max_interval, log_formatter=JsonFormatter(), ) @@ -329,6 +336,33 @@ def test_maintain_worker_pool__restarts_dead_workers(self): class TestWorkerProcess: """Tests for the WorkerProcess class.""" + @pytest.fixture(autouse=True) + def restore_backend_poll_options(self): + """Restore the backend poll options after each test.""" + backend = default_task_backend + saved = (backend.poll_interval, backend.poll_max_interval) + yield + backend.poll_interval, backend.poll_max_interval = saved + + def test_run__applies_poll_overrides_to_backend(self): + """Apply poll overrides to the backend before starting consumer threads.""" + backend = default_task_backend + poll_interval = datetime.timedelta(seconds=0.02) + poll_max_interval = datetime.timedelta(seconds=0.3) + enqueued = backend.enqueue(echo, args=[1]) + worker = _make_worker( + max_tasks=1, + poll_interval=poll_interval, + poll_max_interval=poll_max_interval, + ) + worker.shutdown_requested.set() + # The backend registry resolves one instance per thread; running in this + # thread applies the overrides to the instance asserted on below. + worker.run() + assert backend.poll_interval == poll_interval + assert backend.poll_max_interval == poll_max_interval + assert backend.get_result(enqueued.id).status is TaskResultStatus.SUCCESSFUL + def test_record_task__increments_count(self): """record_task increments task_count.""" worker = _make_worker(max_tasks=5) diff --git a/threadmill/backends/redis.py b/threadmill/backends/redis.py index 8c1a12e..cf4932b 100644 --- a/threadmill/backends/redis.py +++ b/threadmill/backends/redis.py @@ -155,6 +155,13 @@ def __init__(self, alias: str, params: dict) -> None: 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.poll_interval = self.options.get( + "poll_interval", datetime.timedelta(seconds=0.01) + ) + self.poll_max_interval = self.options.get( + "poll_max_interval", datetime.timedelta(seconds=1) + ) + self._miss_count = 0 self._acquire_script = self.client.register_script(self.ACQUIRE_SCRIPT) self._acknowledge_script = self.client.register_script(self.ACKNOWLEDGE_SCRIPT) @@ -270,17 +277,25 @@ def acquire( str(int(self.lease_ttl.total_seconds() * 1000)), ], ): + self._miss_count = 0 return self.deserialize_task_result(data) try: - if deadline - time.monotonic() <= 0: - raise TimeoutError( - "No task available within the specified timeout." - ) + remaining = deadline - time.monotonic() except TypeError: raise queue.Empty("No task available.") - else: - time.sleep(0.01) + if remaining <= 0: + raise TimeoutError("No task available within the specified timeout.") + # Stop doubling once the interval reaches poll_max_interval; larger + # exponents would only overflow the float math. + cap = int(self.poll_max_interval / self.poll_interval).bit_length() + interval_secs = min( + self.poll_interval.total_seconds() * 2 ** min(self._miss_count, cap), + self.poll_max_interval.total_seconds(), + remaining, + ) + self._miss_count += 1 + time.sleep(interval_secs) def acknowledge(self, task_result: TaskResult) -> None: serialized = self.serialize_task_result(task_result) diff --git a/threadmill/executor.py b/threadmill/executor.py index 5285f3b..14df8b2 100644 --- a/threadmill/executor.py +++ b/threadmill/executor.py @@ -96,6 +96,8 @@ class TaskExecutor: threads: int = 1 max_tasks: int = 0 max_tasks_jitter: int = 0 + poll_interval: datetime.timedelta = datetime.timedelta(seconds=0.01) + poll_max_interval: datetime.timedelta = datetime.timedelta(seconds=1) is_publishing: bool = dataclasses.field(default=True, init=False) worker_processes: list[WorkerProcess] = dataclasses.field( default_factory=list, init=False @@ -127,6 +129,8 @@ def create_worker_process(self) -> WorkerProcess: backend_alias=self.backend.alias, queues=self.queues, exit_empty=self.exit_empty, + poll_interval=self.poll_interval, + poll_max_interval=self.poll_max_interval, log_formatter=self.log_formatter, ) worker.start() @@ -188,6 +192,8 @@ def __init__( backend_alias: str = "", queues: tuple[str, ...] = (), exit_empty: bool = False, + poll_interval: datetime.timedelta = datetime.timedelta(seconds=0.01), + poll_max_interval: datetime.timedelta = datetime.timedelta(seconds=1), log_formatter: logging.Formatter, ) -> None: """Create process with dedicated thread pool for task execution.""" @@ -198,6 +204,8 @@ def __init__( self.backend_alias = backend_alias self.queues = queues self.exit_empty = exit_empty + self.poll_interval = poll_interval + self.poll_max_interval = poll_max_interval self.log_formatter = log_formatter self.task_count = 0 self.lock: threading.Lock | None = None @@ -211,6 +219,8 @@ def run(self) -> None: self.lock = threading.Lock() self.expired = threading.Event() backend = task_backends[self.backend_alias] + backend.poll_interval = self.poll_interval + backend.poll_max_interval = self.poll_max_interval consumer_threads = [ WorkerThread(worker=self, index=index, backend=backend) for index in range(self.thread_count) diff --git a/threadmill/management/commands/threadmill.py b/threadmill/management/commands/threadmill.py index 7c40219..9235676 100644 --- a/threadmill/management/commands/threadmill.py +++ b/threadmill/management/commands/threadmill.py @@ -1,3 +1,4 @@ +import datetime import logging import signal import sys @@ -67,6 +68,18 @@ 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( + "--poll-interval", + type=float, + default=0.01, + help="Base wait between idle acquire attempts in seconds.", + ) + parser.add_argument( + "--poll-max-interval", + type=float, + default=1, + help="Maximum wait between idle acquire attempts in seconds.", + ) parser.add_argument( "--exit-empty", action="store_true", @@ -90,6 +103,8 @@ def handle( threads, max_tasks, max_tasks_jitter, + poll_interval, + poll_max_interval, exit_empty, log_format, **options, @@ -127,6 +142,8 @@ def handle( threads=threads, max_tasks=max_tasks, max_tasks_jitter=max_tasks_jitter, + poll_interval=datetime.timedelta(seconds=poll_interval), + poll_max_interval=datetime.timedelta(seconds=poll_max_interval), exit_empty=exit_empty, queues=queues, log_formatter=log_formatter,