sync-v2 step 3: lease-owned account claimer (replaces startup supervisor)

Ownership of an account now moves through a Redis lease so any number of sync
workers can run against the same account set with no double-IDLE.

- leases.py: sync:live registry + sync:lease:{id} (SET NX EX) with owner-checked
  Lua renew/release; mark_live / mark_not_live helpers.
- coordinator.py: SyncCoordinator claims from sync:live up to MAX_ACCOUNTS_PER_
  WORKER, renews every LEASE_RENEW (fences the runner on lease loss), reconciles
  sync:live against Postgres, heartbeats a worker key, and rebalances when it
  holds well above a fair share.
- supervisor start/stop now drive a coordinator; engine._set_state and
  delete_account maintain sync:live membership (+ lease drop = fencing teardown).
- config: MAX_ACCOUNTS_PER_WORKER, LEASE_TTL, LEASE_RENEW.
- tests: acceptance 1-3 (single ownership, crash failover, create/delete
  lifecycle) driving the real coordinator + Redis with faked IMAP loops.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-07-18 20:34:22 +02:00
co-authored by Claude Opus 4.8
parent 23704f323c
commit bcf11fcecb
7 changed files with 549 additions and 30 deletions
+7
View File
@@ -80,6 +80,13 @@ async def delete_account(
await session.execute(delete(Account).where(Account.id == account_id)) await session.execute(delete(Account).where(Account.id == account_id))
await session.flush() 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 # Content-addressed blobs are shared across the user's accounts (dedup): only
# remove a blob when no remaining message/attachment references it. # remove a blob when no remaining message/attachment references it.
for sha in candidate_shas: for sha in candidate_shas:
+4
View File
@@ -40,6 +40,10 @@ class Settings(BaseSettings):
worker_role: str = Field(default="all", alias="WORKER_ROLE") worker_role: str = Field(default="all", alias="WORKER_ROLE")
interactive_concurrency: int = Field(default=20, alias="INTERACTIVE_CONCURRENCY") interactive_concurrency: int = Field(default=20, alias="INTERACTIVE_CONCURRENCY")
bulk_concurrency: int = Field(default=2, alias="BULK_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_id: str = Field(default="", alias="OAUTH_GOOGLE_CLIENT_ID")
oauth_google_client_secret: str = Field(default="", alias="OAUTH_GOOGLE_CLIENT_SECRET") oauth_google_client_secret: str = Field(default="", alias="OAUTH_GOOGLE_CLIENT_SECRET")
+205
View File
@@ -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)
+11
View File
@@ -41,6 +41,17 @@ async def _set_state(account_id: uuid.UUID, state: str, detail: str = "") -> Non
await session.execute( await session.execute(
update(Account).where(Account.id == account_id).values(state=state, state_detail=detail) 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() bus = get_event_bus()
account = None account = None
async with session_scope() as session: async with session_scope() as session:
+83
View File
@@ -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)))
+13 -30
View File
@@ -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` The sync worker's ``on_startup`` starts a :class:`SyncCoordinator`, which claims
per live account. This step (v2 §9.2) only extracts the per-account loops into a accounts from the ``sync:live`` registry via Redis leases and runs one
runner; account discovery is still startup-only and is replaced by the lease :class:`~app.sync.runner.AccountRunner` per owned account. ``on_shutdown`` stops
claimer in a later step. it, releasing every lease so peers can take the accounts over immediately.
""" """
from __future__ import annotations from __future__ import annotations
import asyncio
import logging import logging
import uuid
from typing import Any from typing import Any
from sqlalchemy import select from app.sync.coordinator import SyncCoordinator
from app.db import session_scope
from app.models import Account
from app.sync.runner import AccountRunner
logger = logging.getLogger(__name__) 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: async def start_supervisor(ctx: dict[str, Any]) -> None:
"""Launch an AccountRunner per live account (called from the worker on_startup).""" coordinator = SyncCoordinator()
runners: dict[uuid.UUID, AccountRunner] = {} await coordinator.start()
for account_id in await _live_account_ids(): ctx["coordinator"] = coordinator
runner = AccountRunner(account_id)
runner.start()
runners[account_id] = runner
ctx["runners"] = runners
logger.info("live supervisor started for %d account(s)", len(runners))
async def stop_supervisor(ctx: dict[str, Any]) -> None: async def stop_supervisor(ctx: dict[str, Any]) -> None:
"""Stop and tear down all account runners (called from the worker on_shutdown).""" coordinator: SyncCoordinator | None = ctx.get("coordinator")
runners: dict[uuid.UUID, AccountRunner] = ctx.get("runners", {}) if coordinator is not None:
if runners: await coordinator.stop()
await asyncio.gather(*(r.stop() for r in runners.values()), return_exceptions=True) ctx["coordinator"] = None
ctx["runners"] = {}
+226
View File
@@ -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()