452 lines
14 KiB
Python
452 lines
14 KiB
Python
"""Built-in executable workflow steps."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import signal
|
|
import subprocess
|
|
import tempfile
|
|
import time
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import IO, Mapping
|
|
|
|
import psutil
|
|
|
|
from gyxx_flow.core.context import RunContext
|
|
from gyxx_flow.core.exit_codes import (
|
|
COOKIE_SKIP_EXIT_CODE,
|
|
IDEMPOTENT_REPLAY_EXIT_CODE,
|
|
)
|
|
from gyxx_flow.core.text import bounded_head_tail
|
|
|
|
_CREATE_SUSPENDED = 0x00000004
|
|
_TERMINATION_GRACE_SECONDS = 0.5
|
|
_TERMINATION_CONFIRM_SECONDS = 5.0
|
|
_FAILURE_OUTPUT_MAX_CHARS = 16_000
|
|
|
|
|
|
class _WindowsJob:
|
|
"""A private Job Object that contains exactly one command process tree."""
|
|
|
|
def __init__(self) -> None:
|
|
if os.name != "nt":
|
|
raise OSError("Windows Job Objects are unavailable on this platform")
|
|
|
|
import ctypes
|
|
from ctypes import wintypes
|
|
|
|
class BasicAccountingInformation(ctypes.Structure):
|
|
_fields_ = [
|
|
("TotalUserTime", ctypes.c_longlong),
|
|
("TotalKernelTime", ctypes.c_longlong),
|
|
("ThisPeriodTotalUserTime", ctypes.c_longlong),
|
|
("ThisPeriodTotalKernelTime", ctypes.c_longlong),
|
|
("TotalPageFaultCount", wintypes.DWORD),
|
|
("TotalProcesses", wintypes.DWORD),
|
|
("ActiveProcesses", wintypes.DWORD),
|
|
("TotalTerminatedProcesses", wintypes.DWORD),
|
|
]
|
|
|
|
kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
|
|
kernel32.CreateJobObjectW.argtypes = (wintypes.LPVOID, wintypes.LPCWSTR)
|
|
kernel32.CreateJobObjectW.restype = wintypes.HANDLE
|
|
kernel32.OpenProcess.argtypes = (
|
|
wintypes.DWORD,
|
|
wintypes.BOOL,
|
|
wintypes.DWORD,
|
|
)
|
|
kernel32.OpenProcess.restype = wintypes.HANDLE
|
|
kernel32.AssignProcessToJobObject.argtypes = (
|
|
wintypes.HANDLE,
|
|
wintypes.HANDLE,
|
|
)
|
|
kernel32.AssignProcessToJobObject.restype = wintypes.BOOL
|
|
kernel32.TerminateJobObject.argtypes = (wintypes.HANDLE, wintypes.UINT)
|
|
kernel32.TerminateJobObject.restype = wintypes.BOOL
|
|
kernel32.QueryInformationJobObject.argtypes = (
|
|
wintypes.HANDLE,
|
|
ctypes.c_int,
|
|
wintypes.LPVOID,
|
|
wintypes.DWORD,
|
|
wintypes.LPVOID,
|
|
)
|
|
kernel32.QueryInformationJobObject.restype = wintypes.BOOL
|
|
kernel32.CloseHandle.argtypes = (wintypes.HANDLE,)
|
|
kernel32.CloseHandle.restype = wintypes.BOOL
|
|
|
|
handle = kernel32.CreateJobObjectW(None, None)
|
|
if not handle:
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
self._ctypes = ctypes
|
|
self._kernel32 = kernel32
|
|
self._accounting_type = BasicAccountingInformation
|
|
self._handle = handle
|
|
|
|
def assign(self, pid: int) -> None:
|
|
process_set_quota = 0x0100
|
|
process_terminate = 0x0001
|
|
process_handle = self._kernel32.OpenProcess(
|
|
process_set_quota | process_terminate,
|
|
False,
|
|
pid,
|
|
)
|
|
if not process_handle:
|
|
raise self._ctypes.WinError(self._ctypes.get_last_error())
|
|
try:
|
|
if not self._kernel32.AssignProcessToJobObject(
|
|
self._handle,
|
|
process_handle,
|
|
):
|
|
raise self._ctypes.WinError(self._ctypes.get_last_error())
|
|
finally:
|
|
self._kernel32.CloseHandle(process_handle)
|
|
|
|
def terminate(self, exit_code: int) -> None:
|
|
if not self._kernel32.TerminateJobObject(self._handle, exit_code):
|
|
raise self._ctypes.WinError(self._ctypes.get_last_error())
|
|
|
|
def wait_empty(self, timeout_seconds: float) -> None:
|
|
deadline = time.monotonic() + timeout_seconds
|
|
while True:
|
|
information = self._accounting_type()
|
|
if not self._kernel32.QueryInformationJobObject(
|
|
self._handle,
|
|
1,
|
|
self._ctypes.byref(information),
|
|
self._ctypes.sizeof(information),
|
|
None,
|
|
):
|
|
raise self._ctypes.WinError(self._ctypes.get_last_error())
|
|
if information.ActiveProcesses == 0:
|
|
return
|
|
if time.monotonic() >= deadline:
|
|
raise RuntimeError("command process tree did not terminate")
|
|
time.sleep(0.01)
|
|
|
|
def close(self) -> None:
|
|
if self._handle:
|
|
self._kernel32.CloseHandle(self._handle)
|
|
self._handle = None
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class StepExecution:
|
|
"""One action invocation outcome before retry aggregation."""
|
|
|
|
exit_code: int
|
|
stdout: str = ""
|
|
stderr: str = ""
|
|
error: str | None = None
|
|
skipped: bool = False
|
|
reason: str | None = None
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class CommandStep:
|
|
"""Execute an argv command without shell parsing or implicit process context."""
|
|
|
|
argv: tuple[str, ...]
|
|
cwd: Path
|
|
env: Mapping[str, str]
|
|
|
|
def __post_init__(self) -> None:
|
|
object.__setattr__(self, "argv", tuple(self.argv))
|
|
object.__setattr__(self, "cwd", Path(self.cwd))
|
|
object.__setattr__(self, "env", dict(self.env))
|
|
if not self.argv or not all(
|
|
isinstance(argument, str) and argument and "\x00" not in argument
|
|
for argument in self.argv
|
|
):
|
|
raise ValueError("argv must contain non-empty strings")
|
|
if not all(
|
|
isinstance(key, str)
|
|
and key
|
|
and "\x00" not in key
|
|
and isinstance(value, str)
|
|
and "\x00" not in value
|
|
for key, value in self.env.items()
|
|
):
|
|
raise ValueError("env must contain valid string keys and values")
|
|
|
|
def execute(
|
|
self,
|
|
*,
|
|
context: RunContext,
|
|
timeout_seconds: float | None,
|
|
dry_run: bool,
|
|
) -> StepExecution:
|
|
del context
|
|
if dry_run:
|
|
return StepExecution(exit_code=0, skipped=True, reason="dry-run")
|
|
with (
|
|
tempfile.TemporaryFile(
|
|
mode="w+", encoding="utf-8", errors="replace", newline=""
|
|
) as stdout_file,
|
|
tempfile.TemporaryFile(
|
|
mode="w+", encoding="utf-8", errors="replace", newline=""
|
|
) as stderr_file,
|
|
):
|
|
try:
|
|
process, process_scope = _start_process(
|
|
list(self.argv),
|
|
cwd=self.cwd,
|
|
env=dict(self.env),
|
|
stdout=stdout_file,
|
|
stderr=stderr_file,
|
|
)
|
|
except OSError as exc:
|
|
return StepExecution(exit_code=127, error=f"cannot start command: {exc}")
|
|
|
|
try:
|
|
return_code = process.wait(timeout=timeout_seconds)
|
|
except subprocess.TimeoutExpired as exc:
|
|
_terminate_process_tree(process, process_scope)
|
|
stdout = _read_output(stdout_file)
|
|
stderr = _read_output(stderr_file)
|
|
return StepExecution(
|
|
exit_code=124,
|
|
stdout=stdout,
|
|
stderr=stderr,
|
|
error=f"command timed out after {exc.timeout} seconds",
|
|
)
|
|
except BaseException:
|
|
_terminate_process_tree(process, process_scope)
|
|
raise
|
|
finally:
|
|
if process_scope is not None:
|
|
process_scope.close()
|
|
|
|
stdout = _read_output(stdout_file)
|
|
stderr = _read_output(stderr_file)
|
|
|
|
if return_code == COOKIE_SKIP_EXIT_CODE:
|
|
return StepExecution(
|
|
exit_code=return_code,
|
|
stdout=stdout,
|
|
stderr=stderr,
|
|
skipped=True,
|
|
reason="cookie-preflight",
|
|
)
|
|
|
|
if return_code == IDEMPOTENT_REPLAY_EXIT_CODE:
|
|
return StepExecution(
|
|
exit_code=return_code,
|
|
stdout=stdout,
|
|
stderr=stderr,
|
|
skipped=True,
|
|
reason="idempotent-replay",
|
|
)
|
|
|
|
error = None
|
|
if return_code != 0:
|
|
error = _command_failure_detail(
|
|
stdout,
|
|
stderr,
|
|
return_code=return_code,
|
|
)
|
|
return StepExecution(
|
|
exit_code=return_code,
|
|
stdout=stdout,
|
|
stderr=stderr,
|
|
error=error,
|
|
)
|
|
|
|
|
|
def _start_process(
|
|
argv: list[str],
|
|
*,
|
|
cwd: Path,
|
|
env: dict[str, str],
|
|
stdout: IO[str],
|
|
stderr: IO[str],
|
|
) -> tuple[subprocess.Popen[str], _WindowsJob | None]:
|
|
effective_env = dict(env)
|
|
if os.name == "nt":
|
|
# Python launched with redirected stdio can otherwise inherit the
|
|
# Windows ANSI code page (typically GBK), while the parent streams
|
|
# are opened as UTF-8. Keep child diagnostics readable in journals
|
|
# and the console, especially for Chinese business errors.
|
|
effective_env.setdefault("PYTHONIOENCODING", "utf-8")
|
|
common = {
|
|
"cwd": cwd,
|
|
"env": effective_env,
|
|
"shell": False,
|
|
"stdout": stdout,
|
|
"stderr": stderr,
|
|
"text": True,
|
|
"encoding": "utf-8",
|
|
"errors": "replace",
|
|
}
|
|
if os.name != "nt":
|
|
return subprocess.Popen(argv, start_new_session=True, **common), None
|
|
|
|
job = _WindowsJob()
|
|
process: subprocess.Popen[str] | None = None
|
|
assigned = False
|
|
try:
|
|
process = subprocess.Popen(
|
|
argv,
|
|
creationflags=(
|
|
_CREATE_SUSPENDED | subprocess.CREATE_NEW_PROCESS_GROUP
|
|
),
|
|
**common,
|
|
)
|
|
job.assign(process.pid)
|
|
assigned = True
|
|
psutil.Process(process.pid).resume()
|
|
return process, job
|
|
except BaseException:
|
|
if process is not None:
|
|
if assigned:
|
|
try:
|
|
job.terminate(127)
|
|
job.wait_empty(_TERMINATION_CONFIRM_SECONDS)
|
|
except (OSError, RuntimeError):
|
|
pass
|
|
else:
|
|
try:
|
|
process.kill()
|
|
except OSError:
|
|
pass
|
|
try:
|
|
process.wait(timeout=_TERMINATION_CONFIRM_SECONDS)
|
|
except (OSError, subprocess.TimeoutExpired):
|
|
pass
|
|
job.close()
|
|
raise
|
|
|
|
|
|
def start_managed_process(
|
|
argv: list[str],
|
|
*,
|
|
cwd: Path,
|
|
env: dict[str, str],
|
|
stdout: IO[str],
|
|
stderr: IO[str],
|
|
) -> tuple[subprocess.Popen[str], _WindowsJob | None]:
|
|
"""Start a process in a scope that can terminate its complete child tree."""
|
|
|
|
return _start_process(
|
|
argv,
|
|
cwd=cwd,
|
|
env=env,
|
|
stdout=stdout,
|
|
stderr=stderr,
|
|
)
|
|
|
|
|
|
def _terminate_process_tree(
|
|
process: subprocess.Popen[str],
|
|
process_scope: _WindowsJob | None,
|
|
) -> None:
|
|
tracked = _process_tree_snapshot(process.pid)
|
|
if os.name == "nt":
|
|
if process_scope is None:
|
|
raise RuntimeError("Windows command process is missing its Job Object")
|
|
process_scope.terminate(124)
|
|
process_scope.wait_empty(_TERMINATION_CONFIRM_SECONDS)
|
|
else:
|
|
_terminate_posix_process_group(process.pid)
|
|
|
|
try:
|
|
process.wait(timeout=_TERMINATION_CONFIRM_SECONDS)
|
|
except subprocess.TimeoutExpired as exc:
|
|
raise RuntimeError("command root process did not terminate") from exc
|
|
_confirm_processes_stopped(tracked)
|
|
|
|
|
|
def terminate_managed_process_tree(
|
|
process: subprocess.Popen[str],
|
|
process_scope: _WindowsJob | None,
|
|
) -> None:
|
|
"""Terminate and confirm the complete tree of a managed process."""
|
|
|
|
_terminate_process_tree(process, process_scope)
|
|
|
|
|
|
def _terminate_posix_process_group(process_group_id: int) -> None:
|
|
try:
|
|
os.killpg(process_group_id, signal.SIGTERM)
|
|
except ProcessLookupError:
|
|
return
|
|
except PermissionError as exc:
|
|
raise RuntimeError("cannot terminate command process group") from exc
|
|
|
|
deadline = time.monotonic() + _TERMINATION_GRACE_SECONDS
|
|
while _posix_group_has_live_processes(process_group_id):
|
|
if time.monotonic() >= deadline:
|
|
break
|
|
time.sleep(0.01)
|
|
if _posix_group_has_live_processes(process_group_id):
|
|
try:
|
|
os.killpg(process_group_id, signal.SIGKILL)
|
|
except ProcessLookupError:
|
|
return
|
|
except PermissionError as exc:
|
|
raise RuntimeError("cannot kill command process group") from exc
|
|
|
|
deadline = time.monotonic() + _TERMINATION_CONFIRM_SECONDS
|
|
while _posix_group_has_live_processes(process_group_id):
|
|
if time.monotonic() >= deadline:
|
|
raise RuntimeError("command process group did not terminate")
|
|
time.sleep(0.01)
|
|
|
|
|
|
def _posix_group_has_live_processes(process_group_id: int) -> bool:
|
|
for candidate in psutil.process_iter(("pid", "status")):
|
|
try:
|
|
if (
|
|
os.getpgid(candidate.pid) == process_group_id
|
|
and candidate.info["status"] != psutil.STATUS_ZOMBIE
|
|
):
|
|
return True
|
|
except (OSError, psutil.Error):
|
|
continue
|
|
return False
|
|
|
|
|
|
def _process_tree_snapshot(pid: int) -> tuple[psutil.Process, ...]:
|
|
try:
|
|
root = psutil.Process(pid)
|
|
return (*root.children(recursive=True), root)
|
|
except psutil.Error:
|
|
return ()
|
|
|
|
|
|
def _confirm_processes_stopped(processes: tuple[psutil.Process, ...]) -> None:
|
|
_, alive = psutil.wait_procs(
|
|
list(processes),
|
|
timeout=_TERMINATION_CONFIRM_SECONDS,
|
|
)
|
|
live = []
|
|
for process in alive:
|
|
try:
|
|
if process.is_running() and process.status() != psutil.STATUS_ZOMBIE:
|
|
live.append(process.pid)
|
|
except psutil.Error:
|
|
continue
|
|
if live:
|
|
raise RuntimeError("command descendants did not terminate")
|
|
|
|
|
|
def _read_output(handle) -> str: # type: ignore[no-untyped-def]
|
|
handle.flush()
|
|
handle.seek(0)
|
|
return handle.read()
|
|
|
|
|
|
def _command_failure_detail(
|
|
stdout: str,
|
|
stderr: str,
|
|
*,
|
|
return_code: int,
|
|
) -> str:
|
|
chunks = [chunk for chunk in (stdout.strip(), stderr.strip()) if chunk]
|
|
if not chunks:
|
|
return f"command exited with code {return_code}"
|
|
return bounded_head_tail(
|
|
"\n".join(chunks),
|
|
_FAILURE_OUTPUT_MAX_CHARS,
|
|
)
|