Fixed the timeout retry policy

Author: kushagraThapar

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]

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:

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

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

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()