bench.py

#!/usr/bin/env python3
"""
Framework Benchmark Runner

Benchmarks 10 framework configurations with 5 endpoints each using bombardier:
- Django Bolt
- Django REST Framework (DRF) - default, granian, gunicorn variants
- FastAPI - default, granian variants
- Litestar - default, granian variants
- Django Ninja - default, granian variants

Endpoints:
1. /json-1k  - ~1KB JSON response
2. /json-10k - ~10KB JSON response
3. /db       - 10 database reads (simple query)
4. /articles - Paginated articles with nested author + tags
5. /articles/1 - Single article with nested author + tags + comments

Best practices followed:
- Warmup requests before benchmarking
- Multiple runs for consistency
- Controlled concurrency
- Results saved to JSON for flexible analysis

Usage:
    # Run all frameworks and save to results/ directory (merge with existing)
    python bench.py

    # Run specific frameworks and override all existing data
    python bench.py --frameworks fastapi-uvicorn litestar-uvicorn --override

    # Run single framework and update results (merges with existing)
    python bench.py --frameworks fastapi-uvicorn

    # Generate visualizations from saved results
    python visualize.py
"""

from __future__ import annotations

import argparse
import json
import signal
import subprocess
import sys
import time
import threading
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Optional

import psutil


@dataclass
class BenchConfig:
    """Benchmark configuration."""

    connections: int = 100  # Concurrent connections
    duration: int = 10      # Duration in seconds
    warmup_requests: int = 1000  # Warmup requests
    runs: int = 3           # Number of runs per endpoint


@dataclass
class BenchResult:
    """Single benchmark result."""

    framework: str
    endpoint: str
    rps: float
    latency_avg_ms: float
    latency_p99_ms: float
    errors: int
    duration_s: float
    # Resource usage metrics
    mem_peak_mb: float = 0.0
    mem_avg_mb: float = 0.0
    cpu_peak_percent: float = 0.0
    cpu_avg_percent: float = 0.0


class ResourceMonitor:
    """Monitor CPU and memory usage of a process."""

    def __init__(self, port: int, use_docker: bool = False, container_name: Optional[str] = None):
        self.port = port
        self.use_docker = use_docker
        self.container_name = container_name
        self.samples = []
        self.running = False
        self.thread = None

    def _find_process_by_port(self) -> Optional[psutil.Process]:
        """Find process listening on the given port."""
        for proc in psutil.process_iter(['pid', 'name']):
            try:
                for conn in proc.net_connections(kind='inet'):
                    if conn.laddr.port == self.port and conn.status == 'LISTEN':
                        return proc
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                continue
        return None

    def _get_docker_stats(self) -> Optional[dict]:
        """Get Docker container stats."""
        if not self.container_name:
            return None

        try:
            cmd = [
                "docker", "stats", self.container_name,
                "--no-stream", "--format",
                "{{.MemUsage}}\t{{.CPUPerc}}"
            ]
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=2)
            if result.returncode == 0:
                mem_str, cpu_str = result.stdout.strip().split('\t')
                # Parse memory (e.g., "123.4MiB / 500MiB")
                mem_mb = float(mem_str.split('/')[0].strip().replace('MiB', '').replace('GiB', '000'))
                # Parse CPU (e.g., "12.34%")
                cpu_percent = float(cpu_str.strip().replace('%', ''))
                return {"mem_mb": mem_mb, "cpu_percent": cpu_percent}
        except (subprocess.TimeoutExpired, ValueError, IndexError):
            pass
        return None

    def _monitor_loop(self):
        """Background monitoring loop."""
        process = None
        if not self.use_docker:
            process = self._find_process_by_port()
            if not process:
                print(f"  WARNING: Could not find process for port {self.port}")
                return

        while self.running:
            try:
                if self.use_docker:
                    stats = self._get_docker_stats()
                    if stats:
                        self.samples.append(stats)
                else:
                    if process and process.is_running():
                        mem_mb = process.memory_info().rss / (1024 * 1024)
                        cpu_percent = process.cpu_percent(interval=0.1)
                        self.samples.append({"mem_mb": mem_mb, "cpu_percent": cpu_percent})
                    else:
                        break
                time.sleep(0.5)
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                break

    def start(self):
        """Start monitoring in background thread."""
        self.running = True
        self.samples = []
        self.thread = threading.Thread(target=self._monitor_loop, daemon=True)
        self.thread.start()

    def stop(self) -> dict:
        """Stop monitoring and return statistics."""
        self.running = False
        if self.thread:
            self.thread.join(timeout=2)

        if not self.samples:
            return {
                "mem_peak_mb": 0.0,
                "mem_avg_mb": 0.0,
                "cpu_peak_percent": 0.0,
                "cpu_avg_percent": 0.0,
            }

        mem_values = [s["mem_mb"] for s in self.samples]
        cpu_values = [s["cpu_percent"] for s in self.samples]

        return {
            "mem_peak_mb": max(mem_values),
            "mem_avg_mb": sum(mem_values) / len(mem_values),
            "cpu_peak_percent": max(cpu_values),
            "cpu_avg_percent": sum(cpu_values) / len(cpu_values),
        }


