Per Partition Circuit Breaker

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -3,6 +3,7 @@
 ### 4.10.0b5 (Unreleased)
 
 #### Features Added
+* Per partition circuit breaker support. It can be enabled through the environment variable `AZURE_COSMOS_ENABLE_CIRCUIT_BREAKER`. See [PR 40302](https://github.com/Azure/azure-sdk-for-python/pull/40302).
 
 #### Breaking Changes
 

sdk/cosmos/azure-cosmos/azure/cosmos/_base.py

@@ -63,6 +63,7 @@
     'priority': 'priorityLevel',
     'no_response': 'responsePayloadOnWriteDisabled',
     'max_item_count': 'maxItemCount',
+    'excluded_locations': 'excludedLocations',
 }
 
 # Cosmos resource ID validation regex breakdown:

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

@@ -51,6 +51,18 @@ class _Constants:
     HS_MAX_ITEMS_CONFIG_DEFAULT: int = 1000
     MAX_ITEM_BUFFER_VS_CONFIG: str = "AZURE_COSMOS_MAX_ITEM_BUFFER_VECTOR_SEARCH"
     MAX_ITEM_BUFFER_VS_CONFIG_DEFAULT: int = 50000
+    CIRCUIT_BREAKER_ENABLED_CONFIG: str =  "AZURE_COSMOS_ENABLE_CIRCUIT_BREAKER"
+    CIRCUIT_BREAKER_ENABLED_CONFIG_DEFAULT: str = "False"
+    # 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
+    CONSECUTIVE_ERROR_COUNT_TOLERATED_FOR_WRITE: str = "AZURE_COSMOS_CONSECUTIVE_ERROR_COUNT_TOLERATED_FOR_WRITE"
+    CONSECUTIVE_ERROR_COUNT_TOLERATED_FOR_WRITE_DEFAULT: int = 5
+    FAILURE_PERCENTAGE_TOLERATED = "AZURE_COSMOS_FAILURE_PERCENTAGE_TOLERATED"
+    FAILURE_PERCENTAGE_TOLERATED_DEFAULT: int = 90
+    STALE_PARTITION_UNAVAILABILITY_CHECK = "AZURE_COSMOS_STALE_PARTITION_UNAVAILABILITY_CHECK_IN_SECONDS"
+    STALE_PARTITION_UNAVAILABILITY_CHECK_DEFAULT: int = 120
+    # -------------------------------------------------------------------------
 
     # Error code translations
     ERROR_TRANSLATIONS: Dict[int, str] = {

sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py

@@ -48,7 +48,7 @@
     HttpResponse  # pylint: disable=no-legacy-azure-core-http-response-import
 
 from . import _base as base
-from . import _global_endpoint_manager as global_endpoint_manager
+from ._global_partition_endpoint_manager_circuit_breaker import _GlobalPartitionEndpointManagerForCircuitBreaker
 from . import _query_iterable as query_iterable
 from . import _runtime_constants as runtime_constants
 from . import _session
@@ -164,7 +164,7 @@ def __init__( # pylint: disable=too-many-statements
         self.last_response_headers: CaseInsensitiveDict = CaseInsensitiveDict()
 
         self.UseMultipleWriteLocations = False
-        self._global_endpoint_manager = global_endpoint_manager._GlobalEndpointManager(self)
+        self._global_endpoint_manager = _GlobalPartitionEndpointManagerForCircuitBreaker(self)
 
         retry_policy = None
         if isinstance(self.connection_policy.ConnectionRetryConfiguration, HTTPPolicy):
@@ -2043,7 +2043,8 @@ def PatchItem(
         headers = base.GetHeaders(self, self.default_headers, "patch", path, document_id, resource_type,
                                   documents._OperationType.Patch, options)
         # Patch will use WriteEndpoint since it uses PUT operation
-        request_params = RequestObject(resource_type, documents._OperationType.Patch)
+        request_params = RequestObject(resource_type, documents._OperationType.Patch, headers)
+        request_params.set_excluded_location_from_options(options)
         request_data = {}
         if options.get("filterPredicate"):
             request_data["condition"] = options.get("filterPredicate")
@@ -2131,7 +2132,8 @@ def _Batch(
         base._populate_batch_headers(initial_headers)
         headers = base.GetHeaders(self, initial_headers, "post", path, collection_id, "docs",
                                   documents._OperationType.Batch, options)
-        request_params = RequestObject("docs", documents._OperationType.Batch)
+        request_params = RequestObject("docs", documents._OperationType.Batch, headers)
+        request_params.set_excluded_location_from_options(options)
         return cast(
             Tuple[List[Dict[str, Any]], CaseInsensitiveDict],
             self.__Post(path, request_params, batch_operations, headers, **kwargs)
@@ -2190,8 +2192,9 @@ def DeleteAllItemsByPartitionKey(
         path = '{}{}/{}'.format(path, "operations", "partitionkeydelete")
         collection_id = base.GetResourceIdOrFullNameFromLink(collection_link)
         headers = base.GetHeaders(self, self.default_headers, "post", path, collection_id,
-                                  "partitionkey", documents._OperationType.Delete, options)
-        request_params = RequestObject("partitionkey", documents._OperationType.Delete)
+                                  http_constants.ResourceType.PartitionKey, documents._OperationType.Delete, options)
+        request_params = RequestObject(http_constants.ResourceType.PartitionKey, documents._OperationType.Delete, headers)
+        request_params.set_excluded_location_from_options(options)
         _, last_response_headers = self.__Post(
             path=path,
             request_params=request_params,
@@ -2362,7 +2365,8 @@ def ExecuteStoredProcedure(
                                   documents._OperationType.ExecuteJavaScript, options)
 
         # ExecuteStoredProcedure will use WriteEndpoint since it uses POST operation
-        request_params = RequestObject("sprocs", documents._OperationType.ExecuteJavaScript)
+        request_params = RequestObject("sprocs", documents._OperationType.ExecuteJavaScript, headers)
+        request_params.set_excluded_location_from_options(options)
         result, self.last_response_headers = self.__Post(path, request_params, params, headers, **kwargs)
         return result
 
@@ -2558,7 +2562,7 @@ def GetDatabaseAccount(
 
         headers = base.GetHeaders(self, self.default_headers, "get", "", "", "",
                                   documents._OperationType.Read,{}, client_id=self.client_id)
-        request_params = RequestObject("databaseaccount", documents._OperationType.Read, url_connection)
+        request_params = RequestObject("databaseaccount", documents._OperationType.Read, headers, url_connection)
         result, last_response_headers = self.__Get("", request_params, headers, **kwargs)
         self.last_response_headers = last_response_headers
         database_account = DatabaseAccount()
@@ -2607,7 +2611,7 @@ def _GetDatabaseAccountCheck(
 
         headers = base.GetHeaders(self, self.default_headers, "get", "", "", "",
                                   documents._OperationType.Read,{}, client_id=self.client_id)
-        request_params = RequestObject("databaseaccount", documents._OperationType.Read, url_connection)
+        request_params = RequestObject("databaseaccount", documents._OperationType.Read, headers, url_connection)
         self.__Get("", request_params, headers, **kwargs)
 
 
@@ -2646,7 +2650,8 @@ def Create(
                                   options)
         # Create will use WriteEndpoint since it uses POST operation
 
-        request_params = RequestObject(typ, documents._OperationType.Create)
+        request_params = RequestObject(typ, documents._OperationType.Create, headers)
+        request_params.set_excluded_location_from_options(options)
         result, last_response_headers = self.__Post(path, request_params, body, headers, **kwargs)
         self.last_response_headers = last_response_headers
 
@@ -2692,7 +2697,8 @@ def Upsert(
         headers[http_constants.HttpHeaders.IsUpsert] = True
 
         # Upsert will use WriteEndpoint since it uses POST operation
-        request_params = RequestObject(typ, documents._OperationType.Upsert)
+        request_params = RequestObject(typ, documents._OperationType.Upsert, headers)
+        request_params.set_excluded_location_from_options(options)
         result, last_response_headers = self.__Post(path, request_params, body, headers, **kwargs)
         self.last_response_headers = last_response_headers
         # update session for write request
@@ -2735,7 +2741,8 @@ def Replace(
         headers = base.GetHeaders(self, initial_headers, "put", path, id, typ, documents._OperationType.Replace,
                                   options)
         # Replace will use WriteEndpoint since it uses PUT operation
-        request_params = RequestObject(typ, documents._OperationType.Replace)
+        request_params = RequestObject(typ, documents._OperationType.Replace, headers)
+        request_params.set_excluded_location_from_options(options)
         result, last_response_headers = self.__Put(path, request_params, resource, headers, **kwargs)
         self.last_response_headers = last_response_headers
 
@@ -2776,7 +2783,8 @@ def Read(
         initial_headers = initial_headers or self.default_headers
         headers = base.GetHeaders(self, initial_headers, "get", path, id, typ, documents._OperationType.Read, options)
         # Read will use ReadEndpoint since it uses GET operation
-        request_params = RequestObject(typ, documents._OperationType.Read)
+        request_params = RequestObject(typ, documents._OperationType.Read, headers)
+        request_params.set_excluded_location_from_options(options)
         result, last_response_headers = self.__Get(path, request_params, headers, **kwargs)
         self.last_response_headers = last_response_headers
         if response_hook:
@@ -2815,7 +2823,8 @@ def DeleteResource(
         headers = base.GetHeaders(self, initial_headers, "delete", path, id, typ, documents._OperationType.Delete,
                                   options)
         # Delete will use WriteEndpoint since it uses DELETE operation
-        request_params = RequestObject(typ, documents._OperationType.Delete)
+        request_params = RequestObject(typ, documents._OperationType.Delete, headers)
+        request_params.set_excluded_location_from_options(options)
         result, last_response_headers = self.__Delete(path, request_params, headers, **kwargs)
         self.last_response_headers = last_response_headers
 
@@ -3047,23 +3056,27 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]:
         initial_headers = self.default_headers.copy()
         # Copy to make sure that default_headers won't be changed.
         if query is None:
+            op_typ = documents._OperationType.QueryPlan if is_query_plan else documents._OperationType.ReadFeed
             # Query operations will use ReadEndpoint even though it uses GET(for feed requests)
-            request_params = RequestObject(
-                resource_type,
-                documents._OperationType.QueryPlan if is_query_plan else documents._OperationType.ReadFeed
-            )
             headers = base.GetHeaders(
                 self,
                 initial_headers,
                 "get",
                 path,
                 resource_id,
                 resource_type,
-                request_params.operation_type,
+                op_typ,
                 options,
                 partition_key_range_id
             )
 
+            request_params = RequestObject(
+                resource_type,
+                op_typ,
+                headers
+            )
+            request_params.set_excluded_location_from_options(options)
+
             change_feed_state: Optional[ChangeFeedState] = options.get("changeFeedState")
             if change_feed_state is not None:
                 change_feed_state.populate_request_headers(self._routing_map_provider, headers)
@@ -3089,7 +3102,6 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]:
             raise SystemError("Unexpected query compatibility mode.")
 
         # Query operations will use ReadEndpoint even though it uses POST(for regular query operations)
-        request_params = RequestObject(resource_type, documents._OperationType.SqlQuery)
         req_headers = base.GetHeaders(
             self,
             initial_headers,
@@ -3102,6 +3114,9 @@ def __GetBodiesFromQueryResult(result: Dict[str, Any]) -> List[Dict[str, Any]]:
             partition_key_range_id
         )
 
+        request_params = RequestObject(resource_type, documents._OperationType.SqlQuery, req_headers)
+        request_params.set_excluded_location_from_options(options)
+
         # check if query has prefix partition key
         isPrefixPartitionQuery = kwargs.pop("isPrefixPartitionQuery", None)
         if isPrefixPartitionQuery:
@@ -3256,7 +3271,7 @@ def _AddPartitionKey(
         options: Mapping[str, Any]
     ) -> Dict[str, Any]:
         collection_link = base.TrimBeginningAndEndingSlashes(collection_link)
-        partitionKeyDefinition = self._get_partition_key_definition(collection_link)
+        partitionKeyDefinition = self._get_partition_key_definition(collection_link, options)
         new_options = dict(options)
         # If the collection doesn't have a partition key definition, skip it as it's a legacy collection
         if partitionKeyDefinition:
@@ -3358,15 +3373,19 @@ def _UpdateSessionIfRequired(
             # update session
             self.session.update_session(response_result, response_headers)
 
-    def _get_partition_key_definition(self, collection_link: str) -> Optional[Dict[str, Any]]:
+    def _get_partition_key_definition(
+            self,
+            collection_link: str,
+            options: Mapping[str, Any]
+    ) -> Optional[Dict[str, Any]]:
         partition_key_definition: Optional[Dict[str, Any]]
         # If the document collection link is present in the cache, then use the cached partitionkey definition
         if collection_link in self.__container_properties_cache:
             cached_container: Dict[str, Any] = self.__container_properties_cache.get(collection_link, {})
             partition_key_definition = cached_container.get("partitionKey")
         # Else read the collection from backend and add it to the cache
         else:
-            container = self.ReadContainer(collection_link)
+            container = self.ReadContainer(collection_link, options)
             partition_key_definition = container.get("partitionKey")
             self.__container_properties_cache[collection_link] = _set_properties_cache(container)
         return partition_key_definition

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

@@ -30,8 +30,10 @@
 
 from . import _constants as constants
 from . import exceptions
+from ._request_object import RequestObject
+from ._routing.routing_range import PartitionKeyRangeWrapper
 from .documents import DatabaseAccount
-from ._location_cache import LocationCache
+from ._location_cache import LocationCache, current_time_millis
 
 
 # pylint: disable=protected-access
@@ -50,19 +52,14 @@ def __init__(self, client):
         self.DefaultEndpoint = client.url_connection
         self.refresh_time_interval_in_ms = self.get_refresh_time_interval_in_ms_stub()
         self.location_cache = LocationCache(
-            self.PreferredLocations,
             self.DefaultEndpoint,
-            self.EnableEndpointDiscovery,
-            client.connection_policy.UseMultipleWriteLocations
+            client.connection_policy
         )
         self.refresh_needed = False
         self.refresh_lock = threading.RLock()
         self.last_refresh_time = 0
         self._database_account_cache = None
 
-    def get_use_multiple_write_locations(self):
-        return self.location_cache.can_use_multiple_write_locations()
-
     def get_refresh_time_interval_in_ms_stub(self):
         return constants._Constants.DefaultEndpointsRefreshTime
 
@@ -72,7 +69,11 @@ def get_write_endpoint(self):
     def get_read_endpoint(self):
         return self.location_cache.get_read_regional_routing_context()
 
-    def resolve_service_endpoint(self, request):
+    def resolve_service_endpoint(
+            self,
+            request: RequestObject,
+            pk_range_wrapper: PartitionKeyRangeWrapper # pylint: disable=unused-argument
+    ) -> str:
         return self.location_cache.resolve_service_endpoint(request)
 
     def mark_endpoint_unavailable_for_read(self, endpoint, refresh_cache):
@@ -98,7 +99,7 @@ def update_location_cache(self):
         self.location_cache.update_location_cache()
 
     def refresh_endpoint_list(self, database_account, **kwargs):
-        if self.location_cache.current_time_millis() - self.last_refresh_time > self.refresh_time_interval_in_ms:
+        if current_time_millis() - self.last_refresh_time > self.refresh_time_interval_in_ms:
             self.refresh_needed = True
         if self.refresh_needed:
             with self.refresh_lock:
@@ -114,11 +115,11 @@ def _refresh_endpoint_list_private(self, database_account=None, **kwargs):
         if database_account:
             self.location_cache.perform_on_database_account_read(database_account)
             self.refresh_needed = False
-            self.last_refresh_time = self.location_cache.current_time_millis()
+            self.last_refresh_time = current_time_millis()
         else:
             if self.location_cache.should_refresh_endpoints() or self.refresh_needed:
                 self.refresh_needed = False
-                self.last_refresh_time = self.location_cache.current_time_millis()
+                self.last_refresh_time = current_time_millis()
                 # this will perform getDatabaseAccount calls to check endpoint health
                 self._endpoints_health_check(**kwargs)