Skip to content

Commit

Permalink
Session: Fixes Gateway session scope during merges (#2980)
Browse files Browse the repository at this point in the history 
Fix handling of parent tokens on a partition merge scenario
  • Loading branch information
@ealsur
ealsur committed Jan 21, 2022
1 parent 3e9bab6 commit fb7e398
Show file tree
Hide file tree
Showing 3 changed files with 243 additions and 4 deletions.
16 changes: 12 additions & 4 deletions Microsoft.Azure.Cosmos/src/SessionContainer.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@ namespace Microsoft.Azure.Cosmos.Common
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Globalization;
using System.Text;
using System.Threading;
using Microsoft.Azure.Cosmos.Core.Trace;
Expand All @@ -15,6 +14,7 @@ namespace Microsoft.Azure.Cosmos.Common

internal sealed class SessionContainer : ISessionContainer
{
private static readonly string sessionTokenSeparator = ":";
private volatile SessionContainerState state;

public SessionContainer(string hostName)
Expand Down Expand Up @@ -153,18 +153,26 @@ private static string ResolvePartitionLocalSessionTokenForGateway(SessionContain
{
if (partitionKeyRangeIdToTokenMap.TryGetValue(partitionKeyRangeId, out ISessionToken sessionToken))
{
return partitionKeyRangeId + ":" + sessionToken.ConvertToString();
return partitionKeyRangeId + SessionContainer.sessionTokenSeparator + sessionToken.ConvertToString();
}
else if (request.RequestContext.ResolvedPartitionKeyRange.Parents != null)
{
ISessionToken parentSessionToken = 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();
// A partition can have more than 1 parent (merge). In that case, we apply Merge to generate a token with both parent's max LSNs
parentSessionToken = parentSessionToken != null ? parentSessionToken.Merge(sessionToken) : sessionToken;
}
}

// When we don't have the session token for a partition, we can leverage the session token of the parent(s)
if (parentSessionToken != null)
{
return partitionKeyRangeId + SessionContainer.sessionTokenSeparator + parentSessionToken.ConvertToString();
}
}
}

Expand Down Expand Up @@ -391,7 +399,7 @@ private static string GetSessionTokenString(ConcurrentDictionary<string, ISessio
}

sb.Append(pair.Key);
sb.Append(":");
sb.Append(SessionContainer.sessionTokenSeparator);
sb.Append(pair.Value.ConvertToString());
}

Expand Down
113 changes: 113 additions & 0 deletions Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1029,6 +1029,119 @@ Task<HttpResponseMessage> sendFunc(HttpRequestMessage request)
}
}

/// <summary>
/// Simulating partition split and cache having only the parent token
/// </summary>
[TestMethod]
public async Task GatewayStoreModel_ObtainsSessionFromParent_AfterSplit()
{
SessionContainer sessionContainer = new SessionContainer("testhost");

string collectionResourceId = ResourceId.NewDocumentCollectionId(42, 129).DocumentCollectionId.ToString();
string collectionFullname = "dbs/db1/colls/collName";

// Set token for the parent
string parentPKRangeId = "0";
string parentSession = "1#100#4=90#5=1";
sessionContainer.SetSessionToken(
collectionResourceId,
collectionFullname,
new StoreRequestNameValueCollection() { { HttpConstants.HttpHeaders.SessionToken, $"{parentPKRangeId}:{parentSession}" } }
);

// Create the request for the child
string childPKRangeId = "1";
DocumentServiceRequest documentServiceRequestToChild = DocumentServiceRequest.CreateFromName(OperationType.Read, "dbs/db1/colls/collName/docs/42", ResourceType.Document, AuthorizationTokenType.PrimaryMasterKey, null);

documentServiceRequestToChild.RequestContext.ResolvedPartitionKeyRange = new PartitionKeyRange()
{
Id = childPKRangeId,
MinInclusive = "",
MaxExclusive = "AA",
Parents = new Collection<string>() { parentPKRangeId } // PartitionKeyRange says who is the parent
};

Mock<IGlobalEndpointManager> globalEndpointManager = new Mock<IGlobalEndpointManager>();
await this.GetGatewayStoreModelForConsistencyTest(async (gatewayStoreModel) =>
{
await GatewayStoreModel.ApplySessionTokenAsync(
documentServiceRequestToChild,
ConsistencyLevel.Session,
sessionContainer,
partitionKeyRangeCache: new Mock<PartitionKeyRangeCache>(null, null, null).Object,
clientCollectionCache: new Mock<ClientCollectionCache>(sessionContainer, gatewayStoreModel, null, null).Object,
globalEndpointManager: globalEndpointManager.Object);
Assert.AreEqual($"{childPKRangeId}:{parentSession}", documentServiceRequestToChild.Headers[HttpConstants.HttpHeaders.SessionToken]);
});
}

