[Internal] Circuit Breaker: Adds Code to Implement Per Partition Circuit Breaker (#5023)

# Pull Request Template

## Description

The idea of having a per partition circuit breaker (aka PPCB) is to
optimize **a)** read availability , in a single master account and
**b)** read + write availability in a multi master account, during a
time when a specific partition in one of the regions is experiencing an
outage/ quorum loss. This feature is independent of the partition level
failover triggered by the backend. The per partition circuit breaker is
developed behind a feature flag `AZURE_COSMOS_CIRCUIT_BREAKER_ENABLED`.

However, when the partition level failover is enabled, we will enable
the PPCB by default so that the reads can benefits from it.

## Scope

- For single master, only the read requests will use the circuit breaker
to add the pk-range to region override mapping, and use this mapping as
a source of truth to route the read requests.

- For multi master, both the read and write requests will use the
circuit breaker to add the pk-range to region override mapping and use
this mapping as a source of truth to route both read and write requests.

## Understanding the Configurations exposed by the environment
variables:

- `AZURE_COSMOS_CIRCUIT_BREAKER_ENABLED`: This environment variable is
used to enable/ disable the partition level circuit breaker feature. The
default value is `false`.

-
`AZURE_COSMOS_PPCB_STALE_PARTITION_UNAVAILABILITY_REFRESH_INTERVAL_IN_SECONDS`:
This environment variable is used to set the background periodic address
refresh task interval. The default value for this interval is `60
seconds`.

-
`AZURE_COSMOS_PPCB_ALLOWED_PARTITION_UNAVAILABILITY_DURATION_IN_SECONDS`:
This environment variable is used to set the partition unavailability
time duration in seconds. The unavailability time indicates how long a
partition can remain unhealthy, before it can re-validate it's
connection status. The default value for this property is `5 seconds`.

- `AZURE_COSMOS_PPCB_CONSECUTIVE_FAILURE_COUNT_FOR_READS`: This
environment variable is used to set the consecutive failure count for
reads, before triggering per partition level circuit breaker flow. The
default value for this flag is `10` consecutive failures within 1 min
window.

- `AZURE_COSMOS_PPCB_CONSECUTIVE_FAILURE_COUNT_FOR_WRITES`: This
environment variable is used to set the consecutive failure count for
writes, before triggering per partition level circuit breaker flow. The
default value for this flag is `5` consecutive failures within 1 min
window.

## Understanding the Working Principle:

On a high level, there are three parts of the circuit breaker
implementation:

- **Short Circuit and Failover detection:** The failover detection logic
will reside in the SDK ClientRetryPolicy, just like we have for PPAF.
Ideally the detection logic is based on the below two principles:

- **Status Codes:** The status codes that are indicative of partition
level circuit breaker would be the following: a) `503` Service
Unavailable, b) `408` Request Timeout, c) cancellation token expired.

- **Threshold:** Once the failover condition is met, the SDK will look
for some consecutive failures, until it hits a particular threshold.
Once this threshold is met, the SDK will fail over the read requests to
the next preferred region for that offending partition. For example, if
the threshold value for read requests is `10`, then the SDK will look
for `10` consecutive failures. If the threshold is met/ exceeded, the
SDK will add the region failover information for that partition.

- **Failover a faulty partition to the next preferred region:** Once the
failover conditions are met, the `ClientRetryPolicy` will trigger a
partition level override using
`GlobalPartitionEndpointManagerCore.TryMarkEndpointUnavailableForPartitionKeyRange`
to the next region in the preferred region list. This failover
information will help the current, as well as the subsequent requests
(reads in single master and both reads and writes in multi master) to
route the request to the next region.

- **Failback the faulty partition to it's original first preferred
region:** With PPAF enabled, ideally the write requests will rely on
403.3 (Write Forbidden) signal to fail the partition back to the primary
write region. However, this is not true for reads. That means SDK
doesn’t have a definitive signal to identify when to initiate a failback
for read requests.
Hence, the idea is to create a background task during the time of read
failover, which will keep track of the pk-range and region mapping. The
task will periodically fetch the address from the gateway address cache
for those pk ranges in the faulty region, and it will try to initiate
Rntbd connection to all 4 replicas of that partition. The RNTBD open
connection attempt will be made similar to that of the replica
validation flow. The life cycle of the background task will get
initiated during a failover and will remain until the SDK is disposed.
If the attempt to make the connection to all 4 replicas is successful,
then the task will remove/ override the entry with the primary region,
resulting the SDK to failback the read requests.


