"""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: 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) # --- 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)