FabledSteward/plugins/docker/rekey.py

# plugins/docker/rekey.py
"""Moves history kept under a swarm task's own name, "<host>/<service>.<slot>.
<task>", onto its service's key, "swarm/<service>": samples, hourly rollups,
anomalies and the alert rules and maintenance windows that name it.

The raw table holds millions of rows, too many to move in a migration that the
app has to wait for. A small task moves what it can in a short pass in the
background instead, newest history first, and does nothing once none are left. Network
rates stay where they are: they are kept per host, and age out with retention.
"""
from __future__ import annotations

import asyncio
import logging
import time
from typing import TYPE_CHECKING

from sqlalchemy import delete, select, text, update

from steward.core.scheduler import ScheduledTask
from steward.models.alerts import AlertRule, MaintenanceWindow
from steward.models.metrics import PluginMetric, PluginMetricHourly, PluginMetricLatest
from steward.models.norms import MetricAnomaly, MetricAnomalySpan
from steward.models.tickets import TicketLink
from steward.models.timeline import TimelineEntry
from steward.timeline.record import anomaly_changed

from .service import container_resource, task_service

if TYPE_CHECKING:
    from quart import Quart

logger = logging.getLogger(__name__)

BATCH = 5000            # raw samples moved in one go
EVERY_S = 15
PASS_S = 2              # a key with only a few rows takes milliseconds, and there can be many
TRAFFIC = ("net_rx_bps", "net_tx_bps")
_TASK_KEY = r"^[^/]+/.+\.([0-9]+|[a-z0-9]{25})\.[a-z0-9]{25}$"
_NOT_TRAFFIC = "metric_name NOT IN ('net_rx_bps', 'net_tx_bps')"

# Hours the service has no row for change name; hours it has are combined, as
# two tasks of a service can share an hour.
_RENAME_HOURS = text(f"""
UPDATE plugin_metrics_hourly o SET resource_name = :new
WHERE o.source_module = 'docker' AND o.resource_name = :old AND o.{_NOT_TRAFFIC}
  AND NOT EXISTS (SELECT 1 FROM plugin_metrics_hourly n
                  WHERE n.source_module = 'docker' AND n.resource_name = :new
                    AND n.metric_name = o.metric_name AND n.bucket = o.bucket)
""")
_COMBINE_HOURS = text("""
UPDATE plugin_metrics_hourly n SET
    value_avg = COALESCE((n.value_avg * n.sample_count + o.value_avg * o.sample_count)
                         / NULLIF(n.sample_count + o.sample_count, 0), n.value_avg),
    value_max = GREATEST(n.value_max, o.value_max),
    value_p95 = GREATEST(n.value_p95, o.value_p95),
    sample_count = n.sample_count + o.sample_count
FROM plugin_metrics_hourly o
WHERE o.source_module = 'docker' AND o.resource_name = :old
  AND n.source_module = 'docker' AND n.resource_name = :new
  AND n.metric_name = o.metric_name AND n.bucket = o.bucket
""")
_DROP_HOURS = text(f"""
DELETE FROM plugin_metrics_hourly
WHERE source_module = 'docker' AND resource_name = :old AND {_NOT_TRAFFIC}
""")


def service_key(old: str) -> str:
    """The key a task-named resource belongs under now."""
    return container_resource("", old.split("/", 1)[1])


async def old_keys(session) -> list[str]:
    """Every task-named docker resource still holding something to move, to be
    taken from the end: those with raw samples last, as they are what a chart
    and a normal read first. The latest-value table stands for the raw samples,
    as a key stays in it until its samples have aged out."""
    async def names(column, *where) -> set[str]:
        found = (await session.execute(select(column).distinct().where(*where, column.op("~")(_TASK_KEY)))).scalars()
        return {k for k in found if task_service(k.split("/", 1)[1])}

    recent = await names(PluginMetricLatest.resource_name, PluginMetricLatest.source_module == "docker",
                         PluginMetricLatest.metric_name.notin_(TRAFFIC))
    hourly = await names(PluginMetricHourly.resource_name, PluginMetricHourly.source_module == "docker",
                         PluginMetricHourly.metric_name.notin_(TRAFFIC))
    named = (await names(MetricAnomaly.resource_name, MetricAnomaly.source_module == "docker")
             | await names(AlertRule.resource_name, AlertRule.source_module == "docker")
             | await names(MaintenanceWindow.scope_resource, MaintenanceWindow.scope_source == "docker"))
    logger.info("Swarm history under task names: %d keys with raw samples, %d with hourly rows, "
                "%d named by anomalies, rules or windows", len(recent), len(hourly), len(named))
    return sorted(hourly - recent) + sorted((named - hourly) - recent) + sorted(recent)