## Type of change

- [x] New feature (non-breaking change which adds functionality)

## Closing issues

To automatically close an issue: closes #4981

Author: kundadebdatta

Microsoft.Azure.Cosmos/src/ClientRetryPolicy.cs

@@ -28,7 +28,7 @@ internal sealed class ClientRetryPolicy : IDocumentClientRetryPolicy
         private readonly GlobalEndpointManager globalEndpointManager;
         private readonly GlobalPartitionEndpointManager partitionKeyRangeLocationCache;
         private readonly bool enableEndpointDiscovery;
-        private readonly bool isPertitionLevelFailoverEnabled;
+        private readonly bool isPartitionLevelFailoverEnabled;
         private int failoverRetryCount;
 
         private int sessionTokenRetryCount;
@@ -45,7 +45,7 @@ public ClientRetryPolicy(
             GlobalPartitionEndpointManager partitionKeyRangeLocationCache,
             RetryOptions retryOptions,
             bool enableEndpointDiscovery,
-            bool isPertitionLevelFailoverEnabled)
+            bool isPartitionLevelFailoverEnabled)
         {
             this.throttlingRetry = new ResourceThrottleRetryPolicy(
                 retryOptions.MaxRetryAttemptsOnThrottledRequests,
@@ -59,7 +59,7 @@ public ClientRetryPolicy(
             this.serviceUnavailableRetryCount = 0;
             this.canUseMultipleWriteLocations = false;
             this.isMultiMasterWriteRequest = false;
-            this.isPertitionLevelFailoverEnabled = isPertitionLevelFailoverEnabled;
+            this.isPartitionLevelFailoverEnabled = isPartitionLevelFailoverEnabled;
         }
 
         /// <summary> 
@@ -80,13 +80,9 @@ public async Task<ShouldRetryResult> ShouldRetryAsync(
                     this.documentServiceRequest?.RequestContext?.LocationEndpointToRoute?.ToString() ?? string.Empty,
                     this.documentServiceRequest?.ResourceAddress ?? string.Empty);
 
-                if (this.isPertitionLevelFailoverEnabled)
-                {
-                    // In the event of the routing gateway having outage on region A, mark the partition as unavailable assuming that the
-                    // partition has been failed over to region B, when per partition automatic failover is enabled.
-                    this.partitionKeyRangeLocationCache.TryMarkEndpointUnavailableForPartitionKeyRange(
-                         this.documentServiceRequest);
-                }
+                // In the event of the routing gateway having outage on region A, mark the partition as unavailable assuming that the
+                // partition has been failed over to region B, when per partition automatic failover is enabled.
+                this.TryMarkEndpointUnavailableForPkRange(isSystemResourceUnavailableForWrite: false);
 
                 // Mark both read and write requests because it gateway exception.
                 // This means all requests going to the region will fail.
@@ -113,7 +109,7 @@ public async Task<ShouldRetryResult> ShouldRetryAsync(
                         StatusCodes.TooManyRequests, SubStatusCodes.SystemResourceUnavailable);
 
                     return this.TryMarkEndpointUnavailableForPkRangeAndRetryOnServiceUnavailable(
-                        shouldMarkEndpointUnavailableForPkRange: true);
+                        isSystemResourceUnavailableForWrite: true);
                 }
 
                 ShouldRetryResult shouldRetryResult = await this.ShouldRetryInternalAsync(
@@ -176,7 +172,7 @@ public async Task<ShouldRetryResult> ShouldRetryAsync(
                     StatusCodes.TooManyRequests, SubStatusCodes.SystemResourceUnavailable);
 
                 return this.TryMarkEndpointUnavailableForPkRangeAndRetryOnServiceUnavailable(
-                    shouldMarkEndpointUnavailableForPkRange: true);
+                    isSystemResourceUnavailableForWrite: true);
             }
 
             return await this.throttlingRetry.ShouldRetryAsync(cosmosResponseMessage, cancellationToken);
