From f4bb24950317334044fb3a77ebba3b8de9b09fb2 Mon Sep 17 00:00:00 2001 From: AhmadMasry Date: Tue, 6 Oct 2026 16:33:53 +0300 Subject: [PATCH] fix: share async topology monitor per cluster Every AsyncAwsWrapperConnection.connect() builds its own host list provider, and each provider created its own AsyncClusterTopologyMonitor with its own dedicated monitoring connection. A pool of N connections to one cluster held 2N server connections. Sync shares one monitor per cluster id through monitor_service.run_if_absent. Keep async topology monitors in a module-level registry keyed by (provider class, cluster id, event loop). Providers for the same cluster reuse the running monitor, and the monitor publishes each topology into every provider that uses it. Lifecycle matches sync: closing a connection doesn't stop the monitor; stop() stops it for the cluster, like RdsHostListProvider.stop_monitor, and release_resources_async() stops it at shutdown. Fixes #1284 --- .../aio/cluster_topology_monitor.py | 25 ++- .../aio/host_list_provider.py | 40 +++- .../unit/test_aio_shared_topology_monitor.py | 174 ++++++++++++++++++ 3 files changed, 227 insertions(+), 12 deletions(-) create mode 100644 tests/unit/test_aio_shared_topology_monitor.py diff --git a/aws_advanced_python_wrapper/aio/cluster_topology_monitor.py b/aws_advanced_python_wrapper/aio/cluster_topology_monitor.py index 265f8aa8b..8b1de6057 100644 --- a/aws_advanced_python_wrapper/aio/cluster_topology_monitor.py +++ b/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: diff --git a/aws_advanced_python_wrapper/aio/host_list_provider.py b/aws_advanced_python_wrapper/aio/host_list_provider.py index 38345699c..bec12997c 100644 --- a/aws_advanced_python_wrapper/aio/host_list_provider.py +++ b/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() diff --git a/tests/unit/test_aio_shared_topology_monitor.py b/tests/unit/test_aio_shared_topology_monitor.py new file mode 100644 index 000000000..8b681024b --- /dev/null +++ b/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())