File size: 4,700 Bytes
5655a42 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | """Contain each managed Windows router tree without adopting the owner process."""
from __future__ import annotations
import ctypes
from ctypes import wintypes
import subprocess
import sys
import threading
import psutil
class _BasicLimits(ctypes.Structure):
_fields_ = [
("PerProcessUserTimeLimit", ctypes.c_longlong),
("PerJobUserTimeLimit", ctypes.c_longlong),
("LimitFlags", wintypes.DWORD),
("MinimumWorkingSetSize", ctypes.c_size_t),
("MaximumWorkingSetSize", ctypes.c_size_t),
("ActiveProcessLimit", wintypes.DWORD),
("Affinity", ctypes.c_size_t),
("PriorityClass", wintypes.DWORD),
("SchedulingClass", wintypes.DWORD),
]
class _IoCounters(ctypes.Structure):
_fields_ = [(name, ctypes.c_ulonglong) for name in (
"ReadOperationCount", "WriteOperationCount", "OtherOperationCount",
"ReadTransferCount", "WriteTransferCount", "OtherTransferCount",
)]
class _ExtendedLimits(ctypes.Structure):
_fields_ = [
("BasicLimitInformation", _BasicLimits),
("IoInfo", _IoCounters),
("ProcessMemoryLimit", ctypes.c_size_t),
("JobMemoryLimit", ctypes.c_size_t),
("PeakProcessMemoryUsed", ctypes.c_size_t),
("PeakJobMemoryUsed", ctypes.c_size_t),
]
class _WindowsJob:
def __init__(self):
self._lock = threading.Lock()
self._api = ctypes.WinDLL("kernel32", use_last_error=True)
for name, args, result in (
("CreateJobObjectW", [ctypes.c_void_p, wintypes.LPCWSTR], wintypes.HANDLE),
("SetInformationJobObject", [wintypes.HANDLE, ctypes.c_int,
ctypes.c_void_p, wintypes.DWORD], wintypes.BOOL),
("AssignProcessToJobObject", [wintypes.HANDLE, wintypes.HANDLE], wintypes.BOOL),
("CloseHandle", [wintypes.HANDLE], wintypes.BOOL),
):
fn = getattr(self._api, name)
fn.argtypes = args
fn.restype = result
# NULL security attributes create a non-inheritable, unnamed owner handle.
self._handle = self._api.CreateJobObjectW(None, None)
if not self._handle:
raise ctypes.WinError(ctypes.get_last_error())
try:
limits = _ExtendedLimits()
# Neither BREAKAWAY_OK nor SILENT_BREAKAWAY_OK: descendants stay contained.
limits.BasicLimitInformation.LimitFlags = 0x2000 # KILL_ON_JOB_CLOSE
if not self._api.SetInformationJobObject(
self._handle, 9, ctypes.byref(limits), ctypes.sizeof(limits)):
raise ctypes.WinError(ctypes.get_last_error())
except BaseException:
self.close()
raise
def assign(self, proc: subprocess.Popen) -> None:
# Popen retains the original process handle, avoiding a PID-reuse race.
if not self._api.AssignProcessToJobObject(self._handle, int(proc._handle)):
raise ctypes.WinError(ctypes.get_last_error())
def close(self) -> None:
"""Terminate the contained tree; repeated closes are harmless."""
with self._lock:
if self._handle is not None:
if not self._api.CloseHandle(self._handle):
raise ctypes.WinError(ctypes.get_last_error())
self._handle = None
def spawn_server(cmd, **kwargs) -> tuple[subprocess.Popen, _WindowsJob | None]:
"""Start a router, returning its process and an owner-held containment handle.
Keep the job until shutdown and call close() to terminate the entire tree.
Windows closes it automatically if the owner dies. Other hosts retain Popen's
ordinary behavior. Assignment happens before the child's first instruction.
"""
if sys.platform != "win32":
return subprocess.Popen(cmd, **kwargs), None
job = _WindowsJob()
proc = None
try:
kwargs["creationflags"] = kwargs.get("creationflags", 0) | 0x00000004 # CREATE_SUSPENDED
proc = subprocess.Popen(cmd, **kwargs)
job.assign(proc)
psutil.Process(proc.pid).resume()
return proc, job
except BaseException:
try:
if proc is not None:
# Assignment may have failed: closing an empty job is not enough.
proc.kill()
proc.wait(timeout=10)
finally:
try:
job.close()
finally:
if proc is not None:
for stream in (proc.stdin, proc.stdout, proc.stderr):
if stream is not None:
stream.close()
proc._handle.Close()
raise
|