DatabaseAccountRetryPolicy Improvements (#42525)

* DatabaseAccountRetry - initial commit

* Update sdk/cosmos/azure-cosmos/tests/test_retry_policy_async.py

Co-authored-by: Copilot <[email protected]>

* DatabaseAccountRetry - adding comments

* fixing pylink comments

* updated CHANGELOG.md

* fixing bug for write region unavailability

* fix: comments

* fix: comments

* fix: comments

* fix: change log update

---------

Co-authored-by: Copilot <[email protected]>

Author: dibahlfi

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -8,6 +8,8 @@
 #### Breaking Changes
 
 #### Bugs Fixed
+* Improved the resilience of Database Account Read metadata operation against short-lived network issues by increasing number of retries. See [PR 42525](https://github.com/Azure/azure-sdk-for-python/pull/42525).
+* Fixed bug where during health checks read regions were marked as unavailable for write operations. See [PR 42525](https://github.com/Azure/azure-sdk-for-python/pull/42525).
 
 #### Other Changes
 * Added session token false progress merge logic. See [42393](https://github.com/Azure/azure-sdk-for-python/pull/42393)

sdk/cosmos/azure-cosmos/azure/cosmos/_constants.py

@@ -56,6 +56,13 @@ class _Constants:
     CIRCUIT_BREAKER_ENABLED_CONFIG: str =  "AZURE_COSMOS_ENABLE_CIRCUIT_BREAKER"
     CIRCUIT_BREAKER_ENABLED_CONFIG_DEFAULT: str = "False"
     AAD_SCOPE_OVERRIDE: str = "AZURE_COSMOS_AAD_SCOPE_OVERRIDE"
+
+    # Database Account Retry Policy constants
+    AZURE_COSMOS_HEALTH_CHECK_MAX_RETRIES: str = "AZURE_COSMOS_HEALTH_CHECK_MAX_RETRIES"
+    AZURE_COSMOS_HEALTH_CHECK_MAX_RETRIES_DEFAULT: int = 3
+    AZURE_COSMOS_HEALTH_CHECK_RETRY_AFTER_MS: str = "AZURE_COSMOS_HEALTH_CHECK_RETRY_AFTER_MS"
+    AZURE_COSMOS_HEALTH_CHECK_RETRY_AFTER_MS_DEFAULT: int = 100
+
     # Only applicable when circuit breaker is enabled -------------------------
     CONSECUTIVE_ERROR_COUNT_TOLERATED_FOR_READ: str = "AZURE_COSMOS_CONSECUTIVE_ERROR_COUNT_TOLERATED_FOR_READ"
     CONSECUTIVE_ERROR_COUNT_TOLERATED_FOR_READ_DEFAULT: int = 10

sdk/cosmos/azure-cosmos/azure/cosmos/_database_account_retry_policy.py

@@ -22,28 +22,55 @@
 """Internal class for database account retry policy implementation in the
 Azure Cosmos database service.
 """
+import os
+from azure.core.exceptions import ServiceRequestError, ServiceResponseError
+from azure.cosmos import _constants
+
 
 class DatabaseAccountRetryPolicy(object):
-    """The database account retry policy which should only retry once regardless of errors.
-    """
+    """Implements retry logic for database account reads in Azure Cosmos DB."""
+
+    # List of HTTP status codes considered transient errors for retry logic.
+    transient_status_codes = [502, 503, 504]
+
+    # Tuple of exception types considered transient errors for retry logic.
+    transient_exceptions = (ServiceRequestError, ServiceResponseError)
 
     def __init__(self, connection_policy):
         self.retry_count = 0
-        self.retry_after_in_milliseconds = 0
-        self.max_retry_attempt_count = 1
+        self.retry_after_in_milliseconds = int(os.getenv(
+            _constants._Constants.AZURE_COSMOS_HEALTH_CHECK_RETRY_AFTER_MS,
+            str(_constants._Constants.AZURE_COSMOS_HEALTH_CHECK_RETRY_AFTER_MS_DEFAULT)
+        ))
+        self.max_retry_attempt_count = int(os.getenv(
+            _constants._Constants.AZURE_COSMOS_HEALTH_CHECK_MAX_RETRIES,
+            str(_constants._Constants.AZURE_COSMOS_HEALTH_CHECK_MAX_RETRIES_DEFAULT)
+        ))
         self.connection_policy = connection_policy
 
-    def ShouldRetry(self, exception):  # pylint: disable=unused-argument
-        """Returns true if the request should retry based on the passed-in exception.
+    def ShouldRetry(self, exception):
+        """
+        Determines if the given exception is transient and if a retry should be attempted.
 
-        :param exceptions.CosmosHttpResponseError exception:
-        :returns: a boolean stating whether the request should be retried
+        :param exception: The exception instance to evaluate.
+        :type exception: Exception
+        :return: True if the exception is transient and retry attempts to remain, False otherwise.
         :rtype: bool
         """
 
-        if self.retry_count >= self.max_retry_attempt_count:
-            return False
+        is_transient = False
+
+        # Check for transient HTTP status codes
+        status_code = getattr(exception, "status_code", None)
+        if status_code in self.transient_status_codes:
+            is_transient = True
+
+        # Check for transient exception types
+        if isinstance(exception, self.transient_exceptions):
+            is_transient = True
 
-        self.retry_count += 1
+        if is_transient and self.retry_count < self.max_retry_attempt_count:
+            self.retry_count += 1
+            return True
 
-        return True
+        return False

sdk/cosmos/azure-cosmos/azure/cosmos/_global_endpoint_manager.py

@@ -96,6 +96,15 @@ def force_refresh_on_startup(self, database_account):
     def update_location_cache(self):
         self.location_cache.update_location_cache()
 
+    def _mark_endpoint_unavailable(self, endpoint: str):
+        """Marks an endpoint as unavailable for the appropriate operations.
+        :param str endpoint: The endpoint to mark as unavailable.
+        """
+        write_endpoints = self.location_cache.get_all_write_endpoints()
+        self.mark_endpoint_unavailable_for_read(endpoint, False)
+        if endpoint in write_endpoints:
+            self.mark_endpoint_unavailable_for_write(endpoint, False)
+
     def refresh_endpoint_list(self, database_account, **kwargs):
         if current_time_millis() - self.last_refresh_time > self.refresh_time_interval_in_ms:
             self.refresh_needed = True
@@ -145,8 +154,7 @@ def _GetDatabaseAccount(self, **kwargs) -> Tuple[DatabaseAccount, str]:
         except (exceptions.CosmosHttpResponseError, AzureError):
             # when atm is available, L: 145, 146 should be removed as the global endpoint shouldn't be used
             # for dataplane operations anymore
-            self.mark_endpoint_unavailable_for_read(self.DefaultEndpoint, False)
-            self.mark_endpoint_unavailable_for_write(self.DefaultEndpoint, False)
+            self._mark_endpoint_unavailable(self.DefaultEndpoint)
             for location_name in self.PreferredLocations:
                 locational_endpoint = LocationCache.GetLocationalEndpoint(self.DefaultEndpoint, location_name)
                 try:
@@ -155,8 +163,7 @@ def _GetDatabaseAccount(self, **kwargs) -> Tuple[DatabaseAccount, str]:
                     self.location_cache.mark_endpoint_available(locational_endpoint)
                     return database_account, locational_endpoint
                 except (exceptions.CosmosHttpResponseError, AzureError):
-                    self.mark_endpoint_unavailable_for_read(locational_endpoint, False)
-                    self.mark_endpoint_unavailable_for_write(locational_endpoint, False)
+                    self._mark_endpoint_unavailable(locational_endpoint)
             raise
 
     def _endpoints_health_check(self, **kwargs):
@@ -191,8 +198,8 @@ def _endpoints_health_check(self, **kwargs):
                     success_count += 1
                     self.location_cache.mark_endpoint_available(endpoint)
                 except (exceptions.CosmosHttpResponseError, AzureError):
-                    self.mark_endpoint_unavailable_for_read(endpoint, False)
-                    self.mark_endpoint_unavailable_for_write(endpoint, False)
+                    self._mark_endpoint_unavailable(endpoint)
+
                 finally:
                     # after the health check for that endpoint setting the timeouts back to their original values
                     self.client.connection_policy.override_dba_timeouts(previous_dba_read_timeout,

sdk/cosmos/azure-cosmos/azure/cosmos/_location_cache.py

@@ -203,6 +203,13 @@ def perform_on_database_account_read(self, database_account):
             database_account._EnableMultipleWritableLocations,
         )
 
+    def get_all_write_endpoints(self) -> Set[str]:
+        return {
+            endpoint
+            for context in self.get_write_regional_routing_contexts()
+            for endpoint in (context.get_primary(), context.get_alternate())
+        }
+
     def get_ordered_write_locations(self):
         return self.account_write_locations
 

sdk/cosmos/azure-cosmos/azure/cosmos/aio/_global_endpoint_manager_async.py

@@ -100,6 +100,15 @@ async def force_refresh_on_startup(self, database_account):
     def update_location_cache(self):
         self.location_cache.update_location_cache()
 
+    def _mark_endpoint_unavailable(self, endpoint: str):
+        """Marks an endpoint as unavailable for the appropriate operations.
+        :param str endpoint: The endpoint to mark as unavailable.
+        """
+        write_endpoints = self.location_cache.get_all_write_endpoints()
+        self.mark_endpoint_unavailable_for_read(endpoint, False)
+        if endpoint in write_endpoints:
+            self.mark_endpoint_unavailable_for_write(endpoint, False)
+
     async def refresh_endpoint_list(self, database_account, **kwargs):
         if self.refresh_task and self.refresh_task.done():
             try:
@@ -142,8 +151,7 @@ async def _database_account_check(self, endpoint: str, **kwargs: Dict[str, Any])
             await self.client._GetDatabaseAccountCheck(endpoint, **kwargs)
             self.location_cache.mark_endpoint_available(endpoint)
         except (exceptions.CosmosHttpResponseError, AzureError):
-            self.mark_endpoint_unavailable_for_read(endpoint, False)
-            self.mark_endpoint_unavailable_for_write(endpoint, False)
+            self._mark_endpoint_unavailable(endpoint)
 
     async def _endpoints_health_check(self, **kwargs):
         """Gets the database account for each endpoint.
@@ -185,8 +193,7 @@ async def _GetDatabaseAccount(self, **kwargs) -> Tuple[DatabaseAccount, str]:
         # until we get the database account and return None at the end, if we are not able
         # to get that info from any endpoints
         except (exceptions.CosmosHttpResponseError, AzureError):
-            self.mark_endpoint_unavailable_for_read(self.DefaultEndpoint, False)
-            self.mark_endpoint_unavailable_for_write(self.DefaultEndpoint, False)
+            self._mark_endpoint_unavailable(self.DefaultEndpoint)
             for location_name in self.PreferredLocations:
                 locational_endpoint = LocationCache.GetLocationalEndpoint(self.DefaultEndpoint, location_name)
                 try:
@@ -195,8 +202,7 @@ async def _GetDatabaseAccount(self, **kwargs) -> Tuple[DatabaseAccount, str]:
                     self.location_cache.mark_endpoint_available(locational_endpoint)
                     return database_account, locational_endpoint
                 except (exceptions.CosmosHttpResponseError, AzureError):
-                    self.mark_endpoint_unavailable_for_read(locational_endpoint, False)
-                    self.mark_endpoint_unavailable_for_write(locational_endpoint, False)
+                    self._mark_endpoint_unavailable(locational_endpoint)
             raise
 
     async def _GetDatabaseAccountStub(self, endpoint, **kwargs):

sdk/cosmos/azure-cosmos/tests/test_health_check_async.py

@@ -186,17 +186,34 @@ async def test_health_check_failure_async(self, setup, preferred_location, use_w
         try:
             setup[COLLECTION].client_connection._global_endpoint_manager.startup = False
             setup[COLLECTION].client_connection._global_endpoint_manager.refresh_needed = True
-            for i in range(2):
-                await setup[COLLECTION].create_item(body={'id': 'item' + str(uuid.uuid4()), 'pk': 'pk'})
-                # wait for background task to finish
-                await asyncio.sleep(2)
+            # Trigger the background health check
+            await setup[COLLECTION].create_item(body={'id': 'item' + str(uuid.uuid4()), 'pk': 'pk'})
+
+            # Poll until the background task marks the endpoints as unavailable
+            start_time = time.time()
+            while (len(setup[COLLECTION].client_connection._global_endpoint_manager.location_cache.location_unavailability_info_by_endpoint) < len(
+                    REGIONS) and time.time() - start_time < 10):
+                await asyncio.sleep(0.1)
         finally:
             _global_endpoint_manager_async._GlobalEndpointManager._GetDatabaseAccountStub = self.original_getDatabaseAccountStub
             setup[COLLECTION].client_connection.connection_policy.PreferredLocations = self.original_preferred_locations
 
         num_unavailable_endpoints = len(REGIONS)
         unavailable_endpoint_info = setup[COLLECTION].client_connection._global_endpoint_manager.location_cache.location_unavailability_info_by_endpoint
         assert len(unavailable_endpoint_info) == num_unavailable_endpoints
+        # Allow both global and regional endpoint to be considered write endpoints when global write is enabled
+        write_endpoints = {
+            _location_cache.LocationCache.GetLocationalEndpoint(self.host, REGION_1)
+        }
+        if use_write_global_endpoint:
+            write_endpoints.add(self.host)
+
+        for endpoint, info in unavailable_endpoint_info.items():
+            assert _location_cache.EndpointOperationType.ReadType in info["operationType"]
+            if endpoint in write_endpoints:
+                assert _location_cache.EndpointOperationType.WriteType in info["operationType"]
+            else:
+                assert _location_cache.EndpointOperationType.WriteType not in info["operationType"]
 
     async def mock_health_check(self, **kwargs):
         await asyncio.sleep(100)

sdk/cosmos/azure-cosmos/tests/test_retry_policy.py

@@ -11,8 +11,12 @@
 import azure.cosmos.exceptions as exceptions
 import test_config
 from azure.cosmos import _retry_utility, PartitionKey, documents
-from azure.cosmos.http_constants import HttpHeaders, StatusCodes
+from azure.cosmos.http_constants import HttpHeaders, StatusCodes, ResourceType
 from _fault_injection_transport import FaultInjectionTransport
+import os
+from azure.core.exceptions import ServiceResponseError
+from azure.cosmos._database_account_retry_policy import DatabaseAccountRetryPolicy
+from azure.cosmos._constants import _Constants
 
 
 def setup_method_with_custom_transport(
@@ -483,6 +487,90 @@ def test_patch_replace_no_retry(self):
             container.replace_item(item=doc['id'], body=doc)
         assert connection_retry_policy.counter == 0
 
+
+    def test_database_account_read_retry_policy(self):
+        os.environ['AZURE_COSMOS_HEALTH_CHECK_MAX_RETRIES'] = '5'
+        os.environ['AZURE_COSMOS_HEALTH_CHECK_RETRY_AFTER_MS'] = '100'
+        max_retries = int(os.environ['AZURE_COSMOS_HEALTH_CHECK_MAX_RETRIES'])
+        retry_after_ms = int(os.environ['AZURE_COSMOS_HEALTH_CHECK_RETRY_AFTER_MS'])
+        self.original_execute_function = _retry_utility.ExecuteFunction
+        mock_execute = self.MockExecuteFunctionDBA(self.original_execute_function)
+        _retry_utility.ExecuteFunction = mock_execute
+
+        try:
+            with self.assertRaises(exceptions.CosmosHttpResponseError) as context:
+                cosmos_client.CosmosClient(self.host, self.masterKey)
+                # Client initialization triggers database account read
+
+            self.assertEqual(context.exception.status_code, 503)
+            self.assertEqual(mock_execute.counter, max_retries + 1)
+            policy = DatabaseAccountRetryPolicy(self.connectionPolicy)
+            self.assertEqual(policy.retry_after_in_milliseconds, retry_after_ms)
+        finally:
+            del os.environ["AZURE_COSMOS_HEALTH_CHECK_MAX_RETRIES"]
+            del os.environ["AZURE_COSMOS_HEALTH_CHECK_RETRY_AFTER_MS"]
+            _retry_utility.ExecuteFunction = self.original_execute_function
+
+    def test_database_account_read_retry_policy_defaults(self):
+        self.original_execute_function = _retry_utility.ExecuteFunction
+        mock_execute = self.MockExecuteFunctionDBA(self.original_execute_function)
+        _retry_utility.ExecuteFunction = mock_execute
+
+        try:
+            with self.assertRaises(exceptions.CosmosHttpResponseError) as context:
+                cosmos_client.CosmosClient(self.host, self.masterKey)
+                # Triggers database account read
+
+            self.assertEqual(context.exception.status_code, 503)
+            self.assertEqual(
+                mock_execute.counter,
+                _Constants.AZURE_COSMOS_HEALTH_CHECK_MAX_RETRIES_DEFAULT + 1
+            )
+            policy = DatabaseAccountRetryPolicy(self.connectionPolicy)
+            self.assertEqual(
+                policy.retry_after_in_milliseconds,
+                _Constants.AZURE_COSMOS_HEALTH_CHECK_RETRY_AFTER_MS_DEFAULT
+            )
+        finally:
+            _retry_utility.ExecuteFunction = self.original_execute_function
+
+    def test_database_account_read_retry_with_service_response_error(self):
+        self.original_execute_function = _retry_utility.ExecuteFunction
+        mock_execute = self.MockExecuteFunctionDBAServiceRequestError(self.original_execute_function)
+        _retry_utility.ExecuteFunction = mock_execute
+
+        try:
+            with self.assertRaises(ServiceResponseError):
+                cosmos_client.CosmosClient(self.host, self.masterKey)
+                # Client initialization triggers database account read
+
+            # Should use default retry attempts from _constants.py
+            self.assertEqual(
+                mock_execute.counter,
+                _Constants.AZURE_COSMOS_HEALTH_CHECK_MAX_RETRIES_DEFAULT + 1
+            )
+            policy = DatabaseAccountRetryPolicy(self.connectionPolicy)
+            self.assertEqual(
+                policy.retry_after_in_milliseconds,
+                _Constants.AZURE_COSMOS_HEALTH_CHECK_RETRY_AFTER_MS_DEFAULT
+            )
+        finally:
+            _retry_utility.ExecuteFunction = self.original_execute_function
+
+    class MockExecuteFunctionDBAServiceRequestError(object):
+        def __init__(self, org_func):
+            self.org_func = org_func
+            self.counter = 0
+
+        def __call__(self, func, *args, **kwargs):
+            # The second argument to the internal _request function is the RequestObject.
+            request_object = args[1]
+            if (request_object.operation_type == documents._OperationType.Read and
+                    request_object.resource_type == ResourceType.DatabaseAccount):
+                self.counter += 1
+                raise ServiceResponseError("mocked service response error")
+            return self.org_func(func, *args, **kwargs)
+
     def _MockExecuteFunction(self, function, *args, **kwargs):
         response = test_config.FakeResponse({HttpHeaders.RetryAfterInMilliseconds: self.retry_after_in_milliseconds})
         raise exceptions.CosmosHttpResponseError(
@@ -516,5 +604,21 @@ def __call__(self, func, *args, **kwargs):
                     message="Connection was reset",
                     response=test_config.FakeResponse({}))
 
+    class MockExecuteFunctionDBA(object):
+        def __init__(self, org_func):
+            self.org_func = org_func
+            self.counter = 0
+
+        def __call__(self, func, *args, **kwargs):
+            request_object = args[1]
+            if (request_object.operation_type == documents._OperationType.Read and
+                    request_object.resource_type == ResourceType.DatabaseAccount):
+                self.counter += 1
+                raise exceptions.CosmosHttpResponseError(
+                    status_code=503,
+                    message="Service Unavailable.")
+            return self.org_func(func, *args, **kwargs)
+
+
 if __name__ == '__main__':
     unittest.main()