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

async_task_api_worker/__init__.py

@@ -0,0 +1,42 @@
+"""SDK to write workers for the AsyncTaskAPI platform."""
+
+from async_task_api_worker._version import __version__
+from async_task_api_worker.context import SyncTaskContext, TaskContext
+from async_task_api_worker.errors import (
+    RESERVED_CODES,
+    ErrorCode,
+    FatalError,
+    ProbeError,
+    RetryableError,
+    WorkerConfigurationError,
+    WorkerDefinitionError,
+    WorkerError,
+)
+from async_task_api_worker.protocol.envelope import InputRef  # reached by handlers through ``ctx.inputs``
+from async_task_api_worker.protocol.json_types import JsonObject, JsonValue, ProgressPayload
+from async_task_api_worker.protocol.version import PROTOCOL_VERSION
+from async_task_api_worker.settings import LogLevel, WorkerMode, WorkerSettings
+from async_task_api_worker.worker import Worker
+
+__all__ = [
+    "PROTOCOL_VERSION",
+    "RESERVED_CODES",
+    "ErrorCode",
+    "FatalError",
+    "InputRef",
+    "JsonObject",
+    "JsonValue",
+    "LogLevel",
+    "ProbeError",
+    "ProgressPayload",
+    "RetryableError",
+    "SyncTaskContext",
+    "TaskContext",
+    "Worker",
+    "WorkerConfigurationError",
+    "WorkerDefinitionError",
+    "WorkerError",
+    "WorkerMode",
+    "WorkerSettings",
+    "__version__",
+]

async_task_api_worker/__main__.py

@@ -0,0 +1,5 @@
+import sys
+
+from async_task_api_worker.cli import main
+
+sys.exit(main())

async_task_api_worker/_runtime/__init__.py

@@ -0,0 +1 @@
+"""Private runtime: nothing imported from here is public API."""

async_task_api_worker/_runtime/amqp_transport.py

