Gateway: Adds optimization to reduce gateway header size by scoping session token to specific partitions.

Microsoft.Azure.Cosmos/src/DocumentClient.cs

@@ -943,6 +943,8 @@ private async Task GetInitializationTaskAsync(IStoreClientFactory storeClientFac
             this.partitionKeyRangeCache = new PartitionKeyRangeCache(this, this.GatewayStoreModel, this.collectionCache);
             this.ResetSessionTokenRetryPolicy = new ResetSessionTokenRetryPolicyFactory(this.sessionContainer, this.collectionCache, this.retryPolicy);
 
+            gatewayStoreModel.SetCaches(this.partitionKeyRangeCache, this.collectionCache);
+
             if (this.ConnectionPolicy.ConnectionMode == ConnectionMode.Gateway)
             {
                 this.StoreModel = this.GatewayStoreModel;

Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs

@@ -10,8 +10,10 @@ namespace Microsoft.Azure.Cosmos
     using System.Linq;
     using System.Net;
     using System.Net.Http;
+    using System.Net.Http.Headers;
     using System.Threading;
     using System.Threading.Tasks;
+    using Microsoft.Azure.Cosmos.Common;
     using Microsoft.Azure.Cosmos.Core.Trace;
     using Microsoft.Azure.Cosmos.Routing;
     using Microsoft.Azure.Documents;
@@ -28,6 +30,10 @@ internal class GatewayStoreModel : IStoreModel, IDisposable
 
         private GatewayStoreClient gatewayStoreClient;
 
+        // Caches to resolve the PartitionKeyRange from request. For Session Token Optimization.
+        private ClientCollectionCache ClientCollectionCache;
+        private PartitionKeyRangeCache PartitionKeyRangeCache;
+
         public GatewayStoreModel(
             GlobalEndpointManager endpointManager,
             ISessionContainer sessionContainer,
@@ -49,10 +55,12 @@ public GatewayStoreModel(
 
         public virtual async Task<DocumentServiceResponse> ProcessMessageAsync(DocumentServiceRequest request, CancellationToken cancellationToken = default)
         {
-            GatewayStoreModel.ApplySessionToken(
+            await GatewayStoreModel.ApplySessionTokenAsync(
                 request,
                 this.defaultConsistencyLevel,
-                this.sessionContainer);
+                this.sessionContainer,
+                this.PartitionKeyRangeCache,
+                this.ClientCollectionCache);
 
             DocumentServiceResponse response;
             try
@@ -143,6 +151,12 @@ public virtual async Task<AccountProperties> GetDatabaseAccountAsync(Func<ValueT
             return databaseAccount;
         }
 
+        public void SetCaches(PartitionKeyRangeCache partitionKeyRangeCache, ClientCollectionCache clientCollectionCache)
+        {
+            this.ClientCollectionCache = clientCollectionCache;
+            this.PartitionKeyRangeCache = partitionKeyRangeCache;
+        }
+
         public void Dispose()
         {
             this.Dispose(true);
@@ -193,10 +207,12 @@ private void CaptureSessionToken(
             }
         }
 
-        internal static void ApplySessionToken(
+        internal static async Task ApplySessionTokenAsync(
             DocumentServiceRequest request,
             ConsistencyLevel defaultConsistencyLevel,
-            ISessionContainer sessionContainer)
+            ISessionContainer sessionContainer,
+            PartitionKeyRangeCache partitionKeyRangeCache,
+            ClientCollectionCache clientCollectionCache)
         {
             if (request.Headers == null)
             {
@@ -227,20 +243,134 @@ internal static void ApplySessionToken(
                 (!string.IsNullOrEmpty(requestConsistencyLevel)
                     && string.Equals(requestConsistencyLevel, ConsistencyLevel.Session.ToString(), StringComparison.OrdinalIgnoreCase));
 
-            if (!sessionConsistency)
+            if (!sessionConsistency || !request.IsReadOnlyRequest)
+            {
+                return; // Only apply the session token in case of session consistency and the request is read only
+            }
+
+            string sessionToken = null;
+            bool isSuccess = false;
+            if (clientCollectionCache != null && partitionKeyRangeCache != null)
             {
-                return; // Only apply the session token in case of session consistency
+                (isSuccess, sessionToken) = await TryResolveSessionTokenAsync(request, sessionContainer, partitionKeyRangeCache, clientCollectionCache);
             }
 
-            //Apply the ambient session.
-            string sessionToken = sessionContainer.ResolveGlobalSessionToken(request);
+            if (!isSuccess)
+            {
+                sessionToken = sessionContainer.ResolveGlobalSessionToken(request);
+            }
 
             if (!string.IsNullOrEmpty(sessionToken))
             {
                 request.Headers[HttpConstants.HttpHeaders.SessionToken] = sessionToken;
             }
         }
 
+        internal static async Task<Tuple<bool, string>> TryResolveSessionTokenAsync(DocumentServiceRequest request, 
+                                                                    ISessionContainer sessionContainer, 
+                                                                    PartitionKeyRangeCache partitionKeyRangeCache, 
+                                                                    ClientCollectionCache clientCollectionCache)
+        {
+            if (request == null)
+            {
+                throw new ArgumentNullException(nameof(request));
+            }
+
+            if (sessionContainer == null)
+            {
+                throw new ArgumentNullException(nameof(sessionContainer));
+            }
+
+            if (partitionKeyRangeCache == null)
+            {
+                throw new ArgumentNullException(nameof(partitionKeyRangeCache));
+            }
+
+            if (clientCollectionCache == null)
+            {
+                throw new ArgumentNullException(nameof(clientCollectionCache));
+            }
+
+            if (request.ResourceType.IsPartitioned())
+            {
+                (bool isSuccess, PartitionKeyRange partitionKeyRange) = await TryResolvePartitionKeyRangeAsync(request: request, 
+                                                                        sessionContainer: sessionContainer, 
+                                                                        partitionKeyRangeCache: partitionKeyRangeCache, 
+                                                                        clientCollectionCache: clientCollectionCache, 
+                                                                        refreshCache: false);
+
+                if (isSuccess && sessionContainer is SessionContainer gatewaySessionContainer)
+                {   
+                    string localSessionToken = gatewaySessionContainer.ResolvePartitionLocalSessionTokenForGateway(request, partitionKeyRange.Id);
+                    if (!string.IsNullOrEmpty(localSessionToken))
+                    {
+                        return new Tuple<bool, string>(true, localSessionToken);
+                    }
+                }
+            }
+
+            return new Tuple<bool, string>(false, null);
+        }
+
+        private static async Task<Tuple<bool, PartitionKeyRange>> TryResolvePartitionKeyRangeAsync(DocumentServiceRequest request, 
+                                                                                   ISessionContainer sessionContainer, 
+                                                                                   PartitionKeyRangeCache partitionKeyRangeCache, 
+                                                                                   ClientCollectionCache clientCollectionCache, 
+                                                                                   bool refreshCache)
+        {
+            if (refreshCache)
+            {
+                request.ForceMasterRefresh = true;
+                request.ForceNameCacheRefresh = true;
+            }
+
+            PartitionKeyRange partitonKeyRange = null;
+            ContainerProperties collection = await clientCollectionCache.ResolveCollectionAsync(request, CancellationToken.None);
+
+            string partitionKeyString = request.Headers[HttpConstants.HttpHeaders.PartitionKey];
+            if (partitionKeyString != null)
+            {
+                CollectionRoutingMap collectionRoutingMap = await partitionKeyRangeCache.TryLookupAsync(collectionRid: collection.ResourceId,
+                                                                                                        previousValue: null,
+                                                                                                        request: request,
+                                                                                                        cancellationToken: CancellationToken.None);
+
+                if (refreshCache && collectionRoutingMap != null)
+                {
+                    collectionRoutingMap = await partitionKeyRangeCache.TryLookupAsync(collectionRid: collection.ResourceId,
+                                                                                        previousValue: collectionRoutingMap,
+                                                                                        request: request,
+                                                                                        cancellationToken: CancellationToken.None);
+                }
+
+                partitonKeyRange = AddressResolver.TryResolveServerPartitionByPartitionKey(request: request, 
+                                                                                           partitionKeyString: partitionKeyString, 
+                                                                                           collectionCacheUptoDate: false, 
+                                                                                           collection: collection, 
+                                                                                           routingMap: collectionRoutingMap);
+            }
+            else if (request.PartitionKeyRangeIdentity != null)
+            {
+                PartitionKeyRangeIdentity partitionKeyRangeId = request.PartitionKeyRangeIdentity;
+                partitonKeyRange = await partitionKeyRangeCache.TryGetPartitionKeyRangeByIdAsync(collection.ResourceId, 
+                                                                                                 partitionKeyRangeId.ToString(), 
+                                                                                                 refreshCache);
+            }
+
+            if (partitonKeyRange == null)
+            {
+                if (refreshCache)
+                {
+                    return new Tuple<bool, PartitionKeyRange>(false, null);
+                }
+
+                // need to refresh cache. Maybe split happened.
+                return await TryResolvePartitionKeyRangeAsync(request, sessionContainer, partitionKeyRangeCache, clientCollectionCache, true);
+            }
+
+            return new Tuple<bool, PartitionKeyRange>(true, partitonKeyRange);
+        }
+
         // DEVNOTE: This can be replace with ReplicatedResourceClient.IsMasterOperation on next Direct sync
         internal static bool IsMasterOperation(
             ResourceType resourceType,

Microsoft.Azure.Cosmos/src/Routing/AddressResolver.cs

@@ -447,7 +447,7 @@ private async Task<ResolutionResult> TryResolveServerPartitionAsync(
             object effectivePartitionKeyStringObject = null;
             if (partitionKeyString != null)
             {
-                range = this.TryResolveServerPartitionByPartitionKey(
+                range = TryResolveServerPartitionByPartitionKey(
                     request,
                     partitionKeyString,
                     collectionCacheIsUptodate,
@@ -543,7 +543,7 @@ private PartitionKeyRange TryResolveSinglePartitionCollection(
                     // due to parallel usage of V3 SDK and a possible storage or throughput split
                     // The current client might be legacy and not aware of this.
                     // In such case route the request to the first partition
-                    return this.TryResolveServerPartitionByPartitionKey(
+                    return AddressResolver.TryResolveServerPartitionByPartitionKey(
                                         request,
                                         "[]", // This corresponds to first partition
                                         collectionCacheIsUptoDate,
@@ -624,7 +624,7 @@ private async Task<ResolutionResult> TryResolveServerPartitionByPartitionKeyRang
             return new ResolutionResult(partitionKeyRange, addresses, identity);
         }
 
-        private PartitionKeyRange TryResolveServerPartitionByPartitionKey(
+        internal static PartitionKeyRange TryResolveServerPartitionByPartitionKey(
             DocumentServiceRequest request,
             string partitionKeyString,
             bool collectionCacheUptoDate,

Microsoft.Azure.Cosmos/src/SessionContainer.cs

@@ -49,6 +49,11 @@ public ISessionToken ResolvePartitionLocalSessionToken(DocumentServiceRequest re
             return SessionContainer.ResolvePartitionLocalSessionToken(this.state, request, partitionKeyRangeId);
         }
 
+        public string ResolvePartitionLocalSessionTokenForGateway(DocumentServiceRequest request, string partitionKeyRangeId)
+        {
+            return SessionContainer.ResolvePartitionLocalSessionTokenForGateway(this.state, request, partitionKeyRangeId);
+        }
+
         public void ClearTokenByCollectionFullname(string collectionFullname)
         {
             SessionContainer.ClearTokenByCollectionFullname(this.state, collectionFullname);
@@ -139,6 +144,30 @@ private static ISessionToken ResolvePartitionLocalSessionToken(SessionContainerS
             return SessionTokenHelper.ResolvePartitionLocalSessionToken(request, partitionKeyRangeId, SessionContainer.GetPartitionKeyRangeIdToTokenMap(self, request));
         }
 
+        private static string ResolvePartitionLocalSessionTokenForGateway(SessionContainerState self, DocumentServiceRequest request, string partitionKeyRangeId)
+        {
+            ConcurrentDictionary<string, ISessionToken> partitionKeyRangeIdToTokenMap = SessionContainer.GetPartitionKeyRangeIdToTokenMap(self, request);
+            if (partitionKeyRangeIdToTokenMap != null)
+            {
+                if (partitionKeyRangeIdToTokenMap.TryGetValue(partitionKeyRangeId, out ISessionToken sessionToken))
+                {
+                    return partitionKeyRangeId + ":" + sessionToken.ConvertToString();
+                }
+                else if (request.RequestContext.ResolvedPartitionKeyRange.Parents != null)
+                {
+                    for (int parentIndex = request.RequestContext.ResolvedPartitionKeyRange.Parents.Count - 1; parentIndex >= 0; parentIndex--)
+                    {
+                        if (partitionKeyRangeIdToTokenMap.TryGetValue(request.RequestContext.ResolvedPartitionKeyRange.Parents[parentIndex], out sessionToken))
+                        {
+                            return partitionKeyRangeId + ":" + sessionToken.ConvertToString();
+                        }
+                    }
+                }
+            }
+
+            return null;
+        }
+
         private static void ClearTokenByCollectionFullname(SessionContainerState self, string collectionFullname)
         {
             if (!string.IsNullOrEmpty(collectionFullname))