pyfect-dev 0.1.0__py3-none-any.whl

pyfect/pubsub/_hub.py

@@ -0,0 +1,214 @@
+"""The state a ``PubSub`` shares with its subscriptions.
+
+A parked publish owns its surplus items. The hub admits them into the store
+as space frees, oldest publish first, and resumes the publisher once all of
+them are in (Effect's ``BackPressureStrategy``). An interrupted publisher
+withdraws what was not admitted and has delivered the rest.
+
+As in Effect, the takes waiting on a subscription are served in the order
+they parked, and a message published while one waits is read for it at
+once, so it no longer takes room in the store. It is not handed to the
+waiting fiber: it moves to the subscription's ``claimed`` buffer and the
+fiber is woken to collect it. A take that does not wait never reads a
+claimed message meant for a woken one. If the woken fiber is interrupted
+before it collects, the message passes to the next waiting take (or the
+next take), so nothing is lost.
+"""
+
+from __future__ import annotations
+
+from collections import deque
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+from pyfect._core import Effect, _Succeed, _Sync  # pyright: ignore[reportPrivateUsage]
+from pyfect._waiters import WaitList
+from pyfect.option import Option
+from pyfect.pubsub._atomic import Atomic, BackingSubscription
+from pyfect.pubsub._replay import ReplayWindow
+from pyfect.queue import Strategy
+from pyfect.scope import CloseableScope
+
+
+@dataclass(eq=False, slots=True)
+class Publish[A]:
+    """A parked publish. ``items[admitted:]`` are still waiting."""
+
+    items: list[A]
+    resume: Callable[[Effect[bool]], None]
+    admitted: int = 0
+
+
+@dataclass(eq=False, slots=True)
+class Ticket:
+    """A parked take. ``claimed`` once a message was read for it; it is
+    woken without one when the hub ends or the subscription closes."""
+
+    wake: Callable[[], None]
+    claimed: bool = False
+
+
+@dataclass(eq=False, slots=True)
+class Reader[A]:
+    """One subscription's state. A take that does not wait reads ``replay``,
+    then the spare ``claimed`` messages (those no woken take is owed), then
+    ``backing``."""
+
+    backing: BackingSubscription[A]
+    replay: ReplayWindow[A]
+    waiting: deque[Ticket] = field(default_factory=deque[Ticket])
+    claimed: deque[A] = field(default_factory=deque[A])
+    # Woken tickets that have not collected their message yet.
+    owed: int = 0
+    unsubscribed: bool = False
+    # The fiber that unsubscribed; it interrupts the takes then waiting.
+    closed_by: int | None = None
+
+    def size(self) -> int:
+        """What a take that does not wait can read now."""
+        return len(self.claimed) - self.owed + self.backing.size()
+
+    def read_up_to(self, n: int) -> list[A]:
+        """Read up to ``n`` messages for a take that does not wait. Messages
+        owed to woken takes are not read; while takes wait there is nothing
+        else (``dispatch`` gave it to them)."""
+        if n <= 0:
+            return []
+        spare = min(n, len(self.claimed) - self.owed)
+        out = [self.claimed.popleft() for _ in range(spare)]
+        return out + self.backing.poll_up_to(n - spare)
+
+    def collect(self, ticket: Ticket) -> A:
+        """The message read for a woken ``ticket``."""
+        ticket.claimed = False
+        self.owed -= 1
+        return self.claimed.popleft()
+
+    def release(self, ticket: Ticket) -> bool:
+        """A take ended without collecting: its message is spare again.
+        Whether there was one."""
+        if not ticket.claimed or self.unsubscribed:
+            return False
+        ticket.claimed = False
+        self.owed -= 1
+        return True
+
+    def discard(self, ticket: Ticket) -> None:
+        if ticket in self.waiting:
+            self.waiting.remove(ticket)
+
+    def dispatch(self) -> None:
+        """Give each waiting take, oldest first, a message to collect."""
+        while self.waiting:
+            if len(self.claimed) == self.owed:
+                polled = self.backing.poll()
+                if polled is None:
+                    return
+                self.claimed.append(polled.value)
+            ticket = self.waiting.popleft()
+            ticket.claimed = True
+            self.owed += 1
+            ticket.wake()
+
+    def wake_all(self) -> None:
+        """Wake every waiting take without a message, to look again."""
+        tickets, self.waiting = self.waiting, deque[Ticket]()
+        for ticket in tickets:
+            ticket.wake()
+
+
+class Hub[A]:
+    __slots__ = (
+        "atomic",
+        "ended",
+        "is_shutdown",
+        "publishers",
+        "readers",
+        "scope",
+        "shutdown_waiters",
+        "strategy",
+    )
+
+    def __init__(self, atomic: Atomic[A], strategy: Strategy) -> None:
+        self.atomic = atomic
+        self.strategy: Strategy = strategy
+        self.publishers: deque[Publish[A]] = deque()
+        self.readers: dict[Reader[A], None] = {}
+        self.shutdown_waiters = WaitList()
+        # Subscriptions are closed with this scope.
+        self.scope = CloseableScope()
+        self.ended: Option[A] = None
+        self.is_shutdown = False
+
+    def accepts(self) -> bool:
+        """Whether a publish can still deliver."""
+        return not self.is_shutdown and self.ended is None
+
+    def subscribe(self) -> Reader[A]:
+        reader = Reader(self.atomic.subscribe(), self.atomic.replay_window())
+        self.readers[reader] = None
+        return reader
+
+    def unsubscribe(self, reader: Reader[A], by: int) -> None:
+        if reader.unsubscribed:
+            return
+        reader.unsubscribed = True
+        reader.closed_by = by
+        del self.readers[reader]
+        if reader.owed:
+            # Effect runs a woken take before the unsubscribe goes on, so
+            # it still gets its message and, for a batch, what is left.
+            reader.claimed.extend(reader.backing.poll_up_to(reader.backing.size()))
+        else:
+            reader.claimed.clear()
+        reader.backing.unsubscribe()
+        reader.replay.close()
+        reader.wake_all()
+        self.deliver()
+
+    def deliver(self) -> None:
+        """Read new messages for waiting takes and admit parked publishes
+        into the room that frees, until neither makes progress (Effect's
+        ``completeSubscribers`` and ``onPubSubEmptySpace``); then resume the
+        publishers that are done."""
+        done: list[Publish[A]] = []
+        while True:
+            # A woken take runs at once and may subscribe or unsubscribe.
+            for reader in list(self.readers):
+                if not reader.unsubscribed:
+                    reader.dispatch()
+            if not self._admit(done):
+                break
+        for entry in done:
+            entry.resume(_Succeed(True))
+
+    def wake_takers(self) -> None:
+        """The hub ended: every waiting take looks again."""
+        for reader in list(self.readers):
+            reader.wake_all()
+
+    def _admit(self, done: list[Publish[A]]) -> bool:
+        admitted = False
+        while self.publishers and not self.atomic.is_full():
+            entry = self.publishers[0]
+            self.atomic.publish(entry.items[entry.admitted])
+            entry.admitted += 1
+            admitted = True
+            if entry.admitted == len(entry.items):
+                done.append(self.publishers.popleft())
+        return admitted
+
+    def park(self, entry: Publish[A]) -> Effect[None]:
+        """Queue a publish's surplus; the canceler withdraws it."""
+        self.publishers.append(entry)
+        return _Sync(lambda: self._withdraw(entry))
+
+    def _withdraw(self, entry: Publish[A]) -> None:
+        if entry in self.publishers:
+            self.publishers.remove(entry)
+
+    def reject_publishers(self) -> None:
+        """The hub ended: every parked publish returns ``False``."""
+        entries, self.publishers = self.publishers, deque[Publish[A]]()
+        for entry in entries:
+            entry.resume(_Succeed(False))

