orchestrata_ultimate.gateway.pubsub

Reference › orchestrata_ultimate › gateway › pubsub

pubsub is the in-process topic subscription registry of the gateway (SPEC-014 US-001). PubSubManager stores Subscription records from models behind an asyncio lock. It is used by protocol for broadcasts, by rpc, by several handlers, and by the orchestrator's pipeline kernel. See gateway.

• subscribe validates its inputs, records a subscription for a connection with an fnmatch-style topic pattern and optional filters, updates the subscription gauge in metrics, and returns the subscription id.
• unsubscribe, unsubscribe_by_id and unsubscribe_all remove subscriptions by topic, by id or for a whole connection.
• get_matching_connections takes a snapshot under the lock and returns each connection once if its pattern matches the topic. When a payload is passed it also applies the subscription's filters through evaluate_filters.
• evaluate_filters is the predicate (SPEC-014 FR-009): an empty filter set matches everything, keys combine with AND, subject is a glob match, volk_id matches either the to or from field, and any other key needs equality with the payload field.
• publish returns the number of matching connections; actual delivery is done by the message processor.
• get_subscriptions, get_subscription_by_id and get_total_subscription_count support queries.

At a glance

Class diagram

Dependency diagram

API

Classes

PubSubManager

Line 58

Tracks topic subscriptions per connection.

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/gateway/pubsub.py

"""In-process pub/sub for WebSocket gateway topics."""

import asyncio
import fnmatch
import logging
from typing import Any

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.gateway.metrics import (
    subscription_added,
    subscription_removed,
)
from orchestrata_ultimate.gateway.models import Subscription

logger = logging.getLogger(__name__)


@traces(spec="014", story="US-001", reqs=["REQ-014-003", "REQ-014-005", "REQ-014-007"])
def evaluate_filters(payload: dict[str, Any], filters: dict[str, Any]) -> bool:
    """Evaluate whether *payload* satisfies a subscription's *filters*.

    Implements SPEC-014 FR-009's key-value predicates. An empty (or falsy)
    ``filters`` mapping matches every payload -- subscriptions created
    without filters keep receiving everything on their topic, unchanged.

    Multiple filter keys combine with AND semantics (REQ-014-007): every
    key must match for the payload to pass. Most keys are matched by plain
    equality against the top-level payload field of the same name. Two
    keys get spec-defined richer semantics:

    - ``subject``: fnmatch-style glob match against ``payload["subject"]``
      (REQ-014-005 -- "subject pattern (glob match)").
    - ``volk_id``: matches if *either* ``payload["to"]`` or
      ``payload["from"]`` equals the filter value (REQ-014-003 -- "volk_id
      (TO or FROM)"), since CXP-shaped events carry both fields.

    Args:
        payload: The event payload being considered for delivery.
        filters: The subscription's key-value filter predicates.

    Returns:
        True if payload matches every filter key (or filters is empty).
    """
    if not filters:
        return True
    for key, expected in filters.items():
        if key == "volk_id":
            if payload.get("to") != expected and payload.get("from") != expected:
                return False
        elif key == "subject":
            if not fnmatch.fnmatch(str(payload.get("subject", "")), str(expected)):
                return False
        elif payload.get(key) != expected:
            return False
    return True


