HttpTimeoutPolicy Improvements Phase 2: Refactors Code to Separate out Retry Policy Timeouts for point-reads and non-point-reads on PPAF (#5482)

# Pull Request Template

## Description
This change fixes the aggressive retry timeouts for requests when PPAF
enabled. This change contains distinct retry policy timeouts for
respective operations and allows us to manage the timeouts separately.

## Type of change

Please delete options that are not relevant.

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


## Closing issues

To automatically close an issue: closes #5484

---------

Co-authored-by: Debdatta Kunda <dkunda@microsoft.com>
Co-authored-by: Debdatta Kunda <87335885+kundadebdatta@users.noreply.github.com>

Author: Praveen-Msft

Microsoft.Azure.Cosmos/src/HttpClient/HttpTimeoutPolicy.cs

@@ -55,9 +55,16 @@ public static HttpTimeoutPolicy GetTimeoutPolicy(
                 // Data Plane Reads.
                 else if (documentServiceRequest.IsReadOnlyRequest)
                 {
-                    return isPartitionLevelFailoverEnabled
-                        ? HttpTimeoutPolicyForPartitionFailover.InstanceShouldThrow503OnTimeout
-                        : HttpTimeoutPolicyDefault.InstanceShouldThrow503OnTimeout;
+                    if (isPartitionLevelFailoverEnabled)
+                    {
+                        return documentServiceRequest.OperationType == OperationType.Read 
+                            ? HttpTimeoutPolicyForPartitionFailover.InstanceShouldThrow503OnTimeoutForPointReads
+                            : HttpTimeoutPolicyForPartitionFailover.InstanceShouldThrow503OnTimeoutForNonPointReads;
+                    }
+                    else
+                    {
+                         return HttpTimeoutPolicyDefault.InstanceShouldThrow503OnTimeout;
+                    }
                 }
             }
 

Microsoft.Azure.Cosmos/src/HttpClient/HttpTimeoutPolicyForPartitionFailover.cs

@@ -9,37 +9,47 @@ namespace Microsoft.Azure.Cosmos
 
     internal sealed class HttpTimeoutPolicyForPartitionFailover : HttpTimeoutPolicy
     {
-        public static readonly HttpTimeoutPolicy Instance = new HttpTimeoutPolicyForPartitionFailover(false);
-        public static readonly HttpTimeoutPolicy InstanceShouldThrow503OnTimeout = new HttpTimeoutPolicyForPartitionFailover(true);
-        public bool shouldThrow503OnTimeout;
+        public static readonly HttpTimeoutPolicy InstanceShouldThrow503OnTimeoutForNonPointReads = new HttpTimeoutPolicyForPartitionFailover(isPointRead: false);
+        public static readonly HttpTimeoutPolicy InstanceShouldThrow503OnTimeoutForPointReads = new HttpTimeoutPolicyForPartitionFailover(isPointRead: true);
+        private readonly bool isPointRead;
         private static readonly string Name = nameof(HttpTimeoutPolicyDefault);
 
-        private HttpTimeoutPolicyForPartitionFailover(bool shouldThrow503OnTimeout)
+        private HttpTimeoutPolicyForPartitionFailover(bool isPointRead)
         {
-            this.shouldThrow503OnTimeout = shouldThrow503OnTimeout;
+            this.isPointRead = isPointRead;
         }
 
-        private readonly IReadOnlyList<(TimeSpan requestTimeout, TimeSpan delayForNextRequest)> TimeoutsAndDelays = new List<(TimeSpan requestTimeout, TimeSpan delayForNextRequest)>()
+        // Timeouts and delays are based on the following rationale:
+        // For point reads: 3 attempts with timeouts of 1s, 6s, and 6s respectively.
+        // For non-point reads: 3 attempts with timeouts of 6s, 6s, and 10s respectively.
+        private readonly IReadOnlyList<(TimeSpan requestTimeout, TimeSpan delayForNextRequest)> TimeoutsAndDelaysForPointReads = new List<(TimeSpan requestTimeout, TimeSpan delayForNextRequest)>()
         {
-            (TimeSpan.FromSeconds(.5), TimeSpan.Zero),
-            (TimeSpan.FromSeconds(.5), TimeSpan.Zero),
             (TimeSpan.FromSeconds(1), TimeSpan.Zero),
+            (TimeSpan.FromSeconds(6), TimeSpan.Zero),
+            (TimeSpan.FromSeconds(6), TimeSpan.Zero),
+        };
+
+        private readonly IReadOnlyList<(TimeSpan requestTimeout, TimeSpan delayForNextRequest)> TimeoutsAndDelaysForNonPointReads = new List<(TimeSpan requestTimeout, TimeSpan delayForNextRequest)>()
+        {
+            (TimeSpan.FromSeconds(6), TimeSpan.Zero),
+            (TimeSpan.FromSeconds(6), TimeSpan.Zero),
+            (TimeSpan.FromSeconds(10), TimeSpan.Zero),
         };
 
         public override string TimeoutPolicyName => HttpTimeoutPolicyForPartitionFailover.Name;
 
-        public override int TotalRetryCount => this.TimeoutsAndDelays.Count;
+        public override int TotalRetryCount => this.isPointRead ? this.TimeoutsAndDelaysForPointReads.Count : this.TimeoutsAndDelaysForNonPointReads.Count;
 
         public override IEnumerator<(TimeSpan requestTimeout, TimeSpan delayForNextRequest)> GetTimeoutEnumerator()
         {
-            return this.TimeoutsAndDelays.GetEnumerator();
+            return this.isPointRead ? this.TimeoutsAndDelaysForPointReads.GetEnumerator() : this.TimeoutsAndDelaysForNonPointReads.GetEnumerator();
         }
 
         public override bool ShouldRetryBasedOnResponse(HttpMethod requestHttpMethod, HttpResponseMessage responseMessage)
         {
             return false;
         }
 
-        public override bool ShouldThrow503OnTimeout => this.shouldThrow503OnTimeout;
+        public override bool ShouldThrow503OnTimeout => true;
     }
 }

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

