github/krkn
1
0
Fork 0
mirror of https://github.com/krkn-chaos/krkn.git synced 2026-08-25 09:27:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: add connection pooling for HTTP health checks in cerberus and HealthChecker (#1236)

Browse Source 
Signed-off-by: 1PoPTRoN <vrxn.arp1traj@gmail.com>
Co-authored-by: Paige Patton <64206430+paigerube14@users.noreply.github.com>
This commit is contained in:
Arpit Raj
2026-05-14 11:22:54 -04:00
committed by GitHub
co-authored by Paige Patton
parent 2628665584
commit d70c56aa30
5 changed files with 201 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files
krkn/cerberus/setup.py
+6 -2
View File
@@ -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":
krkn/utils/HealthChecker.py
+15 -1
View File
@@ -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
run_kraken.py
+1
View File
@@ -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:
tests/test_cerberus_setup.py
+96 -59
View File
@@ -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__':
tests/test_health_checker.py
+83 -28
View File
@@ -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()