FabledSteward/plugins/host_agent/maintenance.py

# plugins/host_agent/maintenance.py
"""A host's maintenance job, followed to its end through any restart (#5767).

Steward plans and the host's agent acts: a job is one host's update and, only
when the update says one is needed, its reboot, with the host's swarm node
drained first and made active again after. Each action is one step the agent
hands to its root helper (maint.py, #5766).

A job first asks the helper what the update would change (#6098). Docker's
own packages restart the daemon from their install scripts, stopping every
container on the host, so a swarm node is drained before an update that
upgrades Docker; otherwise only before a reboot. After a drain the node's
tasks get DRAIN_SETTLE to move before the update or reboot goes ahead. The
containers running outside the swarm when the job began must be running again
when it ends, within RECOVER_LIMIT, or the job fails naming them.

Everything lives on the job's row, so Steward restarting, or the host Steward
runs on rebooting, loses nothing:
the next post from the agent carries the step's result, and the `started`
event a host sends after booting is what completes a reboot.

The transitions here are model-free (a job is anything with the attributes),
so they're unit-tested without a database; routes.ingest and the scheduler's
tick apply them.

States: queued → [checking →] [draining → settling →] updating →
[draining → settling →] [rebooting →] [undraining →] [recovering →] done, or
failed / timed_out. A node is drained at most once a job: before the update
when Docker is upgraded, else before the reboot. settling and recovering are
Steward's own waits, moved on by the minute tick. An agent older than
CHECK_AGENT gets no check: its job starts at updating.
"""
from __future__ import annotations

from datetime import datetime, timedelta, timezone

ACTIVE = ("queued", "checking", "draining", "settling", "updating", "rebooting",
          "undraining", "recovering")
FINISHED = ("done", "failed", "timed_out")
# The helper action each handed state asks for.
STEP_ACTION = {"checking": "check", "updating": "update", "draining": "drain",
               "rebooting": "reboot", "undraining": "undrain"}
# How long a state may last before the job is given up. Queued counts from the
# job's creation (the host may be silent); the rest from their hand-off.
DEADLINES = {
    "queued": timedelta(hours=1),
    "checking": timedelta(minutes=30),     # two listings, the helper stops each at 10
    "draining": timedelta(minutes=10),
    "settling": timedelta(minutes=15),     # the tick moves it on at DRAIN_SETTLE
    "updating": timedelta(minutes=90),     # the helper itself stops at 60
    "rebooting": timedelta(minutes=30),
    "undraining": timedelta(minutes=10),
    "recovering": timedelta(minutes=15),   # the tick fails it at RECOVER_LIMIT
}
# After a drain, how long the node's swarm tasks get to move before the
# update or the reboot goes ahead (the operator, #6098: "drain wait 5 minutes
# and then reboot").
DRAIN_SETTLE = timedelta(minutes=5)
# How long the containers that ran outside the swarm before the job get to
# run again after it.
RECOVER_LIMIT = timedelta(minutes=10)
# The first agent whose helper answers a check (#6098).
CHECK_AGENT = (1, 24, 0)
# A restart policy under which Docker starts a container again after a reboot.
RESTARTS = ("always", "unless-stopped")
CONTAINERS_KEPT = 40
# A host's clock and Steward's differ by seconds; a boot this far before the
# reboot was handed out is still taken as that reboot.
BOOT_SLACK = timedelta(seconds=60)
DETAIL_MAX = 500


def step_id(job, state: str) -> str:
    """The helper request id for a job's step: unique per job and step, and in
    the helper's id alphabet (letters, digits, dashes)."""
    return f"{job.id}-{STEP_ACTION[state]}"


def _enter(job, state: str, now: datetime) -> None:
    job.state, job.state_at = state, now
    job.step_id = step_id(job, state) if state in STEP_ACTION else None
    job.step_seen = False               # a new step is handed out until its result arrives
    if state in FINISHED:
        job.finished_at = now


