pydocket 0.5.2__py3-none-any.whl → 0.6.0__py3-none-any.whl

docket/__init__.py

@@ -13,11 +13,13 @@ from .dependencies import (
     CurrentDocket,
     CurrentExecution,
     CurrentWorker,
+    Depends,
     ExponentialRetry,
     Perpetual,
     Retry,
     TaskKey,
     TaskLogger,
+    Timeout,
 )
 from .docket import Docket
 from .execution import Execution
@@ -36,5 +38,7 @@ __all__ = [
     "ExponentialRetry",
     "Logged",
     "Perpetual",
+    "Timeout",
+    "Depends",
     "__version__",
 ]

docket/annotations.py

@@ -28,3 +28,19 @@ class Annotation(abc.ABC):
 
 class Logged(Annotation):
     """Instructs docket to include arguments to this parameter in the log."""
+
+    length_only: bool = False
+
+    def __init__(self, length_only: bool = False) -> None:
+        self.length_only = length_only
+
+    def format(self, argument: Any) -> str:
+        if self.length_only:
+            if isinstance(argument, (dict, set)):
+                return f"{{len {len(argument)}}}"
+            elif isinstance(argument, tuple):
+                return f"(len {len(argument)})"
+            elif hasattr(argument, "__len__"):
+                return f"[len {len(argument)}]"
+
+        return repr(argument)

docket/dependencies.py

@@ -1,35 +1,60 @@
 import abc
-import inspect
 import logging
+import time
+from contextlib import AsyncExitStack, asynccontextmanager
+from contextvars import ContextVar
 from datetime import timedelta
-from typing import Any, Awaitable, Callable, Counter, TypeVar, cast
+from types import TracebackType
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncContextManager,
+    AsyncGenerator,
+    Awaitable,
+    Callable,
+    Counter,
+    Generic,
+    TypeVar,
+    cast,
+)
 
 from .docket import Docket
-from .execution import Execution
-from .worker import Worker
+from .execution import Execution, TaskFunction, get_signature
+
+if TYPE_CHECKING:  # pragma: no cover
+    from .worker import Worker
 
 
 class Dependency(abc.ABC):
     single: bool = False
 
+    docket: ContextVar[Docket] = ContextVar("docket")
+    worker: ContextVar["Worker"] = ContextVar("worker")
+    execution: ContextVar[Execution] = ContextVar("execution")
+
     @abc.abstractmethod
-    def __call__(
-        self, docket: Docket, worker: Worker, execution: Execution
-    ) -> Any: ...  # pragma: no cover
+    async def __aenter__(self) -> Any: ...  # pragma: no cover
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> bool: ...  # pragma: no cover
 
 
 class _CurrentWorker(Dependency):
-    def __call__(self, docket: Docket, worker: Worker, execution: Execution) -> Worker:
-        return worker
+    async def __aenter__(self) -> "Worker":
+        return self.worker.get()
 
 
-def CurrentWorker() -> Worker:
-    return cast(Worker, _CurrentWorker())
+def CurrentWorker() -> "Worker":
+    return cast("Worker", _CurrentWorker())
 
 
 class _CurrentDocket(Dependency):
-    def __call__(self, docket: Docket, worker: Worker, execution: Execution) -> Docket:
-        return docket
+    async def __aenter__(self) -> Docket:
+        return self.docket.get()
 
 
 def CurrentDocket() -> Docket:
@@ -37,10 +62,8 @@ def CurrentDocket() -> Docket:
 
 
 class _CurrentExecution(Dependency):
-    def __call__(
-        self, docket: Docket, worker: Worker, execution: Execution
-    ) -> Execution:
-        return execution
+    async def __aenter__(self) -> Execution:
+        return self.execution.get()
 
 
 def CurrentExecution() -> Execution:
@@ -48,8 +71,8 @@ def CurrentExecution() -> Execution:
 
 
 class _TaskKey(Dependency):
-    def __call__(self, docket: Docket, worker: Worker, execution: Execution) -> str:
-        return execution.key
+    async def __aenter__(self) -> str:
+        return self.execution.get().key
 
 
 def TaskKey() -> str:
@@ -57,15 +80,14 @@ def TaskKey() -> str:
 
 
 class _TaskLogger(Dependency):
-    def __call__(
-        self, docket: Docket, worker: Worker, execution: Execution
-    ) -> logging.LoggerAdapter[logging.Logger]:
+    async def __aenter__(self) -> logging.LoggerAdapter[logging.Logger]:
+        execution = self.execution.get()
         logger = logging.getLogger(f"docket.task.{execution.function.__name__}")
         return logging.LoggerAdapter(
             logger,
             {
-                **docket.labels(),
-                **worker.labels(),
+                **self.docket.get().labels(),
+                **self.worker.get().labels(),
                 **execution.specific_labels(),
             },
         )
@@ -85,7 +107,8 @@ class Retry(Dependency):
         self.delay = delay
         self.attempt = 1
 
