add heartbeat during monitor execution

Author: GabrielSalla

docs/configuration_file.md

@@ -62,15 +62,15 @@ application_queue:
 - `controller_procedures`: Object. Procedures to be executed by the Controller and their settings.
 - `controller_procedures.monitors_stuck`: Object. Settings for the procedure to fix monitors stuck in "queued" or "running" status.
 - `controller_procedures.monitors_stuck.schedule`: String using Cron format. Schedule to execute the `monitors_stuck` procedure.
-- `controller_procedures.monitors_stuck.params.time_tolerance`: Integer. Time tolerance in seconds for a monitor to be considered as stuck.
+- `controller_procedures.monitors_stuck.params.time_tolerance`: Integer. Time tolerance in seconds for a monitor to be considered as stuck. This parameter is directly impacted by the `executor_monitor_heartbeat_time` setting and the recommended value is 2 times the heartbeat time.
 
 ## Executor Settings
 - `executor_concurrency`: Integer. Number of tasks that can be executed at the same time by each Executor.
 - `executor_sleep`: Integer. Time, in seconds, the Executor will sleep when there are no tasks in the queue before trying again.
 - `executor_monitor_timeout`: Integer. Timeout, in seconds, for monitor execution.
 - `executor_reaction_timeout`: Integer. Timeout, in seconds, for reactions execution.
 - `executor_request_timeout`: Integer. Timeout, in seconds, for requests execution.
-- `executor_monitor_heartbeat_time`: Integer. Time, in seconds, between each monitor heartbeat.
+- `executor_monitor_heartbeat_time`: Integer. Time, in seconds, between each monitor heartbeat. This parameter impacts the controller procedure `monitors_stuck.time_tolerance` parameter.
 
 ## Issues Creation
 - `max_issues_creation`: Integer. Maximum number of issues that can be created by each monitor in a single search. Can be overridden by the monitors' configuration.

src/components/executor/monitor_handler.py

@@ -8,8 +8,10 @@
 import prometheus_client
 from pydantic import ValidationError
 
+import components.task_manager as task_manager
 import registry as registry
 from base_exception import BaseSentinelaException
+from configs import configs
 from data_models.process_monitor_payload import ProcessMonitorPayload
 from internal_database import get_session
 from models import Alert, Issue, Monitor
@@ -321,6 +323,13 @@ async def _run_routines(monitor: Monitor, tasks: list[Literal["search", "update"
         await _alerts_routine(monitor)
 
 
+async def _heartbeat_routine(monitor: Monitor) -> None:
+    """Heartbeat routine for the monitor, setting the last heartbeat to the current time"""
+    while True:
+        await monitor.set_last_heartbeat()
+        await asyncio.sleep(configs.executor_monitor_heartbeat_time)
+
+
 async def run(message: dict[Any, Any]) -> None:
     """Process a message with type 'process_monitor', loading the monitor and executing it's
     routines, while also detecting errors and reporting them accordingly"""
@@ -339,12 +348,12 @@ async def run(message: dict[Any, Any]) -> None:
         _logger.error(f"Monitor {monitor_id} not found. Skipping message")
         return
 
-    await registry.wait_monitor_loaded(monitor_id)
-
     # Skip executing the monitor if it's already running
     if monitor.running:
         return
 
+    await registry.wait_monitor_loaded(monitor_id)
+
     prometheus_labels = {
         "monitor_id": monitor.id,
         "monitor_name": monitor.name,
@@ -353,6 +362,10 @@ async def run(message: dict[Any, Any]) -> None:
     monitor_running = prometheus_monitor_running.labels(**prometheus_labels)
     monitor_running.inc()
 
+    heartbeat_task = task_manager.create_task(
+        _heartbeat_routine(monitor), parent_task=asyncio.current_task()
+    )
+
     try:
         await monitor.set_running(True)
 
@@ -376,6 +389,9 @@ async def run(message: dict[Any, Any]) -> None:
         _logger.error(traceback.format_exc().strip())
         _logger.info("Exception caught successfully, going on")
     finally:
+        # Cancel the heartbeat task
+        heartbeat_task.cancel()
+
         # Refresh the monitor before updating to prevent overwriting information that might have
         # changed while the routines were executing
         await monitor.refresh()

tests/components/executor/test_monitor_handler.py

@@ -1,6 +1,6 @@
 import asyncio
 import time
-from datetime import datetime, timezone
+from datetime import datetime, timedelta, timezone
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
@@ -11,7 +11,7 @@
 from data_models.monitor_options import AlertOptions, CountRule, IssueOptions, PriorityLevels
 from models import Alert, AlertPriority, AlertStatus, Issue, IssueStatus, Monitor
 from tests.test_utils import assert_message_in_log, assert_message_not_in_log
-from utils.time import time_since
+from utils.time import now, time_since
 
 pytestmark = pytest.mark.asyncio(loop_scope="session")
 
@@ -1100,6 +1100,27 @@ async def do_nothing(monitor): ...
     assert sample_monitor.active_alerts[0].id == alert.id
 
 
+# Test _heartbeat_routine
+
+
+async def test_heartbeat_routine(monkeypatch, sample_monitor: Monitor):
+    """'_heartbeat_routine' should handle execution timeouts while running the monitor routines"""
+    monkeypatch.setattr(monitor_handler.configs, "executor_monitor_heartbeat_time", 0.5)
+
+    await sample_monitor.refresh()
+    assert sample_monitor.last_heartbeat is None
+
+    heartbeat_task = asyncio.create_task(monitor_handler._heartbeat_routine(sample_monitor))
+
+    await asyncio.sleep(0)
+    for _ in range(4):
+        await sample_monitor.refresh()
+        assert sample_monitor.last_heartbeat > now() - timedelta(seconds=0.1)
+        await asyncio.sleep(0.5)
+
+    heartbeat_task.cancel()
+
+
 # Test run
 
 
@@ -1158,6 +1179,33 @@ async def test_run_monitor_skip_running(caplog, mocker, sample_monitor: Monitor)
     run_routines_spy.assert_not_called()
 
 
+@pytest.mark.flaky(reruns=2)
+async def test_run_monitor_heartbeat(monkeypatch, sample_monitor: Monitor):
+    """'run' should handle execution timeouts while running the monitor routines"""
+
+    async def sleep(monitor, tasks):
+        await asyncio.sleep(2.1)
+
+    monkeypatch.setattr(monitor_handler, "_run_routines", sleep)
+
+    monkeypatch.setattr(monitor_handler.configs, "executor_monitor_heartbeat_time", 0.5)
+
+    await sample_monitor.refresh()
+    assert sample_monitor.last_heartbeat is None
+
+    run_task = asyncio.create_task(
+        monitor_handler.run({"payload": {"monitor_id": sample_monitor.id, "tasks": ["search"]}})
+    )
+
+    await asyncio.sleep(0.05)
+    for _ in range(4):
+        await sample_monitor.refresh()
+        assert sample_monitor.last_heartbeat > now() - timedelta(seconds=0.1)
+        await asyncio.sleep(0.5)
+
+    await run_task
+
+
 @pytest.mark.parametrize("tasks", [["search"], ["update"], ["search", "update"]])
 async def test_run_monitor_set_running(mocker, sample_monitor: Monitor, tasks):
     """'run' should set the monitor's 'running' flag to 'True' while running the routines"""