MathApp/tools/loadtest/loadtest.py

"""
Load test: N students ask the assistant at the same moment (stream=true, exactly like the
frontend), while B other students keep browsing the app (home, quiz, summary, quota).

usage: python loadtest.py BASE_URL users.json N_ASKERS N_BROWSERS ROUND OUT.json [--stagger S]
  users.json: output of mkusers.py; the first N_ASKERS users ask, the next N_BROWSERS browse.
  ROUND: which of each user's prepared questions to ask (0..2) — distinct text, so no answer-cache hits.
Timing per ask: t_headers (HTTP status line), t_first (first `chunk` event), t_done (`done` event).
"""
import asyncio
import json
import sys
import time

import httpx

BASE, USERS, N, B, ROUND, OUT = sys.argv[1], sys.argv[2], int(sys.argv[3]), int(sys.argv[4]), int(sys.argv[5]), sys.argv[6]
WRITES = '--writes' in sys.argv
STAGGER = float(sys.argv[sys.argv.index('--stagger') + 1]) if '--stagger' in sys.argv else 0.0
users = json.load(open(USERS))
askers, browsers = users[:N], users[N:N + B]
T0 = None
stop = asyncio.Event()


async def ask(client, u, idx):
    await asyncio.sleep(idx * STAGGER)
    a = u['asks'][ROUND]
    rec = {'userId': u['userId'], 'grade': u['grade'], 'start': time.monotonic() - T0}
    t = time.monotonic()
    try:
        async with client.stream('POST', f'{BASE}/api/ai/ask', json={
            'userQuestion': a['question'], 'lessonId': str(a['lessonId']), 'stream': True,
        }, headers={'Authorization': f"Bearer {u['token']}"}, timeout=httpx.Timeout(180, connect=30)) as r:
            rec['status'] = r.status_code
            rec['t_headers'] = time.monotonic() - t
            if r.status_code != 200:
                rec['body'] = (await r.aread()).decode(errors='replace')[:300]
                return rec
            buf, nchunks = '', 0
            async for text in r.aiter_text():
                buf += text
                while '\n\n' in buf:
                    part, buf = buf.split('\n\n', 1)
                    ev = next((l[7:] for l in part.split('\n') if l.startswith('event: ')), '')
                    if ev == 'chunk':
                        nchunks += 1
                        rec.setdefault('t_first', time.monotonic() - t)
                    elif ev == 'done':
                        rec['t_done'] = time.monotonic() - t
                        d = json.loads(next(l[6:] for l in part.split('\n') if l.startswith('data: ')))
                        rec['chunksUsed'] = len(d.get('usedChunks') or [])
                        rec['messageId'] = d.get('messageId')
                    elif ev == 'error':
                        rec['error'] = part[-200:]
            rec['nchunks'] = nchunks
            rec['t_end'] = time.monotonic() - t
    except Exception as exc:  # recorded, never hidden: a timeout is a result here
        rec['exc'] = f'{type(exc).__name__}: {exc}'[:200]
        rec['t_end'] = time.monotonic() - t
    return rec


async def browse(client, u):
    recs = []
    paths = ['/api/home', f"/api/nodes/{u['nodeId']}/quiz", '/api/user/summary', '/api/ai/quota']
    if WRITES:
        paths.append('answer')
    quiz = None
    i = 0
    while not stop.is_set():
        p = paths[i % len(paths)]
        i += 1
        t = time.monotonic()
        try:
            if p == 'answer':
                q = quiz['questions'][i % len(quiz['questions'])]
                r = await client.post(f"{BASE}/api/questions/{q['id']}/answer", json={'selectedOptionId': int(q['options'][0]['id'])},
                                      headers={'Authorization': f"Bearer {u['token']}"}, timeout=60)
                p = '/x/answer'
            else:
                r = await client.get(f'{BASE}{p}', headers={'Authorization': f"Bearer {u['token']}"}, timeout=60)
                if '/quiz' in p and r.status_code == 200:
                    quiz = r.json()
            recs.append({'path': p.split('/')[2], 'status': r.status_code, 'dt': time.monotonic() - t, 'at': t - T0})
        except Exception as exc:
            recs.append({'path': p.split('/')[2], 'status': 0, 'exc': type(exc).__name__, 'dt': time.monotonic() - t, 'at': t - T0})
        await asyncio.sleep(1.0)
    return recs


async def main():
    global T0
    limits = httpx.Limits(max_connections=200, max_keepalive_connections=200)
    async with httpx.AsyncClient(limits=limits, http2=False) as client:
        T0 = time.monotonic()
        btasks = [asyncio.create_task(browse(client, u)) for u in browsers]
        await asyncio.sleep(3 if browsers else 0)  # baseline browsing before the burst
        burst_at = time.monotonic() - T0
        results = await asyncio.gather(*(ask(client, u, i) for i, u in enumerate(askers)))
        burst_end = time.monotonic() - T0
        await asyncio.sleep(2 if browsers else 0)
        stop.set()
        brecs = [r for t in btasks for r in await t]
    json.dump({'base': BASE, 'n': N, 'browsers': B, 'round': ROUND, 'stagger': STAGGER,
               'burst_at': burst_at, 'burst_end': burst_end, 'asks': results, 'browse': brecs},
              open(OUT, 'w'), ensure_ascii=False, indent=1)
    ok = [r for r in results if r.get('t_done')]
    td = sorted(r['t_done'] for r in ok)
    tf = sorted(r['t_first'] for r in ok if r.get('t_first'))
    q = lambda a, p: a[min(len(a) - 1, int(p * len(a)))] if a else None
    print(f'N={N} ok={len(ok)}/{N} done p50={q(td, .5) and round(q(td, .5), 1)} p90={q(td, .9) and round(q(td, .9), 1)} '
          f'max={td and round(td[-1], 1)} first p50={q(tf, .5) and round(q(tf, .5), 1)}')
    bad = [r for r in results if not r.get('t_done')]
    for r in bad:
        print('  FAIL', {k: r.get(k) for k in ('userId', 'status', 'body', 'error', 'exc', 't_end')})
    if brecs:
        during = sorted(r['dt'] for r in brecs if burst_at <= r['at'] <= burst_end)
        before = sorted(r['dt'] for r in brecs if r['at'] < burst_at)
        errs = [r for r in brecs if r['status'] != 200]
        w = sorted(r['dt'] for r in brecs if r['path'] == 'answer' and burst_at <= r['at'] <= burst_end)
        if w:
            print(f'  writes during: n={len(w)} p50={round(q(w, .5) * 1000)}ms max={round(w[-1] * 1000)}ms statuses={sorted(set(r["status"] for r in brecs if r["path"] == "answer"))}')
        print(f'  browse before p50={q(before, .5) and round(q(before, .5) * 1000)}ms | during p50={q(during, .5) and round(q(during, .5) * 1000)}ms '
              f'p90={q(during, .9) and round(q(during, .9) * 1000)}ms max={during and round(during[-1] * 1000)}ms n={len(during)} errors={len(errs)}')


asyncio.run(main())