Add deadman: prime a forced reboot before risky PCIe ops

dau_utils/deadman.py

@@ -0,0 +1,153 @@
+"""Prime a forced reboot before a risky PCIe operation, cancel it on success.
+
+The dpv1 wedge class that needs a power cycle is a driver/PCIe hang where the
+kernel and systemd stay alive -- they keep petting the hardware watchdog, so
+systemd's ``RuntimeWatchdogSec`` never fires, yet ``systemctl reboot`` hangs on
+the stuck device. The recovery is ``sysrq-b`` (``emergency_restart()``), which
+resets immediately without touching the wedged driver.
+
+``arm`` schedules that reset as a transient systemd timer so it survives the
+controlling SSH session; ``disarm`` cancels it. Run ``arm`` before a rescan,
+flash, or register probe and ``disarm`` once it returns cleanly. If the box
+wedges before ``disarm``, the timer fires and reboots it. A full kernel lock
+(systemd itself dead) is still caught by systemd's own hardware watchdog.
+"""
+
+from __future__ import annotations
+
+import argparse
+import subprocess
+import sys
+from collections.abc import Iterator, Sequence
+from contextlib import contextmanager
+
+DEFAULT_UNIT = "dau-deadman"
+DEFAULT_TIMEOUT_S = 180
+
+# sysrq 'b' = emergency_restart(): reboot past a wedged driver, where a clean
+# `systemctl reboot` would block on it. Best-effort sync first so a live
+# filesystem lands its journal; if the box is too wedged to sync, the reboot
+# still proceeds.
+RESET_SCRIPT = "echo s > /proc/sysrq-trigger; echo b > /proc/sysrq-trigger"
+
+# systemctl reports these when a unit is not running; anything else (notably
+# "active"/"activating") means the pending reset is still live.
+_INACTIVE_STATES = frozenset({"inactive", "failed", "unknown", "dead"})
+
+
+class DeadmanError(RuntimeError):
+    """A deadman operation could not be confirmed -- treat the host as unsafe."""
+
+
+def arm_command(timeout_s: int = DEFAULT_TIMEOUT_S, *, unit: str = DEFAULT_UNIT) -> tuple[str, ...]:
+    """The ``systemd-run`` invocation that schedules the reset ``timeout_s`` from now."""
+    if timeout_s < 1:
+        raise ValueError(f"deadman timeout must be at least 1 second, got {timeout_s}")
+    return (
+        "sudo",
+        "systemd-run",
+        f"--unit={unit}",
+        f"--on-active={timeout_s}",
+        "--timer-property=AccuracySec=1s",
+        "--collect",
+        "/bin/sh",
+        "-c",
+        RESET_SCRIPT,
+    )
+
+
+def disarm_commands(*, unit: str = DEFAULT_UNIT) -> tuple[tuple[str, ...], ...]:
+    """Stop the pending timer/service and clear any failed state, idempotently."""
+    return (
+        ("sudo", "systemctl", "stop", f"{unit}.timer", f"{unit}.service"),
+        ("sudo", "systemctl", "reset-failed", f"{unit}.timer", f"{unit}.service"),
+    )
+
+
+def status_command(*, unit: str = DEFAULT_UNIT) -> tuple[str, ...]:
+    """List the pending deadman timer, if armed."""
+    return ("systemctl", "list-timers", "--all", f"{unit}.timer")
+
+
+def _is_active(name: str) -> bool:
+    """True only if systemctl positively reports ``name`` running. A failed
+    query (D-Bus down, sudo denied) is treated as active -- we cannot claim a
+    unit is stopped unless systemctl confirms it."""
+    result = subprocess.run(("systemctl", "is-active", name), check=False, capture_output=True, text=True)
+    return result.stdout.strip() not in _INACTIVE_STATES
+
+
+def is_armed(*, unit: str = DEFAULT_UNIT) -> bool:
+    """True if a deadman timer or its service is still live for ``unit``."""
+    return _is_active(f"{unit}.timer") or _is_active(f"{unit}.service")
+
+
+def arm(timeout_s: int = DEFAULT_TIMEOUT_S, *, unit: str = DEFAULT_UNIT) -> None:
+    """Schedule the forced reset. Refuses to stomp an already-armed timer so
+    concurrent callers cannot silently cancel each other's protection; clears
+    only inactive stale state before scheduling."""
+    if is_armed(unit=unit):
+        raise DeadmanError(f"{unit} is already armed; disarm it before arming again")
+    subprocess.run(("sudo", "systemctl", "reset-failed", f"{unit}.timer", f"{unit}.service"), check=False, capture_output=True)
+    subprocess.run(arm_command(timeout_s, unit=unit), check=True)
+
+
+def disarm(*, unit: str = DEFAULT_UNIT) -> None:
+    """Cancel the pending reset and confirm it is gone. Raises ``DeadmanError``
+    if the timer cannot be verified inactive -- the caller must not treat the
+    host as safe until this returns cleanly."""
+    subprocess.run(("sudo", "systemctl", "stop", f"{unit}.timer", f"{unit}.service"), check=False, capture_output=True)
+    if is_armed(unit=unit):
+        raise DeadmanError(f"{unit} still armed after stop; reset may still fire -- intervene before trusting the host")
+    subprocess.run(("sudo", "systemctl", "reset-failed", f"{unit}.timer", f"{unit}.service"), check=False, capture_output=True)
+
+
+@contextmanager
+def armed(timeout_s: int = DEFAULT_TIMEOUT_S, *, unit: str = DEFAULT_UNIT) -> Iterator[None]:
+    """Arm around a risky block; disarm on the way out, success or exception."""
+    arm(timeout_s, unit=unit)
+    try:
+        yield
+    finally:
+        disarm(unit=unit)
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Prime a forced reboot before a risky PCIe op; cancel it on success")
+    parser.add_argument("action", choices=("arm", "disarm", "status"), help="Deadman action")
+    parser.add_argument("--timeout", type=int, default=DEFAULT_TIMEOUT_S, help="Seconds before the reset fires when arming")
+    parser.add_argument("--unit", default=DEFAULT_UNIT, help="Transient systemd unit name")
+    parser.add_argument("--dry-run", action="store_true", help="Print the command(s) without running them")
+    args = parser.parse_args(argv)
+
+    if args.action == "arm":
+        if args.dry_run:
+            print(" ".join(arm_command(args.timeout, unit=args.unit)))
+            return 0
+        try:
+            arm(args.timeout, unit=args.unit)
+        except DeadmanError as error:
+            print(f"deadman NOT armed: {error}", file=sys.stderr)
+            return 1
+        print(f"deadman armed: reset in {args.timeout}s (unit {args.unit}); disarm before then")
+        return 0
+
+    if args.action == "disarm":
+        if args.dry_run:
+            for command in disarm_commands(unit=args.unit):
+                print(" ".join(command))
+            return 0
+        try:
+            disarm(unit=args.unit)
+        except DeadmanError as error:
+            print(f"deadman DISARM FAILED: {error}", file=sys.stderr)
+            return 1
+        print(f"deadman disarmed (unit {args.unit})")
+        return 0
+
+    subprocess.run(status_command(unit=args.unit), check=False)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