@@ -236,8 +232,7 @@ private async Task<ShouldRetryResult> ShouldRetryInternalAsync(
                     this.documentServiceRequest?.ResourceAddress ?? string.Empty);
 
                 // Mark the partition key range as unavailable to retry future request on a new region.
-                this.partitionKeyRangeLocationCache.TryMarkEndpointUnavailableForPartitionKeyRange(
-                     this.documentServiceRequest);
+                this.TryMarkEndpointUnavailableForPkRange(isSystemResourceUnavailableForWrite: false);
             }
 
             // Received 403.3 on write region, initiate the endpoint rediscovery
@@ -313,7 +308,7 @@ private async Task<ShouldRetryResult> ShouldRetryInternalAsync(
             if (statusCode == HttpStatusCode.ServiceUnavailable)
             {
                 return this.TryMarkEndpointUnavailableForPkRangeAndRetryOnServiceUnavailable(
-                    shouldMarkEndpointUnavailableForPkRange: true);
+                    isSystemResourceUnavailableForWrite: false);
             }
 
             return null;
@@ -442,23 +437,18 @@ private ShouldRetryResult ShouldRetryOnSessionNotAvailable(DocumentServiceReques
         /// Service Unavailable response is received, indicating that the service might be temporarily unavailable.
         /// It optionally marks the partition key range as unavailable, which will influence future routing decisions.
         /// </summary>
-        /// <param name="shouldMarkEndpointUnavailableForPkRange">A boolean flag indicating whether the endpoint for the
-        /// current partition key range should be marked as unavailable.</param>
+        /// <param name="isSystemResourceUnavailableForWrite">A boolean flag indicating whether the endpoint for the
+        /// current partition key range should be marked as unavailable, if the failure happened due to system
+        /// resource unavailability.</param>
         /// <returns>An instance of <see cref="ShouldRetryResult"/> indicating whether the operation should be retried.</returns>
         private ShouldRetryResult TryMarkEndpointUnavailableForPkRangeAndRetryOnServiceUnavailable(
-            bool shouldMarkEndpointUnavailableForPkRange)
+            bool isSystemResourceUnavailableForWrite)
         {
             DefaultTrace.TraceWarning("ClientRetryPolicy: ServiceUnavailable. Refresh cache and retry. Failed Location: {0}; ResourceAddress: {1}",
                 this.documentServiceRequest?.RequestContext?.LocationEndpointToRoute?.ToString() ?? string.Empty,
                 this.documentServiceRequest?.ResourceAddress ?? string.Empty);
 
-            if (shouldMarkEndpointUnavailableForPkRange)
-            {
-                // Mark the partition as unavailable.
-                // Let the ClientRetry logic decide if the request should be retried
-                this.partitionKeyRangeLocationCache.TryMarkEndpointUnavailableForPartitionKeyRange(
-                     this.documentServiceRequest);
-            }
+            this.TryMarkEndpointUnavailableForPkRange(isSystemResourceUnavailableForWrite);
 
             return this.ShouldRetryOnServiceUnavailable();
         }
@@ -477,7 +467,7 @@ private ShouldRetryResult ShouldRetryOnServiceUnavailable()
 
             if (!this.canUseMultipleWriteLocations
                     && !this.isReadRequest
-                    && !this.isPertitionLevelFailoverEnabled)
+                    && !this.isPartitionLevelFailoverEnabled)
             {
                 // Write requests on single master cannot be retried if partition level failover is disabled.
                 // This means there are no other regions available to serve the writes.
@@ -506,6 +496,30 @@ private ShouldRetryResult ShouldRetryOnServiceUnavailable()
             return ShouldRetryResult.RetryAfter(TimeSpan.Zero);
         }
 
