pyfect-dev 0.1.0__py3-none-any.whl

pyfect/rpc/client/_serialized.py

@@ -0,0 +1,321 @@
+"""The RPC client over a protocol (Effect's ``RpcClient.make``): encodes
+the requests for the ``ClientProtocol`` and decodes its responses."""
+
+from __future__ import annotations
+
+import itertools
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, Never, cast, overload
+
+from pyfect import exit as Exit
+from pyfect import schema as S
+from pyfect._core import Effect, uninterruptible_mask
+from pyfect._fiber import Fiber
+from pyfect._layer_core import Layer
+from pyfect._layer_for_life import with_layer
+from pyfect.cause import Cause
+from pyfect.context import Context
+from pyfect.deferred import Deferred
+from pyfect.effect import gen, interrupt, provide_context, succeed, suspend
+from pyfect.rpc._rpc import AnyRpc, exit_schema, parse_options
+from pyfect.rpc.client._core import ClientCore, ClientOptions, OnFromClient
+from pyfect.rpc.client._core import make_no_serialization as make_core
+from pyfect.rpc.client._protocol import ClientProtocol, ClientProtocolService, Message
+from pyfect.rpc.client_error import RpcClientDefect, RpcClientError
+from pyfect.rpc.group import RpcGroup
+from pyfect.rpc.message import (
+    Ack,
+    Eof,
+    Interrupt,
+    Request,
+    RequestId,
+    ResponseChunk,
+    ResponseDefect,
+    ResponseExit,
+)
+from pyfect.rpc.schema import get_stream_schemas
+from pyfect.scope import Scope, fork_scoped
+
+_VOID: Effect[None] = succeed(None)
+_client_ids = itertools.count()
+
+type _Parse = Callable[[object], Effect[Any, S.SchemaError, Any]]
+
+
+@dataclass(frozen=True, slots=True)
+class _Schemas:
+    decode_chunk: _Parse | None
+    encode_payload: _Parse
+    decode_exit: _Parse
+
+
+def _schemas(protocol: ClientProtocolService, rpc: AnyRpc) -> _Schemas:
+    codec_for = protocol.codec_for
+    streams = get_stream_schemas(rpc.success_schema)
+    options = parse_options(rpc)
+    chunk = (
+        None
+        if streams is None
+        else S.decode_unknown_effect(codec_for(S.NonEmptyArray(streams.success)), **options)
+    )
+    return _Schemas(
+        chunk,
+        S.encode_unknown_effect(codec_for(rpc.payload_schema), **options),
+        S.decode_unknown_effect(codec_for(exit_schema(rpc)), **options),
+    )
+
+
+@dataclass(frozen=True, slots=True)
+class _Entry:
+    rpc: AnyRpc
+    context: Context[Any]
+    schemas: _Schemas
+
+
+class _Serialized:
+    __slots__ = (
+        "_cache",
+        "_closing",
+        "_decode_defect",
+        "_entries",
+        "_group",
+        "_id",
+        "_protocol",
+        "core",
+    )
+
+    def __init__(self, group: type[RpcGroup], protocol: ClientProtocolService) -> None:
+        self._group = group
+        self._protocol = protocol
+        self._id = next(_client_ids)
+        self._entries: dict[RequestId, _Entry] = {}
+        self._cache: dict[str, _Schemas] = {}
+        self._closing = Deferred[None, Never]()
+        self._decode_defect = S.decode_sync(protocol.codec_for(S.Defect()))
+        self.core: ClientCore | None = None
+
+    def _core(self) -> ClientCore:
+        core = self.core
+        if core is None:
+            raise RuntimeError("the RPC client is not made yet")
+        return core
+
+    def _schemas_of(self, rpc: AnyRpc) -> _Schemas:
+        found = self._cache.get(rpc.tag)
+        if found is None:
+            found = self._cache[rpc.tag] = _schemas(self._protocol, rpc)
+        return found
+
+    def close(self) -> Effect[None]:
+        """Stop sending requests, and tell the server to interrupt the ones
+        still waiting: a call or stream consumed outside the client's scope
+        is resumed by the core, not interrupted, so it sends no ``Interrupt``
+        itself (Effect leaves them running on the server)."""
+
+        @gen
+        def interrupt_all():
+            for request_id in list(self._entries):
+                yield from self._interrupt_request(request_id).exit()
+
+        return self._closing.succeed(None).flat_map(
+            lambda _: interrupt_all.timeout(1.0).exit().as_void()
+        )
+
+    def _defect_of(self, encoded: object) -> BaseException:
+        """The defect a server sent, or why it could not be decoded: this
+        runs in the client's read fiber, which a raise would end."""
+        try:
+            return cast("BaseException", self._decode_defect(encoded))
+        except S.SchemaError as error:
+            return RpcClientDefect("Error decoding defect", error)
+
+    def _interrupt_request(self, request_id: RequestId) -> Effect[None, RpcClientError]:
+        if self._entries.pop(request_id, None) is None:
+            return _VOID
+        return self._protocol.send(self._id, {"_tag": "Interrupt", "requestId": request_id})
+
+    # --- requests -----------------------------------------------------------------------
+
+    def on_from_client(self, sent: OnFromClient) -> Effect[None, Any]:
+        message = sent.message
+        match message:
+            case Request():
+                return self._send_request(message)
+            case Ack(request_id=request_id):
+                if request_id not in self._entries:
+                    return _VOID
+                return self._protocol.send(self._id, {"_tag": "Ack", "requestId": request_id})
+            case Interrupt(request_id=request_id):
+                return self._interrupt_request(request_id)
+            case Eof():
+                return _VOID
+
+    def _send_request(self, message: Request) -> Effect[None, Any]:
+        rpc = self._group.rpcs[message.tag]
+        fiber = Fiber.get_current()
+        context = Context.empty() if fiber is None else fiber._context  # pyright: ignore[reportPrivateUsage]
+        entry = _Entry(rpc, context, self._schemas_of(rpc))
+        self._entries[message.id] = entry
+
+        def send(payload: object) -> Effect[None, Any]:
+            if self._closing.is_done_unsafe():
+                self._entries.pop(message.id, None)
+                return interrupt()
+            encoded: dict[str, Any] = {
+                "_tag": "Request",
+                "id": message.id,
+                "tag": message.tag,
+                "payload": payload,
+                "headers": [[k, v] for k, v in message.headers.items()],
+            }
+            if message.span_id is not None:
+                encoded["traceId"] = message.trace_id
+                encoded["spanId"] = message.span_id
+                encoded["sampled"] = message.sampled
+            return self._protocol.send(self._id, encoded)
+
+        def forget(cause: Cause[Any]) -> Effect[None]:
+            if not cause.has_interrupts_only:
+                self._entries.pop(message.id, None)
+            return _VOID
+
+        encoded = entry.schemas.encode_payload(message.payload).pipe(provide_context(context))
+        return encoded.or_die().flat_map(send).on_error(forget)
+
+    # --- responses ----------------------------------------------------------------------
+
+    def on_response(self, message: Message) -> Effect[None]:
+        core = self._core()
+        match message.get("_tag"):
+            case "Chunk":
+                return self._chunk(message)
+            case "Exit":
+                return self._exit(message)
+            case "Defect":
+                self._entries.clear()
+                return suspend(
+                    lambda: core.write(ResponseDefect(0, self._defect_of(message.get("defect"))))
+                )
+            case "ClientProtocolError":
+                error = cast("RpcClientError", message["error"])
+                ids = list(self._entries)
+                self._entries.clear()
+                failure = Exit.fail(error)
+
+                @gen
+                def deliver():
+                    for request_id in ids:
+                        yield from core.write(ResponseExit(0, request_id, failure))
+
+                return deliver.uninterruptible()
+            case _:
+                return _VOID
+
+    def _chunk(self, message: Message) -> Effect[None]:
+        request_id: RequestId = message["requestId"]
+        entry = self._entries.get(request_id)
+        if entry is None or entry.schemas.decode_chunk is None:
+            return _VOID
+        core = self._core()
+        decoded = entry.schemas.decode_chunk(message.get("values")).pipe(
+            provide_context(entry.context)
+        )
+
+        def failed(cause: Cause[Any]) -> Effect[None]:
+            written = core.write(ResponseExit(0, request_id, Exit.Failure(cause)))
+            if cause.has_interrupts:
+                return written
+            stop = self._interrupt_request(request_id).ignore().race_first(self._closing.wait())
+            return written.flat_map(lambda _: stop)
+
+        written: Effect[None] = decoded.or_die().flat_map(
+            lambda values: core.write(ResponseChunk(0, request_id, tuple(values)))
+        )
+        # Uninterruptible but for the decode and write: an interrupt during
+        # the cleanup lands after the decode failure is recovered.
+        return uninterruptible_mask(
+            lambda restore: (
+                restore(written)
+                .on_error(failed)
+                .catch_cause_if(lambda cause: not cause.has_interrupts, lambda _: _VOID)
+            )
+        )
+
+    def _exit(self, message: Message) -> Effect[None]:
+        request_id: RequestId = message["requestId"]
+        entry = self._entries.pop(request_id, None)
+        if entry is None:
+            return _VOID
+        core = self._core()
+        decoded = entry.schemas.decode_exit(message.get("exit")).pipe(
+            provide_context(entry.context)
+        )
+        return (
+            decoded.or_die()
+            .exit()
+            .flat_map(
+                lambda result: core.write(
+                    ResponseExit(
+                        0,
+                        request_id,
+                        result.value
+                        if isinstance(result, Exit.Success)
+                        else Exit.Failure(result.cause),
+                    )
+                )
+            )
+        )
+
+
+@overload
+def make[G: RpcGroup](
+    group: type[G], options: ClientOptions | None = None
+) -> Effect[G, Never, Scope | ClientProtocol]: ...
+@overload
+def make[G: RpcGroup, E, R](
+    group: type[G], options: ClientOptions | None = None, *, layer: Layer[ClientProtocol, E, R]
+) -> Effect[G, E, R | Scope]: ...
+def make(
+    group: type[RpcGroup],
+    options: ClientOptions | None = None,
+    *,
+    layer: Layer[ClientProtocol, Any, Any] | None = None,
+) -> Effect[Any, Any, Any]:
+    """A client of ``group`` over the ``ClientProtocol`` in the context
+    (Effect's ``RpcClient.make``), or over ``layer``, built into the
+    caller's ``Scope``. The client is valid until that scope closes; a
+    protocol provided to ``make(...)`` alone is released once the client
+    is made, and its calls then fail with an ``RpcClientError``. Client
+    middleware comes from the context too (``middleware.layer_client``)."""
+    return with_layer(lambda: _make(group, options), layer)
+
+
+def _make(
+    group: type[RpcGroup], options: ClientOptions | None
+) -> Effect[RpcGroup, Never, Scope | ClientProtocol]:
+    @gen
+    def run():
+        protocol = yield from ClientProtocol()
+        state = _Serialized(group, protocol)
+        resolved = options or ClientOptions()
+        core_options = ClientOptions(
+            span_prefix=resolved.span_prefix,
+            span_attributes=resolved.span_attributes,
+            disable_tracing=resolved.disable_tracing,
+            supports_ack=protocol.supports_ack,
+            generate_request_id=resolved.generate_request_id,
+        )
+        client, core = yield from make_core(group, state.on_from_client, core_options)
+        state.core = core
+        listen = protocol.run(state._id, state.on_response)  # pyright: ignore[reportPrivateUsage]
+        yield from fork_scoped(listen.catch_cause(lambda _: _VOID).interruptible())
+        scope = yield from Scope()
+        # Finalizers run in reverse order: this runs while the client still runs.
+        yield from scope.add_finalizer(lambda _: state.close())
+        return client
+
+    return run
+
+
+__all__ = ["make"]

