iron/scripts/real_openlist_chaos_acceptance.py

from __future__ import annotations

import hashlib
import os
import subprocess
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from scripts.real_openlist_acceptance import AcceptanceConfig, IronClient, header_value


def main() -> int:
    config = AcceptanceConfig.from_env()
    client = IronClient(config)
    container = os.getenv("OPENLIST_DOCKER_CONTAINER", "openlist")
    file_class = os.getenv("IRON_CHAOS_FILE_CLASS", "other")
    size_bytes = int(os.getenv("IRON_CHAOS_FILE_SIZE_BYTES", str(8 * 1024 * 1024)))
    flap_count = int(os.getenv("IRON_CHAOS_FLAP_COUNT", "2"))

    print(f"[1/10] logging into Iron at {config.iron_base_url}")
    client.login()

    print(f"[2/10] ensuring backend {config.backend_name!r}")
    backend = client.ensure_webdav_backend()
    healthy = client.check_backend(backend["id"])
    if healthy["status"] != "healthy":
        raise RuntimeError(f"backend unhealthy before chaos test: {healthy}")

    original_policy = next(policy for policy in client.list_policies() if policy["file_class"] == file_class)
    client.upsert_policy(
        file_class,
        require_local=True,
        stable_replica_count=1,
        opportunistic_replica_count=0,
        preferred_backend_ids=[backend["id"]],
        excluded_backend_ids=[],
        max_non_local_size_bytes=None,
    )

    try:
        print(f"[3/10] uploading large file ({size_bytes} bytes)")
        content = build_deterministic_bytes(size_bytes)
        filename = f"chaos-large-{int(time.time())}.bin"
        file_id = client.upload_bytes(filename=filename, content=content, directory_id=config.directory_id)

        print("[4/10] waiting for local + remote replicas")
        detail = client.wait_for_ready_replica_set(file_id, {"bkd_local_default", backend["id"]})
        print(f"       ready replicas: {detail['ready_replica_backend_ids']}")

        print("[5/10] validating ranged stream reads")
        validate_stream_ranges(client, file_id, content)
        print("       ranged stream checks passed")

        print("[6/10] forcing remote recovery after deleting local physical file")
        local_replica = next(replica for replica in detail["replicas"] if replica["backend_id"] == "bkd_local_default")
        local_base_path = local_backend_base_path(client)
        local_file_path = Path(local_base_path) / local_replica["storage_key"]
        cache_file_path = default_cache_root() / local_replica["storage_key"]
        if local_file_path.exists():
            local_file_path.unlink()
        if cache_file_path.exists():
            cache_file_path.unlink()
        recovered = client.download_file(file_id)
        if recovered != content:
            raise RuntimeError("remote recovery bytes do not match original content")
        print(f"       remote recovery restored cache at {cache_file_path}")

        print(f"[7/10] flapping OpenList {flap_count} times")
        for index in range(flap_count):
            docker(["stop", container])
            wait_for_unhealthy(client, backend["id"])
            docker(["start", container])
            wait_for_healthy(client, backend["id"])
            print(f"       flap {index + 1}/{flap_count} recovered")

        print("[8/10] deleting file to recycle bin and restoring it")
        client.delete_file(file_id)
        recycle_items = client.list_recycle_bin()
        if not any(item["id"] == file_id for item in recycle_items):
            raise RuntimeError("deleted file not found in recycle bin")
        client.restore_file(file_id)
        detail = client.file_detail(file_id)
        if detail["file"]["id"] != file_id:
            raise RuntimeError("restored file detail did not match original file")
        print("       restore succeeded")

        print("[9/10] reconciling after restore and validating download")
        detail = client.wait_for_ready_replica_set(file_id, {"bkd_local_default", backend["id"]})
        restored_bytes = client.download_file(file_id)
        if restored_bytes != content:
            raise RuntimeError("restored file download does not match original content")
        print(f"       replicas after restore: {detail['ready_replica_backend_ids']}")

        print("[10/10] final integrity digest")
        print(f"       sha256: {hashlib.sha256(content).hexdigest()}")
        print("Chaos acceptance passed.")
        return 0
    finally:
        client.upsert_policy(
            file_class,
            require_local=bool(original_policy["require_local"]),
            stable_replica_count=int(original_policy["stable_replica_count"]),
            opportunistic_replica_count=int(original_policy["opportunistic_replica_count"]),
            preferred_backend_ids=list(original_policy["preferred_backend_ids"]),
            excluded_backend_ids=list(original_policy["excluded_backend_ids"]),
            max_non_local_size_bytes=original_policy["max_non_local_size_bytes"],
        )


def build_deterministic_bytes(size_bytes: int) -> bytes:
    chunk = hashlib.sha256(b"iron-chaos-seed").digest()
    repeats = (size_bytes // len(chunk)) + 1
    return (chunk * repeats)[:size_bytes]


def validate_stream_ranges(client: IronClient, file_id: str, content: bytes) -> None:
    checks = [
        (0, 255),
        (1024, 2047),
        (len(content) // 2, (len(content) // 2) + 511),
        (len(content) - 512, len(content) - 1),
    ]
    for start, end in checks:
        body, headers = client.stream_file(file_id, range_header=f"bytes={start}-{end}")
        if body != content[start : end + 1]:
            raise RuntimeError(f"range body mismatch for {start}-{end}")
        content_range = header_value(headers, "Content-Range")
        if content_range != f"bytes {start}-{end}/{len(content)}":
            raise RuntimeError(f"unexpected content-range for {start}-{end}: {content_range}")


def wait_for_unhealthy(client: IronClient, backend_id: str, timeout_seconds: float = 30) -> None:
    deadline = time.time() + timeout_seconds
    while time.time() < deadline:
        result = client.check_backend(backend_id)
        if result["status"] == "unhealthy":
            return
        time.sleep(1)
    raise RuntimeError("backend did not become unhealthy in time")


def wait_for_healthy(client: IronClient, backend_id: str, timeout_seconds: float = 45) -> None:
    deadline = time.time() + timeout_seconds
    while time.time() < deadline:
        result = client.check_backend(backend_id)
        if result["status"] == "healthy":
            return
        time.sleep(1)
    raise RuntimeError("backend did not recover to healthy in time")


def local_backend_base_path(client: IronClient) -> str:
    for backend in client.list_backends():
        if backend["id"] == "bkd_local_default":
            return str(backend["config"]["base_path"])
    raise RuntimeError("local-default backend not found")


def default_cache_root() -> Path:
    return Path(os.getenv("IRON_CACHE_DIR", ".iron-cache"))


def docker(args: list[str]) -> None:
    subprocess.run(["docker", *args], cwd=ROOT, check=True)


if __name__ == "__main__":
    raise SystemExit(main())