ariadne/ariadne/metrics/metrics.py

from __future__ import annotations

from datetime import datetime

from prometheus_client import Counter, Gauge, Histogram


TASK_RUNS_TOTAL = Counter(
    "ariadne_task_runs_total",
    "Ariadne task runs by status",
    ["task", "status"],
)
TASK_DURATION_SECONDS = Histogram(
    "ariadne_task_duration_seconds",
    "Ariadne task durations in seconds",
    ["task", "status"],
    buckets=(0.5, 1, 2, 5, 10, 30, 60, 120, 300),
)

SCHEDULE_LAST_RUN_TS = Gauge(
    "ariadne_schedule_last_run_timestamp_seconds",
    "Last schedule run timestamp",
    ["task"],
)
SCHEDULE_LAST_SUCCESS_TS = Gauge(
    "ariadne_schedule_last_success_timestamp_seconds",
    "Last successful schedule run timestamp",
    ["task"],
)
SCHEDULE_LAST_ERROR_TS = Gauge(
    "ariadne_schedule_last_error_timestamp_seconds",
    "Last failed schedule run timestamp",
    ["task"],
)
SCHEDULE_NEXT_RUN_TS = Gauge(
    "ariadne_schedule_next_run_timestamp_seconds",
    "Next scheduled run timestamp",
    ["task"],
)
SCHEDULE_STATUS = Gauge(
    "ariadne_schedule_last_status",
    "Last schedule status (1=ok,0=error)",
    ["task"],
)

ACCESS_REQUESTS = Gauge(
    "ariadne_access_requests_total",
    "Access requests by status",
    ["status"],
)

CLUSTER_STATE_LAST_TS = Gauge(
    "ariadne_cluster_state_timestamp_seconds",
    "Last cluster state snapshot timestamp",
)
CLUSTER_STATE_NODES_TOTAL = Gauge(
    "ariadne_cluster_nodes_total",
    "Cluster nodes total",
)
CLUSTER_STATE_NODES_READY = Gauge(
    "ariadne_cluster_nodes_ready",
    "Cluster nodes Ready",
)
CLUSTER_STATE_PODS_RUNNING = Gauge(
    "ariadne_cluster_pods_running",
    "Cluster pods Running",
)
CLUSTER_STATE_KUSTOMIZATIONS_NOT_READY = Gauge(
    "ariadne_cluster_kustomizations_not_ready",
    "Flux kustomizations not Ready",
)
GAME_MODE_ACTIVE = Gauge(
    "ariadne_game_mode_active",
    "Ariadne game mode state (1=active,0=idle)",
    ["node", "game"],
)
GAME_MODE_TRANSITIONS_TOTAL = Counter(
    "ariadne_game_mode_transitions_total",
    "Ariadne game mode transitions by action and status",
    ["action", "status", "game"],
)
GAME_MODE_MANAGED_REPLICAS = Gauge(
    "ariadne_game_mode_managed_replicas",
    "Replicas for workloads managed by Ariadne game mode",
    ["namespace", "deployment"],
)
GAME_MODE_LAST_TRANSITION_TS = Gauge(
    "ariadne_game_mode_last_transition_timestamp_seconds",
    "Last Ariadne game mode transition timestamp",
    ["action", "status", "game"],
)
GAME_STREAM_FIREWALL_UNLOCKS_TOTAL = Counter(
    "ariadne_game_stream_firewall_unlocks_total",
    "Game-stream firewall unlock attempts by status",
    ["status"],
)
GAME_STREAM_FIREWALL_ACTIVE_UNLOCKS = Gauge(
    "ariadne_game_stream_firewall_active_unlocks",
    "Active game-stream firewall unlock count reported by the gatekeeper",
)
GAME_STREAM_PAIRING_ATTEMPTS_TOTAL = Counter(
    "ariadne_game_stream_pairing_attempts_total",
    "Moonlight pairing attempts through Ariadne by status",
    ["status"],
)


