[Cosmos] Fix Service Request Retry Policy (#41804)

* fix request service retry policy

* Update CHANGELOG.md

* Update _service_request_retry_policy.py

Author: bambriz

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -9,6 +9,7 @@
 
 #### Bugs Fixed
 * Fixed bug where replacing manual throughput using `ThroughputProperties` would not work. See [PR 41564](https://github.com/Azure/azure-sdk-for-python/pull/41564)
+* Fixed bug where constantly raising Service Request Error Exceptions would cause the Service Request Retry Policy to indefinitely retry the request during a query or when a request was sent without a request object. See [PR 41804](https://github.com/Azure/azure-sdk-for-python/pull/41804)
 
 #### Other Changes
 

sdk/cosmos/azure-cosmos/azure/cosmos/_service_request_retry_policy.py

@@ -30,7 +30,7 @@ def __init__(self, connection_policy, global_endpoint_manager, pk_range_wrapper,
             else:
                 self.total_retries = len(self.global_endpoint_manager.location_cache.write_regional_routing_contexts)
 
-    def ShouldRetry(self):
+    def ShouldRetry(self):  # pylint: disable=too-many-return-statements
         """Returns true if the request should retry based on preferred regions and retries already done.
 
         :returns: a boolean stating whether the request should be retried
@@ -87,6 +87,11 @@ def ShouldRetry(self):
                     location_endpoint = self.resolve_next_region_service_endpoint()
 
             self.request.route_to_location(location_endpoint)
+            return True
+        # Check if the next retry about to be done is safe
+        if (self.failover_retry_count + 1) >= self.total_retries:
+            return False
+        self.failover_retry_count += 1
         return True
 
     # This function prepares the request to go to the second endpoint in the same region

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

@@ -54,28 +54,57 @@ def test_service_request_retry_policy(self):
                                                                                   self.REGION3: self.REGIONAL_ENDPOINT}
         original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT,
                                                                   self.REGIONAL_ENDPOINT]
+
+        expected_counter = len(original_location_cache.read_regional_routing_contexts)
         try:
             # Mock the function to return the ServiceRequestException we retry
             mf = self.MockExecuteServiceRequestException()
             _retry_utility.ExecuteFunction = mf
             container.read_item(created_item['id'], created_item['pk'])
             pytest.fail("Exception was not raised.")
         except ServiceRequestError:
-            assert mf.counter == 3
+            assert mf.counter == expected_counter
+        finally:
+            _retry_utility.ExecuteFunction = self.original_execute_function
+
+        # Now we test with a query operation, iterating through items sends request without request object
+        # retry policy should eventually raise an exception as it should stop retrying with a max retry attempt
+        # equal to the available read region locations
+
+        # Change the location cache to have 3 preferred read regions and 3 available read endpoints by location
+        original_location_cache = mock_client.client_connection._global_endpoint_manager.location_cache
+        original_location_cache.account_read_locations = [self.REGION1, self.REGION2, self.REGION3]
+        original_location_cache.available_read_regional_endpoints_by_locations = {
+            self.REGION1: self.REGIONAL_ENDPOINT,
+            self.REGION2: self.REGIONAL_ENDPOINT,
+            self.REGION3: self.REGIONAL_ENDPOINT}
+        original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT,
+                                                                      self.REGIONAL_ENDPOINT]
+
+        expected_counter = len(original_location_cache.read_regional_routing_contexts)
+        try:
+            # Mock the function to return the ServiceRequestException we retry
+            mf = self.MockExecuteServiceRequestException()
+            _retry_utility.ExecuteFunction = mf
+            items = list(container.query_items(query="SELECT * FROM c", partition_key=created_item['pk']))
+            pytest.fail("Exception was not raised.")
+        except ServiceRequestError:
+            assert mf.counter == expected_counter
         finally:
             _retry_utility.ExecuteFunction = self.original_execute_function
 
         # Now we change the location cache to have only 1 preferred read region
         original_location_cache.account_read_locations = [self.REGION1]
         original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT]
+        expected_counter = len(original_location_cache.read_regional_routing_contexts)
         try:
             # Reset the function to reset the counter
             mf = self.MockExecuteServiceRequestException()
             _retry_utility.ExecuteFunction = mf
             container.read_item(created_item['id'], created_item['pk'])
             pytest.fail("Exception was not raised.")
         except ServiceRequestError:
-            assert mf.counter == 1
+            assert mf.counter == expected_counter
         finally:
             _retry_utility.ExecuteFunction = self.original_execute_function
 
@@ -84,6 +113,7 @@ def test_service_request_retry_policy(self):
         original_location_cache.write_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT]
         original_location_cache.available_write_regional_endpoints_by_locations = {self.REGION1: self.REGIONAL_ENDPOINT,
                                                                                    self.REGION2: self.REGIONAL_ENDPOINT}
+        expected_counter = len(original_location_cache.write_regional_routing_contexts)
         try:
             # Reset the function to reset the counter
             mf = self.MockExecuteServiceRequestException()
@@ -92,7 +122,7 @@ def test_service_request_retry_policy(self):
             pytest.fail("Exception was not raised.")
         except ServiceRequestError:
             # Should retry twice in each region
-            assert mf.counter == 2
+            assert mf.counter == expected_counter
         finally:
             _retry_utility.ExecuteFunction = self.original_execute_function
 

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

@@ -65,28 +65,57 @@ async def test_service_request_retry_policy_async(self):
                 self.REGION3: self.REGIONAL_ENDPOINT}
             original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT,
                                                                       self.REGIONAL_ENDPOINT]
+            expected_counter = len(original_location_cache.read_regional_routing_contexts)
             try:
                 # Mock the function to return the ServiceRequestException we retry
                 mf = self.MockExecuteServiceRequestException()
                 _retry_utility_async.ExecuteFunctionAsync = mf
                 await container.read_item(created_item['id'], created_item['pk'])
                 pytest.fail("Exception was not raised.")
             except ServiceRequestError:
-                assert mf.counter == 3
+                assert mf.counter == expected_counter
+            finally:
+                _retry_utility_async.ExecuteFunctionAsync = self.original_execute_function
+
+            # Now we test with a query operation, iterating through items sends request without request object
+            # retry policy should eventually raise an exception as it should stop retrying with a max retry attempt
+            # equal to the available read region locations
+
+            # Change the location cache to have 3 preferred read regions and 3 available read endpoints by location
+            original_location_cache = mock_client.client_connection._global_endpoint_manager.location_cache
+            original_location_cache.account_read_locations = [self.REGION1, self.REGION2, self.REGION3]
+            original_location_cache.available_read_regional_endpoints_by_locations = {
+                self.REGION1: self.REGIONAL_ENDPOINT,
+                self.REGION2: self.REGIONAL_ENDPOINT,
+                self.REGION3: self.REGIONAL_ENDPOINT}
+            original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT, self.REGIONAL_ENDPOINT,
+                                                                      self.REGIONAL_ENDPOINT]
+            expected_counter = len(original_location_cache.read_regional_routing_contexts)
+
+            try:
+                # Mock the function to return the ServiceRequestException we retry
+                mf = self.MockExecuteServiceRequestException()
+                _retry_utility_async.ExecuteFunctionAsync = mf
+                items = [item async for item in container.query_items(query="SELECT * FROM c",
+                                                                      partition_key=created_item['pk'])]
+                pytest.fail("Exception was not raised.")
+            except ServiceRequestError:
+                assert mf.counter == expected_counter
             finally:
                 _retry_utility_async.ExecuteFunctionAsync = self.original_execute_function
 
             # Now we change the location cache to have only 1 preferred read region
             original_location_cache.account_read_locations = [self.REGION1]
             original_location_cache.read_regional_routing_contexts = [self.REGIONAL_ENDPOINT]
+            expected_counter = len(original_location_cache.read_regional_routing_contexts)
             try:
                 # Reset the function to reset the counter
                 mf = self.MockExecuteServiceRequestException()
                 _retry_utility_async.ExecuteFunctionAsync = mf
                 await container.read_item(created_item['id'], created_item['pk'])
                 pytest.fail("Exception was not raised.")
             except ServiceRequestError:
-                assert mf.counter == 1
+                assert mf.counter == expected_counter
             finally:
                 _retry_utility_async.ExecuteFunctionAsync = self.original_execute_function
 
@@ -96,14 +125,15 @@ async def test_service_request_retry_policy_async(self):
             original_location_cache.available_write_regional_endpoints_by_locations = {
                 self.REGION1: self.REGIONAL_ENDPOINT,
                 self.REGION2: self.REGIONAL_ENDPOINT}
+            expected_counter = len(original_location_cache.write_regional_routing_contexts)
             try:
                 # Reset the function to reset the counter
                 mf = self.MockExecuteServiceRequestException()
                 _retry_utility_async.ExecuteFunctionAsync = mf
                 await container.create_item({"id": str(uuid.uuid4()), "pk": str(uuid.uuid4())})
                 pytest.fail("Exception was not raised.")
             except ServiceRequestError:
-                assert mf.counter == 2
+                assert mf.counter == expected_counter
             finally:
                 _retry_utility_async.ExecuteFunctionAsync = self.original_execute_function