FabledSteward/plugins/host_agent/scheduler.py

# plugins/host_agent/scheduler.py
from __future__ import annotations
import logging
from datetime import datetime, timezone
from typing import Iterable

from sqlalchemy import select

from steward.core.alerts import MetricBatch
from steward.core.scheduler import ScheduledTask
from steward.core.watch import watch
from steward.hosts.service import agent_stale_after_seconds
from steward.models.hosts import Host
from steward.models.metrics import PluginMetricLatest
from .downtime import announced_stops, restart_grace_seconds
from .models import HostAgentRegistration

logger = logging.getLogger(__name__)


def arrivals(app) -> dict[str, datetime]:
    """{host id: when its agent's newest post reached Steward}, in memory (#5514).

    last_seen_at is written in the ingest's own transaction, so when Steward
    can't store a post (its storage stalls and the request is dropped) the
    agent's report is lost with it, and an agent posting every few seconds read
    as offline. A post that arrived is the agent reaching Steward; storing it
    is Steward's part. Set before any write, so it holds however slow they are.
    """
    got = getattr(app, "_agent_arrivals", None)
    if got is None:
        got = app._agent_arrivals = {}
    return got


def seconds_since_seen(regs: Iterable, host_names: dict[str, str],
                       now: datetime, watching_since: datetime | None = None,
                       arrived: dict[str, datetime] | None = None,
                       already_silent: set[str] | frozenset = frozenset(),
                       explained_until: dict[str, datetime] | None = None,
                       ) -> list[tuple[str, float]]:
    """(host name, seconds since its agent last reported) per registration.

    An agent that has never reported counts from when it was registered. A
    registration whose host is gone is skipped. Nothing counts from before
    `watching_since`, when Steward last started watching without a break: an
    agent can't report while Steward is down, so that time is not its silence
    (#5441). An agent that is really gone reads as offline once it has been
    quiet that long while Steward watched. `arrived` (arrivals(app)) counts a
    post Steward received but couldn't store as the agent seen (#5514).

    `already_silent` names hosts whose last recorded figure was already past
    the offline window: their agent was gone before Steward stopped, so their
    whole silence counts. Counting them from `watching_since` read them as
    just seen, which resolved their alert on every restart and fired it again
    three minutes later (#5719).

    `explained_until` ({host id: time}) is a stop the agent announced, a reboot
    or its own restart, plus its grace (#5756): until then the agent reads as
    just seen. A host still silent past it reads its whole silence, so agent
    offline fires as the grace runs out, for silence nobody explained.
    """
    out = []
    for r in regs:
        name = host_names.get(r.host_id)
        seen = max((t for t in (r.last_seen_at or r.created_at, (arrived or {}).get(r.host_id))
                    if t is not None), default=None)
        if name is None or seen is None:
            continue
        if watching_since is not None and name not in already_silent:
            seen = max(seen, watching_since)
        if r.host_id in (explained_until or {}):
            seen = max(seen, explained_until[r.host_id])
        out.append((name, max(0.0, (now - seen).total_seconds())))
    return out


async def record_seconds_since_seen(app) -> None:
    """Record host_agent/<host>/seconds_since_seen for every registered agent,
    so "agent offline" is an ordinary alert rule (seconds_since_seen > 180)."""
    now = datetime.now(timezone.utc)
    async with app.db_sessionmaker() as session:
        async with session.begin():
            regs = (await session.execute(select(HostAgentRegistration))).scalars().all()
            watching_since = watch(app, "host_agent_seconds_since_seen").tick(now)
            if not regs:
                return
            host_names = dict((await session.execute(
                select(Host.id, Host.name).where(Host.id.in_([r.host_id for r in regs]))
            )).all())
            # Agents whose last figure was already past the offline window keep
            # counting through a restart of Steward (#5719).
            stale_after = agent_stale_after_seconds(app.config)
            already_silent = set((await session.execute(
                select(PluginMetricLatest.resource_name).where(
                    PluginMetricLatest.source_module == "host_agent",
                    PluginMetricLatest.metric_name == "seconds_since_seen",
                    PluginMetricLatest.resource_name.in_(list(host_names.values())),
                    PluginMetricLatest.value > stale_after))).scalars())
            explained = await announced_stops(session, now, restart_grace_seconds(app.config))
            batch = MetricBatch()
            for name, seconds in seconds_since_seen(regs, host_names, now, watching_since,
                                                    arrivals(app), already_silent, explained):
                batch.add(source_module="host_agent", resource_name=name,
                          metric_name="seconds_since_seen", value=seconds, recorded_at=now)
            await batch.record(session)


def make_task(app) -> ScheduledTask:
    async def _record():
        try:
            await record_seconds_since_seen(app)
        except Exception:
            logger.exception("host_agent: recording seconds_since_seen failed")

    return ScheduledTask(
        name="host_agent_seconds_since_seen",
        coro_factory=_record,
        interval_seconds=60,
        run_on_startup=False,
    )