FRAMEWORKS = {
    # Bolt (9001)
    "bolt": {"port": 9001, "prefix": ""},

    # DRF variants (902X)
    "drf-uvicorn": {"port": 9021, "prefix": "/drf"},
    "drf-granian": {"port": 9022, "prefix": "/drf"},
    "drf-gunicorn": {"port": 9023, "prefix": "/drf"},

    # FastAPI variants (903X)
    "fastapi-uvicorn": {"port": 9031, "prefix": ""},
    "fastapi-granian": {"port": 9032, "prefix": ""},

    # Litestar variants (904X)
    "litestar-uvicorn": {"port": 9041, "prefix": ""},
    "litestar-granian": {"port": 9042, "prefix": ""},

    # Ninja variants (905X)
    "ninja-uvicorn": {"port": 9051, "prefix": "/ninja"},
    "ninja-granian": {"port": 9052, "prefix": "/ninja"},
}

ENDPOINTS = [
    "/json-1k",
    "/json-10k",
    "/db",
    "/articles?page=1&page_size=20",
    "/articles/1",
]


def run_bombardier(
    url: str,
    connections: int,
    duration: int,
) -> dict | None:
    """Run bombardier and return parsed results."""
    cmd = [
        "bombardier",
        "-c", str(connections),
        "-d", f"{duration}s",
        "--print", "r",
        "--format", "json",
        url,
    ]

    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=duration + 30,
        )
        if result.returncode != 0:
            print(f"  ERROR: bombardier failed: {result.stderr}")
            return None

        return json.loads(result.stdout)
    except subprocess.TimeoutExpired:
        print(f"  ERROR: bombardier timed out")
        return None
    except json.JSONDecodeError as e:
        print(f"  ERROR: Failed to parse bombardier output: {e}")
        return None
    except FileNotFoundError:
        print("ERROR: bombardier not found. Install it first:")
        print("  go install github.com/codesenberg/bombardier@latest")
        sys.exit(1)


def warmup(url: str, requests: int) -> bool:
    """Send warmup requests using bombardier."""
    cmd = [
        "bombardier",
        "-c", "10",
        "-n", str(requests),
        "--print", "r",
        url,
    ]

    try:
        result = subprocess.run(cmd, capture_output=True, timeout=60)
        return result.returncode == 0
    except (subprocess.TimeoutExpired, FileNotFoundError):
        return False


def check_server(host: str, port: int, prefix: str) -> bool:
    """Check if server is responding."""
    import urllib.request
    import urllib.error

    url = f"http://{host}:{port}{prefix}/json-1k"
    try:
        with urllib.request.urlopen(url, timeout=5) as resp:
            return resp.status == 200
    except (urllib.error.URLError, TimeoutError, ConnectionResetError, ConnectionRefusedError):
        return False


