pydocket 0.5.0__tar.gz → 0.5.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydocket
-Version: 0.5.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.5.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.5.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
@@ -53,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__(
@@ -216,13 +212,17 @@ 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()
+        perpetual_scheduling_task = asyncio.create_task(
+            self._perpetual_scheduling_loop(worker_stopping)
+        )
 
         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 +329,14 @@ class Worker:
                     await asyncio.gather(*active_tasks, return_exceptions=True)
                     await process_completed_tasks()
 
-                should_stop.set()
+                worker_stopping.set()
                 await scheduler_task
+                await perpetual_scheduling_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 +390,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,6 +417,49 @@ class Worker:
 
         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:
@@ -511,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
@@ -532,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
@@ -553,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

{pydocket-0.5.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