pydocket 0.5.1__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

@@ -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,
@@ -37,6 +38,7 @@ from .execution import (
     Strike,
     StrikeInstruction,
     StrikeList,
+    TaskFunction,
 )
 from .instrumentation import (
     REDIS_DISRUPTIONS,
@@ -56,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
@@ -90,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
@@ -110,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]
@@ -196,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)
@@ -256,9 +258,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 +298,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 +354,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 +380,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:

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