Source code for herostools.actor.archiver.base

# coding:utf-8

from typing import Iterable, Any
from abc import abstractmethod
import threading
import queue

from heros import EventObserver

from herostools.helper import log




class HERODataArchiver(EventObserver):
    """Base EventObserver that subscribes to a HERO event and archives its payload.

    The event payload may be a ``(data, metadata)`` tuple or a plain dict.
    Subclasses implement :meth:`_store`.

    Args:
        object_selector: Zenoh object selector for the devices to subscribe to.
        event_name: Name of the event.
        default_metadata: Default metadata merged with every incoming payload's metadata.
        max_retries: Number of times to retry a failed store before dropping the item.
    """

    def __init__(
        self,
        object_selector: str,
        event_name: str,
        default_metadata: dict | None = None,
        max_retries: int = 5,
        *args,
        **kwargs,
    ):
        self.metadata = default_metadata if default_metadata is not None else {}
        self.max_retries = int(max_retries)
        EventObserver.__init__(self, object_selector=object_selector, event_name=event_name, **kwargs)

        self._payload_queue = queue.Queue()
        self._stop_event = threading.Event()
        self._worker_thread = threading.Thread(target=self._process_queue, daemon=True)
        self._worker_thread.start()

        self.register_callback(self.feed)



    def _process_queue(self) -> None:
        """Background worker that consumes the payload queue."""
        while not self._stop_event.is_set():
            try:
                source_name, payload, metadata, retry_count = self._payload_queue.get(timeout=1)
                log.debug(
                    f"Got data from queue source: {source_name}, metadata: {metadata}, retry_count: {retry_count}."
                )
                try:
                    self._store(source_name, payload, metadata)
                except Exception as e:  # noqa: BLE001
                    log.warning(f"Storing data from source: {source_name}, metadata: {metadata} failed with {e}.")
                    if retry_count < self.max_retries:
                        log.info(f"Re-queuing data from source: {source_name}, metadata: {metadata}.")
                        self.feed(source_name, (payload, metadata), retry_count + 1)
                    else:
                        log.error(
                            f"Max-retries exceeded for data from source: {source_name}, metadata: {metadata}. Dropping data!"
                        )
                finally:
                    self._payload_queue.task_done()
            except queue.Empty:
                continue




    def is_queue_drained(self) -> bool:
        """Return True when all queued items have been processed.

        Returns:
            True if every item put on the queue has had task_done() called.
        """
        return self._payload_queue.unfinished_tasks == 0




    def _stop(self):
        """Stop the background thread gracefully."""
        self._stop_event.set()
        self._worker_thread.join()




    def _teardown(self) -> None:
        """Called by boss on shutdown."""
        self._stop()




    def feed(self, source_name: str, data: Iterable, retry_count: int = 0) -> None:
        """Callback registered with the source event.

        ``data`` may be a ``(payload, metadata)`` tuple or a plain dict.
        When a plain dict is received it is used as both payload and metadata
        so downstream ``_store`` implementations can still locate identifier
        keys regardless of which field they are in.

        Args:
            source_name: Name of the event source (the HERO).
            data: Either a ``(payload, metadata)`` tuple or a plain dict payload.
            retry_count: Number of times this item has already been retried.
        """
        if isinstance(data, dict):
            payload = data
            metadata = self.metadata | data
        else:
            payload = data[0]
            metadata = self.metadata | data[1]
        log.debug(f"Feeding data from {source_name} with metadata {metadata}.")
        self._payload_queue.put((source_name, payload, metadata, retry_count))




    @abstractmethod
    def _store(self, source_name: str, payload: Any, metadata: dict) -> None:
        """Store the payload.  Must be implemented by subclasses.

        Args:
            source_name: Name of the event source (the HERO).
            payload: The actual data.
            metadata: Incoming metadata merged with ``default_metadata``.
        """
        pass