Metadata requests: Fixes bug where certain metadata requests are not retried with a client cold start with only query requests (#5108)

# Pull Request Template

## Description

### Context 

Currently, there is a bug in the SDK where upon a cold start of the SDK
and rare edge cases involving online/offline-ing regions, where only
query requests are made, the SDK will not retry certain status code
responses from metadata requests causing the entire request to fail. The
correct behavior would be for the SDK to do cross region retries on
these metadata requests.

This pulls request includes several updates to enhance error handling
and retry logic in the Cosmos DB SDK. The changes mainly focus on
extending support for additional server error types and improving retry
policies for various scenarios.

### Improvements to retry logic:

*
[`ClientRetryPolicy.cs`](diffhunk://#diff-2b056512ca285b1d95e025e31f60345059fa92d958becc38f90a6fb54ce1bbb4R331-R341):
Enhanced retry logic to handle `InternalServerError` ,
`DatabaseAccountNotFound`, and `LeaseNotFound` status codes.
*
[`MetadataRequestThrottleRetryPolicy.cs`](diffhunk://#diff-a5ed5985909c3dcb6e4ce186cdd662d590dac5297ea14e68560c7d1eca307be4L26-R28):
Refactored retry policy to handle additional status codes and renamed
methods and constants to reflect the broader scope of endpoint
unavailability.

### FaultInjection enhancements to error handling testing:

*
[`FaultInjectionRuleBuilder.cs`](diffhunk://#diff-d827164a4a6a0d8737e6598f8132c915ef48a1fc01daaa6422706f770dada5d5L152-R156):
Added support for additional server error types such as
`DatabaseAccountNotFound`, `ServiceUnavailable`, `InternalServerError`,
and `LeaseNotFound` in the for metadata requests.
*
[`FaultInjectionServerErrorType.cs`](diffhunk://#diff-0c89faa9a48c428a7a98662d995474e34295618ac60e677ad9762fd048f33601L75-R82):
Updated the `FaultInjectionServerErrorType` enum to include
`LeaseNotFound` and corrected the status code for
`DatabaseAccountNotFound`.
*
[`FaultInjectionServerErrorResultInternal.cs`](diffhunk://#diff-1ae8256c6d505a8f3b0a350978a0cc9a08f6234f7328f28f1db14302ee691d72L473-R473):
Added handling for `LeaseNotFound` and updated the status code for
`DatabaseAccountNotFound` in the `GetInjectedServerError` method.
* 
### Testing updates:

*
[`CosmosItemIntegrationTests.cs`](diffhunk://#diff-16d429adf686a32936696d2014afab3fc8faf91f10c880850fb8b30f8b96bb33R153-R214):
Added a new test method
`MetadataEndpointUnavailableCrossRegionalRetryTest` to validate the
retry logic for various server error types.
*
[`ClientRetryPolicyTests.cs`](diffhunk://#diff-d3fdfdc5d4f4d8af2c2cc463d928285680e7695861422ef2e3330d1a956807e1L167-R176):
Extended existing tests to cover additional status codes and substatus
codes.

## Type of change

Please delete options that are not relevant.

- [] Bug fix (non-breaking change which fixes an issue)

## Closing issues

To automatically close an issue: closes #4710

---------

Co-authored-by: Kiran Kumar Kolli <kirankk@microsoft.com>

Author: NaluTripician

Microsoft.Azure.Cosmos/FaultInjection/src/FaultInjectionRuleBuilder.cs

@@ -149,7 +149,11 @@ private void ValidateGatewayConnection()
             {
                 if (serverErrorResult?.GetServerErrorType() != FaultInjectionServerErrorType.TooManyRequests
                     && serverErrorResult?.GetServerErrorType() != FaultInjectionServerErrorType.ResponseDelay
-                    && serverErrorResult?.GetServerErrorType() != FaultInjectionServerErrorType.SendDelay)
+                    && serverErrorResult?.GetServerErrorType() != FaultInjectionServerErrorType.SendDelay
+                    && serverErrorResult?.GetServerErrorType() != FaultInjectionServerErrorType.DatabaseAccountNotFound
+                    && serverErrorResult?.GetServerErrorType() != FaultInjectionServerErrorType.ServiceUnavailable
+                    && serverErrorResult?.GetServerErrorType() != FaultInjectionServerErrorType.InternalServerError
+                    && serverErrorResult?.GetServerErrorType() != FaultInjectionServerErrorType.LeaseNotFound)
                 {
                     throw new ArgumentException($"{serverErrorResult?.GetServerErrorType()} is not supported for metadata requests.");
                 }

Microsoft.Azure.Cosmos/FaultInjection/src/FaultInjectionServerErrorType.cs

@@ -72,8 +72,13 @@ public enum FaultInjectionServerErrorType
         ServiceUnavailable,
 
         /// <summary>
-        /// 404:1008 Database account not found from gateway
+        /// 403:1008 Database account not found from gateway
         /// </summary>
         DatabaseAccountNotFound,
+
+        /// <summary>
+        /// 410:1022 Lease not Found
+        /// </summary>
+        LeaseNotFound,
     }
 }

Microsoft.Azure.Cosmos/FaultInjection/src/implementation/FaultInjectionServerErrorResultInternal.cs

@@ -461,7 +461,7 @@ public HttpResponseMessage GetInjectedServerError(DocumentServiceRequest dsr, st
 
                     httpResponse.Headers.Add(
                         WFConstants.BackendHeaders.SubStatus,
-                        ((int)SubStatusCodes.RUBudgetExceeded).ToString(CultureInfo.InvariantCulture));
+                        ((int)SubStatusCodes.Unknown).ToString(CultureInfo.InvariantCulture));
                     httpResponse.Headers.Add(WFConstants.BackendHeaders.LocalLSN, lsn);
 
                     return httpResponse;
@@ -470,7 +470,7 @@ public HttpResponseMessage GetInjectedServerError(DocumentServiceRequest dsr, st
 
                     httpResponse = new HttpResponseMessage
                     {
-                        StatusCode = HttpStatusCode.NotFound,
+                        StatusCode = HttpStatusCode.Forbidden,
                         Content = new FauntInjectionHttpContent(
                             new MemoryStream(
                                 FaultInjectionResponseEncoding.GetBytes($"Fault Injection Server Error: DatabaseAccountNotFound, rule: {ruleId}"))),
@@ -488,6 +488,28 @@ public HttpResponseMessage GetInjectedServerError(DocumentServiceRequest dsr, st
 
                     return httpResponse;
 
+                case FaultInjectionServerErrorType.LeaseNotFound:
+
+                    httpResponse = new HttpResponseMessage
+                    {
+                        StatusCode = HttpStatusCode.Gone,
+                        Content = new FauntInjectionHttpContent(
+                            new MemoryStream(
+                                FaultInjectionResponseEncoding.GetBytes($"Fault Injection Server Error: LeaseNotFound, rule: {ruleId}"))),
+                    };
+
+                    foreach (string header in headers.AllKeys())
+                    {
+                        httpResponse.Headers.Add(header, headers.Get(header));
+                    }
+
+                    httpResponse.Headers.Add(
+                        WFConstants.BackendHeaders.SubStatus,
+                        ((int)SubStatusCodes.LeaseNotFound).ToString(CultureInfo.InvariantCulture));
+                    httpResponse.Headers.Add(WFConstants.BackendHeaders.LocalLSN, lsn);
+
+                    return httpResponse;
+
                 default:
                     throw new ArgumentException($"Server error type {this.serverErrorType} is not supported");
             }

Microsoft.Azure.Cosmos/src/ClientRetryPolicy.cs

@@ -291,11 +291,12 @@ private async Task<ShouldRetryResult> ShouldRetryInternalAsync(
                     this.documentServiceRequest?.RequestContext?.LocationEndpointToRoute?.ToString() ?? string.Empty,
                     this.documentServiceRequest?.ResourceAddress ?? string.Empty);
 
+                //Retry policy will retry on the next preffered region as the original requert region is not accepting requests
                 return await this.ShouldRetryOnEndpointFailureAsync(
                     isReadRequest: this.isReadRequest,
                     markBothReadAndWriteAsUnavailable: false,
                     forceRefresh: false,
-                    retryOnPreferredLocations: false);
+                    retryOnPreferredLocations: true);
             }
 
             if (statusCode == HttpStatusCode.NotFound
@@ -311,6 +312,13 @@ private async Task<ShouldRetryResult> ShouldRetryInternalAsync(
                     isSystemResourceUnavailableForWrite: false);
             }
 
+            // Recieved 500 status code or lease not found
+            if ((statusCode == HttpStatusCode.InternalServerError && this.isReadRequest)
+                || (statusCode == HttpStatusCode.Gone && subStatusCode == SubStatusCodes.LeaseNotFound))
+            {
+                return this.ShouldRetryOnUnavailableEndpointStatusCodes();
+            }
+
             return null;
         }
 
@@ -450,14 +458,15 @@ private ShouldRetryResult TryMarkEndpointUnavailableForPkRangeAndRetryOnServiceU
 
             this.TryMarkEndpointUnavailableForPkRange(isSystemResourceUnavailableForWrite);
 
-            return this.ShouldRetryOnServiceUnavailable();
+            return this.ShouldRetryOnUnavailableEndpointStatusCodes();
         }
 
         /// <summary>
         /// For a ServiceUnavailable (503.0) we could be having a timeout from Direct/TCP locally or a request to Gateway request with a similar response due to an endpoint not yet available.
         /// We try and retry the request only if there are other regions available. The retry logic is applicable for single master write accounts as well.
+        /// Other status codes include InternalServerError (500.0) and LeaseNotFound (410.1022).
         /// </summary>
-        private ShouldRetryResult ShouldRetryOnServiceUnavailable()
+        private ShouldRetryResult ShouldRetryOnUnavailableEndpointStatusCodes()
         {
             if (this.serviceUnavailableRetryCount++ >= ClientRetryPolicy.MaxServiceUnavailableRetryCount)
             {

Microsoft.Azure.Cosmos/src/MetadataRequestThrottleRetryPolicy.cs

@@ -23,9 +23,9 @@ internal sealed class MetadataRequestThrottleRetryPolicy : IDocumentClientRetryP
         private const int DefaultMaxWaitTimeInSeconds = 60;
 
         /// <summary>
-        /// A constant integer defining the default maximum retry count on service unavailable.
+        /// A constant integer defining the default maximum retry count on unavailable endpoint.
         /// </summary>
-        private const int DefaultMaxServiceUnavailableRetryCount = 1;
+        private const int DefaultMaxUnavailableEndpointRetryCount = 1;
 
         /// <summary>
         /// An instance of <see cref="IGlobalEndpointManager"/>.
@@ -38,9 +38,9 @@ internal sealed class MetadataRequestThrottleRetryPolicy : IDocumentClientRetryP
         private readonly IDocumentClientRetryPolicy throttlingRetryPolicy;
 
         /// <summary>
-        /// An integer defining the maximum retry count on service unavailable.
+        /// An integer defining the maximum retry count on unavailable endpoint.
         /// </summary>
-        private readonly int maxServiceUnavailableRetryCount;
+        private readonly int maxUnavailableEndpointRetryCount;
 
         /// <summary>
         /// An instance of <see cref="Uri"/> containing the location endpoint where the partition key
@@ -49,9 +49,9 @@ internal sealed class MetadataRequestThrottleRetryPolicy : IDocumentClientRetryP
         private MetadataRetryContext retryContext;
 
         /// <summary>
-        /// An integer capturing the current retry count on service unavailable.
+        /// An integer capturing the current retry count on unavailable endpoint.
         /// </summary>
-        private int serviceUnavailableRetryCount;
+        private int unavailableEndpointRetryCount;
 
         /// <summary>
         /// The constructor to initialize an instance of <see cref="MetadataRequestThrottleRetryPolicy"/>.
@@ -66,8 +66,8 @@ public MetadataRequestThrottleRetryPolicy(
             int maxRetryWaitTimeInSeconds = DefaultMaxWaitTimeInSeconds)
         {
             this.globalEndpointManager = endpointManager;
-            this.maxServiceUnavailableRetryCount = Math.Max(
-                MetadataRequestThrottleRetryPolicy.DefaultMaxServiceUnavailableRetryCount,
+            this.maxUnavailableEndpointRetryCount = Math.Max(
+                MetadataRequestThrottleRetryPolicy.DefaultMaxUnavailableEndpointRetryCount,
                 this.globalEndpointManager.PreferredLocationCount);
 
             this.throttlingRetryPolicy = new ResourceThrottleRetryPolicy(
@@ -91,11 +91,43 @@ public Task<ShouldRetryResult> ShouldRetryAsync(
             Exception exception,
             CancellationToken cancellationToken)
         {
-            if (exception is CosmosException cosmosException
-                && cosmosException.StatusCode == HttpStatusCode.ServiceUnavailable
-                && cosmosException.Headers.SubStatusCode == SubStatusCodes.TransportGenerated503)
+            if (exception is CosmosException cosmosException)
             {
-                if (this.IncrementRetryIndexOnServiceUnavailableForMetadataRead())
+                return this.ShouldRetryInternalAsync(
+                    cosmosException.StatusCode, 
+                    (SubStatusCodes)cosmosException.SubStatusCode,
+                    exception, 
+                    cancellationToken);
+            }
+
+            if (exception is DocumentClientException clientException)
+            {
+                return this.ShouldRetryInternalAsync(
+                    clientException.StatusCode,
+                    clientException.GetSubStatus(),
+                    exception, cancellationToken);
+            }
+
+            return this.throttlingRetryPolicy.ShouldRetryAsync(exception, cancellationToken);
+        }
+
+        private Task<ShouldRetryResult> ShouldRetryInternalAsync(
+            HttpStatusCode? statusCode, 
+            SubStatusCodes subStatus,
+            Exception exception, 
+            CancellationToken cancellationToken)
+        {
+            if (statusCode == null)
+            {
+                return this.throttlingRetryPolicy.ShouldRetryAsync(exception, cancellationToken);
+            }
+
+            if (statusCode == HttpStatusCode.ServiceUnavailable 
+                || statusCode == HttpStatusCode.InternalServerError
+                || (statusCode == HttpStatusCode.Gone && subStatus == SubStatusCodes.LeaseNotFound)
+                || (statusCode == HttpStatusCode.Forbidden && subStatus == SubStatusCodes.DatabaseAccountNotFound))
+            {
+                if (this.IncrementRetryIndexOnUnavailableEndpointForMetadataRead())
                 {
                     return Task.FromResult(ShouldRetryResult.RetryAfter(TimeSpan.Zero));
                 }
@@ -114,16 +146,36 @@ public Task<ShouldRetryResult> ShouldRetryAsync(
             ResponseMessage cosmosResponseMessage,
             CancellationToken cancellationToken)
         {
-            if (cosmosResponseMessage?.StatusCode == HttpStatusCode.ServiceUnavailable
-                && cosmosResponseMessage?.Headers?.SubStatusCode == SubStatusCodes.TransportGenerated503)
+            return this.ShouldRetryInternalAsync(
+                cosmosResponseMessage.StatusCode,
+                (SubStatusCodes)Convert.ToInt32(cosmosResponseMessage.Headers[WFConstants.BackendHeaders.SubStatus]),
+                cosmosResponseMessage,
+                cancellationToken);
+        }
+
+        private Task<ShouldRetryResult> ShouldRetryInternalAsync(
+            HttpStatusCode? statusCode,
+            SubStatusCodes subStatus,
+            ResponseMessage responseMessage,
+            CancellationToken cancellationToken)
+        {
+            if (statusCode == null)
+            {
+                return this.throttlingRetryPolicy.ShouldRetryAsync(responseMessage, cancellationToken);
+            }
+
+            if (statusCode == HttpStatusCode.ServiceUnavailable 
+                || statusCode == HttpStatusCode.InternalServerError
+                || (statusCode == HttpStatusCode.Gone && subStatus == SubStatusCodes.LeaseNotFound)
+                || (statusCode == HttpStatusCode.Forbidden && subStatus == SubStatusCodes.DatabaseAccountNotFound))
             {
-                if (this.IncrementRetryIndexOnServiceUnavailableForMetadataRead())
+                if (this.IncrementRetryIndexOnUnavailableEndpointForMetadataRead())
                 {
                     return Task.FromResult(ShouldRetryResult.RetryAfter(TimeSpan.Zero));
                 }
             }
 
-            return this.throttlingRetryPolicy.ShouldRetryAsync(cosmosResponseMessage, cancellationToken);
+            return this.throttlingRetryPolicy.ShouldRetryAsync(responseMessage, cancellationToken);
         }
 
         /// <summary>
@@ -146,23 +198,23 @@ public void OnBeforeSendRequest(DocumentServiceRequest request)
         }
 
         /// <summary>
-        /// Increments the location index when a service unavailable exception ocurrs, for any future read requests.
+        /// Increments the location index when a unavailable endpoint exception ocurrs, for any future read requests.
         /// </summary>
         /// <returns>A boolean flag indicating if the operation was successful.</returns>
-        private bool IncrementRetryIndexOnServiceUnavailableForMetadataRead()
+        private bool IncrementRetryIndexOnUnavailableEndpointForMetadataRead()
         {
-            if (this.serviceUnavailableRetryCount++ >= this.maxServiceUnavailableRetryCount)
+            if (this.unavailableEndpointRetryCount++ >= this.maxUnavailableEndpointRetryCount)
             {
-                DefaultTrace.TraceWarning("MetadataRequestThrottleRetryPolicy: Retry count: {0} has exceeded the maximum permitted retry count on service unavailable: {1}.", this.serviceUnavailableRetryCount, this.maxServiceUnavailableRetryCount);
+                DefaultTrace.TraceWarning("MetadataRequestThrottleRetryPolicy: Retry count: {0} has exceeded the maximum permitted retry count on unavailable endpoint: {1}.", this.unavailableEndpointRetryCount, this.maxUnavailableEndpointRetryCount);
                 return false;
             }
 
             // Retrying on second PreferredLocations.
             // RetryCount is used as zero-based index.
-            DefaultTrace.TraceWarning("MetadataRequestThrottleRetryPolicy: Incrementing the metadata retry location index to: {0}.", this.serviceUnavailableRetryCount);
+            DefaultTrace.TraceWarning("MetadataRequestThrottleRetryPolicy: Incrementing the metadata retry location index to: {0}.", this.unavailableEndpointRetryCount);
             this.retryContext = new MetadataRetryContext()
             {
-                RetryLocationIndex = this.serviceUnavailableRetryCount,
+                RetryLocationIndex = this.unavailableEndpointRetryCount,
                 RetryRequestOnPreferredLocations = true,
             };
 

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemIntegrationTests.cs

@@ -8,6 +8,7 @@
     using System.Text.Json.Serialization;
     using System.Threading;
     using System.Threading.Tasks;
+    using Microsoft.Azure.Cosmos.Diagnostics;
     using Microsoft.Azure.Cosmos.FaultInjection;
     using Microsoft.VisualStudio.TestTools.UnitTesting;
     using static Microsoft.Azure.Cosmos.SDK.EmulatorTests.MultiRegionSetupHelpers;
@@ -147,6 +148,85 @@ public async Task ReadMany2UnreachablePartitionsTest()
             }
         }
 
+        [TestMethod]
+        [TestCategory("MultiRegion")]
+        [DataRow(FaultInjectionServerErrorType.ServiceUnavailable)]
+        [DataRow(FaultInjectionServerErrorType.InternalServerError)]
+        [DataRow(FaultInjectionServerErrorType.DatabaseAccountNotFound)]
+        [DataRow(FaultInjectionServerErrorType.LeaseNotFound)]
+        public async Task MetadataEndpointUnavailableCrossRegionalRetryTest(FaultInjectionServerErrorType serverErrorType)
+        {
+            FaultInjectionRule collReadBad = new FaultInjectionRuleBuilder(
+                id: "collread",
+                condition: new FaultInjectionConditionBuilder()
+                    .WithOperationType(FaultInjectionOperationType.MetadataContainer)
+                    .WithRegion(region1)
+                    .Build(),
+                result: new FaultInjectionServerErrorResultBuilder(serverErrorType)
+                    .Build())
+                .Build();
+
+            FaultInjectionRule pkRangeBad = new FaultInjectionRuleBuilder(
+                id: "pkrange",
+                condition: new FaultInjectionConditionBuilder()
+                    .WithOperationType(FaultInjectionOperationType.MetadataPartitionKeyRange)
+                    .WithRegion(region1)
+                    .Build(),
+                result: new FaultInjectionServerErrorResultBuilder(serverErrorType)
+                    .Build())
+                .Build();
+
+            collReadBad.Disable();
+            pkRangeBad.Disable();
+
+            FaultInjector faultInjector = new FaultInjector(new List<FaultInjectionRule> { pkRangeBad, collReadBad });
+
+            CosmosClientOptions cosmosClientOptions = new CosmosClientOptions()
+            {
+                ConsistencyLevel = ConsistencyLevel.Session,
+                ConnectionMode = ConnectionMode.Direct,
+                Serializer = this.cosmosSystemTextJsonSerializer,
+                FaultInjector = faultInjector,
+                ApplicationPreferredRegions = new List<string> { region1, region2,  region3 }
+            };
+
+            using (CosmosClient fiClient = new CosmosClient(
+                connectionString: this.connectionString,
+                clientOptions: cosmosClientOptions))
+            {
+                Database fidb = fiClient.GetDatabase(MultiRegionSetupHelpers.dbName);
+                Container fic = fidb.GetContainer(MultiRegionSetupHelpers.containerName);
+
+                pkRangeBad.Enable();
+                collReadBad.Enable();
+
+                try
+                {
+                    FeedIterator<CosmosIntegrationTestObject> frTest = fic.GetItemQueryIterator<CosmosIntegrationTestObject>("SELECT * FROM c");
+                    while (frTest.HasMoreResults)
+                    {
+                        FeedResponse<CosmosIntegrationTestObject> feedres = await frTest.ReadNextAsync();
+
+                        Assert.AreEqual(HttpStatusCode.OK, feedres.StatusCode);
+                    }
+                }
+                catch (CosmosException ex)
+                {
+                    Assert.Fail(ex.Message);
+                }
+                finally
+                {
+                    //Cross regional retry needs to ocur (could trigger for other metadata call to try on secondary region so rule would not trigger)
+                    Assert.IsTrue(pkRangeBad.GetHitCount() + collReadBad.GetHitCount() >= 1);
+
+                    pkRangeBad.Disable();
+                    collReadBad.Disable();
+
+                    fiClient.Dispose();
+                }
+            }
+        }
+
         [TestMethod]
         [TestCategory("MultiRegion")]
         public async Task AddressRefreshTimeoutTest()