dau_utils/tests/test_deadman.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+import runpy
+import subprocess
+import sys
+
+import pytest
+
+from dau_utils import deadman
+from dau_utils.deadman import (
+    DEFAULT_TIMEOUT_S,
+    DEFAULT_UNIT,
+    DeadmanError,
+    arm,
+    arm_command,
+    disarm,
+    disarm_commands,
+    is_armed,
+    main,
+    status_command,
+)
+
+
+def test_arm_command_schedules_a_transient_sysrq_reset_timer() -> None:
+    command = arm_command(120, unit="dau-deadman")
+
+    assert command[:6] == (
+        "sudo",
+        "systemd-run",
+        "--unit=dau-deadman",
+        "--on-active=120",
+        "--timer-property=AccuracySec=1s",
+        "--collect",
+    )
+    assert command[6:8] == ("/bin/sh", "-c")
+    assert "/proc/sysrq-trigger" in command[8]
+    assert command[8].strip().endswith("echo b > /proc/sysrq-trigger")
+
+
+def test_arm_command_rejects_a_nonpositive_timeout() -> None:
+    for bad in (0, -5):
+        try:
+            arm_command(bad)
+        except ValueError:
+            continue
+        raise AssertionError(f"expected ValueError for timeout {bad}")
+
+
+def test_disarm_stops_the_timer_and_service_then_clears_failed_state() -> None:
+    stop, reset = disarm_commands(unit="dau-deadman")
+
+    assert stop == ("sudo", "systemctl", "stop", "dau-deadman.timer", "dau-deadman.service")
+    assert reset == ("sudo", "systemctl", "reset-failed", "dau-deadman.timer", "dau-deadman.service")
+
+
+def test_status_command_lists_the_named_timer() -> None:
+    assert status_command(unit="dau-deadman") == ("systemctl", "list-timers", "--all", "dau-deadman.timer")
+
+
+def test_cli_arm_dry_run_prints_the_scheduled_reset_command(capsys) -> None:
+    exit_code = main(["arm", "--timeout", "90", "--dry-run"])
+
+    assert exit_code == 0
+    printed = capsys.readouterr().out.strip()
+    assert printed == " ".join(arm_command(90, unit=DEFAULT_UNIT))
+
+
+def test_cli_disarm_dry_run_prints_both_teardown_commands(capsys) -> None:
+    exit_code = main(["disarm", "--dry-run"])
+
+    assert exit_code == 0
+    lines = capsys.readouterr().out.splitlines()
+    assert lines == [" ".join(command) for command in disarm_commands(unit=DEFAULT_UNIT)]
+
+
+def test_cli_arm_dry_run_defaults_to_the_module_timeout(capsys) -> None:
+    main(["arm", "--dry-run"])
+
+    assert f"--on-active={DEFAULT_TIMEOUT_S}" in capsys.readouterr().out
+
+
+class _FakeSystemctl:
+    """Stands in for subprocess.run: is-active returns a scripted state, and
+    every mutating call (sudo/systemd-run/stop) is recorded."""
+
+    def __init__(self, active_states: dict[str, str]) -> None:
+        self.active_states = active_states
+        self.calls: list[tuple[str, ...]] = []
+
+    def __call__(self, command, check=False, capture_output=False, text=False):  # noqa: ANN001
+        command = tuple(command)
+        self.calls.append(command)
+        if command[:2] == ("systemctl", "is-active"):
+            state = self.active_states.get(command[2], "inactive")
+            return subprocess.CompletedProcess(command, 0, stdout=f"{state}\n", stderr="")
+        return subprocess.CompletedProcess(command, 0, stdout="", stderr="")
+
+
+def test_arm_refuses_to_replace_an_already_armed_timer(monkeypatch) -> None:
+    fake = _FakeSystemctl({f"{DEFAULT_UNIT}.timer": "active"})
+    monkeypatch.setattr(deadman.subprocess, "run", fake)
+
+    with pytest.raises(DeadmanError, match="already armed"):
+        arm(120)
+
+    assert not any(call[:2] == ("sudo", "systemd-run") for call in fake.calls)
+
+
+def test_arm_schedules_when_no_timer_is_live(monkeypatch) -> None:
+    fake = _FakeSystemctl({})  # everything inactive
+    monkeypatch.setattr(deadman.subprocess, "run", fake)
+
+    arm(120)
+
+    assert any(call[:2] == ("sudo", "systemd-run") for call in fake.calls)
+
+
+def test_disarm_raises_when_the_timer_survives_the_stop(monkeypatch) -> None:
+    fake = _FakeSystemctl({f"{DEFAULT_UNIT}.timer": "active"})  # stop is a no-op here
+    monkeypatch.setattr(deadman.subprocess, "run", fake)
+
+    with pytest.raises(DeadmanError, match="still armed after stop"):
+        disarm()
+
+
+def test_disarm_succeeds_only_once_the_timer_is_confirmed_inactive(monkeypatch) -> None:
+    fake = _FakeSystemctl({})  # is-active reports inactive
+    monkeypatch.setattr(deadman.subprocess, "run", fake)
+
+    disarm()  # no raise
+
+    assert ("sudo", "systemctl", "stop", f"{DEFAULT_UNIT}.timer", f"{DEFAULT_UNIT}.service") in fake.calls
+
+
+def test_is_armed_treats_a_failed_query_as_still_armed(monkeypatch) -> None:
+    fake = _FakeSystemctl({f"{DEFAULT_UNIT}.timer": "activating"})
+    monkeypatch.setattr(deadman.subprocess, "run", fake)
+
+    assert is_armed() is True
+
+
+def test_cli_disarm_reports_failure_when_timer_cannot_be_confirmed_gone(monkeypatch, capsys) -> None:
+    fake = _FakeSystemctl({f"{DEFAULT_UNIT}.timer": "active"})
+    monkeypatch.setattr(deadman.subprocess, "run", fake)
+
+    exit_code = main(["disarm"])
+
+    assert exit_code == 1
+    assert "DISARM FAILED" in capsys.readouterr().err
+
+
+def test_module_entrypoint_runs_cli_for_uninstalled_checkout(capsys, monkeypatch) -> None:
+    monkeypatch.setattr(sys, "argv", ["deadman", "arm", "--timeout", "42", "--dry-run"])
+    monkeypatch.delitem(sys.modules, "dau_utils.deadman", raising=False)
+
+    try:
+        runpy.run_module("dau_utils.deadman", run_name="__main__")
+    except SystemExit as exc:
+        assert exc.code == 0
+
+    assert "--on-active=42" in capsys.readouterr().out

pyproject.toml

@@ -50,6 +50,7 @@ develop = [
 
 [project.scripts]
 dau-utils-pci-runtime-pm = "dau_utils.pci_runtime_pm:main"
+dau-utils-deadman = "dau_utils.deadman:main"
 
 [project.urls]
 Repository = "https://github.com/dau-dev/dau-utils"