def _finish(job, state: str, now: datetime, detail: str | None = None) -> None:
    if detail:
        job.detail = detail[:DETAIL_MAX]
    _enter(job, state, now)


def can_check(agent_version) -> bool:
    """Whether an agent this version hands the helper a check."""
    try:
        return tuple(int(x) for x in str(agent_version).split(".")[:3]) >= CHECK_AGENT
    except ValueError:
        return False


def start(job, now: datetime, check: bool = False) -> None:
    """The host's agent is about to be handed the job: what the update would
    change first, when its helper can say, else the update."""
    if job.state == "queued":
        _enter(job, "checking" if check else "updating", now)


def drains(job) -> bool:
    """Whether the job drains the host's swarm node: not when the host isn't
    in a swarm, nor when the operator keeps its node drained."""
    return bool(job.node_id and job.drain_host_id)


def outside_swarm(job) -> list:
    """The containers the check found running outside the swarm."""
    return (job.checked or {}).get("containers") or []


def executor(job) -> str:
    """The host whose agent carries out the current step: the job's own host,
    except a drain or undrain, which a swarm manager does."""
    if job.state in ("draining", "undraining") and job.drain_host_id:
        return job.drain_host_id
    return job.host_id


def handoff(job) -> dict | None:
    """What to hand the executing agent now, or None. Handed on every post
    until a result for the step arrives; the agent and the helper both ignore
    an id they've already taken, so a repeat is harmless."""
    if job.state not in STEP_ACTION or job.step_seen:
        return None
    out = {"id": job.step_id, "action": STEP_ACTION[job.state]}
    if job.state in ("draining", "undraining") and job.node_id:
        out["node"] = job.node_id
    return out


REBOOT_POLICIES = ("if_required", "always", "never")


def wants_reboot(job) -> bool:
    """Reboot only when the update says one is needed (the operator's
    default), always (a plan may ask), or never."""
    if job.reboot_policy == "always":
        return True
    return bool(job.reboot_needed) and job.reboot_policy == "if_required"


def _checked(event: dict) -> dict:
    """What a check said, kept on the job: bounded, and only the fields used."""
    containers = []
    for c in event.get("containers") or []:
        if isinstance(c, dict) and isinstance(c.get("name"), str) and c["name"]:
            containers.append({"name": c["name"][:64], "restart": str(c.get("restart") or "no")[:16]})
    packages = [str(p)[:64] for p in event.get("docker_packages") or [] if isinstance(p, str)]
    pending = event.get("pending")
    return {"pending": pending if isinstance(pending, int) else None,
            "docker_packages": packages[:10], "containers": containers[:CONTAINERS_KEPT]}


def _after_check(job, event: dict, now: datetime) -> None:
    job.docker_update = event.get("docker_update") is True
    job.checked = _checked(event)
    _enter(job, "draining" if job.docker_update and drains(job) else "updating", now)


def _after_update(job, now: datetime) -> None:
    if not wants_reboot(job):
        _after_work(job, now)
    elif drains(job) and job.drained_at is None:
        _enter(job, "draining", now)
    else:
        _enter(job, "rebooting", now)


def _after_work(job, now: datetime) -> None:
    """The update (and any reboot) is over: make the node active again if the
    job drained it, then see the containers back."""
    if job.drained_at is not None and drains(job):
        _enter(job, "undraining", now)
    else:
        _recover(job, now)


def _recover(job, now: datetime) -> None:
    _enter(job, "recovering" if outside_swarm(job) else "done", now)


