[INTERNAL] SessionRetry: Adds support to yield immediately after the original attempt when READ_SESSION_NOT_AVAILABLE error is thrown

Microsoft.Azure.Cosmos/src/SessionRetryOptions.cs

@@ -12,6 +12,8 @@ namespace Microsoft.Azure.Cosmos
     /// </summary>
     internal sealed class SessionRetryOptions : ISessionRetryOptions
     {
+        private int maxInRegionRetryCount;
+        private bool remoteRegionPreferred;
         /// <summary>
         /// Initializes a new instance of the <see cref="SessionRetryOptions"/> class.
         /// </summary>
@@ -26,18 +28,40 @@ public SessionRetryOptions()
         /// replication latency between the regions you chose
         /// </summary>
         public TimeSpan MinInRegionRetryTime { get; private set; }
-
+        
         /// <summary>
-        /// Sets the maximum number of retries within each region for read and write operations. The minimum value is 1 - the backoff time for the last in-region retry will ensure that the total retry time within the
+        /// Sets the maximum number of retries within each region for read and write operations - the backoff time for the last in-region retry will ensure that the total retry time within the
         /// region is at least the min. in-region retry time.
         /// </summary>
-        public int MaxInRegionRetryCount { get; private set; }
+        public int MaxInRegionRetryCount 
+        {
+            get => this.maxInRegionRetryCount;
+            internal set
+            {
+                if (value <= 0 && !this.RemoteRegionPreferred)
+                {
+                    throw new ArgumentException("MaxInRegionRetryCount can only be set to 0 or less when RemoteRegionPreferred is true.");
+                }
+                this.maxInRegionRetryCount = value;
+            }
+        }
 
         /// <summary>
         /// hints which guide SDK-internal retry policies on how early to switch retries to a different region. If true, will retry all replicas once and add a minimum delay before switching to the next region.If false, it will
         /// retry in the local region up to 5s
         /// </summary>
-        public bool RemoteRegionPreferred { get; set; } = false;
+        public bool RemoteRegionPreferred 
+        {
+            get => this.remoteRegionPreferred;
+            set
+            {
+                if (!value && this.MaxInRegionRetryCount <= 0)
+                {
+                    throw new ArgumentException("RemoteRegionPreferred cannot be set to false when MaxInRegionRetryCount is 0 or less.");
+                }
+                this.remoteRegionPreferred = value;
+            }
+        } 
 
     }
 }

Microsoft.Azure.Cosmos/src/Util/ConfigurationManager.cs

@@ -80,7 +80,8 @@ internal static class ConfigurationManager
         internal static readonly int MinMinInRegionRetryTimeForWritesInMs = 100;
 
         /// <summary>
-        /// intent is If a client specify a value, we will force it to be atleast 1, otherwise default is going to be 1(right now both the values are 1 but we have the provision to change them in future).
+        /// intent is If a client specify a value, we will force it to be atleast 1, otherwise default is going to be 1(right now both the values are 1
+        /// but we have the provision to change them in future).
         /// </summary>
         internal static readonly string MaxRetriesInLocalRegionWhenRemoteRegionPreferred = "AZURE_COSMOS_MAX_RETRIES_IN_LOCAL_REGION_WHEN_REMOTE_REGION_PREFERRED";
         internal static readonly int DefaultMaxRetriesInLocalRegionWhenRemoteRegionPreferred = 1;
@@ -113,11 +114,11 @@ public static T GetEnvironmentVariable<T>(string variable, T defaultValue)
         public static int GetMaxRetriesInLocalRegionWhenRemoteRegionPreferred()
         {
             return Math.Max(
-                ConfigurationManager
-                    .GetEnvironmentVariable(
-                        variable: MaxRetriesInLocalRegionWhenRemoteRegionPreferred,
-                        defaultValue: DefaultMaxRetriesInLocalRegionWhenRemoteRegionPreferred),
-                MinMaxRetriesInLocalRegionWhenRemoteRegionPreferred);
+               ConfigurationManager
+                   .GetEnvironmentVariable(
+                       variable: MaxRetriesInLocalRegionWhenRemoteRegionPreferred,
+                       defaultValue: DefaultMaxRetriesInLocalRegionWhenRemoteRegionPreferred),
+               MinMaxRetriesInLocalRegionWhenRemoteRegionPreferred);
         }
 
         public static TimeSpan GetMinRetryTimeInLocalRegionWhenRemoteRegionPreferred()

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

@@ -34,6 +34,74 @@ public async Task TestInitAsync()
             Assert.IsTrue(this.writeRegionMap.Count() >= 2);
 
         }