def benchmark_framework(
    framework: str,
    host: str,
    port: int,
    prefix: str,
    config: BenchConfig,
    use_docker: bool = False,
    container_name: Optional[str] = None,
) -> list[BenchResult]:
    """Benchmark a single framework."""
    results = []
    base_url = f"http://{host}:{port}{prefix}"

    print(f"\n{'='*60}")
    print(f"Benchmarking: {framework.upper()}")
    print(f"Base URL: {base_url}")
    print(f"Mode: {'Docker' if use_docker else 'Local'}")
    print(f"{'='*60}")

    # Check server
    if not check_server(host, port, prefix):
        print(f"  ERROR: Server not responding at {base_url}")
        return results

    for endpoint in ENDPOINTS:
        url = f"{base_url}{endpoint}"
        print(f"\n  Endpoint: {endpoint}")

        # Warmup
        print(f"    Warming up ({config.warmup_requests} requests)...")
        if not warmup(url, config.warmup_requests):
            print(f"    WARNING: Warmup failed")

        # Multiple runs
        run_results = []
        for run in range(config.runs):
            print(f"    Run {run + 1}/{config.runs}...", end=" ", flush=True)

            # Start resource monitoring
            monitor = ResourceMonitor(port, use_docker=use_docker, container_name=container_name)
            monitor.start()
            time.sleep(0.5)  # Let monitor stabilize

            data = run_bombardier(url, config.connections, config.duration)

            # Stop resource monitoring
            resource_stats = monitor.stop()

            if data is None:
                print("FAILED")
                continue

            rps = data.get("result", {}).get("rps", {}).get("mean", 0)
            latency = data.get("result", {}).get("latency", {})
            latency_avg = latency.get("mean", 0) / 1_000_000  # ns to ms
            latency_p99 = latency.get("percentiles", {}).get("99", 0) / 1_000_000

            errors = 0
            for code, count in data.get("result", {}).get("statusCodeDistribution", {}).items():
                if not code.startswith("2"):
                    errors += count

            print(f"RPS: {rps:,.0f}, Latency: {latency_avg:.2f}ms (p99: {latency_p99:.2f}ms), "
                  f"Mem: {resource_stats['mem_peak_mb']:.0f}MB, CPU: {resource_stats['cpu_avg_percent']:.1f}%")

            run_results.append({
                "rps": rps,
                "latency_avg": latency_avg,
                "latency_p99": latency_p99,
                "errors": errors,
                "mem_peak_mb": resource_stats["mem_peak_mb"],
                "mem_avg_mb": resource_stats["mem_avg_mb"],
                "cpu_peak_percent": resource_stats["cpu_peak_percent"],
                "cpu_avg_percent": resource_stats["cpu_avg_percent"],
            })

        if run_results:
            # Take the best RPS run
            best = max(run_results, key=lambda x: x["rps"])
            results.append(BenchResult(
                framework=framework,
                endpoint=endpoint,
                rps=best["rps"],
                latency_avg_ms=best["latency_avg"],
                latency_p99_ms=best["latency_p99"],
                errors=best["errors"],
                duration_s=config.duration,
                mem_peak_mb=best["mem_peak_mb"],
                mem_avg_mb=best["mem_avg_mb"],
                cpu_peak_percent=best["cpu_peak_percent"],
                cpu_avg_percent=best["cpu_avg_percent"],
            ))

    return results


