From d70c56aa307f5d077c9bd527ea8c0014124670e3 Mon Sep 17 00:00:00 2001 From: Arpit Raj Date: Thu, 14 May 2026 20:52:54 +0530 Subject: [PATCH] fix: add connection pooling for HTTP health checks in cerberus and HealthChecker (#1236) Signed-off-by: 1PoPTRoN Co-authored-by: Paige Patton <64206430+paigerube14@users.noreply.github.com> --- krkn/cerberus/setup.py | 8 +- krkn/utils/HealthChecker.py | 16 +++- run_kraken.py | 1 + tests/test_cerberus_setup.py | 155 ++++++++++++++++++++++------------- tests/test_health_checker.py | 111 ++++++++++++++++++------- 5 files changed, 201 insertions(+), 90 deletions(-) diff --git a/krkn/cerberus/setup.py b/krkn/cerberus/setup.py index c32d2153..cb9c8ab0 100644 --- a/krkn/cerberus/setup.py +++ b/krkn/cerberus/setup.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import atexit import logging import requests import sys @@ -21,6 +22,9 @@ check_application_routes = "" cerberus_url = None exit_on_failure = False cerberus_enabled = False +http_session = requests.Session() # Singleton for connection pooling +atexit.register(http_session.close) # Cleanup on process exit + def set_url(config): global exit_on_failure @@ -48,7 +52,7 @@ def get_status(start_time, end_time): "is not provided." ) sys.exit(1) - cerberus_status = requests.get(cerberus_url, timeout=60).content + cerberus_status = http_session.get(cerberus_url, timeout=60).content cerberus_status = True if cerberus_status == b"True" else False # Fail if the application routes monitored by cerberus @@ -140,7 +144,7 @@ def application_status( start_time, end_time): try: failed_routes = [] status = True - metrics = requests.get(url, timeout=60).content + metrics = http_session.get(url, timeout=60).content metrics_json = json.loads(metrics) for entry in metrics_json["history"]["failures"]: if entry["component"] == "route": diff --git a/krkn/utils/HealthChecker.py b/krkn/utils/HealthChecker.py index c095db08..839ce2ea 100644 --- a/krkn/utils/HealthChecker.py +++ b/krkn/utils/HealthChecker.py @@ -24,12 +24,26 @@ from krkn_lib.models.telemetry.models import HealthCheck class HealthChecker: current_iterations: int = 0 ret_value = 0 + def __init__(self, iterations): self.iterations = iterations + self.http_session = requests.Session() + + def close(self): + """Close the HTTP session and release resources""" + if self.http_session: + self.http_session.close() + self.http_session = None + + def __enter__(self): + return self + + def __exit__(self, *args): + self.close() def make_request(self, url, auth=None, headers=None, verify=True): response_data = {} - response = requests.get(url, auth=auth, headers=headers, verify=verify, timeout=3) + response = self.http_session.get(url, auth=auth, headers=headers, verify=verify, timeout=3) response_data["url"] = url response_data["status"] = response.status_code == 200 response_data["status_code"] = response.status_code diff --git a/run_kraken.py b/run_kraken.py index 71d1c47b..29c20e47 100644 --- a/run_kraken.py +++ b/run_kraken.py @@ -485,6 +485,7 @@ def main(options, command: Optional[str]) -> int: # to json, and recreate a new object from it. end_time = int(time.time()) health_check_worker.join() + health_checker.close() try: chaos_telemetry.health_checks = health_check_telemetry_queue.get_nowait() except queue.Empty: diff --git a/tests/test_cerberus_setup.py b/tests/test_cerberus_setup.py index b48ab20f..053d0bb7 100644 --- a/tests/test_cerberus_setup.py +++ b/tests/test_cerberus_setup.py @@ -13,6 +13,7 @@ import unittest from unittest.mock import patch, MagicMock, Mock import sys import json +import requests from krkn.cerberus import setup as cerberus_setup @@ -69,46 +70,46 @@ class TestCerberusSetup(unittest.TestCase): self.assertFalse(cerberus_setup.exit_on_failure) self.assertFalse(cerberus_setup.cerberus_enabled) - @patch('krkn.cerberus.setup.requests.get') - def test_get_status_cerberus_disabled(self, mock_get): - """Test get_status when cerberus is disabled""" + @patch.object(cerberus_setup, 'http_session') + def test_get_status_cerberus_disabled(self, mock_session): + """Test get_status when cerberus is disabled makes no HTTP calls""" cerberus_setup.cerberus_enabled = False result = cerberus_setup.get_status(0, 100) self.assertTrue(result) - mock_get.assert_not_called() + mock_session.get.assert_not_called() - @patch('krkn.cerberus.setup.requests.get') - def test_get_status_cerberus_enabled_healthy(self, mock_get): + @patch.object(cerberus_setup, 'http_session') + def test_get_status_cerberus_enabled_healthy(self, mock_session): """Test get_status when cerberus is enabled and cluster is healthy""" cerberus_setup.cerberus_enabled = True cerberus_setup.cerberus_url = "http://cerberus.example.com" - + mock_response = MagicMock() mock_response.content = b"True" - mock_get.return_value = mock_response + mock_session.get.return_value = mock_response result = cerberus_setup.get_status(0, 100) self.assertTrue(result) - mock_get.assert_called_once_with("http://cerberus.example.com", timeout=60) + mock_session.get.assert_called_once_with("http://cerberus.example.com", timeout=60) - @patch('krkn.cerberus.setup.requests.get') - def test_get_status_cerberus_enabled_unhealthy(self, mock_get): + @patch.object(cerberus_setup, 'http_session') + def test_get_status_cerberus_enabled_unhealthy(self, mock_session): """Test get_status when cerberus is enabled and cluster is unhealthy""" cerberus_setup.cerberus_enabled = True cerberus_setup.cerberus_url = "http://cerberus.example.com" - + mock_response = MagicMock() mock_response.content = b"False" - mock_get.return_value = mock_response + mock_session.get.return_value = mock_response with self.assertRaises(SystemExit) as cm: cerberus_setup.get_status(0, 100) - + self.assertEqual(cm.exception.code, 1) - mock_get.assert_called_once_with("http://cerberus.example.com", timeout=60) + mock_session.get.assert_called_once_with("http://cerberus.example.com", timeout=60) def test_get_status_no_url_provided(self): """Test get_status when cerberus is enabled but URL is not provided""" @@ -117,12 +118,14 @@ class TestCerberusSetup(unittest.TestCase): with self.assertRaises(SystemExit) as cm: cerberus_setup.get_status(0, 100) - + self.assertEqual(cm.exception.code, 1) - @patch('krkn.cerberus.setup.requests.get') - def test_get_status_with_application_routes_check_success(self, mock_get): - """Test get_status with application routes check when routes are healthy""" + @patch.object(cerberus_setup, 'http_session') + def test_get_status_cerberus_healthy_returns_true(self, mock_session): + """Test get_status returns True when cerberus reports healthy. + Note: check_application_routes is shadowed locally in get_status() + (pre-existing issue), so route-check branch is not exercised here.""" cerberus_setup.cerberus_enabled = True cerberus_setup.cerberus_url = "http://cerberus.example.com" cerberus_setup.check_application_routes = "route1,route2" @@ -135,26 +138,26 @@ class TestCerberusSetup(unittest.TestCase): mock_response.content = b"True" return mock_response - mock_get.side_effect = mock_get_side_effect + mock_session.get.side_effect = mock_get_side_effect result = cerberus_setup.get_status(0, 100) self.assertTrue(result) - self.assertEqual(mock_get.call_count, 2) + self.assertEqual(mock_session.get.call_count, 2) - @patch('krkn.cerberus.setup.requests.get') - def test_get_status_with_application_routes_check_failure(self, mock_get): - """Test get_status with cerberus returning False (unhealthy)""" + @patch.object(cerberus_setup, 'http_session') + def test_get_status_with_application_routes_check_failure(self, mock_session): + """Test get_status when cerberus returns False (unhealthy)""" cerberus_setup.cerberus_enabled = True cerberus_setup.cerberus_url = "http://cerberus.example.com" - + mock_response = MagicMock() - mock_response.content = b"False" # Cerberus reports unhealthy - mock_get.return_value = mock_response + mock_response.content = b"False" + mock_session.get.return_value = mock_response with self.assertRaises(SystemExit) as cm: cerberus_setup.get_status(0, 100) - + self.assertEqual(cm.exception.code, 1) @patch('krkn.cerberus.setup.get_status') @@ -163,7 +166,6 @@ class TestCerberusSetup(unittest.TestCase): cerberus_setup.exit_on_failure = False mock_get_status.return_value = True - # Should not raise SystemExit cerberus_setup.publish_kraken_status(0, 100) mock_get_status.assert_called_once_with(0, 100) @@ -176,7 +178,7 @@ class TestCerberusSetup(unittest.TestCase): with self.assertRaises(SystemExit) as cm: cerberus_setup.publish_kraken_status(0, 100) - + self.assertEqual(cm.exception.code, 1) mock_get_status.assert_called_once_with(0, 100) @@ -186,7 +188,6 @@ class TestCerberusSetup(unittest.TestCase): cerberus_setup.exit_on_failure = False mock_get_status.return_value = False - # Should not raise SystemExit cerberus_setup.publish_kraken_status(0, 100) mock_get_status.assert_called_once_with(0, 100) @@ -199,58 +200,58 @@ class TestCerberusSetup(unittest.TestCase): with self.assertRaises(SystemExit) as cm: cerberus_setup.publish_kraken_status(0, 100) - + self.assertEqual(cm.exception.code, 1) mock_get_status.assert_called_once_with(0, 100) - @patch('krkn.cerberus.setup.requests.get') - def test_application_status_no_failures(self, mock_get): + @patch.object(cerberus_setup, 'http_session') + def test_application_status_no_failures(self, mock_session): """Test application_status when there are no route failures""" cerberus_setup.cerberus_url = "http://cerberus.example.com" - + mock_response = MagicMock() mock_response.content = json.dumps({ "history": { "failures": [] } }).encode() - mock_get.return_value = mock_response + mock_session.get.return_value = mock_response status, failed_routes = cerberus_setup.application_status(0, 6000) self.assertTrue(status) self.assertEqual(failed_routes, set()) expected_url = "http://cerberus.example.com/history?loopback=100.0" - mock_get.assert_called_once_with(expected_url, timeout=60) + mock_session.get.assert_called_once_with(expected_url, timeout=60) - @patch('krkn.cerberus.setup.requests.get') - def test_application_status_with_route_failures(self, mock_get): + @patch.object(cerberus_setup, 'http_session') + def test_application_status_with_route_failures(self, mock_session): """Test application_status when there are route failures""" cerberus_setup.cerberus_url = "http://cerberus.example.com" - + mock_response = MagicMock() mock_response.content = json.dumps({ "history": { "failures": [ {"component": "route", "name": "route1"}, {"component": "route", "name": "route2"}, - {"component": "pod", "name": "pod1"}, # Should be ignored - {"component": "route", "name": "route1"}, # Duplicate, should only appear once + {"component": "pod", "name": "pod1"}, # Non-route: should be ignored + {"component": "route", "name": "route1"}, # Duplicate: deduped by set() ] } }).encode() - mock_get.return_value = mock_response + mock_session.get.return_value = mock_response status, failed_routes = cerberus_setup.application_status(0, 6000) self.assertFalse(status) self.assertEqual(failed_routes, {"route1", "route2"}) - @patch('krkn.cerberus.setup.requests.get') - def test_application_status_with_non_route_failures(self, mock_get): + @patch.object(cerberus_setup, 'http_session') + def test_application_status_with_non_route_failures(self, mock_session): """Test application_status when there are non-route failures only""" cerberus_setup.cerberus_url = "http://cerberus.example.com" - + mock_response = MagicMock() mock_response.content = json.dumps({ "history": { @@ -260,7 +261,7 @@ class TestCerberusSetup(unittest.TestCase): ] } }).encode() - mock_get.return_value = mock_response + mock_session.get.return_value = mock_response status, failed_routes = cerberus_setup.application_status(0, 6000) @@ -273,35 +274,71 @@ class TestCerberusSetup(unittest.TestCase): with self.assertRaises(SystemExit) as cm: cerberus_setup.application_status(0, 100) - + self.assertEqual(cm.exception.code, 1) - @patch('krkn.cerberus.setup.requests.get') - def test_application_status_request_exception(self, mock_get): + @patch.object(cerberus_setup, 'http_session') + def test_application_status_request_exception(self, mock_session): """Test application_status when request raises an exception""" cerberus_setup.cerberus_url = "http://cerberus.example.com" - - mock_get.side_effect = Exception("Connection error") + + mock_session.get.side_effect = Exception("Connection error") with self.assertRaises(SystemExit) as cm: cerberus_setup.application_status(0, 6000) - + self.assertEqual(cm.exception.code, 1) - @patch('krkn.cerberus.setup.requests.get') - def test_application_status_duration_calculation(self, mock_get): + @patch.object(cerberus_setup, 'http_session') + def test_application_status_duration_calculation(self, mock_session): """Test application_status correctly calculates duration in minutes""" cerberus_setup.cerberus_url = "http://cerberus.example.com" - + mock_response = MagicMock() mock_response.content = json.dumps({"history": {"failures": []}}).encode() - mock_get.return_value = mock_response + mock_session.get.return_value = mock_response - # Duration: (300 - 0) / 60 = 5 minutes cerberus_setup.application_status(0, 300) expected_url = "http://cerberus.example.com/history?loopback=5.0" - mock_get.assert_called_once_with(expected_url, timeout=60) + mock_session.get.assert_called_once_with(expected_url, timeout=60) + + def test_http_session_is_singleton(self): + """Test that http_session is a requests.Session and the same object across accesses""" + session1 = cerberus_setup.http_session + session2 = cerberus_setup.http_session + self.assertIsInstance(session1, requests.Session) + self.assertIs(session1, session2) + + def test_http_session_reused_across_calls(self): + """Test that application_status reuses the module-level http_session""" + cerberus_setup.cerberus_url = "http://cerberus.example.com" + mock_response = MagicMock() + mock_response.content = json.dumps({"history": {"failures": []}}).encode() + original_session = cerberus_setup.http_session + + with patch.object(cerberus_setup.http_session, 'get', return_value=mock_response): + cerberus_setup.application_status(0, 300) + self.assertIs(cerberus_setup.http_session, original_session) + + cerberus_setup.application_status(0, 600) + self.assertIs(cerberus_setup.http_session, original_session) + + def test_http_session_atexit_registered(self): + """Test that http_session.close is registered via atexit for cleanup""" + import atexit + # atexit._run_exitfuncs is internal, so verify registration via the module code + # The atexit handler should have been registered at module import time + # We verify by checking the atexit registry contains our session's close + registered = False + # atexit callbacks are stored internally; verify by re-registering and checking no error + # Best we can do without poking internals: verify the session is closeable + session = cerberus_setup.http_session + self.assertTrue(callable(getattr(session, 'close', None))) + # Verify atexit module was imported and used in setup.py + import inspect + source = inspect.getsource(cerberus_setup) + self.assertIn('atexit.register(http_session.close)', source) if __name__ == '__main__': diff --git a/tests/test_health_checker.py b/tests/test_health_checker.py index 295950e1..28b3e0a6 100644 --- a/tests/test_health_checker.py +++ b/tests/test_health_checker.py @@ -18,6 +18,7 @@ Assisted By: Claude Code import queue import unittest +import requests from datetime import datetime from unittest.mock import MagicMock, patch @@ -39,6 +40,7 @@ class TestHealthChecker(unittest.TestCase): """ Clean up after each test """ + self.checker.close() self.checker.current_iterations = 0 self.checker.ret_value = 0 @@ -51,21 +53,26 @@ class TestHealthChecker(unittest.TestCase): return response_data return side_effect - @patch('requests.get') - def test_make_request_success(self, mock_get): + def _make_mock_session(self): + """Create a mock session with responses""" + mock_session = MagicMock() + mock_response = MagicMock() + mock_response.status_code = 200 + mock_session.get.return_value = mock_response + return mock_session + + def test_make_request_success(self): """ Test make_request returns success for 200 status code """ - mock_response = MagicMock() - mock_response.status_code = 200 - mock_get.return_value = mock_response + self.checker.http_session = self._make_mock_session() result = self.checker.make_request("http://example.com") self.assertEqual(result["url"], "http://example.com") self.assertEqual(result["status"], True) self.assertEqual(result["status_code"], 200) - mock_get.assert_called_once_with( + self.checker.http_session.get.assert_called_once_with( "http://example.com", auth=None, headers=None, @@ -73,20 +80,17 @@ class TestHealthChecker(unittest.TestCase): timeout=3 ) - @patch('requests.get') - def test_make_request_with_auth(self, mock_get): + def test_make_request_with_auth(self): """ Test make_request with basic authentication """ - mock_response = MagicMock() - mock_response.status_code = 200 - mock_get.return_value = mock_response + self.checker.http_session = self._make_mock_session() auth = ("user", "pass") result = self.checker.make_request("http://example.com", auth=auth) self.assertEqual(result["status"], True) - mock_get.assert_called_once_with( + self.checker.http_session.get.assert_called_once_with( "http://example.com", auth=auth, headers=None, @@ -94,20 +98,17 @@ class TestHealthChecker(unittest.TestCase): timeout=3 ) - @patch('requests.get') - def test_make_request_with_bearer_token(self, mock_get): + def test_make_request_with_bearer_token(self): """ Test make_request with bearer token authentication """ - mock_response = MagicMock() - mock_response.status_code = 200 - mock_get.return_value = mock_response + self.checker.http_session = self._make_mock_session() headers = {"Authorization": "Bearer token123"} result = self.checker.make_request("http://example.com", headers=headers) self.assertEqual(result["status"], True) - mock_get.assert_called_once_with( + self.checker.http_session.get.assert_called_once_with( "http://example.com", auth=None, headers=headers, @@ -115,33 +116,31 @@ class TestHealthChecker(unittest.TestCase): timeout=3 ) - @patch('requests.get') - def test_make_request_failure(self, mock_get): + def test_make_request_failure(self): """ Test make_request returns failure for non-200 status code """ + mock_session = MagicMock() mock_response = MagicMock() mock_response.status_code = 500 - mock_get.return_value = mock_response + mock_session.get.return_value = mock_response + self.checker.http_session = mock_session result = self.checker.make_request("http://example.com") self.assertEqual(result["status"], False) self.assertEqual(result["status_code"], 500) - @patch('requests.get') - def test_make_request_with_verify_false(self, mock_get): + def test_make_request_with_verify_false(self): """ Test make_request with SSL verification disabled """ - mock_response = MagicMock() - mock_response.status_code = 200 - mock_get.return_value = mock_response + self.checker.http_session = self._make_mock_session() result = self.checker.make_request("https://example.com", verify=False) self.assertEqual(result["status"], True) - mock_get.assert_called_once_with( + self.checker.http_session.get.assert_called_once_with( "https://example.com", auth=None, headers=None, @@ -494,10 +493,66 @@ class TestHealthChecker(unittest.TestCase): self.checker.iterations = 2 self.checker.run_health_check(config, self.health_check_queue) - # Verify sleep was called with custom interval mock_sleep.assert_called_with(5) + def test_healthchecker_uses_session_for_connection_pooling(self): + """ + Test that HealthChecker uses a session for connection pooling + """ + checker = HealthChecker(iterations=5) + self.assertTrue(hasattr(checker, 'http_session')) + import requests + self.assertIsInstance(checker.http_session, requests.Session) + + def test_healthchecker_close_method(self): + """ + Test that close() closes the session and sets it to None + """ + checker = HealthChecker(iterations=5) + mock_session = MagicMock(spec=requests.Session) + checker.http_session = mock_session + + checker.close() + + mock_session.close.assert_called_once() + self.assertIsNone(checker.http_session) + + def test_healthchecker_session_reused_across_requests(self): + """ + Test that the same session is reused across multiple make_request calls + """ + self.checker.http_session = self._make_mock_session() + original_session = self.checker.http_session + + self.checker.make_request("http://example1.com") + self.checker.make_request("http://example2.com") + + self.assertIs(self.checker.http_session, original_session) + self.assertEqual(self.checker.http_session.get.call_count, 2) + + def test_healthchecker_close_twice_is_safe(self): + """ + Test that calling close() twice does not raise an error + """ + checker = HealthChecker(iterations=5) + checker.close() + # Second close should not raise + checker.close() + self.assertIsNone(checker.http_session) + + def test_healthchecker_context_manager(self): + """ + Test that HealthChecker can be used as a context manager + """ + with HealthChecker(iterations=3) as checker: + mock_session = self._make_mock_session() + checker.http_session = mock_session + checker.make_request("http://example.com") + mock_session.get.assert_called_once() + # After exiting context, session should be closed + self.assertIsNone(checker.http_session) + if __name__ == "__main__": unittest.main()