pyfect/rpc/client/_transports.py

@@ -0,0 +1,328 @@
+"""The socket transport of an RPC client (Effect's ``makeProtocolSocket``)."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, Never, cast
+
+from pyfect import schedule
+from pyfect._core import Effect
+from pyfect.cause import Cause
+from pyfect.clock import Clock
+from pyfect.context import Reference
+from pyfect.duration import Duration, Seconds, nanos_or_infinite, to_seconds
+from pyfect.effect import fail, gen, succeed, suspend
+from pyfect.latch import Latch
+from pyfect.result import Ok
+from pyfect.rpc.client._protocol import (
+    ClientProtocolParts,
+    ClientProtocolService,
+    Message,
+    WriteTo,
+    with_run_client,
+)
+from pyfect.rpc.client_error import RpcClientDefect, RpcClientError
+from pyfect.rpc.message import PING, ServerMessage, decode_from_server
+from pyfect.rpc.serialization import Frame, Parser, RpcSerialization, RpcSerializationService
+from pyfect.schedule import Schedule
+from pyfect.scope import Scope, fork_scoped, scoped
+from pyfect.socket import (
+    Socket,
+    SocketError,
+    SocketOpenError,
+    SocketReadError,
+    SocketService,
+    Writer,
+)
+from pyfect.socket import SocketError as SocketErrorClass
+
+_VOID: Effect[None] = succeed(None)
+
+
+def _key(request_id: object) -> tuple[type[object], object]:
+    """A request id as a dict key: ``1`` and ``"1"`` stay apart."""
+    return (type(request_id), request_id)
+
+
+def _defect(message: str, cause: object) -> RpcClientError:
+    return RpcClientError(RpcClientDefect(message, cause))
+
+
+@dataclass(frozen=True, slots=True)
+class ConnectionHooksService:
+    """Run when a socket protocol connects and when it disconnects."""
+
+    on_connect: Effect[None]
+    on_disconnect: Effect[None]
+
+
+class ConnectionHooks(
+    Reference[ConnectionHooksService | None], key="effect/rpc/RpcClient/ConnectionHooks"
+):
+    """The ``ConnectionHooks`` of the socket protocols made here."""
+
+    @classmethod
+    def default(cls) -> ConnectionHooksService | None:
+        return None
+
+
+def default_retry_policy() -> Schedule[Any, Any, Never, Never]:
+    """Exponential from 0.5s by 1.5, at most 5s apart."""
+    return cast(
+        "Schedule[Any, Any, Never, Never]",
+        schedule.min(schedule.exponential(0.5, 1.5), schedule.spaced(5.0)),
+    )
+
+
+class _Pinger:
+    __slots__ = ("_awaiting", "_clock", "_last_seen", "_timeout", "latch")
+
+    def __init__(self, clock: Any, timeout: Seconds) -> None:
+        self._clock = clock
+        self._timeout = nanos_or_infinite(timeout)
+        self._last_seen = clock.monotonic_nanos_unsafe()
+        self._awaiting = False
+        self.latch = Latch(False)
+
+    def on_frame(self) -> None:
+        self._last_seen = self._clock.monotonic_nanos_unsafe()
+        self._awaiting = False
+
+    def reset(self) -> None:
+        self.on_frame()
+        self.latch.close_unsafe()
+
+    def tick(self, write_ping: Effect[None, Any]) -> Effect[None, Any]:
+        def body() -> Effect[None, Any]:
+            elapsed = self._clock.monotonic_nanos_unsafe() - self._last_seen
+            if self._awaiting and elapsed > self._timeout:
+                return self.latch.open().as_void()
+            self._awaiting = True
+            return write_ping
+
+        return suspend(body)
+
+
+@dataclass(frozen=True, slots=True)
+class SocketOptions:
+    """How a socket protocol keeps its connection: a ping every
+    ``ping_interval`` (dropping the connection when nothing came back
+    within ``ping_timeout``), reconnecting by ``retry_policy``.
+    ``retry_transient_errors`` keeps requests waiting while the socket
+    cannot open (reporting each error to ``on_transient_error``)."""
+
+    ping_interval: Duration = 5.0
+    ping_timeout: Duration | None = None
+    retry_transient_errors: bool = False
+    retry_policy: Schedule[Any, Any, Never, Never] | None = None
+    on_transient_error: Callable[[RpcClientError], Effect[None]] | None = None
+    on_ping_timeout: Effect[None] | None = None
+
+
+class _SocketClient:
+    __slots__ = (
+        "_current_error",
+        "_ids",
+        "_options",
+        "_parser",
+        "_pinger",
+        "_requests",
+        "_serialization",
+        "_write",
+        "_writer",
+    )
+
+    def __init__(
+        self,
+        serialization: RpcSerializationService,
+        *,
+        write: WriteTo,
+        ids: set[int],
+        writer: Writer,
+        pinger: _Pinger,
+        options: SocketOptions,
+    ) -> None:
+        self._serialization = serialization
+        self._write = write
+        self._ids = ids
+        self._writer = writer
+        self._pinger = pinger
+        self._options = options
+        self._parser: Parser = serialization.make_unsafe()
+        self._requests: dict[object, int] = {}
+        self._current_error: RpcClientError | None = None
+
+    def write_ping(self) -> Effect[None, Any]:
+        return suspend(lambda: self._write_frame(self._parser.encode(PING)))
+
+    def _write_frame(self, frame: Frame | None) -> Effect[None, Any]:
+        return _VOID if frame is None else self._writer.write(frame)
+
+    def _broadcast(self, response: Message) -> Effect[None]:
+        @gen
+        def deliver():
+            for client_id in list(self._ids):
+                yield from self._write(client_id, response)
+
+        if response.get("_tag") not in ("Defect", "ClientProtocolError"):
+            return deliver
+        self._requests.clear()
+        return deliver.uninterruptible()
+
+    def broadcast_error(self, error: RpcClientError) -> Effect[None]:
+        self._current_error = error
+        return self._broadcast({"_tag": "ClientProtocolError", "error": error})
+
+    def _process(self, data: Frame) -> Effect[None]:
+        try:
+            responses = decode_from_server(self._parser.decode(data))
+        except Exception as cause:
+            return self._broadcast(
+                {"_tag": "ClientProtocolError", "error": _defect("Error decoding message", cause)}
+            )
+        if not responses:
+            return _VOID
+        self._pinger.on_frame()
+
+        @gen
+        def each():
+            for response in responses:
+                yield from self._route(response)
+
+        return each
+
+    def _route(self, response: ServerMessage) -> Effect[None]:
+        if response["_tag"] == "Pong":
+            return _VOID
+        if response["_tag"] == "Defect":
+            return self._broadcast(response)
+        key = _key(response["requestId"])
+        client_id = self._requests.get(key)
+        if client_id is None:
+            return self._broadcast(response)
+        if response["_tag"] == "Exit":
+            self._requests.pop(key, None)
+        return self._write(client_id, response)
+
+    def connection(
+        self, socket: SocketService, hooks: ConnectionHooksService | None
+    ) -> Effect[Never, SocketError, Any]:
+        """One connection: read until the socket fails or a ping times out.
+        ``hooks`` run when it connects, and when it ends, connected or not."""
+
+        def body() -> Effect[Never, SocketError, Any]:
+            self._parser = self._serialization.make_unsafe()
+            self._pinger.reset()
+            connected: list[bool] = []
+
+            @gen
+            def read():
+                reader = yield from socket.reader
+                self._current_error = None
+                connected.append(True)
+                if hooks is not None:
+                    yield from hooks.on_connect
+                while True:
+                    frames = yield from reader.pull
+                    for frame in frames:
+                        yield from self._process(frame)
+
+            on_timeout = self._options.on_ping_timeout or _VOID
+
+            @gen
+            def timed_out():
+                yield from scoped(read).race_first(self._pinger.latch.wait())
+                if connected:
+                    yield from on_timeout.ignore_cause()
+                return (
+                    yield from fail(SocketErrorClass(SocketReadError(RuntimeError("ping timeout"))))
+                )
+
+            return timed_out
+
+        connection = suspend(body)
+        return connection if hooks is None else connection.ensuring(hooks.on_disconnect)
+
+    def on_cause(self, cause: Cause[SocketError]) -> Effect[None]:
+        """Tell the clients how the connection ended: its socket error, or
+        an ``RpcClientDefect`` for a defect or an interruption."""
+        found = cause.find_error()
+        if not isinstance(found, Ok):
+            defect = _defect("Unknown socket error", cause.squash())
+            return self.broadcast_error(defect)
+        error = found.value
+        rpc_error = RpcClientError(error.reason)
+        if self._options.retry_transient_errors and isinstance(error.reason, SocketOpenError):
+            hook = self._options.on_transient_error
+            return _VOID if hook is None else hook(rpc_error).ignore_cause()
+        return self.broadcast_error(rpc_error)
+
+    def send(self, client_id: int, message: Message) -> Effect[None, RpcClientError]:
+        def body() -> Effect[None, RpcClientError]:
+            if self._current_error is not None:
+                return fail(self._current_error)
+            tag = message.get("_tag")
+            if tag == "Interrupt":
+                request_id = message["requestId"]
+                self._requests.pop(_key(request_id), None)
+            written = suspend(lambda: self._write_frame(self._parser.encode(message))).or_die()
+            if tag != "Request":
+                return written
+            key = _key(message["id"])
+            self._requests[key] = client_id
+            return written.on_error(lambda _: succeed(self._requests.pop(key, None)).as_void())
+
+        return suspend(body)
+
+
+def make_protocol_socket(
+    options: SocketOptions | None = None,
+) -> Effect[ClientProtocolService, Never, Scope | RpcSerialization | Socket]:
+    """A client protocol over the ``Socket`` in the context, reconnecting
+    when it drops."""
+    resolved = options or SocketOptions()
+
+    def parts(write: WriteTo, ids: set[int]) -> Effect[ClientProtocolParts, Never, Any]:
+        @gen
+        def make():
+            socket = yield from Socket()
+            serialization = yield from RpcSerialization()
+            hooks = yield from ConnectionHooks()
+            clock = yield from Clock()
+            writer = yield from socket.writer
+            interval = to_seconds(resolved.ping_interval)
+            timeout = (
+                interval if resolved.ping_timeout is None else to_seconds(resolved.ping_timeout)
+            )
+            pinger = _Pinger(clock, timeout)
+            client = _SocketClient(
+                serialization, write=write, ids=ids, writer=writer, pinger=pinger, options=resolved
+            )
+            yield from fork_scoped(
+                pinger.tick(client.write_ping()).delay(interval).ignore().forever().interruptible()
+            )
+            policy = resolved.retry_policy or default_retry_policy()
+            connection = client.connection(socket, hooks).tap_cause(client.on_cause)
+            yield from fork_scoped(
+                connection.retry(policy).catch_all(
+                    lambda error: client.broadcast_error(RpcClientError(error.reason))
+                )
+            )
+            return ClientProtocolParts(client.send, True, serialization.codec_for)
+
+        return make
+
+    return cast(
+        "Effect[ClientProtocolService, Never, Scope | RpcSerialization | Socket]",
+        with_run_client(parts),
+    )
+
+
+__all__ = [
+    "ConnectionHooks",
+    "ConnectionHooksService",
+    "SocketOptions",
+    "default_retry_policy",
+    "make_protocol_socket",
+]