async-task-api-worker 0.1.0rc2__py3-none-any.whl

async_task_api_worker/_runtime/lifecycle.py

@@ -0,0 +1,463 @@
+"""Process lifecycle: startup checks, signals, graceful drain and shutdown.
+
+Start: validate the settings and the time budget, run the startup hooks, connect (retrying with the reconnect
+interval while the broker is unreachable), consume.
+
+SIGTERM or SIGINT (first): stop consuming (``basic.cancel``), interrupt the retry backoffs, give the in-flight tasks
+``WORKER_SHUTDOWN_GRACE_S`` to finish (they publish and ack normally), then cancel the rest, which requeues their
+deliveries, close the broker and run the shutdown hooks in reverse order. A second signal skips what is left of the
+grace period.
+
+Job mode (``WORKER_MODE=job``) runs the same steps for a one-shot process: the prefetch is 1, consuming stops at the
+``WORKER_JOB_MAX_TASKS``-th delivery (so the broker hands nothing more), and the process drains and exits 0 once those
+deliveries are settled, or when nothing arrives for ``WORKER_JOB_IDLE_TIMEOUT_S``. A broker that stays unreachable for
+``WORKER_JOB_CONNECT_TIMEOUT_S`` ends the process with an error (exit 1) instead of retrying forever. Every outcome of
+a delivery (success, failure, requeue) is a clean exit: the exit code tells the Job controller whether the pod crashed.
+"""
+
+import asyncio
+import contextlib
+import logging
+import os
+import signal
+import sys
+import threading
+import time
+from collections.abc import Awaitable, Callable
+from concurrent.futures import ThreadPoolExecutor
+from enum import Enum
+from typing import TYPE_CHECKING, Final, NoReturn
+
+from async_task_api_worker._runtime.engine import Engine
+from async_task_api_worker._runtime.transport import CLOSE_TIMEOUT_S, Delivery, Transport, TransportUnavailableError
+from async_task_api_worker.errors import WorkerConfigurationError, WorkerDefinitionError
+from async_task_api_worker.observability.health import HealthServer, RuntimeState
+from async_task_api_worker.observability.logging import configure_logging
+from async_task_api_worker.observability.loop_lag import EventLoopLagProbe
+from async_task_api_worker.observability.metrics import WorkerMetrics
+from async_task_api_worker.settings import LogLevel, WorkerSettings
+
+if TYPE_CHECKING:
+    from async_task_api_worker.worker import Worker
+
+_LOGGER: Final[logging.Logger] = logging.getLogger("async_task_api_worker")
+
+EXIT_OK: Final[int] = 0
+EXIT_FAILURE: Final[int] = 1
+EXIT_CONFIGURATION: Final[int] = 2
+# What the broker's consumer timeout must exceed, on top of the task timeout and the longest retry backoff: the time
+# to publish the result and ack.
+ACK_SAFETY_MARGIN_S: Final[float] = 60.0
+# A nack on a half-open connection can block until the heartbeat timeout: do not wait longer than this for them.
+CANCEL_NACK_TIMEOUT_S: Final[float] = 5.0
+# The threads of ``asyncio.to_thread`` and ``run_in_executor(None, ...)``: the loop's default executor.
+DEFAULT_EXECUTOR_THREAD_PREFIX: Final[str] = "async-task-api-worker-default"
+# How long the exit waits for them: ``asyncio.run`` would wait 300 s.
+EXECUTOR_JOIN_TIMEOUT_S: Final[float] = 2.0
+_STOP_SIGNALS: Final[tuple[signal.Signals, ...]] = (signal.SIGTERM, signal.SIGINT)
+
+
+class LifecycleState(Enum):
+    STARTING = "starting"
+    READY = "ready"
+    DRAINING = "draining"
+    STOPPED = "stopped"
+
+
+def check_time_budget(settings: WorkerSettings, task_timeout_s: float) -> None:
+    """The broker takes a delivery back after its consumer timeout: a task must always conclude before that."""
+    needed = task_timeout_s + settings.worker_retry_backoff_max_s + ACK_SAFETY_MARGIN_S
+    if needed >= settings.worker_consumer_timeout_s:
+        msg = (
+            f"The time budget does not fit: task timeout {task_timeout_s:g} s + retry backoff "
+            f"{settings.worker_retry_backoff_max_s:g} s + safety margin {ACK_SAFETY_MARGIN_S:g} s = {needed:g} s "
+            f"must stay under the broker consumer timeout WORKER_CONSUMER_TIMEOUT_S="
+            f"{settings.worker_consumer_timeout_s:g} s"
+        )
+        raise WorkerConfigurationError(msg)
+
+
+class Lifecycle:
+    """Runs one worker on one transport until it is asked to stop."""
+
+    def __init__(
+        self,
+        worker: "Worker",
+        transport: Transport,
+        *,
+        install_signals: bool = True,
+        cancel_timeout_s: float = CANCEL_NACK_TIMEOUT_S,
+        close_timeout_s: float = CLOSE_TIMEOUT_S,
+        serve_probes: bool = False,
+    ) -> None:
+        self._worker = worker
+        self._transport = transport
+        self._install_signals = install_signals
+        self._cancel_timeout_s = cancel_timeout_s
+        self._close_timeout_s = close_timeout_s
+        self._serve_probes = serve_probes
+        self.metrics = WorkerMetrics()
+        self.overrunning: tuple[str, ...] = ()
+        self._stop = asyncio.Event()
+        self._skip_grace = asyncio.Event()
+        self._activity = asyncio.Event()
+        self._accepted = 0
+        self._stop_consuming_task: asyncio.Task[None] | None = None
+        self.state = LifecycleState.STARTING
+        self._engine: Engine | None = None
+
+    def request_stop(self) -> None:
+        """First call: drain. Any later call: skip what is left of the grace period."""
+        if self._stop.is_set():
+            self._skip_grace.set()
+        self._stop.set()
+
+    async def run(self) -> None:
+        """Start, serve until a stop is requested, drain, close. Raises on a configuration or startup failure."""
+        settings = self._worker.settings
+        definition = self._worker._require_definition()
+        timeout_s = (
+            settings.worker_task_timeout_s if settings.worker_task_timeout_s is not None else definition.timeout_s
+        )
+        check_time_budget(settings, timeout_s)
+
+        engine = self._worker._create_engine(self._transport.publisher, observer=self.metrics.observer)
+        self._engine = engine
+        self.metrics.track_in_flight(lambda: len(engine.inflight))
+        health = self._health_server(engine, settings)
+        loop_lag = EventLoopLagProbe(self.metrics.loop_lag.observe)
+        started_hooks = 0
+        loop = asyncio.get_running_loop()
+        handled = self._add_signal_handlers(loop)
+        try:
+            # Served from the very start, in both modes: ``/health`` answers while the startup hooks run and ``/ready``
+            # stays 503 until the broker is connected. A ScaledJob counts a pod only once it is Ready, so job mode
+            # keeps it.
+            if health is not None:
+                await health.start()
+                loop_lag.start()
+            for hook in self._worker._startup_hooks:
+                if not await self._until_stopped(hook()):
+                    break  # stopped during startup: nothing was consumed, drain and exit
+                started_hooks += 1
+            if started_hooks == len(self._worker._startup_hooks) and await self._connect(settings):
+                self.state = LifecycleState.READY
+                _LOGGER.info("Worker %s is consuming %s.", self._worker.service, settings.in_queue_name)
+                await (self._wait_job_done(engine, settings) if settings.is_job else self._stop.wait())
+        finally:
+            self.state = LifecycleState.DRAINING
+            try:
+                await self._drain(engine, settings)
+            finally:
+                for signum in handled:
+                    loop.remove_signal_handler(signum)
+                await self._run_shutdown_hooks(started_hooks)
+                self.state = LifecycleState.STOPPED
+                await loop_lag.stop()
+                if health is not None:
+                    await health.stop()
+
+    def _health_server(self, engine: Engine, settings: WorkerSettings) -> HealthServer | None:
+        if not self._serve_probes:
+            return None
+        state = RuntimeState(
+            ping_broker=self._transport.ping,
+            # Not ready once it no longer takes deliveries: shutting down, or a job whose quota is reached.
+            is_draining=lambda: (
+                self._stop.is_set()
+                or self._stop_consuming_task is not None
+                or self.state in {LifecycleState.DRAINING, LifecycleState.STOPPED}
+            ),
+            overrunning=lambda: tuple(
+                entry.task_id
+                for entry in engine.inflight.overrun(time.monotonic(), settings.worker_task_overrun_grace_s)
+            ),
+            metrics=self.metrics.exposition,
+        )
+        return HealthServer(
+            host=settings.health_check_host,
+            port=settings.health_check_port,
+            state=state,
+            probes=self._worker._probes,
+            probe_timeout_s=settings.worker_probe_timeout_s,
+        )
+
+    def _add_signal_handlers(self, loop: asyncio.AbstractEventLoop) -> list[signal.Signals]:
+        handled: list[signal.Signals] = []
+        if not self._install_signals:
+            return handled
+        for signum in _STOP_SIGNALS:
+            try:
+                loop.add_signal_handler(signum, self.request_stop)
+            except (NotImplementedError, RuntimeError, ValueError):
+                _LOGGER.warning("Cannot install a handler for %s: not in the main thread.", signum.name)
+            else:
+                handled.append(signum)
+        return handled
+
+    async def _connect(self, settings: WorkerSettings) -> bool:
+        """Start the transport, retrying while the broker is unreachable. False when a stop came first.
+
+        A job gives up after ``WORKER_JOB_CONNECT_TIMEOUT_S`` and re-raises the last error: the pod must fail, not wait.
+        """
+        engine = self._engine
+        assert engine is not None
+        loop = asyncio.get_running_loop()
+        give_up_at = loop.time() + settings.worker_job_connect_timeout_s
+        while not self._stop.is_set():
+            # Bound the attempt itself in job mode: a host that drops packets would otherwise hold it for the AMQP
+            # client's own connection timeout, far past the budget.
+            attempt_timeout_s = max(give_up_at - loop.time(), 0.0) if settings.is_job else None
+            try:
+                async with asyncio.timeout(attempt_timeout_s):
+                    if not await self._until_stopped(self._transport.start(self._on_delivery)):
+                        return False
+            except (TimeoutError, TransportUnavailableError) as error:
+                if settings.is_job and loop.time() >= give_up_at:
+                    limit = f"WORKER_JOB_CONNECT_TIMEOUT_S={settings.worker_job_connect_timeout_s:g} s"
+                    msg = f"The broker could not be reached within {limit}"
+                    raise TransportUnavailableError(msg) from error
+                _LOGGER.warning(
+                    "The broker is unreachable, retrying in %s s.",
+                    settings.worker_broker_reconnect_interval_s,
+                )
+                await self._wait_stop(settings.worker_broker_reconnect_interval_s)
+            else:
+                return True
+        return False
+
+    async def _until_stopped(self, work: Awaitable[object]) -> bool:
+        """Run ``work``, cancelling it if a stop is requested first. False when it was cancelled by the stop.
+
+        A SIGTERM during startup must not wait for a connection attempt (up to the AMQP client's own timeout) or a slow
+        startup hook: nothing was consumed yet, so the worker has nothing to drain. A hook that blocks the event loop
+        itself cannot be interrupted.
+        """
+        task = asyncio.ensure_future(work)
+        stop = asyncio.ensure_future(self._stop.wait())
+        try:
+            await asyncio.wait({task, stop}, return_when=asyncio.FIRST_COMPLETED)
+        finally:
+            stop.cancel()
+            if not task.done():
+                task.cancel()
+                await asyncio.wait({task})
+        if task.cancelled():
+            return False
+        task.result()
+        return True
+
+    def _on_delivery(self, delivery: Delivery) -> None:
+        assert self._engine is not None
+        self._accepted += 1
+        if self._worker.settings.is_job and self._accepted >= self._worker.settings.worker_job_max_tasks:
+            # The cancel is scheduled before the engine task, so it goes out before this delivery can be acked: with a
+            # prefetch of 1, that ack is what would let the broker hand out one more message.
+            self._begin_stop_consuming()
+        task = self._engine.spawn(delivery)
+        task.add_done_callback(lambda _: self._activity.set())
+        self._activity.set()
+
+    async def _wait_job_done(self, engine: Engine, settings: WorkerSettings) -> None:
+        """Return when the job is over: quota settled, idle for too long, or a stop was requested."""
+        while not self._stop.is_set():
+            # Cleared first: a wake-up that happens during the cancel below must not be lost.
+            self._activity.clear()
+            quota_reached = self._accepted >= settings.worker_job_max_tasks
+            if quota_reached:
+                await self._stop_consuming_once()
+            busy = bool(engine.tasks())
+            if quota_reached and not busy:
+                _LOGGER.info("Job done: %d delivery(ies) settled.", self._accepted)
+                return
+            # Only a worker with nothing in hand and quota left can be idle.
+            idle_timeout_s = None if busy or quota_reached else settings.worker_job_idle_timeout_s
+            if not await self._wait_activity(idle_timeout_s) and not (self._activity.is_set() or engine.tasks()):
+                _LOGGER.info("Job idle for %s s: no delivery, exiting.", settings.worker_job_idle_timeout_s)
+                return
+
+    async def _wait_activity(self, timeout_s: float | None) -> bool:
+        """Wait for an activity or a stop. False when ``timeout_s`` elapsed first."""
+        activity = asyncio.ensure_future(self._activity.wait())
+        stop = asyncio.ensure_future(self._stop.wait())
+        try:
+            done, _ = await asyncio.wait({activity, stop}, timeout=timeout_s, return_when=asyncio.FIRST_COMPLETED)
+        finally:
+            activity.cancel()
+            stop.cancel()
+        return bool(done)
+
+    def _begin_stop_consuming(self) -> asyncio.Task[None]:
+        if self._stop_consuming_task is None:
+            self._stop_consuming_task = asyncio.ensure_future(self._transport.stop_consuming())
+        return self._stop_consuming_task
+
+    async def _stop_consuming_once(self) -> None:
+        await self._begin_stop_consuming()
+
+    async def _wait_stop(self, seconds: float) -> None:
+        try:
+            await asyncio.wait_for(self._stop.wait(), seconds)
+        except TimeoutError:
+            return
+
+    async def _drain(self, engine: Engine, settings: WorkerSettings) -> None:
+        loop = asyncio.get_running_loop()
+        # The grace starts at the stop: the consumer cancel, which can wait for a reconnect, is part of it.
+        deadline = loop.time() + settings.worker_shutdown_grace_s
+        try:
+            await self._stop_consuming_once()
+        except Exception:
+            # The in-flight tasks still get their grace period, and the transport is still closed.
+            _LOGGER.exception("Stopping the consumer failed, draining anyway.")
+        engine.request_shutdown()
+        while pending := engine.tasks():
+            remaining = deadline - loop.time()
+            if remaining <= 0 or self._skip_grace.is_set():
+                break
+            skip = asyncio.ensure_future(self._skip_grace.wait())
+            try:
+                await asyncio.wait({*pending, skip}, timeout=remaining, return_when=asyncio.FIRST_COMPLETED)
+            finally:
+                skip.cancel()
+        leftover = engine.tasks()
+        if leftover:
+            _LOGGER.warning(
+                "%d task(s) did not finish within the grace period: requeueing their deliveries.",
+                len(leftover),
+            )
+            for task in leftover:
+                task.cancel()
+            _, unsettled = await asyncio.wait(leftover, timeout=self._cancel_timeout_s)
+            if unsettled:
+                _LOGGER.error(
+                    "%d delivery(ies) could not be requeued within %s s: the broker redelivers them when the "
+                    "connection closes.",
+                    len(unsettled),
+                    self._cancel_timeout_s,
+                )
+        await self._close_transport()
+        engine.close()
+        # Sync handlers whose thread is still running keep their registry entry: the process cannot join them.
+        self.overrunning = tuple(entry.task_id for entry in engine.inflight.snapshot())
+
+    async def _close_transport(self) -> None:
+        try:
+            async with asyncio.timeout(self._close_timeout_s):
+                await self._transport.close()
+        except TimeoutError:
+            _LOGGER.error("Closing the broker connection took more than %s s: giving up.", self._close_timeout_s)
+        except Exception:
+            _LOGGER.exception("Closing the broker connection failed.")
+
+    async def _run_shutdown_hooks(self, started_hooks: int) -> None:
+        """Run the shutdown hooks in reverse order, within one budget for all of them.
+
+        Each hook gets what is left of ``WORKER_SHUTDOWN_HOOKS_TIMEOUT_S``; one that runs out of it is cancelled,
+        logged, and the next one goes on. A per-hook timeout would make the bound depend on how many hooks there are.
+        A hook that is out of budget before it starts is skipped. A hook that blocks the event loop, or swallows the
+        cancellation, cannot be bounded.
+        """
+        if started_hooks < len(self._worker._startup_hooks):
+            return
+        loop = asyncio.get_running_loop()
+        deadline = loop.time() + self._worker.settings.worker_shutdown_hooks_timeout_s
+        for hook in reversed(self._worker._shutdown_hooks):
+            if loop.time() >= deadline:
+                _LOGGER.error("A shutdown hook was skipped: the shutdown hooks budget is spent.")
+                continue
+            try:
+                async with asyncio.timeout_at(deadline) as budget:
+                    await hook()
+            except TimeoutError:
+                if budget.expired():
+                    _LOGGER.error("A shutdown hook ran out of the shutdown hooks budget and was cancelled.")
+                else:
+                    _LOGGER.exception("A shutdown hook failed.")  # its own TimeoutError
+            except Exception:
+                _LOGGER.exception("A shutdown hook failed.")
+
+
+def _close_loop(loop: asyncio.AbstractEventLoop, executor: ThreadPoolExecutor, join_timeout_s: float) -> list[str]:
+    """Close the loop like ``asyncio.run`` would, but never wait more than ``join_timeout_s`` in all.
+
+    Returns what is still running after that: the names of the default-executor threads (an ``async def`` handler
+    cancelled while it waits in ``asyncio.to_thread`` leaves its thread behind, and ``asyncio.run`` would join it for up
+    to 300 s), and the tasks that ignored their cancellation.
+    """
+    deadline = time.monotonic() + join_timeout_s
+    stuck: list[str] = []
+    try:
+        pending = [task for task in asyncio.all_tasks(loop) if not task.done()]
+        for task in pending:
+            task.cancel()
+        if pending:
+            _, ignoring = loop.run_until_complete(asyncio.wait(pending, timeout=join_timeout_s))
+            stuck.extend(f"task {task.get_name()}" for task in ignoring)
+        with contextlib.suppress(TimeoutError):
+            loop.run_until_complete(
+                asyncio.wait_for(loop.shutdown_asyncgens(), max(deadline - time.monotonic(), 0.0)),
+            )
+    finally:
+        asyncio.set_event_loop(None)
+        loop.close()
+    executor.shutdown(wait=False, cancel_futures=True)
+    for thread in threading.enumerate():
+        if thread.name.startswith(DEFAULT_EXECUTOR_THREAD_PREFIX):
+            thread.join(max(deadline - time.monotonic(), 0.0))
+            if thread.is_alive():
+                stuck.append(thread.name)
+    return stuck
+
+
+def _leave(reason: str, exit_now: Callable[[int], NoReturn]) -> NoReturn:
+    _LOGGER.error(reason)
+    logging.shutdown()
+    sys.stdout.flush()
+    sys.stderr.flush()
+    exit_now(EXIT_FAILURE)
+
+
+def run_blocking(
+    worker: "Worker",
+    *,
+    exit_now: Callable[[int], NoReturn] = os._exit,
+    executor_join_timeout_s: float = EXECUTOR_JOIN_TIMEOUT_S,
+) -> int:
+    """Run the worker until it stops. Returns the process exit code.
+
+    A thread that never returns would make the interpreter wait for it at exit: a ``def`` handler's, or the thread of an
+    ``async def`` handler cancelled while it waited in ``asyncio.to_thread``. When some are left after the drain, they
+    are logged and the process exits at once through ``exit_now``, with code 1: the drain did not complete.
+    """
+    try:
+        settings = worker.settings
+    except WorkerConfigurationError as error:
+        configure_logging(LogLevel.INFO)
+        _LOGGER.error("Invalid worker configuration: %s", error)
+        return EXIT_CONFIGURATION
+    configure_logging(settings.log_level)
+    # The loop is driven by hand, with a default executor of our own whose threads can be recognised and given up on.
+    loop = asyncio.new_event_loop()
+    executor = ThreadPoolExecutor(thread_name_prefix=DEFAULT_EXECUTOR_THREAD_PREFIX)
+    loop.set_default_executor(executor)
+    asyncio.set_event_loop(loop)
+    overrunning: tuple[str, ...] = ()
+    code = EXIT_OK
+    try:
+        overrunning = loop.run_until_complete(worker._serve_report())
+    except (WorkerConfigurationError, WorkerDefinitionError) as error:
+        _LOGGER.error("Invalid worker configuration: %s", error)
+        code = EXIT_CONFIGURATION
+    except TransportUnavailableError as error:
+        _LOGGER.error("The broker stayed unreachable, giving up: %s", error)
+        code = EXIT_FAILURE
+    except Exception:
+        _LOGGER.exception("The worker stopped on an error.")
+        code = EXIT_FAILURE
+    stuck_threads = _close_loop(loop, executor, executor_join_timeout_s)
+    if overrunning:
+        _leave(f"Exiting with task(s) still running in threads: {', '.join(overrunning)}", exit_now)
+    if stuck_threads:
+        _leave(f"Exiting with thread(s) or task(s) still running after the drain: {', '.join(stuck_threads)}", exit_now)
+    return code

