pyfect-dev 0.1.0__py3-none-any.whl

pyfect/cluster/runner_health.py

@@ -0,0 +1,93 @@
+"""Whether a runner is alive (Effect's ``RunnerHealth``). Sharding keeps
+the shards of a live runner where they are, and moves those of a dead
+one: ``layer_noop`` says every runner is alive, ``layer_ping`` pings it
+through ``Runners``, and ``layer_k8s`` reads its pod's readiness."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Never
+
+from pyfect import layer as layer_
+from pyfect._core import Effect
+from pyfect.cluster.k8s_http_client import K8sHttpClient, Pod, make_get_pods
+from pyfect.cluster.runner_address import RunnerAddress
+from pyfect.cluster.runners import Runners, RunnersService
+from pyfect.context import Tag
+from pyfect.effect import succeed
+from pyfect.exit import is_success
+from pyfect.layer import Layer
+from pyfect.schedule import spaced
+
+
+@dataclass(frozen=True, slots=True)
+class RunnerHealthService:
+    is_alive: Callable[[RunnerAddress], Effect[bool]]
+
+
+class RunnerHealth(Tag[RunnerHealthService], key="effect/cluster/RunnerHealth"):
+    """How this runner checks whether another is alive."""
+
+
+def _alive(_: RunnerAddress) -> Effect[bool]:
+    return succeed(True)
+
+
+layer_noop: Layer[RunnerHealth] = layer_.succeed(RunnerHealth, RunnerHealthService(_alive))
+"""Every runner is alive (for tests and local development)."""
+
+_PING_SCHEDULE = spaced(0.5)
+
+
+def _make_ping(runners: RunnersService) -> RunnerHealthService:
+    def is_alive(address: RunnerAddress) -> Effect[bool]:
+        pinged = runners.ping(address).timeout(10).retry(times=5, schedule=_PING_SCHEDULE)
+        return pinged.exit().map(is_success)
+
+    return RunnerHealthService(is_alive)
+
+
+make_ping: Effect[RunnerHealthService, Never, Runners] = Runners().map(_make_ping)
+"""Alive when a ping answers within ten seconds, retried five times half
+a second apart."""
+
+layer_ping: Layer[RunnerHealth, Never, Runners] = layer_.effect(RunnerHealth, make_ping)
+
+
+def make_k8s(
+    *, namespace: str | None = None, label_selector: str | None = None
+) -> Effect[RunnerHealthService, Never, K8sHttpClient]:
+    """Alive when the pod at the runner's host is ready or initializing;
+    alive too when the Kubernetes API cannot be read."""
+
+    def service(pods: Effect[dict[str, Pod], object]) -> RunnerHealthService:
+        def is_alive(address: RunnerAddress) -> Effect[bool]:
+            def ready(by_ip: dict[str, Pod]) -> bool:
+                pod = by_ip.get(address.host)
+                return pod is not None and pod.is_ready_or_initializing
+
+            return pods.map(ready).catch_cause(lambda _: succeed(True))
+
+        return RunnerHealthService(is_alive)
+
+    return make_get_pods(namespace=namespace, label_selector=label_selector).map(service)
+
+
+def layer_k8s(
+    *, namespace: str | None = None, label_selector: str | None = None
+) -> Layer[RunnerHealth, Never, K8sHttpClient]:
+    """``make_k8s`` as a layer. The ``HttpClient`` must trust the pod CA
+    and the service account must be allowed to list pods."""
+    return layer_.effect(RunnerHealth, make_k8s(namespace=namespace, label_selector=label_selector))
+
+
+__all__ = [
+    "RunnerHealth",
+    "RunnerHealthService",
+    "layer_k8s",
+    "layer_noop",
+    "layer_ping",
+    "make_k8s",
+    "make_ping",
+]

pyfect/cluster/runner_server.py

@@ -0,0 +1,286 @@
+"""The server half of the runner protocol (Effect's ``RunnerServer``):
+handlers for the ``Runners.Rpcs`` that other runners send, handing each
+message to ``Sharding`` and answering persisted requests from
+``MessageStorage``. ``layer`` serves them over the ``ServerProtocol`` in
+the context; ``layer_with_clients`` adds ``Sharding`` and the RPC
+``Runners``, and ``layer_client_only`` gives a process that serves
+nothing and owns no shards its own ``Sharding`` and ``Runners``."""
+# pyright: reportPrivateUsage=false
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, replace
+from typing import Any, Never, cast
+
+from pyfect import layer as layer_
+from pyfect._core import Effect, _WithFiber
+from pyfect._fiber import Fiber
+from pyfect.cause import Cause
+from pyfect.cluster import envelope as envelope_
+from pyfect.cluster import reply as reply_
+from pyfect.cluster import runner_health, runners
+from pyfect.cluster import sharding as sharding_
+from pyfect.cluster.cluster_error import (
+    AlreadyProcessingMessage,
+    EntityNotAssignedToRunner,
+    MalformedMessage,
+)
+from pyfect.cluster.message import Incoming, IncomingEnvelope, IncomingRequest
+from pyfect.cluster.message_storage import MessageStorage, MessageStorageService
+from pyfect.cluster.reply import ReplyWithContext
+from pyfect.cluster.runner_health import RunnerHealth
+from pyfect.cluster.runner_storage import RunnerStorage
+from pyfect.cluster.runners import (
+    EnvelopePayload,
+    NotifyPayload,
+    RequestPayload,
+    RpcClientProtocol,
+    Rpcs,
+    Runners,
+)
+from pyfect.cluster.sharding import Sharding, ShardingService
+from pyfect.cluster.sharding_config import ShardingConfig, ShardingConfigService
+from pyfect.effect import callback, catch_tag, fail_cause, gen, service_option, succeed
+from pyfect.errors import Done
+from pyfect.exit import Exit, Failure
+from pyfect.layer import Layer
+from pyfect.queue import Queue
+from pyfect.rpc import RpcRequest
+from pyfect.rpc.group import HandlersFor
+from pyfect.rpc.serialization import CodecFor
+from pyfect.rpc.server import ServerOptions, ServerProtocol
+from pyfect.rpc.server import layer as server_layer
+from pyfect.scope import CloseableScope, Scope, fork_scoped
+
+_VOID: Effect[None] = succeed(None)
+
+type _Encoded = Effect[reply_.OnWire, EntityNotAssignedToRunner]
+type _Resume = Callable[[_Encoded], None]
+
+
+def _respond_nothing(_: ReplyWithContext) -> Effect[None]:
+    return _VOID
+
+
+def _ignore_already_processing(_: AlreadyProcessingMessage) -> Effect[None]:
+    return _VOID
+
+
+_request_scope: Effect[CloseableScope | None] = service_option(Scope).map(
+    lambda found: None if found is None else found.value
+)
+"""The scope of the request being handled (the server gives each one)."""
+
+
+def _defect_reply(
+    reply: ReplyWithContext, defect: BaseException, codec_for: CodecFor
+) -> Effect[reply_.OnWire]:
+    fallback = ReplyWithContext.from_defect(
+        id=reply.reply.id, request_id=reply.reply.request_id, defect=defect
+    )
+    return reply_.serialize(fallback, codec_for).or_die()
+
+
+@dataclass(frozen=True, slots=True)
+class _Handlers:
+    sharding: ShardingService
+    storage: MessageStorageService
+    codec_for: CodecFor
+
+    def incoming(
+        self,
+        request: envelope_.PartialRequest,
+        caller_scope: CloseableScope | None,
+        respond: Callable[[ReplyWithContext], Effect[None, Any]],
+    ) -> IncomingRequest:
+        return IncomingRequest(
+            envelope=request,
+            last_sent_reply=None,
+            codec_for=self.codec_for,
+            caller_scope=caller_scope,
+            respond=respond,
+        )
+
+    def notify(self, request: RpcRequest[NotifyPayload]) -> Effect[None, Any]:
+        envelope = request.payload["envelope"]
+        message: Incoming = (
+            self.incoming(envelope, None, _respond_nothing)
+            if isinstance(envelope, envelope_.Request)
+            else IncomingEnvelope(envelope=envelope)
+        )
+        if request.payload["persisted"]:
+            return self.sharding.notify(message)
+        return self.sharding.send(message)
+
+    def effect(self, request: RpcRequest[RequestPayload]) -> Effect[reply_.OnWire, Any]:
+        sharding = self.sharding
+
+        @gen
+        def run():
+            caller_scope = yield from _request_scope
+            pending: list[_Encoded] = []
+            resumes: list[_Resume] = []
+
+            def respond(reply: ReplyWithContext) -> Effect[None]:
+                encoded = reply_.serialize_or_defect(reply, self.codec_for)
+                if resumes:
+                    resumes[0](encoded)
+                else:
+                    pending.append(encoded)
+                return _VOID
+
+            message = self.incoming(request.payload["request"], caller_scope, respond)
+            if request.payload["persisted"]:
+                return (yield from self.await_persisted(message, resumes))
+            yield from sharding.send(message)
+
+            def register(resume: _Resume) -> None:
+                if pending:
+                    resume(pending[0])
+                else:
+                    resumes.append(resume)
+
+            return (yield from callback(register))
+
+        return run
+
+    def await_persisted(self, message: IncomingRequest, resumes: list[_Resume]) -> _Encoded:
+        """The reply to a persisted request, through a reply handler on
+        storage, once sharding has read the request."""
+
+        def body(parent: Fiber[Any, Any]) -> _Encoded:
+            def register(resume: _Resume) -> Effect[None]:
+                resumes.append(resume)
+
+                def on_exit(exit: Exit[Any, Any]) -> None:
+                    if isinstance(exit, Failure):
+                        resume(fail_cause(exit.cause))
+
+                handler = parent._fork(
+                    self.storage.register_reply_handler(message),
+                    daemon=True,
+                    start_immediately=True,
+                )
+                handler.add_observer_unsafe(on_exit)
+                notified = self.sharding.notify(message, wait_until_read=True).pipe(
+                    catch_tag(AlreadyProcessingMessage, _ignore_already_processing)
+                )
+                parent._fork(notified, daemon=True, start_immediately=True).add_observer_unsafe(
+                    on_exit
+                )
+                return handler.interrupt()
+
+            return callback(register)
+
+        return _WithFiber(body)
+
+    def stream(
+        self, request: RpcRequest[RequestPayload]
+    ) -> Effect[Queue[reply_.OnWire, EntityNotAssignedToRunner | Done], Any, Scope]:
+        codec_for = self.codec_for
+
+        @gen
+        def run():
+            caller_scope = yield from _request_scope
+            queue = yield from Queue[reply_.OnWire, EntityNotAssignedToRunner | Done].unbounded()
+
+            def offer(encoded: reply_.OnWire, last: bool) -> None:
+                queue.offer_unsafe(encoded)
+                if last:
+                    cast("Queue[reply_.OnWire, Done]", queue).end_unsafe()
+
+            def respond(reply: ReplyWithContext) -> Effect[None]:
+                def malformed(error: MalformedMessage) -> Effect[None]:
+                    # The fallback defect reply is terminal: end the stream.
+                    return _defect_reply(reply, error, codec_for).map(
+                        lambda encoded: offer(encoded, True)
+                    )
+
+                return (
+                    reply_.serialize(reply, codec_for)
+                    .map(lambda encoded: offer(encoded, isinstance(reply.reply, reply_.WithExit)))
+                    .pipe(catch_tag(MalformedMessage, malformed))
+                )
+
+            message = self.incoming(request.payload["request"], caller_scope, respond)
+            if request.payload["persisted"]:
+
+                def failed(cause: Cause[EntityNotAssignedToRunner]) -> Effect[None]:
+                    return queue.fail_cause(cause).as_void()
+
+                handler = self.storage.register_reply_handler(message).on_error(failed)
+                yield from fork_scoped(handler)
+                yield from self.sharding.notify(message, wait_until_read=True)
+            else:
+                yield from self.sharding.send(message)
+            return queue
+
+        return run
+
+    def envelope(self, request: RpcRequest[EnvelopePayload]) -> Effect[None, Any]:
+        return self.sharding.send(IncomingEnvelope(envelope=request.payload["envelope"]))
+
+
+@gen
+def _handlers():
+    sharding = yield from Sharding()
+    storage = yield from MessageStorage()
+    protocol = yield from ServerProtocol()
+    handlers = _Handlers(sharding, storage, protocol.codec_for)
+    return (
+        Rpcs.handlers()
+        .handle(Rpcs.ping)(lambda _: _VOID)
+        .handle(Rpcs.notify)(handlers.notify)
+        .handle(Rpcs.effect)(handlers.effect)
+        .handle(Rpcs.stream)(handlers.stream)
+        .handle(Rpcs.envelope)(handlers.envelope)
+    )
+
+
+layer_handlers: Layer[HandlersFor[Rpcs], Never, Sharding | MessageStorage | ServerProtocol] = (
+    Rpcs.to_layer_effect(_handlers)
+)
+"""The handlers of the runner protocol (Effect's ``layerHandlers``)."""
+
+layer: Layer[Never, Never, ServerProtocol | Sharding | MessageStorage] = server_layer(
+    Rpcs,
+    ServerOptions(span_prefix="RunnerServer", disable_tracing=True, disable_fatal_defects=True),
+).pipe(layer_.provide(layer_handlers))
+"""A runner server over the ``ServerProtocol`` in the context (Effect's
+``RunnerServer.layer``): it answers pings and hands the messages of other
+runners to ``Sharding``."""
+
+
+layer_with_clients: Layer[
+    Sharding | Runners,
+    Never,
+    ServerProtocol
+    | ShardingConfig
+    | RpcClientProtocol
+    | MessageStorage
+    | RunnerStorage
+    | RunnerHealth,
+] = layer.pipe(layer_.provide_merge(sharding_.layer)).pipe(layer_.provide_merge(runners.layer_rpc))
+"""A runner server with the ``Sharding`` and RPC ``Runners`` it serves
+(Effect's ``layerWithClients``)."""
+
+
+def _client_only(config: ShardingConfigService) -> ShardingConfigService:
+    return replace(config, runner_address=None)
+
+
+layer_client_only: Layer[
+    Sharding | Runners,
+    Never,
+    ShardingConfig | RpcClientProtocol | MessageStorage | RunnerStorage,
+] = (
+    sharding_.layer.pipe(layer_.provide_merge(runners.layer_rpc))
+    .pipe(layer_.provide(runner_health.layer_noop))
+    .pipe(layer_.provide(layer_.effect(ShardingConfig, ShardingConfig().map(_client_only))))
+)
+"""``Sharding`` and ``Runners`` for a process that sends to the cluster
+but serves no runner RPCs and owns no shards: its ``runner_address`` is
+cleared (Effect's ``layerClientOnly``)."""
+
+__all__ = ["layer", "layer_client_only", "layer_handlers", "layer_with_clients"]

pyfect/cluster/runner_storage.py

@@ -0,0 +1,212 @@
+"""Where runners register and lock the shards they own (Effect's
+``RunnerStorage``). ``make_memory`` / ``layer_memory`` keep everything in
+the process: one runner, no locking. ``make_encoded`` wraps a driver that
+stores addresses, runners and shard ids as text
+(``pyfect.cluster.sql_runner_storage``)."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable, Sequence
+from dataclasses import dataclass
+from typing import Protocol
+
+from pyfect import layer as layer_
+from pyfect import schema as S
+from pyfect._core import Effect
+from pyfect.cluster import shard_id as shard_id_
+from pyfect.cluster.cluster_error import PersistenceError
+from pyfect.cluster.machine_id import MachineId
+from pyfect.cluster.runner import Runner
+from pyfect.cluster.runner_address import RunnerAddress
+from pyfect.cluster.shard_id import ShardId
+from pyfect.context import Tag
+from pyfect.effect import succeed, sync
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class RunnerStorageService:
+    """``acquire`` takes the locks it can of ``shard_ids`` and gives back
+    the shards now held; ``refresh`` extends the locks still held."""
+
+    register: Callable[[Runner, bool], Effect[MachineId, PersistenceError]]
+    unregister: Callable[[RunnerAddress], Effect[None, PersistenceError]]
+    get_runners: Effect[list[tuple[Runner, bool]], PersistenceError]
+    set_runner_health: Callable[[RunnerAddress, bool], Effect[None, PersistenceError]]
+    acquire: Callable[[RunnerAddress, Iterable[ShardId]], Effect[list[ShardId], PersistenceError]]
+    refresh: Callable[[RunnerAddress, Iterable[ShardId]], Effect[list[ShardId], PersistenceError]]
+    release: Callable[[RunnerAddress, ShardId], Effect[None, PersistenceError]]
+    release_all: Callable[[RunnerAddress], Effect[None, PersistenceError]]
+
+
+class RunnerStorage(Tag[RunnerStorageService], key="effect/cluster/RunnerStorage"):
+    """The runner storage of this cluster."""
+
+
+class EncodedRunnerStorage(Protocol):
+    """What a runner storage driver implements (Effect's
+    ``RunnerStorage.Encoded``): addresses are ``host:port``, runners their
+    JSON text, shards ``group:id``. ``acquire`` gets at least one shard."""
+
+    @property
+    def get_runners(self) -> Effect[list[tuple[str, bool]], PersistenceError]: ...
+
+    def register(
+        self, address: str, runner: str, healthy: bool, /
+    ) -> Effect[int, PersistenceError]: ...
+
+    def unregister(self, address: str, /) -> Effect[None, PersistenceError]: ...
+
+    def set_runner_health(
+        self, address: str, healthy: bool, /
+    ) -> Effect[None, PersistenceError]: ...
+
+    def acquire(
+        self, address: str, shard_ids: Sequence[str], /
+    ) -> Effect[list[str], PersistenceError]: ...
+
+    def refresh(
+        self, address: str, shard_ids: Sequence[str], /
+    ) -> Effect[list[str], PersistenceError]: ...
+
+    def release(self, address: str, shard_id: str, /) -> Effect[None, PersistenceError]: ...
+
+    def release_all(self, address: str, /) -> Effect[None, PersistenceError]: ...
+
+
+_runner_json = S.from_json_string(Runner)
+_encode_runner = S.encode_sync(_runner_json)
+_decode_runner = S.decode_unknown_sync(_runner_json)
+
+
+def encode_runner(runner: Runner) -> str:
+    """A runner as stored: its JSON text."""
+    return _encode_runner(runner)
+
+
+def encode_runner_address(address: RunnerAddress) -> str:
+    return f"{address.host}:{address.port}"
+
+
+def _decode_runners(rows: list[tuple[str, bool]]) -> list[tuple[Runner, bool]]:
+    """The rows that decode (a malformed runner is left out)."""
+    runners: list[tuple[Runner, bool]] = []
+    for text, healthy in rows:
+        try:
+            runners.append((_decode_runner(text), healthy))
+        except S.SchemaError:
+            continue
+    return runners
+
+
+def make_encoded(encoded: EncodedRunnerStorage) -> RunnerStorageService:
+    """A runner storage on the driver ``encoded``."""
+
+    def acquire(
+        address: RunnerAddress, shard_ids: Iterable[ShardId]
+    ) -> Effect[list[ShardId], PersistenceError]:
+        shards = [str(s) for s in shard_ids]
+        if not shards:
+            return succeed([])
+        return encoded.acquire(encode_runner_address(address), shards).map(_shard_ids)
+
+    def refresh(
+        address: RunnerAddress, shard_ids: Iterable[ShardId]
+    ) -> Effect[list[ShardId], PersistenceError]:
+        shards = [str(s) for s in shard_ids]
+        return encoded.refresh(encode_runner_address(address), shards).map(_shard_ids)
+
+    return RunnerStorageService(
+        register=lambda runner, healthy: encoded.register(
+            encode_runner_address(runner.address), encode_runner(runner), healthy
+        ).map(MachineId),
+        unregister=lambda address: encoded.unregister(encode_runner_address(address)),
+        get_runners=encoded.get_runners.map(_decode_runners),
+        set_runner_health=lambda address, healthy: encoded.set_runner_health(
+            encode_runner_address(address), healthy
+        ),
+        acquire=acquire,
+        refresh=refresh,
+        release=lambda address, shard: encoded.release(encode_runner_address(address), str(shard)),
+        release_all=lambda address: encoded.release_all(encode_runner_address(address)),
+    )
+
+
+def _shard_ids(shards: list[str]) -> list[ShardId]:
+    return [shard_id_.from_string(s) for s in shards]
+
+
+def _make_memory() -> RunnerStorageService:
+    runners: dict[RunnerAddress, tuple[Runner, bool]] = {}
+    acquired: list[ShardId] = []
+    next_id = 0
+
+    def register(runner: Runner, healthy: bool) -> Effect[MachineId, PersistenceError]:
+        def run() -> MachineId:
+            nonlocal next_id
+            runners[runner.address] = (runner, healthy)
+            next_id += 1
+            return MachineId(next_id - 1)
+
+        return sync(run)
+
+    def unregister(address: RunnerAddress) -> Effect[None, PersistenceError]:
+        return sync(lambda: runners.pop(address, None)).as_void()
+
+    def set_runner_health(address: RunnerAddress, healthy: bool) -> Effect[None, PersistenceError]:
+        def run() -> None:
+            entry = runners.get(address)
+            if entry is not None:
+                runners[address] = (entry[0], healthy)
+
+        return sync(run)
+
+    def acquire(
+        _address: RunnerAddress, shard_ids: Iterable[ShardId]
+    ) -> Effect[list[ShardId], PersistenceError]:
+        def run() -> list[ShardId]:
+            acquired[:] = shard_ids
+            return list(acquired)
+
+        return sync(run)
+
+    def refresh(
+        _address: RunnerAddress, _shard_ids: Iterable[ShardId]
+    ) -> Effect[list[ShardId], PersistenceError]:
+        return sync(lambda: list(acquired))
+
+    def release(_address: RunnerAddress, _shard_id: ShardId) -> Effect[None, PersistenceError]:
+        return succeed(None)
+
+    def release_all(_address: RunnerAddress) -> Effect[None, PersistenceError]:
+        return succeed(None)
+
+    return RunnerStorageService(
+        register=register,
+        unregister=unregister,
+        get_runners=sync(lambda: list(runners.values())),
+        set_runner_health=set_runner_health,
+        acquire=acquire,
+        refresh=refresh,
+        release=release,
+        release_all=release_all,
+    )
+
+
+def make_memory() -> Effect[RunnerStorageService]:
+    """Runner storage in memory (Effect's ``makeMemory``)."""
+    return sync(_make_memory)
+
+
+layer_memory: layer_.Layer[RunnerStorage] = layer_.effect(RunnerStorage, make_memory())
+
+__all__ = [
+    "EncodedRunnerStorage",
+    "PersistenceError",
+    "RunnerStorage",
+    "RunnerStorageService",
+    "encode_runner",
+    "encode_runner_address",
+    "layer_memory",
+    "make_encoded",
+    "make_memory",
+]

pyfect/cluster/runners/__init__.py

@@ -0,0 +1,55 @@
+"""How runners talk to each other (Effect's ``Runners``): ping a runner,
+send it a message, tell it a persisted message is waiting, and deliver to
+this runner's entities through storage.
+
+``make`` completes a ``RunnersTransport`` with local delivery;
+``layer_noop`` serves a single runner, and ``layer_rpc`` speaks the
+``Rpcs`` protocol to other runners through an ``RpcClientProtocol``.
+"""
+
+from pyfect.cluster.runners._local import make
+from pyfect.cluster.runners._noop import layer_noop, make_noop
+from pyfect.cluster.runners._rpc import (
+    EnvelopePayload,
+    NotifyPayload,
+    ReplyEncodedSchema,
+    RequestPayload,
+    Rpcs,
+    layer_rpc,
+    make_rpc_client,
+    make_rpc_runners,
+)
+from pyfect.cluster.runners._service import (
+    LocalSendError,
+    NotifyLocalFn,
+    RpcClientProtocol,
+    RpcClientProtocolService,
+    Runners,
+    RunnersService,
+    RunnersTransport,
+    SendError,
+    SendLocalFn,
+)
+
+__all__ = [
+    "EnvelopePayload",
+    "LocalSendError",
+    "NotifyLocalFn",
+    "NotifyPayload",
+    "ReplyEncodedSchema",
+    "RequestPayload",
+    "RpcClientProtocol",
+    "RpcClientProtocolService",
+    "Rpcs",
+    "Runners",
+    "RunnersService",
+    "RunnersTransport",
+    "SendError",
+    "SendLocalFn",
+    "layer_noop",
+    "layer_rpc",
+    "make",
+    "make_noop",
+    "make_rpc_client",
+    "make_rpc_runners",
+]