orchestrata_ultimate.cxp.direct_transport

Reference › orchestrata_ultimate › cxp › direct_transport

direct_transport is a thin asyncio wrapper over cxpt, the free-threaded direct transport that carries CXP traffic. direct_cxp_adapter is the only importer. It moves opaque method names and byte payloads and knows nothing of packets. See cxp.

AsyncDirectClient wraps cxpt.Client. Each request gets a correlation id that keys a pending future (unary calls) or a queue (streams). The transport's completion file descriptor is watched with the event loop's reader, so replies are resolved on the loop thread by _drain. call sends a unary request and returns the response payload, raising RpcError for a non-OK status and TimeoutError when the timeout expires; the timeout matters because a peer that dies after the request was submitted would otherwise leave the future pending forever. call_status returns the status and payload without raising on a bad status, call_stream yields stream items, and close and aclose fail any waiters. DirectServer runs the server side: it takes a target and a handler, serves requests on worker threads, turns a handler exception into an internal-error status, and drains stream generators into item and end frames; it works as a context manager. RpcError carries the status code and payload.

At a glance

Class diagram

Dependency diagram

API

Classes

RpcError

Bases: Exception · Line 93

Raised by :meth:AsyncDirectClient.call when the peer returns non-OK.

AsyncDirectClient

Line 102

Asyncio client over cxpt.Client.

DirectServer

Line 309

Server harness over cxpt.Server.

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/cxp/direct_transport.py

"""Async wrapper over the ``cxpt`` free-threaded direct transport.

This is the GIL-off-safe replacement for ``grpc.aio`` on the internal CXP RPC
path (kernel / gateway / KG, all on the trusted Incus fleet). It drops HTTP/2
entirely in favour of raw length-prefixed frames over Unix domain sockets
(co-located peers) or TCP (cross-host):

* No broker, no HTTP/2, no streams / flow-control -> the "CANCELLED under load"
  failure class is structurally impossible.
* The native worker / reader / acceptor threads inside ``cxpt`` are pure C and
  never touch a Python object, so they cannot race the free-threaded CPython
  (mimalloc) heap the way grpc's EventEngine threads do (that race is what hangs
  the KG suite at interpreter exit; see
  ``tests/knowledge_graph/conftest.py`` and the ``python314t-heap-corruption``
  investigation). The only threads that touch Python are real Python threads
  with a valid tstate: the asyncio loop thread (client) and the handler-pool
  threads (server).

Both unary and server-streaming RPCs are supported. The response frame carries
a ``flags`` byte (``UNARY`` / ``STREAM_ITEM`` / ``STREAM_END``); the **server**
decides which by what it emits, so the request path is identical for both:

* ``await client.call(method, payload)`` -> a single response (unary).
* ``async for chunk in client.call_stream(method, payload)`` -> a stream.

On the server, a handler returning a ``(status, bytes)`` tuple is unary; a
handler returning an *iterable* of ``(status, bytes)`` is a stream (each item is
a ``STREAM_ITEM`` frame, followed by a terminating ``STREAM_END``).

A client that stops consuming a stream early (breaks out of the ``async for``,
or otherwise closes the async generator ``call_stream`` returns) sends a
CANCEL control frame for that stream; the server stops producing into it and
frees its stream-table entry instead of running the handler to natural
completion for a peer that is no longer listening (backlog #23, the
bounded-leak follow-up to server-streaming). See ``AsyncDirectClient.call_stream``
and ``DirectServer._serve_stream``.

Wire status codes reuse the gRPC canonical numbers so call sites can migrate
without relearning semantics.
"""

from __future__ import annotations

import asyncio
import contextlib
import functools
import itertools
import logging
import threading
from collections.abc import AsyncIterator, Callable, Iterator

import cxpt

from orchestrata_ultimate.core.base.traces import traces

logger = logging.getLogger(__name__)

# Canonical status codes (subset of the gRPC set) carried in the response frame.
STATUS_OK = 0
STATUS_UNKNOWN = 2
STATUS_INVALID_ARGUMENT = 3
STATUS_NOT_FOUND = 5
STATUS_UNIMPLEMENTED = 12
STATUS_INTERNAL = 13
STATUS_UNAVAILABLE = 14

# Response-frame flags (server-driven). UNARY is a single terminal response;
# a stream is zero or more STREAM_ITEM frames terminated by one STREAM_END.
_FLAG_UNARY = 0
_FLAG_STREAM_ITEM = 1
_FLAG_STREAM_END = 2

_U32 = 0xFFFFFFFF

