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

async_task_api_worker/_runtime/engine.py

@@ -0,0 +1,543 @@
+"""The per-delivery state machine: one ack or one nack for every delivery, whatever happens.
+
+Transport-agnostic: it runs on a :class:`Delivery` source and a :class:`Publisher`. The outcome of a delivery:
+
+- poison envelope: ack, nothing published;
+- ``started`` publish fails: nack (requeue), the handler does not run;
+- the handler concludes (success, ``INPUT_INVALID``, ``OUTPUT_INVALID``, ``TASK_TIMEOUT``, a ``WorkerError``, any
+  other exception as ``INTERNAL_ERROR``): the terminal message is published, **then** the delivery is acked;
+- a terminal publish that fails: nack (requeue), never left pending;
+- ``RetryableError`` before the last attempt: backoff, then nack (requeue); on the last attempt it concludes with
+  its own code;
+- cancellation (shutdown after the grace period): nack (requeue), then the cancellation is re-raised.
+"""
+
+import asyncio
+import hashlib
+import logging
+import random
+import threading
+import time
+from collections.abc import Awaitable, Callable, Mapping
+from contextvars import Token
+from dataclasses import dataclass
+from enum import StrEnum
+from functools import partial
+from typing import Final, Literal, cast
+
+from pydantic import BaseModel, ValidationError
+
+from async_task_api_worker._runtime.inflight import InFlightRegistry, InFlightTask
+from async_task_api_worker._runtime.log_context import TaskLogContext, current_task
+from async_task_api_worker._runtime.observer import MessageType, NoopObserver, RuntimeObserver, TaskStatus
+from async_task_api_worker._runtime.progress import Clock, ProgressDropLog, ProgressReporter, Sleep
+from async_task_api_worker._runtime.retry import Uniform, attempt_of, backoff_seconds
+from async_task_api_worker._runtime.sync_bridge import SyncBridge, make_threadsafe_progress
+from async_task_api_worker._runtime.transport import Delivery, Publisher
+from async_task_api_worker.context import SyncTaskContext, TaskContext
+from async_task_api_worker.errors import MAX_MESSAGE_LENGTH, ErrorCode, RetryableError, WorkerError
+from async_task_api_worker.protocol.envelope import InputRef, PoisonDelivery, TaskEnvelope, decode_envelope
+from async_task_api_worker.protocol.messages import MessageEncoder
+from async_task_api_worker.settings import WorkerSettings
+
+_LOGGER: Final[logging.Logger] = logging.getLogger("async_task_api_worker")
+# Enough to correlate a poison delivery with the broker or the API, and nothing of its content.
+_BODY_DIGEST_HEX_LENGTH: Final[int] = 12
+_MAX_REPORTED_ISSUES: Final[int] = 10
+_ELLIPSIS: Final[str] = "..."
+
+type AsyncHandler = Callable[[BaseModel, TaskContext], Awaitable[BaseModel]]
+type SyncHandler = Callable[[BaseModel, SyncTaskContext], BaseModel]
+
+
+@dataclass(frozen=True, slots=True)
+class TaskDefinition:
+    """What ``Worker.task`` registers: the handler and the models read from its annotations."""
+
+    handler: AsyncHandler | SyncHandler
+    is_async: bool
+    input_model: type[BaseModel]
+    output_model: type[BaseModel]
+    timeout_s: float
+
+
+@dataclass(frozen=True, slots=True)
+class _Terminal:
+    """A terminal message to publish. ``code`` is ``None`` for a success."""
+
+    body: bytes
+    code: str | None
+
+
+@dataclass(frozen=True, slots=True)
+class _Retry:
+    error: RetryableError
+
+
+class _DeadlineExceededError(Exception):
+    """The handler did not return before the deadline."""
+
+
+@dataclass(slots=True)
+class _Run:
+    """Per-delivery state that must survive an exception: what the concurrency slot waits for."""
+
+    entry: InFlightTask | None = None
+    thread_future: asyncio.Future[BaseModel] | None = None
+    log_token: Token[TaskLogContext | None] | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class _Attempt:
+    """Everything one handler run needs."""
+
+    envelope: TaskEnvelope
+    number: int
+    deadline: float
+    reporter: ProgressReporter
+    log: logging.LoggerAdapter[logging.Logger]
+    run: _Run
+    counted: bool = True
+
+    @property
+    def task_id(self) -> str:
+        return self.envelope.task_id
+
+
+class DeliveryOutcome(StrEnum):
+    """What the engine decided about a delivery: the ``outcome`` of its decision log line."""
+
+    SUCCESS = "success"
+    FAILURE = "failure"
+    POISON = "poison"
+    RETRY = "retry"
+    PUBLISH_FAILED = "publish_failed"
+    SHUTDOWN_REQUEUED = "shutdown_requeued"
+    CANCELLED = "cancelled"
+    ENGINE_ERROR = "engine_error"
+
+
+class SettleAction(StrEnum):
+    """What was done to the broker delivery: the ``action`` of its decision log line."""
+
+    ACK = "ack"
+    CONCLUDE = "conclude"
+    NACK_REQUEUE = "nack-requeue"
+
+
+class _Settler:
+    """Settles one delivery exactly once, however many code paths reach for it.
+
+    Every settlement writes one INFO "decision" line with the task id, the attempt over the limit, the outcome, the
+    failure code and the action, so that ``grep <task_id>`` tells the whole story of a task. Never a payload.
+    """
+
+    def __init__(self, delivery: Delivery, observer: RuntimeObserver, service: str, delivery_limit: int) -> None:
+        self._delivery = delivery
+        self._observer = observer
+        self._service = service
+        self._delivery_limit = delivery_limit
+        self._attempt = attempt_of(delivery)
+        self.task_id: str | None = None
+        self.settled = False
+
+    def _log_decision(self, outcome: DeliveryOutcome, action: SettleAction, code: str | None) -> None:
+        _LOGGER.info(
+            "Delivery settled: outcome=%s action=%s code=%s attempt=%d/%d",
+            outcome,
+            action,
+            code,
+            self._attempt,
+            self._delivery_limit,
+            extra={
+                "task_id": self.task_id,
+                "service": self._service,
+                "attempt": self._attempt,
+                "delivery_limit": self._delivery_limit,
+                "outcome": outcome.value,
+                "action": action.value,
+                "code": code,
+            },
+        )
+
+    async def ack(self, outcome: DeliveryOutcome, *, code: str | None = None, concluded: bool = False) -> None:
+        """Ack the delivery. ``concluded``: a terminal message was published for it first."""
+        self.settled = True
+        self._log_decision(outcome, SettleAction.CONCLUDE if concluded else SettleAction.ACK, code)
+        try:
+            await self._delivery.ack()
+        except Exception:
+            # The result is published, but the broker will redeliver (typically a closed channel after the
+            # consumer timeout): the work is redone. Make it loud.
+            self._observer.ack_failed()
+            _LOGGER.exception(
+                "ACK FAILED for task %s: the result is published but the message will be redelivered.",
+                self.task_id,
+            )
+
+    async def nack_requeue(self, outcome: DeliveryOutcome, *, code: str | None = None) -> None:
+        self.settled = True
+        self._log_decision(outcome, SettleAction.NACK_REQUEUE, code)
+        try:
+            await self._delivery.nack(requeue=True)
+        except Exception:
+            # Dead channel: the broker redelivers when the channel closes.
+            _LOGGER.exception("NACK FAILED for task %s: the broker redelivers when the channel closes.", self.task_id)
+
+
+def describe_validation_error(error: ValidationError) -> str:
+    """Public text of an input validation error: locations and messages, never the input values."""
+    issues = error.errors(include_url=False, include_context=False, include_input=False)
+    parts = [
+        f"{'.'.join(str(part) for part in issue['loc']) or '<root>'}: {issue['msg']}"
+        for issue in issues[:_MAX_REPORTED_ISSUES]
+    ]
+    text = "Invalid task input: " + "; ".join(parts)
+    if len(issues) > _MAX_REPORTED_ISSUES:
+        text += f"; and {len(issues) - _MAX_REPORTED_ISSUES} more"
+    text = " ".join(text.split())
+    if len(text) > MAX_MESSAGE_LENGTH:
+        text = text[: MAX_MESSAGE_LENGTH - len(_ELLIPSIS)] + _ELLIPSIS
+    return text
+
+
+class Engine:
+    """Processes deliveries for one task definition."""
+
+    def __init__(
+        self,
+        *,
+        definition: TaskDefinition,
+        service: str,
+        settings: WorkerSettings,
+        publisher: Publisher,
+        encoder: MessageEncoder | None = None,
+        observer: RuntimeObserver | None = None,
+        clock: Clock = time.monotonic,
+        sleep: Sleep = asyncio.sleep,
+        uniform: Uniform = random.uniform,
+    ) -> None:
+        self._definition = definition
+        self._service = service
+        self._settings = settings
+        self._publisher = publisher
+        self._encoder = encoder or MessageEncoder()
+        self._observer: RuntimeObserver = observer or NoopObserver()
+        self._clock = clock
+        self._sleep = sleep
+        self._uniform = uniform
+        self._timeout_s = (
+            settings.worker_task_timeout_s if settings.worker_task_timeout_s is not None else definition.timeout_s
+        )
+        self._slots = asyncio.Semaphore(settings.worker_concurrency)
+        self._inflight = InFlightRegistry()
+        self._bridge = SyncBridge(settings.worker_concurrency)
+        self._drops = ProgressDropLog(self._observer)
+        self._shutdown = asyncio.Event()
+        self._tasks: set[asyncio.Task[None]] = set()
+        self._warned_uncounted = False
+
+    @property
+    def inflight(self) -> InFlightRegistry:
+        return self._inflight
+
+    def tasks(self) -> tuple[asyncio.Task[None], ...]:
+        """Snapshot of the tasks spawned and not finished yet (safe to iterate while they complete)."""
+        return tuple(self._tasks)
+
+    def request_shutdown(self) -> None:
+        """Interrupt the retry backoffs: their deliveries are requeued at once."""
+        self._shutdown.set()
+
+    def close(self) -> None:
+        self._bridge.shutdown()
+
+    def spawn(self, delivery: Delivery) -> asyncio.Task[None]:
+        """Process ``delivery`` in a task the engine owns, so that a drain can wait for it or cancel it."""
+        task = asyncio.get_running_loop().create_task(self.process_delivery(delivery))
+        self._tasks.add(task)
+        task.add_done_callback(self._tasks.discard)
+        return task
+
+    async def process_delivery(self, delivery: Delivery) -> None:
+        """Run one delivery to its single ack or nack. Re-raises cancellation after requeueing."""
+        settler = _Settler(delivery, self._observer, self._service, self._settings.worker_delivery_limit)
+        try:
+            acquired = await self._acquire_slot()
+        except BaseException as error:
+            # Cancelled while waiting for a slot: the delivery must be requeued too.
+            await self._abort(error, settler)
+            raise
+        if not acquired:
+            # A shutdown was requested while this delivery waited for a slot: it must not run during the grace
+            # period, the broker hands it to another worker.
+            self._observer.shutdown_requeued()
+            await asyncio.shield(settler.nack_requeue(DeliveryOutcome.SHUTDOWN_REQUEUED))
+            return
+        run = _Run()
+        try:
+            await self._process(delivery, settler, run)
+        except BaseException as error:
+            await self._abort(error, settler)
+            raise
+        finally:
+            if run.log_token is not None:
+                current_task.reset(run.log_token)
+            self._release_slot(run)
+
+    async def _acquire_slot(self) -> bool:
+        """Wait for a concurrency slot. False, with no slot held, when a shutdown is requested first."""
+        if self._shutdown.is_set():
+            return False
+        if not self._slots.locked():
+            await self._slots.acquire()
+            return True
+        acquire = asyncio.ensure_future(self._slots.acquire())
+        stop = asyncio.ensure_future(self._shutdown.wait())
+        try:
+            await asyncio.wait({acquire, stop}, return_when=asyncio.FIRST_COMPLETED)
+        except BaseException:
+            await self._give_up(acquire)
+            raise
+        finally:
+            stop.cancel()
+        if acquire.done() and not self._shutdown.is_set():
+            return True
+        await self._give_up(acquire)
+        return False
+
+    async def _give_up(self, acquire: asyncio.Task[Literal[True]]) -> None:
+        """Stop waiting for a slot, and hand back the one that was granted at the very same moment, if any."""
+        acquire.cancel()
+        await asyncio.wait({acquire})
+        if not acquire.cancelled():
+            self._slots.release()
+
+    async def _abort(self, error: BaseException, settler: _Settler) -> None:
+        """Cancellation (shutdown after the grace period) or a bug: the delivery must not stay pending."""
+        cancelled = isinstance(error, asyncio.CancelledError)
+        if not cancelled:
+            _LOGGER.exception("Unexpected error escaped the engine for task %s.", settler.task_id)
+        if not settler.settled:
+            if cancelled:
+                self._observer.shutdown_requeued()
+            outcome = DeliveryOutcome.CANCELLED if cancelled else DeliveryOutcome.ENGINE_ERROR
+            await asyncio.shield(settler.nack_requeue(outcome))
+
+    def _release_slot(self, run: _Run) -> None:
+        """Free the concurrency slot. A sync handler that outlived its task keeps it until its thread returns."""
+
+        def release(_: object = None) -> None:
+            if run.entry is not None:
+                self._inflight.unregister(run.entry)
+            self._slots.release()
+
+        if run.thread_future is not None and not run.thread_future.done():
+            run.thread_future.add_done_callback(release)
+        else:
+            release()
+
+    async def _process(self, delivery: Delivery, settler: _Settler, run: _Run) -> None:
+        received = time.monotonic()
+        decoded = decode_envelope(delivery.body)
+        if isinstance(decoded, PoisonDelivery):
+            settler.task_id = decoded.task_id
+            _LOGGER.error(
+                "Poison delivery (%s), task_id=%s, body_bytes=%d, body_sha256=%s",
+                decoded.reason,
+                decoded.task_id,
+                len(delivery.body),
+                hashlib.sha256(delivery.body).hexdigest()[:_BODY_DIGEST_HEX_LENGTH],
+            )
+            self._observer.poison(decoded.reason)
+            await settler.ack(DeliveryOutcome.POISON, code=decoded.reason.value)
+            return
+
+        task_id = decoded.task_id
+        settler.task_id = task_id
+        attempt = attempt_of(delivery)
+        if attempt > self._settings.worker_delivery_limit:
+            await self._conclude(settler, self._failure(task_id, ErrorCode.TASK_ABANDONED_AFTER_REDELIVERIES), received)
+            return
+        if not await self._publish(self._encoder.started(task_id), MessageType.STARTED):
+            await settler.nack_requeue(DeliveryOutcome.PUBLISH_FAILED)
+            return
+
+        run.log_token = current_task.set(TaskLogContext(task_id, self._service, attempt))
+        deadline = received + self._timeout_s
+        run.entry = self._inflight.register(task_id, deadline, is_sync=not self._definition.is_async)
+        log = logging.LoggerAdapter(_LOGGER, {"task_id": task_id, "service": self._service, "attempt": attempt})
+        reporter = ProgressReporter(
+            task_id=task_id,
+            encoder=self._encoder,
+            publish=self._publish_progress,
+            min_interval_s=self._settings.worker_progress_min_interval_s,
+            drops=self._drops,
+            clock=self._clock,
+            sleep=self._sleep,
+        )
+        try:
+            # A first delivery legitimately carries no counter. A redelivery without one (a classic queue) cannot be
+            # counted, so a retry could never be bounded.
+            counted = delivery.delivery_count is not None or not delivery.redelivered
+            verdict = await self._run(_Attempt(decoded, attempt, deadline, reporter, log, run, counted))
+        finally:
+            # Before any terminal message: nothing is published as progress after it.
+            await reporter.close()
+
+        if isinstance(verdict, _Retry):
+            self._observer.task_retried()
+            delay = backoff_seconds(
+                attempt,
+                base_s=self._settings.worker_retry_backoff_base_s,
+                max_s=self._settings.worker_retry_backoff_max_s,
+                uniform=self._uniform,
+            )
+            await self._backoff(delay)
+            await settler.nack_requeue(DeliveryOutcome.RETRY, code=verdict.error.code)
+            return
+        await self._conclude(settler, verdict, received)
+
+    async def _conclude(self, settler: _Settler, terminal: _Terminal, received: float) -> None:
+        """Publish the terminal message, then ack. A failed publish requeues the delivery instead."""
+        message_type = MessageType.SUCCESS if terminal.code is None else MessageType.FAILURE
+        if not await self._publish(terminal.body, message_type):
+            await settler.nack_requeue(DeliveryOutcome.PUBLISH_FAILED, code=terminal.code)
+            return
+        outcome = DeliveryOutcome.SUCCESS if terminal.code is None else DeliveryOutcome.FAILURE
+        await settler.ack(outcome, code=terminal.code, concluded=True)
+        status = TaskStatus.SUCCESS if terminal.code is None else TaskStatus.FAILURE
+        self._observer.task_concluded(status, terminal.code, time.monotonic() - received)
+
+    async def _run(self, attempt: _Attempt) -> _Terminal | _Retry:
+        try:
+            task_input = self._definition.input_model.model_validate(attempt.envelope.data.body)
+        except ValidationError as error:
+            return self._failure(attempt.task_id, ErrorCode.INPUT_INVALID, describe_validation_error(error))
+        try:
+            result = await self._invoke(attempt, task_input)
+        except asyncio.CancelledError as error:
+            task = asyncio.current_task()
+            if task is not None and task.cancelling() > 0:
+                raise  # a real cancellation (shutdown): the delivery is requeued by the caller
+            # The handler's own code was cancelled (a cancelled child task or future): it is its failure, and a
+            # redelivery would fail the same way.
+            return self._on_error(attempt, error)
+        except Exception as error:
+            return self._on_error(attempt, error)
+        return self._success(attempt, result)
+
+    def _on_error(self, attempt: _Attempt, error: BaseException) -> _Terminal | _Retry:
+        """Map what the handler (or the deadline) raised to the message the client gets."""
+        log, task_id = attempt.log, attempt.task_id
+        if isinstance(error, _DeadlineExceededError):
+            log.error("Task exceeded its time limit of %s s.", self._timeout_s)
+            message = f"The task exceeded its time limit of {self._timeout_s:g} s."
+            return self._failure(task_id, ErrorCode.TASK_TIMEOUT, message)
+        if isinstance(error, RetryableError):
+            if not attempt.counted:
+                self._warn_uncounted()
+                log.warning("Retryable error %s on an uncounted redelivery, concluding the task.", error.code)
+                return self._failure(task_id, error)
+            if attempt.number >= self._settings.worker_delivery_limit:
+                log.warning("Retryable error on the last attempt, concluding with %s.", error.code)
+                return self._failure(task_id, error)
+            log.warning("Retryable error %s, requeueing after a backoff.", error.code)
+            return _Retry(error)
+        if isinstance(error, WorkerError):
+            log.warning("Task failed with %s.", error.code)
+            return self._failure(task_id, error)
+        # ``str(error)`` is never published: it could leak anything. The traceback stays in the logs.
+        log.error("Unexpected error while processing the task.", exc_info=error)
+        return self._failure(task_id, ErrorCode.INTERNAL_ERROR)
+
+    def _warn_uncounted(self) -> None:
+        """Once per process: the queue gives no delivery counter, so retries cannot be bounded."""
+        if self._warned_uncounted:
+            return
+        self._warned_uncounted = True
+        _LOGGER.warning(
+            "The input queue gives no delivery counter (it is not a quorum queue): a RetryableError concludes the "
+            "task on its first redelivery, and a crash loop is not bounded. Retries and crash-loop protection need "
+            "quorum queues.",
+        )
+
+    async def _invoke(self, attempt: _Attempt, task_input: BaseModel) -> BaseModel:
+        remaining = max(0.0, attempt.deadline - time.monotonic())
+        context = attempt.envelope.data.context
+        client_id = context.client_id if context else None
+        inputs: Mapping[str, InputRef] = context.input_refs if context else {}
+        identity = (attempt.task_id, client_id, attempt.number, attempt.deadline, attempt.log, inputs)
+
+        if self._definition.is_async:
+            handler = cast("AsyncHandler", self._definition.handler)
+            timeout = asyncio.timeout(remaining)
+            try:
+                async with timeout:
+                    return await handler(task_input, TaskContext(*identity, attempt.reporter.report))
+            except TimeoutError:
+                if timeout.expired():
+                    raise _DeadlineExceededError from None
+                raise
+
+        loop = asyncio.get_running_loop()
+        cancel_event = threading.Event()
+        progress = make_threadsafe_progress(loop, attempt.reporter)
+        sync_handler = cast("SyncHandler", self._definition.handler)
+        call = partial(sync_handler, task_input, SyncTaskContext(*identity, progress, cancel_event))
+        future = self._bridge.submit(loop, call)
+        attempt.run.thread_future = future
+        # A thread that outlives the task must not leave an unretrieved exception behind.
+        future.add_done_callback(lambda done: done.cancelled() or done.exception())
+        try:
+            return await asyncio.wait_for(asyncio.shield(future), remaining)
+        except TimeoutError:
+            if future.done():
+                # Finished at the same instant as the deadline (or raised a TimeoutError itself): the outcome of the
+                # thread wins, not the timer.
+                return future.result()
+            cancel_event.set()
+            raise _DeadlineExceededError from None
+        except asyncio.CancelledError:
+            cancel_event.set()
+            raise
+
+    def _success(self, attempt: _Attempt, result: BaseModel) -> _Terminal:
+        expected = self._definition.output_model
+        if not isinstance(result, expected):
+            attempt.log.error("The handler returned %s, expected %s.", type(result).__name__, expected.__name__)
+            return self._failure(attempt.task_id, ErrorCode.OUTPUT_INVALID)
+        try:
+            return _Terminal(self._encoder.success(attempt.task_id, result.model_dump(mode="json")), None)
+        except Exception:
+            attempt.log.exception("The handler output cannot be serialized.")
+            return self._failure(attempt.task_id, ErrorCode.OUTPUT_INVALID)
+
+    def _failure(self, task_id: str, error: WorkerError | ErrorCode, message: str | None = None) -> _Terminal:
+        code = error.code if isinstance(error, WorkerError) else error.value
+        return _Terminal(self._encoder.failure(task_id, error, message), code)
+
+    async def _publish(self, body: bytes, message_type: MessageType) -> bool:
+        try:
+            await self._publisher.publish(body)
+        except Exception:
+            self._observer.publish_failed(message_type)
+            _LOGGER.exception("Publishing a %s message failed.", message_type)
+            return False
+        self._observer.message_published(message_type)
+        return True
+
+    async def _publish_progress(self, body: bytes) -> None:
+        """Progress is best effort: a failed publish is counted and logged, never raised."""
+        await self._publish(body, MessageType.PROGRESS)
+
+    async def _backoff(self, delay_s: float) -> None:
+        """Wait ``delay_s``, or until a shutdown is requested."""
+        sleeper = asyncio.ensure_future(self._sleep(delay_s))
+        waiter = asyncio.ensure_future(self._shutdown.wait())
+        try:
+            await asyncio.wait({sleeper, waiter}, return_when=asyncio.FIRST_COMPLETED)
+        finally:
+            for pending in (sleeper, waiter):
+                pending.cancel()
+            await asyncio.wait({sleeper, waiter})

