diff --git a/backend/app/accounts/service.py b/backend/app/accounts/service.py index 3ca5b5f..4818195 100644 --- a/backend/app/accounts/service.py +++ b/backend/app/accounts/service.py @@ -80,6 +80,13 @@ async def delete_account( await session.execute(delete(Account).where(Account.id == account_id)) await session.flush() + # Deregister from sync:live and drop any lease so the owning sync worker + # stops the account's runner on its next renewal (refactor §2 fencing). + from app.redis_client import get_redis + from app.sync import leases + + await leases.mark_not_live(get_redis(), account_id) + # Content-addressed blobs are shared across the user's accounts (dedup): only # remove a blob when no remaining message/attachment references it. for sha in candidate_shas: diff --git a/backend/app/config.py b/backend/app/config.py index 483cafa..24b900d 100644 --- a/backend/app/config.py +++ b/backend/app/config.py @@ -40,6 +40,10 @@ class Settings(BaseSettings): worker_role: str = Field(default="all", alias="WORKER_ROLE") interactive_concurrency: int = Field(default=20, alias="INTERACTIVE_CONCURRENCY") bulk_concurrency: int = Field(default=2, alias="BULK_CONCURRENCY") + # Lease protocol (§2): account ownership across sync workers. + max_accounts_per_worker: int = Field(default=400, alias="MAX_ACCOUNTS_PER_WORKER") + lease_ttl: int = Field(default=60, alias="LEASE_TTL") + lease_renew: int = Field(default=20, alias="LEASE_RENEW") oauth_google_client_id: str = Field(default="", alias="OAUTH_GOOGLE_CLIENT_ID") oauth_google_client_secret: str = Field(default="", alias="OAUTH_GOOGLE_CLIENT_SECRET") diff --git a/backend/app/sync/coordinator.py b/backend/app/sync/coordinator.py new file mode 100644 index 0000000..cf86f87 --- /dev/null +++ b/backend/app/sync/coordinator.py @@ -0,0 +1,205 @@ +"""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) diff --git a/backend/app/sync/engine.py b/backend/app/sync/engine.py index 09a1fbc..68b1e84 100644 --- a/backend/app/sync/engine.py +++ b/backend/app/sync/engine.py @@ -41,6 +41,17 @@ async def _set_state(account_id: uuid.UUID, state: str, detail: str = "") -> Non await session.execute( update(Account).where(Account.id == account_id).values(state=state, state_detail=detail) ) + # Keep the sync:live registry in step with the account's state so the sync + # coordinator claims it once live and drops it otherwise (refactor §2). + from app.redis_client import get_redis + from app.sync import leases + + redis = get_redis() + if state == "live": + await leases.mark_live(redis, account_id) + else: + await leases.mark_not_live(redis, account_id) + bus = get_event_bus() account = None async with session_scope() as session: diff --git a/backend/app/sync/leases.py b/backend/app/sync/leases.py new file mode 100644 index 0000000..2918bfe --- /dev/null +++ b/backend/app/sync/leases.py @@ -0,0 +1,83 @@ +"""Redis lease primitives for account ownership (refactor §2). + +A sync worker must hold the lease ``sync:lease:{account_id}`` before it runs an +account's sync loops, and no two workers may hold the same lease. Claims are +atomic (``SET NX EX``); renewal and release are owner-checked Lua so a worker can +never renew or delete a lease that has been taken over by another worker. + +The set ``sync:live`` is the registry of accounts that *should* be synced; the +claimer only ever claims ids drawn from it. +""" + +from __future__ import annotations + +import uuid + +import redis.asyncio as aioredis + +from app.redis_client import aw + +LIVE_SET = "sync:live" +_LEASE_PREFIX = "sync:lease:" +WORKERS_PREFIX = "sync:workers:" + + +def lease_key(account_id: uuid.UUID) -> str: + return f"{_LEASE_PREFIX}{account_id}" + + +def worker_key(worker_id: str) -> str: + return f"{WORKERS_PREFIX}{worker_id}" + + +# Renew only if we still own the lease (guards against a stale worker extending a +# lease another worker has already claimed after expiry). Returns 1 on success. +_RENEW_LUA = """ +if redis.call('GET', KEYS[1]) == ARGV[1] then + return redis.call('EXPIRE', KEYS[1], ARGV[2]) +else + return 0 +end +""" + +# Delete only if we still own the lease. Returns 1 if we released it, else 0. +_RELEASE_LUA = """ +if redis.call('GET', KEYS[1]) == ARGV[1] then + return redis.call('DEL', KEYS[1]) +else + return 0 +end +""" + + +class LeaseScripts: + """Server-registered Lua for owner-checked renew/release.""" + + def __init__(self, redis: aioredis.Redis) -> None: + self._renew = redis.register_script(_RENEW_LUA) + self._release = redis.register_script(_RELEASE_LUA) + + async def renew(self, account_id: uuid.UUID, worker_id: str, ttl: int) -> bool: + result = await self._renew(keys=[lease_key(account_id)], args=[worker_id, ttl]) + return bool(result) + + async def release(self, account_id: uuid.UUID, worker_id: str) -> bool: + result = await self._release(keys=[lease_key(account_id)], args=[worker_id]) + return bool(result) + + +async def claim(redis: aioredis.Redis, account_id: uuid.UUID, worker_id: str, ttl: int) -> bool: + """Atomically claim the lease for ``account_id``; True iff we now own it.""" + ok = await redis.set(lease_key(account_id), worker_id, nx=True, ex=ttl) + return bool(ok) + + +async def mark_live(redis: aioredis.Redis, account_id: uuid.UUID) -> None: + """Register an account as one that should be synced (idempotent).""" + await aw(redis.sadd(LIVE_SET, str(account_id))) + + +async def mark_not_live(redis: aioredis.Redis, account_id: uuid.UUID) -> None: + """Deregister an account and drop any lease so its owner stops it (fencing).""" + await aw(redis.srem(LIVE_SET, str(account_id))) + await aw(redis.delete(lease_key(account_id))) diff --git a/backend/app/sync/supervisor.py b/backend/app/sync/supervisor.py index 357bcd3..4bb5ec8 100644 --- a/backend/app/sync/supervisor.py +++ b/backend/app/sync/supervisor.py @@ -1,46 +1,29 @@ -"""Live-sync supervisor (refactor §4, transitional). +"""Live-sync entrypoint (refactor §2, step 3). -Runs inside the sync worker: on startup it launches one :class:`AccountRunner` -per live account. This step (v2 §9.2) only extracts the per-account loops into a -runner; account discovery is still startup-only and is replaced by the lease -claimer in a later step. +The sync worker's ``on_startup`` starts a :class:`SyncCoordinator`, which claims +accounts from the ``sync:live`` registry via Redis leases and runs one +:class:`~app.sync.runner.AccountRunner` per owned account. ``on_shutdown`` stops +it, releasing every lease so peers can take the accounts over immediately. """ from __future__ import annotations -import asyncio import logging -import uuid from typing import Any -from sqlalchemy import select - -from app.db import session_scope -from app.models import Account -from app.sync.runner import AccountRunner +from app.sync.coordinator import SyncCoordinator logger = logging.getLogger(__name__) -async def _live_account_ids() -> list[uuid.UUID]: - async with session_scope() as session: - return list(await session.scalars(select(Account.id).where(Account.state == "live"))) - - async def start_supervisor(ctx: dict[str, Any]) -> None: - """Launch an AccountRunner per live account (called from the worker on_startup).""" - runners: dict[uuid.UUID, AccountRunner] = {} - for account_id in await _live_account_ids(): - runner = AccountRunner(account_id) - runner.start() - runners[account_id] = runner - ctx["runners"] = runners - logger.info("live supervisor started for %d account(s)", len(runners)) + coordinator = SyncCoordinator() + await coordinator.start() + ctx["coordinator"] = coordinator async def stop_supervisor(ctx: dict[str, Any]) -> None: - """Stop and tear down all account runners (called from the worker on_shutdown).""" - runners: dict[uuid.UUID, AccountRunner] = ctx.get("runners", {}) - if runners: - await asyncio.gather(*(r.stop() for r in runners.values()), return_exceptions=True) - ctx["runners"] = {} + coordinator: SyncCoordinator | None = ctx.get("coordinator") + if coordinator is not None: + await coordinator.stop() + ctx["coordinator"] = None diff --git a/backend/tests/sync/test_acceptance.py b/backend/tests/sync/test_acceptance.py new file mode 100644 index 0000000..bf12a46 --- /dev/null +++ b/backend/tests/sync/test_acceptance.py @@ -0,0 +1,226 @@ +"""Sync-engine v2 acceptance tests 1-3 (refactor §10). + +These drive the real :class:`SyncCoordinator` against the test Redis with the +per-account IMAP loops replaced by fakes, so the lease protocol — single +ownership, crash failover, and create/delete lifecycle — is exercised without a +live mail server. Intervals are shortened so the tests run in seconds. + + 1. Two workers, N accounts: every account owned by exactly one worker, exactly + one live (IDLE) connection each. + 2. Kill a worker: its orphaned accounts are re-owned by the survivor once the + leases expire, still one connection each (no double ownership). + 3. A newly-live account is claimed without a restart; a removed account's + connection is torn down (lease fencing), with no further ownership. +""" + +from __future__ import annotations + +import asyncio +import uuid +from collections import defaultdict +from collections.abc import AsyncIterator, Callable + +import pytest_asyncio +from sqlalchemy import delete +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine +from sqlalchemy.pool import NullPool + +from app.config import get_settings +from app.identity.passwords import hash_password +from app.models import Account, User +from app.redis_client import get_redis +from app.sync import coordinator as coordinator_mod +from app.sync import leases, runner +from app.sync.coordinator import SyncCoordinator + +# --- fakes: replace the real IMAP loops with connection-counting stand-ins ---- + + +@pytest_asyncio.fixture +async def connections(monkeypatch) -> AsyncIterator[dict[uuid.UUID, int]]: # type: ignore[no-untyped-def] + """Patch the runner's loops; yield a live-connection count per account.""" + conns: dict[uuid.UUID, int] = defaultdict(int) + + async def fake_live_loop(account_id: uuid.UUID, folder: str, stop: asyncio.Event) -> None: + conns[account_id] += 1 # one IDLE connection while the loop is alive + try: + await stop.wait() + finally: + conns[account_id] -= 1 + + async def fake_secondary(account_id: uuid.UUID, stop: asyncio.Event | None = None) -> None: + assert stop is not None + await stop.wait() + + monkeypatch.setattr(runner, "live_loop", fake_live_loop) + monkeypatch.setattr(runner, "secondary_poll_loop", fake_secondary) + yield conns + + +@pytest_asyncio.fixture +async def make_accounts() -> AsyncIterator[Callable[[int, str], object]]: + """Factory creating N live accounts; cleans up its users afterwards.""" + engine = create_async_engine(get_settings().database_url, poolclass=NullPool) + sm = async_sessionmaker(engine, expire_on_commit=False) + user_ids: list[uuid.UUID] = [] + + async def _make(n: int, state: str = "live") -> list[uuid.UUID]: + async with sm() as session: + user = User( + handle=f"acc{uuid.uuid4().hex[:8]}@test.local", password_hash=hash_password("pw") + ) + session.add(user) + await session.flush() + user_ids.append(user.id) + ids: list[uuid.UUID] = [] + for _ in range(n): + account = Account( + user_id=user.id, + email=f"{uuid.uuid4().hex[:10]}@example.com", + display_name="Acct", + imap_config={}, + smtp_config={}, + secret_ct=b"", + discovery_src="manual", + state=state, + ) + session.add(account) + await session.flush() + ids.append(account.id) + await session.commit() + return ids + + try: + yield _make + finally: + async with sm() as session: + for uid in user_ids: + await session.execute(delete(User).where(User.id == uid)) + await session.commit() + await engine.dispose() + + +async def _no_reconcile() -> None: + """Stub: keep the shared dev DB's real live accounts out of the test.""" + return None + + +def _fast(c: SyncCoordinator) -> SyncCoordinator: + """Compress cadences and drive sync:live explicitly (no DB reconcile). + + The tests share the dev Postgres, which holds real live accounts; reconciling + from it would pull them into the test. We seed ``sync:live`` ourselves instead. + """ + c.claim_interval = 0.05 + c.claim_jitter = 0.0 + c.lease_renew = 0.3 + c.lease_ttl = 2 # short so a crashed worker's leases expire quickly + c.reconcile_interval = 1000.0 + c.heartbeat_refresh = 1000.0 + c.rebalance_interval = 1000.0 + c._reconcile_once = _no_reconcile # type: ignore[method-assign] + return c + + +async def _seed(ids: set[uuid.UUID]) -> None: + for account_id in ids: + await leases.mark_live(get_redis(), account_id) + + +async def _until(pred: Callable[[], bool], timeout_s: float = 6.0) -> None: + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout_s + while loop.time() < deadline: + if pred(): + return + await asyncio.sleep(0.05) + assert pred(), "condition not met within timeout" + + +# --- acceptance 1: single ownership, one connection per account --------------- + + +async def test_acceptance_1_single_ownership(connections, make_accounts) -> None: # type: ignore[no-untyped-def] + ids = set(await make_accounts(20)) + await _seed(ids) + c1 = _fast(SyncCoordinator("w1")) + c2 = _fast(SyncCoordinator("w2")) + c1.rebalance_interval = c2.rebalance_interval = 0.3 # let the load even out + # Start concurrently so the two workers race for the leases. + await asyncio.gather(c1.start(), c2.start()) + try: + # Both workers end up holding accounts (via the initial race + rebalance), + # together covering all 20 with no account owned twice. + await _until(lambda: c1.owned | c2.owned == ids and bool(c1.owned and c2.owned)) + assert c1.owned & c2.owned == set() + # At steady state each account has exactly one live IDLE connection. + await _until(lambda: all(connections[a] == 1 for a in ids)) + finally: + await c1.stop() + await c2.stop() + await _until(lambda: all(connections[a] == 0 for a in ids)) + + +# --- acceptance 2: crash failover -------------------------------------------- + + +async def test_acceptance_2_failover(connections, make_accounts) -> None: # type: ignore[no-untyped-def] + ids = set(await make_accounts(10)) + await _seed(ids) + c1 = _fast(SyncCoordinator("w1")) + await c1.start() + try: + await _until(lambda: c1.owned == ids) # sole worker owns everything + + c2 = _fast(SyncCoordinator("w2")) + await c2.start() + + # Simulate `kill -9` on w1: connections drop and renewals stop, but the + # leases are NOT released — they must expire before w2 can take over. + for task in c1._tasks: + task.cancel() + await asyncio.gather(*c1._tasks, return_exceptions=True) + c1._tasks = [] + for r in list(c1._runners.values()): + await r.stop() + c1._runners.clear() + + try: + # w2 re-owns every orphan once the (short) leases expire. + await _until(lambda: c2.owned == ids, timeout_s=8.0) + assert all(connections[a] == 1 for a in ids) # one connection each, no dupes + finally: + await c2.stop() + finally: + # w1 already crashed; clear any residual leases it left behind. + await get_redis().flushdb() + await _until(lambda: all(connections[a] == 0 for a in ids)) + + +# --- acceptance 3: create / delete lifecycle without restart ------------------ + + +async def test_acceptance_3_create_delete(connections, make_accounts) -> None: # type: ignore[no-untyped-def] + c = _fast(SyncCoordinator("w1")) + await c.start() # starts with no live accounts + try: + assert c.owned == set() + + # A newly-live account (as the bootstrap job's _set_state('live') would do): + # register it in sync:live; the running coordinator must claim it, no restart. + (aid,) = await make_accounts(1) + await leases.mark_live(get_redis(), aid) + await _until(lambda: aid in c.owned) + assert connections[aid] == 1 + + # Removing it (as delete_account does) fences the runner: lease gone -> + # next renewal fails -> connection torn down, no further ownership. + await leases.mark_not_live(get_redis(), aid) + await _until(lambda: aid not in c.owned) + await _until(lambda: connections[aid] == 0) + finally: + await c.stop() + + +def test_worker_id_is_unique() -> None: + assert coordinator_mod._new_worker_id() != coordinator_mod._new_worker_id()