def apply_result(job, event: dict, now: datetime) -> bool:
    """Apply a `maintenance` event for this job's current step; True when the
    job changed. A result for a step the job has moved past is ignored, so a
    repeated or late report does nothing."""
    if job.state not in STEP_ACTION or event.get("id") != job.step_id:
        return False
    state = event.get("state")
    if state in ("running", "rebooting"):
        if job.step_seen:
            return False
        job.step_seen = True
        return True
    if state not in ("done", "failed"):
        return False
    if state == "failed" or event.get("ok") is False:
        what = STEP_ACTION[job.state]
        detail = f"{what} failed: {event.get('detail') or 'no detail'}"
        if job.state == "undraining":
            detail += " (the node is still drained)"
        _finish(job, "failed", now, detail)
        return True
    if job.state == "checking":
        _after_check(job, event, now)
    elif job.state == "updating":
        if isinstance(event.get("packages"), int):
            job.packages = event["packages"]
        job.reboot_needed = event.get("reboot_required") is True
        _after_update(job, now)
    elif job.state == "draining":
        job.drained_at = now
        _enter(job, "settling", now)
    elif job.state == "undraining":
        _recover(job, now)
    else:                               # rebooting: the boot completes it, not this
        return False
    return True


def apply_boot(job, boot_at: datetime | None, now: datetime) -> bool:
    """A `started` event from the job's host. A boot after the reboot was
    handed out is that reboot: the host is back."""
    if job.state != "rebooting" or boot_at is None or job.state_at is None:
        return False
    if boot_at < job.state_at - BOOT_SLACK:
        return False                    # the agent restarted; the host didn't
    job.back_at = now
    _after_work(job, now)
    return True


def settle(job, now: datetime) -> bool:
    """A drained node has had DRAIN_SETTLE for its tasks to move: on to the
    update, or to the reboot when the update is done (reboot_needed is set)."""
    if job.state != "settling" or job.state_at is None or now - job.state_at < DRAIN_SETTLE:
        return False
    _enter(job, "updating" if job.reboot_needed is None else "rebooting", now)
    return True


def recovered(job, running: set, now: datetime) -> bool:
    """Done once every container that ran outside the swarm before the job
    runs again (`running`: the host's containers seen running since the job
    came to this state); failed, naming the rest, past RECOVER_LIMIT."""
    if job.state != "recovering" or job.state_at is None:
        return False
    missing = [c for c in outside_swarm(job) if c["name"] not in running]
    if not missing:
        _enter(job, "done", now)
        return True
    if now - job.state_at < RECOVER_LIMIT:
        return False
    names = [c["name"] + ("" if c["restart"] in RESTARTS else " (no restart policy)")
             for c in missing[:6]]
    more = f" and {len(missing) - 6} more" if len(missing) > 6 else ""
    _finish(job, "failed", now, f"not running again {int(RECOVER_LIMIT.total_seconds() // 60)} min "
                                f"after the job: {', '.join(names)}{more}")
    return True


def check_deadline(job, now: datetime) -> bool:
    """Give up on a job whose state has lasted past its deadline."""
    limit = DEADLINES.get(job.state)
    since = job.state_at or job.created_at
    if limit is None or since is None or now - since <= limit:
        return False
    waiting = {"queued": "the host's agent didn't check in",
               "checking": "the check of what the update would change didn't finish",
               "draining": "the node wasn't drained",
               "settling": "the wait after the drain didn't end",
               "updating": "the update didn't finish",
               "rebooting": "the host didn't come back",
               "undraining": "the node wasn't made active again",
               "recovering": "its containers weren't seen again"}[job.state]
    _finish(job, "timed_out", now, f"{waiting} within {int(limit.total_seconds() // 60)} min")
    return True


def window_end(job, now: datetime) -> datetime:
    """When the job's maintenance window should end: the longest the rest of
    the job could take, or now once it has finished."""
    if job.state in FINISHED:
        return now
    left = sum((DEADLINES[s] for s in ACTIVE[ACTIVE.index(job.state):]), timedelta())
    return now + left


def utcnow() -> datetime:
    return datetime.now(timezone.utc)


# ── applied to the database: ingest, the tick, and starting a job ────────────

