Genarrative/scripts/export-gitea-rust-cache.py

#!/usr/bin/env python3
"""只把 master push 的任务内 sccache 对象导出为 Gitea Actions 产物。"""

import base64
from datetime import datetime, timedelta, timezone
import hashlib
import io
import json
import os
from pathlib import Path
import re
import stat
import subprocess
import sys
import tarfile
import tempfile
import time
import urllib.error
import urllib.parse
import urllib.request
import zipfile


MAX_OBJECT_BYTES = 4 * 1024**3
CHUNK_BYTES = 8 * 1024**2
JOBS = {
    "ai-game-creator-shell-rust-lane-1",
    "ai-game-creator-shell-rust-lane-2",
    "ai-game-creator-shell-rust-crates",
    "backend-tests",
    "native-shell-tests",
}


def object_path(name):
    parts = name.split("/")
    return (
        len(parts) == 3
        and re.fullmatch(r"[a-f0-9]{64}", parts[2]) is not None
        and parts[0] == parts[2][0]
        and parts[1] == parts[2][1]
    )


class HashingReader:
    def __init__(self, stream):
        self.stream = stream
        self.digest = hashlib.sha256()

    def read(self, size=-1):
        data = self.stream.read(size)
        self.digest.update(data)
        return data


def scan_objects(root):
    objects = root / "objects"
    if objects.is_symlink() or not objects.is_dir():
        raise ValueError("cache object directory is unavailable")
    candidates = []
    for directory, directories, files in os.walk(objects, followlinks=False):
        directories[:] = [
            name for name in directories if not (Path(directory) / name).is_symlink()
        ]
        for filename in files:
            path = Path(directory) / filename
            relative = path.relative_to(objects).as_posix()
            info = path.lstat()
            if object_path(relative) and stat.S_ISREG(info.st_mode):
                candidates.append((path, relative, info))
    return candidates


def save_baseline(root, destination):
    # sccache 命中会更新 mtime，不能把时间变化当成内容变化，否则又变成全量上传。
    baseline = {
        relative: {"size": info.st_size, "mtime_ns": info.st_mtime_ns}
        for _, relative, info in scan_objects(root)
    }
    destination.write_text(json.dumps(baseline), encoding="utf-8")


def pack_snapshot(root, destination, metadata, baseline, limit=MAX_OBJECT_BYTES):
    """仅归档新 key；命中的已有对象只传递新近使用时间，不传输对象内容。"""
    candidates = scan_objects(root)
    candidates.sort(key=lambda item: (-item[2].st_mtime_ns, item[1]))
    entries = []
    touched = []
    total = 0
    with zipfile.ZipFile(destination, "w", compression=zipfile.ZIP_STORED) as bundle, \
            bundle.open("snapshot.tar", "w", force_zip64=True) as tar_stream, \
            tarfile.open(fileobj=tar_stream, mode="w|", format=tarfile.USTAR_FORMAT) as archive:
        for path, relative, info in candidates:
            previous = baseline.get(relative)
            if previous is not None and previous["size"] == info.st_size:
                if previous["mtime_ns"] != info.st_mtime_ns:
                    touched.append({"path": "objects/" + relative, "mtime_ns": info.st_mtime_ns})
                continue
            if info.st_size <= 0 or total + info.st_size > limit:
                continue
            # report 已成功停止 daemon；打开后再核实对象，避免导出变化中的文件。
            descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0))
            with os.fdopen(descriptor, "rb") as stream:
                actual = os.fstat(stream.fileno())
                if not stat.S_ISREG(actual.st_mode) or (
                    actual.st_size, actual.st_mtime_ns, actual.st_ino
                ) != (info.st_size, info.st_mtime_ns, info.st_ino):
                    raise ValueError("cache object changed during export")
                member = tarfile.TarInfo("objects/" + relative)
                member.size = info.st_size
                member.mode = 0o644
                member.mtime = int(info.st_mtime)
                reader = HashingReader(stream)
                archive.addfile(member, reader)
                entries.append({
                    "path": member.name,
                    "size": info.st_size,
                    "sha256": reader.digest.hexdigest(),
                    "mtime_ns": info.st_mtime_ns,
                })
                total += info.st_size
        manifest = dict(metadata, schema=1, mode="delta", objects=entries, touched=touched)
        data = json.dumps(manifest, ensure_ascii=False, sort_keys=True).encode("utf-8")
        member = tarfile.TarInfo("manifest.json")
        member.size = len(data)
        member.mode = 0o644
        archive.addfile(member, io.BytesIO(data))
    return len(entries), total


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        # runtime token 只允许发往明确配置的 Gitea 地址。
        return None


