webmail-client/backend/loadtest/soak_sync.py

"""Sync-engine soak / load test — acceptance criterion §10.7 (mock provider).

Simulates ``SOAK_ACCOUNTS`` accounts (default 1000), a fraction with active
sessions, sharded across ``SOAK_WORKERS`` in-process ``SyncCoordinator`` instances
using a mock ``MailProvider`` (no real IMAP). It samples, over the run:

* persistent sockets vs hot-account count   (target: persistent ≤ hot, total ≤ hot×1.05)
* process RSS                                (target: stable — no per-account leak)
* oldest sync lag across all accounts        (target: < POLL_COLD)

**Fully isolated**: it creates its own Postgres database (``webmail_soak``) and uses
a separate Redis DB, so it never touches the dev database, the dev Redis, or a
running worker. Both are torn down at the end.

Run (defaults = the §10.7 spec, ~1h):

    cd backend && python -m loadtest.soak_sync

Fast smoke run (a couple of minutes, shortened cadences):

    cd backend && SOAK_ACCOUNTS=200 SOAK_DURATION_S=150 POLL_COLD=30 POLL_WARM=15 \
        SOAK_TIER_EVAL_S=5 SOAK_IDLE_S=10 python -m loadtest.soak_sync

Exit code 0 = all targets met, 1 = a target missed (usable in CI-nightly).
"""

from __future__ import annotations

import os
import re
import subprocess
import sys

# --- isolation: point DB + Redis at throwaway locations BEFORE importing app ---

_DEV_DB = os.environ.get(
    "DATABASE_URL", "postgresql+asyncpg://webmail:webmail_dev@127.0.0.1:5433/webmail_dev"
)
_SOAK_DB_NAME = os.environ.get("SOAK_DB_NAME", "webmail_soak")
_SOAK_DB_URL = re.sub(r"/[^/?]+(\?|$)", f"/{_SOAK_DB_NAME}\\1", _DEV_DB)
os.environ["DATABASE_URL"] = _SOAK_DB_URL
os.environ.setdefault("SOAK_REDIS_URL", "redis://127.0.0.1:6380/14")
os.environ["REDIS_URL"] = os.environ["SOAK_REDIS_URL"]

import asyncio  # noqa: E402
import logging  # noqa: E402
import time  # noqa: E402
import uuid  # noqa: E402
from dataclasses import dataclass, field  # noqa: E402

import asyncpg  # noqa: E402
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine  # noqa: E402
from sqlalchemy.pool import NullPool  # noqa: E402

from app.config import get_settings  # noqa: E402
from app.identity.passwords import hash_password  # noqa: E402
from app.models import Account, Mailbox, User  # noqa: E402
from app.redis_client import close_redis, get_redis  # noqa: E402
from app.sync import leases  # noqa: E402
from app.sync.coordinator import SyncCoordinator  # noqa: E402
from loadtest.mock_provider import MockProvider, MockStats  # noqa: E402

# --- config ------------------------------------------------------------------


def _int(name: str, default: int) -> int:
    return int(os.environ.get(name, default))


@dataclass
class Config:
    accounts: int = field(default_factory=lambda: _int("SOAK_ACCOUNTS", 1000))
    session_pct: int = field(default_factory=lambda: _int("SOAK_SESSION_PCT", 20))
    duration_s: int = field(default_factory=lambda: _int("SOAK_DURATION_S", 3600))
    workers: int = field(default_factory=lambda: _int("SOAK_WORKERS", 3))
    sample_s: int = field(default_factory=lambda: _int("SOAK_SAMPLE_S", 30))
    idle_s: int = field(default_factory=lambda: _int("SOAK_IDLE_S", 60))
    claim_s: int = field(default_factory=lambda: _int("SOAK_CLAIM_S", 5))
    tier_eval_s: int = field(default_factory=lambda: _int("SOAK_TIER_EVAL_S", 120))

    def __post_init__(self) -> None:
        s = get_settings()
        self.poll_cold = s.poll_cold
        self.poll_warm = s.poll_warm
        self.session_grace_s = s.tier_session_grace_min * 60
        # Warm-up must outlast one cold poll so cold accounts have synced once, but
        # never eat the whole run on a short smoke test.
        default_warmup = min(self.poll_cold + 90, int(self.duration_s * 0.4))
        self.warmup_s = _int("SOAK_WARMUP_S", default_warmup)


# --- database lifecycle ------------------------------------------------------


def _admin_dsn() -> str:
    # asyncpg DSN for the maintenance DB (strip the +asyncpg driver + swap db name).
    dsn = _DEV_DB.replace("postgresql+asyncpg://", "postgresql://")
    return re.sub(r"/[^/?]+(\?|$)", "/postgres\\1", dsn)


async def _create_database() -> None:
    conn = await asyncpg.connect(_admin_dsn())
    try:
        await conn.execute(f'DROP DATABASE IF EXISTS "{_SOAK_DB_NAME}" WITH (FORCE)')
        await conn.execute(f'CREATE DATABASE "{_SOAK_DB_NAME}"')
    finally:
        await conn.close()


