Services¶

class NATSService(ServiceBase):
    """Service implementation using NATS as the messaging backend.

    Subscribes to ``commands.{name}.req`` for inbound commands, publishes
    responses to ``events.params``, sensor data to ``events.data``, and
    OTA/status updates to ``events.infos``. A heartbeat is published to
    ``heartbeat.{name}.service`` at the configured cadence, which defaults to
    every ten seconds.

    Incoming messages are typically SenML-formatted JSON lists. Each record
    must have an ``n`` field (command name) and a ``vs`` field (a
    JSON-encoded string containing the command payload).

    Example:
        Subclass `NATSService` to build a complete microservice, naming the
        plugin it drives on the inner ``Meta``::

            class TemperatureService(NATSService):
                class Meta:
                    plugin_class = TemperaturePlugin

    """

    class Meta:
        """Service metadata. Subclasses declare their own ``plugin_class``."""

        strategy_class = NATSMessageStrategy

    class ServiceSubject(str, Enum):
        """NATS subject constants for outbound publications.

        Attributes:
            MESSAGES: Subject for raw sensor data (SenML records).
            PARAMS: Subject for command responses and parameter updates.
            INFO: Subject for OTA status events and service info.

        """

        MESSAGES = "events.data"
        PARAMS = "events.params"
        INFO = "events.infos"

    def __init__(self, name: str, logger: logging.Logger | None = None) -> None:
        super().__init__(name, logger)
        self._params_sync = ParamsSyncContext()

    # --- Properties ---

    @property
    def params_sync(self) -> ParamsSyncContext:
        """The sync context guarding the params-publication loop."""
        return self._params_sync

    # --- Concrete Methods ---

    async def connect(self) -> None:
        """Connect to NATS and subscribe to this service's command subject."""
        await super().connect()
        await self.register_handler(f"commands.{self.name}.req", self.remote_handler)

    async def setup(self) -> None:
        """Set the service up, then announce its command catalog once.

        The ``commands`` record goes out after the plugins are acquired, so
        it reflects every command the service and its plugins expose. A
        failure to build or publish it is logged rather than raised, so it
        cannot abort start-up.

        """
        await super().setup()
        try:
            records = [CommandsRecordFactory.create(self.base_name, self.command_catalog()).as_dict()]
            await self.publish(records, self.ServiceSubject.PARAMS.value)
        except Exception as e:
            self._logger.error("Failed to send commands: %s", e)

    async def on_plugins_changed(self) -> None:
        """Wake ``params_sync_loop`` so the new plugin set gets announced."""
        await self.notify_params_changed()

    async def notify_params_changed(self) -> None:
        """Wake ``params_sync_loop`` so the plugin topology and params get published.

        Called by `on_plugins_changed` after every attach and detach, and
        safe to call directly. Registration must already be complete, because
        the loop reads ``self._plugins`` as soon as it wakes. Wake-ups that
        arrive before the loop runs are merged into one publication.

        """
        async with self.params_sync.condition:
            self.params_sync.is_ready = True
            self.params_sync.condition.notify_all()

    async def params_sync_loop(self) -> None:
        """Publish the plugin topology and params whenever the plugin registry changes.

        Runs until cancelled, woken by `notify_params_changed`. Each wake-up
        calls `publish_plugins`, then `sync_params`.

        """
        while True:
            self._logger.debug("Waiting for new params to become available...")
            async with self.params_sync.condition:
                await self.params_sync.condition.wait_for(lambda: self.params_sync.is_ready)
                self.params_sync.is_ready = False
                self._logger.debug("New params confirmed; proceeding to send.")

            await self.publish_plugins()
            await self.sync_params()

    async def publish_plugins(self) -> None:
        """Publish the plugin topology as a ``plugins`` record.

        Published also when no plugin is attached, so a consumer learns the
        set emptied. A failed publish is logged rather than raised, so it
        cannot take the loop down.

        """
        records = [PluginsRecordFactory.create(self.base_name, self.plugin_topology()).as_dict()]
        try:
            await self.publish(records, self.ServiceSubject.PARAMS.value)
        except Exception as e:
            self._logger.error("Failed to send plugins: %s", e)

    async def sync_params(self) -> None:
        """Publish service-specific params after the plugin topology.

        A hook for subclasses, called by `params_sync_loop` after
        `publish_plugins`. Does nothing by default.

        """

    async def start(self) -> None:
        """Announce the plugins, schedule the params loop and heartbeat, and start the service."""
        # The startup record; the loop below publishes it once it runs.
        await self.notify_params_changed()
        await self.task_manager.add_task("params_sync_loop", self.params_sync_loop)
        # Cadence comes from configuration; the method's own default is left
        # alone so heartbeat() stays independently callable.
        interval_ms = get_config().service.heartbeat_interval_ms
        await self.task_manager.add_task("heartbeat", self.heartbeat, interval_ms=interval_ms)
        await super().start()

    async def heartbeat(self, interval_ms: int = 10_000) -> None:
        """Publish a heartbeat message at the configured interval.

        Publishes an empty message to ``heartbeat.{name}.service`` every
        ``interval_ms`` milliseconds. Intended to run as a managed
        background task.

        Args:
            interval_ms: Heartbeat interval in milliseconds. Defaults to 10 000.
                ``start`` passes the configured cadence; this default applies
                only when the method is called directly.

        """
        while True:
            self._logger.info("Publishing heartbeat...")
            await self.strategy.publish(f"heartbeat.{self.name}.service", None)
            await asyncio.sleep(interval_ms / 1000.0)

    async def _reply_error(self, msg: msg.Msg, text: str) -> None:
        """Publish an error record to ``events.params``, passing ``msg.reply`` through.

        Args:
            msg: The message being answered.
            text: The error text sent as the record's value.

        """
        error_response = [{"bn": self.base_name, "n": "error", "vs": text}]
        await self.publish(error_response, self.ServiceSubject.PARAMS.value, msg.reply)

    async def remote_handler(self, msg: msg.Msg) -> None:
        """Handle an incoming NATS message from ``commands.{name}.req``.

        Decodes the SenML envelope, extracts the command name and payload,
        dispatches via ``command_handler``, and publishes the response to
        ``events.params``. When ``msg.reply`` is set it is passed through as the
        NATS reply-to (request/reply correlation), not as an alternative subject.
        A malformed envelope (an empty list, a non-object record, or a missing
        or non-string ``n`` or ``vs``) gets the "Invalid message format" reply.

        Args:
            msg: The raw NATS message.

        """
        subject = msg.subject
        msg_data = msg.data.decode()
        self._logger.info("Received message on subject [%s] with reply-to: %s", subject, msg.reply or "None")
        try:
            msg_content = json.loads(msg_data)

            record = msg_content[0] if isinstance(msg_content, list) and msg_content else msg_content
            if not isinstance(record, dict):
                raise MessageFormatError("n", detail="the command record must be a JSON object")
            if "n" not in record:
                raise MessageFormatError("n")
            if not isinstance(record["n"], str):
                raise MessageFormatError("n", detail="'n' must be a string")
            if "vs" not in record:
                raise MessageFormatError("vs")
            if not isinstance(record["vs"], str):
                raise MessageFormatError("vs", detail="'vs' must be a string")

            cmd_name = record["n"]
            payload = json.loads(record["vs"])

            result: Dict = await self.command_handler(cmd_name, payload)
            msg_response = [{"bn": self.base_name, "n": cmd_name, "vs": json.dumps(result)}]
            await self.publish(msg_response, self.ServiceSubject.PARAMS.value, msg.reply)

        except MessageFormatError as message_format_exc:
            self._logger.error("%s", message_format_exc)
            await self._reply_error(msg, "Invalid message format")

        except CommandError as command_exc:
            # Expected client error (unknown command, bad plugin id): log
            # cleanly without a stack trace and return the message.
            self._logger.error("Command rejected: %s", command_exc)
            await self._reply_error(msg, str(command_exc))

        except json.JSONDecodeError as json_decode_exc:
            self._logger.error("JSON decode error at position %d: %s", json_decode_exc.pos, json_decode_exc.msg)
            await self._reply_error(msg, "Invalid JSON")

        except Exception as exc:
            self._logger.exception("Command handling failed: %s", exc)
            await self._reply_error(msg, str(exc))