+        /// <summary>
+        /// Attempts to mark the endpoint associated with the current partition key range as unavailable
+        /// which will influence future routing decisions.
+        /// </summary>
+        /// <param name="isSystemResourceUnavailableForWrite">A boolean flag indicating if the system resource was unavailable. If true,
+        /// the endpoint will be marked unavailable for the pk-range of a multi master write request, bypassing the circuit breaker check.</param>
+        /// <returns>A boolean flag indicating whether the endpoint was marked as unavailable.</returns>
+        private bool TryMarkEndpointUnavailableForPkRange(
+            bool isSystemResourceUnavailableForWrite)
+        {
+            if (this.documentServiceRequest != null
+                && (isSystemResourceUnavailableForWrite
+                || this.IsRequestEligibleForPerPartitionAutomaticFailover()
+                || this.IsRequestEligibleForPartitionLevelCircuitBreaker()))
+            {
+                // Mark the partition as unavailable.
+                // Let the ClientRetry logic decide if the request should be retried
+                return this.partitionKeyRangeLocationCache.TryMarkEndpointUnavailableForPartitionKeyRange(
+                    request: this.documentServiceRequest);
+            }
+
+            return false;
+        }
+
         /// <summary>
         /// Returns a boolean flag indicating if the endpoint should be marked as unavailable
         /// due to a 429 response with a sub status code of 3092 (system resource unavailable).
@@ -524,6 +538,32 @@ private bool ShouldMarkEndpointUnavailableOnSystemResourceUnavailableForWrite(
                 && subStatusCode == SubStatusCodes.SystemResourceUnavailable;
         }
 
+        /// <summary>
+        /// Determines if a request is eligible for per-partition automatic failover.
+        /// A request is eligible if it is a write request, partition level failover is enabled,
+        /// and the global endpoint manager cannot use multiple write locations for the request.
+        /// </summary>
+        /// <returns>True if the request is eligible for per-partition automatic failover, otherwise false.</returns>
+        private bool IsRequestEligibleForPerPartitionAutomaticFailover()
+        {
+            return this.partitionKeyRangeLocationCache.IsRequestEligibleForPerPartitionAutomaticFailover(
+                this.documentServiceRequest);
+        }
+
+        /// <summary>
+        /// Determines if a request is eligible for partition-level circuit breaker.
+        /// This method checks if the request is a read-only request or a multi master write request, if partition-level circuit breaker is enabled,
+        /// and if the partition key range location cache indicates that the partition can fail over based on the number of request failures.
+        /// </summary>
+        /// <returns>
+        /// True if the read request is eligible for partition-level circuit breaker, otherwise false.
+        /// </returns>
+        private bool IsRequestEligibleForPartitionLevelCircuitBreaker()
+        {
+            return this.partitionKeyRangeLocationCache.IsRequestEligibleForPartitionLevelCircuitBreaker(this.documentServiceRequest)
+                        && this.partitionKeyRangeLocationCache.IncrementRequestFailureCounterAndCheckIfPartitionCanFailover(this.documentServiceRequest);
+        }
+
         private sealed class RetryContext
         {
             public int RetryLocationIndex { get; set; }

Microsoft.Azure.Cosmos/src/ConnectionPolicy.cs

@@ -332,6 +332,12 @@ public bool EnablePartitionLevelFailover
             set;
         }
 
+        public bool EnablePartitionLevelCircuitBreaker
+        {
+            get;
+            set;
+        }
+
         /// <summary>
         /// Gets or sets the certificate validation callback.
         /// </summary>

Microsoft.Azure.Cosmos/src/CosmosClientOptions.cs

@@ -729,7 +729,14 @@ public Func<HttpClient> HttpClientFactory
         /// <summary>
         /// Enable partition key level failover
         /// </summary>
-        internal bool EnablePartitionLevelFailover { get; set; } = ConfigurationManager.IsPartitionLevelFailoverEnabled(defaultValue: false);
+        internal bool EnablePartitionLevelFailover { get; set; } = ConfigurationManager.IsPartitionLevelFailoverEnabled(defaultValue: false);
+
+        /// <summary>
+        /// Enable partition level circuit breaker (aka PPCB). For compute gateway use case, by default per partition automatic failover will be disabled, so does the PPCB.
+        /// If compute gateway chooses to enable PPAF, then the .NET SDK will enable PPCB by default, which will improve the read availability and latency. This would mean
+        /// when PPAF is enabled, the SDK will automatically enable PPCB as well.
+        /// </summary>
+        internal bool EnablePartitionLevelCircuitBreaker { get; set; } = ConfigurationManager.IsPartitionLevelCircuitBreakerEnabled(defaultValue: false);
 
         /// <summary>
         /// Quorum Read allowed with eventual consistency account or consistent prefix account.
