74 lines
3.0 KiB
Python
74 lines
3.0 KiB
Python
"""Windows Job Object ownership for suite-launched native servers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
|
|
|
|
class WindowsKillOnCloseJob:
|
|
"""Kill assigned children when the owning Python process loses this handle."""
|
|
|
|
def __init__(self) -> None:
|
|
self.handle = None
|
|
def assign(self, process_handle: int) -> None:
|
|
if os.name != "nt":
|
|
return
|
|
import ctypes
|
|
from ctypes import wintypes
|
|
|
|
class IO_COUNTERS(ctypes.Structure):
|
|
_fields_ = [(name, ctypes.c_ulonglong) for name in (
|
|
"ReadOperationCount", "WriteOperationCount", "OtherOperationCount",
|
|
"ReadTransferCount", "WriteTransferCount", "OtherTransferCount",
|
|
)]
|
|
|
|
class JOBOBJECT_BASIC_LIMIT_INFORMATION(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 JOBOBJECT_EXTENDED_LIMIT_INFORMATION(ctypes.Structure):
|
|
_fields_ = [
|
|
("BasicLimitInformation", JOBOBJECT_BASIC_LIMIT_INFORMATION),
|
|
("IoInfo", IO_COUNTERS),
|
|
("ProcessMemoryLimit", ctypes.c_size_t),
|
|
("JobMemoryLimit", ctypes.c_size_t),
|
|
("PeakProcessMemoryUsed", ctypes.c_size_t),
|
|
("PeakJobMemoryUsed", ctypes.c_size_t),
|
|
]
|
|
|
|
kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
|
|
kernel32.CreateJobObjectW.restype = wintypes.HANDLE
|
|
kernel32.SetInformationJobObject.argtypes = [
|
|
wintypes.HANDLE, ctypes.c_int, ctypes.c_void_p, wintypes.DWORD
|
|
]
|
|
kernel32.AssignProcessToJobObject.argtypes = [wintypes.HANDLE, wintypes.HANDLE]
|
|
handle = kernel32.CreateJobObjectW(None, None)
|
|
if not handle:
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
info = JOBOBJECT_EXTENDED_LIMIT_INFORMATION()
|
|
info.BasicLimitInformation.LimitFlags = 0x00002000 # KILL_ON_JOB_CLOSE
|
|
if not kernel32.SetInformationJobObject(handle, 9, ctypes.byref(info), ctypes.sizeof(info)):
|
|
kernel32.CloseHandle(handle)
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
if not kernel32.AssignProcessToJobObject(handle, wintypes.HANDLE(process_handle)):
|
|
kernel32.CloseHandle(handle)
|
|
raise ctypes.WinError(ctypes.get_last_error())
|
|
self.handle = handle
|
|
|
|
def close(self) -> None:
|
|
if self.handle is None or os.name != "nt":
|
|
return
|
|
import ctypes
|
|
|
|
ctypes.WinDLL("kernel32", use_last_error=True).CloseHandle(self.handle)
|
|
self.handle = None
|