83 lines
2.4 KiB
Python
83 lines
2.4 KiB
Python
#!/usr/bin/env python3
|
|
"""Run a command in its own process group and clean it up on timeout/cancel."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
|
|
|
|
GRACE_SECONDS = 10
|
|
|
|
|
|
def main() -> int:
|
|
if len(sys.argv) < 3:
|
|
print("usage: run-with-timeout.py <seconds> <command> [args...]", file=sys.stderr)
|
|
return 2
|
|
|
|
try:
|
|
timeout_seconds = float(sys.argv[1])
|
|
except ValueError:
|
|
print("timeout must be a number of seconds", file=sys.stderr)
|
|
return 2
|
|
|
|
process = subprocess.Popen(sys.argv[2:], start_new_session=True)
|
|
stop_signal: int | None = None
|
|
stop_started: float | None = None
|
|
timed_out = False
|
|
|
|
def request_stop(signum: int, _frame: object) -> None:
|
|
nonlocal stop_signal, stop_started
|
|
if stop_signal is None:
|
|
stop_signal = signum
|
|
stop_started = time.monotonic()
|
|
try:
|
|
os.killpg(process.pid, signal.SIGTERM)
|
|
except ProcessLookupError:
|
|
pass
|
|
|
|
for signum in (signal.SIGTERM, signal.SIGINT, signal.SIGHUP):
|
|
signal.signal(signum, request_stop)
|
|
|
|
deadline = time.monotonic() + timeout_seconds
|
|
while True:
|
|
now = time.monotonic()
|
|
if stop_signal is None and not timed_out and now >= deadline:
|
|
timed_out = True
|
|
stop_started = now
|
|
try:
|
|
os.killpg(process.pid, signal.SIGTERM)
|
|
except ProcessLookupError:
|
|
pass
|
|
elif stop_started is not None and now - stop_started >= GRACE_SECONDS:
|
|
try:
|
|
os.killpg(process.pid, signal.SIGKILL)
|
|
except ProcessLookupError:
|
|
pass
|
|
stop_started = None
|
|
|
|
try:
|
|
result = process.wait(timeout=0.2)
|
|
if stop_signal is not None:
|
|
try:
|
|
os.killpg(process.pid, signal.SIGKILL)
|
|
except ProcessLookupError:
|
|
pass
|
|
return 128 + stop_signal
|
|
if timed_out:
|
|
try:
|
|
os.killpg(process.pid, signal.SIGKILL)
|
|
except ProcessLookupError:
|
|
pass
|
|
return 124
|
|
return result
|
|
except subprocess.TimeoutExpired:
|
|
continue
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|