fix(aio): async monitoring fixes (#1284, #1276, #1287, #1288)

aws_advanced_python_wrapper/aio/aurora_connection_tracker.py

@@ -270,6 +270,8 @@ def __init__(
         self._plugin_service = plugin_service
         self._tracker = tracker or AsyncOpenedConnectionTracker()
         self._current_writer: Optional[HostInfo] = None
+        # The instance a cluster-endpoint connection landed on.
+        self._connected_instance: Optional[HostInfo] = None
         self._pending_invalidations: Set[asyncio.Task] = set()
 
         # Telemetry counter -- the sync tracker plugin doesn't emit this,
@@ -342,6 +344,10 @@ async def _pin_current_writer(self, host_info: HostInfo, conn: Any) -> None:
             role = await self._plugin_service.get_host_role(conn)
             if role == HostRole.WRITER:
                 host = self._plugin_service.current_host_info
+                # Topology names instances, so pin the instance behind a
+                # cluster endpoint; if it's unknown, keep the topology pin.
+                if host is not None and not rds.is_rds_instance(host.host):
+                    host = self._connected_instance
                 if host is not None and (
                         self._current_writer is None
                         or not self._same_host(host, self._current_writer)):
@@ -376,6 +382,7 @@ async def _fill_instance_alias(
             identified = await self._plugin_service.identify_connection(conn)
             if identified is None:
                 return
+            self._connected_instance = identified
             host_info.reset_aliases()
             host_info.add_alias(host_info.as_alias())
             host_info.add_alias(*identified.as_aliases())
@@ -505,9 +512,14 @@ async def _invalidate_departed_host(
         """
         if pre_failover_host is None:
             return
+        # A cluster-endpoint connection departed the instance it landed on.
+        if (self._connected_instance is not None
+                and not RdsUtils().is_rds_instance(pre_failover_host.host)):
+            pre_failover_host = self._connected_instance
         post = self._plugin_service.current_host_info
         if post is None or self._same_host(pre_failover_host, post):
             return
+        self._connected_instance = post
         if (self._current_writer is not None
                 and not self._same_host(pre_failover_host, self._current_writer)
                 and not self._same_host(post, self._current_writer)):

aws_advanced_python_wrapper/aio/cluster_topology_monitor.py

@@ -18,7 +18,8 @@
 :meth:`AsyncAuroraHostListProvider.force_refresh`, and sleeps. Replaces
 the sync :class:`ClusterTopologyMonitor`'s thread-based loop.
 
-The monitor is deliberately minimal: one task per provider instance, fixed
+The monitor is deliberately minimal: one task per cluster (shared by every
+provider for that cluster, see ``host_list_provider._topology_monitors``), fixed
 interval, no suggestions feedback loop (sync EFM uses that; the async EFM
 plugin runs its own monitors). Cancellation is clean -- ``stop()`` cancels
 the task and awaits its exit.
@@ -45,6 +46,7 @@
 import asyncio
 import copy
 import time
+import weakref
 from typing import (TYPE_CHECKING, Any, Awaitable, Callable, Dict, Optional,
                     Set, Tuple)
 
@@ -106,6 +108,9 @@ def __init__(
             When ``None``, falls back to ``connection_getter`` (test/back-compat).
         """
         self._provider = provider
+        # Providers that receive published topologies. Weak so the monitor
+        # doesn't keep a dropped provider alive.
+        self._subscribers: weakref.WeakSet = weakref.WeakSet()
         self._connection_getter = connection_getter
         self._connection_factory = connection_factory
         self._owned_conn: Optional[Any] = None
@@ -148,6 +153,10 @@ def last_topology(self) -> Topology:
     def is_running(self) -> bool:
         return self._task is not None and not self._task.done()
 
+    def subscribe(self, provider: Any) -> None:
+        """Add ``provider`` to the set that receives published topologies."""
+        self._subscribers.add(provider)
+
     def is_in_panic_mode(self) -> bool:
         """True when panic-mode probe tasks are currently running."""
         return bool(self._probe_tasks)
@@ -369,12 +378,13 @@ def _publish_topology(self, topology: Topology) -> None:
         topology = self._role_corrected(topology)
         self._last_topology = topology
         self._check_for_writer_change(topology)
-        adopt = getattr(self._provider, "adopt_topology", None)
-        if adopt is not None:
-            try:
-                adopt(topology)
-            except Exception:  # noqa: BLE001 - cache publication is best-effort
-                pass
+        for provider in list(self._subscribers) or [self._provider]:
+            adopt = getattr(provider, "adopt_topology", None)
+            if adopt is not None:
+                try:
+                    adopt(topology)
+                except Exception:  # noqa: BLE001 - cache publication is best-effort
+                    pass
         self._topology_updated.set()
 
     async def _wait_for_tick(self, interval: float) -> None:

aws_advanced_python_wrapper/aio/host_list_provider.py

@@ -58,6 +58,12 @@
 
 Topology = Tuple[HostInfo, ...]
 
+# One topology monitor per cluster, shared by all its providers (each connect
+# builds its own provider). The event loop is part of the key because a monitor
+# task only runs on the loop that started it.
+_MonitorKey = Tuple[type, str, asyncio.AbstractEventLoop]
+_topology_monitors: Dict[_MonitorKey, Any] = {}
+
 
 def _is_programming_error(ex: BaseException) -> bool:
     """True when ``ex`` is a PEP-249 ``ProgrammingError`` (bad SQL / not an
@@ -239,6 +245,7 @@ def __init__(
         # Monitor wiring (N.1b). Lazy-created on first refresh so
         # provider construction stays cheap and test-friendly.
         self._monitor: Optional[Any] = None  # AsyncClusterTopologyMonitor
+        self._monitor_key: Optional[_MonitorKey] = None
         self._last_conn: Optional[Any] = None
         # Panic-mode probe: opens a pipeline connection to a host and
         # classifies its role, used by the monitor to discover the new
@@ -375,7 +382,7 @@ async def force_monitoring_refresh(
         return topology or ()
 
     def _get_or_create_monitor(self) -> Optional[Any]:
-        """Lazy-construct the per-provider topology monitor.
+        """Return this cluster's shared topology monitor, creating it if needed.
 
         Returns ``None`` if monitor construction fails or if the
         topology_query is absent (e.g., static/unsupported dialects).
@@ -385,6 +392,17 @@ def _get_or_create_monitor(self) -> Optional[Any]:
             return self._monitor
         if not self._topology_query:
             return None
+        key: _MonitorKey = (
+            type(self), self._cluster_id, asyncio.get_running_loop())
+        for stale_key in [
+                k for k, m in _topology_monitors.items() if not m.is_running()]:
+            del _topology_monitors[stale_key]
+        shared = _topology_monitors.get(key)
+        if shared is not None:
+            shared.subscribe(self)
+            self._monitor = shared
+            self._monitor_key = key
+            return shared
         try:
             from aws_advanced_python_wrapper.aio.cluster_topology_monitor import \
                 AsyncClusterTopologyMonitor
@@ -420,7 +438,10 @@ def _get_or_create_monitor(self) -> Optional[Any]:
             connection_factory=self._monitor_connection_factory,
         )
         monitor.start()
+        monitor.subscribe(self)
         self._monitor = monitor
+        self._monitor_key = key
+        _topology_monitors[key] = monitor
         # Register monitor teardown with the global cleanup hook so
         # release_resources_async() tears it down.
         try:
@@ -651,14 +672,18 @@ def _host_from_server_id(self, server_id: str) -> str:
         return server_id
 
     async def stop(self) -> None:
-        """Tear down the provider's background topology monitor.
+        """Stop this cluster's shared topology monitor.
 
-        Sync parity: sync ``RdsHostListProvider`` stops its monitor on release.
-        The async provider previously no-op'd here, leaking the monitor task
-        (it only stopped via the global cleanup hook).
+        As in sync, closing a connection doesn't stop the monitor; it runs
+        until this, ``release_resources_async()``, or its event loop closes.
         """
         monitor = self._monitor
+        key = self._monitor_key
         self._monitor = None
+        self._monitor_key = None
+        if key is not None and monitor is not None \
+                and _topology_monitors.get(key) is monitor:
+            del _topology_monitors[key]
         if monitor is not None:
             try:
                 await monitor.stop()

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,15 @@
 # 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 (a monitor task only runs on the loop that started it).
 # 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] = {}
+# via the single module shutdown hook. The lock guards it across threads.
+_RegistryKey = Tuple[str, asyncio.AbstractEventLoop]
+_monitors: Dict[_RegistryKey, AsyncHostMonitorV2] = {}
+_monitors_lock = threading.RLock()
 _shutdown_hook_registered: bool = False
 
 
@@ -142,11 +147,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 +163,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 +203,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 +221,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 +329,10 @@ def can_dispose(self) -> bool:
         return not self._active_contexts and not self._new_contexts
 
     def is_usable(self) -> bool:
+        # Other loops have their own monitors, so only a closed loop matters.
         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
@@ -659,11 +668,11 @@ async def _get_monitoring_host_info(self) -> HostInfo:
         current_host_info = self._plugin_service.current_host_info
         if current_host_info is None:
             raise AwsWrapperError(Messages.get("HostMonitoringV2Plugin.HostInfoNone"))
-        self._monitoring_host_info = current_host_info
-        rds_url_type = self._rds_utils.identify_rds_type(self._monitoring_host_info.host)
+        rds_url_type = self._rds_utils.identify_rds_type(current_host_info.host)
 
         try:
             if not rds_url_type.is_rds_cluster:
+                self._monitoring_host_info = current_host_info
                 return self._monitoring_host_info
             logger.debug("HostMonitoringV2Plugin.ClusterEndpointHostInfo")
             current_connection = self._plugin_service.current_connection

aws_advanced_python_wrapper/aio/wrapper.py

@@ -67,6 +67,14 @@
     "aurora_connection_tracker",
 })
 
+# Plugins that use the topology when the dialect has one (Aurora, Multi-AZ,
+# Global), to resolve the instance behind a cluster endpoint. Sync gets a
+# topology provider from those dialects whatever the plugins.
+_TOPOLOGY_WHEN_AVAILABLE_PLUGINS = frozenset({
+    "host_monitoring",
+    "host_monitoring_v2",
+})
+
 
 def _build_host_list_provider(
         props: Properties,
@@ -76,7 +84,8 @@ def _build_host_list_provider(
     """Pick an async host list provider based on plugins + DB dialect.
 
     Selection order, matching sync database_dialect.py:
-      * ``plugins`` references no topology-requiring plugin -> static.
+      * ``plugins`` references no topology-requiring plugin, and no host
+        monitoring plugin on a topology-aware dialect -> static.
       * database_dialect is a MultiAz dialect -> MultiAz provider.
       * database_dialect is a GlobalAurora dialect -> GlobalAurora provider.
       * Otherwise -> Aurora provider (the default topology path).
@@ -90,9 +99,14 @@ def _build_host_list_provider(
         AsyncMultiAzHostListProvider, AsyncStaticHostListProvider)
     from aws_advanced_python_wrapper.aio.plugin_factory import \
         parse_plugins_property
-
-    codes = parse_plugins_property(props) or []
-    if not any(c.strip() in _TOPOLOGY_REQUIRING_PLUGINS for c in codes):
+    from aws_advanced_python_wrapper.database_dialect import \
+        TopologyAwareDatabaseDialect
+
+    codes = {c.strip() for c in parse_plugins_property(props) or []}
+    needs_topology = bool(codes & _TOPOLOGY_REQUIRING_PLUGINS) or (
+        bool(codes & _TOPOLOGY_WHEN_AVAILABLE_PLUGINS)
+        and isinstance(database_dialect, TopologyAwareDatabaseDialect))
+    if not needs_topology:
         return AsyncStaticHostListProvider(props)
 
     # Topology-aware path. Pick the provider matching the dialect.