async def _window(session, job, host_name: str, now: datetime) -> None:
    """Keep the job's maintenance window over the host until the job ends: one
    scoped to this host's agent metrics, so an alert the job causes waits.
    A swarm node's job keeps a second over the swarm-wide figures (#5788):
    draining and rebooting a node takes it down and can leave a service
    pinned to it short, which is the job, not news."""
    from steward.models.alerts import MaintenanceWindow
    if job.window_id:
        w = await session.get(MaintenanceWindow, job.window_id)
        if w is not None:
            w.end_at = max(w.start_at, window_end(job, now))
    elif job.created_by is None or job.state in FINISHED:
        return                          # a window needs an owner; none for a finished job
    else:
        w = MaintenanceWindow(name=f"Maintenance: {host_name}"[:128], start_at=now,
                              end_at=window_end(job, now), scope_source="host_agent",
                              scope_resource=host_name, created_by=job.created_by)
        session.add(w)
        await session.flush()
        job.window_id = w.id
    if job.node_id:
        await _swarm_window(session, job, host_name, now)


SWARM_SCOPE = ("docker", "swarm")       # the swarm-wide figures' source and resource


async def _swarm_window(session, job, host_name: str, now: datetime) -> None:
    """The swarm half of a swarm node's job window, found again by its name
    and the job's start, so it needs no column of its own."""
    from sqlalchemy import select
    from steward.models.alerts import MaintenanceWindow
    name = f"Maintenance: {host_name} (swarm)"[:128]
    w = (await session.execute(select(MaintenanceWindow).where(
        MaintenanceWindow.name == name, MaintenanceWindow.scope_source == SWARM_SCOPE[0],
        MaintenanceWindow.scope_resource == SWARM_SCOPE[1],
        MaintenanceWindow.start_at >= job.created_at))).scalars().first()
    if w is not None:
        w.end_at = max(w.start_at, window_end(job, now))
    elif job.created_by is not None and job.state not in FINISHED:
        session.add(MaintenanceWindow(name=name, start_at=now, end_at=window_end(job, now),
                                      scope_source=SWARM_SCOPE[0], scope_resource=SWARM_SCOPE[1],
                                      created_by=job.created_by))


async def _host_name(session, host_id: str) -> str:
    from steward.models.hosts import Host
    host = await session.get(Host, host_id)
    return host.name if host is not None else host_id


async def apply_events(session, host_id: str, events, now: datetime) -> int:
    """Apply one post's `maintenance` results and `started` boots to the jobs
    they belong to; returns how many jobs changed. A result is matched by its
    step id, so a drain reported by a manager reaches the job of the node it
    drained."""
    from sqlalchemy import select
    from .downtime import _at
    from .models import HostMaintenanceJob
    if not isinstance(events, list):
        return 0
    changed = set()
    for event in events:
        if not isinstance(event, dict):
            continue
        if event.get("kind") == "maintenance" and isinstance(event.get("id"), str):
            job = (await session.execute(select(HostMaintenanceJob).where(
                HostMaintenanceJob.step_id == event["id"],
                HostMaintenanceJob.state.in_(ACTIVE)))).scalar_one_or_none()
            if job is not None and apply_result(job, event, now):
                changed.add(job)
        elif event.get("kind") == "started":
            boot_at = _at(event.get("boot_at"))
            for job in (await session.execute(select(HostMaintenanceJob).where(
                    HostMaintenanceJob.host_id == host_id,
                    HostMaintenanceJob.state == "rebooting"))).scalars():
                if apply_boot(job, boot_at, now):
                    changed.add(job)
    for job in changed:
        await _window(session, job, await _host_name(session, job.host_id), now)
    return len(changed)


async def handoffs(session, host_id: str, host_name: str, now: datetime,
                   agent_version: str | None = None) -> list[dict]:
    """The steps this host's agent should carry out now. A queued job for the
    host starts here, when its agent is known to be listening, with a check
    when the agent is new enough to hand one over."""
    from sqlalchemy import or_, select
    from .models import HostMaintenanceJob
    jobs = (await session.execute(select(HostMaintenanceJob).where(
        HostMaintenanceJob.state.in_(ACTIVE),
        or_(HostMaintenanceJob.host_id == host_id,
            HostMaintenanceJob.drain_host_id == host_id)
    ).order_by(HostMaintenanceJob.created_at))).scalars().all()
    out = []
    for job in jobs:
        if job.state == "queued" and job.host_id == host_id:
            start(job, now, check=can_check(agent_version))
            await _window(session, job, host_name, now)
        if executor(job) == host_id:
            step = handoff(job)
            if step is not None:
                out.append(step)
    return out