pyfect/pubsub/_pubsub.py

@@ -0,0 +1,214 @@
+"""``PubSub[A]`` itself. See the package docstring for its behavior."""
+# pyright: reportPrivateUsage=false
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable
+from typing import Any, Never, overload
+
+from pyfect._core import Effect, _Callback, _Gen, _Succeed, _Suspend, _Sync, _WithFiber
+from pyfect._fiber import Fiber
+from pyfect._unsubscripted import Unsubscripted
+from pyfect.cause import interrupt as interrupt_cause
+from pyfect.exit import Failure
+from pyfect.option import Some
+from pyfect.pubsub._atomic import Atomic
+from pyfect.pubsub._hub import Hub, Publish
+from pyfect.pubsub._subscription import Subscription
+from pyfect.queue import Strategy
+from pyfect.scope import Scope
+
+
+class PubSub[A = Unsubscripted]:
+    __slots__ = ("_hub",)
+
+    def __init__(self, atomic: Atomic[A], strategy: Strategy) -> None:
+        self._hub = Hub(atomic, strategy)
+
+    @overload
+    @classmethod
+    def make[T](
+        cls: type[PubSub[Unsubscripted]], atomic: Callable[[], Atomic[T]], strategy: Strategy
+    ) -> Effect[PubSub[T]]: ...
+    @overload
+    @classmethod
+    def make(cls, atomic: Callable[[], Atomic[A]], strategy: Strategy) -> Effect[PubSub[A]]: ...
+    @classmethod
+    def make[T](cls, atomic: Callable[[], Atomic[T]], strategy: Strategy) -> Effect[PubSub[T]]:
+        """A hub over the store ``atomic()`` builds, handling a full store by
+        ``strategy`` (Effect's ``make``, with the built-in strategies)."""
+        return _Sync(lambda: PubSub(atomic(), strategy))
+
+    @classmethod
+    def bounded(cls, capacity: int | None, *, replay: int | None = None) -> Effect[PubSub[A]]:
+        """Publishers wait while ``capacity`` messages are unread. ``replay``
+        messages are kept for new subscribers. A capacity below 1 is a
+        defect; ``None`` is unbounded."""
+        return _Sync(lambda: cls(Atomic(capacity, replay), "suspend"))
+
+    @classmethod
+    def dropping(cls, capacity: int | None, *, replay: int | None = None) -> Effect[PubSub[A]]:
+        """A publish to a full hub drops the message and returns ``False``."""
+        return _Sync(lambda: cls(Atomic(capacity, replay), "dropping"))
+
+    @classmethod
+    def sliding(cls, capacity: int | None, *, replay: int | None = None) -> Effect[PubSub[A]]:
+        """A publish to a full hub drops the oldest unread message."""
+        return _Sync(lambda: cls(Atomic(capacity, replay), "sliding"))
+
+    @classmethod
+    def unbounded(cls, *, replay: int | None = None) -> Effect[PubSub[A]]:
+        return _Sync(lambda: cls(Atomic(None, replay), "dropping"))
+
+    @property
+    def capacity(self) -> int | None:
+        """How many unread messages the hub holds; ``None`` if unbounded."""
+        return self._hub.atomic.capacity
+
+    # --- publishing -------------------------------------------------------
+
+    def publish(self, value: A) -> Effect[bool]:
+        """Publish to every current subscriber. ``False`` if the hub ended or
+        shut down, or a dropping hub is full; a bounded hub waits for room."""
+
+        def register(resume: Callable[[Effect[bool]], None]) -> Effect[None] | None:
+            hub = self._hub
+            if self.publish_unsafe(value):
+                resume(_Succeed(True))
+            elif not hub.accepts() or hub.strategy == "dropping":
+                resume(_Succeed(False))
+            elif hub.strategy == "sliding":
+                resume(_Succeed(self._slide_in([value])))
+            else:
+                return self._park([value], resume)
+            return None
+
+        return _Callback(register)
+
+    def publish_unsafe(self, value: A) -> bool:
+        """Publish without waiting, sliding or dropping: ``False`` if the
+        store is full. Call on the scheduler thread of the hub's fibers."""
+        hub = self._hub
+        if not hub.accepts() or not hub.atomic.publish(value):
+            return False
+        hub.deliver()
+        return True
+
+    def publish_all(self, items: Iterable[A]) -> Effect[bool]:
+        """Publish ``items`` in order; ``True`` if every one was published.
+        A bounded hub publishes what fits and waits to publish the rest."""
+
+        def register(resume: Callable[[Effect[bool]], None]) -> Effect[None] | None:
+            hub = self._hub
+            if not hub.accepts():
+                resume(_Succeed(False))
+                return None
+            surplus = hub.atomic.publish_all(items)
+            hub.deliver()
+            if not surplus:
+                resume(_Succeed(True))
+            elif hub.strategy == "dropping":
+                resume(_Succeed(False))
+            elif hub.strategy == "sliding":
+                resume(_Succeed(self._slide_in(surplus)))
+            else:
+                return self._park(surplus, resume)
+            return None
+
+        return _Callback(register)
+
+    def end(self, value: A) -> Effect[bool]:
+        """Publish nothing more. Subscribers take what they have not read,
+        then ``value`` from every take; waiting publishes return ``False``.
+        ``False`` if the hub already ended or shut down."""
+        return _Sync(lambda: self.end_unsafe(value))
+
+    def end_unsafe(self, value: A) -> bool:
+        hub = self._hub
+        if not hub.accepts():
+            return False
+        hub.ended = Some(value)
+        hub.reject_publishers()
+        hub.wake_takers()
+        return True
+
+    # --- subscribing ------------------------------------------------------
+
+    def subscribe(self) -> Effect[Subscription[A], Never, Scope]:
+        """Subscribe until the current scope closes (or the hub shuts down).
+        The subscription reads the replayed messages, then everything
+        published from now on."""
+        hub = self._hub
+
+        def body():
+            local = yield from Scope()
+            scope = yield from hub.scope.fork()
+            reader = hub.subscribe()
+            yield from scope.add_finalizer(
+                lambda _: _WithFiber(lambda fiber: _Sync(lambda: hub.unsubscribe(reader, fiber.id)))
+            )
+            yield from local.add_finalizer(scope.close)
+            return Subscription[A](hub, reader)
+
+        return _Gen(body).uninterruptible()
+
+    # --- lifecycle --------------------------------------------------------
+
+    def shutdown(self) -> Effect[None]:
+        """Unsubscribe everyone, interrupting their takes. Later publishes
+        return ``False``. Waiting publishes complete: with nobody subscribed
+        the store empties and admits them (as in Effect, whose strategy
+        shutdown then finds no publisher left to interrupt)."""
+
+        def body(fiber: Fiber[Any, Any]) -> Effect[None]:
+            hub = self._hub
+            if hub.is_shutdown:
+                return _Succeed(None)
+            hub.is_shutdown = True
+            hub.shutdown_waiters.notify_all()
+            return hub.scope.close(Failure(interrupt_cause(fiber.id)))
+
+        return _WithFiber(body).uninterruptible()
+
+    def is_shutdown(self) -> Effect[bool]:
+        return _Sync(self.is_shutdown_unsafe)
+
+    def is_shutdown_unsafe(self) -> bool:
+        return self._hub.is_shutdown
+
+    def await_shutdown(self) -> Effect[None]:
+        """Wait until the hub is shut down."""
+        return _Suspend(lambda: self._hub.shutdown_waiters.wait(self.is_shutdown_unsafe))
+
+    def size(self) -> Effect[int]:
+        """How many messages are stored unread (0 once shut down)."""
+        return _Sync(self.size_unsafe)
+
+    def size_unsafe(self) -> int:
+        return 0 if self._hub.is_shutdown else self._hub.atomic.size()
+
+    def is_full(self) -> Effect[bool]:
+        return _Sync(lambda: self.size_unsafe() == self.capacity)
+
+    def is_empty(self) -> Effect[bool]:
+        return _Sync(lambda: self.size_unsafe() == 0)
+
+    # --- internals --------------------------------------------------------
+
+    def _park(self, surplus: list[A], resume: Callable[[Effect[bool]], None]) -> Effect[None]:
+        """Wait to publish ``surplus``. Room a waiting subscriber freed is
+        filled at once, as Effect's ``handleSurplus`` does."""
+        canceler = self._hub.park(Publish(surplus, resume))
+        self._hub.deliver()
+        return canceler
+
+    def _slide_in(self, items: list[A]) -> bool:
+        """Publish each item after dropping the oldest stored message, as
+        Effect's ``SlidingStrategy`` does (also when a waiting subscriber
+        has just read one and made room)."""
+        atomic = self._hub.atomic
+        for item in items:
+            atomic.slide()
+            atomic.publish(item)
+        self._hub.deliver()
+        return True

