From 43f23662c88a46499fc766f3a96c06e64aa4d795 Mon Sep 17 00:00:00 2001 From: Asket Agarwal Date: Fri, 29 Jan 2021 20:13:34 +0530 Subject: [PATCH 01/12] Session Token Optimization for Gateway Mode. --- Microsoft.Azure.Cosmos/src/DocumentClient.cs | 2 + .../src/GatewayStoreModel.cs | 113 ++++++++++++++++-- .../src/Routing/AddressResolver.cs | 6 +- .../GatewaySessionTokenTests.cs | 78 ++++++++++++ .../GatewayStoreModelTest.cs | 14 +-- 5 files changed, 196 insertions(+), 17 deletions(-) create mode 100644 Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs diff --git a/Microsoft.Azure.Cosmos/src/DocumentClient.cs b/Microsoft.Azure.Cosmos/src/DocumentClient.cs index 0d0fba6533..8078263931 100644 --- a/Microsoft.Azure.Cosmos/src/DocumentClient.cs +++ b/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; diff --git a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs index f6c34d1083..86a8c80ce4 100644 --- a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs +++ b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs @@ -10,6 +10,7 @@ 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.Core.Trace; @@ -28,6 +29,10 @@ internal class GatewayStoreModel : IStoreModel, IDisposable private GatewayStoreClient gatewayStoreClient; + // Caches to resolve the PartitionKeyRange from request. For Session Token Optimization. + internal ClientCollectionCache ClientCollectionCache { get; set; } + internal PartitionKeyRangeCache PartitionKeyRangeCache { get; set; } + public GatewayStoreModel( GlobalEndpointManager endpointManager, ISessionContainer sessionContainer, @@ -49,10 +54,12 @@ public GatewayStoreModel( public virtual async Task 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 +150,12 @@ public virtual async Task GetDatabaseAccountAsync(Func ResolveSessionTokenAsync(DocumentServiceRequest request, ISessionContainer sessionContainer, PartitionKeyRangeCache partitionKeyRangeCache, ClientCollectionCache clientCollectionCache) + { + if (request == null) + { + throw new ArgumentNullException(nameof(request)); + } + + PartitionKeyRange partitionKeyRange = null; + if (request.ResourceType.IsPartitioned()) + { + partitionKeyRange = await ResolvePartitionKeyRangeAsync(request, sessionContainer, partitionKeyRangeCache, clientCollectionCache, false); + } + + ISessionToken localSessionToken = null; + if (partitionKeyRange != null) + { + localSessionToken = sessionContainer.ResolvePartitionLocalSessionToken(request, partitionKeyRange.Id); + } + + return localSessionToken.ConvertToString(); + } + + private static async Task ResolvePartitionKeyRangeAsync(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); + + if (request.Headers[HttpConstants.HttpHeaders.PartitionKey] != 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.Headers[HttpConstants.HttpHeaders.PartitionKey], false, collection, 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 null; + } + + // need to refresh cache. Maybe split happened. + return await ResolvePartitionKeyRangeAsync(request, sessionContainer, partitionKeyRangeCache, clientCollectionCache, true); + } + + return partitonKeyRange; + } + // DEVNOTE: This can be replace with ReplicatedResourceClient.IsMasterOperation on next Direct sync internal static bool IsMasterOperation( ResourceType resourceType, diff --git a/Microsoft.Azure.Cosmos/src/Routing/AddressResolver.cs b/Microsoft.Azure.Cosmos/src/Routing/AddressResolver.cs index b1c419924b..ebf7bb4e5e 100644 --- a/Microsoft.Azure.Cosmos/src/Routing/AddressResolver.cs +++ b/Microsoft.Azure.Cosmos/src/Routing/AddressResolver.cs @@ -447,7 +447,7 @@ private async Task 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 TryResolveServerPartitionByPartitionKey( request, "[]", // This corresponds to first partition collectionCacheIsUptoDate, @@ -624,7 +624,7 @@ private async Task TryResolveServerPartitionByPartitionKeyRang return new ResolutionResult(partitionKeyRange, addresses, identity); } - private PartitionKeyRange TryResolveServerPartitionByPartitionKey( + internal static PartitionKeyRange TryResolveServerPartitionByPartitionKey( DocumentServiceRequest request, string partitionKeyString, bool collectionCacheUptoDate, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs new file mode 100644 index 0000000000..ad1f1dbbd6 --- /dev/null +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs @@ -0,0 +1,78 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests +{ + using System; + using System.Net; + using System.Threading; + using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Common; + using Microsoft.Azure.Cosmos.Routing; + using Microsoft.Azure.Documents; + using Microsoft.VisualStudio.TestTools.UnitTesting; + + [TestClass] + public class GatewaySessionTokenTests : BaseCosmosClientHelper + { + private ContainerInternal Container = null; + private const string PartitionKey = "/pk"; + + [TestInitialize] + public async Task TestInitialize() + { + this.cosmosClient = TestCommon.CreateCosmosClient(useGateway: true); + this.database = await this.cosmosClient.CreateDatabaseAsync( + id: Guid.NewGuid().ToString()); + ContainerResponse response = await this.database.CreateContainerAsync( + new ContainerProperties(id: Guid.NewGuid().ToString(), partitionKeyPath: PartitionKey), + throughput: 20000, + cancellationToken: this.cancellationToken); + Assert.IsNotNull(response); + Assert.IsNotNull(response.Container); + Assert.IsNotNull(response.Resource); + this.Container = (ContainerInlineCore)response; + } + + [TestCleanup] + public async Task Cleanup() + { + await base.TestCleanup(); + } + + [TestMethod] + public async Task TestGatewayModelSession() + { + // Create items with different + for (int i = 0; i < 1000; i++) + { + ToDoActivity item = ToDoActivity.CreateRandomToDoActivity(); + item.pk = "Status" + i.ToString(); + item.id = i.ToString(); + ItemResponse itemResponse = await this.Container.CreateItemAsync(item); + Assert.AreEqual(HttpStatusCode.Created, itemResponse.StatusCode); + } + + ContainerProperties containerProperties = await this.Container.GetCachedContainerPropertiesAsync(); + + ISessionContainer sessionContainer = this.cosmosClient.DocumentClient.sessionContainer; + string docLink = "dbs/" + this.database.Id + "/colls/" + containerProperties.Id + "/docs/3"; + Documents.Collections.INameValueCollection headers = new StoreRequestNameValueCollection(); + headers.Set(HttpConstants.HttpHeaders.PartitionKey, "[\"Status3\"]"); + + DocumentServiceRequest request = DocumentServiceRequest.Create(OperationType.Read, ResourceType.Document, docLink, AuthorizationTokenType.PrimaryMasterKey, headers); + string globalSessionToken = sessionContainer.ResolveGlobalSessionToken(request); + Assert.IsTrue(globalSessionToken.Split(',').Length > 1); + + await GatewayStoreModel.ApplySessionTokenAsync(request, + Cosmos.ConsistencyLevel.Session, + sessionContainer, + await this.cosmosClient.DocumentClient.GetPartitionKeyRangeCacheAsync(), + await this.cosmosClient.DocumentClient.GetCollectionCacheAsync()); + + string sessionToken = request.Headers[HttpConstants.HttpHeaders.SessionToken]; + Assert.IsTrue(!string.IsNullOrEmpty(sessionToken) && sessionToken.Split(',').Length == 1); + } + } +} diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs index 503ee03f81..de9b123b82 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs @@ -147,7 +147,7 @@ public async Task TestRetries() } [TestMethod] - public void TestApplySessionForMasterOperation() + public async Task TestApplySessionForMasterOperation() { List resourceTypes = new List() { @@ -190,7 +190,7 @@ public void TestApplySessionForMasterOperation() dsr.Headers.Add(HttpConstants.HttpHeaders.SessionToken, Guid.NewGuid().ToString()); - GatewayStoreModel.ApplySessionToken( + await GatewayStoreModel.ApplySessionTokenAsync( dsr, ConsistencyLevel.Session, new Mock().Object); @@ -211,7 +211,7 @@ public void TestApplySessionForMasterOperation() dsrQueryPlan.Headers.Add(HttpConstants.HttpHeaders.SessionToken, Guid.NewGuid().ToString()); - GatewayStoreModel.ApplySessionToken( + await GatewayStoreModel.ApplySessionTokenAsync( dsrQueryPlan, ConsistencyLevel.Session, new Mock().Object); @@ -220,7 +220,7 @@ public void TestApplySessionForMasterOperation() } [TestMethod] - public void TestApplySessionForDataOperation() + public async Task TestApplySessionForDataOperation() { List resourceTypes = new List() { @@ -257,7 +257,7 @@ public void TestApplySessionForDataOperation() string dsrSessionToken = Guid.NewGuid().ToString(); dsr.Headers.Add(HttpConstants.HttpHeaders.SessionToken, dsrSessionToken); - GatewayStoreModel.ApplySessionToken( + await GatewayStoreModel.ApplySessionTokenAsync( dsr, ConsistencyLevel.Session, new Mock().Object); @@ -274,7 +274,7 @@ public void TestApplySessionForDataOperation() Mock sMock = new Mock(); sMock.Setup(x => x.ResolveGlobalSessionToken(dsrNoSessionToken)).Returns(dsrSessionToken); - GatewayStoreModel.ApplySessionToken( + await GatewayStoreModel.ApplySessionTokenAsync( dsrNoSessionToken, ConsistencyLevel.Session, sMock.Object); @@ -297,7 +297,7 @@ public void TestApplySessionForDataOperation() string sessionToken = Guid.NewGuid().ToString(); dsrSprocExecute.Headers.Add(HttpConstants.HttpHeaders.SessionToken, sessionToken); - GatewayStoreModel.ApplySessionToken( + await GatewayStoreModel.ApplySessionTokenAsync( dsrSprocExecute, ConsistencyLevel.Session, new Mock().Object); From d436acf162d90e5cc3c368a26d3454b47243b8b2 Mon Sep 17 00:00:00 2001 From: Asket Agarwal Date: Mon, 1 Feb 2021 16:11:35 +0530 Subject: [PATCH 02/12] Adding Test for Reading session token after write. --- .../src/GatewayStoreModel.cs | 34 ++++++-- .../GatewaySessionTokenTests.cs | 87 ++++++++++++++++--- .../GatewayStoreModelTest.cs | 9 +- 3 files changed, 113 insertions(+), 17 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs index 86a8c80ce4..a227408f27 100644 --- a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs +++ b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs @@ -264,26 +264,47 @@ internal static async Task ApplySessionTokenAsync( } } - internal static async Task ResolveSessionTokenAsync(DocumentServiceRequest request, ISessionContainer sessionContainer, PartitionKeyRangeCache partitionKeyRangeCache, ClientCollectionCache clientCollectionCache) + internal static async Task ResolveSessionTokenAsync(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)); + } + PartitionKeyRange partitionKeyRange = null; if (request.ResourceType.IsPartitioned()) { partitionKeyRange = await ResolvePartitionKeyRangeAsync(request, sessionContainer, partitionKeyRangeCache, clientCollectionCache, false); } - ISessionToken localSessionToken = null; if (partitionKeyRange != null) { - localSessionToken = sessionContainer.ResolvePartitionLocalSessionToken(request, partitionKeyRange.Id); + string localSessionToken = sessionContainer.ResolvePartitionLocalSessionToken(request, partitionKeyRange.Id).ConvertToString(); + if (!string.IsNullOrEmpty(localSessionToken)) + { + return partitionKeyRange.Id + ":" + localSessionToken; + } } - return localSessionToken.ConvertToString(); + return null; } private static async Task ResolvePartitionKeyRangeAsync(DocumentServiceRequest request, @@ -301,7 +322,8 @@ private static async Task ResolvePartitionKeyRangeAsync(Docum PartitionKeyRange partitonKeyRange = null; ContainerProperties collection = await clientCollectionCache.ResolveCollectionAsync(request, CancellationToken.None); - if (request.Headers[HttpConstants.HttpHeaders.PartitionKey] != null) + string partitionKeyString = request.Headers[HttpConstants.HttpHeaders.PartitionKey]; + if (partitionKeyString != null) { CollectionRoutingMap collectionRoutingMap = await partitionKeyRangeCache.TryLookupAsync(collectionRid: collection.ResourceId, previousValue: null, @@ -316,7 +338,7 @@ private static async Task ResolvePartitionKeyRangeAsync(Docum cancellationToken: CancellationToken.None); } - partitonKeyRange = AddressResolver.TryResolveServerPartitionByPartitionKey(request, request.Headers[HttpConstants.HttpHeaders.PartitionKey], false, collection, collectionRoutingMap); + partitonKeyRange = AddressResolver.TryResolveServerPartitionByPartitionKey(request, partitionKeyString, false, collection, collectionRoutingMap); } else if (request.PartitionKeyRangeIdentity != null) { diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs index ad1f1dbbd6..334318cb54 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs @@ -5,7 +5,9 @@ namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests { using System; + using System.Collections.Generic; using System.Net; + using System.Net.Http; using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Common; @@ -33,6 +35,16 @@ public async Task TestInitialize() Assert.IsNotNull(response.Container); Assert.IsNotNull(response.Resource); this.Container = (ContainerInlineCore)response; + + // Create items with different + for (int i = 0; i < 500; i++) + { + ToDoActivity item = ToDoActivity.CreateRandomToDoActivity(); + item.pk = "Status" + i.ToString(); + item.id = i.ToString(); + ItemResponse itemResponse = await this.Container.CreateItemAsync(item); + Assert.AreEqual(HttpStatusCode.Created, itemResponse.StatusCode); + } } [TestCleanup] @@ -44,16 +56,6 @@ public async Task Cleanup() [TestMethod] public async Task TestGatewayModelSession() { - // Create items with different - for (int i = 0; i < 1000; i++) - { - ToDoActivity item = ToDoActivity.CreateRandomToDoActivity(); - item.pk = "Status" + i.ToString(); - item.id = i.ToString(); - ItemResponse itemResponse = await this.Container.CreateItemAsync(item); - Assert.AreEqual(HttpStatusCode.Created, itemResponse.StatusCode); - } - ContainerProperties containerProperties = await this.Container.GetCachedContainerPropertiesAsync(); ISessionContainer sessionContainer = this.cosmosClient.DocumentClient.sessionContainer; @@ -74,5 +76,70 @@ await this.cosmosClient.DocumentClient.GetPartitionKeyRangeCacheAsync(), string sessionToken = request.Headers[HttpConstants.HttpHeaders.SessionToken]; Assert.IsTrue(!string.IsNullOrEmpty(sessionToken) && sessionToken.Split(',').Length == 1); } + + [TestMethod] + public async Task GatewaySameSessionTokenTest() + { + string createSessionToken = null; + GatewaySessionTokenTests.HttpClientHandlerHelper httpClientHandler = new HttpClientHandlerHelper + { + ResponseCallBack = (result) => + { + HttpResponseMessage response = result.Result; + if (response.StatusCode != HttpStatusCode.Created) + { + return response; + } + + response.Headers.TryGetValues("x-ms-session-token", out IEnumerable sessionTokens); + foreach (string singleToken in sessionTokens) + { + createSessionToken = singleToken; + break; + } + return response; + } + }; + + using (CosmosClient client = TestCommon.CreateCosmosClient(builder => builder + .WithConnectionModeGateway() + .WithConsistencyLevel(Cosmos.ConsistencyLevel.Session) + .WithHttpClientFactory(() => new HttpClient(httpClientHandler)))) + { + Container container = client.GetContainer(this.database.Id, this.Container.Id); + + ToDoActivity item = ToDoActivity.CreateRandomToDoActivity("Status1001", "1001"); + ItemResponse itemResponse = await container.CreateItemAsync(item); + + // Read back the created Item and check if the session token is identical. + string docLink = "dbs/" + this.database.Id + "/colls/" + this.Container.Id + "/docs/1001"; + Documents.Collections.INameValueCollection headers = new StoreRequestNameValueCollection(); + headers.Set(HttpConstants.HttpHeaders.PartitionKey, "[\"Status1001\"]"); + + DocumentServiceRequest request = DocumentServiceRequest.Create(OperationType.Read, ResourceType.Document, docLink, AuthorizationTokenType.PrimaryMasterKey, headers); + await GatewayStoreModel.ApplySessionTokenAsync(request, + Cosmos.ConsistencyLevel.Session, + client.DocumentClient.sessionContainer, + await client.DocumentClient.GetPartitionKeyRangeCacheAsync(), + await client.DocumentClient.GetCollectionCacheAsync()); + + string readSessionToken = request.Headers[HttpConstants.HttpHeaders.SessionToken]; + Assert.AreEqual(readSessionToken, createSessionToken); + } + } + + private class HttpClientHandlerHelper : DelegatingHandler + { + public HttpClientHandlerHelper() : base(new HttpClientHandler()) + { + } + + public Func, HttpResponseMessage> ResponseCallBack { get; set; } + + protected override Task SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) + { + return base.SendAsync(request, cancellationToken).ContinueWith(this.ResponseCallBack); + } + } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs index de9b123b82..681db4291e 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs @@ -279,7 +279,14 @@ await GatewayStoreModel.ApplySessionTokenAsync( ConsistencyLevel.Session, sMock.Object); - Assert.AreEqual(dsrSessionToken, dsrNoSessionToken.Headers[HttpConstants.HttpHeaders.SessionToken]); + if (dsrNoSessionToken.IsReadOnlyRequest) + { + Assert.AreEqual(dsrSessionToken, dsrNoSessionToken.Headers[HttpConstants.HttpHeaders.SessionToken]); + } + else + { + Assert.IsNull(dsrNoSessionToken.Headers[HttpConstants.HttpHeaders.SessionToken]); + } } } From e3c916bb110eaac13c84d98564330efa4a6b47e5 Mon Sep 17 00:00:00 2001 From: Asket Agarwal Date: Tue, 2 Feb 2021 18:37:06 +0530 Subject: [PATCH 03/12] Changing the Properties to private --- .../src/GatewayStoreModel.cs | 24 ++++++++++++------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs index a227408f27..f6e8f2e349 100644 --- a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs +++ b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs @@ -30,8 +30,8 @@ internal class GatewayStoreModel : IStoreModel, IDisposable private GatewayStoreClient gatewayStoreClient; // Caches to resolve the PartitionKeyRange from request. For Session Token Optimization. - internal ClientCollectionCache ClientCollectionCache { get; set; } - internal PartitionKeyRangeCache PartitionKeyRangeCache { get; set; } + private ClientCollectionCache ClientCollectionCache { get; set; } + private PartitionKeyRangeCache PartitionKeyRangeCache { get; set; } public GatewayStoreModel( GlobalEndpointManager endpointManager, @@ -292,7 +292,11 @@ internal static async Task ResolveSessionTokenAsync(DocumentServiceReque PartitionKeyRange partitionKeyRange = null; if (request.ResourceType.IsPartitioned()) { - partitionKeyRange = await ResolvePartitionKeyRangeAsync(request, sessionContainer, partitionKeyRangeCache, clientCollectionCache, false); + partitionKeyRange = await ResolvePartitionKeyRangeAsync(request: request, + sessionContainer: sessionContainer, + partitionKeyRangeCache: partitionKeyRangeCache, + clientCollectionCache: clientCollectionCache, + refreshCache: false); } if (partitionKeyRange != null) @@ -308,10 +312,10 @@ internal static async Task ResolveSessionTokenAsync(DocumentServiceReque } private static async Task ResolvePartitionKeyRangeAsync(DocumentServiceRequest request, - ISessionContainer sessionContainer, - PartitionKeyRangeCache partitionKeyRangeCache, - ClientCollectionCache clientCollectionCache, - bool refreshCache) + ISessionContainer sessionContainer, + PartitionKeyRangeCache partitionKeyRangeCache, + ClientCollectionCache clientCollectionCache, + bool refreshCache) { if (refreshCache) { @@ -338,7 +342,11 @@ private static async Task ResolvePartitionKeyRangeAsync(Docum cancellationToken: CancellationToken.None); } - partitonKeyRange = AddressResolver.TryResolveServerPartitionByPartitionKey(request, partitionKeyString, false, collection, collectionRoutingMap); + partitonKeyRange = AddressResolver.TryResolveServerPartitionByPartitionKey(request: request, + partitionKeyString: partitionKeyString, + collectionCacheUptoDate: false, + collection: collection, + routingMap: collectionRoutingMap); } else if (request.PartitionKeyRangeIdentity != null) { From 08daac3f01bc8400c855b77492145cea45aad1ab Mon Sep 17 00:00:00 2001 From: Asket Agarwal Date: Tue, 2 Feb 2021 20:07:35 +0530 Subject: [PATCH 04/12] Avoiding null returning methods. --- .../src/GatewayStoreModel.cs | 26 ++++++++++--------- 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs index f6e8f2e349..2320536e96 100644 --- a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs +++ b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs @@ -244,16 +244,17 @@ internal static async Task ApplySessionTokenAsync( if (!sessionConsistency || !request.IsReadOnlyRequest) { - 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 } string sessionToken = null; + bool isSuccess = false; if (clientCollectionCache != null && partitionKeyRangeCache != null) { - sessionToken = await ResolveSessionTokenAsync(request, sessionContainer, partitionKeyRangeCache, clientCollectionCache); + (isSuccess, sessionToken) = await TryResolveSessionTokenAsync(request, sessionContainer, partitionKeyRangeCache, clientCollectionCache); } - if (string.IsNullOrEmpty(sessionToken)) + if (!isSuccess) { sessionToken = sessionContainer.ResolveGlobalSessionToken(request); } @@ -264,7 +265,7 @@ internal static async Task ApplySessionTokenAsync( } } - internal static async Task ResolveSessionTokenAsync(DocumentServiceRequest request, + internal static async Task> TryResolveSessionTokenAsync(DocumentServiceRequest request, ISessionContainer sessionContainer, PartitionKeyRangeCache partitionKeyRangeCache, ClientCollectionCache clientCollectionCache) @@ -290,28 +291,29 @@ internal static async Task ResolveSessionTokenAsync(DocumentServiceReque } PartitionKeyRange partitionKeyRange = null; + bool isSuccess = false; if (request.ResourceType.IsPartitioned()) { - partitionKeyRange = await ResolvePartitionKeyRangeAsync(request: request, + (isSuccess, partitionKeyRange) = await TryResolvePartitionKeyRangeAsync(request: request, sessionContainer: sessionContainer, partitionKeyRangeCache: partitionKeyRangeCache, clientCollectionCache: clientCollectionCache, refreshCache: false); } - if (partitionKeyRange != null) + if (isSuccess) { string localSessionToken = sessionContainer.ResolvePartitionLocalSessionToken(request, partitionKeyRange.Id).ConvertToString(); if (!string.IsNullOrEmpty(localSessionToken)) { - return partitionKeyRange.Id + ":" + localSessionToken; + return new Tuple(true, partitionKeyRange.Id + ":" + localSessionToken); } } - return null; + return new Tuple(false, null); } - private static async Task ResolvePartitionKeyRangeAsync(DocumentServiceRequest request, + private static async Task> TryResolvePartitionKeyRangeAsync(DocumentServiceRequest request, ISessionContainer sessionContainer, PartitionKeyRangeCache partitionKeyRangeCache, ClientCollectionCache clientCollectionCache, @@ -360,14 +362,14 @@ private static async Task ResolvePartitionKeyRangeAsync(Docum { if (refreshCache) { - return null; + return new Tuple(false, null); } // need to refresh cache. Maybe split happened. - return await ResolvePartitionKeyRangeAsync(request, sessionContainer, partitionKeyRangeCache, clientCollectionCache, true); + return await TryResolvePartitionKeyRangeAsync(request, sessionContainer, partitionKeyRangeCache, clientCollectionCache, true); } - return partitonKeyRange; + return new Tuple(true, partitonKeyRange); } // DEVNOTE: This can be replace with ReplicatedResourceClient.IsMasterOperation on next Direct sync From 1f82dbc9a4d337453f7a9959c9c2748a9da719e7 Mon Sep 17 00:00:00 2001 From: Asket Agarwal Date: Wed, 3 Feb 2021 20:13:12 +0530 Subject: [PATCH 05/12] Improving Readability --- .../src/GatewayStoreModel.cs | 20 +++++++++---------- .../src/Routing/AddressResolver.cs | 2 +- .../GatewayStoreModelTest.cs | 20 ++++++++++++++----- 3 files changed, 25 insertions(+), 17 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs index 2320536e96..709b662cbe 100644 --- a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs +++ b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs @@ -210,8 +210,8 @@ internal static async Task ApplySessionTokenAsync( DocumentServiceRequest request, ConsistencyLevel defaultConsistencyLevel, ISessionContainer sessionContainer, - PartitionKeyRangeCache partitionKeyRangeCache = null, - ClientCollectionCache clientCollectionCache = null) + PartitionKeyRangeCache partitionKeyRangeCache, + ClientCollectionCache clientCollectionCache) { if (request.Headers == null) { @@ -290,23 +290,21 @@ internal static async Task> TryResolveSessionTokenAsync(Docu throw new ArgumentNullException(nameof(clientCollectionCache)); } - PartitionKeyRange partitionKeyRange = null; - bool isSuccess = false; if (request.ResourceType.IsPartitioned()) { - (isSuccess, partitionKeyRange) = await TryResolvePartitionKeyRangeAsync(request: request, + (bool isSuccess, PartitionKeyRange partitionKeyRange) = await TryResolvePartitionKeyRangeAsync(request: request, sessionContainer: sessionContainer, partitionKeyRangeCache: partitionKeyRangeCache, clientCollectionCache: clientCollectionCache, refreshCache: false); - } - if (isSuccess) - { - string localSessionToken = sessionContainer.ResolvePartitionLocalSessionToken(request, partitionKeyRange.Id).ConvertToString(); - if (!string.IsNullOrEmpty(localSessionToken)) + if (isSuccess) { - return new Tuple(true, partitionKeyRange.Id + ":" + localSessionToken); + string localSessionToken = sessionContainer.ResolvePartitionLocalSessionToken(request, partitionKeyRange.Id).ConvertToString(); + if (!string.IsNullOrEmpty(localSessionToken)) + { + return new Tuple(true, partitionKeyRange.Id + ":" + localSessionToken); + } } } diff --git a/Microsoft.Azure.Cosmos/src/Routing/AddressResolver.cs b/Microsoft.Azure.Cosmos/src/Routing/AddressResolver.cs index ebf7bb4e5e..2c469e6078 100644 --- a/Microsoft.Azure.Cosmos/src/Routing/AddressResolver.cs +++ b/Microsoft.Azure.Cosmos/src/Routing/AddressResolver.cs @@ -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 TryResolveServerPartitionByPartitionKey( + return AddressResolver.TryResolveServerPartitionByPartitionKey( request, "[]", // This corresponds to first partition collectionCacheIsUptoDate, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs index 681db4291e..229d0abc2e 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs @@ -193,7 +193,9 @@ public async Task TestApplySessionForMasterOperation() await GatewayStoreModel.ApplySessionTokenAsync( dsr, ConsistencyLevel.Session, - new Mock().Object); + new Mock().Object, + partitionKeyRangeCache: null, + clientCollectionCache: null); Assert.IsNull(dsr.Headers[HttpConstants.HttpHeaders.SessionToken]); } @@ -214,7 +216,9 @@ await GatewayStoreModel.ApplySessionTokenAsync( await GatewayStoreModel.ApplySessionTokenAsync( dsrQueryPlan, ConsistencyLevel.Session, - new Mock().Object); + new Mock().Object, + partitionKeyRangeCache: null, + clientCollectionCache: null); Assert.IsNull(dsrQueryPlan.Headers[HttpConstants.HttpHeaders.SessionToken]); } @@ -260,7 +264,9 @@ public async Task TestApplySessionForDataOperation() await GatewayStoreModel.ApplySessionTokenAsync( dsr, ConsistencyLevel.Session, - new Mock().Object); + new Mock().Object, + partitionKeyRangeCache: null, + clientCollectionCache: null); Assert.AreEqual(dsrSessionToken, dsr.Headers[HttpConstants.HttpHeaders.SessionToken]); @@ -277,7 +283,9 @@ await GatewayStoreModel.ApplySessionTokenAsync( await GatewayStoreModel.ApplySessionTokenAsync( dsrNoSessionToken, ConsistencyLevel.Session, - sMock.Object); + sMock.Object, + partitionKeyRangeCache: null, + clientCollectionCache: null); if (dsrNoSessionToken.IsReadOnlyRequest) { @@ -307,7 +315,9 @@ await GatewayStoreModel.ApplySessionTokenAsync( await GatewayStoreModel.ApplySessionTokenAsync( dsrSprocExecute, ConsistencyLevel.Session, - new Mock().Object); + new Mock().Object, + partitionKeyRangeCache: null, + clientCollectionCache: null); Assert.AreEqual(sessionToken, dsrSprocExecute.Headers[HttpConstants.HttpHeaders.SessionToken]); } From 5951f33ed5646ab38604214b1a8b36fc94ad4a2e Mon Sep 17 00:00:00 2001 From: Asket Agarwal Date: Mon, 8 Feb 2021 18:36:21 +0530 Subject: [PATCH 06/12] Having an explicit resolve session token method for gateway --- .../src/GatewayStoreModel.cs | 13 +++++---- .../src/SessionContainer.cs | 29 +++++++++++++++++++ 2 files changed, 36 insertions(+), 6 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs index 709b662cbe..1f21354422 100644 --- a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs +++ b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs @@ -13,6 +13,7 @@ namespace Microsoft.Azure.Cosmos 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; @@ -30,8 +31,8 @@ internal class GatewayStoreModel : IStoreModel, IDisposable private GatewayStoreClient gatewayStoreClient; // Caches to resolve the PartitionKeyRange from request. For Session Token Optimization. - private ClientCollectionCache ClientCollectionCache { get; set; } - private PartitionKeyRangeCache PartitionKeyRangeCache { get; set; } + private ClientCollectionCache ClientCollectionCache; + private PartitionKeyRangeCache PartitionKeyRangeCache; public GatewayStoreModel( GlobalEndpointManager endpointManager, @@ -298,12 +299,12 @@ internal static async Task> TryResolveSessionTokenAsync(Docu clientCollectionCache: clientCollectionCache, refreshCache: false); - if (isSuccess) - { - string localSessionToken = sessionContainer.ResolvePartitionLocalSessionToken(request, partitionKeyRange.Id).ConvertToString(); + if (isSuccess && sessionContainer is SessionContainer gatewaySessionContainer) + { + string localSessionToken = gatewaySessionContainer.ResolvePartitionLocalSessionTokenForGateway(request, partitionKeyRange.Id); if (!string.IsNullOrEmpty(localSessionToken)) { - return new Tuple(true, partitionKeyRange.Id + ":" + localSessionToken); + return new Tuple(true, localSessionToken); } } } diff --git a/Microsoft.Azure.Cosmos/src/SessionContainer.cs b/Microsoft.Azure.Cosmos/src/SessionContainer.cs index 93f1de9fdc..84563014c0 100644 --- a/Microsoft.Azure.Cosmos/src/SessionContainer.cs +++ b/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 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)) From dd4b7e9218a32ee62d8630eb103ba01a58f12edf Mon Sep 17 00:00:00 2001 From: Asket Agarwal Date: Tue, 16 Feb 2021 20:15:11 +0530 Subject: [PATCH 07/12] Improving styling and readability --- .../src/GatewayStoreModel.cs | 27 ++++++++--------- .../src/Routing/AddressResolver.cs | 2 +- .../src/SessionContainer.cs | 7 +++-- .../GatewaySessionTokenTests.cs | 2 -- .../GatewayStoreModelTest.cs | 30 ++++++++++++------- 5 files changed, 38 insertions(+), 30 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs index 1f21354422..5ab037be20 100644 --- a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs +++ b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs @@ -31,8 +31,8 @@ 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; + private ClientCollectionCache clientCollectionCache; + private PartitionKeyRangeCache partitionKeyRangeCache; public GatewayStoreModel( GlobalEndpointManager endpointManager, @@ -59,8 +59,8 @@ await GatewayStoreModel.ApplySessionTokenAsync( request, this.defaultConsistencyLevel, this.sessionContainer, - this.PartitionKeyRangeCache, - this.ClientCollectionCache); + this.partitionKeyRangeCache, + this.clientCollectionCache); DocumentServiceResponse response; try @@ -151,10 +151,11 @@ public virtual async Task GetDatabaseAccountAsync(Func TryResolveServerPartitionAsync( object effectivePartitionKeyStringObject = null; if (partitionKeyString != null) { - range = TryResolveServerPartitionByPartitionKey( + range = AddressResolver.TryResolveServerPartitionByPartitionKey( request, partitionKeyString, collectionCacheIsUptodate, diff --git a/Microsoft.Azure.Cosmos/src/SessionContainer.cs b/Microsoft.Azure.Cosmos/src/SessionContainer.cs index 84563014c0..9fbd3f88c8 100644 --- a/Microsoft.Azure.Cosmos/src/SessionContainer.cs +++ b/Microsoft.Azure.Cosmos/src/SessionContainer.cs @@ -144,7 +144,9 @@ private static ISessionToken ResolvePartitionLocalSessionToken(SessionContainerS return SessionTokenHelper.ResolvePartitionLocalSessionToken(request, partitionKeyRangeId, SessionContainer.GetPartitionKeyRangeIdToTokenMap(self, request)); } - private static string ResolvePartitionLocalSessionTokenForGateway(SessionContainerState self, DocumentServiceRequest request, string partitionKeyRangeId) + private static string ResolvePartitionLocalSessionTokenForGateway(SessionContainerState self, + DocumentServiceRequest request, + string partitionKeyRangeId) { ConcurrentDictionary partitionKeyRangeIdToTokenMap = SessionContainer.GetPartitionKeyRangeIdToTokenMap(self, request); if (partitionKeyRangeIdToTokenMap != null) @@ -157,7 +159,8 @@ private static string ResolvePartitionLocalSessionTokenForGateway(SessionContain { for (int parentIndex = request.RequestContext.ResolvedPartitionKeyRange.Parents.Count - 1; parentIndex >= 0; parentIndex--) { - if (partitionKeyRangeIdToTokenMap.TryGetValue(request.RequestContext.ResolvedPartitionKeyRange.Parents[parentIndex], out sessionToken)) + if (partitionKeyRangeIdToTokenMap.TryGetValue(request.RequestContext.ResolvedPartitionKeyRange.Parents[parentIndex], + out sessionToken)) { return partitionKeyRangeId + ":" + sessionToken.ConvertToString(); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs index 334318cb54..0fa751819b 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs @@ -10,8 +10,6 @@ namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests using System.Net.Http; using System.Threading; using System.Threading.Tasks; - using Microsoft.Azure.Cosmos.Common; - using Microsoft.Azure.Cosmos.Routing; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs index 229d0abc2e..0bee04db2b 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs @@ -194,8 +194,8 @@ await GatewayStoreModel.ApplySessionTokenAsync( dsr, ConsistencyLevel.Session, new Mock().Object, - partitionKeyRangeCache: null, - clientCollectionCache: null); + partitionKeyRangeCache: new Mock(null, null, null).Object, + clientCollectionCache: new Mock(new SessionContainer("testhost"), this.GetGatewayStoreModelForConsistencyTest(), null, null).Object); Assert.IsNull(dsr.Headers[HttpConstants.HttpHeaders.SessionToken]); } @@ -217,8 +217,9 @@ await GatewayStoreModel.ApplySessionTokenAsync( dsrQueryPlan, ConsistencyLevel.Session, new Mock().Object, - partitionKeyRangeCache: null, - clientCollectionCache: null); + partitionKeyRangeCache: new Mock(null, null, null).Object, + clientCollectionCache: new Mock(new SessionContainer("testhost"), this.GetGatewayStoreModelForConsistencyTest(), null, null).Object); + Assert.IsNull(dsrQueryPlan.Headers[HttpConstants.HttpHeaders.SessionToken]); } @@ -265,8 +266,9 @@ await GatewayStoreModel.ApplySessionTokenAsync( dsr, ConsistencyLevel.Session, new Mock().Object, - partitionKeyRangeCache: null, - clientCollectionCache: null); + partitionKeyRangeCache: new Mock(null, null, null).Object, + clientCollectionCache: new Mock(new SessionContainer("testhost"), this.GetGatewayStoreModelForConsistencyTest(), null, null).Object); + Assert.AreEqual(dsrSessionToken, dsr.Headers[HttpConstants.HttpHeaders.SessionToken]); @@ -284,10 +286,11 @@ await GatewayStoreModel.ApplySessionTokenAsync( dsrNoSessionToken, ConsistencyLevel.Session, sMock.Object, - partitionKeyRangeCache: null, - clientCollectionCache: null); + partitionKeyRangeCache: new Mock(null, null, null).Object, + clientCollectionCache: new Mock(new SessionContainer("testhost"), this.GetGatewayStoreModelForConsistencyTest(), null, null).Object); + - if (dsrNoSessionToken.IsReadOnlyRequest) + if (dsrNoSessionToken.IsReadOnlyRequest || dsrNoSessionToken.OperationType == OperationType.Batch) { Assert.AreEqual(dsrSessionToken, dsrNoSessionToken.Headers[HttpConstants.HttpHeaders.SessionToken]); } @@ -316,8 +319,9 @@ await GatewayStoreModel.ApplySessionTokenAsync( dsrSprocExecute, ConsistencyLevel.Session, new Mock().Object, - partitionKeyRangeCache: null, - clientCollectionCache: null); + partitionKeyRangeCache: new Mock(null, null, null).Object, + clientCollectionCache: new Mock(new SessionContainer("testhost"), this.GetGatewayStoreModelForConsistencyTest(), null, null).Object); + Assert.AreEqual(sessionToken, dsrSprocExecute.Headers[HttpConstants.HttpHeaders.SessionToken]); } @@ -768,6 +772,10 @@ private GatewayStoreModel GetGatewayStoreModelForConsistencyTest() null, MockCosmosUtil.CreateCosmosHttpClient(() => new HttpClient(httpMessageHandler))); + ClientCollectionCache clientCollectionCache = new Mock(new SessionContainer("testhost"), storeModel, null, null).Object; + PartitionKeyRangeCache partitionKeyRangeCache = new Mock(null, storeModel, clientCollectionCache).Object; + storeModel.SetCaches(partitionKeyRangeCache, clientCollectionCache); + return storeModel; } From e59f07e417ffe883ce84bb0269300080f391151a Mon Sep 17 00:00:00 2001 From: Asket Agarwal Date: Tue, 16 Feb 2021 21:16:52 +0530 Subject: [PATCH 08/12] MErging latest changes --- .../GatewaySessionTokenTests.cs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs index 0fa751819b..33b0b82933 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs @@ -10,6 +10,7 @@ namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests using System.Net.Http; using System.Threading; using System.Threading.Tasks; + using Microsoft.Azure.Cosmos.Tracing; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; @@ -54,7 +55,7 @@ public async Task Cleanup() [TestMethod] public async Task TestGatewayModelSession() { - ContainerProperties containerProperties = await this.Container.GetCachedContainerPropertiesAsync(); + ContainerProperties containerProperties = await this.Container.GetCachedContainerPropertiesAsync(false, NoOpTrace.Singleton, CancellationToken.None); ISessionContainer sessionContainer = this.cosmosClient.DocumentClient.sessionContainer; string docLink = "dbs/" + this.database.Id + "/colls/" + containerProperties.Id + "/docs/3"; @@ -69,7 +70,7 @@ await GatewayStoreModel.ApplySessionTokenAsync(request, Cosmos.ConsistencyLevel.Session, sessionContainer, await this.cosmosClient.DocumentClient.GetPartitionKeyRangeCacheAsync(), - await this.cosmosClient.DocumentClient.GetCollectionCacheAsync()); + await this.cosmosClient.DocumentClient.GetCollectionCacheAsync(NoOpTrace.Singleton)); string sessionToken = request.Headers[HttpConstants.HttpHeaders.SessionToken]; Assert.IsTrue(!string.IsNullOrEmpty(sessionToken) && sessionToken.Split(',').Length == 1); @@ -119,7 +120,7 @@ await GatewayStoreModel.ApplySessionTokenAsync(request, Cosmos.ConsistencyLevel.Session, client.DocumentClient.sessionContainer, await client.DocumentClient.GetPartitionKeyRangeCacheAsync(), - await client.DocumentClient.GetCollectionCacheAsync()); + await client.DocumentClient.GetCollectionCacheAsync(NoOpTrace.Singleton)); string readSessionToken = request.Headers[HttpConstants.HttpHeaders.SessionToken]; Assert.AreEqual(readSessionToken, createSessionToken); From f1d56e50248676c7a591eead4dc2bb80c1d07d05 Mon Sep 17 00:00:00 2001 From: Asket Agarwal Date: Wed, 17 Mar 2021 18:57:55 +0530 Subject: [PATCH 09/12] nitpicks --- .../src/GatewayStoreModel.cs | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs index 5ab037be20..db14d25d35 100644 --- a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs +++ b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs @@ -16,6 +16,7 @@ namespace Microsoft.Azure.Cosmos using Microsoft.Azure.Cosmos.Common; using Microsoft.Azure.Cosmos.Core.Trace; using Microsoft.Azure.Cosmos.Routing; + using Microsoft.Azure.Cosmos.Tracing; using Microsoft.Azure.Documents; using Microsoft.Azure.Documents.Collections; using Newtonsoft.Json; @@ -324,7 +325,7 @@ private static async Task> TryResolvePartitionKey } PartitionKeyRange partitonKeyRange = null; - ContainerProperties collection = await clientCollectionCache.ResolveCollectionAsync(request, CancellationToken.None); + ContainerProperties collection = await clientCollectionCache.ResolveCollectionAsync(request, CancellationToken.None, NoOpTrace.Singleton); string partitionKeyString = request.Headers[HttpConstants.HttpHeaders.PartitionKey]; if (partitionKeyString != null) @@ -332,14 +333,16 @@ private static async Task> TryResolvePartitionKey CollectionRoutingMap collectionRoutingMap = await partitionKeyRangeCache.TryLookupAsync(collectionRid: collection.ResourceId, previousValue: null, request: request, - cancellationToken: CancellationToken.None); + cancellationToken: CancellationToken.None, + NoOpTrace.Singleton); if (refreshCache && collectionRoutingMap != null) { collectionRoutingMap = await partitionKeyRangeCache.TryLookupAsync(collectionRid: collection.ResourceId, previousValue: collectionRoutingMap, request: request, - cancellationToken: CancellationToken.None); + cancellationToken: CancellationToken.None, + NoOpTrace.Singleton); } partitonKeyRange = AddressResolver.TryResolveServerPartitionByPartitionKey(request: request, @@ -352,7 +355,8 @@ private static async Task> TryResolvePartitionKey { PartitionKeyRangeIdentity partitionKeyRangeId = request.PartitionKeyRangeIdentity; partitonKeyRange = await partitionKeyRangeCache.TryGetPartitionKeyRangeByIdAsync(collection.ResourceId, - partitionKeyRangeId.ToString(), + partitionKeyRangeId.ToString(), + NoOpTrace.Singleton, refreshCache); } @@ -364,7 +368,11 @@ private static async Task> TryResolvePartitionKey } // need to refresh cache. Maybe split happened. - return await TryResolvePartitionKeyRangeAsync(request, sessionContainer, partitionKeyRangeCache, clientCollectionCache, true); + return await GatewayStoreModel.TryResolvePartitionKeyRangeAsync(request: request, + sessionContainer: sessionContainer, + partitionKeyRangeCache: partitionKeyRangeCache, + clientCollectionCache: clientCollectionCache, + refreshCache: true); } return new Tuple(true, partitonKeyRange); From 13474de88ece6b85eba836a2ba42c273d0348505 Mon Sep 17 00:00:00 2001 From: Asket Agarwal Date: Tue, 23 Mar 2021 18:38:38 +0530 Subject: [PATCH 10/12] Fixing build --- .../GatewaySessionTokenTests.cs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs index 33b0b82933..94b1602675 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs @@ -69,7 +69,7 @@ public async Task TestGatewayModelSession() await GatewayStoreModel.ApplySessionTokenAsync(request, Cosmos.ConsistencyLevel.Session, sessionContainer, - await this.cosmosClient.DocumentClient.GetPartitionKeyRangeCacheAsync(), + await this.cosmosClient.DocumentClient.GetPartitionKeyRangeCacheAsync(NoOpTrace.Singleton), await this.cosmosClient.DocumentClient.GetCollectionCacheAsync(NoOpTrace.Singleton)); string sessionToken = request.Headers[HttpConstants.HttpHeaders.SessionToken]; @@ -119,7 +119,7 @@ public async Task GatewaySameSessionTokenTest() await GatewayStoreModel.ApplySessionTokenAsync(request, Cosmos.ConsistencyLevel.Session, client.DocumentClient.sessionContainer, - await client.DocumentClient.GetPartitionKeyRangeCacheAsync(), + await client.DocumentClient.GetPartitionKeyRangeCacheAsync(NoOpTrace.Singleton), await client.DocumentClient.GetCollectionCacheAsync(NoOpTrace.Singleton)); string readSessionToken = request.Headers[HttpConstants.HttpHeaders.SessionToken]; From aa7f1b905812eed383409c339b9326dcc566b817 Mon Sep 17 00:00:00 2001 From: Asket Agarwal Date: Mon, 29 Mar 2021 10:57:38 +0530 Subject: [PATCH 11/12] Adding Resolved PArtition to request --- Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs index db14d25d35..f18af38a98 100644 --- a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs +++ b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs @@ -300,7 +300,8 @@ internal static async Task> TryResolveSessionTokenAsync(Docu refreshCache: false); if (isSuccess && sessionContainer is SessionContainer gatewaySessionContainer) - { + { + request.RequestContext.ResolvedPartitionKeyRange = partitionKeyRange; string localSessionToken = gatewaySessionContainer.ResolvePartitionLocalSessionTokenForGateway(request, partitionKeyRange.Id); if (!string.IsNullOrEmpty(localSessionToken)) { From 5a47553b06ec63565ed0a6da221016d4775869a2 Mon Sep 17 00:00:00 2001 From: Asket Agarwal Date: Mon, 29 Mar 2021 10:59:05 +0530 Subject: [PATCH 12/12] adding class name to static method --- Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs index f18af38a98..a0ecdfc258 100644 --- a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs +++ b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs @@ -250,7 +250,7 @@ internal static async Task ApplySessionTokenAsync( return; // Only apply the session token in case of session consistency and the request is read only } - (bool isSuccess, string sessionToken) = await TryResolveSessionTokenAsync(request, + (bool isSuccess, string sessionToken) = await GatewayStoreModel.TryResolveSessionTokenAsync(request, sessionContainer, partitionKeyRangeCache, clientCollectionCache);