Merge pull request #8359 from fstagni/90_use_psutil

[9.0] Subprocess: use psutil for killing processes

Author: fstagni

src/DIRAC/Core/Utilities/Subprocess.py

@@ -1,7 +1,7 @@
 """
 DIRAC Wrapper to execute python and system commands with a wrapper, that might
 set a timeout.
-3 FUNCTIONS are provided:
+3 functions are provided:
 
      - shellCall( iTimeOut, cmdSeq, callbackFunction = None, env = None ):
        it uses subprocess.Popen class with "shell = True".
@@ -26,6 +26,7 @@
        should be used to wrap third party python functions
 
 """
+
 import os
 import selectors
 import signal
@@ -161,7 +162,6 @@ def __init__(self, timeout=False, bufferLimit=52428800):
 
         self.child = None
         self.childPID = 0
-        self.childKilled = False
         self.callback = None
         self.bufferList = []
         self.cmdSeq = []
@@ -193,7 +193,7 @@ def __readFromFD(self, fd, baseLength=0):
                     f"First and last data in buffer: \n{dataString[:100]} \n....\n {dataString[-100:]} ",
                 )
                 retDict = S_ERROR(
-                    "Reached maximum allowed length (%d bytes) " "for called function return value" % self.bufferLimit
+                    "Reached maximum allowed length (%d bytes) for called function return value" % self.bufferLimit
                 )
                 retDict["Value"] = dataString
                 return retDict
@@ -241,60 +241,71 @@ def __selectFD(self, readSeq, timeout=False):
         events = sel.select(timeout=timeout or self.timeout or None)
         return [key.fileobj for key, event in events if event & selectors.EVENT_READ]
 
-    def __killPid(self, pid, sig=9):
-        """send signal :sig: to process :pid:
-
-        :param int pid: process id
-        :param int sig: signal to send, default 9 (SIGKILL)
-        """
+    def __terminateProcess(self, process):
+        """Tries to terminate a process with SIGTERM. Returns a (gone, alive) tuple"""
+        self.log.verbose(f"Sending SIGTERM signal to PID {process.pid}")
         try:
-            os.kill(pid, sig)
-        except Exception as x:
-            if str(x) != "[Errno 3] No such process":
-                self.log.exception("Exception while killing timed out process")
-                raise x
-
-    def __poll(self, pid):
-        """wait for :pid:"""
+            process.terminate()
+        except psutil.NoSuchProcess:
+            return ([], [])
+        return psutil.wait_procs([process], timeout=60)
+
+    def __poll(self, process):
+        """Non-blocking check of whether process `pid` is still alive.
+        Returns:
+        - (0, 0) if process is still running  (like os.waitpid(pid, os.WNOHANG))
+        - (pid, exitcode) if process has terminated
+        - None if process info cannot be retrieved
+        """
         try:
-            return os.waitpid(pid, os.WNOHANG)
-        except os.error:
-            if self.childKilled:
-                return False
+            exitcode = process.wait(timeout=0)
+            return (process.pid, exitcode)  # exited
+        except psutil.TimeoutExpired:
+            return (0, 0)  # still running
+        except psutil.NoSuchProcess:
             return None
 
     def killChild(self, recursive=True):
