lukasandClaude Opus 4.8
19335116da
Each owned account is tiered by the coordinator; the runner shapes its work to
match — correctness identical, only latency/socket cost differ.
- hot: persistent INBOX IDLE + secondary poll every POLL_WARM.
- warm: no persistent socket; transient INBOX+special poll every POLL_WARM.
- cold: same transient poll every POLL_COLD.
- engine: poll_all_loop (transient full poll); secondary_poll_loop interval is
now a parameter; POLL_INTERVAL removed in favor of POLL_WARM/POLL_COLD.
- coordinator: tier eval every 120s from session heat (sess:active:{user}) +
recent mail (max internal_date within TIER_HOT_ACTIVITY_H), with the
HOT_BUDGET_PER_WORKER guard shedding least-recently-active session-less hot
accounts; user.active event promotes a user's accounts to hot at once (≤5s).
- API: login + WS connect/ping publish user.active and refresh sess:active
(TTL TIER_SESSION_GRACE_MIN). New sync:tier:{id} key, default warm on claim.
- config: POLL_WARM/POLL_COLD/TIER_*/HOT_BUDGET_PER_WORKER.
- tests: 5 tier tests (session/recent-mail promotion, cold, user.active,
budget demotion) alongside the lease/lifecycle acceptance tests.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
349 lines
15 KiB
Python
349 lines
15 KiB
Python
"""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)
|