pyfect-dev 0.1.0__py3-none-any.whl

pyfect/rpc/client/_worker.py

@@ -0,0 +1,283 @@
+"""The worker transport of an RPC client (Effect's ``makeProtocolWorker``):
+requests go to a pool of workers, each serving the group with
+``rpc.server.layer_protocol_worker_runner``."""
+
+from __future__ import annotations
+
+import itertools
+from dataclasses import dataclass
+from typing import Any, cast, overload
+
+from pyfect import schedule
+from pyfect import schema as S
+from pyfect._core import Effect
+from pyfect.cause import Cause
+from pyfect.duration import Duration
+from pyfect.effect import gen, service_option, succeed, suspend, sync
+from pyfect.latch import Latch
+from pyfect.layer import Layer
+from pyfect.layer import effect as layer_effect
+from pyfect.pool import Pool
+from pyfect.result import Ok
+from pyfect.rpc.client._protocol import (
+    ClientProtocol,
+    ClientProtocolParts,
+    ClientProtocolService,
+    Message,
+    WriteTo,
+    with_run_client,
+)
+from pyfect.rpc.client._transports import ConnectionHooks
+from pyfect.rpc.client_error import RpcClientDefect, RpcClientError
+from pyfect.rpc.message import MalformedMessage, ServerMessage, decode_from_server
+from pyfect.rpc.serialization import CodecFor
+from pyfect.rpc.worker import InitialMessage
+from pyfect.scope import Scope, add_finalizer, fork_scoped, scoped
+from pyfect.workers.worker import Spawner, Worker, WorkerPlatform, WorkerPlatformService
+from pyfect.workers.worker_error import WorkerError
+
+_VOID: Effect[None] = succeed(None)
+
+type _Backing = Worker[Message, Any]
+
+
+@dataclass(frozen=True, slots=True)
+class _Entry:
+    client_id: int
+    worker: _Backing
+    latch: Latch
+
+
+def _codec_for(schema: Any) -> Any:
+    # Worker messages are values, so they do not depend on RpcSerialization.
+    return S.to_codec_json(schema)
+
+
+class _WorkerClient:
+    __slots__ = ("_client_ids", "_entries", "_ids", "_initial", "_platform", "_write")
+
+    def __init__(
+        self,
+        write: WriteTo,
+        client_ids: set[int],
+        platform: WorkerPlatformService,
+        initial: Effect[tuple[object, tuple[object, ...]]] | None,
+    ) -> None:
+        self._write = write
+        self._client_ids = client_ids
+        self._platform = platform
+        self._initial = initial
+        self._entries: dict[object, _Entry] = {}
+        self._ids = itertools.count()
+
+    def _broadcast(self, response: Message) -> Effect[None]:
+        @gen
+        def deliver():
+            for client_id in list(self._client_ids):
+                yield from self._write(client_id, response)
+
+        return deliver
+
+    def _on_message(self, backing: _Backing, message: object) -> Effect[None]:
+        try:
+            responses = decode_from_server([message])
+        except MalformedMessage as cause:
+            return self._fail(backing, RpcClientDefect("Error decoding message", cause))
+        return _VOID if not responses else self._on_response(responses[0])
+
+    def _on_response(self, response: ServerMessage) -> Effect[None]:
+        if response["_tag"] == "Pong":
+            return _VOID
+        if response["_tag"] == "Defect":
+            for entry in self._entries.values():
+                entry.latch.open_unsafe()
+            self._entries.clear()
+            return self._broadcast(response)
+        request_id = response["requestId"]
+        if response["_tag"] == "Exit":
+            entry = self._entries.pop(request_id, None)
+            if entry is not None:
+                entry.latch.open_unsafe()
+        else:
+            entry = self._entries.get(request_id)
+        if entry is not None:
+            return self._write(entry.client_id, response)
+        return self._broadcast(response)
+
+    def _fail(self, backing: _Backing, reason: Any) -> Effect[None]:
+        """Release the requests on ``backing`` and fail the clients'."""
+        for request_id, entry in list(self._entries.items()):
+            if entry.worker is backing:
+                del self._entries[request_id]
+                entry.latch.open_unsafe()
+        return self._broadcast({"_tag": "ClientProtocolError", "error": RpcClientError(reason)})
+
+    def _on_cause(self, backing: _Backing, cause: Cause[Any]) -> Effect[None]:
+        found = cause.find_error()
+        if isinstance(found, Ok) and isinstance(found.value, WorkerError):
+            reason: Any = found.value.reason
+        else:
+            reason = RpcClientDefect("Error in worker", cause.squash())
+        return self._fail(backing, reason)
+
+    def acquire(self) -> Effect[_Backing, WorkerError, Any]:
+        @gen
+        def acquire():
+            backing: _Backing = yield from self._platform.spawn(next(self._ids))
+            initial = self._initial
+            on_spawn = (
+                None
+                if initial is None
+                else initial.flat_map(
+                    lambda made: backing.send(
+                        {"_tag": "InitialMessage", "value": made[0]}, made[1]
+                    ).or_die()
+                )
+            )
+            running = (
+                backing.run(lambda message: self._on_message(backing, message), on_spawn=on_spawn)
+                .tap_cause(lambda cause: self._on_cause(backing, cause))
+                .retry(schedule.spaced(1.0))
+                .annotate_logs({"module": "RpcClient", "method": "makeProtocolWorker"})
+                .interruptible()
+            )
+            yield from fork_scoped(running)
+            return backing
+
+        return acquire
+
+    def close(self) -> None:
+        for entry in self._entries.values():
+            entry.latch.open_unsafe()
+        self._entries.clear()
+
+    def send_with(self, pool: Pool[_Backing, WorkerError]) -> Any:
+        def send(client_id: int, request: Message) -> Effect[None, RpcClientError]:
+            def body() -> Effect[None, RpcClientError]:
+                tag = request.get("_tag")
+                if tag == "Request":
+
+                    def lent(worker: _Backing) -> Effect[None, WorkerError]:
+                        latch = Latch(False)
+                        self._entries[request["id"]] = _Entry(client_id, worker, latch)
+                        return worker.send(request).flat_map(lambda _: latch.wait())
+
+                    return scoped(pool.get().flat_map(lent)).or_die()
+                if tag == "Interrupt":
+                    entry = self._entries.pop(request["requestId"], None)
+                    if entry is None:
+                        return _VOID
+                    entry.latch.open_unsafe()
+                    return entry.worker.send(request).or_die()
+                if tag == "Ack":
+                    entry = self._entries.get(request["requestId"])
+                    if entry is None:
+                        return _VOID
+                    return entry.worker.send(request).or_die()
+                return _VOID
+
+            return suspend(body)
+
+        return send
+
+
+type _Requirements = Scope | WorkerPlatform | Spawner
+
+
+@overload
+def make_protocol_worker(
+    *, size: int, concurrency: int = 1, target_utilization: float = 1.0
+) -> Effect[ClientProtocolService, WorkerError, _Requirements]: ...
+@overload
+def make_protocol_worker(
+    *,
+    min_size: int,
+    max_size: int,
+    time_to_live: Duration,
+    concurrency: int = 1,
+    target_utilization: float = 1.0,
+) -> Effect[ClientProtocolService, WorkerError, _Requirements]: ...
+def make_protocol_worker(
+    *,
+    size: int | None = None,
+    min_size: int | None = None,
+    max_size: int | None = None,
+    time_to_live: Duration | None = None,
+    concurrency: int = 1,
+    target_utilization: float = 1.0,
+) -> Effect[ClientProtocolService, WorkerError, _Requirements]:
+    """A client protocol over a pool of workers from the ``WorkerPlatform``:
+    ``size`` of them, or between ``min_size`` and ``max_size`` with idle
+    ones above ``min_size`` released after ``time_to_live``. Each worker
+    serves ``concurrency`` requests at a time; a worker that fails is
+    respawned a second later, and its requests get a
+    ``ClientProtocolError``. An ``InitialMessage`` in the context is sent
+    to each worker first."""
+    ranged = (min_size, max_size, time_to_live)
+    sized = size is not None and all(option is None for option in ranged)
+    bounded = size is None and all(option is not None for option in ranged)
+    if not (sized or bounded):
+        raise TypeError(
+            "make_protocol_worker takes size=, or min_size=, max_size= and time_to_live="
+        )
+
+    def parts(
+        write: WriteTo, client_ids: set[int]
+    ) -> Effect[ClientProtocolParts, WorkerError, Any]:
+        @gen
+        def make():
+            platform = yield from WorkerPlatform()
+            initial = yield from service_option(InitialMessage)
+            hooks = yield from ConnectionHooks()
+            client = _WorkerClient(
+                write, client_ids, platform, None if initial is None else initial.value
+            )
+            acquire = client.acquire()
+            if size is not None:
+                pool = yield from Pool.make(
+                    acquire,
+                    size=size,
+                    concurrency=concurrency,
+                    target_utilization=target_utilization,
+                )
+            else:
+                pool = yield from Pool.make_with_ttl(
+                    acquire,
+                    min=cast("int", min_size),
+                    max=cast("int", max_size),
+                    time_to_live=cast("Duration", time_to_live),
+                    concurrency=concurrency,
+                    target_utilization=target_utilization,
+                )
+            yield from add_finalizer(lambda _: sync(client.close))
+            yield from scoped(pool.get())
+            if hooks is not None:
+                yield from hooks.on_connect
+            return ClientProtocolParts(client.send_with(pool), True, cast("CodecFor", _codec_for))
+
+        return make
+
+    return cast("Effect[ClientProtocolService, WorkerError, _Requirements]", with_run_client(parts))
+
+
+@overload
+def layer_protocol_worker(
+    *, size: int, concurrency: int = 1, target_utilization: float = 1.0
+) -> Layer[ClientProtocol, WorkerError, WorkerPlatform | Spawner]: ...
+@overload
+def layer_protocol_worker(
+    *,
+    min_size: int,
+    max_size: int,
+    time_to_live: Duration,
+    concurrency: int = 1,
+    target_utilization: float = 1.0,
+) -> Layer[ClientProtocol, WorkerError, WorkerPlatform | Spawner]: ...
+def layer_protocol_worker(
+    **options: Any,
+) -> Layer[ClientProtocol, WorkerError, WorkerPlatform | Spawner]:
+    """``make_protocol_worker`` as a layer."""
+    return layer_effect(ClientProtocol, make_protocol_worker(**options))
+
+
+__all__ = ["layer_protocol_worker", "make_protocol_worker"]