@@ -2,6 +2,7 @@
 {
     using System;
     using System.Collections.Generic;
+    using System.Diagnostics;
     using System.IO;
     using System.Linq;
     using System.Net;
@@ -12,7 +13,7 @@
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.Diagnostics;
-    using Microsoft.Azure.Cosmos.FaultInjection;
+    using Microsoft.Azure.Cosmos.FaultInjection;
     using Microsoft.VisualStudio.TestTools.UnitTesting;
     using Newtonsoft.Json.Linq;
     using static Microsoft.Azure.Cosmos.Routing.GlobalPartitionEndpointManagerCore;
@@ -80,7 +81,8 @@ public void TestCleanup()
             finally
             {
                 //Do not delete the resources (except MM Write test object), georeplication is slow and we want to reuse the resources
-                this.client?.Dispose();
+                this.client?.Dispose();
+                Environment.SetEnvironmentVariable(ConfigurationManager.StalePartitionUnavailabilityRefreshIntervalInSeconds, null);
             }
         }
 
@@ -467,7 +469,7 @@ await this.container.DeleteItemAsync<CosmosIntegrationTestObject>(
         }
 
         [TestMethod]
-        [TestCategory("MultiRegion")]
+        [TestCategory("MultiRegion")]
         [DataRow(ConnectionMode.Direct, "15", "10", DisplayName = "Direct Mode - Scenario when the total iteration count is 15 and circuit breaker consecutive failure threshold is set to 10.")]
         [DataRow(ConnectionMode.Direct, "25", "20", DisplayName = "Direct Mode - Scenario when the total iteration count is 25 and circuit breaker consecutive failure threshold is set to 20.")]
         [DataRow(ConnectionMode.Direct, "35", "30", DisplayName = "Direct Mode - Scenario when the total iteration count is 35 and circuit breaker consecutive failure threshold is set to 30.")]
@@ -602,7 +604,7 @@ public async Task ReadItemAsync_WithCircuitBreakerEnabledAndSingleMasterAccountA
         }
 
         [TestMethod]
-        [TestCategory("MultiRegion")]
+        [TestCategory("MultiRegion")]
         [DataRow(ConnectionMode.Direct, DisplayName ="Direct Mode")]
         [DataRow(ConnectionMode.Gateway, DisplayName = "Gateway Mode")]
         [Owner("nalutripician")]
