MoShiro/ci/guard-step.py

#!/usr/bin/env python3
"""Run a CI step so that one failure stops the whole pipeline and nothing is left hanging.

usage: guard-step.py <name> <seconds> <command> [args...]

All guarded steps of one pipeline (across workflows) share a state directory. The first step that fails, is
cancelled or times out writes a `failed` marker there. Every other guarded step then
  - exits at once if it has not started yet (code 125), or
  - is stopped if it is running: SIGTERM to its process group so traps run and flock/Buildah state is released,
    SIGKILL after a grace period, and the lingering group is always killed when the step ends (code 143).
A pipeline deadline is the last resort for a step that waits on something that never happens.
"""

from __future__ import annotations

import json
import os
import shutil
import signal
import subprocess
import sys
import time
from pathlib import Path

GRACE_SECONDS = 15
POLL_SECONDS = 1.0
STATE_RETENTION_SECONDS = 3 * 86400
ABORTED = 125
INTERRUPTED = 143
TIMED_OUT = 124


def state_dir() -> Path:
    base = Path(os.environ.get("MOSHIRO_CI_STATE", "/home/justme/woodpecker/cache/pipeline-state"))
    key = f"{os.environ.get('CI_REPO_ID', '0')}-{os.environ.get('CI_PIPELINE_NUMBER', '0')}"
    return base / key


def prune(base: Path, keep: Path) -> None:
    cutoff = time.time() - STATE_RETENTION_SECONDS
    try:
        for entry in base.iterdir():
            if entry != keep and entry.is_dir() and entry.stat().st_mtime < cutoff:
                shutil.rmtree(entry, ignore_errors=True)
    except OSError:
        pass


def write_once(path: Path, content: str) -> bool:
    """Create the file atomically; False when somebody else created it first."""
    try:
        descriptor = os.open(path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o644)
    except FileExistsError:
        return False
    with os.fdopen(descriptor, "w") as handle:
        handle.write(content)
    return True


def kill_group(pid: int, signum: int) -> None:
    try:
        os.killpg(pid, signum)
    except (ProcessLookupError, PermissionError):
        pass


def main() -> int:
    if len(sys.argv) < 4:
        print("usage: guard-step.py <name> <seconds> <command> [args...]", file=sys.stderr)
        return 2
    name = sys.argv[1]
    try:
        timeout_seconds = float(sys.argv[2])
    except ValueError:
        print("timeout must be a number of seconds", file=sys.stderr)
        return 2
    command = sys.argv[3:]

    directory = state_dir()
    try:
        directory.mkdir(parents=True, exist_ok=True)
        prune(directory.parent, directory)
    except OSError as error:
        # 没有状态目录就退化成普通的超时包装，不能因为守卫本身让步骤失败
        print(f"guard-step: no pipeline state ({error}); running without cross-step abort", file=sys.stderr)
        directory = None

    failed = directory / "failed" if directory else None
    pipeline_deadline = float(os.environ.get("MOSHIRO_PIPELINE_DEADLINE", "6000"))
    started_file = directory / "started" if directory else None
    if started_file is not None:
        write_once(started_file, str(time.time()))

    def mark(reason: str, code: int) -> None:
        if failed is not None and write_once(failed, json.dumps({"step": name, "reason": reason, "code": code, "at": time.time()})):
            print(f"guard-step: {name} {reason}; stopping the rest of the pipeline", file=sys.stderr)

    def earlier_failure() -> str | None:
        if failed is None or not failed.exists():
            return None
        try:
            info = json.loads(failed.read_text())
            return f"{info.get('step')} {info.get('reason')}"
        except (OSError, ValueError):
            return "an earlier step failed"

    def pipeline_expired() -> bool:
        if started_file is None:
            return False
        try:
            return time.time() - float(started_file.read_text()) > pipeline_deadline
        except (OSError, ValueError):
            return False

    reason = earlier_failure()
    if reason is not None:
        print(f"guard-step: {name} not started: {reason}", file=sys.stderr)
        return ABORTED
    if pipeline_expired():
        mark("pipeline deadline exceeded", TIMED_OUT)
        return TIMED_OUT

    process = subprocess.Popen(command, start_new_session=True)
    interrupted: int | None = None
    stop_reason: str | None = None
    stop_started: float | None = None

    def request_stop(why: str) -> None:
        nonlocal stop_reason, stop_started
        if stop_reason is None:
            stop_reason = why
            stop_started = time.monotonic()
            kill_group(process.pid, signal.SIGTERM)

    def on_signal(signum: int, _frame: object) -> None:
        nonlocal interrupted
        interrupted = signum
        request_stop(f"received signal {signum}")

    for signum in (signal.SIGTERM, signal.SIGINT, signal.SIGHUP):
        signal.signal(signum, on_signal)

    deadline = time.monotonic() + timeout_seconds
    result: int | None = None
    try:
        while result is None:
            now = time.monotonic()
            if stop_reason is None:
                other = earlier_failure()
                if other is not None:
                    request_stop(f"aborted by {other}")
                elif now >= deadline:
                    request_stop(f"timed out after {int(timeout_seconds)}s")
                elif pipeline_expired():
                    request_stop("pipeline deadline exceeded")
            elif stop_started is not None and now - stop_started >= GRACE_SECONDS:
                kill_group(process.pid, signal.SIGKILL)
                stop_started = None
            try:
                result = process.wait(timeout=POLL_SECONDS)
            except subprocess.TimeoutExpired:
                continue
    finally:
        # 不管怎么结束，这个步骤自己的进程组里不留任何东西（它们可能还握着 flock 或 Buildah 的存储）
        kill_group(process.pid, signal.SIGKILL)

    if stop_reason is None:
        if result != 0:
            mark(f"failed with exit code {result}", result)
        return result
    if stop_reason.startswith("aborted by "):
        print(f"guard-step: {name} {stop_reason}", file=sys.stderr)
        return INTERRUPTED
    if interrupted is not None:
        mark(stop_reason, 128 + interrupted)
        return 128 + interrupted
    mark(stop_reason, TIMED_OUT)
    return TIMED_OUT


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