Harden stress test kill on exit (#267)

* Harden stress test kill on exit

* only import psutil in the kill child process function

* Apply formatting

* Address Review comments, remove unused import

Author: amanusk

s_tui/helper_functions.py

@@ -20,6 +20,7 @@
 import os
 import logging
 import platform
+import signal
 import subprocess
 import re
 import csv
@@ -74,17 +75,53 @@ def get_processor_name():
     return platform.processor()
 
 
-def kill_child_processes(parent_proc):
-    """Kills a process and all its children"""
-    logging.debug("Killing stress process")
+def kill_child_processes(parent_proc, timeout=3):
+    """Kills a process and all its children gracefully.
+
+    Attempts SIGTERM via process group first, falls back to per-process
+    terminate, then SIGKILL after timeout.
+    """
+    import psutil
+
+    if parent_proc is None:
+        logging.debug("No stress process to kill")
+        return
+
+    logging.debug("Killing stress process %s", parent_proc)
+
+    # Try process-group SIGTERM first (covers the whole tree)
     try:
-        for proc in parent_proc.children(recursive=True):
-            logging.debug("Killing %s", proc)
-            proc.kill()
-        parent_proc.kill()
-    except AttributeError:
-        logging.debug("No such process")
-        logging.debug("Could not kill process")
+        pgid = os.getpgid(parent_proc.pid)
+        os.killpg(pgid, signal.SIGTERM)
+        logging.debug("Sent SIGTERM to process group %s", pgid)
+    except (OSError, ProcessLookupError, AttributeError):
+        # Process group kill failed — fall back to per-process terminate
+        logging.debug("Process group kill failed, falling back to per-process")
+        try:
+            for proc in parent_proc.children(recursive=True):
+                try:
+                    proc.terminate()
+                except (psutil.NoSuchProcess, ProcessLookupError):
+                    pass
+            parent_proc.terminate()
+        except (psutil.NoSuchProcess, AttributeError):
+            logging.debug("Process already gone during terminate")
+            return
+
+    # Wait for graceful exit, then SIGKILL stragglers
+    try:
+        gone, alive = psutil.wait_procs(
+            [parent_proc] + parent_proc.children(recursive=True),
+            timeout=timeout,
+        )
+        for proc in alive:
+            logging.debug("Sending SIGKILL to straggler %s", proc)
+            try:
+                proc.kill()
+            except (psutil.NoSuchProcess, ProcessLookupError):
+                pass
+    except (psutil.NoSuchProcess, AttributeError):
+        logging.debug("Process already gone during wait")
 
 
 def output_to_csv(sources, csv_writeable_file):

s_tui/s_tui.py

@@ -23,6 +23,7 @@
 from __future__ import absolute_import
 
 import argparse
+import atexit
 import signal
 import logging
 import os
@@ -122,6 +123,7 @@ def unhandled_input(self, uinput):
             graph_controller.view.on_menu_close()
 
     signal.signal(signal.SIGINT, signal_handler)
+    signal.signal(signal.SIGTERM, signal_handler)
 
 
 class StressController:
@@ -175,7 +177,10 @@ def start_stress(self, stress_cmd):
         with open(os.devnull, "w") as dev_null:
             try:
                 stress_proc = subprocess.Popen(
-                    stress_cmd, stdout=dev_null, stderr=dev_null
+                    stress_cmd,
+                    stdout=dev_null,
+                    stderr=dev_null,
+                    start_new_session=True,
                 )
                 self.set_stress_process(psutil.Process(stress_proc.pid))
             except OSError:
@@ -957,6 +962,7 @@ def main():
 
     global graph_controller
     graph_controller = GraphController(args)
+    atexit.register(graph_controller.stress_controller.kill_stress_process)
     graph_controller.main()
 
 

tests/test_helper_functions.py

@@ -1,11 +1,14 @@
 """Tests for s_tui.helper_functions module."""
 
 import os
+import signal
 import sys
 import json
 import tempfile
 
+import psutil
 import pytest
+from unittest.mock import MagicMock, patch
 
 from s_tui.helper_functions import (
     __version__,
@@ -201,9 +204,59 @@ def test_make_user_config_dir(self, tmp_path, monkeypatch):
 
 class TestKillChildProcesses:
     def test_none_parent(self):
-        """Should handle None gracefully (AttributeError caught)."""
+        """Should handle None gracefully."""
         kill_child_processes(None)  # should not raise
 
+    def test_sigterm_via_process_group(self):
+        """Should send SIGTERM to the process group first."""
+        parent = MagicMock()
+        parent.pid = 12345
+        parent.children.return_value = []
+
+        with patch("os.getpgid", return_value=12345) as mock_getpgid, patch(
+            "os.killpg"
+        ) as mock_killpg, patch("psutil.wait_procs", return_value=([], [])):
+            kill_child_processes(parent, timeout=1)
+            mock_getpgid.assert_called_once_with(12345)
+            mock_killpg.assert_called_once_with(12345, signal.SIGTERM)
+
+    def test_fallback_to_per_process_terminate(self):
+        """When process group kill fails, falls back to per-process terminate."""
+        child = MagicMock()
+        parent = MagicMock()
+        parent.pid = 100
+        parent.children.return_value = [child]
+
+        with patch("os.getpgid", side_effect=OSError("no pgid")), patch(
+            "psutil.wait_procs", return_value=([], [])
+        ):
+            kill_child_processes(parent, timeout=1)
+            child.terminate.assert_called_once()
+            parent.terminate.assert_called_once()
+
+    def test_sigkill_stragglers_after_timeout(self):
+        """Processes still alive after timeout get SIGKILL."""
+        straggler = MagicMock()
+        parent = MagicMock()
+        parent.pid = 200
+        parent.children.return_value = []
+
+        with patch("os.getpgid", return_value=200), patch("os.killpg"), patch(
+            "psutil.wait_procs", return_value=([], [straggler])
+        ):
+            kill_child_processes(parent, timeout=1)
+            straggler.kill.assert_called_once()
+
+    def test_already_dead_process(self):
+        """Should not raise when process is already dead."""
+        parent = MagicMock()
+        parent.pid = 300
+
+        with patch("os.getpgid", side_effect=ProcessLookupError()), patch.object(
+            parent, "children", side_effect=psutil.NoSuchProcess(300)
+        ):
+            kill_child_processes(parent, timeout=1)  # should not raise
+
 
 # ---------------------------------------------------------------------------
 # output_to_terminal / output_to_json

tests/test_stress_controller.py

@@ -94,6 +94,17 @@ def test_start_stress(self, mocker):
         sc.start_stress(["stress", "-c", "4"])
         assert sc.get_stress_process() is mock_psutil_proc
 
+    def test_start_stress_uses_new_session(self, mocker):
+        """start_stress passes start_new_session=True for process group isolation."""
+        mock_popen = mocker.patch("subprocess.Popen")
+        mock_popen.return_value.pid = 99999
+        mocker.patch("psutil.Process", return_value=MagicMock())
+
+        sc = StressController(True, False)
+        sc.start_stress(["stress", "-c", "4"])
+        _, kwargs = mock_popen.call_args
+        assert kwargs.get("start_new_session") is True
+
     def test_start_stress_oserror(self, mocker):
         """start_stress handles OSError gracefully."""
         mocker.patch("subprocess.Popen", side_effect=OSError("not found"))