From 1bec1f9051ab30c2cb8c71193edea591f17adb27 Mon Sep 17 00:00:00 2001 From: Sahil Lenka <76817449+Sahil-u07@users.noreply.github.com> Date: Thu, 30 Jul 2026 20:37:30 +0530 Subject: [PATCH] multiprocess_nodes: clean up pool in finally block, fix log level (#1415) ThreadPool.close() was placed after starmap/map but outside a finally block. If the pool operation throws, close() never runs and worker threads hang around until process exit. Moved pool.close() and added pool.join() into a finally block. Changed logging.info to logging.error in the exception handler so pool failures are visible at the default log level. Wrote two tests before touching code. The first verifies that close() and join() are called even when map raises. The second verifies that logging.error is used instead of logging.info. Signed-off-by: Sahil Lenka Co-authored-by: Paige Patton <64206430+paigerube14@users.noreply.github.com> --- .../shut_down/shut_down_scenario_plugin.py | 11 +++-- tests/test_shut_down_scenario_plugin.py | 40 ++++++++++++++++++- 2 files changed, 46 insertions(+), 5 deletions(-) diff --git a/krkn/scenario_plugins/shut_down/shut_down_scenario_plugin.py b/krkn/scenario_plugins/shut_down/shut_down_scenario_plugin.py index 88ff3c3d..c410bb47 100644 --- a/krkn/scenario_plugins/shut_down/shut_down_scenario_plugin.py +++ b/krkn/scenario_plugins/shut_down/shut_down_scenario_plugin.py @@ -62,9 +62,9 @@ class ShutDownScenarioPlugin(AbstractScenarioPlugin): return 1 def multiprocess_nodes(self, cloud_object_function, nodes, processes=0): + # pool object with number of element + pool = None try: - # pool object with number of element - if processes == 0: pool = ThreadPool(processes=len(nodes)) else: @@ -82,9 +82,12 @@ class ShutDownScenarioPlugin(AbstractScenarioPlugin): else: logging.info("pool type" + str(type(nodes))) pool.map(cloud_object_function, nodes) - pool.close() except Exception as e: - logging.info("Error on pool multiprocessing: " + str(e)) + logging.error("Error on pool multiprocessing: " + str(e)) + finally: + if pool: + pool.close() + pool.join() # Inject the cluster shut down scenario # krkn_lib diff --git a/tests/test_shut_down_scenario_plugin.py b/tests/test_shut_down_scenario_plugin.py index aad0ff82..44e2976d 100644 --- a/tests/test_shut_down_scenario_plugin.py +++ b/tests/test_shut_down_scenario_plugin.py @@ -409,7 +409,7 @@ class TestShutDownScenarioPlugin(unittest.TestCase): self.assertEqual(call_args[0], mock_cloud_function) mock_pool_instance.close.assert_called_once() - @patch('logging.info') + @patch('logging.error') @patch('krkn.scenario_plugins.shut_down.shut_down_scenario_plugin.ThreadPool') def test_multiprocess_nodes_with_exception(self, mock_threadpool, mock_logging): """ @@ -425,6 +425,44 @@ class TestShutDownScenarioPlugin(unittest.TestCase): mock_logging.assert_called() logged_args, logged_kwargs = mock_logging.call_args self.assertIn("Error on pool multiprocessing", logged_args[0]) + + @patch('logging.error') + @patch('krkn.scenario_plugins.shut_down.shut_down_scenario_plugin.ThreadPool') + def test_multiprocess_nodes_pool_close_on_exception(self, mock_threadpool, mock_logging): + """ + Test pool.close() and pool.join() are called even when map raises + """ + mock_pool_instance = Mock() + mock_threadpool.return_value = mock_pool_instance + mock_pool_instance.map.side_effect = Exception("map failed") + + nodes = ["node1", "node2"] + mock_cloud_function = Mock() + + self.plugin.multiprocess_nodes(mock_cloud_function, nodes, processes=0) + + mock_pool_instance.close.assert_called_once() + mock_pool_instance.join.assert_called_once() + + @patch('logging.error') + @patch('krkn.scenario_plugins.shut_down.shut_down_scenario_plugin.ThreadPool') + def test_multiprocess_nodes_logs_error_on_exception(self, mock_threadpool, mock_logging): + """ + Test logging.error is used when an exception occurs + """ + mock_pool_instance = Mock() + mock_threadpool.return_value = mock_pool_instance + mock_pool_instance.map.side_effect = Exception("map failed") + + nodes = ["node1", "node2"] + mock_cloud_function = Mock() + + self.plugin.multiprocess_nodes(mock_cloud_function, nodes, processes=0) + + mock_logging.assert_called() + logged_message = mock_logging.call_args[0][0] + self.assertIn("Error on pool multiprocessing", logged_message) + @patch('krkn.scenario_plugins.shut_down.shut_down_scenario_plugin.AWS') @patch('time.sleep') @patch('time.time')