randogoth/pneumatube
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
pneumatube/scripts/stress_test.py

532 lines
23 KiB
Python
Raw Blame History

#!/usr/bin/env python3
"""Concurrent load test for pneumatube.
Two targets:
direct Talks to a `pneumatube` subprocess directly over LMTP/SMTP,
bypassing Postfix. Isolates pneumatube's own concurrency
behaviour -- in particular, SmolWorker (worker.py) serializes
every smolmail operation onto one thread, so this is where that
design's throughput ceiling actually shows up.
postfix Drives the same load through the real docker-compose.yml rig
(Postfix -> LMTP -> pneumatube), to confirm nothing changes with
a real MTA in front. Needs podman (or docker) + compose.
Both targets spin up a throwaway smolmaild, register a gateway account plus
N recipient/sender accounts against it, generate concurrent load, and verify
every expected message actually arrived exactly once -- throughput numbers
are meaningless if messages were silently dropped or duplicated.
Usage:
uv run python scripts/stress_test.py --target direct
uv run python scripts/stress_test.py --target postfix
uv run python scripts/stress_test.py --target both --concurrency 40
Needs a `smolmail` checkout next to this repo (same default pneumatube
itself uses: PNEUMATUBE_SMOLMAIL_REPO, default "../smolmail").
"""
from __future__ import annotations
import argparse
import contextlib
import os
import shutil
import smtplib
import socket
import subprocess
import sys
import tempfile
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from pathlib import Path
from aiosmtpd.controller import Controller
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
from pneumatube import vendor # noqa: E402
REPO_ROOT = Path(__file__).resolve().parent.parent
def free_port() -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
def wait_for_port(host: str, port: int, timeout: float = 15.0) -> None:
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
try:
with socket.create_connection((host, port), timeout=0.5):
return
except OSError:
time.sleep(0.1)
raise TimeoutError(f"nothing listening on {host}:{port} after {timeout}s")
# --- reporting ----------------------------------------------------------------
@dataclass
class Report:
label: str
attempted: int = 0
succeeded: int = 0
rejected: int = 0
errors: int = 0
latencies: list[float] = field(default_factory=list)
def record(self, outcome: str, latency: float) -> None:
self.attempted += 1
self.latencies.append(latency)
if outcome == "ok":
self.succeeded += 1
elif outcome == "rejected":
self.rejected += 1
else:
self.errors += 1
def print(self) -> None:
lat = sorted(self.latencies)
p50 = lat[len(lat) // 2] if lat else 0.0
p99 = lat[int(len(lat) * 0.99)] if lat else 0.0
print(f"-- {self.label} --")
print(f" attempted={self.attempted} succeeded={self.succeeded} "
f"rejected={self.rejected} errors={self.errors}")
print(f" latency p50={p50 * 1000:.1f}ms p99={p99 * 1000:.1f}ms "
f"max={max(lat, default=0) * 1000:.1f}ms")
# --- smolmail test fixtures ----------------------------------------------------
@dataclass
class SmolServer:
host: str
port: int
pubkey_b32: str
process: subprocess.Popen | None = None
def parse_pubkey(text: str) -> str:
for line in text.splitlines():
if "public key" in line:
return line.strip().rsplit(" ", 1)[-1]
raise RuntimeError(f"no public key line in output:\n{text}")
def start_smolmaild(smolmail_repo: Path, tmp_dir: Path, rate_connections: int,
rate_sends: int) -> SmolServer:
port = free_port()
key_path = tmp_dir / "server.key"
subprocess.run(
["uv", "run", str(smolmail_repo / "smolmaild.py"), "keygen", "--key", str(key_path)],
cwd=smolmail_repo, check=True, capture_output=True, text=True)
process = subprocess.Popen(
["uv", "run", str(smolmail_repo / "smolmaild.py"), "serve",
"--key", str(key_path), "--db", str(tmp_dir / "server.db"),
"--host", "127.0.0.1", "--port", str(port),
# smolmaild's connection/send rate limits are keyed by source IP
# (SPEC.md S10); every sender in this script shares one IP, so the
# defaults (120 conn/min, 60 sends/min) would measure the rate
# limiter, not pneumatube, unless raised for a many-sender test.
"--rate-connections", str(rate_connections), "--rate-sends", str(rate_sends)],
cwd=smolmail_repo, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
pubkey = None
deadline = time.monotonic() + 15
while time.monotonic() < deadline:
line = process.stdout.readline()
if not line:
break
if "server public key" in line:
pubkey = line.strip().rsplit(" ", 1)[-1]
break
if pubkey is None:
process.kill()
raise RuntimeError("smolmaild never printed its public key")
wait_for_port("127.0.0.1", port)
return SmolServer("127.0.0.1", port, pubkey, process)
@dataclass
class SmolAccount:
username: str
key_path: Path
db_path: Path
address: str # "user@host:port"
def make_account(smol, tmp_dir: Path, name: str, server: SmolServer) -> SmolAccount:
"""Keygen + trust + register, reusing smolmail.py's own CLI commands
directly (the same functions the gateway itself calls) rather than
reimplementing the protocol or shelling out per account.
A bare TCP connect (wait_for_port) can succeed slightly before a
containerized smolmaild can actually complete a Noise handshake, so
register is retried a few times rather than treating that race as fatal.
"""
key_path, db_path = tmp_dir / f"{name}.key", tmp_dir / f"{name}.db"
store = smol.Store(str(db_path))
smol.cmd_keygen(argparse.Namespace(key=str(key_path), force=True), store)
smol.cmd_trust(argparse.Namespace(host=f"{server.host}:{server.port}",
key_b32=server.pubkey_b32, force=True), store)
address = f"{name}@{server.host}:{server.port}"
register_args = argparse.Namespace(key=str(key_path), address=address, token=None)
attempts = 20
for attempt in range(attempts):
try:
smol.cmd_register(register_args, store)
break
except (smol.SmolError, ConnectionError, OSError):
if attempt == attempts - 1:
raise
time.sleep(1.0)
return SmolAccount(name, key_path, db_path, address)
def repoint_gateway_account(smol, gateway: SmolAccount, host: str, port: int,
pubkey_b32: str) -> None:
"""Point the gateway's own remembered server connection at a different
host string (e.g. a compose service name) without a real round trip --
`trust` only ever stores a pin locally (smolmail.py's own cmd_trust)."""
store = smol.Store(str(gateway.db_path))
smol.cmd_trust(argparse.Namespace(host=host, key_b32=pubkey_b32, force=True), store)
store.set_account(smol.Address(gateway.username, host, port))
# --- outbound SMTP sink (stands in for the real relay) -------------------------
class CountingSink:
def __init__(self) -> None:
self.lock = threading.Lock()
self.received: list[tuple[str, tuple[str, ...]]] = []
async def handle_DATA(self, server, session, envelope):
with self.lock:
self.received.append((envelope.mail_from, tuple(envelope.rcpt_tos)))
return "250 OK"
def count(self) -> int:
with self.lock:
return len(self.received)
# --- load generators (each measures its own latency) ---------------------------
def send_one_lmtp(host: str, port: int, mail_from: str, rcpt: str, body: bytes) -> tuple[str, float]:
t0 = time.monotonic()
try:
with smtplib.LMTP(host, port, timeout=10) as client:
client.ehlo("stress-test")
client.mail(mail_from)
code, _ = client.rcpt(rcpt)
if code >= 500:
return "rejected", time.monotonic() - t0
if code >= 400:
return "error", time.monotonic() - t0
client.data(body)
return "ok", time.monotonic() - t0
except Exception:
return "error", time.monotonic() - t0
def send_one_smtp(host: str, port: int, mail_from: str, rcpt: str, body: bytes) -> tuple[str, float]:
t0 = time.monotonic()
try:
with smtplib.SMTP(host, port, timeout=10) as client:
client.ehlo("stress-test")
client.mail(mail_from)
code, _ = client.rcpt(rcpt)
if code >= 500:
return "rejected", time.monotonic() - t0
if code >= 400:
return "error", time.monotonic() - t0
client.data(body)
return "ok", time.monotonic() - t0
except Exception:
return "error", time.monotonic() - t0
def send_one_smol(smol, account: SmolAccount, to_address: str, email_to: str,
subject: str) -> tuple[str, float]:
t0 = time.monotonic()
store = smol.Store(str(account.db_path))
send_args = argparse.Namespace(
key=str(account.key_path), address=to_address, body=f"load from {account.username}",
subject=subject, reply_to=None, header=[f"Email-To: {email_to}"],
anonymous=True, no_pad=False)
try:
smol.cmd_send(send_args, store)
return "ok", time.monotonic() - t0
except Exception:
return "error", time.monotonic() - t0
def run_concurrent(report: Report, concurrency: int, jobs: list) -> list[bool]:
"""jobs: list of (fn, args, expected_valid). Returns per-job valid flags
in submission order so the caller can cross-check expected vs actual."""
valid_flags = [valid for _, _, valid in jobs]
with ThreadPoolExecutor(max_workers=concurrency) as pool:
start = time.monotonic()
futures = [pool.submit(fn, *fn_args) for fn, fn_args, _ in jobs]
for future in futures:
outcome, latency = future.result()
report.record(outcome, latency)
elapsed = time.monotonic() - start
report.print()
print(f" throughput={len(jobs) / elapsed:.1f} msg/s over {elapsed:.1f}s")
return valid_flags
# --- direct target ---------------------------------------------------------------
def run_direct(args: argparse.Namespace) -> None:
smolmail_repo = Path(args.smolmail_repo).resolve()
with tempfile.TemporaryDirectory(prefix="pneumatube-stress-") as tmp:
tmp_dir = Path(tmp)
print(f"[direct] starting smolmaild in {tmp_dir}")
server = start_smolmaild(smolmail_repo, tmp_dir, args.server_rate_connections,
args.server_rate_sends)
smol = vendor.load(smolmail_repo)
gateway = make_account(smol, tmp_dir, "gatewayuser", server)
recipients = [make_account(smol, tmp_dir, f"recipient{i}", server)
for i in range(args.recipients)]
senders = [make_account(smol, tmp_dir, f"sender{i}", server)
for i in range(args.concurrency)]
print(f"[direct] {len(recipients)} recipient account(s), "
f"{len(senders)} sender account(s) ready")
sink = CountingSink()
sink_port = free_port()
sink_controller = Controller(sink, hostname="127.0.0.1", port=sink_port)
sink_controller.start()
listen_port, lookup_port = free_port(), free_port()
env = {**os.environ,
"PNEUMATUBE_SMOLMAIL_REPO": str(smolmail_repo),
"PNEUMATUBE_SMOLMAIL_KEY": str(gateway.key_path),
"PNEUMATUBE_SMOLMAIL_DB": str(gateway.db_path),
"PNEUMATUBE_SMOLMAIL_DOMAIN": f"{server.host}:{server.port}",
"PNEUMATUBE_EMAIL_DOMAIN": "gateway.example",
"PNEUMATUBE_FETCH_INTERVAL": "2",
# Throughput, not the accept-gate features, is this script's
# concern -- the throwaway recipient/sender accounts here
# never go through the accept dance with the gateway.
"PNEUMATUBE_REQUIRE_ACCEPTED_SENDER": "0",
"PNEUMATUBE_REQUIRE_ACCEPTED_RECIPIENT": "0",
"PNEUMATUBE_LISTEN_PROTOCOL": "lmtp",
"PNEUMATUBE_LISTEN_HOST": "127.0.0.1",
"PNEUMATUBE_LISTEN_PORT": str(listen_port),
"PNEUMATUBE_LOOKUP_HOST": "127.0.0.1",
"PNEUMATUBE_LOOKUP_PORT": str(lookup_port),
"PNEUMATUBE_SMTP_RELAY_HOST": "127.0.0.1",
"PNEUMATUBE_SMTP_RELAY_PORT": str(sink_port),
"PNEUMATUBE_SMTP_RELAY_STARTTLS": "0",
"PNEUMATUBE_GATEWAY_FROM_ADDRESS": "gateway@gateway.example"}
print("[direct] starting pneumatube")
proc = subprocess.Popen(["uv", "run", "pneumatube"], cwd=REPO_ROOT, env=env,
stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
try:
wait_for_port("127.0.0.1", listen_port)
wait_for_port("127.0.0.1", lookup_port)
# Phase A: email -> smolmail, concurrent LMTP, mixed valid/invalid.
total = args.concurrency * args.messages
jobs = []
for i in range(total):
valid = i % 3 != 0 # 2/3 valid, 1/3 unknown-user
rcpt = (f"{recipients[i % len(recipients)].username}@gateway.example"
if valid else "nobody-at-all@gateway.example")
body = f"Subject: load {i}\r\n\r\nmessage {i}\r\n".encode()
jobs.append((send_one_lmtp, ("127.0.0.1", listen_port,
"stress@sender.example", rcpt, body), valid))
report_a = Report("email -> smolmail (LMTP, direct)")
valid_flags = run_concurrent(report_a, args.concurrency, jobs)
expected_ok, expected_rejected = sum(valid_flags), valid_flags.count(False)
ok = report_a.succeeded == expected_ok and report_a.rejected == expected_rejected
print(f" correctness: {'PASS' if ok else 'FAIL'} "
f"(expected ok={expected_ok} rejected={expected_rejected})")
# Phase B: smolmail -> email, concurrent SEND, then drain via the
# gateway's own poll loop (already running in the subprocess).
jobs = [(send_one_smol, (smol, senders[i % len(senders)], gateway.address,
f"dest{i}@example.org", f"load {i}"), True)
for i in range(total)]
report_b = Report("smolmail -> email (SEND + poll drain)")
run_concurrent(report_b, args.concurrency, jobs)
drain_deadline = time.monotonic() + max(30.0, total * 0.05)
while sink.count() < report_b.succeeded and time.monotonic() < drain_deadline:
time.sleep(0.5)
drained = sink.count()
print(f" drained {drained}/{report_b.succeeded} into the outbound sink "
f"({'PASS' if drained == report_b.succeeded else 'FAIL: lost or stuck messages'})")
finally:
proc.terminate()
with contextlib.suppress(Exception):
proc.wait(timeout=5)
sink_controller.stop()
server.process.terminate()
with contextlib.suppress(Exception):
server.process.wait(timeout=5)
# --- postfix target ---------------------------------------------------------------
def podman(*cmd: str, **kwargs) -> subprocess.CompletedProcess:
return subprocess.run(["podman", *cmd], text=True, capture_output=True, **kwargs)
def run_postfix(args: argparse.Namespace) -> None:
if shutil.which("podman") is None:
print("[postfix] podman not found; skipping", file=sys.stderr)
return
smolmail_repo = Path(args.smolmail_repo).resolve()
smol = vendor.load(smolmail_repo)
network = "pneumatube-stress-net"
env = {**os.environ,
"PNEUMATUBE_SMOLMAIL_DOMAIN": "smolmail-dev:1961",
"PNEUMATUBE_EMAIL_DOMAIN": "gateway.example",
# Throughput, not the accept-gate features, is this script's
# concern -- the throwaway recipient accounts here never go
# through the accept dance with the gateway.
"PNEUMATUBE_REQUIRE_ACCEPTED_SENDER": "0",
"PNEUMATUBE_REQUIRE_ACCEPTED_RECIPIENT": "0",
"PNEUMATUBE_SMTP_RELAY_HOST": "smolmail-dev",
"PNEUMATUBE_SMTP_RELAY_PORT": "2626",
"PNEUMATUBE_SMTP_RELAY_STARTTLS": "0",
"PNEUMATUBE_GATEWAY_FROM_ADDRESS": "gateway@gateway.example"}
with tempfile.TemporaryDirectory(prefix="pneumatube-stress-pf-") as tmp:
tmp_dir = Path(tmp)
data_dir = REPO_ROOT / "data"
data_dir.mkdir(exist_ok=True)
try:
print("[postfix] building images")
subprocess.run(["podman", "compose", "build"], cwd=REPO_ROOT, env=env, check=True)
print("[postfix] starting throwaway smolmaild-dev container")
podman("network", "create", network)
server_dir = tmp_dir / "server"
server_dir.mkdir()
keygen = podman("run", "--rm", "-v", f"{smolmail_repo}:/smolmail:ro",
"-v", f"{server_dir}:/data",
"localhost/pneumatube_pneumatube:latest",
"uv", "run", "/smolmail/smolmaild.py", "keygen", "--key", "/data/server.key")
keygen.check_returncode()
pubkey = parse_pubkey(keygen.stdout)
host_port = free_port()
podman("run", "-d", "--name", "pneumatube-stress-smolmaild",
"--network", network, "--network-alias", "smolmail-dev",
"-p", f"{host_port}:1961",
"-v", f"{smolmail_repo}:/smolmail:ro", "-v", f"{server_dir}:/data",
"localhost/pneumatube_pneumatube:latest",
"uv", "run", "/smolmail/smolmaild.py", "serve",
"--key", "/data/server.key", "--db", "/data/server.db",
"--host", "0.0.0.0", "--port", "1961",
"--rate-connections", str(args.server_rate_connections),
"--rate-sends", str(args.server_rate_sends)).check_returncode()
wait_for_port("127.0.0.1", host_port)
server = SmolServer("127.0.0.1", host_port, pubkey)
gateway = make_account(smol, tmp_dir, "gatewayuser", server)
recipients = [make_account(smol, tmp_dir, f"recipient{i}", server)
for i in range(args.recipients)]
repoint_gateway_account(smol, gateway, "smolmail-dev", 1961, pubkey)
(data_dir / "gateway.key").write_bytes(gateway.key_path.read_bytes())
(data_dir / "gateway.db").write_bytes(gateway.db_path.read_bytes())
print("[postfix] starting pneumatube + postfix")
subprocess.run(["podman", "compose", "up", "-d"], cwd=REPO_ROOT, env=env, check=True)
# Join the compose-started containers to the network carrying
# smolmail-dev (they start on compose's own default network,
# which smolmail-dev was never part of) -- not the other way
# around, which would just reconnect smolmail-dev to a network
# it's already on and leave pneumatube unable to resolve it.
podman("network", "connect", network, "pneumatube_pneumatube_1")
podman("network", "connect", network, "pneumatube_postfix_1")
wait_for_port("127.0.0.1", 2525, timeout=60.0)
# Port 25 accepting doesn't mean postconf's chroot/transport_maps
# setup (entrypoint.sh) has finished; give it a moment to settle.
time.sleep(2)
total = args.concurrency * args.messages
jobs = []
for i in range(total):
valid = i % 3 != 0
rcpt = (f"{recipients[i % len(recipients)].username}@gateway.example"
if valid else "nobody-at-all@gateway.example")
body = f"Subject: pf-load {i}\r\n\r\nmessage {i}\r\n".encode()
jobs.append((send_one_smtp, ("127.0.0.1", 2525,
"stress@sender.example", rcpt, body), valid))
report = Report("email -> smolmail (through Postfix)")
valid_flags = run_concurrent(report, args.concurrency, jobs)
expected_ok, expected_rejected = sum(valid_flags), valid_flags.count(False)
print(f" correctness (live reject, not bounce): "
f"{'PASS' if report.rejected == expected_rejected else 'FAIL'} "
f"(expected ok={expected_ok} rejected={expected_rejected})")
deadline = time.monotonic() + max(30.0, total * 0.1)
delivered = 0
while time.monotonic() < deadline:
delivered = 0
for r in recipients:
store = smol.Store(str(r.db_path))
smol.cmd_fetch(argparse.Namespace(key=str(r.key_path), db=str(r.db_path),
reset=False, keep=False), store)
delivered += len(store.mail("all"))
if delivered >= expected_ok:
break
time.sleep(1)
print(f" delivered {delivered}/{expected_ok} across recipient inboxes "
f"({'PASS' if delivered == expected_ok else 'FAIL: lost or stuck messages'})")
finally:
print("[postfix] tearing down")
subprocess.run(["podman", "compose", "down"], cwd=REPO_ROOT, env=env)
podman("rm", "-f", "pneumatube-stress-smolmaild")
podman("network", "rm", network)
with contextlib.suppress(Exception):
(data_dir / "gateway.key").unlink()
(data_dir / "gateway.db").unlink()
data_dir.rmdir()
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--target", choices=["direct", "postfix", "both"], default="direct")
parser.add_argument("--concurrency", type=int, default=20,
help="concurrent sender workers per phase")
parser.add_argument("--messages", type=int, default=10,
help="messages sent by each worker")
parser.add_argument("--recipients", type=int, default=10,
help="number of pre-registered valid recipient accounts")
parser.add_argument("--smolmail-repo", default="../smolmail")
parser.add_argument("--server-rate-connections", type=int, default=6000,
help="smolmaild --rate-connections for the throwaway test server "
"(default 120/min would rate-limit this whole test, since every "
"sender shares one source IP; pass 120 to test against that "
"realistic default instead)")
parser.add_argument("--server-rate-sends", type=int, default=6000,
help="smolmaild --rate-sends for the throwaway test server "
"(see --server-rate-connections; realistic default is 60)")
args = parser.parse_args()
if args.target in ("direct", "both"):
run_direct(args)
if args.target in ("postfix", "both"):
run_postfix(args)
if __name__ == "__main__":
main()
Reference in a new issue View git blame Copy permalink