-    def __call__(self, docket: Docket, worker: Worker, execution: Execution) -> "Retry":
+    async def __aenter__(self) -> "Retry":
+        execution = self.execution.get()
         retry = Retry(attempts=self.attempts, delay=self.delay)
         retry.attempt = execution.attempt
         return retry
@@ -104,9 +127,9 @@ class ExponentialRetry(Retry):
         self.minimum_delay = minimum_delay
         self.maximum_delay = maximum_delay
 
-    def __call__(
-        self, docket: Docket, worker: Worker, execution: Execution
-    ) -> "ExponentialRetry":
+    async def __aenter__(self) -> "ExponentialRetry":
+        execution = self.execution.get()
+
         retry = ExponentialRetry(
             attempts=self.attempts,
             minimum_delay=self.minimum_delay,
@@ -155,9 +178,8 @@ class Perpetual(Dependency):
         self.automatic = automatic
         self.cancelled = False
 
-    def __call__(
-        self, docket: Docket, worker: Worker, execution: Execution
-    ) -> "Perpetual":
+    async def __aenter__(self) -> "Perpetual":
+        execution = self.execution.get()
         perpetual = Perpetual(every=self.every)
         perpetual.args = execution.args
         perpetual.kwargs = execution.kwargs
@@ -171,27 +193,121 @@ class Perpetual(Dependency):
         self.kwargs = kwargs
 
 
+class Timeout(Dependency):
+    single = True
+
+    base: timedelta
+
+    _deadline: float
+
+    def __init__(self, base: timedelta) -> None:
+        self.base = base
+
+    async def __aenter__(self) -> "Timeout":
+        timeout = Timeout(base=self.base)
+        timeout.start()
+        return timeout
+
+    def start(self) -> None:
+        self._deadline = time.monotonic() + self.base.total_seconds()
+
+    def expired(self) -> bool:
+        return time.monotonic() >= self._deadline
+
+    def remaining(self) -> timedelta:
+        return timedelta(seconds=self._deadline - time.monotonic())
+
+    def extend(self, by: timedelta | None = None) -> None:
+        if by is None:
+            by = self.base
+        self._deadline += by.total_seconds()
+
+
+R = TypeVar("R")
+
+DependencyFunction = Callable[..., Awaitable[R] | AsyncContextManager[R]]
+
+
+_parameter_cache: dict[
+    TaskFunction | DependencyFunction[Any],
+    dict[str, Dependency],
+] = {}
+
+
 def get_dependency_parameters(
-    function: Callable[..., Awaitable[Any]],
+    function: TaskFunction | DependencyFunction[Any],
 ) -> dict[str, Dependency]:
-    dependencies: dict[str, Any] = {}
+    if function in _parameter_cache:
+        return _parameter_cache[function]
+
+    dependencies: dict[str, Dependency] = {}
 
-    signature = inspect.signature(function)
+    signature = get_signature(function)
 
-    for param_name, param in signature.parameters.items():
+    for parameter, param in signature.parameters.items():
         if not isinstance(param.default, Dependency):
             continue
 
-        dependencies[param_name] = param.default
+        dependencies[parameter] = param.default
 
+    _parameter_cache[function] = dependencies
     return dependencies
 
 
+class _Depends(Dependency, Generic[R]):
+    dependency: DependencyFunction[R]
+
+    cache: ContextVar[dict[DependencyFunction[Any], Any]] = ContextVar("cache")
+    stack: ContextVar[AsyncExitStack] = ContextVar("stack")
+
+    def __init__(
+        self, dependency: Callable[[], Awaitable[R] | AsyncContextManager[R]]
+    ) -> None:
+        self.dependency = dependency
+
+    async def _resolve_parameters(
+        self,
+        function: TaskFunction | DependencyFunction[Any],
+    ) -> dict[str, Any]:
+        stack = self.stack.get()
+
+        arguments: dict[str, Any] = {}
+        parameters = get_dependency_parameters(function)
+
+        for parameter, dependency in parameters.items():
+            arguments[parameter] = await stack.enter_async_context(dependency)
+
+        return arguments
+
+    async def __aenter__(self) -> R:
+        cache = self.cache.get()
+
+        if self.dependency in cache:
+            return cache[self.dependency]
+
+        stack = self.stack.get()
+        arguments = await self._resolve_parameters(self.dependency)
+
+        value = self.dependency(**arguments)
+
+        if isinstance(value, AsyncContextManager):
+            value = await stack.enter_async_context(value)
+        else:
+            value = await value
+
+        cache[self.dependency] = value
+        return value
+
+
+def Depends(dependency: DependencyFunction[R]) -> R:
+    return cast(R, _Depends(dependency))
+
+
 D = TypeVar("D", bound=Dependency)
 
 
 def get_single_dependency_parameter_of_type(
-    function: Callable[..., Awaitable[Any]], dependency_type: type[D]
+    function: TaskFunction, dependency_type: type[D]
 ) -> D | None:
     assert dependency_type.single, "Dependency must be single"
     for _, dependency in get_dependency_parameters(function).items():
@@ -210,7 +326,7 @@ def get_single_dependency_of_type(
     return None
 
 
-def validate_dependencies(function: Callable[..., Awaitable[Any]]) -> None:
+def validate_dependencies(function: TaskFunction) -> None:
     parameters = get_dependency_parameters(function)
 
     counts = Counter(type(dependency) for dependency in parameters.values())
@@ -220,3 +336,31 @@ def validate_dependencies(function: Callable[..., Awaitable[Any]]) -> None:
             raise ValueError(
                 f"Only one {dependency_type.__name__} dependency is allowed per task"
             )
+
+
+@asynccontextmanager
+async def resolved_dependencies(
+    worker: "Worker", execution: Execution
+) -> AsyncGenerator[dict[str, Any], None]:
+    # Set context variables once at the beginning
+    Dependency.docket.set(worker.docket)
+    Dependency.worker.set(worker)
+    Dependency.execution.set(execution)
+
+    _Depends.cache.set({})
+
+    async with AsyncExitStack() as stack:
+        _Depends.stack.set(stack)
+
+        arguments: dict[str, Any] = {}
+
+        parameters = get_dependency_parameters(execution.function)
+        for parameter, dependency in parameters.items():
+            kwargs = execution.kwargs
+            if parameter in kwargs:
+                arguments[parameter] = kwargs[parameter]
+                continue
+
+            arguments[parameter] = await stack.enter_async_context(dependency)
+
+        yield arguments

docket/docket.py

@@ -38,6 +38,7 @@ from .execution import (
     Strike,
     StrikeInstruction,
     StrikeList,
+    TaskFunction,
 )
 from .instrumentation import (
     REDIS_DISRUPTIONS,
@@ -57,7 +58,7 @@ tracer: trace.Tracer = trace.get_tracer(__name__)
 P = ParamSpec("P")
 R = TypeVar("R")
 
-TaskCollection = Iterable[Callable[..., Awaitable[Any]]]
+TaskCollection = Iterable[TaskFunction]
 
 RedisStreamID = bytes
 RedisMessageID = bytes
@@ -91,7 +92,7 @@ class RunningExecution(Execution):
         worker: str,
         started: datetime,
     ) -> None:
-        self.function: Callable[..., Awaitable[Any]] = execution.function
+        self.function: TaskFunction = execution.function
         self.args: tuple[Any, ...] = execution.args
         self.kwargs: dict[str, Any] = execution.kwargs
         self.when: datetime = execution.when
@@ -111,7 +112,7 @@ class DocketSnapshot:
 
 
 class Docket:
-    tasks: dict[str, Callable[..., Awaitable[Any]]]
+    tasks: dict[str, TaskFunction]
     strike_list: StrikeList
 
     _monitor_strikes_task: asyncio.Task[None]
@@ -197,7 +198,7 @@ class Docket:
         finally:
             await asyncio.shield(r.__aexit__(None, None, None))
 
-    def register(self, function: Callable[..., Awaitable[Any]]) -> None:
+    def register(self, function: TaskFunction) -> None:
         from .dependencies import validate_dependencies
 
         validate_dependencies(function)

docket/execution.py

@@ -7,23 +7,40 @@ from typing import Any, Awaitable, Callable, Hashable, Literal, Mapping, Self, c
 
 import cloudpickle  # type: ignore[import]
 
+from opentelemetry import propagate
+import opentelemetry.context
 
 from .annotations import Logged
+from docket.instrumentation import message_getter
 
 logger: logging.Logger = logging.getLogger(__name__)
 
+TaskFunction = Callable[..., Awaitable[Any]]
 Message = dict[bytes, bytes]
 
 
+_signature_cache: dict[Callable[..., Any], inspect.Signature] = {}
+
+
+def get_signature(function: Callable[..., Any]) -> inspect.Signature:
+    if function in _signature_cache:
+        return _signature_cache[function]
+
+    signature = inspect.signature(function)
+    _signature_cache[function] = signature
+    return signature
+
+
 class Execution:
     def __init__(
         self,
-        function: Callable[..., Awaitable[Any]],
+        function: TaskFunction,
         args: tuple[Any, ...],
         kwargs: dict[str, Any],
         when: datetime,
         key: str,
         attempt: int,
+        trace_context: opentelemetry.context.Context | None = None,
     ) -> None:
         self.function = function
         self.args = args
@@ -31,6 +48,7 @@ class Execution:
         self.when = when
         self.key = key
         self.attempt = attempt
+        self.trace_context = trace_context
 
     def as_message(self) -> Message:
         return {
@@ -43,9 +61,7 @@ class Execution:
         }
 
     @classmethod
-    def from_message(
-        cls, function: Callable[..., Awaitable[Any]], message: Message
-    ) -> Self:
+    def from_message(cls, function: TaskFunction, message: Message) -> Self:
         return cls(
             function=function,
             args=cloudpickle.loads(message[b"args"]),
@@ -53,6 +69,7 @@ class Execution:
             when=datetime.fromisoformat(message[b"when"].decode()),
             key=message[b"key"].decode(),
             attempt=int(message[b"attempt"].decode()),
+            trace_context=propagate.extract(message, getter=message_getter),
         )
 
     def general_labels(self) -> Mapping[str, str]:
@@ -68,7 +85,7 @@ class Execution:
 
     def call_repr(self) -> str:
         arguments: list[str] = []
-        signature = inspect.signature(self.function)
+        signature = get_signature(self.function)
         function_name = self.function.__name__
 
         logged_parameters = Logged.annotated_parameters(signature)
@@ -77,14 +94,14 @@ class Execution:
 
         for i, argument in enumerate(self.args[: len(parameter_names)]):
             parameter_name = parameter_names[i]
-            if parameter_name in logged_parameters:
-                arguments.append(repr(argument))
+            if logged := logged_parameters.get(parameter_name):
+                arguments.append(logged.format(argument))
             else:
                 arguments.append("...")
 
         for parameter_name, argument in self.kwargs.items():
-            if parameter_name in logged_parameters:
-                arguments.append(f"{parameter_name}={repr(argument)}")
+            if logged := logged_parameters.get(parameter_name):
+                arguments.append(f"{parameter_name}={logged.format(argument)}")
             else:
                 arguments.append(f"{parameter_name}=...")
 
@@ -217,10 +234,10 @@ class StrikeList:
         if function_name in self.task_strikes and not task_strikes:
             return True
 
-        sig = inspect.signature(execution.function)
+        signature = get_signature(execution.function)
 
         try:
-            bound_args = sig.bind(*execution.args, **execution.kwargs)
+            bound_args = signature.bind(*execution.args, **execution.kwargs)
             bound_args.apply_defaults()
         except TypeError:
             # If we can't make sense of the arguments, just assume the task is fine
@@ -265,6 +282,8 @@ class StrikeList:
                 case "between":  # pragma: no branch
                     lower, upper = strike_value
                     return lower <= value <= upper
+                case _:  # pragma: no cover
+                    raise ValueError(f"Unknown operator: {operator}")
         except (ValueError, TypeError):
             # If we can't make the comparison due to incompatible types, just log the
             # error and assume the task is not stricken

docket/worker.py

@@ -1,11 +1,10 @@
 import asyncio
-import inspect
 import logging
 import sys
 from datetime import datetime, timedelta, timezone
 from types import TracebackType
 from typing import (
-    TYPE_CHECKING,
+    Coroutine,
     Mapping,
     Protocol,
     Self,
@@ -13,18 +12,27 @@ from typing import (
 )
 from uuid import uuid4
 
-import redis.exceptions
-from opentelemetry import propagate, trace
+from opentelemetry import trace
 from opentelemetry.trace import Tracer
 from redis.asyncio import Redis
-from redis.exceptions import LockError
-
+from redis.exceptions import ConnectionError, LockError
+
+from docket.execution import get_signature
+
+from .dependencies import (
+    Dependency,
+    Perpetual,
+    Retry,
+    Timeout,
+    get_single_dependency_of_type,
+    get_single_dependency_parameter_of_type,
+    resolved_dependencies,
+)
 from .docket import (
     Docket,
     Execution,
     RedisMessage,
     RedisMessageID,
-    RedisMessages,
     RedisReadGroupResponse,
 )
 from .instrumentation import (
@@ -41,7 +49,6 @@ from .instrumentation import (
     TASKS_STARTED,
     TASKS_STRICKEN,
     TASKS_SUCCEEDED,
-    message_getter,
     metrics_server,
 )
 
@@ -49,10 +56,6 @@ logger: logging.Logger = logging.getLogger(__name__)
 tracer: Tracer = trace.get_tracer(__name__)
 
 
-if TYPE_CHECKING:  # pragma: no cover
-    from .dependencies import Dependency
-
-
 class _stream_due_tasks(Protocol):
     async def __call__(
         self, keys: list[str], args: list[str | float]
@@ -75,7 +78,7 @@ class Worker:
         concurrency: int = 10,
         redelivery_timeout: timedelta = timedelta(minutes=5),
         reconnection_delay: timedelta = timedelta(seconds=5),
-        minimum_check_interval: timedelta = timedelta(milliseconds=100),
+        minimum_check_interval: timedelta = timedelta(milliseconds=250),
         scheduling_resolution: timedelta = timedelta(milliseconds=250),
     ) -> None:
         self.docket = docket
@@ -197,13 +200,14 @@ class Worker:
     async def _run(self, forever: bool = False) -> None:
         logger.info("Starting worker %r with the following tasks:", self.name)
         for task_name, task in self.docket.tasks.items():
-            signature = inspect.signature(task)
+            signature = get_signature(task)
             logger.info("* %s%s", task_name, signature)
 
         while True:
             try:
-                return await self._worker_loop(forever=forever)
-            except redis.exceptions.ConnectionError:
+                async with self.docket.redis() as redis:
+                    return await self._worker_loop(redis, forever=forever)
+            except ConnectionError:
                 REDIS_DISRUPTIONS.add(1, self.labels())
                 logger.warning(
                     "Error connecting to redis, retrying in %s...",
@@ -212,123 +216,133 @@ class Worker:
                 )
                 await asyncio.sleep(self.reconnection_delay.total_seconds())
 
-    async def _worker_loop(self, forever: bool = False):
+    async def _worker_loop(self, redis: Redis, forever: bool = False):
         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, worker_stopping)
+        scheduler_task = asyncio.create_task(
+            self._scheduler_loop(redis, worker_stopping)
+        )
+
+        active_tasks: dict[asyncio.Task[None], RedisMessageID] = {}
+        available_slots = self.concurrency
+
+        async def check_for_work() -> bool:
+            logger.debug("Checking for work", extra=self._log_context())
+            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 get_redeliveries(redis: Redis) -> RedisReadGroupResponse:
+            logger.debug("Getting redeliveries", extra=self._log_context())
+            _, redeliveries, *_ = await redis.xautoclaim(
+                name=self.docket.stream_key,
+                groupname=self.docket.worker_group_name,
+                consumername=self.name,
+                min_idle_time=int(self.redelivery_timeout.total_seconds() * 1000),
+                start_id="0-0",
+                count=available_slots,
+            )
+            return [(b"__redelivery__", redeliveries)]
+
+        async def get_new_deliveries(redis: Redis) -> RedisReadGroupResponse:
+            logger.debug("Getting new deliveries", extra=self._log_context())
+            return await redis.xreadgroup(
+                groupname=self.docket.worker_group_name,
+                consumername=self.name,
+                streams={self.docket.stream_key: ">"},
+                block=int(self.minimum_check_interval.total_seconds() * 1000),
+                count=available_slots,
             )
-            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:
-                    message_id = active_tasks.pop(task)
-
-                    await task
-
-                    async with redis.pipeline() as pipeline:
-                        pipeline.xack(
-                            self.docket.stream_key,
-                            self.docket.worker_group_name,
-                            message_id,
-                        )
-                        pipeline.xdel(
-                            self.docket.stream_key,
-                            message_id,
-                        )
-                        await pipeline.execute()
 
-            has_work: bool = True
+        def start_task(message_id: RedisMessageID, message: RedisMessage) -> bool:
+            if not message:  # pragma: no cover
+                return False
 
-            if not forever:  # pragma: no branch
-                has_work = await check_for_work()
+            function_name = message[b"function"].decode()
+            if not (function := self.docket.tasks.get(function_name)):
+                logger.warning(
+                    "Task function %r not found",
+                    function_name,
+                    extra=self._log_context(),
+                )
+                return False
 
-            try:
-                while forever or has_work or active_tasks:
-                    await process_completed_tasks()
+            execution = Execution.from_message(function, message)
 
-                    available_slots = self.concurrency - len(active_tasks)
+            task = asyncio.create_task(self._execute(execution))
+            active_tasks[task] = message_id
 
-                    def start_task(
-                        message_id: RedisMessageID, message: RedisMessage
-                    ) -> None:
-                        if not message:  # pragma: no cover
-                            return
+            nonlocal available_slots
+            available_slots -= 1
 
-                        task = asyncio.create_task(self._execute(message))
-                        active_tasks[task] = message_id
+            return True
 
-                        nonlocal available_slots
-                        available_slots -= 1
+        async def ack_message(redis: Redis, message_id: RedisMessageID) -> None:
+            logger.debug("Acknowledging message", extra=self._log_context())
+            async with redis.pipeline() as pipeline:
+                pipeline.xack(
+                    self.docket.stream_key,
+                    self.docket.worker_group_name,
+                    message_id,
+                )
+                pipeline.xdel(
+                    self.docket.stream_key,
+                    message_id,
+                )
+                await pipeline.execute()
 
-                    if available_slots <= 0:
-                        await asyncio.sleep(self.minimum_check_interval.total_seconds())
-                        continue
-
-                    redeliveries: RedisMessages
-                    _, redeliveries, *_ = await redis.xautoclaim(
-                        name=self.docket.stream_key,
-                        groupname=self.docket.worker_group_name,
-                        consumername=self.name,
-                        min_idle_time=int(
-                            self.redelivery_timeout.total_seconds() * 1000
-                        ),
-                        start_id="0-0",
-                        count=available_slots,
-                    )
+        async def process_completed_tasks() -> None:
+            completed_tasks = {task for task in active_tasks if task.done()}
+            for task in completed_tasks:
+                message_id = active_tasks.pop(task)
+                await task
+                await ack_message(redis, message_id)
 
-                    for message_id, message in redeliveries:
-                        start_task(message_id, message)
+        has_work: bool = True
 
-                    if available_slots <= 0:
-                        continue
-
-                    new_deliveries: RedisReadGroupResponse = await redis.xreadgroup(
-                        groupname=self.docket.worker_group_name,
-                        consumername=self.name,
-                        streams={self.docket.stream_key: ">"},
-                        block=(
-                            int(self.minimum_check_interval.total_seconds() * 1000)
-                            if forever or active_tasks
-                            else None
-                        ),
-                        count=available_slots,
-                    )
+        try:
+            while forever or has_work or active_tasks:
+                await process_completed_tasks()
+
+                available_slots = self.concurrency - len(active_tasks)
 
-                    for _, messages in new_deliveries:
+                if available_slots <= 0:
+                    await asyncio.sleep(self.minimum_check_interval.total_seconds())
+                    continue
+
+                for source in [get_redeliveries, get_new_deliveries]:
+                    for _, messages in await source(redis):
                         for message_id, message in messages:
-                            start_task(message_id, message)
+                            if not start_task(message_id, message):
+                                await self._delete_known_task(redis, message)
+                                await ack_message(redis, message_id)
 
-                    if not forever and not active_tasks and not new_deliveries:
-                        has_work = await check_for_work()
+                    if available_slots <= 0:
+                        break
 
-            except asyncio.CancelledError:
-                if active_tasks:  # pragma: no cover
-                    logger.info(
-                        "Shutdown requested, finishing %d active tasks...",
-                        len(active_tasks),
-                        extra=self._log_context(),
-                    )
-            finally:
-                if active_tasks:
-                    await asyncio.gather(*active_tasks, return_exceptions=True)
-                    await process_completed_tasks()
+                if not forever and not active_tasks:
+                    has_work = await check_for_work()
 
-                worker_stopping.set()
-                await scheduler_task
+        except asyncio.CancelledError:
+            if active_tasks:  # pragma: no cover
+                logger.info(
+                    "Shutdown requested, finishing %d active tasks...",
+                    len(active_tasks),
+                    extra=self._log_context(),
+                )
+        finally:
+            if active_tasks:
+                await asyncio.gather(*active_tasks, return_exceptions=True)
+                await process_completed_tasks()
+
+            worker_stopping.set()
+            await scheduler_task
 
     async def _scheduler_loop(
         self,
@@ -389,6 +403,7 @@ class Worker:
 
         while not worker_stopping.is_set() or total_work:
             try:
+                logger.debug("Scheduling due tasks", extra=self._log_context())
                 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],
@@ -415,8 +430,6 @@ 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(
@@ -438,25 +451,22 @@ class Worker:
             except LockError:  # pragma: no cover
                 return
 
-    async def _execute(self, message: RedisMessage) -> None:
-        key = message[b"key"].decode()
-
-        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
-            )
+    async def _delete_known_task(
+        self, redis: Redis, execution_or_message: Execution | RedisMessage
+    ) -> None:
+        if isinstance(execution_or_message, Execution):
+            key = execution_or_message.key
+        elif bytes_key := execution_or_message.get(b"key"):
+            key = bytes_key.decode()
+        else:  # pragma: no cover
             return
 
-        execution = Execution.from_message(function, message)
-        dependencies = self._get_dependencies(execution)
+        logger.debug("Deleting known task", extra=self._log_context())
+        known_task_key = self.docket.known_task_key(key)
+        await redis.delete(known_task_key)
 
-        log_context = {**log_context, **execution.specific_labels()}
+    async def _execute(self, execution: Execution) -> None:
+        log_context = {**self._log_context(), **execution.specific_labels()}
         counter_labels = {**self.labels(), **execution.general_labels()}
 
         arrow = "↬" if execution.attempt > 1 else "↪"
@@ -464,7 +474,7 @@ class Worker:
 
         if self.docket.strike_list.is_stricken(execution):
             async with self.docket.redis() as redis:
-                await redis.delete(self.docket.known_task_key(key))
+                await self._delete_known_task(redis, execution)
 
             arrow = "🗙"
             logger.warning("%s %s", arrow, call, extra=log_context)
@@ -474,20 +484,16 @@ class Worker:
         if execution.key in self._execution_counts:
             self._execution_counts[execution.key] += 1
 
-        # 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()
+        initiating_span = trace.get_current_span(execution.trace_context)
+        initiating_context = initiating_span.get_span_context()
         links = [trace.Link(initiating_context)] if initiating_context.is_valid else []
 
         start = datetime.now(timezone.utc)
         punctuality = start - execution.when
-        log_context = {**log_context, "punctuality": punctuality.total_seconds()}
+        log_context = {
+            **log_context,
+            "punctuality": punctuality.total_seconds(),
+        }
         duration = timedelta(0)
 
         TASKS_STARTED.add(1, counter_labels)
@@ -496,77 +502,103 @@ class Worker:
 
         logger.info("%s [%s] %s", arrow, punctuality, call, extra=log_context)
 
-        try:
-            with tracer.start_as_current_span(
-                execution.function.__name__,
-                kind=trace.SpanKind.CONSUMER,
-                attributes={
-                    **self.labels(),
-                    **execution.specific_labels(),
-                    "code.function.name": execution.function.__name__,
-                },
-                links=links,
-            ):
-                await execution.function(
-                    *execution.args,
-                    **{
-                        **execution.kwargs,
-                        **dependencies,
-                    },
+        with tracer.start_as_current_span(
+            execution.function.__name__,
+            kind=trace.SpanKind.CONSUMER,
+            attributes={
+                **self.labels(),
+                **execution.specific_labels(),
+                "code.function.name": execution.function.__name__,
+            },
+            links=links,
+        ):
+            async with resolved_dependencies(self, execution) as dependencies:
+                # 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 self._delete_known_task(redis, execution)
+
+                try:
+                    if timeout := get_single_dependency_of_type(dependencies, Timeout):
+                        await self._run_function_with_timeout(
+                            execution, dependencies, timeout
+                        )
+                    else:
+                        await execution.function(
+                            *execution.args,
+                            **{
+                                **execution.kwargs,
+                                **dependencies,
+                            },
+                        )
 
-            TASKS_SUCCEEDED.add(1, counter_labels)
-            duration = datetime.now(timezone.utc) - start
-            log_context["duration"] = duration.total_seconds()
-            rescheduled = await self._perpetuate_if_requested(
-                execution, dependencies, duration
-            )
-            arrow = "↫" if rescheduled else "↩"
-            logger.info("%s [%s] %s", arrow, duration, call, extra=log_context)
-        except Exception:
-            TASKS_FAILED.add(1, counter_labels)
-            duration = datetime.now(timezone.utc) - start
-            log_context["duration"] = duration.total_seconds()
-            retried = await self._retry_if_requested(execution, dependencies)
-            if not retried:
-                retried = await self._perpetuate_if_requested(
-                    execution, dependencies, duration
-                )
-            arrow = "↫" if retried else "↩"
-            logger.exception("%s [%s] %s", arrow, duration, call, extra=log_context)
-        finally:
-            TASKS_RUNNING.add(-1, counter_labels)
-            TASKS_COMPLETED.add(1, counter_labels)
-            TASK_DURATION.record(duration.total_seconds(), counter_labels)
+                    TASKS_SUCCEEDED.add(1, counter_labels)
+                    duration = datetime.now(timezone.utc) - start
+                    log_context["duration"] = duration.total_seconds()
+                    rescheduled = await self._perpetuate_if_requested(
+                        execution, dependencies, duration
+                    )
+                    arrow = "↫" if rescheduled else "↩"
+                    logger.info("%s [%s] %s", arrow, duration, call, extra=log_context)
+                except Exception:
+                    TASKS_FAILED.add(1, counter_labels)
+                    duration = datetime.now(timezone.utc) - start
+                    log_context["duration"] = duration.total_seconds()
+                    retried = await self._retry_if_requested(execution, dependencies)
+                    if not retried:
+                        retried = await self._perpetuate_if_requested(
+                            execution, dependencies, duration
+                        )
+                    arrow = "↫" if retried else "↩"
+                    logger.exception(
+                        "%s [%s] %s", arrow, duration, call, extra=log_context
+                    )
+                finally:
+                    TASKS_RUNNING.add(-1, counter_labels)
+                    TASKS_COMPLETED.add(1, counter_labels)
+                    TASK_DURATION.record(duration.total_seconds(), counter_labels)
 
-    def _get_dependencies(
+    async def _run_function_with_timeout(
         self,
         execution: Execution,
-    ) -> dict[str, "Dependency"]:
-        from .dependencies import get_dependency_parameters
-
-        parameters = get_dependency_parameters(execution.function)
-
-        dependencies: dict[str, "Dependency"] = {}
-
-        for parameter_name, dependency in parameters.items():
-            # If the argument is already provided, skip it, which allows users to call
-            # the function directly with the arguments they want.
-            if parameter_name in execution.kwargs:
-                dependencies[parameter_name] = execution.kwargs[parameter_name]
-                continue
-
-            dependencies[parameter_name] = dependency(self.docket, self, execution)
+        dependencies: dict[str, Dependency],
+        timeout: Timeout,
+    ) -> None:
+        task_coro = cast(
+            Coroutine[None, None, None],
+            execution.function(*execution.args, **execution.kwargs, **dependencies),
+        )
+        task = asyncio.create_task(task_coro)
+        try:
+            while not task.done():  # pragma: no branch
+                remaining = timeout.remaining().total_seconds()
+                if timeout.expired():
+                    task.cancel()
+                    break
+
+                try:
+                    await asyncio.wait_for(asyncio.shield(task), timeout=remaining)
+                    return
+                except asyncio.TimeoutError:
+                    continue
+        finally:
+            if not task.done():
+                task.cancel()
 
-        return dependencies
+            try:
+                await task
+            except asyncio.CancelledError:
+                raise asyncio.TimeoutError
 
     async def _retry_if_requested(
         self,
         execution: Execution,
-        dependencies: dict[str, "Dependency"],
+        dependencies: dict[str, Dependency],
     ) -> bool:
-        from .dependencies import Retry, get_single_dependency_of_type
-
         retry = get_single_dependency_of_type(dependencies, Retry)
         if not retry:
             return False
@@ -584,11 +616,9 @@ class Worker:
     async def _perpetuate_if_requested(
         self,
         execution: Execution,
-        dependencies: dict[str, "Dependency"],
+        dependencies: dict[str, Dependency],
         duration: timedelta | None = None,
     ) -> bool:
-        from .dependencies import Perpetual, get_single_dependency_of_type
-
         perpetual = get_single_dependency_of_type(dependencies, Perpetual)
         if not perpetual:
             return False
@@ -667,7 +697,7 @@ class Worker:
 
             except asyncio.CancelledError:  # pragma: no cover
                 return
-            except redis.exceptions.ConnectionError:
+            except ConnectionError:
                 REDIS_DISRUPTIONS.add(1, self.labels())
                 logger.exception(
                     "Error sending worker heartbeat",

{pydocket-0.5.2.dist-info → pydocket-0.6.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydocket
-Version: 0.5.2
+Version: 0.6.0
 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.6.0.dist-info/RECORD

@@ -0,0 +1,16 @@
+docket/__init__.py,sha256=124XWbyQQHO1lhCoLQ-oheZnu4vNDHIaq4Whb7z3ogI,831
+docket/__main__.py,sha256=Vkuh7aJ-Bl7QVpVbbkUksAd_hn05FiLmWbc-8kbhZQ4,34
+docket/annotations.py,sha256=I00zB32BYWOQSNEjjCkc5n5DwTnT277I_BRYUJPS7w4,1474
+docket/cli.py,sha256=OWql6QFthSbvRCGkIg-ufo26F48z0eCmzRXJYOdyAEc,20309
+docket/dependencies.py,sha256=pkjseBZjdSpgW9g2H4cZ_RXIRZ2ZfdngBCXJGUcbmao,10052
+docket/docket.py,sha256=KJxgiyOskEHsRQOmfgLpJCYDNNleHI-vEKK3uBPL_K8,21420
+docket/execution.py,sha256=da1uYxSNAfz5FuNyCzX4I_PglHiMaf1oEv--K5TkjXc,13297
+docket/instrumentation.py,sha256=bZlGA02JoJcY0J1WGm5_qXDfY0AXKr0ZLAYu67wkeKY,4611
+docket/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+docket/tasks.py,sha256=RIlSM2omh-YDwVnCz6M5MtmK8T_m_s1w2OlRRxDUs6A,1437
+docket/worker.py,sha256=3sMcwGfSJ0Q4y5AuaqdgiGniDhJ21nM2PQmroJi_Q-A,26430
+pydocket-0.6.0.dist-info/METADATA,sha256=ktk1hqLmP_VSqYmdRtHFDPbEeRQD1J66ZAHEqaDXejk,13092
+pydocket-0.6.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+pydocket-0.6.0.dist-info/entry_points.txt,sha256=4WOk1nUlBsUT5O3RyMci2ImuC5XFswuopElYcLHtD5k,47
+pydocket-0.6.0.dist-info/licenses/LICENSE,sha256=YuVWU_ZXO0K_k2FG8xWKe5RGxV24AhJKTvQmKfqXuyk,1087
+pydocket-0.6.0.dist-info/RECORD,,

pydocket-0.5.2.dist-info/RECORD

@@ -1,16 +0,0 @@
-docket/__init__.py,sha256=7oruGALDoU6W_ntF-mMxxv3FFtO970DVzj3lUgoVIiM,775
-docket/__main__.py,sha256=Vkuh7aJ-Bl7QVpVbbkUksAd_hn05FiLmWbc-8kbhZQ4,34
-docket/annotations.py,sha256=GZwOPtPXyeIhnsLh3TQMBnXrjtTtSmF4Ratv4vjPx8U,950
-docket/cli.py,sha256=OWql6QFthSbvRCGkIg-ufo26F48z0eCmzRXJYOdyAEc,20309
-docket/dependencies.py,sha256=0P8GJTMWrzm9uZkQejCiRfT6IBisY7Hp1-4HAGTWv6w,6326
-docket/docket.py,sha256=p2G7QNn4H0sUhDlAI5BO5C6cRTy1ZWUZmFEuohX3RM8,21470
-docket/execution.py,sha256=PDrlAr8VzmB6JvqKO71YhXUcTcGQW7eyXrSKiTcAexE,12508
-docket/instrumentation.py,sha256=bZlGA02JoJcY0J1WGm5_qXDfY0AXKr0ZLAYu67wkeKY,4611
-docket/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-docket/tasks.py,sha256=RIlSM2omh-YDwVnCz6M5MtmK8T_m_s1w2OlRRxDUs6A,1437
-docket/worker.py,sha256=gqY_N7H9Jxh_0YIYQk0mucj_UrZNKItkT1xkuhwYmlY,25301
-pydocket-0.5.2.dist-info/METADATA,sha256=VbNbGmDdseQkzH64LFmsPNtw6kwbIc8cL73jlhS0vck,13092
-pydocket-0.5.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-pydocket-0.5.2.dist-info/entry_points.txt,sha256=4WOk1nUlBsUT5O3RyMci2ImuC5XFswuopElYcLHtD5k,47
-pydocket-0.5.2.dist-info/licenses/LICENSE,sha256=YuVWU_ZXO0K_k2FG8xWKe5RGxV24AhJKTvQmKfqXuyk,1087
-pydocket-0.5.2.dist-info/RECORD,,