sync-v2 step 5: activity tiers (hot/warm/cold)

Each owned account is tiered by the coordinator; the runner shapes its work to
match — correctness identical, only latency/socket cost differ.

- hot: persistent INBOX IDLE + secondary poll every POLL_WARM.
- warm: no persistent socket; transient INBOX+special poll every POLL_WARM.
- cold: same transient poll every POLL_COLD.
- engine: poll_all_loop (transient full poll); secondary_poll_loop interval is
  now a parameter; POLL_INTERVAL removed in favor of POLL_WARM/POLL_COLD.
- coordinator: tier eval every 120s from session heat (sess:active:{user}) +
  recent mail (max internal_date within TIER_HOT_ACTIVITY_H), with the
  HOT_BUDGET_PER_WORKER guard shedding least-recently-active session-less hot
  accounts; user.active event promotes a user's accounts to hot at once (≤5s).
- API: login + WS connect/ping publish user.active and refresh sess:active
  (TTL TIER_SESSION_GRACE_MIN). New sync:tier:{id} key, default warm on claim.
- config: POLL_WARM/POLL_COLD/TIER_*/HOT_BUDGET_PER_WORKER.
- tests: 5 tier tests (session/recent-mail promotion, cold, user.active,
  budget demotion) alongside the lease/lifecycle acceptance tests.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-07-19 07:03:54 +02:00
