From b51fb1dd63cb4a7f06c6be8907276ab9bf41e080 Mon Sep 17 00:00:00 2001 From: Dmitry Kropachev Date: Thu, 16 Jul 2026 15:56:55 -0400 Subject: [PATCH] Fix BatchWrite affinity and IPv6 routing URLs --- alternator/core/handlers.py | 7 +- alternator/core/key_affinity.py | 10 ++- alternator/core/live_nodes.py | 26 +++++++- tests/unit/test_async_manager.py | 26 ++++++++ tests/unit/test_key_affinity.py | 108 +++++++++++++++++++++++++++---- tests/unit/test_sync_manager.py | 26 ++++++++ 6 files changed, 181 insertions(+), 22 deletions(-) diff --git a/alternator/core/handlers.py b/alternator/core/handlers.py index 75b9a8b..500166f 100644 --- a/alternator/core/handlers.py +++ b/alternator/core/handlers.py @@ -22,6 +22,7 @@ create_header_filter_handler, create_user_agent_header_handler, ) +from alternator.core.live_nodes import _format_host_port from alternator.core.query_plan import LazyQueryPlan from alternator.core.request import extract_operation_name, extract_request_params from alternator.exceptions import NoNodesAvailableError @@ -83,20 +84,20 @@ def create_query_plan( """Create a URI iterator for a single request.""" node_addresses = nodes.nodes if preferred_node is not None and preferred_node in node_addresses: - yield f"{scheme}://{preferred_node}:{port}" + yield f"{scheme}://{_format_host_port(preferred_node, port)}" remaining_nodes = tuple( node for node in node_addresses if node != preferred_node ) seed = _stable_seed(preferred_node) plan = LazyQueryPlan(nodes=remaining_nodes, seed=seed) for node in plan: - yield f"{scheme}://{node}:{port}" + yield f"{scheme}://{_format_host_port(node, port)}" return seed = random.getrandbits(64) plan = LazyQueryPlan(nodes=node_addresses, seed=seed) for node in plan: - yield f"{scheme}://{node}:{port}" + yield f"{scheme}://{_format_host_port(node, port)}" # Register event handler to update endpoint per-request def update_endpoint( diff --git a/alternator/core/key_affinity.py b/alternator/core/key_affinity.py index 6108733..b5303fd 100644 --- a/alternator/core/key_affinity.py +++ b/alternator/core/key_affinity.py @@ -13,6 +13,7 @@ from alternator._constants import PK_DISCOVERY_TIMEOUT_SECONDS from alternator.core.hashing import hash_attribute_value +from alternator.core.query_plan import LazyQueryPlan if TYPE_CHECKING: from mypy_boto3_dynamodb import DynamoDBClient @@ -205,7 +206,6 @@ def _select_batch_write_affinity_node( get_pk_name: Callable[[str], str | None], ) -> str | None: votes: Counter[str] = Counter() - selector = AffinitySelector() for candidate in _iter_batch_write_candidates(params): pk_name = get_pk_name(candidate.table_name) @@ -226,7 +226,7 @@ def _select_batch_write_affinity_node( except (ValueError, TypeError, UnicodeEncodeError): continue - node = selector.select(nodes, hash_value) + node = _select_query_plan_first_node(nodes, hash_value) if node is not None: votes[node] += 1 @@ -240,6 +240,12 @@ def _select_batch_write_affinity_node( return winners[0] +def _select_query_plan_first_node(nodes: NodeList, hash_value: int) -> str | None: + if not nodes: + return None + return next(LazyQueryPlan(nodes=tuple(sorted(nodes.nodes)), seed=hash_value)) + + def _iter_batch_write_candidates( params: dict[str, Any], ) -> Iterable[_BatchWriteCandidate]: diff --git a/alternator/core/live_nodes.py b/alternator/core/live_nodes.py index 65de23a..74dace3 100644 --- a/alternator/core/live_nodes.py +++ b/alternator/core/live_nodes.py @@ -4,6 +4,7 @@ import asyncio import contextlib +import ipaddress import logging import threading import time @@ -137,7 +138,8 @@ def build_localnodes_url(self, scope: RoutingScope, seed: str) -> str: Returns: Full URL with query string """ - base = f"{self._config.scheme}://{seed}:{self._config.port}/localnodes" + authority = _format_host_port(seed, self._config.port) + base = f"{self._config.scheme}://{authority}/localnodes" query = scope.get_localnodes_query() return f"{base}?{query}" if query else base @@ -236,6 +238,22 @@ def _no_nodes_for_scope_error(scope: RoutingScope) -> ConfigurationError: return ConfigurationError(f"No nodes found for routing scope: {scope.description}") +def _format_host_port(host: str, port: int) -> str: + """Return a URL authority for a host-only address and shared port.""" + if host.startswith("[") and host.endswith("]"): + return f"{host}:{port}" + if _is_ipv6_literal(host): + return f"[{host}]:{port}" + return f"{host}:{port}" + + +def _is_ipv6_literal(host: str) -> bool: + try: + return ipaddress.ip_address(host).version == 6 + except ValueError: + return False + + class SyncLiveNodesManager: """ Synchronous LiveNodesManager with background refresh thread. @@ -319,7 +337,8 @@ def next_node_uri(self) -> str: scope_name=self._config.routing_scope.name, attempted_hosts=list(self._config.seed_hosts), ) - return f"{self._config.scheme}://{node}:{self._config.port}" + authority = _format_host_port(node, self._config.port) + return f"{self._config.scheme}://{authority}" def refresh_nodes(self) -> bool: """ @@ -509,7 +528,8 @@ def next_node_uri(self) -> str: scope_name=self._config.routing_scope.name, attempted_hosts=list(self._config.seed_hosts), ) - return f"{self._config.scheme}://{node}:{self._config.port}" + authority = _format_host_port(node, self._config.port) + return f"{self._config.scheme}://{authority}" async def refresh_nodes(self) -> bool: """ diff --git a/tests/unit/test_async_manager.py b/tests/unit/test_async_manager.py index 45c47ea..80bb1fa 100644 --- a/tests/unit/test_async_manager.py +++ b/tests/unit/test_async_manager.py @@ -162,6 +162,32 @@ async def empty_fetch(url: str) -> list[str]: url = manager._core.build_localnodes_url(ClusterScope(), "192.168.1.1") assert url == "http://192.168.1.1:8000/localnodes" + @pytest.mark.asyncio + async def test_url_construction_brackets_ipv6_seed(self) -> None: + """Test raw IPv6 seed hosts are bracketed in discovery URLs.""" + + async def empty_fetch(url: str) -> list[str]: + return [] + + config = Config(seed_hosts=["2001:db8::1"], port=8000) + manager = AsyncLiveNodesManager(config, empty_fetch) + + url = manager._core.build_localnodes_url(ClusterScope(), "2001:db8::1") + assert url == "http://[2001:db8::1]:8000/localnodes" + + @pytest.mark.asyncio + async def test_next_node_uri_brackets_ipv6_node(self) -> None: + """Test raw IPv6 live nodes are bracketed in operation URLs.""" + + async def mock_fetch(url: str) -> list[str]: + return ["2001:db8::2"] + + config = Config(seed_hosts=["seed"], port=8000) + manager = AsyncLiveNodesManager(config, mock_fetch) + await manager.refresh_nodes() + + assert manager.next_node_uri() == "http://[2001:db8::2]:8000" + @pytest.mark.asyncio async def test_url_construction_with_dc_scope(self, config: Config) -> None: """Test URL construction with datacenter scope.""" diff --git a/tests/unit/test_key_affinity.py b/tests/unit/test_key_affinity.py index fff7aa9..ab1bf48 100644 --- a/tests/unit/test_key_affinity.py +++ b/tests/unit/test_key_affinity.py @@ -24,6 +24,7 @@ should_use_affinity, ) from alternator.core.live_nodes import NodeList +from alternator.core.query_plan import LazyQueryPlan from alternator.core.request import extract_request_params Params = dict[str, Any] @@ -51,6 +52,19 @@ def _pk_value_for_node(nodes: NodeList, target_node: str, prefix: str) -> str: raise AssertionError(f"could not find value for node {target_node}") +def _query_plan_first_node(nodes: NodeList, hash_value: int) -> str: + return next(LazyQueryPlan(nodes=tuple(sorted(nodes.nodes)), seed=hash_value)) + + +def _pk_value_for_batch_node(nodes: NodeList, target_node: str, prefix: str) -> str: + for index in range(1000): + value = f"{prefix}-{index}" + hash_value = hash_attribute_value("S", value) + if _query_plan_first_node(nodes, hash_value) == target_node: + return value + raise AssertionError(f"could not find batch value for node {target_node}") + + class TestIsRmwOperation: """Tests for is_rmw_operation function.""" @@ -360,7 +374,7 @@ def capture_request(request: AWSPreparedRequest, **_: object) -> None: def test_batch_write_single_put_selects_node(self) -> None: """Test BatchWriteItem with a single PutRequest selects its node.""" nodes = NodeList(nodes=("a", "b", "c"), scope_name="test") - value = _pk_value_for_node(nodes, "a", "batch-put") + value = _pk_value_for_batch_node(nodes, "a", "batch-put") params = { "RequestItems": { "orders": [{"PutRequest": {"Item": {"pk": {"S": value}}}}], @@ -378,10 +392,49 @@ def test_batch_write_single_put_selects_node(self) -> None: == "a" ) + def test_batch_write_vote_uses_query_plan_first_pick(self) -> None: + """Test BatchWriteItem votes use canonical query-plan first pick.""" + nodes = NodeList( + nodes=("node1", "node2", "node3", "node4", "node5", "node6"), + scope_name="test", + ) + value = None + expected = None + modulo = None + for index in range(1000): + candidate = f"canonical-batch-{index}" + hash_value = hash_attribute_value("S", candidate) + query_plan_node = _query_plan_first_node(nodes, hash_value) + modulo_node = AffinitySelector().select(nodes, hash_value) + if query_plan_node != modulo_node: + value = candidate + expected = query_plan_node + modulo = modulo_node + break + + assert value is not None + params = { + "RequestItems": { + "orders": [{"PutRequest": {"Item": {"pk": {"S": value}}}}], + } + } + + assert ( + select_affinity_node( + mode="ANY_WRITE", + operation_name="BatchWriteItem", + params=params, + nodes=nodes, + get_pk_name={"orders": "pk"}.get, + ) + == expected + ) + assert expected != modulo + def test_batch_write_single_delete_selects_node(self) -> None: """Test BatchWriteItem with a single DeleteRequest selects its node.""" nodes = NodeList(nodes=("a", "b", "c"), scope_name="test") - value = _pk_value_for_node(nodes, "c", "batch-delete") + value = _pk_value_for_batch_node(nodes, "c", "batch-delete") params = { "RequestItems": { "orders": [{"DeleteRequest": {"Key": {"pk": {"S": value}}}}], @@ -402,9 +455,9 @@ def test_batch_write_single_delete_selects_node(self) -> None: def test_batch_write_mixed_put_delete_unique_winner(self) -> None: """Test BatchWriteItem votes for the unique preferred node.""" nodes = NodeList(nodes=("a", "b", "c"), scope_name="test") - b1 = _pk_value_for_node(nodes, "b", "b1") - b2 = _pk_value_for_node(nodes, "b", "b2") - c1 = _pk_value_for_node(nodes, "c", "c1") + b1 = _pk_value_for_batch_node(nodes, "b", "b1") + b2 = _pk_value_for_batch_node(nodes, "b", "b2") + c1 = _pk_value_for_batch_node(nodes, "c", "c1") params = { "RequestItems": { "orders": [ @@ -429,9 +482,9 @@ def test_batch_write_mixed_put_delete_unique_winner(self) -> None: def test_batch_write_multi_table_reversed_order_same_winner(self) -> None: """Test batch voting is independent of table and request order.""" nodes = NodeList(nodes=("a", "b", "c"), scope_name="test") - b1 = _pk_value_for_node(nodes, "b", "orders-b1") - b2 = _pk_value_for_node(nodes, "b", "sessions-b2") - a1 = _pk_value_for_node(nodes, "a", "orders-a1") + b1 = _pk_value_for_batch_node(nodes, "b", "orders-b1") + b2 = _pk_value_for_batch_node(nodes, "b", "sessions-b2") + a1 = _pk_value_for_batch_node(nodes, "a", "orders-a1") orders = [ {"PutRequest": {"Item": {"pk": {"S": b1}}}}, {"PutRequest": {"Item": {"pk": {"S": a1}}}}, @@ -543,8 +596,8 @@ def test_batch_write_no_nodes_falls_back(self) -> None: def test_batch_write_tied_votes_fall_back(self) -> None: """Test tied preferred-node votes produce no preferred node.""" nodes = NodeList(nodes=("a", "b", "c"), scope_name="test") - a1 = _pk_value_for_node(nodes, "a", "tie-a") - b1 = _pk_value_for_node(nodes, "b", "tie-b") + a1 = _pk_value_for_batch_node(nodes, "a", "tie-a") + b1 = _pk_value_for_batch_node(nodes, "b", "tie-b") params = { "RequestItems": { "orders": [ @@ -569,9 +622,8 @@ def test_batch_write_binary_pk_selects_stable_node(self) -> None: """Test binary partition-key values use stable hashing.""" nodes = NodeList(nodes=("a", "b", "c"), scope_name="test") binary_value = b"\x00\x01stable" - expected = AffinitySelector().select( - nodes, - hash_attribute_value("B", binary_value), + expected = _query_plan_first_node( + nodes, hash_attribute_value("B", binary_value) ) params = { "RequestItems": { @@ -595,7 +647,7 @@ def test_batch_write_binary_pk_selects_stable_node(self) -> None: def test_batch_write_selection_does_not_mutate_params(self) -> None: """Test BatchWriteItem affinity selection leaves request params unchanged.""" nodes = NodeList(nodes=("a", "b", "c"), scope_name="test") - value = _pk_value_for_node(nodes, "b", "no-mutate") + value = _pk_value_for_batch_node(nodes, "b", "no-mutate") params = { "RequestItems": { "orders": [{"PutRequest": {"Item": {"pk": {"S": value}}}}], @@ -661,6 +713,34 @@ def compute_affinity_node( assert first_url == "http://b:8000/" assert {second_url, third_url} == {"http://a:8000/", "http://c:8000/"} + def test_query_plan_brackets_ipv6_nodes(self) -> None: + """Test request handler formats raw IPv6 node addresses as URL authorities.""" + config = Config(seed_hosts=["seed"], port=8000) + manager = MagicMock() + manager.nodes = NodeList( + nodes=("2001:db8::1", "2001:db8::2"), scope_name="cluster" + ) + events = MagicMock() + + _register_alternator_handlers(events, manager, config) + handlers = { + call[0][1].__name__: call[0][1] for call in events.register.call_args_list + } + + request = MagicMock() + request.url = "http://seed:8000/" + request.headers = {"X-Amz-Target": "DynamoDB_20120810.ListTables"} + request.body = b"{}" + request._alternator_query_plan = None + + update_endpoint = handlers["update_endpoint"] + update_endpoint(request) + + assert request.url in { + "http://[2001:db8::1]:8000/", + "http://[2001:db8::2]:8000/", + } + class TestExtractPartitionKey: """Tests for extract_partition_key function.""" diff --git a/tests/unit/test_sync_manager.py b/tests/unit/test_sync_manager.py index ed21537..7e7473f 100644 --- a/tests/unit/test_sync_manager.py +++ b/tests/unit/test_sync_manager.py @@ -159,6 +159,32 @@ def mock_fetch(url: str) -> Sequence[str]: assert captured_url[0] == "http://192.168.1.1:8000/localnodes" + def test_url_construction_brackets_ipv6_seed(self) -> None: + """Test raw IPv6 seed hosts are bracketed in discovery URLs.""" + config = Config(seed_hosts=["2001:db8::1"], port=8000) + captured_url: list[str] = [] + + def mock_fetch(url: str) -> Sequence[str]: + captured_url.append(url) + return ["node1"] + + manager = SyncLiveNodesManager(config, mock_fetch) + manager.refresh_nodes() + + assert captured_url[0] == "http://[2001:db8::1]:8000/localnodes" + + def test_next_node_uri_brackets_ipv6_node(self) -> None: + """Test raw IPv6 live nodes are bracketed in operation URLs.""" + config = Config(seed_hosts=["seed"], port=8000) + + def mock_fetch(url: str) -> Sequence[str]: + return ["2001:db8::2"] + + manager = SyncLiveNodesManager(config, mock_fetch) + manager.refresh_nodes() + + assert manager.next_node_uri() == "http://[2001:db8::2]:8000" + def test_url_construction_with_dc_scope(self) -> None: """Test URL includes query string for datacenter scope.""" config = Config(