fix: Release all stopped monitors (#397)

Author: joyc-bq

aws_advanced_python_wrapper/host_monitoring_plugin.py

@@ -432,7 +432,7 @@ def run(self):
                             self._monitor_container.release_monitor(self)
                             break
 
-                        sleep(Monitor._INACTIVE_SLEEP_MS / 1000)
+                        self.sleep(Monitor._INACTIVE_SLEEP_MS / 1000)
                         continue
 
                     status_check_start_time_ns = perf_counter_ns()
@@ -482,7 +482,7 @@ def run(self):
                         # Use this delay for all active contexts
                         self._host_check_timeout_ms = delay_ms
 
-                    sleep(delay_ms / 1000)
+                    self.sleep(delay_ms / 1000)
                 except InterruptedError as e:
                     raise e
                 except Exception as e:
@@ -494,6 +494,7 @@ def run(self):
             logger.debug("Monitor.StoppingMonitorUnhandledException", self._host_info.host)
             logger.debug(e, exc_info=True)
         finally:
+            self._monitor_container.release_monitor(self)
             self.stop()
             if self._monitoring_conn is not None:
                 try:
@@ -543,6 +544,10 @@ def _execute_conn_check(self, conn: Connection, timeout_sec: float):
             driver_dialect.execute("Cursor.execute", lambda: cursor.execute(query), query, exec_timeout=timeout_sec)
             cursor.fetchone()
 
+    # Used to help with testing
+    def sleep(self, duration: int):
+        sleep(duration)
+
 
 class MonitoringThreadContainer:
     """

tests/unit/test_monitor.py

@@ -153,6 +153,35 @@ def test_run_host_available(
     assert container._tasks_map.get(monitor) is None
 
 
+def test_ensure_stopped_monitor_removed_from_map(
+        mocker,
+        monitor,
+        host_info,
+        monitoring_conn_props,
+        mock_plugin_service,
+        mock_conn,
+        mock_driver_dialect):
+    remove_delays()
+    host_alias = "host-1"
+    container = MonitoringThreadContainer()
+    container._monitor_map.put_if_absent(host_alias, monitor)
+    container._tasks_map.put_if_absent(monitor, mocker.MagicMock())
+
+    executor = ThreadPoolExecutor()
+    context = MonitoringContext(monitor, mock_conn, mock_driver_dialect, 10, 1, 3)
+
+    mocker.patch(
+        "aws_advanced_python_wrapper.host_monitoring_plugin.Monitor.sleep", side_effect=InterruptedError())
+    monitor.start_monitoring(context)
+
+    future = executor.submit(monitor.run)
+    sleep(0.1)  # Allow some time for the monitor to loop
+    wait([future], 3)
+
+    assert container._monitor_map.get(host_alias) is None
+    assert container._tasks_map.get(monitor) is None
+
+
 def test_run_host_unavailable(
         mocker,
         monitor,