webmail-client/backend/app/sync/coordinator.py

"""Sync coordinator: lease-owned account claimer (refactor §2, step 3).

Replaces the startup-only supervisor. Each sync worker runs one coordinator that:

* claims accounts from ``sync:live`` (up to ``MAX_ACCOUNTS_PER_WORKER``),
  starting an :class:`AccountRunner` for each lease it wins;
* renews its leases every ``LEASE_RENEW`` seconds and, on losing one (a renewal
  that returns 0), stops that account's runner immediately (fencing);
* reconciles ``sync:live`` against Postgres periodically as a safety net;
* heartbeats a worker key so peers can compute a fair share and rebalance.

All ownership is via :mod:`app.sync.leases`; two workers can never run the same
account because claims are ``SET NX`` and renew/release are owner-checked.
"""

from __future__ import annotations

import asyncio
import contextlib
import json
import logging
import os
import random
import secrets
import socket
import time
import uuid
from datetime import UTC, datetime, timedelta

from sqlalchemy import func, select

from app.config import get_settings
from app.db import session_scope
from app.events import channel_for
from app.models import Account, Message
from app.redis_client import aw, get_redis
from app.sync import leases
from app.sync.runner import AccountRunner

logger = logging.getLogger(__name__)


def _new_worker_id() -> str:
    return f"{socket.gethostname()}:{os.getpid()}:{secrets.token_hex(4)}"


