feat: complete production workflow migration

2026-08-06 14:29:57 +08:00
1 parent 7f215e79c4
commit 8df5266abb
@@ -2,12 +2,129 @@

 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 Mapping
+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
+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)
@@ -59,43 +176,257 @@ class CommandStep:
        del context
        if dry_run:
            return StepExecution(exit_code=0, skipped=True, reason="dry-run")
-        try:
-            completed = subprocess.run(
-                list(self.argv),
-                cwd=self.cwd,
-                env=dict(self.env),
-                shell=False,
-                capture_output=True,
-                text=True,
-                timeout=timeout_seconds,
-                check=False,
-            )
-        except subprocess.TimeoutExpired as exc:
-            stdout = _to_text(exc.stdout)
-            stderr = _to_text(exc.stderr)
+        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=124,
+                exit_code=return_code,
                stdout=stdout,
                stderr=stderr,
-                error=f"command timed out after {exc.timeout} seconds",
+                skipped=True,
+                reason="cookie-preflight",
            )
-        except OSError as exc:
-            return StepExecution(exit_code=127, error=f"cannot start command: {exc}")

        error = None
-        if completed.returncode != 0:
-            error = completed.stderr.strip() or f"command exited with code {completed.returncode}"
+        if return_code != 0:
+            error = _command_failure_detail(
+                stdout,
+                stderr,
+                return_code=return_code,
+            )
        return StepExecution(
-            exit_code=completed.returncode,
-            stdout=completed.stdout,
-            stderr=completed.stderr,
+            exit_code=return_code,
+            stdout=stdout,
+            stderr=stderr,
            error=error,
        )


-def _to_text(value: str | bytes | None) -> str:
-    if value is None:
-        return ""
-    if isinstance(value, bytes):
-        return value.decode(errors="replace")
-    return value
+def _start_process(
+    argv: list[str],
+    *,
+    cwd: Path,
+    env: dict[str, str],
+    stdout: IO[str],
+    stderr: IO[str],
+) -> tuple[subprocess.Popen[str], _WindowsJob | None]:
+    common = {
+        "cwd": cwd,
+        "env": 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,
+    )