pydocket 0.5.0__tar.gz → 0.5.2__tar.gz

{pydocket-0.5.0 → pydocket-0.5.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydocket
-Version: 0.5.0
+Version: 0.5.2
 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.5.0 → pydocket-0.5.2}/chaos/driver.py

@@ -23,12 +23,6 @@ from .tasks import toxic
 
 logging.getLogger().setLevel(logging.INFO)
 
-# Quiets down the testcontainers logger
-testcontainers_logger = logging.getLogger("testcontainers.core.container")
-testcontainers_logger.setLevel(logging.ERROR)
-testcontainers_logger = logging.getLogger("testcontainers.core.waiting_utils")
-testcontainers_logger.setLevel(logging.ERROR)
-
 console = logging.StreamHandler(stream=sys.stdout)
 console.setFormatter(
     logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")

pydocket-0.5.2/examples/common.py

@@ -0,0 +1,69 @@
+import asyncio
+import os
+import socket
+from contextlib import asynccontextmanager
+from typing import AsyncGenerator
+
+from docker import DockerClient
+
+
+@asynccontextmanager
+async def run_redis(version: str) -> AsyncGenerator[str, None]:
+    def get_free_port() -> int:
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind(("", 0))
+            return s.getsockname()[1]
+
+    port = get_free_port()
+
+    client = DockerClient.from_env()
+    container = client.containers.run(
+        f"redis:{version}",
+        detach=True,
+        ports={"6379/tcp": port},
+        auto_remove=True,
+    )
+
+    # Wait for Redis to be ready
+    for line in container.logs(stream=True):
+        if b"Ready to accept connections" in line:
+            break
+
+    url = f"redis://localhost:{port}/0"
+    print("***** Redis is running on %s *****", url)
+    try:
+        yield url
+    finally:
+        container.stop()
+
+
+async def run_example_workers(workers: int, concurrency: int, tasks: str):
+    async with run_redis("7.4.2") as redis_url:
+        processes = [
+            await asyncio.create_subprocess_exec(
+                "docket",
+                "worker",
+                "--name",
+                f"worker-{i}",
+                "--url",
+                redis_url,
+                "--tasks",
+                tasks,
+                "--concurrency",
+                str(concurrency),
+                env={
+                    **os.environ,
+                    "PYTHONPATH": os.path.abspath(
+                        os.path.join(os.path.dirname(__file__), "..")
+                    ),
+                },
+            )
+            for i in range(workers)
+        ]
+        try:
+            await asyncio.gather(*[p.wait() for p in processes])
+        except asyncio.CancelledError:
+            for p in processes:
+                p.kill()
+        finally:
+            await asyncio.gather(*[p.wait() for p in processes])

pydocket-0.5.2/examples/find_and_flood.py

@@ -0,0 +1,39 @@
+import asyncio
+from datetime import timedelta
+from logging import Logger, LoggerAdapter
+from typing import Annotated
+
+from docket import Docket
+from docket.annotations import Logged
+from docket.dependencies import CurrentDocket, Perpetual, TaskLogger
+
+from .common import run_example_workers
+
+
+async def find(
+    docket: Docket = CurrentDocket(),
+    logger: LoggerAdapter[Logger] = TaskLogger(),
+    perpetual: Perpetual = Perpetual(every=timedelta(seconds=3), automatic=True),
+) -> None:
+    for i in range(1, 10 + 1):
+        await docket.add(flood, key=str(i))(i)
+
+
+async def flood(
+    item: Annotated[int, Logged],
+    logger: LoggerAdapter[Logger] = TaskLogger(),
+) -> None:
+    logger.info("Working on %s", item)
+
+
+tasks = [find, flood]
+
+
+if __name__ == "__main__":
+    asyncio.run(
+        run_example_workers(
+            workers=3,
+            concurrency=8,
+            tasks="examples.find_and_flood:tasks",
+        )
+    )

{pydocket-0.5.0 → pydocket-0.5.2}/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.5.0 → pydocket-0.5.2}/src/docket/docket.py