async def _running(session, job) -> set:
    """The job's host's containers seen running since the job came to its
    state; without the Docker plugin, the ones it waits for, as nothing can
    say otherwise."""
    from sqlalchemy import select
    models = docker_models()
    if models is None:
        return {c["name"] for c in outside_swarm(job)}
    DockerContainer = models.DockerContainer
    return set((await session.execute(select(DockerContainer.name).where(
        DockerContainer.host_id == job.host_id, DockerContainer.status == "running",
        DockerContainer.scraped_at >= job.state_at))).scalars())


async def tick(session, now: datetime) -> int:
    """Move on the jobs waiting on Steward (the wait after a drain, the
    containers seen again) and give up on jobs past their deadline; returns
    how many were given up."""
    from sqlalchemy import select
    from .models import HostMaintenanceJob
    gave_up = 0
    for job in (await session.execute(select(HostMaintenanceJob).where(
            HostMaintenanceJob.state.in_(ACTIVE)))).scalars():
        moved = settle(job, now)
        if job.state == "recovering":
            moved = recovered(job, await _running(session, job), now) or moved
        if check_deadline(job, now):
            gave_up += 1
            moved = True
        if moved:
            await _window(session, job, await _host_name(session, job.host_id), now)
    return gave_up


class CannotStart(Exception):
    """Why a host can't take a maintenance job now, in words for the page."""


def pick_drainer(node_hostname_to_id: dict, managers: dict, host_name: str,
                 host_id: str, helpers: set, availability: str = "active") -> tuple:
    """(node_id, drain_host_id) for a host in a swarm, (None, None) for one
    that isn't. `managers` maps a manager's host id to its swarm node id;
    `helpers` is the hosts whose agent has the maintenance helper. The host
    drains itself when it is a manager, else another manager does. A node the
    operator already set to drain or pause (a manager kept free of work) is
    left as it is: no drain before, and no making it active after."""
    node_id = node_hostname_to_id.get(host_name.lower())
    if node_id is None:
        return None, None
    if availability != "active":
        return node_id, None
    if host_id in managers and host_id in helpers:
        return node_id, host_id
    for mgr in sorted(managers):
        if mgr != host_id and mgr in helpers:
            return node_id, mgr
    raise CannotStart("it is a swarm node and no swarm manager has the maintenance "
                      "helper to drain it: run Update agent on a manager")


async def plan_job(session, host, reg, reboot_policy: str = "if_required") -> tuple:
    """(node_id, drain_host_id) for a job on `host`, or CannotStart saying why
    it can't have one now. The panel asks this before offering the button."""
    from sqlalchemy import select
    from .models import HostAgentRegistration, HostMaintenanceJob
    if reg is None or not reg.maint_helper:
        raise CannotStart("its agent has no maintenance helper: run Update agent first")
    if reboot_policy not in REBOOT_POLICIES:
        raise CannotStart(f"unknown reboot policy {reboot_policy!r}")
    busy = (await session.execute(select(HostMaintenanceJob.id).where(
        HostMaintenanceJob.host_id == host.id,
        HostMaintenanceJob.state.in_(ACTIVE)))).first()
    if busy is not None:
        raise CannotStart("a maintenance job is already running on it")
    nodes, managers, availability = await _swarm(session)
    helpers = set((await session.execute(select(HostAgentRegistration.host_id).where(
        HostAgentRegistration.maint_helper.is_(True)))).scalars())
    node_id = nodes.get(host.name.lower())
    try:
        return pick_drainer(nodes, managers, host.name, host.id, helpers,
                            availability.get(node_id, "active"))
    except CannotStart:
        # A job that never reboots needs a drain only if Docker is upgraded;
        # with no manager to drain, it runs undrained rather than not at all.
        if reboot_policy == "never":
            return None, None
        raise


