packmate/stress/network/api_control.py

428 lines
16 KiB
Python

import argparse
import base64
import binascii
import json
import math
import os
import struct
import time
import urllib.error
import urllib.request
from pathlib import Path
from network_load import PAYLOAD_MARKER, build_payload
BASE_URL = os.getenv("STRESS_API_URL", "http://172.31.250.10:65000/api").rstrip("/")
USER = os.getenv("STRESS_API_USER", "stress")
PASSWORD = os.getenv("STRESS_API_PASSWORD", "stress")
AUTHORIZATION = "Basic " + base64.b64encode(f"{USER}:{PASSWORD}".encode()).decode()
EXPECTED_LOCAL_IP = os.getenv("STRESS_TARGET_HOST", "172.31.250.10")
def request(method: str, path: str, body=None, timeout: float = 10):
payload = None if body is None else json.dumps(body).encode()
headers = {"Authorization": AUTHORIZATION}
if payload is not None:
headers["Content-Type"] = "application/json"
req = urllib.request.Request(BASE_URL + path, data=payload, headers=headers, method=method)
with urllib.request.urlopen(req, timeout=timeout) as response:
data = response.read()
return json.loads(data) if data else None
def wait_and_prepare(timeout: int) -> None:
deadline = time.monotonic() + timeout
while True:
try:
request("GET", "/service/")
break
except (urllib.error.URLError, TimeoutError):
if time.monotonic() >= deadline:
raise RuntimeError("Packmate API did not become ready")
time.sleep(1)
services = request("GET", "/service/")
if not any(service["port"] == 18080 for service in services):
request("POST", "/service/", {
"port": 18080,
"name": "network-stress",
"decryptTls": False,
"http": False,
"urldecodeHttpRequests": False,
"mergeAdjacentPackets": False,
"parseWebSockets": False,
})
patterns = request("GET", "/pattern/")
if not any(pattern["name"] == "network-stress-marker" for pattern in patterns):
request("POST", "/pattern/", {
"name": "network-stress-marker",
"value": PAYLOAD_MARKER.hex(),
"color": "#ff00ff",
"searchType": "SUBBYTES",
"searchTarget": "CONTENT",
"directionType": "BOTH",
"actionType": "FIND",
"serviceId": 18080,
})
if not request("GET", "/pcap/started"):
request("POST", "/pcap/start")
time.sleep(2)
def probe(
requests: int,
pause: float,
report: str,
max_p99_ms: float,
stop_file: str | None,
min_requests: int,
max_queue: int,
abort_file: str | None,
) -> None:
latencies = []
failures = []
max_capture_queue = 0
capture_drops = 0
interface_drops = 0
application_drops = 0
body = {"startingFrom": None, "pageSize": 20, "pattern": None, "favorites": False}
for index in range(requests):
if stop_file and index >= min_requests and Path(stop_file).exists():
break
started = time.perf_counter()
try:
request("POST", "/stream/all", body)
latencies.append((time.perf_counter() - started) * 1000)
metrics = request("GET", "/pcap/metrics")
if not metrics["captureOpen"] or not metrics["statsAvailable"]:
raise RuntimeError(f"pcap statistics unavailable: {metrics.get('statsError')}")
max_capture_queue = max(max_capture_queue, metrics["maximumQueuedPackets"])
capture_drops = max(capture_drops, metrics["droppedPackets"])
interface_drops = max(interface_drops, metrics["interfaceDroppedPackets"])
application_drops = max(application_drops, metrics["applicationDroppedPackets"])
if (
metrics["maximumQueuedPackets"] > max_queue
or metrics["droppedPackets"] > 0
or metrics["interfaceDroppedPackets"] > 0
or metrics["applicationDroppedPackets"] > 0
):
raise RuntimeError(
"capture limit exceeded: "
f"maxQueueObserved={metrics['maximumQueuedPackets']} maxQueue={max_queue} "
f"drops={metrics['droppedPackets']} interfaceDrops={metrics['interfaceDroppedPackets']} "
f"applicationDrops={metrics['applicationDroppedPackets']}"
)
except Exception as exc:
failures.append(f"request={index}: {type(exc).__name__}: {exc}")
if abort_file:
Path(abort_file).touch()
break
time.sleep(pause)
latencies.sort()
p99 = latencies[max(0, math.ceil(len(latencies) * 0.99) - 1)] if latencies else math.inf
result = {
"maximumRequests": requests,
"requests": len(latencies) + len(failures),
"successfulRequests": len(latencies),
"failedRequests": len(failures),
"latencyP99Millis": round(p99, 3) if latencies else None,
"latencyMaxMillis": round(max(latencies, default=0), 3),
"maxCaptureQueue": max_capture_queue,
"captureDroppedPackets": capture_drops,
"interfaceDroppedPackets": interface_drops,
"applicationDroppedPackets": application_drops,
"sampleErrors": failures[:20],
"stoppedByWorkloadCompletion": bool(stop_file and Path(stop_file).exists()),
}
Path(report).write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
print(json.dumps(result), flush=True)
if (
failures
or p99 > max_p99_ms
or capture_drops > 0
or interface_drops > 0
or application_drops > 0
or max_capture_queue > max_queue
or len(latencies) < min_requests
or (stop_file and not Path(stop_file).exists())
):
raise SystemExit(1)
def monitor_capture(
report: str,
stop_file: str,
abort_file: str,
max_queue: int,
pause: float,
min_samples: int,
) -> None:
samples = 0
maximum_queue = 0
maximum_drops = 0
maximum_interface_drops = 0
maximum_application_drops = 0
failures = []
last_metrics = None
started_at = int(time.time() * 1000)
while True:
try:
metrics = request("GET", "/pcap/metrics")
last_metrics = metrics
samples += 1
maximum_queue = max(maximum_queue, metrics["maximumQueuedPackets"])
maximum_drops = max(maximum_drops, metrics["droppedPackets"])
maximum_interface_drops = max(maximum_interface_drops, metrics["interfaceDroppedPackets"])
maximum_application_drops = max(
maximum_application_drops, metrics["applicationDroppedPackets"]
)
if not metrics["captureOpen"] or not metrics["statsAvailable"]:
raise RuntimeError(f"pcap statistics unavailable: {metrics.get('statsError')}")
if (
maximum_queue > max_queue
or maximum_drops > 0
or maximum_interface_drops > 0
or maximum_application_drops > 0
):
raise RuntimeError(
"capture limit exceeded: "
f"maxQueue={maximum_queue} limit={max_queue} "
f"drops={maximum_drops} interfaceDrops={maximum_interface_drops} "
f"applicationDrops={maximum_application_drops}"
)
except Exception as exc:
failures.append(f"{type(exc).__name__}: {exc}")
Path(abort_file).touch()
break
if samples >= min_samples and Path(stop_file).exists():
break
time.sleep(pause)
result = {
"startedAtUnixMillis": started_at,
"finishedAtUnixMillis": int(time.time() * 1000),
"samples": samples,
"maxCaptureQueue": maximum_queue,
"captureDroppedPackets": maximum_drops,
"interfaceDroppedPackets": maximum_interface_drops,
"applicationDroppedPackets": maximum_application_drops,
"queueLimit": max_queue,
"stoppedByWorkloadCompletion": Path(stop_file).exists(),
"abortedWorkload": Path(abort_file).exists(),
"lastMetrics": last_metrics,
"sampleErrors": failures[:20],
}
Path(report).write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
print(json.dumps(result), flush=True)
if (
failures
or samples < min_samples
or not Path(stop_file).exists()
or maximum_queue > max_queue
or maximum_drops > 0
or maximum_interface_drops > 0
or maximum_application_drops > 0
):
raise SystemExit(1)
def capture_metrics(report: str, require_empty_queue: bool, wait_seconds: int, max_queue: int | None) -> None:
deadline = time.monotonic() + wait_seconds
while True:
metrics = request("GET", "/pcap/metrics")
if not require_empty_queue or (metrics["queuedPackets"] == 0 and metrics["activeProcessors"] == 0):
break
if time.monotonic() >= deadline:
break
time.sleep(1)
Path(report).write_text(json.dumps(metrics, indent=2) + "\n", encoding="utf-8")
print(json.dumps(metrics), flush=True)
if (
not metrics["captureOpen"]
or not metrics["statsAvailable"]
or metrics["receivedPackets"] <= 0
or metrics["droppedPackets"] > 0
or metrics["interfaceDroppedPackets"] > 0
or metrics["applicationDroppedPackets"] > 0
or (max_queue is not None and metrics["maximumQueuedPackets"] > max_queue)
):
raise SystemExit(1)
if require_empty_queue and (metrics["queuedPackets"] != 0 or metrics["activeProcessors"] != 0):
raise SystemExit(1)
def parse_frames(content: bytes) -> list[bytes]:
frames = []
offset = 0
while offset < len(content):
if offset + 4 > len(content):
raise ValueError("truncated frame header")
size = struct.unpack_from("!I", content, offset)[0]
offset += 4
if size <= 0 or offset + size > len(content):
raise ValueError(f"invalid frame size {size}")
frames.append(content[offset:offset + size])
offset += size
return frames
def get_all_packets(stream_id: int) -> list[dict]:
packets = []
cursor = None
page_size = 1000
while True:
page = request("POST", f"/packet/{stream_id}", {
"startingFrom": cursor,
"pageSize": page_size,
})
packets.extend(page)
if len(page) < page_size:
return packets
new_cursor = page[-1]["id"]
if new_cursor == cursor:
raise ValueError("packet pagination cursor did not advance")
cursor = new_cursor
def verify_payloads(
run_id: int,
sample_size: int,
frames_per_stream: int,
payload_bytes: int,
report: str,
) -> None:
streams = request("POST", "/stream/all", {
"startingFrom": None,
"pageSize": sample_size,
"pattern": None,
"favorites": False,
})
failures = []
connection_ids = set()
verified = 0
for stream in streams:
try:
packets = get_all_packets(stream["id"])
incoming = b"".join(base64.b64decode(packet["content"]) for packet in packets if packet["incoming"])
outgoing = b"".join(base64.b64decode(packet["content"]) for packet in packets if not packet["incoming"])
incoming_frames = parse_frames(incoming)
outgoing_frames = parse_frames(outgoing)
if len(incoming_frames) != frames_per_stream or len(outgoing_frames) != frames_per_stream:
raise ValueError(
f"frame count incoming={len(incoming_frames)} outgoing={len(outgoing_frames)}"
)
identities = []
for sequence, payload in enumerate(incoming_frames):
if len(payload) < 24:
raise ValueError("short request payload")
frame_run, connection, shard, frame_sequence = struct.unpack("!QQII", payload[:24])
if frame_run != run_id or frame_sequence != sequence:
raise ValueError(
f"unexpected marker run={frame_run} connection={connection} sequence={frame_sequence}"
)
identities.append((connection, shard))
expected = build_payload(run_id, connection, shard, sequence, payload_bytes)
if payload != expected:
raise ValueError(f"request payload differs at sequence {sequence}")
if outgoing_frames[sequence] != expected:
raise ValueError(f"response payload differs at sequence {sequence}")
if len(set(identities)) != 1:
raise ValueError("one captured TCP stream contains multiple generator connections")
if identities[0] in connection_ids:
raise ValueError("duplicate generator connection across captured streams")
for packet in packets:
if not packet.get("sourceIp") or not packet.get("destinationIp"):
raise ValueError("missing packet IP metadata")
local_ip = packet["destinationIp"] if packet["incoming"] else packet["sourceIp"]
if local_ip != EXPECTED_LOCAL_IP:
raise ValueError(f"unexpected local IP {local_ip}")
connection_ids.add(identities[0])
verified += 1
except (ValueError, KeyError, struct.error, binascii.Error) as exc:
failures.append(f"stream={stream.get('id')}: {type(exc).__name__}: {exc}")
result = {
"runId": run_id,
"requestedSample": sample_size,
"returnedStreams": len(streams),
"verifiedStreams": verified,
"failedStreams": len(failures),
"sampleErrors": failures[:50],
}
Path(report).write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
print(json.dumps(result), flush=True)
if len(streams) < sample_size or failures:
raise SystemExit(1)
def main() -> None:
parser = argparse.ArgumentParser()
subparsers = parser.add_subparsers(dest="command", required=True)
prepare = subparsers.add_parser("prepare")
prepare.add_argument("--timeout", type=int, default=300)
latency = subparsers.add_parser("probe")
latency.add_argument("--requests", type=int, default=300)
latency.add_argument("--pause", type=float, default=0.1)
latency.add_argument("--report", required=True)
latency.add_argument("--max-p99-ms", type=float, default=500)
latency.add_argument("--stop-file")
latency.add_argument("--min-requests", type=int, default=1)
latency.add_argument("--max-queue", type=int, required=True)
latency.add_argument("--abort-file")
monitor = subparsers.add_parser("monitor")
monitor.add_argument("--report", required=True)
monitor.add_argument("--stop-file", required=True)
monitor.add_argument("--abort-file", required=True)
monitor.add_argument("--max-queue", type=int, required=True)
monitor.add_argument("--pause", type=float, default=0.25)
monitor.add_argument("--min-samples", type=int, default=2)
metrics = subparsers.add_parser("metrics")
metrics.add_argument("--report", required=True)
metrics.add_argument("--require-empty-queue", action="store_true")
metrics.add_argument("--wait-seconds", type=int, default=0)
metrics.add_argument("--max-queue", type=int)
verify = subparsers.add_parser("verify")
verify.add_argument("--run-id", type=int, required=True)
verify.add_argument("--sample-size", type=int, default=200)
verify.add_argument("--frames", type=int, required=True)
verify.add_argument("--payload-bytes", type=int, required=True)
verify.add_argument("--report", required=True)
args = parser.parse_args()
if args.command == "prepare":
wait_and_prepare(args.timeout)
elif args.command == "probe":
probe(
args.requests,
args.pause,
args.report,
args.max_p99_ms,
args.stop_file,
args.min_requests,
args.max_queue,
args.abort_file,
)
elif args.command == "monitor":
monitor_capture(
args.report,
args.stop_file,
args.abort_file,
args.max_queue,
args.pause,
args.min_samples,
)
elif args.command == "metrics":
capture_metrics(args.report, args.require_empty_queue, args.wait_seconds, args.max_queue)
else:
verify_payloads(args.run_id, args.sample_size, args.frames, args.payload_bytes, args.report)
if __name__ == "__main__":
main()