diff --git a/python/constellation/core/async_experimental/__init__.py b/python/constellation/core/async_experimental/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/constellation/core/async_experimental/async_basecontroller.py b/python/constellation/core/async_experimental/async_basecontroller.py new file mode 100644 index 000000000..9d31bf16b --- /dev/null +++ b/python/constellation/core/async_experimental/async_basecontroller.py @@ -0,0 +1,146 @@ +""" +SPDX-FileCopyrightText: 2026 DESY and the Constellation authors +SPDX-License-Identifier: EUPL-1.2 + +Async base and scriptable controller framework classes. +""" + +import asyncio +from typing import Any +from uuid import UUID + +import zmq +import zmq.asyncio + +from constellation.core.async_experimental.async_chirp import CHIRPEvent, DiscoveredService +from constellation.core.async_experimental.async_heartbeat import AsyncHeartbeatChecker +from constellation.core.async_experimental.async_monitoringlistener import AsyncMonitoringListener +from constellation.core.chirp import CHIRPServiceIdentifier +from constellation.core.controller import SatelliteUpdate +from constellation.core.logging import setup_cli_logging +from constellation.core.message.cscp1 import CSCP1Message + + +class AsyncBaseController(AsyncMonitoringListener, AsyncHeartbeatChecker): + """Async equivalent of BaseController. + + CSCP commands use zmq.asyncio with await send/recv directly, eliminating + the dedicated executor and synchronous context. _async_ctx from + BaseSatelliteFrame is used for all ZMQ sockets including CSCP REQ sockets. + Discovery and heartbeat run as coroutines. + """ + + def __init__( + self, + loop: asyncio.AbstractEventLoop | None = None, + **kwds: Any, + ) -> None: + super().__init__(**kwds) + self._transmitters: dict[str, zmq.asyncio.Socket] = {} + self._transmitter_uuids: dict[UUID, str] = {} + self._cscp_lock = asyncio.Lock() + + self.register_chirp_callback("basecontroller_control", self._on_control_service) + self.register_chirp_callback("basecontroller_heartbeat", self._on_heartbeat_service) + + def _on_control_service(self, event: CHIRPEvent, service: DiscoveredService) -> None: + """Handle CONTROL service connect/disconnect.""" + if service.service_id != CHIRPServiceIdentifier.CONTROL: + return + if event == CHIRPEvent.SERVICE_CONNECTED: + asyncio.ensure_future(self._setup_transmitter(service.host_id, service.addresses[0], service.port)) + elif event == CHIRPEvent.SERVICE_DISCONNECTED: + self._cleanup_transmitter(service.host_id) + + def _on_heartbeat_service(self, event: CHIRPEvent, service: DiscoveredService) -> None: + """Handle HEARTBEAT service connect/disconnect.""" + if service.service_id != CHIRPServiceIdentifier.HEARTBEAT: + return + if event == CHIRPEvent.SERVICE_CONNECTED: + name = self._transmitter_uuids.get(service.host_id, f"Unknown-{str(service.host_id)[:8]}") + self.register_heartbeat_host(service.host_id, service.addresses[0], service.port, name) + elif event == CHIRPEvent.SERVICE_DISCONNECTED: + self.unregister_heartbeat_host(service.host_id) + + async def _cscp_request( + self, + sock: zmq.asyncio.Socket, + command: str, + payload: Any = None, + ) -> CSCP1Message: + """Send a CSCP1 request and await the response. + + _cscp_lock serialises all REQ/REP exchanges so no two coroutines + share a socket simultaneously. + """ + request = CSCP1Message(self.name, (CSCP1Message.Type.REQUEST, command)) + if payload is not None: + request.payload = payload + async with self._cscp_lock: + await sock.send_multipart(request.assemble().frames) + response_frames = await sock.recv_multipart() + return CSCP1Message.disassemble(response_frames) + + async def command( + self, + cmd: str, + sat: str = "", + satcls: str = "", + payload: Any = None, + ) -> str: + """Send a CSCP command on the event loop.""" + key = f"{satcls}.{sat}".strip(".") + sock = self._transmitters.get(key) + if sock is None: + return f"ERROR: No transmitter for {key}" + try: + result = await self._cscp_request(sock, cmd, payload) + return f"{result.verb_msg}" + except Exception as e: + return f"ERROR: {e}" + + def _on_satellite_update(self, name: str, update_type: SatelliteUpdate) -> None: + """Called on satellite connect/disconnect. Override in subclass.""" + + async def _setup_transmitter(self, uuid: UUID, address: str, port: int) -> None: + """Connect to a satellite CSCP port on the event loop.""" + if uuid in self._transmitter_uuids: + return + sock = self._async_ctx.socket(zmq.REQ) + sock.connect(f"tcp://{address}:{port}") + sock.setsockopt(zmq.LINGER, 2000) + try: + msg = await asyncio.wait_for(self._cscp_request(sock, "get_commands"), timeout=5.0) + except Exception: + sock.close() + return + canonical_name = msg.sender + self._transmitters[canonical_name] = sock + self._transmitter_uuids[uuid] = canonical_name + self._on_satellite_update(canonical_name, SatelliteUpdate.ADDED) + + def _cleanup_transmitter(self, uuid: UUID) -> None: + """Remove and close a CSCP socket on satellite departure.""" + canonical_name = self._transmitter_uuids.pop(uuid, None) + if canonical_name is None: + return + sock = self._transmitters.pop(canonical_name, None) + if sock is not None: + sock.close() + self._on_satellite_update(canonical_name, SatelliteUpdate.REMOVED) + + def shutdown(self) -> None: + """Shut down all components.""" + for sock in self._transmitters.values(): + sock.close() + self._hb_receiver.close() + self._cmdp_pool.close() + self._async_ctx.term() + + +class AsyncScriptableController(AsyncBaseController): + """Async equivalent of ScriptableController.""" + + def __init__(self, log_level: str = "INFO", **kwds: Any) -> None: + setup_cli_logging(log_level) + super().__init__(**kwds) diff --git a/python/constellation/core/async_experimental/async_chirp.py b/python/constellation/core/async_experimental/async_chirp.py new file mode 100644 index 000000000..7c2e7566a --- /dev/null +++ b/python/constellation/core/async_experimental/async_chirp.py @@ -0,0 +1,402 @@ +""" +SPDX-FileCopyrightText: 2026 DESY and the Constellation authors +SPDX-License-Identifier: EUPL-1.2 + +Async CHIRP listener and manager. +""" + +import asyncio +import socket +import sys +from collections.abc import Callable +from dataclasses import dataclass, field +from enum import Enum, auto +from typing import Any +from uuid import UUID + +from constellation.core.base import BaseSatelliteFrame +from constellation.core.chirp import ( + CHIRP_MULTICAST_ADDRESS, + CHIRP_PORT, + CHIRPMessage, + CHIRPMessageType, + CHIRPServiceIdentifier, + get_uuid, +) +from constellation.core.multicast import MULTICAST_TTL +from constellation.core.network import get_interface_addresses + + +class CHIRPEvent(Enum): + """Event type for CHIRP service callbacks.""" + + SERVICE_CONNECTED = auto() + SERVICE_DISCONNECTED = auto() + + +@dataclass(eq=False) +class DiscoveredService: + """A service discovered via CHIRP. + + addresses is a list because a host may be reachable on multiple network + interfaces. Additional addresses are appended as OFFERs arrive from the + same host on different interfaces. + """ + + group_id: UUID + host_id: UUID + service_id: CHIRPServiceIdentifier + port: int + addresses: list[str] = field(default_factory=list) + + def __eq__(self, other: object) -> bool: + """Equality on host_id + service_id only, matching chirpmanager.py.""" + if isinstance(other, DiscoveredService): + return self.host_id == other.host_id and self.service_id == other.service_id + return NotImplemented + + +class AsyncMulticastSocket(asyncio.DatagramProtocol): + """Async-compatible multicast socket and DatagramProtocol. + + Owns the raw socket setup and the asyncio transport layer. + Incoming datagrams are forwarded to datagram_callback when set. + + Send sockets remain synchronous since UDP sendto writes to the + kernel buffer without blocking. + """ + + def __init__( + self, + interface_addresses: list[str], + multicast_address: str, + multicast_port: int, + ) -> None: + self._multicast_endpoint = (multicast_address, multicast_port) + self._datagram_callback: Callable[[bytes, tuple[str, int]], None] | None = None + self._transport: asyncio.DatagramTransport | None = None + self._send_sockets: list[socket.socket] = [] + for interface_address in interface_addresses: + sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, MULTICAST_TTL) + sock.setsockopt( + socket.IPPROTO_IP, + socket.IP_MULTICAST_LOOP, + 0 if interface_address != "127.0.0.1" else 1, + ) + sock.setsockopt( + socket.IPPROTO_IP, + socket.IP_MULTICAST_IF, + socket.inet_aton(interface_address), + ) + self._send_sockets.append(sock) + + self._recv_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP) + self._recv_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + if sys.platform == "darwin": + self._recv_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) + self._recv_socket.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, MULTICAST_TTL) + self._recv_socket.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_LOOP, 1) + self._recv_socket.bind(("0.0.0.0", multicast_port)) + for interface_address in interface_addresses: + ip_mreq = socket.inet_aton(multicast_address) + socket.inet_aton(interface_address) + self._recv_socket.setsockopt(socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP, ip_mreq) + self._recv_socket.setblocking(False) + + @property + def datagram_callback(self) -> Callable[[bytes, tuple[str, int]], None] | None: + """Callback invoked by datagram_received for each incoming packet.""" + return self._datagram_callback + + @datagram_callback.setter + def datagram_callback(self, callback: Callable[[bytes, tuple[str, int]], None] | None) -> None: + self._datagram_callback = callback + + def connection_made(self, transport: asyncio.DatagramTransport) -> None: + self._transport = transport + + def datagram_received(self, data: bytes, addr: tuple[str, int]) -> None: + if self._datagram_callback is not None: + self._datagram_callback(data, addr) + + async def start(self, loop: asyncio.AbstractEventLoop) -> None: + """Hand the receive socket to the event loop.""" + await loop.create_datagram_endpoint( + lambda: self, + sock=self._recv_socket, + ) + + def send(self, data: bytes) -> None: + """Send data on all interfaces.""" + for sock in self._send_sockets: + sock.sendto(data, self._multicast_endpoint) + + def close(self) -> None: + """Close all sockets and the transport.""" + for sock in self._send_sockets: + sock.close() + if self._transport is not None: + self._transport.close() + + +class AsyncCHIRPListener: + """Async CHIRP listener. + + Owns an AsyncMulticastSocket, decodes incoming CHIRP messages, + maintains a list of discovered services (mirroring CHIRPManager logic), + and dispatches to registered callbacks. + + Callbacks receive (CHIRPEvent, DiscoveredService). + REQUEST messages are forwarded to request_callback if set. + """ + + def __init__( + self, + name: str, + group: str, + interface_addresses: list[str], + ) -> None: + self._host_uuid = get_uuid(name) + self._group_uuid = get_uuid(group) + self._socket = AsyncMulticastSocket( + interface_addresses, + CHIRP_MULTICAST_ADDRESS, + CHIRP_PORT, + ) + self._socket.datagram_callback = self._on_datagram_received + self._callbacks: dict[str, Callable[[CHIRPEvent, DiscoveredService], None]] = {} + self._request_callback: Callable[[CHIRPServiceIdentifier], None] | None = None + self._discovered_services: list[DiscoveredService] = [] + + def register_callback( + self, + callback_id: str, + callback_func: Callable[[CHIRPEvent, DiscoveredService], None], + ) -> None: + """Register a callback for CHIRP service events.""" + self._callbacks[callback_id] = callback_func + + def unregister_callback(self, callback_id: str) -> None: + """Remove a previously registered callback.""" + self._callbacks.pop(callback_id, None) + + def send_chirp_message(self, message: CHIRPMessage) -> None: + """Send a CHIRP message on all interfaces.""" + self._socket.send(message.pack()) + + def get_discovered(self, service_id: CHIRPServiceIdentifier) -> list[DiscoveredService]: + """Return all discovered services matching service_id.""" + return [s for s in self._discovered_services if s.service_id == service_id] + + @property + def request_callback(self) -> Callable[[CHIRPServiceIdentifier], None] | None: + """Callback invoked when a CHIRP REQUEST is received.""" + return self._request_callback + + @request_callback.setter + def request_callback( + self, + callback: Callable[[CHIRPServiceIdentifier], None] | None, + ) -> None: + self._request_callback = callback + + async def start(self, loop: asyncio.AbstractEventLoop) -> None: + """Start the multicast socket on the event loop.""" + await self._socket.start(loop) + + def close(self) -> None: + """Close the socket and transport.""" + self._socket.close() + + def _on_datagram_received(self, data: bytes, addr: tuple[str, int]) -> None: + """Decode and dispatch an incoming CHIRP datagram.""" + msg = CHIRPMessage() + try: + msg.unpack(data) + except Exception: + return + + if msg.host_uuid == self._host_uuid: + return + if msg.group_uuid != self._group_uuid: + return + + msg.from_address = addr[0] + + if msg.msgtype == CHIRPMessageType.REQUEST: + if self._request_callback is not None: + self._request_callback(msg.serviceid) + elif msg.msgtype == CHIRPMessageType.OFFER: + self._discover_service(msg, addr[0]) + elif msg.msgtype == CHIRPMessageType.DEPART and msg.port != 0: + self._depart_service(msg) + + def _discover_service(self, msg: CHIRPMessage, from_address: str) -> None: + """Mirror CHIRPManager._discover_service logic.""" + already_discovered = False + for svc in list(self._discovered_services): + if svc.host_id == msg.host_uuid and svc.service_id == msg.serviceid: + if svc.port != msg.port: + self._discovered_services.remove(svc) + for cb in self._callbacks.values(): + cb(CHIRPEvent.SERVICE_DISCONNECTED, svc) + else: + if from_address not in svc.addresses: + svc.addresses.append(from_address) + already_discovered = True + break + + if not already_discovered: + svc = DiscoveredService( + group_id=msg.group_uuid, + host_id=msg.host_uuid, + service_id=msg.serviceid, + port=msg.port, + addresses=[from_address], + ) + self._discovered_services.append(svc) + for cb in self._callbacks.values(): + cb(CHIRPEvent.SERVICE_CONNECTED, svc) + + def _depart_service(self, msg: CHIRPMessage) -> None: + """Mirror CHIRPManager._depart_service logic.""" + for svc in list(self._discovered_services): + if svc.host_id == msg.host_uuid and svc.service_id == msg.serviceid: + self._discovered_services.remove(svc) + for cb in self._callbacks.values(): + cb(CHIRPEvent.SERVICE_DISCONNECTED, svc) + break + + +class AsyncCHIRPManager(BaseSatelliteFrame): + """Async equivalent of CHIRPManager. + + Owns an AsyncCHIRPListener and directly implements the service registry. + Replaces both the former async_chirpmanager.AsyncCHIRPManager wrapper and + the former AsyncCHIRPManager(AsyncCHIRPListener) registry class. + Registers CHIRP discovery as a coroutine via _add_com_task. + """ + + def __init__( + self, + name: str, + group: str, + interface: list[str] | None, + **kwds: Any, + ) -> None: + super().__init__(name=name, **kwds) + self.group = group + interface_addresses = get_interface_addresses(interface) + self._host_uuid = get_uuid(self.name) + self._group_uuid = get_uuid(group) + self._listener = AsyncCHIRPListener( + name=self.name, + group=group, + interface_addresses=interface_addresses, + ) + self._registered_services: dict[CHIRPServiceIdentifier, int] = {} + self._listener.request_callback = self._handle_request + + def _add_com_task(self) -> None: + """Register the async CHIRP discovery coroutine.""" + super()._add_com_task() + self._com_task_factories.append(self._run_chirp) + + async def _run_chirp(self, stop: asyncio.Event) -> None: + """Start CHIRP discovery and run until stop is set.""" + loop = asyncio.get_running_loop() + await self._listener.start(loop) + await stop.wait() + self._listener.close() + + def register_chirp_callback( + self, + callback_id: str, + callback_func: Callable[[CHIRPEvent, DiscoveredService], None], + ) -> None: + """Register a callback for CHIRP service events.""" + self._listener.register_callback(callback_id, callback_func) + + def unregister_chirp_callback(self, callback_id: str) -> None: + """Remove a previously registered CHIRP callback.""" + self._listener.unregister_callback(callback_id) + + def get_discovered(self, service_id: CHIRPServiceIdentifier) -> list[DiscoveredService]: + """Return discovered services matching service_id.""" + return self._listener.get_discovered(service_id) + + def request(self, service_id: CHIRPServiceIdentifier) -> None: + """Send a CHIRP REQUEST for a specific service type.""" + msg = CHIRPMessage( + CHIRPMessageType.REQUEST, + self._group_uuid, + self._host_uuid, + service_id, + 0, + ) + self._listener.send_chirp_message(msg) + + def register_service(self, service_id: CHIRPServiceIdentifier, port: int) -> None: + """Register a service to offer via CHIRP. + + Sends DEPART for the old port if already registered. + Sends OFFER immediately upon registration. + """ + if service_id in self._registered_services: + old_port = self._registered_services[service_id] + depart = CHIRPMessage( + CHIRPMessageType.DEPART, + self._group_uuid, + self._host_uuid, + service_id, + old_port, + ) + self._listener.send_chirp_message(depart) + self._registered_services[service_id] = port + offer = CHIRPMessage( + CHIRPMessageType.OFFER, + self._group_uuid, + self._host_uuid, + service_id, + port, + ) + self._listener.send_chirp_message(offer) + + def unregister_service(self, service_id: CHIRPServiceIdentifier) -> None: + """Unregister a service and send DEPART.""" + port = self._registered_services.pop(service_id, None) + if port is not None: + msg = CHIRPMessage( + CHIRPMessageType.DEPART, + self._group_uuid, + self._host_uuid, + service_id, + port, + ) + self._listener.send_chirp_message(msg) + + def emit_offers(self) -> None: + """Send OFFER for all registered services.""" + for service_id, port in self._registered_services.items(): + msg = CHIRPMessage( + CHIRPMessageType.OFFER, + self._group_uuid, + self._host_uuid, + service_id, + port, + ) + self._listener.send_chirp_message(msg) + + def _handle_request(self, service_id: CHIRPServiceIdentifier) -> None: + """Respond to incoming CHIRP requests with matching offers.""" + if service_id in self._registered_services: + port = self._registered_services[service_id] + msg = CHIRPMessage( + CHIRPMessageType.OFFER, + self._group_uuid, + self._host_uuid, + service_id, + port, + ) + self._listener.send_chirp_message(msg) diff --git a/python/constellation/core/async_experimental/async_heartbeat.py b/python/constellation/core/async_experimental/async_heartbeat.py new file mode 100644 index 000000000..5f9b4dff0 --- /dev/null +++ b/python/constellation/core/async_experimental/async_heartbeat.py @@ -0,0 +1,243 @@ +""" +SPDX-FileCopyrightText: 2026 DESY and the Constellation authors +SPDX-License-Identifier: EUPL-1.2 + +Async heartbeat receiver and checker. +""" + +import asyncio +import time +from collections.abc import Callable +from dataclasses import dataclass, field +from datetime import datetime +from typing import Any +from uuid import UUID + +import zmq.asyncio + +from constellation.core.async_experimental.async_pools import AsyncSubscriberPool +from constellation.core.base import BaseSatelliteFrame +from constellation.core.chp import CHPRole, chp_decode_message +from constellation.core.protocol.cscp1 import SatelliteState +from constellation.core.util import case_insensitive_dict + + +@dataclass +class HeartbeatState: + """Tracked state for a single satellite.""" + + host: UUID + name: str + state: SatelliteState = SatelliteState.DEAD + last_refresh: float = field(default_factory=time.monotonic) + last_statechange: datetime = field(default_factory=datetime.now) + interval_ms: int = 2000 + lives: int = 3 + role: CHPRole = CHPRole.DYNAMIC + status: str = "" + + def refresh(self) -> None: + self.last_refresh = time.monotonic() + + def seconds_since_refresh(self) -> float: + return time.monotonic() - self.last_refresh + + +class AsyncHeartbeatReceiver: + """Async heartbeat receiver backed by AsyncSubscriberPool. + + Socket management is delegated to an internal AsyncSubscriberPool. + CHP has no topic filtering so the pool subscribes to "" for all sockets, + applied automatically to every socket added via add_satellite(). + + Call add_satellite() and remove_satellite() from the event loop. + When calling from another thread use loop.call_soon_threadsafe(). + """ + + INIT_LIVES = 3 + INIT_INTERVAL = 2000 + + def __init__( + self, + ctx: zmq.asyncio.Context, + on_state_change: Callable[[str, SatelliteState, SatelliteState], None] | None = None, + on_satellite_dead: Callable[[str], None] | None = None, + ) -> None: + self._on_state_change = on_state_change + self._on_satellite_dead = on_satellite_dead + + self._states: dict[UUID, HeartbeatState] = {} + self._name_to_uuid: case_insensitive_dict[UUID] = case_insensitive_dict() + + self._pool = AsyncSubscriberPool(ctx, self._on_frames) + self._pool.set_topics([""]) + + def add_satellite(self, uuid: UUID, address: str, port: int, name: str) -> None: + """Register a satellite for heartbeat tracking.""" + if uuid in self._states: + return + self._pool.add_socket(uuid, address, port) + self._states[uuid] = HeartbeatState( + host=uuid, + name=name, + lives=self.INIT_LIVES, + interval_ms=self.INIT_INTERVAL, + ) + self._name_to_uuid[name] = uuid + + def remove_satellite(self, uuid: UUID) -> None: + """Deregister a satellite and close its socket.""" + self._pool.remove_socket(uuid) + hb = self._states.pop(uuid, None) + if hb is not None: + self._name_to_uuid.pop(hb.name, None) + + @property + def states(self) -> case_insensitive_dict[SatelliteState]: + """Current states keyed by canonical name.""" + return case_insensitive_dict({hb.name: hb.state for hb in self._states.values()}) + + @property + def state_changes(self) -> case_insensitive_dict[datetime]: + """Last state change timestamps keyed by canonical name.""" + return case_insensitive_dict({hb.name: hb.last_statechange for hb in self._states.values()}) + + @property + def statuses(self) -> case_insensitive_dict[str]: + """Last status message keyed by canonical name.""" + return case_insensitive_dict({hb.name: hb.status for hb in self._states.values()}) + + async def run(self, stop: asyncio.Event) -> None: + """Run pool polling and stale check loop concurrently until stop is set.""" + await asyncio.gather( + self._pool.run(stop), + self._stale_check_loop(stop), + ) + + async def _stale_check_loop(self, stop: asyncio.Event) -> None: + """Check for missed heartbeats every 300ms.""" + while not stop.is_set(): + await asyncio.sleep(0.3) + self._check_stale_connections() + + def _on_frames(self, uuid: UUID, frames: list[bytes]) -> None: + """Process a CHP frame delivered by the pool.""" + try: + name, _timestamp, state_val, flags, interval, status = chp_decode_message(frames) + except Exception: + return + + hb = self._states.get(uuid) + if hb is None: + return + + if hb.name.casefold() != name.casefold(): + self._name_to_uuid.pop(hb.name, None) + hb.name = name + self._name_to_uuid[name] = uuid + + state = SatelliteState(state_val) + + if state != hb.state: + old_state = hb.state + hb.state = state + hb.last_statechange = datetime.now() + if self._on_state_change: + self._on_state_change(name, old_state, state) + + hb.refresh() + hb.interval_ms = interval + hb.role = CHPRole.from_flags(flags) + + if status: + hb.status = status + + if hb.lives != self.INIT_LIVES: + hb.lives = self.INIT_LIVES + + def _check_stale_connections(self) -> None: + """Decrement lives for satellites that have missed heartbeats.""" + for _uuid, hb in list(self._states.items()): + if hb.lives <= 0: + continue + expected = (hb.interval_ms / 1000) * 1.5 + if hb.seconds_since_refresh() > expected: + hb.lives -= 1 + if hb.lives == 0: + if self._on_satellite_dead: + self._on_satellite_dead(hb.name) + hb.state = SatelliteState.DEAD + else: + hb.refresh() + + def close(self) -> None: + """Close all sockets and clear state.""" + self._pool.close() + self._states.clear() + self._name_to_uuid.clear() + + +class AsyncHeartbeatChecker(BaseSatelliteFrame): + """Async equivalent of HeartbeatChecker. + + Owns an AsyncHeartbeatReceiver internally. + Uses self._async_ctx from BaseSatelliteFrame. + Callbacks fire on the event loop — no call_soon_threadsafe needed. + """ + + def __init__(self, **kwds: Any) -> None: + super().__init__(**kwds) + self._hb_receiver = AsyncHeartbeatReceiver( + self._async_ctx, + on_state_change=self._on_state_change, + on_satellite_dead=self._on_satellite_dead, + ) + + def _add_com_task(self) -> None: + """Register the async heartbeat receiver coroutine.""" + super()._add_com_task() + self._com_task_factories.append(self._hb_receiver.run) + + def register_heartbeat_host( + self, + uuid: UUID, + address: str, + port: int, + name: str, + ) -> None: + """Register a satellite for heartbeat tracking.""" + self._hb_receiver.add_satellite(uuid, address, port, name) + + def unregister_heartbeat_host(self, uuid: UUID) -> None: + """Deregister a satellite from heartbeat tracking.""" + self._hb_receiver.remove_satellite(uuid) + + @property + def heartbeat_states(self) -> case_insensitive_dict[SatelliteState]: + """Current states keyed by canonical name.""" + return self._hb_receiver.states + + @property + def heartbeat_state_changes(self) -> case_insensitive_dict[datetime]: + """Last state change timestamps keyed by canonical name.""" + return self._hb_receiver.state_changes + + @property + def heartbeat_statuses(self) -> case_insensitive_dict[str]: + """Last status messages keyed by canonical name.""" + return self._hb_receiver.statuses + + def _on_state_change( + self, + name: str, + old_state: SatelliteState, + new_state: SatelliteState, + ) -> None: + """Called when a satellite changes state. Override in subclass.""" + + def _on_satellite_dead(self, name: str) -> None: + """Called when a satellite stops sending heartbeats. Override in subclass.""" + + def close(self) -> None: + """Close all heartbeat sockets.""" + self._hb_receiver.close() diff --git a/python/constellation/core/async_experimental/async_monitoringlistener.py b/python/constellation/core/async_experimental/async_monitoringlistener.py new file mode 100644 index 000000000..ae3617371 --- /dev/null +++ b/python/constellation/core/async_experimental/async_monitoringlistener.py @@ -0,0 +1,82 @@ +""" +SPDX-FileCopyrightText: 2026 DESY and the Constellation authors +SPDX-License-Identifier: EUPL-1.2 + +Async monitoring listener framework class. +""" + +import logging +from typing import Any +from uuid import UUID + +from constellation.core.async_experimental.async_chirp import ( + AsyncCHIRPManager, + CHIRPEvent, + DiscoveredService, +) +from constellation.core.async_experimental.async_pools import AsyncSubscriberPool +from constellation.core.chirp import CHIRPServiceIdentifier +from constellation.core.message.cmdp1 import ( + CMDP1LogMessage, + CMDP1Message, + CMDP1StatMessage, +) + + +class AsyncMonitoringListener(AsyncCHIRPManager): + """Async equivalent of MonitoringListener. + + Inherits AsyncCHIRPManager and adds CMDP log/metric receiving + via AsyncSubscriberPool. + _async_ctx is provided by BaseSatelliteFrame.__init__ via the MRO chain. + """ + + def __init__(self, **kwds: Any) -> None: + super().__init__(**kwds) + self._cmdp_pool = AsyncSubscriberPool( + self._async_ctx, + callback=self._on_cmdp_message, + ) + self.register_chirp_callback("monitoring_listener", self._on_monitoring_service) + + def _add_com_task(self) -> None: + """Register the async CMDP pool coroutine.""" + super()._add_com_task() + self._com_task_factories.append(self._cmdp_pool.run) + + def set_topics(self, topics: list[str]) -> None: + """Set CMDP topic subscriptions across all sockets.""" + self._cmdp_pool.set_topics(topics) + + def _on_monitoring_service(self, event: CHIRPEvent, service: DiscoveredService) -> None: + """Handle MONITORING service connect/disconnect.""" + if service.service_id != CHIRPServiceIdentifier.MONITORING: + return + if event == CHIRPEvent.SERVICE_CONNECTED: + self._cmdp_pool.add_socket(service.host_id, service.addresses[0], service.port) + elif event == CHIRPEvent.SERVICE_DISCONNECTED: + self._cmdp_pool.remove_socket(service.host_id) + + def _on_cmdp_message(self, uuid: UUID, frames: list[bytes]) -> None: + """Decode and dispatch an incoming CMDP message.""" + try: + msg = CMDP1Message.disassemble(frames) + if msg.is_log_message(): + log_msg = CMDP1LogMessage.from_cmdp_message(msg) + self.receive_log(log_msg.to_log_record()) + elif msg.is_stat_message(): + stat_msg = CMDP1StatMessage.from_cmdp_message(msg) + self.receive_metric( + stat_msg.sender, + stat_msg.metric, + stat_msg.time, + stat_msg.value, + ) + except Exception: + pass + + def receive_log(self, record: logging.LogRecord) -> None: + """Called when a log message arrives. Override in subclass.""" + + def receive_metric(self, sender: str, metric, time, value) -> None: + """Called when a metric arrives. Override in subclass.""" diff --git a/python/constellation/core/async_experimental/async_pools.py b/python/constellation/core/async_experimental/async_pools.py new file mode 100644 index 000000000..e0967f7a2 --- /dev/null +++ b/python/constellation/core/async_experimental/async_pools.py @@ -0,0 +1,128 @@ +""" +SPDX-FileCopyrightText: 2026 DESY and the Constellation authors +SPDX-License-Identifier: EUPL-1.2 +""" + +import asyncio +from collections.abc import Callable +from uuid import UUID + +import zmq +import zmq.asyncio + + +class AsyncSubscriberPool: + """Async ZMQ SUB socket pool using zmq.asyncio. + + All methods must be called from the event loop. Use + loop.call_soon_threadsafe() when calling from other threads. + """ + + def __init__( + self, + ctx: zmq.asyncio.Context, + callback: Callable[[UUID, list[bytes]], None], + ) -> None: + self._ctx = ctx + self._callback = callback + self._sockets: dict[UUID, zmq.asyncio.Socket] = {} + self._socket_to_uuid: dict[zmq.asyncio.Socket, UUID] = {} + self._poller = zmq.asyncio.Poller() + self._topics: list[str] = [] + + def add_socket(self, uuid: UUID, address: str, port: int) -> None: + """Add a subscriber socket and apply all current global topics.""" + if uuid in self._sockets: + return + sock = self._ctx.socket(zmq.SUB) + sock.connect(f"tcp://{address}:{port}") + sock.setsockopt(zmq.LINGER, 0) + for topic in self._topics: + sock.setsockopt_string(zmq.SUBSCRIBE, topic) + self._poller.register(sock, zmq.POLLIN) + self._sockets[uuid] = sock + self._socket_to_uuid[sock] = uuid + + def remove_socket(self, uuid: UUID) -> None: + """Remove and close the subscriber socket for a departing satellite.""" + sock = self._sockets.pop(uuid, None) + if sock is not None: + self._socket_to_uuid.pop(sock, None) + self._poller.unregister(sock) + sock.close() + + def subscribe(self, topic: str, uuid: UUID | None = None) -> None: + """Subscribe to a topic on one or all sockets. + + When uuid is None, topic is added to _topics so future sockets added + via add_socket() are automatically subscribed. When uuid is given, + only that socket is subscribed and _topics is not modified. + """ + if uuid is None: + for sock in self._sockets.values(): + sock.setsockopt_string(zmq.SUBSCRIBE, topic) + if topic not in self._topics: + self._topics.append(topic) + else: + sock = self._sockets.get(uuid) + if sock is not None: + sock.setsockopt_string(zmq.SUBSCRIBE, topic) + + def unsubscribe(self, topic: str, uuid: UUID | None = None) -> None: + """Unsubscribe from a topic on one or all sockets. + + When uuid is None, topic is removed from _topics. When uuid is given, + only that socket is unsubscribed and _topics is not modified. + """ + if uuid is None: + for sock in self._sockets.values(): + sock.setsockopt_string(zmq.UNSUBSCRIBE, topic) + if topic in self._topics: + self._topics.remove(topic) + else: + sock = self._sockets.get(uuid) + if sock is not None: + sock.setsockopt_string(zmq.UNSUBSCRIBE, topic) + + def set_topics(self, new_topics: list[str]) -> None: + """Replace all global subscriptions across all sockets.""" + old_set = set(self._topics) + new_set = set(new_topics) + for topic in old_set - new_set: + for sock in self._sockets.values(): + sock.setsockopt_string(zmq.UNSUBSCRIBE, topic) + for topic in new_set - old_set: + for sock in self._sockets.values(): + sock.setsockopt_string(zmq.SUBSCRIBE, topic) + self._topics = list(new_topics) + + async def run(self, stop: asyncio.Event) -> None: + """Poll all sockets until stop is set.""" + while not stop.is_set(): + if not self._sockets: + await asyncio.sleep(0.05) + continue + events = dict(await self._poller.poll(timeout=50)) + for sock in events: + uuid = self._socket_to_uuid.get(sock) + if uuid is None: + continue + try: + msg = await sock.recv_multipart() + self._callback(uuid, msg) + while sock.getsockopt(zmq.EVENTS) & zmq.POLLIN: + msg = await sock.recv_multipart() + self._callback(uuid, msg) + except zmq.ZMQError: + pass + + def close(self) -> None: + """Close all sockets.""" + for sock in self._sockets.values(): + try: + self._poller.unregister(sock) + except Exception: + pass + sock.close() + self._sockets.clear() + self._socket_to_uuid.clear() diff --git a/python/constellation/core/base.py b/python/constellation/core/base.py index 6647b4863..7e1c25d1e 100644 --- a/python/constellation/core/base.py +++ b/python/constellation/core/base.py @@ -5,6 +5,7 @@ This module provides a base class for Constellation Satellite modules. """ +import asyncio import atexit import logging import socket @@ -14,6 +15,7 @@ from typing import Any, cast import zmq +import zmq.asyncio from . import __version__, __version_code_name__ from .logging import ConstellationLogger @@ -105,6 +107,7 @@ def __init__(self, name: str, **_kwds: Any): # add type name to create the canonical name self.name = f"{self.type}.{name}" self.context = zmq.Context() + self._async_ctx = zmq.asyncio.Context() self.log = self.get_logger(self.type.upper()) @@ -121,6 +124,8 @@ def __init__(self, name: str, **_kwds: Any): self._com_thread_pool: dict[str, threading.Thread] = {} # Event to indicate to communication service threads to stop self._com_thread_evt: threading.Event | None = None + # List of async coroutine factories for async communication tasks + self._com_task_factories: list = [] # add self to list of satellites to destroy on shutdown global SATELLITE_LIST # noqa @@ -146,6 +151,20 @@ def _add_com_thread(self) -> None: """ self.log.debug("Satellite Base class _add_thread called") + def _add_com_task(self) -> None: + """Method to add async coroutine factories to the task pool. + + Does nothing in the base class. Override in subclass and call + super()._add_com_task() first, then append to _com_task_factories. + + """ + + async def _start_com_tasks(self, stop: asyncio.Event) -> None: + """Run all registered async tasks concurrently until stop is set.""" + if not self._com_task_factories: + return + await asyncio.gather(*[factory(stop) for factory in self._com_task_factories]) + def _start_com_threads(self) -> None: """Start all background communication threads.""" self._com_thread_evt = threading.Event() @@ -180,6 +199,7 @@ def terminate(self) -> None: self.reentry() self.log.debug("Terminating ZMQ context.") self.context.term() + self._async_ctx.term() SATELLITE_LIST: list[BaseSatelliteFrame] = [] diff --git a/python/constellation/meson.build b/python/constellation/meson.build index 296341976..cdfd3361c 100644 --- a/python/constellation/meson.build +++ b/python/constellation/meson.build @@ -86,3 +86,15 @@ version_code_name_file = configure_file( py.install_sources(version_code_name_file, subdir: 'constellation/core', ) + +core_async_experimental_files = files( + 'core/async_experimental/__init__.py', + 'core/async_experimental/async_basecontroller.py', + 'core/async_experimental/async_chirp.py', + 'core/async_experimental/async_heartbeat.py', + 'core/async_experimental/async_monitoringlistener.py', + 'core/async_experimental/async_pools.py', +) + +py.install_sources(core_async_experimental_files, + subdir: 'constellation/core/async_experimental')