use Windows JobObjects to have all children killed when the parent pr…

…ocess dies unexpectedly

pyptlib/test/util_subproc_main.py

@@ -13,6 +13,7 @@ def startChild(subcmd, report=False, stdout=SINK, **kwargs):
     proc = Popen(
         ["python", "-m", "pyptlib.test.util_subproc_child", subcmd],
         stdout = stdout,
+        usesnewjobobject = True,
         **kwargs
     )
     if report:

pyptlib/util/subproc.py

@@ -25,11 +25,27 @@
     """
     _use_ctrlbreak_as_sigterm = bool(os.getenv("USE_CTRLBREAK_AS_SIGTERM", True))
 
+    """
+    Set this to true to automatically kill all descendents when this process
+    dies. NOTE: if any direct children uses this module too, and sets this flag
+    to autokill their children (i.e. this process's grandchildren), you MUST
+    pass "usesnewjobobject = True" when creating that child, otherwise it will
+    encounter an Access Denied error when creating the grandchildren.
+    """
+    # TODO(infinity0): write a test for this, similar to test_killall_kill
+    # Note: setting this to True defeats the point of some of the tests, so
+    # keep the default value as False. Perhaps we could make that work better.
+    _kill_children_on_death = bool(os.getenv("KILL_CHILDREN_ON_DEATH", False))
+
+
     from ctypes import byref, windll, WinError
     from ctypes.wintypes import DWORD
     import win32api, win32job
 
+    _CREATE_BREAKAWAY_FROM_JOB = 0x01000000 # see _kill_children_on_death
     _SYNCHRONIZE = 0x00100000 # generically useful
+    _PROCESS_SET_QUOTA = 0x0100 # required for AssignProcessToJobObject
+    _PROCESS_TERMINATE = 0x0001 # required for AssignProcessToJobObject
     _PROCESS_QUERY_INFORMATION = 0x0400 # required for GetExitCodeProcess
     _STILL_ACTIVE = 259 # GetExitCodeProcess returns this for still-running process
 
@@ -51,6 +67,22 @@
     _Popen_defaults = tmp.items()
     del tmp
 
+if mswindows and _kill_children_on_death:
+    _chJob = win32job.CreateJobObject(None, "")
+    if not _chJob:
+        raise WinError()
+
+    chJeli = win32job.QueryInformationJobObject(
+        _chJob, win32job.JobObjectExtendedLimitInformation)
+    chJeli['BasicLimitInformation']['LimitFlags'] |= (
+        win32job.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE |
+        win32job.JOB_OBJECT_LIMIT_BREAKAWAY_OK)
+
+    if win32job.SetInformationJobObject(
+        _chJob, win32job.JobObjectExtendedLimitInformation, chJeli) == 0:
+        raise WinError()
+    del chJeli
+
 class Popen(subprocess.Popen):
     """Wrapper for subprocess.Popen that tracks every child process.
 
@@ -66,6 +98,11 @@ class Popen(subprocess.Popen):
 
     def __init__(self, *args, **kwargs):
         kwargs = dict(_Popen_defaults + kwargs.items())
+        if 'usesnewjobobject' in kwargs:
+            if mswindows and _kill_children_on_death:
+                kwargs['creationflags'] = (
+                    kwargs.get('creationflags', 0) | _CREATE_BREAKAWAY_FROM_JOB)
+            del kwargs['usesnewjobobject']
         if 'creationflagsmerge' in kwargs:
             kwargs['creationflags'] = (
                 kwargs.get('creationflags', 0) | kwargs['creationflagsmerge'])
@@ -78,6 +115,12 @@ def __init__(self, *args, **kwargs):
         subprocess.Popen.__init__(self, *args, **kwargs)
         _CHILD_PROCS.append(self)
 
+        if mswindows and _kill_children_on_death:
+            handle = windll.kernel32.OpenProcess(
+                _SYNCHRONIZE | _PROCESS_SET_QUOTA | _PROCESS_TERMINATE, 0, self.pid)
+            if win32job.AssignProcessToJobObject(_chJob, handle) == 0:
+                raise WinError()
+
     if mswindows and _use_ctrlbreak_as_sigterm:
         def send_signal(self, sig):
             if sig == signal.SIGTERM: