Skip to content

refactor: identify hosts by host id and host-and-port instead of aliases #1297

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
karenc-bq wants to merge 1 commit into from
+1,982 −1,179
Open

refactor: identify hosts by host id and host-and-port instead of aliases #1297

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .flake8
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
[flake8]
max-line-length = 150
extend-select = TC, TC1
exclude = venv/
exclude = venv/, .venv/
226 changes: 79 additions & 147 deletions aws_advanced_python_wrapper/aio/aurora_connection_tracker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,11 +16,11 @@

Port of sync ``aurora_connection_tracker_plugin.py``. Tracks every
connection opened through the plugin pipeline in a class-level
:class:`AsyncOpenedConnectionTracker` registry keyed by the RDS
instance-endpoint alias (mirroring sync normalization). On writer
change (detected via ``AsyncPluginService.all_hosts`` or a raised
:class:`FailoverError`), closes every tracked connection to the old
writer so pooled consumers don't reuse a stale demoted-writer socket.
:class:`AsyncOpenedConnectionTracker` registry keyed by host id and by
host and port. On writer change (detected via
``AsyncPluginService.all_hosts`` or a raised :class:`FailoverError`),
closes every tracked connection to the old writer so pooled consumers
don't reuse a stale demoted-writer socket.

Simpler than the sync implementation:

Expand All @@ -47,7 +47,7 @@
import threading
from time import perf_counter_ns
from typing import (TYPE_CHECKING, Any, Awaitable, Callable, ClassVar, Dict,
Optional, Set)
Optional, Set, Tuple)
from weakref import WeakSet

from aws_advanced_python_wrapper.aio.cleanup import register_shutdown_hook
Expand All @@ -58,6 +58,7 @@
from aws_advanced_python_wrapper.utils.log import Logger
from aws_advanced_python_wrapper.utils.notifications import HostEvent
from aws_advanced_python_wrapper.utils.rds_utils import RdsUtils
from aws_advanced_python_wrapper.utils.utils import Utils

if TYPE_CHECKING:
from aws_advanced_python_wrapper.aio.driver_dialect.base import \
Expand All @@ -70,24 +71,17 @@


class AsyncOpenedConnectionTracker:
"""Alias-keyed registry of open connections, shared across plugin instances.
"""Registry of open connections, shared across plugin instances.

Keys by the RDS **instance** endpoint alias when one is present in
``host_info.as_aliases()`` -- matching sync's
:class:`OpenedConnectionTracker` at
``aurora_connection_tracker_plugin.py:47``. Falls back to
``host_info.as_alias()`` when no RDS instance alias is
discoverable (non-Aurora or custom endpoint connections).
An instance endpoint names the instance already, so its connections are
tracked under one key. Any other endpoint may be a custom domain name,
which is tracked both under the host id and under the endpoint the caller
used, so a later notification naming either one finds the connections.

This lets a cluster-endpoint-connected session still be
invalidated when a later ``CONVERTED_TO_READER`` notification
carries the instance alias.

Class-level ``_tracked`` mirrors sync ``OpenedConnectionTracker`` at
``aurora_connection_tracker_plugin.py:47`` so multiple pooled
:class:`AsyncAuroraConnectionTrackerPlugin` instances can coordinate:
one plugin's :meth:`invalidate_all` closes every peer connection to
the same host.
``_tracked`` is class-level, as sync's registry is, so multiple pooled
:class:`AsyncAuroraConnectionTrackerPlugin` instances coordinate: one
plugin's :meth:`invalidate_all` closes every peer connection to the same
host.

Still simpler than sync:

Expand All @@ -101,74 +95,49 @@ def __init__(self) -> None:
# Instance-level nothing -- all state lives on the class dict.
pass

@staticmethod
def _canonical_key(host_info: HostInfo) -> str:
"""Return the instance-endpoint alias if available, else host:port.
_rds_utils: ClassVar[RdsUtils] = RdsUtils()

Mirrors sync's multi-alias normalization in
``populate_opened_connection_set``/``invalidate_all_connections``/
``remove_connection_tracking`` at
``aurora_connection_tracker_plugin.py:116-137``.
"""
rds_utils = RdsUtils()
if rds_utils.is_rds_instance(host_info.host):
return host_info.as_alias()
for alias in host_info.as_aliases():
if rds_utils.is_rds_instance(rds_utils.remove_port(alias)):
return alias
return host_info.as_alias()

