fix: share async topology monitor per cluster

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,10 @@ def __init__(
             When ``None``, falls back to ``connection_getter`` (test/back-compat).
         """
         self._provider = provider
+        # Every provider sharing this monitor; each one's cache receives the
+        # topologies the monitor publishes. Weak so a provider whose connection
+        # was dropped without close() doesn't stay alive through the monitor.
+        self._subscribers: weakref.WeakSet = weakref.WeakSet()
         self._connection_getter = connection_getter
         self._connection_factory = connection_factory
         self._owned_conn: Optional[Any] = None
@@ -148,6 +154,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 +379,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,15 @@
 
 Topology = Tuple[HostInfo, ...]
 
+# Topology monitors shared by every provider of the same cluster. Every async
+# connect builds its own provider, so a per-provider monitor meant one
+# dedicated monitoring connection per application connection (issue #1284).
+# Sync parity: monitor_service.run_if_absent(ClusterTopologyMonitorImpl,
+# cluster_id, ...). Keyed by event loop as well because an asyncio monitor
+# task is bound to 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 +248,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 +385,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 +395,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 +441,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 +675,20 @@ 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).
+        Sync parity: ``RdsHostListProvider.stop_monitor`` stops the cluster's
+        monitor for every connection (blue/green switchover). Closing a
+        connection does not stop it, as in sync: the monitor lives 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()

tests/unit/test_aio_shared_topology_monitor.py

@@ -0,0 +1,174 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License").
+#  You may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""Async topology monitors are shared per cluster, like sync's
+``monitor_service.run_if_absent(ClusterTopologyMonitorImpl, cluster_id, ...)``.
+
+Every async connect builds its own host list provider, so a per-provider
+monitor meant one dedicated monitoring connection per application connection
+(issue #1284).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import gc
+from typing import List
+from unittest.mock import MagicMock
+
+import pytest
+
+from aws_advanced_python_wrapper.aio import cleanup as aio_cleanup
+from aws_advanced_python_wrapper.aio.host_list_provider import \
+    AsyncAuroraHostListProvider
+from aws_advanced_python_wrapper.utils.properties import Properties
+
+
+@pytest.fixture(autouse=True)
+def _clear_hooks():
+    aio_cleanup.clear_shutdown_hooks()
+    yield
+    aio_cleanup.clear_shutdown_hooks()
+
+
+class _ConnFactory:
+    """Counts dedicated monitoring connections opened."""
+
+    def __init__(self) -> None:
+        self.opened = 0
+
+    async def __call__(self):
+        self.opened += 1
+        return object()
+
+
+def _provider(
+        factory: _ConnFactory,
+        host: str = "cluster.example",
+        rows=(("srv-1", True), ("srv-2", False))) -> AsyncAuroraHostListProvider:
+    provider = AsyncAuroraHostListProvider(
+        Properties({"host": host, "port": "5432"}),
+        MagicMock(),
+        monitor_connection_factory=factory,
+    )
+
+    async def _rows(_conn):  # noqa: ARG001 - signature fixed
+        return list(rows)
+
+    provider._run_topology_query = _rows  # type: ignore[method-assign]
+    return provider
+
+
+async def _stop_all(providers: List[AsyncAuroraHostListProvider]) -> None:
+    for p in providers:
+        await p.stop()
+
+
+def test_providers_for_same_cluster_share_one_monitor() -> None:
+    factory = _ConnFactory()
+
+    async def _body():
+        providers = [_provider(factory) for _ in range(10)]
+        for p in providers:
+            await p.force_refresh(object())
+        await asyncio.sleep(0.05)  # let the background task open its connection
+
+        monitors = {id(p._get_or_create_monitor()) for p in providers}
+        assert len(monitors) == 1
+        assert factory.opened == 1
+        await _stop_all(providers)
+
+    asyncio.run(_body())
+
+
+def test_shared_monitor_publishes_topology_to_every_provider() -> None:
+    factory = _ConnFactory()
+
+    async def _body():
+        first, second = _provider(factory), _provider(factory)
+        await first.force_refresh(object())
+        await second.force_refresh(object())
+        second._topology_cache = None
+        # Simulate a background tick discovering a new topology.
+        monitor = first._get_or_create_monitor()
+        monitor._publish_topology(await first._fetch_from_db(object()))
+
+        assert second._topology_cache is not None
+        assert {h.host for h in second._topology_cache} == \
+            {h.host for h in first._topology_cache}
+        await _stop_all([first, second])
+
+    asyncio.run(_body())
+
+
+def test_providers_for_different_clusters_get_separate_monitors() -> None:
+    factory = _ConnFactory()
+
+    async def _body():
+        a = _provider(factory, host="cluster-a.example")
+        b = _provider(factory, host="cluster-b.example")
+        await a.force_refresh(object())
+        await b.force_refresh(object())
+
+        assert a._get_or_create_monitor() is not b._get_or_create_monitor()
+        await _stop_all([a, b])
+
+    asyncio.run(_body())
+
+
+def test_monitor_outlives_its_connections_and_is_reused() -> None:
+    # Sync parity: closing a connection doesn't stop the cluster's monitor, so
+    # connections that come and go (NullPool, a pool draining to zero) reuse
+    # one monitoring connection instead of opening a new one each time.
+    factory = _ConnFactory()
+
+    async def _body():
+        first = _provider(factory)
+        await first.force_refresh(object())
+        await asyncio.sleep(0.05)
+        monitor = first._get_or_create_monitor()
+        del first
+        gc.collect()
+
+        later = _provider(factory)
+        await later.force_refresh(object())
+        await asyncio.sleep(0.05)
+        assert later._get_or_create_monitor() is monitor
+        assert monitor.is_running()
+        assert factory.opened == 1
+        await later.stop()
+
+    asyncio.run(_body())
+
+
+def test_stop_stops_the_shared_monitor_for_the_cluster() -> None:
+    # Sync parity: RdsHostListProvider.stop_monitor stops the cluster's monitor
+    # (blue/green switchover); the next provider starts a fresh one.
+    factory = _ConnFactory()
+
+    async def _body():
+        first, second = _provider(factory), _provider(factory)
+        await first.force_refresh(object())
+        await second.force_refresh(object())
+        monitor = first._get_or_create_monitor()
+
+        await first.stop()
+        assert not monitor.is_running()
+
+        later = _provider(factory)
+        await later.force_refresh(object())
+        assert later._get_or_create_monitor() is not monitor
+        await later.stop()
+
+    asyncio.run(_body())