"""Long-running commands: start one, read what it has printed, feed it input, stop it.""" from __future__ import annotations import os import select import signal import subprocess import threading import time from dataclasses import dataclass, field from pathlib import Path OUTPUT_BUFFER_BYTES = 1 << 20 OUTPUT_CAPTURE_BYTES = 8 << 20 STOP_GRACE_SECONDS = 5.0 MAX_LIVE_PROCESSES = 8 STDIN_WRITE_TIMEOUT_SECONDS = 5.0 def _kill_group(pgid: int) -> None: for sig in (signal.SIGTERM, signal.SIGKILL): try: os.killpg(pgid, sig) except ProcessLookupError: return time.sleep(0.2) def run_bounded( argv: list[str], *, cwd: Path, env: dict[str, str], timeout: float, max_bytes: int = OUTPUT_CAPTURE_BYTES ) -> tuple[int | None, str, bool]: """Run to completion, capping wall time AND captured bytes. `capture_output=True` would buffer without limit: a command printing `yes` reaches gigabytes of RSS in seconds, which is a one-call denial of service against the host.""" popen = subprocess.Popen( argv, cwd=cwd, env=env, stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, start_new_session=True, ) pgid = os.getpgid(popen.pid) stream = popen.stdout assert stream is not None os.set_blocking(stream.fileno(), False) collected = bytearray() truncated = False deadline = time.monotonic() + timeout timed_out = False try: while True: if popen.poll() is not None and not select.select([stream], [], [], 0)[0]: break if time.monotonic() >= deadline: timed_out = True break if not select.select([stream], [], [], 0.05)[0]: continue chunk = stream.read(65536) if chunk is None: continue if not chunk: if popen.poll() is not None: break continue room = max_bytes - len(collected) if room <= 0: truncated = True break collected += chunk[:room] finally: if popen.poll() is None: _kill_group(pgid) stream.close() popen.wait() exit_code = None if timed_out or truncated else popen.returncode return exit_code, collected.decode("utf-8", "replace"), truncated @dataclass class ManagedProcess: process_id: str command: str cwd: str popen: subprocess.Popen[bytes] started_at: float # Captured at start: reading it later races the reaper, and after PID reuse # os.getpgid would name a group this registry never started. pgid: int = 0 _lock: threading.Lock = field(default_factory=threading.Lock) _buffer: bytearray = field(default_factory=bytearray) _dropped: int = 0 @property def running(self) -> bool: return self.popen.poll() is None @property def exit_code(self) -> int | None: return self.popen.poll() def _pump(self) -> None: stream = self.popen.stdout assert stream is not None # read1 returns whatever has arrived; read(n) would block until n bytes, # hiding a running process's output until it produced 4 KiB or exited. for chunk in iter(lambda: stream.read1(4096), b""): with self._lock: self._buffer += chunk overflow = len(self._buffer) - OUTPUT_BUFFER_BYTES if overflow > 0: del self._buffer[:overflow] self._dropped += overflow stream.close() self.popen.wait() def read(self, since: int) -> tuple[str, int, int]: """Output from absolute offset `since`, the next offset, and how many bytes were discarded before the returned slice because the buffer overflowed.""" with self._lock: start = max(since, self._dropped, 0) data = bytes(self._buffer[start - self._dropped :]) return data.decode("utf-8", "replace"), start + len(data), start - since def write_stdin(self, text: str) -> None: if not self.running: raise ValueError(f"process {self.process_id} has already exited") stdin = self.popen.stdin if stdin is None or stdin.closed: raise ValueError(f"process {self.process_id} has no open stdin") # A process that never reads stdin fills the 64 KiB pipe and a blocking # write hangs forever. On one event loop that freezes every other tool, # including the one that would stop this process. payload = text.encode("utf-8") fd = stdin.fileno() os.set_blocking(fd, False) deadline = time.monotonic() + STDIN_WRITE_TIMEOUT_SECONDS while payload: if time.monotonic() >= deadline: raise ValueError( f"process {self.process_id} is not reading stdin; " f"{len(payload)} bytes were not delivered" ) if not select.select([], [fd], [], 0.05)[1]: continue try: payload = payload[os.write(fd, payload) :] except BlockingIOError: continue except BrokenPipeError: raise ValueError(f"process {self.process_id} closed its stdin") from None def stop(self) -> str: if not self.running: return f"process {self.process_id} had already exited ({self.exit_code})" os.killpg(self.pgid, signal.SIGTERM) deadline = time.monotonic() + STOP_GRACE_SECONDS while time.monotonic() < deadline: if not self.running: return f"process {self.process_id} stopped" time.sleep(0.1) os.killpg(self.pgid, signal.SIGKILL) self.popen.wait(timeout=STOP_GRACE_SECONDS) return f"process {self.process_id} killed after {STOP_GRACE_SECONDS:.0f}s grace" class ProcessRegistry: def __init__(self) -> None: self._processes: dict[str, ManagedProcess] = {} self._next = 1 self._lock = threading.Lock() def start(self, argv: list[str], *, command: str, cwd: Path, env: dict[str, str]) -> ManagedProcess: live = sum(1 for managed in self.all() if managed.running) if live >= MAX_LIVE_PROCESSES: raise ValueError( f"{live} background processes are already running (limit {MAX_LIVE_PROCESSES}); " "stop one before starting another" ) popen = subprocess.Popen( argv, cwd=cwd, env=env, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, start_new_session=True, ) with self._lock: process_id = f"proc-{self._next}" self._next += 1 managed = ManagedProcess( process_id=process_id, command=command, cwd=str(cwd), popen=popen, started_at=time.time(), pgid=os.getpgid(popen.pid), ) self._processes[process_id] = managed threading.Thread(target=managed._pump, daemon=True).start() return managed def get(self, process_id: str) -> ManagedProcess: try: return self._processes[process_id] except KeyError: raise ValueError(f"no such process: {process_id}") from None def all(self) -> list[ManagedProcess]: with self._lock: return list(self._processes.values()) def stop_all(self) -> None: for managed in self.all(): if managed.running: managed.stop()