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

Microsoft.Azure.Cosmos/src/DocumentClient.cs

@@ -946,6 +946,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,13 @@ 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 +208,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 +244,132 @@ internal static void ApplySessionToken(
  (!string.IsNullOrEmpty(requestConsistencyLevel)
  && string.Equals(requestConsistencyLevel, ConsistencyLevel.Session.ToString(), StringComparison.OrdinalIgnoreCase));
 
- if (!sessionConsistency)
+ if (!sessionConsistency || (!request.IsReadOnlyRequest && request.OperationType != OperationType.Batch))
  {
- return; // Only apply the session token in case of session consistency
+ return; // Only apply the session token in case of session consistency and the request is read only
  }
 
- //Apply the ambient session.
- string sessionToken = sessionContainer.ResolveGlobalSessionToken(request);
+ (bool isSuccess, string sessionToken) = await TryResolveSessionTokenAsync(request, 
+ sessionContainer, 
+ partitionKeyRangeCache, 
+ clientCollectionCache);
+
+ 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 = AddressResolver.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,33 @@ 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))