#!/usr/bin/env python3 """Plan and publish the R2 catalog mirror for the CPU workers. The CPU workers (London) fetch replays and maps on demand from R2 through the public CDN host, never from AWS. R2 is content-addressed: replays/corpus/.w3g originals AND c28 conversions maps/. every map the catalog can reference catalog/current.json -> {"sha256","key","replays",...} catalog/serving-index-.json compact immutable serving index Subcommands (run from the developer machine; AWS creds from the mastering-aws profile, R2 creds from the repo .env): plan write a prioritized plan.jsonl of S3 objects still missing from R2, each with a time-limited presigned GET URL, so a box with no AWS credentials can do the transfer (r2-mirror-worker.py) index compact the verified serving index (schema 2, the AWS layer index) into the schema-3 index the workers read publish upload the compact index to R2 and flip catalog/current.json recache stamp the immutable cache header on objects that lack it """ from __future__ import annotations import argparse import gzip import hashlib import json import os import re import sys import time from pathlib import Path SHA1_HEX = 40 def r2_client(): import boto3 env = {} env_path = Path(__file__).resolve().parents[3] / ".env" for line in env_path.read_text().splitlines(): if "=" in line and not line.startswith("#"): key, _, value = line.partition("=") env[key.strip()] = value.strip().strip('"').strip("'") session = boto3.session.Session( aws_access_key_id=env["R2_ACCESS_KEY_ID"], aws_secret_access_key=env["R2_SECRET_ACCESS_KEY"], region_name="auto") # boto3 >= 1.36 adds CRC32 checksum trailers by default; R2 rejects or # stalls on them. Only compute checksums where the API requires them. from botocore.config import Config config = Config(request_checksum_calculation="when_required", response_checksum_validation="when_required", retries={"max_attempts": 3}) return (session.client("s3", endpoint_url=env["R2_ENDPOINT"], config=config), env["R2_BUCKET"]) def s3_client(profile: str, region: str): import boto3 return boto3.session.Session(profile_name=profile, region_name=region).client("s3") def list_keys(client, bucket: str, prefix: str) -> set[str]: keys: set[str] = set() paginator = client.get_paginator("list_objects_v2") for page in paginator.paginate(Bucket=bucket, Prefix=prefix): for obj in page.get("Contents", []): keys.add(obj["Key"]) return keys def map_extension(name: str) -> str: ext = name.rsplit(".", 1)[-1].lower() if "." in name else "" if ext not in ("w3m", "w3x", "tmp"): raise ValueError(f"unsupported map extension in {name!r}") return ext # Engine profiles the London seats serve (w3cs-criu-build-images.sh and # lab-game-session.sh know these names). A layer-index entry whose AWS base # has no London profile is left out of the serving index, so the catalog # routes it to the AWS video path; adding a profile here (after the engine # is installed on every box) brings its replays to London at the next # publish. LONDON_PROFILES = frozenset(( "native-1100", "native-1110", "native-1120", "native-1140", "native-1150", "native-1160", "native-1170", "native-1180", "native-120e", "native-121b", "native-1220", "native-1230", "native-124ab", "native-124cde", "native-125b", "native-1260", "native-1285", "native-1311-d3d9")) ENGINE_PROFILES = {"1.14": "native-1140", "1.28.5": "native-1285", "1.31.1": "native-1311-d3d9"} BASE_ID = re.compile(r"classic-([0-9]+[a-z]*)-") def engine_profile(entry: dict) -> str: """London profile name for a layer-index entry: the AWS layer base `classic--...` names the engine (`classic-1120-tft-open1-audio-v1` -> `native-1120`); the engine string is the fallback for bases without one. Returns "" when London has no such profile.""" base = str(entry.get("baseId") or "") engine = str(entry.get("engine") or "") match = BASE_ID.match(base) if match: profile = "native-" + match.group(1) if profile == "native-1311": profile = "native-1311-d3d9" else: profile = ENGINE_PROFILES.get(engine, "") return profile if profile in LONDON_PROFILES else "" def load_profile_overrides(path) -> dict[str, str]: """Per-replay London profile overrides (see serving-profile-overrides.json): `{"schema": 1, "overrides": {"": "", ...}}`. The AWS layer base names the engine family, and for raw version 24 the pipeline has only a 1.24c/d/e base, so replays recorded on 1.24a/b (the 2009 events) would run on the wrong family: the map loads and the clock runs, but no recorded action executes. The override routes them to the London profile that actually plays them.""" data = json.loads(Path(path).read_text()) if data.get("schema") != 1 or not isinstance(data.get("overrides"), dict): raise ValueError("expected a schema-1 profile override file") overrides = {} for source, profile in data["overrides"].items(): source = str(source).lower() if len(source) != SHA1_HEX or profile not in LONDON_PROFILES: raise ValueError(f"bad override {source!r} -> {profile!r}") overrides[source] = str(profile) return overrides def compact_index(layer_index: dict, profiles=LONDON_PROFILES, overrides: dict[str, str] | None = None) -> dict: """Schema 3: only what a worker needs to materialize a replay from R2. Entries on engines London does not serve are skipped and counted in `skipped` (by engine) so a publish says what stays AWS-only. An entry in `overrides` (source sha1 -> London profile) replaces the profile the layer base implies; overrides are counted in `overridden` (by profile).""" if layer_index.get("schema") != 2: raise ValueError("expected the schema-2 layer index") overrides = overrides or {} out = {} skipped: dict[str, int] = {} overridden: dict[str, int] = {} for source, entry in sorted(layer_index["replays"].items()): profile = engine_profile(entry) if source in overrides: profile = overrides[source] overridden[profile] = overridden.get(profile, 0) + 1 if profile not in profiles: engine = str(entry.get("engine") or entry.get("baseId") or "?") skipped[engine] = skipped.get(engine, 0) + 1 continue played = str(entry["playedReplaySha1"]).lower() map_sha1 = str(entry["mapContentSha1"]).lower() if len(source) != SHA1_HEX or len(played) != SHA1_HEX or len(map_sha1) != SHA1_HEX: raise ValueError(f"bad identity in {source}") map_path = str(entry["mapPath"]) mode = str(entry.get("classicGameMode") or "").lower() if mode not in ("roc", "tft"): raise ValueError(f"bad game mode in {source}") out[source] = { "played": played, "mapSha1": map_sha1, "mapExt": map_extension(map_path.replace("/", "\\").split("\\")[-1]), "mapPath": map_path, "profile": profile, "mode": mode, "family": str(entry.get("family") or mode.upper()), "players": [str(p) for p in (entry.get("players") or [])][:12], "durationMs": int(entry.get("durationMs") or 0), } return {"schema": 3, "generatedAt": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), "replayCount": len(out), "skipped": skipped, "overridden": overridden, "replays": out} def cmd_index(args): layer = json.loads(Path(args.layer_index).read_text()) overrides = load_profile_overrides(args.profile_overrides) \ if args.profile_overrides else {} compact = compact_index(layer, overrides=overrides) missing = sorted(set(overrides) - set(layer["replays"])) if missing: print(json.dumps({"warning": "overrides not in the layer index", "count": len(missing), "first": missing[:5]}), file=sys.stderr) body = json.dumps(compact, sort_keys=True, separators=(",", ":")).encode() Path(args.output).write_bytes(body) print(json.dumps({"replays": compact["replayCount"], "skipped": compact["skipped"], "overridden": compact["overridden"], "bytes": len(body), "sha256": hashlib.sha256(body).hexdigest(), "output": args.output})) def cmd_publish(args): body = Path(args.index).read_bytes() data = json.loads(body) if data.get("schema") != 3: raise SystemExit("publish expects a schema-3 index") sha256 = hashlib.sha256(body).hexdigest() client, bucket = r2_client() key = f"catalog/serving-index-{sha256}.json" client.put_object(Bucket=bucket, Key=key, Body=body, ContentType="application/json", CacheControl="public, max-age=31536000, immutable") current = {"schema": 1, "sha256": sha256, "key": key, "replays": data["replayCount"], "generatedAt": data["generatedAt"], "publishedAt": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())} client.put_object(Bucket=bucket, Key="catalog/current.json", Body=json.dumps(current, sort_keys=True).encode(), ContentType="application/json", CacheControl="public, max-age=60") print(json.dumps(current)) def cmd_plan(args): s3 = s3_client(args.profile, args.region) r2, r2_bucket = r2_client() layer = json.loads(Path(args.layer_index).read_text())["replays"] manifest = json.loads(Path(args.maps_manifest).read_text()) readiness = [json.loads(line) for line in gzip.open(args.readiness, "rt")] print("listing R2 ...", file=sys.stderr) r2_replays = {k.split("/")[-1][:-4] for k in list_keys(r2, r2_bucket, "replays/corpus/") if k.endswith(".w3g")} r2_maps = {k.split("/")[-1].split(".")[0] for k in list_keys(r2, r2_bucket, "maps/")} print(f"R2 has {len(r2_replays)} replays, {len(r2_maps)} maps", file=sys.stderr) def presign(key: str) -> str: return s3.generate_presigned_url( "get_object", Params={"Bucket": args.bucket, "Key": key}, ExpiresIn=args.expires) plan: list[dict] = [] seen: set[str] = set() def add(kind: str, s3_key: str, r2_key: str | None, sha1: str | None, priority: int): ident = r2_key or f"auto:{s3_key}" if ident in seen: return seen.add(ident) plan.append({"priority": priority, "kind": kind, "s3Key": s3_key, "key": r2_key, "sha1": sha1, "url": presign(s3_key)}) # 1. maps the verified index references, 2. their conversions. index_map_keys: dict[str, str] = {} for source, entry in layer.items(): index_map_keys.setdefault(entry["mapContentSha1"].lower(), entry["mapObjectKey"]) for map_sha1, s3_key in index_map_keys.items(): if map_sha1 in r2_maps: continue add("map", s3_key, f"maps/{map_sha1}.{map_extension(s3_key)}", map_sha1, 1) for source, entry in layer.items(): played = entry["playedReplaySha1"].lower() if played in r2_replays: continue s3_key = (f"replays/archive/{source}-c28.w3g" if played != source else f"replays/archive/{source}.w3g") if entry.get("engine") == "1.14": s3_key = f"worker-image/classic/replaykit/1.14/v3/replays/{played}.w3g" add("replay", s3_key, f"replays/corpus/{played}.w3g", played, 2) # 3. every other map content (the full corpus, so the next release needs # no map transfer), 4. every other conversion and original. for map_sha1, keys in manifest["content"].items(): if map_sha1 in r2_maps or map_sha1 in index_map_keys: continue s3_key = sorted(keys)[0] try: ext = map_extension(s3_key) except ValueError: continue add("map", s3_key, f"maps/{map_sha1}.{ext}", map_sha1, 3) for entry in readiness: source = entry["sha1"].lower() if entry.get("convertedPresent") and entry.get("convertedKey") and source not in layer: add("replay", entry["convertedKey"], None, None, 4) # sha1 known after download if entry.get("originalPresent") and source not in r2_replays: add("replay", entry["originalKey"], f"replays/corpus/{source}.w3g", source, 5) plan.sort(key=lambda item: item["priority"]) with open(args.output, "w") as out: for item in plan: out.write(json.dumps(item) + "\n") counts = {} for item in plan: counts[item["priority"]] = counts.get(item["priority"], 0) + 1 print(json.dumps({"planned": len(plan), "byPriority": counts, "expiresInHours": args.expires / 3600, "output": args.output})) IMMUTABLE = "public, max-age=31536000, immutable" def cmd_recache(args): """Stamp the immutable cache header on content-addressed objects that lack it (objects uploaded before the mirror existed). In-place metadata rewrite (CopyObject onto itself), no data transfer; idempotent.""" from concurrent.futures import ThreadPoolExecutor client, bucket = r2_client() keys = [] for prefix in args.prefix: keys.extend(sorted(list_keys(client, bucket, prefix))) print(f"{len(keys)} objects under {args.prefix}", file=sys.stderr) def fix(key: str) -> str: head = client.head_object(Bucket=bucket, Key=key) if head.get("CacheControl") == IMMUTABLE: return "ok" client.copy_object( Bucket=bucket, Key=key, CopySource={"Bucket": bucket, "Key": key}, MetadataDirective="REPLACE", CacheControl=IMMUTABLE, ContentType=head.get("ContentType") or "application/octet-stream") return "fixed" counts = {"ok": 0, "fixed": 0, "error": 0} with ThreadPoolExecutor(max_workers=args.threads) as pool: for index, result in enumerate(pool.map(lambda k: _safe(fix, k), keys), 1): counts[result] += 1 if index % 2000 == 0: print(json.dumps({"done": index, **counts}), file=sys.stderr, flush=True) print(json.dumps({"objects": len(keys), **counts})) def _safe(function, key): try: return function(key) except Exception as error: # noqa: BLE001 print(f"error {key}: {error}", file=sys.stderr) return "error" def main(): parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) sub = parser.add_subparsers(dest="command", required=True) p = sub.add_parser("plan") p.add_argument("--layer-index", required=True) p.add_argument("--maps-manifest", required=True) p.add_argument("--readiness", required=True) p.add_argument("--output", required=True) p.add_argument("--bucket", default="war3replays-assets-440964995465") p.add_argument("--profile", default="mastering-aws") p.add_argument("--region", default="us-east-1") p.add_argument("--expires", type=int, default=48 * 3600) p.set_defaults(func=cmd_plan) p = sub.add_parser("index") p.add_argument("--layer-index", required=True) p.add_argument("--output", required=True) p.add_argument("--profile-overrides", default=str(Path(__file__).with_name("serving-profile-overrides.json")), help="per-replay London profile overrides (schema 1); " "'' disables") p.set_defaults(func=cmd_index) p = sub.add_parser("publish") p.add_argument("--index", required=True) p.set_defaults(func=cmd_publish) p = sub.add_parser("recache") p.add_argument("--prefix", action="append", default=None, help="object prefix to stamp (repeatable; default replays/corpus/ and maps/)") p.add_argument("--threads", type=int, default=32) p.set_defaults(func=cmd_recache) args = parser.parse_args() if getattr(args, "prefix", 0) is None: args.prefix = ["replays/corpus/", "maps/"] args.func(args) if __name__ == "__main__": main()