class SensorService(NATSService):
    """Sensor service publishing plugin readings and params over NATS.

    Declares no ``plugin_class``: `SensorPlugin` is abstract, so defaulting to it
    would only defer the failure from construction to the first ``new_plugin``
    call. A deployment names its own concrete plugin.

    """

    class Meta:
        """Service metadata. Subclasses must declare a concrete ``plugin_class``."""

        commands = SENSOR_COMMANDS

    def __init__(self, service_name: str, logger: logging.Logger | None = None) -> None:
        super().__init__(service_name, logger)
        self._data_sync = DataSyncContext()
        self._data_queue: asyncio.Queue = asyncio.Queue()

    # --- Properties ---

    @property
    def data_sync(self) -> DataSyncContext:
        """The sync context guarding the reading-publication loop."""
        return self._data_sync

    # --- Concrete Methods ---

    def wire_plugin(self, plugin: PluginBase) -> None:
        """Route the plugin's data and info callbacks to this service.

        Args:
            plugin: The plugin about to be attached.

        """
        plugin.on_data_available = self.async_data_handler
        plugin.on_info_available = self.async_info_handler

    async def start(self) -> None:
        """Start the reading loop, then the base loops, and block until the tasks finish."""
        await self.task_manager.add_task("data_sync_loop", self.data_sync_loop)
        await super().start()

    async def stop(self) -> None:
        """Disconnect every attached sensor, then stop the service."""
        for device_id, sensor in self._plugins.items():
            self._logger.info("Disconnecting sensor '%s'...", device_id)
            await sensor.disconnect()
        await super().stop()

    async def data_sync_loop(self) -> None:
        """Publish queued readings whenever a plugin signals new data.

        Runs until cancelled. Waking is edge-triggered by `async_data_handler`
        rather than polled, so an idle sensor costs nothing.

        """
        while True:
            self._logger.debug("Waiting for new data to become available...")
            async with self.data_sync.condition:
                await self.data_sync.condition.wait_for(lambda: self.data_sync.is_ready)
                self.data_sync.is_ready = False
                self._logger.debug("New data confirmed; proceeding to send.")

            await self.sync_data()

    async def sync_data(self) -> None:
        """Drain the data queue, publishing each batch as SenML records.

        A batch that fails to publish is logged and dropped rather than
        retried, so one bad batch cannot stall the queue behind it.

        """
        builder: SenMLRecordBuilder = SenMLRecordBuilder(self.base_name, DatapointRecordFactory)
        while not self._data_queue.empty():
            batch_data: List = await self._data_queue.get()
            if not batch_data:
                self._logger.debug("Skipped empty batch.")
                continue

            start = time.perf_counter()
            records: List = [builder.build(data).as_dict() for data in batch_data]
            self._logger.info("Records prepared for batch: \n%s", json.dumps(records, indent=2))
            self._logger.debug("Preparing to send %d SenML records (batch size: %d)", len(records), len(batch_data))
            try:
                await self.publish(records, self.ServiceSubject.MESSAGES.value)
                elapsed = time.perf_counter() - start
                self._logger.info("Data batch sent successfully in %.3f seconds.", elapsed)
            except Exception as e:
                self._logger.error("Failed to send data batch: %s", e)

    async def sync_params(self) -> None:
        """Publish the output params of every plugin.

        Skipped when no plugin is attached. A failed publish is logged rather
        than raised, so it cannot take the loop down.

        """
        builder: SenMLRecordBuilder = SenMLRecordBuilder(self.base_name, MetricRecordFactory)

        if not self._plugins:
            self._logger.debug("Skipped empty plugins.")
            return

        start = time.perf_counter()

        # Collect all sensor parameters
        all_params: List = []
        for _, sensor in self._plugins.items():
            all_params.extend(sensor.out_params)

        # Build records from collected parameters
        records: List = [builder.build(all_params).as_dict()]

        self._logger.info("Records prepared for params: \n%s", json.dumps(records, indent=2))
        self._logger.debug("Preparing to send %d SenML records (plugins count: %d)", len(records), len(self._plugins))

        try:
            await self.publish(records, self.ServiceSubject.PARAMS.value)
            elapsed = time.perf_counter() - start
            self._logger.info("Params sent successfully in %.3f seconds.", elapsed)
        except Exception as e:
            self._logger.error("Failed to send params: %s", e)

    async def async_data_handler(self, data: List) -> None:
        """Queue a plugin's readings and wake `data_sync_loop`.

        Args:
            data: The datapoints pushed by the plugin.

        """
        data_size = len(data)
        self._logger.info("Handling incoming sensor data: %d datapoints.", data_size)

        async with self.data_sync.condition:
            self._logger.debug("Data lock acquired; queuing %d datapoints.", data_size)
            try:
                await self._data_queue.put(data)
                self.data_sync.is_ready = True
                self.data_sync.condition.notify_all()
                self._logger.info("Data handling complete; %d datapoints queued and consumers notified.", data_size)
            except Exception as e:
                self._logger.exception("Error while handling sensor data: %s", e)

    async def async_info_handler(self, info: dict) -> None:
        """Publish a device-info dict to the INFO subject.

        Unlike sensor data (SenML-wrapped in ``sync_data``), ``info`` is
        published as-is to ``events.infos``.

        """
        self._logger.info("Handling incoming device info: %s", info)
        try:
            await self.publish(info, self.ServiceSubject.INFO.value)
        except Exception as e:
            self._logger.exception("Error while handling device info: %s", e)

