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 a batch at a time in the background
instead, one key after another, 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
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
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. The
    latest-value table stands for the raw samples, as a key stays in it until
    its samples have aged out."""
    found: set[str] = set()
    for model in (PluginMetricLatest, PluginMetricHourly):
        found.update((await session.execute(
            select(model.resource_name).distinct().where(
                model.source_module == "docker", model.metric_name.notin_(TRAFFIC),
                model.resource_name.op("~")(_TASK_KEY)))).scalars())
    for model in (MetricAnomaly, AlertRule):
        found.update((await session.execute(
            select(model.resource_name).distinct().where(
                model.source_module == "docker", model.resource_name.op("~")(_TASK_KEY)))).scalars())
    found.update((await session.execute(
        select(MaintenanceWindow.scope_resource).distinct().where(
            MaintenanceWindow.scope_source == "docker",
            MaintenanceWindow.scope_resource.op("~")(_TASK_KEY)))).scalars())
    return sorted(k for k in found if task_service(k.split("/", 1)[1]))


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))
    await session.execute(update(MetricAnomaly).where(
        MetricAnomaly.source_module == "docker", MetricAnomaly.resource_name == old)
        .values(resource_name=new).execution_options(synchronize_session=False))
    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:
            async with session.begin():
                if remaining is None:
                    remaining = await old_keys(session)
                    logger.info("Swarm history under task names: %d keys to move", len(remaining))
                if remaining and await move_key(session, remaining[-1]):
                    remaining.pop()

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