A tiny keep-alive JSON POST load tester for checking API headroom before overbuilding

pulsepost.py

#!/usr/bin/env python3
"""PulsePost: tiny keep-alive JSON POST load tester.

Example:
    python pulsepost.py \
        --url https://api.example.com/v1/events/value \
        --api-key "$API_KEY" \
        --stat-name api.random_load_test \
        --workers 50 \
        --duration 60
"""

from __future__ import annotations

import argparse
import http.client
import json
import random
import ssl
import statistics
import threading
import time
from collections import Counter
from dataclasses import dataclass
from urllib.parse import ParseResult, urlparse


@dataclass
class Config:
    url: str
    api_key: str
    stat_name: str
    min_value: int
    max_value: int
    timeout: float

    @property
    def parsed_url(self) -> ParseResult:
        return urlparse(self.url)


@dataclass
class Result:
    status: int | None
    latency_ms: float
    error: str | None = None


class KeepAliveClient:
    """Keeps one TCP/TLS connection open per worker unless --no-keep-alive is used."""

    def __init__(self, config: Config) -> None:
        self.config = config
        self.connection: http.client.HTTPConnection | None = None

    def close(self) -> None:
        if self.connection:
            self.connection.close()
            self.connection = None

    def _connect(self) -> http.client.HTTPConnection:
        if self.connection:
            return self.connection

        parsed_url = self.config.parsed_url
        if parsed_url.scheme == "https":
            self.connection = http.client.HTTPSConnection(
                parsed_url.hostname,
                parsed_url.port or 443,
                timeout=self.config.timeout,
                context=ssl.create_default_context(),
            )
        else:
            self.connection = http.client.HTTPConnection(
                parsed_url.hostname,
                parsed_url.port or 80,
                timeout=self.config.timeout,
            )
        return self.connection

    def post_event(self) -> Result:
        started = time.perf_counter()
        try:
            return self._post_event()
        except Exception as error:
            self.close()
            return Result(None, (time.perf_counter() - started) * 1000, str(error))

    def _post_event(self) -> Result:
        # Shape this payload to match your API. This version sends a simple metric event.
        parsed_url = self.config.parsed_url
        payload = {
            "apiKey": self.config.api_key,
            "statName": self.config.stat_name,
            "value": random.randint(self.config.min_value, self.config.max_value),
        }
        data = json.dumps(payload).encode("utf-8")
        path = parsed_url.path or "/"
        if parsed_url.query:
            path = f"{path}?{parsed_url.query}"

        connection = self._connect()
        started = time.perf_counter()
        connection.request(
            "POST",
            path,
            body=data,
            headers={
                "Content-Type": "application/json",
                "Content-Length": str(len(data)),
                "Host": parsed_url.netloc,
            },
        )
        response = connection.getresponse()
        response.read()
        latency_ms = (time.perf_counter() - started) * 1000

        if response.getheader("Connection", "").lower() == "close":
            self.close()

        error = None if 200 <= response.status < 300 else response.reason
        return Result(response.status, latency_ms, error)


def post_event(config: Config) -> Result:
    client = KeepAliveClient(config)
    try:
        return client.post_event()
    finally:
        client.close()


def worker(
    config: Config,
    deadline: float,
    results: list[Result],
    lock: threading.Lock,
    keep_alive: bool,
) -> None:
    client = KeepAliveClient(config)
    while time.monotonic() < deadline:
        result = client.post_event() if keep_alive else post_event(config)
        with lock:
            results.append(result)
    client.close()


def percentile(values: list[float], pct: float) -> float:
    if not values:
        return 0.0
    index = min(round((len(values) - 1) * pct), len(values) - 1)
    return sorted(values)[index]


def print_summary(results: list[Result], elapsed_seconds: float) -> None:
    """Print enough signal for a quick API headroom check: rate, statuses, p50/p95."""

    statuses = Counter(
        result.status if result.status is not None else "error" for result in results
    )
    successes = sum(
        count
        for status, count in statuses.items()
        if isinstance(status, int) and 200 <= status < 300
    )
    failures = len(results) - successes
    latencies = [result.latency_ms for result in results]
    errors = Counter(result.error for result in results if result.error)

    print("\nSummary")
    print("-------")
    print(f"Duration: {elapsed_seconds:.1f}s")
    print(f"Total requests: {len(results)}")
    print(f"Successful requests: {successes}")
    print(f"Failed requests: {failures}")
    print(f"Requests/sec: {len(results) / elapsed_seconds:.2f}")
    print(f"Status counts: {dict(statuses)}")

    if latencies:
        print(f"Latency min: {min(latencies):.1f}ms")
        print(f"Latency avg: {statistics.fmean(latencies):.1f}ms")
        print(f"Latency p50: {percentile(latencies, 0.50):.1f}ms")
        print(f"Latency p95: {percentile(latencies, 0.95):.1f}ms")
        print(f"Latency max: {max(latencies):.1f}ms")

    if errors:
        print(f"Top errors: {errors.most_common(5)}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="PulsePost: tiny keep-alive JSON POST load tester."
    )
    parser.add_argument("--url", required=True, help="JSON POST endpoint to test.")
    parser.add_argument("--api-key", required=True, help="API key sent as payload.apiKey.")
    parser.add_argument("--stat-name", required=True, help="Stat name sent as payload.statName.")
    parser.add_argument("--duration", type=int, default=60, help="Run duration in seconds.")
    parser.add_argument("--workers", type=int, default=5, help="Concurrent request workers.")
    parser.add_argument("--timeout", type=float, default=10.0, help="Request timeout in seconds.")
    parser.add_argument("--min-value", type=int, default=1, help="Minimum random value.")
    parser.add_argument("--max-value", type=int, default=50, help="Maximum random value.")
    parser.add_argument(
        "--no-keep-alive",
        action="store_true",
        help="Open a fresh TCP/TLS connection for every request.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    config = Config(
        url=args.url,
        api_key=args.api_key,
        stat_name=args.stat_name,
        min_value=args.min_value,
        max_value=args.max_value,
        timeout=args.timeout,
    )
    results: list[Result] = []
    lock = threading.Lock()
    deadline = time.monotonic() + args.duration
    started = time.perf_counter()
    keep_alive = not args.no_keep_alive

    threads = [
        threading.Thread(
            target=worker,
            args=(config, deadline, results, lock, keep_alive),
            daemon=True,
        )
        for _ in range(args.workers)
    ]
    for thread in threads:
        thread.start()
    for thread in threads:
        thread.join()

    print_summary(results, time.perf_counter() - started)


if __name__ == "__main__":
    main()