async_task_api_worker/_runtime/log_context.py

@@ -0,0 +1,20 @@
+"""The task a log record belongs to, carried by a context variable.
+
+The engine binds it around a delivery, and a logging filter (see ``observability.logging``) copies it onto the records
+that do not carry it already, so that the lines of a library used inside a handler (an HTTP client, a storage client)
+are attributable to the task too. A context variable follows the ``await`` chain of the delivery's task; for ``def``
+handlers the sync bridge runs the thread in a copy of the context.
+"""
+
+from contextvars import ContextVar
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True, slots=True)
+class TaskLogContext:
+    task_id: str
+    service: str
+    attempt: int
+
+
+current_task: ContextVar[TaskLogContext | None] = ContextVar("async_task_api_worker_current_task", default=None)

async_task_api_worker/_runtime/observer.py

@@ -0,0 +1,78 @@
+"""Events the engine emits. The engine never imports a metrics library: an observer turns them into metrics."""
+
+from enum import StrEnum
+from typing import Protocol
+
+from async_task_api_worker.protocol.envelope import PoisonReason
+
+
+class TaskStatus(StrEnum):
+    SUCCESS = "success"
+    FAILURE = "failure"
+
+
+class MessageType(StrEnum):
+    """The messages a worker publishes, as the metric label values of the observer events."""
+
+    STARTED = "started"
+    PROGRESS = "progress"
+    SUCCESS = "success"
+    FAILURE = "failure"
+
+
+class ProgressDropReason(StrEnum):
+    """Why a progress update was not published (or not published as such)."""
+
+    INVALID_FRACTION = "invalid_fraction"
+    INVALID_PAYLOAD = "invalid_payload"
+    NOT_INCREASING = "not_increasing"
+    COALESCED = "coalesced"
+    AFTER_CONCLUSION = "after_conclusion"
+
+
+class RuntimeObserver(Protocol):
+    def poison(self, reason: PoisonReason) -> None: ...
+
+    def task_retried(self) -> None: ...
+
+    def task_concluded(self, status: TaskStatus, code: str | None, duration_s: float) -> None:
+        """Called after the ack, like the broker sees it. ``code`` is the failure code, ``None`` on success."""
+        ...
+
+    def message_published(self, message_type: MessageType) -> None: ...
+
+    def publish_failed(self, message_type: MessageType) -> None: ...
+
+    def ack_failed(self) -> None: ...
+
+    def progress_dropped(self, reason: ProgressDropReason) -> None: ...
+
+    def shutdown_requeued(self) -> None: ...
+
+
+class NoopObserver:
+    """Observer that ignores every event."""
+
+    def poison(self, reason: PoisonReason) -> None:
+        pass
+
+    def task_retried(self) -> None:
+        pass
+
+    def task_concluded(self, status: TaskStatus, code: str | None, duration_s: float) -> None:
+        pass
+
+    def message_published(self, message_type: MessageType) -> None:
+        pass
+
+    def publish_failed(self, message_type: MessageType) -> None:
+        pass
+
+    def ack_failed(self) -> None:
+        pass
+
+    def progress_dropped(self, reason: ProgressDropReason) -> None:
+        pass
+
+    def shutdown_requeued(self) -> None:
+        pass