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

…ocess dies unexpectedly

pyptlib/util/subproc.py

@@ -26,11 +26,23 @@
     """
     _use_ctrlbreak_as_sigterm = bool(os.getenv("USE_CTRLBREAK_AS_SIGTERM", 1))
 
+    """
+    Set KILL_CHILDREN_ON_DEATH=1 in the environment to automatically kill all
+    descendents when this process dies.
+    """
+    # 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", 0))
+
     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
 
@@ -47,10 +59,37 @@
 _Popen_defaults = zip(a.args[-len(a.defaults):],a.defaults); del a
 if mswindows:
     # required for os.kill() to work
+    _Popen_creationflags = subprocess.CREATE_NEW_PROCESS_GROUP
+
+    if _kill_children_on_death:
+        _chJob = win32job.CreateJobObject(None, "")
+        if not _chJob:
+            raise WinError()
+
+        chJeli = win32job.QueryInformationJobObject(
+            _chJob, win32job.JobObjectExtendedLimitInformation)
+        # JOB_OBJECT_LIMIT_BREAKAWAY_OK allows children to assign grandchildren
+        # to their own jobs
+        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
+
+        # If we already belong to a JobObject, our children are auto-assigned
+        # to that and AssignProcessToJobObject(ch, _chJob) fails. This flag
+        # prevents this auto-assignment, as long as the parent JobObject has
+        # JOB_OBJECT_LIMIT_BREAKAWAY_OK set on it as well.
+        _Popen_creationflags |= _CREATE_BREAKAWAY_FROM_JOB
+
     tmp = dict(_Popen_defaults)
-    tmp['creationflags'] |= subprocess.CREATE_NEW_PROCESS_GROUP
+    tmp['creationflags'] |= _Popen_creationflags
     _Popen_defaults = tmp.items()
-    del tmp
+    del tmp, _Popen_creationflags
+
 
 class Popen(subprocess.Popen):
     """Wrapper for subprocess.Popen that tracks every child process.
@@ -79,6 +118,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: