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