fix(aio): keep host_monitoring_v2 monitors per event loop

aws_advanced_python_wrapper/aio/host_monitoring_plugin.py

@@ -22,7 +22,8 @@
   ``plugin_service.force_connect`` and probes THAT (never the in-band
   application connection). Sync ref: ``HostMonitorV2.check_connection_status``.
 * **Shared monitors.** One monitor per ``"{time}:{interval}:{count}:{url}"`` key
-  is shared across every connection with the same parameters, held in a
+  and event loop is shared across every connection with the same parameters on
+  that loop, held in a
   module-level registry with sliding idle expiry. Sync ref:
   ``MonitorServiceV2.get_monitor`` + ``_CACHE_CLEANUP_NANO``.
 * **Execute-window-bounded monitoring.** Each network-bound execute registers a
@@ -51,6 +52,7 @@
 from __future__ import annotations
 
 import asyncio
+import threading
 import time
 import weakref
 from typing import (TYPE_CHECKING, Any, Awaitable, Callable, Dict, List,
@@ -96,12 +98,21 @@
 # connects to an unreachable host hang without a bound, so always supply one.
 _DEFAULT_MONITORING_CONNECT_TIMEOUT_SEC = 10
 
-# Module-level shared-monitor registry, keyed "{time}:{interval}:{count}:{url}".
+# Module-level shared-monitor registry, keyed by "{time}:{interval}:{count}:{url}"
+# and event loop.
 # Holds monitors (which reference the plugin service, NOT the plugin), so a
 # closed connection's plugin stays collectable. Sliding-expiry disposal is
 # opportunistic (on next monitor request) + on host-deleted notifications +
 # via the single module shutdown hook.
-_monitors: Dict[str, AsyncHostMonitorV2] = {}
+#
+# The key also holds the event loop the monitor runs on: an asyncio task only
+# runs on the loop that started it, so an application with one loop per thread
+# gets one monitor per loop, and a loop never stops another live loop's
+# monitors (#1287). The lock guards the registry across those threads, like
+# the RLock in sync's monitor service.
+_RegistryKey = Tuple[str, asyncio.AbstractEventLoop]
+_monitors: Dict[_RegistryKey, AsyncHostMonitorV2] = {}
+_monitors_lock = threading.RLock()
 _shutdown_hook_registered: bool = False
 
 
@@ -142,11 +153,12 @@ def _cleanup_idle_monitors() -> None:
     task/loop is no longer usable. Opportunistic: called on each monitor
     request so cleanup piggybacks on activity rather than a standing task."""
     now = time.monotonic()
-    for key, monitor in list(_monitors.items()):
-        if not monitor.is_usable() or (
-                monitor.can_dispose and now - monitor.last_used >= _MONITOR_EXPIRATION_SEC):
-            monitor.stop()
-            _monitors.pop(key, None)
+    with _monitors_lock:
+        for key, monitor in list(_monitors.items()):
+            if not monitor.is_usable() or (
+                    monitor.can_dispose and now - monitor.last_used >= _MONITOR_EXPIRATION_SEC):
+                monitor.stop()
+                _monitors.pop(key, None)
 
 
 def _get_or_create_monitor(
@@ -157,34 +169,37 @@ def _get_or_create_monitor(
         failure_detection_interval_ms: int,
         failure_detection_count: int,
         aborted_counter: Optional[TelemetryCounter]) -> AsyncHostMonitorV2:
-    _cleanup_idle_monitors()
     key = _monitor_key(
         failure_detection_time_ms, failure_detection_interval_ms,
         failure_detection_count, host_info.url)
-    existing = _monitors.get(key)
-    if existing is not None and existing.is_usable():
-        existing.touch()
-        return existing
-    if existing is not None:
-        # Stale (task done or bound to a closed loop): replace it.
-        existing.stop()
-    monitor = AsyncHostMonitorV2(
-        plugin_service, host_info, props, failure_detection_time_ms,
-        failure_detection_interval_ms, failure_detection_count,
-        aborted_counter, key)
-    monitor.start()
-    _monitors[key] = monitor
-    _register_shutdown_hook_once()
-    return monitor
+    registry_key = (key, asyncio.get_running_loop())
+    with _monitors_lock:
+        _cleanup_idle_monitors()
+        existing = _monitors.get(registry_key)
+        if existing is not None and existing.is_usable():
+            existing.touch()
+            return existing
+        if existing is not None:
+            # Stale (task done or loop closed): replace it.
+            existing.stop()
+        monitor = AsyncHostMonitorV2(
+            plugin_service, host_info, props, failure_detection_time_ms,
+            failure_detection_interval_ms, failure_detection_count,
+            aborted_counter, key)
+        monitor.start()
+        _monitors[registry_key] = monitor
+        _register_shutdown_hook_once()
+        return monitor
 
 
 def _stop_monitors_for_hosts(host_urls: Set[str]) -> None:
     """Stop + drop shared monitors whose host was deleted from the topology."""
-    for key, monitor in list(_monitors.items()):
-        if monitor.host_url in host_urls:
-            logger.debug("HostMonitorV2.StopMonitoringThread", monitor.host)
-            monitor.stop()
-            _monitors.pop(key, None)
+    with _monitors_lock:
+        for key, monitor in list(_monitors.items()):
+            if monitor.host_url in host_urls:
+                logger.debug("HostMonitorV2.StopMonitoringThread", monitor.host)
+                monitor.stop()
+                _monitors.pop(key, None)
 
 
 async def _stop_all_monitors() -> None:
@@ -194,9 +209,10 @@ async def _stop_all_monitors() -> None:
     plugin or connection. Idempotent.
     """
     global _shutdown_hook_registered
-    monitors = list(_monitors.values())
-    _monitors.clear()
-    _shutdown_hook_registered = False
+    with _monitors_lock:
+        monitors = list(_monitors.values())
+        _monitors.clear()
+        _shutdown_hook_registered = False
     for monitor in monitors:
         monitor.stop()
     if monitors:
@@ -211,10 +227,11 @@ def _reset_monitor_registry() -> None:
     creates a fresh loop), so cancelling and forgetting them is sufficient.
     """
     global _shutdown_hook_registered
-    for monitor in list(_monitors.values()):
-        monitor.stop()
-    _monitors.clear()
-    _shutdown_hook_registered = False
+    with _monitors_lock:
+        for monitor in list(_monitors.values()):
+            monitor.stop()
+        _monitors.clear()
+        _shutdown_hook_registered = False
 
 
 class _AsyncMonitoringContext:
@@ -318,12 +335,12 @@ def can_dispose(self) -> bool:
         return not self._active_contexts and not self._new_contexts
 
     def is_usable(self) -> bool:
+        # Usable from the loop that owns it; other loops get their own monitor
+        # (registry key), so a different running loop isn't a reason to stop
+        # this one (#1287). Only a closed loop is.
         if self._stopped or self._task is None or self._task.done():
             return False
-        try:
-            return self._loop is asyncio.get_running_loop()
-        except RuntimeError:
-            return False
+        return self._loop is not None and not self._loop.is_closed()
 
     def stop(self) -> None:
         """Signal the monitor to stop (sync, fire-and-forget). The task drains