class SingleSensorService(SensorService):
    """Sensor service driving exactly one plugin, attached at setup."""

    class Meta:
        """Service metadata. Subclasses must declare a concrete ``plugin_class``."""

    async def acquire_plugins(self) -> None:
        """Attach the single declared plugin.

        A successful attach wakes the params loop through `on_plugins_changed`,
        so a plugin that failed to connect does not advertise metrics it
        cannot produce.

        """
        await self.attach_plugin(self.new_plugin())

class OTAService(NATSService):
    """Service driving a single OTA plugin and publishing its status.

    The plugin pushes status as it changes rather than being polled, so an
    update's progress reaches subscribers without a sync loop.

    """

    # --- Concrete Methods ---

    def wire_plugin(self, plugin: PluginBase) -> None:
        """Route the plugin's data callback to `notify_ota_status`.

        Args:
            plugin: The plugin about to be attached.

        """
        plugin.on_data_available = self.notify_ota_status

    async def acquire_plugins(self) -> None:
        """Attach the single plugin declared in ``Meta.plugin_class``."""
        await self.attach_plugin(self.new_plugin())

    async def notify_ota_status(self, data: List) -> None:
        """Publish an OTA status update on the service's info subject.

        Args:
            data: The status records pushed by the plugin.

        """
        self._logger.info("OTAService received data: %s", data)
        await self.publish(data, subject=self.ServiceSubject.INFO.value)