@staticmethod
def _all_keys(host_info: HostInfo) -> Set[str]:
"""Return every alias worth tracking ``host_info`` under.

Sync's ``OpenedConnectionTracker`` tracks connections keyed on
every alias in ``host_info.as_aliases()`` so a later invalidation
via any alias closes the same underlying connection set. We
mirror that here: track under each alias plus the canonical
RDS-instance endpoint.
"""
keys: Set[str] = set()
canonical = AsyncOpenedConnectionTracker._canonical_key(host_info)
if canonical:
keys.add(canonical)
try:
for alias in host_info.as_aliases():
if alias:
keys.add(alias)
except Exception: # noqa: BLE001 - as_aliases() is best-effort
pass
return keys
@classmethod
def _tracking_keys(cls, host_info: HostInfo) -> Tuple[str, ...]:
"""The keys ``host_info``'s connections are tracked under."""
if cls._rds_utils.is_rds_instance(host_info.host):
return (host_info.get_host_and_port(),)
return tuple(
key for key in (host_info.host_id, host_info.get_host_and_port()) if key)

def track(self, host_info: HostInfo, conn: Any) -> None:
for key in self._all_keys(host_info):
for key in self._tracking_keys(host_info):
conn_set = AsyncOpenedConnectionTracker._tracked.setdefault(
key, WeakSet())
conn_set.add(conn)

def remove(self, host_info: HostInfo, conn: Any) -> None:
for key in self._all_keys(host_info):
for key in self._tracking_keys(host_info):
conn_set = AsyncOpenedConnectionTracker._tracked.get(key)
if conn_set is not None:
conn_set.discard(conn)

async def invalidate_all(self, host_info: HostInfo) -> None:
"""Close every tracked connection reachable via any alias of ``host_info``.

Best-effort: individual close failures are swallowed so one
broken connection doesn't block closing the rest. Fixes the
prior canonical-key-only semantics that could miss invalidations
when the caller held a non-canonical alias.
"""
# Gather the full connection set from every alias; close each
# connection exactly once, then purge the alias entries.
"""Close every connection tracked for ``host_info``."""
logger.debug(
"OpenedConnectionTracker.InvalidatingConnections", host_info.url)
await self.invalidate_all_for_keys(
host_info.get_host_and_port(), host_info.host, host_info.host_id)

async def invalidate_all_for_keys(self, *keys: Optional[str]) -> None:
"""Close every connection tracked under any of ``keys``.

Best-effort: individual close failures are swallowed so one broken
connection does not block closing the rest.
"""
seen_conns: Set[int] = set()
to_close: list = []
keys = self._all_keys(host_info)
for key in keys:
conn_set = AsyncOpenedConnectionTracker._tracked.get(key)
present = tuple(key for key in keys if key)
# Drop each key before the first await. Collecting everything first
# and popping afterwards would discard a connection a concurrent
# track() added in between, leaving it open and untracked.
for key in present:
conn_set = AsyncOpenedConnectionTracker._tracked.pop(key, None)
if conn_set is None:
continue
for c in list(conn_set):
Expand Down Expand Up @@ -205,18 +174,14 @@ async def invalidate_all(self, host_info: HostInfo) -> None:
# is_closed=False to its owner.
logger.debug(
f"[AsyncOpenedConnectionTracker] best-effort close of a "
f"tracked connection to '{host_info.url}' failed: {e}")

for key in keys:
try:
del AsyncOpenedConnectionTracker._tracked[key]
except KeyError:
pass
f"tracked connection failed: {e}")

def _tracked_for(self, host_info: HostInfo) -> WeakSet[Any]:
"""Test-only accessor using the canonical key derivation."""
key = self._canonical_key(host_info)
return AsyncOpenedConnectionTracker._tracked.get(key, WeakSet())
"""Test-only accessor using the first tracking key."""
keys = self._tracking_keys(host_info)
if not keys:
return WeakSet()
return AsyncOpenedConnectionTracker._tracked.get(keys[0], WeakSet())

def log_opened_connections(self) -> None:
"""Debug-log the registry contents (sync parity:
Expand Down Expand Up @@ -295,11 +260,34 @@ async def connect(
connect_func: Callable[..., Awaitable[Any]]) -> Any:
conn = await connect_func()
if conn is not None:
await self._fill_instance_alias(host_info, conn)
self._tracker.track(host_info, conn)
await self._pin_current_writer(host_info, conn)
connection_host_info = await self._identify_connection_host(host_info, conn)
self._tracker.track(connection_host_info, conn)
await self._pin_current_writer(connection_host_info, conn)
return conn

async def _identify_connection_host(
self, host_info: HostInfo, conn: Any) -> HostInfo:
"""The host this connection actually landed on.

