Skip to content

Commit 63fac64

Browse files
simplynaveen20simplynaveen20
authored
Scoping session token per partition level for gateway call and enabling session token for multi master create request (Azure#25555)
* Scoping session token per partition level for gateway call and enabling session token for multi master create request * checkstyle fix * fixing existing test in RxGatewayStoreModelTest.java * fixing test case * adding feed range query test * Removing unwanted change from SessionTokenHelper.java * Handling scenerio of partition split with scoped session token * fixing test cases * Fixing test cases * fixing test cases * adding multimaster group in session test
1 parent 8d215ce commit 63fac64

File tree

11 files changed

+452
-74
lines changed

11 files changed

+452
-74
lines changed

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReplicatedResourceClientUtils.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,8 @@ public static boolean isMasterResource(ResourceType resourceType) {
3636
resourceType == ResourceType.PartitionKeyRange ||
3737
resourceType == ResourceType.DocumentCollection ||
3838
resourceType == ResourceType.Trigger ||
39-
resourceType == ResourceType.UserDefinedFunction) {
39+
resourceType == ResourceType.UserDefinedFunction ||
40+
resourceType == ResourceType.ClientEncryptionKey) {
4041
return true;
4142
}
4243

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -407,6 +407,13 @@ private void initializeGatewayConfigurationReader() {
407407
// https://msdata.visualstudio.com/CosmosDB/_workitems/edit/332589
408408
}
409409

410+
private void updateGatewayProxy() {
411+
((RxGatewayStoreModel)this.gatewayProxy).setGatewayServiceConfigurationReader(this.gatewayConfigurationReader);
412+
((RxGatewayStoreModel)this.gatewayProxy).setCollectionCache(this.collectionCache);
413+
((RxGatewayStoreModel)this.gatewayProxy).setPartitionKeyRangeCache(this.partitionKeyRangeCache);
414+
((RxGatewayStoreModel)this.gatewayProxy).setUseMultipleWriteLocations(this.useMultipleWriteLocations);
415+
}
416+
410417
public void init(CosmosClientMetadataCachesSnapshot metadataCachesSnapshot, Function<HttpClient, HttpClient> httpClientInterceptor) {
411418
try {
412419
// TODO: add support for openAsync
@@ -448,6 +455,7 @@ public void init(CosmosClientMetadataCachesSnapshot metadataCachesSnapshot, Func
448455
this.partitionKeyRangeCache = new RxPartitionKeyRangeCache(RxDocumentClientImpl.this,
449456
collectionCache);
450457

458+
updateGatewayProxy();
451459
if (this.connectionPolicy.getConnectionMode() == ConnectionMode.GATEWAY) {
452460
this.storeModel = this.gatewayProxy;
453461
} else {
@@ -460,7 +468,6 @@ public void init(CosmosClientMetadataCachesSnapshot metadataCachesSnapshot, Func
460468
clientTelemetry.init();
461469
this.queryPlanCache = new ConcurrentHashMap<>();
462470
this.retryPolicy.setRxCollectionCache(this.collectionCache);
463-
((RxGatewayStoreModel)this.gatewayProxy).setCollectionCache(this.collectionCache);
464471
} catch (Exception e) {
465472
logger.error("unexpected failure in initializing client.", e);
466473
close();
@@ -565,7 +572,7 @@ private void createStoreModel(boolean subscribeRntbdStatus) {
565572
this.sessionContainer,
566573
this.gatewayConfigurationReader,
567574
this,
568-
false
575+
this.useMultipleWriteLocations
569576
);
570577

571578
this.storeModel = new ServerStoreModel(storeClient);

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java‎

Lines changed: 146 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -7,15 +7,20 @@
77
import com.azure.cosmos.CosmosException;
88
import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
99
import com.azure.cosmos.implementation.caches.RxClientCollectionCache;
10+
import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache;
1011
import com.azure.cosmos.implementation.directconnectivity.DirectBridgeInternal;
12+
import com.azure.cosmos.implementation.directconnectivity.GatewayServiceConfigurationReader;
1113
import com.azure.cosmos.implementation.directconnectivity.HttpUtils;
14+
import com.azure.cosmos.implementation.directconnectivity.RequestHelper;
1215
import com.azure.cosmos.implementation.directconnectivity.StoreResponse;
1316
import com.azure.cosmos.implementation.directconnectivity.WebExceptionUtility;
1417
import com.azure.cosmos.implementation.http.HttpClient;
1518
import com.azure.cosmos.implementation.http.HttpHeaders;
1619
import com.azure.cosmos.implementation.http.HttpRequest;
1720
import com.azure.cosmos.implementation.http.HttpResponse;
1821
import com.azure.cosmos.implementation.http.ReactorNettyRequestRecord;
22+
import com.azure.cosmos.implementation.routing.PartitionKeyInternal;
23+
import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper;
1924
import com.azure.cosmos.implementation.throughputControl.ThroughputControlStore;
2025
import io.netty.handler.codec.http.HttpMethod;
2126
import io.netty.handler.codec.http.HttpResponseStatus;
@@ -53,6 +58,9 @@ class RxGatewayStoreModel implements RxStoreModel {
5358
private ConsistencyLevel defaultConsistencyLevel;
5459
private ISessionContainer sessionContainer;
5560
private ThroughputControlStore throughputControlStore;
61+
private boolean useMultipleWriteLocations;
62+
private RxPartitionKeyRangeCache partitionKeyRangeCache;
63+
private GatewayServiceConfigurationReader gatewayServiceConfigurationReader;
5664
private RxClientCollectionCache collectionCache;
5765

5866
public RxGatewayStoreModel(
@@ -94,7 +102,35 @@ public RxGatewayStoreModel(
94102
this.sessionContainer = sessionContainer;
95103
}
96104

97-
void setCollectionCache(RxClientCollectionCache collectionCache) {
105+
void setGatewayServiceConfigurationReader(GatewayServiceConfigurationReader gatewayServiceConfigurationReader) {
106+
this.gatewayServiceConfigurationReader = gatewayServiceConfigurationReader;
107+
}
108+
109+
public void setPartitionKeyRangeCache(RxPartitionKeyRangeCache partitionKeyRangeCache) {
110+
this.partitionKeyRangeCache = partitionKeyRangeCache;
111+
}
112+
113+
public void setUseMultipleWriteLocations(boolean useMultipleWriteLocations) {
114+
this.useMultipleWriteLocations = useMultipleWriteLocations;
115+
}
116+
117+
boolean isUseMultipleWriteLocations() {
118+
return useMultipleWriteLocations;
119+
}
120+
121+
RxPartitionKeyRangeCache getPartitionKeyRangeCache() {
122+
return partitionKeyRangeCache;
123+
}
124+
125+
GatewayServiceConfigurationReader getGatewayServiceConfigurationReader() {
126+
return gatewayServiceConfigurationReader;
127+
}
128+
129+
RxClientCollectionCache getCollectionCache() {
130+
return collectionCache;
131+
}
132+
133+
public void setCollectionCache(RxClientCollectionCache collectionCache) {
98134
this.collectionCache = collectionCache;
99135
}
100136

@@ -473,11 +509,8 @@ public Mono<RxDocumentServiceResponse> processMessage(RxDocumentServiceRequest r
473509

474510
return Mono.error(dce);
475511
}
476-
).map(response ->
477-
{
478-
this.captureSessionToken(request, response.getResponseHeaders());
479-
return response;
480-
}
512+
).flatMap(response ->
513+
this.captureSessionTokenAndHandlePartitionSplit(request, response.getResponseHeaders()).then(Mono.just(response))
481514
);
482515
}
483516

@@ -503,12 +536,29 @@ private void captureSessionToken(RxDocumentServiceRequest request, Map<String, S
503536
}
504537
}
505538

539+
private Mono<Void> captureSessionTokenAndHandlePartitionSplit(RxDocumentServiceRequest request,
540+
Map<String, String> responseHeaders) {
541+
this.captureSessionToken(request, responseHeaders);
542+
if (request.requestContext.resolvedPartitionKeyRange != null &&
543+
StringUtils.isNotEmpty(request.requestContext.resolvedCollectionRid) &&
544+
StringUtils.isNotEmpty(responseHeaders.get(HttpConstants.HttpHeaders.PARTITION_KEY_RANGE_ID)) &&
545+
!responseHeaders.get(HttpConstants.HttpHeaders.PARTITION_KEY_RANGE_ID).equals(request.requestContext.resolvedPartitionKeyRange.getId())) {
546+
return this.partitionKeyRangeCache.refreshAsync(BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), request.requestContext.resolvedCollectionRid)
547+
.flatMap(collectionRoutingMapValueHolder -> Mono.empty());
548+
}
549+
return Mono.empty();
550+
}
551+
506552
private Mono<Void> addIntendedCollectionRidAndSessionToken(RxDocumentServiceRequest request) {
507-
applySessionToken(request);
508-
if(this.collectionCache != null && request.getResourceType().equals(ResourceType.Document)) {
553+
return applySessionToken(request).then(addIntendedCollectionRid(request));
554+
}
555+
556+
private Mono<Void> addIntendedCollectionRid(RxDocumentServiceRequest request) {
557+
if (this.collectionCache != null && request.getResourceType().equals(ResourceType.Document)) {
509558
return this.collectionCache.resolveCollectionAsync(BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), request).flatMap(documentCollectionValueHolder -> {
510-
if(StringUtils.isEmpty(request.getHeaders().get(INTENDED_COLLECTION_RID_HEADER))) {
511-
request.getHeaders().put(INTENDED_COLLECTION_RID_HEADER, request.requestContext.resolvedCollectionRid);
559+
if (StringUtils.isEmpty(request.getHeaders().get(INTENDED_COLLECTION_RID_HEADER))) {
560+
request.getHeaders().put(INTENDED_COLLECTION_RID_HEADER,
561+
request.requestContext.resolvedCollectionRid);
512562
} else {
513563
request.intendedCollectionRidPassedIntoSDK = true;
514564
}
@@ -518,40 +568,105 @@ private Mono<Void> addIntendedCollectionRidAndSessionToken(RxDocumentServiceRequ
518568
return Mono.empty();
519569
}
520570

521-
private void applySessionToken(RxDocumentServiceRequest request) {
571+
private Mono<Void> applySessionToken(RxDocumentServiceRequest request) {
522572
Map<String, String> headers = request.getHeaders();
523573
Objects.requireNonNull(headers, "RxDocumentServiceRequest::headers is required and cannot be null");
524574

525-
String requestConsistencyLevel = headers.get(HttpConstants.HttpHeaders.CONSISTENCY_LEVEL);
575+
// Master resource operations don't require session token.
576+
if (isMasterOperation(request.getResourceType(), request.getOperationType())) {
577+
if (!Strings.isNullOrEmpty(request.getHeaders().get(HttpConstants.HttpHeaders.SESSION_TOKEN))) {
578+
request.getHeaders().remove(HttpConstants.HttpHeaders.SESSION_TOKEN);
579+
}
580+
return Mono.empty();
581+
}
526582

527-
boolean sessionTokenApplicable =
528-
Strings.areEqual(requestConsistencyLevel, ConsistencyLevel.SESSION.toString()) ||
529-
(this.defaultConsistencyLevel == ConsistencyLevel.SESSION &&
530-
// skip applying the session token when Eventual Consistency is explicitly requested
531-
// on request-level for data plane operations.
532-
// The session token is ignored on teh backend/gateway in this case anyway
533-
// and the session token can be rather large (even run in the 16 KB header length problem
534-
// on the gateway - so not worth sending when not needed
535-
(!request.isReadOnlyRequest() ||
536-
request.getResourceType() != ResourceType.Document ||
537-
!Strings.areEqual(requestConsistencyLevel, ConsistencyLevel.EVENTUAL.toString())));
583+
boolean sessionConsistency = RequestHelper.getConsistencyLevelToUse(this.gatewayServiceConfigurationReader,
584+
request) == ConsistencyLevel.SESSION;
538585

539586
if (!Strings.isNullOrEmpty(request.getHeaders().get(HttpConstants.HttpHeaders.SESSION_TOKEN))) {
540-
if (!sessionTokenApplicable || isMasterOperation(request.getResourceType(), request.getOperationType())) {
587+
if (!sessionConsistency ||
588+
(!request.isReadOnlyRequest() && request.getOperationType() != OperationType.Batch && !this.useMultipleWriteLocations)){
541589
request.getHeaders().remove(HttpConstants.HttpHeaders.SESSION_TOKEN);
542590
}
543-
return; //User is explicitly controlling the session.
591+
return Mono.empty(); //User is explicitly controlling the session.
544592
}
545593

546-
if (!sessionTokenApplicable || isMasterOperation(request.getResourceType(), request.getOperationType())) {
547-
return; // Only apply the session token in case of session consistency and when resource is not a master resource
594+
if (!sessionConsistency ||
595+
(!request.isReadOnlyRequest() && request.getOperationType() != OperationType.Batch && !this.useMultipleWriteLocations)) {
596+
return Mono.empty();
597+
// Only apply the session token in case of session consistency and if request is read only,
598+
// apply token for write request only if batch operation or multi master
548599
}
549600

550-
//Apply the ambient session.
551-
String sessionToken = this.sessionContainer.resolveGlobalSessionToken(request);
601+
if (this.collectionCache != null && this.partitionKeyRangeCache != null) {
602+
return this.collectionCache.resolveCollectionAsync(BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), request).
603+
flatMap(collectionValueHolder -> {
552604

553-
if (!Strings.isNullOrEmpty(sessionToken)) {
554-
headers.put(HttpConstants.HttpHeaders.SESSION_TOKEN, sessionToken);
605+
if(collectionValueHolder== null || collectionValueHolder.v == null) {
606+
//Apply the ambient session.
607+
String sessionToken = this.sessionContainer.resolveGlobalSessionToken(request);
608+
609+
if (!Strings.isNullOrEmpty(sessionToken)) {
610+
headers.put(HttpConstants.HttpHeaders.SESSION_TOKEN, sessionToken);
611+
}
612+
return Mono.empty();
613+
}
614+
return partitionKeyRangeCache.tryLookupAsync(BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics),
615+
collectionValueHolder.v.getResourceId(),
616+
null,
617+
null).flatMap(collectionRoutingMapValueHolder -> {
618+
if (collectionRoutingMapValueHolder == null || collectionRoutingMapValueHolder.v == null) {
619+
//Apply the ambient session.
620+
String sessionToken = this.sessionContainer.resolveGlobalSessionToken(request);
621+
622+
if (!Strings.isNullOrEmpty(sessionToken)) {
623+
headers.put(HttpConstants.HttpHeaders.SESSION_TOKEN, sessionToken);
624+
}
625+
return Mono.empty();
626+
}
627+
String partitionKeyRangeId =
628+
request.getHeaders().get(HttpConstants.HttpHeaders.PARTITION_KEY_RANGE_ID);
629+
PartitionKeyInternal partitionKeyInternal = request.getPartitionKeyInternal();
630+
631+
if (StringUtils.isNotEmpty(partitionKeyRangeId)) {
632+
PartitionKeyRange range =
633+
collectionRoutingMapValueHolder.v.getRangeByPartitionKeyRangeId(partitionKeyRangeId);
634+
request.requestContext.resolvedPartitionKeyRange = range;
635+
if (request.requestContext.resolvedPartitionKeyRange == null) {
636+
SessionTokenHelper.setPartitionLocalSessionToken(request, partitionKeyRangeId,
637+
sessionContainer);
638+
} else {
639+
SessionTokenHelper.setPartitionLocalSessionToken(request, sessionContainer);
640+
}
641+
} else if (partitionKeyInternal != null) {
642+
String effectivePartitionKeyString = PartitionKeyInternalHelper
643+
.getEffectivePartitionKeyString(
644+
partitionKeyInternal,
645+
collectionValueHolder.v.getPartitionKey());
646+
PartitionKeyRange range =
647+
collectionRoutingMapValueHolder.v.getRangeByEffectivePartitionKey(effectivePartitionKeyString);
648+
request.requestContext.resolvedPartitionKeyRange = range;
649+
SessionTokenHelper.setPartitionLocalSessionToken(request, sessionContainer);
650+
} else {
651+
//Apply the ambient session.
652+
String sessionToken = this.sessionContainer.resolveGlobalSessionToken(request);
653+
654+
if (!Strings.isNullOrEmpty(sessionToken)) {
655+
headers.put(HttpConstants.HttpHeaders.SESSION_TOKEN, sessionToken);
656+
}
657+
}
658+
659+
return Mono.empty();
660+
});
661+
});
662+
} else {
663+
//Apply the ambient session.
664+
String sessionToken = this.sessionContainer.resolveGlobalSessionToken(request);
665+
666+
if (!Strings.isNullOrEmpty(sessionToken)) {
667+
headers.put(HttpConstants.HttpHeaders.SESSION_TOKEN, sessionToken);
668+
}
669+
return Mono.empty();
555670
}
556671
}
557672

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/SessionTokenHelper.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,9 +32,11 @@ public static void setOriginalSessionToken(RxDocumentServiceRequest request, Str
3232
}
3333

3434
public static void setPartitionLocalSessionToken(RxDocumentServiceRequest request, ISessionContainer sessionContainer) {
35-
String originalSessionToken = request.getHeaders().get(HttpConstants.HttpHeaders.SESSION_TOKEN);
36-
String partitionKeyRangeId = request.requestContext.resolvedPartitionKeyRange.getId();
35+
setPartitionLocalSessionToken(request, request.requestContext.resolvedPartitionKeyRange.getId(), sessionContainer);
36+
}
3737

38+
public static void setPartitionLocalSessionToken(RxDocumentServiceRequest request, String partitionKeyRangeId, ISessionContainer sessionContainer) {
39+
String originalSessionToken = request.getHeaders().get(HttpConstants.HttpHeaders.SESSION_TOKEN);
3840

3941
if (Strings.isNullOrEmpty(partitionKeyRangeId)) {
4042
// AddressCache/address resolution didn't produce partition key range id.

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/caches/RxPartitionKeyRangeCache.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,16 @@ public Mono<Utils.ValueHolder<PartitionKeyRange>> tryGetRangeByPartitionKeyRange
168168
});
169169
}
170170

171+
public Mono<Utils.ValueHolder<CollectionRoutingMap>> refreshAsync(MetadataDiagnosticsContext metaDataDiagnosticsContext, String collectionRid) {
172+
return this.tryLookupAsync(
173+
metaDataDiagnosticsContext,
174+
collectionRid,
175+
null,
176+
null
177+
).flatMap(collectionRoutingMapValueHolder -> tryLookupAsync(metaDataDiagnosticsContext, collectionRid,
178+
collectionRoutingMapValueHolder.v, null));
179+
}
180+
171181
private Mono<CollectionRoutingMap> getRoutingMapForCollectionAsync(
172182
MetadataDiagnosticsContext metaDataDiagnosticsContext,
173183
String collectionRid,

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,7 @@ private static <T extends Resource> Mono<Pair<List<Range<String>>,QueryInfo>> ge
115115
Instant endTime = Instant.now(); // endTime for query plan diagnostics
116116
PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = queryPlanCache.get(query.getQueryText());
117117
if (partitionedQueryExecutionInfo != null) {
118-
logger.info("Skipping query plan round trip by using the cached plan");
118+
logger.debug("Skipping query plan round trip by using the cached plan");
119119
return getTargetRangesFromQueryPlan(cosmosQueryRequestOptions, collection, queryExecutionContext,
120120
partitionedQueryExecutionInfo, startTime, endTime);
121121
}

0 commit comments

Comments
 (0)