+
+        [TestMethod]
+        [DataRow(FaultInjectionOperationType.ReadItem, DisplayName = "No retries for ReadItem when MaxInRegionRetryCount is 0 and EnableRemoteRegionPreferredForSessionRetry is true")]
+        [DataRow(FaultInjectionOperationType.QueryItem, DisplayName = "No retries for QueryItem when MaxInRegionRetryCount is 0 and and EnableRemoteRegionPreferredForSessionRetry is true")]
+        [TestCategory("MultiMaster")]
+        public async Task ReadOrQueryOperationWithMaxInRegionRetryCountZero(FaultInjectionOperationType faultInjectionOperationType)
+        {
+            string[] preferredRegions = this.writeRegionMap.Keys.ToArray();
+
+            FaultInjectionRule badSessionTokenRule = new FaultInjectionRuleBuilder(
+                    id: "badSessionTokenRule",
+                    condition:
+                        new FaultInjectionConditionBuilder()
+                            .WithOperationType(faultInjectionOperationType)
+                            .WithRegion(preferredRegions[0])
+                            .Build(),
+                    result:
+                        FaultInjectionResultBuilder.GetResultBuilder(FaultInjectionServerErrorType.ReadSessionNotAvailable)
+                            .Build())
+                    .WithDuration(TimeSpan.FromMinutes(10))
+                    .Build();
+
+            List<FaultInjectionRule> rules = new List<FaultInjectionRule>() { badSessionTokenRule };
+            FaultInjector faultInjector = new FaultInjector(rules);
+            Assert.IsNotNull(faultInjector);
+
+            CosmosClientOptions clientOptions = new CosmosClientOptions()
+            {
+                EnableRemoteRegionPreferredForSessionRetry = true,
+                ConsistencyLevel = ConsistencyLevel.Session,
+                ApplicationPreferredRegions = preferredRegions,
+                ConnectionMode = ConnectionMode.Direct,
+            };
+
+            // Explicitly set MaxInRegionRetryCount to 0 in SessionRetryOptions
+            clientOptions.SessionRetryOptions.MaxInRegionRetryCount = 0;
+            using (CosmosClient faultInjectionClient = new CosmosClient(
+                        connectionString: this.connectionString,
+                        clientOptions: faultInjector.GetFaultInjectionClientOptions(clientOptions)))
+            {
+                Database database = faultInjectionClient.GetDatabase(MultiRegionSetupHelpers.dbName);
+                Container container = await database.CreateContainerIfNotExistsAsync("sessionRetryPolicy", "/id");
+                string GUID = Guid.NewGuid().ToString();
+                dynamic testObject = new
+                {
+                    id = GUID,
+                    name = "customer one",
+                    address = new
+                    {
+                        line1 = "45 new street",
+                        city = "mckinney",
+                        postalCode = "98989",
+                    }
+                };
+
+                ItemResponse<dynamic> response = await container.CreateItemAsync<dynamic>(testObject);
+                Assert.IsNotNull(response);
+
+                OperationExecutionResult executionResult = await this.PerformDocumentOperation(faultInjectionOperationType, container, testObject);
+                this.ValidateOperationExecutionResult(executionResult, true);
+
+                // Assert that only the original attempt happened (no retries)
+                long hitCount = badSessionTokenRule.GetHitCount();
+                Assert.AreEqual(4, hitCount, $"There should be only one attempt (no retries) for {faultInjectionOperationType} when MaxInRegionRetryCount is 0 and RemotePreferredRegion is set to true.");
+            }
+
+        }
+
         [TestMethod]
         [DataRow(FaultInjectionOperationType.ReadItem, 2, true, DisplayName = "Validate Read Item operation with remote region preferred.")]
         [DataRow(FaultInjectionOperationType.QueryItem, 1, true, DisplayName = "Validate Query Item operation with remote region preferred.")]
@@ -46,6 +114,7 @@ public async Task ReadOperationWithReadSessionUnavailableTest(FaultInjectionOper
             string[] preferredRegions = this.writeRegionMap.Keys.ToArray();
             Environment.SetEnvironmentVariable(ConfigurationManager.MinInRegionRetryTimeForWritesInMs, "100");
             Environment.SetEnvironmentVariable(ConfigurationManager.MaxRetriesInLocalRegionWhenRemoteRegionPreferred, Convert.ToString(sessionTokenMismatchRetryAttempts));
+
             try
             {
                 // if I go to first region for reading an item, I should get a 404/2002 response for 10 minutes

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

@@ -35,6 +35,17 @@ public void SessionRetryOptionsValidValuesTest()
             }
 
         }
+        [TestMethod]
+        public void SessionRetryOptionsValidateMaxInRegionRetryCount()
+        {
+            SessionRetryOptions options = new SessionRetryOptions
+            {
+                RemoteRegionPreferred = false
+            };
+
+            Assert.ThrowsException<ArgumentException>(() => options.MaxInRegionRetryCount = 0);
+            Assert.ThrowsException<ArgumentException>(() => options.MaxInRegionRetryCount = -5);
+        }
 
         [TestMethod]
         public void SessionRetryOptionsDefaultValuesTest()
@@ -43,7 +54,6 @@ public void SessionRetryOptionsDefaultValuesTest()
             {
                 EnableRemoteRegionPreferredForSessionRetry = true,
             };
-
             Assert.IsTrue(clientOptions.SessionRetryOptions.MinInRegionRetryTime == TimeSpan.FromMilliseconds(500));
             Assert.IsTrue(clientOptions.SessionRetryOptions.MaxInRegionRetryCount == 1);