- Heartbeat JSON now carries {owned, sockets_open, tiers:{hot,warm,cold}, ts} per
§2 so peers and operators can see each worker's load and socket footprint.
- run.py runs the sync coordinator directly (not via a placeholder arq worker) so
SIGINT/SIGTERM can stop the claimer, cancel runners, and release every lease
explicitly within a 10s drain budget — peers adopt the accounts immediately
instead of waiting out the lease TTL. Removed the sync:control-queue hack and
the supervisor indirection (start/stop_supervisor, sync_on_startup/shutdown).
- Acceptance: graceful-shutdown test proves a peer adopts released accounts within
seconds under a 30s TTL (i.e. via release, not expiry).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
416 lines
16 KiB
Python
416 lines
16 KiB
Python
"""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
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
import pytest_asyncio
|
|
from sqlalchemy import delete, select
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
from sqlalchemy.pool import NullPool
|
|
|
|
from app.config import get_settings
|
|
from app.db import session_scope
|
|
from app.events import get_event_bus
|
|
from app.identity.passwords import hash_password
|
|
from app.models import Account, Mailbox, Message, 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 (hot-tier IDLE) connection count."""
|
|
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 persistent IDLE connection while alive (hot)
|
|
try:
|
|
await stop.wait()
|
|
finally:
|
|
conns[account_id] -= 1
|
|
|
|
async def fake_secondary(
|
|
account_id: uuid.UUID, stop: asyncio.Event | None = None, interval: float | None = None
|
|
) -> None:
|
|
assert stop is not None
|
|
await stop.wait()
|
|
|
|
async def fake_poll_all(
|
|
account_id: uuid.UUID, interval: float, stop: asyncio.Event | None = None
|
|
) -> None:
|
|
assert stop is not None # warm/cold: no persistent connection, just wait
|
|
await stop.wait()
|
|
|
|
monkeypatch.setattr(runner, "live_loop", fake_live_loop)
|
|
monkeypatch.setattr(runner, "secondary_poll_loop", fake_secondary)
|
|
monkeypatch.setattr(runner, "poll_all_loop", fake_poll_all)
|
|
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.tier_eval_interval = 1000.0 # lease tests pin the tier; no eval churn
|
|
c._reconcile_once = _no_reconcile # type: ignore[method-assign]
|
|
return c
|
|
|
|
|
|
async def _seed(ids: set[uuid.UUID], tier: str = "hot") -> None:
|
|
"""Register accounts as live and pin their tier (hot => persistent IDLE)."""
|
|
redis = get_redis()
|
|
for account_id in ids:
|
|
await leases.mark_live(redis, account_id)
|
|
await redis.set(leases.tier_key(account_id), tier)
|
|
|
|
|
|
async def _user_of(account_id: uuid.UUID) -> uuid.UUID:
|
|
async with session_scope() as session:
|
|
user_id = await session.scalar(select(Account.user_id).where(Account.id == account_id))
|
|
assert user_id is not None
|
|
return user_id
|
|
|
|
|
|
async def _add_recent_mail(account_id: uuid.UUID, when: datetime) -> None:
|
|
"""Insert one INBOX message dated ``when`` so the account reads as active."""
|
|
user_id = await _user_of(account_id)
|
|
async with session_scope() as session:
|
|
mailbox = Mailbox(account_id=account_id, user_id=user_id, path="INBOX", special_use="inbox")
|
|
session.add(mailbox)
|
|
await session.flush()
|
|
session.add(
|
|
Message(
|
|
user_id=user_id,
|
|
account_id=account_id,
|
|
mailbox_id=mailbox.id,
|
|
uid=1,
|
|
internal_date=when,
|
|
blob_sha256="0" * 64,
|
|
)
|
|
)
|
|
await session.commit()
|
|
|
|
|
|
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 get_redis().set(leases.tier_key(aid), "hot") # persistent IDLE
|
|
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()
|
|
|
|
|
|
async def test_event_driven_teardown(connections, make_accounts) -> None: # type: ignore[no-untyped-def]
|
|
"""§3: an account.state_changed event tears the runner down at once.
|
|
|
|
Leases are kept long (no fencing) and the verify loop slow, so the only path
|
|
that can stop the runner in time is the cross-process event subscriber.
|
|
"""
|
|
(aid,) = await make_accounts(1)
|
|
user_id = uuid.uuid4()
|
|
await _seed({aid})
|
|
c = _fast(SyncCoordinator("w1"))
|
|
c.lease_ttl = 60 # lease stays valid: fencing cannot explain the teardown
|
|
c.lease_renew = 30
|
|
await c.start()
|
|
try:
|
|
await _until(lambda: aid in c.owned)
|
|
assert connections[aid] == 1
|
|
|
|
# As delete_account does: drop it from sync:live (so it isn't re-claimed),
|
|
# then emit the event. The lease is left intact, so only the event — not
|
|
# fencing — can explain the teardown.
|
|
await get_redis().srem(leases.LIVE_SET, str(aid))
|
|
await get_event_bus().publish(
|
|
"account.state_changed",
|
|
{"user_id": str(user_id), "account_id": str(aid), "state": "deleted"},
|
|
)
|
|
await _until(lambda: aid not in c.owned, timeout_s=3.0)
|
|
await _until(lambda: connections[aid] == 0, timeout_s=3.0)
|
|
finally:
|
|
await c.stop()
|
|
|
|
|
|
# --- acceptance 6: graceful shutdown -----------------------------------------
|
|
|
|
|
|
async def test_acceptance_6_graceful_shutdown(connections, make_accounts) -> None: # type: ignore[no-untyped-def]
|
|
"""A graceful stop releases every lease so a peer adopts the accounts fast.
|
|
|
|
Leases are long-lived, so adoption within a couple of seconds proves the
|
|
accounts were handed over by explicit release — not by waiting out the TTL.
|
|
"""
|
|
ids = set(await make_accounts(6))
|
|
await _seed(ids) # hot => persistent IDLE
|
|
c1 = _fast(SyncCoordinator("w1"))
|
|
c2 = _fast(SyncCoordinator("w2"))
|
|
c1.lease_ttl = c2.lease_ttl = 30 # adoption must come from release, not expiry
|
|
await c1.start()
|
|
await _until(lambda: c1.owned == ids) # sole owner
|
|
await c2.start() # nothing to claim yet — all leased by w1
|
|
try:
|
|
await c1.stop() # graceful: stops claimer, cancels runners, releases leases
|
|
assert c1.owned == set()
|
|
assert await get_redis().exists(leases.worker_key("w1")) == 0 # heartbeat gone
|
|
# w2 adopts every account well within the 30s TTL, one connection each.
|
|
await _until(lambda: c2.owned == ids, timeout_s=4.0)
|
|
await _until(lambda: all(connections[a] == 1 for a in ids))
|
|
finally:
|
|
await c2.stop()
|
|
await _until(lambda: all(connections[a] == 0 for a in ids))
|
|
|
|
|
|
# --- §5 activity tiers -------------------------------------------------------
|
|
|
|
|
|
async def test_tier_session_promotes_to_hot(connections, make_accounts) -> None: # type: ignore[no-untyped-def]
|
|
(aid,) = await make_accounts(1)
|
|
user_id = await _user_of(aid)
|
|
await leases.mark_live(get_redis(), aid) # claimed as warm by default
|
|
await leases.mark_user_active(get_redis(), user_id, 300) # active session -> hot
|
|
c = _fast(SyncCoordinator("w1"))
|
|
c.tier_eval_interval = 0.1
|
|
await c.start()
|
|
try:
|
|
await _until(lambda: aid in c.owned)
|
|
await _until(lambda: connections[aid] == 1) # hot => persistent IDLE
|
|
assert c._runners[aid].tier == "hot"
|
|
finally:
|
|
await c.stop()
|
|
|
|
|
|
async def test_tier_no_session_no_mail_goes_cold(connections, make_accounts) -> None: # type: ignore[no-untyped-def]
|
|
(aid,) = await make_accounts(1)
|
|
await leases.mark_live(get_redis(), aid) # no session, no mail
|
|
c = _fast(SyncCoordinator("w1"))
|
|
c.tier_eval_interval = 0.1
|
|
await c.start()
|
|
try:
|
|
await _until(lambda: aid in c.owned)
|
|
await _until(lambda: c._runners[aid].tier == "cold")
|
|
assert connections[aid] == 0 # cold holds no persistent connection
|
|
finally:
|
|
await c.stop()
|
|
|
|
|
|
async def test_tier_recent_mail_promotes_to_hot(connections, make_accounts) -> None: # type: ignore[no-untyped-def]
|
|
(aid,) = await make_accounts(1)
|
|
await _add_recent_mail(aid, datetime.now(UTC)) # activity within TIER_HOT_ACTIVITY_H
|
|
await leases.mark_live(get_redis(), aid)
|
|
c = _fast(SyncCoordinator("w1"))
|
|
c.tier_eval_interval = 0.1
|
|
await c.start()
|
|
try:
|
|
await _until(lambda: aid in c.owned)
|
|
await _until(lambda: connections[aid] == 1) # recent mail => hot
|
|
assert c._runners[aid].tier == "hot"
|
|
finally:
|
|
await c.stop()
|
|
|
|
|
|
async def test_user_active_event_promotes_immediately(connections, make_accounts) -> None: # type: ignore[no-untyped-def]
|
|
aid0, aid1 = await make_accounts(2) # both under one user
|
|
user_id = await _user_of(aid0)
|
|
await leases.mark_live(get_redis(), aid0)
|
|
await leases.mark_live(get_redis(), aid1)
|
|
c = _fast(SyncCoordinator("w1")) # tier eval disabled: only the event can promote
|
|
await c.start()
|
|
try:
|
|
await _until(lambda: {aid0, aid1} <= c.owned)
|
|
assert connections[aid0] == 0 and connections[aid1] == 0 # warm to start
|
|
await get_event_bus().publish("user.active", {"user_id": str(user_id)})
|
|
await _until(lambda: connections[aid0] == 1 and connections[aid1] == 1, timeout_s=3.0)
|
|
finally:
|
|
await c.stop()
|
|
|
|
|
|
async def test_hot_budget_demotes_least_active(connections, make_accounts) -> None: # type: ignore[no-untyped-def]
|
|
aid_old, aid_new = await make_accounts(2) # activity-hot, same user, NO session
|
|
now = datetime.now(UTC)
|
|
await _add_recent_mail(aid_old, now - timedelta(hours=1)) # least recently active
|
|
await _add_recent_mail(aid_new, now)
|
|
await leases.mark_live(get_redis(), aid_old)
|
|
await leases.mark_live(get_redis(), aid_new)
|
|
c = _fast(SyncCoordinator("w1"))
|
|
c.tier_eval_interval = 0.1
|
|
c.hot_budget = 1 # only one may be hot
|
|
await c.start()
|
|
try:
|
|
await _until(lambda: {aid_old, aid_new} <= c.owned)
|
|
# The newer account stays hot; the older is shed to warm by the budget.
|
|
await _until(lambda: c._runners[aid_new].tier == "hot")
|
|
await _until(lambda: c._runners[aid_old].tier == "warm")
|
|
assert connections[aid_new] == 1 and connections[aid_old] == 0
|
|
finally:
|
|
await c.stop()
|
|
|
|
|
|
def test_worker_id_is_unique() -> None:
|
|
assert coordinator_mod._new_worker_id() != coordinator_mod._new_worker_id()
|