Nexus/server/background/heartbeat_flush.py

"""Nexus — Heartbeat Flush Background Task
Every 10 minutes: batch-flush Redis heartbeat data to MySQL servers table.
MySQL only stores historical snapshots; frontend reads live data from Redis.
"""

import asyncio
import logging
from datetime import datetime, timezone

from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import AsyncSession

from server.infrastructure.database.session import AsyncSessionLocal
from server.infrastructure.redis.client import get_redis
from server.domain.models import Server
from server.infrastructure.database.fleet_metric_repo import FleetMetricRepositoryImpl
from server.infrastructure.database.server_metric_repo import ServerMetricRepositoryImpl
from server.utils.fleet_metrics import aggregate_fleet_metrics
from server.utils.server_metrics import sample_from_redis_heartbeat

logger = logging.getLogger("nexus.heartbeat_flush")

FLUSH_INTERVAL = 600  # 10 minutes (seconds)
REDIS_KEY_PREFIX = "heartbeat:"
# Redis TTL for live keys: server.api.agent.REDIS_KEY_EXPIRE (= FLUSH_INTERVAL * 1.5)


async def _write_server_metric_samples(
    session: AsyncSession,
    redis_by_id: dict[int, dict],
    recorded_at: datetime,
) -> int:
    """Insert per-server CPU/mem samples for all agent-installed servers."""
    result = await session.execute(
        select(Server.id).where(
            Server.agent_version.isnot(None),
            Server.agent_version != "",
        )
    )
    agent_ids = [row[0] for row in result.all()]
    if not agent_ids:
        return 0

    repo = ServerMetricRepositoryImpl(session)
    written = 0
    for server_id in agent_ids:
        data = redis_by_id.get(server_id)
        if data:
            is_online, cpu_pct, mem_pct = sample_from_redis_heartbeat(data)
        else:
            is_online, cpu_pct, mem_pct = False, None, None
        await repo.insert_sample(
            server_id=server_id,
            recorded_at=recorded_at,
            is_online=is_online,
            cpu_pct=cpu_pct,
            mem_pct=mem_pct,
        )
        written += 1
    return written


async def _flush_redis_keys(session: AsyncSession, keys: list[str], redis) -> tuple[int, int, list[dict], dict[int, dict]]:
    """Update servers from Redis keys; return fleet rows and id→heartbeat map."""
    flushed = 0
    skipped = 0
    heartbeat_rows: list[dict] = []
    redis_by_id: dict[int, dict] = {}

    for key in keys:
        try:
            server_id = int(key.replace(REDIS_KEY_PREFIX, ""))
            data = await redis.hgetall(key)

            if not data:
                continue

            row = dict(data)
            heartbeat_rows.append(row)
            redis_by_id[server_id] = row
            is_online = data.get("is_online") == "True"
            system_info = data.get("system_info", "{}")
            agent_version = data.get("agent_version", "")

            last_hb_str = data.get("last_heartbeat", "")
            try:
                last_heartbeat = datetime.fromisoformat(last_hb_str)
                if last_heartbeat.tzinfo is None:
                    last_heartbeat = last_heartbeat.replace(tzinfo=timezone.utc)
            except (ValueError, TypeError):
                last_heartbeat = datetime.now(timezone.utc)

            result = await session.execute(
                update(Server)
                .where(Server.id == server_id)
                .values(
                    is_online=is_online,
                    system_info=system_info,
                    last_heartbeat=last_heartbeat,
                    agent_version=agent_version,
                )
            )
            if (result.rowcount or 0) > 0:
                flushed += 1
            else:
                skipped += 1
        except Exception as e:
            logger.warning(f"Failed to flush server {key}: {e}")
            skipped += 1

    return flushed, skipped, heartbeat_rows, redis_by_id


async def _commit_flush_metrics(
    session: AsyncSession,
    heartbeat_rows: list[dict],
    redis_by_id: dict[int, dict],
    recorded_at: datetime,
) -> None:
    fleet_repo = FleetMetricRepositoryImpl(session)
    if heartbeat_rows:
        agg = aggregate_fleet_metrics(heartbeat_rows)
        await fleet_repo.insert_sample(
            recorded_at=recorded_at,
            online_count=agg.online_count,
            sample_count=agg.sample_count,
            cpu_avg=agg.cpu_avg,
            mem_avg=agg.mem_avg,
            disk_avg=agg.disk_avg,
        )
    purged_fleet = await fleet_repo.purge_older_than()

    server_samples = await _write_server_metric_samples(session, redis_by_id, recorded_at)
    server_repo = ServerMetricRepositoryImpl(session)
    purged_server = await server_repo.purge_older_than()

    await session.commit()
    if purged_fleet:
        logger.info("Fleet metric purge: removed %s rows older than retention", purged_fleet)
    if purged_server:
        logger.info("Server metric purge: removed %s rows older than retention", purged_server)
    if server_samples:
        logger.info("Server metric samples: wrote %s rows", server_samples)


async def heartbeat_flush_loop():
    """Background task: every 10min, flush Redis heartbeat data to MySQL.

    Flow:
    1. Scan all heartbeat:* keys in Redis
    2. For each key, extract server_id + heartbeat data
    3. Batch UPDATE MySQL servers table (is_online, system_info, last_heartbeat, agent_version)
    4. Write fleet + per-server metric samples; purge retention
    """
    logger.info("Heartbeat flush loop started (interval: 10min)")
    while True:
        await asyncio.sleep(FLUSH_INTERVAL)

        redis = get_redis()
        recorded_at = datetime.now(timezone.utc)

        try:
            keys: list[str] = []
            cursor = 0
            while True:
                cursor, batch = await redis.scan(cursor, match=f"{REDIS_KEY_PREFIX}*", count=100)
                keys.extend(batch)
                if cursor == 0:
                    break
        except Exception as e:
            logger.error(f"Redis scan failed: {e}")
            continue

        if not keys:
            logger.info("No heartbeat keys in Redis — writing offline server metric samples only")
            async with AsyncSessionLocal() as session:
                try:
                    await _commit_flush_metrics(session, [], {}, recorded_at)
                except Exception as e:
                    logger.error(f"Heartbeat flush commit failed (no Redis keys): {e}")
                    await session.rollback()
            continue

        flushed = 0
        skipped = 0
        async with AsyncSessionLocal() as session:
            try:
                flushed, skipped, heartbeat_rows, redis_by_id = await _flush_redis_keys(
                    session, keys, redis
                )
                await _commit_flush_metrics(session, heartbeat_rows, redis_by_id, recorded_at)
            except Exception as e:
                logger.error(f"Heartbeat flush commit failed: {e}")
                await session.rollback()
                continue

        logger.info(
            f"Heartbeat flush: {flushed}/{len(keys)} servers updated to MySQL"
            + (f", {skipped} failed" if skipped else "")
        )