co-authored by Claude Opus 4.8
parent 245672d3f2
commit 19335116da
9 changed files with 365 additions and 47 deletions
+7 -1
View File
@@ -29,7 +29,6 @@ class Settings(BaseSettings):
meili_key: str = Field(default="", alias="MEILI_KEY") meili_key: str = Field(default="", alias="MEILI_KEY")
undo_seconds: int = Field(default=10, alias="UNDO_SECONDS") undo_seconds: int = Field(default=10, alias="UNDO_SECONDS")
poll_interval: int = Field(default=120, alias="POLL_INTERVAL")
session_ttl_days: int = Field(default=30, alias="SESSION_TTL_DAYS") session_ttl_days: int = Field(default=30, alias="SESSION_TTL_DAYS")
max_attachment_mb: int = Field(default=25, alias="MAX_ATTACHMENT_MB") max_attachment_mb: int = Field(default=25, alias="MAX_ATTACHMENT_MB")
log_level: str = Field(default="INFO", alias="LOG_LEVEL") log_level: str = Field(default="INFO", alias="LOG_LEVEL")
@@ -44,6 +43,13 @@ class Settings(BaseSettings):
max_accounts_per_worker: int = Field(default=400, alias="MAX_ACCOUNTS_PER_WORKER") max_accounts_per_worker: int = Field(default=400, alias="MAX_ACCOUNTS_PER_WORKER")
lease_ttl: int = Field(default=60, alias="LEASE_TTL") lease_ttl: int = Field(default=60, alias="LEASE_TTL")
lease_renew: int = Field(default=20, alias="LEASE_RENEW") lease_renew: int = Field(default=20, alias="LEASE_RENEW")
# Activity tiers (§5): hot = persistent IDLE, warm/cold = transient polls.
poll_warm: int = Field(default=180, alias="POLL_WARM")
poll_cold: int = Field(default=900, alias="POLL_COLD")
tier_hot_activity_h: int = Field(default=24, alias="TIER_HOT_ACTIVITY_H")
tier_session_grace_min: int = Field(default=30, alias="TIER_SESSION_GRACE_MIN")
tier_cold_after_d: int = Field(default=3, alias="TIER_COLD_AFTER_D")
hot_budget_per_worker: int = Field(default=250, alias="HOT_BUDGET_PER_WORKER")
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")
+6 -1
View File
@@ -51,6 +51,8 @@ from app.identity.sessions import (
revoke_session, revoke_session,
) )
from app.models import User from app.models import User
from app.redis_client import get_redis
from app.sync import leases
router = APIRouter(prefix="/auth", tags=["auth"]) router = APIRouter(prefix="/auth", tags=["auth"])
@@ -117,7 +119,10 @@ async def login(body: LoginRequest, request: Request, response: Response, db: Db
if user is None: if user is None:
raise_api_error(401, "invalid_credentials", "incorrect handle or password") raise_api_error(401, "invalid_credentials", "incorrect handle or password")
await _establish_session(request, response, user) await _establish_session(request, response, user)
# Sync v2 §3: a login marks the user active, promoting their accounts' tier. # Sync v2 §3/§5: a login marks the user active (session heat) and promotes
# their accounts' tier to hot.
settings = get_settings()
await leases.mark_user_active(get_redis(), user.id, settings.tier_session_grace_min * 60)
await get_event_bus().publish("user.active", {"user_id": str(user.id)}) await get_event_bus().publish("user.active", {"user_id": str(user.id)})
return _user_out(user) return _user_out(user)
+1 -1
View File
@@ -85,7 +85,7 @@ class MailProvider(Protocol):
async def idle(self, mailbox: str) -> None: async def idle(self, mailbox: str) -> None:
"""Blocks until the server signals change or 25 min elapse (re-issue IDLE). May raise """Blocks until the server signals change or 25 min elapse (re-issue IDLE). May raise
ProviderNotSupported => sync loop falls back to poll() every POLL_INTERVAL (default 120 s). ProviderNotSupported => sync loop falls back to poll() every POLL_WARM (default 180 s).
""" """
... ...
+103 -7
View File
@@ -25,13 +25,14 @@ import secrets
import socket import socket
import time import time
import uuid import uuid
from datetime import UTC, datetime, timedelta
from sqlalchemy import select from sqlalchemy import func, select
from app.config import get_settings from app.config import get_settings
from app.db import session_scope from app.db import session_scope
from app.events import channel_for from app.events import channel_for
from app.models import Account from app.models import Account, Message
from app.redis_client import aw, get_redis from app.redis_client import aw, get_redis
from app.sync import leases from app.sync import leases
from app.sync.runner import AccountRunner from app.sync.runner import AccountRunner
@@ -65,6 +66,11 @@ class SyncCoordinator:
self.heartbeat_ttl = 30 self.heartbeat_ttl = 30
self.heartbeat_refresh = 10.0 self.heartbeat_refresh = 10.0
self.rebalance_interval = 60.0 self.rebalance_interval = 60.0
self.tier_eval_interval = 120.0
self.session_grace_s = s.tier_session_grace_min * 60
self.hot_activity = timedelta(hours=s.tier_hot_activity_h)
self.cold_after = timedelta(days=s.tier_cold_after_d)
self.hot_budget = s.hot_budget_per_worker
# --- lifecycle ----------------------------------------------------------- # --- lifecycle -----------------------------------------------------------
@@ -80,6 +86,7 @@ class SyncCoordinator:
asyncio.create_task(self._loop(self._reconcile_once, self.reconcile_interval)), 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._heartbeat_once, self.heartbeat_refresh)),
asyncio.create_task(self._loop(self._rebalance_once, self.rebalance_interval)), asyncio.create_task(self._loop(self._rebalance_once, self.rebalance_interval)),
asyncio.create_task(self._loop(self._evaluate_tiers_once, self.tier_eval_interval)),
asyncio.create_task(self._event_subscriber()), asyncio.create_task(self._event_subscriber()),
] ]
logger.info("sync coordinator %s started", self.worker_id) logger.info("sync coordinator %s started", self.worker_id)
@@ -131,7 +138,7 @@ class SyncCoordinator:
if account_id in self._runners: if account_id in self._runners:
continue continue
if await leases.claim(self._redis, account_id, self.worker_id, self.lease_ttl): if await leases.claim(self._redis, account_id, self.worker_id, self.lease_ttl):
self._start_runner(account_id) await self._start_runner(account_id)
async def _renew_once(self) -> None: async def _renew_once(self) -> None:
for account_id in list(self._runners): for account_id in list(self._runners):
@@ -160,8 +167,11 @@ class SyncCoordinator:
if runner is not None: if runner is not None:
await runner.stop() await runner.stop()
def _start_runner(self, account_id: uuid.UUID) -> None: async def _start_runner(self, account_id: uuid.UUID) -> None:
runner = AccountRunner(account_id) # Inherit any tier a previous owner recorded; default warm on first claim.
tier = await aw(self._redis.get(leases.tier_key(account_id))) or "warm"
await aw(self._redis.set(leases.tier_key(account_id), tier))
runner = AccountRunner(account_id, tier)
runner.start() runner.start()
self._runners[account_id] = runner self._runners[account_id] = runner
@@ -225,7 +235,9 @@ class SyncCoordinator:
the lease to lapse. the lease to lapse.
""" """
pubsub = self._redis.pubsub() pubsub = self._redis.pubsub()
await pubsub.subscribe(channel_for("account.state_changed")) state_ch = channel_for("account.state_changed")
active_ch = channel_for("user.active")
await pubsub.subscribe(state_ch, active_ch)
try: try:
async for message in pubsub.listen(): async for message in pubsub.listen():
if self._stop.is_set(): if self._stop.is_set():
@@ -236,7 +248,10 @@ class SyncCoordinator:
payload = json.loads(message["data"]).get("payload", {}) payload = json.loads(message["data"]).get("payload", {})
except (ValueError, TypeError): except (ValueError, TypeError):
continue continue
await self._on_state_changed(payload) if message.get("channel") == active_ch:
await self._on_user_active(payload)
else:
await self._on_state_changed(payload)
finally: finally:
with contextlib.suppress(Exception): with contextlib.suppress(Exception):
await pubsub.aclose() # type: ignore[no-untyped-call] await pubsub.aclose() # type: ignore[no-untyped-call]
@@ -250,3 +265,84 @@ class SyncCoordinator:
await leases.mark_live(self._redis, account_id) # claimable next cycle await leases.mark_live(self._redis, account_id) # claimable next cycle
elif account_id in self._runners: elif account_id in self._runners:
await self._release(account_id) # disabled/deleted/error -> stop now await self._release(account_id) # disabled/deleted/error -> stop now
async def _on_user_active(self, payload: dict[str, object]) -> None:
"""Promote all this user's owned accounts to hot at once (≤5s, §5)."""
raw = payload.get("user_id")
if not isinstance(raw, str):
return
user_id = uuid.UUID(raw)
owned = list(self._runners)
if not owned:
return
async with session_scope() as session:
aids = list(
await session.scalars(
select(Account.id).where(Account.id.in_(owned), Account.user_id == user_id)
)
)
for account_id in aids:
await self._set_tier(account_id, "hot")
# --- tier evaluation (§5) ------------------------------------------------
async def _set_tier(self, account_id: uuid.UUID, tier: str) -> None:
runner = self._runners.get(account_id)
if runner is None:
return
await aw(self._redis.set(leases.tier_key(account_id), tier))
await runner.apply_tier(tier)
async def _account_activity(
self, ids: list[uuid.UUID]
) -> dict[uuid.UUID, tuple[uuid.UUID, datetime | None]]:
"""Return {account_id: (user_id, last_mail_at)} for the owned accounts."""
async with session_scope() as session:
rows = (
await session.execute(
select(Account.id, Account.user_id, func.max(Message.internal_date))
.select_from(Account)
.outerjoin(Message, Message.account_id == Account.id)
.where(Account.id.in_(ids))
.group_by(Account.id, Account.user_id)
)
).all()
return {row[0]: (row[1], row[2]) for row in rows}
async def _evaluate_tiers_once(self) -> None:
"""Recompute each owned account's tier from session heat + recent mail."""
owned = list(self._runners)
if not owned:
return
activity = await self._account_activity(owned)
now = datetime.now(UTC)
sessions: dict[uuid.UUID, bool] = {}
desired: dict[uuid.UUID, str] = {}
demotable: list[tuple[datetime, uuid.UUID]] = [] # activity-hot, no session
hot_count = 0
for account_id in owned:
user_id, last_mail = activity.get(account_id, (None, None))
session_active = False
if user_id is not None:
if user_id not in sessions:
sessions[user_id] = bool(
await aw(self._redis.exists(leases.session_key(user_id)))
)
session_active = sessions[user_id]
recent = last_mail is not None and (now - last_mail) < self.hot_activity
if session_active or recent:
desired[account_id] = "hot"
hot_count += 1
if not session_active and last_mail is not None:
demotable.append((last_mail, account_id))
elif last_mail is not None and (now - last_mail) < self.cold_after:
desired[account_id] = "warm"
else:
desired[account_id] = "cold"
# Global hot budget: shed the least-recently-active session-less hot accounts.
if hot_count > self.hot_budget:
demotable.sort(key=lambda item: item[0]) # oldest activity first
for _, account_id in demotable[: hot_count - self.hot_budget]:
desired[account_id] = "warm"
for account_id, tier in desired.items():
await self._set_tier(account_id, tier)
+26 -6
View File
@@ -255,7 +255,7 @@ async def live_loop(
from app.db import session_scope from app.db import session_scope
stop = stop or asyncio.Event() stop = stop or asyncio.Event()
poll_interval = get_settings().poll_interval idle_fallback = get_settings().poll_warm
async with session_scope() as session: async with session_scope() as session:
account = await _load_account(session, account_id) account = await _load_account(session, account_id)
@@ -267,7 +267,7 @@ async def live_loop(
try: try:
await provider.idle(mailbox_path) await provider.idle(mailbox_path)
except ProviderNotSupported: except ProviderNotSupported:
await asyncio.sleep(poll_interval) await asyncio.sleep(idle_fallback)
if stop.is_set(): if stop.is_set():
break break
async with session_scope() as session: async with session_scope() as session:
@@ -341,19 +341,39 @@ async def refresh_account_job(ctx: dict[str, Any], account_id: str) -> None:
await poll_folders_once(uuid.UUID(account_id), include_inbox=True) await poll_folders_once(uuid.UUID(account_id), include_inbox=True)
async def secondary_poll_loop(account_id: uuid.UUID, stop: asyncio.Event | None = None) -> None: async def secondary_poll_loop(
account_id: uuid.UUID, stop: asyncio.Event | None = None, interval: float | None = None
) -> None:
"""Periodically CONDSTORE-poll non-INBOX folders (Sent/Drafts/Archive/Trash/Junk). """Periodically CONDSTORE-poll non-INBOX folders (Sent/Drafts/Archive/Trash/Junk).
INBOX has its own IDLE loop (real-time). This keeps the other folders in step INBOX has its own IDLE loop (real-time). This keeps the other folders in step
with server-side changes (e.g. Gmail-filed sent mail) and changes made by other with server-side changes (e.g. Gmail-filed sent mail) and changes made by other
clients. CONDSTORE/QRESYNC makes each pass a cheap MODSEQ delta. clients. CONDSTORE/QRESYNC makes each pass a cheap MODSEQ delta. Used by hot-tier
runners alongside the INBOX IDLE loop.
""" """
stop = stop or asyncio.Event() stop = stop or asyncio.Event()
interval = get_settings().poll_interval period = interval if interval is not None else get_settings().poll_warm
while not stop.is_set(): while not stop.is_set():
# Wait first: the initial state was already covered by the backfill. # Wait first: the initial state was already covered by the backfill.
with contextlib.suppress(TimeoutError):
await asyncio.wait_for(stop.wait(), timeout=period)
if stop.is_set():
break
await _poll_secondary_once(account_id)
async def poll_all_loop(
account_id: uuid.UUID, interval: float, stop: asyncio.Event | None = None
) -> None:
"""Transient CONDSTORE poll of INBOX + special folders every ``interval`` seconds.
The warm/cold tiers use this instead of a persistent IDLE connection: no socket
is held between passes. Correctness matches the hot tier — only latency differs.
"""
stop = stop or asyncio.Event()
while not stop.is_set():
with contextlib.suppress(TimeoutError): with contextlib.suppress(TimeoutError):
await asyncio.wait_for(stop.wait(), timeout=interval) await asyncio.wait_for(stop.wait(), timeout=interval)
if stop.is_set(): if stop.is_set():
break break
await _poll_secondary_once(account_id) await poll_folders_once(account_id, include_inbox=True)
+16
View File
@@ -20,6 +20,8 @@ from app.redis_client import aw
LIVE_SET = "sync:live" LIVE_SET = "sync:live"
_LEASE_PREFIX = "sync:lease:" _LEASE_PREFIX = "sync:lease:"
WORKERS_PREFIX = "sync:workers:" WORKERS_PREFIX = "sync:workers:"
_TIER_PREFIX = "sync:tier:"
_SESSION_PREFIX = "sess:active:"
def lease_key(account_id: uuid.UUID) -> str: def lease_key(account_id: uuid.UUID) -> str:
@@ -30,6 +32,20 @@ def worker_key(worker_id: str) -> str:
return f"{WORKERS_PREFIX}{worker_id}" return f"{WORKERS_PREFIX}{worker_id}"
def tier_key(account_id: uuid.UUID) -> str:
return f"{_TIER_PREFIX}{account_id}"
def session_key(user_id: uuid.UUID | str) -> str:
"""Presence key the API refreshes while a user has an active session (§5)."""
return f"{_SESSION_PREFIX}{user_id}"
async def mark_user_active(redis: aioredis.Redis, user_id: uuid.UUID | str, grace_s: int) -> None:
"""Refresh the user's session-heat key; workers read it for tier promotion."""
await aw(redis.set(session_key(user_id), "1", ex=grace_s))
# Renew only if we still own the lease (guards against a stale worker extending a # 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. # lease another worker has already claimed after expiry). Returns 1 on success.
_RENEW_LUA = """ _RENEW_LUA = """
+64 -23
View File
@@ -1,9 +1,16 @@
"""Per-account sync runner (refactor §4). """Per-account sync runner (refactor §4, §5).
An ``AccountRunner`` owns the long-lived sync loops for a single account and An ``AccountRunner`` owns the long-lived sync work for a single account and
exposes a small ``start`` / ``stop`` lifecycle. This is a pure extraction of the exposes a ``start`` / ``stop`` / ``apply_tier`` lifecycle. The task set is shaped
loops that previously lived inline in the supervisor — behavior is unchanged by the account's activity tier (§5):
(INBOX IDLE + secondary CONDSTORE poll, each self-healing with backoff).
* **hot** — a persistent INBOX IDLE loop plus a secondary CONDSTORE poll of the
special folders every ``POLL_WARM`` seconds.
* **warm** — no persistent socket; a single transient poll of INBOX + special
folders every ``POLL_WARM`` seconds.
* **cold** — the same single transient poll, every ``POLL_COLD`` seconds.
Correctness is identical across tiers — only latency and socket cost differ.
""" """
from __future__ import annotations from __future__ import annotations
@@ -16,54 +23,88 @@ from collections.abc import Awaitable, Callable
from sqlalchemy import select from sqlalchemy import select
from app.config import get_settings
from app.db import session_scope from app.db import session_scope
from app.models import Account from app.models import Account
from app.sync.engine import live_loop, secondary_poll_loop from app.sync.engine import live_loop, poll_all_loop, secondary_poll_loop
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
_RECONNECT_BACKOFF = 30.0 _RECONNECT_BACKOFF = 30.0
_VERIFY_INTERVAL = 60.0 _VERIFY_INTERVAL = 60.0
Tier = str # "hot" | "warm" | "cold"
class AccountRunner: class AccountRunner:
"""Owns the sync loops for one account until stopped.""" """Owns the tier-shaped sync work for one account until stopped."""
def __init__(self, account_id: uuid.UUID) -> None: def __init__(self, account_id: uuid.UUID, tier: Tier = "warm") -> None:
self.account_id = account_id self.account_id = account_id
self.tier = tier
self._stop = asyncio.Event() self._stop = asyncio.Event()
self._tasks: list[asyncio.Task[None]] = [] self._tasks: list[asyncio.Task[None]] = [] # tier-independent (verify)
self._work: list[asyncio.Task[None]] = [] # tier-shaped loops
# Set when the account is found no longer live (§3 defense in depth); the # Set when the account is found no longer live (§3 defense in depth); the
# coordinator reaps such runners and releases their lease. # coordinator reaps such runners and releases their lease.
self.dead = asyncio.Event() self.dead = asyncio.Event()
self.verify_interval = _VERIFY_INTERVAL self.verify_interval = _VERIFY_INTERVAL
def start(self) -> None: def start(self) -> None:
"""Launch the account's sync loops (idempotent-ish: call once).""" """Launch the verify task and the current tier's work loops."""
self._stop.clear() self._stop.clear()
self._tasks = [ self._tasks = [asyncio.create_task(self._verify_loop())]
asyncio.create_task(self._supervise(self._run_live, "live")), self._start_work()
asyncio.create_task(self._supervise(self._run_secondary, "secondary")),
asyncio.create_task(self._verify_loop()),
]
def running(self) -> bool: def running(self) -> bool:
return any(not t.done() for t in self._tasks) return any(not t.done() for t in (*self._tasks, *self._work))
async def stop(self) -> None: async def stop(self) -> None:
"""Signal and cancel the loops, waiting for teardown.""" """Signal and cancel every loop, waiting for teardown."""
self._stop.set() self._stop.set()
for task in self._tasks: await self._cancel(self._work)
task.cancel() await self._cancel(self._tasks)
if self._tasks: self._work = []
await asyncio.gather(*self._tasks, return_exceptions=True)
self._tasks = [] self._tasks = []
async def apply_tier(self, tier: Tier) -> None:
"""Reshape the work loops for a new tier (no-op if unchanged)."""
if tier == self.tier and self._work:
return
logger.info("account %s tier %s -> %s", self.account_id, self.tier, tier)
self.tier = tier
await self._cancel(self._work)
self._start_work()
# --- internals -----------------------------------------------------------
def _start_work(self) -> None:
settings = get_settings()
if self.tier == "hot":
self._work = [
asyncio.create_task(self._supervise(self._run_live, "live")),
asyncio.create_task(self._supervise(self._run_secondary, "secondary")),
]
else:
interval = float(settings.poll_warm if self.tier == "warm" else settings.poll_cold)
self._work = [
asyncio.create_task(self._supervise(lambda: self._run_poll(interval), "poll"))
]
async def _cancel(self, tasks: list[asyncio.Task[None]]) -> None:
for task in tasks:
task.cancel()
if tasks:
await asyncio.gather(*tasks, return_exceptions=True)
async def _run_live(self) -> None: async def _run_live(self) -> None:
await live_loop(self.account_id, "INBOX", self._stop) await live_loop(self.account_id, "INBOX", self._stop)
async def _run_secondary(self) -> None: async def _run_secondary(self) -> None:
await secondary_poll_loop(self.account_id, self._stop) await secondary_poll_loop(self.account_id, self._stop, get_settings().poll_warm)
async def _run_poll(self, interval: float) -> None:
await poll_all_loop(self.account_id, interval, self._stop)
async def _verify_loop(self) -> None: async def _verify_loop(self) -> None:
"""Periodically confirm the account is still live; flag ``dead`` if not. """Periodically confirm the account is still live; flag ``dead`` if not.
@@ -88,7 +129,7 @@ class AccountRunner:
return state == "live" return state == "live"
async def _supervise(self, run: Callable[[], Awaitable[None]], label: str) -> None: async def _supervise(self, run: Callable[[], Awaitable[None]], label: str) -> None:
"""Run a loop, restarting it (with backoff) until stopped.""" """Run a loop, restarting it (with backoff) until stopped or cancelled."""
while not self._stop.is_set(): while not self._stop.is_set():
try: try:
await run() await run()
+7 -1
View File
@@ -13,10 +13,12 @@ import logging
from fastapi import APIRouter, WebSocket, WebSocketDisconnect from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from app.config import get_settings
from app.events import get_event_bus from app.events import get_event_bus
from app.identity.cookies import SESSION_COOKIE_INSECURE, SESSION_COOKIE_SECURE from app.identity.cookies import SESSION_COOKIE_INSECURE, SESSION_COOKIE_SECURE
from app.identity.sessions import resolve_session from app.identity.sessions import resolve_session
from app.redis_client import get_redis from app.redis_client import get_redis
from app.sync import leases
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -35,7 +37,9 @@ async def ws_endpoint(websocket: WebSocket) -> None:
await websocket.accept() await websocket.accept()
user_id = str(session.user_id) user_id = str(session.user_id)
# Sync v2 §3: an open WS means the user is active -> promote their tier. grace_s = get_settings().tier_session_grace_min * 60
# Sync v2 §3/§5: an open WS means the user is active -> session heat + promote.
await leases.mark_user_active(get_redis(), user_id, grace_s)
await get_event_bus().publish("user.active", {"user_id": user_id}) await get_event_bus().publish("user.active", {"user_id": user_id})
pubsub = get_redis().pubsub() pubsub = get_redis().pubsub()
await pubsub.psubscribe("evt:*") await pubsub.psubscribe("evt:*")
@@ -57,6 +61,8 @@ async def ws_endpoint(websocket: WebSocket) -> None:
raw = await websocket.receive_text() raw = await websocket.receive_text()
try: try:
if json.loads(raw).get("type") == "ping": if json.loads(raw).get("type") == "ping":
# A ping keeps the user's session heat alive (§5).
await leases.mark_user_active(get_redis(), user_id, grace_s)
await websocket.send_json({"type": "pong"}) await websocket.send_json({"type": "pong"})
except (ValueError, TypeError): except (ValueError, TypeError):
continue continue
+135 -7
View File
@@ -19,16 +19,18 @@ import asyncio
import uuid import uuid
from collections import defaultdict from collections import defaultdict
from collections.abc import AsyncIterator, Callable from collections.abc import AsyncIterator, Callable
from datetime import UTC, datetime, timedelta
import pytest_asyncio import pytest_asyncio
from sqlalchemy import delete from sqlalchemy import delete, select
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from sqlalchemy.pool import NullPool from sqlalchemy.pool import NullPool
from app.config import get_settings from app.config import get_settings
from app.db import session_scope
from app.events import get_event_bus from app.events import get_event_bus
from app.identity.passwords import hash_password from app.identity.passwords import hash_password
from app.models import Account, User from app.models import Account, Mailbox, Message, User
from app.redis_client import get_redis from app.redis_client import get_redis
from app.sync import coordinator as coordinator_mod from app.sync import coordinator as coordinator_mod
from app.sync import leases, runner from app.sync import leases, runner
@@ -39,22 +41,31 @@ from app.sync.coordinator import SyncCoordinator
@pytest_asyncio.fixture @pytest_asyncio.fixture
async def connections(monkeypatch) -> AsyncIterator[dict[uuid.UUID, int]]: # type: ignore[no-untyped-def] 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.""" """Patch the runner's loops; yield a live (hot-tier IDLE) connection count."""
conns: dict[uuid.UUID, int] = defaultdict(int) conns: dict[uuid.UUID, int] = defaultdict(int)
async def fake_live_loop(account_id: uuid.UUID, folder: str, stop: asyncio.Event) -> None: 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 conns[account_id] += 1 # one persistent IDLE connection while alive (hot)
try: try:
await stop.wait() await stop.wait()
finally: finally:
conns[account_id] -= 1 conns[account_id] -= 1
async def fake_secondary(account_id: uuid.UUID, stop: asyncio.Event | None = None) -> None: async def fake_secondary(
account_id: uuid.UUID, stop: asyncio.Event | None = None, interval: float | None = None
) -> None:
assert stop is not None assert stop is not None
await stop.wait() 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, "live_loop", fake_live_loop)
monkeypatch.setattr(runner, "secondary_poll_loop", fake_secondary) monkeypatch.setattr(runner, "secondary_poll_loop", fake_secondary)
monkeypatch.setattr(runner, "poll_all_loop", fake_poll_all)
yield conns yield conns
@@ -119,13 +130,44 @@ def _fast(c: SyncCoordinator) -> SyncCoordinator:
c.reconcile_interval = 1000.0 c.reconcile_interval = 1000.0
c.heartbeat_refresh = 1000.0 c.heartbeat_refresh = 1000.0
c.rebalance_interval = 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] c._reconcile_once = _no_reconcile # type: ignore[method-assign]
return c return c
async def _seed(ids: set[uuid.UUID]) -> None: 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: for account_id in ids:
await leases.mark_live(get_redis(), account_id) 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: async def _until(pred: Callable[[], bool], timeout_s: float = 6.0) -> None:
@@ -210,6 +252,7 @@ async def test_acceptance_3_create_delete(connections, make_accounts) -> None:
# A newly-live account (as the bootstrap job's _set_state('live') would do): # 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. # register it in sync:live; the running coordinator must claim it, no restart.
(aid,) = await make_accounts(1) (aid,) = await make_accounts(1)
await get_redis().set(leases.tier_key(aid), "hot") # persistent IDLE
await leases.mark_live(get_redis(), aid) await leases.mark_live(get_redis(), aid)
await _until(lambda: aid in c.owned) await _until(lambda: aid in c.owned)
assert connections[aid] == 1 assert connections[aid] == 1
@@ -254,5 +297,90 @@ async def test_event_driven_teardown(connections, make_accounts) -> None: # typ
await c.stop() await c.stop()
# --- §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: def test_worker_id_is_unique() -> None:
assert coordinator_mod._new_worker_id() != coordinator_mod._new_worker_id() assert coordinator_mod._new_worker_id() != coordinator_mod._new_worker_id()