@@ -0,0 +1,251 @@
+"""The broker transport, on aio-pika. This and ``declare.py`` are the only modules that import it.
+
+One robust connection (it reconnects by itself and restores what it opened), with two channels:
+
+- the consume channel: a prefetch equal to the concurrency and manual acknowledgements, the engine acking only after
+  the terminal publish. The consumer is cancelled by this module, with the tag it owns, so that the order of the cancel
+  and of the last deliveries is ours: every delivery sent before the cancel is handed to the engine, and a cancelled
+  consumer is never restored after a reconnect;
+- the publish channel: publisher confirms on, and ``mandatory`` publishes that raise when nothing is routed.
+
+Queues are declared by :mod:`declare` (durable, no argument: the broker decides the type and the policies) and only
+attached here with a passive declare.
+
+Published messages are persistent, with a JSON content type and no other property of our own: no correlation id, no
+headers. aio-pika adds a random ``message_id``, which means nothing to anyone; the listener reads the body only (a task
+is identified by its ``task_id`` inside it).
+"""
+
+import asyncio
+import contextlib
+import logging
+from collections.abc import Awaitable, Callable
+from typing import Final
+
+from aio_pika import DeliveryMode, Message
+from aio_pika.abc import (
+    AbstractChannel,
+    AbstractIncomingMessage,
+    AbstractQueue,
+    AbstractRobustConnection,
+    ConsumerTag,
+)
+from aio_pika.connection import make_url
+from aio_pika.exceptions import AMQPConnectionError, ChannelInvalidStateError, ProbableAuthenticationError
+from aio_pika.robust_connection import RobustConnection
+
+from async_task_api_worker._runtime.declare import authentication_failed, declare_queues
+from async_task_api_worker._runtime.transport import CLOSE_TIMEOUT_S, OnDelivery, Publisher, TransportUnavailableError
+
+# Quorum queues count a delivery that was returned after a failure (channel closed, consumer gone) in
+# ``x-delivery-count``, and every delivery of a message handed out again, including a voluntary
+# ``nack(requeue=True)``, in ``x-acquired-count``. Only the failures can reach the broker's delivery limit.
+_DELIVERY_COUNT_HEADERS: Final[tuple[str, ...]] = ("x-delivery-count", "x-acquired-count")
+# How long a publish waits for the broker's confirmation before it counts as failed.
+_PUBLISH_CONFIRM_TIMEOUT_S: Final[float] = 30.0
+# How long ``basic.cancel`` may take. During an outage the robust channel waits for the reconnect before it sends it.
+CANCEL_TIMEOUT_S: Final[float] = 5.0
+_CONTENT_TYPE: Final[str] = "application/json"
+# The share of the close timeout the channels may take; the connection gets the rest.
+_CHANNELS_SHARE: Final[float] = 0.5
+_LOGGER: Final[logging.Logger] = logging.getLogger("async_task_api_worker")
+_UNREACHABLE: Final[tuple[type[Exception], ...]] = (OSError, AMQPConnectionError)
+
+type NewConnection = Callable[[str, float], AbstractRobustConnection]
+
+
+def _robust_connection(url: str, reconnect_interval_s: float) -> AbstractRobustConnection:
+    """A robust connection, not connected yet: the caller owns it from here and can close it whatever happens."""
+    return RobustConnection(make_url(url), reconnect_interval=reconnect_interval_s)
+
+
+class _AmqpDelivery:
+    def __init__(self, message: AbstractIncomingMessage) -> None:
+        self._message = message
+
+    @property
+    def body(self) -> bytes:
+        return self._message.body
+
+    @property
+    def delivery_count(self) -> int | None:
+        """How many times the broker already handed this message out: the largest of the two broker counters."""
+        headers = self._message.headers or {}
+        counts = [value for name in _DELIVERY_COUNT_HEADERS if isinstance(value := headers.get(name), int)]
+        return max(counts) if counts else None
+
+    @property
+    def redelivered(self) -> bool:
+        return bool(self._message.redelivered)
+
+    async def ack(self) -> None:
+        await self._message.ack()
+
+    async def nack(self, *, requeue: bool) -> None:
+        # Never ``reject``: the broker counts it as a failure, like a crash.
+        await self._message.nack(requeue=requeue)
+
+
+class _AmqpPublisher:
+    def __init__(self, transport: "AmqpTransport") -> None:
+        self._transport = transport
+
+    async def publish(self, body: bytes) -> None:
+        await self._transport.publish(body)
+
+
+class AmqpTransport:
+    """RabbitMQ through aio-pika: one consume channel (prefetch = concurrency), one publish channel with confirms."""
+
+    def __init__(
+        self,
+        *,
+        url: str,
+        in_queue: str,
+        out_queue: str,
+        prefetch: int,
+        reconnect_interval_s: float,
+        new_connection: NewConnection = _robust_connection,
+        cancel_timeout_s: float = CANCEL_TIMEOUT_S,
+        close_timeout_s: float = CLOSE_TIMEOUT_S,
+        declare: Callable[[str, tuple[str, ...]], Awaitable[None]] = declare_queues,
+    ) -> None:
+        self._new_connection = new_connection
+        self._cancel_timeout_s = cancel_timeout_s
+        self._close_timeout_s = close_timeout_s
+        self._declare = declare
+        self._url = url
+        self._in_queue = in_queue
+        self._out_queue = out_queue
+        self._prefetch = prefetch
+        self._reconnect_interval_s = reconnect_interval_s
+        self._on_delivery: OnDelivery | None = None
+        self._connection: AbstractRobustConnection | None = None
+        self._consume_channel: AbstractChannel | None = None
+        self._publish_channel: AbstractChannel | None = None
+        self._queue: AbstractQueue | None = None
+        self._consumer_tag: ConsumerTag | None = None
+        self._stopped = False
+        self._publisher = _AmqpPublisher(self)
+
+    @property
+    def publisher(self) -> Publisher:
+        return self._publisher
+
+    async def start(self, on_delivery: OnDelivery) -> None:
+        await self._declare(self._url, (self._in_queue, self._out_queue))
+        self._on_delivery = on_delivery
+        self._stopped = False
+        connection = self._new_connection(self._url, self._reconnect_interval_s)
+        try:
+            await connection.connect()
+            consume_channel = await connection.channel()
+            await consume_channel.set_qos(prefetch_count=self._prefetch)
+            publish_channel = await connection.channel(publisher_confirms=True, on_return_raises=True)
+            # Passive: the queue was declared above, and a passive declare sends no argument.
+            queue = await consume_channel.declare_queue(self._in_queue, passive=True)
+            self._consumer_tag = await queue.consume(self._on_message)
+        except ProbableAuthenticationError as error:
+            await self._discard(connection)
+            raise authentication_failed() from error
+        except _UNREACHABLE as error:
+            await self._discard(connection)
+            msg = "The broker is unreachable"
+            raise TransportUnavailableError(msg) from error
+        except BaseException:
+            # Also a cancellation (the job connect timeout): the half-built robust connection would keep reconnecting.
+            await self._discard(connection)
+            raise
+        self._connection = connection
+        self._consume_channel = consume_channel
+        self._publish_channel = publish_channel
+        self._queue = queue
+
+    async def _discard(self, connection: AbstractRobustConnection) -> None:
+        with contextlib.suppress(Exception):
+            async with asyncio.timeout(self._close_timeout_s):
+                await connection.close()
+
+    async def _on_message(self, message: AbstractIncomingMessage) -> None:
+        assert self._on_delivery is not None
+        if self._stopped:
+            # A delivery that still arrives after the stop (the consumer comes back after a reconnect, say): the engine
+            # is draining and must not take it. Give it back at once.
+            with contextlib.suppress(Exception):
+                await message.nack(requeue=True)
+            return
+        self._on_delivery(_AmqpDelivery(message))
+
+    async def stop_consuming(self) -> None:
+        """``basic.cancel``: once it returns, no new delivery comes, and every one sent before it was handed over.
+
+        Always sent, even while the connection is down: the robust queue forgets the consumer before its first await,
+        so that it is not restored after the reconnect, and the cancel itself then waits for the channel, hence the
+        timeout. Whatever still arrives after the stop is given back (see ``_on_message``).
+        """
+        queue, tag = self._queue, self._consumer_tag
+        self._consumer_tag = None
+        try:
+            if queue is not None and tag is not None:
+                with contextlib.suppress(ChannelInvalidStateError, TimeoutError):
+                    async with asyncio.timeout(self._cancel_timeout_s):
+                        await queue.cancel(tag)
+        finally:
+            # After the cancel, not before: what was sent before it still reaches the engine.
+            self._stopped = True
+
+    async def ping(self, timeout_s: float) -> bool:
+        del timeout_s  # the robust connection tracks its own state: no round trip is needed
+        connection = self._connection
+        return connection is not None and not connection.is_closed and connection.connected.is_set()
+
+    async def close(self) -> None:
+        """Close the channels, wait for the broker to confirm it, then close the connection.
+
+        Closing the connection alone writes ``Connection.Close`` and drops the socket without waiting. The broker's
+        reader then shuts the channel process down with an exit signal that takes priority over a ``basic.ack`` still in
+        its mailbox, and the quorum queue gives the delivery that was just concluded back: it is processed twice. A
+        ``Channel.Close`` is an RPC answered after what the channel already received, the ack included.
+
+        Each step is bounded (the channels get half of ``close_timeout_s``, the connection the rest) and a failure of
+        one does not prevent the next: unacked deliveries are then requeued by the broker, the correct fallback.
+        """
+        connection, self._connection = self._connection, None
+        channels = [channel for channel in (self._consume_channel, self._publish_channel) if channel is not None]
+        self._consume_channel = None
+        self._publish_channel = None
+        self._queue = None
+        if connection is None:
+            return
+        channels_timeout_s = self._close_timeout_s * _CHANNELS_SHARE
+        try:
+            async with asyncio.timeout(channels_timeout_s):
+                await asyncio.gather(*(self._close_channel(channel) for channel in channels))
+        except TimeoutError:
+            _LOGGER.warning(
+                "Closing the channels took more than %s s: closing the connection anyway.",
+                channels_timeout_s,
+            )
+        async with asyncio.timeout(self._close_timeout_s - channels_timeout_s):
+            await connection.close()
+
+    @staticmethod
+    async def _close_channel(channel: AbstractChannel) -> None:
+        try:
+            await channel.close()
+        except Exception:
+            _LOGGER.warning("Closing a channel failed: closing the connection anyway.", exc_info=True)
+
+    async def publish(self, body: bytes) -> None:
+        """Persistent, publisher-confirmed, mandatory publish on the output queue (default exchange)."""
+        channel = self._publish_channel
+        if channel is None:
+            msg = "The transport is not started"
+            raise TransportUnavailableError(msg)
+        message = Message(body, delivery_mode=DeliveryMode.PERSISTENT, content_type=_CONTENT_TYPE)
+        await channel.default_exchange.publish(
+            message,
+            routing_key=self._out_queue,
+            mandatory=True,
+            timeout=_PUBLISH_CONFIRM_TIMEOUT_S,
+        )

