Added support for multiple destinations in healthcheck (#2704)

Author: sbylica-splunk

package/sbin/healthcheck.py

@@ -1,6 +1,7 @@
 import logging
 import os
 import subprocess
+import re
 
 from flask_wtf.csrf import CSRFProtect
 from flask import Flask, jsonify
@@ -18,11 +19,20 @@ def str_to_bool(value):
         'yes'
     }
 
+def get_list_of_destinations():
+    found_destinations = []
+    regex = r"^SC4S_DEST_SPLUNK_HEC_(.*)_URL$"
+
+    for var_key, var_variable in os.environ.items():
+        if re.search(regex, var_key):
+            found_destinations.append(var_variable)
+    return set(found_destinations)
+
 class Config:
-    SC4S_DEST_SPLUNK_HEC_DEFAULT_URL = os.getenv('SC4S_DEST_SPLUNK_HEC_DEFAULT_URL')
     HEALTHCHECK_PORT = int(os.getenv('SC4S_LISTEN_STATUS_PORT', '8080'))
     CHECK_QUEUE_SIZE = str_to_bool(os.getenv('HEALTHCHECK_CHECK_QUEUE_SIZE', "false"))
     MAX_QUEUE_SIZE = int(os.getenv('HEALTHCHECK_MAX_QUEUE_SIZE', '10000'))
+    DESTINATIONS = get_list_of_destinations()
 
 logging.basicConfig(
     format="%(asctime)s - healthcheck.py - %(levelname)s - %(message)s",
@@ -52,11 +62,11 @@ def check_syslog_ng_health() -> bool:
         return False
 
 def check_queue_size(
-        sc4s_dest_splunk_hec_default=Config.SC4S_DEST_SPLUNK_HEC_DEFAULT_URL,
+        sc4s_dest_splunk_hec_destinations=Config.DESTINATIONS,
         max_queue_size=Config.MAX_QUEUE_SIZE
     ) -> bool:
     """Check syslog-ng queue size and compare it against the configured maximum limit."""
-    if not sc4s_dest_splunk_hec_default:
+    if not sc4s_dest_splunk_hec_destinations:
         logger.error(
             "SC4S_DEST_SPLUNK_HEC_DEFAULT_URL not configured. "
             "Ensure the default HEC destination is set, or disable HEALTHCHECK_CHECK_QUEUE_SIZE."
@@ -75,15 +85,22 @@ def check_queue_size(
             return False
 
         stats = result.stdout.splitlines()
-        destination_stat = next(
-            (s for s in stats if ";queued;" in s and sc4s_dest_splunk_hec_default in s),
-            None
-        )
-        if not destination_stat:
-            logger.error("No matching queue stats found for the destination URL.")
-            return False
 
-        queue_size = int(destination_stat.split(";")[-1])
+        queue_sizes_all_destinations = []
+
+        for destination in sc4s_dest_splunk_hec_destinations:
+            destination_stat = next(
+                (s for s in stats if ";queued;" in s and destination in s),
+                None
+            )
+
+            if not destination_stat:
+                logger.error(f"No matching queue stats found for the destination URL {destination}.")
+                return False
+
+            queue_sizes_all_destinations.append(int(destination_stat.split(";")[-1]))
+
+        queue_size = max(queue_sizes_all_destinations)
         if queue_size > max_queue_size:
             logger.warning(
                 f"Queue size {queue_size} exceeds the maximum limit of {max_queue_size}."

tests/test_healthcheck_unit_tests.py

@@ -8,6 +8,7 @@
     check_syslog_ng_health,
     subprocess,
     check_queue_size,
+    get_list_of_destinations,
 )
 
 # str_to_bool
@@ -52,9 +53,9 @@ def test_check_syslog_ng_health_exception(mock_run):
 # check_queue_size
 def test_check_queue_size_no_url():
     """
-    If sc4s_dest_splunk_hec_default is not set, check_queue_size should fail.
+    If sc4s_dest_splunk_hec_destinations is not set, check_queue_size should fail.
     """
-    assert check_queue_size(sc4s_dest_splunk_hec_default=None, max_queue_size=1000) is False
+    assert check_queue_size(sc4s_dest_splunk_hec_destinations=None, max_queue_size=1000) is False
 
 @patch("subprocess.run")
 def test_check_queue_size_stats_fail(mock_run):
@@ -63,7 +64,7 @@ def test_check_queue_size_stats_fail(mock_run):
     """
     mock_run.return_value.returncode = 1
     mock_run.return_value.stderr = "stats error"
-    assert check_queue_size(sc4s_dest_splunk_hec_default="http://example.com:8088", max_queue_size=1000) is False
+    assert check_queue_size(sc4s_dest_splunk_hec_destinations={"http://example.com:8088"}, max_queue_size=1000) is False
 
 @patch("subprocess.run")
 def test_check_queue_size_no_matching_stats(mock_run):
@@ -72,7 +73,7 @@ def test_check_queue_size_no_matching_stats(mock_run):
     """
     mock_run.return_value.returncode = 0
     mock_run.return_value.stdout = "some;other;stat;line\nanother;stat"
-    assert check_queue_size(sc4s_dest_splunk_hec_default="http://example.com:8088", max_queue_size=1000) is False
+    assert check_queue_size(sc4s_dest_splunk_hec_destinations={"http://example.com:8088"}, max_queue_size=1000) is False
 
 @patch("subprocess.run")
 def test_check_queue_size_exceeds_limit(mock_run):
@@ -84,7 +85,7 @@ def test_check_queue_size_exceeds_limit(mock_run):
         "destination;queued;http://example.com:8088;2000\n"
         "another;queued;http://other-url.com;1234"
     )
-    assert check_queue_size(sc4s_dest_splunk_hec_default="http://example.com:8088", max_queue_size=1000) is False
+    assert check_queue_size(sc4s_dest_splunk_hec_destinations={"http://example.com:8088"}, max_queue_size=1000) is False
 
 @patch("subprocess.run")
 def test_check_queue_size_under_limit(mock_run):
@@ -96,7 +97,7 @@ def test_check_queue_size_under_limit(mock_run):
         "destination;queued;http://example.com:8088;500\n"
         "another;queued;http://other-url.com;1234"
     )
-    assert check_queue_size(sc4s_dest_splunk_hec_default="http://example.com:8088", max_queue_size=1000) is True
+    assert check_queue_size(sc4s_dest_splunk_hec_destinations={"http://example.com:8088"}, max_queue_size=1000) is True
 
 @patch("subprocess.run")
 def test_check_queue_size_equals_limit(mock_run):
@@ -108,7 +109,62 @@ def test_check_queue_size_equals_limit(mock_run):
         "destination;queued;http://example.com:8088;1000\n"
         "another;queued;http://other-url.com;1234"
     )
