""" Minimal HTTP/1.1 keep-alive load generator (stdlib only, raw asyncio sockets). Each of N connections is an asyncio.Protocol that sends one GET, parses the response (Content-Length or chunked), records latency, and immediately sends the next request on the same connection (closed-loop, like wrk -c N). Usage: python load.py --port 8801 --path /-/versions.json -c 8 -d 10 [--warmup 1] [--cookie 'ds_actor=...'] [--header 'If-None-Match: "abc"'] [--procs 2] [--json] --procs P splits the N connections over P worker processes so that the load generator itself is not the bottleneck. """ import argparse import asyncio import json import multiprocessing import statistics import time HEADERS, BODY_CL, CHUNK_SIZE, CHUNK_DATA, CHUNK_END = range(5) class Conn(asyncio.Protocol): def __init__(self, request, stats, deadline_holder, done): self.request = request self.stats = stats self.deadline = deadline_holder self.done = done self.buf = bytearray() self.state = HEADERS self.remaining = 0 self.status = 0 self.body_bytes = 0 self.t0 = 0 def connection_made(self, transport): self.transport = transport self.send() def send(self): self.state = HEADERS self.body_bytes = 0 self.t0 = time.perf_counter_ns() self.transport.write(self.request) def connection_lost(self, exc): if not self.done.done(): self.stats["errors"] += 1 self.done.set_result(None) def data_received(self, data): self.buf += data while self._parse(): now = time.perf_counter_ns() st = self.stats st["all"] += 1 if now >= st["measure_from"]: st["lat"].append(now - self.t0) st["status"][self.status] = st["status"].get(self.status, 0) + 1 st["bytes"] += self.body_bytes if now >= self.deadline[0]: if not self.done.done(): self.done.set_result(None) self.transport.close() return self.send() def _parse(self): """Advance the state machine. Return True when a full response is in.""" buf = self.buf while True: if self.state == HEADERS: idx = buf.find(b"\r\n\r\n") if idx < 0: return False head = bytes(buf[:idx]) del buf[: idx + 4] self.status = int(head[9:12]) lower = head.lower() cl = lower.find(b"\r\ncontent-length:") if cl >= 0: end = lower.find(b"\r\n", cl + 2) if end < 0: end = len(lower) self.remaining = int(lower[cl + 17 : end].strip()) self.state = BODY_CL elif b"\r\ntransfer-encoding: chunked" in lower: self.state = CHUNK_SIZE else: # 304 / 204 / no body self.state = HEADERS return True elif self.state == BODY_CL: n = min(self.remaining, len(buf)) del buf[:n] self.remaining -= n self.body_bytes += n if self.remaining: return False self.state = HEADERS return True elif self.state == CHUNK_SIZE: idx = buf.find(b"\r\n") if idx < 0: return False size = int(bytes(buf[:idx]).split(b";")[0], 16) del buf[: idx + 2] if size == 0: self.state = CHUNK_END else: self.remaining = size + 2 self.body_bytes += size self.state = CHUNK_DATA elif self.state == CHUNK_DATA: n = min(self.remaining, len(buf)) del buf[:n] self.remaining -= n if self.remaining: return False self.state = CHUNK_SIZE elif self.state == CHUNK_END: # assume no trailers if len(buf) < 2: return False del buf[:2] self.state = HEADERS return True def build_request(args): lines = [ f"GET {args.path} HTTP/1.1", f"Host: {args.host}:{args.port}", "User-Agent: load.py", "Accept: */*", "Connection: keep-alive", ] if args.cookie: lines.append(f"Cookie: {args.cookie}") lines.extend(args.header or []) return ("\r\n".join(lines) + "\r\n\r\n").encode("latin-1") async def run_worker(args, conns): loop = asyncio.get_running_loop() request = build_request(args) start = time.perf_counter_ns() stats = { "lat": [], "status": {}, "bytes": 0, "errors": 0, "all": 0, "measure_from": start + int(args.warmup * 1e9), } deadline = [start + int((args.warmup + args.duration) * 1e9)] dones = [] for _ in range(conns): done = loop.create_future() dones.append(done) await loop.create_connection( lambda d=done: Conn(request, stats, deadline, d), args.host, args.port ) await asyncio.gather(*dones) end = time.perf_counter_ns() stats["elapsed"] = (end - stats["measure_from"]) / 1e9 return stats def worker_entry(payload): args, conns = payload return asyncio.run(run_worker(args, conns)) def pct(sorted_vals, p): if not sorted_vals: return float("nan") k = min(len(sorted_vals) - 1, round(p / 100 * (len(sorted_vals) - 1))) return sorted_vals[k] def main(): p = argparse.ArgumentParser() p.add_argument("--host", default="127.0.0.1") p.add_argument("--port", type=int, required=True) p.add_argument("--path", required=True) p.add_argument("-c", "--connections", type=int, default=8) p.add_argument("-d", "--duration", type=float, default=10) p.add_argument("--warmup", type=float, default=1.0) p.add_argument("--procs", type=int, default=1) p.add_argument("--cookie") p.add_argument("--header", action="append") p.add_argument("--json", action="store_true") p.add_argument("--label", default="") args = p.parse_args() procs = max(1, min(args.procs, args.connections)) if procs == 1: results = [asyncio.run(run_worker(args, args.connections))] else: split = [args.connections // procs] * procs for i in range(args.connections % procs): split[i] += 1 ctx = multiprocessing.get_context("fork") with ctx.Pool(procs) as pool: results = pool.map(worker_entry, [(args, n) for n in split]) lat = sorted(x for r in results for x in r["lat"]) status = {} for r in results: for k, v in r["status"].items(): status[k] = status.get(k, 0) + v elapsed = statistics.mean(r["elapsed"] for r in results) n = len(lat) out = { "label": args.label, "path": args.path, "connections": args.connections, "procs": procs, "requests": n, "elapsed_s": round(elapsed, 3), "rps": round(n / elapsed, 1) if elapsed else 0, "lat_us": { "mean": round(statistics.mean(lat) / 1e3, 1) if lat else None, "p50": round(pct(lat, 50) / 1e3, 1), "p90": round(pct(lat, 90) / 1e3, 1), "p99": round(pct(lat, 99) / 1e3, 1), "p99.9": round(pct(lat, 99.9) / 1e3, 1), "max": round(lat[-1] / 1e3, 1) if lat else None, }, "status": status, "avg_body_bytes": round(sum(r["bytes"] for r in results) / n) if n else 0, "errors": sum(r["errors"] for r in results), "requests_total": sum(r["all"] for r in results), } if args.json: print(json.dumps(out)) else: l = out["lat_us"] print( f"{args.label or args.path}: c={args.connections} {out['rps']} req/s " f"n={n} p50={l['p50']}us p90={l['p90']}us p99={l['p99']}us " f"max={l['max']}us status={status} body={out['avg_body_bytes']}B " f"errors={out['errors']}" ) if __name__ == "__main__": main()