remove infra.md.example, infra.md is the source of truth

ayn-antivirus/ayn_antivirus/dashboard/collector.py

@@ -0,0 +1,181 @@
+"""Background metrics collector for the AYN Antivirus dashboard."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import random
+from datetime import datetime
+from typing import Any, Dict, Optional
+
+import psutil
+
+from ayn_antivirus.constants import DASHBOARD_COLLECTOR_INTERVAL
+
+logger = logging.getLogger("ayn_antivirus.dashboard.collector")
+
+
+class MetricsCollector:
+    """Periodically sample system metrics and store them in the dashboard DB.
+
+    Parameters
+    ----------
+    store:
+        A :class:`DashboardStore` instance to write metrics into.
+    interval:
+        Seconds between samples.
+    """
+
+    def __init__(self, store: Any, interval: int = DASHBOARD_COLLECTOR_INTERVAL) -> None:
+        self.store = store
+        self.interval = interval
+        self._task: Optional[asyncio.Task] = None
+        self._running = False
+
+    async def start(self) -> None:
+        """Begin collecting metrics on a background asyncio task."""
+        self._running = True
+        self._task = asyncio.create_task(self._collect_loop())
+        logger.info("Metrics collector started (interval=%ds)", self.interval)
+
+    async def stop(self) -> None:
+        """Cancel the background task and wait for it to finish."""
+        self._running = False
+        if self._task:
+            self._task.cancel()
+            try:
+                await self._task
+            except asyncio.CancelledError:
+                pass
+        logger.info("Metrics collector stopped")
+
+    # ------------------------------------------------------------------
+    # Internal loop
+    # ------------------------------------------------------------------
+
+    async def _collect_loop(self) -> None:
+        while self._running:
+            try:
+                await asyncio.to_thread(self._sample)
+            except Exception as exc:
+                logger.error("Collector error: %s", exc)
+            await asyncio.sleep(self.interval)
+
+    def _sample(self) -> None:
+        """Take a single metric snapshot and persist it."""
+        cpu = psutil.cpu_percent(interval=1)
+        mem = psutil.virtual_memory()
+
+        disks = []
+        for part in psutil.disk_partitions(all=False):
+            try:
+                usage = psutil.disk_usage(part.mountpoint)
+                disks.append({
+                    "mount": part.mountpoint,
+                    "device": part.device,
+                    "total": usage.total,
+                    "used": usage.used,
+                    "free": usage.free,
+                    "percent": usage.percent,
+                })
+            except (PermissionError, OSError):
+                continue
+
+        try:
+            load = list(os.getloadavg())
+        except (OSError, AttributeError):
+            load = [0.0, 0.0, 0.0]
+
+        try:
+            net_conns = len(psutil.net_connections(kind="inet"))
+        except (psutil.AccessDenied, OSError):
+            net_conns = 0
+
+        self.store.record_metric(
+            cpu=cpu,
+            mem_pct=mem.percent,
+            mem_used=mem.used,
+            mem_total=mem.total,
+            disk_usage=disks,
+            load_avg=load,
+            net_conns=net_conns,
+        )
+
+        # Periodic cleanup (~1 in 100 samples).
+        if random.randint(1, 100) == 1:
+            self.store.cleanup_old_metrics()
+
+    # ------------------------------------------------------------------
+    # One-shot snapshot (no storage)
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def get_snapshot() -> Dict[str, Any]:
+        """Return a live system snapshot without persisting it."""
+        cpu = psutil.cpu_percent(interval=0.1)
+        cpu_per_core = psutil.cpu_percent(interval=0.1, percpu=True)
+        cpu_freq = psutil.cpu_freq(percpu=False)
+        mem = psutil.virtual_memory()
+        swap = psutil.swap_memory()
+
+        disks = []
+        for part in psutil.disk_partitions(all=False):
+            try:
+                usage = psutil.disk_usage(part.mountpoint)
+                disks.append({
+                    "mount": part.mountpoint,
+                    "device": part.device,
+                    "total": usage.total,
+                    "used": usage.used,
+                    "percent": usage.percent,
+                })
+            except (PermissionError, OSError):
+                continue
+
+        try:
+            load = list(os.getloadavg())
+        except (OSError, AttributeError):
+            load = [0.0, 0.0, 0.0]
+
+        try:
+            net_conns = len(psutil.net_connections(kind="inet"))
+        except (psutil.AccessDenied, OSError):
+            net_conns = 0
+
+        # Top processes by CPU
+        top_procs = []
+        try:
+            for p in sorted(psutil.process_iter(['pid', 'name', 'cpu_percent', 'memory_percent']),
+                            key=lambda x: x.info.get('cpu_percent', 0) or 0, reverse=True)[:8]:
+                info = p.info
+                if (info.get('cpu_percent') or 0) > 0.1:
+                    top_procs.append({
+                        "pid": info['pid'],
+                        "name": info['name'] or '?',
+                        "cpu": round(info.get('cpu_percent', 0) or 0, 1),
+                        "mem": round(info.get('memory_percent', 0) or 0, 1),
+                    })
+        except Exception:
+            pass
+
+        return {
+            "cpu_percent": cpu,
+            "cpu_per_core": cpu_per_core,
+            "cpu_cores": psutil.cpu_count(logical=True),
+            "cpu_freq_mhz": round(cpu_freq.current) if cpu_freq else 0,
+            "mem_percent": mem.percent,
+            "mem_used": mem.used,
+            "mem_total": mem.total,
+            "mem_available": mem.available,
+            "mem_cached": getattr(mem, 'cached', 0),
+            "mem_buffers": getattr(mem, 'buffers', 0),
+            "swap_percent": swap.percent,
+            "swap_used": swap.used,
+            "swap_total": swap.total,
+            "disk_usage": disks,
+            "load_avg": load,
+            "net_connections": net_conns,
+            "top_processes": top_procs,
+            "timestamp": datetime.utcnow().isoformat(),
+        }