pyfect/pubsub/_replay.py

@@ -0,0 +1,107 @@
+"""The replay buffer of a ``PubSub`` and the window a new subscriber reads.
+
+The buffer keeps the last ``capacity`` published messages, whether or not
+anyone is subscribed. A subscriber starts with a ``ReplayWindow``: a
+snapshot of those messages, read before anything published after it
+subscribed.
+
+While the window is not drained, the subscriber has read nothing else, so a
+message slid out of the hub before the subscriber read it would leave a gap
+between the window and the subscription. The buffer logs each slid message
+with its publish index, and the window shifts forward over the ones newer
+than its own newest message (Effect's ``ReplayWindowImpl.sync``): the
+subscriber still reads one unbroken run of messages.
+"""
+
+from __future__ import annotations
+
+from collections import deque
+
+from pyfect.option import Option, Some
+
+
+class ReplayBuffer[A]:
+    __slots__ = ("_next_index", "capacity", "items", "slide_count", "slides")
+
+    def __init__(self, capacity: int) -> None:
+        self.capacity = capacity
+        self.items: deque[tuple[int, A]] = deque(maxlen=capacity)
+        self._next_index = 0
+        self.slides: deque[tuple[int, A]] = deque(maxlen=capacity)
+        self.slide_count = 0
+
+    def offer(self, value: A) -> int:
+        """Keep ``value`` (dropping the oldest past capacity); its index."""
+        index = self._next_index
+        self._next_index += 1
+        self.items.append((index, value))
+        return index
+
+    def slide(self, value: A, index: int) -> None:
+        """Log that the message published at ``index`` slid out of the hub."""
+        self.slides.append((index, value))
+        self.slide_count += 1
+
+
+class ReplayWindow[A]:
+    """The replayed messages a subscription reads first."""
+
+    __slots__ = ("_buffer", "_newest", "_slides_seen", "_values")
+
+    def __init__(self, buffer: ReplayBuffer[A] | None) -> None:
+        self._buffer = buffer
+        self._values: deque[A] = deque()
+        self._newest = -1
+        self._slides_seen = 0
+        if buffer is not None:
+            for index, value in buffer.items:
+                self._values.append(value)
+                self._newest = index
+            self._slides_seen = buffer.slide_count
+
+    @property
+    def remaining(self) -> int:
+        return len(self._values)
+
+    def take(self) -> Option[A]:
+        """The next replayed message; ``None`` once drained."""
+        if not self._values:
+            return None
+        self._sync()
+        value = self._values.popleft()
+        if not self._values:
+            self.close()
+        return Some(value)
+
+    def take_n(self, n: int) -> list[A]:
+        """Up to ``n`` replayed messages (none if ``n`` is 0 or less)."""
+        if not self._values or n <= 0:
+            return []
+        self._sync()
+        values = [self._values.popleft() for _ in range(min(n, len(self._values)))]
+        if not self._values:
+            self.close()
+        return values
+
+    def take_all(self) -> list[A]:
+        return self.take_n(len(self._values))
+
+    def close(self) -> None:
+        """Drop what is left; a drained window keeps nothing alive."""
+        self._values.clear()
+        self._buffer = None
+
+    def _sync(self) -> None:
+        buffer = self._buffer
+        if buffer is None:
+            return
+        slides = buffer.slide_count - self._slides_seen
+        self._slides_seen = buffer.slide_count
+        if slides == 0:
+            return
+        logged = list(buffer.slides)
+        for index, value in logged[len(logged) - min(slides, len(logged)) :]:
+            if index > self._newest:
+                self._values.popleft()
+                self._values.append(value)
+                self._newest = index