pyfect/rpc/client_error.py

@@ -0,0 +1,35 @@
+"""The errors of RPC clients (Effect's ``RpcClientError``)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from pyfect.errors import TaggedError
+from pyfect.http.error import HttpClientError
+from pyfect.socket import SocketErrorReason
+
+
+@dataclass(eq=False)
+class RpcClientDefect(TaggedError):
+    """The client could not make sense of what the protocol gave it."""
+
+    message: str
+    cause: object = None
+
+
+type RpcClientErrorReason = SocketErrorReason | HttpClientError | RpcClientDefect
+
+
+@dataclass(eq=False)
+class RpcClientError(TaggedError):
+    """A request could not be sent or answered; ``reason`` says why: the
+    socket's error, the HTTP client's, or an ``RpcClientDefect``."""
+
+    reason: RpcClientErrorReason
+
+    @property
+    def message(self) -> str:
+        return f"{type(self.reason).__name__}: {self.reason.message}"
+
+
+__all__ = ["RpcClientDefect", "RpcClientError", "RpcClientErrorReason"]

pyfect/rpc/group.py

@@ -0,0 +1,296 @@
+"""Groups of RPCs (Effect's ``RpcGroup``) and their handlers.
+
+A group is a class whose attributes are RPCs::
+
+    class Users(RpcGroup):
+        get_user = rpc.make("GetUser", payload=S.Struct(GetUser), success=User)
+
+Its handlers are built one RPC at a time, then made into a layer the
+server needs::
+
+    UsersLive = Users.to_layer(
+        Users.handlers()
+        .handle(Users.get_user)(lambda request: find_user(request.payload["id"]))
+    )
+
+A client of the group is an instance of it (``RpcClient.make(Users)``),
+on which each RPC attribute is the method calling it.
+"""
+
+from __future__ import annotations
+
+import types
+from collections.abc import Callable, Mapping
+from typing import Any, ClassVar, Never, Protocol, Self, cast, get_args, get_origin, overload
+
+from pyfect._core import Effect
+from pyfect._effect_args import require_effect
+from pyfect.context import Context, ServiceMarker
+from pyfect.deferred import Deferred
+from pyfect.effect import context, die, provide_context, succeed
+from pyfect.errors import Done
+from pyfect.layer import Layer
+from pyfect.layer import effect_context as layer_effect_context
+from pyfect.queue import Dequeue
+from pyfect.rpc._rpc import AnyRpc, HandlerEntry, Rpc, RpcRequest, StreamRpc, handler_tag
+from pyfect.rpc.middleware import RpcMiddleware
+from pyfect.scope import Scope
+from pyfect.stream import Stream
+
+
+class HandlersFor[G](ServiceMarker):
+    """What a server of the group ``G`` needs: a handler for each of its
+    RPCs (type-level marker; Effect's ``Rpc.ToHandler``)."""
+
+    @property
+    def group(self) -> G:
+        raise TypeError("HandlersFor is a type-level marker")
+
+
+type HandlerFn = Callable[[RpcRequest[Any]], Any]
+
+
+class Handlers[G, R = Never]:
+    """The handlers of the group ``G`` so far, needing ``R``. Immutable:
+    ``handle`` gives a new one."""
+
+    __slots__ = ("_by_tag", "group")
+
+    def __init__(self, group: type[RpcGroup], by_tag: Mapping[str, HandlerFn]) -> None:
+        self.group = group
+        self._by_tag = dict(by_tag)
+
+    @property
+    def by_tag(self) -> Mapping[str, HandlerFn]:
+        return self._by_tag
+
+    @overload
+    def handle[P, S, E](
+        self, rpc: Rpc[P, S, E, Never], /
+    ) -> Bound[G, RpcRequest[P], S, E, Never, R]: ...
+    @overload
+    def handle[P, S, E, MP, ME, MT](
+        self, rpc: Rpc[P, S, E, RpcMiddleware[MP, ME, Any], MT], /
+    ) -> Bound[G, RpcRequest[P], S, E | ME, MP, R | MT]: ...
+    @overload
+    def handle[P, A, E](
+        self, rpc: StreamRpc[P, A, E, Never], /
+    ) -> StreamBound[G, RpcRequest[P], A, E, Never, R]: ...
+    @overload
+    def handle[P, A, E, MP, ME, MT](
+        self, rpc: StreamRpc[P, A, E, RpcMiddleware[MP, ME, Any], MT], /
+    ) -> StreamBound[G, RpcRequest[P], A, E | ME, MP, R | MT]: ...
+    def handle(self, rpc: AnyRpc, /) -> Any:
+        """The slot of ``rpc``'s handler: call it with the handler, which
+        gets the request (payload, headers, client, ...) and returns the
+        effect to run (for a stream RPC, a stream or a scoped queue). The
+        services ``rpc``'s middleware provides are removed from the
+        handler's requirements, and the middleware is required instead."""
+        if self.group.rpcs.get(rpc.tag) is not rpc:
+            raise ValueError(f"{rpc.tag!r} is not an RPC of {self.group.__name__}")
+        return _Slot(self, rpc)
+
+    def _register(self, rpc: AnyRpc, handler: HandlerFn) -> Handlers[Any, Any]:
+        return Handlers(self.group, {**self._by_tag, rpc.tag: handler})
+
+
+class _Slot:
+    __slots__ = ("handlers", "rpc")
+
+    def __init__(self, handlers: Handlers[Any, Any], rpc: AnyRpc) -> None:
+        self.handlers = handlers
+        self.rpc = rpc
+
+    def __call__(self, handler: HandlerFn, /) -> Handlers[Any, Any]:
+        return self.handlers._register(self.rpc, handler)  # pyright: ignore[reportPrivateUsage]
+
+
+class Bound[G, X, S, E, MP, RR](Protocol):
+    """The handler slot of a unary RPC: request ``X``, success ``S``,
+    errors ``E``; its middleware provides ``MP``; ``RR`` is what the
+    handlers need so far. They are fixed by ``handle``, before the handler
+    is seen (in one call ``MP`` would be solved from the handler too)."""
+
+    @overload
+    def __call__(
+        self, handler: Callable[[X], Effect[S | Deferred[S, E], E, MP]], /
+    ) -> Handlers[G, RR]: ...
+    @overload
+    def __call__[R1](
+        self, handler: Callable[[X], Effect[S | Deferred[S, E], E, R1 | MP]], /
+    ) -> Handlers[G, RR | R1]: ...
+
+
+class StreamBound[G, X, A, E, MP, RR](Protocol):
+    """The handler slot of a stream RPC (see ``Bound``): the handler gives
+    a stream, or a scoped queue that ends with ``Done``."""
+
+    @overload
+    def __call__(self, handler: Callable[[X], Stream[A, E, MP]], /) -> Handlers[G, RR]: ...
+    @overload
+    def __call__[R1](
+        self, handler: Callable[[X], Stream[A, E, R1 | MP]], /
+    ) -> Handlers[G, RR | R1]: ...
+    @overload
+    def __call__(
+        self, handler: Callable[[X], Effect[Dequeue[A, E | Done], E, MP | Scope]], /
+    ) -> Handlers[G, RR]: ...
+    @overload
+    def __call__[R1](
+        self, handler: Callable[[X], Effect[Dequeue[A, E | Done], E, R1 | MP | Scope]], /
+    ) -> Handlers[G, RR | R1]: ...
+
+
+class RpcGroup:
+    """A group of RPCs: subclass it with the RPCs as attributes."""
+
+    rpcs: ClassVar[Mapping[str, AnyRpc]] = {}
+    _wraps_rpcs: ClassVar[bool] = False
+    _rpc_methods: dict[str, Any]
+
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        super().__init_subclass__(**kwargs)
+        for base in types.get_original_bases(cls):
+            origin = get_origin(base)
+            if isinstance(origin, type) and issubclass(origin, RpcGroup) and origin._wraps_rpcs:
+                _add_middleware(cls, get_args(base)[0])
+        found: dict[str, AnyRpc] = {}
+        for klass in reversed(cls.__mro__):
+            for value in vars(klass).values():
+                if isinstance(value, AnyRpc):
+                    found[value.tag] = value
+        cls.rpcs = found
+
+    @classmethod
+    def handlers(cls) -> Handlers[Self]:
+        """No handlers yet: add one per RPC with ``handle``."""
+        return Handlers(cls, {})
+
+    @classmethod
+    def to_layer[R](cls, handlers: Handlers[Self, R], /) -> Layer[HandlersFor[Self], Never, R]:
+        """The layer of the group's handlers (Effect's ``toLayer``). Each
+        handler runs with the services of the layer's context. An RPC left
+        without a handler is a ``ValueError`` here (Effect rejects it at
+        compile time; pyright cannot name the RPCs not yet handled)."""
+        message = cls._missing_handlers(handlers)
+        if message is not None:
+            raise ValueError(message)
+        return cls.to_layer_effect(succeed(handlers))
+
+    @classmethod
+    def _missing_handlers(cls, handlers: Handlers[Any, Any]) -> str | None:
+        missing = [
+            f"{tag} ({cls.__name__}.{_attribute_of(cls, rpc)})"
+            for tag, rpc in cls.rpcs.items()
+            if tag not in handlers.by_tag
+        ]
+        if not missing:
+            return None
+        return f"{cls.__name__} has no handler for {', '.join(missing)}"
+
+    @classmethod
+    def to_layer_effect[EX, RR = Never](
+        cls, handlers: Effect[Handlers[Self, RR], EX, RR], /
+    ) -> Layer[HandlersFor[Self], EX, RR]:
+        """``to_layer`` with the handlers made by an effect."""
+        require_effect(handlers, "handlers")
+        return layer_effect_context(handlers.flat_map(cls._handler_context))
+
+    @classmethod
+    def _handler_context(cls, handlers: Handlers[Any, Any]) -> Effect[Context[Any]]:
+        message = cls._missing_handlers(handlers)
+        if message is not None:
+            return die(ValueError(message))
+
+        def entries(services: Context[Any]) -> Context[Any]:
+            kept = Context({t: s for t, s in services.map_unsafe.items() if t is not Scope})
+            return Context(
+                {
+                    handler_tag(rpc): HandlerEntry(rpc.tag, handlers.by_tag[rpc.tag], kept)
+                    for rpc in cls.rpcs.values()
+                }
+            )
+
+        return context().map(entries)
+
+
+def _attribute_of(group: type, rpc: AnyRpc) -> str:
+    """The attribute name ``rpc`` has on ``group`` (its tag if none)."""
+    for klass in group.__mro__:
+        for name, value in vars(klass).items():
+            if value is rpc:
+                return name
+    return rpc.tag
+
+
+class HandlerOf[T](ServiceMarker):
+    """A handler of the RPC ``T`` in the context (type-level marker)."""
+
+    @property
+    def rpc(self) -> T:
+        raise TypeError("HandlerOf is a type-level marker")
+
+
+def to_layer_handler[P, S, E, R2](
+    rpc: Rpc[P, S, E, Never],
+    handler: Callable[[RpcRequest[P]], Effect[S | Deferred[S, E], E, R2]],
+) -> Layer[HandlerOf[Rpc[P, S, E]], Never, R2]:
+    """The layer of one RPC's handler (Effect's ``toLayerHandler``); it
+    runs with the services of the layer's context."""
+
+    def entry(services: Context[Any]) -> Context[Any]:
+        kept = Context({t: s for t, s in services.map_unsafe.items() if t is not Scope})
+        return Context({handler_tag(rpc): HandlerEntry(rpc.tag, handler, kept)})
+
+    return layer_effect_context(context().map(entry))
+
+
+def access_handler[P, S, E](
+    rpc: Rpc[P, S, E, Never],
+) -> Effect[Callable[[RpcRequest[P]], Effect[S, E]], Never, HandlerOf[Rpc[P, S, E]]]:
+    """The handler of ``rpc`` from the context, running with its layer's
+    services (Effect's ``accessHandler``)."""
+
+    def found(services: Context[Any]) -> Callable[[RpcRequest[P]], Effect[S, E]]:
+        entry: HandlerEntry = services.get_unsafe(handler_tag(rpc))
+
+        def run(request: RpcRequest[P]) -> Effect[S, E]:
+            return cast("Effect[S, E]", entry.handler(request)).pipe(provide_context(entry.context))
+
+        return run
+
+    return cast(
+        "Effect[Callable[[RpcRequest[P]], Effect[S, E]], Never, HandlerOf[Rpc[P, S, E]]]",
+        context().map(found),
+    )
+
+
+def _add_middleware(group: type[RpcGroup], middleware: Any) -> None:
+    for name, value in list(vars(group).items()):
+        if isinstance(value, AnyRpc):
+            setattr(group, name, value.with_middleware(middleware))
+
+
+class RpcGroupWith[GM](RpcGroup):
+    """A group whose RPCs are all wrapped by ``GM`` (Effect's
+    ``RpcGroup.middleware``)::
+
+        class Users(RpcGroupWith[Auth]):
+            get = rpc.make("Users.get", ...)
+    """
+
+    _wraps_rpcs: ClassVar[bool] = True
+
+
+__all__ = [
+    "Bound",
+    "HandlerFn",
+    "HandlerOf",
+    "Handlers",
+    "HandlersFor",
+    "RpcGroup",
+    "RpcGroupWith",
+    "StreamBound",
+    "access_handler",
+    "to_layer_handler",
+]