From 37a1146b7e95e1c9b9ddd0f28d4c70cd6cf10b80 Mon Sep 17 00:00:00 2001 From: Asket Agarwal Date: Tue, 11 Jan 2022 02:28:32 +0530 Subject: [PATCH 1/2] Avoid Global Session Token on GW if we can't resolve scoped session token --- .../src/GatewayStoreModel.cs | 7 +--- .../GatewaySessionTokenTests.cs | 34 +++++++++++++++++++ 2 files changed, 35 insertions(+), 6 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.EmulatorTests/GatewaySessionTokenTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs index 9f5a353cb3..c6dffeb475 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs @@ -147,5 +147,39 @@ await client.DocumentClient.GetCollectionCacheAsync(NoOpTrace.Singleton), Assert.AreEqual(readSessionToken, createSessionToken); } } + + [TestMethod] + public async Task FirstCallToPartitionAvoidsSessionTokenTestAsync() + { + bool isFirstRequestToPartition = true; + HttpClientHandlerHelper httpClientHandlerHelper = new HttpClientHandlerHelper + { + RequestCallBack = (request, cancellationToken) => + { + if (isFirstRequestToPartition) + { + Assert.IsFalse(request.Headers.Contains(HttpConstants.HttpHeaders.SessionToken)); + } + else + { + Assert.IsTrue(request.Headers.Contains(HttpConstants.HttpHeaders.SessionToken)); + } + + return null; + } + }; + + using (CosmosClient client = TestCommon.CreateCosmosClient(builder => builder + .WithConnectionModeGateway() + .WithConsistencyLevel(Cosmos.ConsistencyLevel.Session) + .WithHttpClientFactory(() => new HttpClient(httpClientHandlerHelper)))) + { + Container container = client.GetContainer(this.database.Id, this.Container.Id); + await container.ReadItemAsync("1", new Cosmos.PartitionKey("Status1")); + await container.ReadItemAsync("3", new Cosmos.PartitionKey("Status3")); + isFirstRequestToPartition = false; + await container.ReadItemAsync("3", new Cosmos.PartitionKey("Status3")); + } + } } } From 7b916e62171b3b795987bcf523444ecfc7043b43 Mon Sep 17 00:00:00 2001 From: Asket Agarwal Date: Thu, 13 Jan 2022 22:07:07 +0530 Subject: [PATCH 2/2] fixing gatewaystoremodel tests --- .../GatewaySessionTokenTests.cs | 34 ---- .../GatewayStoreModelTest.cs | 162 +++++++++++++----- 2 files changed, 123 insertions(+), 73 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 c6dffeb475..9f5a353cb3 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/GatewaySessionTokenTests.cs @@ -147,39 +147,5 @@ await client.DocumentClient.GetCollectionCacheAsync(NoOpTrace.Singleton), Assert.AreEqual(readSessionToken, createSessionToken); } } - - [TestMethod] - public async Task FirstCallToPartitionAvoidsSessionTokenTestAsync() - { - bool isFirstRequestToPartition = true; - HttpClientHandlerHelper httpClientHandlerHelper = new HttpClientHandlerHelper - { - RequestCallBack = (request, cancellationToken) => - { - if (isFirstRequestToPartition) - { - Assert.IsFalse(request.Headers.Contains(HttpConstants.HttpHeaders.SessionToken)); - } - else - { - Assert.IsTrue(request.Headers.Contains(HttpConstants.HttpHeaders.SessionToken)); - } - - return null; - } - }; - - using (CosmosClient client = TestCommon.CreateCosmosClient(builder => builder - .WithConnectionModeGateway() - .WithConsistencyLevel(Cosmos.ConsistencyLevel.Session) - .WithHttpClientFactory(() => new HttpClient(httpClientHandlerHelper)))) - { - Container container = client.GetContainer(this.database.Id, this.Container.Id); - await container.ReadItemAsync("1", new Cosmos.PartitionKey("Status1")); - await container.ReadItemAsync("3", new Cosmos.PartitionKey("Status3")); - isFirstRequestToPartition = false; - await container.ReadItemAsync("3", new Cosmos.PartitionKey("Status3")); - } - } } } 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); } }