Fix zombie health check watcher threads on krkn runner failures

krkn_ai/chaos_engines/krkn_runner.py

@@ -115,17 +115,19 @@ def run(self, scenario: BaseScenario, generation_id: int) -> CommandRunResult:
             # Start watching application urls for health checks
             health_check_watcher.run()
 
-            # Run command (show logs when verbose mode is enabled)
-            log, returncode = run_shell(
-                self.process_es_env_string(command, True), do_not_log=not is_verbose()
-            )
+            try:
+                # Run command (show logs when verbose mode is enabled)
+                log, returncode = run_shell(
+                    self.process_es_env_string(command, True), do_not_log=not is_verbose()
+                )
 
-            # Extract return code from run log which is part of telemetry data present in the log
-            returncode, run_uuid = self.__extract_returncode_from_run(log, returncode)
-            logger.info("Krkn scenario return code: %d", returncode)
+                # Extract return code from run log which is part of telemetry data present in the log
+                returncode, run_uuid = self.__extract_returncode_from_run(log, returncode)
+                logger.info("Krkn scenario return code: %d", returncode)
 
-            # Stop watching application urls for health checks
-            health_check_watcher.stop()
+            finally:
+                # Stop watching application urls for health checks
+                health_check_watcher.stop()
 
         end_time = datetime.datetime.now()
 

tests/test_krkn_runner_leak.py

@@ -0,0 +1,58 @@
+import unittest
+from unittest.mock import MagicMock, patch
+import sys
+
+
+from krkn_ai.chaos_engines.krkn_runner import KrknRunner
+from krkn_ai.models.config import ConfigFile, FitnessFunction, HealthCheckConfig
+
+class TestKrknRunnerThreadLeak(unittest.TestCase):
+    @patch('krkn_ai.chaos_engines.krkn_runner.create_prometheus_client')
+    @patch('krkn_ai.chaos_engines.krkn_runner.HealthCheckWatcher')
+    @patch('krkn_ai.chaos_engines.krkn_runner.run_shell') 
+    def test_run_shell_exception_cleanup(self, mock_run_shell, mock_watcher_cls, mock_create_prom):
+        # Setup mocks
+        mock_watcher = MagicMock()
+        mock_watcher_cls.return_value = mock_watcher
+
+        # Configure run_shell to return success by default (for init checks)
+        mock_run_shell.return_value = ("output", 0)
+
+        # Setup Config
+        config = MagicMock(spec=ConfigFile)
+        config.kubeconfig_file_path = "fake_path"
+        config.fitness_function = MagicMock(spec=FitnessFunction)
+        config.health_checks = MagicMock(spec=HealthCheckConfig)
+        config.elastic = None
+        config.wait_duration = 10
+
+        from krkn_ai.models.scenario.base import Scenario
+
+        runner = KrknRunner(config, "output_dir")
+
+        # Create a dummy scenario
+        scenario = MagicMock(spec=Scenario)
+        scenario.parameters = []
+        scenario.krknctl_name = "test-scenario"
+
+        # Now make run_shell crash
+        mock_run_shell.side_effect = Exception("Simulated Shell Crash")
+
+        print("\n[TEST] Executing runner.run() with failing run_shell...")
+        try:
+            runner.run(scenario, 1)
+        except Exception as e:
+            print(f"[INFO] Caught expected exception: {e}")
+
+        # Verification
+        print("[VERIFY] Checking if HealthCheckWatcher.stop() was called...")
+
+        # Assert that stop was called
+        mock_watcher.stop.assert_called_once()
+        print("[SUCCESS] stop() WAS called. Thread leak prevented.")
+
+        # Verify run() was also called
+        mock_watcher.run.assert_called()
+
+if __name__ == '__main__':
+    unittest.main()