from __future__ import annotations
|
|
from dataclasses import dataclass
|
from datetime import datetime, timezone
|
import os
|
from pathlib import Path
|
import signal
|
import subprocess
|
import time
|
from typing import Sequence
|
|
from .models import ErrorCode
|
|
|
@dataclass(frozen=True)
|
class ProcessResult:
|
argv_redacted: tuple[str, ...]
|
started: bool
|
pid: int | None
|
exit_code: int | None
|
timed_out: bool
|
terminate_issued: bool
|
exited: bool
|
liveness_unknown: bool
|
stdout_bytes: bytes
|
stderr_bytes: bytes
|
stdout_truncated: bool
|
stderr_truncated: bool
|
started_at_utc: str | None
|
ended_at_utc: str | None
|
elapsed_ms: int | None
|
error_code: ErrorCode | None
|
|
|
class ProcessSupervisor:
|
def __init__(self, *, max_stream_bytes: int = 1_048_576, kill_wait_ms: int = 5_000):
|
self.max_stream_bytes = max_stream_bytes
|
self.kill_wait_ms = kill_wait_ms
|
|
@staticmethod
|
def _utc_now() -> str:
|
return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
|
|
def run(self, argv: Sequence[str], *, timeout_ms: int, cwd: Path | None = None,
|
stdin_bytes: bytes | None = None, redacted_argv: Sequence[str] | None = None) -> ProcessResult:
|
if not argv or timeout_ms <= 0:
|
raise ValueError("argv and positive timeout required")
|
shown = tuple(redacted_argv or argv)
|
started_at = self._utc_now()
|
begin = time.monotonic()
|
deadline = begin + timeout_ms / 1000
|
creationflags = subprocess.CREATE_NEW_PROCESS_GROUP if os.name == "nt" else 0
|
try:
|
proc = subprocess.Popen(
|
tuple(argv), cwd=str(cwd) if cwd else None, stdin=subprocess.PIPE if stdin_bytes is not None else subprocess.DEVNULL,
|
stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=False, creationflags=creationflags,
|
start_new_session=(os.name != "nt"),
|
)
|
except OSError:
|
return ProcessResult(shown, False, None, None, False, False, True, False, b"", b"", False,
|
False, started_at, self._utc_now(), round((time.monotonic()-begin)*1000),
|
ErrorCode.PROCESS_START_FAILED)
|
timed_out = False
|
terminate_issued = False
|
liveness_unknown = False
|
try:
|
termination_reserve_ms = min(self.kill_wait_ms, max(100, min(1_000, timeout_ms // 5)),
|
max(1, timeout_ms - 1))
|
normal_wait = max(0.001, (timeout_ms - termination_reserve_ms) / 1000)
|
out, err = proc.communicate(stdin_bytes, timeout=normal_wait)
|
except subprocess.TimeoutExpired:
|
timed_out = True
|
terminate_issued = True
|
self._terminate_tree(proc, max(0.001, deadline - time.monotonic()))
|
try:
|
kill_remaining = max(0.0, deadline - time.monotonic())
|
if kill_remaining <= 0:
|
raise subprocess.TimeoutExpired(tuple(argv), timeout_ms / 1000)
|
out, err = proc.communicate(timeout=min(self.kill_wait_ms / 1000, kill_remaining))
|
except subprocess.TimeoutExpired:
|
out, err = b"", b""
|
liveness_unknown = True
|
try:
|
proc.kill()
|
except OSError:
|
pass
|
elapsed = round((time.monotonic() - begin) * 1000)
|
out_trim, out_trunc = self._cap(out)
|
err_trim, err_trunc = self._cap(err)
|
exited = proc.poll() is not None
|
code = None
|
if liveness_unknown:
|
code = ErrorCode.PROCESS_LIVENESS_UNKNOWN
|
elif timed_out:
|
code = ErrorCode.PROCESS_TIMEOUT
|
elif out_trunc or err_trunc:
|
code = ErrorCode.PROCESS_OUTPUT_LIMIT
|
return ProcessResult(shown, True, proc.pid, proc.returncode if exited else None, timed_out, terminate_issued,
|
exited, liveness_unknown, out_trim, err_trim, out_trunc, err_trunc,
|
started_at, self._utc_now(), elapsed, code)
|
|
def _cap(self, data: bytes) -> tuple[bytes, bool]:
|
if len(data) <= self.max_stream_bytes:
|
return data, False
|
return data[:self.max_stream_bytes], True
|
|
@staticmethod
|
def _terminate_tree(proc: subprocess.Popen[bytes], timeout_seconds: float = 5.0) -> None:
|
try:
|
if os.name == "nt":
|
subprocess.run(("taskkill", "/PID", str(proc.pid), "/T", "/F"), shell=False,
|
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
|
timeout=max(0.001, min(5.0, timeout_seconds)))
|
else:
|
os.killpg(proc.pid, signal.SIGKILL)
|
except (OSError, subprocess.SubprocessError):
|
try:
|
proc.kill()
|
except OSError:
|
pass
|