/// <summary>
/// Simulating partition merge and cache having only the parents tokens
/// </summary>
[TestMethod]
public async Task GatewayStoreModel_ObtainsSessionFromParents_AfterMerge()
{
SessionContainer sessionContainer = new SessionContainer("testhost");

string collectionResourceId = ResourceId.NewDocumentCollectionId(42, 129).DocumentCollectionId.ToString();
string collectionFullname = "dbs/db1/colls/collName";

// Set tokens for the parents
string parentPKRangeId = "0";
int maxGlobalLsn = 100;
int maxLsnRegion1 = 200;
int maxLsnRegion2 = 300;
int maxLsnRegion3 = 400;

// Generate 2 tokens, one has max global but lower regional, the other lower global but higher regional
// Expect the merge to contain all the maxes
string parentSession = $"1#{maxGlobalLsn}#1={maxLsnRegion1 - 1}#2={maxLsnRegion2}#3={maxLsnRegion3 - 1}";
sessionContainer.SetSessionToken(
collectionResourceId,
collectionFullname,
new StoreRequestNameValueCollection() { { HttpConstants.HttpHeaders.SessionToken, $"{parentPKRangeId}:{parentSession}" } }
);

string parent2PKRangeId = "1";
string parent2Session = $"1#{maxGlobalLsn - 1}#1={maxLsnRegion1}#2={maxLsnRegion2 - 1}#3={maxLsnRegion3}";
sessionContainer.SetSessionToken(
collectionResourceId,
collectionFullname,
new StoreRequestNameValueCollection() { { HttpConstants.HttpHeaders.SessionToken, $"{parent2PKRangeId}:{parent2Session}" } }
);

string tokenWithAllMax = $"1#{maxGlobalLsn}#1={maxLsnRegion1}#2={maxLsnRegion2}#3={maxLsnRegion3}";

// Create the request for the child
// Request for a child from both parents
string childPKRangeId = "2";

DocumentServiceRequest documentServiceRequestToChild = DocumentServiceRequest.CreateFromName(OperationType.Read, "dbs/db1/colls/collName/docs/42", ResourceType.Document, AuthorizationTokenType.PrimaryMasterKey, null);

documentServiceRequestToChild.RequestContext.ResolvedPartitionKeyRange = new PartitionKeyRange()
{
Id = childPKRangeId,
MinInclusive = "",
MaxExclusive = "FF",
Parents = new Collection<string>() { parentPKRangeId, parent2PKRangeId } // PartitionKeyRange says who are the parents
};

Mock<IGlobalEndpointManager> globalEndpointManager = new Mock<IGlobalEndpointManager>();
await this.GetGatewayStoreModelForConsistencyTest(async (gatewayStoreModel) =>
{
await GatewayStoreModel.ApplySessionTokenAsync(
documentServiceRequestToChild,
ConsistencyLevel.Session,
sessionContainer,
partitionKeyRangeCache: new Mock<PartitionKeyRangeCache>(null, null, null).Object,
clientCollectionCache: new Mock<ClientCollectionCache>(sessionContainer, gatewayStoreModel, null, null).Object,
globalEndpointManager: globalEndpointManager.Object);
Assert.AreEqual($"{childPKRangeId}:{tokenWithAllMax}", documentServiceRequestToChild.Headers[HttpConstants.HttpHeaders.SessionToken]);
});
}

private class MockMessageHandler : HttpMessageHandler
{
private readonly Func<HttpRequestMessage, Task<HttpResponseMessage>> sendFunc;
Expand Down
118 changes: 118 additions & 0 deletions Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/SessionContainerTest.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ namespace Microsoft.Azure.Cosmos
using Client;
using Microsoft.Azure.Documents;
using Microsoft.Azure.Documents.Collections;
using System.Threading.Tasks;

/// <summary>
/// Tests for <see cref="SessionContainer"/> class.
Expand Down Expand Up @@ -806,5 +807,122 @@ public void TestNewCollectionResourceIdInvalidatesOldCollectionResourceId()
Assert.IsTrue(string.IsNullOrEmpty(sessionContainer.GetSessionToken(string.Format("dbs/{0}/colls/{1}", dbResourceId, oldCollectionResourceId))));
Assert.IsFalse(string.IsNullOrEmpty(sessionContainer.GetSessionToken(string.Format("dbs/{0}/colls/{1}", dbResourceId, newCollectionResourceId))));
}