def request(method, url, token, data=None, headers=None, attempts=3):
    opener = urllib.request.build_opener(NoRedirect())
    for attempt in range(attempts):
        req = urllib.request.Request(url, data=data, method=method, headers={
            **({"Authorization": "Bearer " + token} if token else {}),
            "Content-Type": "application/json",
            **(headers or {}),
        })
        try:
            with opener.open(req, timeout=120) as response:
                body = response.read(1024 * 1024)
                return json.loads(body) if body else {}
        except urllib.error.HTTPError as error:
            if error.code not in (408, 429, 500, 502, 503, 504) or attempt == attempts - 1:
                raise RuntimeError(f"artifact {method} failed: HTTP {error.code}") from None
        except (urllib.error.URLError, TimeoutError, ConnectionError):
            if attempt == attempts - 1:
                raise RuntimeError(f"artifact {method} connection failed") from None
        time.sleep(2 ** attempt)


def artifact_base_url(env):
    # .runner.address 可指向仅支持 RPC 的领取网关，不能用 GITHUB_SERVER_URL。
    repository = env["GITHUB_REPOSITORY"]
    clone_url = env["GENARRATIVE_GITEA_REPOSITORY_URL"]
    suffix = "/" + repository + ".git"
    parsed = urllib.parse.urlsplit(clone_url)
    if (
        parsed.scheme not in ("http", "https")
        or not parsed.netloc
        or parsed.username or parsed.password or parsed.query or parsed.fragment
        or not parsed.path.endswith(suffix)
    ):
        raise ValueError("a credential-free HTTP Gitea repository URL is required")
    return urllib.parse.urlunsplit((
        parsed.scheme, parsed.netloc, parsed.path[:-len(suffix)], "", ""
    ))


def upload_snapshot(path, name, run_id, base_url, token):
    """Gitea 1.26.4 原生 v4：只有完成全部块并校验 SHA256 后才发布。"""
    producer_match = re.fullmatch(r"rust-cache-v1-([a-z0-9-]+)-attempt-(0|[1-9][0-9]*)", name)
    if not producer_match or producer_match[1] not in JOBS:
        raise ValueError("unexpected cache artifact name")
    producer = f"{producer_match[1]}:{producer_match[2]}"
    endpoint = f"{base_url}/twirp/github.actions.results.api.v1.ArtifactService"
    # Gitea 从任务凭据确定 job，只要求请求中的 run ID 与任务所属 run 一致。
    identity = {"workflowRunBackendId": str(run_id), "name": name}
    created = request("POST", endpoint + "/CreateArtifact", token, json.dumps({
        **identity, "version": 4,
        # Gitea 将剩余小时向下取整成天；多给一天，服务端才实际保留至少七天。
        "expiresAt": (datetime.now(timezone.utc) + timedelta(days=8)).isoformat(),
    }).encode())
    received = urllib.parse.urlsplit(created.get("signedUploadUrl", ""))
    expected_path = urllib.parse.urlsplit(endpoint).path + "/UploadArtifact"
    if not created.get("ok") or received.path != expected_path or not received.query:
        raise ValueError("unexpected artifact upload path")
    # 服务可能返回外部 AppURL；保留签名，但只连接配置的内部 Gitea 地址。
    upload_url = endpoint + "/UploadArtifact?" + received.query
    size = path.stat().st_size
    digest = hashlib.sha256()
    blocks = []
    with path.open("rb") as source:
        while chunk := source.read(CHUNK_BYTES):
            # 双层编码中的内层保留所属任务；宿主只回收带此前缀的过期未完成块。
            block = base64.b64encode(
                f"genarrative-rust-cache-v1:{producer}:{len(blocks):08d}".encode()
            ).decode()
            blocks.append(block)
            digest.update(chunk)
            request("PUT", upload_url + "&" + urllib.parse.urlencode({
                "comp": "block", "blockid": block,
            }), None, chunk, {
                "Content-Type": "application/octet-stream",
            })
    blocklist = "<BlockList>" + "".join(f"<Latest>{block}</Latest>" for block in blocks) + "</BlockList>"
    request("PUT", upload_url + "&comp=blocklist", None, blocklist.encode(), {
        "Content-Type": "application/xml",
    })
    # Finalize 会消耗服务端块列表；响应不确定时不盲目重试、也不宣告成功。
    finalized = request("POST", endpoint + "/FinalizeArtifact", token, json.dumps({
        **identity, "size": str(size), "hash": "sha256:" + digest.hexdigest(),
    }).encode(), attempts=1)
    if not finalized.get("ok"):
        raise ValueError("artifact was not finalized")