# handler(method, payload, cancelled) -> either a unary (status, response)
# tuple, or an Iterator of (status, chunk) tuples for a server-streaming
# reply. Iterator (not Iterable) keeps the arms disjoint for the type checker
# — a bare tuple is Iterable but not Iterator — so isinstance(result, tuple)
# narrows cleanly. ``cancelled`` is a zero-arg predicate a streaming handler
# MAY poll from deep inside a single, otherwise-uninterruptible blocking pull
# (e.g. an async generator awaiting a queue with a long internal timeout) to
# notice a client CANCEL promptly instead of only between yielded items --
# see DirectCXPServer._drive_stream for the canonical consumer (backlog #23).
# Unary handlers and simple synchronous-generator streaming handlers can
# ignore it: DirectServer._serve_stream already checks cancellation between
# items regardless.
Handler = Callable[
    [bytes, bytes, "Callable[[], bool]"],
    "tuple[int, bytes] | Iterator[tuple[int, bytes]]",
]


class RpcError(Exception):
    """Raised by :meth:`AsyncDirectClient.call` when the peer returns non-OK."""

    def __init__(self, status: int, payload: bytes) -> None:
        super().__init__(f"RPC failed with status {status}")
        self.status = status
        self.payload = payload


class AsyncDirectClient:
    """Asyncio client over ``cxpt.Client``.

    Correlation ids key an in-flight :class:`asyncio.Future` map (unary) or an
    :class:`asyncio.Queue` map (streaming); the transport's completion fd is
    watched with ``loop.add_reader`` so responses resolve on the loop thread.
    Construct inside a running event loop (or call any coroutine from one) — the
    loop is bound lazily on first use.
    """

    def __init__(self, target: str, nworkers: int = 4) -> None:
        self._client = cxpt.Client(target, nworkers)
        self._pending: dict[int, asyncio.Future[tuple[int, bytes]]] = {}
        self._streams: dict[int, asyncio.Queue[tuple[int, int, bytes]]] = {}
        self._corr = itertools.count(1)
        self._loop: asyncio.AbstractEventLoop | None = None
        self._fd = self._client.wakeup_fd()
        self._closed = False

    def _ensure_reader(self) -> asyncio.AbstractEventLoop:
        loop = asyncio.get_running_loop()
        if self._loop is None:
            self._loop = loop
            loop.add_reader(self._fd, self._drain)
        return loop

    def _drain(self) -> None:
        # Runs on the loop thread. Resolve unary futures; feed stream queues.
        for corr, flags, status, payload in self._client.poll():
            if flags == _FLAG_UNARY:
                fut = self._pending.pop(corr, None)
                if fut is not None and not fut.done():
                    fut.set_result((status, payload))
            else:
                queue = self._streams.get(corr)
                if queue is not None:
                    queue.put_nowait((flags, status, payload))

    async def call(
        self, method: str, payload: bytes, timeout: float | None = None
    ) -> bytes:
        """Send a unary request and await the response.

        Returns the response payload on ``STATUS_OK``; raises :class:`RpcError`
        otherwise, or :class:`TimeoutError` if *timeout* elapses first.

        A timeout is not optional in practice. Replies arrive by correlation id
        through :meth:`_drain`, so a peer that dies (or simply never answers)
        after the request was submitted leaves the future pending FOREVER --
        there is no connection-level event that completes it. Measured
        2026-08-07: a gateway talking to a kernel that went away blocked until
        pytest-timeout killed it at 120s, and would have blocked indefinitely
        in production. Callers pass their configured
        ``CXPClientConfig.timeout_seconds``.
        """
        if self._closed:
            raise RuntimeError("client is closed")
        loop = self._ensure_reader()
        corr = next(self._corr) & _U32
        fut: asyncio.Future[tuple[int, bytes]] = loop.create_future()
        self._pending[corr] = fut
        try:
            self._client.submit(corr, method, payload)
        except Exception:
            self._pending.pop(corr, None)
            raise
        try:
            status, resp = await self._await_reply(fut, timeout)
        finally:
            # Drop the correlation entry on every exit path: a timed-out
            # request whose reply lands later must not accumulate in
            # `_pending` (an unbounded leak on a flapping peer).
            self._pending.pop(corr, None)
        if status != STATUS_OK:
            raise RpcError(status, resp)
        return resp

    @staticmethod
    async def _await_reply(
        fut: asyncio.Future[tuple[int, bytes]], timeout: float | None
    ) -> tuple[int, bytes]:
        """Await *fut*, bounded by *timeout* when one is configured."""
        if timeout is None:
            return await fut
        return await asyncio.wait_for(fut, timeout=timeout)

    async def call_status(
        self, method: str, payload: bytes, timeout: float | None = None
    ) -> tuple[int, bytes]:
        """Like :meth:`call` but return ``(status, payload)`` without raising.

        Still raises :class:`TimeoutError` if *timeout* elapses -- see
        :meth:`call` for why an unbounded wait is not survivable.
        """
        if self._closed:
            raise RuntimeError("client is closed")
        loop = self._ensure_reader()
        corr = next(self._corr) & _U32
        fut: asyncio.Future[tuple[int, bytes]] = loop.create_future()
        self._pending[corr] = fut
        try:
            self._client.submit(corr, method, payload)
        except Exception:
            self._pending.pop(corr, None)
            raise
        try:
            return await self._await_reply(fut, timeout)
        finally:
            self._pending.pop(corr, None)

    @traces(spec="001", story="US-002", reqs=["REQ-001-US-002-012"])
    async def call_stream(self, method: str, payload: bytes) -> AsyncIterator[bytes]:
        """Send a request and yield response payloads until the stream ends.

        The server drives the stream: each ``STREAM_ITEM`` frame yields its
        payload; a ``STREAM_END`` frame terminates the iteration. A non-OK status
        on any frame raises :class:`RpcError` (and ends the stream).

        Backlog #23: if this async generator is closed before it observes
        ``STREAM_END`` -- an early ``break`` out of the consuming ``async for``,
        or an explicit ``.aclose()`` -- a CANCEL control frame is sent for
        ``corr`` on the SAME worker connection ``submit()`` used (required so
        the server's per-connection stream lookup finds it; see ``cxpt.Client
        .cancel`` and the wire-format note in ``_native.c``). A stream that
        already ended (normally or via :class:`RpcError`) needs no cancel --
        the server already unregistered it.
        """
        if self._closed:
            raise RuntimeError("client is closed")
        self._ensure_reader()
        corr = next(self._corr) & _U32
        queue: asyncio.Queue[tuple[int, int, bytes]] = asyncio.Queue()
        self._streams[corr] = queue
        worker_idx: int | None = None
        ended = False
        try:
            worker_idx = self._client.submit(corr, method, payload)
            while True:
                flags, status, chunk = await queue.get()
                if status != STATUS_OK:
                    ended = True
                    raise RpcError(status, chunk)
                if flags == _FLAG_STREAM_END:
                    ended = True
                    return
                yield chunk
        finally:
            self._streams.pop(corr, None)
            if not ended and worker_idx is not None and not self._closed:
                with contextlib.suppress(OSError):
                    self._client.cancel(corr, worker_idx)

    def _remove_reader(self) -> None:
        if self._loop is not None:
            try:
                self._loop.remove_reader(self._fd)
            except ValueError, OSError:  # loop already closed
                pass

    def _fail_waiters(self) -> None:
        for fut in self._pending.values():
            if not fut.done():
                fut.set_exception(RpcError(STATUS_UNAVAILABLE, b""))
        self._pending.clear()
        # Wake any in-flight stream consumers so they raise instead of hanging.
        for queue in self._streams.values():
            queue.put_nowait((_FLAG_STREAM_END, STATUS_UNAVAILABLE, b""))
        self._streams.clear()

    def close(self) -> None:
        if self._closed:
            return
        self._closed = True
        self._remove_reader()
        self._client.close()
        self._fail_waiters()

    async def aclose(self) -> None:
        """Close without blocking the event loop.

        ``cxpt.Client.close()`` joins native worker threads (blocking), so it
        runs in a thread; ``remove_reader`` and future/queue wakeups are
        loop-affine and stay on the calling loop thread.
        """
        if self._closed:
            return
        self._closed = True
        self._remove_reader()
        await asyncio.to_thread(self._client.close)
        self._fail_waiters()

    async def __aenter__(self) -> AsyncDirectClient:
        return self

    async def __aexit__(self, *exc: object) -> None:
        await self.aclose()


