from collections import defaultdict, deque
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, fields, MISSING
from enum import Enum
from logging import Logger, getLogger
from multiprocessing.connection import wait, Connection
from multiprocessing.queues import Queue
from re import findall
from threading import Lock
from typing import Any, Callable, DefaultDict, NamedTuple, Type, cast
from uuid import UUID
try:
from markdown import markdown
except ImportError as e: # pragma: no cover - exercised only without the extra
raise ImportError(
"edri.api.broker requires markdown. "
"Install it with: pip install edri[api]"
) from e
from edri.abstract.manager.manager_base import ManagerBase
from edri.api.dataclass import Client, Cookie, Header, Scope
from edri.api.dataclass.api_event import api_events, ApiEvent
from edri.api.middleware import Middleware, MiddlewareControl
from edri.config.constant import ApiType
from edri.config.setting import API_RESPONSE_WRAPPED, API_ENABLE_DOCUMENTATION, API_ENABLE_TOPOLOGY, API_SENDER_THREADS, \
HEALTH_CHECK_LOG
from edri.dataclass.directive.base import InternalServerErrorResponseDirective
from edri.dataclass.directive.http import NotFoundResponseDirective
from edri.dataclass.event import Event, EventHandlingType, ApiInfo
from edri.dataclass.injection import Injection
from edri.dataclass.response import ResponseStatus, Response
from edri.events.api import group, manage, client
from edri.events.edri.router import HealthCheck, Subscribe, Topology
from edri.utility import Storage
from edri.utility.coercion import TypeCoercer
from edri.utility.validation import StringValidator, IntegerValidator, FloatValidator, DateValidator, \
TimeValidator, DateTimeValidator
__all__ = ["EventInfo"]
class _SendPump:
"""Bounded thread pool that performs the broker's outbound sends off the
single-threaded event loop, with per-client FIFO ordering.
The broker is single-threaded; it used to send responses with a blocking
``Connection.send``. A slow or stuck client (e.g. a few large concurrent
responses whose reader has fallen behind) then blocked the whole event loop,
so even trivial requests timed out and registrations failed (head-of-line
blocking).
Here the broker only ``submit``s (non-blocking). A fixed pool of worker
threads performs the blocking sends, and **at most one thread services a
given client at a time** (per-client queue + in-flight flag), so:
* per-client send order is preserved (matters for WebSocket streams),
* a slow client ties up a single pool thread, not the broker — and not the
other clients, up to ``max_workers`` simultaneously stuck clients.
A failed send is reported via ``on_dead(key, type)``; the broker turns that
into an ``Unregister`` processed on its own thread, so client/subscription
state is never mutated from a pool thread.
"""
def __init__(self, logger: Logger, on_dead: Callable[[str, ApiType], None], max_workers: int) -> None:
self._logger = logger
self._on_dead = on_dead
self._executor = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="broker-sender")
self._lock = Lock()
self._queues: dict[Client, deque[tuple[str, ApiType, Event]]] = {}
self._inflight: set[Client] = set()
def submit(self, client: "Client", key: str, type: ApiType, event: "Event") -> None:
with self._lock:
queue = self._queues.get(client)
if queue is None:
queue = deque()
self._queues[client] = queue
queue.append((key, type, event))
if client in self._inflight:
return
self._inflight.add(client)
self._executor.submit(self._drain, client)
def _drain(self, client: "Client") -> None:
while True:
with self._lock:
queue = self._queues.get(client)
if not queue:
self._inflight.discard(client)
self._queues.pop(client, None)
return
key, type, event = queue.popleft()
try:
client.socket.send(event)
except (BrokenPipeError, ConnectionResetError, OSError) as e:
# The client connection is gone: drop it (the broker unregisters
# it on its own thread).
self._logger.warning("Client send failed; dropping client", exc_info=e)
with self._lock:
self._queues.pop(client, None)
self._inflight.discard(client)
self._on_dead(key, type)
return
except Exception as e:
# The connection is fine but this event could not be sent — most
# often the response is not picklable (e.g. it carries a lambda,
# an ORM attribute or a DB session). Don't silently swallow it in
# the pool thread's unread future; log it and reply 500 so the
# client gets an immediate error instead of timing out.
self._logger.error("Failed to send response; replying 500: %s", event, exc_info=e)
try:
self._send_failure(client, event)
except Exception as failure_error:
self._logger.error("Failed to send 500 fallback; dropping client",
exc_info=failure_error)
with self._lock:
self._queues.pop(client, None)
self._inflight.discard(client)
self._on_dead(key, type)
return
@staticmethod
def _send_failure(client: "Client", event: "Event") -> None:
"""Replace the (unsendable) response with a clean 500 and send that, so a
client whose response failed to serialize gets an error rather than a
hang. Clearing the non-underscore response fields drops whatever made the
original unpicklable."""
response = event.get_response()
if response is not None:
for response_field in fields(response):
if not response_field.name.startswith("_"):
try:
setattr(response, response_field.name, None)
except Exception:
pass
response.get_directives().clear()
response.set_status(ResponseStatus.FAILED)
response.add_directive(InternalServerErrorResponseDirective("Response could not be serialized"))
client.socket.send(event)
def forget(self, client: "Client") -> None:
with self._lock:
self._queues.pop(client, None)
self._inflight.discard(client)
def shutdown(self) -> None:
self._executor.shutdown(wait=False, cancel_futures=True)
def _json_safe(value: Any) -> Any:
"""Recursively coerce a value into something picklable and JSON-serializable
for the documentation payload. Containers are walked; anything that is not a
plain primitive (regexes, ORM attributes, DB sessions, lambda transforms, …)
becomes its ``str()`` — which is what a docs page wants to display anyway, and
keeps the response sendable across the broker→listener pipe."""
if value is None or isinstance(value, (str, int, float, bool)):
return value
if isinstance(value, dict):
return {str(key): _json_safe(item) for key, item in value.items()}
if isinstance(value, (list, tuple, set)):
return [_json_safe(item) for item in value]
return str(value)
class TrieNode:
def __init__(self, key: Any = None) -> None:
self.key = key
self.children: dict[tuple[str, Any], TrieNode] = {}
self.is_end: bool = False
self.clients: set[Client] = set()
def __repr__(self) -> str:
"""
Returns a tree-like string representation of the trie starting from this node.
Returns:
str: String representation of the trie starting from this node.
"""
return self._build_repr(level=0)
def _build_repr(self, level: int = 0, max_depth: int | None = None, max_children: int | None = None) -> str:
"""
Recursively builds the string representation of the trie.
Parameters:
level (int): Current depth level in the trie.
max_depth (int | None): Maximum depth to traverse for the representation.
max_children (int | None): Maximum number of children to display per node.
Returns:
str: String representation of the trie starting from this node.
"""
indent = " " * level
node_label = f"TrieNode(key={self.key}, is_end={self.is_end}, clients={len(self.clients)})"
representation = f"{indent}{node_label}\n"
if max_depth is not None and level >= max_depth:
return representation
children_keys = sorted(self.children.keys())
displayed_children = cast("list[tuple[str, Any] | str]", children_keys)
if max_children is not None and len(children_keys) > max_children:
displayed_children = children_keys[:max_children] + ['...']
for child_key in displayed_children:
if isinstance(child_key, str):
representation += f"{indent} ...\n"
break
child_node = self.children[child_key]
representation += child_node._build_repr(level=level + 1, max_depth=max_depth, max_children=max_children)
return representation
[docs]
@dataclass
class EventInfo:
resource: str | None
handling_type: EventHandlingType
# Comparison operators a WS subscription filter may use. Equality (``eq``) is also
# the shorthand: a bare scalar value is treated as ``{"eq": value}``.
FILTER_OPERATORS = frozenset({"eq", "ne", "in", "gt", "gte", "lt", "lte"})
class Predicate(NamedTuple):
"""One coerced subscription filter constraint: ``field <op> value``.
For ``in`` the value is a tuple of allowed values; for every other operator
it is a single coerced scalar.
"""
field: str
op: str
value: Any
class Broker(ManagerBase):
"""
The ApiBroker class acts as an intermediary between the central router and all API handlers,
managing client registrations, event dispatching, and communication in a multiprocessing
environment. It extends ManagerBase to leverage common management functionalities and
adds capabilities specific to API interactions.
Attributes:
api_broker_queue (Queue): A multiprocessing queue used to receive events specifically for the API broker.
clients (Optional[Storage[Client]]): A storage mechanism for registered clients, allowing for quick access and management.
db: Placeholder attribute for potential database connections or operations, not implemented by default.
events (dict): A mapping of event resources to event classes, derived from the global api_events definition.
middlewares_request (list[Middleware]): List of middlewares applied to incoming requests.
middlewares_response (list[Middleware]): List of middlewares applied to outgoing responses.
Methods:
__init__(router_queue: "Queue[Event]", api_broker_queue: "Queue[Event]", middlewares: list[Middleware]):
Constructor for the ApiBroker class.
name (property): Returns the name of the ApiBroker instance.
send_subscribed(event: Event): Sends an event to clients subscribed to matching parameters.
send_all(event: Event): Broadcasts an event to all clients of a specific API type.
client_register(event: client.Register): Registers a new client and sends a confirmation event.
client_unregister(event: Event): Unregisters a client based on the event's key (the dead-client/queue-only path).
event_register(event: manage.EdriRegister, source: Client): Registers a client for specific events and parameters.
solve(event: Event): Determines the appropriate method to handle an incoming event.
_prepare(): Prepares the ApiBroker for operation by initializing storage for clients.
additional_pipes (property): Returns a set of additional pipes to monitor for incoming events.
run_resolver(): Main loop for processing incoming events and dispatching them to the appropriate handlers.
resolve_callback(event: Event, source: Client): Callback method for handling events received from a registered client's own pipe.
_prepare_resolvers(): Prepares resolvers for handling specific types of events.
resolve(event: Event): Processes an event through the appropriate middlewares and forwards it to its destination.
Usage:
This class is instantiated with specific queues and is used within a multiprocessing
application to manage communications between the router and API handlers. It handles
client registrations, event dispatching, and utilizes middlewares for request and
response processing.
"""
def __init__(self, router_queue: Queue[Event], api_broker_queue: Queue[Event],
middlewares: list[Middleware]) -> None:
"""
Initializes an ApiBroker instance with specific queues for routing and
API broker events.
Parameters:
router_queue (Queue): The queue for receiving routed events from other
components of the system.
api_broker_queue (Queue): The queue specifically for the API broker to
receive events.
Returns:
None
"""
super().__init__(router_queue)
self.api_broker_queue = api_broker_queue
self.logger = getLogger(__name__)
self.clients: Storage[Client]
self._pump: _SendPump
self._client_keys: dict[Client, str] = {}
self.db = None
self.events = {api_event.event: EventInfo(api_event.resource, api_event.handling) for api_event in api_events}
self.resources = {api_event.resource: api_event.event for api_event in api_events if api_event.resource}
self.middlewares_request = []
self.middlewares_response = []
for middleware in middlewares:
if middleware.is_request:
self.middlewares_request.append(middleware)
if middleware.is_response:
self.middlewares_response.append(middleware)
requests, responses = middleware.register_events()
for event in api_events:
if event.event in requests or event.event in responses:
self.events[event.event] = EventInfo(event.resource, event.handling)
self.roots: DefaultDict[Type[Event], TrieNode] = defaultdict(TrieNode)
# Subscriptions that use any non-equality operator do not fit the equality
# trie; they are matched by a linear scan of compiled predicates per type.
self.predicates: DefaultDict[Type[Event], list[tuple[tuple[Predicate, ...], Client]]] = defaultdict(list)
self._coercer = TypeCoercer()
# ManagerBase.name is itself a read-only property overriding Process's
# writeable attribute (see its own type: ignore[override]); overriding it
# again here hits the same mypy limitation.
@property
def name(self) -> str: # type: ignore[override]
"""
Property returning the name of the ApiBroker instance. Primarily used for
logging and identification purposes.
Returns:
str: The name "ApiBroker".
"""
return "ApiBroker"
def _send(self, client: Client, event: Event) -> None:
"""
Attempts to send an event to a specific client. If the connection is broken
or reset during sending, it logs a warning and unregisters the client.
Parameters:
client (Client): The client object representing the target for the event.
event (Event): The event to be sent to the client.
Returns:
None
"""
key = self._client_keys.get(client)
if key is None:
self.logger.warning("No key for client; dropping event %s", event)
return
# Non-blocking: hand off to the sender pool. A failed send is reported
# asynchronously and unregisters the client on the broker thread (see
# _SendPump / _on_dead_client).
self._pump.submit(client, key, client.type, event)
def _on_dead_client(self, key: str, type: ApiType) -> None:
"""Pool-thread callback for a failed send: post an Unregister so the
broker drops the client on its own thread."""
self.api_broker_queue.put(self._dead_client_event(key, type))
def _refused_by_middleware(self, event: Event, client: Client) -> bool:
"""Run the request middlewares over one subscription-management event.
Subscription management is answered by the Broker and never reaches
the Router, so without this it would be the one part of an
application's API that no middleware ever sees -- and a client denied
every query could simply subscribe to the same data instead.
Denial follows the same convention as an ordinary request: a
middleware marks the response and continues, and the marked event is
the answer. Here that answer goes straight back to the client, because
a Manage event has nowhere else to go.
"""
if not self.middlewares_request:
return False
resolved = self.handle_middlewares(event, self.middlewares_request)
if resolved is None:
return True
if resolved.has_response():
self._send(client, resolved)
return True
if resolved is not event:
# Replaced with something that is no longer subscription
# management; it is routed like any other event instead.
self.send_to_router(resolved)
return True
return False
def send_subscribed(self, event: Event) -> None:
"""
Sends an event to all clients that are subscribed to the parameters
matching the event's attributes.
Parameters:
event (Event): The event to be sent to subscribed clients.
Returns:
None
"""
event_fields = {field.name: getattr(event, field.name) for field in fields(event)}
matched: set[Client] = set(self._search_trie(self.roots[event.__class__], event_fields))
for predicates, client in self.predicates.get(event.__class__, []):
if client not in matched and self._predicate_match(predicates, event_fields):
matched.add(client)
for client in matched:
self.logger.debug("->? %s", event)
self._send(client, event)
def send_all(self, event: Event) -> None:
"""
Broadcasts an event to every client holding a persistent, pushable
connection - WebSocket and gRPC's bidi Connect stream. REST/HTML
clients are one-shot request/response and have no open connection
to push to; a pooled/idle gRPC unary pipe (ClientRegistry.acquire_pipe)
is also a Client of type GRPC but nothing is reading it between
calls, so in practice only WS connections and gRPC clients actively
inside a Connect stream ever receive one of these.
Parameters:
event (Event): The event to be broadcasted.
Returns:
None
"""
self.logger.debug("->* %s", event)
for client in self.clients.values():
if client.type in (ApiType.WS, ApiType.GRPC):
self._send(client, event)
def client_register(self, event: client.Register) -> None:
"""
Handles the registration of a new client based on the received register event.
It assigns a unique key to the client, stores it, and acknowledges the registration.
Parameters:
event (client.Register): The event containing the client's registration information.
Returns:
None
"""
client_obj = Client(event.socket, event.type)
key = self.clients.append(client_obj)
self._client_keys[client_obj] = key
self._pipe_to_client[client_obj.socket] = client_obj
event._api = ApiInfo(key, event.type)
# `Register.socket` is annotated `Connection` (non-Optional); cleared here
# once the socket has been stored on `client_obj` so it isn't carried
# (and re-pickled) on the event any longer.
event.socket = None # type: ignore[assignment]
event.response.set_status(ResponseStatus.OK)
self._send(client_obj, event)
def _dead_client_event(self, key: str, type: ApiType) -> "client.Unregister":
"""Builds an Unregister the broker can process on its own thread to drop a
client whose sender thread reported a dead connection."""
dead = client.Unregister()
dead._api = ApiInfo(key, type)
return dead
def _do_unregister(self, client: Client) -> None:
"""
Canonical cleanup for a departing client, given the resolved Client
object directly. Used both by client_unregister() (the dead-client/
queue-only path, resolved via event._api.key) and by
resolve_callback() (an Unregister arriving on the client's own live
pipe, resolved via the pipe itself — see _pipe_to_client).
Returns:
None
"""
self._pump.forget(client)
key = self._client_keys.pop(client, None)
self._pipe_to_client.pop(client.socket, None)
try:
client.socket.close()
except OSError:
pass
if key is not None:
self.clients.pop(key, None)
self._remove_client_from_tries(client)
self._remove_client_from_predicates(client)
def client_unregister(self, event: Event) -> None:
"""
Handles the unregistration of a client identified by the event's key.
Kept specifically for the dead-client signal (posted internally via
api_broker_queue when a send fails on a pool thread) — that path has
no live pipe to resolve identity from, so it is the one place that
still reads event._api.key rather than a pipe-resolved Client.
Parameters:
event (Event): The event containing the key of the client to be unregistered.
Returns:
None
"""
if not (event._api and event._api.key):
self.logger.warning("Client cant be unregistered because key is missing!")
return
try:
client = self.clients[event._api.key]
except KeyError as e:
self.logger.debug("Client was not found!", exc_info=e)
return
self._do_unregister(client)
def event_register(self, event: manage.EdriRegister, source: Client) -> None:
"""
Registers a client for specific events and their parameters, allowing for
customized event handling. This method updates the client's event subscriptions
and their associated parameters based on the provided event registration details.
Args:
event (manage.Register): An object representing the event registration request:
a list of ``events`` the client wishes to subscribe to, and corresponding
``parameters`` and ``values`` for those events.
source (Client): The client that sent this request, already resolved by
resolve_callback() from the pipe it arrived on.
Returns:
None
Raises:
ValueError: If the ``events`` list in the ``event`` object is empty, indicating
that no events have been specified for registration.
LookupError: If any of the events specified in the ``event.events`` list do not
exist in the ``self.events`` mapping, indicating an attempt to register for
undefined events.
"""
client = source
if self._refused_by_middleware(event, client):
return
event_type = self.resources.get(event.event)
if event_type is None:
self.logger.error("Event was not found: %s", event.event)
event.response.set_status(ResponseStatus.FAILED)
self._send(client, event)
return
try:
predicates = self._compile_param_set(event_type, event.param_set)
except ValueError as e:
self.logger.error("Invalid subscription filter for %s: %s", event.event, e)
event.response.set_status(ResponseStatus.FAILED)
self._send(client, event)
return
root = self.roots[event_type]
if not root.clients:
# Subscribe's `response` field is non-init (default_factory via the
# @event decorator), which mypy cannot see through.
self.send_to_router(Subscribe(pipe=None, name=self.name, event_type=event_type, request=False,
from_time=self._from_time)) # type: ignore[call-arg]
if all(predicate.op == "eq" for predicate in predicates):
# Pure-equality subscription: store on the fast equality trie.
self._insert_client_to_trie(root, {p.field: p.value for p in predicates}, client)
else:
# Any non-equality operator (ne/in/gt/gte/lt/lte) goes to the predicate scan.
self.predicates[event_type].append((predicates, client))
event.response.set_status(ResponseStatus.OK)
self._send(client, event)
def event_unregister(self, event: manage.EdriUnregister, source: Client) -> None:
"""
Unregisters a client from a specific subscription based on the provided param_set.
Removes the client from the trie along the path defined by param_set.
Parameters:
event (manage.Unregister): The event containing unregistration details, including param_set.
source (Client): The client that sent this request, already resolved by
resolve_callback() from the pipe it arrived on.
Returns:
None
"""
client = source
if self._refused_by_middleware(event, client):
return
event_type = self.resources.get(event.event)
if event_type is None:
self.logger.error("Event was not found: %s", event.event)
event.response.set_status(ResponseStatus.FAILED)
self._send(client, event)
return
# Compile the filter exactly as registration did, so the removal targets the
# same structure (trie path or predicate subscription) that was inserted.
try:
predicates = self._compile_param_set(event_type, event.param_set)
except ValueError as e:
self.logger.error("Invalid subscription filter for %s: %s", event.event, e)
event.response.set_status(ResponseStatus.FAILED)
self._send(client, event)
return
root = self.roots[event_type]
if all(predicate.op == "eq" for predicate in predicates):
param_items = sorted((p.field, p.value) for p in predicates)
removed = self._remove_client_from_trie_by_param_set(root, param_items, 0, client)
else:
removed = self._remove_predicate_subscription(event_type, predicates, client)
if removed:
self.logger.debug("Client %s successfully unregistered from %s", client, event.param_set)
event.response.set_status(ResponseStatus.OK)
else:
self.logger.warning("Client %s was not registered for %s", client, event.param_set)
event.response.set_status(ResponseStatus.FAILED)
self._send(client, event)
def documentation(self, event: client.EdriDocumentation, source: Client) -> None:
def split_camel(name: str) -> list[str]:
"""Split a CamelCase name into its component words, preserving acronyms."""
return findall(r'[A-Z]+(?=[A-Z][a-z0-9])|[A-Z][a-z0-9]+|[A-Z]+|[a-z0-9]+', name)
def detect_endpoints(api_event: ApiEvent) -> list[str]:
return [api_type.name for api_type in ApiType if api_type.value not in api_event.exclude]
def extract_parameters_response(event_cls: Type[Event]) -> tuple[dict[str, Any], dict[str, Any]]:
VALIDATION_MAP: dict[type, tuple[str, type]] = {
StringValidator: ("string", StringValidator.__bases__[0]),
IntegerValidator: ("integer", IntegerValidator.__bases__[0]),
FloatValidator: ("float", FloatValidator.__bases__[0]),
DateValidator: ("date", DateValidator.__bases__[0]),
TimeValidator: ("time", TimeValidator.__bases__[0]),
DateTimeValidator: ("datetime", DateTimeValidator.__bases__[0]),
}
def resolve_validation(field_type: Injection) -> tuple[str | None, dict[str, Any]]:
parameters = dict(field_type.parameters)
for injected_callable in field_type.classes:
derived_parameters = getattr(
injected_callable,
"__edri_validator_parameters__",
None,
)
if isinstance(derived_parameters, dict):
parameters.update(derived_parameters)
candidate = getattr(
injected_callable,
"__edri_validator__",
injected_callable,
)
for validation_cls, (type_str, base_type) in VALIDATION_MAP.items():
if candidate is validation_cls:
return str(base_type), {
"type": type_str,
"parameters": parameters,
}
# No dedicated EDRI validator matched (UUID, Decimal, bool, a
# Python enum): still report a usable documentation type from a
# callable's declared base type rather than null.
for injected_callable in field_type.classes:
base = getattr(injected_callable, "__edri_base_type__", None)
if not isinstance(base, type):
continue
if issubclass(base, Enum):
doc_type = "string"
elif issubclass(base, bool):
doc_type = "boolean"
elif issubclass(base, int):
doc_type = "integer"
elif issubclass(base, float):
doc_type = "float"
else:
doc_type = "string"
return str(base), {"type": doc_type, "parameters": parameters}
return None, {"parameters": parameters}
def extract_response_info(response_type: Type[Response]) -> dict[str, Any]:
params = {}
for field in fields(response_type):
if field.name.startswith("_"):
continue
param = {"type": str(field.type)}
if field.default is not MISSING:
param["default"] = field.default
params[field.name] = param
return {
"name": response_type.__name__,
# __doc__ is `str | None`; markdown() is a 3rd-party stub typed
# as requiring `str`.
"description": markdown(response_type.__doc__) or "", # type: ignore[arg-type]
"parameters": params,
}
params = {}
response = {}
for field in fields(event_cls):
if field.name.startswith("_") or field.name == "method":
continue
if field.name == "response":
# dataclasses.Field.type is typed as `type | str | Any` (it can
# be an unresolved string when `from __future__ import
# annotations` is active); the "response" field is always a
# resolved Response subclass by construction (see @event).
response = extract_response_info(cast(Type[Response], field.type))
param: dict[str, Any] = {"type": str(field.type)}
if isinstance(field.type, Injection):
param["type"], param["validation"] = resolve_validation(field.type)
if field.default is not MISSING:
if isinstance(field.default, Cookie):
param["note"] = "Value will be load from cookie"
elif isinstance(field.default, Header):
param["note"] = "Value will be load from header"
elif isinstance(field.default, Scope):
param["note"] = "Value will be load from scope"
else:
param["default"] = field.default
params[field.name] = param
return params, response
def sort_api_tree(node: dict[str, Any]) -> dict[str, Any]:
leaf_nodes: dict[str, Any] = {}
group_nodes: dict[str, Any] = {}
for key, value in node.items():
if isinstance(value, dict) and "_value" in value:
leaf_nodes[key] = value
elif isinstance(value, dict):
group_nodes[key] = sort_api_tree(value) # Recurse into subgroups
# Combine leaf nodes first, then group nodes
sorted_node = {**dict(sorted(leaf_nodes.items())), **dict(sorted(group_nodes.items()))}
return sorted_node
api_resources: dict[str, Any] = {}
for api_event in api_events:
if api_event.event is Topology and not API_ENABLE_TOPOLOGY:
continue
class_name = api_event.event.__name__
parts = split_camel(class_name)
node = api_resources
for part in parts[:-1]:
node = node.setdefault(part, {})
leaf = parts[-1]
parameters, response = extract_parameters_response(api_event.event)
node[leaf] = {
"_value": {
"name": class_name,
"description": markdown(api_event.event.__doc__) or "", # type: ignore[arg-type]
"url": api_event.url,
"method": api_event.method.name if api_event.method else None,
"endpoints": detect_endpoints(api_event),
"command": api_event.resource,
"parameters": parameters,
"cookies": api_event.cookies,
"headers": api_event.headers,
"response": response,
}
}
# Validator parameters, field defaults, etc. can hold objects that are not
# picklable/JSON-serializable (regexes, ORM attributes, DB sessions, lambda
# transforms). Coerce the whole tree to safe values so the response can be
# sent across the broker→listener pipe and serialized for the client.
event.response.api_resources = _json_safe(sort_api_tree(api_resources))
event.response.wrapped = API_RESPONSE_WRAPPED
self._send(source, event)
def solve(self, event: Event) -> None:
"""
Determines the appropriate method to handle an incoming event based on its type and
handling type specified in the events mapping.
Parameters:
event (Event): The event to be handled.
Returns:
None
"""
if (event_info := self.events.get(event.__class__)) is None:
self.logger.error("Unknown event type: %s", event.__class__)
else:
# Response middlewares unwind in reverse of the request order (onion
# model): given middlewares=[A, B, C], requests run A→B→C and
# responses run C→B→A.
resolved_event = self.handle_middlewares(event, list(reversed(self.middlewares_response)))
if resolved_event:
handling_type = event_info.handling_type
if handling_type & EventHandlingType.SPECIFIC:
target = self._pending_outbound.pop(resolved_event._id, None)
if target is None:
if resolved_event._api is None:
# Nothing came in through the API: an @api event can
# equally be raised inside the application (by the
# scheduler, by a manager), and then there is no
# client to answer and nothing wrong.
self.logger.debug("No client waiting for internal event %s", resolved_event)
else:
self.logger.warning("No pending client for event %s", resolved_event)
else:
self._send(target, resolved_event)
elif handling_type & EventHandlingType.SUBSCRIBED:
self.send_subscribed(resolved_event)
elif handling_type & EventHandlingType.ALL:
self.send_all(resolved_event)
else:
self.logger.warning("Event %s with unexpected handling type %s", resolved_event, handling_type)
def _prepare(self) -> None:
"""
Prepares the ApiBroker for operation by initializing necessary components,
such as client storage.
Returns:
None
"""
super()._prepare()
self.clients = Storage[Client]()
self._client_keys = {}
self._pump = _SendPump(self.logger, self._on_dead_client, API_SENDER_THREADS)
# Inbound correlation: which registered client a live pipe belongs to,
# so an incoming event is identified by *which pipe delivered it*
# rather than a key read off the event payload. Populated only in
# client_register(); drained in _do_unregister().
self._pipe_to_client: dict[Connection, Client] = {}
# Outbound correlation: which client is waiting for the reply to a
# given request, keyed on the request's own event._id (set
# in resolve_callback() before running request middleware, read
# back in solve() when the reply returns) — avoids trusting
# event._api.key to have survived the Router/Manager/Worker round
# trip unmodified.
self._pending_outbound: dict[UUID, Client] = {}
@property
def additional_pipes(self) -> set[Connection]:
"""
Provides a set of additional pipes that should be monitored for incoming
events: the API broker queue's reader (registration bootstrap and the
internal dead-client signal) plus every registered client's own pipe.
Returns:
Set[Connection]: Pipes to watch for inbound client-originated events.
"""
# `_reader` is a real attribute of multiprocessing.queues.Queue at
# runtime, just not exposed on the typeshed stub.
return {self.api_broker_queue._reader, *self._pipe_to_client.keys()} # type: ignore[attr-defined]
def run_resolver(self) -> None:
"""
The main event loop of the ApiBroker, continuously checking for and processing incoming events from various sources.
The loop iterates indefinitely until a KeyboardInterrupt is raised, at which point it returns None.
Three kinds of pipe fire here: the router pipe (replies/broadcasts
coming back from the rest of the system), the broker queue's reader
(registration bootstrap and the internal dead-client signal — see
_resolve_bootstrap), and every registered client's own pipe (ordinary
requests from an already-registered connection, identified by *which*
pipe delivered them via _pipe_to_client, not by a key on the event).
Returns:
None: Raises a KeyboardInterrupt exception upon user termination.
"""
while True:
pipes: set[Connection] = set()
if self.router_pipe:
pipes.add(self.router_pipe)
pipes.update(self.workers.pipes())
pipes.update(self.additional_pipes)
try:
active_pipes = wait(pipes, timeout=self.workers._next_timeout(10.0))
worker_pipes = self.workers.pipes()
ready_workers = [pipe for pipe in active_pipes if pipe in worker_pipes]
for active_pipe in ready_workers:
# wait()'s return type is broader than what these pipes actually
# are (Connection | socket | int); everything this broker waits
# on is a Connection.
self.resolve_non_router_pipe(cast(Connection, active_pipe))
self.workers._process_timeouts()
for active_pipe in active_pipes:
if active_pipe in worker_pipes:
continue
if active_pipe == self.router_pipe:
event: Event = self.router_pipe.recv()
event._timing.stamp(self.__class__.__name__, "Received")
if HEALTH_CHECK_LOG or not isinstance(event, HealthCheck):
self.logger.debug("Received event: %s", event)
self.resolve(event)
elif active_pipe is self.api_broker_queue._reader: # type: ignore[attr-defined]
try:
event = active_pipe.recv()
except (EOFError, OSError) as e:
self.logger.error("Communication problems", exc_info=e)
continue
event._timing.stamp(self.__class__.__name__, "Received")
self._resolve_bootstrap(event)
else:
# Look up the owning client *before* recv() — if it's
# already gone (unregistered earlier in this same
# wait() batch, e.g. its dead-client Unregister came
# off the queue reader first), skip without touching
# a Connection that may already be closed.
source = self._pipe_to_client.get(cast(Connection, active_pipe))
if source is None:
continue
try:
event = cast(Connection, active_pipe).recv()
except (EOFError, OSError) as e:
self.logger.warning("Client pipe closed unexpectedly", exc_info=e)
self._do_unregister(source)
continue
event._timing.stamp(self.__class__.__name__, "Received")
self.resolve_callback(event, source)
self.workers._process_timeouts()
except KeyboardInterrupt:
return None
def _resolve_bootstrap(self, event: Event) -> None:
"""
Handles the two event kinds that still travel over api_broker_queue:
a new connection's Register (no pipe exists yet to route it through)
and the internal dead-client Unregister (a control signal with no
live pipe to trust — see _SendPump/_on_dead_client).
Returns:
None
"""
if isinstance(event, client.Register):
self.client_register(event)
elif isinstance(event, client.Unregister):
self.client_unregister(event)
else:
self.logger.error("Unexpected event on bootstrap queue: %s", event)
def resolve_callback(self, event: Event, source: Client) -> None:
"""
Handles incoming events received from an already-registered client's
own pipe, dispatching them to the appropriate handling method based
on their type.
Parameters:
event (Event): The incoming event to be handled.
source (Client): The client this event was received from,
resolved by run_resolver() from *which pipe* delivered it.
Returns:
None
"""
self.logger.debug("<- Client %s", event)
if isinstance(event, group.Client):
# client.Register never reaches here — it only ever arrives via
# the bootstrap queue, before any pipe is registered to route it
# through (see _resolve_bootstrap). Any other group.Client
# subtype not explicitly handled below is silently ignored here,
# exactly as before this redesign.
if isinstance(event, client.Unregister):
self._do_unregister(source)
elif isinstance(event, client.EdriDocumentation):
if API_ENABLE_DOCUMENTATION:
self.documentation(event, source)
else:
event.response.set_status(ResponseStatus.FAILED)
event.response.add_directive(NotFoundResponseDirective())
self._send(source, event)
elif isinstance(event, group.Manage):
if isinstance(event, manage.EdriRegister):
self.event_register(event, source)
elif isinstance(event, manage.EdriUnregister):
self.event_unregister(event, source)
elif isinstance(event, Topology) and not API_ENABLE_TOPOLOGY:
event.response.set_status(ResponseStatus.FAILED)
event.response.add_directive(NotFoundResponseDirective())
self._send(source, event)
else:
self.logger.debug("Router <- API Broker %s", event)
event_info = self.events.get(event.__class__)
if event_info and event_info.handling_type & EventHandlingType.SPECIFIC:
# Middleware can answer synchronously via REDIRECT (or its
# error path), before anything has been sent to the Router.
self._pending_outbound[event._id] = source
try:
resolved_event = self.handle_middlewares(event, self.middlewares_request)
except Exception as e:
self._pending_outbound.pop(event._id, None)
self.logger.exception("Exception during middleware running", exc_info=e)
else:
if resolved_event:
event_info = self.events.get(resolved_event.__class__)
if event_info and event_info.handling_type & EventHandlingType.SPECIFIC:
self._pending_outbound[resolved_event._id] = source
self.send_to_router(resolved_event)
def _transfer_pending_client(self, event: Event, replacement: Event) -> None:
"""Keep reply correlation when middleware substitutes a different event."""
if replacement._id != event._id:
source = self._pending_outbound.pop(event._id, None)
event_info = self.events.get(replacement.__class__)
if source is not None and event_info and event_info.handling_type & EventHandlingType.SPECIFIC:
self._pending_outbound[replacement._id] = source
def handle_middlewares(self, event: Event, middlewares: list[Middleware]) -> Event | None:
event_info = self.events.get(event.__class__, None)
for middleware in middlewares:
try:
event._timing.stamp(middleware.__class__.__name__, "Started")
# Snapshot the direction before middleware mutates the event:
# a request rejection commonly sets event.response in place.
is_response = event.has_response()
if not is_response:
control = middleware.process_request(event, event_info)
else:
control = middleware.process_response(event, event_info)
event._timing.stamp(middleware.__class__.__name__, "Finished")
if not isinstance(control, MiddlewareControl):
self.logger.error("Middleware %s suppose to return MiddlewareControl object",
middleware.__class__.__name__)
# `event: Event` here (not a subtype), so `.response` is the base
# class's `Optional[Response]`; this handler assumes it is set.
cast(Response, event.response).set_status(ResponseStatus.FAILED)
cast(Response, event.response).add_directive(InternalServerErrorResponseDirective(
f"Middleware {middleware.__class__.__name__} suppose to return MiddlewareControl object"))
self.resolve(event)
raise TypeError(
"Middleware %s suppose to return MiddlewareControl object" % middleware.__class__.__name__)
if control.action == MiddlewareControl.ActionType.CONTINUE:
continue
elif control.action == MiddlewareControl.ActionType.STOP:
self.logger.info("Processing stopped: %s", middleware.__class__.__name__)
self._pending_outbound.pop(event._id, None)
return None
elif control.action == MiddlewareControl.ActionType.REPLACE_EVENT:
self.logger.debug("Event replaced by middleware: %s. New event: %s", middleware.__class__.__name__,
control.event)
if control.event is not None:
self._transfer_pending_client(event, control.event)
return control.event
elif control.action == MiddlewareControl.ActionType.REDIRECT:
self.logger.debug("Event redirected by middleware to router: %s. New event: %s",
middleware.__class__.__name__,
control.event)
# MiddlewareControl.event is Optional in general, but REDIRECT
# is only ever constructed via redirect_event(), which requires
# a real Event.
redirected = cast(Event, control.event)
self._transfer_pending_client(event, redirected)
if not is_response:
self.solve(redirected)
else:
self.send_to_router(redirected)
return None
except Exception as e:
self.logger.error("Unknown error while processing request in middleware", exc_info=e)
cast(Response, event.response).set_status(ResponseStatus.FAILED)
cast(Response, event.response).add_directive(InternalServerErrorResponseDirective(
f"Unknown error while processing request in middleware: {middleware.__class__.__name__}"))
self.resolve(event)
raise e
return event
def _prepare_resolvers(self) -> None:
"""
Prepares the resolvers for handling specific types of events, enhancing
the base implementation with API-specific event handling.
Returns:
None
"""
for event in api_events:
if "response" not in event.event.__annotations__:
self._requests[event.event] = self.solve
else:
self._responses[event.event] = self.solve
super()._prepare_resolvers()
def _compile_param_set(self, event_type: Type[Event], param_set: dict[str, Any]) -> tuple[Predicate, ...]:
"""
Compile a raw filter map into canonical, coerced :class:`Predicate`s.
Each value is either a bare scalar (shorthand for ``eq``) or a single-key
operator object ``{"<op>": operand}`` with ``op`` in :data:`FILTER_OPERATORS`.
Operands are coerced to the matching event field's declared type using the
same machinery the Listener uses for an HTTP request, so the filter and the
event attribute compare correctly even for UUID/Enum/datetime fields.
A filter that could never match — an unknown field, an unknown operator, a
value that cannot be coerced, or a non-scalar value — is rejected loudly with
a ``ValueError`` instead of being stored as a subscription that never fires.
The result is sorted by field name so the same filter always compiles to the
same tuple regardless of JSON key order (used for trie paths and for matching
a subscription on unregister).
"""
field_types = {field.name: field.type for field in fields(event_type)}
predicates: list[Predicate] = []
for name, raw in param_set.items():
if name not in field_types:
raise ValueError(f"Unknown field '{name}' for event '{event_type.__name__}'")
annotation = field_types[name]
op, operand = self._parse_operator(raw)
if op == "in":
if not isinstance(operand, list):
raise ValueError(f"Operator 'in' for field '{name}' requires a list of values")
values = tuple(self._coerce_scalar(name, annotation, item) for item in operand)
predicates.append(Predicate(name, op, values))
else:
predicates.append(Predicate(name, op, self._coerce_scalar(name, annotation, operand)))
predicates.sort(key=lambda predicate: predicate.field)
return tuple(predicates)
@staticmethod
def _parse_operator(raw: Any) -> tuple[str, Any]:
"""Resolve a raw filter value into ``(operator, operand)``.
A single-key dict whose key is a known operator is that operator; anything
else (incl. bare scalars and other dicts) is shorthand for ``eq``.
"""
if isinstance(raw, dict) and len(raw) == 1:
(key, operand), = raw.items()
if key in FILTER_OPERATORS:
return key, operand
return "eq", raw
def _coerce_scalar(self, name: str, annotation: Any, value: Any) -> Any:
"""Coerce a single operand to the field type and require it to be scalar."""
try:
converted = self._coercer.convert_type(value, annotation)
except (TypeError, ValueError) as e:
raise ValueError(
f"Filter value {value!r} for field '{name}' cannot be coerced to {annotation}"
) from e
try:
hash(converted)
except TypeError:
raise ValueError(
f"Filter value for field '{name}' must be a scalar, got {type(converted).__name__}"
)
return converted
@staticmethod
def _predicate_match(predicates: tuple[Predicate, ...], event_fields: dict[str, Any]) -> bool:
"""True if every predicate holds against the event's field values (AND)."""
for field, op, value in predicates:
if field not in event_fields:
return False
actual = event_fields[field]
try:
if op == "eq":
matched = actual == value
elif op == "ne":
matched = actual != value
elif op == "in":
matched = actual in value
elif op == "gt":
matched = actual > value
elif op == "gte":
matched = actual >= value
elif op == "lt":
matched = actual < value
else: # "lte"
matched = actual <= value
except TypeError:
# Incomparable types (e.g. ordering against None) never match.
return False
if not matched:
return False
return True
def _remove_predicate_subscription(self, event_type: Type[Event],
predicates: tuple[Predicate, ...], client: Client) -> bool:
"""Remove one predicate subscription for ``client`` matching ``predicates``."""
subscriptions = self.predicates.get(event_type)
if not subscriptions:
return False
for index, (stored, stored_client) in enumerate(subscriptions):
if stored_client is client and stored == predicates:
del subscriptions[index]
if not subscriptions:
del self.predicates[event_type]
return True
return False
def _remove_client_from_predicates(self, client: Client) -> None:
"""Drop every predicate subscription belonging to a fully unregistered client."""
for event_type in list(self.predicates.keys()):
subscriptions = [sub for sub in self.predicates[event_type] if sub[1] is not client]
if subscriptions:
self.predicates[event_type] = subscriptions
else:
del self.predicates[event_type]
@staticmethod
def _insert_client_to_trie(node: TrieNode, param_set: dict[str, Any], client: Client) -> None:
"""
Inserts a client into the trie based on the provided parameter set.
Parameters:
node (TrieNode): The current node in the trie.
param_set (dict[str, Any]): The parameter set defining the subscription path.
client (Client): The client to be inserted.
Returns:
None
"""
# Canonical key order so two clients filtering on the same fields share a
# trie path regardless of the order the keys arrived in the JSON request.
for param, value in sorted(param_set.items(), key=lambda item: item[0]):
key = (param, value)
if key not in node.children:
node.children[key] = TrieNode(key=key)
node = node.children[key]
node.is_end = True
node.clients.add(client)
@classmethod
def _search_trie(cls, node: TrieNode, event_fields: dict[str, Any],
path: set[tuple[str, Any]] | None = None) -> list[Client]:
"""
Searches the trie to find all clients subscribed to parameters matching the event's attributes.
Parameters:
node (TrieNode): The current node in the trie.
event_fields (dict[str, Any]): The event's field name -> value map (computed once by the caller).
path (set): A set to keep track of the current traversal path.
Returns:
list[Client]: A list of clients that match the event's parameters.
"""
if path is None:
path = set()
results: list[Client] = []
if node.is_end:
results.extend(node.clients)
for (param, value), child in node.children.items():
if param in event_fields and event_fields[param] == value and (param, value) not in path:
path.add((param, value))
results.extend(cls._search_trie(child, event_fields, path))
path.remove((param, value))
return results
def _remove_client_from_tries(self, client: Client) -> bool:
return any([self._remove_client_from_trie(root, client) for root in self.roots.values()])
@classmethod
def _remove_client_from_trie(cls, node: TrieNode, client: Client) -> bool:
"""
Recursively removes the client from the trie. Used when a client unregisters entirely.
Returns True if the node should be deleted (no clients and no children), False otherwise.
Parameters:
node (TrieNode): The current node in the trie.
client (Client): The client to be removed.
Returns:
bool: True if the node can be deleted, False otherwise.
"""
# Remove the client from the current node if present
if client in node.clients:
node.clients.remove(client)
# Make a list of keys to avoid modifying the dictionary while iterating
keys = list(node.children.keys())
for key in keys:
child = node.children[key]
# Recursively remove the client from child nodes
should_delete_child = cls._remove_client_from_trie(child, client)
if should_delete_child:
# Delete the child node if it has no clients and no children
del node.children[key]
# Determine if the current node should be deleted
if not node.clients and not node.children:
return True # Node can be deleted
else:
return False # Node should remain
def _remove_client_from_trie_by_param_set(
self, node: TrieNode, param_items: list[tuple[str, Any]], index: int, client: Client
) -> bool:
"""
Recursive function to remove client from a specific path in trie.
Returns True if the client was removed, False otherwise.
Parameters:
node (TrieNode): The current node in the trie.
param_items (list[tuple[str, Any]]): The list of parameter items defining the path.
index (int): The current index in the param_items list.
client (Client): The client to be removed.
Returns:
bool: True if the client was removed, False otherwise.
"""
if index < len(param_items):
param, value = param_items[index]
key = (param, value)
if key in node.children:
child = node.children[key]
client_removed = self._remove_client_from_trie_by_param_set(child, param_items, index + 1, client)
# After recursion, check if the child node should be deleted
if client_removed:
if not child.clients and not child.children:
del node.children[key]
return True
else:
return False
else:
# The path does not exist; client was not registered for this param_set
return False
else:
# At the leaf node corresponding to param_set
if client in node.clients:
node.clients.remove(client)
return True
else:
# Client was not registered for this param_set
return False
def after_start(self) -> None:
for middleware in self.middlewares_request:
requests, responses = middleware.register_events()
# Subscribe's `response` field is non-init (default_factory via the
# @event decorator), which mypy cannot see through.
for request in requests:
self.send_to_router(
Subscribe(pipe=None, name=self.name, event_type=request, request=True, from_time=self._from_time) # type: ignore[call-arg]
)
for response in responses:
self.send_to_router(
Subscribe(pipe=None, name=self.name, event_type=response, request=False, from_time=self._from_time) # type: ignore[call-arg]
)
def resolve_unknown(self, event: Event) -> None:
if isinstance(event, Subscribe) and event.has_response():
return
if not event.has_response():
for middleware in self.middlewares_request:
requests, _ = middleware.register_events()
if event.__class__ in requests:
resolved_event = self.handle_middlewares(event, [middleware])
if resolved_event:
self.send_to_router(resolved_event)
return
else:
for middleware in reversed(self.middlewares_response):
_, responses = middleware.register_events()
if event.__class__ in responses:
resolved_event = self.handle_middlewares(event, [middleware])
if resolved_event:
self.solve(resolved_event)
return
super().resolve_unknown(event)
def _subscribe_static_resolvers(self) -> None:
"""
Registers static resolver methods for specific event types, such as health checks and store get requests.
"""
self.logger.debug("Registering static resolvers")
self._requests[HealthCheck] = self.health_check