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

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,17 +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 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 (
@@ -40,7 +49,6 @@ from .instrumentation import (
     TASKS_STARTED,
     TASKS_STRICKEN,
     TASKS_SUCCEEDED,
-    message_getter,
     metrics_server,
 )
 
@@ -48,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]
@@ -74,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
@@ -196,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...",
@@ -211,127 +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()
-        perpetual_scheduling_task = asyncio.create_task(
-            self._perpetual_scheduling_loop(worker_stopping)
+
+        scheduler_task = asyncio.create_task(
+            self._scheduler_loop(redis, worker_stopping)
         )
 
-        async with self.docket.redis() as redis:
-            scheduler_task = asyncio.create_task(
-                self._scheduler_loop(redis, worker_stopping)
+        active_tasks: dict[asyncio.Task[None], RedisMessageID] = {}
+        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)
+
+                if available_slots <= 0:
+                    await asyncio.sleep(self.minimum_check_interval.total_seconds())
+                    continue
 
-                    for _, messages in new_deliveries:
+                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
-                await perpetual_scheduling_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,
@@ -392,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],
@@ -417,73 +429,53 @@ 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
+    async def _schedule_all_automatic_perpetual_tasks(self) -> None:
+        async with self.docket.redis() as redis:
             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)
+                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
 
-        return minimum_interval
+                        if not perpetual.automatic:
+                            continue
 
-    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))
+                        key = task_function.__name__
 
-        log_context: Mapping[str, str | float] = self._log_context()
+                        await self.docket.add(task_function, key=key)()
+            except LockError:  # pragma: no cover
+                return
 
-        function_name = message[b"function"].decode()
-        function = self.docket.tasks.get(function_name)
-        if function is None:
-            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)
+        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 "↪"
         call = execution.call_repr()
 
         if self.docket.strike_list.is_stricken(execution):
+            async with self.docket.redis() as redis:
+                await self._delete_known_task(redis, execution)
+
             arrow = "🗙"
             logger.warning("%s %s", arrow, call, extra=log_context)
             TASKS_STRICKEN.add(1, counter_labels | {"docket.where": "worker"})
@@ -492,15 +484,16 @@ class Worker:
         if execution.key in self._execution_counts:
             self._execution_counts[execution.key] += 1
 
-        dependencies = self._get_dependencies(execution)
-
-        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)
@@ -509,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
@@ -597,26 +616,28 @@ class Worker:
     async def _perpetuate_if_requested(
         self,
         execution: Execution,
-        dependencies: dict[str, "Dependency"],
-        duration: timedelta,
+        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
 
         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
@@ -676,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",