@@ -983,6 +990,7 @@ internal virtual ConnectionPolicy GetConnectionPolicy(int clientId)
                 MaxTcpConnectionsPerEndpoint = this.MaxTcpConnectionsPerEndpoint,
                 EnableEndpointDiscovery = !this.LimitToEndpoint,
                 EnablePartitionLevelFailover = this.EnablePartitionLevelFailover,
+                EnablePartitionLevelCircuitBreaker = this.EnablePartitionLevelFailover || this.EnablePartitionLevelCircuitBreaker,
                 PortReuseMode = this.portReuseMode,
                 EnableTcpConnectionEndpointRediscovery = this.EnableTcpConnectionEndpointRediscovery,
                 EnableAdvancedReplicaSelectionForTcp = this.EnableAdvancedReplicaSelectionForTcp,
@@ -1221,6 +1229,11 @@ internal string GetUserAgentSuffix()
                 featureFlag += (int)UserAgentFeatureFlags.PerPartitionAutomaticFailover;
             }
 
+            if (this.EnablePartitionLevelFailover || this.EnablePartitionLevelCircuitBreaker)
+            {
+                featureFlag += (int)UserAgentFeatureFlags.PerPartitionCircuitBreaker;
+            }
+
             if (featureFlag == 0)
             {
                 return this.ApplicationName;

Microsoft.Azure.Cosmos/src/Diagnostics/UserAgentFeatureFlags.cs

@@ -16,5 +16,7 @@ namespace Microsoft.Azure.Cosmos
     internal enum UserAgentFeatureFlags
     {
         PerPartitionAutomaticFailover = 1,
+
+        PerPartitionCircuitBreaker = 2,
     }
 }

Microsoft.Azure.Cosmos/src/DocumentClient.cs

@@ -939,8 +939,11 @@ internal virtual void Initialize(Uri serviceEndpoint,
 #endif
 
             this.GlobalEndpointManager = new GlobalEndpointManager(this, this.ConnectionPolicy);
-            this.PartitionKeyRangeLocation = this.ConnectionPolicy.EnablePartitionLevelFailover
-                ? new GlobalPartitionEndpointManagerCore(this.GlobalEndpointManager)
+            this.PartitionKeyRangeLocation = this.ConnectionPolicy.EnablePartitionLevelFailover || this.ConnectionPolicy.EnablePartitionLevelCircuitBreaker
+                ? new GlobalPartitionEndpointManagerCore(
+                    this.GlobalEndpointManager,
+                    this.ConnectionPolicy.EnablePartitionLevelFailover,
+                    this.ConnectionPolicy.EnablePartitionLevelCircuitBreaker)
                 : GlobalPartitionEndpointManagerNoOp.Instance;
 
             this.httpClient = CosmosHttpClientCore.CreateWithConnectionPolicy(

Microsoft.Azure.Cosmos/src/RetryPolicy.cs

@@ -13,7 +13,7 @@ internal sealed class RetryPolicy : IRetryPolicyFactory
         private readonly GlobalPartitionEndpointManager partitionKeyRangeLocationCache;
         private readonly GlobalEndpointManager globalEndpointManager;
         private readonly bool enableEndpointDiscovery;
-        private readonly bool isPertitionLevelFailoverEnabled;
+        private readonly bool isPartitionLevelFailoverEnabled;
         private readonly RetryOptions retryOptions;
 
         /// <summary>
@@ -25,7 +25,7 @@ public RetryPolicy(
             GlobalPartitionEndpointManager partitionKeyRangeLocationCache)
         {
             this.enableEndpointDiscovery = connectionPolicy.EnableEndpointDiscovery;
-            this.isPertitionLevelFailoverEnabled = connectionPolicy.EnablePartitionLevelFailover;
+            this.isPartitionLevelFailoverEnabled = connectionPolicy.EnablePartitionLevelFailover;
             this.globalEndpointManager = globalEndpointManager;
             this.retryOptions = connectionPolicy.RetryOptions;
             this.partitionKeyRangeLocationCache = partitionKeyRangeLocationCache;
@@ -41,7 +41,7 @@ public IDocumentClientRetryPolicy GetRequestPolicy()
                 this.partitionKeyRangeLocationCache,
                 this.retryOptions,
                 this.enableEndpointDiscovery,
-                this.isPertitionLevelFailoverEnabled);
+                this.isPartitionLevelFailoverEnabled);
 
             return clientRetryPolicy;
         }