pyfect-dev 0.1.0__py3-none-any.whl

pyfect/cluster/sharding/_clients.py

@@ -0,0 +1,271 @@
+"""Entity clients: one RPC client per entity type, whose messages are
+routed by ``send_outgoing`` and whose replies come back through each
+request's ``respond``."""
+# pyright: reportPrivateUsage=false
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import replace
+from typing import TYPE_CHECKING, Any, Never
+
+from pyfect._concurrency import for_each
+from pyfect._core import Effect, _WithFiber
+from pyfect._fiber import Fiber
+from pyfect.cluster import _interruptors as active_teardown
+from pyfect.cluster._resource_map import ResourceMap
+from pyfect.cluster.cluster_schema import (
+    ClientTracingEnabled,
+    Persisted,
+    is_uninterruptible_for_client,
+)
+from pyfect.cluster.entity._services import CurrentAddress
+from pyfect.cluster.entity_address import EntityAddress
+from pyfect.cluster.entity_id import EntityId
+from pyfect.cluster.envelope import AckChunk, Interrupt, make_request
+from pyfect.cluster.message import OutgoingEnvelope, OutgoingRequest, dynamic_annotations
+from pyfect.cluster.reply import Chunk, Reply
+from pyfect.cluster.snowflake import Snowflake
+from pyfect.context import Context, Tag
+from pyfect.effect import gen, succeed, suspend, sync
+from pyfect.rpc import message as rpc_message
+from pyfect.rpc._methods import CallOptions, StreamMethod, UnaryMethod
+from pyfect.rpc.client._core import ClientCore, ClientOptions, OnFromClient
+from pyfect.rpc.client._core import make_no_serialization as make_client_core
+from pyfect.scope import Scope
+
+from ._routing import Routing
+from ._singletons import SingletonShard
+from ._state import ClientRequestEntry, ShardingState
+
+if TYPE_CHECKING:
+    from pyfect.cluster.entity import Entity
+
+_VOID: Effect[None] = succeed(None)
+
+
+class ClientAddress(Tag[EntityAddress], key="effect/cluster/Sharding/ClientAddress"):
+    """The address a client call is for."""
+
+
+def _respond_discard(_: Reply) -> Effect[None]:
+    return _VOID
+
+
+class Clients:
+    __slots__ = ("_clients", "_routing", "_state")
+
+    def __init__(
+        self,
+        state: ShardingState,
+        routing: Routing,
+        clients: ResourceMap[type[Entity], Callable[[str], Any], Never],
+    ) -> None:
+        self._state = state
+        self._routing = routing
+        self._clients = clients
+
+    @staticmethod
+    def make(state: ShardingState, routing: Routing) -> Effect[Clients, Never, Scope]:
+        holder: list[Clients] = []
+
+        def lookup(entity: type[Entity]) -> Effect[Callable[[str], Any], Never, Scope]:
+            return holder[0]._make_client(entity)
+
+        def build(clients: ResourceMap[type[Entity], Callable[[str], Any], Never]) -> Clients:
+            made = Clients(state, routing, clients)
+            holder.append(made)
+            return made
+
+        return ResourceMap.make(lookup).map(build)
+
+    def get(self, entity: type[Entity]) -> Effect[Callable[[str], Any]]:
+        return self._clients.get(entity)
+
+    # --- abandoning callers ---------------------------------------------------------------
+
+    def abandon_undelivered(
+        self, address: EntityAddress, is_delivered: Callable[[Snowflake], bool]
+    ) -> Effect[None]:
+        """During shutdown, abandon local callers parked on persisted
+        requests a retiring entity never started: the next owner serves
+        them. On a runner that keeps running they stay waiting."""
+        state = self._state
+
+        def run() -> Effect[None]:
+            if not state.is_shutdown:
+                return _VOID
+            abandon: list[Effect[None]] = []
+            for request_id, entry in list(state.client_requests.items()):
+                if (
+                    not entry.message.annotations.get(Persisted)
+                    or entry.message.envelope.address != address
+                    or is_delivered(request_id)
+                ):
+                    continue
+                abandon.append(state.storage.unregister_reply_handler(request_id, interrupt=True))
+            return for_each(abandon, lambda effect: effect).as_void()
+
+        return suspend(run)
+
+    # --- one entity type ------------------------------------------------------------------
+
+    def _on_from_client(  # noqa: PLR0911
+        self, entity: type[Entity], cores: list[ClientCore], sent: OnFromClient
+    ) -> Effect[None, Any]:
+        state = self._state
+        routing = self._routing
+        address = sent.context.get_unsafe(ClientAddress)
+        message = sent.message
+        if isinstance(message, rpc_message.Request):
+            request = message
+
+            def body(fiber: Fiber[Any, Any]) -> Effect[None, Any]:
+                request_id = Snowflake(int(request.id))
+                rpc = entity.rpcs[request.tag]
+                envelope = make_request(
+                    request_id=request_id,
+                    address=address,
+                    tag=request.tag,
+                    payload=request.payload,
+                    headers=request.headers,
+                    trace_id=request.trace_id,
+                    span_id=request.span_id,
+                    sampled=request.sampled,
+                )
+                entry: ClientRequestEntry
+
+                def respond(reply: Reply) -> Effect[None]:
+                    if isinstance(reply, Chunk):
+                        entry.last_chunk_id = reply.id
+                        return cores[0].write(
+                            rpc_message.ResponseChunk(0, reply.request_id, reply.values)
+                        )
+                    state.client_requests.pop(reply.request_id, None)
+                    return cores[0].write(rpc_message.ResponseExit(0, reply.request_id, reply.exit))
+
+                outgoing = OutgoingRequest(
+                    envelope=envelope,
+                    last_received_reply=None,
+                    rpc=rpc,
+                    context=fiber._context,
+                    respond=_respond_discard if sent.discard else respond,
+                    annotations=dynamic_annotations(rpc, envelope),
+                )
+                if not sent.discard:
+                    entry = ClientRequestEntry(rpc, fiber._context, outgoing)
+                    state.client_requests[request_id] = entry
+                return routing.send_outgoing(outgoing, sent.discard)
+
+            return _WithFiber(body)
+        if isinstance(message, rpc_message.Ack):
+            request_id = Snowflake(int(message.request_id))
+            entry = state.client_requests.get(request_id)
+            if entry is None or entry.last_chunk_id is None:
+                return _VOID
+            ack = OutgoingEnvelope(
+                envelope=AckChunk(
+                    id=state.snowflakes.next_unsafe(),
+                    address=address,
+                    request_id=request_id,
+                    reply_id=entry.last_chunk_id,
+                ),
+                rpc=entry.rpc,
+            )
+            return routing.send_outgoing(ack, False, None, entry.runner)
+        if isinstance(message, rpc_message.Interrupt):
+            request_id = Snowflake(int(message.request_id))
+            entry = state.client_requests.pop(request_id, None)
+            if entry is None:
+                return _VOID
+            if is_uninterruptible_for_client(entry.message.annotations):
+                return _VOID
+            # Durable messages ignore interrupts on shutdown or from a shard
+            # being reassigned.
+            caller = entry.context.get_option(CurrentAddress)
+            singleton_shard = entry.context.get_option(SingletonShard)
+            transient = (
+                state.is_shutdown
+                or (caller is not None and active_teardown.is_active(caller.value))
+                or (
+                    singleton_shard is not None
+                    and active_teardown.is_shard_active(singleton_shard.value)
+                )
+            )
+            if transient and entry.message.annotations.get(Persisted):
+                return _VOID
+            interrupt = OutgoingEnvelope(
+                envelope=Interrupt(
+                    id=state.snowflakes.next_unsafe(), address=address, request_id=request_id
+                ),
+                rpc=entry.rpc,
+            )
+            return routing.send_outgoing(interrupt, False, 3, entry.runner).ignore()
+        return _VOID
+
+    def _make_client(self, entity: type[Entity]) -> Effect[Callable[[str], Any], Never, Scope]:
+        state = self._state
+        clients = self
+
+        def run():
+            scope = yield from Scope()
+            yield from scope.add_finalizer(
+                lambda _: sync(lambda: active_teardown.release_entity_type(entity.type))
+            )
+            cores: list[ClientCore] = []
+            _, core = yield from make_client_core(
+                entity.protocol,
+                lambda sent: clients._on_from_client(entity, cores, sent),
+                ClientOptions(
+                    span_prefix=f"{entity.type}.client",
+                    disable_tracing=not entity.annotations.get(ClientTracingEnabled),
+                    supports_ack=True,
+                    generate_request_id=state.snowflakes.next_unsafe,
+                ),
+            )
+            cores.append(core)
+            yield from scope.add_finalizer(
+                lambda _: sync(lambda: active_teardown.acquire_entity_type(entity.type))
+            )
+
+            def client_for(entity_id: str) -> Any:
+                id = EntityId(entity_id)
+                address = EntityAddress(
+                    shard_id=state.get_shard_id(id, entity.get_shard_group(id)),
+                    entity_id=id,
+                    entity_type=entity.type,
+                )
+                address_context = Context.make(ClientAddress, address)
+
+                def with_address(call: CallOptions) -> CallOptions:
+                    merged = (
+                        address_context
+                        if call.context is None
+                        else call.context.merge(address_context)
+                    )
+                    return replace(call, context=merged)
+
+                methods: dict[str, Any] = {}
+                for rpc in entity.rpcs.values():
+                    if rpc.is_stream:
+                        methods[rpc.tag] = StreamMethod[Any, Any, Any](
+                            lambda payload, call, rpc=rpc: core.open_stream(
+                                rpc, payload, with_address(call)
+                            )
+                        )
+                    else:
+                        methods[rpc.tag] = UnaryMethod[Any, Any, Any](
+                            lambda payload, call, rpc=rpc: core.call(
+                                rpc, payload, with_address(call)
+                            )
+                        )
+                client = entity()
+                client._rpc_methods = methods
+                return client
+
+            return client_for
+
+        return gen(run)
+
+
+__all__ = ["ClientAddress", "Clients"]

