293 lines
10 KiB
Python
293 lines
10 KiB
Python
"""Windows-specific functionality for stdio transport operations."""
|
|
|
|
import logging
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import weakref
|
|
from contextlib import suppress
|
|
from pathlib import Path
|
|
from typing import BinaryIO, TextIO, TypeAlias, cast
|
|
|
|
import anyio
|
|
from anyio.abc import Process
|
|
from anyio.streams.file import FileReadStream, FileWriteStream
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Windows-specific imports for Job Objects
|
|
if sys.platform == "win32":
|
|
import msvcrt
|
|
|
|
import pywintypes
|
|
import win32api
|
|
import win32con
|
|
import win32job
|
|
else:
|
|
# Type stubs for non-Windows platforms
|
|
win32api = None
|
|
win32con = None
|
|
msvcrt = None
|
|
win32job = None
|
|
pywintypes = None
|
|
|
|
|
|
def rebind_std_handle_to_fd(fd: int) -> None:
|
|
"""Points the Win32 standard-handle slot for fd 0, 1, or 2 at fd's current OS handle.
|
|
|
|
os.dup2 updates only the CRT descriptor table; subprocess handle inheritance
|
|
reads the Win32 slot, so it must be repointed too.
|
|
|
|
Raises:
|
|
OSError: The slot could not be set.
|
|
"""
|
|
if sys.platform != "win32" or not win32api or not msvcrt or not pywintypes:
|
|
return
|
|
std_ids = {0: win32api.STD_INPUT_HANDLE, 1: win32api.STD_OUTPUT_HANDLE, 2: win32api.STD_ERROR_HANDLE}
|
|
try:
|
|
win32api.SetStdHandle(std_ids[fd], msvcrt.get_osfhandle(fd))
|
|
except pywintypes.error as exc:
|
|
# Normalized so callers' OSError-based best-effort handling covers it.
|
|
raise OSError(f"SetStdHandle failed for fd {fd}") from exc
|
|
|
|
|
|
# How often FallbackProcess polls the underlying Popen for exit.
|
|
_EXIT_POLL_INTERVAL = 0.01
|
|
|
|
# Job Object handle per spawned process, for tree termination at shutdown.
|
|
# Values stay pywin32 PyHANDLEs: if no pop site ever runs, the dying weak entry
|
|
# drops the last reference and the PyHANDLE destructor closes the handle, which
|
|
# is what makes KILL_ON_JOB_CLOSE reap an abandoned tree.
|
|
_process_jobs: "weakref.WeakKeyDictionary[Process | FallbackProcess, object]" = weakref.WeakKeyDictionary()
|
|
|
|
|
|
def get_windows_executable_command(command: str) -> str:
|
|
"""Resolves the command to a Windows executable path.
|
|
|
|
Tries the bare name first, then the common script extensions (.cmd, .bat,
|
|
.exe, .ps1).
|
|
"""
|
|
try:
|
|
if command_path := shutil.which(command):
|
|
return command_path
|
|
|
|
for ext in [".cmd", ".bat", ".exe", ".ps1"]:
|
|
ext_version = f"{command}{ext}"
|
|
if ext_path := shutil.which(ext_version):
|
|
return ext_path
|
|
|
|
return command
|
|
except OSError:
|
|
return command # path probing failed (permissions, broken symlinks)
|
|
|
|
|
|
class FallbackProcess:
|
|
"""Async wrapper around subprocess.Popen for SelectorEventLoop.
|
|
|
|
Windows event loops without async subprocess support get this Popen-backed
|
|
fallback, with anyio file streams wrapping the pipes.
|
|
"""
|
|
|
|
def __init__(self, popen_obj: subprocess.Popen[bytes]) -> None:
|
|
self.popen: subprocess.Popen[bytes] = popen_obj
|
|
stdin = popen_obj.stdin
|
|
stdout = popen_obj.stdout
|
|
|
|
self.stdin = FileWriteStream(cast(BinaryIO, stdin)) if stdin else None
|
|
self.stdout = FileReadStream(cast(BinaryIO, stdout)) if stdout else None
|
|
|
|
async def wait(self) -> int:
|
|
"""Waits for exit by polling the Popen.
|
|
|
|
A thread blocked in Popen.wait() cannot be cancelled by anyio, which
|
|
would defeat every timeout placed around this call.
|
|
"""
|
|
while (returncode := self.popen.poll()) is None:
|
|
await anyio.sleep(_EXIT_POLL_INTERVAL)
|
|
return returncode
|
|
|
|
def terminate(self) -> None:
|
|
"""Terminates the subprocess."""
|
|
self.popen.terminate()
|
|
|
|
def kill(self) -> None:
|
|
"""Kills the subprocess (on Windows the same hard kill as terminate)."""
|
|
self.popen.kill()
|
|
|
|
@property
|
|
def pid(self) -> int:
|
|
"""Returns the process ID."""
|
|
return self.popen.pid
|
|
|
|
@property
|
|
def returncode(self) -> int | None:
|
|
"""The exit code, or None while the process is still running.
|
|
|
|
Polls the Popen so death is observable without anyone calling wait().
|
|
"""
|
|
return self.popen.poll()
|
|
|
|
|
|
# The process handle stdio_client drives: anyio's Process, or the Popen-backed
|
|
# fallback used on Windows event loops without async subprocess support.
|
|
ServerProcess: TypeAlias = Process | FallbackProcess
|
|
|
|
|
|
async def create_windows_process(
|
|
command: str,
|
|
args: list[str],
|
|
env: dict[str, str] | None = None,
|
|
errlog: TextIO | None = sys.stderr,
|
|
cwd: Path | str | None = None,
|
|
) -> Process | FallbackProcess:
|
|
"""Creates a subprocess with Job Object support for tree termination.
|
|
|
|
Spawns via anyio's open_process; event loops without async subprocess
|
|
support (notably the SelectorEventLoop) raise NotImplementedError, in which
|
|
case the spawn falls back to a Popen-backed FallbackProcess. Either way the
|
|
process is then assigned to a Job Object so its children can be terminated
|
|
with it; children spawned before the assignment completes are not captured
|
|
(see the inline note below).
|
|
|
|
Returns:
|
|
Process | FallbackProcess: The spawned process with async stdin/stdout streams.
|
|
"""
|
|
try:
|
|
process = await anyio.open_process(
|
|
[command, *args],
|
|
env=env,
|
|
# Ensure we don't create console windows for each process
|
|
creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
|
|
stderr=errlog,
|
|
cwd=cwd,
|
|
)
|
|
except NotImplementedError:
|
|
# Windows event loops without async subprocess support (SelectorEventLoop)
|
|
process = await _create_windows_fallback_process(command, args, env, errlog, cwd)
|
|
|
|
# Children spawned before the assignment completes land outside the job
|
|
# (membership is inherited at CreateProcess, never acquired retroactively);
|
|
# if that ever bites, the fix is a CREATE_SUSPENDED spawn -> assign -> resume.
|
|
job = _create_job_object()
|
|
_maybe_assign_process_to_job(process, job)
|
|
return process
|
|
|
|
|
|
async def _create_windows_fallback_process(
|
|
command: str,
|
|
args: list[str],
|
|
env: dict[str, str] | None = None,
|
|
errlog: TextIO | None = sys.stderr,
|
|
cwd: Path | str | None = None,
|
|
) -> FallbackProcess:
|
|
"""Spawns via subprocess.Popen and wraps it in FallbackProcess."""
|
|
popen_obj = subprocess.Popen(
|
|
[command, *args],
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
stderr=errlog,
|
|
env=env,
|
|
cwd=cwd,
|
|
bufsize=0, # Unbuffered output
|
|
creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
|
|
)
|
|
return FallbackProcess(popen_obj)
|
|
|
|
|
|
def _create_job_object() -> object | None:
|
|
"""Creates a Windows Job Object configured to terminate all its processes when closed."""
|
|
if sys.platform != "win32" or not win32api or not win32job:
|
|
return None
|
|
|
|
job = None
|
|
try:
|
|
job = win32job.CreateJobObject(None, "")
|
|
extended_info = win32job.QueryInformationJobObject(job, win32job.JobObjectExtendedLimitInformation)
|
|
|
|
extended_info["BasicLimitInformation"]["LimitFlags"] |= win32job.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
|
|
win32job.SetInformationJobObject(job, win32job.JobObjectExtendedLimitInformation, extended_info)
|
|
return job
|
|
except pywintypes.error:
|
|
logger.warning("Failed to create Job Object for process tree management", exc_info=True)
|
|
# If creation succeeded but configuration failed, close the handle now.
|
|
if job is not None:
|
|
_close_job_handle(job)
|
|
return None
|
|
|
|
|
|
def _maybe_assign_process_to_job(process: Process | FallbackProcess, job: object | None) -> None:
|
|
"""Assigns the process to the job and records it for tree termination.
|
|
|
|
On any failure the job handle is closed instead.
|
|
"""
|
|
if job is None:
|
|
return
|
|
|
|
if sys.platform != "win32" or not win32api or not win32con or not win32job:
|
|
return
|
|
|
|
try:
|
|
process_handle = win32api.OpenProcess(
|
|
win32con.PROCESS_SET_QUOTA | win32con.PROCESS_TERMINATE, False, process.pid
|
|
)
|
|
if not process_handle:
|
|
raise pywintypes.error(0, "OpenProcess", "Failed to open process handle")
|
|
|
|
try:
|
|
win32job.AssignProcessToJobObject(job, process_handle)
|
|
finally:
|
|
win32api.CloseHandle(process_handle)
|
|
# Record only after the CloseHandle above succeeded: had it failed, the
|
|
# except below would close the job and KILL_ON_JOB_CLOSE takes the server.
|
|
_process_jobs[process] = job
|
|
except pywintypes.error:
|
|
logger.warning("Failed to assign process %d to Job Object", process.pid, exc_info=True)
|
|
_close_job_handle(job)
|
|
|
|
|
|
def close_process_job(process: Process | FallbackProcess) -> None:
|
|
"""Closes the process's Job Object handle, if it still has one.
|
|
|
|
KILL_ON_JOB_CLOSE makes the close also kill any members still alive,
|
|
deterministically rather than at GC time; a deliberate divergence from
|
|
POSIX, where a graceful server's children are left alive.
|
|
"""
|
|
if sys.platform != "win32":
|
|
return
|
|
|
|
job = _process_jobs.pop(process, None)
|
|
if job is not None:
|
|
_close_job_handle(job)
|
|
|
|
|
|
async def terminate_windows_process_tree(process: Process | FallbackProcess) -> None:
|
|
"""Terminates the process's job, or just the process if it has no job.
|
|
|
|
Job termination is an immediate hard kill of every member. Windows has no
|
|
tree-wide SIGTERM; the stdin-close grace period is the server's chance to
|
|
exit cleanly.
|
|
"""
|
|
if sys.platform != "win32":
|
|
return
|
|
|
|
job = _process_jobs.pop(process, None)
|
|
if job is not None and win32job:
|
|
try:
|
|
with suppress(pywintypes.error): # the job might already be terminated
|
|
win32job.TerminateJobObject(job, 1)
|
|
finally:
|
|
_close_job_handle(job)
|
|
|
|
# The process may have no job (creation or assignment failed); kill it directly too.
|
|
try:
|
|
process.terminate()
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def _close_job_handle(job: object) -> None:
|
|
"""Closes a Job Object handle, tolerating one that is already closed."""
|
|
if win32api and pywintypes:
|
|
with suppress(pywintypes.error):
|
|
win32api.CloseHandle(job)
|