tests/unit/test_aio_host_monitoring_plugin.py

@@ -26,6 +26,7 @@
 import asyncio
 import gc
 import socket as socket_mod
+import threading
 import time
 import weakref
 from unittest.mock import AsyncMock, MagicMock
@@ -666,3 +667,63 @@ def test_async_dialects_support_abort_connection():
 
     assert AsyncPsycopgDriverDialect().supports_abort_connection() is True
     assert AsyncAiomysqlDriverDialect().supports_abort_connection() is True
+
+
+# ---- Multiple event loops (#1287) --------------------------------------
+
+
+class _LoopThread:
+    """An event loop running forever on its own thread, like one loop per
+    worker thread in an application."""
+
+    def __init__(self) -> None:
+        self.loop = asyncio.new_event_loop()
+        self._thread = threading.Thread(target=self.loop.run_forever, daemon=True)
+        self._thread.start()
+
+    def run(self, coro):
+        return asyncio.run_coroutine_threadsafe(coro, self.loop).result(timeout=5)
+
+    def close(self) -> None:
+        self.loop.call_soon_threadsafe(self.loop.stop)
+        self._thread.join(timeout=5)
+        self.loop.close()
+
+
+def test_event_loops_do_not_stop_each_others_monitors():
+    """A statement on one loop must not stop another live loop's monitor for
+    the same host, and each loop keeps reusing its own monitor."""
+    other = _LoopThread()
+    try:
+        p_other, svc_other, *_ = _build(grace_ms=100000)
+        svc_other.force_connect = AsyncMock()
+
+        async def _body():
+            p_here, svc_here, *_ = _build(grace_ms=100000)
+            svc_here.force_connect = AsyncMock()
+            seen = set()
+            for _ in range(3):
+                other.run(_execute_once(p_other))
+                await _execute_once(p_here)
+                seen.update(id(m) for m in list(_monitors.values()))
+            assert len(_monitors) == 2
+            assert len(seen) == 2
+            assert all(not m._stopped for m in list(_monitors.values()))
+
+        asyncio.run(_body())
+    finally:
+        efm._reset_monitor_registry()
+        other.close()
+
+
+def test_monitor_whose_loop_closed_is_replaced():
+    plugin, svc, *_ = _build(grace_ms=100000)
+    svc.force_connect = AsyncMock()
+
+    asyncio.run(_execute_once(plugin))
+    (old,) = list(_monitors.values())
+
+    asyncio.run(_execute_once(plugin))
+    (new,) = list(_monitors.values())
+    assert new is not old
+    assert old._stopped is True