Fixed the timeout logic

Author: kushagraThapar

sdk/cosmos/azure-cosmos/azure/cosmos/_timeout_failover_retry_policy.py

@@ -18,7 +18,6 @@ def __init__(self, connection_policy, global_endpoint_manager, *args):
 
         self.global_endpoint_manager = global_endpoint_manager
         self.retry_count = 0
-        self.location_index = 0
         self.connection_policy = connection_policy
         self.request = args[0] if args else None
 
@@ -29,14 +28,13 @@ def ShouldRetry(self, _exception):
         :returns: a boolean stating whether the request should be retried
         :rtype: bool
         """
-        if self.request:
-            if not _OperationType.IsReadOnlyOperation(self.request.operation_type):
-                return False
+        # we don't retry on write operations for timeouts or service unavailable
+        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:
@@ -47,46 +45,19 @@ def ShouldRetry(self, _exception):
             return False
 
         if self.request:
-            # Update the last routed location to where this request was routed previously.
-            # So that we can check in location cache if we need to return the current or previous
-            # based on where the request was routed previously.
-            self.request.last_routed_location_endpoint_within_region = self.request.location_endpoint_to_route
-
-            if _OperationType.IsReadOnlyOperation(self.request.operation_type):
-                # We just directly got to the next location in case of read requests
-                # We don't retry again on the same region for regional endpoint
-                location_endpoint = self.resolve_next_region_service_endpoint()
-            else:
-                self.global_endpoint_manager.swap_regional_endpoint_values(self.request)
-                location_endpoint = self.resolve_current_region_service_endpoint()
-                # This is the case where both current and previous point to the same writable endpoint
-                # In this case we don't want to retry again, rather failover to the next region
-                if self.request.last_routed_location_endpoint_within_region == location_endpoint:
-                    location_endpoint = self.resolve_next_region_service_endpoint()
-
+            location_endpoint = self.resolve_next_region_service_endpoint()
             self.request.route_to_location(location_endpoint)
         return True
 
-
-    # This function prepares the request to go to the second endpoint in the same region
-    def resolve_current_region_service_endpoint(self):
-        # clear previous location-based routing directive
-        self.request.clear_route_to_location()
-        # resolve the next service endpoint in the same region
-        # since we maintain 2 endpoints per region for write operations
-        self.request.route_to_location_with_preferred_location_flag(self.location_index, True)
-        return self.global_endpoint_manager.resolve_service_endpoint(self.request)
-
     # This function prepares the request to go to the next region
     def resolve_next_region_service_endpoint(self):
-        self.location_index += 1
         # clear previous location-based routing directive
         self.request.clear_route_to_location()
         # clear the last routed endpoint within same region since we are going to a new region now
         self.request.last_routed_location_endpoint_within_region = None
         # set location-based routing directive based on retry count
         # ensuring usePreferredLocations is set to True for retry
-        self.request.route_to_location_with_preferred_location_flag(self.location_index, True)
+        self.request.route_to_location_with_preferred_location_flag(self.retry_count, True)
         # Resolve the endpoint for the request and pin the resolution to the resolved endpoint
         # This enables marking the endpoint unavailability on endpoint failover/unreachability
         return self.global_endpoint_manager.resolve_service_endpoint(self.request)