async def _merge_anomalies(session, new: str) -> None:
    """One row per metric still to review, as for any key: the rows waiting on
    the service become the newest of them, with all their occurrences."""
    rows = (await session.execute(
        select(MetricAnomaly).where(
            MetricAnomaly.source_module == "docker", MetricAnomaly.resource_name == new,
            MetricAnomaly.review == "to_review")
        .order_by(MetricAnomaly.last_started_at.desc()))).scalars().all()
    keep: dict[str, MetricAnomaly] = {}
    for row in rows:
        into = keep.setdefault(row.metric_name, row)
        if into is row:
            continue
        await session.execute(update(MetricAnomalySpan).where(
            MetricAnomalySpan.anomaly_id == row.id).values(anomaly_id=into.id))
        into.occurrences += row.occurrences
        into.started_at = min(into.started_at, row.started_at)
        if row.ended_at is None:
            into.ended_at = None
        if row.peak_level > into.peak_level:
            into.peak_level, into.peak_at = row.peak_level, row.peak_at
            into.usual_hold, into.usual_peak = row.usual_hold, row.usual_peak
        linked = (await session.execute(select(TicketLink.ticket_id).where(
            TicketLink.ref_type == "metric_anomaly", TicketLink.ref_id == into.id))).scalars().all()
        await session.execute(update(TicketLink).where(
            TicketLink.ref_type == "metric_anomaly", TicketLink.ref_id == row.id,
            TicketLink.ticket_id.notin_(linked)).values(ref_id=into.id))
        await session.execute(delete(TicketLink).where(
            TicketLink.ref_type == "metric_anomaly", TicketLink.ref_id == row.id))
        await session.execute(delete(TimelineEntry).where(
            TimelineEntry.ref_type == "metric_anomaly", TimelineEntry.ref_id == row.id))
        await session.delete(row)


async def move_key(session, old: str) -> bool:
    """Move the next batch of `old`, and the rest of what it holds once its raw
    samples are done. Returns whether it is finished. Runs inside the caller's
    transaction."""
    new = service_key(old)
    batch = select(PluginMetric.id).where(
        PluginMetric.source_module == "docker", PluginMetric.resource_name == old,
        PluginMetric.metric_name.notin_(TRAFFIC)).limit(BATCH)
    moved = await session.execute(
        update(PluginMetric).where(PluginMetric.id.in_(batch)).values(resource_name=new)
        .execution_options(synchronize_session=False))
    if moved.rowcount >= BATCH:
        return False
    keys = {"old": old, "new": new}
    for statement in (_RENAME_HOURS, _COMBINE_HOURS, _DROP_HOURS):
        await session.execute(statement, keys)
    await session.execute(delete(PluginMetricLatest).where(
        PluginMetricLatest.source_module == "docker", PluginMetricLatest.resource_name == old,
        PluginMetricLatest.metric_name.notin_(TRAFFIC)))
    await session.execute(update(AlertRule).where(
        AlertRule.source_module == "docker", AlertRule.resource_name == old).values(resource_name=new))
    await session.execute(update(MaintenanceWindow).where(
        MaintenanceWindow.scope_source == "docker", MaintenanceWindow.scope_resource == old)
        .values(scope_resource=new))
    renamed = await session.execute(update(MetricAnomaly).where(
        MetricAnomaly.source_module == "docker", MetricAnomaly.resource_name == old)
        .values(resource_name=new).execution_options(synchronize_session=False))
    if not renamed.rowcount:
        return True
    session.expire_all()
    await _merge_anomalies(session, new)
    await session.flush()
    for row in (await session.execute(select(MetricAnomaly).where(
            MetricAnomaly.source_module == "docker", MetricAnomaly.resource_name == new))).scalars().all():
        await anomaly_changed(session, row)
    return True


def make_rekey_task(app: "Quart") -> ScheduledTask:
    remaining: list[str] | None = None
    busy = asyncio.Lock()

    async def rekey() -> None:
        nonlocal remaining
        if remaining == [] or busy.locked():
            return
        async with busy, app.db_sessionmaker() as session:
            if remaining is None:
                async with session.begin():
                    remaining = await old_keys(session)
            deadline = time.monotonic() + PASS_S
            while remaining and time.monotonic() < deadline:
                async with session.begin():
                    finished = await move_key(session, remaining[-1])
                if not finished:        # a full batch of samples is enough for one pass
                    break
                remaining.pop()
            if not remaining:
                logger.info("Swarm history under task names: all moved")

    return ScheduledTask(name="docker_rekey", coro_factory=rekey, interval_seconds=EVERY_S)