class PubSubManager:
    """Tracks topic subscriptions per connection.

    Pattern matching uses fnmatch-style wildcards
    (e.g. ``agent.*`` matches ``agent.status``).

    Subscriptions are identified by a unique subscription_id
    and support optional key-value filters for event matching.
    """

    def __init__(self) -> None:
        """Initialise empty subscription store."""
        self._subs: dict[str, set[str]] = {}
        self._subscriptions: dict[str, Subscription] = {}
        self._lock = asyncio.Lock()

    async def subscribe(
        self,
        connection_id: str,
        topic_pattern: str,
        filters: dict[str, Any] | None = None,
    ) -> str:
        """Subscribe *connection_id* to *topic_pattern*.

        Args:
            connection_id: Connection to subscribe.
            topic_pattern: fnmatch-style topic pattern.
            filters: Optional key-value filters for events.

        Returns:
            The subscription_id for this subscription.

        Raises:
            TypeError: If arguments are not strings.
            ValueError: If arguments are empty.
        """
        if not isinstance(connection_id, str):
            raise TypeError("connection_id must be a str")
        if not connection_id:
            raise ValueError("connection_id must not be empty")
        if not isinstance(topic_pattern, str):
            raise TypeError("topic_pattern must be a str")
        if not topic_pattern:
            raise ValueError("topic_pattern must not be empty")
        async with self._lock:
            sub = Subscription(
                connection_id=connection_id,
                topic=topic_pattern,
                filters=filters or {},
            )
            self._subscriptions[sub.subscription_id] = sub
            patterns = self._subs.setdefault(connection_id, set())
            patterns.add(topic_pattern)
            logger.info(
                "%s subscribed to %s (sub=%s)",
                connection_id,
                topic_pattern,
                sub.subscription_id,
            )
            # REQ-014-147: live subscription count, labeled by topic.
            subscription_added(topic_pattern)
            return sub.subscription_id

    async def unsubscribe(
        self,
        connection_id: str,
        topic_pattern: str,
    ) -> bool:
        """Remove a single subscription by topic pattern.

        Args:
            connection_id: Connection to unsubscribe.
            topic_pattern: Exact pattern to remove.

        Returns:
            True if the pattern was present.

        Raises:
            TypeError: If arguments are not strings.
            ValueError: If arguments are empty.
        """
        if not isinstance(connection_id, str):
            raise TypeError("connection_id must be a str")
        if not connection_id:
            raise ValueError("connection_id must not be empty")
        if not isinstance(topic_pattern, str):
            raise TypeError("topic_pattern must be a str")
        if not topic_pattern:
            raise ValueError("topic_pattern must not be empty")
        async with self._lock:
            patterns = self._subs.get(connection_id)
            if not patterns:
                return False
            if topic_pattern not in patterns:
                return False
            patterns.discard(topic_pattern)
            if not patterns:
                del self._subs[connection_id]
            self._remove_sub_records(connection_id, topic_pattern)
            return True

    async def unsubscribe_by_id(
        self,
        subscription_id: str,
    ) -> bool:
        """Remove a subscription by its subscription_id.

        Args:
            subscription_id: ID returned from subscribe().

        Returns:
            True if the subscription existed and was removed.

        Raises:
            TypeError: If subscription_id is not a str.
            ValueError: If subscription_id is empty.
        """
        if not isinstance(subscription_id, str):
            raise TypeError("subscription_id must be a str")
        if not subscription_id:
            raise ValueError("subscription_id must not be empty")
        async with self._lock:
            sub = self._subscriptions.pop(subscription_id, None)
            if not sub:
                return False
            patterns = self._subs.get(sub.connection_id)
            if patterns:
                patterns.discard(sub.topic)
                if not patterns:
                    del self._subs[sub.connection_id]
            # REQ-014-147: one fewer live subscription on this topic.
            subscription_removed(sub.topic)
            return True

    @traces(
        spec="014",
        story="US-012",
        reqs=["REQ-014-196"],
    )
    async def unsubscribe_all(self, connection_id: str) -> int:
        """Remove every subscription for a connection.

        Args:
            connection_id: Connection to clear.

        Returns:
            Number of patterns removed.

        Raises:
            TypeError: If connection_id is not a str.
            ValueError: If connection_id is empty.
        """
        if not isinstance(connection_id, str):
            raise TypeError("connection_id must be a str")
        if not connection_id:
            raise ValueError("connection_id must not be empty")
        async with self._lock:
            patterns = self._subs.pop(connection_id, set())
            to_remove = [
                sid
                for sid, s in self._subscriptions.items()
                if s.connection_id == connection_id
            ]
            # REQ-014-147: one decrement per removed subscription record's
            # topic (not per pattern -- a connection can hold more than one
            # subscription record for the same topic pattern).
            for sid in to_remove:
                subscription_removed(self._subscriptions[sid].topic)
                del self._subscriptions[sid]
            return len(patterns)

    @traces(spec="014", story="US-001", reqs=["REQ-014-007"])
    async def get_matching_connections(
        self,
        topic: str,
        payload: dict[str, Any] | None = None,
    ) -> list[str]:
        """Return connection IDs whose pattern (and filters) match *topic*.

        Args:
            topic: Concrete topic name to match against.
            payload: Optional event payload. When given, a connection is
                included only if at least one of its subscriptions whose
                topic pattern matches *topic* also satisfies that
                subscription's ``filters`` against *payload* (SPEC-014
                FR-009). Subscriptions with no filters always pass once the
                pattern matches. When omitted, filters are not evaluated
                (topic-pattern matching only) -- preserves the historic
                behaviour for callers that only need the recipient set.

        Returns:
            List of matching connection IDs.

        Raises:
            TypeError: If topic is not a str.
            ValueError: If topic is empty.
        """
        if not isinstance(topic, str):
            raise TypeError("topic must be a str")
        if not topic:
            raise ValueError("topic must not be empty")
        # Snapshot under the lock — writers hold self._lock, so reading
        # _subscriptions without it can raise "dict changed size during
        # iteration".
        async with self._lock:
            snapshot = list(self._subscriptions.values())
        matches: list[str] = []
        seen: set[str] = set()
        for sub in snapshot:
            if sub.connection_id in seen:
                continue
            if not fnmatch.fnmatch(topic, sub.topic):
                continue
            if payload is not None and not evaluate_filters(payload, sub.filters):
                continue
            matches.append(sub.connection_id)
            seen.add(sub.connection_id)
        return matches

    async def get_subscriptions(self, connection_id: str) -> list[str]:
        """Return topic patterns for *connection_id*.

        Args:
            connection_id: Connection to query.

        Returns:
            Sorted list of subscribed topic patterns.

        Raises:
            TypeError: If connection_id is not a str.
            ValueError: If connection_id is empty.
        """
        if not isinstance(connection_id, str):
            raise TypeError("connection_id must be a str")
        if not connection_id:
            raise ValueError("connection_id must not be empty")
        # Copy under the lock — a concurrent (un)subscribe may mutate the
        # underlying set while sorted() iterates it otherwise.
        async with self._lock:
            patterns = set(self._subs.get(connection_id, set()))
        return sorted(patterns)

    def get_subscription_by_id(
        self,
        subscription_id: str,
    ) -> Subscription | None:
        """Look up a subscription by ID.

        Args:
            subscription_id: The subscription to find.

        Returns:
            The Subscription or None if not found.
        """
        return self._subscriptions.get(subscription_id)

    @traces(spec="014", story="US-012", reqs=["REQ-014-199"])
    def get_total_subscription_count(self) -> int:
        """Return total active subscriptions across clients."""
        return len(self._subscriptions)

    async def publish(
        self,
        topic: str,
        payload: dict[str, Any] | None = None,
    ) -> int:
        """Publish an event to all matching subscriptions.

        Args:
            topic: Event topic name.
            payload: Optional event payload.

        Returns:
            Number of matching connections.

        Raises:
            TypeError: If topic is not a str.
            ValueError: If topic is empty.
        """
        if not isinstance(topic, str):
            raise TypeError("topic must be a str")
        if not topic:
            raise ValueError("topic must not be empty")
        return len(await self.get_matching_connections(topic))

    def _remove_sub_records(
        self,
        connection_id: str,
        topic: str,
    ) -> None:
        """Remove subscription records matching connection+topic."""
        to_remove = [
            sid
            for sid, s in self._subscriptions.items()
            if s.connection_id == connection_id and s.topic == topic
        ]
        for sid in to_remove:
            del self._subscriptions[sid]
            # REQ-014-147: one fewer live subscription on this topic.
            subscription_removed(topic)