-    assert check_queue_size(sc4s_dest_splunk_hec_default="http://example.com:8088", max_queue_size=1000) is True
+    assert check_queue_size(sc4s_dest_splunk_hec_destinations={"http://example.com:8088"}, max_queue_size=1000) is True
+
+@patch("subprocess.run")
+def test_check_queue_size_multiple_destinations(mock_run):
+    """
+    If queue size for all destinations is <= HEALTHCHECK_MAX_QUEUE_SIZE, check_queue_size should pass.
+    """
+    mock_run.return_value.returncode = 0
+    mock_run.return_value.stdout = (
+        "destination;queued;http://example.com:8088;300\n"
+        "destination;queued;http://another.com:8088;500\n"
+        "another;queued;http://other-url.com;1234"
+    )
+    assert check_queue_size(sc4s_dest_splunk_hec_destinations={"http://example.com:8088", "http://another.com:8088"},
+                            max_queue_size=1000) is True
+
+@patch("subprocess.run")
+def test_check_queue_size_multiple_destinations_over_limit(mock_run):
+    """
+    If queue size for at least one destination is > HEALTHCHECK_MAX_QUEUE_SIZE, check_queue_size should fail.
+    """
+    mock_run.return_value.returncode = 0
+    mock_run.return_value.stdout = (
+        "destination;queued;http://example.com:8088;1300\n"
+        "destination;queued;http://another.com:8088;500\n"
+        "another;queued;http://other-url.com;1234"
+    )
+    assert check_queue_size(sc4s_dest_splunk_hec_destinations={"http://example.com:8088", "http://another.com:8088"},
+                            max_queue_size=1000) is False
+
+@patch("subprocess.run")
+def test_check_queue_size_multiple_destinations_all_over_limit(mock_run):
+    """
+    If queue size for all destinations is > HEALTHCHECK_MAX_QUEUE_SIZE, check_queue_size should fail.
+    """
+    mock_run.return_value.returncode = 0
+    mock_run.return_value.stdout = (
+        "destination;queued;http://example.com:8088;1300\n"
+        "destination;queued;http://another.com:8088;1500\n"
+        "another;queued;http://other-url.com;1234"
+    )
+    assert check_queue_size(sc4s_dest_splunk_hec_destinations={"http://example.com:8088", "http://another.com:8088"},
+                            max_queue_size=1000) is False
+
+@patch("subprocess.run")
+def test_check_queue_size_multiple_incomplete_info(mock_run):
+    """
+    If stats run successfully but do not contain stats for one of the desired destinations, it should fail.
+    """
+    mock_run.return_value.returncode = 0
+    mock_run.return_value.stdout = (
+        "destination;queued;http://example.com:8088;300\n"
+        "another;queued;http://other-url.com;1234"
+    )
+    assert check_queue_size(sc4s_dest_splunk_hec_destinations={"http://example.com:8088", "http://another.com:8088"},
+                            max_queue_size=1000) is False
 
 @patch("subprocess.run", side_effect=Exception("some exception"))
 def test_check_queue_size_exception(mock_run):
@@ -139,4 +195,26 @@ def test_health_endpoint_no_queue_check(mock_run, client):
 
     response = client.get("/health")
     assert response.status_code == 200
-    assert response.json["status"] == "healthy"
+    assert response.json["status"] == "healthy"
+
+@patch.dict(
+    os.environ,
+    {
+        "SC4S_DEST_SPLUNK_HEC_DEFAULT_URL": "http://my_test_url:1234",
+        "SC4S_DEST_SPLUNK_HEC_OTHER_URL": "http://my_hec:1234",
+        "SOME_OTHER_URL": "http://my_url/test_url",
+        "SOME_OTHER_ENV_VARIABLE": "my_variable",
+        "SC4S_LISTEN_STATUS_PORT": "1234",
+    },
+    clear=True
+)
+def test_get_destinations():
+    """
+    Check if get_list_of_destinations method parses and returns the expected
+    destinations from environment variables.
+    """
+    destinations = get_list_of_destinations()
+
+    assert len(destinations) == 2
+    assert "http://my_test_url:1234" in destinations
+    assert "http://my_hec:1234" in destinations