280 lines
11 KiB
Python
280 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""Stress ImageFind WebDAV reception without retaining generated large files."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import threading
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from urllib.parse import quote
|
|
|
|
import httpx
|
|
|
|
BLOCK = bytes(1024 * 1024)
|
|
|
|
|
|
def padded_hash(prefix: bytes, size: int) -> str:
|
|
value = hashlib.sha256(prefix)
|
|
remaining = size - len(prefix)
|
|
while remaining > 0:
|
|
length = min(len(BLOCK), remaining)
|
|
value.update(BLOCK[:length])
|
|
remaining -= length
|
|
return value.hexdigest()
|
|
|
|
|
|
def padded_body(prefix: bytes, size: int):
|
|
yield prefix
|
|
remaining = size - len(prefix)
|
|
while remaining > 0:
|
|
length = min(len(BLOCK), remaining)
|
|
yield BLOCK[:length]
|
|
remaining -= length
|
|
|
|
|
|
def request(client: httpx.Client, method: str, path: str, **kwargs: Any) -> httpx.Response:
|
|
response = client.request(method, path, **kwargs)
|
|
if response.is_error:
|
|
raise RuntimeError(f"{method} {path} -> {response.status_code}: {response.text[:1000]}")
|
|
return response
|
|
|
|
|
|
def encoded(parts: list[str]) -> str:
|
|
return "/webdav/" + "/".join(quote(part, safe="") for part in parts)
|
|
|
|
|
|
def wait_uploads(api: httpx.Client, ids: set[str], timeout: int) -> dict[str, dict[str, Any]]:
|
|
deadline = time.monotonic() + timeout
|
|
terminal: dict[str, dict[str, Any]] = {}
|
|
while time.monotonic() < deadline:
|
|
rows = request(api, "GET", "/api/v1/uploads?limit=500").json()
|
|
by_id = {item["id"]: item for item in rows if item["id"] in ids}
|
|
terminal = {
|
|
key: value
|
|
for key, value in by_id.items()
|
|
if value["status"] in {"completed", "failed", "cancelled"}
|
|
}
|
|
if len(terminal) == len(ids):
|
|
failures = [item for item in terminal.values() if item["status"] != "completed"]
|
|
if failures:
|
|
raise AssertionError(f"committed WebDAV uploads failed: {failures}")
|
|
return terminal
|
|
time.sleep(1)
|
|
raise TimeoutError(f"WebDAV transfers did not finish: {ids - set(terminal)}")
|
|
|
|
|
|
def wait_video(api: httpx.Client, filename: str, timeout: int = 120) -> dict[str, Any]:
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
rows = request(api, "GET", "/api/v1/videos?limit=500").json()
|
|
match = next((item for item in rows if item.get("source_key", "").endswith("/" + filename)), None)
|
|
if match:
|
|
return match
|
|
time.sleep(1)
|
|
raise TimeoutError(f"committed stress video did not appear: {filename}")
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--base-url", required=True)
|
|
parser.add_argument("--token-file", type=Path, required=True)
|
|
parser.add_argument("--run-dir", type=Path, required=True)
|
|
parser.add_argument("--case-id", required=True)
|
|
parser.add_argument("--size-mb", type=int, required=True)
|
|
parser.add_argument("--count", type=int, required=True)
|
|
parser.add_argument("--concurrency", type=int, default=4)
|
|
parser.add_argument("--commit-count", type=int, default=1)
|
|
parser.add_argument("--timeout", type=int, default=1800)
|
|
args = parser.parse_args()
|
|
|
|
state = json.loads((args.run_dir / "state.json").read_text())
|
|
token = args.token_file.read_text().strip()
|
|
prefix = f"{state['run_id']}-{args.case_id}-stress"
|
|
collection = f"{prefix}-collection"
|
|
group = "uploads"
|
|
size = args.size_mb * 1024**2
|
|
fixture = (args.run_dir / "e2e-positive.mp4").read_bytes()
|
|
if size < len(fixture):
|
|
raise ValueError("stress size must be at least the fixture size")
|
|
checksum = padded_hash(fixture, size)
|
|
auth = httpx.BasicAuth("imagefind", token)
|
|
api = httpx.Client(
|
|
base_url=args.base_url.rstrip("/"),
|
|
headers={"Authorization": f"Bearer {token}"},
|
|
timeout=httpx.Timeout(60, connect=10),
|
|
)
|
|
dav = httpx.Client(
|
|
base_url=args.base_url.rstrip("/"),
|
|
auth=auth,
|
|
timeout=httpx.Timeout(args.timeout, connect=15, read=args.timeout, write=args.timeout),
|
|
)
|
|
for parts in ([collection], [collection, group]):
|
|
response = dav.request("MKCOL", encoded(parts))
|
|
if response.status_code not in {201, 405}:
|
|
raise RuntimeError(f"MKCOL {parts} -> {response.status_code}: {response.text[:1000]}")
|
|
|
|
stop = threading.Event()
|
|
samples: list[dict[str, Any]] = []
|
|
|
|
def monitor() -> None:
|
|
with httpx.Client(
|
|
base_url=args.base_url.rstrip("/"),
|
|
headers={"Authorization": f"Bearer {token}"},
|
|
timeout=10,
|
|
) as client:
|
|
while not stop.wait(1):
|
|
started = time.perf_counter()
|
|
try:
|
|
status = client.get("/api/v1/status")
|
|
status.raise_for_status()
|
|
latency = round((time.perf_counter() - started) * 1000, 1)
|
|
resources = client.get("/api/v1/system/resources")
|
|
resources.raise_for_status()
|
|
payload = resources.json()
|
|
samples.append(
|
|
{
|
|
"api_ms": latency,
|
|
"cpu": payload.get("cpu_percent"),
|
|
"memory_available": payload.get("memory_available_bytes"),
|
|
"writer_wait_ms": (payload.get("database") or {}).get("last_wait_ms"),
|
|
"ok": True,
|
|
}
|
|
)
|
|
except Exception as exc:
|
|
samples.append({"ok": False, "error": type(exc).__name__})
|
|
|
|
monitor_thread = threading.Thread(target=monitor, name="imagefind-webdav-monitor", daemon=True)
|
|
monitor_thread.start()
|
|
|
|
files = [f"{prefix}-{index + 1:03d}.mp4" for index in range(args.count)]
|
|
temporary_paths = {name: encoded([collection, group, name + ".part"]) for name in files}
|
|
|
|
def put_one(filename: str) -> dict[str, Any]:
|
|
path = temporary_paths[filename]
|
|
started = time.perf_counter()
|
|
with httpx.Client(
|
|
base_url=args.base_url.rstrip("/"),
|
|
auth=auth,
|
|
timeout=httpx.Timeout(args.timeout, connect=15, read=args.timeout, write=args.timeout),
|
|
) as client:
|
|
response = client.put(
|
|
path,
|
|
content=padded_body(fixture, size),
|
|
headers={
|
|
"Content-Length": str(size),
|
|
"X-Content-SHA256": checksum,
|
|
},
|
|
)
|
|
if response.is_error:
|
|
raise RuntimeError(f"PUT {filename} -> {response.status_code}: {response.text[:1000]}")
|
|
if response.status_code not in {201, 204}:
|
|
raise RuntimeError(f"PUT {filename} returned {response.status_code}")
|
|
return {
|
|
"filename": filename,
|
|
"status": response.status_code,
|
|
"elapsed_ms": round((time.perf_counter() - started) * 1000, 1),
|
|
"offset": int(response.headers.get("upload-offset", 0)),
|
|
}
|
|
|
|
started_all = time.perf_counter()
|
|
results: list[dict[str, Any]] = []
|
|
error: Exception | None = None
|
|
committed: list[tuple[str, str]] = []
|
|
videos: list[dict[str, Any]] = []
|
|
try:
|
|
with ThreadPoolExecutor(max_workers=args.concurrency) as executor:
|
|
futures = {executor.submit(put_one, filename): filename for filename in files}
|
|
for future in as_completed(futures):
|
|
results.append(future.result())
|
|
if any(item["offset"] != size for item in results):
|
|
raise AssertionError("one or more WebDAV PUT responses reported an incomplete offset")
|
|
|
|
for filename in files[: min(args.commit_count, len(files))]:
|
|
source_path = temporary_paths[filename]
|
|
destination = encoded([collection, group, filename])
|
|
moved = request(
|
|
dav,
|
|
"MOVE",
|
|
source_path,
|
|
headers={
|
|
"Destination": args.base_url.rstrip("/") + destination,
|
|
"Overwrite": "F",
|
|
},
|
|
)
|
|
upload_id = moved.headers.get("x-imagefind-upload-id")
|
|
if not upload_id:
|
|
raise AssertionError(f"MOVE omitted upload id for {filename}")
|
|
committed.append((filename, upload_id))
|
|
|
|
if committed:
|
|
wait_uploads(api, {upload_id for _, upload_id in committed}, args.timeout)
|
|
for filename, _ in committed:
|
|
video = wait_video(api, filename)
|
|
ranged = request(
|
|
dav,
|
|
"GET",
|
|
encoded([collection, group, filename]),
|
|
headers={"Range": "bytes=0-63"},
|
|
)
|
|
if ranged.status_code != 206 or ranged.content != fixture[:64]:
|
|
raise AssertionError(f"committed range read failed: {filename}")
|
|
videos.append(video)
|
|
except Exception as exc:
|
|
error = exc
|
|
finally:
|
|
for filename in files[len(committed) :]:
|
|
try:
|
|
dav.delete(temporary_paths[filename])
|
|
except Exception:
|
|
pass
|
|
for video in videos:
|
|
try:
|
|
removed = request(api, "DELETE", f"/api/v1/videos/{video['id']}?delete_source=true").json()
|
|
if removed.get("trash_id"):
|
|
request(api, "DELETE", f"/api/v1/trash/{removed['trash_id']}")
|
|
except Exception:
|
|
pass
|
|
stop.set()
|
|
monitor_thread.join(timeout=15)
|
|
if error is not None:
|
|
raise error
|
|
|
|
elapsed = time.perf_counter() - started_all
|
|
successful_samples = [item for item in samples if item.get("ok")]
|
|
report = {
|
|
"case_id": args.case_id,
|
|
"size_mb": args.size_mb,
|
|
"count": args.count,
|
|
"concurrency": args.concurrency,
|
|
"committed": len(committed),
|
|
"elapsed_seconds": round(elapsed, 2),
|
|
"throughput_mib_s": round(args.size_mb * args.count / max(elapsed, 0.001), 2),
|
|
"put_latency_ms": {
|
|
"minimum": min(item["elapsed_ms"] for item in results),
|
|
"maximum": max(item["elapsed_ms"] for item in results),
|
|
},
|
|
"monitor": {
|
|
"samples": len(samples),
|
|
"failed_samples": sum(1 for item in samples if not item.get("ok")),
|
|
"max_api_ms": max((item["api_ms"] for item in successful_samples), default=0),
|
|
"max_cpu_percent": max((item["cpu"] for item in successful_samples), default=0),
|
|
"min_memory_available_gb": round(
|
|
min((item["memory_available"] for item in successful_samples), default=0) / 1024**3,
|
|
2,
|
|
),
|
|
"max_writer_wait_ms": max((item["writer_wait_ms"] for item in successful_samples), default=0),
|
|
},
|
|
}
|
|
report_path = args.run_dir / f"webdav-stress-{args.case_id}.json"
|
|
report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n")
|
|
print(json.dumps(report, ensure_ascii=False, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|