def record_task_run(task: str, status: str, duration_sec: float | None) -> None:
    """Increment task counters and duration histograms for one run."""

    TASK_RUNS_TOTAL.labels(task=task, status=status).inc()
    if duration_sec is not None:
        TASK_DURATION_SECONDS.labels(task=task, status=status).observe(duration_sec)


def record_schedule_state(task: str, last_run_ts: float | None, last_success_ts: float | None, next_run_ts: float | None, ok: bool | None) -> None:
    """Publish the latest scheduler timestamps and status for a task."""

    if last_run_ts:
        SCHEDULE_LAST_RUN_TS.labels(task=task).set(last_run_ts)
    if last_success_ts:
        SCHEDULE_LAST_SUCCESS_TS.labels(task=task).set(last_success_ts)
    if next_run_ts:
        SCHEDULE_NEXT_RUN_TS.labels(task=task).set(next_run_ts)
    if ok is not None:
        SCHEDULE_STATUS.labels(task=task).set(1 if ok else 0)
        if ok is False and last_run_ts:
            SCHEDULE_LAST_ERROR_TS.labels(task=task).set(last_run_ts)


def set_access_request_counts(counts: dict[str, int]) -> None:
    """Set access-request gauges grouped by lifecycle status."""

    for status, count in counts.items():
        ACCESS_REQUESTS.labels(status=status).set(count)


def set_cluster_state_metrics(collected_at: datetime, nodes_total: int | None, nodes_ready: int | None, pods_running: float | None, kustomizations_not_ready: int | None) -> None:
    """Set cluster-state gauges from the most recent collector snapshot."""

    CLUSTER_STATE_LAST_TS.set(collected_at.timestamp())
    if nodes_total is not None:
        CLUSTER_STATE_NODES_TOTAL.set(nodes_total)
    if nodes_ready is not None:
        CLUSTER_STATE_NODES_READY.set(nodes_ready)
    if pods_running is not None:
        CLUSTER_STATE_PODS_RUNNING.set(pods_running)
    if kustomizations_not_ready is not None:
        CLUSTER_STATE_KUSTOMIZATIONS_NOT_READY.set(kustomizations_not_ready)


def set_game_mode_state(node: str, game: str, active: bool) -> None:
    """Publish the current game-mode state."""

    GAME_MODE_ACTIVE.labels(node=node or "unknown", game=game or "unknown").set(1 if active else 0)


def record_game_mode_transition(action: str, status: str, game: str) -> None:
    """Increment game-mode transition metrics."""

    labels = {"action": action or "unknown", "status": status or "unknown", "game": game or "unknown"}
    GAME_MODE_TRANSITIONS_TOTAL.labels(**labels).inc()
    GAME_MODE_LAST_TRANSITION_TS.labels(**labels).set(datetime.now().timestamp())


def set_game_mode_managed_replicas(namespace: str, deployment: str, replicas: int | None) -> None:
    """Publish desired replicas for a workload that game mode can move."""

    if replicas is not None:
        GAME_MODE_MANAGED_REPLICAS.labels(namespace=namespace, deployment=deployment).set(replicas)


def record_game_stream_firewall_unlock(status: str) -> None:
    """Increment the game-stream firewall unlock counter."""

    GAME_STREAM_FIREWALL_UNLOCKS_TOTAL.labels(status=status or "unknown").inc()


def set_game_stream_firewall_active_unlocks(count: int | None) -> None:
    """Publish the number of active Moonlight firewall unlocks."""

    if count is not None:
        GAME_STREAM_FIREWALL_ACTIVE_UNLOCKS.set(count)


def record_game_stream_pairing_attempt(status: str) -> None:
    """Increment the game-stream pairing attempt counter."""

    GAME_STREAM_PAIRING_ATTEMPTS_TOTAL.labels(status=status or "unknown").inc()