@@ -721,15 +723,14 @@ public async Task ReadItemAsync_WithCircuitBreakerEnabledAndTimeoutCounterOverwr
             }
             finally
             {
-                Environment.SetEnvironmentVariable(ConfigurationManager.PartitionLevelCircuitBreakerEnabled, null);
-                Environment.SetEnvironmentVariable(ConfigurationManager.CircuitBreakerConsecutiveFailureCountForReads, null);
-
+                Environment.SetEnvironmentVariable(ConfigurationManager.PartitionLevelCircuitBreakerEnabled, null);
+                Environment.SetEnvironmentVariable(ConfigurationManager.CircuitBreakerTimeoutCounterResetWindowInMinutes, null);
                 await this.TryDeleteItems(itemsList);
             }
         }
 
         [TestMethod]
-        [TestCategory("MultiRegion")]
+        [TestCategory("MultiRegion")]
         [Owner("dkunda")]
         [Timeout(70000)]
         public async Task ReadItemAsync_WithCircuitBreakerEnabledAndSingleMasterAccountAndServiceUnavailableReceivedFromTwoRegions_ShouldApplyPartitionLevelOverrideToThridRegion()
@@ -893,7 +894,7 @@ public async Task ReadItemAsync_WithCircuitBreakerEnabledAndSingleMasterAccountA
         }
 
         [TestMethod]
-        [TestCategory("MultiRegion")]
+        [TestCategory("MultiRegion")]
         [Owner("dkunda")]
         [Timeout(70000)]
         public async Task ReadItemAsync_WithNoPreferredRegionsAndCircuitBreakerEnabledAndSingleMasterAccountAndServiceUnavailableReceived_ShouldApplyPartitionLevelOverride()
@@ -1011,7 +1012,7 @@ public async Task ReadItemAsync_WithNoPreferredRegionsAndCircuitBreakerEnabledAn
 
         [TestMethod]
         [Owner("dkunda")]
-        [TestCategory("MultiRegion")]
+        [TestCategory("MultiRegion")]
         [Timeout(70000)]
         public async Task ReadItemAsync_WithCircuitBreakerDisabledAndSingleMasterAccountAndServiceUnavailableReceived_ShouldNotApplyPartitionLevelOverride()
         {
@@ -1095,14 +1096,12 @@ public async Task ReadItemAsync_WithCircuitBreakerDisabledAndSingleMasterAccount
             finally
             {
                 Environment.SetEnvironmentVariable(ConfigurationManager.PartitionLevelCircuitBreakerEnabled, null);
-                Environment.SetEnvironmentVariable(ConfigurationManager.CircuitBreakerConsecutiveFailureCountForReads, null);
-
                 await this.TryDeleteItems(itemsList);
             }
         }
 
         [TestMethod]
-        [Owner("dkunda")]
+        [Owner("dkunda")]
         [TestCategory("MultiRegion")]
         [Timeout(70000)]
         public async Task CreateItemAsync_WithCircuitBreakerEnabledAndSingleMasterAccountAndServiceUnavailableReceived_ShouldNotApplyPartitionLevelOverride()
@@ -1182,7 +1181,7 @@ public async Task CreateItemAsync_WithCircuitBreakerEnabledAndSingleMasterAccoun
 
         [TestMethod]
         [Owner("dkunda")]
-        [TestCategory("MultiMaster")]
+        [TestCategory("MultiMaster")]
         [DataRow(ConnectionMode.Direct, "15", "10", DisplayName = "Direct Mode - Scenario whtn the total iteration count is 15 and circuit breaker consecutive failure threshold is set to 10.")]
         [DataRow(ConnectionMode.Direct, "25", "20", DisplayName = "Direct Mode - Scenario whtn the total iteration count is 25 and circuit breaker consecutive failure threshold is set to 20.")]
         [DataRow(ConnectionMode.Direct, "35", "30", DisplayName = "Direct Mode - Scenario whtn the total iteration count is 35 and circuit breaker consecutive failure threshold is set to 30.")]
@@ -2259,7 +2258,132 @@ public async Task ClinetOverrides0msRequestTimeoutValueForPPAF()
             Assert.IsNotNull(strat);
             Assert.AreNotEqual(0, strat.Threshold);
         }
