pydocket 0.4.0__tar.gz → 0.5.1__tar.gz

{pydocket-0.4.0 → pydocket-0.5.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydocket
-Version: 0.4.0
+Version: 0.5.1
 Summary: A distributed background task system for Python functions
 Project-URL: Homepage, https://github.com/chrisguidry/docket
 Project-URL: Bug Tracker, https://github.com/chrisguidry/docket/issues

{pydocket-0.4.0 → pydocket-0.5.1}/chaos/driver.py

@@ -76,9 +76,9 @@ async def run_redis(version: str) -> AsyncGenerator[tuple[str, Container], None]
 
 async def main(
     mode: Literal["performance", "chaos"] = "chaos",
-    tasks: int = 5000,
-    producers: int = 4,
-    workers: int = 7,
+    tasks: int = 20000,
+    producers: int = 5,
+    workers: int = 10,
 ):
     async with (
         run_redis("7.4.2") as (redis_url, redis_container),
@@ -97,9 +97,7 @@ async def main(
         # Add in some random strikes to performance test
         for _ in range(100):
             parameter = f"param_{random.randint(1, 100)}"
-            operator: Operator = random.choice(
-                ["==", "!=", ">", ">=", "<", "<=", "between"]
-            )
+            operator = random.choice(list(Operator))
             value = f"val_{random.randint(1, 1000)}"
             await docket.strike("rando", parameter, operator, value)
 
@@ -141,11 +139,9 @@ async def main(
                 redis_url,
                 "--tasks",
                 "chaos.tasks:chaos_tasks",
-                env=environment
-                | {
-                    "OTEL_SERVICE_NAME": "chaos-worker",
-                    "DOCKET_WORKER_REDELIVERY_TIMEOUT": "5s",
-                },
+                "--redelivery-timeout",
+                "5s",
+                env=environment | {"OTEL_SERVICE_NAME": "chaos-worker"},
                 stdout=subprocess.DEVNULL,
                 stderr=subprocess.DEVNULL,
             )

{pydocket-0.4.0 → pydocket-0.5.1}/chaos/producer.py

@@ -1,8 +1,11 @@
 import asyncio
+import datetime
 import logging
 import os
+import random
 import sys
 import time
+from datetime import timedelta
 
 import redis.exceptions
 
@@ -14,6 +17,10 @@ logging.getLogger().setLevel(logging.INFO)
 logger = logging.getLogger("chaos.producer")
 
 
+def now() -> datetime.datetime:
+    return datetime.datetime.now(datetime.timezone.utc)
+
+
 async def main(tasks_to_produce: int):
     docket = Docket(
         name=os.environ["DOCKET_NAME"],
@@ -25,7 +32,9 @@ async def main(tasks_to_produce: int):
             async with docket:
                 async with docket.redis() as r:
                     for _ in range(tasks_sent, tasks_to_produce):
-                        execution = await docket.add(hello)()
+                        jitter = 5 * ((random.random() * 2) - 1)
+                        when = now() + timedelta(seconds=jitter)
+                        execution = await docket.add(hello, when=when)()
                         await r.zadd("hello:sent", {execution.key: time.time()})
                         logger.info("Added task %s", execution.key)
                         tasks_sent += 1

{pydocket-0.4.0 → pydocket-0.5.1}/src/docket/cli.py

@@ -162,6 +162,7 @@ def worker(
                 "This can be specified multiple times.  A task collection is any "
                 "iterable of async functions."
             ),
+            envvar="DOCKET_TASKS",
         ),
     ] = ["docket.tasks:standard_tasks"],
     docket_: Annotated[
@@ -236,6 +237,14 @@ def worker(
             envvar="DOCKET_WORKER_MINIMUM_CHECK_INTERVAL",
         ),
     ] = timedelta(milliseconds=100),
+    scheduling_resolution: Annotated[
+        timedelta,
+        typer.Option(
+            parser=duration,
+            help="How frequently to check for future tasks to be scheduled",
+            envvar="DOCKET_WORKER_SCHEDULING_RESOLUTION",
+        ),
+    ] = timedelta(milliseconds=250),
     until_finished: Annotated[
         bool,
         typer.Option(
@@ -260,6 +269,7 @@ def worker(
             redelivery_timeout=redelivery_timeout,
             reconnection_delay=reconnection_delay,
             minimum_check_interval=minimum_check_interval,
+            scheduling_resolution=scheduling_resolution,
             until_finished=until_finished,
             metrics_port=metrics_port,
             tasks=tasks,
@@ -542,6 +552,18 @@ def relative_time(now: datetime, when: datetime) -> str:
 
 @app.command(help="Shows a snapshot of what's on the docket right now")
 def snapshot(
+    tasks: Annotated[
+        list[str],
+        typer.Option(
+            "--tasks",
+            help=(
+                "The dotted path of a task collection to register with the docket. "
+                "This can be specified multiple times.  A task collection is any "
+                "iterable of async functions."
+            ),
+            envvar="DOCKET_TASKS",
+        ),
+    ] = ["docket.tasks:standard_tasks"],
     docket_: Annotated[
         str,
         typer.Option(
@@ -560,6 +582,9 @@ def snapshot(
 ) -> None:
     async def run() -> DocketSnapshot:
         async with Docket(name=docket_, url=url) as docket:
+            for task_path in tasks:
+                docket.register_collection(task_path)
+
             return await docket.snapshot()
 
     snapshot = asyncio.run(run())

{pydocket-0.4.0 → pydocket-0.5.1}/src/docket/dependencies.py

@@ -2,7 +2,7 @@ import abc
 import inspect
 import logging
 from datetime import timedelta
-from typing import Any, Awaitable, Callable, Counter, cast
+from typing import Any, Awaitable, Callable, Counter, TypeVar, cast
 
 from .docket import Docket
 from .execution import Execution
@@ -130,12 +130,29 @@ class Perpetual(Dependency):
     single = True
 
     every: timedelta
+    automatic: bool
+
     args: tuple[Any, ...]
     kwargs: dict[str, Any]
+
     cancelled: bool
 
-    def __init__(self, every: timedelta = timedelta(0)) -> None:
+    def __init__(
+        self,
+        every: timedelta = timedelta(0),
+        automatic: bool = False,
+    ) -> None:
+        """Declare a task that should be run perpetually.
+
+        Args:
+            every: The target interval between task executions.
+            automatic: If set, this task will be automatically scheduled during worker
+                startup and continually through the worker's lifespan.  This ensures
+                that the task will always be scheduled despite crashes and other
+                adverse conditions.  Automatic tasks must not require any arguments.
+        """
         self.every = every
+        self.automatic = automatic
         self.cancelled = False
 
     def __call__(
@@ -170,6 +187,29 @@ def get_dependency_parameters(
     return dependencies
 
 
+D = TypeVar("D", bound=Dependency)
+
+
+def get_single_dependency_parameter_of_type(
+    function: Callable[..., Awaitable[Any]], dependency_type: type[D]
+) -> D | None:
+    assert dependency_type.single, "Dependency must be single"
+    for _, dependency in get_dependency_parameters(function).items():
+        if isinstance(dependency, dependency_type):
+            return dependency
+    return None
+
+
+def get_single_dependency_of_type(
+    dependencies: dict[str, Dependency], dependency_type: type[D]
+) -> D | None:
+    assert dependency_type.single, "Dependency must be single"
+    for _, dependency in dependencies.items():
+        if isinstance(dependency, dependency_type):
+            return dependency
+    return None
+
+
 def validate_dependencies(function: Callable[..., Awaitable[Any]]) -> None:
     parameters = get_dependency_parameters(function)
 

{pydocket-0.4.0 → pydocket-0.5.1}/src/docket/worker.py

@@ -6,11 +6,9 @@ from datetime import datetime, timedelta, timezone
 from types import TracebackType
 from typing import (
     TYPE_CHECKING,
-    Any,
     Mapping,
     Protocol,
     Self,
-    TypeVar,
     cast,
 )
 from uuid import uuid4
@@ -18,6 +16,7 @@ from uuid import uuid4
 import redis.exceptions
 from opentelemetry import propagate, trace
 from opentelemetry.trace import Tracer
+from redis.asyncio import Redis
 
 from .docket import (
     Docket,
@@ -52,8 +51,6 @@ tracer: Tracer = trace.get_tracer(__name__)
 if TYPE_CHECKING:  # pragma: no cover
     from .dependencies import Dependency
 
-D = TypeVar("D", bound="Dependency")
-
 
 class _stream_due_tasks(Protocol):
     async def __call__(
@@ -68,6 +65,7 @@ class Worker:
     redelivery_timeout: timedelta
     reconnection_delay: timedelta
     minimum_check_interval: timedelta
+    scheduling_resolution: timedelta
 
     def __init__(
         self,
@@ -77,6 +75,7 @@ class Worker:
         redelivery_timeout: timedelta = timedelta(minutes=5),
         reconnection_delay: timedelta = timedelta(seconds=5),
         minimum_check_interval: timedelta = timedelta(milliseconds=100),
+        scheduling_resolution: timedelta = timedelta(milliseconds=250),
     ) -> None:
         self.docket = docket
         self.name = name or f"worker:{uuid4()}"
@@ -84,6 +83,7 @@ class Worker:
         self.redelivery_timeout = redelivery_timeout
         self.reconnection_delay = reconnection_delay
         self.minimum_check_interval = minimum_check_interval
+        self.scheduling_resolution = scheduling_resolution
 
     async def __aenter__(self) -> Self:
         self._heartbeat_task = asyncio.create_task(self._heartbeat())
@@ -128,6 +128,7 @@ class Worker:
         redelivery_timeout: timedelta = timedelta(minutes=5),
         reconnection_delay: timedelta = timedelta(seconds=5),
         minimum_check_interval: timedelta = timedelta(milliseconds=100),
+        scheduling_resolution: timedelta = timedelta(milliseconds=250),
         until_finished: bool = False,
         metrics_port: int | None = None,
         tasks: list[str] = ["docket.tasks:standard_tasks"],
@@ -144,6 +145,7 @@ class Worker:
                     redelivery_timeout=redelivery_timeout,
                     reconnection_delay=reconnection_delay,
                     minimum_check_interval=minimum_check_interval,
+                    scheduling_resolution=scheduling_resolution,
                 ) as worker:
                     if until_finished:
                         await worker.run_until_finished()
@@ -210,57 +212,28 @@ class Worker:
                 await asyncio.sleep(self.reconnection_delay.total_seconds())
 
     async def _worker_loop(self, forever: bool = False):
-        async with self.docket.redis() as redis:
-            stream_due_tasks: _stream_due_tasks = cast(
-                _stream_due_tasks,
-                redis.register_script(
-                    # Lua script to atomically move scheduled tasks to the stream
-                    # KEYS[1]: queue key (sorted set)
-                    # KEYS[2]: stream key
-                    # ARGV[1]: current timestamp
-                    # ARGV[2]: docket name prefix
-                    """
-                local total_work = redis.call('ZCARD', KEYS[1])
-                local due_work = 0
-
-                if total_work > 0 then
-                    local tasks = redis.call('ZRANGEBYSCORE', KEYS[1], 0, ARGV[1])
-
-                    for i, key in ipairs(tasks) do
-                        local hash_key = ARGV[2] .. ":" .. key
-                        local task_data = redis.call('HGETALL', hash_key)
-
-                        if #task_data > 0 then
-                            local task = {}
-                            for j = 1, #task_data, 2 do
-                                task[task_data[j]] = task_data[j+1]
-                            end
-
-                            redis.call('XADD', KEYS[2], '*',
-                                'key', task['key'],
-                                'when', task['when'],
-                                'function', task['function'],
-                                'args', task['args'],
-                                'kwargs', task['kwargs'],
-                                'attempt', task['attempt']
-                            )
-                            redis.call('DEL', hash_key)
-                            due_work = due_work + 1
-                        end
-                    end
-                end
+        worker_stopping = asyncio.Event()
 
-                if due_work > 0 then
-                    redis.call('ZREMRANGEBYSCORE', KEYS[1], 0, ARGV[1])
-                end
+        await self._schedule_all_automatic_perpetual_tasks()
+        perpetual_scheduling_task = asyncio.create_task(
+            self._perpetual_scheduling_loop(worker_stopping)
+        )
 
-                return {total_work, due_work}
-                """
-                ),
+        async with self.docket.redis() as redis:
+            scheduler_task = asyncio.create_task(
+                self._scheduler_loop(redis, worker_stopping)
             )
-
             active_tasks: dict[asyncio.Task[None], RedisMessageID] = {}
 
+            async def check_for_work() -> bool:
+                async with redis.pipeline() as pipeline:
+                    pipeline.xlen(self.docket.stream_key)
+                    pipeline.zcard(self.docket.queue_key)
+                    results: list[int] = await pipeline.execute()
+                    stream_len = results[0]
+                    queue_len = results[1]
+                    return stream_len > 0 or queue_len > 0
+
             async def process_completed_tasks() -> None:
                 completed_tasks = {task for task in active_tasks if task.done()}
                 for task in completed_tasks:
@@ -280,10 +253,13 @@ class Worker:
                         )
                         await pipeline.execute()
 
-            future_work, due_work = sys.maxsize, 0
+            has_work: bool = True
+
+            if not forever:  # pragma: no branch
+                has_work = await check_for_work()
 
             try:
-                while forever or future_work or active_tasks:
+                while forever or has_work or active_tasks:
                     await process_completed_tasks()
 
                     available_slots = self.concurrency - len(active_tasks)
@@ -297,28 +273,13 @@ class Worker:
                         task = asyncio.create_task(self._execute(message))
                         active_tasks[task] = message_id
 
-                        nonlocal available_slots, future_work
+                        nonlocal available_slots
                         available_slots -= 1
-                        future_work += 1
 
                     if available_slots <= 0:
                         await asyncio.sleep(self.minimum_check_interval.total_seconds())
                         continue
 
-                    future_work, due_work = await stream_due_tasks(
-                        keys=[self.docket.queue_key, self.docket.stream_key],
-                        args=[datetime.now(timezone.utc).timestamp(), self.docket.name],
-                    )
-                    if due_work > 0:
-                        logger.debug(
-                            "Moved %d/%d due tasks from %s to %s",
-                            due_work,
-                            future_work,
-                            self.docket.queue_key,
-                            self.docket.stream_key,
-                            extra=self._log_context(),
-                        )
-
                     redeliveries: RedisMessages
                     _, redeliveries, *_ = await redis.xautoclaim(
                         name=self.docket.stream_key,
@@ -348,10 +309,14 @@ class Worker:
                         ),
                         count=available_slots,
                     )
+
                     for _, messages in new_deliveries:
                         for message_id, message in messages:
                             start_task(message_id, message)
 
+                    if not forever and not active_tasks and not new_deliveries:
+                        has_work = await check_for_work()
+
             except asyncio.CancelledError:
                 if active_tasks:  # pragma: no cover
                     logger.info(
@@ -364,7 +329,142 @@ class Worker:
                     await asyncio.gather(*active_tasks, return_exceptions=True)
                     await process_completed_tasks()
 
+                worker_stopping.set()
+                await scheduler_task
+                await perpetual_scheduling_task
+
+    async def _scheduler_loop(
+        self,
+        redis: Redis,
+        worker_stopping: asyncio.Event,
+    ) -> None:
+        """Loop that moves due tasks from the queue to the stream."""
+
+        stream_due_tasks: _stream_due_tasks = cast(
+            _stream_due_tasks,
+            redis.register_script(
+                # Lua script to atomically move scheduled tasks to the stream
+                # KEYS[1]: queue key (sorted set)
+                # KEYS[2]: stream key
+                # ARGV[1]: current timestamp
+                # ARGV[2]: docket name prefix
+                """
+            local total_work = redis.call('ZCARD', KEYS[1])
+            local due_work = 0
+
+            if total_work > 0 then
+                local tasks = redis.call('ZRANGEBYSCORE', KEYS[1], 0, ARGV[1])
+
+                for i, key in ipairs(tasks) do
+                    local hash_key = ARGV[2] .. ":" .. key
+                    local task_data = redis.call('HGETALL', hash_key)
+
+                    if #task_data > 0 then
+                        local task = {}
+                        for j = 1, #task_data, 2 do
+                            task[task_data[j]] = task_data[j+1]
+                        end
+
+                        redis.call('XADD', KEYS[2], '*',
+                            'key', task['key'],
+                            'when', task['when'],
+                            'function', task['function'],
+                            'args', task['args'],
+                            'kwargs', task['kwargs'],
+                            'attempt', task['attempt']
+                        )
+                        redis.call('DEL', hash_key)
+                        due_work = due_work + 1
+                    end
+                end
+            end
+
+            if due_work > 0 then
+                redis.call('ZREMRANGEBYSCORE', KEYS[1], 0, ARGV[1])
+            end
+
+            return {total_work, due_work}
+            """
+            ),
+        )
+
+        total_work: int = sys.maxsize
+
+        while not worker_stopping.is_set() or total_work:
+            try:
+                total_work, due_work = await stream_due_tasks(
+                    keys=[self.docket.queue_key, self.docket.stream_key],
+                    args=[datetime.now(timezone.utc).timestamp(), self.docket.name],
+                )
+
+                if due_work > 0:
+                    logger.debug(
+                        "Moved %d/%d due tasks from %s to %s",
+                        due_work,
+                        total_work,
+                        self.docket.queue_key,
+                        self.docket.stream_key,
+                        extra=self._log_context(),
+                    )
+            except Exception:  # pragma: no cover
+                logger.exception(
+                    "Error in scheduler loop",
+                    exc_info=True,
+                    extra=self._log_context(),
+                )
+            finally:
+                await asyncio.sleep(self.scheduling_resolution.total_seconds())
+
+        logger.debug("Scheduler loop finished", extra=self._log_context())
+
+    async def _perpetual_scheduling_loop(self, worker_stopping: asyncio.Event) -> None:
+        """Loop that ensures that automatic perpetual tasks are always scheduled."""
+
+        while not worker_stopping.is_set():
+            minimum_interval = self.scheduling_resolution
+            try:
+                minimum_interval = await self._schedule_all_automatic_perpetual_tasks()
+            except Exception:  # pragma: no cover
+                logger.exception(
+                    "Error in perpetual scheduling loop",
+                    exc_info=True,
+                    extra=self._log_context(),
+                )
+            finally:
+                # Wait until just before the next time any task would need to be
+                # scheduled (one scheduling_resolution before the lowest interval)
+                interval = max(
+                    minimum_interval - self.scheduling_resolution,
+                    self.scheduling_resolution,
+                )
+                assert interval <= self.scheduling_resolution
+                await asyncio.sleep(interval.total_seconds())
+
+    async def _schedule_all_automatic_perpetual_tasks(self) -> timedelta:
+        from .dependencies import Perpetual, get_single_dependency_parameter_of_type
+
+        minimum_interval = self.scheduling_resolution
+        for task_function in self.docket.tasks.values():
+            perpetual = get_single_dependency_parameter_of_type(
+                task_function, Perpetual
+            )
+            if perpetual is None:
+                continue
+
+            if not perpetual.automatic:
+                continue
+
+            key = task_function.__name__
+            await self.docket.add(task_function, key=key)()
+            minimum_interval = min(minimum_interval, perpetual.every)
+
+        return minimum_interval
+
     async def _execute(self, message: RedisMessage) -> None:
+        key = message[b"key"].decode()
+        async with self.docket.redis() as redis:
+            await redis.delete(self.docket.known_task_key(key))
+
         log_context: Mapping[str, str | float] = self._log_context()
 
         function_name = message[b"function"].decode()
@@ -377,9 +477,6 @@ class Worker:
 
         execution = Execution.from_message(function, message)
 
-        async with self.docket.redis() as redis:
-            await redis.delete(self.docket.known_task_key(execution.key))
-
         log_context = {**log_context, **execution.specific_labels()}
         counter_labels = {**self.labels(), **execution.general_labels()}
 
@@ -458,12 +555,12 @@ class Worker:
     def _get_dependencies(
         self,
         execution: Execution,
-    ) -> dict[str, Any]:
+    ) -> dict[str, "Dependency"]:
         from .dependencies import get_dependency_parameters
 
         parameters = get_dependency_parameters(execution.function)
 
-        dependencies: dict[str, Any] = {}
+        dependencies: dict[str, "Dependency"] = {}
 
         for parameter_name, dependency in parameters.items():
             # If the argument is already provided, skip it, which allows users to call
@@ -479,16 +576,14 @@ class Worker:
     async def _retry_if_requested(
         self,
         execution: Execution,
-        dependencies: dict[str, Any],
+        dependencies: dict[str, "Dependency"],
     ) -> bool:
-        from .dependencies import Retry
+        from .dependencies import Retry, get_single_dependency_of_type
 
-        retries = [retry for retry in dependencies.values() if isinstance(retry, Retry)]
-        if not retries:
+        retry = get_single_dependency_of_type(dependencies, Retry)
+        if not retry:
             return False
 
-        retry = retries[0]
-
         if retry.attempts is None or execution.attempt < retry.attempts:
             execution.when = datetime.now(timezone.utc) + retry.delay
             execution.attempt += 1
@@ -500,19 +595,16 @@ class Worker:
         return False
 
     async def _perpetuate_if_requested(
-        self, execution: Execution, dependencies: dict[str, Any], duration: timedelta
+        self,
+        execution: Execution,
+        dependencies: dict[str, "Dependency"],
+        duration: timedelta,
     ) -> bool:
-        from .dependencies import Perpetual
-
-        perpetuals = [
-            perpetual
-            for perpetual in dependencies.values()
-            if isinstance(perpetual, Perpetual)
-        ]
-        if not perpetuals:
-            return False
+        from .dependencies import Perpetual, get_single_dependency_of_type
 
-        perpetual = perpetuals[0]
+        perpetual = get_single_dependency_of_type(dependencies, Perpetual)
+        if not perpetual:
+            return False
 
         if perpetual.cancelled:
             return False
@@ -572,11 +664,10 @@ class Worker:
                         pipeline.zcount(self.docket.queue_key, 0, now)
                         pipeline.zcount(self.docket.queue_key, now, "+inf")
 
-                        (
-                            stream_depth,
-                            overdue_depth,
-                            schedule_depth,
-                        ) = await pipeline.execute()
+                        results: list[int] = await pipeline.execute()
+                        stream_depth = results[0]
+                        overdue_depth = results[1]
+                        schedule_depth = results[2]
 
                         QUEUE_DEPTH.set(
                             stream_depth + overdue_depth, self.docket.labels()

{pydocket-0.4.0 → pydocket-0.5.1}/tests/conftest.py

@@ -155,7 +155,9 @@ async def docket(redis_url: str, aiolib: str) -> AsyncGenerator[Docket, None]:
 @pytest.fixture
 async def worker(docket: Docket) -> AsyncGenerator[Worker, None]:
     async with Worker(
-        docket, minimum_check_interval=timedelta(milliseconds=10)
+        docket,
+        minimum_check_interval=timedelta(milliseconds=10),
+        scheduling_resolution=timedelta(milliseconds=10),
     ) as worker:
         yield worker
 

{pydocket-0.4.0 → pydocket-0.5.1}/tests/test_fundamentals.py

@@ -1037,3 +1037,31 @@ async def test_perpetual_tasks_perpetuate_even_after_errors(
     await worker.run_at_most({execution.key: 3})
 
     assert calls == 3
+
+
+async def test_perpetual_tasks_can_be_automatically_scheduled(
+    docket: Docket, worker: Worker
+):
+    """Perpetual tasks can be automatically scheduled"""
+
+    calls = 0
+
+    async def my_automatic_task(
+        perpetual: Perpetual = Perpetual(
+            every=timedelta(milliseconds=50), automatic=True
+        ),
+    ):
+        assert isinstance(perpetual, Perpetual)
+
+        assert perpetual.every == timedelta(milliseconds=50)
+
+        nonlocal calls
+        calls += 1
+
+    # Note we never add this task to the docket, we just register it.
+    docket.register(my_automatic_task)
+
+    # The automatic key will be the task function's name
+    await worker.run_at_most({"my_automatic_task": 3})
+
+    assert calls == 3