async_task_api_worker/_runtime/declare.py

@@ -0,0 +1,40 @@
+"""Queue declaration, done by the SDK itself and never with arguments.
+
+The broker decides the queue type (the virtual host default) and the policies add the delivery limit and the dead
+letter exchange. Declaring with an explicit ``x-queue-type`` (``classic``, say) would fail with
+``PRECONDITION_FAILED`` against an existing quorum queue. Consumers and publishers are then attached with a passive
+declare only.
+"""
+
+from collections.abc import Awaitable, Callable
+
+import aio_pika
+from aio_pika.abc import AbstractConnection
+
+from async_task_api_worker._runtime.transport import TransportUnavailableError
+from async_task_api_worker.errors import WorkerConfigurationError
+
+_UNREACHABLE: tuple[type[Exception], ...] = (OSError, aio_pika.exceptions.AMQPConnectionError)
+
+
+def authentication_failed() -> WorkerConfigurationError:
+    """The broker refused the credentials: retrying would loop forever, it is a configuration error."""
+    return WorkerConfigurationError("Broker authentication failed: check the credentials in BROKER_URL")
+
+
+type Connect = Callable[[str], Awaitable[AbstractConnection]]
+
+
+async def declare_queues(url: str, queue_names: tuple[str, ...], *, connect: Connect = aio_pika.connect) -> None:
+    """Declare each queue durable, with no arguments. Idempotent when the queue already exists."""
+    try:
+        connection = await connect(url)
+    except aio_pika.exceptions.ProbableAuthenticationError as error:
+        raise authentication_failed() from error
+    except _UNREACHABLE as error:
+        msg = "The broker is unreachable"
+        raise TransportUnavailableError(msg) from error
+    async with connection:
+        channel = await connection.channel()
+        for name in queue_names:
+            await channel.declare_queue(name, durable=True)