"""Own one worker's process tree, including children left after its exit.""" from __future__ import annotations import os import subprocess import threading import time class ProcessTree: @classmethod def launch(cls, command, **kwargs): if os.name == "nt": kwargs["creationflags"] = kwargs.get("creationflags", 0) | 0x4 # CREATE_SUSPENDED else: kwargs["start_new_session"] = True process = subprocess.Popen(command, **kwargs) try: tree = cls(process) process.ytdl_process_tree = tree if os.name == "nt": tree.resume() return process except BaseException: tree = getattr(process, "ytdl_process_tree", None) if tree: tree.close() if process.poll() is None: process.kill() process.wait(timeout=10) for stream in (process.stdin, process.stdout, process.stderr): if stream: stream.close() raise def __init__(self, process): self.process = process self.lock = threading.Lock() self.handle = None if os.name != "nt": return import ctypes from ctypes import wintypes class BasicLimits(ctypes.Structure): _fields_ = [("ProcessTime", ctypes.c_longlong), ("JobTime", ctypes.c_longlong), ("Flags", wintypes.DWORD), ("MinWorkingSet", ctypes.c_size_t), ("MaxWorkingSet", ctypes.c_size_t), ("ActiveLimit", wintypes.DWORD), ("Affinity", ctypes.c_size_t), ("Priority", wintypes.DWORD), ("Scheduling", wintypes.DWORD)] class Limits(ctypes.Structure): _fields_ = [("Basic", BasicLimits), ("Io", ctypes.c_ulonglong * 6), ("ProcessMemory", ctypes.c_size_t), ("JobMemory", ctypes.c_size_t), ("PeakProcessMemory", ctypes.c_size_t), ("PeakJobMemory", ctypes.c_size_t)] class Accounting(ctypes.Structure): _fields_ = [("Times", ctypes.c_longlong * 4), ("PageFaults", wintypes.DWORD), ("TotalProcesses", wintypes.DWORD), ("ActiveProcesses", wintypes.DWORD), ("TerminatedProcesses", wintypes.DWORD)] self.ctypes, self.Accounting = ctypes, Accounting 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), ("QueryInformationJobObject", [wintypes.HANDLE, ctypes.c_int, ctypes.c_void_p, wintypes.DWORD, ctypes.c_void_p], wintypes.BOOL), ("AssignProcessToJobObject", [wintypes.HANDLE, wintypes.HANDLE], wintypes.BOOL), ("TerminateJobObject", [wintypes.HANDLE, wintypes.UINT], wintypes.BOOL), ("OpenProcess", [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD], wintypes.HANDLE), ("CloseHandle", [wintypes.HANDLE], wintypes.BOOL), ("CreateToolhelp32Snapshot", [wintypes.DWORD, wintypes.DWORD], wintypes.HANDLE), ("Thread32First", [wintypes.HANDLE, ctypes.c_void_p], wintypes.BOOL), ("Thread32Next", [wintypes.HANDLE, ctypes.c_void_p], wintypes.BOOL), ("OpenThread", [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD], wintypes.HANDLE), ("ResumeThread", [wintypes.HANDLE], wintypes.DWORD), ("WaitForSingleObject", [wintypes.HANDLE, wintypes.DWORD], wintypes.DWORD), ]: function = getattr(self.api, name) function.argtypes, function.restype = args, result handle = self.api.CreateJobObjectW(None, None) if not handle: raise ctypes.WinError(ctypes.get_last_error()) try: limits = Limits() limits.Basic.Flags = 0x2000 # JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE if not self.api.SetInformationJobObject(handle, 9, ctypes.byref(limits), ctypes.sizeof(limits)): raise ctypes.WinError(ctypes.get_last_error()) process_handle = self.api.OpenProcess(0x100 | 0x1, False, process.pid) if not process_handle: raise ctypes.WinError(ctypes.get_last_error()) try: if not self.api.AssignProcessToJobObject(handle, process_handle): raise ctypes.WinError(ctypes.get_last_error()) finally: self.api.CloseHandle(process_handle) self.handle = handle except BaseException: self.api.CloseHandle(handle) raise def resume(self): from ctypes import wintypes ctypes = self.ctypes class ThreadEntry(ctypes.Structure): _fields_ = [("Size", wintypes.DWORD), ("Usage", wintypes.DWORD), ("Id", wintypes.DWORD), ("Owner", wintypes.DWORD), ("BasePriority", wintypes.LONG), ("DeltaPriority", wintypes.LONG), ("Flags", wintypes.DWORD)] snapshot = self.api.CreateToolhelp32Snapshot(0x4, 0) # TH32CS_SNAPTHREAD if snapshot == ctypes.c_void_p(-1).value: raise ctypes.WinError(ctypes.get_last_error()) try: entry = ThreadEntry() entry.Size = ctypes.sizeof(entry) found = self.api.Thread32First(snapshot, ctypes.byref(entry)) while found: if entry.Owner == self.process.pid: thread = self.api.OpenThread(0x2, False, entry.Id) if not thread: raise ctypes.WinError(ctypes.get_last_error()) try: if self.api.ResumeThread(thread) == 0xffffffff: raise ctypes.WinError(ctypes.get_last_error()) return finally: self.api.CloseHandle(thread) found = self.api.Thread32Next(snapshot, ctypes.byref(entry)) raise RuntimeError("Не найден поток запуска помощника.") finally: self.api.CloseHandle(snapshot) def terminate(self): with self.lock: if os.name == "nt": if not self.handle: return from ctypes import wintypes capacity = 16 while True: class ProcessIds(self.ctypes.Structure): _fields_ = [("Assigned", wintypes.DWORD), ("Count", wintypes.DWORD), ("Ids", self.ctypes.c_size_t * capacity)] members = ProcessIds() if self.api.QueryInformationJobObject(self.handle, 3, self.ctypes.byref(members), self.ctypes.sizeof(members), None): break if self.ctypes.get_last_error() != 234: # ERROR_MORE_DATA raise self.ctypes.WinError(self.ctypes.get_last_error()) capacity = max(capacity * 2, members.Assigned) handles = [] deadline = time.monotonic() + 10 try: for pid in members.Ids[:members.Count]: handle = self.api.OpenProcess(0x100000, False, pid) # SYNCHRONIZE if handle: handles.append(handle) if not self.api.TerminateJobObject(self.handle, 1): raise self.ctypes.WinError(self.ctypes.get_last_error()) # Job accounting can reach zero before individual process # handles signal. Wait for FFmpeg to release its files too. for handle in handles: timeout = max(0, int((deadline - time.monotonic()) * 1000)) result = self.api.WaitForSingleObject(handle, timeout) if result == 258: raise TimeoutError("Не удалось дождаться остановки процессов загрузки.") if result != 0: raise self.ctypes.WinError(self.ctypes.get_last_error()) finally: for handle in handles: self.api.CloseHandle(handle) else: import signal try: os.killpg(self.process.pid, signal.SIGKILL) except ProcessLookupError: pass self.process.wait(timeout=10) def close(self): self.terminate() with self.lock: if self.handle: self.api.CloseHandle(self.handle) self.handle = None