From ac78da9632bf406e6ac38b4d9ea79cc03eb7d41d Mon Sep 17 00:00:00 2001 From: Kushagra Thapar Date: Wed, 5 Feb 2025 23:13:52 -0800 Subject: [PATCH] Fixed the timeout retry policy --- .../azure-cosmos/azure/cosmos/_retry_utility.py | 7 +++++-- .../azure/cosmos/_timeout_failover_retry_policy.py | 13 ++++--------- .../azure/cosmos/aio/_retry_utility_async.py | 5 +++-- .../azure-cosmos/azure/cosmos/http_constants.py | 1 - sdk/cosmos/azure-cosmos/test/test_globaldb_mock.py | 2 +- 5 files changed, 13 insertions(+), 15 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_retry_utility.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_retry_utility.py index 99784facecc4..927ed7a41baa 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_retry_utility.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_retry_utility.py @@ -131,7 +131,6 @@ def Execute(client, global_endpoint_manager, function, *args, **kwargs): sub_status_code=SubStatusCodes.THROUGHPUT_OFFER_NOT_FOUND) return result except exceptions.CosmosHttpResponseError as e: - retry_policy = defaultRetry_policy if request and _has_database_account_header(request.headers): retry_policy = database_account_retry_policy # Re-assign retry policy based on error code @@ -173,8 +172,12 @@ def Execute(client, global_endpoint_manager, function, *args, **kwargs): retry_policy.container_rid = cached_container["_rid"] request.headers[retry_policy._intended_headers] = retry_policy.container_rid - elif e.status_code in [StatusCodes.REQUEST_TIMEOUT, StatusCodes.SERVICE_UNAVAILABLE]: + elif e.status_code == StatusCodes.REQUEST_TIMEOUT: retry_policy = timeout_failover_retry_policy + elif e.status_code >= StatusCodes.INTERNAL_SERVER_ERROR: + retry_policy = timeout_failover_retry_policy + else: + retry_policy = defaultRetry_policy # If none of the retry policies applies or there is no retry needed, set the # throttle related response headers and re-throw the exception back arg[0] diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_timeout_failover_retry_policy.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_timeout_failover_retry_policy.py index aa66cd4f76e7..60f0208e6351 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_timeout_failover_retry_policy.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_timeout_failover_retry_policy.py @@ -5,18 +5,17 @@ Cosmos database service. """ from azure.cosmos.documents import _OperationType -from . import http_constants class _TimeoutFailoverRetryPolicy(object): def __init__(self, connection_policy, global_endpoint_manager, *args): - self._max_retry_attempt_count = 120 - self._max_service_unavailable_retry_count = 1 - self.retry_after_in_milliseconds = 0 + self.retry_after_in_milliseconds = 500 self.args = args self.global_endpoint_manager = global_endpoint_manager + # If an account only has 1 region, then we still want to retry once on the same region + self._max_retry_attempt_count = len(self.global_endpoint_manager.location_cache.read_regional_endpoints) + 1 self.retry_count = 0 self.connection_policy = connection_policy self.request = args[0] if args else None @@ -28,17 +27,13 @@ def ShouldRetry(self, _exception): :returns: a boolean stating whether the request should be retried :rtype: bool """ - # we don't retry on write operations for timeouts or service unavailable + # we don't retry on write operations for timeouts or any internal server errors if self.request and (not _OperationType.IsReadOnlyOperation(self.request.operation_type)): return False if not self.connection_policy.EnableEndpointDiscovery: return False - # Check if the next retry about to be done is safe - if _exception.status_code == http_constants.StatusCodes.SERVICE_UNAVAILABLE and \ - self.retry_count >= self._max_service_unavailable_retry_count: - return False self.retry_count += 1 # Check if the next retry about to be done is safe if self.retry_count >= self._max_retry_attempt_count: diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_retry_utility_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_retry_utility_async.py index 74df8ea9479f..c4be5a1afc2c 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_retry_utility_async.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_retry_utility_async.py @@ -130,7 +130,6 @@ async def ExecuteAsync(client, global_endpoint_manager, function, *args, **kwarg return result except exceptions.CosmosHttpResponseError as e: - retry_policy = None if request and _has_database_account_header(request.headers): retry_policy = database_account_retry_policy elif e.status_code == StatusCodes.FORBIDDEN and e.sub_status in \ @@ -171,7 +170,9 @@ async def ExecuteAsync(client, global_endpoint_manager, function, *args, **kwarg retry_policy.container_rid = cached_container["_rid"] request.headers[retry_policy._intended_headers] = retry_policy.container_rid - elif e.status_code in [StatusCodes.REQUEST_TIMEOUT, StatusCodes.SERVICE_UNAVAILABLE]: + elif e.status_code == StatusCodes.REQUEST_TIMEOUT: + retry_policy = timeout_failover_retry_policy + elif e.status_code >= StatusCodes.INTERNAL_SERVER_ERROR: retry_policy = timeout_failover_retry_policy else: retry_policy = defaultRetry_policy diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/http_constants.py b/sdk/cosmos/azure-cosmos/azure/cosmos/http_constants.py index 8a7b57b7c93f..31a95d2600d6 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/http_constants.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/http_constants.py @@ -400,7 +400,6 @@ class StatusCodes: RETRY_WITH = 449 INTERNAL_SERVER_ERROR = 500 - SERVICE_UNAVAILABLE = 503 # Operation pause and cancel. These are FAKE status codes for QOS logging purpose only. OPERATION_PAUSED = 1200 diff --git a/sdk/cosmos/azure-cosmos/test/test_globaldb_mock.py b/sdk/cosmos/azure-cosmos/test/test_globaldb_mock.py index 5548b51839b3..83d5e2603e9f 100644 --- a/sdk/cosmos/azure-cosmos/test/test_globaldb_mock.py +++ b/sdk/cosmos/azure-cosmos/test/test_globaldb_mock.py @@ -166,7 +166,7 @@ def MockExecuteFunction(self, function, *args, **kwargs): def MockGetDatabaseAccountStub(self, endpoint): raise exceptions.CosmosHttpResponseError( - status_code=StatusCodes.SERVICE_UNAVAILABLE, message="Service unavailable") + status_code=StatusCodes.INTERNAL_SERVER_ERROR, message="Internal Server Error") def test_global_db_endpoint_discovery_retry_policy(self): connection_policy = documents.ConnectionPolicy()