class SyncCoordinator:
    """Owns the leased account runners for a single sync worker."""

    def __init__(self, worker_id: str | None = None) -> None:
        self.worker_id = worker_id or _new_worker_id()
        self._redis = get_redis()
        self._scripts = leases.LeaseScripts(self._redis)
        self._runners: dict[uuid.UUID, AccountRunner] = {}
        self._stop = asyncio.Event()
        self._tasks: list[asyncio.Task[None]] = []

        s = get_settings()
        self.lease_ttl = s.lease_ttl
        self.lease_renew = s.lease_renew
        self.max_accounts = s.max_accounts_per_worker
        # Operational cadences (not operator-tunable; §2 protocol timings).
        self.claim_interval = 10.0
        self.claim_jitter = 5.0
        self.reconcile_interval = 300.0
        self.heartbeat_ttl = 30
        self.heartbeat_refresh = 10.0
        self.rebalance_interval = 60.0
        self.tier_eval_interval = 120.0
        self.session_grace_s = s.tier_session_grace_min * 60
        self.hot_activity = timedelta(hours=s.tier_hot_activity_h)
        self.cold_after = timedelta(days=s.tier_cold_after_d)
        self.hot_budget = s.hot_budget_per_worker

    # --- lifecycle -----------------------------------------------------------

    async def start(self) -> None:
        self._stop.clear()
        await self._reconcile_once()  # seed sync:live so the first claim isn't delayed
        await self._heartbeat_once()
        self._tasks = [
            asyncio.create_task(
                self._loop(self._claim_once, self.claim_interval, self.claim_jitter)
            ),
            asyncio.create_task(self._loop(self._renew_once, self.lease_renew)),
            asyncio.create_task(self._loop(self._reconcile_once, self.reconcile_interval)),
            asyncio.create_task(self._loop(self._heartbeat_once, self.heartbeat_refresh)),
            asyncio.create_task(self._loop(self._rebalance_once, self.rebalance_interval)),
            asyncio.create_task(self._loop(self._evaluate_tiers_once, self.tier_eval_interval)),
            asyncio.create_task(self._event_subscriber()),
        ]
        logger.info("sync coordinator %s started", self.worker_id)

    async def stop(self) -> None:
        self._stop.set()
        for task in self._tasks:
            task.cancel()
        if self._tasks:
            await asyncio.gather(*self._tasks, return_exceptions=True)
        self._tasks = []
        # Release every lease so peers can pick the accounts up without waiting for TTL.
        await asyncio.gather(
            *(self._release(aid) for aid in list(self._runners)), return_exceptions=True
        )
        await self._redis.delete(leases.worker_key(self.worker_id))
        logger.info("sync coordinator %s stopped", self.worker_id)

    @property
    def owned(self) -> set[uuid.UUID]:
        return set(self._runners)

    # --- periodic loop driver ------------------------------------------------

    async def _loop(self, step: object, interval: float, jitter: float = 0.0) -> None:
        while not self._stop.is_set():
            delay = interval + random.uniform(0.0, jitter)
            try:
                await asyncio.wait_for(self._stop.wait(), timeout=delay)
                return  # stop requested
            except TimeoutError:
                pass
            try:
                await step()  # type: ignore[operator]
            except asyncio.CancelledError:
                raise
            except Exception:
                logger.exception("coordinator step failed on worker %s", self.worker_id)

    # --- claim / renew / release ---------------------------------------------

    async def _claim_once(self) -> None:
        if len(self._runners) >= self.max_accounts:
            return
        async for raw in self._redis.sscan_iter(leases.LIVE_SET):
            if self._stop.is_set() or len(self._runners) >= self.max_accounts:
                break
            account_id = uuid.UUID(raw)
            if account_id in self._runners:
                continue
            if await leases.claim(self._redis, account_id, self.worker_id, self.lease_ttl):
                await self._start_runner(account_id)

    async def _renew_once(self) -> None:
        for account_id in list(self._runners):
            runner = self._runners.get(account_id)
            if runner is None:
                continue
            # Reap a runner that self-verified it is no longer live (§3).
            if runner.dead.is_set():
                logger.info("reaping dead runner for account %s", account_id)
                await self._release(account_id)
                continue
            if not await self._scripts.renew(account_id, self.worker_id, self.lease_ttl):
                logger.warning("lost lease for account %s; fencing runner", account_id)
                await self._drop(account_id)

    async def _release(self, account_id: uuid.UUID) -> None:
        """Graceful hand-off: stop the runner, then delete the lease if still ours."""
        runner = self._runners.pop(account_id, None)
        if runner is not None:
            await runner.stop()
        await self._scripts.release(account_id, self.worker_id)

    async def _drop(self, account_id: uuid.UUID) -> None:
        """Fencing: stop the runner without touching the lease (we no longer own it)."""
        runner = self._runners.pop(account_id, None)
        if runner is not None:
            await runner.stop()

    async def _start_runner(self, account_id: uuid.UUID) -> None:
        # Inherit any tier a previous owner recorded; default warm on first claim.
        tier = await aw(self._redis.get(leases.tier_key(account_id))) or "warm"
        await aw(self._redis.set(leases.tier_key(account_id), tier))
        runner = AccountRunner(account_id, tier)
        runner.start()
        self._runners[account_id] = runner

    # --- reconcile / heartbeat / rebalance -----------------------------------

    async def _reconcile_once(self) -> None:
        """Make ``sync:live`` match the set of live accounts in Postgres.

        Never touches leases — only set membership. Accounts that disappear here
        stop being claimed; their runners are dropped when their lease is removed.
        """
        async with session_scope() as session:
            live_db = {
                str(a)
                for a in await session.scalars(select(Account.id).where(Account.state == "live"))
            }
        live_redis = set(await aw(self._redis.smembers(leases.LIVE_SET)))
        to_add = live_db - live_redis
        to_remove = live_redis - live_db
        if to_add:
            await aw(self._redis.sadd(leases.LIVE_SET, *to_add))
        if to_remove:
            await aw(self._redis.srem(leases.LIVE_SET, *to_remove))
            # Stop any runners for accounts we own that are no longer live.
            for raw in to_remove:
                account_id = uuid.UUID(raw)
                if account_id in self._runners:
                    await self._release(account_id)

    async def _heartbeat_once(self) -> None:
        """Publish the worker's load for peer rebalancing + observability (§2)."""
        tiers = {"hot": 0, "warm": 0, "cold": 0}
        sockets_open = 0
        for runner in self._runners.values():
            tiers[runner.tier] = tiers.get(runner.tier, 0) + 1
            if runner.tier == "hot":
                sockets_open += 1  # one persistent IDLE connection per hot account
        payload = json.dumps(
            {
                "owned": len(self._runners),
                "sockets_open": sockets_open,
                "tiers": tiers,
                "ts": time.time(),
            }
        )
        await self._redis.set(leases.worker_key(self.worker_id), payload, ex=self.heartbeat_ttl)

    async def _alive_worker_count(self) -> int:
        count = 0
        async for _ in self._redis.scan_iter(match=f"{leases.WORKERS_PREFIX}*"):
            count += 1
        return max(count, 1)

    async def _rebalance_once(self) -> None:
        """Shed accounts if we hold well above a fair share of the live total."""
        total = int(await aw(self._redis.scard(leases.LIVE_SET)))
        workers = await self._alive_worker_count()
        target = total / workers
        threshold = int(target * 1.2) + 5
        excess = len(self._runners) - threshold
        if excess <= 0:
            return
        # Release the newest claims first (least likely to be mid-backfill on us).
        for account_id in list(self._runners)[-excess:]:
            await self._release(account_id)
        logger.info("worker %s shed %d account(s) to rebalance", self.worker_id, excess)

    # --- event-driven lifecycle (§3) -----------------------------------------

    async def _event_subscriber(self) -> None:
        """React to ``account.state_changed`` across processes (Redis pub/sub).

        Going live makes an account claimable at once; going away tears the
        runner down immediately (≤30s from the API call) instead of waiting for
        the lease to lapse.
        """
        pubsub = self._redis.pubsub()
        state_ch = channel_for("account.state_changed")
        active_ch = channel_for("user.active")
        await pubsub.subscribe(state_ch, active_ch)
        try:
            async for message in pubsub.listen():
                if self._stop.is_set():
                    break
                if message.get("type") != "message":
                    continue
                try:
                    payload = json.loads(message["data"]).get("payload", {})
                except (ValueError, TypeError):
                    continue
                if message.get("channel") == active_ch:
                    await self._on_user_active(payload)
                else:
                    await self._on_state_changed(payload)
        finally:
            with contextlib.suppress(Exception):
                await pubsub.aclose()  # type: ignore[no-untyped-call]

    async def _on_state_changed(self, payload: dict[str, object]) -> None:
        raw = payload.get("account_id")
        if not isinstance(raw, str):
            return
        account_id = uuid.UUID(raw)
        if payload.get("state") == "live":
            await leases.mark_live(self._redis, account_id)  # claimable next cycle
        elif account_id in self._runners:
            await self._release(account_id)  # disabled/deleted/error -> stop now

    async def _on_user_active(self, payload: dict[str, object]) -> None:
        """Promote all this user's owned accounts to hot at once (≤5s, §5)."""
        raw = payload.get("user_id")
        if not isinstance(raw, str):
            return
        user_id = uuid.UUID(raw)
        owned = list(self._runners)
        if not owned:
            return
        async with session_scope() as session:
            aids = list(
                await session.scalars(
                    select(Account.id).where(Account.id.in_(owned), Account.user_id == user_id)
                )
            )
        for account_id in aids:
            await self._set_tier(account_id, "hot")

    # --- tier evaluation (§5) ------------------------------------------------

    async def _set_tier(self, account_id: uuid.UUID, tier: str) -> None:
        runner = self._runners.get(account_id)
        if runner is None:
            return
        await aw(self._redis.set(leases.tier_key(account_id), tier))
        await runner.apply_tier(tier)

    async def _account_activity(
        self, ids: list[uuid.UUID]
    ) -> dict[uuid.UUID, tuple[uuid.UUID, datetime | None]]:
        """Return {account_id: (user_id, last_mail_at)} for the owned accounts."""
        async with session_scope() as session:
            rows = (
                await session.execute(
                    select(Account.id, Account.user_id, func.max(Message.internal_date))
                    .select_from(Account)
                    .outerjoin(Message, Message.account_id == Account.id)
                    .where(Account.id.in_(ids))
                    .group_by(Account.id, Account.user_id)
                )
            ).all()
        return {row[0]: (row[1], row[2]) for row in rows}

    async def _evaluate_tiers_once(self) -> None:
        """Recompute each owned account's tier from session heat + recent mail."""
        owned = list(self._runners)
        if not owned:
            return
        activity = await self._account_activity(owned)
        now = datetime.now(UTC)
        sessions: dict[uuid.UUID, bool] = {}
        desired: dict[uuid.UUID, str] = {}
        demotable: list[tuple[datetime, uuid.UUID]] = []  # activity-hot, no session
        hot_count = 0
        for account_id in owned:
            user_id, last_mail = activity.get(account_id, (None, None))
            session_active = False
            if user_id is not None:
                if user_id not in sessions:
                    sessions[user_id] = bool(
                        await aw(self._redis.exists(leases.session_key(user_id)))
                    )
                session_active = sessions[user_id]
            recent = last_mail is not None and (now - last_mail) < self.hot_activity
            if session_active or recent:
                desired[account_id] = "hot"
                hot_count += 1
                if not session_active and last_mail is not None:
                    demotable.append((last_mail, account_id))
            elif last_mail is not None and (now - last_mail) < self.cold_after:
                desired[account_id] = "warm"
            else:
                desired[account_id] = "cold"
        # Global hot budget: shed the least-recently-active session-less hot accounts.
        if hot_count > self.hot_budget:
            demotable.sort(key=lambda item: item[0])  # oldest activity first
            for _, account_id in demotable[: hot_count - self.hot_budget]:
                desired[account_id] = "warm"
        for account_id, tier in desired.items():
            await self._set_tier(account_id, tier)