async def create_job(session, host, reg, user_id: str | None,
                     reboot_policy: str = "if_required", now: datetime | None = None,
                     run_id: str | None = None):
    """Queue a maintenance job for `host`; it starts with the agent's next post."""
    from .models import HostMaintenanceJob
    now = now or utcnow()
    node_id, drain_host_id = await plan_job(session, host, reg, reboot_policy)
    job = HostMaintenanceJob(host_id=host.id, state="queued", state_at=now,
                             reboot_policy=reboot_policy, node_id=node_id,
                             drain_host_id=drain_host_id, created_by=user_id, created_at=now,
                             run_id=run_id)
    session.add(job)
    await session.flush()
    return job


async def latest_job(session, host_id: str, now: datetime, days: int = 366):
    """The host's running job, else its last one (within a year), for the panel."""
    from sqlalchemy import select
    from .models import HostMaintenanceJob
    return (await session.execute(select(HostMaintenanceJob).where(
        HostMaintenanceJob.host_id == host_id,
        HostMaintenanceJob.created_at >= now - timedelta(days=days),
    ).order_by(HostMaintenanceJob.state.in_(ACTIVE).desc(),
               HostMaintenanceJob.created_at.desc()).limit(1))).scalar_one_or_none()


LABELS = {
    "queued": "Waiting for the agent",
    "checking": "Checking what the update will change",
    "draining": "Draining its swarm node",
    "settling": "Giving its swarm tasks 5 min to move",
    "updating": "Updating",
    "rebooting": "Rebooting",
    "undraining": "Making its swarm node active again",
    "recovering": "Waiting for its containers to run again",
    "done": "Done",
    "failed": "Failed",
    "timed_out": "Timed out",
}


def summary(job) -> str:
    """One line on what the job did, for the panel."""
    if job.state not in FINISHED:
        return LABELS[job.state]
    parts = []
    if job.docker_update:
        parts.append("Docker upgraded")
    if job.packages is not None:
        parts.append(f"{job.packages} package{'s' if job.packages != 1 else ''} updated")
    if job.back_at is not None:
        parts.append("rebooted")
    elif job.reboot_needed and job.reboot_policy == "never":
        parts.append("reboot needed, not allowed by this job")
    elif job.state == "done" and job.reboot_needed is False:
        parts.append("no reboot needed")
    head = LABELS[job.state]
    return f"{head}: {', '.join(parts)}" if parts else head


def docker_models():
    """The loaded Docker plugin's models module, or None when it isn't loaded."""
    from steward.core.plugin_manager import loaded_plugin_submodule
    return loaded_plugin_submodule("docker", "models")


async def _swarm(session) -> tuple:
    """({node hostname: node id}, {manager host id: its node id},
    {node id: availability}) from what the Docker plugin last stored; empty
    when it isn't installed. A manager is a host whose agent reported the
    swarm's nodes."""
    models = docker_models()
    if models is None:
        return {}, {}, {}
    DockerSwarmNode = models.DockerSwarmNode
    from sqlalchemy import select
    from steward.models.hosts import Host
    rows = (await session.execute(select(DockerSwarmNode))).scalars().all()
    nodes = {(n.hostname or "").lower(): n.node_id for n in rows if n.hostname}
    reporters = {n.host_id for n in rows}
    names = dict((await session.execute(select(Host.id, Host.name).where(
        Host.id.in_(reporters)))).all()) if reporters else {}
    managers = {hid: nodes.get(name.lower()) for hid, name in names.items()}
    availability = {n.node_id: (n.availability or "active")
                    for n in sorted(rows, key=lambda n: n.updated_at)}
    return nodes, managers, availability