pyfect-dev 0.1.0__py3-none-any.whl

pyfect/rpc/server/_transports.py

@@ -0,0 +1,427 @@
+"""The transports of an RPC server: sockets (TCP, WebSocket) and HTTP
+(Effect's ``makeSocketProtocol`` and ``makeProtocolWithHttpEffect``)."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, Literal, Never, cast
+
+from pyfect._core import Effect
+from pyfect.effect import fail, gen, succeed, suspend
+from pyfect.errors import Done
+from pyfect.http.origin import origin_allowed
+from pyfect.http.server_request import HttpServerRequest
+from pyfect.http.server_response import HttpServerResponse, empty, text
+from pyfect.http.server_response import bytes_ as bytes_response
+from pyfect.http.server_response import stream as stream_response
+from pyfect.queue import Queue
+from pyfect.rpc.message import RequestId, response_defect_encoded
+from pyfect.rpc.serialization import (
+    Frame,
+    MaxBufferSizeExceeded,
+    Parser,
+    RpcSerialization,
+    RpcSerializationService,
+    encode_text,
+    is_binary,
+)
+from pyfect.rpc.server._protocol import (
+    Message,
+    ProtocolParts,
+    ServerProtocolService,
+    WriteRequest,
+    with_run,
+)
+from pyfect.rpc.server._serialized import encode_defect_for
+from pyfect.scope import Scope, fork_scoped, scoped
+from pyfect.socket import CloseEvent, SocketCloseError, SocketError, SocketService, Writer
+from pyfect.socket.server import SocketServer
+from pyfect.stream import Stream, from_iterable, from_queue
+
+_VOID: Effect[None] = succeed(None)
+
+type BufferSize = int | Literal["unbounded"]
+type RequestHeaders = tuple[tuple[str, str], ...]
+
+
+def _with_headers(message: Message, headers: RequestHeaders) -> Message:
+    """A request with the transport's headers added to its own. A name the
+    transport carries replaces the request's: a proxy vouches for the
+    transport's headers, while the message is the caller's (Effect lets the
+    message's win)."""
+    if message.get("_tag") != "Request" or not headers:
+        return message
+    carried = {name.lower() for name, _ in headers}
+    own = [
+        pair
+        for pair in cast("list[Any]", list(message.get("headers") or ()))
+        if not (_is_pair(pair) and str(pair[0]).lower() in carried)
+    ]
+    return {**message, "headers": [*own, *headers]}
+
+
+def _is_pair(pair: object) -> bool:
+    return isinstance(pair, list | tuple) and len(cast("tuple[object, ...]", pair)) == 2
+
+
+# --- sockets ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True, slots=True)
+class SocketProtocol:
+    """A server protocol over sockets: ``on_socket`` serves one connection
+    (until it closes); the ``protocol`` is the server's side."""
+
+    protocol: ServerProtocolService
+    on_socket: Callable[[SocketService, RequestHeaders], Effect[None, Never, Scope]]
+
+
+def make_socket_protocol() -> Effect[SocketProtocol, Never, RpcSerialization]:
+    @gen
+    def make():
+        serialization = yield from RpcSerialization()
+        encode_defect = encode_defect_for(serialization.codec_for)
+        disconnects = yield from Queue[int, Never].unbounded()
+        clients: dict[int, Callable[[Message], Effect[None]]] = {}
+        ids = iter(range(1 << 62))
+        writes: list[WriteRequest] = []
+
+        def on_socket(socket: SocketService, headers: RequestHeaders) -> Effect[None, Never, Scope]:
+            @gen
+            def serve():
+                scope = yield from Scope()
+                parser = serialization.make_unsafe()
+                client_id = next(ids)
+
+                def leave(_: object) -> Effect[None]:
+                    clients.pop(client_id, None)
+                    return disconnects.offer(client_id).as_void()
+
+                yield from scope.add_finalizer(leave)
+                writer = yield from socket.writer
+                clients[client_id] = lambda response: _socket_write(
+                    writer, parser, encode_defect, response
+                )
+
+                def process(data: Frame) -> Effect[None, SocketError]:
+                    try:
+                        decoded = parser.decode(data)
+                    except MaxBufferSizeExceeded as cause:
+                        return writer.write(CloseEvent(1009, str(cause)))
+                    except Exception as cause:
+                        defect = parser.encode(response_defect_encoded(encode_defect(cause)))
+                        return _VOID if defect is None else writer.write(defect)
+
+                    @gen
+                    def each():
+                        for message in decoded:
+                            yield from writes[0](
+                                client_id, _with_headers(cast("Message", message), headers)
+                            )
+
+                    return each
+
+                @gen
+                def read():
+                    reader = yield from socket.reader
+                    while True:
+                        frames = yield from reader.pull
+                        for frame in frames:
+                            yield from process(frame)
+
+                def closed(error: SocketError) -> Effect[None, SocketError]:
+                    return _VOID if isinstance(error.reason, SocketCloseError) else fail(error)
+
+                yield from read.catch_all(closed).or_die()
+
+            return serve
+
+        def send(client_id: int, response: Message) -> Effect[None]:
+            write = clients.get(client_id)
+            return _VOID if write is None else write(response)
+
+        def parts(write: WriteRequest) -> Effect[ProtocolParts]:
+            writes.append(write)
+            return succeed(
+                ProtocolParts(
+                    disconnects=disconnects,
+                    send=send,
+                    end=lambda _: _VOID,
+                    supports_ack=True,
+                    supports_span_propagation=True,
+                    codec_for=serialization.codec_for,
+                )
+            )
+
+        protocol = yield from with_run(parts)
+        return SocketProtocol(protocol, on_socket)
+
+    return make
+
+
+def _socket_write(
+    writer: Writer, parser: Parser, encode_defect: Callable[[object], object], response: Message
+) -> Effect[None]:
+    try:
+        encoded = parser.encode(response)
+    except Exception as cause:
+        encoded = parser.encode(response_defect_encoded(encode_defect(cause)))
+    if encoded is None:
+        return _VOID
+    return writer.write(encoded).or_die()
+
+
+def make_protocol_socket_server() -> Effect[
+    ServerProtocolService, Never, RpcSerialization | SocketServer | Scope
+]:
+    """A server protocol over the connections of the ``SocketServer``."""
+
+    @gen
+    def make():
+        server = yield from SocketServer()
+        made = yield from make_socket_protocol()
+        yield from fork_scoped(server.run(lambda socket: scoped(made.on_socket(socket, ()))))
+        return made.protocol
+
+    return make
+
+
+# --- WebSocket and HTTP ----------------------------------------------------------------
+
+
+type HttpEffect = Effect[HttpServerResponse, Never, Scope | HttpServerRequest]
+
+
+@dataclass(frozen=True, slots=True)
+class HttpProtocol:
+    """A server protocol and the HTTP handler feeding it."""
+
+    protocol: ServerProtocolService
+    http_effect: HttpEffect
+
+
+def _request_headers(request: Any) -> RequestHeaders:
+    return tuple((k, v) for k, v in request.headers.items())
+
+
+def make_protocol_with_http_effect_websocket() -> Effect[HttpProtocol, Never, RpcSerialization]:
+    """A server protocol over WebSockets, upgraded by its HTTP handler."""
+
+    @gen
+    def make():
+        made = yield from make_socket_protocol()
+
+        @gen
+        def handle():
+            request = yield from HttpServerRequest()
+            if not (yield from origin_allowed(request)):
+                return empty(status=403)
+            socket = yield from request.upgrade.or_die()
+            yield from made.on_socket(socket, _request_headers(request))
+            return empty()
+
+        return HttpProtocol(made.protocol, handle)
+
+    return make
+
+
+def _read_frame(request: Any, binary: bool) -> Effect[Frame]:
+    return request.bytes_.or_die() if binary else request.text.or_die()
+
+
+@dataclass(slots=True, eq=False)
+class _HttpClient:
+    write: Callable[[Message], Effect[None]]
+    end: Effect[None]
+
+
+def make_protocol_with_http_effect(
+    *, stream_buffer_size: BufferSize = 16
+) -> Effect[HttpProtocol, Never, RpcSerialization]:
+    """A server protocol over HTTP POSTs: each request body is a client's
+    messages; the response holds its responses (streamed as they come for
+    a framed serialization)."""
+
+    @gen
+    def make():
+        serialization = yield from RpcSerialization()
+        encode_defect = encode_defect_for(serialization.codec_for)
+        disconnects = yield from Queue[int, Never].unbounded()
+        clients: dict[int, _HttpClient] = {}
+        ids = iter(range(1 << 62))
+        writes: list[WriteRequest] = []
+
+        def parts(write: WriteRequest) -> Effect[ProtocolParts]:
+            writes.append(write)
+
+            def send(client_id: int, response: Message) -> Effect[None]:
+                client = clients.get(client_id)
+                return _VOID if client is None else client.write(response)
+
+            def end(client_id: int) -> Effect[None]:
+                client = clients.get(client_id)
+                return _VOID if client is None else client.end
+
+            return succeed(
+                ProtocolParts(
+                    disconnects=disconnects,
+                    send=send,
+                    end=end,
+                    supports_ack=False,
+                    supports_span_propagation=False,
+                    codec_for=serialization.codec_for,
+                )
+            )
+
+        protocol = yield from with_run(parts)
+        handler = _HttpHandler(
+            serialization,
+            encode_defect=encode_defect,
+            disconnects=disconnects,
+            clients=clients,
+            ids=ids,
+            writes=writes,
+            buffer=stream_buffer_size,
+        )
+        return HttpProtocol(protocol, handler.handle())
+
+    return make
+
+
+class _HttpHandler:
+    __slots__ = (
+        "_buffer",
+        "_clients",
+        "_disconnects",
+        "_encode_defect",
+        "_ids",
+        "_serialization",
+        "_writes",
+    )
+
+    def __init__(
+        self,
+        serialization: RpcSerializationService,
+        *,
+        encode_defect: Callable[[object], object],
+        disconnects: Queue[int, Never],
+        clients: dict[int, _HttpClient],
+        ids: Any,
+        writes: list[WriteRequest],
+        buffer: BufferSize,
+    ) -> None:
+        self._serialization = serialization
+        self._encode_defect = encode_defect
+        self._disconnects = disconnects
+        self._clients = clients
+        self._ids = ids
+        self._writes = writes
+        self._buffer = buffer
+
+    def handle(self) -> HttpEffect:
+        serialization = self._serialization
+        framed = serialization.includes_framing
+
+        @gen
+        def run():
+            request = yield from HttpServerRequest()
+            if not (yield from origin_allowed(request)):
+                return empty(status=403)
+            scope = yield from Scope()
+            binary = is_binary(serialization)
+            data = yield from _read_frame(request, binary)
+            client_id: int = next(self._ids)
+            capacity = self._buffer if framed and isinstance(self._buffer, int) else None
+            queue = yield from Queue[Any, Done].make(capacity)
+            parser = serialization.make_unsafe()
+            request_ids: list[RequestId] = []
+            client = _HttpClient(self._writer(queue, parser), queue.end().as_void())
+
+            def leave(_: object) -> Effect[None]:
+                self._clients.pop(client_id, None)
+                self._disconnects.offer_unsafe(client_id)
+                if queue.is_done_unsafe():
+                    return _VOID
+                write = self._writes[0]
+
+                @gen
+                def interrupt_all():
+                    yield from queue.shutdown()
+                    for request_id in request_ids:
+                        yield from write(client_id, {"_tag": "Interrupt", "requestId": request_id})
+
+                return interrupt_all
+
+            yield from scope.add_finalizer(leave)
+            self._clients[client_id] = client
+            headers = _request_headers(request)
+            try:
+                decoded = [cast("Message", m) for m in parser.decode(data)]
+            except Exception as cause:
+                decoded = []
+                yield from client.write(response_defect_encoded(self._encode_defect(cause)))
+            for message in decoded:
+                if message.get("_tag") == "Request":
+                    request_ids.append(message["id"])
+                yield from self._writes[0](client_id, _with_headers(message, headers))
+            yield from self._writes[0](client_id, {"_tag": "Eof"})
+            content_type = serialization.content_type
+            if not framed:
+                collected: Effect[list[Any], Never] = queue.collect()
+                responses: list[Any] = yield from collected.interruptible()
+                encoded = parser.encode(responses) if responses else None
+                return text("" if encoded is None else _as_text(encoded), content_type=content_type)
+            take_all: Effect[list[Any], Never] = queue.take_all().catch_all(
+                lambda _: succeed(list[Any]())
+            )
+            initial: list[Any] = yield from take_all.interruptible()
+            chunks = [encode_text(c) for c in initial]
+            if queue.is_done_unsafe():
+                joined = b"".join(chunks)
+                if binary:
+                    return bytes_response(joined, content_type=content_type)
+                return text(joined.decode(), content_type=content_type)
+            body: Stream[bytes, Never] = from_iterable(chunks).concat(
+                from_queue(queue).map(encode_text)
+            )
+            return stream_response(body, content_type=content_type)
+
+        return run
+
+    def _writer(self, queue: Queue[Any, Done], parser: Parser) -> Callable[[Message], Effect[None]]:
+        if not self._serialization.includes_framing:
+
+            def buffered(response: Message) -> Effect[None]:
+                if response.get("_tag") == "Request" and response.get("isNotification"):
+                    return _VOID
+                return queue.offer(response).as_void()
+
+            return buffered
+
+        def framed(response: Message) -> Effect[None]:
+            def body() -> Effect[None]:
+                try:
+                    encoded = parser.encode(response)
+                except Exception as cause:
+                    encoded = parser.encode(response_defect_encoded(self._encode_defect(cause)))
+                return _VOID if encoded is None else queue.offer(encoded).as_void()
+
+            return suspend(body)
+
+        return framed
+
+
+def _as_text(frame: Frame) -> str:
+    return frame if isinstance(frame, str) else frame.decode()
+
+
+__all__ = [
+    "HttpEffect",
+    "HttpProtocol",
+    "SocketProtocol",
+    "make_protocol_socket_server",
+    "make_protocol_with_http_effect",
+    "make_protocol_with_http_effect_websocket",
+    "make_socket_protocol",
+]

pyfect/rpc/server/_worker.py

@@ -0,0 +1,108 @@
+"""The worker transport of an RPC server (Effect's
+``makeProtocolWorkerRunner``): the server runs inside a worker and serves
+the parent's requests through the ``WorkerRunner``."""
+
+from __future__ import annotations
+
+from typing import Any, cast
+
+from pyfect import schema as S
+from pyfect._core import Effect
+from pyfect.deferred import Deferred
+from pyfect.effect import current_fiber, gen, sync
+from pyfect.layer import Layer
+from pyfect.layer import effect as layer_effect
+from pyfect.logger import log_error
+from pyfect.option import Option, Some
+from pyfect.queue import Queue
+from pyfect.rpc.serialization import CodecFor
+from pyfect.rpc.server._protocol import (
+    Message,
+    ProtocolParts,
+    ServerProtocol,
+    ServerProtocolService,
+    WriteRequest,
+    with_run,
+)
+from pyfect.scheduler import CurrentScheduler
+from pyfect.scope import Scope, fork_scoped
+from pyfect.workers.worker_error import WorkerError
+from pyfect.workers.worker_runner import WorkerRunnerPlatform
+
+
+def _codec_for(schema: Any) -> Any:
+    # Worker messages are values, so they do not depend on RpcSerialization.
+    return S.to_codec_json(schema)
+
+
+def _some(value: object) -> Option[object]:
+    return Some(value)
+
+
+def make_protocol_worker_runner() -> Effect[
+    ServerProtocolService, WorkerError, WorkerRunnerPlatform | Scope
+]:
+    """A server protocol over the ``WorkerRunner`` of the worker this runs
+    in. When the parent closes the worker, the fiber that made the
+    protocol (a ``layer.launch``) is interrupted, so the worker's program
+    ends."""
+
+    def parts(write: WriteRequest) -> Effect[ProtocolParts, WorkerError, Any]:
+        @gen
+        def make():
+            fiber = yield from current_fiber()
+            scheduler = yield from CurrentScheduler()
+            platform = yield from WorkerRunnerPlatform()
+            backing = yield from platform.start()
+            initial = Deferred[object]()
+            client_ids: set[int] = set()
+            disconnects = yield from Queue[int].unbounded()
+
+            def on_message(client_id: int, message: Message) -> Effect[Any]:
+                client_ids.add(client_id)
+                if message.get("_tag") == "InitialMessage":
+                    return initial.succeed(message.get("value"))
+                return write(client_id, message)
+
+            def end_program(_: object) -> Effect[None]:
+                return sync(lambda: scheduler.schedule(lambda: fiber.interrupt_as_unsafe(fiber.id)))
+
+            yield from fork_scoped(
+                backing.run(on_message).tap_cause(log_error).on_exit(end_program)
+            )
+            if backing.disconnects is not None:
+                dequeue = backing.disconnects
+
+                def left(client_id: int) -> Effect[None]:
+                    client_ids.discard(client_id)
+                    return disconnects.offer(client_id).as_void()
+
+                yield from fork_scoped(dequeue.take().flat_map(left).forever())
+
+            def send(client_id: int, response: Message) -> Effect[None]:
+                return backing.send(client_id, response)
+
+            return ProtocolParts(
+                disconnects=disconnects,
+                send=send,
+                end=lambda _: sync(lambda: None),
+                supports_ack=True,
+                supports_span_propagation=True,
+                codec_for=cast("CodecFor", _codec_for),
+                initial_message=initial.wait().map(_some),
+            )
+
+        return make
+
+    return cast(
+        "Effect[ServerProtocolService, WorkerError, WorkerRunnerPlatform | Scope]",
+        with_run(parts),
+    )
+
+
+layer_protocol_worker_runner: Layer[ServerProtocol, WorkerError, WorkerRunnerPlatform] = (
+    layer_effect(ServerProtocol, make_protocol_worker_runner())
+)
+"""A server protocol over the current ``WorkerRunnerPlatform``."""
+
+__all__ = ["layer_protocol_worker_runner", "make_protocol_worker_runner"]

pyfect/rpc/test.py

@@ -0,0 +1,74 @@
+"""An in-memory client of a group's handlers (Effect's ``RpcTest``).
+
+``make_client(Group, layer=GroupLive)`` builds the handlers' layer into the
+caller's ``Scope``, so the handlers' resources (a database, forked workers)
+stay open while the client is used::
+
+    client = yield from rpc_test.make_client(Todos, layer=TodosLive)
+
+Without ``layer`` the handlers come from the context, which must then
+outlive the client: a layer provided to ``make_client(...)`` alone is
+released once the client is made, as in Effect.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Never, overload
+
+from pyfect._core import Effect
+from pyfect._layer_core import Layer
+from pyfect._layer_for_life import with_layer
+from pyfect.effect import gen
+from pyfect.rpc.client._core import ClientCore, OnFromClient
+from pyfect.rpc.client._core import make_no_serialization as make_client_core
+from pyfect.rpc.group import HandlersFor, RpcGroup
+from pyfect.rpc.message import FromServer
+from pyfect.rpc.server._core import ServerOptions
+from pyfect.rpc.server._core import make_no_serialization as make_server_core
+from pyfect.scope import Scope
+
+
+@overload
+def make_client[G: RpcGroup](
+    group: type[G], *, disable_fatal_defects: bool = False
+) -> Effect[G, Never, Scope | HandlersFor[G]]: ...
+@overload
+def make_client[G: RpcGroup, E, R](
+    group: type[G], *, layer: Layer[HandlersFor[G], E, R], disable_fatal_defects: bool = False
+) -> Effect[G, E, R | Scope]: ...
+def make_client(
+    group: type[RpcGroup],
+    *,
+    layer: Layer[Any, Any, Any] | None = None,
+    disable_fatal_defects: bool = False,
+) -> Effect[Any, Any, Any]:
+    """A client of ``group`` calling its handlers directly, without
+    serialization. The handlers (and middleware) are ``layer``, built into
+    the caller's ``Scope``, or else come from the context; the client is
+    valid until that scope closes."""
+    return with_layer(lambda: _make_client(group, disable_fatal_defects), layer)
+
+
+def _make_client(group: type[RpcGroup], disable_fatal_defects: bool) -> Effect[Any, Never, Any]:
+    cores: list[ClientCore] = []
+
+    def on_from_server(response: FromServer) -> Effect[None]:
+        return cores[0].write(response)
+
+    @gen
+    def make():
+        server = yield from make_server_core(
+            group, on_from_server, ServerOptions(disable_fatal_defects=disable_fatal_defects)
+        )
+
+        def on_from_client(sent: OnFromClient) -> Effect[None]:
+            return server.write(0, sent.message)
+
+        client, core = yield from make_client_core(group, on_from_client)
+        cores.append(core)
+        return client
+
+    return make
+
+
+__all__ = ["make_client"]

pyfect/rpc/worker.py

@@ -0,0 +1,92 @@
+"""The initial message of RPC workers (Effect's ``RpcWorker``).
+
+A worker client (``rpc.client.layer_protocol_worker``) can send each
+worker one value before any request: provide ``InitialMessage`` with
+``layer_initial_message(schema, build)``, and read it in the worker with
+``initial_message(schema)``, which waits for it.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from pyfect import schema as S
+from pyfect._core import Effect
+from pyfect._effect_args import require_effect
+from pyfect.context import Context, Tag
+from pyfect.effect import context, fail, provide_context, provide_service, succeed
+from pyfect.errors import NoSuchElementException
+from pyfect.layer import Layer
+from pyfect.layer import effect as layer_effect
+from pyfect.option import Option
+from pyfect.rpc.server._protocol import ServerProtocol
+from pyfect.workers.transferable import Collector, make_collector_unsafe
+
+type Encoded = tuple[object, tuple[object, ...]]
+"""An encoded initial message and the values to transfer with it."""
+
+
+class InitialMessage(Tag[Effect[Encoded]], key="effect/rpc/RpcWorker/InitialMessage"):
+    """Makes the value a worker client sends each worker first."""
+
+
+def make_initial_message[A, E = Any, R = Any](
+    schema: S.Codec[A, Any, Any, Any], eff: Effect[A, E, R]
+) -> Effect[Encoded, E | S.SchemaError, R]:
+    """Run ``eff`` and encode its value with ``schema``'s JSON codec,
+    collecting its transferables."""
+    require_effect(eff, "eff")
+    encode = S.encode_effect(S.to_codec_json(schema))
+
+    def encoded(value: A) -> Effect[Encoded, S.SchemaError, Any]:
+        collector = make_collector_unsafe()
+        return (
+            encode(value)
+            .pipe(provide_service(Collector, collector))
+            .map(lambda data: (data, tuple(collector.clear_unsafe())))
+        )
+
+    return eff.flat_map(encoded)
+
+
+def layer_initial_message[A, R = Any](
+    schema: S.Codec[A, Any, Any, Any], build: Effect[A, Any, R]
+) -> Layer[InitialMessage, Any, R]:
+    """Provide ``InitialMessage`` from ``build``, run with the layer's
+    services each time a worker spawns; an encoding failure is a defect."""
+    require_effect(build, "build")
+
+    def made(services: Context[Any]) -> Effect[Encoded]:
+        return make_initial_message(schema, build).or_die().pipe(provide_context(services))
+
+    return layer_effect(InitialMessage, context().map(made))
+
+
+def _from_option(found: Option[object]) -> Effect[object, NoSuchElementException]:
+    if found is None:
+        return fail(NoSuchElementException())
+    return succeed(found.value)
+
+
+def initial_message[A](
+    schema: S.Codec[A, Any, Any, Any],
+) -> Effect[A, NoSuchElementException | S.SchemaError, ServerProtocol]:
+    """The initial message the client sent this worker, decoded with
+    ``schema``'s JSON codec; fails when the protocol has none."""
+    decode = S.decode_unknown_effect(S.to_codec_json(schema))
+    return (
+        ServerProtocol()
+        .flat_map(lambda protocol: protocol.initial_message)
+        .flat_map(_from_option)
+        .flat_map(decode)
+    )
+
+
+__all__ = [
+    "Encoded",
+    "InitialMessage",
+    "NoSuchElementException",
+    "initial_message",
+    "layer_initial_message",
+    "make_initial_message",
+]