async def _drop_database() -> None:
    conn = await asyncpg.connect(_admin_dsn())
    try:
        await conn.execute(f'DROP DATABASE IF EXISTS "{_SOAK_DB_NAME}" WITH (FORCE)')
    finally:
        await conn.close()


def _migrate() -> None:
    env = {**os.environ, "DATABASE_URL": _SOAK_DB_URL}
    backend_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    subprocess.run(
        [sys.executable, "-c", "from alembic.config import main; main(argv=['upgrade', 'head'])"],
        cwd=backend_dir,
        env=env,
        check=True,
    )


_MAILBOXES = [
    ("INBOX", "inbox"),
    ("Sent", "sent"),
    ("Drafts", "drafts"),
    ("Archive", "archive"),
    ("Trash", "trash"),
]


async def _seed(cfg: Config) -> list[uuid.UUID]:
    """Create N users, each with one live account and a handful of mailboxes."""
    engine = create_async_engine(_SOAK_DB_URL, poolclass=NullPool)
    sm = async_sessionmaker(engine, expire_on_commit=False)
    user_ids: list[uuid.UUID] = []
    pw = hash_password("soak")
    async with sm() as session:
        for i in range(cfg.accounts):
            user = User(handle=f"soak-{i:05d}@load.local", password_hash=pw)
            session.add(user)
            await session.flush()
            user_ids.append(user.id)
            account = Account(
                user_id=user.id,
                email=f"soak-{i:05d}@load.local",
                display_name="Soak",
                imap_config={"host": "mock", "port": 993, "tls": "ssl"},
                smtp_config={"host": "mock", "port": 465, "tls": "ssl"},
                secret_ct=b"",
                discovery_src="manual",
                state="live",
            )
            session.add(account)
            await session.flush()
            for path, special in _MAILBOXES:
                session.add(
                    Mailbox(account_id=account.id, user_id=user.id, path=path, special_use=special)
                )
            if i % 200 == 0:
                await session.flush()
        await session.commit()
    await engine.dispose()
    return user_ids


# --- metrics -----------------------------------------------------------------


def _rss_mb() -> float:
    with open("/proc/self/status") as fh:
        for line in fh:
            if line.startswith("VmRSS:"):
                return int(line.split()[1]) / 1024.0
    return 0.0


@dataclass
class Sample:
    t: float
    owned: int
    hot: int
    warm: int
    cold: int
    persistent: int
    total_open: int
    max_lag: float
    rss_mb: float
    synced: int  # accounts that have completed >=1 poll (steady-state gate)


async def _tier_counts(coords: list[SyncCoordinator]) -> tuple[int, int, int, int]:
    owned = hot = warm = cold = 0
    for c in coords:
        for runner in list(c._runners.values()):
            owned += 1
            if runner.tier == "hot":
                hot += 1
            elif runner.tier == "warm":
                warm += 1
            else:
                cold += 1
    return owned, hot, warm, cold


def _max_lag(now: float, seeded: int, start: float) -> float:
    """Oldest sync lag across accounts. Never-synced accounts count from run start."""
    worst = 0.0
    synced = MockStats.last_sync
    for ts in synced.values():
        worst = max(worst, now - ts)
    if len(synced) < seeded:  # some accounts have not completed a first poll yet
        worst = max(worst, now - start)
    return worst


# --- run ---------------------------------------------------------------------


def _percentile(values: list[float], pct: float) -> float:
    if not values:
        return 0.0
    ordered = sorted(values)
    idx = min(len(ordered) - 1, int(pct / 100.0 * len(ordered)))
    return ordered[idx]


