184 lines
8.0 KiB
Python
Executable File
184 lines
8.0 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Concurrent persistent-connection HTTP load + optional Linux RustProxy RSS sampler.
|
|
|
|
This script issues REAL HTTP requests only when pointed at a RUNNING service.
|
|
A Python test origin can verify the load generator but cannot validate RustProxy
|
|
memory. Use an isolated staging instance, not a production endpoint by default.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from collections import Counter
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
import http.client
|
|
import json
|
|
from pathlib import Path
|
|
import ssl
|
|
import threading
|
|
import time
|
|
from urllib.parse import urlsplit
|
|
|
|
|
|
def proc_sample(pid: int) -> dict | None:
|
|
path = Path(f"/proc/{pid}")
|
|
try:
|
|
status = (path / "status").read_text(encoding="utf-8")
|
|
counters: dict[str, int] = {}
|
|
for row in status.splitlines():
|
|
if row.startswith(("VmRSS:", "VmHWM:", "Threads:")):
|
|
name, value = row.split(":", 1)
|
|
counters[name.lower()] = int(value.split()[0]) * (1024 if name != "Threads" else 1)
|
|
counters["open_fds"] = len(list((path / "fd").iterdir()))
|
|
return counters
|
|
except (OSError, ValueError, IndexError):
|
|
return None
|
|
|
|
|
|
def run_load(
|
|
*, url: str, requests: int, concurrency: int, unique_paths: int,
|
|
timeout: float, host_header: str | None = None, pid: int | None = None,
|
|
sample_interval: float = 1.0, cooldown: float = 0.0, insecure: bool = False,
|
|
) -> dict:
|
|
parsed = urlsplit(url)
|
|
if parsed.scheme not in {"http", "https"} or not parsed.hostname or parsed.username or parsed.password:
|
|
raise ValueError("URL must be http(s)://host[:port]/path without credentials")
|
|
if requests < 1 or not (1 <= concurrency <= 2048) or unique_paths < 0:
|
|
raise ValueError("requests >= 1, concurrency 1..2048 and unique_paths >= 0 are required")
|
|
if not (0 < timeout <= 120) or sample_interval <= 0 or cooldown < 0:
|
|
raise ValueError("invalid timeout/sampling/cooldown")
|
|
base_target = parsed.path or "/"
|
|
if parsed.query:
|
|
base_target += "?" + parsed.query
|
|
if unique_paths and parsed.query:
|
|
raise ValueError("unique-path mode expects a URL without query parameters")
|
|
port = parsed.port or (443 if parsed.scheme == "https" else 80)
|
|
tls = ssl._create_unverified_context() if insecure else ssl.create_default_context()
|
|
connection_type = http.client.HTTPSConnection if parsed.scheme == "https" else http.client.HTTPConnection
|
|
headers = {"user-agent": "RustProxy-Memory-Stress/0.31.4", "accept": "*/*"}
|
|
if host_header:
|
|
headers["host"] = host_header
|
|
|
|
sampler_stop = threading.Event()
|
|
trace: list[dict] = []
|
|
started = time.monotonic()
|
|
def sampler() -> None:
|
|
while not sampler_stop.is_set():
|
|
point = proc_sample(pid) if pid is not None else None
|
|
if point is not None:
|
|
trace.append({"elapsed_seconds": round(time.monotonic() - started, 2), **point})
|
|
sampler_stop.wait(sample_interval)
|
|
|
|
sampler_thread = threading.Thread(target=sampler, daemon=True)
|
|
if pid is not None:
|
|
sampler_thread.start()
|
|
|
|
index_lock = threading.Lock()
|
|
cursor = 0
|
|
def next_index() -> int | None:
|
|
nonlocal cursor
|
|
with index_lock:
|
|
if cursor >= requests:
|
|
return None
|
|
index = cursor
|
|
cursor += 1
|
|
return index
|
|
|
|
def worker() -> tuple[Counter, int]:
|
|
statuses: Counter = Counter()
|
|
total_bytes = 0
|
|
connection = None
|
|
while True:
|
|
index = next_index()
|
|
if index is None:
|
|
break
|
|
if unique_paths:
|
|
target = base_target.rstrip("/") + f"/memory-{index % unique_paths}"
|
|
else:
|
|
target = base_target
|
|
try:
|
|
if connection is None:
|
|
options = {"timeout": timeout}
|
|
if parsed.scheme == "https":
|
|
options["context"] = tls
|
|
connection = connection_type(parsed.hostname, port, **options)
|
|
connection.request("GET", target, headers=headers)
|
|
response = connection.getresponse()
|
|
statuses[str(response.status)] += 1
|
|
while chunk := response.read(64 * 1024):
|
|
total_bytes += len(chunk)
|
|
if response.will_close:
|
|
connection.close()
|
|
connection = None
|
|
except (OSError, TimeoutError, http.client.HTTPException) as error:
|
|
statuses["error:" + type(error).__name__] += 1
|
|
if connection is not None:
|
|
connection.close()
|
|
connection = None
|
|
if connection is not None:
|
|
connection.close()
|
|
return statuses, total_bytes
|
|
|
|
try:
|
|
with ThreadPoolExecutor(max_workers=concurrency) as executor:
|
|
results = list(executor.map(lambda _: worker(), range(concurrency)))
|
|
traffic_end = time.monotonic()
|
|
if cooldown:
|
|
time.sleep(cooldown)
|
|
end = time.monotonic()
|
|
finally:
|
|
sampler_stop.set()
|
|
if pid is not None:
|
|
sampler_thread.join(timeout=sample_interval + 1)
|
|
|
|
statuses: Counter = Counter()
|
|
received = 0
|
|
for codes, response_bytes in results:
|
|
statuses.update(codes)
|
|
received += response_bytes
|
|
report = {
|
|
"target": f"{parsed.scheme}://{parsed.hostname}:{port}{base_target}",
|
|
"requests_sent": sum(statuses.values()),
|
|
"concurrency": concurrency,
|
|
"unique_paths": unique_paths,
|
|
"statuses": dict(sorted(statuses.items())),
|
|
"response_bytes_read": received,
|
|
"traffic_seconds": round(traffic_end - started, 3),
|
|
"requests_per_second": round(requests / max(0.001, traffic_end - started), 1),
|
|
"wall_seconds": round(end - started, 3),
|
|
"sampler_pid": pid,
|
|
"rss_start_mib": round(trace[0]["vmrss"] / (1024 * 1024), 2) if trace else None,
|
|
"rss_peak_mib": round(max(t["vmrss"] for t in trace) / (1024 * 1024), 2) if trace else None,
|
|
"rss_end_mib": round(trace[-1]["vmrss"] / (1024 * 1024), 2) if trace else None,
|
|
"open_fds_peak": max((t["open_fds"] for t in trace), default=None),
|
|
"threads_peak": max((t.get("threads", 0) for t in trace), default=None),
|
|
"rss_samples": trace,
|
|
"note": "Process RSS values apply only to --pid. No PID means no RustProxy memory measurement.",
|
|
}
|
|
return report
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--url", required=True, help="Staging URL (use your configured vhost route)")
|
|
parser.add_argument("--host-header", help="HTTP Host for a named vhost")
|
|
parser.add_argument("--pid", type=int, help="Linux PID of the RustProxy process to sample from /proc")
|
|
parser.add_argument("--requests", type=int, default=100_000)
|
|
parser.add_argument("--concurrency", type=int, default=64)
|
|
parser.add_argument("--unique-paths", type=int, default=0, help="Cycle this many distinct URL paths; 0 repeats the same path")
|
|
parser.add_argument("--timeout", type=float, default=8.0)
|
|
parser.add_argument("--sample-interval", type=float, default=1.0)
|
|
parser.add_argument("--cooldown", type=float, default=30.0, help="Observe RSS after traffic stops")
|
|
parser.add_argument("--insecure", action="store_true", help="Ignore TLS verification for self-signed STAGING only")
|
|
parser.add_argument("--output", type=Path, help="Optional machine-readable JSON report path")
|
|
args = parser.parse_args()
|
|
report = run_load(url=args.url, host_header=args.host_header, pid=args.pid, requests=args.requests,
|
|
concurrency=args.concurrency, unique_paths=args.unique_paths, timeout=args.timeout,
|
|
sample_interval=args.sample_interval, cooldown=args.cooldown, insecure=args.insecure)
|
|
print(json.dumps({k: v for k, v in report.items() if k != "rss_samples"}, indent=2))
|
|
if args.output:
|
|
args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|