-
+
+        [TestMethod]
+        [TestCategory("MultiRegion")]
+        [Owner("pkolluri")]
+        [Timeout(70000)]
+        public async Task QueryItemAsync_WithCircuitBreakerEnabledMultiRegionAndServiceResponseDelay_ShouldFailOverToNextRegionAsync()
+        {
+            // Arrange.
+            Environment.SetEnvironmentVariable(ConfigurationManager.PartitionLevelCircuitBreakerEnabled, "True");
+            Environment.SetEnvironmentVariable(ConfigurationManager.CircuitBreakerConsecutiveFailureCountForReads, "1");
+
+            // Enabling fault injection rule to simulate a 503 service unavailable scenario.
+            string serviceResponseDelayRuleId = "response-delay-rule-" + Guid.NewGuid().ToString();
+            FaultInjectionRule serviceResponseDelayRuleFromRegion1 = new FaultInjectionRuleBuilder(
+                id: serviceResponseDelayRuleId,
+                condition:
+                    new FaultInjectionConditionBuilder()
+                        .WithOperationType(FaultInjectionOperationType.QueryItem)
+                        .WithConnectionType(FaultInjectionConnectionType.Gateway)
+                        .WithRegion(region1)
+                        .Build(),
+                result:
+                    FaultInjectionResultBuilder.GetResultBuilder(FaultInjectionServerErrorType.ResponseDelay)
+                        .WithDelay(TimeSpan.FromSeconds(70))
+                        .Build())
+                .Build();
+
+            serviceResponseDelayRuleFromRegion1.Disable();
+
+            List<FaultInjectionRule> rules = new List<FaultInjectionRule> { serviceResponseDelayRuleFromRegion1};
+            FaultInjector faultInjector = new FaultInjector(rules);
+
+            List<string> preferredRegions = new List<string> { region1, region2, region3 };
+            CosmosClientOptions cosmosClientOptions = new CosmosClientOptions()
+            {
+                ConsistencyLevel = ConsistencyLevel.Session,
+                FaultInjector = faultInjector,
+                ApplicationPreferredRegions = preferredRegions,
+                ConnectionMode = ConnectionMode.Gateway,
+            };
+
+            List<CosmosIntegrationTestObject> itemsList = new()
+            {
+                new() { Id = "smTestId2", Pk = "smpk1" },
+            };
+
+            try
+            {
+                using CosmosClient cosmosClient = new(connectionString: this.connectionString, clientOptions: cosmosClientOptions);
+                Database database = cosmosClient.GetDatabase(MultiRegionSetupHelpers.dbName);
+                Container container = database.GetContainer(MultiRegionSetupHelpers.containerName);
+
+                // Act and Assert.
+                await this.TryCreateItems(itemsList);
+
+                //Must Ensure the data is replicated to all regions
+                await Task.Delay(3000);
+
+                bool isRegion1Available = true;
+                bool isRegion2Available = true;
+
+                int thresholdCounter = 0;
+                int totalIterations = 7;
+                int ppcbDefaultThreshold = 1;
+                int firstRegionServiceUnavailableAttempt = 1;
+
+                for (int attemptCount = 1; attemptCount <= totalIterations; attemptCount++)
+                {
+                    try
+                    {
+                        string sqlQueryText = $"SELECT * FROM c WHERE c.id = '{itemsList[0].Id}'";
+                        using FeedIterator<CosmosIntegrationTestObject> feedIterator = container.GetItemQueryIterator<CosmosIntegrationTestObject>(sqlQueryText, requestOptions: new QueryRequestOptions());
+
+                        while (feedIterator.HasMoreResults)
+                        {
+                            FeedResponse<CosmosIntegrationTestObject> response = await feedIterator.ReadNextAsync();
+                            Assert.AreEqual(System.Net.HttpStatusCode.OK, response.StatusCode);
+                            IReadOnlyList<(string regionName, Uri uri)> contactedRegionMapping = response.Diagnostics.GetContactedRegions();
+                            HashSet<string> contactedRegions = new(contactedRegionMapping.Select(r => r.regionName));
+
+                            if (isRegion1Available && isRegion2Available)
+                            {
+                                Assert.IsTrue(contactedRegions.Count == 1, "Assert that, when no failure happened, the query request is being served from region 1.");
+                                Assert.IsTrue(contactedRegions.Contains(region1));
+
+                                // Simulating service unavailable on region 1.
+                                if (attemptCount == firstRegionServiceUnavailableAttempt)
+                                {
+                                    isRegion1Available = false;
+                                    serviceResponseDelayRuleFromRegion1.Enable();
+                                }
+                            }
+                            else if (isRegion2Available)
+                            {
+                                if (thresholdCounter <= ppcbDefaultThreshold)
+                                {
+                                    Assert.IsTrue(contactedRegions.Count == 2, "Asserting that when the query request succeeds before the consecutive failure count reaches the threshold, the partition didn't fail over to the next region, and the request was retried.");
+                                    Assert.IsTrue(contactedRegions.Contains(region1) && contactedRegions.Contains(region2), "Asserting that both region 1 and region 2 were contacted.");
+                                    thresholdCounter++;
+                                }
+                                else
+                                {
+                                    Assert.IsTrue(contactedRegions.Count == 1, "Asserting that when the consecutive failure count reaches the threshold, the partition was failed over to the next region, and the subsequent query request/s were successful on the next region");
+                                }
+                            }
+                        }
+                    }
+                    catch (CosmosException ce)
+                    {
+                        Assert.Fail("Query operation should succeed with successful failover to next region." + ce.Diagnostics.ToString());
+                    }
+                    catch (Exception ex)
+                    {
+                        Assert.Fail($"Unhandled Exception was thrown during Query operation call. Message: {ex.Message}");
+                    }
+                }
+            }
+            finally
+            {
+                Environment.SetEnvironmentVariable(ConfigurationManager.PartitionLevelCircuitBreakerEnabled, null);
+                Environment.SetEnvironmentVariable(ConfigurationManager.CircuitBreakerConsecutiveFailureCountForReads, null);
+
+                await this.TryDeleteItems(itemsList);
+            }
+        }
+
         private async Task TryCreateItems(List<CosmosIntegrationTestObject> testItems)
         {
             foreach (CosmosIntegrationTestObject item in testItems)

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosHttpClientCoreTests.cs

@@ -592,6 +592,53 @@ await TestScenarioAsync(
                 expectedNumberOfRetrys: 3);
         }
 