async def _soak(cfg: Config) -> int:
    import app.sync.engine as engine_mod

    # Route every provider build to the mock (no real IMAP for the whole run).
    engine_mod.build_provider = lambda account: MockProvider(account.id, cfg.idle_s)  # type: ignore[assignment]
    MockStats.reset()

    redis = get_redis()
    await redis.flushdb()

    user_ids = await _seed(cfg)
    active_users = user_ids[: cfg.accounts * cfg.session_pct // 100]

    async def refresh_sessions() -> None:
        for uid in active_users:
            await leases.mark_user_active(redis, uid, cfg.session_grace_s)

    await refresh_sessions()

    coords = [SyncCoordinator(f"soak-w{i}") for i in range(cfg.workers)]
    for c in coords:
        c.claim_interval = cfg.claim_s
        c.claim_jitter = 1.0
        c.tier_eval_interval = cfg.tier_eval_s
        c.heartbeat_refresh = 5.0
        c.rebalance_interval = 20.0
        c.reconcile_interval = 60.0
    await asyncio.gather(*(c.start() for c in coords))

    start = time.time()
    deadline = start + cfg.duration_s
    next_session_refresh = start + cfg.session_grace_s / 2
    samples: list[Sample] = []
    print(
        f"soak: {cfg.accounts} accounts, {len(active_users)} active sessions, "
        f"{cfg.workers} workers, duration {cfg.duration_s}s, warmup {cfg.warmup_s}s, "
        f"POLL_COLD={cfg.poll_cold}s POLL_WARM={cfg.poll_warm}s",
        flush=True,
    )
    print(
        f"{'t':>6} {'owned':>6} {'hot':>5} {'warm':>5} {'cold':>5} "
        f"{'persist':>8} {'open':>5} {'lag_s':>7} {'rss_mb':>8}",
        flush=True,
    )
    try:
        while time.time() < deadline:
            await asyncio.sleep(cfg.sample_s)
            now = time.time()
            if now >= next_session_refresh:
                await refresh_sessions()
                next_session_refresh = now + cfg.session_grace_s / 2
            owned, hot, warm, cold = await _tier_counts(coords)
            s = Sample(
                t=now - start,
                owned=owned,
                hot=hot,
                warm=warm,
                cold=cold,
                persistent=MockStats.persistent_open,
                total_open=MockStats.open,
                max_lag=_max_lag(now, cfg.accounts, start),
                rss_mb=_rss_mb(),
                synced=len(MockStats.last_sync),
            )
            samples.append(s)
            print(
                f"{s.t:6.0f} {s.owned:6d} {s.hot:5d} {s.warm:5d} {s.cold:5d} "
                f"{s.persistent:8d} {s.total_open:5d} {s.max_lag:7.0f} {s.rss_mb:8.1f}",
                flush=True,
            )
    finally:
        await asyncio.gather(*(c.stop() for c in coords), return_exceptions=True)

    return _report(cfg, samples)


def _report(cfg: Config, samples: list[Sample]) -> int:
    # Steady state begins once every account has completed its first sync (and at
    # least the configured warm-up has elapsed) — before that, lag legitimately
    # reflects accounts still ramping up, not a starved steady state.
    all_synced_t = next((s.t for s in samples if s.synced >= cfg.accounts), None)
    if all_synced_t is None:
        worst = max((s.synced for s in samples), default=0)
        print(
            f"\nFAIL: not all accounts synced (best {worst}/{cfg.accounts}); "
            "increase SOAK_DURATION_S",
            flush=True,
        )
        return 1
    floor = max(cfg.warmup_s, all_synced_t)
    steady = [s for s in samples if s.t >= floor]
    if not steady:
        print("\nFAIL: no steady-state samples (increase SOAK_DURATION_S)", flush=True)
        return 1

    owned = max(s.owned for s in steady)
    hot = max(s.hot for s in steady)
    persistent_p95 = _percentile([s.persistent for s in steady], 95)
    open_p95 = _percentile([float(s.total_open) for s in steady], 95)
    lag_max = max(s.max_lag for s in steady)
    rss_base = steady[0].rss_mb
    rss_max = max(s.rss_mb for s in steady)
    rss_growth = (rss_max - rss_base) / rss_base * 100 if rss_base else 0.0

    socket_budget = hot * 1.05  # §10.7: persistent sockets <= hot count + 5%
    lag_budget = cfg.poll_cold * 1.15
    rss_budget_pct = 15.0

    # §10.7 hard criteria.
    checks = [
        (
            "all accounts owned",
            owned >= cfg.accounts,
            f"owned={owned}/{cfg.accounts}",
        ),
        (
            "persistent sockets <= hot+5%",
            persistent_p95 <= socket_budget,
            f"persist_p95={persistent_p95:.0f} budget={socket_budget:.0f} (hot={hot})",
        ),
        (
            "oldest lag < POLL_COLD*1.15",
            lag_max < lag_budget,
            f"lag_max={lag_max:.0f}s budget={lag_budget:.0f}s",
        ),
        (
            "RSS stable (<15% growth)",
            rss_growth <= rss_budget_pct,
            f"rss {rss_base:.0f}->{rss_max:.0f}MB (+{rss_growth:.1f}%)",
        ),
    ]

    print("\n=== soak result (steady-state) ===", flush=True)
    ok = True
    for name, passed, detail in checks:
        ok = ok and passed
        print(f"  [{'PASS' if passed else 'FAIL'}] {name:30} {detail}", flush=True)
    # Observational (transient poll connections, not a §10.7 gate).
    print(
        f"  [info] total open incl. transient   open_p95={open_p95:.0f}",
        flush=True,
    )
    print(f"\nsteady samples: {len(steady)}  hot~{hot}  owned~{owned}", flush=True)
    print("RESULT:", "PASS" if ok else "FAIL", flush=True)
    return 0 if ok else 1


async def _main() -> int:
    if not os.environ.get("SOAK_DEBUG"):
        logging.disable(logging.WARNING)  # quiet the per-account loops; keep ERROR+
    else:
        logging.basicConfig(level=logging.INFO)
    print(f"provisioning isolated soak DB '{_SOAK_DB_NAME}' + Redis {os.environ['REDIS_URL']} ...")
    await _create_database()
    _migrate()
    try:
        rc = await _soak(Config())
    finally:
        await get_redis().flushdb()
        await close_redis()
        from app.db import dispose_engine

        await dispose_engine()
        await _drop_database()
        print(f"torn down soak DB '{_SOAK_DB_NAME}'", flush=True)
    return rc


if __name__ == "__main__":
    raise SystemExit(asyncio.run(_main()))