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.

Cancelling a pipeline in Woodpecker (local backend) sends SIGKILL to the step's process group and nothing else, so
this script cannot catch it. The command runs in its own session (so it is not in that group) and would be left
running. A watcher process in a third session holds the read end of a pipe whose write end only this script has; when
the script dies for any reason without saying "done", the watcher stops the command's group (SIGTERM, then SIGKILL).
"""

from __future__ import annotations

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

GRACE_SECONDS = float(os.environ.get("MOSHIRO_GUARD_GRACE", "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 watch(group: int, descriptor: int) -> int:
    """Wait for the guard to finish; if it vanished without saying so, stop the command's process group."""
    with os.fdopen(descriptor, "rb") as pipe:
        if pipe.read() == b"done":
            return 0
    kill_group(group, signal.SIGTERM)
    time.sleep(GRACE_SECONDS)
    kill_group(group, signal.SIGKILL)
    return 0


def start_watcher(group: int) -> int | None:
    """Returns the write end of the watcher's pipe, or None when no watcher could be started."""
    try:
        read_end, write_end = os.pipe()
        subprocess.Popen(
            [sys.executable, os.path.abspath(__file__), "--watch", str(group), str(read_end)],
            pass_fds=(read_end,), start_new_session=True,
            stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
        )
        os.close(read_end)
        return write_end
    except OSError:
        return None


def main() -> int:
    if len(sys.argv) == 4 and sys.argv[1] == "--watch":
        return watch(int(sys.argv[2]), int(sys.argv[3]))
    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)
    watcher_pipe = start_watcher(process.pid)
    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 watcher_pipe is not None:
            try:
                os.write(watcher_pipe, b"done")
                os.close(watcher_pipe)
            except OSError:
                pass

    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())