377 lines
14 KiB
Python
Executable File
377 lines
14 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Orthogonal ImageFind WebDAV upload stress test with isolated cleanup.
|
|
|
|
The default profile is CI-friendly. Pass ``--full`` to execute the requested
|
|
10 MiB / 1 GiB / 5 GiB singles, ten-file batch and fifty 10 MiB-file batch.
|
|
The API token is used for both WebDAV Basic authentication and cleanup.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import base64
|
|
import concurrent.futures
|
|
import hashlib
|
|
import os
|
|
import sys
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from collections.abc import Iterator
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from urllib.parse import quote
|
|
|
|
import httpx
|
|
|
|
MIB = 1024**2
|
|
GIB = 1024**3
|
|
CHUNK = MIB
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Case:
|
|
virtual_path: str
|
|
size: int
|
|
resume: bool = False
|
|
|
|
|
|
@dataclass
|
|
class TimedBody:
|
|
name: str
|
|
size: int
|
|
prefix: bytes = b""
|
|
start: int = 0
|
|
end: int | None = None
|
|
finished_at: float | None = None
|
|
|
|
def __iter__(self) -> Iterator[bytes]:
|
|
yield from blocks(
|
|
self.name,
|
|
self.size,
|
|
prefix=self.prefix,
|
|
start=self.start,
|
|
end=self.end,
|
|
)
|
|
# httpx resumes the iterator after it has handed the final request
|
|
# bytes to the transport. The remaining interval is server response
|
|
# latency, including any accidental post-PUT whole-file work.
|
|
self.finished_at = time.monotonic()
|
|
|
|
|
|
def blocks(
|
|
name: str,
|
|
size: int,
|
|
*,
|
|
prefix: bytes = b"",
|
|
start: int = 0,
|
|
end: int | None = None,
|
|
) -> Iterator[bytes]:
|
|
end = size if end is None else end
|
|
seed = hashlib.sha256(name.encode()).digest()
|
|
block = (seed * (CHUNK // len(seed) + 1))[:CHUNK]
|
|
position = start
|
|
while position < end:
|
|
if position < len(prefix):
|
|
length = min(len(prefix) - position, end - position)
|
|
yield prefix[position : position + length]
|
|
position += length
|
|
continue
|
|
offset = position % CHUNK
|
|
length = min(CHUNK - offset, end - position)
|
|
yield block[offset : offset + length]
|
|
position += length
|
|
|
|
|
|
def encoded_path(path: str) -> str:
|
|
return "/".join(quote(part, safe="") for part in path.split("/"))
|
|
|
|
|
|
class StressRun:
|
|
def __init__(
|
|
self,
|
|
base_url: str,
|
|
token: str,
|
|
full: bool,
|
|
large: bool,
|
|
workers: int,
|
|
timeout: float,
|
|
probe_timeout: float,
|
|
media_prefix: bytes = b"",
|
|
):
|
|
self.base_url = base_url.rstrip("/")
|
|
self.token = token
|
|
self.full = full
|
|
self.large = large
|
|
self.workers = workers
|
|
self.timeout = timeout
|
|
self.probe_timeout = probe_timeout
|
|
self.media_prefix = media_prefix
|
|
self.run_id = f"ifstress-{uuid.uuid4().hex[:10]}"
|
|
basic = base64.b64encode(f"imagefind:{token}".encode()).decode()
|
|
self.dav_headers = {"Authorization": f"Basic {basic}"}
|
|
self.api_headers = {"Authorization": f"Bearer {token}"}
|
|
self.upload_ids: set[str] = set()
|
|
self.request_timeout = httpx.Timeout(connect=15, write=60, read=120, pool=30)
|
|
self.probe_latencies: list[float] = []
|
|
self.probe_errors: list[str] = []
|
|
|
|
def probe(self, stopping: threading.Event) -> None:
|
|
timeout = httpx.Timeout(connect=3, write=3, read=self.probe_timeout, pool=3)
|
|
while not stopping.is_set():
|
|
started = time.monotonic()
|
|
try:
|
|
response = httpx.get(f"{self.base_url}/api/v1/status", timeout=timeout)
|
|
response.raise_for_status()
|
|
self.probe_latencies.append(time.monotonic() - started)
|
|
except Exception as exc:
|
|
if len(self.probe_errors) < 10:
|
|
self.probe_errors.append(f"{type(exc).__name__}: {exc}")
|
|
stopping.wait(0.25)
|
|
|
|
def cases(self) -> list[Case]:
|
|
if self.large:
|
|
return [
|
|
Case(f"{self.run_id}/large-{index}.mp4", 256 * MIB)
|
|
for index in range(4)
|
|
]
|
|
single_sizes = [10 * MIB, GIB, 5 * GIB] if self.full else [10 * MIB]
|
|
cases = [
|
|
Case(f"{self.run_id}/single-{size}.mp4", size, resume=index == 0)
|
|
for index, size in enumerate(single_sizes)
|
|
]
|
|
cases.append(Case(f"root-{self.run_id}.mp4", 10 * MIB if self.full else MIB))
|
|
for index in range(10):
|
|
depth = (1, 5, 10)[index % 3]
|
|
groups = "/".join(f"level-{level}" for level in range(1, depth + 1))
|
|
size = 10 * MIB if self.full else MIB
|
|
cases.append(Case(f"{self.run_id}/{groups}/{self.run_id}-ten-{index:02d}.mp4", size))
|
|
for index in range(50):
|
|
size = 10 * MIB if self.full else 256 * 1024
|
|
cases.append(Case(f"{self.run_id}/fifty/{self.run_id}-fifty-{index:02d}.mp4", size))
|
|
return cases
|
|
|
|
def put(self, case: Case) -> tuple[str, int, str | None, float]:
|
|
url = f"{self.base_url}/webdav/{encoded_path(case.virtual_path)}"
|
|
headers = dict(self.dav_headers)
|
|
upload_id = None
|
|
response_body: TimedBody
|
|
with httpx.Client(timeout=self.request_timeout) as client:
|
|
if case.resume:
|
|
split = case.size // 2
|
|
first_body = TimedBody(
|
|
case.virtual_path, case.size, self.media_prefix, end=split
|
|
)
|
|
first = client.put(
|
|
url,
|
|
headers={
|
|
**headers,
|
|
"Content-Length": str(split),
|
|
"Content-Range": f"bytes 0-{split - 1}/{case.size}",
|
|
},
|
|
content=first_body,
|
|
)
|
|
first.raise_for_status()
|
|
if first.headers.get("upload-offset") != str(split):
|
|
raise RuntimeError(f"{case.virtual_path}: first offset was not persisted")
|
|
head = client.head(url, headers=headers)
|
|
head.raise_for_status()
|
|
if head.headers.get("upload-offset") != str(split):
|
|
raise RuntimeError(f"{case.virtual_path}: HEAD did not report the partial offset")
|
|
response_body = TimedBody(
|
|
case.virtual_path, case.size, self.media_prefix, start=split
|
|
)
|
|
response = client.put(
|
|
url,
|
|
headers={
|
|
**headers,
|
|
"Content-Length": str(case.size - split),
|
|
"Content-Range": f"bytes {split}-{case.size - 1}/{case.size}",
|
|
},
|
|
content=response_body,
|
|
)
|
|
else:
|
|
response_body = TimedBody(case.virtual_path, case.size, self.media_prefix)
|
|
response = client.put(
|
|
url,
|
|
headers={**headers, "Content-Length": str(case.size)},
|
|
content=response_body,
|
|
)
|
|
response.raise_for_status()
|
|
upload_id = response.headers.get("x-imagefind-upload-id")
|
|
response_latency = time.monotonic() - (response_body.finished_at or time.monotonic())
|
|
return case.virtual_path, case.size, upload_id, response_latency
|
|
|
|
def wait(self, expected: int) -> None:
|
|
deadline = time.monotonic() + self.timeout
|
|
while time.monotonic() < deadline:
|
|
response = httpx.get(
|
|
f"{self.base_url}/api/v1/uploads?limit=500",
|
|
headers=self.api_headers,
|
|
timeout=30,
|
|
)
|
|
response.raise_for_status()
|
|
rows = [row for row in response.json() if row.get("id") in self.upload_ids]
|
|
failures = [row for row in rows if row.get("status") == "failed"]
|
|
if failures:
|
|
raise RuntimeError(f"upload failures: {failures}")
|
|
if len(rows) == expected and all(row.get("status") == "completed" for row in rows):
|
|
return
|
|
time.sleep(2)
|
|
raise TimeoutError("uploads did not complete before the stress timeout")
|
|
|
|
def upload_rows(self) -> list[dict]:
|
|
response = httpx.get(
|
|
f"{self.base_url}/api/v1/uploads?limit=500",
|
|
headers=self.api_headers,
|
|
timeout=self.request_timeout,
|
|
)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
|
|
def verify_dedup(self, case: Case) -> None:
|
|
url = f"{self.base_url}/webdav/{encoded_path(case.virtual_path)}"
|
|
before = {row["id"] for row in self.upload_rows()}
|
|
body = TimedBody(case.virtual_path, case.size, self.media_prefix)
|
|
response = httpx.put(
|
|
url,
|
|
headers={
|
|
**self.dav_headers,
|
|
"Content-Length": str(case.size),
|
|
},
|
|
content=body,
|
|
timeout=self.request_timeout,
|
|
)
|
|
response.raise_for_status()
|
|
if response.headers.get("x-imagefind-deduplicated") != "true":
|
|
raise RuntimeError("full retry after a simulated lost response was not deduplicated")
|
|
after = {row["id"] for row in self.upload_rows()}
|
|
if after != before:
|
|
raise RuntimeError("full retry created another upload task")
|
|
latency = time.monotonic() - (body.finished_at or time.monotonic())
|
|
print(
|
|
f"deduplicated lost-response retry response_latency={latency:.3f}s",
|
|
flush=True,
|
|
)
|
|
|
|
def cleanup(self) -> None:
|
|
try:
|
|
videos = httpx.get(
|
|
f"{self.base_url}/api/v1/videos?limit=500", headers=self.api_headers, timeout=30
|
|
).json()
|
|
for video in videos:
|
|
if self.run_id in str(video.get("display_name")) or self.run_id in str(video.get("source_key")):
|
|
response = httpx.delete(
|
|
f"{self.base_url}/api/v1/videos/{video['id']}?delete_source=true",
|
|
headers=self.api_headers,
|
|
timeout=60,
|
|
)
|
|
if response.status_code != 404:
|
|
response.raise_for_status()
|
|
collections = httpx.get(
|
|
f"{self.base_url}/api/v1/collections", headers=self.api_headers, timeout=30
|
|
).json()
|
|
for collection in collections:
|
|
if collection.get("name") == self.run_id:
|
|
httpx.delete(
|
|
f"{self.base_url}/api/v1/collections/{collection['id']}",
|
|
headers=self.api_headers,
|
|
timeout=30,
|
|
).raise_for_status()
|
|
trash = httpx.get(
|
|
f"{self.base_url}/api/v1/trash", headers=self.api_headers, timeout=30
|
|
).json()
|
|
for item in trash:
|
|
if self.run_id in str(item.get("display_name")) or self.run_id in str(item.get("original_key")):
|
|
response = httpx.delete(
|
|
f"{self.base_url}/api/v1/trash/{item['id']}",
|
|
headers=self.api_headers,
|
|
timeout=60,
|
|
)
|
|
if response.status_code != 404:
|
|
response.raise_for_status()
|
|
except Exception as exc:
|
|
print(f"cleanup warning: {exc}", file=sys.stderr)
|
|
|
|
def execute(self, keep: bool) -> None:
|
|
cases = self.cases()
|
|
started = time.monotonic()
|
|
probe_stopping = threading.Event()
|
|
probe_thread = threading.Thread(
|
|
target=self.probe,
|
|
args=(probe_stopping,),
|
|
name="imagefind-stress-probe",
|
|
daemon=True,
|
|
)
|
|
probe_thread.start()
|
|
try:
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=self.workers) as executor:
|
|
futures = [executor.submit(self.put, case) for case in cases]
|
|
for future in concurrent.futures.as_completed(futures):
|
|
path, size, upload_id, response_latency = future.result()
|
|
if upload_id:
|
|
self.upload_ids.add(upload_id)
|
|
print(
|
|
f"received {path} ({size} bytes) "
|
|
f"response_latency={response_latency:.3f}s",
|
|
flush=True,
|
|
)
|
|
self.wait(len(self.upload_ids))
|
|
self.verify_dedup(cases[0])
|
|
if self.probe_errors:
|
|
raise RuntimeError(
|
|
"ImageFind API became unreachable during WebDAV transfer: "
|
|
+ "; ".join(self.probe_errors)
|
|
)
|
|
maximum_probe = max(self.probe_latencies, default=0)
|
|
print(
|
|
f"PASS run={self.run_id} files={len(cases)} bytes={sum(case.size for case in cases)} "
|
|
f"seconds={time.monotonic() - started:.1f} max_probe_latency={maximum_probe:.3f}s",
|
|
flush=True,
|
|
)
|
|
finally:
|
|
probe_stopping.set()
|
|
probe_thread.join(timeout=self.probe_timeout + 5)
|
|
if not keep:
|
|
self.cleanup()
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--base-url", required=True, help="ImageFind native URL, e.g. http://nas:8765")
|
|
parser.add_argument("--token", default=os.getenv("IMAGEFIND_WEBDAV_TOKEN"))
|
|
profile = parser.add_mutually_exclusive_group()
|
|
profile.add_argument("--full", action="store_true", help="include 1 GiB, 5 GiB and full-size batches")
|
|
profile.add_argument("--large", action="store_true", help="run four concurrent 256 MiB uploads")
|
|
parser.add_argument("--workers", type=int, default=8)
|
|
parser.add_argument("--timeout", type=float, default=7200)
|
|
parser.add_argument(
|
|
"--probe-timeout",
|
|
type=float,
|
|
default=10,
|
|
help="fail when the API is unreachable for this many seconds during transfers",
|
|
)
|
|
parser.add_argument("--keep", action="store_true", help="keep isolated stress data for inspection")
|
|
parser.add_argument(
|
|
"--media-prefix",
|
|
help="optional valid media file prepended before deterministic padding",
|
|
)
|
|
args = parser.parse_args()
|
|
if not args.token:
|
|
parser.error("--token or IMAGEFIND_WEBDAV_TOKEN is required")
|
|
media_prefix = Path(args.media_prefix).read_bytes() if args.media_prefix else b""
|
|
StressRun(
|
|
args.base_url,
|
|
args.token,
|
|
args.full,
|
|
args.large,
|
|
min(10, max(1, args.workers)),
|
|
args.timeout,
|
|
max(1, args.probe_timeout),
|
|
media_prefix,
|
|
).execute(args.keep)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|