diff --git a/aws_advanced_python_wrapper/aio/aurora_connection_tracker.py b/aws_advanced_python_wrapper/aio/aurora_connection_tracker.py index 93e287ee..455b32cd 100644 --- a/aws_advanced_python_wrapper/aio/aurora_connection_tracker.py +++ b/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)): diff --git a/aws_advanced_python_wrapper/aio/cluster_topology_monitor.py b/aws_advanced_python_wrapper/aio/cluster_topology_monitor.py index 265f8aa8..24656720 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,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: diff --git a/aws_advanced_python_wrapper/aio/host_list_provider.py b/aws_advanced_python_wrapper/aio/host_list_provider.py index 38345699..68718c75 100644 --- a/aws_advanced_python_wrapper/aio/host_list_provider.py +++ b/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() diff --git a/aws_advanced_python_wrapper/aio/host_monitoring_plugin.py b/aws_advanced_python_wrapper/aio/host_monitoring_plugin.py index 7458f91d..400282db 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,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 diff --git a/aws_advanced_python_wrapper/aio/wrapper.py b/aws_advanced_python_wrapper/aio/wrapper.py index 4ec42b83..7aba09b9 100644 --- a/aws_advanced_python_wrapper/aio/wrapper.py +++ b/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. diff --git a/tests/unit/test_aio_aurora_connection_tracker.py b/tests/unit/test_aio_aurora_connection_tracker.py index 9056d959..62260d14 100644 --- a/tests/unit/test_aio_aurora_connection_tracker.py +++ b/tests/unit/test_aio_aurora_connection_tracker.py @@ -1005,3 +1005,123 @@ async def _raising(): assert plugin._current_writer.host == new_writer.host asyncio.run(_run()) + + +def test_cluster_endpoint_connection_survives_its_first_execute(): + """The pin is the instance, not the cluster endpoint, so the first + execute doesn't see a writer change and close the connection.""" + plugin, svc, driver_dialect, tracker = _build() + cluster = "mydb.cluster-xyz.us-east-1.rds.amazonaws.com" + instance = "inst-1.xyz.us-east-1.rds.amazonaws.com" + writer = HostInfo(host=instance, port=5432, role=HostRole.WRITER) + + async def _refresh(*a, **k): + svc._all_hosts = (writer,) + svc.refresh_host_list = AsyncMock(side_effect=_refresh) + svc.identify_connection = AsyncMock( + return_value=HostInfo(host=instance, port=5432)) + svc.get_host_role = AsyncMock(return_value=HostRole.WRITER) + + conn = _plain_conn("conn") + + async def _connect_func(): + return conn + + async def _run(): + svc._current_host_info = HostInfo(host=cluster, port=5432) + await plugin.connect( + target_driver_func=MagicMock(), driver_dialect=driver_dialect, + host_info=HostInfo(host=cluster, port=5432), props=svc.props, + is_initial_connection=True, connect_func=_connect_func) + assert plugin._current_writer is not None + assert plugin._current_writer.host == instance + + result = await plugin.execute( + MagicMock(), "Cursor.execute", AsyncMock(return_value="ok")) + assert result == "ok" + conn.close.assert_not_called() + + asyncio.run(_run()) + + +def _connect_via(plugin, svc, driver_dialect, endpoint, conn): + async def _connect_func(): + return conn + + host_info = HostInfo(host=endpoint, port=5432) + svc._current_host_info = host_info + return plugin.connect( + target_driver_func=MagicMock(), driver_dialect=driver_dialect, + host_info=host_info, props=svc.props, + is_initial_connection=True, connect_func=_connect_func) + + +def test_failover_invalidates_departed_instance_of_cluster_endpoint_connection(): + """With lagged topology, failover still closes idle connections to the + instance a cluster-endpoint connection departed.""" + plugin, svc, driver_dialect, tracker = _build() + cluster = "mydb.cluster-xyz.us-east-1.rds.amazonaws.com" + old = "inst-1.xyz.us-east-1.rds.amazonaws.com" + new = "inst-2.xyz.us-east-1.rds.amazonaws.com" + lagged = (HostInfo(old, 5432, role=HostRole.WRITER), + HostInfo(new, 5432, role=HostRole.READER)) + + async def _refresh(*a, **k): + svc._all_hosts = lagged + svc.refresh_host_list = AsyncMock(side_effect=_refresh) + svc.force_refresh_host_list = AsyncMock(side_effect=_refresh) + svc.identify_connection = AsyncMock(return_value=HostInfo(old, 5432)) + svc.get_host_role = AsyncMock(return_value=HostRole.WRITER) + + idle_conn = _plain_conn("idle") + active_conn = _plain_conn("active") + + async def _run(): + await _connect_via(plugin, svc, driver_dialect, cluster, idle_conn) + await _connect_via(plugin, svc, driver_dialect, cluster, active_conn) + + async def _raising(): + svc._current_host_info = HostInfo(new, 5432, role=HostRole.WRITER) + raise FailoverSuccessError("failover") + with pytest.raises(FailoverSuccessError): + await plugin.execute(MagicMock(), "Cursor.execute", _raising) + idle_conn.close.assert_called() + + asyncio.run(_run()) + + +def test_reader_endpoint_failover_does_not_invalidate_the_writer(): + """A reader-endpoint connection that fails over between readers departs + a reader, so writer connections must survive.""" + plugin, svc, driver_dialect, tracker = _build() + reader_ep = "mydb.cluster-ro-xyz.us-east-1.rds.amazonaws.com" + writer = "inst-1.xyz.us-east-1.rds.amazonaws.com" + reader_1 = "inst-2.xyz.us-east-1.rds.amazonaws.com" + reader_2 = "inst-3.xyz.us-east-1.rds.amazonaws.com" + topology = (HostInfo(writer, 5432, role=HostRole.WRITER), + HostInfo(reader_1, 5432, role=HostRole.READER), + HostInfo(reader_2, 5432, role=HostRole.READER)) + + async def _refresh(*a, **k): + svc._all_hosts = topology + svc.refresh_host_list = AsyncMock(side_effect=_refresh) + svc.force_refresh_host_list = AsyncMock(side_effect=_refresh) + svc.identify_connection = AsyncMock(return_value=HostInfo(reader_1, 5432)) + svc.get_host_role = AsyncMock(return_value=HostRole.READER) + + writer_conn = _plain_conn("writer") + reader_conn = _plain_conn("reader") + + async def _run(): + tracker.track(HostInfo(writer, 5432), writer_conn) + await _connect_via(plugin, svc, driver_dialect, reader_ep, reader_conn) + assert plugin._current_writer.host == writer + + async def _raising(): + svc._current_host_info = HostInfo(reader_2, 5432, role=HostRole.READER) + raise FailoverSuccessError("failover") + with pytest.raises(FailoverSuccessError): + await plugin.execute(MagicMock(), "Cursor.execute", _raising) + writer_conn.close.assert_not_called() + + asyncio.run(_run()) diff --git a/tests/unit/test_aio_host_list_provider_autodetect.py b/tests/unit/test_aio_host_list_provider_autodetect.py index a3bf6996..80de2c69 100644 --- a/tests/unit/test_aio_host_list_provider_autodetect.py +++ b/tests/unit/test_aio_host_list_provider_autodetect.py @@ -110,3 +110,27 @@ class AuroraPgDialect: provider = _build_host_list_provider( Properties({"plugins": "failover"}), dd, AuroraPgDialect()) assert type(provider) is AsyncAuroraHostListProvider + + +def test_host_monitoring_on_topology_aware_dialect_returns_topology_provider() -> None: + # Host monitoring needs the topology to identify the instance behind a + # cluster endpoint. + from aws_advanced_python_wrapper.database_dialect import ( + AuroraMysqlDialect, AuroraPgDialect) + dd = AsyncPsycopgDriverDialect() + for code in ("host_monitoring_v2", "host_monitoring"): + for dialect in (AuroraPgDialect(), AuroraMysqlDialect()): + provider = _build_host_list_provider( + Properties({"plugins": code}), dd, dialect) + assert isinstance(provider, AsyncAuroraHostListProvider), (code, dialect) + + +def test_host_monitoring_on_non_topology_dialect_stays_static() -> None: + # Plain PostgreSQL / RDS instances have no topology to query. + from aws_advanced_python_wrapper.database_dialect import ( + PgDatabaseDialect, RdsPgDialect) + dd = AsyncPsycopgDriverDialect() + for dialect in (RdsPgDialect(), PgDatabaseDialect(), None): + provider = _build_host_list_provider( + Properties({"plugins": "host_monitoring_v2"}), dd, dialect) + assert isinstance(provider, AsyncStaticHostListProvider), dialect diff --git a/tests/unit/test_aio_host_monitoring_plugin.py b/tests/unit/test_aio_host_monitoring_plugin.py index 88aa91b3..37fc86eb 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,80 @@ def test_async_dialects_support_abort_connection(): assert AsyncPsycopgDriverDialect().supports_abort_connection() is True assert AsyncAiomysqlDriverDialect().supports_abort_connection() is True + + +# ---- Multiple event loops ---------------------------------------------- + + +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 + + +def test_failed_instance_identification_is_not_cached(): + """An unidentified cluster endpoint isn't cached and monitored later.""" + plugin, svc, _dd, _conn = _build() + cluster = HostInfo(host="mydb.cluster-xyz.us-east-1.rds.amazonaws.com", port=5432) + svc._current_host_info = cluster + svc.identify_connection = AsyncMock(return_value=None) + + async def _body(): + for _ in range(2): + with pytest.raises(AwsWrapperError): + await plugin._get_monitoring_host_info() + assert svc.identify_connection.await_count == 2 + assert plugin._monitoring_host_info is None + + asyncio.run(_body()) 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 00000000..808f97f1 --- /dev/null +++ b/tests/unit/test_aio_shared_topology_monitor.py @@ -0,0 +1,165 @@ +# 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, as in sync.""" + +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: + # As in sync, closing a connection doesn't stop the cluster's monitor. + 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: + # Like sync's stop_monitor; the next provider starts a fresh monitor. + 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())