pyfect-dev 0.1.0__py3-none-any.whl

pyfect/cluster/sharding/_routing.py

@@ -0,0 +1,296 @@
+"""Sending messages: to a local entity, through storage, or to the runner
+that owns the shard, retrying while shards move and abandoning messages
+the runner can no longer route once it shuts down."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+from pyfect._core import Effect
+from pyfect.clock import sleep
+from pyfect.cluster import _cluster_abandon, deliver_at
+from pyfect.cluster.cluster_error import (
+    AlreadyProcessingMessage,
+    EntityNotAssignedToRunner,
+    MailboxFull,
+    MalformedMessage,
+    PersistenceError,
+    RunnerUnavailable,
+)
+from pyfect.cluster.cluster_schema import Persisted
+from pyfect.cluster.envelope import AckChunk, Request
+from pyfect.cluster.message import (
+    Incoming,
+    IncomingEnvelope,
+    IncomingLocal,
+    IncomingRequest,
+    Outgoing,
+    OutgoingRequest,
+)
+from pyfect.cluster.runner_address import RunnerAddress
+from pyfect.duration import Seconds
+from pyfect.effect import callback, catch_tag, die, fail, succeed, suspend
+from pyfect.latch import Latch
+from pyfect.logger import log_debug, log_warning
+
+from ._state import PendingNotification, ShardingState, entity_not_registered
+
+_VOID: Effect[None] = succeed(None)
+
+type SendError = EntityNotAssignedToRunner | MailboxFull | AlreadyProcessingMessage
+type OutgoingError = (
+    MailboxFull | AlreadyProcessingMessage | PersistenceError | EntityNotAssignedToRunner
+)
+
+
+class Routing:
+    __slots__ = ("_state",)
+
+    def __init__(self, state: ShardingState) -> None:
+        self._state = state
+
+    # --- waiting for an entity type -------------------------------------------------------
+
+    def _await_registration_deadline(self, entity_type: str) -> Effect[Any]:
+        """Sleep until the registration deadline shared with the storage
+        loop, checking again in case registration started meanwhile."""
+
+        def run() -> Effect[Any]:
+            remaining = self._state.registration_time_remaining()
+            if remaining > 0:
+                return sleep(Seconds(remaining / 1000)).and_then(
+                    self._await_registration_deadline(entity_type)
+                )
+            return entity_not_registered(entity_type)
+
+        return suspend(run)
+
+    def wait_for_entity_manager(self, entity_type: str) -> Effect[None]:
+        state = self._state
+        latch = state.entity_manager_latches.get(entity_type)
+        if latch is None:
+            latch = Latch()
+            state.entity_manager_latches[entity_type] = latch
+        return latch.wait().race_first(self._await_registration_deadline(entity_type))
+
+    # --- local delivery -------------------------------------------------------------------
+
+    def send_local(self, message: Outgoing | Incoming) -> Effect[None, Any]:
+        state = self._state
+        address = message.envelope.address
+        is_request = isinstance(message, IncomingRequest | OutgoingRequest)
+
+        def loop() -> Effect[None, Any]:
+            # Entities draining from a released shard still take interrupts
+            # and acks.
+            if not state.is_entity_on_local_shards(address) and (
+                is_request or address.shard_id not in state.releasing_shards
+            ):
+                return fail(EntityNotAssignedToRunner(address=address))
+            found = state.entity_managers.get(address.entity_type)
+            if found is None:
+                return self.wait_for_entity_manager(address.entity_type).and_then(suspend(loop))
+            if found.status == "closed" or (
+                found.status == "closing" and isinstance(message, IncomingRequest)
+            ):
+                # Shutting down: no new requests.
+                return fail(EntityNotAssignedToRunner(address=address))
+            if isinstance(message, IncomingRequest | IncomingEnvelope):
+                return found.manager.send(message)
+            return state.runners.send_local(
+                message,
+                send=found.manager.send_local,
+                simulate_remote_serialization=state.config.simulate_remote_serialization,
+            )
+
+        send = suspend(loop)
+        return state.await_shard_lock(address.shard_id).and_then(send) if is_request else send
+
+    def notify_local(
+        self, message: Outgoing | Incoming, discard: bool, *, wait_until_read: bool = False
+    ) -> Effect[None, Any]:
+        state = self._state
+
+        def notify_now() -> Effect[None, EntityNotAssignedToRunner]:
+            if not state.storage_enabled:
+                return die(RuntimeError("Sharding.notifyLocal: storage is disabled"))
+            _read_when_due(state, message)
+            return state.storage_read_latch.open().as_void()
+
+        def notify_storage(_: IncomingLocal) -> Effect[None, EntityNotAssignedToRunner]:
+            return notify_now()
+
+        def loop() -> Effect[None, Any]:  # noqa: PLR0911
+            address = message.envelope.address
+            found = state.entity_managers.get(address.entity_type)
+            if found is None:
+                return self.wait_for_entity_manager(address.entity_type).and_then(suspend(loop))
+            if found.status == "closed":
+                return fail(EntityNotAssignedToRunner(address=address))
+            is_local = state.is_entity_on_local_shards(address)
+            if isinstance(message, IncomingRequest | IncomingEnvelope):
+                if not is_local:
+                    return fail(EntityNotAssignedToRunner(address=address))
+                if isinstance(message, IncomingRequest) and found.manager.is_processing_for(
+                    message, exclude_replies=True
+                ):
+                    return fail(
+                        AlreadyProcessingMessage(
+                            address=address, envelope_id=message.envelope.request_id
+                        )
+                    )
+                if isinstance(message, IncomingRequest) and wait_until_read:
+                    if not state.storage_enabled:
+                        return notify_now()
+                    request = message
+                    return callback(lambda resume: self._park(request, resume))
+                return notify_now()
+            return state.runners.notify_local(
+                message, notify=notify_storage, discard=discard, storage_only=not is_local
+            )
+
+        return suspend(loop)
+
+    def _park(
+        self,
+        message: IncomingRequest,
+        resume: Callable[[Effect[None, EntityNotAssignedToRunner]], None],
+    ) -> None:
+        state = self._state
+        request_id = message.envelope.request_id
+        entry = state.pending_notifications.get(request_id)
+        if entry is not None:
+            previous = entry.resume
+
+            def both(effect: Effect[None, EntityNotAssignedToRunner]) -> None:
+                previous(effect)
+                resume(effect)
+
+            entry.resume = both
+            return
+        state.pending_notifications[request_id] = PendingNotification(resume, message)
+        state.storage_read_latch.open_unsafe()
+
+    # --- outgoing messages ----------------------------------------------------------------
+
+    def send_outgoing(
+        self,
+        message: Outgoing,
+        discard: bool,
+        retries: int | None = None,
+        runner: RunnerAddress | None = None,
+    ) -> Effect[None, OutgoingError]:
+        """Route ``message``; ``runner`` is the runner serving the request,
+        for its interrupts and acks (later attempts use the current shard
+        assignment)."""
+        state = self._state
+        annotations = (
+            message.annotations if isinstance(message, OutgoingRequest) else message.rpc.annotations
+        )
+        persisted = annotations.get(Persisted)
+        should_fail = not discard and (
+            isinstance(message, OutgoingRequest) or isinstance(message.envelope, AckChunk)
+        )
+
+        def abandon(error: EntityNotAssignedToRunner) -> Effect[None, OutgoingError]:
+            if not persisted:
+                if should_fail:
+                    return fail(error)
+                return log_debug(
+                    "Abandoning outgoing message during shutdown", message.envelope.address
+                )
+            persist: Effect[Any, Any] = (
+                state.storage.save_request(message)
+                if isinstance(message, OutgoingRequest)
+                else state.storage.save_envelope(message)
+            )
+            then: Effect[None, Any]
+            if should_fail:
+                # The message is durable, so a transient routing state is
+                # never an error for the caller; the runner can no longer
+                # observe the reply, so the caller is interrupted and the
+                # request is served under the next owner.
+                then = (
+                    _cluster_abandon.interrupt
+                    if isinstance(message, OutgoingRequest)
+                    else fail(error)
+                )
+            else:
+                then = log_warning(
+                    "Persisting outgoing message abandoned during shutdown",
+                    message.envelope.address,
+                )
+            return persist.pipe(catch_tag(MalformedMessage, _die)).and_then(then)
+
+        def attempt() -> Effect[None, Any]:  # noqa: PLR0911
+            address = message.envelope.address
+            if persisted and not state.storage_enabled:
+                return die(
+                    RuntimeError("Sharding.sendOutgoing: Persisted messages require MessageStorage")
+                )
+            if should_fail and state.is_shutdown:
+                return fail(EntityNotAssignedToRunner(address=address))
+            target = runner if runner is not None else state.shard_assignments.get(address.shard_id)
+            is_local = target is not None and state.is_local_runner(target)
+            if persisted:
+                if is_local:
+                    return self.notify_local(message, discard)
+                return state.runners.notify(target, message, discard=discard)
+            if target is None:
+                return fail(EntityNotAssignedToRunner(address=address))
+            if isinstance(message, OutgoingRequest):
+                entry = state.client_requests.get(message.envelope.request_id)
+                if entry is not None:
+                    entry.runner = target
+            if is_local:
+                return self.send_local(message)
+            if discard:
+                return state.runners.notify(target, message, discard=discard)
+            return state.runners.send(target, message)
+
+        def recover(error: Any) -> Effect[None, OutgoingError]:
+            if not isinstance(error, EntityNotAssignedToRunner | RunnerUnavailable):
+                return fail(error)
+            # Abandon the message during teardown: retrying would loop
+            # forever once the runner is shutting down.
+            if isinstance(error, EntityNotAssignedToRunner):
+                target = state.entity_managers.get(message.envelope.address.entity_type)
+                cannot_recover = state.is_shutdown or (
+                    target is not None and target.status != "alive"
+                )
+                if cannot_recover:
+                    return abandon(error)
+            if retries == 0:
+                return die(error)
+            return self.send_outgoing(
+                message, discard, None if retries is None else retries - 1
+            ).delay(state.config.send_retry_interval)
+
+        return suspend(attempt).catch_all(recover)
+
+    def reset(self, request_id: Any, expected_reply_id: Any | None) -> Effect[bool]:
+        return self._state.storage.clear_replies(
+            request_id, expected_reply_id=expected_reply_id
+        ).match_cause(on_failure=lambda _: False, on_success=lambda _: True)
+
+
+def _read_when_due(state: ShardingState, message: Outgoing | Incoming) -> None:
+    """For a request due later (``DeliverAt``, such as a durable clock's
+    wake-up), also read storage when it falls due, on the ``Clock``: the
+    storage poll alone finds it up to ``entity_message_poll_interval`` late
+    (pyfect-only; the poll stays the fallback for another runner's messages)."""
+    request = message.envelope
+    if not isinstance(request, Request):
+        return
+    due = deliver_at.to_millis(request.payload)
+    wait = None if due is None else due - state.now_millis()
+    if wait is not None and wait > 0:
+        state.run_fork(sleep(Seconds(wait / 1000)).and_then(state.storage_read_latch.open()))
+
+
+def _die(error: MalformedMessage) -> Effect[Any]:
+    return die(error)
+
+
+__all__ = ["Routing"]

pyfect/cluster/sharding/_runners_sync.py

@@ -0,0 +1,198 @@
+"""Keeps this runner's view of the runners and shard assignments in sync
+with ``RunnerStorage``, and runs the health checks of the other runners
+as a singleton. It shuts down after the clients, which may still need
+shard assignments while entities shut down."""
+# pyright: reportPrivateUsage=false
+
+from __future__ import annotations
+
+from typing import Any, Never
+
+from pyfect import cause as Cause
+from pyfect._concurrency import for_each
+from pyfect._core import Effect
+from pyfect.clock import sleep
+from pyfect.cluster import cluster_metrics
+from pyfect.cluster.runner import Runner
+from pyfect.cluster.runner_address import RunnerAddress
+from pyfect.cluster.shard_id import ShardId
+from pyfect.context import Context
+from pyfect.duration import Seconds
+from pyfect.effect import gen, succeed
+from pyfect.hash_ring import HashRing
+from pyfect.logger import log_debug, log_warning
+from pyfect.schedule import spaced
+
+from ._state import ShardingState
+
+_VOID: Effect[None] = succeed(None)
+
+
+def _with_timeout[A, E](effect: Effect[A, E]) -> Effect[A, Any]:
+    return effect.timeout(5)
+
+
+def sync_loop(state: ShardingState) -> Effect[Never]:  # noqa: PLR0915
+    self_runner = state.self_runner
+
+    @gen
+    def run():  # noqa: PLR0912, PLR0915
+        rings: dict[str, HashRing[RunnerAddress]] = {}
+        next_runners: dict[Runner, bool] = {}
+        healthy_runners: set[Runner] = set()
+
+        def register_self():
+            if state.is_shutdown or self_runner is None:
+                return
+            yield from log_debug("Registering runner", self_runner)
+            machine_id = yield from _with_timeout(state.runner_storage.register(self_runner, True))
+            yield from state.snowflakes.set_machine_id(machine_id)
+
+        while True:
+            # Make sure this runner is registered.
+            if (
+                self_runner is not None
+                and not state.is_shutdown
+                and self_runner not in state.all_runners
+            ):
+                yield from state.registration_lock.with_permits(gen(register_self))
+
+            runners = yield from _with_timeout(state.runner_storage.get_runners)
+            changed = False
+            for runner, healthy in runners:
+                next_runners[runner] = healthy
+                was_healthy = runner in healthy_runners
+                if not healthy or was_healthy:
+                    if healthy == was_healthy or not was_healthy:
+                        # No change.
+                        state.all_runners.pop(runner, None)
+                    continue
+                changed = True
+                healthy_runners.add(runner)
+                state.all_runners.pop(runner, None)
+                for group in runner.groups:
+                    ring = rings.get(group)
+                    if ring is None:
+                        ring = HashRing[RunnerAddress]()
+                        rings[group] = ring
+                    ring.add(runner.address, weight=runner.weight)
+
+            # Remove the runners no longer present or healthy.
+            for runner in list(state.all_runners):
+                changed = True
+                del state.all_runners[runner]
+                healthy_runners.discard(runner)
+                state.run_fork(state.runners.on_runner_unavailable(runner.address))
+                for group in runner.groups:
+                    found = rings.get(group)
+                    if found is not None:
+                        found.remove(runner.address)
+
+            previous = state.all_runners
+            state.all_runners = next_runners
+            next_runners = previous
+            next_runners.clear()
+            state.healthy_runner_count = len(healthy_runners)
+
+            # Make sure this runner is registered.
+            if (
+                self_runner is not None
+                and not state.is_shutdown
+                and self_runner not in state.all_runners
+            ):
+                continue
+
+            # Recompute the shard assignments when the healthy runners changed.
+            if changed:
+                state.self_shards.clear()
+                for group, ring in rings.items():
+                    assignments = ring.get_shards(state.config.shards_per_group)
+                    for i in range(state.config.shards_per_group):
+                        shard = ShardId(group=group, id=i + 1)
+                        if assignments is not None:
+                            owner = assignments[i]
+                            state.shard_assignments[shard] = owner
+                            if state.shard_locks_healthy_latch.is_open() and state.is_local_runner(
+                                owner
+                            ):
+                                state.self_shards.add(shard)
+                        else:
+                            state.shard_assignments.pop(shard, None)
+                yield from log_debug("New shard assignments", sorted(map(str, state.self_shards)))
+                state.active_shards_latch.open_unsafe()
+                yield from state.notify_shard_lock_waiters()
+                if self_runner is not None:
+                    cluster_metrics.runners_healthy.update_unsafe(
+                        1 if self_runner in healthy_runners else 0, Context.empty()
+                    )
+
+            if self_runner is not None and not healthy_runners:
+                yield from log_warning("No healthy runners available")
+                # To prevent a deadlock, mark this runner healthy to start the
+                # health check singleton again.
+                yield from _with_timeout(
+                    state.runner_storage.set_runner_health(self_runner.address, True)
+                )
+
+            yield from sleep(state.config.refresh_assignments_interval)
+
+    def failed(cause: Cause.Cause[Any]) -> Effect[None]:
+        return log_debug(cause)
+
+    return (
+        run.catch_cause(failed)
+        .repeat(spaced(Seconds(1)))
+        .annotate_logs(
+            {
+                "module": "effect/cluster/Sharding",
+                "fiber": "RunnerStorage sync",
+                "runner": state.runner_address,
+            }
+        )
+        .forever()
+    )
+
+
+def health_check(state: ShardingState) -> Effect[Never]:
+    """Check the other runners' health and record it in storage; the
+    last healthy runner is never marked unhealthy."""
+
+    def check(item: tuple[Runner, bool]) -> Effect[None, Any]:
+        runner, healthy = item
+
+        def decided(alive: bool) -> Effect[None, Any]:
+            if healthy == alive:
+                return _VOID
+            if alive:
+                state.healthy_runner_count += 1
+                return log_debug("Runner is healthy", runner).and_then(
+                    state.runner_storage.set_runner_health(runner.address, alive)
+                )
+            if state.healthy_runner_count <= 1:
+                return _VOID
+            state.healthy_runner_count -= 1
+            return log_debug("Runner is unhealthy", runner).and_then(
+                state.runner_storage.set_runner_health(runner.address, alive)
+            )
+
+        return state.runner_health.is_alive(runner.address).flat_map(decided)
+
+    @gen
+    def run():
+        while True:
+            # Skip health checks when this is the only runner.
+            if len(state.all_runners) > 1:
+                yield from for_each(list(state.all_runners.items()), check, concurrency=10)
+            yield from sleep(state.config.runner_health_check_interval)
+
+    def failed(cause: Cause.Cause[Any]) -> Effect[None]:
+        return log_debug("Runner health check failed", cause)
+
+    return (
+        run.catch_cause(failed)
+        .forever()
+        .annotate_logs({"module": "effect/cluster/Sharding", "fiber": "Runner health check"})
+    )
+
+
+__all__ = ["health_check", "sync_loop"]

pyfect/cluster/sharding/_service.py

@@ -0,0 +1,91 @@
+"""The ``Sharding`` service interface and tag."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, Never, Protocol
+
+from pyfect._core import Effect
+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.shard_id import ShardId
+from pyfect.cluster.snowflake import Snowflake
+from pyfect.context import Tag
+
+if TYPE_CHECKING:
+    from pyfect.cluster.entity import Entity, EntityHandlers, EntityOptions
+    from pyfect.cluster.sharding_registration_event import ShardingRegistrationEvent
+    from pyfect.stream import Stream
+
+
+class ShardingService(Protocol):
+    """What ``Sharding`` provides. ``send`` delivers a message to a local
+    entity, ``notify`` tells it a persisted message is waiting, and
+    ``send_outgoing`` routes a message from a client."""
+
+    @property
+    def get_registration_events(self) -> Stream[ShardingRegistrationEvent]:
+        """The entities and singletons registered from now on."""
+        ...
+
+    def get_shard_id(self, entity_id: EntityId, group: str, /) -> ShardId: ...
+
+    def has_shard_id(self, shard_id: ShardId, /) -> bool: ...
+
+    @property
+    def get_snowflake(self) -> Effect[Snowflake]: ...
+
+    @property
+    def is_shutdown(self) -> Effect[bool]: ...
+
+    def make_client[T: Entity](self, entity: type[T], /) -> Effect[Callable[[str], T]]: ...
+
+    def register_entity(
+        self,
+        entity: type[Entity],
+        build: Effect[EntityHandlers[Any, Any], Never, Any],
+        options: EntityOptions,
+        /,
+    ) -> Effect[None, Never, Any]: ...
+
+    def register_singleton(
+        self, name: str, run: Effect[None, Any, Any], shard_group: str | None, /
+    ) -> Effect[None, Never, Any]: ...
+
+    def send(
+        self, message: Incoming, /
+    ) -> Effect[None, EntityNotAssignedToRunner | MailboxFull | AlreadyProcessingMessage]: ...
+
+    def send_outgoing(
+        self, message: Outgoing, discard: bool, /
+    ) -> Effect[
+        None,
+        MailboxFull | AlreadyProcessingMessage | PersistenceError | EntityNotAssignedToRunner,
+    ]: ...
+
+    def notify(
+        self, message: Incoming, /, *, wait_until_read: bool = False
+    ) -> Effect[None, EntityNotAssignedToRunner | AlreadyProcessingMessage]: ...
+
+    def reset(
+        self, request_id: Snowflake, /, *, expected_reply_id: Snowflake | None = None
+    ) -> Effect[bool]: ...
+
+    @property
+    def poll_storage(self) -> Effect[None]: ...
+
+    @property
+    def active_entity_count(self) -> Effect[int]: ...
+
+
+class Sharding(Tag[ShardingService], key="effect/cluster/Sharding"):
+    """The sharding service of this runner."""
+
+
+__all__ = ["Sharding", "ShardingService"]