Something went wrong. Try again.
My version of https://github.com/vercel-labs/fx
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395from __future__ import annotations
import contextlibimport dataclassesimport jsonimport osimport pathlibimport shleximport signalimport subprocessimport sysimport threadingimport timefrom collections import dequefrom collections.abc import Iterator, Mapping, Sequencefrom typing import TextIO
from scripts.pgso.model import PgsoError, require_empty_stderr as ensure_empty_stderr
COMMAND_HEARTBEAT_SECONDS = 30.0MAX_CAPTURED_OUTPUT_CHARS = 1024 * 1024INHERITED_ENVIRONMENT_KEYS = ( "LANG", "LC_ALL", "LC_CTYPE", "LOGNAME", "PATH", "SHELL", "SYSTEM_VERSION_COMPAT", "TEMP", "TERM", "TMP", "TMPDIR", "TZ", "USER", "__CF_USER_TEXT_ENCODING",)
@dataclasses.dataclass(frozen=True)class CommandResult: argv: tuple[str, ...] returncode: int stdout: str stderr: str elapsed_seconds: float stdout_total_chars: int = 0 stderr_total_chars: int = 0 stdout_truncated: bool = False stderr_truncated: bool = False
class _BoundedCapture: def __init__(self, maximum_chars: int) -> None: self.maximum_chars = maximum_chars self.chunks: deque[str] = deque() self.retained_chars = 0 self.total_chars = 0
@property def truncated(self) -> bool: return self.total_chars > self.retained_chars
def append(self, chunk: str) -> None: self.total_chars += len(chunk) self.chunks.append(chunk) self.retained_chars += len(chunk) while self.retained_chars > self.maximum_chars and self.chunks: excess = self.retained_chars - self.maximum_chars oldest = self.chunks[0] if len(oldest) <= excess: self.chunks.popleft() self.retained_chars -= len(oldest) else: self.chunks[0] = oldest[excess:] self.retained_chars -= excess
def text(self) -> str: retained = "".join(self.chunks) if not self.truncated: return retained omitted = self.total_chars - self.retained_chars return f"[pgso truncated {omitted} earlier characters]\n{retained}"
def emit_progress(message: str) -> None: print(f"[pgso] {message}", flush=True)
def hermetic_environment(home: pathlib.Path) -> dict[str, str]: environment = { key: os.environ[key] for key in INHERITED_ENVIRONMENT_KEYS if key in os.environ } environment["HOME"] = str(home) return environment
@contextlib.contextmanagerdef cancellation_guard() -> Iterator[None]: if threading.current_thread() is not threading.main_thread(): yield return
handled = (signal.SIGINT, signal.SIGTERM) previous = {item: signal.getsignal(item) for item in handled}
def cancel(signum, _frame) -> None: name = signal.Signals(signum).name raise PgsoError(f"pipeline cancelled by {name}")
try: for item in handled: signal.signal(item, cancel) yield finally: for item, handler in previous.items(): signal.signal(item, handler)
def _write_log( path: pathlib.Path, *, argv: tuple[str, ...], elapsed_seconds: float, returncode: int | None, stdout: str, stderr: str, status: str, stdout_total_chars: int | None = None, stderr_total_chars: int | None = None, stdout_truncated: bool = False, stderr_truncated: bool = False,) -> None: payload = { "argv": argv, "command_class": pathlib.Path(argv[0]).name, "elapsed_seconds": elapsed_seconds, "returncode": returncode, "status": status, "stderr": stderr, "stderr_total_chars": ( len(stderr) if stderr_total_chars is None else stderr_total_chars ), "stderr_truncated": stderr_truncated, "stdout": stdout, "stdout_total_chars": ( len(stdout) if stdout_total_chars is None else stdout_total_chars ), "stdout_truncated": stdout_truncated, } path.parent.mkdir(parents=True, exist_ok=True) path.write_text( json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8", )
def _terminate_process_group(process: subprocess.Popen[str]) -> None: try: os.killpg(process.pid, signal.SIGTERM) except ProcessLookupError: pass
try: process.wait(timeout=5) except subprocess.TimeoutExpired: try: os.killpg(process.pid, signal.SIGKILL) except ProcessLookupError: pass process.wait()
def _stream_pipe( pipe: TextIO, sink: TextIO, capture: _BoundedCapture,) -> None: try: while True: chunk = pipe.read(4096) if not chunk: return capture.append(chunk) sink.write(chunk) sink.flush() finally: pipe.close()
def run_checked( argv: Sequence[str | os.PathLike[str]], *, cwd: pathlib.Path, env: Mapping[str, str] | None, timeout_s: float, log_path: pathlib.Path, require_empty_stderr: bool = False, max_capture_chars: int | None = None,) -> CommandResult: argv_tuple = tuple(os.fspath(argument) for argument in argv) if not argv_tuple: raise PgsoError("cannot run an empty command") if timeout_s <= 0: raise PgsoError("command timeout must be positive") capture_limit = MAX_CAPTURED_OUTPUT_CHARS if max_capture_chars is None else max_capture_chars if capture_limit <= 0: raise PgsoError("command capture limit must be positive")
command_display = shlex.join(argv_tuple) emit_progress(f"command started: {command_display} (log: {log_path})") started = time.monotonic() try: process = subprocess.Popen( argv_tuple, cwd=cwd, env=None if env is None else dict(env), stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, encoding="utf-8", errors="replace", start_new_session=True, ) except OSError as error: elapsed_seconds = time.monotonic() - started _write_log( log_path, argv=argv_tuple, elapsed_seconds=elapsed_seconds, returncode=None, stdout="", stderr=str(error), status="spawn_failed", ) emit_progress( f"command spawn failed in {elapsed_seconds:.3f}s: " f"{command_display}: {error}" ) if isinstance(error, FileNotFoundError): raise PgsoError(f"executable not found: {argv_tuple[0]}") from error raise PgsoError(f"could not start command: {error}") from error
if process.stdout is None or process.stderr is None: _terminate_process_group(process) raise PgsoError("command output pipes were not created")
stdout_capture = _BoundedCapture(capture_limit) stderr_capture = _BoundedCapture(capture_limit) stdout_thread = threading.Thread( target=_stream_pipe, args=(process.stdout, sys.stdout, stdout_capture), daemon=True, ) stderr_thread = threading.Thread( target=_stream_pipe, args=(process.stderr, sys.stderr, stderr_capture), daemon=True, ) stdout_thread.start() stderr_thread.start()
# POSIX timed waits poll; observe exit before supervision and log-drain delays. completed = threading.Event() completed_at = started wait_error: BaseException | None = None
def wait_for_exit() -> None: nonlocal completed_at, wait_error try: process.wait() except BaseException as error: wait_error = error finally: completed_at = time.monotonic() completed.set()
wait_thread = threading.Thread(target=wait_for_exit, daemon=True) timed_out = False deadline = time.monotonic() + timeout_s next_heartbeat = time.monotonic() + COMMAND_HEARTBEAT_SECONDS try: wait_thread.start() while process.poll() is None: now = time.monotonic() if now >= deadline: timed_out = True _terminate_process_group(process) break wait_seconds = max(0.001, min(deadline, next_heartbeat) - now) if completed.wait(wait_seconds): break now = time.monotonic() if now >= deadline: timed_out = True _terminate_process_group(process) break if now >= next_heartbeat: emit_progress( f"command still running after {now - started:.1f}s: " f"{command_display} (log: {log_path})" ) while next_heartbeat <= now: next_heartbeat += COMMAND_HEARTBEAT_SECONDS wait_thread.join() if wait_error is not None: raise wait_error except BaseException: _terminate_process_group(process) if wait_thread.ident is not None: wait_thread.join() stdout_thread.join() stderr_thread.join() elapsed_seconds = time.monotonic() - started _write_log( log_path, argv=argv_tuple, elapsed_seconds=elapsed_seconds, returncode=process.returncode, stdout=stdout_capture.text(), stderr=stderr_capture.text(), status="cancelled", stdout_total_chars=stdout_capture.total_chars, stderr_total_chars=stderr_capture.total_chars, stdout_truncated=stdout_capture.truncated, stderr_truncated=stderr_capture.truncated, ) emit_progress(f"command cancelled after {elapsed_seconds:.3f}s: {command_display}") raise
elapsed_seconds = completed_at - started stdout_thread.join() stderr_thread.join() stdout = stdout_capture.text() stderr = stderr_capture.text()
def write_status(status: str) -> None: _write_log( log_path, argv=argv_tuple, elapsed_seconds=elapsed_seconds, returncode=process.returncode, stdout=stdout, stderr=stderr, status=status, stdout_total_chars=stdout_capture.total_chars, stderr_total_chars=stderr_capture.total_chars, stdout_truncated=stdout_capture.truncated, stderr_truncated=stderr_capture.truncated, )
if timed_out: write_status("timed_out") emit_progress( f"command timed out in {elapsed_seconds:.3f}s: {command_display}" ) raise PgsoError( f"command timed out after {timeout_s:g} seconds: {argv_tuple[0]}" )
result = CommandResult( argv=argv_tuple, returncode=process.returncode, stdout=stdout, stderr=stderr, elapsed_seconds=elapsed_seconds, stdout_total_chars=stdout_capture.total_chars, stderr_total_chars=stderr_capture.total_chars, stdout_truncated=stdout_capture.truncated, stderr_truncated=stderr_capture.truncated, ) if result.returncode != 0: write_status("failed") emit_progress( f"command failed in {elapsed_seconds:.3f}s: {command_display} " f"(exit: {result.returncode})" ) raise PgsoError( f"command failed with exit code {result.returncode}: {argv_tuple[0]}" )
if require_empty_stderr and stderr: write_status("unexpected_stderr") emit_progress( f"command rejected stderr in {elapsed_seconds:.3f}s: {command_display}" ) ensure_empty_stderr("command", stderr)
write_status("passed") emit_progress(f"command passed in {elapsed_seconds:.3f}s: {command_display}") return result