-        """kill child process
-
-        :param boolean recursive: flag to kill all descendants
+        """Kills a process tree (including children) with signal SIGTERM. If that fails, escalate to SIGKILL
+        returns (gone, alive) tuple.
         """
-        pgid = os.getpgid(self.childPID)
-        if pgid != os.getpgrp():
-            try:
-                # Child is in its own group: kill the group
-                os.killpg(pgid, signal.SIGTERM)
-            except OSError:
-                # Process is already dead
-                pass
-        else:
-            # No separate group: walk the tree
-            parent = psutil.Process(self.childPID)
-            procs = parent.children(recursive=recursive)
-            procs.append(parent)
-            for p in procs:
+
+        self.log.info(f"Killing childPID {self.childPID}")
+
+        gone, alive = [], []
+        try:
+            child_process = psutil.Process(self.childPID)
+        except psutil.NoSuchProcess:
+            self.log.warn(f"Child PID {self.childPID} no longer exists")
+            return (gone, alive)
+
+        if recursive:
+            # grandchildren
+            children = child_process.children(recursive=True)
+            self.log.info(f"Sending kill signal to {len(children)} children PIDs")
+            for p in children:
                 try:
                     p.terminate()
                 except psutil.NoSuchProcess:
-                    pass
-            _gone, alive = psutil.wait_procs(procs, timeout=10)
-            # Escalate any survivors
+                    continue
+            g, a = psutil.wait_procs(children, timeout=60)
+            gone.extend(g)
+            alive.extend(a)
+
+        # now killing the child_process
+        g, a = self.__terminateProcess(child_process)
+        gone.extend(g)
+        alive.extend(a)
+
+        # if there's something still alive, use SIGKILL
+        if alive:
             for p in alive:
                 try:
                     p.kill()
                 except psutil.NoSuchProcess:
                     pass
 
-        self.childKilled = True
+        return psutil.wait_procs(alive, timeout=60)
 
     def pythonCall(self, function, *stArgs, **stKeyArgs):
         """call python function :function: with :stArgs: and :stKeyArgs:"""
@@ -309,8 +320,6 @@ def pythonCall(self, function, *stArgs, **stKeyArgs):
         if pid == 0:
             os.close(readFD)
             self.__executePythonFunction(function, writeFD, *stArgs, **stKeyArgs)
-            # FIXME: the close it is done at __executePythonFunction, do we need it here?
-            os.close(writeFD)
         else:
             os.close(writeFD)
             readSeq = self.__selectFD([readFD])
@@ -319,14 +328,13 @@ def pythonCall(self, function, *stArgs, **stKeyArgs):
             try:
                 if len(readSeq) == 0:
                     self.log.debug("Timeout limit reached for pythonCall", function.__name__)
-                    self.__killPid(pid)
-
-                    # HACK to avoid python bug
-                    # self.wait()
-                    retries = 10000
-                    while os.waitpid(pid, 0) == -1 and retries > 0:
-                        time.sleep(0.001)
-                        retries -= 1
+                    gone, alive = self.__terminateProcess(psutil.Process(pid))
+                    if alive:
+                        for p in alive:
+                            try:
+                                p.kill()
+                            except psutil.NoSuchProcess:
+                                continue
 
                     return S_ERROR('%d seconds timeout for "%s" call' % (self.timeout, function.__name__))
                 elif readSeq[0] == readFD:
@@ -400,7 +408,7 @@ def __readFromFile(self, fd, baseLength):
         if len(dataString) + baseLength > self.bufferLimit:
             self.log.error("Maximum output buffer length reached")
             retDict = S_ERROR(
-                "Reached maximum allowed length (%d bytes) for called " "function return value" % self.bufferLimit
+                "Reached maximum allowed length (%d bytes) for called function return value" % self.bufferLimit
             )
             retDict["Value"] = dataString
             return retDict
@@ -446,6 +454,7 @@ def systemCall(
                 start_new_session=start_new_session,
             )
             self.childPID = self.child.pid
+            child_process = psutil.Process(self.childPID)
         except OSError as v:
             retDict = S_ERROR(repr(v))
             retDict["Value"] = (-1, "", str(v))
@@ -464,9 +473,9 @@ def systemCall(
             self.bufferList = [["", 0], ["", 0]]
             initialTime = time.time()
 
-            exitStatus = self.__poll(self.child.pid)
+            exitStatus = self.__poll(child_process)
 
-            while (0, 0) == exitStatus or exitStatus is None:
+            while (0, 0) == exitStatus:  # This means that the process is still alive
                 retDict = self.__readFromCommand()
                 if not retDict["OK"]:
                     return retDict
@@ -478,14 +487,14 @@ def systemCall(
                         1, "Timeout (%d seconds) for '%s' call" % (self.timeout, cmdSeq)
                     )
                 time.sleep(0.01)
-                exitStatus = self.__poll(self.child.pid)
+                exitStatus = self.__poll(child_process)
 
             self.__readFromCommand()
 
             if exitStatus:
                 exitStatus = exitStatus[1]
 
-            if exitStatus >= 256:
+            if exitStatus and exitStatus >= 256:
                 exitStatus = int(exitStatus / 256)
             return S_OK((exitStatus, self.bufferList[0][0], self.bufferList[1][0]))
         finally:

src/DIRAC/Core/Utilities/test/Test_Subprocess.py

@@ -3,23 +3,25 @@
 # Date: 2012/12/11 18:04:25
 ########################################################################
 
-""" :mod: SubprocessTests
-    =======================
+""":mod: SubprocessTests
+=======================
 
-    .. module: SubprocessTests
-    :synopsis: unittest for Subprocess module
-    .. moduleauthor:: [email protected]
+.. module: SubprocessTests
+:synopsis: unittest for Subprocess module
+.. moduleauthor:: [email protected]
 
-    unittest for Subprocess module
+unittest for Subprocess module
 """
-import time
+
 import platform
+import time
 from os.path import dirname, join
 from subprocess import Popen
 
+import psutil
 import pytest
 
-from DIRAC.Core.Utilities.Subprocess import systemCall, shellCall, pythonCall, getChildrenPIDs, Subprocess
+from DIRAC.Core.Utilities.Subprocess import Subprocess, getChildrenPIDs, pythonCall, shellCall, systemCall
 
 # Mark this entire module as slow
 pytestmark = pytest.mark.slow
@@ -72,3 +74,51 @@ def test_decodingCommandOutput():
     retVal = sp.systemCall(r"""python -c 'import os; os.fdopen(2, "wb").write(b"\xdf")'""", shell=True)
     assert retVal["OK"]
     assert retVal["Value"] == (0, "", "\ufffd")
+
+
+@pytest.fixture
+def subprocess_instance():
+    """Provides a Subprocess instance for testing."""
+    subp = Subprocess()
+    return subp
+
+
+@pytest.fixture
+def dummy_child():
+    """Spawn a dummy process tree: parent -> child."""
+    # Start a shell that sleeps, with a subprocess child
+    parent = Popen(["bash", "-c", "sleep 10 & wait"])
+    time.sleep(0.2)  # give it a moment to start
+    yield parent
+    # Ensure cleanup
+    try:
+        parent.terminate()
+        parent.wait(timeout=1)
+    except Exception:
+        pass
+
+
+def test_kill_child_process_tree(subprocess_instance, dummy_child):
+    """Test that killChild kills both parent and its children."""
+    subprocess_instance.childPID = dummy_child.pid
+    parent_proc = psutil.Process(subprocess_instance.childPID)
+
+    # Sanity check: parent should exist
+    assert parent_proc.is_running()
+
+    # It should have at least one sleeping child
+    children = parent_proc.children(recursive=True)
+    assert children, "Expected dummy process to have at least one child"
+
+    # Kill the tree
+    gone, alive = subprocess_instance.killChild(recursive=True)
+
+    # Verify the parent and children are terminated
+    for p in gone:
+        assert not p.is_running(), f"Process {p.pid} still alive"
+    for p in alive:
+        assert not p.is_running(), f"Process {p.pid} still alive"
+
+    # Verify parent is gone
+    with pytest.raises(psutil.NoSuchProcess):
+        psutil.Process(subprocess_instance.childPID)