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 json
import logging
import os
import random
import secrets
import socket
import time
import uuid

from sqlalchemy import select

from app.config import get_settings
from app.db import session_scope
from app.models import Account
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

    # --- 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)),
        ]
        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):
                self._start_runner(account_id)

    async def _renew_once(self) -> None:
        for account_id in list(self._runners):
            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()

    def _start_runner(self, account_id: uuid.UUID) -> None:
        runner = AccountRunner(account_id)
        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:
        payload = json.dumps({"owned": len(self._runners), "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)