stackpilot/backend/services/stats_service.py

"""Live CPU/memory stats for running containers, aggregated per compose stack.

Reads a one-shot ``docker stats`` sample per running container (the daemon
includes ``precpu_stats`` so a single read yields a usable CPU delta) and sums
them by ``com.docker.compose.project`` label, which equals the stack id.

Sampling is not free: it is one blocking call to the daemon *per running
container*, and the dashboard and the stacks list both poll this every five
seconds. Two open tabs on a 40-container host meant a sustained ~16 samples a
second. Results are therefore cached for :data:`CACHE_TTL`, the same shape
``dashboard_service`` already uses for its fleet aggregate — one sweep serves
every reader in the window, and the numbers stay well inside what a
five-second poll can show.
"""
from __future__ import annotations

import threading
import time
from concurrent.futures import ThreadPoolExecutor

from docker_client import DockerError, get_client, safe_call

COMPOSE_PROJECT_LABEL = "com.docker.compose.project"

#: Slightly under the frontend's 5s poll, so a refresh usually gets fresh
#: numbers while concurrent readers still share one sweep.
CACHE_TTL = 4.0

_cache: dict = {"data": None, "ts": 0.0}
# Held across the sample so N simultaneous callers trigger one sweep, not N.
_lock = threading.Lock()


def _container_stats(container) -> dict | None:
    try:
        s = container.stats(stream=False)
    except Exception:  # noqa: BLE001 - skip containers that fail to report
        return None

    cpu = s.get("cpu_stats", {}) or {}
    pre = s.get("precpu_stats", {}) or {}
    cpu_delta = cpu.get("cpu_usage", {}).get("total_usage", 0) - pre.get("cpu_usage", {}).get(
        "total_usage", 0
    )
    sys_delta = cpu.get("system_cpu_usage", 0) - pre.get("system_cpu_usage", 0)
    online = (
        cpu.get("online_cpus")
        or len(cpu.get("cpu_usage", {}).get("percpu_usage") or [])
        or 1
    )
    cores_used = (cpu_delta / sys_delta) * online if sys_delta > 0 and cpu_delta > 0 else 0.0

    mem = s.get("memory_stats", {}) or {}
    usage = mem.get("usage", 0) or 0
    detail = mem.get("stats", {}) or {}
    # Match `docker stats`: exclude reclaimable page cache from "used".
    inactive = detail.get("inactive_file") or detail.get("total_inactive_file") or 0
    mem_used = max(usage - inactive, 0)

    hc = container.attrs.get("HostConfig", {}) or {}
    nano = hc.get("NanoCpus") or 0
    quota = hc.get("CpuQuota") or 0
    period = hc.get("CpuPeriod") or 0
    if nano:
        cpu_limit = nano / 1e9
    elif quota and period:
        cpu_limit = quota / period
    else:
        cpu_limit = None
    mem_limit = hc.get("Memory") or 0

    return {
        "project": (container.labels or {}).get(COMPOSE_PROJECT_LABEL),
        "cores_used": cores_used,
        "cpu_limit": cpu_limit,
        "mem_used": mem_used,
        "mem_limit": mem_limit or None,
    }


def stack_stats(refresh: bool = False) -> dict:
    """Return {stack_id: {cpu_used, cpu_limit, mem_used, mem_limit, containers}}.

    Limits are the summed assigned limits across the stack's containers, or null
    when none of them have that limit set. Served from a short-lived cache
    unless ``refresh`` is set.
    """
    if not refresh and _cache["data"] is not None:
        if time.monotonic() - _cache["ts"] < CACHE_TTL:
            return _cache["data"]

    with _lock:
        # Somebody may have refreshed it while we waited for the lock.
        if not refresh and _cache["data"] is not None:
            if time.monotonic() - _cache["ts"] < CACHE_TTL:
                return _cache["data"]
        data = _sample()
        _cache["data"] = data
        _cache["ts"] = time.monotonic()
        return data


def _sample() -> dict:
    """One full sweep across every running container."""
    try:
        client = get_client()
        containers = safe_call(client.containers.list)  # running only
    except DockerError:
        return {}

    with ThreadPoolExecutor(max_workers=8) as pool:
        samples = list(pool.map(_container_stats, containers))

    agg: dict[str, dict] = {}
    for st in samples:
        if not st or not st["project"]:
            continue
        a = agg.setdefault(
            st["project"],
            {"cpu_used": 0.0, "cpu_limit": 0.0, "has_cpu": False,
             "mem_used": 0, "mem_limit": 0, "has_mem": False, "containers": 0},
        )
        a["cpu_used"] += st["cores_used"]
        a["mem_used"] += st["mem_used"]
        a["containers"] += 1
        if st["cpu_limit"]:
            a["cpu_limit"] += st["cpu_limit"]
            a["has_cpu"] = True
        if st["mem_limit"]:
            a["mem_limit"] += st["mem_limit"]
            a["has_mem"] = True

    return {
        proj: {
            "cpu_used": round(a["cpu_used"], 3),
            "cpu_limit": round(a["cpu_limit"], 3) if a["has_cpu"] else None,
            "mem_used": a["mem_used"],
            "mem_limit": a["mem_limit"] if a["has_mem"] else None,
            "containers": a["containers"],
        }
        for proj, a in agg.items()
    }