Source code for herostools.flow._subscriber

"""Queued per-source subscriber for flow data."""

import queue
import threading
from collections.abc import Callable
from typing import Any

from heros import serdes

from herostools.helper import log

_STOP = object()




class FlowSubscriber:
    """Subscribe to the `flow_data` event of several HEROs and process messages on one worker.

    The zenoh callback only enqueues the raw bytes. Decoding happens on the worker, so a payload
    that cannot be decoded (for example one referencing a HERO that is gone) neither blocks zenoh
    nor kills the worker. Decoding a payload that references a dead HERO builds a `RemoteHERO`,
    which blocks the single worker for the discovery timeout (about 2 s) and stalls every source
    of this subscriber. Control items can be put on the same queue with `submit`, which keeps
    message handling and control work strictly ordered (FIFO).

    Args:
        session: Zenoh session used to declare the subscribers.
        realm: HEROS realm of the observed objects.
        handler: Called as `handler(source, payload)` on the worker for each decoded message.
        queue_warn_threshold: Queue depth at which a warning is logged. The warning is logged
            once per crossing and re-armed when the depth falls below the threshold again.
        name: Name of the owning node. It appears in log messages and in the worker thread name
            `flow-subscriber-<name>`.
    """

    def __init__(
        self,
        session: Any,
        realm: str,
        handler: Callable[[str, Any], None],
        queue_warn_threshold: int = 100,
        name: str = "unnamed",
    ):
        self._name = name
        self._session = session
        self._realm = realm
        self._handler = handler
        self._queue_warn_threshold = int(queue_warn_threshold)
        if self._queue_warn_threshold <= 0:
            raise ValueError(f"queue_warn_threshold must be positive, got {queue_warn_threshold}.")
        self._queue: queue.Queue[Any] = queue.Queue()
        self._subscribers: dict[str, Any] = {}
        self._lock = threading.Lock()
        self._warn_lock = threading.Lock()
        self._warned = False
        self._stopped = False
        self._worker = threading.Thread(target=self._run, name=f"flow-subscriber-{name}", daemon=True)
        self._worker.start()



    def _key(self, source: str) -> str:
        return f"@object/{self._realm}/{source}/flow_data"




    def add(self, source: str) -> None:
        """Subscribe to the `flow_data` event of a source. Does nothing if already subscribed.

        Args:
            source: Name of the source HERO.
        """
        with self._lock:
            if self._stopped or source in self._subscribers:
                return
            self._subscribers[source] = self._session.declare_subscriber(
                self._key(source), lambda sample, _source=source: self._on_sample(_source, sample)
            )




    def submit(self, func: Callable[[], None]) -> None:
        """Run a callable on the worker, ordered with the messages.

        Args:
            func: Callable without arguments. Exceptions are logged and do not stop the worker.
        """
        self._put(func)




    def stop(self, timeout: float = 5.0) -> None:
        """Undeclare all subscribers and stop the worker. Safe to call more than once.

        Queued items are discarded. Only an item already running is finished. If it does not
        finish within the timeout, a warning is logged and the worker exits afterwards.

        Args:
            timeout: Seconds to wait for the worker to finish.
        """
        with self._lock:
            if self._stopped:
                return
            self._stopped = True
            subscribers = list(self._subscribers.values())
            self._subscribers.clear()
        for subscriber in subscribers:
            subscriber.undeclare()
        self._queue.put(_STOP)
        if threading.current_thread() is not self._worker:
            self._worker.join(timeout)
            if self._worker.is_alive():
                log.warning(f"Flow subscriber {self._name}: worker did not stop within {timeout} s")




    def _on_sample(self, source: str, sample: Any) -> None:
        """Zenoh callback: enqueue the raw bytes without decoding."""
        self._put((source, sample.payload.to_bytes()))




    def _put(self, item: Any) -> None:
        if self._stopped:
            log.debug(f"Flow subscriber {self._name}: stopped, ignoring item")
            return
        with self._warn_lock:
            self._queue.put(item)
            depth = self._queue.qsize()
            crossed = depth >= self._queue_warn_threshold and not self._warned
            if crossed:
                self._warned = True
        if crossed:
            log.warning(
                f"Flow subscriber {self._name}: queue depth {depth} reached threshold {self._queue_warn_threshold}"
            )




    def _run(self) -> None:
        """Worker loop: decode messages, call the handler, run submitted callables."""
        while True:
            item = self._queue.get()
            if item is _STOP or self._stopped:
                return
            with self._warn_lock:
                if self._queue.qsize() < self._queue_warn_threshold:
                    self._warned = False
            if callable(item):
                self._run_item(item)
                continue
            self._dispatch(*item)




    def _dispatch(self, source: str, raw: bytes) -> None:
        try:
            payload = serdes.deserialize(raw)
        except Exception:  # noqa: BLE001
            log.error(
                f"Flow subscriber {self._name}: could not decode flow_data message from {source}, dropping it",
                exc_info=True,
            )
            return
        self._run_item(lambda: self._handler(source, payload))




    def _run_item(self, func: Callable[[], None]) -> None:
        try:
            func()
        except Exception:  # noqa: BLE001
            log.exception(f"Flow subscriber {self._name}: item failed")