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:
@@ -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:
|
||||||
|
|||||||
@@ -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")
|
||||||
|
|||||||
@@ -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)
|
||||||
@@ -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:
|
||||||
|
|||||||
@@ -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)))
|
||||||
@@ -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"] = {}
|
|
||||||
|
|||||||
@@ -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()
|
||||||
Reference in New Issue
Block a user