/// <summary>
/// Use the session token of the parent if request comes for a child
/// </summary>
[TestMethod]
public void TestResolveSessionTokenFromParent_Gateway_AfterSplit()
{
SessionContainer sessionContainer = new SessionContainer("127.0.0.1");

string collectionResourceId = ResourceId.NewDocumentCollectionId(42, 129).DocumentCollectionId.ToString();
string collectionFullname = "dbs/db1/colls/collName";

// Set token for the parent
string parentPKRangeId = "0";
string parentSession = "1#100#4=90#5=1";
sessionContainer.SetSessionToken(
collectionResourceId,
collectionFullname,
new StoreRequestNameValueCollection() { { HttpConstants.HttpHeaders.SessionToken, $"{parentPKRangeId}:{parentSession}" } }
);

// We send requests for the children

string childPKRangeId = "1";
string childPKRangeId2 = "1";

DocumentServiceRequest documentServiceRequestToChild1 = DocumentServiceRequest.CreateFromName(OperationType.Read, "dbs/db1/colls/collName/docs/42", ResourceType.Document, AuthorizationTokenType.PrimaryMasterKey, null);

documentServiceRequestToChild1.RequestContext.ResolvedPartitionKeyRange = new PartitionKeyRange()
{
Id = childPKRangeId,
MinInclusive = "",
MaxExclusive = "AA",
Parents = new Collection<string>() { parentPKRangeId } // PartitionKeyRange says who is the parent
};

string resolvedToken = sessionContainer.ResolvePartitionLocalSessionTokenForGateway(
documentServiceRequestToChild1,
childPKRangeId);// For one of the children

Assert.AreEqual($"{childPKRangeId}:{parentSession}", resolvedToken);

DocumentServiceRequest documentServiceRequestToChild2 = DocumentServiceRequest.CreateFromName(OperationType.Read, "dbs/db1/colls/collName/docs/42", ResourceType.Document, AuthorizationTokenType.PrimaryMasterKey, null);

documentServiceRequestToChild2.RequestContext.ResolvedPartitionKeyRange = new PartitionKeyRange()
{
Id = childPKRangeId2,
MinInclusive = "AA",
MaxExclusive = "FF",
Parents = new Collection<string>() { parentPKRangeId } // PartitionKeyRange says who is the parent
};

resolvedToken = sessionContainer.ResolvePartitionLocalSessionTokenForGateway(
documentServiceRequestToChild2,
childPKRangeId2);// For the other child

Assert.AreEqual($"{childPKRangeId2}:{parentSession}", resolvedToken);
}

// <summary>
/// Use the session token of the parent if request comes for a child when 2 parents are present
/// </summary>
[TestMethod]
public void TestResolveSessionTokenFromParent_Gateway_AfterMerge()
{
SessionContainer sessionContainer = new SessionContainer("127.0.0.1");

string collectionResourceId = ResourceId.NewDocumentCollectionId(42, 129).DocumentCollectionId.ToString();
string collectionFullname = "dbs/db1/colls/collName";

// Set tokens for the parents
string parentPKRangeId = "0";
int maxGlobalLsn = 100;
int maxLsnRegion1 = 200;
int maxLsnRegion2 = 300;
int maxLsnRegion3 = 400;

// Generate 2 tokens, one has max global but lower regional, the other lower global but higher regional
// Expect the merge to contain all the maxes
string parentSession = $"1#{maxGlobalLsn}#1={maxLsnRegion1 - 1}#2={maxLsnRegion2}#3={maxLsnRegion3 - 1}";
sessionContainer.SetSessionToken(
collectionResourceId,
collectionFullname,
new StoreRequestNameValueCollection() { { HttpConstants.HttpHeaders.SessionToken, $"{parentPKRangeId}:{parentSession}" } }
);

string parent2PKRangeId = "1";
string parent2Session = $"1#{maxGlobalLsn - 1}#1={maxLsnRegion1}#2={maxLsnRegion2 - 1}#3={maxLsnRegion3}";
sessionContainer.SetSessionToken(
collectionResourceId,
collectionFullname,
new StoreRequestNameValueCollection() { { HttpConstants.HttpHeaders.SessionToken, $"{parent2PKRangeId}:{parent2Session}" } }
);

string tokenWithAllMax = $"1#{maxGlobalLsn}#1={maxLsnRegion1}#2={maxLsnRegion2}#3={maxLsnRegion3}";

// Request for a child from both parents
string childPKRangeId = "2";

DocumentServiceRequest documentServiceRequestToChild1 = DocumentServiceRequest.CreateFromName(OperationType.Read, "dbs/db1/colls/collName/docs/42", ResourceType.Document, AuthorizationTokenType.PrimaryMasterKey, null);

documentServiceRequestToChild1.RequestContext.ResolvedPartitionKeyRange = new PartitionKeyRange()
{
Id = childPKRangeId,
MinInclusive = "",
MaxExclusive = "FF",
Parents = new Collection<string>() { parentPKRangeId, parent2PKRangeId } // PartitionKeyRange says who are the parents
};

string resolvedToken = sessionContainer.ResolvePartitionLocalSessionTokenForGateway(
documentServiceRequestToChild1,
childPKRangeId);// For one of the children

// Expect the resulting token is for the child partition but containing all maxes of the lsn of the parents
Assert.AreEqual($"{childPKRangeId}:{tokenWithAllMax}", resolvedToken);
}

}
}

0 comments on commit fb7e398

Please sign in to comment.