diff --git a/aws_advanced_python_wrapper/aio/host_monitoring_plugin.py b/aws_advanced_python_wrapper/aio/host_monitoring_plugin.py index 7458f91d..76114141 100644 --- a/aws_advanced_python_wrapper/aio/host_monitoring_plugin.py +++ b/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 diff --git a/tests/unit/test_aio_host_monitoring_plugin.py b/tests/unit/test_aio_host_monitoring_plugin.py index 88aa91b3..6daae155 100644 --- a/tests/unit/test_aio_host_monitoring_plugin.py +++ b/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