def save_results_to_json(results: list[BenchResult], results_dir: Path, override: bool = False) -> None:
    """Save benchmark results to JSON files (one file per framework).

    Args:
        results: List of benchmark results
        results_dir: Directory to save framework JSON files
        override: If True, delete all existing files first. If False, only update specified frameworks.
    """
    # Create results directory
    results_dir.mkdir(exist_ok=True)

    # If override mode, delete all existing framework files
    if override:
        for existing_file in results_dir.glob("*.json"):
            existing_file.unlink()
        print(f"\nOverride mode: Cleared all existing files in {results_dir}/")

    # Group results by framework
    frameworks_data = {}
    for result in results:
        if result.framework not in frameworks_data:
            frameworks_data[result.framework] = {}

        frameworks_data[result.framework][result.endpoint] = {
            "rps": result.rps,
            "latency_avg_ms": result.latency_avg_ms,
            "latency_p99_ms": result.latency_p99_ms,
            "errors": result.errors,
            "duration_s": result.duration_s,
            "mem_peak_mb": result.mem_peak_mb,
            "mem_avg_mb": result.mem_avg_mb,
            "cpu_peak_percent": result.cpu_peak_percent,
            "cpu_avg_percent": result.cpu_avg_percent,
        }

    # Save each framework to its own file
    for framework, endpoints_data in frameworks_data.items():
        framework_file = results_dir / f"{framework}.json"
        with open(framework_file, 'w') as f:
            json.dump(endpoints_data, f, indent=2)
        print(f"  Saved: {framework_file}")

    print(f"\nResults saved to: {results_dir}/")
    print(f"  Frameworks updated: {len(frameworks_data)}")
    if override:
        print("  Mode: Override (replaced all existing data)")
    else:
        print("  Mode: Merge (updated specified frameworks only)")


def start_docker_service(framework: str, port: int, prefix: str = "") -> bool:
    """Start a framework in Docker and wait for it to be ready."""
    # Container name mapping
    container_name = f"benchmark-{framework}"

    # Framework-specific commands matching run_*.sh scripts
    commands = {
        # Bolt (9001)
        "bolt": ["sh", "-c", "DJANGO_SETTINGS_MODULE=django_project.settings_bolt uv run python manage.py runbolt --host 0.0.0.0 --port 9001 --processes 1"],

        # DRF variants (902X)
        "drf-uvicorn": ["uv", "run", "uvicorn", "django_project.asgi:application", "--host", "0.0.0.0", "--port", "9021", "--workers", "1", "--no-access-log"],
        "drf-granian": ["uv", "run", "granian", "--interface", "wsgi", "--host", "0.0.0.0", "--port", "9022", "django_project.wsgi:application"],
        "drf-gunicorn": ["uv", "run", "gunicorn", "-c", "gunicorn_drf.conf.py", "django_project.wsgi:application"],

        # FastAPI variants (903X)
        "fastapi-uvicorn": ["uv", "run", "uvicorn", "fastapi_app:app", "--host", "0.0.0.0", "--port", "9031", "--workers", "1", "--no-access-log"],
        "fastapi-granian": ["uv", "run", "granian", "--interface", "asgi", "--host", "0.0.0.0", "--port", "9032", "fastapi_app:app"],

        # Litestar variants (904X)
        "litestar-uvicorn": ["uv", "run", "uvicorn", "litestar_app:app", "--host", "0.0.0.0", "--port", "9041", "--workers", "1", "--no-access-log"],
        "litestar-granian": ["uv", "run", "granian", "--interface", "asgi", "--host", "0.0.0.0", "--port", "9042", "litestar_app:app"],

        # Ninja variants (905X)
        "ninja-uvicorn": ["uv", "run", "uvicorn", "django_project.asgi:application", "--host", "0.0.0.0", "--port", "9051", "--workers", "1", "--no-access-log"],
        "ninja-granian": ["uv", "run", "granian", "--interface", "asgi", "--host", "0.0.0.0", "--port", "9052", "django_project.asgi:application"],
    }

    # Django settings module mapping
    django_settings = {
        "bolt": "django_project.settings_bolt",
        "drf-uvicorn": "django_project.settings_drf",
        "drf-granian": "django_project.settings_drf",
        "drf-gunicorn": "django_project.settings_drf_gunicorn",  # No pool for gevent
        "ninja-uvicorn": "django_project.settings_ninja",
        "ninja-granian": "django_project.settings_ninja",
    }

    # Gunicorn environment variables
    gunicorn_env = {
        "drf-gunicorn": {"GUNICORN_PORT": "9023", "GUNICORN_WORKERS": "1", "GUNICORN_WORKER_CONNECTIONS": "1000"},
    }

    try:
        # Build docker run command
        cmd = [
            "docker", "run", "-d",
            "--rm",
            "--name", container_name,
            "-p", f"{port}:{port}",
            "--cpus=1",
            "--memory=750m",
            "--add-host=host.docker.internal:host-gateway",
            "--env-file", "docker.env",
        ]

        # Django-specific settings
        if framework in django_settings:
            cmd.extend(["-e", f"DJANGO_SETTINGS_MODULE={django_settings[framework]}"])

        # Gunicorn-specific environment variables
        if framework in gunicorn_env:
            for key, value in gunicorn_env[framework].items():
                cmd.extend(["-e", f"{key}={value}"])

        cmd.extend([
            "benchmark-framework:latest"
        ] + commands[framework])

        result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
        if result.returncode != 0:
            print(f"  ERROR: Failed to start {container_name}: {result.stderr}")
            return False

        print(f"  Started {container_name} container")
        print(f"  Waiting for {framework} to be ready...", end="", flush=True)

        # Wait for service to respond (up to 30 seconds)
        for i in range(30):
            if check_server("127.0.0.1", port, prefix):
                print(" Ready!")
                return True
            print(".", end="", flush=True)
            time.sleep(1)

        print(" Timeout!")
        print(f"  ERROR: Service {framework} did not respond in time")
        # Clean up failed container
        subprocess.run(["docker", "stop", container_name], capture_output=True, timeout=10)
        return False

    except subprocess.TimeoutExpired:
        print(f"  ERROR: Timeout starting {framework}")
        return False


