diff --git a/.vscode/cspell.json b/.vscode/cspell.json index f9c517b24066..f16a91b0be67 100644 --- a/.vscode/cspell.json +++ b/.vscode/cspell.json @@ -1678,6 +1678,7 @@ { "filename": "sdk/cosmos/azure-cosmos/**", "words": [ + "PPAF", "colls", "pkranges", "Upserts", diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index b8b1c451ad10..05648b5863fe 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -3,6 +3,7 @@ ### 4.16.2 (Unreleased) #### Features Added +* Added metadata cache cross-region hedging. When enabled, the SDK hedges the container and partition-key-range metadata cache reads (both initial population and cache refresh) to a second region if the primary region is slow, returning the first acceptable response. Controlled by the new `enable_metadata_hedging` client keyword (tri-state: `None` follows the account's PPAF state, `True` forces it on, `False` disables it). The `AZURE_COSMOS_METADATA_HEDGING_ENABLED` environment variable acts as an operator kill-switch that overrides the client keyword and PPAF state. For partition-key-range reads only the first change-feed page is hedged; later pages are pinned to the winning region to preserve continuation integrity. #### Breaking Changes diff --git a/sdk/cosmos/azure-cosmos/api.md b/sdk/cosmos/azure-cosmos/api.md index f19133480576..9268cddf2ff6 100644 --- a/sdk/cosmos/azure-cosmos/api.md +++ b/sdk/cosmos/azure-cosmos/api.md @@ -1697,6 +1697,7 @@ namespace azure.cosmos.aio availability_strategy: Union[bool, dict[str, Any]] = False, availability_strategy_max_concurrency: Optional[int] = ..., consistency_level: Optional[str] = ..., + enable_metadata_hedging: Optional[bool] = ..., **kwargs: Any ) -> None: ... @@ -3754,6 +3755,7 @@ namespace azure.cosmos.http_constants CROSS_PARTITION_QUERY_NOT_SERVABLE = 1004 DATABASE_ACCOUNT_NOT_FOUND = 1008 INSUFFICIENT_BINDABLE_PARTITIONS = 1007 + LEASE_NOT_FOUND = 1022 NAME_CACHE_IS_STALE = 1000 OWNER_RESOURCE_NOT_FOUND = 1003 PARTITION_KEY_MISMATCH = 1001 @@ -3883,6 +3885,7 @@ namespace azure.cosmos.scripts consistency_level: Optional[str] = None, availability_strategy: Union[bool, dict[str, Any]] = False, availability_strategy_executor: Optional[ThreadPoolExecutor] = None, + enable_metadata_hedging: Optional[bool] = None, **kwargs: Any ) -> None: ... @@ -4677,6 +4680,7 @@ namespace azure.cosmos.user consistency_level: Optional[str] = None, availability_strategy: Union[bool, dict[str, Any]] = False, availability_strategy_executor: Optional[ThreadPoolExecutor] = None, + enable_metadata_hedging: Optional[bool] = None, **kwargs: Any ) -> None: ... diff --git a/sdk/cosmos/azure-cosmos/api.metadata.yml b/sdk/cosmos/azure-cosmos/api.metadata.yml index ae4780d7163d..cce73c6b6d21 100644 --- a/sdk/cosmos/azure-cosmos/api.metadata.yml +++ b/sdk/cosmos/azure-cosmos/api.metadata.yml @@ -1,3 +1,3 @@ -apiMdSha256: de3b8c0f2a0405fac7152d562c982bf98bf7bb546c46f05e39164574e47f7f7a +apiMdSha256: 600817d370f2bae76b5ffda0ac83b670ef973df33840a56213a172f51f6bbcf7 parserVersion: 0.3.28 pythonVersion: 3.13.14 diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_availability_strategy_config.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_availability_strategy_config.py index d2b9ed1e2283..838ee70c5fa9 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_availability_strategy_config.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_availability_strategy_config.py @@ -21,12 +21,28 @@ """Configuration types for Azure Cosmos DB availability strategies.""" +import os from typing import Optional, Any, Union # Default values for cross-region hedging strategy DEFAULT_THRESHOLD_MS = 500 DEFAULT_THRESHOLD_STEPS_MS = 100 +# Defaults for cold-start metadata cache hedging. These are SDK-derived and not +# customer-configurable. The threshold is an aggressive tail-latency trigger: metadata +# (control-plane) reads use a short read timeout (DBAReadTimeout, 3s by default), so the +# 1.5s hedge fires well before the primary's own timeout, giving a slow region a chance +# to be raced by a second region while the primary attempt is still outstanding. The +# concurrency budget caps the number of in-flight metadata hedges per client. +DEFAULT_METADATA_HEDGING_THRESHOLD_MS = 1500 +DEFAULT_METADATA_HEDGING_THRESHOLD_STEPS_MS = 500 +DEFAULT_METADATA_HEDGING_CONCURRENCY_BUDGET = 8 + +# Operator kill-switch (parity with the .NET AZURE_COSMOS_METADATA_HEDGING_ENABLED env var). +# When set to a recognized truthy/falsy value it overrides both the customer opt-in and the +# account PPAF state. +METADATA_HEDGING_ENABLED_ENV_VAR = "AZURE_COSMOS_METADATA_HEDGING_ENABLED" + class CrossRegionHedgingStrategy: """Configuration for cross-region request hedging strategy. @@ -34,7 +50,7 @@ class CrossRegionHedgingStrategy: :param config: Dictionary containing configuration values, defaults to None :type config: Optional[Dict[str, Any]] :raises ValueError: If configuration values are invalid - + The config dictionary can contain: - threshold_ms: Time in ms before routing to alternate region (default: 500) - threshold_steps_ms: Time interval between routing attempts (default: 100) @@ -53,11 +69,71 @@ def __init__(self, config: Optional[dict[str, Any]] = None) -> None: raise ValueError("threshold_steps_ms must be positive") +class MetadataCrossRegionHedgingStrategy(CrossRegionHedgingStrategy): + """Cold-start metadata cache cross-region hedging configuration. + + Unlike :class:`CrossRegionHedgingStrategy`, the metadata hedging threshold is a + fixed, SDK-derived value and is not customer-configurable. The strategy is bounded + to the primary request plus a single cross-region hedge. + """ + + def __init__(self) -> None: + super().__init__( + { + "threshold_ms": DEFAULT_METADATA_HEDGING_THRESHOLD_MS, + "threshold_steps_ms": DEFAULT_METADATA_HEDGING_THRESHOLD_STEPS_MS, + } + ) + + +def _parse_metadata_hedging_env_override() -> Optional[bool]: + """Parse the ``AZURE_COSMOS_METADATA_HEDGING_ENABLED`` env var into a tri-state bool. + + :returns: ``True``/``False`` when the env var is set to a recognized value, else ``None``. + :rtype: Optional[bool] + """ + raw = os.environ.get(METADATA_HEDGING_ENABLED_ENV_VAR) + if raw is None: + return None + value = raw.strip().lower() + if value in ("true", "1", "yes"): + return True + if value in ("false", "0", "no"): + return False + return None + + +def resolve_metadata_hedging_opt_in(opt_in: Optional[bool], ppaf_enabled: bool) -> bool: + """Resolve the tri-state cold-start metadata hedging opt-in to a concrete bool. + + When the customer leaves the opt-in ``None`` (the default), cold-start metadata + hedging follows the account's PPAF (Per-Partition Automatic Failover) state: it is + enabled when PPAF is enabled and disabled otherwise. An explicit ``True`` enables + hedging even when PPAF is disabled, and an explicit ``False`` disables it regardless + of PPAF. The ``AZURE_COSMOS_METADATA_HEDGING_ENABLED`` environment variable, when set to a + recognized truthy/falsy value, overrides both the customer opt-in and the account PPAF + state (operator kill-switch, parity with the .NET SDK). + + :param opt_in: The customer-supplied tri-state opt-in (True/False/None). + :type opt_in: Optional[bool] + :param ppaf_enabled: Whether Per-Partition Automatic Failover is enabled for the account. + :type ppaf_enabled: bool + :returns: True if cold-start metadata hedging should be applied, False otherwise. + :rtype: bool + """ + env_override = _parse_metadata_hedging_env_override() + if env_override is not None: + return env_override + if opt_in is None: + return ppaf_enabled + return opt_in + + def _validate_request_hedging_strategy( config: Optional[Union[bool, dict[str, Any]]] ) -> Union[CrossRegionHedgingStrategy, bool, None]: """Validate and create a CrossRegionHedgingStrategy for a request. - + :param config: Configuration for availability strategy. Can be: - None: Returns None (no strategy, uses client default if available) - True: Returns strategy with default values (threshold_ms=500, threshold_steps_ms=100) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_availability_strategy_handler_base.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_availability_strategy_handler_base.py index be2b8399d6cd..95657b516eec 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_availability_strategy_handler_base.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_availability_strategy_handler_base.py @@ -21,7 +21,7 @@ """Module containing base classes and mixins for availability strategy handlers.""" -from typing import List, Optional, Union, TYPE_CHECKING +from typing import Any, List, Optional, Union, TYPE_CHECKING from . import exceptions from ._location_cache import RegionalRoutingContext @@ -69,6 +69,40 @@ def _create_excluded_regions_for_hedging( return excluded + def _record_winner( + self, + winner_sink: Optional[List[Any]], + is_primary: bool, + available_locations: List[str], + request_params: RequestObject, + ) -> None: + """Record the winning branch for PartitionKeyRange continuation pinning. + + Populates ``winner_sink[0]`` with the winning region and, when a hedge won, the + set of regions to exclude so later change-feed pages pin to the winning region + (continuation-etag integrity). No-op when ``winner_sink`` is None. Shared by the + sync and async metadata hedging handlers. + + :param winner_sink: Length-1 sink list to receive the winner descriptor, or None. + :type winner_sink: Optional[List[Any]] + :param is_primary: Whether the primary branch produced the returned outcome. + :type is_primary: bool + :param available_locations: Ordered applicable region names (index 0 = primary). + :type available_locations: List[str] + :param request_params: The originating request parameters. + :type request_params: ~azure.cosmos._request_object.RequestObject + :rtype: None + """ + if winner_sink is None: + return + winning_index = 0 if is_primary else 1 + winner_sink[0] = { + "hedge_won": not is_primary, + "winning_region": available_locations[winning_index], + "pin_excluded_locations": self._create_excluded_regions_for_hedging( + winning_index, available_locations, request_params.excluded_locations), + } + def _is_non_transient_error(self, result: BaseException) -> bool: """Check if exception represents a non-transient error. diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py index a5d01a7a12c9..6bd08f7208fd 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -62,6 +62,7 @@ from . import http_constants, exceptions from ._auth_policy import CosmosBearerTokenCredentialPolicy from ._availability_strategy_config import validate_client_hedging_strategy, CrossRegionHedgingStrategy +from ._metadata_hedging import MetadataCrossRegionHedgingHandler from ._base import _build_properties_cache from ._change_feed.change_feed_iterable import ChangeFeedIterable from ._change_feed.change_feed_state import ChangeFeedState @@ -71,7 +72,7 @@ from ._cosmos_responses import CosmosDict, CosmosList, CosmosItemPaged from ._range_partition_resolver import RangePartitionResolver from ._read_items_helper import ReadItemsHelperSync -from ._request_object import RequestObject +from ._request_object import RequestObject, METADATA_CACHE_POPULATION_OPTION from ._retry_utility import ConnectionRetryPolicy from ._routing import routing_map_provider from ._query_advisor import get_query_advice_info @@ -135,7 +136,7 @@ class _QueryCompatibilityMode: _DefaultStringHashPrecision = 3 _DefaultStringRangePrecision = -1 - def __init__( # pylint: disable=too-many-statements + def __init__( # pylint: disable=too-many-statements,too-many-locals self, url_connection: str, auth: CredentialDict, @@ -143,6 +144,7 @@ def __init__( # pylint: disable=too-many-statements consistency_level: Optional[str] = None, availability_strategy: Union[bool, dict[str, Any]] = False, availability_strategy_executor: Optional[ThreadPoolExecutor] = None, + enable_metadata_hedging: Optional[bool] = None, **kwargs: Any ) -> None: """ @@ -170,6 +172,12 @@ def __init__( # pylint: disable=too-many-statements self.availability_strategy: Union[CrossRegionHedgingStrategy, None] =\ validate_client_hedging_strategy(availability_strategy) self.availability_strategy_executor: Optional[ThreadPoolExecutor] = availability_strategy_executor + # Tri-state opt-in for metadata cache cross-region hedging. None follows the + # account's PPAF state, True forces it on, False disables it (no handler is created). + self._metadata_hedging_opt_in: Optional[bool] = enable_metadata_hedging + self._metadata_hedging_handler: Optional[MetadataCrossRegionHedgingHandler] = ( + None if enable_metadata_hedging is False else MetadataCrossRegionHedgingHandler() + ) self.master_key: Optional[str] = None self.resource_tokens: Optional[Mapping[str, Any]] = None @@ -2975,6 +2983,7 @@ def Read( headers, options.get("partitionKey", None)) request_params.set_excluded_location_from_options(options) + request_params.set_metadata_cache_population_from_options(options) request_params.set_availability_strategy(options, self.availability_strategy) request_params.availability_strategy_executor = self.availability_strategy_executor base.set_session_token_header(self, headers, path, request_params, options) @@ -3310,6 +3319,7 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]: options.get("partitionKey", None) ) request_params.set_excluded_location_from_options(options) + request_params.set_metadata_cache_population_from_options(options) request_params.set_availability_strategy(options, self.availability_strategy) request_params.availability_strategy_executor = self.availability_strategy_executor @@ -3948,7 +3958,8 @@ def refresh_routing_map_provider( def _refresh_container_properties_cache(self, container_link: str): # If container properties cache is stale, refresh it by reading the container. - container = self.ReadContainer(container_link, options=None) + container = self.ReadContainer( + container_link, options={METADATA_CACHE_POPULATION_OPTION: True}) # Only cache Container Properties that will not change in the lifetime of the container self._set_container_properties_cache(container_link, _build_properties_cache(container, container_link)) @@ -3995,7 +4006,9 @@ def _get_partition_key_definition( partition_key_definition = cached_container.get("partitionKey") # Else read the collection from backend and add it to the cache else: - container = self.ReadContainer(collection_link, options) + read_options = {**options} if options else {} + read_options[METADATA_CACHE_POPULATION_OPTION] = True + container = self.ReadContainer(collection_link, read_options) partition_key_definition = container.get("partitionKey") self._set_container_properties_cache(collection_link, _build_properties_cache(container, collection_link)) return partition_key_definition diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_metadata_hedging.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_metadata_hedging.py new file mode 100644 index 000000000000..df0c850df65d --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_metadata_hedging.py @@ -0,0 +1,441 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Cold-start metadata cache cross-region hedging for Azure Cosmos DB. + +This is the synchronous port of the .NET ``MetadataHedgingStrategy`` (PR +Azure/azure-cosmos-dotnet-v3#5999). It provides bounded cross-region hedging for +cold-start metadata cache reads (container/Collection reads and PartitionKeyRange +read-feed reads): the primary request is dispatched immediately and, if it has not +produced an acceptable response within a fixed SDK-derived threshold, a single hedge +request is dispatched to a second region. The first acceptable winner is returned. +""" + +import copy +import logging +import os +import time +from concurrent.futures import CancelledError, Future, ThreadPoolExecutor, as_completed +from threading import Event, Semaphore +from typing import Any, Callable, Dict, List, Optional, Tuple + +from azure.core.exceptions import ServiceRequestError, ServiceResponseError # pylint: disable=no-legacy-azure-core-http-response-import +from azure.core.pipeline.transport import HttpRequest # pylint: disable=no-legacy-azure-core-http-response-import + +from . import exceptions +from ._availability_strategy_config import ( + DEFAULT_METADATA_HEDGING_CONCURRENCY_BUDGET, + MetadataCrossRegionHedgingStrategy, +) +from ._availability_strategy_handler_base import AvailabilityStrategyHandlerMixin +from ._global_partition_endpoint_manager_circuit_breaker import _GlobalPartitionEndpointManagerForCircuitBreaker +from ._request_object import RequestObject +from .documents import _OperationType +from .http_constants import ResourceType, StatusCodes, SubStatusCodes + +logger = logging.getLogger("azure.cosmos.metadata_hedging") + +ResponseType = Tuple[Dict[str, Any], Dict[str, Any]] + + +class MetadataHedgeSkipReason: + """Reason a cold-start metadata hedge was not dispatched (for diagnostics/logging).""" + NONE = "None" + NOT_SUPPORTED_RESOURCE = "ResourceTypeNotSupported" + ALREADY_HEDGING = "AlreadyHedgedThisOperation" + SINGLE_REGION = "SingleRegion" + BUDGET_EXHAUSTED = "BudgetExhausted" + + +def is_supported_metadata_request(request_params: RequestObject) -> bool: + """Return True if the request is a metadata read eligible for cold-start hedging. + + Supported reads mirror the .NET design: a Collection read or a PartitionKeyRange + read-feed. PartitionKeyRange reads issued as a plain ``Read`` are also accepted + defensively. + + :param request_params: The request parameters. + :type request_params: ~azure.cosmos._request_object.RequestObject + :returns: True if the request is a supported metadata read. + :rtype: bool + """ + resource_type = request_params.resource_type + operation_type = request_params.operation_type + if resource_type == ResourceType.Collection and operation_type == _OperationType.Read: + return True + if resource_type == ResourceType.PartitionKeyRange and operation_type in ( + _OperationType.ReadFeed, + _OperationType.Read, + ): + return True + return False + + +def is_regional_failure( + status_code: Optional[int], + sub_status: Optional[int], + exception: Optional[BaseException], +) -> bool: + """Return True if the response/exception is a regional failure for metadata hedging. + + A regional failure is one that should advance a metadata read to a different region + (so it must not be accepted as a winner). This mirrors the .NET ``IsRegionalFailure`` + classification: transport-level failures and timeouts, ``503``, ``500``, + ``403`` with sub-status ``DatabaseAccountNotFound``, and ``410`` with sub-status + ``LeaseNotFound`` (a regional/backend-lease failure, not a request-definitive error). + + :param status_code: HTTP status code from the response, or None for a transport failure. + :type status_code: Optional[int] + :param sub_status: Cosmos sub-status code from the response. + :type sub_status: Optional[int] + :param exception: Exception observed instead of (or in addition to) the response, or None. + :type exception: Optional[BaseException] + :returns: True if the failure is regional. + :rtype: bool + """ + if isinstance(exception, (ServiceRequestError, ServiceResponseError)): + return True + + if status_code is None: + return False + + if status_code in (StatusCodes.SERVICE_UNAVAILABLE, StatusCodes.INTERNAL_SERVER_ERROR): + return True + if status_code == StatusCodes.FORBIDDEN and sub_status == SubStatusCodes.DATABASE_ACCOUNT_NOT_FOUND: + return True + if status_code == StatusCodes.GONE and sub_status == SubStatusCodes.LEASE_NOT_FOUND: + return True + return False + + +def _status_codes_from_exception(exception: BaseException) -> Tuple[Optional[int], Optional[int]]: + if isinstance(exception, exceptions.CosmosHttpResponseError): + return exception.status_code, exception.sub_status + return None, None + + +class _BranchOutcome: + """Classification of a settled hedging branch (parity with the .NET ``BranchOutcome``). + + Keeping the primary authoritative depends on this split: only a ``REGIONAL_FAILURE`` is + worth hedging, and a hedge can win *only* by producing a ``SUCCESS`` -- it can never + override a primary ``SUCCESS`` or ``DEFINITIVE`` outcome. + """ + SUCCESS = "Success" + REGIONAL_FAILURE = "RegionalFailure" + DEFINITIVE = "Definitive" + CANCELLED = "Cancelled" + + +def classify_branch_outcome(exception: Optional[BaseException]) -> str: + """Classify a settled (non-cancelled) branch outcome. + + A missing exception is a ``SUCCESS`` (these metadata reads throw for every status + ``>= 400``, so an exceptionless return is always a good response). A regional-failure + exception is ``REGIONAL_FAILURE`` -- the region, not the request, is at fault, so + another region is worth trying. Any other terminal error -- a non-regional definitive + error such as ``404`` / ``409`` / ``412`` or an auth failure -- is ``DEFINITIVE`` and + authoritative; a hedge must never win with it. Internal loser cancellation is handled + by the caller (which knows its own ``CancelledError`` type) and never reaches here. + + :param exception: The exception raised by the branch, or None on success. + :type exception: Optional[BaseException] + :returns: One of the :class:`_BranchOutcome` string constants. + :rtype: str + """ + if exception is None: + return _BranchOutcome.SUCCESS + status_code, sub_status = _status_codes_from_exception(exception) + if is_regional_failure(status_code, sub_status, exception): + return _BranchOutcome.REGIONAL_FAILURE + return _BranchOutcome.DEFINITIVE + + +class MetadataCrossRegionHedgingHandler(AvailabilityStrategyHandlerMixin): + """Bounded cross-region hedging handler for cold-start metadata cache reads. + + One instance per client. The per-client concurrency budget caps the number of + in-flight metadata hedges; when it is exhausted, eligible requests fall back to a + primary-only send. + + :param concurrency_budget: Max number of in-flight metadata hedges for this client. + :type concurrency_budget: int + """ + + def __init__(self, concurrency_budget: int = DEFAULT_METADATA_HEDGING_CONCURRENCY_BUDGET) -> None: + budget = max(1, concurrency_budget) + self._budget = Semaphore(budget) + # Long-lived executor shared across this client's metadata hedges. Every in-flight + # hedge needs two worker threads (primary + hedge), so size the pool to at least + # 2 * budget; otherwise hedge tasks can starve behind primaries on low-core hosts. + max_workers = max(os.cpu_count() or 1, 2 * budget) + self._executor = ThreadPoolExecutor(max_workers=max_workers) # pylint: disable=consider-using-with + + def close(self) -> None: + """Shut down the shared hedge executor, releasing its worker threads. + + Invoked from the client/connection teardown path (``CosmosClient.__exit__`` / + ``close``) so a disposed client does not leak its metadata-hedge worker threads. + Idempotent -- safe to call multiple times. + + :rtype: None + """ + self._executor.shutdown(wait=False, cancel_futures=True) + + def _classify_future(self, future: "Future") -> str: + """Classify a COMPLETED future into a :class:`_BranchOutcome`. + + :param future: A future that is guaranteed to be done. + :type future: ~concurrent.futures.Future + :returns: One of the :class:`_BranchOutcome` string constants. + :rtype: str + """ + if future.cancelled(): + return _BranchOutcome.CANCELLED + exception = future.exception() + if isinstance(exception, CancelledError): + return _BranchOutcome.CANCELLED + return classify_branch_outcome(exception) + + @staticmethod + def _wait_settled(future: "Future") -> None: + """Block until ``future`` settles, swallowing its outcome (classified separately). + + :param future: The future to wait for. + :type future: ~concurrent.futures.Future + :rtype: None + """ + try: + future.exception() + except CancelledError: + pass + + def _finish( + self, + future: "Future", + winner_sink: Optional[List[Any]], + is_primary: bool, + available_locations: List[str], + request_params: RequestObject, + completion_status: Event, + ) -> ResponseType: + """Signal completion, record the winner, and return/raise the winning branch's outcome. + + :param future: The winning (completed) future whose outcome is returned. + :type future: ~concurrent.futures.Future + :param winner_sink: Optional length-1 sink for the winner descriptor. + :type winner_sink: Optional[List[Any]] + :param is_primary: Whether the winning branch is the primary. + :type is_primary: bool + :param available_locations: Ordered applicable region names (index 0 = primary). + :type available_locations: List[str] + :param request_params: The originating request parameters. + :type request_params: ~azure.cosmos._request_object.RequestObject + :param completion_status: The shared completion event to signal the loser. + :type completion_status: ~threading.Event + :returns: The winning response tuple. + :rtype: Tuple[dict, dict] + """ + completion_status.set() + self._record_winner(winner_sink, is_primary, available_locations, request_params) + if future.cancelled(): + raise CancelledError("The request has been cancelled") + exception = future.exception() + if exception is not None: + raise exception + return future.result() # type: ignore[return-value] + + def _send_with_delay( + self, + request_params: RequestObject, + request: HttpRequest, + execute_request_fn: Callable[..., ResponseType], + location_index: int, + available_locations: List[str], + complete_status: Event, + ) -> ResponseType: + strategy = request_params.availability_strategy + if strategy is None: + raise ValueError("availability_strategy should not be null for metadata hedging") + + delay = 0 if location_index == 0 else strategy.threshold_ms + if delay > 0: + time.sleep(delay / 1000) + + params = copy.deepcopy(request_params) + params.is_hedging_request = location_index > 0 + params.completion_status = complete_status + params.excluded_locations = self._create_excluded_regions_for_hedging( + location_index, + available_locations, + request_params.excluded_locations, + ) + + req = copy.deepcopy(request) + + if complete_status.is_set(): + raise CancelledError("The request has been cancelled") + + return execute_request_fn(params, req) + + def _resolve_winner( + self, + first_future: "Future", + primary_future: "Future", + hedge_future: "Future", + ) -> Tuple["Future", bool]: + """Resolve which settled branch wins, keeping the primary authoritative. + + A hedge can win only by settling first with a ``SUCCESS``, and even then never over + a primary that has already produced a definitive (non-regional) outcome. A primary + regional failure yields to a successful hedge; otherwise the primary's outcome (its + exception rethrown by :meth:`_finish`) is authoritative. Mirrors the .NET + ``ResolveWinnerAsync``. + + :param first_future: The branch that settled first. + :type first_future: ~concurrent.futures.Future + :param primary_future: The primary branch future. + :type primary_future: ~concurrent.futures.Future + :param hedge_future: The hedge branch future. + :type hedge_future: ~concurrent.futures.Future + :returns: The winning future and whether it is the primary branch. + :rtype: Tuple[~concurrent.futures.Future, bool] + """ + if first_future is hedge_future: + if self._classify_future(hedge_future) == _BranchOutcome.SUCCESS: + # A successful hedge wins unless the primary has ALREADY settled with a + # definitive (non-regional) outcome, which the hedge must never override. + if primary_future.done() and \ + self._classify_future(primary_future) != _BranchOutcome.REGIONAL_FAILURE: + return primary_future, True + return hedge_future, False + # A non-success hedge can never win; wait for the primary's authoritative outcome. + self._wait_settled(primary_future) + return primary_future, True + + # Primary settled first. Success or a definitive error is authoritative. + if self._classify_future(primary_future) != _BranchOutcome.REGIONAL_FAILURE: + return primary_future, True + + # Primary regional failure -> a successful hedge may now win; otherwise the primary's + # (authoritative) regional failure is returned. + self._wait_settled(hedge_future) + if self._classify_future(hedge_future) == _BranchOutcome.SUCCESS: + return hedge_future, False + return primary_future, True + + def execute_request( + self, + request_params: RequestObject, + global_endpoint_manager: _GlobalPartitionEndpointManagerForCircuitBreaker, + request: HttpRequest, + execute_request_fn: Callable[..., ResponseType], + winner_sink: Optional[List[Any]] = None, + ) -> ResponseType: + """Execute a metadata read with bounded primary + single-hedge cross-region hedging. + + :param request_params: Request parameters for the metadata read. + :type request_params: ~azure.cosmos._request_object.RequestObject + :param global_endpoint_manager: Manager for endpoint routing and health tracking. + :type global_endpoint_manager: + ~azure.cosmos._GlobalPartitionEndpointManagerForCircuitBreaker + :param request: The HTTP request to be executed. + :type request: ~azure.core.pipeline.transport.HttpRequest + :param execute_request_fn: Function that executes the actual request. + :type execute_request_fn: Callable[..., Tuple[dict, dict]] + :param winner_sink: Optional length-1 list to receive the winner descriptor for + PartitionKeyRange continuation pinning (``hedge_won`` / ``winning_region`` / + ``pin_excluded_locations``). + :type winner_sink: Optional[List[Any]] + :returns: The winning response tuple. + :rtype: Tuple[dict, dict] + """ + available_locations = self._get_applicable_endpoints(request_params, global_endpoint_manager) + if len(available_locations) <= 1: + logger.debug("Metadata hedge skipped: %s", MetadataHedgeSkipReason.SINGLE_REGION) + return execute_request_fn(request_params, request) + + acquired = self._budget.acquire(blocking=False) # pylint: disable=consider-using-with + if not acquired: + logger.debug("Metadata hedge skipped: %s", MetadataHedgeSkipReason.BUDGET_EXHAUSTED) + return execute_request_fn(request_params, request) + + completion_status = Event() + try: + primary_future = self._executor.submit( + self._send_with_delay, request_params, request, execute_request_fn, + 0, available_locations, completion_status) + hedge_future = self._executor.submit( + self._send_with_delay, request_params, request, execute_request_fn, + 1, available_locations, completion_status) + + # Resolve the primary-vs-hedge race, keeping the primary authoritative: a hedge + # can win only by settling first with a SUCCESS, and even then only while the + # primary has not already produced a definitive (non-regional) answer. Mirrors + # the .NET ``ResolveWinnerAsync``. + # + # Note: no per-region failure is recorded when the hedge wins. Metadata cache + # reads are account-global, not per-partition, so the circuit-breaker health + # tracker (keyed on partition-key ranges) does not apply here; a slow primary is + # not necessarily an unhealthy one. + first_future = next(as_completed([primary_future, hedge_future])) + winner_future, is_primary = self._resolve_winner( + first_future, primary_future, hedge_future) + return self._finish( + winner_future, winner_sink, is_primary, + available_locations, request_params, completion_status) + finally: + completion_status.set() + self._budget.release() + + +def execute_metadata_hedging( + handler: MetadataCrossRegionHedgingHandler, + request_params: RequestObject, + global_endpoint_manager: _GlobalPartitionEndpointManagerForCircuitBreaker, + request: HttpRequest, + execute_request_fn: Callable[..., ResponseType], + winner_sink: Optional[List[Any]] = None, +) -> ResponseType: + """Execute a metadata read with cold-start cross-region hedging. + + :param handler: The per-client metadata hedging handler. + :type handler: ~azure.cosmos._metadata_hedging.MetadataCrossRegionHedgingHandler + :param request_params: Request parameters for the metadata read. + :type request_params: ~azure.cosmos._request_object.RequestObject + :param global_endpoint_manager: Manager for endpoint routing and health tracking. + :type global_endpoint_manager: + ~azure.cosmos._GlobalPartitionEndpointManagerForCircuitBreaker + :param request: The HTTP request to be executed. + :type request: ~azure.core.pipeline.transport.HttpRequest + :param execute_request_fn: Function that executes the actual request. + :type execute_request_fn: Callable[..., Tuple[dict, dict]] + :param winner_sink: Optional length-1 list to receive the winner descriptor + (``hedge_won`` / ``winning_region`` / ``pin_excluded_locations``) so a + PartitionKeyRange drain can pin later pages to the winning region. + :type winner_sink: Optional[List[Any]] + :returns: The winning response tuple. + :rtype: Tuple[dict, dict] + """ + if request_params.availability_strategy is None: + request_params.availability_strategy = MetadataCrossRegionHedgingStrategy() + return handler.execute_request( + request_params, global_endpoint_manager, request, execute_request_fn, winner_sink=winner_sink) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_request_object.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_request_object.py index b27b5486cdcd..63470aa9ec88 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_request_object.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_request_object.py @@ -31,6 +31,11 @@ from .documents import _OperationType from .http_constants import ResourceType +# Internal options key used to mark a request as a cold-start metadata-cache +# population read (container / partition-key-range cache). Only set by the +# cache-population callsites; gates cold-start metadata hedging. +METADATA_CACHE_POPULATION_OPTION = "metadataCachePopulation" + class RequestObject(object): # pylint: disable=too-many-instance-attributes def __init__( @@ -59,6 +64,7 @@ def __init__( self.pk_val = pk_val self.retry_write: int = 0 self.is_hedging_request: bool = False # Flag to track if this is a hedged request + self.is_metadata_cache_population: bool = False # Cold-start metadata-cache read self.completion_status: Optional[Union[threading.Event, asyncio.Event]] = None def route_to_location_with_preferred_location_flag( # pylint: disable=name-too-long @@ -97,6 +103,17 @@ def set_excluded_location_from_options(self, options: Mapping[str, Any]) -> None if self._can_set_excluded_location(options): self.excluded_locations = options['excludedLocations'] + def set_metadata_cache_population_from_options(self, options: Mapping[str, Any]) -> None: # pylint: disable=name-too-long + """Mark this request as a cold-start metadata-cache population read when the + internal ``metadataCachePopulation`` options flag is set. + + :param options: The request options that may contain the internal flag. + :type options: Mapping[str, Any] + :return: None + """ + if options and options.get(METADATA_CACHE_POPULATION_OPTION): + self.is_metadata_cache_population = True + def set_retry_write(self, request_options: Mapping[str, Any], client_retry_write: int) -> None: if self.resource_type == ResourceType.Document: if request_options and request_options.get(Constants.Kwargs.RETRY_WRITE): @@ -149,7 +166,7 @@ def set_availability_strategy( def should_cancel_request(self) -> bool: """Check if this request should be cancelled due to parallel request completion. - + :return: True if request should be cancelled, False otherwise :rtype: bool """ diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/aio/routing_map_provider.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/aio/routing_map_provider.py index 668adbae90d2..a3f2f7288e56 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/aio/routing_map_provider.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/aio/routing_map_provider.py @@ -28,6 +28,7 @@ from typing import Dict, Any, Optional, List, TYPE_CHECKING from azure.core.utils import CaseInsensitiveDict from ... import _base, http_constants +from ..._request_object import METADATA_CACHE_POPULATION_OPTION from ..collection_routing_map import CollectionRoutingMap from ...exceptions import CosmosHttpResponseError from .._routing_map_provider_common import ( @@ -408,6 +409,12 @@ async def _fetch_routing_map( ) base_headers: Dict[str, Any] = base_kwargs_for_headers['headers'] + # Cross-region metadata hedging is applied to the FIRST drain page only. + # Pages 2..N are pinned to whichever region won page 1 (via excludedLocations) + # so the change-feed continuation etag stays consistent with a single region. + is_first_drain_page = True + pin_excluded_locations: Optional[List[str]] = None + while True: request_kwargs = dict(kwargs) # Shallow-copy ``base_headers`` so the per-iter @@ -421,6 +428,19 @@ async def _fetch_routing_map( status_capture: List[Optional[int]] = [None] request_kwargs['_internal_response_status_capture'] = status_capture + # Per-page copy so the first-page-only hedging flag and later-page region + # pinning do not bleed across iterations of the shared options dict. + page_options = dict(change_feed_options) + winner_capture: List[Optional[Dict[str, Any]]] = [None] + if is_first_drain_page: + # Mark this read as cold-start metadata-cache population so the request + # layer can hedge it cross-region; other PartitionKeyRange readers + # (e.g. hybrid search) are not flagged and so are not hedged. + page_options[METADATA_CACHE_POPULATION_OPTION] = True + request_kwargs['_internal_hedge_winner_capture'] = winner_capture + elif pin_excluded_locations is not None: + page_options['excludedLocations'] = pin_excluded_locations + # Override If-None-Match with the running etag from the drain # so each page advances. ``prepare_fetch_options_and_headers`` # only sets it from ``current_previous_map.change_feed_etag`` @@ -434,7 +454,7 @@ async def _fetch_routing_map( try: pk_range_generator = self._document_client._ReadPartitionKeyRanges( collection_link, - change_feed_options, + page_options, **request_kwargs ) ranges.extend([item async for item in pk_range_generator]) @@ -444,6 +464,14 @@ async def _fetch_routing_map( collection_link, e) raise + if is_first_drain_page: + # If a hedge won page 1, pin subsequent pages to the winning region so + # the change-feed continuation etag is honored by the same region. + winner = winner_capture[0] + if winner and winner.get("hedge_won"): + pin_excluded_locations = winner.get("pin_excluded_locations") + is_first_drain_page = False + decision, new_etag, current_if_none_match, seen_any_etag = evaluate_drain_page( page_new_etag=response_headers.get(http_constants.HttpHeaders.ETag), current_if_none_match=current_if_none_match, diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_map_provider.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_map_provider.py index c2cca7bb2ec1..3dca5a2601d2 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_map_provider.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_map_provider.py @@ -28,6 +28,7 @@ from typing import Dict, Any, Optional, List, TYPE_CHECKING from azure.core.utils import CaseInsensitiveDict from .. import _base, http_constants +from .._request_object import METADATA_CACHE_POPULATION_OPTION from .collection_routing_map import CollectionRoutingMap from ..exceptions import CosmosHttpResponseError from ._routing_map_provider_common import ( @@ -375,6 +376,12 @@ def _fetch_routing_map( ) base_headers: Dict[str, Any] = base_kwargs_for_headers['headers'] + # Cross-region metadata hedging is applied to the FIRST drain page only. + # Pages 2..N are pinned to whichever region won page 1 (via excludedLocations) + # so the change-feed continuation etag stays consistent with a single region. + is_first_drain_page = True + pin_excluded_locations: Optional[List[str]] = None + while True: request_kwargs = dict(kwargs) # Shallow-copy ``base_headers`` so the per-iter @@ -388,6 +395,19 @@ def _fetch_routing_map( status_capture: List[Optional[int]] = [None] request_kwargs['_internal_response_status_capture'] = status_capture + # Per-page copy so the first-page-only hedging flag and later-page region + # pinning do not bleed across iterations of the shared options dict. + page_options = dict(change_feed_options) + winner_capture: List[Optional[Dict[str, Any]]] = [None] + if is_first_drain_page: + # Mark this read as cold-start metadata-cache population so the request + # layer can hedge it cross-region; other PartitionKeyRange readers + # (e.g. hybrid search) are not flagged and so are not hedged. + page_options[METADATA_CACHE_POPULATION_OPTION] = True + request_kwargs['_internal_hedge_winner_capture'] = winner_capture + elif pin_excluded_locations is not None: + page_options['excludedLocations'] = pin_excluded_locations + # Override If-None-Match with the running etag from the drain # so each page advances. ``prepare_fetch_options_and_headers`` # only sets it from ``current_previous_map.change_feed_etag`` @@ -402,7 +422,7 @@ def _fetch_routing_map( try: pk_range_generator = self._document_client._ReadPartitionKeyRanges( collection_link, - change_feed_options, + page_options, **request_kwargs ) page_ranges.extend(list(pk_range_generator)) @@ -414,6 +434,14 @@ def _fetch_routing_map( ranges.extend(page_ranges) + if is_first_drain_page: + # If a hedge won page 1, pin subsequent pages to the winning region so + # the change-feed continuation etag is honored by the same region. + winner = winner_capture[0] + if winner and winner.get("hedge_won"): + pin_excluded_locations = winner.get("pin_excluded_locations") + is_first_drain_page = False + decision, new_etag, current_if_none_match, seen_any_etag = evaluate_drain_page( page_new_etag=response_headers.get(http_constants.HttpHeaders.ETag), current_if_none_match=current_if_none_match, diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_synchronized_request.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_synchronized_request.py index 1a7e24e5feba..258c88b53b83 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_synchronized_request.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_synchronized_request.py @@ -30,8 +30,9 @@ from azure.core.exceptions import DecodeError # type: ignore from . import exceptions, http_constants, _retry_utility -from ._availability_strategy_config import CrossRegionHedgingStrategy +from ._availability_strategy_config import CrossRegionHedgingStrategy, resolve_metadata_hedging_opt_in from ._availability_strategy_handler import execute_with_hedging +from ._metadata_hedging import execute_metadata_hedging, is_supported_metadata_request from ._constants import _Constants from ._response_decoding import decode_response_body_for_status from ._request_object import RequestObject @@ -246,6 +247,32 @@ def _is_availability_strategy_applicable(request_params: RequestObject) -> bool: request_params.retry_write > 0)) +def _is_metadata_hedging_applicable(client, request_params: RequestObject, global_endpoint_manager) -> bool: + """Determine if cold-start metadata cache hedging should be applied to the request. + + :param client: Document client instance. + :type client: object + :param request_params: Request parameters containing operation details. + :type request_params: ~azure.cosmos._request_object.RequestObject + :param global_endpoint_manager: Manager for endpoint routing and PPAF state. + :type global_endpoint_manager: object + :returns: True if metadata hedging should be applied, False otherwise. + :rtype: bool + """ + handler = getattr(client, "_metadata_hedging_handler", None) + if handler is None: + return False + if request_params.is_hedging_request or request_params.availability_strategy is not None: + return False + if not request_params.is_metadata_cache_population: + return False + if not is_supported_metadata_request(request_params): + return False + opt_in = getattr(client, "_metadata_hedging_opt_in", None) + ppaf_enabled = global_endpoint_manager.is_per_partition_automatic_failover_enabled() + return resolve_metadata_hedging_opt_in(opt_in, ppaf_enabled) + + def _replace_url_prefix(original_url, new_prefix): parts = original_url.split('/', 3) @@ -294,6 +321,31 @@ def SynchronizedRequest( elif request.data is None: request.headers[http_constants.HttpHeaders.ContentLength] = 0 + # Sidecar (length-1 list) used by the /pkranges change-feed drain to learn which + # region won the first hedged page, so it can pin later pages to that region for + # change-feed continuation integrity. Popped here so it never reaches the pipeline. + hedge_winner_capture = kwargs.pop("_internal_hedge_winner_capture", None) + + # Cold-start metadata cache cross-region hedging (Collection / PartitionKeyRange reads). + if _is_metadata_hedging_applicable(client, request_params, global_endpoint_manager): + return execute_metadata_hedging( + client._metadata_hedging_handler, # pylint: disable=protected-access + request_params, + global_endpoint_manager, + request, + lambda req_param, r: _retry_utility.Execute( + client, + global_endpoint_manager, + _Request, + req_param, + connection_policy, + pipeline_client, + r, + **kwargs + ), + winner_sink=hedge_winner_capture, + ) + if request_params.availability_strategy is None: # if ppaf is enabled, then hedging is enabled by default if global_endpoint_manager.is_per_partition_automatic_failover_enabled(): diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_asynchronous_request.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_asynchronous_request.py index ce7d5a44536c..d1e942c822a4 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_asynchronous_request.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_asynchronous_request.py @@ -30,9 +30,11 @@ from . import _retry_utility_async from ._asynchronous_availability_strategy_handler import execute_with_availability_strategy +from ._metadata_hedging import execute_metadata_hedging as execute_metadata_hedging_async from .. import exceptions from .. import http_constants -from .._availability_strategy_config import CrossRegionHedgingStrategy +from .._availability_strategy_config import CrossRegionHedgingStrategy, resolve_metadata_hedging_opt_in +from .._metadata_hedging import is_supported_metadata_request from .._constants import _Constants from .._request_object import RequestObject from .._response_decoding import decode_response_body_for_status @@ -215,6 +217,33 @@ def _is_availability_strategy_applicable(request_params: RequestObject) -> bool: (not _OperationType.IsWriteOperation(request_params.operation_type) or request_params.retry_write > 0)) + +def _is_metadata_hedging_applicable(client, request_params: RequestObject, global_endpoint_manager) -> bool: + """Determine if cold-start metadata cache hedging should be applied to the request. + + :param client: Document client instance. + :type client: object + :param request_params: Request parameters containing operation details. + :type request_params: ~azure.cosmos._request_object.RequestObject + :param global_endpoint_manager: Manager for endpoint routing and PPAF state. + :type global_endpoint_manager: object + :returns: True if metadata hedging should be applied, False otherwise. + :rtype: bool + """ + handler = getattr(client, "_metadata_hedging_handler", None) + if handler is None: + return False + if request_params.is_hedging_request or request_params.availability_strategy is not None: + return False + if not request_params.is_metadata_cache_population: + return False + if not is_supported_metadata_request(request_params): + return False + opt_in = getattr(client, "_metadata_hedging_opt_in", None) + ppaf_enabled = global_endpoint_manager.is_per_partition_automatic_failover_enabled() + return resolve_metadata_hedging_opt_in(opt_in, ppaf_enabled) + + async def AsynchronousRequest( client, request_params, @@ -247,6 +276,31 @@ async def AsynchronousRequest( elif request.data is None: request.headers[http_constants.HttpHeaders.ContentLength] = 0 + # Sidecar (length-1 list) used by the /pkranges change-feed drain to learn which + # region won the first hedged page, so it can pin later pages to that region for + # change-feed continuation integrity. Popped here so it never reaches the pipeline. + hedge_winner_capture = kwargs.pop("_internal_hedge_winner_capture", None) + + # Cold-start metadata cache cross-region hedging (Collection / PartitionKeyRange reads). + if _is_metadata_hedging_applicable(client, request_params, global_endpoint_manager): + return await execute_metadata_hedging_async( + client._metadata_hedging_handler, # pylint: disable=protected-access + request_params, + global_endpoint_manager, + request, + lambda req_param, r: _retry_utility_async.ExecuteAsync( + client, + global_endpoint_manager, + _Request, + req_param, + connection_policy, + pipeline_client, + r, + **kwargs + ), + winner_sink=hedge_winner_capture, + ) + if request_params.availability_strategy is None: # if ppaf is enabled, then hedging is enabled by default if global_endpoint_manager.is_per_partition_automatic_failover_enabled(): diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client.py index fa109d594c31..98cf02816c5c 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client.py @@ -192,6 +192,14 @@ class CosmosClient: # pylint: disable=client-accepts-api-version-keyword Default value is False (hedging disabled). :paramtype availability_strategy: Union[bool, dict[str, Any]] :keyword int availability_strategy_max_concurrency: The max concurrency for parallel requests. + :keyword Optional[bool] enable_metadata_hedging: + Tri-state opt-in for metadata cache cross-region hedging. When ``None`` + (the default), it follows the account's Per-Partition Automatic Failover (PPAF) + state. When ``True``, the SDK hedges the container and partition-key-range metadata + cache reads (both initial population and cache refresh) across regions even when PPAF + is disabled. When ``False``, metadata hedging is suppressed regardless of PPAF. The + threshold and other tuning knobs are SDK-derived defaults and are not customer-configurable. + :paramtype enable_metadata_hedging: Optional[bool] .. admonition:: Example: @@ -212,6 +220,7 @@ def __init__( consistency_level: Optional[str] = None, availability_strategy: Union[bool, dict[str, Any]] = False, availability_strategy_max_concurrency: Optional[int] = None, + enable_metadata_hedging: Optional[bool] = None, **kwargs: Any ) -> None: """Instantiate a new CosmosClient.""" @@ -224,6 +233,7 @@ def __init__( connection_policy=connection_policy, availability_strategy=availability_strategy, availability_strategy_max_concurrency=availability_strategy_max_concurrency, + enable_metadata_hedging=enable_metadata_hedging, **kwargs ) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py index 1f9f2ca87369..40b627cf441e 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py @@ -49,6 +49,7 @@ _GlobalPartitionEndpointManagerForPerPartitionAutomaticFailoverAsync) from .. import _base as base from .._availability_strategy_config import CrossRegionHedgingStrategy, validate_client_hedging_strategy +from ._metadata_hedging import MetadataCrossRegionAsyncHedgingHandler from .._base import _build_properties_cache from .. import documents from .._change_feed.aio.change_feed_iterable import ChangeFeedIterable @@ -139,6 +140,7 @@ def __init__( # pylint: disable=too-many-statements consistency_level: Optional[str] = None, availability_strategy: Union[bool, dict[str, Any]] = False, availability_strategy_max_concurrency: Optional[int] = None, + enable_metadata_hedging: Optional[bool] = None, **kwargs: Any ) -> None: """ @@ -162,6 +164,12 @@ def __init__( # pylint: disable=too-many-statements self.availability_strategy: Union[CrossRegionHedgingStrategy, None] =\ validate_client_hedging_strategy(availability_strategy) self.availability_strategy_max_concurrency: Optional[int] = availability_strategy_max_concurrency + # Tri-state opt-in for metadata cache cross-region hedging. None follows the + # account's PPAF state, True forces it on, False disables it (no handler is created). + self._metadata_hedging_opt_in: Optional[bool] = enable_metadata_hedging + self._metadata_hedging_handler: Optional[MetadataCrossRegionAsyncHedgingHandler] = ( + None if enable_metadata_hedging is False else MetadataCrossRegionAsyncHedgingHandler() + ) self.master_key: Optional[str] = None self.resource_tokens: Optional[Mapping[str, Any]] = None self.aad_credentials: Optional[AsyncTokenCredential] = None @@ -1304,6 +1312,7 @@ async def Read( headers, options.get("partitionKey", None)) request_params.set_excluded_location_from_options(options) + request_params.set_metadata_cache_population_from_options(options) await base.set_session_token_header_async(self, headers, path, request_params, options) request_params.set_availability_strategy(options, self.availability_strategy) request_params.availability_strategy_max_concurrency = self.availability_strategy_max_concurrency @@ -3108,6 +3117,7 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]: options.get("partitionKey", None), ) request_params.set_excluded_location_from_options(options) + request_params.set_metadata_cache_population_from_options(options) request_params.set_availability_strategy(options, self.availability_strategy) request_params.availability_strategy_max_concurrency = self.availability_strategy_max_concurrency headers = base.GetHeaders(self, initial_headers, "get", path, id_, resource_type, @@ -3811,7 +3821,8 @@ async def refresh_routing_map_provider( async def _refresh_container_properties_cache(self, container_link: str): # If container properties cache is stale, refresh it by reading the container. - container = await self.ReadContainer(container_link, options=None) + container = await self.ReadContainer( + container_link, options={_request_object.METADATA_CACHE_POPULATION_OPTION: True}) # Only cache Container Properties that will not change in the lifetime of the container self._set_container_properties_cache(container_link, _build_properties_cache(container, container_link)) @@ -3908,7 +3919,9 @@ async def _get_partition_key_definition( partition_key_definition = cached_container.get("partitionKey") # Else read the collection from backend and add it to the cache else: - container = await self.ReadContainer(collection_link, options) + read_options = {**options} if options else {} + read_options[_request_object.METADATA_CACHE_POPULATION_OPTION] = True + container = await self.ReadContainer(collection_link, read_options) partition_key_definition = container.get("partitionKey") self._set_container_properties_cache(collection_link, _build_properties_cache(container, collection_link)) return partition_key_definition diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_metadata_hedging.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_metadata_hedging.py new file mode 100644 index 000000000000..d32f0dc56b00 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_metadata_hedging.py @@ -0,0 +1,306 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Asynchronous cold-start metadata cache cross-region hedging for Azure Cosmos DB. + +Async port of the .NET ``MetadataHedgingStrategy`` (PR +Azure/azure-cosmos-dotnet-v3#5999). See :mod:`azure.cosmos._metadata_hedging` for the +synchronous counterpart and design notes. +""" + +import asyncio # pylint: disable=do-not-import-asyncio +import copy +import logging +from asyncio import CancelledError, Event, Task # pylint: disable=do-not-import-asyncio +from typing import Any, Awaitable, Callable, Dict, List, Optional, Tuple + +from azure.core.pipeline.transport import HttpRequest # pylint: disable=no-legacy-azure-core-http-response-import + +from .._availability_strategy_config import ( + DEFAULT_METADATA_HEDGING_CONCURRENCY_BUDGET, + MetadataCrossRegionHedgingStrategy, +) +from .._availability_strategy_handler_base import AvailabilityStrategyHandlerMixin +from .._metadata_hedging import _BranchOutcome, classify_branch_outcome +from .._request_object import RequestObject +from ._global_partition_endpoint_manager_circuit_breaker_async import \ + _GlobalPartitionEndpointManagerForCircuitBreakerAsync + +logger = logging.getLogger("azure.cosmos.metadata_hedging") + +ResponseType = Tuple[Dict[str, Any], Dict[str, Any]] + + +class MetadataCrossRegionAsyncHedgingHandler(AvailabilityStrategyHandlerMixin): + """Bounded async cross-region hedging handler for cold-start metadata cache reads. + + One instance per client. The per-client concurrency budget caps the number of + in-flight metadata hedges; when it is exhausted, eligible requests fall back to a + primary-only send. + + :param concurrency_budget: Max number of in-flight metadata hedges for this client. + :type concurrency_budget: int + """ + + def __init__(self, concurrency_budget: int = DEFAULT_METADATA_HEDGING_CONCURRENCY_BUDGET) -> None: + self._budget = asyncio.Semaphore(max(1, concurrency_budget)) + + def _classify_task(self, task: Task) -> str: + """Classify a COMPLETED task into a :class:`_BranchOutcome`. + + :param task: A task that is guaranteed to be done. + :type task: ~asyncio.Task + :returns: One of the :class:`_BranchOutcome` string constants. + :rtype: str + """ + if task.cancelled(): + return _BranchOutcome.CANCELLED + exception = task.exception() + if isinstance(exception, CancelledError): + return _BranchOutcome.CANCELLED + return classify_branch_outcome(exception) + + @staticmethod + async def _wait_settled(task: Task) -> None: + """Wait until ``task`` settles, without propagating its outcome (classified separately). + + :param task: The task to wait for. + :type task: ~asyncio.Task + :rtype: None + """ + await asyncio.wait({task}) + + def _finish( + self, + task: Task, + winner_sink: Optional[List[Any]], + is_primary: bool, + available_locations: List[str], + request_params: RequestObject, + ) -> ResponseType: + """Record the winner and return/raise the winning branch's outcome. + + :param task: The winning (completed) task whose outcome is returned. + :type task: ~asyncio.Task + :param winner_sink: Optional length-1 sink for the winner descriptor. + :type winner_sink: Optional[List[Any]] + :param is_primary: Whether the winning branch is the primary. + :type is_primary: bool + :param available_locations: Ordered applicable region names (index 0 = primary). + :type available_locations: List[str] + :param request_params: The originating request parameters. + :type request_params: ~azure.cosmos._request_object.RequestObject + :returns: The winning response tuple. + :rtype: Tuple[dict, dict] + """ + self._record_winner(winner_sink, is_primary, available_locations, request_params) + if task.cancelled(): + raise CancelledError("The request has been cancelled") + exception = task.exception() + if exception is not None: + raise exception + return task.result() # type: ignore[return-value] + + async def _send_with_delay( + self, + request_params: RequestObject, + request: HttpRequest, + execute_request_fn: Callable[..., Awaitable[ResponseType]], + location_index: int, + available_locations: List[str], + complete_status: Event, + ) -> ResponseType: + strategy = request_params.availability_strategy + if strategy is None: + raise ValueError("availability_strategy should not be null for metadata hedging") + + delay = 0 if location_index == 0 else strategy.threshold_ms + if delay > 0: + await asyncio.sleep(delay / 1000) + + params = copy.deepcopy(request_params) + params.is_hedging_request = location_index > 0 + params.completion_status = complete_status + params.excluded_locations = self._create_excluded_regions_for_hedging( + location_index, + available_locations, + request_params.excluded_locations, + ) + + req = copy.deepcopy(request) + + if complete_status.is_set(): + raise CancelledError("The request has been cancelled") + + return await execute_request_fn(params, req) + + async def _resolve_winner( + self, + done: set, + primary_task: Task, + hedge_task: Task, + ) -> Tuple[Task, bool]: + """Resolve which settled branch wins, keeping the primary authoritative. + + A hedge can win only by settling first with a ``SUCCESS``, and even then never over + a primary that has already produced a definitive (non-regional) outcome. A primary + regional failure yields to a successful hedge; otherwise the primary's outcome (its + exception rethrown by :meth:`_finish`) is authoritative. Mirrors the .NET + ``ResolveWinnerAsync``. + + :param done: The set of tasks completed by the first ``asyncio.wait``. + :type done: set + :param primary_task: The primary branch task. + :type primary_task: ~asyncio.Task + :param hedge_task: The hedge branch task. + :type hedge_task: ~asyncio.Task + :returns: The winning task and whether it is the primary branch. + :rtype: Tuple[~asyncio.Task, bool] + """ + # If both settled together, treat the primary as first so it stays authoritative. + if primary_task in done: + if self._classify_task(primary_task) != _BranchOutcome.REGIONAL_FAILURE: + # Success or a definitive error -> authoritative; never overridden. + return primary_task, True + # Primary regional failure -> a successful hedge may now win; otherwise the + # primary's (authoritative) regional failure is returned. + await self._wait_settled(hedge_task) + if self._classify_task(hedge_task) == _BranchOutcome.SUCCESS: + return hedge_task, False + return primary_task, True + + # Hedge settled first. + if self._classify_task(hedge_task) == _BranchOutcome.SUCCESS: + # A successful hedge wins unless the primary has ALREADY settled with a + # definitive (non-regional) outcome, which the hedge must never override. + if primary_task.done() and \ + self._classify_task(primary_task) != _BranchOutcome.REGIONAL_FAILURE: + return primary_task, True + return hedge_task, False + # A non-success hedge can never win; wait for the primary's authoritative outcome. + await self._wait_settled(primary_task) + return primary_task, True + + async def execute_request( + self, + request_params: RequestObject, + global_endpoint_manager: _GlobalPartitionEndpointManagerForCircuitBreakerAsync, + request: HttpRequest, + execute_request_fn: Callable[..., Awaitable[ResponseType]], + winner_sink: Optional[List[Any]] = None, + ) -> ResponseType: + """Execute a metadata read with bounded primary + single-hedge cross-region hedging. + + :param request_params: Request parameters for the metadata read. + :type request_params: ~azure.cosmos._request_object.RequestObject + :param global_endpoint_manager: Manager for endpoint routing and health tracking. + :type global_endpoint_manager: + ~azure.cosmos.aio._GlobalPartitionEndpointManagerForCircuitBreakerAsync + :param request: The HTTP request to be executed. + :type request: ~azure.core.pipeline.transport.HttpRequest + :param execute_request_fn: Async function that executes the actual request. + :type execute_request_fn: Callable[..., Awaitable[Tuple[dict, dict]]] + :param winner_sink: Optional length-1 list to receive the winner descriptor for + PartitionKeyRange continuation pinning (``hedge_won`` / ``winning_region`` / + ``pin_excluded_locations``). + :type winner_sink: Optional[List[Any]] + :returns: The winning response tuple. + :rtype: Tuple[dict, dict] + """ + available_locations = self._get_applicable_endpoints(request_params, global_endpoint_manager) + if len(available_locations) <= 1: + logger.debug("Metadata hedge skipped: SingleRegion") + return await execute_request_fn(request_params, request) + + # Non-blocking budget check: locked() is True only when the budget is fully + # exhausted (value == 0). There is no await between this check and acquire(), + # so in single-threaded asyncio no other coroutine can consume the slot in + # between and acquire() completes immediately without blocking. + if self._budget.locked(): + logger.debug("Metadata hedge skipped: BudgetExhausted") + return await execute_request_fn(request_params, request) + + await self._budget.acquire() + completion_status = Event() + active_tasks: List[Task] = [] + try: + primary_task = asyncio.create_task(self._send_with_delay( + request_params, request, execute_request_fn, + 0, available_locations, completion_status)) + hedge_task = asyncio.create_task(self._send_with_delay( + request_params, request, execute_request_fn, + 1, available_locations, completion_status)) + active_tasks = [primary_task, hedge_task] + + # Resolve the primary-vs-hedge race, keeping the primary authoritative: a hedge + # can win only by settling first with a SUCCESS, and even then only while the + # primary has not already produced a definitive (non-regional) answer. Mirrors + # the .NET ``ResolveWinnerAsync``. + # + # Note: no per-region failure is recorded when the hedge wins. Metadata cache + # reads are account-global, not per-partition, so the circuit-breaker health + # tracker (keyed on partition-key ranges) does not apply here. + done, _pending = await asyncio.wait( + {primary_task, hedge_task}, return_when=asyncio.FIRST_COMPLETED) + winner_task, is_primary = await self._resolve_winner(done, primary_task, hedge_task) + return self._finish( + winner_task, winner_sink, is_primary, available_locations, request_params) + finally: + completion_status.set() + for task in active_tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*active_tasks, return_exceptions=True) + self._budget.release() + + +async def execute_metadata_hedging( + handler: MetadataCrossRegionAsyncHedgingHandler, + request_params: RequestObject, + global_endpoint_manager: _GlobalPartitionEndpointManagerForCircuitBreakerAsync, + request: HttpRequest, + execute_request_fn: Callable[..., Awaitable[ResponseType]], + winner_sink: Optional[List[Any]] = None, +) -> ResponseType: + """Execute a metadata read with cold-start cross-region hedging. + + :param handler: The per-client metadata hedging handler. + :type handler: ~azure.cosmos.aio._metadata_hedging.MetadataCrossRegionAsyncHedgingHandler + :param request_params: Request parameters for the metadata read. + :type request_params: ~azure.cosmos._request_object.RequestObject + :param global_endpoint_manager: Manager for endpoint routing and health tracking. + :type global_endpoint_manager: + ~azure.cosmos.aio._GlobalPartitionEndpointManagerForCircuitBreakerAsync + :param request: The HTTP request to be executed. + :type request: ~azure.core.pipeline.transport.HttpRequest + :param execute_request_fn: Async function that executes the actual request. + :type execute_request_fn: Callable[..., Awaitable[Tuple[dict, dict]]] + :param winner_sink: Optional length-1 list to receive the winner descriptor + (``hedge_won`` / ``winning_region`` / ``pin_excluded_locations``) so a + PartitionKeyRange drain can pin later pages to the winning region. + :type winner_sink: Optional[List[Any]] + :returns: The winning response tuple. + :rtype: Tuple[dict, dict] + """ + if request_params.availability_strategy is None: + request_params.availability_strategy = MetadataCrossRegionHedgingStrategy() + return await handler.execute_request( + request_params, global_endpoint_manager, request, execute_request_fn, winner_sink=winner_sink) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py b/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py index 360bdca53a63..e377474dca05 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/cosmos_client.py @@ -215,6 +215,14 @@ class CosmosClient: # pylint: disable=client-accepts-api-version-keyword :paramtype availability_strategy: Union[bool, dict[str, Any]] :keyword ~concurrent.futures.thread.ThreadPoolExecutor availability_strategy_executor: Optional ThreadPoolExecutor for handling concurrent operations. + :keyword Optional[bool] enable_metadata_hedging: + Tri-state opt-in for metadata cache cross-region hedging. When ``None`` + (the default), it follows the account's Per-Partition Automatic Failover (PPAF) + state. When ``True``, the SDK hedges the container and partition-key-range metadata + cache reads (both initial population and cache refresh) across regions even when PPAF + is disabled. When ``False``, metadata hedging is suppressed regardless of PPAF. The + threshold and other tuning knobs are SDK-derived defaults and are not customer-configurable. + :paramtype enable_metadata_hedging: Optional[bool] .. admonition:: Example: @@ -245,6 +253,7 @@ def __init__( connection_policy=connection_policy, availability_strategy=kwargs.pop("availability_strategy", False), availability_strategy_executor=kwargs.pop("availability_strategy_executor", None), + enable_metadata_hedging=kwargs.pop("enable_metadata_hedging", None), **kwargs ) @@ -263,6 +272,12 @@ def __exit__(self, *args): self.client_connection._routing_map_provider.release() # pylint: disable=protected-access except Exception: # pylint: disable=broad-except pass + try: + handler = self.client_connection._metadata_hedging_handler # pylint: disable=protected-access + if handler is not None: + handler.close() + except Exception: # pylint: disable=broad-except + pass def close(self) -> None: """Close this instance of CosmosClient. diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/http_constants.py b/sdk/cosmos/azure-cosmos/azure/cosmos/http_constants.py index 4bfd7a574524..660265dd78a7 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/http_constants.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/http_constants.py @@ -430,6 +430,7 @@ class SubStatusCodes: PARTITION_KEY_RANGE_GONE = 1002 COMPLETING_SPLIT = 1007 COMPLETING_PARTITION_MIGRATION = 1008 + LEASE_NOT_FOUND = 1022 # 403: Forbidden Substatus. WRITE_FORBIDDEN = 3 diff --git a/sdk/cosmos/azure-cosmos/tests/test_metadata_hedging.py b/sdk/cosmos/azure-cosmos/tests/test_metadata_hedging.py new file mode 100644 index 000000000000..6fcf9b14262c --- /dev/null +++ b/sdk/cosmos/azure-cosmos/tests/test_metadata_hedging.py @@ -0,0 +1,290 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +"""Unit tests for cold-start metadata cache cross-region hedging (sync).""" + +import threading +import time +import unittest + +from azure.core.exceptions import ServiceRequestError +from azure.core.pipeline.transport import HttpRequest + +from azure.cosmos._availability_strategy_config import ( + MetadataCrossRegionHedgingStrategy, + resolve_metadata_hedging_opt_in, +) +from azure.cosmos._metadata_hedging import ( + MetadataCrossRegionHedgingHandler, + execute_metadata_hedging, + is_regional_failure, + is_supported_metadata_request, +) +from azure.cosmos._request_object import RequestObject +from azure.cosmos.documents import _OperationType +from azure.cosmos.exceptions import CosmosHttpResponseError +from azure.cosmos.http_constants import ResourceType, StatusCodes, SubStatusCodes + + +class _FakeContext: + def __init__(self, endpoint): + self._endpoint = endpoint + + def get_primary(self): + return self._endpoint + + +class _FakeGlobalEndpointManager: + def __init__(self, regions, ppaf=False): + self._contexts = [_FakeContext(r) for r in regions] + self._ppaf = ppaf + self.recorded_failures = [] + + def get_applicable_read_regional_routing_contexts(self, request): # noqa: ARG002 + return self._contexts + + def get_applicable_write_regional_routing_contexts(self, request): # noqa: ARG002 + return self._contexts + + def get_region_name(self, endpoint, is_write): # noqa: ARG002 + return endpoint + + def is_per_partition_automatic_failover_enabled(self): + return self._ppaf + + def record_failure(self, request_params): + self.recorded_failures.append(request_params) + + +def _metadata_request(): + req = RequestObject(ResourceType.Collection, _OperationType.Read, {}) + req.availability_strategy = MetadataCrossRegionHedgingStrategy() + # Shorten the threshold so tests don't wait the full 1.5s. + req.availability_strategy.threshold_ms = 150 + return req + + +def _http_request(): + return HttpRequest("GET", "https://primary.documents.azure.com/") + + +class TestMetadataHedgingHelpers(unittest.TestCase): + def test_resolve_opt_in_tri_state(self): + self.assertTrue(resolve_metadata_hedging_opt_in(None, True)) + self.assertFalse(resolve_metadata_hedging_opt_in(None, False)) + self.assertTrue(resolve_metadata_hedging_opt_in(True, False)) + self.assertFalse(resolve_metadata_hedging_opt_in(False, True)) + + def test_is_regional_failure(self): + self.assertTrue(is_regional_failure(StatusCodes.SERVICE_UNAVAILABLE, None, None)) + self.assertTrue(is_regional_failure(StatusCodes.INTERNAL_SERVER_ERROR, None, None)) + self.assertTrue(is_regional_failure( + StatusCodes.FORBIDDEN, SubStatusCodes.DATABASE_ACCOUNT_NOT_FOUND, None)) + self.assertTrue(is_regional_failure(None, None, ServiceRequestError(message="boom"))) + self.assertFalse(is_regional_failure(StatusCodes.NOT_FOUND, None, None)) + self.assertFalse(is_regional_failure(StatusCodes.FORBIDDEN, None, None)) + self.assertFalse(is_regional_failure(None, None, None)) + + def test_is_supported_metadata_request(self): + self.assertTrue(is_supported_metadata_request( + RequestObject(ResourceType.Collection, _OperationType.Read, {}))) + self.assertTrue(is_supported_metadata_request( + RequestObject(ResourceType.PartitionKeyRange, _OperationType.ReadFeed, {}))) + self.assertTrue(is_supported_metadata_request( + RequestObject(ResourceType.PartitionKeyRange, _OperationType.Read, {}))) + self.assertFalse(is_supported_metadata_request( + RequestObject(ResourceType.Document, _OperationType.Read, {}))) + self.assertFalse(is_supported_metadata_request( + RequestObject(ResourceType.Collection, _OperationType.Create, {}))) + + def test_set_metadata_cache_population_from_options(self): + from azure.cosmos._request_object import METADATA_CACHE_POPULATION_OPTION + req = RequestObject(ResourceType.Collection, _OperationType.Read, {}) + self.assertFalse(req.is_metadata_cache_population) + req.set_metadata_cache_population_from_options({METADATA_CACHE_POPULATION_OPTION: True}) + self.assertTrue(req.is_metadata_cache_population) + # Absent or false flag leaves it unset. + req2 = RequestObject(ResourceType.Collection, _OperationType.Read, {}) + req2.set_metadata_cache_population_from_options({}) + self.assertFalse(req2.is_metadata_cache_population) + req2.set_metadata_cache_population_from_options(None) + self.assertFalse(req2.is_metadata_cache_population) + + +class TestMetadataHedgingHandler(unittest.TestCase): + def setUp(self): + self.handler = MetadataCrossRegionHedgingHandler(concurrency_budget=8) + self.gem = _FakeGlobalEndpointManager(["region-1", "region-2"]) + + def test_primary_wins_fast_no_hedge(self): + calls = [] + + def execute_fn(params, _req): + calls.append(params.is_hedging_request) + if params.is_hedging_request: + # Hedge should not win; sleep beyond test horizon. + time.sleep(5) + return ({"source": "hedge"}, {}) + return ({"source": "primary"}, {}) + + result, _ = self.handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn) + self.assertEqual(result["source"], "primary") + self.assertEqual(self.gem.recorded_failures, []) + + def test_hedge_wins_when_primary_slow(self): + def execute_fn(params, _req): + if params.is_hedging_request: + return ({"source": "hedge"}, {}) + time.sleep(5) + return ({"source": "primary"}, {}) + + result, _ = self.handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn) + self.assertEqual(result["source"], "hedge") + + def test_hedge_auth_reject_not_accepted(self): + def execute_fn(params, _req): + if params.is_hedging_request: + raise CosmosHttpResponseError(status_code=StatusCodes.FORBIDDEN, message="auth") + time.sleep(0.3) + return ({"source": "primary"}, {}) + + result, _ = self.handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn) + # The hedge's 401/403 must never win; primary's success is returned. + self.assertEqual(result["source"], "primary") + + def test_regional_failure_hedge_lets_primary_decide(self): + def execute_fn(params, _req): + if params.is_hedging_request: + raise CosmosHttpResponseError( + status_code=StatusCodes.SERVICE_UNAVAILABLE, message="503") + time.sleep(0.3) + return ({"source": "primary"}, {}) + + result, _ = self.handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn) + self.assertEqual(result["source"], "primary") + + def test_single_region_primary_only(self): + gem = _FakeGlobalEndpointManager(["region-1"]) + count = {"n": 0} + + def execute_fn(params, _req): # noqa: ARG001 + count["n"] += 1 + return ({"source": "primary"}, {}) + + result, _ = self.handler.execute_request( + _metadata_request(), gem, _http_request(), execute_fn) + self.assertEqual(result["source"], "primary") + self.assertEqual(count["n"], 1) + + def test_budget_exhausted_primary_only(self): + handler = MetadataCrossRegionHedgingHandler(concurrency_budget=1) + # Exhaust the budget so the next request must fall back to primary-only. + self.assertTrue(handler._budget.acquire(blocking=False)) # pylint: disable=protected-access + try: + count = {"n": 0} + + def execute_fn(params, _req): # noqa: ARG001 + count["n"] += 1 + return ({"source": "primary"}, {}) + + result, _ = handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn) + self.assertEqual(result["source"], "primary") + self.assertEqual(count["n"], 1) + finally: + handler._budget.release() # pylint: disable=protected-access + + def test_both_branches_fail_prefers_primary(self): + def execute_fn(params, _req): + if params.is_hedging_request: + raise CosmosHttpResponseError( + status_code=StatusCodes.SERVICE_UNAVAILABLE, message="hedge-503") + raise CosmosHttpResponseError( + status_code=StatusCodes.SERVICE_UNAVAILABLE, message="primary-503") + + with self.assertRaises(CosmosHttpResponseError) as ctx: + self.handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn) + self.assertIn("primary-503", str(ctx.exception)) + + def test_execute_metadata_hedging_sets_strategy(self): + req = RequestObject(ResourceType.Collection, _OperationType.Read, {}) + self.assertIsNone(req.availability_strategy) + + def execute_fn(params, _req): # noqa: ARG001 + return ({"source": "primary"}, {}) + + gem = _FakeGlobalEndpointManager(["region-1"]) + execute_metadata_hedging(self.handler, req, gem, _http_request(), execute_fn) + self.assertIsInstance(req.availability_strategy, MetadataCrossRegionHedgingStrategy) + + def test_close_shuts_down_executor_idempotently(self): + handler = MetadataCrossRegionHedgingHandler(concurrency_budget=2) + handler.close() + self.assertTrue(handler._executor._shutdown) # pylint: disable=protected-access + # Idempotent: a second close must not raise. + handler.close() + + +class _FakeClient: + def __init__(self, handler, opt_in): + self._metadata_hedging_handler = handler + self._metadata_hedging_opt_in = opt_in + + +class TestMetadataHedgingApplicability(unittest.TestCase): + def setUp(self): + from azure.cosmos._synchronized_request import _is_metadata_hedging_applicable + self._is_applicable = _is_metadata_hedging_applicable + self.handler = MetadataCrossRegionHedgingHandler() + + def _request(self, resource_type=ResourceType.Collection, operation=_OperationType.Read, + cache_population=True): + req = RequestObject(resource_type, operation, {}) + req.is_metadata_cache_population = cache_population + return req + + def test_applicable_when_opt_in_true(self): + client = _FakeClient(self.handler, True) + gem = _FakeGlobalEndpointManager(["r1", "r2"], ppaf=False) + self.assertTrue(self._is_applicable(client, self._request(), gem)) + + def test_not_applicable_without_cache_population_flag(self): + client = _FakeClient(self.handler, True) + gem = _FakeGlobalEndpointManager(["r1", "r2"], ppaf=True) + # A supported metadata read that is NOT a cache-population read (e.g. a public + # container.read()) must not be hedged. + self.assertFalse(self._is_applicable(client, self._request(cache_population=False), gem)) + + def test_follows_ppaf_when_opt_in_none(self): + client = _FakeClient(self.handler, None) + self.assertTrue(self._is_applicable( + client, self._request(), _FakeGlobalEndpointManager(["r1", "r2"], ppaf=True))) + self.assertFalse(self._is_applicable( + client, self._request(), _FakeGlobalEndpointManager(["r1", "r2"], ppaf=False))) + + def test_not_applicable_without_handler(self): + client = _FakeClient(None, True) + gem = _FakeGlobalEndpointManager(["r1", "r2"], ppaf=True) + self.assertFalse(self._is_applicable(client, self._request(), gem)) + + def test_not_applicable_for_unsupported_resource(self): + client = _FakeClient(self.handler, True) + gem = _FakeGlobalEndpointManager(["r1", "r2"], ppaf=True) + req = self._request(resource_type=ResourceType.Document) + self.assertFalse(self._is_applicable(client, req, gem)) + + def test_not_applicable_for_hedging_request(self): + client = _FakeClient(self.handler, True) + gem = _FakeGlobalEndpointManager(["r1", "r2"], ppaf=True) + req = self._request() + req.is_hedging_request = True + self.assertFalse(self._is_applicable(client, req, gem)) + + +if __name__ == "__main__": + unittest.main() diff --git a/sdk/cosmos/azure-cosmos/tests/test_metadata_hedging_async.py b/sdk/cosmos/azure-cosmos/tests/test_metadata_hedging_async.py new file mode 100644 index 000000000000..4c9b8cd0c81b --- /dev/null +++ b/sdk/cosmos/azure-cosmos/tests/test_metadata_hedging_async.py @@ -0,0 +1,179 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +"""Unit tests for cold-start metadata cache cross-region hedging (async).""" + +import asyncio +import unittest + +from azure.core.pipeline.transport import HttpRequest + +from azure.cosmos._availability_strategy_config import MetadataCrossRegionHedgingStrategy +from azure.cosmos._request_object import RequestObject +from azure.cosmos.aio._metadata_hedging import ( + MetadataCrossRegionAsyncHedgingHandler, + execute_metadata_hedging, +) +from azure.cosmos.documents import _OperationType +from azure.cosmos.exceptions import CosmosHttpResponseError +from azure.cosmos.http_constants import ResourceType, StatusCodes + + +class _FakeContext: + def __init__(self, endpoint): + self._endpoint = endpoint + + def get_primary(self): + return self._endpoint + + +class _FakeGlobalEndpointManagerAsync: + def __init__(self, regions, ppaf=False): + self._contexts = [_FakeContext(r) for r in regions] + self._ppaf = ppaf + self.recorded_failures = [] + + def get_applicable_read_regional_routing_contexts(self, request): # noqa: ARG002 + return self._contexts + + def get_applicable_write_regional_routing_contexts(self, request): # noqa: ARG002 + return self._contexts + + def get_region_name(self, endpoint, is_write): # noqa: ARG002 + return endpoint + + def is_per_partition_automatic_failover_enabled(self): + return self._ppaf + + async def record_failure(self, request_params): + self.recorded_failures.append(request_params) + + +def _metadata_request(): + req = RequestObject(ResourceType.Collection, _OperationType.Read, {}) + req.availability_strategy = MetadataCrossRegionHedgingStrategy() + req.availability_strategy.threshold_ms = 150 + return req + + +def _http_request(): + return HttpRequest("GET", "https://primary.documents.azure.com/") + + +class TestMetadataHedgingHandlerAsync(unittest.TestCase): + def setUp(self): + self.handler = MetadataCrossRegionAsyncHedgingHandler(concurrency_budget=8) + self.gem = _FakeGlobalEndpointManagerAsync(["region-1", "region-2"]) + + def test_primary_wins_fast_no_hedge(self): + async def run(): + async def execute_fn(params, _req): + if params.is_hedging_request: + await asyncio.sleep(5) + return ({"source": "hedge"}, {}) + return ({"source": "primary"}, {}) + + result, _ = await self.handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn) + self.assertEqual(result["source"], "primary") + self.assertEqual(self.gem.recorded_failures, []) + + asyncio.run(run()) + + def test_hedge_wins_when_primary_slow(self): + async def run(): + async def execute_fn(params, _req): + if params.is_hedging_request: + return ({"source": "hedge"}, {}) + await asyncio.sleep(5) + return ({"source": "primary"}, {}) + + result, _ = await self.handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn) + self.assertEqual(result["source"], "hedge") + + asyncio.run(run()) + + def test_hedge_auth_reject_not_accepted(self): + async def run(): + async def execute_fn(params, _req): + if params.is_hedging_request: + raise CosmosHttpResponseError(status_code=StatusCodes.FORBIDDEN, message="auth") + await asyncio.sleep(0.3) + return ({"source": "primary"}, {}) + + result, _ = await self.handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn) + self.assertEqual(result["source"], "primary") + + asyncio.run(run()) + + def test_single_region_primary_only(self): + async def run(): + gem = _FakeGlobalEndpointManagerAsync(["region-1"]) + count = {"n": 0} + + async def execute_fn(params, _req): # noqa: ARG001 + count["n"] += 1 + return ({"source": "primary"}, {}) + + result, _ = await self.handler.execute_request( + _metadata_request(), gem, _http_request(), execute_fn) + self.assertEqual(result["source"], "primary") + self.assertEqual(count["n"], 1) + + asyncio.run(run()) + + def test_budget_exhausted_primary_only(self): + async def run(): + handler = MetadataCrossRegionAsyncHedgingHandler(concurrency_budget=1) + await handler._budget.acquire() # pylint: disable=protected-access + try: + count = {"n": 0} + + async def execute_fn(params, _req): # noqa: ARG001 + count["n"] += 1 + return ({"source": "primary"}, {}) + + result, _ = await handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn) + self.assertEqual(result["source"], "primary") + self.assertEqual(count["n"], 1) + finally: + handler._budget.release() # pylint: disable=protected-access + + asyncio.run(run()) + + def test_both_branches_fail_prefers_primary(self): + async def run(): + async def execute_fn(params, _req): + if params.is_hedging_request: + raise CosmosHttpResponseError( + status_code=StatusCodes.SERVICE_UNAVAILABLE, message="hedge-503") + raise CosmosHttpResponseError( + status_code=StatusCodes.SERVICE_UNAVAILABLE, message="primary-503") + + with self.assertRaises(CosmosHttpResponseError) as ctx: + await self.handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn) + self.assertIn("primary-503", str(ctx.exception)) + + asyncio.run(run()) + + def test_execute_metadata_hedging_sets_strategy(self): + async def run(): + req = RequestObject(ResourceType.Collection, _OperationType.Read, {}) + self.assertIsNone(req.availability_strategy) + + async def execute_fn(params, _req): # noqa: ARG001 + return ({"source": "primary"}, {}) + + gem = _FakeGlobalEndpointManagerAsync(["region-1"]) + await execute_metadata_hedging(self.handler, req, gem, _http_request(), execute_fn) + self.assertIsInstance(req.availability_strategy, MetadataCrossRegionHedgingStrategy) + + asyncio.run(run()) + + +if __name__ == "__main__": + unittest.main() diff --git a/sdk/cosmos/azure-cosmos/tests/test_metadata_hedging_gaps.py b/sdk/cosmos/azure-cosmos/tests/test_metadata_hedging_gaps.py new file mode 100644 index 000000000000..9586b391e153 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/tests/test_metadata_hedging_gaps.py @@ -0,0 +1,251 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +"""Regression tests for the metadata-hedging parity gaps closed as part of the +.NET PR #5999 port: 410+LeaseNotFound classification (SE-002), the threshold-safety +invariant (SE-003), the AZURE_COSMOS_METADATA_HEDGING_ENABLED env kill-switch (SE-005), +and first-page-only continuation pinning (SE-001 / SE-006).""" + +import os +import time +import unittest + +from azure.core.pipeline.transport import HttpRequest + +from azure.cosmos._availability_strategy_config import ( + DEFAULT_METADATA_HEDGING_THRESHOLD_MS, + METADATA_HEDGING_ENABLED_ENV_VAR, + MetadataCrossRegionHedgingStrategy, + resolve_metadata_hedging_opt_in, +) +from azure.cosmos._metadata_hedging import ( + MetadataCrossRegionHedgingHandler, + is_regional_failure, +) +from azure.cosmos._request_object import RequestObject +from azure.cosmos.documents import ConnectionPolicy, _OperationType +from azure.cosmos.exceptions import CosmosHttpResponseError +from azure.cosmos.http_constants import ResourceType, StatusCodes, SubStatusCodes + + +class _FakeContext: + def __init__(self, endpoint): + self._endpoint = endpoint + + def get_primary(self): + return self._endpoint + + +class _FakeGlobalEndpointManager: + def __init__(self, regions): + self._contexts = [_FakeContext(r) for r in regions] + + def get_applicable_read_regional_routing_contexts(self, request): # noqa: ARG002 + return self._contexts + + def get_applicable_write_regional_routing_contexts(self, request): # noqa: ARG002 + return self._contexts + + def get_region_name(self, endpoint, is_write): # noqa: ARG002 + return endpoint + + +def _metadata_request(): + req = RequestObject(ResourceType.Collection, _OperationType.Read, {}) + req.availability_strategy = MetadataCrossRegionHedgingStrategy() + req.availability_strategy.threshold_ms = 150 + return req + + +def _http_request(): + return HttpRequest("GET", "https://primary.documents.azure.com/") + + +class TestRegionalFailureLeaseNotFound(unittest.TestCase): + """SE-002: 410 + LeaseNotFound must be classified as a regional failure.""" + + def test_gone_lease_not_found_is_regional(self): + self.assertTrue( + is_regional_failure(StatusCodes.GONE, SubStatusCodes.LEASE_NOT_FOUND, None)) + + def test_gone_other_substatus_is_not_regional(self): + # A 410 with a different (definitive) sub-status is NOT regional. + self.assertFalse( + is_regional_failure(StatusCodes.GONE, SubStatusCodes.PARTITION_KEY_RANGE_GONE, None)) + self.assertFalse(is_regional_failure(StatusCodes.GONE, SubStatusCodes.UNKNOWN, None)) + + +class TestThresholdInvariant(unittest.TestCase): + """SE-003: threshold must sit strictly below the control-plane read timeout.""" + + def test_threshold_below_dba_read_timeout(self): + dba_read_timeout_ms = ConnectionPolicy().DBAReadTimeout * 1000 + self.assertGreater(DEFAULT_METADATA_HEDGING_THRESHOLD_MS, 0) + self.assertLess(DEFAULT_METADATA_HEDGING_THRESHOLD_MS, dba_read_timeout_ms) + + +class TestEnvKillSwitch(unittest.TestCase): + """SE-005: AZURE_COSMOS_METADATA_HEDGING_ENABLED overrides opt-in and PPAF.""" + + def setUp(self): + self._saved = os.environ.pop(METADATA_HEDGING_ENABLED_ENV_VAR, None) + + def tearDown(self): + if self._saved is None: + os.environ.pop(METADATA_HEDGING_ENABLED_ENV_VAR, None) + else: + os.environ[METADATA_HEDGING_ENABLED_ENV_VAR] = self._saved + + def test_env_false_disables_even_with_opt_in_true(self): + os.environ[METADATA_HEDGING_ENABLED_ENV_VAR] = "false" + self.assertFalse(resolve_metadata_hedging_opt_in(True, True)) + + def test_env_true_enables_even_with_opt_in_false(self): + os.environ[METADATA_HEDGING_ENABLED_ENV_VAR] = "true" + self.assertTrue(resolve_metadata_hedging_opt_in(False, False)) + + def test_unrecognized_env_is_ignored(self): + os.environ[METADATA_HEDGING_ENABLED_ENV_VAR] = "not-a-bool" + self.assertFalse(resolve_metadata_hedging_opt_in(None, False)) + self.assertTrue(resolve_metadata_hedging_opt_in(None, True)) + + +class TestWinnerPinning(unittest.TestCase): + """SE-001 / SE-006: the winner sink drives first-page continuation pinning.""" + + def setUp(self): + self.handler = MetadataCrossRegionHedgingHandler(concurrency_budget=8) + self.gem = _FakeGlobalEndpointManager(["region-1", "region-2"]) + + def test_hedge_win_pins_to_hedge_region(self): + sink = [None] + + def execute_fn(params, _req): + if params.is_hedging_request: + return ({"source": "hedge"}, {}) + time.sleep(5) + return ({"source": "primary"}, {}) + + result, _ = self.handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn, winner_sink=sink) + self.assertEqual(result["source"], "hedge") + self.assertIsNotNone(sink[0]) + self.assertTrue(sink[0]["hedge_won"]) + self.assertEqual(sink[0]["winning_region"], "region-2") + # Excluding region-1 pins the remaining pages to the winning region-2. + self.assertEqual(sink[0]["pin_excluded_locations"], ["region-1"]) + + def test_primary_win_records_no_pin(self): + sink = [None] + + def execute_fn(params, _req): + if params.is_hedging_request: + time.sleep(5) + return ({"source": "hedge"}, {}) + return ({"source": "primary"}, {}) + + result, _ = self.handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn, winner_sink=sink) + self.assertEqual(result["source"], "primary") + self.assertIsNotNone(sink[0]) + self.assertFalse(sink[0]["hedge_won"]) + self.assertEqual(sink[0]["winning_region"], "region-1") + # Primary won -> no additional regions excluded (drain stays on primary). + self.assertEqual(sink[0]["pin_excluded_locations"], []) + + def test_none_sink_is_noop(self): + def execute_fn(params, _req): + if params.is_hedging_request: + return ({"source": "hedge"}, {}) + time.sleep(5) + return ({"source": "primary"}, {}) + + # Default winner_sink=None must not raise. + result, _ = self.handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn) + self.assertEqual(result["source"], "hedge") + + +class TestHedgePrimaryAuthoritative(unittest.TestCase): + """SE-004: the primary is authoritative. A hedge may win ONLY by settling first with a + 2xx success, and even then never over a primary that has already produced a definitive + (non-regional) answer. A hedge definitive error (e.g. 404) must never win. This is the + core .NET ``ResolveWinnerAsync`` invariant; the previous symmetric acceptance let a + hedge win with a non-2xx definitive response.""" + + def setUp(self): + self.handler = MetadataCrossRegionHedgingHandler(concurrency_budget=8) + self.gem = _FakeGlobalEndpointManager(["region-1", "region-2"]) + + def test_hedge_definitive_error_does_not_win_over_slow_primary_success(self): + # Hedge settles first with a definitive 404; the primary is slower but succeeds. + # The hedge must NOT win -- the primary's success is authoritative. + def execute_fn(params, _req): + if params.is_hedging_request: + raise CosmosHttpResponseError(status_code=StatusCodes.NOT_FOUND, message="hedge-404") + time.sleep(0.4) + return ({"source": "primary"}, {}) + + result, _ = self.handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn) + self.assertEqual(result["source"], "primary") + + def test_hedge_definitive_error_defers_to_primary_definitive(self): + # Hedge 404 settles first; the primary is slower and returns a DIFFERENT definitive + # error (409). The primary's authoritative outcome is surfaced, not the hedge's. + def execute_fn(params, _req): + if params.is_hedging_request: + raise CosmosHttpResponseError(status_code=StatusCodes.NOT_FOUND, message="hedge-404") + time.sleep(0.4) + raise CosmosHttpResponseError(status_code=StatusCodes.CONFLICT, message="primary-409") + + with self.assertRaises(CosmosHttpResponseError) as ctx: + self.handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn) + self.assertEqual(ctx.exception.status_code, StatusCodes.CONFLICT) + self.assertIn("primary-409", str(ctx.exception)) + + def test_primary_regional_failure_hedge_definitive_returns_primary(self): + # Primary region is at fault (503) and the hedge returns a definitive 404. Only a + # SUCCESSFUL hedge may win, so the primary's authoritative regional failure is + # returned (the retry policy then classifies the real failure). + def execute_fn(params, _req): + if params.is_hedging_request: + raise CosmosHttpResponseError(status_code=StatusCodes.NOT_FOUND, message="hedge-404") + raise CosmosHttpResponseError( + status_code=StatusCodes.SERVICE_UNAVAILABLE, message="primary-503") + + with self.assertRaises(CosmosHttpResponseError) as ctx: + self.handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn) + self.assertEqual(ctx.exception.status_code, StatusCodes.SERVICE_UNAVAILABLE) + self.assertIn("primary-503", str(ctx.exception)) + + def test_primary_regional_failure_hedge_success_hedge_wins(self): + # Primary regional failure (503) + a successful hedge -> the hedge legitimately wins. + def execute_fn(params, _req): + if params.is_hedging_request: + return ({"source": "hedge"}, {}) + raise CosmosHttpResponseError( + status_code=StatusCodes.SERVICE_UNAVAILABLE, message="primary-503") + + result, _ = self.handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn) + self.assertEqual(result["source"], "hedge") + + def test_primary_definitive_wins_over_pending_hedge_success(self): + # The primary returns a definitive 404 immediately; a hedge success would arrive + # later. The definitive primary is authoritative and is returned without override. + def execute_fn(params, _req): + if params.is_hedging_request: + return ({"source": "hedge"}, {}) + raise CosmosHttpResponseError(status_code=StatusCodes.NOT_FOUND, message="primary-404") + + with self.assertRaises(CosmosHttpResponseError) as ctx: + self.handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn) + self.assertEqual(ctx.exception.status_code, StatusCodes.NOT_FOUND) + + +if __name__ == "__main__": + unittest.main() diff --git a/sdk/cosmos/azure-cosmos/tests/test_metadata_hedging_gaps_async.py b/sdk/cosmos/azure-cosmos/tests/test_metadata_hedging_gaps_async.py new file mode 100644 index 000000000000..97569c1db339 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/tests/test_metadata_hedging_gaps_async.py @@ -0,0 +1,167 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +"""Regression tests for the async metadata-hedging first-page continuation pinning +(SE-001 / SE-006) closed as part of the .NET PR #5999 port.""" + +import asyncio +import unittest + +from azure.core.pipeline.transport import HttpRequest + +from azure.cosmos._availability_strategy_config import MetadataCrossRegionHedgingStrategy +from azure.cosmos._request_object import RequestObject +from azure.cosmos.aio._metadata_hedging import MetadataCrossRegionAsyncHedgingHandler +from azure.cosmos.documents import _OperationType +from azure.cosmos.exceptions import CosmosHttpResponseError +from azure.cosmos.http_constants import ResourceType, StatusCodes + + +class _FakeContext: + def __init__(self, endpoint): + self._endpoint = endpoint + + def get_primary(self): + return self._endpoint + + +class _FakeGlobalEndpointManagerAsync: + def __init__(self, regions): + self._contexts = [_FakeContext(r) for r in regions] + + def get_applicable_read_regional_routing_contexts(self, request): # noqa: ARG002 + return self._contexts + + def get_applicable_write_regional_routing_contexts(self, request): # noqa: ARG002 + return self._contexts + + def get_region_name(self, endpoint, is_write): # noqa: ARG002 + return endpoint + + +def _metadata_request(): + req = RequestObject(ResourceType.Collection, _OperationType.Read, {}) + req.availability_strategy = MetadataCrossRegionHedgingStrategy() + req.availability_strategy.threshold_ms = 150 + return req + + +def _http_request(): + return HttpRequest("GET", "https://primary.documents.azure.com/") + + +class TestWinnerPinningAsync(unittest.TestCase): + def setUp(self): + self.handler = MetadataCrossRegionAsyncHedgingHandler(concurrency_budget=8) + self.gem = _FakeGlobalEndpointManagerAsync(["region-1", "region-2"]) + + def test_hedge_win_pins_to_hedge_region(self): + async def run(): + sink = [None] + + async def execute_fn(params, _req): + if params.is_hedging_request: + return ({"source": "hedge"}, {}) + await asyncio.sleep(5) + return ({"source": "primary"}, {}) + + result, _ = await self.handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn, winner_sink=sink) + self.assertEqual(result["source"], "hedge") + self.assertIsNotNone(sink[0]) + self.assertTrue(sink[0]["hedge_won"]) + self.assertEqual(sink[0]["winning_region"], "region-2") + self.assertEqual(sink[0]["pin_excluded_locations"], ["region-1"]) + + asyncio.run(run()) + + def test_primary_win_records_no_pin(self): + async def run(): + sink = [None] + + async def execute_fn(params, _req): + if params.is_hedging_request: + await asyncio.sleep(5) + return ({"source": "hedge"}, {}) + return ({"source": "primary"}, {}) + + result, _ = await self.handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn, winner_sink=sink) + self.assertEqual(result["source"], "primary") + self.assertIsNotNone(sink[0]) + self.assertFalse(sink[0]["hedge_won"]) + self.assertEqual(sink[0]["pin_excluded_locations"], []) + + asyncio.run(run()) + + def test_none_sink_is_noop(self): + async def run(): + async def execute_fn(params, _req): + if params.is_hedging_request: + return ({"source": "hedge"}, {}) + await asyncio.sleep(5) + return ({"source": "primary"}, {}) + + result, _ = await self.handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn) + self.assertEqual(result["source"], "hedge") + + asyncio.run(run()) + + +class TestHedgePrimaryAuthoritativeAsync(unittest.TestCase): + """SE-004 (async): the primary is authoritative. A hedge may win ONLY by settling first + with a 2xx success; a hedge definitive error (e.g. 404) must never win over the + primary. Mirrors the .NET ``ResolveWinnerAsync`` invariant.""" + + def setUp(self): + self.handler = MetadataCrossRegionAsyncHedgingHandler(concurrency_budget=8) + self.gem = _FakeGlobalEndpointManagerAsync(["region-1", "region-2"]) + + def test_hedge_definitive_error_does_not_win_over_slow_primary_success(self): + async def run(): + async def execute_fn(params, _req): + if params.is_hedging_request: + raise CosmosHttpResponseError(status_code=StatusCodes.NOT_FOUND, message="hedge-404") + await asyncio.sleep(0.4) + return ({"source": "primary"}, {}) + + result, _ = await self.handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn) + self.assertEqual(result["source"], "primary") + + asyncio.run(run()) + + def test_primary_regional_failure_hedge_definitive_returns_primary(self): + async def run(): + async def execute_fn(params, _req): + if params.is_hedging_request: + raise CosmosHttpResponseError(status_code=StatusCodes.NOT_FOUND, message="hedge-404") + raise CosmosHttpResponseError( + status_code=StatusCodes.SERVICE_UNAVAILABLE, message="primary-503") + + with self.assertRaises(CosmosHttpResponseError) as ctx: + await self.handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn) + self.assertEqual(ctx.exception.status_code, StatusCodes.SERVICE_UNAVAILABLE) + self.assertIn("primary-503", str(ctx.exception)) + + asyncio.run(run()) + + def test_primary_regional_failure_hedge_success_hedge_wins(self): + async def run(): + async def execute_fn(params, _req): + if params.is_hedging_request: + return ({"source": "hedge"}, {}) + raise CosmosHttpResponseError( + status_code=StatusCodes.SERVICE_UNAVAILABLE, message="primary-503") + + result, _ = await self.handler.execute_request( + _metadata_request(), self.gem, _http_request(), execute_fn) + self.assertEqual(result["source"], "hedge") + + asyncio.run(run()) + + +if __name__ == "__main__": + unittest.main()