metro-warden/plugins/system/plugin.py

"""
System Plugin — monitors CPU, memory, disk, and load via psutil.

Publishes to:
  system.cpu      — per-core and overall CPU usage %
  system.memory   — RAM and swap usage
  system.disk     — disk partition usage
  system.load     — 1/5/15 minute load averages
  system.snapshot — combined full snapshot
"""

from __future__ import annotations

import asyncio
import logging
import os
from typing import Any, Dict, List

try:
    import psutil
    _PSUTIL_AVAILABLE = True
except ImportError:
    _PSUTIL_AVAILABLE = False

from plugins.base import BasePlugin

log = logging.getLogger(__name__)

DEFAULT_POLL_INTERVAL = 3.0


def _collect_cpu() -> Dict:
    if not _PSUTIL_AVAILABLE:
        return {}
    per_core = psutil.cpu_percent(interval=None, percpu=True)
    freq = psutil.cpu_freq()
    return {
        "percent": psutil.cpu_percent(interval=None),
        "per_core": per_core,
        "core_count": psutil.cpu_count(logical=False),
        "logical_count": psutil.cpu_count(logical=True),
        "freq_mhz": round(freq.current, 1) if freq else 0,
        "freq_max_mhz": round(freq.max, 1) if freq else 0,
    }


def _collect_memory() -> Dict:
    if not _PSUTIL_AVAILABLE:
        return {}
    vm = psutil.virtual_memory()
    sw = psutil.swap_memory()
    return {
        "total": vm.total,
        "available": vm.available,
        "used": vm.used,
        "percent": vm.percent,
        "swap_total": sw.total,
        "swap_used": sw.used,
        "swap_percent": sw.percent,
    }


def _collect_disk() -> List[Dict]:
    if not _PSUTIL_AVAILABLE:
        return []
    partitions = []
    for part in psutil.disk_partitions(all=False):
        try:
            usage = psutil.disk_usage(part.mountpoint)
            partitions.append({
                "device": part.device,
                "mountpoint": part.mountpoint,
                "fstype": part.fstype,
                "total": usage.total,
                "used": usage.used,
                "free": usage.free,
                "percent": usage.percent,
            })
        except PermissionError:
            continue
    return partitions


def _collect_load() -> Dict:
    try:
        avg = os.getloadavg()
        return {"load1": avg[0], "load5": avg[1], "load15": avg[2]}
    except AttributeError:
        return {"load1": 0.0, "load5": 0.0, "load15": 0.0}


def _collect_snapshot() -> Dict:
    return {
        "cpu": _collect_cpu(),
        "memory": _collect_memory(),
        "disk": _collect_disk(),
        "load": _collect_load(),
    }


class SystemPlugin(BasePlugin):
    """Monitors CPU, memory, disk, and load averages via psutil."""

    name = "system"
    version = "1.0.0"
    description = "Monitors system resources: CPU, memory, disk, and load averages"
    tags = ["system", "monitoring"]

    def __init__(
        self,
        bus=None,
        state=None,
        poll_interval: float = DEFAULT_POLL_INTERVAL,
    ) -> None:
        super().__init__(bus=bus, state=state)
        self._poll_interval = poll_interval
        self._task: asyncio.Task | None = None
        self._running = False

    def on_load(self) -> None:
        super().on_load()
        if not _PSUTIL_AVAILABLE:
            self._log.warning("psutil not available — system monitoring degraded")
        self._running = True
        try:
            loop = asyncio.get_running_loop()
            self._task = loop.create_task(self._poll_loop())
        except RuntimeError:
            self._log.debug("no running event loop at load time; task deferred")

    def on_unload(self) -> None:
        self._running = False
        if self._task and not self._task.done():
            self._task.cancel()
        super().on_unload()

    async def _poll_loop(self) -> None:
        self._log.debug("system poll loop started (interval=%.1fs)", self._poll_interval)
        while self._running:
            try:
                snapshot = await asyncio.to_thread(_collect_snapshot)

                self.state_set("system.cpu", snapshot["cpu"])
                self.state_set("system.memory", snapshot["memory"])
                self.state_set("system.disk", snapshot["disk"])
                self.state_set("system.load", snapshot["load"])

                if self._bus:
                    await self._bus.publish("system.cpu", snapshot["cpu"])
                    await self._bus.publish("system.memory", snapshot["memory"])
                    await self._bus.publish("system.disk", snapshot["disk"])
                    await self._bus.publish("system.load", snapshot["load"])
                    await self._bus.publish("system.snapshot", snapshot)

            except asyncio.CancelledError:
                break
            except Exception as exc:
                self._log.error("system poll error: %s", exc)
            await asyncio.sleep(self._poll_interval)
        self._log.debug("system poll loop stopped")

    def on_event(self, topic: str, data: Any) -> None:
        if topic == "system.refresh":
            asyncio.ensure_future(self._poll_loop())