Session Consistency: Adds SessionTokenMismatchRetryPolicy optimization through customer supplied region switch hints

Microsoft.Azure.Cosmos/src/ConnectionPolicy.cs

@@ -536,6 +536,15 @@ internal CosmosClientTelemetryOptions CosmosClientTelemetryOptions
             set;
         }
 
+        /// <summary>
+        /// provides SessionTokenMismatchRetryPolicy optimization through customer supplied region switch hints
+        /// </summary>
+        internal SessionRetryOptions SessionRetryOptions
+        {
+            get;
+            set;
+        }
+
         /// <summary>
         /// GlobalEndpointManager will subscribe to this event if user updates the preferredLocations list in the Azure Cosmos DB service.
         /// </summary>

Microsoft.Azure.Cosmos/src/CosmosClientOptions.cs

@@ -89,7 +89,8 @@ public CosmosClientOptions()
             this.ConnectionProtocol = CosmosClientOptions.DefaultProtocol;
             this.ApiType = CosmosClientOptions.DefaultApiType;
             this.CustomHandlers = new Collection<RequestHandler>();
-            this.CosmosClientTelemetryOptions = new CosmosClientTelemetryOptions();
+            this.CosmosClientTelemetryOptions = new CosmosClientTelemetryOptions();
+            this.SessionRetryOptions = new SessionRetryOptions();
         }
 
         /// <summary>
@@ -120,7 +121,12 @@ public string ApplicationName
         /// <summary>
         /// Get or set session container for the client
         /// </summary>
-        internal ISessionContainer SessionContainer { get; set; }
+        internal ISessionContainer SessionContainer { get; set; }
+
+        /// <summary>
+        /// hint which guide SDK-internal retry policies on how early to switch retries to a different region. 
+        /// </summary>
+        internal SessionRetryOptions SessionRetryOptions { get; private set; }
 
         /// <summary>
         /// Gets or sets the location where the application is running. This will influence the SDK's choice for the Azure Cosmos DB service interaction.
@@ -740,6 +746,20 @@ public Func<HttpClient> HttpClientFactory
         /// after the threshold step time, the SDK will hedge to the third region and so on.
         /// </remarks>
         public AvailabilityStrategy AvailabilityStrategy { get; set; }
+
+        /// <summary>
+        /// provides SessionTokenMismatchRetryPolicy optimization through customer supplied region switch hints
+        /// </summary>
+#if PREVIEW
+        public
+#else
+        internal
+#endif
+        bool EnableRemoteRegionPreferredForSessionRetry
+        {
+            get => this.SessionRetryOptions.RemoteRegionPreferred;
+            set => this.SessionRetryOptions.RemoteRegionPreferred = value;
+        }
 
         /// <summary>
         /// Enable partition key level failover
@@ -1004,7 +1024,8 @@ internal virtual ConnectionPolicy GetConnectionPolicy(int clientId)
                 ConnectionProtocol = this.ConnectionProtocol,
                 UserAgentContainer = this.CreateUserAgentContainerWithFeatures(clientId),
                 UseMultipleWriteLocations = true,
-                IdleTcpConnectionTimeout = this.IdleTcpConnectionTimeout,
+                IdleTcpConnectionTimeout = this.IdleTcpConnectionTimeout,
+                SessionRetryOptions = this.SessionRetryOptions,
                 OpenTcpConnectionTimeout = this.OpenTcpConnectionTimeout,
                 MaxRequestsPerTcpConnection = this.MaxRequestsPerTcpConnection,
                 MaxTcpConnectionsPerEndpoint = this.MaxTcpConnectionsPerEndpoint,

Microsoft.Azure.Cosmos/src/DocumentClient.cs

@@ -6800,7 +6800,8 @@ private void CreateStoreModel(bool subscribeRntbdStatus)
                 !this.enableRntbdChannel,
                 this.UseMultipleWriteLocations && (this.accountServiceConfiguration.DefaultConsistencyLevel != Documents.ConsistencyLevel.Strong),
                 true,
