173 lines
8.5 KiB
Python
173 lines
8.5 KiB
Python
"""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
|