from __future__ import annotations
|
|
import ctypes
|
import json
|
import os
|
import pathlib
|
import sys
|
import subprocess
|
import tempfile
|
import time
|
import unittest
|
from ctypes import wintypes
|
|
PROJECT_DEV = pathlib.Path(__file__).resolve().parents[2]
|
if str(PROJECT_DEV) not in sys.path:
|
sys.path.insert(0, str(PROJECT_DEV))
|
|
from bili_authenticated_extension.job import WindowsJob, close_handles, open_nul_handles # noqa: E402
|
|
HELPER = pathlib.Path(__file__).with_name("sleep_tree_helper.py")
|
|
|
def wait_records(path: pathlib.Path, count: int, timeout: float = 8.0) -> list[dict]:
|
deadline = time.monotonic() + timeout
|
while time.monotonic() < deadline:
|
if path.exists():
|
try:
|
records = [json.loads(line) for line in path.read_text(encoding="ascii").splitlines() if line]
|
except (OSError, json.JSONDecodeError):
|
records = []
|
if len(records) >= count:
|
return records
|
time.sleep(0.05)
|
raise AssertionError(f"expected {count} PID records")
|
|
|
def pid_signaled(pid: int, timeout: float = 5.0) -> bool:
|
kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
|
synchronize = 0x00100000
|
handle = kernel32.OpenProcess(synchronize, False, pid)
|
if not handle:
|
return True
|
try:
|
return kernel32.WaitForSingleObject(handle, int(timeout * 1000)) == 0
|
finally:
|
kernel32.CloseHandle(handle)
|
|
|
@unittest.skipUnless(os.name == "nt", "Windows Job Object test")
|
class JobObjectTests(unittest.TestCase):
|
def test_suspended_assign_resume_and_all_descendants_die_for_seven_terminal_causes(self) -> None:
|
causes = ("cancel", "eof", "disconnect", "keyboard_interrupt", "system_exit", "timeout", "host_kill")
|
for cause in causes:
|
with self.subTest(cause=cause), tempfile.TemporaryDirectory() as temporary:
|
pid_file = pathlib.Path(temporary) / "pids.jsonl"
|
nul_in, nul_out, nul_error = open_nul_handles()
|
job = WindowsJob()
|
process = None
|
try:
|
process = job.launch_suspended(
|
[sys.executable, str(HELPER), "--role", "worker", "--pid-file", str(pid_file)],
|
stdin_handle=nul_in,
|
stdout_handle=nul_out,
|
stderr_handle=nul_error,
|
)
|
records = wait_records(pid_file, 3)
|
self.assertEqual({"worker", "child", "grandchild"}, {item["role"] for item in records})
|
if cause == "host_kill":
|
job.close()
|
else:
|
job.terminate()
|
self.assertTrue(process.wait(5))
|
for item in records:
|
self.assertTrue(pid_signaled(item["pid"]), item)
|
finally:
|
if process is not None:
|
process.close()
|
job.close()
|
close_handles(nul_in, nul_out, nul_error)
|
|
def test_abrupt_job_owner_death_kills_real_worker_child_and_grandchild(self) -> None:
|
with tempfile.TemporaryDirectory() as temporary:
|
pid_file = pathlib.Path(temporary) / "owner-pids.jsonl"
|
owner = subprocess.Popen(
|
[sys.executable, str(HELPER), "--role", "owner", "--pid-file", str(pid_file)],
|
stdin=subprocess.DEVNULL,
|
stdout=subprocess.DEVNULL,
|
stderr=subprocess.DEVNULL,
|
)
|
try:
|
records = wait_records(pid_file, 4)
|
self.assertEqual(
|
{"owner", "worker", "child", "grandchild"},
|
{item["role"] for item in records},
|
)
|
descendants = [item for item in records if item["role"] != "owner"]
|
owner.kill()
|
owner.wait(timeout=5)
|
for item in descendants:
|
self.assertTrue(pid_signaled(item["pid"]), item)
|
finally:
|
if owner.poll() is None:
|
owner.kill()
|
owner.wait(timeout=5)
|
|
|
if __name__ == "__main__":
|
unittest.main()
|