def export(env):
    # 双重限定：手工调用、PR 和其它分支不会扫描、打包或上传对象。
    if env.get("GITHUB_EVENT_NAME") != "push" or env.get("GITHUB_REF") != "refs/heads/master":
        return
    if env.get("GENARRATIVE_CI_RUST_CACHE_EXPORT_READY") != "1":
        print("[rust-cache] export skipped: cache daemon did not stop cleanly")
        return
    job = env.get("GITHUB_JOB")
    if job not in JOBS:
        raise ValueError("unexpected Rust cache job")
    sha = env["GITHUB_SHA"]
    if not re.fullmatch(r"[a-f0-9]{40}", sha):
        raise ValueError("a complete source SHA is required")
    run_id, attempt = int(env["GITHUB_RUN_ID"]), int(env["GITHUB_RUN_ATTEMPT"])
    if run_id <= 0 or attempt < 0:
        raise ValueError("invalid run identity")
    token = env.get("ACTIONS_RUNTIME_TOKEN") or env["GENARRATIVE_GITEA_TOKEN"]
    base_url = artifact_base_url(env)
    root = Path(env.get("GENARRATIVE_CI_RUST_CACHE_ROOT", "/opt/genarrative-ci/rust-cache"))
    baseline_path = root / "export-baseline.json"
    baseline = json.loads(baseline_path.read_text(encoding="utf-8"))
    rustc = subprocess.check_output(["rustc", "-vV"], text=True)
    workspace = str(Path.cwd().resolve())
    if rustc != (root / "rustc.txt").read_text() or workspace != (root / "workspace.txt").read_text().strip():
        raise ValueError("cache provenance changed after prepare")
    metadata = {
        "repository": env["GITHUB_REPOSITORY"],
        "run_id": run_id, "run_attempt": attempt, "job": job, "source_sha": sha,
        "rustc": rustc, "workspace": workspace,
        "inherited_source_sha": (root / "source-commit.txt").read_text().strip(),
        "sccache_version": subprocess.check_output([str(root / "sccache"), "--version"], text=True).strip(),
    }
    if (root / "base-image.txt").is_file():
        metadata["base_image"] = (root / "base-image.txt").read_text().strip()
    name = f"rust-cache-v1-{job}-attempt-{attempt}"
    try:
        with tempfile.TemporaryDirectory(prefix="ci-rust-cache-export-") as directory:
            path = Path(directory) / "snapshot.zip"
            count, size = pack_snapshot(root, path, metadata, baseline)
            upload_snapshot(path, name, run_id, base_url, token)
    finally:
        baseline_path.unlink(missing_ok=True)
    print(f"[rust-cache] artifact={name} objects={count} bytes={size} complete=true")


if __name__ == "__main__":
    if len(sys.argv) == 3 and sys.argv[1] == "baseline":
        save_baseline(Path(os.environ.get("GENARRATIVE_CI_RUST_CACHE_ROOT", "/opt/genarrative-ci/rust-cache")), Path(sys.argv[2]))
    else:
        export(os.environ)