jcc-blockchain/scripts/load_test.py

#!/usr/bin/env python3
"""
load_test.py — stdlib-only RPC load-test harness.
Targets any JSON-RPC endpoint (mock or real Besu).

Usage:
    python3 load_test.py --url http://localhost:8545 \
                         --method eth_sendTransaction \
                         --duration 30 \
                         --workers 8 \
                         --rps 500
"""

import argparse
import json
import multiprocessing
import signal
import sys
import time
import urllib.request
from collections import deque
from typing import Any, Dict, List


def _rpc_request(url: str, payload: Dict[str, Any], timeout: float = 10.0) -> Dict[str, Any]:
    req = urllib.request.Request(
        url,
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return json.loads(resp.read())


def _worker_process(
    worker_id: int,
    url: str,
    method: str,
    duration: float,
    rps: float,
    result_queue,
):
    """Target function for each worker process."""
    end_time = time.monotonic() + duration
    interval = 1.0 / rps if rps > 0 else 0.0
    count = 0
    dropped = 0
    latencies: List[float] = []
    errors: List[str] = []

    # Deterministic account for eth_sendTransaction
    addr = f"0x{'%040x' % worker_id}"
    to_addr = f"0x{'%040x' % ((worker_id + 1) % 10)}"

    while time.monotonic() < end_time:
        start = time.monotonic()
        if method == "eth_sendTransaction":
            payload = {
                "jsonrpc": "2.0",
                "method": "eth_sendTransaction",
                "params": [
                    {
                        "from": addr,
                        "to": to_addr,
                        "value": "0x1",
                        "gas": "0x5208",
                        "gasPrice": "0x1",
                    }
                ],
                "id": count,
            }
        elif method == "eth_getBalance":
            payload = {
                "jsonrpc": "2.0",
                "method": "eth_getBalance",
                "params": [addr, "latest"],
                "id": count,
            }
        else:
            payload = {
                "jsonrpc": "2.0",
                "method": method,
                "params": [],
                "id": count,
            }

        try:
            _rpc_request(url, payload, timeout=10.0)
            latencies.append((time.monotonic() - start) * 1000)
            count += 1
        except Exception as e:
            dropped += 1
            errors.append(str(e))

        # Rate-limit sleep
        elapsed = time.monotonic() - start
        if interval > elapsed:
            time.sleep(interval - elapsed)

    result_queue.put(
        {
            "worker_id": worker_id,
            "requests": count,
            "dropped": dropped,
            "latencies": latencies,
            "errors": errors[:5],  # cap error detail
        }
    )


def run_load_test(url: str, method: str, duration: float, workers: int, rps: float):
    result_queue = multiprocessing.Queue()
    processes: List[multiprocessing.Process] = []

    def sigint_handler(_sig, _frame):
        for p in processes:
            p.terminate()
        sys.exit(0)

    signal.signal(signal.SIGINT, sigint_handler)

    start = time.monotonic()
    for i in range(workers):
        p = multiprocessing.Process(
            target=_worker_process,
            args=(i, url, method, duration, rps / workers, result_queue),
        )
        p.start()
        processes.append(p)

    for p in processes:
        p.join()

    elapsed = time.monotonic() - start
    total_requests = 0
    total_dropped = 0
    all_latencies: List[float] = []
    for _ in range(workers):
        result = result_queue.get()
        total_requests += result["requests"]
        total_dropped += result["dropped"]
        all_latencies.extend(result["latencies"])

    all_latencies.sort()
    mean_latency = sum(all_latencies) / len(all_latencies) if all_latencies else 0.0
    p95 = all_latencies[int(len(all_latencies) * 0.95)] if all_latencies else 0.0
    p99 = all_latencies[int(len(all_latencies) * 0.99)] if all_latencies else 0.0
    actual_tps = total_requests / elapsed if elapsed > 0 else 0.0
    drop_rate = total_dropped / (total_requests + total_dropped) if (total_requests + total_dropped) > 0 else 0.0

    print("\n=== Load Test Results ===")
    print(f"URL          : {url}")
    print(f"Method       : {method}")
    print(f"Workers      : {workers}")
    print(f"Target RPS   : {rps}")
    print(f"Duration     : {elapsed:.1f}s")
    print(f"Total reqs   : {total_requests}")
    print(f"Dropped      : {total_dropped} ({drop_rate*100:.2f}%)")
    print(f"Actual TPS   : {actual_tps:.1f}")
    print(f"Latency avg  : {mean_latency:.2f} ms")
    print(f"Latency p95  : {p95:.2f} ms")
    print(f"Latency p99  : {p99:.2f} ms")

    return {
        "total_requests": total_requests,
        "total_dropped": total_dropped,
        "drop_rate": drop_rate,
        "actual_tps": actual_tps,
        "mean_latency_ms": mean_latency,
        "p95_latency_ms": p95,
        "p99_latency_ms": p99,
        "duration_s": elapsed,
        "pass": drop_rate <= 0.01 and actual_tps >= rps * 0.95,
    }


def main():
    parser = argparse.ArgumentParser(description="RPC load-test harness")
    parser.add_argument("--url", default="http://localhost:8545")
    parser.add_argument("--method", default="eth_getBalance", choices=["eth_sendTransaction", "eth_getBalance", "eth_blockNumber"])
    parser.add_argument("--duration", type=float, default=30.0, help="Test duration in seconds")
    parser.add_argument("--workers", type=int, default=8, help="Number of concurrent worker processes")
    parser.add_argument("--rps", type=float, default=500.0, help="Target aggregate requests per second")
    parser.add_argument("--latency-budget", type=float, default=200.0, help="Latency budget in ms for pass/fail")
    args = parser.parse_args()

    result = run_load_test(args.url, args.method, args.duration, args.workers, args.rps)

    if result["pass"]:
        print("\nPASS: TPS and drop-rate within acceptance thresholds.")
        sys.exit(0)
    else:
        print("\nFAIL: Did not meet TPS/drop-rate thresholds.")
        sys.exit(1)


if __name__ == "__main__":
    main()