async_task_api_worker/_runtime/inflight.py

@@ -0,0 +1,38 @@
+"""Registry of the tasks being processed: what a drain waits for and what overrun detection looks at."""
+
+from collections.abc import Iterator
+from dataclasses import dataclass, field
+from itertools import count
+
+
+@dataclass(frozen=True, slots=True)
+class InFlightTask:
+    token: int
+    task_id: str
+    deadline: float
+    is_sync: bool
+
+
+@dataclass(slots=True)
+class InFlightRegistry:
+    _entries: dict[int, InFlightTask] = field(default_factory=dict)
+    _tokens: Iterator[int] = field(default_factory=count)
+
+    def register(self, task_id: str, deadline: float, *, is_sync: bool) -> InFlightTask:
+        entry = InFlightTask(next(self._tokens), task_id, deadline, is_sync)
+        self._entries[entry.token] = entry
+        return entry
+
+    def unregister(self, entry: InFlightTask) -> None:
+        self._entries.pop(entry.token, None)
+
+    def snapshot(self) -> tuple[InFlightTask, ...]:
+        """A copy: callers may iterate it while tasks come and go."""
+        return tuple(self._entries.values())
+
+    def __len__(self) -> int:
+        return len(self._entries)
+
+    def overrun(self, now: float, grace_s: float) -> tuple[InFlightTask, ...]:
+        """Tasks still registered ``grace_s`` after their deadline (a sync handler that never returns)."""
+        return tuple(entry for entry in self._entries.values() if now > entry.deadline + grace_s)