NAS-133058 / 25.04 / simplify service.terminate_process

src/middlewared/middlewared/plugins/reporting/events.py

@@ -1,4 +1,3 @@
-import psutil
 import time
 
 from middlewared.event import EventSource

src/middlewared/middlewared/plugins/service.py

@@ -1,17 +1,16 @@
 import asyncio
 import errno
+import os
 
-import psutil
-
+import middlewared.sqlalchemy as sa
 from middlewared.plugins.service_.services.all import all_services
 from middlewared.plugins.service_.services.base import IdentifiableServiceInterface
 from middlewared.plugins.service_.utils import app_has_write_privilege_for_service
-
 from middlewared.schema import accepts, Bool, Dict, Int, List, Ref, returns, Str
 from middlewared.service import filterable, CallError, CRUDService, pass_app, periodic, private
-from middlewared.service_exception import MatchNotFound
-import middlewared.sqlalchemy as sa
+from middlewared.service_exception import MatchNotFound, ValidationError
 from middlewared.utils import filter_list, filter_getattrs
+from middlewared.utils.os import terminate_pid
 
 
 class ServiceModel(sa.Model):
@@ -429,33 +428,23 @@ async def become_standby(self, service):
         service_object = await self.middleware.call('service.object', service)
         return await service_object.become_standby()
 
-    @accepts(Int("pid"), Int("timeout", default=10))
-    @returns(Bool(
-        "process_terminated_nicely",
-        description="`true` is process has been successfully terminated with `TERM` and `false` if we had to use `KILL`"
-    ))
-    def terminate_process(self, pid, timeout):
-        """
-        Terminate process by `pid`.
-
-        First send `TERM` signal, then, if was not terminated in `timeout` seconds, send `KILL` signal.
+    @private
+    def terminate_process(self, pid, timeout = 10):
         """
-        try:
-            process = psutil.Process(pid)
-            process.terminate()
-            gone, alive = psutil.wait_procs([process], timeout)
-        except psutil.NoSuchProcess:
-            raise CallError("Process does not exist")
+        Terminate the process with the given `pid`.
 
-        if not alive:
-            return True
+        Send SIGTERM, wait up to `timeout` seconds for the process to terminate.
+        If the process is still running after the timeout, send SIGKILL.
 
+        Returns:
+            boolean: True if process was terminated with SIGTERM, false if SIGKILL was used
+        """
+        if pid == 0 or pid == os.getpid():
+            raise ValidationError('terminate_process.pid', 'Invalid PID')
         try:
-            alive[0].kill()
-        except psutil.NoSuchProcess:
-            return True
-
-        return False
+            return terminate_pid(pid, timeout)
+        except ProcessLookupError:
+            raise ValidationError('terminate_process.pid', f'No such process with PID: {pid}')
 
     @periodic(3600, run_on_start=False)
     @private

src/middlewared/middlewared/utils/os.py

@@ -1,17 +1,53 @@
-import logging
-import os
-import resource
+from os import closerange, kill
+from resource import getrlimit, RLIMIT_NOFILE, RLIM_INFINITY
+from signal import SIGKILL, SIGTERM
+from time import sleep, time
 
-logger = logging.getLogger(__name__)
+__all__ = ['close_fds', 'terminate_pid']
 
-__all__ = ['close_fds']
+ALIVE_SIGNAL = 0
 
 
 def close_fds(low_fd, max_fd=None):
     if max_fd is None:
-        max_fd = resource.getrlimit(resource.RLIMIT_NOFILE)[1]
-        # Avoid infinity as thats not practical
-        if max_fd == resource.RLIM_INFINITY:
+        max_fd = getrlimit(RLIMIT_NOFILE)[1]
+        if max_fd == RLIM_INFINITY:
+            # Avoid infinity as thats not practical
             max_fd = 8192
 
-    os.closerange(low_fd, max_fd)
+    closerange(low_fd, max_fd)
+
+
+def terminate_pid(pid: int, timeout: int = 10) -> bool:
+    # Send SIGTERM to request the process to terminate
+    kill(pid, SIGTERM)
+
+    try:
+        kill(pid, ALIVE_SIGNAL)
+    except ProcessLookupError:
+        # SIGTERM was honored
+        return True
+
+    # process still alive (could take awhile)
+    start_time = time()
+    while True:
+        try:
+            kill(pid, ALIVE_SIGNAL)
+        except ProcessLookupError:
+            # SIGTERM was honored (eventually)
+            return True
+
+        if time() - start_time >= timeout:
+            # Timeout reached; break out of the loop to send SIGKILL
+            break
+
+        # Wait a bit before checking again
+        sleep(0.1)
+
+    try:
+        # Send SIGKILL to forcefully terminate the process
+        kill(pid, SIGKILL)
+        return False
+    except ProcessLookupError:
+        # Process may have terminated between checks
+        return True

tests/unit/test_terminate_pid.py

@@ -0,0 +1,65 @@
+from os import kill
+from signal import signal, SIGCHLD, SIG_IGN
+from subprocess import Popen
+from textwrap import dedent
+from time import sleep
+
+import pytest
+
+from middlewared.utils.os import terminate_pid
+
+
+@pytest.fixture(scope="module", autouse=True)
+def signal_handler():
+    # ignore SIGCHLD so child process is removed
+    # from process table immediately without parent
+    # process having to do any (formal) clean-up
+    orig = signal(SIGCHLD, SIG_IGN)
+    yield
+    signal(SIGCHLD, orig)
+
+
+def test_sigterm():
+    p = Popen(['python', '-c', 'import time; time.sleep(60)'])
+    # Allow process to start
+    sleep(0.2)
+
+    # Call terminate_pid with timeout
+    terminate_pid(p.pid, timeout=5)
+    # Wait a bit to ensure process has time to terminate
+    sleep(0.5)
+
+    # Check if process has terminated
+    try:
+        kill(p.pid, 0)
+        assert False, f"{p.pid!r} still running"
+    except ProcessLookupError:
+        # Process has terminated
+        pass
+
+
+def test_sigkill():
+    script = dedent("""
+        import signal
+        import os
+        signal.signal(signal.SIGTERM, signal.SIG_IGN)
+        max_sleep = 60
+        slept_time = 0
+        while slept_time < max_sleep:
+            time.sleep(1)
+            slept_time += 1
+    """)
+    p = Popen(['python', '-c', script])
+
+    # Call terminate_pid with short timeout
+    terminate_pid(p.pid, timeout=1)
+    # Wait a bit to ensure process has time to terminate
+    sleep(0.5)
+
+    # Check if process has terminated
+    try:
+        kill(p.pid, 0)
+        assert False, f"{p.pid!r} still running"
+    except ProcessLookupError:
+        # Process has terminated
+        pass