@@ -28,6 +28,7 @@ from uuid import uuid4
 import redis.exceptions
 from opentelemetry import propagate, trace
 from redis.asyncio import ConnectionPool, Redis
+from redis.asyncio.client import Pipeline
 
 from .execution import (
     Execution,
@@ -256,9 +257,14 @@ class Docket:
 
         async def scheduler(*args: P.args, **kwargs: P.kwargs) -> Execution:
             execution = Execution(function, args, kwargs, when, key, attempt=1)
-            await self.schedule(execution)
+
+            async with self.redis() as redis:
+                async with redis.pipeline() as pipeline:
+                    await self._schedule(redis, pipeline, execution, replace=False)
+                    await pipeline.execute()
 
             TASKS_ADDED.add(1, {**self.labels(), **execution.general_labels()})
+            TASKS_SCHEDULED.add(1, {**self.labels(), **execution.general_labels()})
 
             return execution
 
@@ -291,15 +297,48 @@ class Docket:
 
         async def scheduler(*args: P.args, **kwargs: P.kwargs) -> Execution:
             execution = Execution(function, args, kwargs, when, key, attempt=1)
-            await self.cancel(key)
-            await self.schedule(execution)
+
+            async with self.redis() as redis:
+                async with redis.pipeline() as pipeline:
+                    await self._schedule(redis, pipeline, execution, replace=True)
+                    await pipeline.execute()
 
             TASKS_REPLACED.add(1, {**self.labels(), **execution.general_labels()})
+            TASKS_CANCELLED.add(1, {**self.labels(), **execution.general_labels()})
+            TASKS_SCHEDULED.add(1, {**self.labels(), **execution.general_labels()})
 
             return execution
 
         return scheduler
 
+    async def schedule(self, execution: Execution) -> None:
+        with tracer.start_as_current_span(
+            "docket.schedule",
+            attributes={
+                **self.labels(),
+                **execution.specific_labels(),
+                "code.function.name": execution.function.__name__,
+            },
+        ):
+            async with self.redis() as redis:
+                async with redis.pipeline() as pipeline:
+                    await self._schedule(redis, pipeline, execution, replace=False)
+                    await pipeline.execute()
+
+        TASKS_SCHEDULED.add(1, {**self.labels(), **execution.general_labels()})
+
+    async def cancel(self, key: str) -> None:
+        with tracer.start_as_current_span(
+            "docket.cancel",
+            attributes={**self.labels(), "docket.key": key},
+        ):
+            async with self.redis() as redis:
+                async with redis.pipeline() as pipeline:
+                    await self._cancel(pipeline, key)
+                    await pipeline.execute()
+
+        TASKS_CANCELLED.add(1, self.labels())
+
     @property
     def queue_key(self) -> str:
         return f"{self.name}:queue"
@@ -314,7 +353,13 @@ class Docket:
     def parked_task_key(self, key: str) -> str:
         return f"{self.name}:{key}"
 
-    async def schedule(self, execution: Execution) -> None:
+    async def _schedule(
+        self,
+        redis: Redis,
+        pipeline: Pipeline,
+        execution: Execution,
+        replace: bool = False,
+    ) -> None:
         if self.strike_list.is_stricken(execution):
             logger.warning(
                 "%r is stricken, skipping schedule of %r",
@@ -334,53 +379,35 @@ class Docket:
         message: dict[bytes, bytes] = execution.as_message()
         propagate.inject(message, setter=message_setter)
 
-        with tracer.start_as_current_span(
-            "docket.schedule",
-            attributes={
-                **self.labels(),
-                **execution.specific_labels(),
-                "code.function.name": execution.function.__name__,
-            },
-        ):
-            key = execution.key
-            when = execution.when
+        key = execution.key
+        when = execution.when
+        known_task_key = self.known_task_key(key)
 
-            async with self.redis() as redis:
+        async with redis.lock(f"{known_task_key}:lock", timeout=10):
+            if replace:
+                await self._cancel(pipeline, key)
+            else:
                 # if the task is already in the queue or stream, retain it
-                if await redis.exists(self.known_task_key(key)):
+                if await redis.exists(known_task_key):
                     logger.debug(
-                        "Task %r is already in the queue or stream, skipping schedule",
+                        "Task %r is already in the queue or stream, not scheduling",
                         key,
                         extra=self.labels(),
                     )
                     return
 
-                async with redis.pipeline() as pipe:
-                    pipe.set(self.known_task_key(key), when.timestamp())
-
-                    if when <= datetime.now(timezone.utc):
-                        pipe.xadd(self.stream_key, message)  # type: ignore[arg-type]
-                    else:
-                        pipe.hset(self.parked_task_key(key), mapping=message)  # type: ignore[arg-type]
-                        pipe.zadd(self.queue_key, {key: when.timestamp()})
+            pipeline.set(known_task_key, when.timestamp())
 
-                    await pipe.execute()
-
-        TASKS_SCHEDULED.add(1, {**self.labels(), **execution.general_labels()})
-
-    async def cancel(self, key: str) -> None:
-        with tracer.start_as_current_span(
-            "docket.cancel",
-            attributes={**self.labels(), "docket.key": key},
-        ):
-            async with self.redis() as redis:
-                async with redis.pipeline() as pipe:
-                    pipe.delete(self.known_task_key(key))
-                    pipe.delete(self.parked_task_key(key))
-                    pipe.zrem(self.queue_key, key)
-                    await pipe.execute()
+            if when <= datetime.now(timezone.utc):
+                pipeline.xadd(self.stream_key, message)  # type: ignore[arg-type]
+            else:
+                pipeline.hset(self.parked_task_key(key), mapping=message)  # type: ignore[arg-type]
+                pipeline.zadd(self.queue_key, {key: when.timestamp()})
 
-        TASKS_CANCELLED.add(1, self.labels())
+    async def _cancel(self, pipeline: Pipeline, key: str) -> None:
+        pipeline.delete(self.known_task_key(key))
+        pipeline.delete(self.parked_task_key(key))
+        pipeline.zrem(self.queue_key, key)
 
     @property
     def strike_key(self) -> str:

{pydocket-0.5.0 → pydocket-0.5.2}/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
@@ -19,6 +17,7 @@ import redis.exceptions
 from opentelemetry import propagate, trace
 from opentelemetry.trace import Tracer
 from redis.asyncio import Redis
+from redis.exceptions import LockError
 
 from .docket import (
     Docket,
@@ -53,8 +52,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__(
@@ -216,13 +213,14 @@ class Worker:
                 await asyncio.sleep(self.reconnection_delay.total_seconds())
 
     async def _worker_loop(self, forever: bool = False):
-        should_stop = asyncio.Event()
+        worker_stopping = asyncio.Event()
+
+        await self._schedule_all_automatic_perpetual_tasks()
 
         async with self.docket.redis() as redis:
             scheduler_task = asyncio.create_task(
-                self._scheduler_loop(redis, should_stop)
+                self._scheduler_loop(redis, worker_stopping)
             )
-
             active_tasks: dict[asyncio.Task[None], RedisMessageID] = {}
 
             async def check_for_work() -> bool:
@@ -329,13 +327,13 @@ class Worker:
                     await asyncio.gather(*active_tasks, return_exceptions=True)
                     await process_completed_tasks()
 
-                should_stop.set()
+                worker_stopping.set()
                 await scheduler_task
 
     async def _scheduler_loop(
         self,
         redis: Redis,
-        should_stop: asyncio.Event,
+        worker_stopping: asyncio.Event,
     ) -> None:
         """Loop that moves due tasks from the queue to the stream."""
 
@@ -389,7 +387,7 @@ class Worker:
 
         total_work: int = sys.maxsize
 
-        while not should_stop.is_set() or total_work:
+        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],
@@ -416,22 +414,47 @@ class Worker:
 
         logger.debug("Scheduler loop finished", extra=self._log_context())
 
+    async def _schedule_all_automatic_perpetual_tasks(self) -> None:
+        from .dependencies import Perpetual, get_single_dependency_parameter_of_type
+
+        async with self.docket.redis() as redis:
+            try:
+                async with redis.lock(
+                    f"{self.docket.name}:perpetual:lock", timeout=10, blocking=False
+                ):
+                    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)()
+            except LockError:  # pragma: no cover
+                return
+
     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()
         function = self.docket.tasks.get(function_name)
         if function is None:
+            async with self.docket.redis() as redis:
+                await redis.delete(self.docket.known_task_key(key))
             logger.warning(
                 "Task function %r not found", function_name, extra=log_context
             )
             return
 
         execution = Execution.from_message(function, message)
+        dependencies = self._get_dependencies(execution)
 
         log_context = {**log_context, **execution.specific_labels()}
         counter_labels = {**self.labels(), **execution.general_labels()}
@@ -440,6 +463,9 @@ class Worker:
         call = execution.call_repr()
 
         if self.docket.strike_list.is_stricken(execution):
+            async with self.docket.redis() as redis:
+                await redis.delete(self.docket.known_task_key(key))
+
             arrow = "🗙"
             logger.warning("%s %s", arrow, call, extra=log_context)
             TASKS_STRICKEN.add(1, counter_labels | {"docket.where": "worker"})
@@ -448,7 +474,12 @@ class Worker:
         if execution.key in self._execution_counts:
             self._execution_counts[execution.key] += 1
 
-        dependencies = self._get_dependencies(execution)
+        # Preemptively reschedule the perpetual task for the future, or clear the
+        # known task key for this task
+        rescheduled = await self._perpetuate_if_requested(execution, dependencies)
+        if not rescheduled:
+            async with self.docket.redis() as redis:
+                await redis.delete(self.docket.known_task_key(key))
 
         context = propagate.extract(message, getter=message_getter)
         initiating_context = trace.get_current_span(context).get_span_context()
@@ -511,12 +542,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
@@ -532,16 +563,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
@@ -553,31 +582,32 @@ 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 | None = None,
     ) -> 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:
+            await self.docket.cancel(execution.key)
             return False
 
         now = datetime.now(timezone.utc)
-        execution.when = max(now, now + perpetual.every - duration)
-        execution.args = perpetual.args
-        execution.kwargs = perpetual.kwargs
+        when = max(now, now + perpetual.every - (duration or timedelta(0)))
+
+        await self.docket.replace(execution.function, when, execution.key)(
+            *perpetual.args,
+            **perpetual.kwargs,
+        )
 
-        await self.docket.schedule(execution)
+        if duration is not None:
+            TASKS_PERPETUATED.add(1, {**self.labels(), **execution.specific_labels()})
 
-        TASKS_PERPETUATED.add(1, {**self.labels(), **execution.specific_labels()})
         return True
 
     @property

{pydocket-0.5.0 → pydocket-0.5.2}/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

{pydocket-0.5.0 → pydocket-0.5.2}/tests/test_worker.py

@@ -374,9 +374,6 @@ async def test_perpetual_tasks_are_scheduled_close_to_target_time(
     ):
         timestamps.append(datetime.now(timezone.utc))
 
-        if len(timestamps) % 2 == 0:
-            await asyncio.sleep(0.05)
-
     await docket.add(perpetual_task, key="my-key")(a="a", b=2)
 
     await worker.run_at_most({"my-key": 8})
@@ -384,11 +381,15 @@ async def test_perpetual_tasks_are_scheduled_close_to_target_time(
     assert len(timestamps) == 8
 
     intervals = [next - previous for previous, next in zip(timestamps, timestamps[1:])]
-    total = timedelta(seconds=sum(i.total_seconds() for i in intervals))
-    average = total / len(intervals)
+    minimum = min(intervals)
+    maximum = max(intervals)
+
+    debug = ", ".join([f"{i.total_seconds() * 1000:.2f}ms" for i in intervals])
 
-    # even with a variable duration, Docket attempts to schedule them equally
-    assert timedelta(milliseconds=45) <= average <= timedelta(milliseconds=75)
+    # even with a variable duration, Docket attempts to schedule them equally and to
+    # abide by the target interval
+    assert minimum >= timedelta(milliseconds=50), debug
+    assert maximum <= timedelta(milliseconds=75), debug
 
 
 async def test_worker_can_exit_from_perpetual_tasks_that_queue_further_tasks(