def _close_stream_items(items: Iterator[tuple[int, bytes]]) -> None:
    """Best-effort ``items.close()`` -- ``Iterator`` doesn't guarantee it, but
    every real handler (a generator or generator-backed adapter) has one; this
    is what actually stops a producer parked mid-``yield`` (backlog #23)."""
    close = getattr(items, "close", None)
    if close is not None:
        close()


class DirectServer:
    """Server harness over ``cxpt.Server``.

    Runs ``workers`` real Python threads (GIL off -> genuine parallelism), each
    blocking on ``next_request`` and dispatching to ``handler``. A handler that
    returns a ``(status, bytes)`` tuple produces a unary reply; a handler that
    returns an *iterable* of ``(status, bytes)`` produces a server-streamed reply
    (one ``STREAM_ITEM`` frame per item, then a terminating ``STREAM_END``).
    Handler exceptions become ``STATUS_INTERNAL`` rather than killing the worker.
    """

    def __init__(self, target: str, handler: Handler, workers: int = 8) -> None:
        self._server = cxpt.Server(target)
        self._handler = handler
        self._closed = False
        self._threads = [
            threading.Thread(target=self._serve, name=f"cxpt-handler-{i}", daemon=True)
            for i in range(workers)
        ]
        for t in self._threads:
            t.start()

    def _serve(self) -> None:
        server = self._server
        handler = self._handler
        while True:
            req = server.next_request()
            if req is None:  # server closing
                return
            conn_id, corr, method, payload = req
            # functools.partial binds conn_id/corr by VALUE now, not by late-
            # binding closure -- a plain lambda here would capture the loop
            # variables themselves and see whatever request _serve is on by
            # the time a streaming handler actually calls this, long after
            # this iteration has moved on.
            cancelled = functools.partial(server.stream_is_cancelled, conn_id, corr)
            try:
                result = handler(method, payload, cancelled)
            except Exception:  # noqa: BLE001 — never let one bad call kill a worker
                logger.exception("cxpt handler raised for method %r", method)
                server.send_response(conn_id, corr, STATUS_INTERNAL, b"", _FLAG_UNARY)
                continue
            if isinstance(result, tuple):
                status, response = result
                server.send_response(conn_id, corr, status, response, _FLAG_UNARY)
            else:
                self._serve_stream(conn_id, corr, method, result)

    @traces(spec="001", story="US-002", reqs=["REQ-001-US-002-012"])
    def _serve_stream(
        self,
        conn_id: int,
        corr: int,
        method: bytes,
        items: Iterator[tuple[int, bytes]],
    ) -> None:
        """Drain a streaming handler's items onto the wire.

        Registers ``(conn_id, corr)`` in the cxpt server's stream table (see
        ``cxpt.Server.stream_register``/``.stream_is_cancelled``/
        ``.stream_unregister``) so an inbound CANCEL frame from the client
        (``AsyncDirectClient.call_stream``) has something to mark. Cancellation
        is checked BEFORE pulling each item -- not just before sending it -- so
        a cancel stops the handler from being asked to produce anything further
        (backlog #23): a dead connection (``send_response`` returning ``False``)
        gets the same treatment, since spinning a handler thread for a peer
        that can never read is the identical bounded-leak class. Either way
        ``items.close()`` releases whatever the handler holds (a generator
        frame, or -- via ``DirectCXPServer._drive_stream``'s own ``finally`` --
        the underlying servicer async generator) instead of it running
        unattended to natural completion. The table entry is always freed in
        ``finally`` so the registry returns to baseline once nothing is in
        flight, whether the stream ended normally, errored, or was cancelled.

        Every terminal ``send_response`` (STREAM_END / the STATUS_INTERNAL
        fallback) is wrapped in ``contextlib.suppress(OSError)``: a
        genuinely-dead connection can flip from "alive" to a failing
        ``write()`` between the boolean-``delivered`` check above and the
        very next send -- ``Server_send_response`` raises ``OSError`` in that
        race rather than returning ``False`` (that only covers the
        already-reaped case). Pre-existing gap, surfaced by cancel/teardown
        tests exercising this path far more densely: previously an OSError
        here from the STATUS_INTERNAL fallback (itself reached because the
        FIRST send already failed) went uncaught and crashed the handler
        thread.
        """
        server = self._server
        server.stream_register(conn_id, corr)
        try:
            while True:
                if server.stream_is_cancelled(conn_id, corr):
                    _close_stream_items(items)
                    return
                try:
                    status, chunk = next(items)
                except StopIteration:
                    with contextlib.suppress(OSError):
                        server.send_response(
                            conn_id, corr, STATUS_OK, b"", _FLAG_STREAM_END
                        )
                    return
                try:
                    delivered = server.send_response(
                        conn_id, corr, status, chunk, _FLAG_STREAM_ITEM
                    )
                except OSError:
                    _close_stream_items(items)
                    return
                if not delivered:
                    _close_stream_items(items)
                    return
        except Exception:  # noqa: BLE001 — surface as a stream-terminating error
            logger.exception("cxpt stream handler raised for method %r", method)
            with contextlib.suppress(OSError):
                server.send_response(
                    conn_id, corr, STATUS_INTERNAL, b"", _FLAG_STREAM_END
                )
        finally:
            server.stream_unregister(conn_id, corr)

    def close(self) -> None:
        if self._closed:
            return
        self._closed = True
        self._server.close()  # wakes next_request -> None, workers exit
        for t in self._threads:
            t.join(timeout=2.0)

    def __enter__(self) -> DirectServer:
        return self

    def __exit__(self, *exc: object) -> None:
        self.close()