scripts/benchmark-bundles.py
Rama main · réplica de lectura del repositorio original.
#!/usr/bin/env python3
"""Paired Linux benchmarks on an isolated, real TCP path with fixed RTT.
Requires root, iproute2, two native binaries, and enough disk for one received
fixture. Network namespaces and qdiscs are private to this run. Public routing
and host qdiscs are never modified. Results include independent SHA-256 checks,
temporary-storage samples, socket counts, and the original debug logs.
"""
from bundle_bench_common import manifest, sha256, transfer_secret, stop_process
import argparse
import json
import os
from pathlib import Path
import secrets
import shutil
import statistics
import subprocess
import threading
import time
def command(*args, **kwargs):
return subprocess.run([str(arg) for arg in args], check=True, **kwargs)
def drain_log(stream, logfile, started, markers):
# Timestamp existing debug events when pipe bytes arrive. read1 drains
# available bytes immediately instead of waiting to fill a 64 KiB buffer.
needles = {
"metadata_ready": b"startup milestone file-metadata-ready",
"first_data": b"startup milestone first-decrypted-data-byte",
"data_complete": b"finished receiving!",
"sender_start": b"beginning sending comms",
"sender_complete": b"closing file",
}
tail = b""
while block := stream.read1(64 << 10):
observed = time.monotonic() - started
combined = tail + block
for name, needle in needles.items():
if name not in markers and needle in combined:
markers[name] = observed
tail = combined[-128:]
logfile.write(block)
class Monitor(threading.Thread):
def __init__(self, sender_temp, receiver_temp, destination):
super().__init__(daemon=True)
self.sender_temp, self.receiver_temp, self.destination = (
sender_temp,
receiver_temp,
destination,
)
self.stopping = threading.Event()
self.sender_peak = self.receiver_peak = 0
self.samples = 0
self.pid = None
self.connections_peak = 0
@staticmethod
def sizes(paths):
total = 0
for path in paths:
try:
if path.is_file():
total += path.stat().st_size
except FileNotFoundError:
pass
return total
def run(self):
while not self.stopping.is_set():
self.sender_peak = max(
self.sender_peak,
self.sizes(self.sender_temp.glob("croc-bundle-send-*")),
)
staged = (
file
for directory in self.destination.glob(".croc-bundle-*")
for file in directory.rglob("*")
)
received = self.sizes(
self.receiver_temp.glob("croc-bundle-receive-*")
) + self.sizes(staged)
self.receiver_peak = max(self.receiver_peak, received)
self.samples += 1
if self.pid and self.samples % 20 == 0:
try:
sockets = (
Path(f"/proc/{self.pid}/net/tcp").read_text().splitlines()[1:]
)
connections = sum(
line.split()[3] == "01"
and 32171 <= int(line.split()[2].split(":")[1], 16) <= 32178
for line in sockets
)
self.connections_peak = max(self.connections_peak, connections)
except FileNotFoundError:
pass
self.stopping.wait(0.005)
def transfer(
args,
binary,
source,
expected,
send_namespace,
receive_namespace,
output,
relay,
zipped=False,
):
output.mkdir()
destination = output / "received"
destination.mkdir()
sender_temp, receiver_temp = output / "sender-temp", output / "receiver-temp"
sender_temp.mkdir()
receiver_temp.mkdir()
secret = transfer_secret()
common = [
"--debug",
"--relay",
relay,
"--relay6",
"",
"--yes",
"--overwrite",
"--ignore-stdin",
"--disable-clipboard",
]
if source.is_file():
common.append("--no-compress")
send = [
"ip",
"netns",
"exec",
send_namespace,
binary,
*common,
"send",
"--no-local",
"--transport",
"relay",
]
if zipped:
send.append("--zip")
send.append(source)
receive = [
"ip",
"netns",
"exec",
receive_namespace,
binary,
*common,
"--out",
destination,
]
monitor = Monitor(sender_temp, receiver_temp, destination)
processes = []
monitor.start()
started = time.monotonic()
try:
drains = []
timing_markers = {}
# A real terminal/pipe cannot be fsynced. Giving croc a regular log file
# makes progressbar fsync every render and measures the log disk instead.
with (output / "sender.log").open("wb") as send_log, (
output / "receiver.log"
).open("wb") as receive_log:
try:
for name, cmd, log, temp in [
("sender", send, send_log, sender_temp),
("receiver", receive, receive_log, receiver_temp),
]:
env = dict(
os.environ,
CROC_SECRET=secret,
CROC_CONFIG_DIR=str(output / f"config-{name}"),
TMPDIR=str(temp),
)
processes.append(
subprocess.Popen(
[str(arg) for arg in cmd],
env=env,
cwd=output,
stdin=subprocess.DEVNULL,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
)
)
timing_markers[name] = {}
drain = threading.Thread(
target=drain_log,
args=(processes[-1].stdout, log, started, timing_markers[name]),
daemon=True,
)
drain.start()
drains.append(drain)
if name == "sender":
monitor.pid = processes[-1].pid
time.sleep(0.15)
for process in reversed(processes):
if process.wait(timeout=args.timeout) != 0:
raise RuntimeError(f"Transfer failed; inspect {output}")
finally:
for process in processes:
stop_process(process)
for drain in drains:
drain.join()
elapsed = time.monotonic() - started
monitor.stopping.set()
monitor.join()
actual = manifest(destination / source.name)
if actual != expected:
raise RuntimeError(
f"Independent pathname, size, or SHA-256 mismatch: {output}"
)
if monitor.sender_peak > 32 << 20 or monitor.receiver_peak > 64 << 20:
raise RuntimeError(f"Bundle storage exceeded its bound: {output}")
leaked = (
list(sender_temp.glob("croc-bundle-*"))
+ list(receiver_temp.glob("croc-bundle-*"))
+ list(destination.rglob(".croc-bundle-*"))
)
if leaked:
raise RuntimeError(f"Bundle temporary files leaked: {leaked}")
send_log = (output / "sender.log").read_text()
result = {
"seconds": elapsed,
"binary": str(binary),
"files": len(actual),
"bytes": sum(file["size"] for file in actual.values()),
"zip": zipped,
"bundle_requests": send_log.count("read bundle-request message"),
"sender_peak_bundle_storage": monitor.sender_peak,
"receiver_peak_bundle_storage": monitor.receiver_peak,
"storage_samples": monitor.samples,
"independent_sha256": True,
"established_data_connections_peak": monitor.connections_peak,
"timing_markers": timing_markers,
"bundle_timing": [
line for line in send_log.splitlines() if "bundle timing" in line
]
+ [
line
for line in (output / "receiver.log").read_text().splitlines()
if "bundle timing" in line
],
}
receive_markers = timing_markers["receiver"]
if (
source.is_file()
and "first_data" in receive_markers
and "data_complete" in receive_markers
):
result["data_seconds"] = (
receive_markers["data_complete"] - receive_markers["first_data"]
)
result["data_MBps"] = result["bytes"] / result["data_seconds"] / 1e6
(output / "result.json").write_text(json.dumps(result, indent=2) + "\n")
# Every trial begins with an empty destination; keep evidence, not copies
# of completed video payloads. This also bounds total benchmark storage.
shutil.rmtree(destination)
return result
finally:
monitor.stopping.set()
monitor.join()
for process in processes:
stop_process(process)
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--base", type=Path, required=True)
parser.add_argument("--new", type=Path, required=True)
parser.add_argument("--artifacts", type=Path, required=True)
parser.add_argument("--source-dir", type=Path)
parser.add_argument("--fixtures-dir", type=Path)
parser.add_argument("--video", type=Path)
parser.add_argument(
"--datasets",
nargs="+",
choices=["small", "source", "video"],
default=["small", "source", "video"],
)
parser.add_argument("--repeats", type=int, default=3)
parser.add_argument("--rtt-ms", type=float, default=75)
parser.add_argument("--timeout", type=int, default=300)
parser.add_argument(
"--resume",
action="store_true",
help="Reuse verified cases only when binaries and frozen inputs still match",
)
parser.add_argument(
"--group",
help="Repeat one dataset/direction/mode group while retaining every frozen input check",
)
args = parser.parse_args()
args.artifacts = args.artifacts.resolve()
args.artifacts.mkdir(parents=True, exist_ok=True)
binary_checks = {str(path): sha256(path) for path in [args.base, args.new]}
binary_record = args.artifacts / "binary-sha256.json"
if (
binary_record.exists()
and json.loads(binary_record.read_text()) != binary_checks
):
raise RuntimeError("Benchmark binaries changed; use a fresh artifact directory")
binary_record.write_text(json.dumps(binary_checks, indent=2))
fixtures = (
args.fixtures_dir.resolve()
if args.fixtures_dir
else args.artifacts / "fixtures"
)
fixtures.mkdir(exist_ok=True)
small = fixtures / "small-256"
small.mkdir(exist_ok=True)
for index in range(256):
fixture = small / f"{index:04}.bin"
if not fixture.exists():
fixture.write_bytes(
bytes(
(position * 31 + index) % 251
for position in range(1024 + index % 16)
)
)
sources = {"small": small}
for name, source in [("source", args.source_dir), ("video", args.video)]:
if name not in args.datasets:
continue
if not source:
parser.error(f"--{name if name == 'video' else 'source-dir'} is required")
target = fixtures / ("video.bin" if name == "video" else "source")
if not target.exists():
if source.is_dir():
shutil.copytree(source, target)
else:
shutil.copy2(source, target)
sources[name] = target
manifests = {name: manifest(source) for name, source in sources.items()}
input_record = args.artifacts / "input-manifests.json"
if input_record.exists() and json.loads(input_record.read_text()) != manifests:
raise RuntimeError("Frozen benchmark inputs changed")
input_record.write_text(json.dumps(manifests, indent=2))
identity = secrets.token_hex(3)
left, right = f"croc-bundle-{identity}-a", f"croc-bundle-{identity}-b"
link_a, link_b = f"cb{identity}a", f"cb{identity}b"
created = []
relay = None
results = []
try:
for namespace in [left, right]:
command("ip", "netns", "add", namespace)
created.append(namespace)
command("ip", "-n", namespace, "link", "set", "lo", "up")
command("ip", "link", "add", link_a, "type", "veth", "peer", "name", link_b)
for namespace, link, address in [
(left, link_a, "10.203.0.1/24"),
(right, link_b, "10.203.0.2/24"),
]:
command("ip", "link", "set", link, "netns", namespace)
command("ip", "-n", namespace, "addr", "add", address, "dev", link)
command("ip", "-n", namespace, "link", "set", link, "up")
command(
"ip",
"netns",
"exec",
namespace,
"tc",
"qdisc",
"add",
"dev",
link,
"root",
"netem",
"delay",
f"{args.rtt_ms / 2}ms",
)
ping = subprocess.check_output(
["ip", "netns", "exec", left, "ping", "-c", "5", "10.203.0.2"], text=True
)
(args.artifacts / "controlled-path.txt").write_text(ping)
relay_log = (args.artifacts / "relay.log").open("w")
relay = subprocess.Popen(
[
"ip",
"netns",
"exec",
right,
str(args.new),
"--debug",
"relay",
"--host",
"0.0.0.0",
"--ports",
"32170,32171,32172,32173,32174,32175,32176,32177,32178",
"--source-join-limit",
"10000",
"--room-join-limit",
"10000",
],
stdout=relay_log,
stderr=subprocess.STDOUT,
)
time.sleep(0.5)
for dataset in args.datasets:
source = sources[dataset]
for direction, sender, receiver in [
("up", left, right),
("down", right, left),
]:
for zipped in ([False, True] if dataset == "source" else [False]):
if (
args.group
and args.group
!= f"{dataset}/{direction}/{'zip' if zipped else 'raw'}"
):
continue
for trial in range(args.repeats):
profiles = [("base", args.base), ("new", args.new)]
if trial % 2:
profiles.reverse()
for label, binary in profiles:
case = f"{dataset}-{direction}-{'zip' if zipped else 'raw'}-{trial}-{label}"
output = args.artifacts / case
if args.resume and (output / "result.json").exists():
result = json.loads(
(output / "result.json").read_text()
)
result.update(
dataset=dataset,
direction=direction,
trial=trial,
profile=label,
)
results.append(result)
continue
if args.resume and output.exists():
output.rename(
args.artifacts
/ (case + f"-failed-{time.time_ns()}")
)
result = transfer(
args,
binary,
source,
manifests[dataset],
sender,
receiver,
output,
"10.203.0.2:32170",
zipped,
)
result.update(
dataset=dataset,
direction=direction,
trial=trial,
profile=label,
)
results.append(result)
with (args.artifacts / "results.jsonl").open("a") as stream:
stream.write(json.dumps(result) + "\n")
print(
f"{case}: {result['seconds']:.3f}s, {result['bundle_requests']} bundles, SHA-256 verified",
flush=True,
)
summary = {}
for result in results:
key = f"{result['dataset']}/{result['direction']}/{'zip' if result['zip'] else 'raw'}/{result['profile']}"
summary.setdefault(key, []).append(result["seconds"])
summary = {
key: {
"median": statistics.median(times),
"min": min(times),
"max": max(times),
"trials": len(times),
}
for key, times in summary.items()
}
(args.artifacts / "summary.json").write_text(
json.dumps(summary, indent=2) + "\n"
)
print(json.dumps(summary, indent=2))
finally:
if relay:
relay.terminate()
relay.wait(timeout=10)
for namespace in reversed(created):
subprocess.run(["ip", "netns", "del", namespace], check=False)
subprocess.run(
["ip", "link", "del", link_a],
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
check=False,
)
if __name__ == "__main__":
main()