-                enableReplicaValidation: this.isReplicaAddressValidationEnabled);
+                enableReplicaValidation: this.isReplicaAddressValidationEnabled,
+                sessionRetryOptions: this.ConnectionPolicy.SessionRetryOptions);
 
             if (subscribeRntbdStatus)
             {

Microsoft.Azure.Cosmos/src/Fluent/CosmosClientBuilder.cs

@@ -581,8 +581,24 @@ public CosmosClientBuilder WithSerializerOptions(CosmosSerializationOptions cosm
         {
             this.clientOptions.SerializerOptions = cosmosSerializerOptions;
             return this;
-        }
-
+        }
+
+        /// <summary>
+        /// provides SessionTokenMismatchRetryPolicy optimization through customer supplied region switch hints
+        /// </summary>
+        /// <param name="enableRemoteRegionPreferredForSessionRetry"></param>
+        /// <returns>The <see cref="CosmosClientBuilder"/> object</returns>
+#if PREVIEW
+        public
+#else
+        internal
+#endif
+        CosmosClientBuilder WithEnableRemoteRegionPreferredForSessionRetry(bool enableRemoteRegionPreferredForSessionRetry)
+        {
+            this.clientOptions.EnableRemoteRegionPreferredForSessionRetry = enableRemoteRegionPreferredForSessionRetry;
+            return this;
+        }
+
         /// <summary>
         /// Set a custom JSON serializer. 
         /// </summary>

Microsoft.Azure.Cosmos/src/SessionRetryOptions.cs

@@ -0,0 +1,43 @@
+// ------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+// ------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos
+{
+    using System;
+    using Microsoft.Azure.Documents;
+
+    /// <summary>
+    /// Implementation of ISessionRetryOptions interface, do not want clients to subclass.
+    /// </summary>
+    internal sealed class SessionRetryOptions : ISessionRetryOptions
+    {
+        /// <summary>
+        /// Initializes a new instance of the <see cref="SessionRetryOptions"/> class.
+        /// </summary>
+        public SessionRetryOptions()
+        {
+            this.MinInRegionRetryTime = ConfigurationManager.GetMinRetryTimeInLocalRegionWhenRemoteRegionPreferred();
+            this.MaxInRegionRetryCount = ConfigurationManager.GetMaxRetriesInLocalRegionWhenRemoteRegionPreferred();
+        }
+        /// <summary>
+        /// Sets the minimum retry time for 404/1002 retries within each region for read and write operations. 
+        /// The minimum value is 100ms - this minimum is enforced to provide a way for the local region to catch-up on replication lag. The default value is 500ms - as a recommendation ensure that this value is higher than the steady-state
+        /// 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
+        /// region is at least the min. in-region retry time.
+        /// </summary>
+        public int MaxInRegionRetryCount { get; private set; }
+
+        /// <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;
+
+    }
+}

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

@@ -72,6 +72,20 @@ internal static class ConfigurationManager
         /// </summary>
         internal static readonly string DistributedQueryGatewayModeEnabled = "AZURE_COSMOS_DISTRIBUTED_QUERY_GATEWAY_ENABLED";
 
+        /// <summary>
+        /// intent is If a client specify a value, we will force it to be atleast 100ms, otherwise default is going to be 500ms
+        /// </summary>
+        internal static readonly string MinInRegionRetryTimeForWritesInMs = "AZURE_COSMOS_SESSION_TOKEN_MISMATCH_IN_REGION_RETRY_TIME_IN_MILLISECONDS";
+        internal static readonly int DefaultMinInRegionRetryTimeForWritesInMs = 500;
+        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).
+        /// </summary>
+        internal static readonly string MaxRetriesInLocalRegionWhenRemoteRegionPreferred = "AZURE_COSMOS_MAX_RETRIES_IN_LOCAL_REGION_WHEN_REMOTE_REGION_PREFERRED";
+        internal static readonly int DefaultMaxRetriesInLocalRegionWhenRemoteRegionPreferred = 1;
+        internal static readonly int MinMaxRetriesInLocalRegionWhenRemoteRegionPreferred = 1;
+
         /// <summary>
         /// A read-only string containing the environment variable name for enabling binary encoding. This will eventually
         /// be removed once binary encoding is enabled by default for both preview
@@ -96,6 +110,26 @@ public static T GetEnvironmentVariable<T>(string variable, T defaultValue)
             return (T)Convert.ChangeType(value, typeof(T));
         }
 
+        public static int GetMaxRetriesInLocalRegionWhenRemoteRegionPreferred()
+        {
+            return Math.Max(
+                ConfigurationManager
+                    .GetEnvironmentVariable(
+                        variable: MaxRetriesInLocalRegionWhenRemoteRegionPreferred,
+                        defaultValue: DefaultMaxRetriesInLocalRegionWhenRemoteRegionPreferred),
+                MinMaxRetriesInLocalRegionWhenRemoteRegionPreferred);
+        }
+
+        public static TimeSpan GetMinRetryTimeInLocalRegionWhenRemoteRegionPreferred()
+        {
+            return TimeSpan.FromMilliseconds(Math.Max(
+                ConfigurationManager
+                    .GetEnvironmentVariable(
+                        variable: MinInRegionRetryTimeForWritesInMs,
+                        defaultValue: DefaultMinInRegionRetryTimeForWritesInMs),
+                MinMinInRegionRetryTimeForWritesInMs));
+        }
+
         /// <summary>
         /// Gets the boolean value of the replica validation environment variable. Note that, replica validation
         /// is enabled by default for the preview package and disabled for GA at the moment. The user can set the