+        [TestMethod]
+        public void HttpTimeoutPolicyForParitionFailoverForQueries()
+        {
+            HttpTimeoutPolicy httpTimeoutPolicyForQuery = HttpTimeoutPolicy.GetTimeoutPolicy(
+                    documentServiceRequest: CosmosHttpClientCoreTests.CreateDocumentServiceRequestByOperation(ResourceType.Document, OperationType.Query),
+                    isPartitionLevelFailoverEnabled: true,
+                    isThinClientEnabled: false);
+            IEnumerator<(TimeSpan requestTimeout, TimeSpan delayForNextRequest)> availableRetries = httpTimeoutPolicyForQuery.GetTimeoutEnumerator();
+
+            int count = 0;
+            while (availableRetries.MoveNext())
+            {
+                if (count <= 1)
+                {
+                    Assert.AreEqual(new TimeSpan(0,0,6), availableRetries.Current.requestTimeout);
+                }
+                else if (count == 2)
+                {
+                    Assert.AreEqual(new TimeSpan(0, 0, 10), availableRetries.Current.requestTimeout);
+                }
+                count++;
+            }
+        }
+
+        [TestMethod]
+        public void HttpTimeoutPolicyForParitionFailoverForReads()
+        {
+            HttpTimeoutPolicy httpTimeoutPolicyForQuery = HttpTimeoutPolicy.GetTimeoutPolicy(
+                    documentServiceRequest: CosmosHttpClientCoreTests.CreateDocumentServiceRequestByOperation(ResourceType.Document, OperationType.Read),
+                    isPartitionLevelFailoverEnabled: true,
+                    isThinClientEnabled: false);
+            IEnumerator<(TimeSpan requestTimeout, TimeSpan delayForNextRequest)> availableRetries = httpTimeoutPolicyForQuery.GetTimeoutEnumerator();
+
+            int count = 0;
+            while (availableRetries.MoveNext())
+            {
+                if (count == 0)
+                {
+                    Assert.AreEqual(new TimeSpan(0, 0, 1), availableRetries.Current.requestTimeout);
+                }
+                else if (count == 1 || count ==2 )
+                {
+                    Assert.AreEqual(new TimeSpan(0, 0, 6), availableRetries.Current.requestTimeout);
+                }
+                count++;
+            }
+        }
 
         private static DocumentServiceRequest CreateDocumentServiceRequestByOperation(
             ResourceType resourceType,