def stop_docker_service(framework: str) -> None:
    """Stop a framework Docker container."""
    container_name = f"benchmark-{framework}"
    try:
        cmd = ["docker", "stop", container_name]
        subprocess.run(cmd, capture_output=True, timeout=30)
        print(f"  Stopped {container_name} container")
    except subprocess.TimeoutExpired:
        print(f"  WARNING: Timeout stopping {container_name}")


def main():
    parser = argparse.ArgumentParser(description="Framework Benchmark Runner")
    parser.add_argument("--host", default="127.0.0.1", help="Server host")
    parser.add_argument("-c", "--connections", type=int, default=100, help="Concurrent connections")
    parser.add_argument("-d", "--duration", type=int, default=10, help="Duration per endpoint (seconds)")
    parser.add_argument("-w", "--warmup", type=int, default=1000, help="Warmup requests")
    parser.add_argument("-r", "--runs", type=int, default=3, help="Runs per endpoint")
    parser.add_argument("-o", "--results-dir", default="results", help="Directory for framework JSON files")
    parser.add_argument("--override", action="store_true",
                        help="Delete all existing result files before saving (default: merge with existing)")
    parser.add_argument("--frameworks", nargs="+", choices=list(FRAMEWORKS.keys()),
                        default=list(FRAMEWORKS.keys()), help="Frameworks to benchmark")
    parser.add_argument("--docker", action="store_true", help="Monitor Docker containers instead of local processes")
    parser.add_argument("--sequential", action="store_true",
                        help="Run frameworks one at a time (start, benchmark, stop). Recommended for fair comparison.")

    args = parser.parse_args()

    config = BenchConfig(
        connections=args.connections,
        duration=args.duration,
        warmup_requests=args.warmup,
        runs=args.runs,
    )

    # Track active containers for cleanup on interrupt
    active_containers = []

    def cleanup_containers(signum=None, frame=None):
        """Clean up active containers on interrupt."""
        if active_containers:
            print("\n\nCleaning up containers...")
            for container in active_containers:
                try:
                    subprocess.run(["docker", "rm", "-f", container],
                                 capture_output=True, timeout=10)
                    print(f"  Removed {container}")
                except Exception as e:
                    print(f"  Failed to remove {container}: {e}")
        sys.exit(0)

    # Register signal handlers for graceful cleanup
    if args.docker and args.sequential:
        signal.signal(signal.SIGINT, cleanup_containers)
        signal.signal(signal.SIGTERM, cleanup_containers)

    print("Framework Benchmark")
    print("=" * 60)
    print(f"Host: {args.host}")
    print(f"Connections: {config.connections}")
    print(f"Duration: {config.duration}s")
    print(f"Warmup: {config.warmup_requests} requests")
    print(f"Runs: {config.runs}")
    print(f"Frameworks: {', '.join(args.frameworks)}")
    print(f"Mode: {'Docker' if args.docker else 'Local'}")
    print(f"Execution: {'Sequential (one at a time)' if args.sequential else 'Parallel (all running)'}")
    print(f"Output: {args.results_dir}/ ({'override all' if args.override else 'merge'})")

    all_results = []

    # Docker container name mapping
    container_names = {
        # Bolt (9001)
        "bolt": "benchmark-bolt",

        # DRF variants (902X)
        "drf-uvicorn": "benchmark-drf-uvicorn",
        "drf-granian": "benchmark-drf-granian",
        "drf-gunicorn": "benchmark-drf-gunicorn",

        # FastAPI variants (903X)
        "fastapi-uvicorn": "benchmark-fastapi-uvicorn",
        "fastapi-granian": "benchmark-fastapi-granian",

        # Litestar variants (904X)
        "litestar-uvicorn": "benchmark-litestar-uvicorn",
        "litestar-granian": "benchmark-litestar-granian",

        # Ninja variants (905X)
        "ninja-uvicorn": "benchmark-ninja-uvicorn",
        "ninja-granian": "benchmark-ninja-granian",
    }

    for framework in args.frameworks:
        fw_config = FRAMEWORKS[framework]

        # Sequential mode: start framework, benchmark, stop
        if args.sequential and args.docker:
            print(f"\n{'='*60}")
            print(f"Starting {framework} container...")
            print(f"{'='*60}")

            container_name = container_names.get(framework)
            if not start_docker_service(framework, fw_config["port"], fw_config["prefix"]):
                print(f"Skipping {framework} due to startup failure")
                continue

            # Track active container for cleanup
            if container_name:
                active_containers.append(container_name)
        results = benchmark_framework(
            framework=framework,
            host=args.host,
            port=fw_config["port"],
            prefix=fw_config["prefix"],
            config=config,
            use_docker=args.docker,
            container_name=container_names.get(framework) if args.docker else None,
        )
        all_results.extend(results)

        # Sequential mode: stop framework after benchmarking
        if args.sequential and args.docker:
            print(f"\n{'='*60}")
            print(f"Stopping {framework} container...")
            print(f"{'='*60}")
            stop_docker_service(framework)

            # Remove from active containers tracking
            container_name = container_names.get(framework)
            if container_name and container_name in active_containers:
                active_containers.remove(container_name)

            time.sleep(2)  # Brief pause before next framework

    if all_results:
        # Save results to JSON files
        results_dir = Path(args.results_dir)
        save_results_to_json(all_results, results_dir, override=args.override)

        # Print summary
        print("\n" + "="*60)
        print("Benchmark Complete!")
        print("="*60)
        print(f"Frameworks benchmarked: {len(set(r.framework for r in all_results))}")
        print(f"Total results: {len(all_results)}")
        print("\nTo generate visualizations, run:")
        print(f"  python visualize.py -i {results_dir}")
    else:
        print("\nNo results collected!")


if __name__ == "__main__":
    main()