pyfect/cluster/sharding/_make.py

@@ -0,0 +1,388 @@
+"""``Sharding.make`` and ``layer``: wires the loops together in the order
+Effect does, so the scope's finalizers run in its shutdown order (hand
+the shards off first, then the singletons, clients, runner sync, storage
+inbox and shard acquisition, and release every lock last)."""
+# pyright: reportPrivateUsage=false
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, Never
+
+from pyfect import layer as layer_
+from pyfect._core import Effect, _Locally
+from pyfect.clock import Clock
+from pyfect.cluster import _entity_manager, _entity_reaper, cluster_metrics, snowflake
+from pyfect.cluster._entity_reaper import EntityReaper
+from pyfect.cluster._interruptors import around_entity_type
+from pyfect.cluster.cluster_error import (
+    AlreadyProcessingMessage,
+    EntityNotAssignedToRunner,
+    MailboxFull,
+    PersistenceError,
+)
+from pyfect.cluster.entity_id import EntityId
+from pyfect.cluster.message import Incoming, Outgoing
+from pyfect.cluster.message_storage import NOOP, MessageStorage
+from pyfect.cluster.runner import Runner
+from pyfect.cluster.runner_health import RunnerHealth
+from pyfect.cluster.runner_storage import RunnerStorage
+from pyfect.cluster.runners import Runners
+from pyfect.cluster.shard_id import ShardId
+from pyfect.cluster.sharding_config import ShardingConfig, shard_group_config
+from pyfect.cluster.sharding_registration_event import EntityRegistered, ShardingRegistrationEvent
+from pyfect.cluster.snowflake import Generator, Snowflake
+from pyfect.context import Context
+from pyfect.effect import context, gen, succeed, suspend, sync
+from pyfect.exit import Exit, Failure
+from pyfect.fiber_map import FiberMap
+from pyfect.latch import Latch
+from pyfect.layer import Layer
+from pyfect.logger import log_debug
+from pyfect.pubsub import PubSub
+from pyfect.scope import Scope, fork_in
+from pyfect.stream import Stream, from_pubsub
+
+from ._acquisition import ShardAcquisition
+from ._clients import Clients
+from ._routing import Routing
+from ._runners_sync import health_check, sync_loop
+from ._service import Sharding, ShardingService
+from ._singletons import Singletons
+from ._state import EntityManagerState, Residency, ShardingState
+from ._storage_loop import StorageInbox
+
+if TYPE_CHECKING:
+    from pyfect.cluster.entity import Entity, EntityHandlers, EntityOptions
+
+_VOID: Effect[None] = succeed(None)
+
+
+class _Sharding:
+    """The ``Sharding`` service of one runner."""
+
+    __slots__ = ("_clients", "_routing", "_shutdown", "_singletons", "_state")
+
+    def __init__(
+        self,
+        state: ShardingState,
+        routing: Routing,
+        clients: Clients,
+        singletons: Singletons,
+        shutdown: Callable[[Exit[Any, Any] | None], Effect[None]],
+    ) -> None:
+        self._state = state
+        self._routing = routing
+        self._clients = clients
+        self._singletons = singletons
+        self._shutdown = shutdown
+
+    @property
+    def get_registration_events(self) -> Stream[ShardingRegistrationEvent]:
+        return from_pubsub(self._state.events)
+
+    def get_shard_id(self, entity_id: EntityId, group: str, /) -> ShardId:
+        return self._state.get_shard_id(entity_id, group)
+
+    def has_shard_id(self, shard_id: ShardId, /) -> bool:
+        state = self._state
+        return not state.is_shutdown and shard_id in state.acquired_shards
+
+    @property
+    def get_snowflake(self) -> Effect[Snowflake]:
+        return sync(self._state.snowflakes.next_unsafe)
+
+    @property
+    def is_shutdown(self) -> Effect[bool]:
+        return sync(lambda: self._state.is_shutdown)
+
+    def make_client[T: Entity](self, entity: type[T], /) -> Effect[Callable[[str], T]]:
+        return self._clients.get(entity)
+
+    def register_entity(
+        self,
+        entity: type[Entity],
+        build: Effect[EntityHandlers[Any, Any], Never, Any],
+        options: EntityOptions,
+        /,
+    ) -> Effect[None, Never, Any]:
+        return _register_entity(self, entity, build, options)
+
+    def register_singleton(
+        self, name: str, run: Effect[None, Any, Any], shard_group: str | None, /
+    ) -> Effect[None, Never, Any]:
+        return self._singletons.register(name, run, shard_group)
+
+    def send(
+        self, message: Incoming, /
+    ) -> Effect[None, EntityNotAssignedToRunner | MailboxFull | AlreadyProcessingMessage]:
+        return self._routing.send_local(message)
+
+    def send_outgoing(
+        self, message: Outgoing, discard: bool, /
+    ) -> Effect[
+        None,
+        MailboxFull | AlreadyProcessingMessage | PersistenceError | EntityNotAssignedToRunner,
+    ]:
+        return self._routing.send_outgoing(message, discard)
+
+    def notify(
+        self, message: Incoming, /, *, wait_until_read: bool = False
+    ) -> Effect[None, EntityNotAssignedToRunner | AlreadyProcessingMessage]:
+        return self._routing.notify_local(message, False, wait_until_read=wait_until_read)
+
+    def reset(
+        self, request_id: Snowflake, /, *, expected_reply_id: Snowflake | None = None
+    ) -> Effect[bool]:
+        return self._routing.reset(request_id, expected_reply_id)
+
+    @property
+    def poll_storage(self) -> Effect[None]:
+        return self._state.storage_read_latch.open().as_void()
+
+    @property
+    def active_entity_count(self) -> Effect[int]:
+        def count() -> int:
+            return sum(
+                len(m.manager._core.active_servers) for m in self._state.entity_managers.values()
+            )
+
+        return sync(count)
+
+
+def _register_entity(
+    sharding: _Sharding,
+    entity: type[Entity],
+    build: Effect[EntityHandlers[Any, Any], Never, Any],
+    options: EntityOptions,
+) -> Effect[None, Never, Any]:
+    state = sharding._state
+
+    @gen
+    def run():
+        runner_address = state.runner_address
+        if runner_address is None or entity.type in state.entity_managers:
+            return
+        scope = yield from Scope()
+        registration = yield from context()
+        holder: list[EntityManagerState] = []
+
+        def closed(_: Exit[Any, Any]) -> Effect[None]:
+            def mark() -> None:
+                if holder:
+                    holder[0].status = "closed"
+
+            return sync(mark)
+
+        yield from scope.add_finalizer(closed)
+        residency = state.residency
+        manager_options = _entity_manager.ManagerOptions(
+            sharding=sharding,
+            storage=state.storage,
+            runner_address=runner_address,
+            residency=residency,
+            abandon_undelivered=sharding._clients.abandon_undelivered,
+            max_idle_time=options.max_idle_time,
+            concurrency=options.concurrency,
+            mailbox_capacity=options.mailbox_capacity,
+            mailbox_metrics=options.mailbox_metrics,
+            disable_fatal_defects=options.disable_fatal_defects,
+            defect_retry_policy=options.defect_retry_policy,
+            span_attributes=options.span_attributes,
+        )
+        # Registration overrides the construction fallbacks, except for the
+        # services the runner owns.
+        manager_context = (
+            state.services.merge(registration)
+            .add(ShardingConfig, state.config)
+            .add(Clock, state.clock)
+            .add(EntityReaper, state.reaper)
+            .add(Scope, scope)
+            .add(Generator, state.snowflakes)
+        )
+        manager = yield from _Locally(
+            _entity_manager.make(entity, build, manager_options), lambda _: manager_context
+        )
+        entry = EntityManagerState(entity, manager)
+        holder.append(entry)
+
+        def closing(_: Exit[Any, Any]) -> Effect[None]:
+            def run_closing() -> Effect[None]:
+                entry.status = "closing"
+                stop = sharding._shutdown(None) if state.config.preemptive_shutdown else _VOID
+                return around_entity_type(entity.type, stop)
+
+            return suspend(run_closing)
+
+        yield from scope.add_finalizer(closing)
+
+        # Register while storage is idle, so message order is kept.
+        def register_now() -> None:
+            if state.registration_start_millis is None:
+                state.registration_start_millis = state.now_millis()
+            state.entity_managers[entity.type] = entry
+            # Persisted messages read before this registration were reset:
+            # read storage again instead of waiting for the next poll.
+            state.storage_read_latch.open_unsafe()
+            latch = state.entity_manager_latches.pop(entity.type, None)
+            if latch is not None:
+                latch.open_unsafe()
+
+        yield from state.storage_read_lock.with_permits(sync(register_now))
+        yield from state.events.publish(EntityRegistered(entity))
+
+    return run
+
+
+@gen
+def _make():  # noqa: PLR0915
+    config = yield from ShardingConfig()
+    clock = yield from Clock()
+    runners = yield from Runners()
+    runner_health = yield from RunnerHealth()
+    snowflakes = yield from Generator()
+    sharding_scope = yield from Scope()
+    services = yield from context()
+    services = Context({t: s for t, s in services.map_unsafe.items() if t is not Scope})
+    storage = yield from MessageStorage()
+    runner_storage = yield from RunnerStorage()
+    events = yield from PubSub[ShardingRegistrationEvent].unbounded()
+    shard_groups = shard_group_config(config)
+    initial_address = config.runner_address
+    self_runner = (
+        None
+        if initial_address is None
+        else Runner(
+            address=initial_address,
+            groups=shard_groups.assigned,
+            weight=config.runner_shard_weight,
+        )
+    )
+    active_shards_latch = Latch()
+    storage_read_latch = Latch(True)
+    reaper = yield from EntityReaper()
+    state = ShardingState(
+        config=config,
+        shard_groups=shard_groups,
+        clock=clock,
+        runners=runners,
+        runner_health=runner_health,
+        snowflakes=snowflakes,
+        sharding_scope=sharding_scope,
+        services=services,
+        storage=storage,
+        storage_enabled=storage is not NOOP,
+        runner_storage=runner_storage,
+        events=events,
+        reaper=reaper,
+        active_shards_latch=active_shards_latch,
+        storage_read_latch=storage_read_latch,
+        self_runner=self_runner,
+        residency=Residency(config.max_resident_entities, storage_read_latch),
+    )
+
+    yield from sharding_scope.add_finalizer(
+        lambda _: log_debug("Shutdown complete").annotate_logs("module", "effect/cluster/Sharding")
+    )
+
+    # --- shard acquisition ---
+    if initial_address is not None:
+        self_address = initial_address
+        state.update_releasing_gauge()
+        # The locks expire over time, so a failure here is ignored.
+        yield from sharding_scope.add_finalizer(
+            lambda _: state.runner_storage.release_all(self_address).ignore()
+        )
+        # Shard releases fork into a child scope closed after the shutdown
+        # handoff, as forks into the closing sharding scope never run.
+        release_scope = yield from sharding_scope.fork()
+        release_map = yield from FiberMap[ShardId, Any, Any].make()
+        acquisition = ShardAcquisition(state, self_address, release_scope, release_map)
+        yield from fork_in(acquisition.acquisition_loop(), sharding_scope)
+        yield from fork_in(acquisition.refresh_loop(), sharding_scope)
+        # Open the shard latch every poll interval.
+        yield from fork_in(
+            state.active_shards_latch.open().delay(config.entity_message_poll_interval).forever(),
+            sharding_scope,
+        )
+        state.handoff = acquisition.make_handoff()
+
+    # --- storage inbox ---
+    if state.storage_enabled and initial_address is not None:
+        inbox = StorageInbox(state, initial_address)
+        yield from fork_in(
+            state.storage_read_latch.open().delay(config.entity_message_poll_interval).forever(),
+            sharding_scope,
+        )
+        # Forked once the state the loop reads is ready.
+        yield from fork_in(inbox.loop(), sharding_scope)
+
+    routing = Routing(state)
+
+    # --- runner storage sync ---
+    if self_runner is not None:
+        cluster_metrics.runners.update_unsafe(1, Context.empty())
+        cluster_metrics.runners_healthy.update_unsafe(1, Context.empty())
+    yield from fork_in(sync_loop(state), sharding_scope)
+
+    # --- clients ---
+    clients = yield from Clients.make(state, routing)
+
+    # --- singletons ---
+    singleton_fibers = yield from FiberMap[Any, Any, Any].make()
+    singletons = Singletons(state, singleton_fibers)
+    state.singletons = singletons
+    state.singletons_ready.open_unsafe()
+
+    def shutdown(exit: Exit[Any, Any] | None) -> Effect[None]:
+        @gen
+        def run():
+            if exit is not None:
+                yield from log_debug(
+                    "Shutting down", exit.cause if isinstance(exit, Failure) else {}
+                ).annotate_logs("module", "effect/cluster/Sharding")
+            if state.is_shutdown:
+                return
+            state.is_shutdown = True
+            yield from state.notify_shard_lock_waiters()
+            if self_runner is not None:
+                yield from state.registration_lock.with_permits(
+                    state.runner_storage.unregister(self_runner.address)
+                ).ignore()
+
+        return run
+
+    sharding = _Sharding(state, routing, clients, singletons, shutdown)
+
+    # --- runner health checks ---
+    if self_runner is not None:
+        yield from singletons.register(
+            "effect/cluster/Sharding/RunnerHealth", health_check(state), None
+        )
+
+    # --- finalization ---
+    # Hand the shards off before the loops stop.
+    def finalize(exit: Exit[Any, Any]) -> Effect[None]:
+        handoff = state.handoff
+        return shutdown(exit).and_then(_VOID if handoff is None else handoff)
+
+    yield from sharding_scope.add_finalizer(finalize)
+    service: ShardingService = sharding
+    return service
+
+
+make: Effect[ShardingService, Never, Any] = _make
+"""The ``Sharding`` service (Effect's ``Sharding.make``)."""
+
+layer: Layer[
+    Sharding,
+    Never,
+    ShardingConfig | Runners | MessageStorage | RunnerStorage | RunnerHealth,
+] = layer_.effect(Sharding, make).pipe(
+    layer_.provide(layer_.merge(snowflake.layer_generator, _entity_reaper.layer))
+)
+"""The ``Sharding`` service, with its own snowflake generator and entity
+reaper (Effect's ``Sharding.layer``). ``MessageStorage.layer_noop`` makes
+persisted sends defect."""
+
+__all__ = ["layer", "make"]