diff --git a/dau_utils/deadman.py b/dau_utils/deadman.py new file mode 100644 index 0000000..b2e7c31 --- /dev/null +++ b/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()) diff --git a/dau_utils/tests/test_deadman.py b/dau_utils/tests/test_deadman.py new file mode 100644 index 0000000..7edccac --- /dev/null +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 936a664..3e7afb5 100644 --- a/pyproject.toml +++ b/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"