From 67cbd89757b778c6d6bd832c5b55e25cc909699d Mon Sep 17 00:00:00 2001 From: Asket Agarwal Date: Sat, 15 Jan 2022 02:50:30 +0530 Subject: [PATCH] =?UTF-8?q?Session:=20Removes=20Global=20Session=20Token?= =?UTF-8?q?=20on=20GW=20request=20if=20we=20can't=20resolve=20the=20scoped?= =?UTF-8?q?=20session=20t=E2=80=A6=20(#2975)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The .NET SDK is sending the global session token if it does not have the partition key range id in the session token cache. This can cause the request to fail from the header being too large. Once the session token is added to the session token cache it is properly filtered for future requests. This impacts both point operations and query for the first operation going to that partition key range id. Solution: If there is no session token for the specificied partition key range id or the parent range ids then no session token should be sent because it is the first request to that partition range id. --- .../src/GatewayStoreModel.cs | 7 +- .../GatewayStoreModelTest.cs | 162 +++++++++++++----- 2 files changed, 124 insertions(+), 45 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs index d5623e9a8e..e299af088a 100644 --- a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs +++ b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs @@ -288,12 +288,7 @@ internal static async Task ApplySessionTokenAsync( partitionKeyRangeCache, clientCollectionCache); - if (!isSuccess) - { - sessionToken = sessionContainer.ResolveGlobalSessionToken(request); - } - - if (!string.IsNullOrEmpty(sessionToken)) + if (isSuccess && !string.IsNullOrEmpty(sessionToken)) { request.Headers[HttpConstants.HttpHeaders.SessionToken] = sessionToken; } 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 d832c70cbe..f5a435e873 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs @@ -23,6 +23,7 @@ namespace Microsoft.Azure.Cosmos using Microsoft.Azure.Cosmos.Tests; using Microsoft.Azure.Cosmos.Tracing; using Newtonsoft.Json; + using System.Collections.ObjectModel; /// /// Tests for . @@ -253,8 +254,7 @@ public async Task TestApplySessionForDataOperation(bool multiMaster) List resourceTypes = new List() { ResourceType.Document, - ResourceType.Conflict, - ResourceType.Batch + ResourceType.Conflict }; List operationTypes = new List() @@ -263,7 +263,8 @@ public async Task TestApplySessionForDataOperation(bool multiMaster) OperationType.Delete, OperationType.Read, OperationType.Upsert, - OperationType.Replace + OperationType.Replace, + OperationType.Batch }; foreach (ResourceType resourceType in resourceTypes) @@ -301,31 +302,33 @@ await GatewayStoreModel.ApplySessionTokenAsync( { // Verify when user does not set session token - DocumentServiceRequest dsrNoSessionToken = DocumentServiceRequest.CreateFromName( - operationType, - "Test", - resourceType, - AuthorizationTokenType.PrimaryMasterKey); + DocumentServiceRequest dsrNoSessionToken = DocumentServiceRequest.Create(operationType, + resourceType, + new Uri("https://foo.com/dbs/db1/colls/coll1", UriKind.Absolute), + AuthorizationTokenType.PrimaryMasterKey); - string dsrSessionToken = Guid.NewGuid().ToString(); - Mock sMock = new Mock(); - sMock.Setup(x => x.ResolveGlobalSessionToken(dsrNoSessionToken)).Returns(dsrSessionToken); + SessionContainer sessionContainer = new SessionContainer(string.Empty); + sessionContainer.SetSessionToken( + ResourceId.NewDocumentCollectionId(42, 129).DocumentCollectionId.ToString(), + "dbs/db1/colls/coll1", + new StoreRequestNameValueCollection() { { HttpConstants.HttpHeaders.SessionToken, "range_1:1#9#4=8#5=7" } }); Mock globalEndpointManager = new Mock(); globalEndpointManager.Setup(gem => gem.CanUseMultipleWriteLocations(It.Is(drs => drs == dsrNoSessionToken))).Returns(multiMaster); await this.GetGatewayStoreModelForConsistencyTest(async (gatewayStoreModel) => { + dsrNoSessionToken.RequestContext.ResolvedPartitionKeyRange = new PartitionKeyRange { Id = "range_1" }; await GatewayStoreModel.ApplySessionTokenAsync( dsrNoSessionToken, ConsistencyLevel.Session, - sMock.Object, + sessionContainer, partitionKeyRangeCache: new Mock(null, null, null).Object, clientCollectionCache: new Mock(new SessionContainer("testhost"), gatewayStoreModel, null, null).Object, globalEndpointManager: globalEndpointManager.Object); if (dsrNoSessionToken.IsReadOnlyRequest || dsrNoSessionToken.OperationType == OperationType.Batch || multiMaster) { - Assert.AreEqual(dsrSessionToken, dsrNoSessionToken.Headers[HttpConstants.HttpHeaders.SessionToken]); + Assert.AreEqual("range_1:1#9#4=8#5=7", dsrNoSessionToken.Headers[HttpConstants.HttpHeaders.SessionToken]); } else { @@ -336,18 +339,19 @@ await GatewayStoreModel.ApplySessionTokenAsync( { // Verify when partition key range is configured - DocumentServiceRequest dsr = DocumentServiceRequest.CreateFromName( - operationType, - "Test", - resourceType, - AuthorizationTokenType.PrimaryMasterKey); - - string partitionKeyRangeId = "1"; - dsr.Headers[WFConstants.BackendHeaders.PartitionKeyRangeId] = new PartitionKeyRangeIdentity(partitionKeyRangeId).ToHeader(); - - string dsrSessionToken = Guid.NewGuid().ToString(); - Mock sMock = new Mock(); - sMock.Setup(x => x.ResolveGlobalSessionToken(dsr)).Returns(dsrSessionToken); + string partitionKeyRangeId = "range_1"; + DocumentServiceRequest dsr = DocumentServiceRequest.Create(operationType, + resourceType, + new Uri("https://foo.com/dbs/db1/colls/coll1", UriKind.Absolute), + AuthorizationTokenType.PrimaryMasterKey, + new StoreRequestNameValueCollection() { { WFConstants.BackendHeaders.PartitionKeyRangeId, new PartitionKeyRangeIdentity(partitionKeyRangeId).ToHeader() } }); + + + SessionContainer sessionContainer = new SessionContainer(string.Empty); + sessionContainer.SetSessionToken( + ResourceId.NewDocumentCollectionId(42, 129).DocumentCollectionId.ToString(), + "dbs/db1/colls/coll1", + new StoreRequestNameValueCollection() { { HttpConstants.HttpHeaders.SessionToken, "range_1:1#9#4=8#5=7" } }); ContainerProperties containerProperties = ContainerProperties.CreateWithResourceId("ccZ1ANCszwk="); containerProperties.Id = "TestId"; @@ -364,19 +368,19 @@ await GatewayStoreModel.ApplySessionTokenAsync( containerProperties.ResourceId, partitionKeyRangeId, It.IsAny(), - false)).Returns(Task.FromResult(new PartitionKeyRange())); + false)).Returns(Task.FromResult(new PartitionKeyRange { Id = "range_1" })); await GatewayStoreModel.ApplySessionTokenAsync( dsr, ConsistencyLevel.Session, - sMock.Object, + sessionContainer, partitionKeyRangeCache: mockPartitionKeyRangeCache.Object, clientCollectionCache: mockCollectionCahce.Object, globalEndpointManager: Mock.Of()); if (dsr.IsReadOnlyRequest || dsr.OperationType == OperationType.Batch) { - Assert.AreEqual(dsrSessionToken, dsr.Headers[HttpConstants.HttpHeaders.SessionToken]); + Assert.AreEqual("range_1:1#9#4=8#5=7", dsr.Headers[HttpConstants.HttpHeaders.SessionToken]); } else { @@ -424,25 +428,27 @@ public async Task TestRequestOverloadRemovesSessionToken(bool multiMaster, bool INameValueCollection headers = new StoreRequestNameValueCollection(); headers.Set(HttpConstants.HttpHeaders.ConsistencyLevel, ConsistencyLevel.Eventual.ToString()); - DocumentServiceRequest dsrNoSessionToken = DocumentServiceRequest.CreateFromName( - isWriteRequest ? OperationType.Create : OperationType.Read, - "Test", - ResourceType.Document, - AuthorizationTokenType.PrimaryMasterKey, - headers); + DocumentServiceRequest dsrNoSessionToken = DocumentServiceRequest.Create(isWriteRequest ? OperationType.Create : OperationType.Read, + ResourceType.Document, + new Uri("https://foo.com/dbs/db1/colls/coll1/docs/doc1", UriKind.Absolute), + AuthorizationTokenType.PrimaryMasterKey, + headers); - string dsrSessionToken = Guid.NewGuid().ToString(); - Mock sMock = new Mock(); - sMock.Setup(x => x.ResolveGlobalSessionToken(dsrNoSessionToken)).Returns(dsrSessionToken); + SessionContainer sessionContainer = new SessionContainer(string.Empty); + sessionContainer.SetSessionToken( + ResourceId.NewDocumentCollectionId(42, 129).DocumentCollectionId.ToString(), + "dbs/db1/colls/coll1", + new StoreRequestNameValueCollection() { { HttpConstants.HttpHeaders.SessionToken, "range_1:1#9#4=8#5=7" } }); Mock globalEndpointManager = new Mock(); globalEndpointManager.Setup(gem => gem.CanUseMultipleWriteLocations(It.Is(drs => drs == dsrNoSessionToken))).Returns(multiMaster); await this.GetGatewayStoreModelForConsistencyTest(async (gatewayStoreModel) => { + dsrNoSessionToken.RequestContext.ResolvedPartitionKeyRange = new PartitionKeyRange { Id = "range_1" }; await GatewayStoreModel.ApplySessionTokenAsync( dsrNoSessionToken, ConsistencyLevel.Session, - sMock.Object, + sessionContainer, partitionKeyRangeCache: new Mock(null, null, null).Object, clientCollectionCache: new Mock(new SessionContainer("testhost"), gatewayStoreModel, null, null).Object, globalEndpointManager: globalEndpointManager.Object); @@ -450,7 +456,7 @@ await GatewayStoreModel.ApplySessionTokenAsync( if (isWriteRequest && multiMaster) { // Multi master write requests should not lower the consistency and remove the session token - Assert.AreEqual(dsrSessionToken, dsrNoSessionToken.Headers[HttpConstants.HttpHeaders.SessionToken]); + Assert.AreEqual("range_1:1#9#4=8#5=7", dsrNoSessionToken.Headers[HttpConstants.HttpHeaders.SessionToken]); } else { @@ -881,6 +887,83 @@ private async Task GatewayStoreModel_Exceptionless_NotUpdateSessionTokenOnKnownR } } + [TestMethod] + [Owner("askagarw")] + public async Task GatewayStoreModel_AvoidGlobalSessionToken() + { + Mock mockDocumentClient = new Mock(); + mockDocumentClient.Setup(client => client.ServiceEndpoint).Returns(new Uri("https://foo")); + using GlobalEndpointManager endpointManager = new GlobalEndpointManager(mockDocumentClient.Object, new ConnectionPolicy()); + SessionContainer sessionContainer = new SessionContainer(string.Empty); + DocumentClientEventSource eventSource = DocumentClientEventSource.Instance; + using GatewayStoreModel storeModel = new GatewayStoreModel( + endpointManager, + sessionContainer, + ConsistencyLevel.Session, + eventSource, + null, + MockCosmosUtil.CreateCosmosHttpClient(() => new HttpClient())); + Mock clientCollectionCache = new Mock(new SessionContainer("testhost"), storeModel, null, null); + Mock partitionKeyRangeCache = new Mock(null, storeModel, clientCollectionCache.Object); + + sessionContainer.SetSessionToken( + ResourceId.NewDocumentCollectionId(42, 129).DocumentCollectionId.ToString(), + "dbs/db1/colls/coll1", + new StoreRequestNameValueCollection() { { HttpConstants.HttpHeaders.SessionToken, "range_1:1#9#4=8#5=7" } }); + + using (DocumentServiceRequest dsr = DocumentServiceRequest.Create(OperationType.Query, + ResourceType.Document, + new Uri("https://foo.com/dbs/db1/colls/coll1", UriKind.Absolute), + AuthorizationTokenType.PrimaryMasterKey)) + { + // pkrange 1 : which has session token + dsr.RequestContext.ResolvedPartitionKeyRange = new PartitionKeyRange { Id = "range_1" }; + await GatewayStoreModel.ApplySessionTokenAsync(dsr, + ConsistencyLevel.Session, + sessionContainer, + partitionKeyRangeCache.Object, + clientCollectionCache.Object, + endpointManager); + Assert.AreEqual(dsr.Headers[HttpConstants.HttpHeaders.SessionToken], "range_1:1#9#4=8#5=7"); + } + + using (DocumentServiceRequest dsr = DocumentServiceRequest.Create(OperationType.Query, + ResourceType.Document, + new Uri("https://foo.com/dbs/db1/colls/coll1", UriKind.Absolute), + AuthorizationTokenType.PrimaryMasterKey)) + { + // pkrange 2 : which has no session token + dsr.RequestContext.ResolvedPartitionKeyRange = new PartitionKeyRange { Id = "range_2" }; + await GatewayStoreModel.ApplySessionTokenAsync(dsr, + ConsistencyLevel.Session, + sessionContainer, + partitionKeyRangeCache.Object, + clientCollectionCache.Object, + endpointManager); + Assert.AreEqual(dsr.Headers[HttpConstants.HttpHeaders.SessionToken], null); + + // There exists global session token, but we do not use it + Assert.AreEqual(sessionContainer.ResolveGlobalSessionToken(dsr), "range_1:1#9#4=8#5=7"); + } + + using (DocumentServiceRequest dsr = DocumentServiceRequest.Create(OperationType.Query, + ResourceType.Document, + new Uri("https://foo.com/dbs/db1/colls/coll1", UriKind.Absolute), + AuthorizationTokenType.PrimaryMasterKey)) + { + // pkrange 3 : Split scenario where session token exists for parent of pk range + Collection parents = new Collection { "range_1" }; + dsr.RequestContext.ResolvedPartitionKeyRange = new PartitionKeyRange { Id = "range_3", Parents = parents }; + await GatewayStoreModel.ApplySessionTokenAsync(dsr, + ConsistencyLevel.Session, + sessionContainer, + partitionKeyRangeCache.Object, + clientCollectionCache.Object, + endpointManager); + Assert.AreEqual(dsr.Headers[HttpConstants.HttpHeaders.SessionToken], "range_3:1#9#4=8#5=7"); + } + } + /// /// When the response contains a PKRangeId header different than the one targeted with the session token, trigger a refresh of the PKRange cache /// @@ -1023,6 +1106,7 @@ private async Task TestGatewayStoreModelProcessMessageAsync(GatewayStoreModel st await storeModel.ProcessMessageAsync(request); request.Headers.Remove("x-ms-session-token"); request.Headers["x-ms-consistency-level"] = "Session"; + request.RequestContext.ResolvedPartitionKeyRange = new PartitionKeyRange { Id = "range_0" }; await storeModel.ProcessMessageAsync(request); } }