A cluster, custom-domain or IP endpoint does not name the instance the
connection reached, so ask the database and record the answer as the
routed host. On failure the endpoint the caller used is kept.
"""
routed = self._plugin_service.routed_host_info
connection_host_info = routed if routed is not None else host_info
url_type = AsyncOpenedConnectionTracker._rds_utils.identify_rds_type(
connection_host_info.host)
if not Utils.requires_host_identification(url_type):
return connection_host_info
try:
identified = await self._plugin_service.identify_connection(conn)
except Exception: # noqa: BLE001 - identification is best-effort
return connection_host_info
if identified is None:
return connection_host_info
self._plugin_service.routed_host_info = identified
return identified

async def _pin_current_writer(self, host_info: HostInfo, conn: Any) -> None:
"""Pin the current writer at connect time so a later failover is
detected as a *transition*.
Expand Down Expand Up @@ -349,39 +337,6 @@ async def _pin_current_writer(self, host_info: HostInfo, conn: Any) -> None:
except Exception: # noqa: BLE001 - writer pinning is best-effort
pass

async def _fill_instance_alias(
self, host_info: HostInfo, conn: Any) -> None:
"""Add the connection's resolved instance endpoint as an alias.

A cluster-endpoint connection lands on whichever instance the
cluster DNS currently points at, but ``host_info.host`` stays the
cluster endpoint -- so the tracker would key it only under the
cluster alias. On writer failover the change is detected against
the *instance* host (``_current_writer_from_hosts`` returns a
topology row), so ``invalidate_all(writer_instance)`` would never
match the cluster-keyed entry and the stale idle connection would
survive the demotion (``test_writer_failover_in_idle_connections``).

Mirrors sync ``AuroraConnectionTrackerPlugin.connect`` which calls
``plugin_service.fill_aliases`` for ``is_rds_cluster`` hosts. Only
cluster-endpoint connections pay the extra identify round-trip;
instance-endpoint connections already key correctly. Best-effort:
on any failure we fall back to cluster-key-only tracking (the
prior behavior), so a transient identify failure never blocks the
connect.
"""
try:
if not RdsUtils().identify_rds_type(host_info.host).is_rds_cluster:
return
identified = await self._plugin_service.identify_connection(conn)
if identified is None:
return
host_info.reset_aliases()
host_info.add_alias(host_info.as_alias())
host_info.add_alias(*identified.as_aliases())
except Exception: # noqa: BLE001 - alias enrichment is best-effort
pass

async def execute(
self,
target: object,
Expand Down Expand Up @@ -543,14 +498,14 @@ def _current_writer_from_hosts(self) -> Optional[HostInfo]:
def _same_host(a: HostInfo, b: HostInfo) -> bool:
return a.host == b.host and a.port == b.port

def _spawn_invalidation(self, old_writer: HostInfo) -> None:
def _spawn_invalidation(self, key: str) -> None:
"""Fire-and-forget invalidation, ONLY for the sync-signature
:meth:`notify_host_list_changed` path (cannot await there). The
execute/failover path awaits :meth:`_invalidate_writer_change`
inline instead -- a spawned task dies un-awaited if the event loop
closes before it runs. Pending tasks from this path are drained by
:meth:`_shutdown` on ``release_resources_async``."""
task = asyncio.create_task(self._tracker.invalidate_all(old_writer))
task = asyncio.create_task(self._tracker.invalidate_all_for_keys(key))
self._pending_invalidations.add(task)
task.add_done_callback(self._pending_invalidations.discard)

Expand Down Expand Up @@ -579,38 +534,15 @@ def notify_host_list_changed(
:attr:`_current_writer` so the next ``execute`` re-detects the
new writer via :meth:`_update_writer_from_topology`.

With the tracker's canonical-key normalization (instance alias
over cluster alias), a notify carrying the instance alias for a
cluster-connected session now correctly finds and invalidates
the tracked entries.
"""
for alias, events in changes.items():
for host, events in changes.items():
if HostEvent.CONVERTED_TO_READER in events:
host_part, port_part = self._parse_alias(alias)
h = HostInfo(host=host_part, port=port_part)
self._spawn_invalidation(h)
# Change keys are urls; tracking keys are not.
self._spawn_invalidation(host.rstrip("/"))
if HostEvent.CONVERTED_TO_WRITER in events:
# Force writer recheck on next execute via
# _update_writer_from_topology.
self._current_writer = None

@staticmethod
def _parse_alias(alias: str) -> tuple:
"""Split a ``host:port`` alias into ``(host, port)``.

Uses ``rpartition(':')`` so IPv6 aliases of the form
``[::1]:5432`` split correctly on the LAST colon. Falls back to
port 5432 when the alias has no colon (or the port part isn't
numeric).
"""
host_part, sep, port_part = alias.rpartition(":")
if not sep:
# No colon -> treat entire alias as host, port defaults.
return alias, 5432
try:
return host_part, int(port_part)
except ValueError:
return host_part, 5432


__all__ = ["AsyncAuroraConnectionTrackerPlugin", "AsyncOpenedConnectionTracker"]
Loading
Loading