chronaxis.batch.process_tree
源码基线：8904c2f744951857659abd40aaa3535d57faf53c
1"""管理本次工作进程树；Windows使用作业对象，Linux使用已确认身份的pidfd进程组。23父进程持有本次任务的作业句柄或pidfd；工作进程完成登记后才运行用户计划。4关闭时回收本任务的剩余后代，不扫描或接管其他任务的进程。5"""67import ctypes8from ctypes import wintypes9import os10import signal11import sys12from time import monotonic, sleep1314from ..results.budget import _cleanup_owner151617class _BasicLimits(ctypes.Structure):18 """对应Windows基本作业限制结构；仅设置关闭句柄时结束整个作业。"""1920 _fields_ = [('per_process', ctypes.c_longlong), ('per_job', ctypes.c_longlong),21 ('flags', wintypes.DWORD), ('min_ws', ctypes.c_size_t),22 ('max_ws', ctypes.c_size_t), ('active_limit', wintypes.DWORD),23 ('affinity', ctypes.c_size_t), ('priority', wintypes.DWORD),24 ('scheduling', wintypes.DWORD)]252627class _ExtendedLimits(ctypes.Structure):28 """对应Windows扩展作业限制，未使用的资源额度保持零值。"""2930 _fields_ = [('basic', _BasicLimits), ('io', ctypes.c_ulonglong * 6),31 ('process_memory', ctypes.c_size_t), ('job_memory', ctypes.c_size_t),32 ('peak_process', ctypes.c_size_t), ('peak_job', ctypes.c_size_t)]333435class ProcessTree:36 """持有一个工作进程树的关闭责任；Linux在GO前完成双向身份核验。"""3738 def __init__(self, pid, *, registration=None, process=None):39 """登记实际PID。Windows沿原pid接口；Linux另需registration和借用process。返回新对象；平台、握手或系统错误保原异常。"""40 self._linux = sys.platform.startswith('linux')41 if self._linux:42 self._register_linux(pid, registration, process)43 return44 if os.name != 'nt':45 raise NotImplementedError('进程树管理需要Windows或Linux 6.9及以上')46 api = ctypes.WinDLL('kernel32', use_last_error=True)47 api.CreateJobObjectW.argtypes = [ctypes.c_void_p, wintypes.LPCWSTR]48 api.CreateJobObjectW.restype = wintypes.HANDLE49 api.SetInformationJobObject.argtypes = [wintypes.HANDLE, ctypes.c_int, ctypes.c_void_p, wintypes.DWORD]50 api.SetInformationJobObject.restype = wintypes.BOOL51 api.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]52 api.OpenProcess.restype = wintypes.HANDLE53 api.AssignProcessToJobObject.argtypes = [wintypes.HANDLE, wintypes.HANDLE]54 api.AssignProcessToJobObject.restype = wintypes.BOOL55 api.CloseHandle.argtypes = [wintypes.HANDLE]56 api.CloseHandle.restype = wintypes.BOOL57 # 两个实际句柄先登记；构造尚未返回也必须能从原异常恢复关闭。58 self._api, self._handle, self._process_handle = api, None, None59 self._closed = False60 try:61 self._handle = api.CreateJobObjectW(None, None)62 if not self._handle:63 raise ctypes.WinError(ctypes.get_last_error())64 limits = _ExtendedLimits()65 limits.basic.flags = 0x2000 # JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE；不授予后代脱离权限。66 if not api.SetInformationJobObject(self._handle, 9, ctypes.byref(limits), ctypes.sizeof(limits)):67 raise ctypes.WinError(ctypes.get_last_error())68 self._process_handle = api.OpenProcess(0x0101, False, pid) # SET_QUOTA | TERMINATE。69 if not self._process_handle or not api.AssignProcessToJobObject(self._handle, self._process_handle):70 raise ctypes.WinError(ctypes.get_last_error())71 # 登记成功后只关本次借助的进程句柄，保留作业句柄直到工作池结束。72 self._close_process_handle()73 except BaseException as error:74 _cleanup_owner(self, error, 'process_tree_cleanup_owner', 'Windows进程树构造')75 raise7677 def _register_linux(self, pid, registration, process):78 """通过本次进程的双向握手登记pidfd。7980 参数: pid为实际PID；registration拥有启动通道；process为借用的Process。81 返回: None；能力与身份核验完成后才发GO。82 异常: 无效身份、握手或系统错误原样传播；未确认fd只关闭，不发送组信号。83 """84 self._fd = None85 self._confirmed = False86 self._closed = False87 self._process = process88 try:89 if registration is None or process is None or type(pid) is not int or (pid != process.pid) or (pid <= 0):90 raise ValueError('Linux进程树只能登记本次启动握手对应的Process')91 registration.close_child()92 expected = b'READY' + registration.token + pid.to_bytes(8, 'big') + pid.to_bytes(8, 'big')93 if registration.receive(process) != expected or pid == os.getpgrp():94 raise ValueError('工作进程组READY身份不匹配')95 # 只打开临时fd；收到打开之后的ACK以前不能向该fd发送组信号。96 self._fd = os.pidfd_open(pid)97 registration.parent.send_bytes(b'PROBE' + registration.token)98 if registration.receive(process) != b'ACK' + registration.token:99 raise ValueError('工作进程组ACK身份不匹配')100 if os.getpgid(pid) != pid or pid == os.getpgrp():101 raise ValueError('实际工作进程未建立独立进程组')102 try:103 signal.pidfd_send_signal(self._fd, 0, None, 4)104 except OSError as error:105 if error.errno in (22, 38):106 raise NotImplementedError('并行任务需要Linux 6.9及以上的pidfd进程组信号') from error107 raise108 # ACK证明原worker仍活着，实际组身份和内核能力通过后才授予业务执行。109 self._confirmed = True110 registration.parent.send_bytes(b'GO' + registration.token)111 except BaseException as error:112 _cleanup_owner(self, error, 'process_tree_cleanup_owner', 'Linux进程树登记')113 raise114115 def _close_linux(self):116 """关闭已确认的进程组并核对实际消失。117118 参数: 无。119 返回: None；未确认fd仅关闭。120 异常: 组未消失或信号失败时保留责任；Linux close本身报错也不重试已释放编号。121 """122 if self._fd is None:123 self._process = None124 self._closed = True125 return126 if self._confirmed:127 try:128 signal.pidfd_send_signal(self._fd, signal.SIGKILL, None, 4)129 except ProcessLookupError:130 pass131 # 回收本次组长僵尸；不接管或扫描其他进程。132 process = self._process133 if process is not None and (not getattr(process, '_closed', False)):134 process.join(timeout=5)135 if process.is_alive():136 raise TimeoutError('进程树组长尚未退出')137 deadline = monotonic() + 5138 while True:139 try:140 signal.pidfd_send_signal(self._fd, 0, None, 4)141 except ProcessLookupError:142 break143 if monotonic() >= deadline:144 raise TimeoutError('进程组尚未全部退出，保留清理责任')145 sleep(0.01)146 # 组已不存在；Linux close可能先释放编号再报错，因此只调用一次。147 descriptor = self._fd148 self._fd = None149 self._process = None150 self._closed = True151 os.close(descriptor)152153 def _close_process_handle(self):154 """关闭登记期间真实打开的进程句柄。参数: 无。返回: None。异常: 失败仍保留句柄。"""155 if self._process_handle:156 if not self._api.CloseHandle(self._process_handle):157 raise ctypes.WinError(ctypes.get_last_error())158 self._process_handle = None159160 def close(self):161 """关闭本次平台资源；Windows失败保留实际句柄，Linux只重试尚未结束的组。162163 参数: 无。返回: None；重复关闭安全。异常: 第一个关闭错误，其余写入notes。164 Windows作业关闭后结束残余后代；登记句柄失败仍尝试结束作业。165 Linux组确认消失后只close一次pidfd；错误向上传播，不重试已释放编号。166 """167 if self._closed:168 return169 if getattr(self, '_linux', False) is True:170 self._close_linux()171 return172 errors = []173 try:174 self._close_process_handle()175 except BaseException as error:176 errors.append(error)177 if self._handle:178 try:179 if not self._api.CloseHandle(self._handle):180 raise ctypes.WinError(ctypes.get_last_error())181 self._handle = None182 except BaseException as error:183 errors.append(error)184 self._closed = not self._handle and not self._process_handle185 if errors:186 for error in errors[1:]:187 errors[0].add_note('Windows进程树另有句柄关闭失败: ' + type(error).__name__)188 raise errors[0]189