diff --git a/sdk/cosmos/azure-cosmos-test/src/main/java/com/azure/cosmos/test/implementation/faultinjection/FaultInjectionRuleProcessor.java b/sdk/cosmos/azure-cosmos-test/src/main/java/com/azure/cosmos/test/implementation/faultinjection/FaultInjectionRuleProcessor.java index 0b221ce659c2..85a5843a304f 100644 --- a/sdk/cosmos/azure-cosmos-test/src/main/java/com/azure/cosmos/test/implementation/faultinjection/FaultInjectionRuleProcessor.java +++ b/sdk/cosmos/azure-cosmos-test/src/main/java/com/azure/cosmos/test/implementation/faultinjection/FaultInjectionRuleProcessor.java @@ -24,6 +24,7 @@ import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdUtils; import com.azure.cosmos.implementation.feedranges.FeedRangeInternal; import com.azure.cosmos.implementation.routing.PartitionKeyRangeIdentity; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import com.azure.cosmos.test.faultinjection.FaultInjectionCondition; import com.azure.cosmos.test.faultinjection.FaultInjectionConnectionErrorResult; import com.azure.cosmos.test.faultinjection.FaultInjectionConnectionType; @@ -446,7 +447,7 @@ private Mono> resolvePhysicalAddresses( ResourceType.Document, null); - faultInjectionAddressRequest.requestContext.locationEndpointToRoute = regionEndpoint; + faultInjectionAddressRequest.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(regionEndpoint); faultInjectionAddressRequest.setPartitionKeyRangeIdentity(new PartitionKeyRangeIdentity(pkRangeId)); if (isWriteOnly) { diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/GlobalPartitionEndpointManagerForCircuitBreakerTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/GlobalPartitionEndpointManagerForCircuitBreakerTests.java index 9e7246b06225..d3b4d9b8134d 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/GlobalPartitionEndpointManagerForCircuitBreakerTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/GlobalPartitionEndpointManagerForCircuitBreakerTests.java @@ -4,7 +4,6 @@ package com.azure.cosmos; import com.azure.cosmos.implementation.GlobalEndpointManager; -import com.azure.cosmos.implementation.MetadataDiagnosticsContext; import com.azure.cosmos.implementation.OperationType; import com.azure.cosmos.implementation.PartitionKeyRange; import com.azure.cosmos.implementation.PointOperationContextForCircuitBreaker; @@ -17,6 +16,7 @@ import com.azure.cosmos.implementation.circuitBreaker.LocationSpecificHealthContext; import com.azure.cosmos.implementation.circuitBreaker.PartitionKeyRangeWrapper; import com.azure.cosmos.implementation.guava25.collect.ImmutableList; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import org.apache.commons.lang3.tuple.Pair; import org.mockito.Mockito; import org.slf4j.Logger; @@ -142,11 +142,11 @@ public void recordHealthyStatus(String partitionLevelCircuitBreakerConfigAsJsonS Object partitionAndLocationSpecificUnavailabilityInfo = partitionKeyRangeToLocationSpecificUnavailabilityInfo.get(new PartitionKeyRangeWrapper(request.requestContext.resolvedPartitionKeyRange, collectionResourceId)); - ConcurrentHashMap locationEndpointToLocationSpecificContextForPartition - = (ConcurrentHashMap) locationEndpointToLocationSpecificContextForPartitionField.get(partitionAndLocationSpecificUnavailabilityInfo); + ConcurrentHashMap locationEndpointToLocationSpecificContextForPartition + = (ConcurrentHashMap) locationEndpointToLocationSpecificContextForPartitionField.get(partitionAndLocationSpecificUnavailabilityInfo); LocationSpecificHealthContext locationSpecificHealthContext - = locationEndpointToLocationSpecificContextForPartition.get(LocationEastUs2EndpointToLocationPair.getKey()); + = locationEndpointToLocationSpecificContextForPartition.get(new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getKey())); assertThat(locationSpecificHealthContext.isRegionAvailableToProcessRequests()).isTrue(); assertThat(locationSpecificHealthContext.isExceptionThresholdBreached()).isFalse(); @@ -167,12 +167,12 @@ public void recordHealthyToHealthyWithFailuresStatusTransition(String partitionL String maxExclusive = "BB"; String collectionResourceId = "dbs/db1/colls/coll1"; - List applicableReadWriteEndpoints = ImmutableList.of( + List applicableReadWriteEndpoints = ImmutableList.of( LocationEastUs2EndpointToLocationPair, LocationEastUsEndpointToLocationPair, LocationCentralUsEndpointToLocationPair) .stream() - .map(Pair::getLeft) + .map(uriStringPair -> new RegionalRoutingContext(uriStringPair.getLeft())) .collect(Collectors.toList()); RxDocumentServiceRequest request = constructRxDocumentServiceRequestInstance( @@ -185,11 +185,11 @@ public void recordHealthyToHealthyWithFailuresStatusTransition(String partitionL maxExclusive, LocationEastUs2EndpointToLocationPair.getKey()); - Mockito.when(this.globalEndpointManagerMock.getReadEndpoints()).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); - Mockito.when(this.globalEndpointManagerMock.getWriteEndpoints()).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); + Mockito.when(this.globalEndpointManagerMock.getReadEndpoints()).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); + Mockito.when(this.globalEndpointManagerMock.getWriteEndpoints()).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); globalPartitionEndpointManagerForCircuitBreaker - .handleLocationExceptionForPartitionKeyRange(request, LocationEastUs2EndpointToLocationPair.getKey()); + .handleLocationExceptionForPartitionKeyRange(request, new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getKey())); Class[] enclosedClasses = GlobalPartitionEndpointManagerForCircuitBreaker.class.getDeclaredClasses(); Class partitionLevelUnavailabilityInfoClass @@ -210,11 +210,11 @@ public void recordHealthyToHealthyWithFailuresStatusTransition(String partitionL Object partitionAndLocationSpecificUnavailabilityInfo = partitionKeyRangeToLocationSpecificUnavailabilityInfo.get(new PartitionKeyRangeWrapper(request.requestContext.resolvedPartitionKeyRange, collectionResourceId)); - ConcurrentHashMap locationEndpointToLocationSpecificContextForPartition - = (ConcurrentHashMap) locationEndpointToLocationSpecificContextForPartitionField.get(partitionAndLocationSpecificUnavailabilityInfo); + ConcurrentHashMap locationEndpointToLocationSpecificContextForPartition + = (ConcurrentHashMap) locationEndpointToLocationSpecificContextForPartitionField.get(partitionAndLocationSpecificUnavailabilityInfo); LocationSpecificHealthContext locationSpecificHealthContext - = locationEndpointToLocationSpecificContextForPartition.get(LocationEastUs2EndpointToLocationPair.getKey()); + = locationEndpointToLocationSpecificContextForPartition.get(new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getKey())); assertThat(locationSpecificHealthContext.isRegionAvailableToProcessRequests()).isTrue(); assertThat(locationSpecificHealthContext.isExceptionThresholdBreached()).isFalse(); @@ -235,12 +235,12 @@ public void recordHealthyWithFailuresToUnavailableStatusTransition(String partit String maxExclusive = "BB"; String collectionResourceId = "dbs/db1/colls/coll1"; - List applicableReadWriteEndpoints = ImmutableList.of( + List applicableReadWriteEndpoints = ImmutableList.of( LocationEastUs2EndpointToLocationPair, LocationEastUsEndpointToLocationPair, LocationCentralUsEndpointToLocationPair) .stream() - .map(Pair::getLeft) + .map(uriStringPair -> new RegionalRoutingContext(uriStringPair.getLeft())) .collect(Collectors.toList()); RxDocumentServiceRequest request = constructRxDocumentServiceRequestInstance( @@ -253,15 +253,15 @@ public void recordHealthyWithFailuresToUnavailableStatusTransition(String partit maxExclusive, LocationEastUs2EndpointToLocationPair.getKey()); - Mockito.when(this.globalEndpointManagerMock.getApplicableWriteEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); - Mockito.when(this.globalEndpointManagerMock.getApplicableReadEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); + Mockito.when(this.globalEndpointManagerMock.getApplicableWriteEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); + Mockito.when(this.globalEndpointManagerMock.getApplicableReadEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); int exceptionCountToHandle = globalPartitionEndpointManagerForCircuitBreaker.getConsecutiveExceptionBasedCircuitBreaker().getAllowedExceptionCountToMaintainStatus(LocationHealthStatus.HealthyWithFailures, readOperationTrue); for (int i = 1; i <= exceptionCountToHandle; i++) { globalPartitionEndpointManagerForCircuitBreaker - .handleLocationExceptionForPartitionKeyRange(request, LocationEastUs2EndpointToLocationPair.getKey()); + .handleLocationExceptionForPartitionKeyRange(request, new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getKey())); } Class[] enclosedClasses = GlobalPartitionEndpointManagerForCircuitBreaker.class.getDeclaredClasses(); @@ -283,11 +283,11 @@ public void recordHealthyWithFailuresToUnavailableStatusTransition(String partit Object partitionAndLocationSpecificUnavailabilityInfo = partitionKeyRangeToLocationSpecificUnavailabilityInfo.get(new PartitionKeyRangeWrapper(request.requestContext.resolvedPartitionKeyRange, collectionResourceId)); - ConcurrentHashMap locationEndpointToLocationSpecificContextForPartition - = (ConcurrentHashMap) locationEndpointToLocationSpecificContextForPartitionField.get(partitionAndLocationSpecificUnavailabilityInfo); + ConcurrentHashMap locationEndpointToLocationSpecificContextForPartition + = (ConcurrentHashMap) locationEndpointToLocationSpecificContextForPartitionField.get(partitionAndLocationSpecificUnavailabilityInfo); LocationSpecificHealthContext locationSpecificHealthContext - = locationEndpointToLocationSpecificContextForPartition.get(LocationEastUs2EndpointToLocationPair.getKey()); + = locationEndpointToLocationSpecificContextForPartition.get(new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getKey())); assertThat(locationSpecificHealthContext.isRegionAvailableToProcessRequests()).isFalse(); assertThat(locationSpecificHealthContext.isExceptionThresholdBreached()).isTrue(); @@ -310,12 +310,12 @@ public void recordUnavailableToHealthyTentativeStatusTransition(String partition String maxExclusive = "BB"; String collectionResourceId = "dbs/db1/colls/coll1"; - List applicableReadWriteEndpoints = ImmutableList.of( + List applicableReadWriteEndpoints = ImmutableList.of( LocationEastUs2EndpointToLocationPair, LocationEastUsEndpointToLocationPair, LocationCentralUsEndpointToLocationPair) .stream() - .map(Pair::getLeft) + .map(uriStringPair -> new RegionalRoutingContext(uriStringPair.getLeft())) .collect(Collectors.toList()); RxDocumentServiceRequest request = constructRxDocumentServiceRequestInstance( @@ -328,15 +328,15 @@ public void recordUnavailableToHealthyTentativeStatusTransition(String partition maxExclusive, LocationEastUs2EndpointToLocationPair.getKey()); - Mockito.when(this.globalEndpointManagerMock.getApplicableWriteEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); - Mockito.when(this.globalEndpointManagerMock.getApplicableReadEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); + Mockito.when(this.globalEndpointManagerMock.getApplicableWriteEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); + Mockito.when(this.globalEndpointManagerMock.getApplicableReadEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); int exceptionCountToHandle = globalPartitionEndpointManagerForCircuitBreaker.getConsecutiveExceptionBasedCircuitBreaker().getAllowedExceptionCountToMaintainStatus(LocationHealthStatus.HealthyWithFailures, readOperationTrue); for (int i = 1; i <= exceptionCountToHandle; i++) { globalPartitionEndpointManagerForCircuitBreaker - .handleLocationExceptionForPartitionKeyRange(request, LocationEastUs2EndpointToLocationPair.getKey()); + .handleLocationExceptionForPartitionKeyRange(request, new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getKey())); } Class[] enclosedClasses = GlobalPartitionEndpointManagerForCircuitBreaker.class.getDeclaredClasses(); @@ -358,11 +358,11 @@ public void recordUnavailableToHealthyTentativeStatusTransition(String partition Object partitionAndLocationSpecificUnavailabilityInfo = partitionKeyRangeToLocationSpecificUnavailabilityInfo.get(new PartitionKeyRangeWrapper(request.requestContext.resolvedPartitionKeyRange, collectionResourceId)); - ConcurrentHashMap locationEndpointToLocationSpecificContextForPartition - = (ConcurrentHashMap) locationEndpointToLocationSpecificContextForPartitionField.get(partitionAndLocationSpecificUnavailabilityInfo); + ConcurrentHashMap locationEndpointToLocationSpecificContextForPartition + = (ConcurrentHashMap) locationEndpointToLocationSpecificContextForPartitionField.get(partitionAndLocationSpecificUnavailabilityInfo); LocationSpecificHealthContext locationSpecificHealthContext - = locationEndpointToLocationSpecificContextForPartition.get(LocationEastUs2EndpointToLocationPair.getKey()); + = locationEndpointToLocationSpecificContextForPartition.get(new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getKey())); assertThat(locationSpecificHealthContext.isRegionAvailableToProcessRequests()).isFalse(); assertThat(locationSpecificHealthContext.isExceptionThresholdBreached()).isTrue(); @@ -373,7 +373,7 @@ public void recordUnavailableToHealthyTentativeStatusTransition(String partition throw new RuntimeException(ex); } - locationSpecificHealthContext = locationEndpointToLocationSpecificContextForPartition.get(LocationEastUs2EndpointToLocationPair.getKey()); + locationSpecificHealthContext = locationEndpointToLocationSpecificContextForPartition.get(new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getKey())); assertThat(locationSpecificHealthContext.isRegionAvailableToProcessRequests()).isTrue(); assertThat(locationSpecificHealthContext.isExceptionThresholdBreached()).isFalse(); @@ -396,12 +396,12 @@ public void recordHealthyTentativeToHealthyStatusTransition(String partitionLeve String maxExclusive = "BB"; String collectionResourceId = "dbs/db1/colls/coll1"; - List applicableReadWriteEndpoints = ImmutableList.of( + List applicableReadWriteEndpoints = ImmutableList.of( LocationEastUs2EndpointToLocationPair, LocationEastUsEndpointToLocationPair, LocationCentralUsEndpointToLocationPair) .stream() - .map(Pair::getLeft) + .map(uriStringPair -> new RegionalRoutingContext(uriStringPair.getLeft())) .collect(Collectors.toList()); RxDocumentServiceRequest request = constructRxDocumentServiceRequestInstance( @@ -414,15 +414,15 @@ public void recordHealthyTentativeToHealthyStatusTransition(String partitionLeve maxExclusive, LocationEastUs2EndpointToLocationPair.getKey()); - Mockito.when(this.globalEndpointManagerMock.getApplicableWriteEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); - Mockito.when(this.globalEndpointManagerMock.getApplicableReadEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); + Mockito.when(this.globalEndpointManagerMock.getApplicableWriteEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); + Mockito.when(this.globalEndpointManagerMock.getApplicableReadEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); int exceptionCountToHandle = globalPartitionEndpointManagerForCircuitBreaker.getConsecutiveExceptionBasedCircuitBreaker().getAllowedExceptionCountToMaintainStatus(LocationHealthStatus.HealthyWithFailures, readOperationTrue); for (int i = 1; i <= exceptionCountToHandle; i++) { globalPartitionEndpointManagerForCircuitBreaker - .handleLocationExceptionForPartitionKeyRange(request, LocationEastUs2EndpointToLocationPair.getKey()); + .handleLocationExceptionForPartitionKeyRange(request, new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getKey())); } Class[] enclosedClasses = GlobalPartitionEndpointManagerForCircuitBreaker.class.getDeclaredClasses(); @@ -444,11 +444,11 @@ public void recordHealthyTentativeToHealthyStatusTransition(String partitionLeve Object partitionAndLocationSpecificUnavailabilityInfo = partitionKeyRangeToLocationSpecificUnavailabilityInfo.get(new PartitionKeyRangeWrapper(request.requestContext.resolvedPartitionKeyRange, collectionResourceId)); - ConcurrentHashMap locationEndpointToLocationSpecificContextForPartition - = (ConcurrentHashMap) locationEndpointToLocationSpecificContextForPartitionField.get(partitionAndLocationSpecificUnavailabilityInfo); + ConcurrentHashMap locationEndpointToLocationSpecificContextForPartition + = (ConcurrentHashMap) locationEndpointToLocationSpecificContextForPartitionField.get(partitionAndLocationSpecificUnavailabilityInfo); LocationSpecificHealthContext locationSpecificHealthContext - = locationEndpointToLocationSpecificContextForPartition.get(LocationEastUs2EndpointToLocationPair.getKey()); + = locationEndpointToLocationSpecificContextForPartition.get(new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getKey())); assertThat(locationSpecificHealthContext.isRegionAvailableToProcessRequests()).isFalse(); assertThat(locationSpecificHealthContext.isExceptionThresholdBreached()).isTrue(); @@ -459,7 +459,7 @@ public void recordHealthyTentativeToHealthyStatusTransition(String partitionLeve throw new RuntimeException(ex); } - locationSpecificHealthContext = locationEndpointToLocationSpecificContextForPartition.get(LocationEastUs2EndpointToLocationPair.getKey()); + locationSpecificHealthContext = locationEndpointToLocationSpecificContextForPartition.get(new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getKey())); int successCountToUpgradeStatus = globalPartitionEndpointManagerForCircuitBreaker.getConsecutiveExceptionBasedCircuitBreaker().getMinimumSuccessCountForStatusUpgrade(LocationHealthStatus.HealthyTentative, readOperationTrue); @@ -489,12 +489,12 @@ public void recordHealthyTentativeToUnavailableTransition(String partitionLevelC String maxExclusive = "BB"; String collectionResourceId = "dbs/db1/colls/coll1"; - List applicableReadWriteEndpoints = ImmutableList.of( + List applicableReadWriteEndpoints = ImmutableList.of( LocationEastUs2EndpointToLocationPair, LocationEastUsEndpointToLocationPair, LocationCentralUsEndpointToLocationPair) .stream() - .map(Pair::getLeft) + .map(uriStringPair -> new RegionalRoutingContext(uriStringPair.getLeft())) .collect(Collectors.toList()); RxDocumentServiceRequest request = constructRxDocumentServiceRequestInstance( @@ -507,15 +507,15 @@ public void recordHealthyTentativeToUnavailableTransition(String partitionLevelC maxExclusive, LocationEastUs2EndpointToLocationPair.getKey()); - Mockito.when(this.globalEndpointManagerMock.getApplicableWriteEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); - Mockito.when(this.globalEndpointManagerMock.getApplicableReadEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); + Mockito.when(this.globalEndpointManagerMock.getApplicableWriteEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); + Mockito.when(this.globalEndpointManagerMock.getApplicableReadEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); int exceptionCountToHandle = globalPartitionEndpointManagerForCircuitBreaker.getConsecutiveExceptionBasedCircuitBreaker().getAllowedExceptionCountToMaintainStatus(LocationHealthStatus.HealthyWithFailures, readOperationTrue); for (int i = 1; i <= exceptionCountToHandle; i++) { globalPartitionEndpointManagerForCircuitBreaker - .handleLocationExceptionForPartitionKeyRange(request, LocationEastUs2EndpointToLocationPair.getKey()); + .handleLocationExceptionForPartitionKeyRange(request, new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getKey())); } Class[] enclosedClasses = GlobalPartitionEndpointManagerForCircuitBreaker.class.getDeclaredClasses(); @@ -537,11 +537,11 @@ public void recordHealthyTentativeToUnavailableTransition(String partitionLevelC Object partitionAndLocationSpecificUnavailabilityInfo = partitionKeyRangeToLocationSpecificUnavailabilityInfo.get(new PartitionKeyRangeWrapper(request.requestContext.resolvedPartitionKeyRange, collectionResourceId)); - ConcurrentHashMap locationEndpointToLocationSpecificContextForPartition - = (ConcurrentHashMap) locationEndpointToLocationSpecificContextForPartitionField.get(partitionAndLocationSpecificUnavailabilityInfo); + ConcurrentHashMap locationEndpointToLocationSpecificContextForPartition + = (ConcurrentHashMap) locationEndpointToLocationSpecificContextForPartitionField.get(partitionAndLocationSpecificUnavailabilityInfo); LocationSpecificHealthContext locationSpecificHealthContext - = locationEndpointToLocationSpecificContextForPartition.get(LocationEastUs2EndpointToLocationPair.getKey()); + = locationEndpointToLocationSpecificContextForPartition.get(new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getKey())); assertThat(locationSpecificHealthContext.isRegionAvailableToProcessRequests()).isFalse(); assertThat(locationSpecificHealthContext.isExceptionThresholdBreached()).isTrue(); @@ -556,10 +556,10 @@ public void recordHealthyTentativeToUnavailableTransition(String partitionLevelC for (int i = 1; i <= exceptionCountToHandle; i++) { globalPartitionEndpointManagerForCircuitBreaker - .handleLocationExceptionForPartitionKeyRange(request, LocationEastUs2EndpointToLocationPair.getKey()); + .handleLocationExceptionForPartitionKeyRange(request, new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getKey())); } - locationSpecificHealthContext = locationEndpointToLocationSpecificContextForPartition.get(LocationEastUs2EndpointToLocationPair.getKey()); + locationSpecificHealthContext = locationEndpointToLocationSpecificContextForPartition.get(new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getKey())); assertThat(locationSpecificHealthContext.isRegionAvailableToProcessRequests()).isFalse(); assertThat(locationSpecificHealthContext.isExceptionThresholdBreached()).isTrue(); @@ -581,12 +581,12 @@ public void allRegionsUnavailableHandling(String partitionLevelCircuitBreakerCon String maxExclusive = "BB"; String collectionResourceId = "dbs/db1/colls/coll1"; - List applicableReadWriteEndpoints = ImmutableList.of( + List applicableReadWriteEndpoints = ImmutableList.of( LocationEastUs2EndpointToLocationPair, LocationEastUsEndpointToLocationPair, LocationCentralUsEndpointToLocationPair) .stream() - .map(Pair::getLeft) + .map(uriStringPair -> new RegionalRoutingContext(uriStringPair.getLeft())) .collect(Collectors.toList()); RxDocumentServiceRequest request = constructRxDocumentServiceRequestInstance( @@ -599,8 +599,8 @@ public void allRegionsUnavailableHandling(String partitionLevelCircuitBreakerCon maxExclusive, LocationEastUs2EndpointToLocationPair.getKey()); - Mockito.when(this.globalEndpointManagerMock.getApplicableWriteEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); - Mockito.when(this.globalEndpointManagerMock.getApplicableReadEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); + Mockito.when(this.globalEndpointManagerMock.getApplicableWriteEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); + Mockito.when(this.globalEndpointManagerMock.getApplicableReadEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); int exceptionCountToHandle = globalPartitionEndpointManagerForCircuitBreaker @@ -609,11 +609,11 @@ public void allRegionsUnavailableHandling(String partitionLevelCircuitBreakerCon for (int i = 1; i <= exceptionCountToHandle; i++) { globalPartitionEndpointManagerForCircuitBreaker - .handleLocationExceptionForPartitionKeyRange(request, LocationEastUs2EndpointToLocationPair.getKey()); + .handleLocationExceptionForPartitionKeyRange(request, new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getKey())); globalPartitionEndpointManagerForCircuitBreaker - .handleLocationExceptionForPartitionKeyRange(request, LocationEastUsEndpointToLocationPair.getKey()); + .handleLocationExceptionForPartitionKeyRange(request, new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getKey())); globalPartitionEndpointManagerForCircuitBreaker - .handleLocationExceptionForPartitionKeyRange(request, LocationCentralUsEndpointToLocationPair.getKey()); + .handleLocationExceptionForPartitionKeyRange(request, new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getKey())); } Class[] enclosedClasses = GlobalPartitionEndpointManagerForCircuitBreaker.class.getDeclaredClasses(); @@ -653,12 +653,12 @@ public void multiContainerBothWithSinglePartitionHealthyToUnavailableHandling(St String collectionResourceId1 = "dbs/db1/colls/coll1"; String collectionResourceId2 = "dbs/db1/colls/coll2"; - List applicableReadWriteEndpoints = ImmutableList.of( + List applicableReadWriteEndpoints = ImmutableList.of( LocationEastUs2EndpointToLocationPair, LocationEastUsEndpointToLocationPair, LocationCentralUsEndpointToLocationPair) .stream() - .map(Pair::getLeft) + .map(uriStringPair -> new RegionalRoutingContext(uriStringPair.getLeft())) .collect(Collectors.toList()); RxDocumentServiceRequest request1 = constructRxDocumentServiceRequestInstance( @@ -681,15 +681,15 @@ public void multiContainerBothWithSinglePartitionHealthyToUnavailableHandling(St maxExclusive, LocationEastUs2EndpointToLocationPair.getKey()); - Mockito.when(this.globalEndpointManagerMock.getApplicableWriteEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); - Mockito.when(this.globalEndpointManagerMock.getApplicableReadEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); + Mockito.when(this.globalEndpointManagerMock.getApplicableWriteEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); + Mockito.when(this.globalEndpointManagerMock.getApplicableReadEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); int exceptionCountToHandle = globalPartitionEndpointManagerForCircuitBreaker.getConsecutiveExceptionBasedCircuitBreaker().getAllowedExceptionCountToMaintainStatus(LocationHealthStatus.HealthyWithFailures, readOperationTrue); for (int i = 1; i <= exceptionCountToHandle; i++) { globalPartitionEndpointManagerForCircuitBreaker - .handleLocationExceptionForPartitionKeyRange(request1, LocationEastUs2EndpointToLocationPair.getKey()); + .handleLocationExceptionForPartitionKeyRange(request1, new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getKey())); } globalPartitionEndpointManagerForCircuitBreaker.handleLocationSuccessForPartitionKeyRange(request2); @@ -714,21 +714,21 @@ public void multiContainerBothWithSinglePartitionHealthyToUnavailableHandling(St = partitionKeyRangeToLocationSpecificUnavailabilityInfo.get(new PartitionKeyRangeWrapper( new PartitionKeyRange(pkRangeId, minInclusive, maxExclusive), collectionResourceId1)); - ConcurrentHashMap locationEndpointToLocationSpecificContextForPartitionForColl1 - = (ConcurrentHashMap) locationEndpointToLocationSpecificContextForPartitionField.get(partitionLevelLocationUnavailabilityInfoSnapshotForColl1); + ConcurrentHashMap locationEndpointToLocationSpecificContextForPartitionForColl1 + = (ConcurrentHashMap) locationEndpointToLocationSpecificContextForPartitionField.get(partitionLevelLocationUnavailabilityInfoSnapshotForColl1); Object partitionLevelLocationUnavailabilityInfoSnapshotForColl2 = partitionKeyRangeToLocationSpecificUnavailabilityInfo.get(new PartitionKeyRangeWrapper( new PartitionKeyRange(pkRangeId, minInclusive, maxExclusive), collectionResourceId2)); - ConcurrentHashMap locationEndpointToLocationSpecificContextForPartitionForColl2 - = (ConcurrentHashMap) locationEndpointToLocationSpecificContextForPartitionField.get(partitionLevelLocationUnavailabilityInfoSnapshotForColl2); + ConcurrentHashMap locationEndpointToLocationSpecificContextForPartitionForColl2 + = (ConcurrentHashMap) locationEndpointToLocationSpecificContextForPartitionField.get(partitionLevelLocationUnavailabilityInfoSnapshotForColl2); LocationSpecificHealthContext locationSpecificHealthContext1 - = locationEndpointToLocationSpecificContextForPartitionForColl1.get(LocationEastUs2EndpointToLocationPair.getKey()); + = locationEndpointToLocationSpecificContextForPartitionForColl1.get(new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getKey())); LocationSpecificHealthContext locationSpecificHealthContext2 - = locationEndpointToLocationSpecificContextForPartitionForColl2.get(LocationEastUs2EndpointToLocationPair.getKey()); + = locationEndpointToLocationSpecificContextForPartitionForColl2.get(new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getKey())); assertThat(locationSpecificHealthContext1.isRegionAvailableToProcessRequests()).isFalse(); assertThat(locationSpecificHealthContext1.isExceptionThresholdBreached()).isTrue(); @@ -766,16 +766,16 @@ public void allRegionsUnavailableHandlingWithMultiThreading(String partitionLeve String collectionResourceId = "dbs/db1/colls/coll1"; PartitionKeyRange partitionKeyRange = new PartitionKeyRange(pkRangeId, minInclusive, maxExclusive); - List applicableReadWriteEndpoints = ImmutableList.of( + List applicableReadWriteEndpoints = ImmutableList.of( LocationEastUs2EndpointToLocationPair, LocationEastUsEndpointToLocationPair, LocationCentralUsEndpointToLocationPair) .stream() - .map(Pair::getLeft) + .map(uriStringPair -> new RegionalRoutingContext(uriStringPair.getLeft())) .collect(Collectors.toList()); - Mockito.when(this.globalEndpointManagerMock.getApplicableWriteEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); - Mockito.when(this.globalEndpointManagerMock.getApplicableReadEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); + Mockito.when(this.globalEndpointManagerMock.getApplicableWriteEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); + Mockito.when(this.globalEndpointManagerMock.getApplicableReadEndpoints(Mockito.anyList())).thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableReadWriteEndpoints)); RxDocumentServiceRequest requestCentralUs = constructRxDocumentServiceRequestInstance( readOperationTrue ? OperationType.Read : OperationType.Create, @@ -883,10 +883,10 @@ private static void validateAllRegionsAreNotUnavailableAfterExceptionInLocation( URI locationWithFailure, String collectionResourceId, PartitionKeyRange partitionKeyRange, - List applicableReadWriteLocations) { + List applicableReadWriteLocations) { logger.warn("Handling exception for {}", locationWithFailure.getPath()); - globalPartitionEndpointManagerForCircuitBreaker.handleLocationExceptionForPartitionKeyRange(request, locationWithFailure); + globalPartitionEndpointManagerForCircuitBreaker.handleLocationExceptionForPartitionKeyRange(request, new RegionalRoutingContext(locationWithFailure)); List unavailableRegions = globalPartitionEndpointManagerForCircuitBreaker.getUnavailableRegionsForPartitionKeyRange(collectionResourceId, partitionKeyRange, request.getOperationType()); @@ -914,7 +914,9 @@ private RxDocumentServiceRequest constructRxDocumentServiceRequestInstance( request.requestContext.resolvedPartitionKeyRange = new PartitionKeyRange(partitionKeyRangeId, minInclusive, maxExclusive); request.requestContext.resolvedPartitionKeyRangeForCircuitBreaker = request.requestContext.resolvedPartitionKeyRange; - request.requestContext.locationEndpointToRoute = locationEndpointToRoute; + + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointToRoute); + request.requestContext.setExcludeRegions(Collections.emptyList()); request.requestContext.setPointOperationContext( new PointOperationContextForCircuitBreaker( diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PartitionLevelCircuitBreakerTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PartitionLevelCircuitBreakerTests.java index 52c1360d21dc..3d6963ad3cf0 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PartitionLevelCircuitBreakerTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PartitionLevelCircuitBreakerTests.java @@ -26,6 +26,7 @@ import com.azure.cosmos.implementation.feedranges.FeedRangeEpkImpl; import com.azure.cosmos.implementation.feedranges.FeedRangePartitionKeyImpl; import com.azure.cosmos.implementation.guava25.base.Function; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import com.azure.cosmos.models.CosmosBatch; import com.azure.cosmos.models.CosmosBatchResponse; import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; @@ -60,7 +61,6 @@ import reactor.core.publisher.Mono; import java.lang.reflect.Field; -import java.net.URI; import java.time.Duration; import java.util.ArrayList; import java.util.Collections; @@ -5144,8 +5144,8 @@ private static double getEstimatedFailureCountSeenPerRegionPerPartitionKeyRange( return 0d; } - ConcurrentHashMap locationEndpointToLocationSpecificContextForPartition - = (ConcurrentHashMap) locationEndpointToLocationSpecificContextForPartitionField.get(partitionAndLocationSpecificUnavailabilityInfo); + ConcurrentHashMap locationEndpointToLocationSpecificContextForPartition + = (ConcurrentHashMap) locationEndpointToLocationSpecificContextForPartitionField.get(partitionAndLocationSpecificUnavailabilityInfo); int count = 0; boolean failuresExist = false; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ProactiveConnectionManagementTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ProactiveConnectionManagementTest.java index a105b2da7ed4..fbfeff9897eb 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ProactiveConnectionManagementTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ProactiveConnectionManagementTest.java @@ -29,6 +29,7 @@ import com.azure.cosmos.implementation.routing.CollectionRoutingMap; import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper; import com.azure.cosmos.implementation.routing.PartitionKeyRangeIdentity; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import com.azure.cosmos.models.CosmosContainerIdentity; import com.azure.cosmos.rx.TestSuiteBase; import org.testng.annotations.BeforeClass; @@ -180,11 +181,13 @@ public void openConnectionsAndInitCachesWithContainer(ProactiveConnectionManagem cosmosAsyncContainer.openConnectionsAndInitCaches(proactiveConnectionRegionCount).block(); - UnmodifiableList readEndpoints = + UnmodifiableList readEndpoints = globalEndpointManager.getReadEndpoints(); + List proactiveConnectionEndpoints = readEndpoints.subList( 0, - Math.min(readEndpoints.size(),proactiveContainerInitConfig.getProactiveConnectionRegionsCount())); + Math.min(readEndpoints.size(),proactiveContainerInitConfig.getProactiveConnectionRegionsCount())) + .stream().map(RegionalRoutingContext::getGatewayRegionalEndpoint).collect(Collectors.toList()); Mono asyncContainerMono = Mono.just(cosmosAsyncContainer); @@ -342,10 +345,14 @@ public void openConnectionsAndInitCachesWithCosmosClient_And_PerContainerConnect ConcurrentHashMap routingMap = getRoutingMap(rxDocumentClient); ConcurrentHashMap collectionInfoByNameMap = getCollectionInfoByNameMap(rxDocumentClient); Set endpoints = ConcurrentHashMap.newKeySet(); - UnmodifiableList readEndpoints = globalEndpointManager.getReadEndpoints(); + UnmodifiableList readEndpoints = globalEndpointManager.getReadEndpoints(); + List proactiveConnectionEndpoints = readEndpoints.subList( 0, - Math.min(readEndpoints.size(), proactiveContainerInitConfig.getProactiveConnectionRegionsCount())); + Math.min(readEndpoints.size(), proactiveContainerInitConfig.getProactiveConnectionRegionsCount())) + .stream() + .map(RegionalRoutingContext::getGatewayRegionalEndpoint) + .collect(Collectors.toList()); Flux asyncContainerFlux = Flux.fromIterable(asyncContainers); Flux>> partitionKeyRangeFlux = @@ -488,10 +495,13 @@ public void openConnectionsAndInitCachesWithCosmosClient_And_PerContainerConnect ConcurrentHashMap routingMap = getRoutingMap(rxDocumentClient); ConcurrentHashMap collectionInfoByNameMap = getCollectionInfoByNameMap(rxDocumentClient); Set endpoints = ConcurrentHashMap.newKeySet(); - UnmodifiableList readEndpoints = globalEndpointManager.getReadEndpoints(); + UnmodifiableList readEndpoints = globalEndpointManager.getReadEndpoints(); List proactiveConnectionEndpoints = readEndpoints.subList( 0, - Math.min(readEndpoints.size(), proactiveContainerInitConfig.getProactiveConnectionRegionsCount())); + Math.min(readEndpoints.size(), proactiveContainerInitConfig.getProactiveConnectionRegionsCount())) + .stream() + .map(RegionalRoutingContext::getGatewayRegionalEndpoint) + .collect(Collectors.toList());; Flux asyncContainerFlux = Flux.fromIterable(asyncContainers); Flux>> partitionKeyRangeFlux = @@ -656,10 +666,13 @@ public void openConnectionsAndInitCachesWithCosmosClient_And_PerContainerConnect ConcurrentHashMap routingMap = getRoutingMap(rxDocumentClient); ConcurrentHashMap collectionInfoByNameMap = getCollectionInfoByNameMap(rxDocumentClient); Set endpoints = ConcurrentHashMap.newKeySet(); - UnmodifiableList readEndpoints = globalEndpointManager.getReadEndpoints(); + UnmodifiableList readEndpoints = globalEndpointManager.getReadEndpoints(); List proactiveConnectionEndpoints = readEndpoints.subList( 0, - Math.min(readEndpoints.size(), proactiveContainerInitConfig.getProactiveConnectionRegionsCount())); + Math.min(readEndpoints.size(), proactiveContainerInitConfig.getProactiveConnectionRegionsCount())) + .stream() + .map(RegionalRoutingContext::getGatewayRegionalEndpoint) + .collect(Collectors.toList());; Flux asyncContainerFlux = Flux.fromIterable(asyncContainers); Flux>> partitionKeyRangeFlux = diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ClientRetryPolicyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ClientRetryPolicyTest.java index 31aa3abe2349..13f6ca4cd0ca 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ClientRetryPolicyTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ClientRetryPolicyTest.java @@ -8,6 +8,7 @@ import com.azure.cosmos.ThrottlingRetryOptions; import com.azure.cosmos.implementation.circuitBreaker.GlobalPartitionEndpointManagerForCircuitBreaker; import com.azure.cosmos.implementation.directconnectivity.ChannelAcquisitionException; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import io.netty.handler.timeout.ReadTimeoutException; import io.reactivex.subscribers.TestSubscriber; import org.mockito.Mockito; @@ -66,7 +67,7 @@ public void networkFailureOnRead() throws Exception { ThrottlingRetryOptions throttlingRetryOptions = new ThrottlingRetryOptions(); GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new URI("http://localhost")).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); + Mockito.doReturn(new RegionalRoutingContext(new URI("http://localhost"))).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(Mockito.eq(null), Mockito.eq(false)); ClientRetryPolicy clientRetryPolicy = new ClientRetryPolicy(mockDiagnosticsClientContext(), endpointManager, true, throttlingRetryOptions, null, globalPartitionEndpointManager); @@ -107,7 +108,7 @@ public void shouldRetryOnGatewayTimeout( GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new URI("http://localhost")).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); + Mockito.doReturn(new RegionalRoutingContext(new URI("http://localhost"))).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(Mockito.eq(null), Mockito.eq(true)); ClientRetryPolicy clientRetryPolicy = new ClientRetryPolicy( @@ -150,7 +151,7 @@ public void tcpNetworkFailureOnRead() throws Exception { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new URI("http://localhost")).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); + Mockito.doReturn(new RegionalRoutingContext(new URI("http://localhost"))).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(Mockito.eq(null), Mockito.eq(false)); Mockito.doReturn(2).when(endpointManager).getPreferredLocationCount(); ClientRetryPolicy clientRetryPolicy = new ClientRetryPolicy(mockDiagnosticsClientContext(), endpointManager, true, retryOptions, null, globalPartitionEndpointManager); @@ -198,7 +199,7 @@ public void networkFailureOnWrite() throws Exception { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new URI("http://localhost")).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); + Mockito.doReturn(new RegionalRoutingContext(new URI("http://localhost"))).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(Mockito.eq(null), Mockito.eq(false)); ClientRetryPolicy clientRetryPolicy = new ClientRetryPolicy(mockDiagnosticsClientContext(), endpointManager, true, throttlingRetryOptions, null, globalPartitionEndpointManager); @@ -233,7 +234,7 @@ public void tcpNetworkFailureOnWrite( GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new URI("http://localhost")).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); + Mockito.doReturn(new RegionalRoutingContext(new URI("http://localhost"))).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(Mockito.eq(null), Mockito.eq(false)); Mockito.doReturn(2).when(endpointManager).getPreferredLocationCount(); ClientRetryPolicy clientRetryPolicy = new ClientRetryPolicy(mockDiagnosticsClientContext(), endpointManager, true, retryOptions, null, globalPartitionEndpointManager); @@ -293,7 +294,7 @@ public void networkFailureOnUpsert() throws Exception { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new URI("http://localhost")).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); + Mockito.doReturn(new RegionalRoutingContext(new URI("http://localhost"))).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(Mockito.eq(null), Mockito.eq(false)); ClientRetryPolicy clientRetryPolicy = new ClientRetryPolicy(mockDiagnosticsClientContext(), endpointManager, true, throttlingRetryOptions, null, globalPartitionEndpointManager); @@ -326,7 +327,7 @@ public void tcpNetworkFailureOnUpsert() throws Exception { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new URI("http://localhost")).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); + Mockito.doReturn(new RegionalRoutingContext(new URI("http://localhost"))).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(Mockito.eq(null), Mockito.eq(false)); Mockito.doReturn(2).when(endpointManager).getPreferredLocationCount(); ClientRetryPolicy clientRetryPolicy = new ClientRetryPolicy(mockDiagnosticsClientContext(), endpointManager, true, retryOptions, null, globalPartitionEndpointManager); @@ -362,7 +363,7 @@ public void networkFailureOnDelete() throws Exception { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new URI("http://localhost")).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); + Mockito.doReturn(new RegionalRoutingContext(new URI("http://localhost"))).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(Mockito.eq(null), Mockito.eq(false)); ClientRetryPolicy clientRetryPolicy = new ClientRetryPolicy(mockDiagnosticsClientContext(), endpointManager, true, throttlingRetryOptions, null, globalPartitionEndpointManager); @@ -396,7 +397,7 @@ public void tcpNetworkFailureOnDelete() throws Exception { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new URI("http://localhost")).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); + Mockito.doReturn(new RegionalRoutingContext(new URI("http://localhost"))).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(Mockito.eq(null), Mockito.eq(false)); Mockito.doReturn(2).when(endpointManager).getPreferredLocationCount(); ClientRetryPolicy clientRetryPolicy = new ClientRetryPolicy(mockDiagnosticsClientContext(), endpointManager, true, retryOptions, null, globalPartitionEndpointManager); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java index 26c34a1c423c..8245f322aa5e 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java @@ -6,8 +6,6 @@ import com.azure.cosmos.implementation.clienttelemetry.MetricCategory; import com.azure.cosmos.implementation.clienttelemetry.TagName; import com.azure.cosmos.implementation.directconnectivity.Protocol; -import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils; -import io.netty.handler.ssl.SslContext; import org.testng.annotations.Test; import java.net.URI; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RegionScopedSessionContainerTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RegionScopedSessionContainerTest.java index 22bee3a04860..359a3a701ec8 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RegionScopedSessionContainerTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RegionScopedSessionContainerTest.java @@ -8,6 +8,7 @@ import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList; import com.azure.cosmos.implementation.guava25.collect.ImmutableList; import com.azure.cosmos.implementation.guava25.collect.ImmutableMap; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import com.azure.cosmos.models.ModelBridgeInternal; import com.azure.cosmos.models.ModelBridgeUtils; import com.azure.cosmos.models.PartitionKey; @@ -374,7 +375,11 @@ public void sessionContainer() throws Exception { int numPartitionKeyRangeIds = 5; String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); - UnmodifiableList endpoints = new UnmodifiableList<>(ImmutableList.of(LocationEastUsEndpointToLocationPair.getLeft(), LocationEastUs2EndpointToLocationPair.getLeft(), LocationCentralUsEndpointToLocationPair.getLeft())); + UnmodifiableList endpoints = new UnmodifiableList<>( + ImmutableList.of( + new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(LocationEastUsEndpointToLocationPair.getLeft()), Mockito.any())).thenReturn(regionContacted); @@ -391,7 +396,7 @@ public void sessionContainer() throws Exception { String resultantSessionToken = partitionKeyRangeId + ":" + lsn; RxDocumentServiceRequest request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document); - request.requestContext.locationEndpointToRoute = LocationEastUsEndpointToLocationPair.getKey(); + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getKey()); sessionContainer.setSessionToken( request, @@ -404,7 +409,7 @@ public void sessionContainer() throws Exception { RxDocumentServiceRequest request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.ReadFeed, ResourceType.DocumentCollection, "dbs/db1/colls/collName_1", Utils.getUTF8Bytes("content1"), new HashMap<>()); - request.requestContext.locationEndpointToRoute = LocationEastUsEndpointToLocationPair.getLeft(); + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()); ISessionToken sessionToken = sessionContainer.resolvePartitionLocalSessionToken(request, "range_1"); assertThat(sessionToken.getLSN()).isEqualTo(1); @@ -415,7 +420,7 @@ public void sessionContainer() throws Exception { GatewayTestUtils.setParent(resolvedPKRange, ImmutableList.of("range_2", "range_x")); dsrContext.resolvedPartitionKeyRange = resolvedPKRange; request.requestContext = dsrContext; - request.requestContext.locationEndpointToRoute = LocationEastUsEndpointToLocationPair.getLeft(); + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()); sessionToken = sessionContainer.resolvePartitionLocalSessionToken(request, resolvedPKRange.getId()); assertThat(sessionToken.getLSN()).isEqualTo(2); @@ -434,14 +439,18 @@ public void setSessionToken_NoSessionTokenForPartitionKeyRangeId() throws Except GlobalEndpointManager globalEndpointManagerMock = Mockito.mock(GlobalEndpointManager.class); ISessionContainer sessionContainer = new RegionScopedSessionContainer("127.0.0.1", false, globalEndpointManagerMock); - UnmodifiableList endpoints = new UnmodifiableList<>(ImmutableList.of(LocationEastUsEndpointToLocationPair.getLeft(), LocationEastUs2EndpointToLocationPair.getLeft(), LocationCentralUsEndpointToLocationPair.getLeft())); + UnmodifiableList endpoints = new UnmodifiableList<>( + ImmutableList.of( + new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(endpointContacted), Mockito.any())).thenReturn(regionContacted); RxDocumentServiceRequest request1 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Create, ResourceType.Document, collectionName + "/docs", Utils.getUTF8Bytes("content1"), new HashMap<>()); - request1.requestContext.locationEndpointToRoute = endpointContacted; + request1.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(endpointContacted); String sessionTokenWithPkRangeIdForRequest1 = partitionKeyRangeId + ":" + sessionToken; @@ -474,7 +483,7 @@ public void setSessionToken_NoSessionTokenForPartitionKeyRangeId() throws Except RxDocumentServiceRequest request2 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document, collectionName + "/docs", Utils.getUTF8Bytes(""), new HashMap<>()); - request2.requestContext.locationEndpointToRoute = endpointContacted; + request2.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(endpointContacted); ISessionToken resolvedSessionToken = sessionContainer.resolvePartitionLocalSessionToken(request2, partitionKeyRangeId); assertThat(resolvedSessionToken).isNotNull(); @@ -497,14 +506,18 @@ public void setSessionToken_MergeOldWithNew() throws Exception { GlobalEndpointManager globalEndpointManagerMock = Mockito.mock(GlobalEndpointManager.class); RegionScopedSessionContainer sessionContainer = new RegionScopedSessionContainer("127.0.0.1", false, globalEndpointManagerMock); - UnmodifiableList endpoints = new UnmodifiableList<>(ImmutableList.of(LocationEastUsEndpointToLocationPair.getLeft(), LocationEastUs2EndpointToLocationPair.getLeft(), LocationCentralUsEndpointToLocationPair.getLeft())); + UnmodifiableList endpoints = new UnmodifiableList<>( + ImmutableList.of( + new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); RxDocumentServiceRequest request1 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Create, ResourceType.Document, collectionName + "/docs", Utils.getUTF8Bytes("content1"), new HashMap<>()); - request1.requestContext.locationEndpointToRoute = locationEndpointContacted; + request1.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String initialSessionTokenWithPkRangeId = partitionKeyRangeId + ":" + initialSessionToken; @@ -526,7 +539,7 @@ public void setSessionToken_MergeOldWithNew() throws Exception { RxDocumentServiceRequest request2 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, ResourceType.Document, collectionName + "/docs", Utils.getUTF8Bytes(""), new HashMap<>()); - request2.requestContext.locationEndpointToRoute = locationEndpointContacted; + request2.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); ISessionToken resolvedSessionToken = sessionContainer.resolvePartitionLocalSessionToken(request2, partitionKeyRangeId); assertThat(resolvedSessionToken).isNotNull(); @@ -550,25 +563,28 @@ public void resolveGlobalSessionTokenReturnsEmptyStringOnCacheMiss() { String initialSessionToken = "1#100#1=20#2=5#3=30"; String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - String resultantSessionToken = partitionKeyRangeId + ":" + initialSessionToken; GlobalEndpointManager globalEndpointManagerMock = Mockito.mock(GlobalEndpointManager.class); RegionScopedSessionContainer sessionContainer = new RegionScopedSessionContainer("127.0.0.1", false, globalEndpointManagerMock); - UnmodifiableList endpoints = new UnmodifiableList<>(ImmutableList.of(LocationEastUsEndpointToLocationPair.getLeft(), LocationEastUs2EndpointToLocationPair.getLeft(), LocationCentralUsEndpointToLocationPair.getLeft())); + UnmodifiableList endpoints = new UnmodifiableList<>( + ImmutableList.of( + new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); RxDocumentServiceRequest collectionCreateRequest = RxDocumentServiceRequest.create( mockDiagnosticsClientContext(), OperationType.Create, ResourceType.DocumentCollection); - collectionCreateRequest.requestContext.locationEndpointToRoute = locationEndpointContacted; + collectionCreateRequest.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); sessionContainer.setSessionToken(collectionCreateRequest, documentCollectionId, "dbs/db1/colls1/collName", ImmutableMap.of(HttpConstants.HttpHeaders.SESSION_TOKEN, partitionKeyRangeId + ":" + initialSessionToken)); RxDocumentServiceRequest request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, ResourceType.Document, "dbs/db1/colls1/collName2/docs/doc1", new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); assertThat(StringUtils.EMPTY).isEqualTo(sessionContainer.resolveGlobalSessionToken(request)); } @@ -583,13 +599,17 @@ public void resolveGlobalSessionTokenReturnsTokenMapUsingName() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>(ImmutableList.of(LocationEastUsEndpointToLocationPair.getLeft(), LocationEastUs2EndpointToLocationPair.getLeft(), LocationCentralUsEndpointToLocationPair.getLeft())); + UnmodifiableList endpoints = new UnmodifiableList<>( + ImmutableList.of( + new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); RxDocumentServiceRequest request1 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document); - request1.requestContext.locationEndpointToRoute = locationEndpointContacted; + request1.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String sessionTokenForRequest1 = "range_0:1#100#1=20#2=5#3=30"; @@ -597,7 +617,7 @@ public void resolveGlobalSessionTokenReturnsTokenMapUsingName() { ImmutableMap.of(HttpConstants.HttpHeaders.SESSION_TOKEN, sessionTokenForRequest1)); RxDocumentServiceRequest request2 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document); - request2.requestContext.locationEndpointToRoute = locationEndpointContacted; + request2.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String sessionTokenForRequest2 = "range_1:1#101#1=20#2=5#3=30"; @@ -625,13 +645,17 @@ public void resolveGlobalSessionTokenReturnsTokenMapUsingResourceId() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>(ImmutableList.of(LocationEastUsEndpointToLocationPair.getLeft(), LocationEastUs2EndpointToLocationPair.getLeft(), LocationCentralUsEndpointToLocationPair.getLeft())); + UnmodifiableList endpoints = new UnmodifiableList<>( + ImmutableList.of( + new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); RxDocumentServiceRequest request1 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document); - request1.requestContext.locationEndpointToRoute = locationEndpointContacted; + request1.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String sessionTokenForRequest1 = "range_0:1#100#1=20#2=5#3=30"; @@ -639,7 +663,7 @@ public void resolveGlobalSessionTokenReturnsTokenMapUsingResourceId() { ImmutableMap.of(HttpConstants.HttpHeaders.SESSION_TOKEN, sessionTokenForRequest1)); RxDocumentServiceRequest request2 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document); - request2.requestContext.locationEndpointToRoute = locationEndpointContacted; + request2.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String sessionTokenForRequest2 = "range_1:1#101#1=20#2=5#3=30"; @@ -648,7 +672,7 @@ public void resolveGlobalSessionTokenReturnsTokenMapUsingResourceId() { RxDocumentServiceRequest request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, documentCollectionId, ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String sessionToken = sessionContainer.resolveGlobalSessionToken(request); @@ -669,13 +693,17 @@ public void resolveLocalSessionTokenReturnsTokenMapUsingName() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>(ImmutableList.of(LocationEastUsEndpointToLocationPair.getLeft(), LocationEastUs2EndpointToLocationPair.getLeft(), LocationCentralUsEndpointToLocationPair.getLeft())); + UnmodifiableList endpoints = new UnmodifiableList<>( + ImmutableList.of( + new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); RxDocumentServiceRequest request1 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document); - request1.requestContext.locationEndpointToRoute = locationEndpointContacted; + request1.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String sessionTokenForRequest1 = "range_0:1#100#1=20#2=5#3=30"; @@ -683,7 +711,7 @@ public void resolveLocalSessionTokenReturnsTokenMapUsingName() { ImmutableMap.of(HttpConstants.HttpHeaders.SESSION_TOKEN, sessionTokenForRequest1)); RxDocumentServiceRequest request2 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document); - request2.requestContext.locationEndpointToRoute = locationEndpointContacted; + request2.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String sessionTokenForRequest2 = "range_1:1#101#1=20#2=5#3=30"; @@ -709,13 +737,17 @@ public void resolveLocalSessionTokenReturnsTokenMapUsingResourceId() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>(ImmutableList.of(LocationEastUsEndpointToLocationPair.getLeft(), LocationEastUs2EndpointToLocationPair.getLeft(), LocationCentralUsEndpointToLocationPair.getLeft())); + UnmodifiableList endpoints = new UnmodifiableList<>( + ImmutableList.of( + new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); RxDocumentServiceRequest request1 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document); - request1.requestContext.locationEndpointToRoute = locationEndpointContacted; + request1.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String sessionTokenWithPkRangeIdForRequest1 = "range_0:1#100#1=20#2=5#3=30"; @@ -723,7 +755,7 @@ public void resolveLocalSessionTokenReturnsTokenMapUsingResourceId() { ImmutableMap.of(HttpConstants.HttpHeaders.SESSION_TOKEN, sessionTokenWithPkRangeIdForRequest1)); RxDocumentServiceRequest request2 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document); - request2.requestContext.locationEndpointToRoute = locationEndpointContacted; + request2.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String sessionTokenWithPkRangeIdForRequest2 = "range_1:1#101#1=20#2=5#3=30"; @@ -732,7 +764,7 @@ public void resolveLocalSessionTokenReturnsTokenMapUsingResourceId() { RxDocumentServiceRequest request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, documentCollectionId, ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); ISessionToken sessionToken = sessionContainer.resolvePartitionLocalSessionToken(request, "range_0"); assertThat(sessionToken).isNotNull(); @@ -754,13 +786,17 @@ public void resolveLocalSessionTokenReturnsNullOnPartitionMiss() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>(ImmutableList.of(LocationEastUsEndpointToLocationPair.getLeft(), LocationEastUs2EndpointToLocationPair.getLeft(), LocationCentralUsEndpointToLocationPair.getLeft())); + UnmodifiableList endpoints = new UnmodifiableList<>( + ImmutableList.of( + new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); RxDocumentServiceRequest request1 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document); - request1.requestContext.locationEndpointToRoute = locationEndpointContacted; + request1.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String sessionTokenWithPkRangeIdForRequest1 = "range_0:1#100#1=20#2=5#3=30"; @@ -768,7 +804,7 @@ public void resolveLocalSessionTokenReturnsNullOnPartitionMiss() { ImmutableMap.of(HttpConstants.HttpHeaders.SESSION_TOKEN, sessionTokenWithPkRangeIdForRequest1)); RxDocumentServiceRequest request2 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document); - request2.requestContext.locationEndpointToRoute = locationEndpointContacted; + request2.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String sessionTokenWithPkRangeIdForRequest2 = "range_1:1#101#1=20#2=5#3=30"; @@ -777,7 +813,7 @@ public void resolveLocalSessionTokenReturnsNullOnPartitionMiss() { RxDocumentServiceRequest requestToResultInPkRangeIdBasedMiss = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, documentCollectionId, ResourceType.Document, new HashMap<>()); - requestToResultInPkRangeIdBasedMiss.requestContext.locationEndpointToRoute = locationEndpointContacted; + requestToResultInPkRangeIdBasedMiss.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); requestToResultInPkRangeIdBasedMiss.requestContext.resolvedPartitionKeyRange = new PartitionKeyRange(); ISessionToken sessionToken = sessionContainer.resolvePartitionLocalSessionToken(requestToResultInPkRangeIdBasedMiss, "range_2"); @@ -795,13 +831,17 @@ public void resolveLocalSessionTokenReturnsNullOnCollectionMiss() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>(ImmutableList.of(LocationEastUsEndpointToLocationPair.getLeft(), LocationEastUs2EndpointToLocationPair.getLeft(), LocationCentralUsEndpointToLocationPair.getLeft())); + UnmodifiableList endpoints = new UnmodifiableList<>( + ImmutableList.of( + new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); RxDocumentServiceRequest request1 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document); - request1.requestContext.locationEndpointToRoute = locationEndpointContacted; + request1.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String sessionTokenWithPkRangeIdForRequest1 = "range_0:1#100#1=20#2=5#3=30"; @@ -809,7 +849,7 @@ public void resolveLocalSessionTokenReturnsNullOnCollectionMiss() { ImmutableMap.of(HttpConstants.HttpHeaders.SESSION_TOKEN, sessionTokenWithPkRangeIdForRequest1)); RxDocumentServiceRequest request2 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document); - request2.requestContext.locationEndpointToRoute = locationEndpointContacted; + request2.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String sessionTokenWithPkRangeIdForRequest2 = "range_1:1#101#1=20#2=5#3=30"; @@ -834,13 +874,17 @@ public void resolvePartitionLocalSessionTokenReturnsTokenOnParentMatch() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>(ImmutableList.of(LocationEastUsEndpointToLocationPair.getLeft(), LocationEastUs2EndpointToLocationPair.getLeft(), LocationCentralUsEndpointToLocationPair.getLeft())); + UnmodifiableList endpoints = new UnmodifiableList<>( + ImmutableList.of( + new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); RxDocumentServiceRequest request1 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document); - request1.requestContext.locationEndpointToRoute = locationEndpointContacted; + request1.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String sessionTokenWithPkRangeIdForRequest1 = "range_0:1#100#1=20#2=5#3=30"; @@ -848,7 +892,7 @@ public void resolvePartitionLocalSessionTokenReturnsTokenOnParentMatch() { ImmutableMap.of(HttpConstants.HttpHeaders.SESSION_TOKEN, sessionTokenWithPkRangeIdForRequest1)); RxDocumentServiceRequest request2 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document); - request2.requestContext.locationEndpointToRoute = locationEndpointContacted; + request2.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String sessionTokenWithPkRangeIdForRequest2 = "range_1:1#101#1=20#2=5#3=30"; @@ -876,13 +920,17 @@ public void clearTokenByCollectionFullNameRemovesToken() { URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); String unparsedSessionToken = "range_0:1#100#1=20#2=5#3=30"; - UnmodifiableList endpoints = new UnmodifiableList<>(ImmutableList.of(LocationEastUsEndpointToLocationPair.getLeft(), LocationEastUs2EndpointToLocationPair.getLeft(), LocationCentralUsEndpointToLocationPair.getLeft())); + UnmodifiableList endpoints = new UnmodifiableList<>( + ImmutableList.of( + new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); RxDocumentServiceRequest documentCollectionCreateRequest = createRequestEntity(OperationType.Create, ResourceType.DocumentCollection, LocationEastUsEndpointToLocationPair.getLeft()); - documentCollectionCreateRequest.requestContext.locationEndpointToRoute = locationEndpointContacted; + documentCollectionCreateRequest.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); sessionContainer.setSessionToken(documentCollectionCreateRequest, documentCollectionId, collectionFullName, ImmutableMap.of(HttpConstants.HttpHeaders.SESSION_TOKEN, unparsedSessionToken)); @@ -890,7 +938,7 @@ public void clearTokenByCollectionFullNameRemovesToken() { // Test getResourceId based RxDocumentServiceRequest request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, documentCollectionId, ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); ISessionToken sessionToken = sessionContainer.resolvePartitionLocalSessionToken(request, "range_0"); assertThat(sessionToken.getLSN()).isEqualTo(100); @@ -898,7 +946,7 @@ public void clearTokenByCollectionFullNameRemovesToken() { // Test names based request = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(),OperationType.Read, collectionFullName + "/docs/doc1", ResourceType.Document); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); sessionToken = sessionContainer.resolvePartitionLocalSessionToken(request, "range_0"); assertThat(sessionToken.getLSN()).isEqualTo(100); @@ -908,7 +956,7 @@ public void clearTokenByCollectionFullNameRemovesToken() { // Test resourceId based request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, documentCollectionId, ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); sessionToken = sessionContainer.resolvePartitionLocalSessionToken(request, "range_0"); assertThat(sessionToken).isNull(); @@ -916,7 +964,7 @@ public void clearTokenByCollectionFullNameRemovesToken() { // Test names based request = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(),OperationType.Read, collectionFullName + "/docs/doc1", ResourceType.Document); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); sessionToken = sessionContainer.resolvePartitionLocalSessionToken(request, "range_0"); assertThat(sessionToken).isNull(); @@ -935,13 +983,17 @@ public void clearTokenByResourceIdRemovesToken() { String unparsedSessionToken = "range_0:1#100#1=20#2=5#3=30"; - UnmodifiableList endpoints = new UnmodifiableList<>(ImmutableList.of(LocationEastUsEndpointToLocationPair.getLeft(), LocationEastUs2EndpointToLocationPair.getLeft(), LocationCentralUsEndpointToLocationPair.getLeft())); + UnmodifiableList endpoints = new UnmodifiableList<>( + ImmutableList.of( + new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); RxDocumentServiceRequest documentCollectionCreateRequest = createRequestEntity(OperationType.Create, ResourceType.DocumentCollection, LocationEastUsEndpointToLocationPair.getLeft()); - documentCollectionCreateRequest.requestContext.locationEndpointToRoute = locationEndpointContacted; + documentCollectionCreateRequest.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); sessionContainer.setSessionToken(documentCollectionCreateRequest, documentCollectionId, collectionFullName, ImmutableMap.of(HttpConstants.HttpHeaders.SESSION_TOKEN, unparsedSessionToken)); @@ -949,7 +1001,7 @@ public void clearTokenByResourceIdRemovesToken() { // Test resourceId based RxDocumentServiceRequest request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, documentCollectionId, ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); ISessionToken sessionToken = sessionContainer.resolvePartitionLocalSessionToken(request, "range_0"); assertThat(sessionToken).isNotNull(); @@ -958,7 +1010,7 @@ public void clearTokenByResourceIdRemovesToken() { // Test names based request = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(),OperationType.Read, collectionFullName + "/docs/doc1", ResourceType.Document); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); sessionToken = sessionContainer.resolvePartitionLocalSessionToken(request, "range_0"); assertThat(sessionToken).isNotNull(); @@ -969,7 +1021,7 @@ public void clearTokenByResourceIdRemovesToken() { // Test resourceId based request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, documentCollectionId, ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); sessionToken = sessionContainer.resolvePartitionLocalSessionToken(request, "range_0"); assertThat(sessionToken).isNull(); @@ -977,7 +1029,7 @@ public void clearTokenByResourceIdRemovesToken() { // Test names based request = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(),OperationType.Read, collectionFullName + "/docs/doc1", ResourceType.Document); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); sessionToken = sessionContainer.resolvePartitionLocalSessionToken(request, "range_0"); assertThat(sessionToken).isNull(); @@ -995,13 +1047,17 @@ public void clearTokenKeepsUnmatchedCollection() { URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); String unparsedSessionToken = "range_0:1#100#1=20#2=5#3=30"; - UnmodifiableList endpoints = new UnmodifiableList<>(ImmutableList.of(LocationEastUsEndpointToLocationPair.getLeft(), LocationEastUs2EndpointToLocationPair.getLeft(), LocationCentralUsEndpointToLocationPair.getLeft())); + UnmodifiableList endpoints = new UnmodifiableList<>( + ImmutableList.of( + new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); RxDocumentServiceRequest documentCollectionCreateRequest = createRequestEntity(OperationType.Create, ResourceType.DocumentCollection, LocationEastUsEndpointToLocationPair.getLeft()); - documentCollectionCreateRequest.requestContext.locationEndpointToRoute = locationEndpointContacted; + documentCollectionCreateRequest.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); sessionContainer.setSessionToken(documentCollectionCreateRequest, documentCollectionId1, collectionFullName1, ImmutableMap.of(HttpConstants.HttpHeaders.SESSION_TOKEN, unparsedSessionToken)); @@ -1009,7 +1065,7 @@ public void clearTokenKeepsUnmatchedCollection() { // Test resourceId based RxDocumentServiceRequest request1 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, documentCollectionId1, ResourceType.Document, new HashMap<>()); - request1.requestContext.locationEndpointToRoute = locationEndpointContacted; + request1.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String documentCollectionId2 = ResourceId.newDocumentCollectionId(getRandomDbId(), randomCollectionId - 1).getDocumentCollectionId().toString(); String collectionFullName2 = "dbs/db1/colls1/collName2"; @@ -1017,7 +1073,7 @@ public void clearTokenKeepsUnmatchedCollection() { // Test resourceId based RxDocumentServiceRequest request2 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, documentCollectionId2, ResourceType.Document, new HashMap<>()); - request2.requestContext.locationEndpointToRoute = locationEndpointContacted; + request2.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); sessionContainer.setSessionToken(request2, documentCollectionId2, collectionFullName2, ImmutableMap.of(HttpConstants.HttpHeaders.SESSION_TOKEN, unparsedSessionToken)); @@ -1050,14 +1106,18 @@ public void setSessionTokenSetsTokenWhenRequestIsntNameBased() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>(ImmutableList.of(LocationEastUsEndpointToLocationPair.getLeft(), LocationEastUs2EndpointToLocationPair.getLeft(), LocationCentralUsEndpointToLocationPair.getLeft())); + UnmodifiableList endpoints = new UnmodifiableList<>( + ImmutableList.of( + new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); RxDocumentServiceRequest request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, collectionFullName + "/docs/doc1", ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String sessionTokenWithPkRangeIdForRequest = "range_0:1#100#4=90#5=1"; request.setResourceId(documentCollectionId); @@ -1065,7 +1125,7 @@ public void setSessionTokenSetsTokenWhenRequestIsntNameBased() { assertThat(request.getIsNameBased()).isFalse(); sessionContainer.setSessionToken(request, ImmutableMap.of(HttpConstants.HttpHeaders.SESSION_TOKEN, sessionTokenWithPkRangeIdForRequest)); request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, documentCollectionId, ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); ISessionToken sessionToken = sessionContainer.resolvePartitionLocalSessionToken(request, "range_0"); assertThat(sessionToken).isNotNull(); @@ -1088,14 +1148,18 @@ public void setSessionTokenGivesPriorityToOwnerFullNameOverResourceAddress() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>(ImmutableList.of(LocationEastUsEndpointToLocationPair.getLeft(), LocationEastUs2EndpointToLocationPair.getLeft(), LocationCentralUsEndpointToLocationPair.getLeft())); + UnmodifiableList endpoints = new UnmodifiableList<>( + ImmutableList.of( + new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); RxDocumentServiceRequest request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, collectionFullName1 + "/docs/doc1", ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String sessionTokenWithPkRangeIdForRequest = "range_0:1#100#4=90#5=1"; request.setResourceId(documentCollectionId); @@ -1105,13 +1169,13 @@ public void setSessionTokenGivesPriorityToOwnerFullNameOverResourceAddress() { HttpConstants.HttpHeaders.OWNER_FULL_NAME, collectionFullName2)); request = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(),OperationType.Read, collectionFullName1 + "/docs/doc1", ResourceType.Document); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); ISessionToken sessionToken = sessionContainer.resolvePartitionLocalSessionToken(request, "range_0"); assertThat(sessionToken).isNull(); request = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(),OperationType.Read, collectionFullName2 + "/docs/doc1", ResourceType.Document); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); sessionToken = sessionContainer.resolvePartitionLocalSessionToken(request, "range_0"); assertThat(sessionToken).isNotNull(); @@ -1131,14 +1195,18 @@ public void setSessionTokenIgnoresOwnerIdWhenRequestIsntNameBased() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>(ImmutableList.of(LocationEastUsEndpointToLocationPair.getLeft(), LocationEastUs2EndpointToLocationPair.getLeft(), LocationCentralUsEndpointToLocationPair.getLeft())); + UnmodifiableList endpoints = new UnmodifiableList<>( + ImmutableList.of( + new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); RxDocumentServiceRequest request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, collectionFullName + "/docs/doc1", ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); request.setResourceId(documentCollectionId1); @@ -1152,7 +1220,7 @@ public void setSessionTokenIgnoresOwnerIdWhenRequestIsntNameBased() { request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, documentCollectionId1, ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); ISessionToken sessionToken = sessionContainer.resolvePartitionLocalSessionToken(request, "range_0"); assertThat(sessionToken).isNotNull(); @@ -1161,7 +1229,7 @@ public void setSessionTokenIgnoresOwnerIdWhenRequestIsntNameBased() { request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, documentCollectionId2, ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); sessionToken = sessionContainer.resolvePartitionLocalSessionToken(request, "range_0"); assertThat(sessionToken).isNull(); @@ -1181,14 +1249,18 @@ public void setSessionTokenGivesPriorityToOwnerIdOverResourceIdWhenRequestIsName String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>(ImmutableList.of(LocationEastUsEndpointToLocationPair.getLeft(), LocationEastUs2EndpointToLocationPair.getLeft(), LocationCentralUsEndpointToLocationPair.getLeft())); + UnmodifiableList endpoints = new UnmodifiableList<>( + ImmutableList.of( + new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(LocationEastUsEndpointToLocationPair.getLeft()), Mockito.any())).thenReturn(regionContacted); RxDocumentServiceRequest request = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(),OperationType.Read, collectionFullName + "/docs/doc1", ResourceType.Document); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); request.setResourceId(documentCollectionId1); @@ -1201,14 +1273,14 @@ public void setSessionTokenGivesPriorityToOwnerIdOverResourceIdWhenRequestIsName request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, documentCollectionId1, ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); ISessionToken sessionToken = sessionContainer.resolvePartitionLocalSessionToken(request, "range_0"); assertThat(sessionToken).isNull(); request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, documentCollectionId2, ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); sessionToken = sessionContainer.resolvePartitionLocalSessionToken(request, "range_0"); assertThat(sessionToken.getLSN()).isEqualTo(100); @@ -1245,14 +1317,18 @@ public void setSessionTokenDoesntOverwriteHigherLSN() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>(ImmutableList.of(LocationEastUsEndpointToLocationPair.getLeft(), LocationEastUs2EndpointToLocationPair.getLeft(), LocationCentralUsEndpointToLocationPair.getLeft())); + UnmodifiableList endpoints = new UnmodifiableList<>( + ImmutableList.of( + new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); RxDocumentServiceRequest request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, collectionFullName + "/docs/doc1", ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String sessionTokenWithPkRangeIdForRequest = "range_0:1#100#4=90#5=1"; @@ -1262,7 +1338,7 @@ public void setSessionTokenDoesntOverwriteHigherLSN() { request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, collectionFullName + "/docs/doc1", ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String sessionTokenWithPkRangeIdToBeOverwritten = "range_0:1#105#4=90#5=1"; @@ -1271,7 +1347,7 @@ public void setSessionTokenDoesntOverwriteHigherLSN() { request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, documentCollectionId, ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); request.setResourceId(documentCollectionId); ISessionToken sessionToken = sessionContainer.resolvePartitionLocalSessionToken(request, "range_0"); @@ -1289,14 +1365,18 @@ public void setSessionTokenOverwriteLowerLSN() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>(ImmutableList.of(LocationEastUsEndpointToLocationPair.getLeft(), LocationEastUs2EndpointToLocationPair.getLeft(), LocationCentralUsEndpointToLocationPair.getLeft())); + UnmodifiableList endpoints = new UnmodifiableList<>( + ImmutableList.of( + new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); RxDocumentServiceRequest request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, collectionFullName + "/docs/doc1", ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String sessionTokenWithPkRangeIdForRequest = "range_0:1#105#4=90#5=1"; request.setResourceId(documentCollectionId); @@ -1305,7 +1385,7 @@ public void setSessionTokenOverwriteLowerLSN() { request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, collectionFullName + "/docs/doc1", ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String sessionTokenWithPkRangeIdToBeOverwritten = "range_0:1#100#4=90#5=1"; @@ -1315,7 +1395,7 @@ public void setSessionTokenOverwriteLowerLSN() { request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, documentCollectionId, ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); request.setResourceId(documentCollectionId); ISessionToken sessionToken = sessionContainer.resolvePartitionLocalSessionToken(request, "range_0"); @@ -1332,13 +1412,17 @@ public void setSessionTokenDoesNothingOnEmptySessionTokenHeader() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>(ImmutableList.of(LocationEastUsEndpointToLocationPair.getLeft(), LocationEastUs2EndpointToLocationPair.getLeft(), LocationCentralUsEndpointToLocationPair.getLeft())); + UnmodifiableList endpoints = new UnmodifiableList<>( + ImmutableList.of( + new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); RxDocumentServiceRequest docReadRequest1 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document); - docReadRequest1.requestContext.locationEndpointToRoute = locationEndpointContacted; + docReadRequest1.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String sessionTokenWithPkRangeIdForDocReadRequest1 = "range_0:1#100#4=90#5=1"; @@ -1346,7 +1430,7 @@ public void setSessionTokenDoesNothingOnEmptySessionTokenHeader() { ImmutableMap.of(HttpConstants.HttpHeaders.SESSION_TOKEN, sessionTokenWithPkRangeIdForDocReadRequest1)); RxDocumentServiceRequest request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, documentCollectionId, ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); String sessionToken = sessionContainer.resolveGlobalSessionToken(request); Set tokens = Sets.newSet(sessionToken.split(",")); @@ -1354,12 +1438,12 @@ public void setSessionTokenDoesNothingOnEmptySessionTokenHeader() { assertThat(tokens.contains(sessionTokenWithPkRangeIdForDocReadRequest1)).isTrue(); RxDocumentServiceRequest docReadRequest2 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document); - docReadRequest2.requestContext.locationEndpointToRoute = locationEndpointContacted; + docReadRequest2.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); sessionContainer.setSessionToken(docReadRequest2, documentCollectionId, collectionFullName, new HashMap<>()); request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, documentCollectionId, ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); sessionToken = sessionContainer.resolveGlobalSessionToken(request); tokens = Sets.newSet(sessionToken.split(",")); @@ -1424,7 +1508,11 @@ public void useParentSessionTokenAfterSplit() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>(ImmutableList.of(LocationEastUsEndpointToLocationPair.getLeft(), LocationEastUs2EndpointToLocationPair.getLeft(), LocationCentralUsEndpointToLocationPair.getLeft())); + UnmodifiableList endpoints = new UnmodifiableList<>( + ImmutableList.of( + new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); @@ -1435,7 +1523,7 @@ public void useParentSessionTokenAfterSplit() { String resultantParentSessionToken = parentPKRangeId + ":" + parentSession; RxDocumentServiceRequest request1 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document); - request1.requestContext.locationEndpointToRoute = locationEndpointContacted; + request1.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); sessionContainer.setSessionToken( request1, @@ -1447,7 +1535,7 @@ public void useParentSessionTokenAfterSplit() { String childPKRangeId = "1"; RxDocumentServiceRequest request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, documentCollectionId1, ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); request.requestContext.resolvedPartitionKeyRange = new PartitionKeyRange( @@ -1472,7 +1560,11 @@ public void useParentSessionTokenAfterMerge() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>(ImmutableList.of(LocationEastUsEndpointToLocationPair.getLeft(), LocationEastUs2EndpointToLocationPair.getLeft(), LocationCentralUsEndpointToLocationPair.getLeft())); + UnmodifiableList endpoints = new UnmodifiableList<>( + ImmutableList.of( + new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getLeft()), + new RegionalRoutingContext(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); @@ -1497,7 +1589,7 @@ public void useParentSessionTokenAfterMerge() { String parent1SessionToken = parent1PKRangeId + ":" + parent1Session; RxDocumentServiceRequest request1 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document); - request1.requestContext.locationEndpointToRoute = locationEndpointContacted; + request1.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); sessionContainer.setSessionToken( request1, @@ -1515,7 +1607,7 @@ public void useParentSessionTokenAfterMerge() { String parent2SessionToken = parent2PKRangeId + ":" + parent2Session; RxDocumentServiceRequest request2 = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document); - request2.requestContext.locationEndpointToRoute = locationEndpointContacted; + request2.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); sessionContainer.setSessionToken( request2, @@ -1527,7 +1619,7 @@ public void useParentSessionTokenAfterMerge() { String childPKRangeId = "2"; RxDocumentServiceRequest request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(),OperationType.Read, documentCollectionId1, ResourceType.Document, new HashMap<>()); - request.requestContext.locationEndpointToRoute = locationEndpointContacted; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointContacted); request.requestContext.resolvedPartitionKeyRange = new PartitionKeyRange( @@ -1567,9 +1659,14 @@ public void resolvePartitionLocalSessionToken( GlobalEndpointManager globalEndpointManagerMock = null; RegionScopedSessionContainer sessionContainer = null; - List writableURIs = writableURIToLocationMappings + List consolidatedWriteRegionalEndpointRoutingContexts = writableURIToLocationMappings .stream() - .map(uriToLocationMappings -> uriToLocationMappings.getLeft()) + .map(uriToLocationMappings -> new RegionalRoutingContext(uriToLocationMappings.getLeft())) + .collect(Collectors.toList()); + + List consolidatedReadRegionalEndpointRoutingContexts = readEndpoints + .stream() + .map(readEndpoint -> new RegionalRoutingContext(readEndpoint)) .collect(Collectors.toList()); DatabaseAccount databaseAccount = ModelBridgeUtils.createDatabaseAccount( @@ -1584,7 +1681,7 @@ public void resolvePartitionLocalSessionToken( .when(globalEndpointManagerMock.getLatestDatabaseAccount()) .thenReturn(databaseAccount); - UnmodifiableList readEndpointsInUnmodifiableList = new UnmodifiableList<>(readEndpoints); + UnmodifiableList readEndpointsInUnmodifiableList = new UnmodifiableList<>(consolidatedReadRegionalEndpointRoutingContexts); Mockito .when(globalEndpointManagerMock.getReadEndpoints()) @@ -1592,7 +1689,7 @@ public void resolvePartitionLocalSessionToken( Mockito .when(globalEndpointManagerMock.getApplicableWriteEndpoints(Mockito.anyList())) - .thenReturn(new UnmodifiableList<>(writableURIs)); + .thenReturn(new UnmodifiableList<>(consolidatedWriteRegionalEndpointRoutingContexts)); Mockito .when(globalEndpointManagerMock.canUseMultipleWriteLocations(Mockito.any())) @@ -1613,7 +1710,7 @@ public void resolvePartitionLocalSessionToken( RxDocumentServiceRequest request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document); // doesn't matter for a request for which the session token has to be resolved - request.requestContext.locationEndpointToRoute = LocationEastUsEndpointToLocationPair.getLeft(); + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getLeft()); request.setResourceId(documentCollectionId1); if (pkToBeUsedForSessionTokenResolution != null) { @@ -1664,7 +1761,7 @@ private static RequestMetadata constructRequestInstance( request.setResourceId(collectionResourceId); request.setPartitionKeyInternal(BridgeInternal.getPartitionKeyInternal(partitionKey)); request.setPartitionKeyDefinition(partitionKeyDefinition); - request.requestContext.locationEndpointToRoute = locationEndpointToRoute; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointToRoute); return new RequestMetadata(request, responseHeaders); } @@ -1676,7 +1773,7 @@ private static RxDocumentServiceRequest createRequestEntity(OperationType operat operationType, resourceType); - request.requestContext.locationEndpointToRoute = locationEndpointToRoute; + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointToRoute); return request; } diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RenameCollectionAwareClientRetryPolicyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RenameCollectionAwareClientRetryPolicyTest.java index 5a352678a334..b2a9feeb879f 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RenameCollectionAwareClientRetryPolicyTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RenameCollectionAwareClientRetryPolicyTest.java @@ -6,11 +6,15 @@ import com.azure.cosmos.implementation.caches.RxClientCollectionCache; import com.azure.cosmos.implementation.circuitBreaker.GlobalPartitionEndpointManagerForCircuitBreaker; import com.azure.cosmos.implementation.directconnectivity.WFConstants; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import io.netty.handler.timeout.ReadTimeoutException; import org.mockito.Mockito; import org.testng.annotations.Test; import reactor.core.publisher.Mono; +import java.net.URI; +import java.net.URISyntaxException; + import static com.azure.cosmos.implementation.ClientRetryPolicyTest.validateSuccess; import static com.azure.cosmos.implementation.TestUtils.mockDiagnosticsClientContext; import static org.assertj.core.api.Assertions.assertThat; @@ -51,10 +55,16 @@ public void onBeforeSendRequestNotInvoked() { } @Test(groups = "unit", timeOut = TIMEOUT) - public void shouldRetryWithNotFoundStatusCode() { + public void shouldRetryWithNotFoundStatusCode() throws URISyntaxException { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(eq(null), eq(false)); + + URI locationEndToRoute = new URI("https://location1.documents.com"); + RegionalRoutingContext consolidatedLocationEndpointToRoute = new RegionalRoutingContext(locationEndToRoute); + + Mockito.when(endpointManager.resolveServiceEndpoint(Mockito.any())).thenReturn(consolidatedLocationEndpointToRoute); + IRetryPolicyFactory retryPolicyFactory = new RetryPolicy(mockDiagnosticsClientContext(), endpointManager, ConnectionPolicy.getDefaultPolicy(), globalPartitionEndpointManager); RxClientCollectionCache rxClientCollectionCache = Mockito.mock(RxClientCollectionCache.class); @@ -78,11 +88,17 @@ public void shouldRetryWithNotFoundStatusCode() { } @Test(groups = "unit", timeOut = TIMEOUT) - public void shouldRetryWithNotFoundStatusCodeAndReadSessionNotAvailableSubStatusCode() { + public void shouldRetryWithNotFoundStatusCodeAndReadSessionNotAvailableSubStatusCode() throws URISyntaxException { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(eq(null), eq(false)); + + URI locationEndToRoute = new URI("https://location1.documents.com"); + RegionalRoutingContext consolidatedLocationEndpointToRoute = new RegionalRoutingContext(locationEndToRoute); + + Mockito.when(endpointManager.resolveServiceEndpoint(Mockito.any())).thenReturn(consolidatedLocationEndpointToRoute); + IRetryPolicyFactory retryPolicyFactory = new RetryPolicy(mockDiagnosticsClientContext(), endpointManager, ConnectionPolicy.getDefaultPolicy(), globalPartitionEndpointManager); RxClientCollectionCache rxClientCollectionCache = Mockito.mock(RxClientCollectionCache.class); @@ -117,10 +133,15 @@ public void shouldRetryWithNotFoundStatusCodeAndReadSessionNotAvailableSubStatus * No retry on bad request exception */ @Test(groups = "unit", timeOut = TIMEOUT) - public void shouldRetryWithGenericException() { + public void shouldRetryWithGenericException() throws URISyntaxException { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); + URI locationEndToRoute = new URI("https://location1.documents.com"); + RegionalRoutingContext consolidatedLocationEndpointToRoute = new RegionalRoutingContext(locationEndToRoute); + + Mockito.when(endpointManager.resolveServiceEndpoint(Mockito.any())).thenReturn(consolidatedLocationEndpointToRoute); + Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(eq(null), eq(false)); IRetryPolicyFactory retryPolicyFactory = new RetryPolicy(mockDiagnosticsClientContext(), endpointManager, ConnectionPolicy.getDefaultPolicy(), globalPartitionEndpointManager); RxClientCollectionCache rxClientCollectionCache = Mockito.mock(RxClientCollectionCache.class); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxGatewayStoreModelTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxGatewayStoreModelTest.java index bd26bb53a0b5..889f09b14314 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxGatewayStoreModelTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxGatewayStoreModelTest.java @@ -11,8 +11,8 @@ import com.azure.cosmos.implementation.http.HttpClient; import com.azure.cosmos.implementation.http.HttpHeaders; import com.azure.cosmos.implementation.http.HttpRequest; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import io.netty.channel.ConnectTimeoutException; -import io.netty.handler.codec.http.HttpMethod; import io.netty.handler.timeout.ReadTimeoutException; import io.reactivex.subscribers.TestSubscriber; import org.mockito.ArgumentCaptor; @@ -80,9 +80,10 @@ public void readTimeout() throws Exception { QueryCompatibilityMode queryCompatibilityMode = QueryCompatibilityMode.Default; UserAgentContainer userAgentContainer = new UserAgentContainer(); GlobalEndpointManager globalEndpointManager = Mockito.mock(GlobalEndpointManager.class); - GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new URI("https://localhost")) + RegionalRoutingContext regionalRoutingContext = new RegionalRoutingContext(new URI("https://localhost")); + + Mockito.doReturn(regionalRoutingContext) .when(globalEndpointManager).resolveServiceEndpoint(any()); HttpClient httpClient = Mockito.mock(HttpClient.class); Mockito.doReturn(Mono.error(ReadTimeoutException.INSTANCE)) @@ -105,7 +106,7 @@ public void readTimeout() throws Exception { OperationType.Read, "/dbs/db/colls/col/docs/docId", ResourceType.Document); dsr.getHeaders().put("key", "value"); dsr.requestContext = new DocumentServiceRequestContext(); - + dsr.requestContext.regionalRoutingContextToRoute = regionalRoutingContext; Mono resp = storeModel.processMessage(dsr); validateFailure(resp, FailureValidator.builder() @@ -125,7 +126,9 @@ public void serviceUnavailable() throws Exception { UserAgentContainer userAgentContainer = new UserAgentContainer(); GlobalEndpointManager globalEndpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new URI("https://localhost")) + + RegionalRoutingContext regionalRoutingContext = new RegionalRoutingContext(new URI("https://localhost")); + Mockito.doReturn(regionalRoutingContext) .when(globalEndpointManager).resolveServiceEndpoint(any()); HttpClient httpClient = Mockito.mock(HttpClient.class); Mockito.doReturn(Mono.error(new SocketException("Dummy SocketException"))) @@ -148,7 +151,7 @@ public void serviceUnavailable() throws Exception { OperationType.Read, "/dbs/db/colls/col/docs/docId", ResourceType.Document); dsr.getHeaders().put("key", "value"); dsr.requestContext = new DocumentServiceRequestContext(); - + dsr.requestContext.regionalRoutingContextToRoute = regionalRoutingContext; Mono resp = storeModel.processMessage(dsr); validateFailure(resp, FailureValidator.builder() @@ -177,9 +180,10 @@ public void applySessionToken( Mockito.doReturn(sdkGlobalSessionToken).when(sessionContainer).resolveGlobalSessionToken(any()); GlobalEndpointManager globalEndpointManager = Mockito.mock(GlobalEndpointManager.class); - GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); + URI locationEndpointToRoute = new URI("https://localhost"); + RegionalRoutingContext regionalRoutingContext = new RegionalRoutingContext(locationEndpointToRoute); - Mockito.doReturn(new URI("https://localhost")) + Mockito.doReturn(new RegionalRoutingContext(new URI("https://localhost"))) .when(globalEndpointManager).resolveServiceEndpoint(any()); HttpClient httpClient = Mockito.mock(HttpClient.class); @@ -207,6 +211,13 @@ public void applySessionToken( operationType, "/fakeResourceFullName", resourceType); + + if (resourceType != ResourceType.DatabaseAccount) { + dsr.requestContext.regionalRoutingContextToRoute = regionalRoutingContext; + } else { + dsr.setEndpointOverride(locationEndpointToRoute); + } + if (sessionTokenFromUser) { dsr.getHeaders().put(HttpConstants.HttpHeaders.SESSION_TOKEN, userControlledSessionToken); } @@ -248,9 +259,8 @@ public void validateApiType() throws Exception { Mockito.doReturn(sdkGlobalSessionToken).when(sessionContainer).resolveGlobalSessionToken(any()); GlobalEndpointManager globalEndpointManager = Mockito.mock(GlobalEndpointManager.class); - GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new URI("https://localhost")) + Mockito.doReturn(new RegionalRoutingContext(new URI("https://localhost"))) .when(globalEndpointManager).resolveServiceEndpoint(any()); HttpClient httpClient = Mockito.mock(HttpClient.class); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ThinClientStoreModelTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ThinClientStoreModelTest.java index 772be371b0c8..64cd7fe37115 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ThinClientStoreModelTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ThinClientStoreModelTest.java @@ -1,23 +1,15 @@ package com.azure.cosmos.implementation; import com.azure.cosmos.ConsistencyLevel; -import com.azure.cosmos.implementation.circuitBreaker.GlobalPartitionEndpointManagerForCircuitBreaker; -import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils; -import com.azure.cosmos.implementation.http.Http2ConnectionConfig; import com.azure.cosmos.implementation.http.HttpClient; -import com.azure.cosmos.implementation.http.HttpClientConfig; -import com.azure.cosmos.implementation.http.HttpHeaders; -import com.azure.cosmos.implementation.http.HttpRequest; -import com.azure.cosmos.implementation.http.ReactorNettyClient; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import io.netty.channel.ConnectTimeoutException; -import org.mockito.ArgumentCaptor; import org.mockito.Mockito; import org.testng.annotations.Test; import reactor.core.publisher.Mono; import java.net.URI; -import static org.assertj.core.api.Assertions.fail; import static org.mockito.ArgumentMatchers.any; public class ThinClientStoreModelTest { @@ -38,7 +30,7 @@ public void testThinClientStoreModel() throws Exception { GlobalEndpointManager globalEndpointManager = Mockito.mock(GlobalEndpointManager.class); - Mockito.doReturn(new URI("https://localhost")) + Mockito.doReturn(new RegionalRoutingContext(new URI("https://localhost:8080"))) .when(globalEndpointManager).resolveServiceEndpoint(any()); // mocking with HTTP/1.1 client, just using this test as basic store model validation. e2e request flow diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/WebExceptionRetryPolicyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/WebExceptionRetryPolicyTest.java index 0a72963fe795..41db724bf295 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/WebExceptionRetryPolicyTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/WebExceptionRetryPolicyTest.java @@ -5,6 +5,7 @@ import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.CosmosException; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import io.netty.handler.timeout.ReadTimeoutException; import io.reactivex.subscribers.TestSubscriber; import org.mockito.Mockito; @@ -37,7 +38,10 @@ public static Object[][] operationTypeProvider() { @Test(groups = {"unit"}) public void shouldRetryOnTimeoutForReadOperations() throws Exception { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); - Mockito.doReturn(new URI("http://localhost:")).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); + + RegionalRoutingContext regionalRoutingContext = new RegionalRoutingContext(new URI("http://localhost:")); + + Mockito.doReturn(regionalRoutingContext).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(Mockito.eq(null), Mockito.eq(true)); RetryContext retryContext = new RetryContext(); @@ -52,6 +56,7 @@ public void shouldRetryOnTimeoutForReadOperations() throws Exception { //Default HttpTimeout Policy dsr = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(), OperationType.Read, "/dbs/db/colls/col/docs/doc", ResourceType.Document); + dsr.requestContext.regionalRoutingContextToRoute = regionalRoutingContext; // 1st Attempt webExceptionRetryPolicy.onBeforeSendRequest(dsr); @@ -89,7 +94,10 @@ public void shouldRetryOnTimeoutForReadOperations() throws Exception { @Test(groups = {"unit"}) public void shouldRetryOnTimeoutForMetaDataReadOperations() throws Exception { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); - Mockito.doReturn(new URI("http://localhost:")).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); + + RegionalRoutingContext regionalRoutingContext = new RegionalRoutingContext(new URI("http://localhost:")); + + Mockito.doReturn(regionalRoutingContext).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(Mockito.eq(null), Mockito.eq(true)); RetryContext retryContext = new RetryContext(); @@ -104,6 +112,7 @@ public void shouldRetryOnTimeoutForMetaDataReadOperations() throws Exception { //Default HttpTimeout Policy dsr = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(), OperationType.Read, "/dbs/db", ResourceType.DatabaseAccount); + dsr.requestContext.regionalRoutingContextToRoute = regionalRoutingContext; // 1st Attempt webExceptionRetryPolicy.onBeforeSendRequest(dsr); @@ -143,7 +152,10 @@ public void shouldRetryOnTimeoutForMetaDataReadOperations() throws Exception { @Test(groups = {"unit"}) public void shouldRetryOnTimeoutForQueryPlanOperations() throws Exception { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); - Mockito.doReturn(new URI("http://localhost:")).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); + + RegionalRoutingContext regionalRoutingContext = new RegionalRoutingContext(new URI("http://localhost:")); + + Mockito.doReturn(regionalRoutingContext).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(Mockito.eq(null), Mockito.eq(true)); RetryContext retryContext = new RetryContext(); @@ -158,6 +170,7 @@ public void shouldRetryOnTimeoutForQueryPlanOperations() throws Exception { //Default HttpTimeout Policy dsr = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(), OperationType.QueryPlan, "/dbs/db/colls/col/docs/doc", ResourceType.Document); + dsr.requestContext.regionalRoutingContextToRoute = regionalRoutingContext; // 1st Attempt webExceptionRetryPolicy.onBeforeSendRequest(dsr); @@ -195,7 +208,10 @@ public void shouldRetryOnTimeoutForQueryPlanOperations() throws Exception { @Test(groups = "unit") public void shouldNotRetryOnTimeoutForWriteOperations() throws Exception { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); - Mockito.doReturn(new URI("http://localhost:")).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); + + RegionalRoutingContext regionalRoutingContext = new RegionalRoutingContext(new URI("http://localhost:")); + + Mockito.doReturn(regionalRoutingContext).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(Mockito.eq(null), Mockito.eq(true)); @@ -209,6 +225,7 @@ public void shouldNotRetryOnTimeoutForWriteOperations() throws Exception { //Data Plane Write - Should not retry dsr = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(), OperationType.Create, "/dbs/db/colls/col/docs/doc", ResourceType.Document); + dsr.requestContext.regionalRoutingContextToRoute = regionalRoutingContext; WebExceptionRetryPolicy webExceptionRetryPolicy = new WebExceptionRetryPolicy(new RetryContext()); webExceptionRetryPolicy.onBeforeSendRequest(dsr); @@ -222,6 +239,7 @@ public void shouldNotRetryOnTimeoutForWriteOperations() throws Exception { //Metadata Write - Should not Retry dsr = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(), OperationType.Create, "/dbs/db", ResourceType.DatabaseAccount); + dsr.requestContext.regionalRoutingContextToRoute = regionalRoutingContext; webExceptionRetryPolicy = new WebExceptionRetryPolicy(new RetryContext()); webExceptionRetryPolicy.onBeforeSendRequest(dsr); @@ -236,7 +254,10 @@ public void shouldNotRetryOnTimeoutForWriteOperations() throws Exception { @Test(groups = "unit", dataProvider = "operationTypeProvider") public void httpNetworkFailureOnAddressRefresh(OperationType operationType) throws Exception { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); - Mockito.doReturn(new URI("http://localhost:")).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); + + RegionalRoutingContext regionalRoutingContext = new RegionalRoutingContext(new URI("http://localhost:")); + + Mockito.doReturn(regionalRoutingContext).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(Mockito.eq(null), Mockito.eq(false)); Mockito.doReturn(2).when(endpointManager).getPreferredLocationCount(); @@ -250,6 +271,7 @@ public void httpNetworkFailureOnAddressRefresh(OperationType operationType) thro operationType, "/dbs/db/colls/col/docs/", ResourceType.Document); dsr.setAddressRefresh(true, false); dsr.requestContext = new DocumentServiceRequestContext(); + dsr.requestContext.regionalRoutingContextToRoute = regionalRoutingContext; // 1st Attempt webExceptionRetryPolicy.onBeforeSendRequest(dsr); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/ConnectionStateListenerTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/ConnectionStateListenerTest.java index d41b688ded2e..97f86bb795d3 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/ConnectionStateListenerTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/ConnectionStateListenerTest.java @@ -26,6 +26,7 @@ import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdEndpoint; import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdRequestManager; import com.azure.cosmos.implementation.routing.PartitionKeyRangeIdentity; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import io.netty.handler.ssl.SslContext; import org.mockito.Mockito; import org.slf4j.Logger; @@ -34,6 +35,8 @@ import org.testng.annotations.Test; import java.io.IOException; +import java.net.URI; +import java.net.URISyntaxException; import java.nio.channels.ClosedChannelException; import java.util.ArrayList; import java.util.Arrays; @@ -85,7 +88,7 @@ public void connectionStateListener_OnConnectionEvent( boolean isTcpConnectionEndpointRediscoveryEnabled, RequestResponseType responseType, boolean markUnhealthy, - boolean markUnhealthyWhenServerShutdown) throws ExecutionException, InterruptedException { + boolean markUnhealthyWhenServerShutdown) throws ExecutionException, InterruptedException, URISyntaxException { ConnectionPolicy connectionPolicy = new ConnectionPolicy(DirectConnectionConfig.getDefaultConfig()); connectionPolicy.setTcpConnectionEndpointRediscoveryEnabled(isTcpConnectionEndpointRediscoveryEnabled); @@ -124,6 +127,8 @@ public void connectionStateListener_OnConnectionEvent( RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Create, ResourceType.Document, "dbs/fakedb/colls/fakeColls", getDocumentDefinition(), new HashMap<>()); + req.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(new URI("https://localhost:8080")); + req.setPartitionKeyRangeIdentity(new PartitionKeyRangeIdentity("fakeCollectionId","fakePartitionKeyRangeId")); // Validate connectionStateListener will always track the latest Uri diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/GlobalAddressResolverTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/GlobalAddressResolverTest.java index 73c4110cea1c..331be53cc7af 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/GlobalAddressResolverTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/GlobalAddressResolverTest.java @@ -11,7 +11,6 @@ import com.azure.cosmos.implementation.ConnectionPolicy; import com.azure.cosmos.implementation.DocumentCollection; import com.azure.cosmos.implementation.GlobalEndpointManager; -import com.azure.cosmos.implementation.circuitBreaker.GlobalPartitionEndpointManagerForCircuitBreaker; import com.azure.cosmos.implementation.IAuthorizationTokenProvider; import com.azure.cosmos.implementation.OpenConnectionResponse; import com.azure.cosmos.implementation.OperationType; @@ -27,6 +26,7 @@ import com.azure.cosmos.implementation.http.HttpClient; import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper; import com.azure.cosmos.implementation.routing.PartitionKeyRangeIdentity; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import com.azure.cosmos.models.CosmosContainerIdentity; import org.mockito.ArgumentMatchers; import org.mockito.Mockito; @@ -78,17 +78,17 @@ public void before_GlobalAddressResolverTest() throws Exception { httpClient = Mockito.mock(HttpClient.class); endpointManager = Mockito.mock(GlobalEndpointManager.class); - List readEndPointList = new ArrayList<>(); - readEndPointList.add(urlforRead1); - readEndPointList.add(urlforRead2); - readEndPointList.add(urlforRead3); - UnmodifiableList readList = new UnmodifiableList<>(readEndPointList); + List readEndPointList = new ArrayList<>(); + readEndPointList.add(new RegionalRoutingContext(urlforRead1)); + readEndPointList.add(new RegionalRoutingContext(urlforRead2)); + readEndPointList.add(new RegionalRoutingContext(urlforRead3)); + UnmodifiableList readList = new UnmodifiableList<>(readEndPointList); - List writeEndPointList = new ArrayList<>(); - writeEndPointList.add(urlforWrite1); - writeEndPointList.add(urlforWrite2); - writeEndPointList.add(urlforWrite3); - UnmodifiableList writeList = new UnmodifiableList<>(writeEndPointList); + List writeEndPointList = new ArrayList<>(); + writeEndPointList.add(new RegionalRoutingContext(urlforWrite1)); + writeEndPointList.add(new RegionalRoutingContext(urlforWrite2)); + writeEndPointList.add(new RegionalRoutingContext(urlforWrite3)); + UnmodifiableList writeList = new UnmodifiableList<>(writeEndPointList); Mockito.when(endpointManager.getReadEndpoints()).thenReturn(readList); Mockito.when(endpointManager.getWriteEndpoints()).thenReturn(writeList); @@ -122,13 +122,13 @@ public void resolveAsync() throws Exception { assertThat(urlsBeforeResolve.contains(urlforRead3)).isFalse();//Last read will be removed from addressCacheByEndpoint after 5 endpoints assertThat(urlsBeforeResolve.contains(urlforRead2)).isTrue(); - URI testUrl = new URI("http://Test.com/"); + RegionalRoutingContext testUrl = new RegionalRoutingContext(new URI("http://Test.com/")); Mockito.when(endpointManager.resolveServiceEndpoint(ArgumentMatchers.any())).thenReturn(testUrl); globalAddressResolver.resolveAsync(request, true); Set urlsAfterResolve = globalAddressResolver.addressCacheByEndpoint.keySet(); assertThat(urlsAfterResolve.size()).isEqualTo(5); assertThat(urlsAfterResolve.contains(urlforRead2)).isFalse();//Last read will be removed from addressCacheByEndpoint after 5 endpoints - assertThat(urlsBeforeResolve.contains(testUrl)).isTrue();//New endpoint will be added in addressCacheByEndpoint + assertThat(urlsBeforeResolve.contains(testUrl.getGatewayRegionalEndpoint())).isTrue();//New endpoint will be added in addressCacheByEndpoint } @Test(groups = "unit") @@ -156,8 +156,9 @@ public void submitOpenConnectionTasksAndInitCaches() { AddressInformation addressInformation = new AddressInformation(true, true, "https://be1.west-us.com:8080", Protocol.TCP); Mockito - .when(endpointManager.getReadEndpoints()) - .thenReturn(new UnmodifiableList(Arrays.asList(urlforRead1, urlforRead2))); + .when(endpointManager.getReadEndpoints()) + .thenReturn(new UnmodifiableList<>( + Arrays.asList(new RegionalRoutingContext(urlforRead1), new RegionalRoutingContext(urlforRead2)))); DocumentCollection documentCollection = new DocumentCollection(); documentCollection.setId("TestColl"); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/GlobalEndPointManagerTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/GlobalEndPointManagerTest.java index a91163ae264c..734a3695a2c9 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/GlobalEndPointManagerTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/GlobalEndPointManagerTest.java @@ -11,6 +11,7 @@ import com.azure.cosmos.implementation.GlobalEndpointManager; import com.azure.cosmos.implementation.LifeCycleUtils; import com.azure.cosmos.implementation.routing.LocationCache; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import org.mockito.ArgumentMatchers; import org.mockito.Mockito; import org.testng.Assert; @@ -99,7 +100,7 @@ public void refreshLocationAsyncForConnectivityIssue() throws Exception { LocationCache locationCache = this.getLocationCache(globalEndPointManager); Assert.assertEquals(locationCache.getReadEndpoints().size(), 2, "Read endpoints should have 2 values"); - Map availableReadEndpointByLocation = this.getAvailableReadEndpointByLocation(locationCache); + Map availableReadEndpointByLocation = this.getAvailableReadEndpointByLocation(locationCache); Assert.assertEquals(availableReadEndpointByLocation.size(), 2); Assert.assertTrue(availableReadEndpointByLocation.keySet().contains("East Asia")); @@ -155,7 +156,7 @@ public void refreshLocationAsyncForConnectivityIssueWithPreferredRegions() throw locationCache = this.getLocationCache(globalEndPointManager); Assert.assertEquals(locationCache.getReadEndpoints().size(), 2); //Cache will not refresh immediately, other preferred region East Asia is still active - Map availableReadEndpointByLocation = this.getAvailableReadEndpointByLocation(locationCache); + Map availableReadEndpointByLocation = this.getAvailableReadEndpointByLocation(locationCache); Assert.assertEquals(availableReadEndpointByLocation.size(), 2); Assert.assertTrue(availableReadEndpointByLocation.keySet().iterator().next().equalsIgnoreCase("East Asia")); @@ -195,7 +196,7 @@ public void refreshLocationAsyncForWriteForbidden() throws Exception { LocationCache locationCache = this.getLocationCache(globalEndPointManager); Assert.assertEquals(locationCache.getReadEndpoints().size(), 1); - Map availableWriteEndpointByLocation = this.getAvailableWriteEndpointByLocation(locationCache); + Map availableWriteEndpointByLocation = this.getAvailableWriteEndpointByLocation(locationCache); Assert.assertTrue(availableWriteEndpointByLocation.keySet().contains("East Asia")); AtomicBoolean isRefreshing = getIsRefreshing(globalEndPointManager); @@ -261,7 +262,7 @@ public void startRefreshLocationTimerAsync() throws Exception { LocationCache locationCache = this.getLocationCache(globalEndPointManager); Assert.assertEquals(locationCache.getReadEndpoints().size(), 1); - Map availableReadEndpointByLocation = this.getAvailableReadEndpointByLocation(locationCache); + Map availableReadEndpointByLocation = this.getAvailableReadEndpointByLocation(locationCache); Assert.assertEquals(availableReadEndpointByLocation.size(), 1); Assert.assertTrue(availableReadEndpointByLocation.keySet().iterator().next().equalsIgnoreCase("East Asia")); @@ -290,33 +291,33 @@ private LocationCache getLocationCache(GlobalEndpointManager globalEndPointManag return locationCache; } - private Map getAvailableWriteEndpointByLocation(LocationCache locationCache) throws Exception { + private Map getAvailableWriteEndpointByLocation(LocationCache locationCache) throws Exception { Field locationInfoField = LocationCache.class.getDeclaredField("locationInfo"); locationInfoField.setAccessible(true); Object locationInfo = locationInfoField.get(locationCache); Class DatabaseAccountLocationsInfoClass = Class.forName("com.azure.cosmos.implementation.routing.LocationCache$DatabaseAccountLocationsInfo"); - Field availableWriteEndpointByLocationField = DatabaseAccountLocationsInfoClass.getDeclaredField("availableWriteEndpointByLocation"); + Field availableWriteEndpointByLocationField = DatabaseAccountLocationsInfoClass.getDeclaredField("availableWriteEndpointsByLocation"); availableWriteEndpointByLocationField.setAccessible(true); - Field availableReadEndpointByLocationField = DatabaseAccountLocationsInfoClass.getDeclaredField("availableReadEndpointByLocation"); + Field availableReadEndpointByLocationField = DatabaseAccountLocationsInfoClass.getDeclaredField("availableReadEndpointsByLocation"); availableReadEndpointByLocationField.setAccessible(true); @SuppressWarnings("unchecked") - Map map = (Map) availableWriteEndpointByLocationField.get(locationInfo); + Map map = (Map) availableWriteEndpointByLocationField.get(locationInfo); return map; } - private Map getAvailableReadEndpointByLocation(LocationCache locationCache) throws Exception { + private Map getAvailableReadEndpointByLocation(LocationCache locationCache) throws Exception { Field locationInfoField = LocationCache.class.getDeclaredField("locationInfo"); locationInfoField.setAccessible(true); Object locationInfo = locationInfoField.get(locationCache); Class DatabaseAccountLocationsInfoClass = Class.forName("com.azure.cosmos.implementation.routing.LocationCache$DatabaseAccountLocationsInfo"); - Field availableReadEndpointByLocationField = DatabaseAccountLocationsInfoClass.getDeclaredField("availableReadEndpointByLocation"); + Field availableReadEndpointByLocationField = DatabaseAccountLocationsInfoClass.getDeclaredField("availableReadEndpointsByLocation"); availableReadEndpointByLocationField.setAccessible(true); @SuppressWarnings("unchecked") - Map map = (Map) availableReadEndpointByLocationField.get(locationInfo); + Map map = (Map) availableReadEndpointByLocationField.get(locationInfo); return map; } @@ -353,8 +354,8 @@ private GlobalEndpointManager getGlobalEndPointManager() throws Exception { LocationCache locationCache = getLocationCache(globalEndPointManager); Assert.assertEquals(locationCache.getReadEndpoints().size(), 2, "Read endpoints should have 2 values"); - Map availableWriteEndpointByLocation = this.getAvailableWriteEndpointByLocation(locationCache); - Map availableReadEndpointByLocation = this.getAvailableReadEndpointByLocation(locationCache); + Map availableWriteEndpointByLocation = this.getAvailableWriteEndpointByLocation(locationCache); + Map availableReadEndpointByLocation = this.getAvailableReadEndpointByLocation(locationCache); Assert.assertEquals(availableWriteEndpointByLocation.size(), 1); Assert.assertEquals(availableReadEndpointByLocation.size(), 2); Assert.assertTrue(availableWriteEndpointByLocation.keySet().contains("East US")); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/MetadataRequestRetryPolicyTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/MetadataRequestRetryPolicyTests.java index f747de6ffa54..13b8cb738289 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/MetadataRequestRetryPolicyTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/MetadataRequestRetryPolicyTests.java @@ -32,6 +32,7 @@ import com.azure.cosmos.implementation.http.HttpClient; import com.azure.cosmos.implementation.http.HttpClientConfig; import com.azure.cosmos.implementation.http.HttpTimeoutPolicyControlPlaneHotPath; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import com.azure.cosmos.implementation.throughputControl.TestItem; import com.azure.cosmos.models.CosmosItemRequestOptions; import com.azure.cosmos.models.CosmosPatchOperations; @@ -234,14 +235,14 @@ public void forceBackgroundAddressRefresh_onConnectionTimeoutAndRequestCancellat GlobalAddressResolver globalAddressResolver = ReflectionUtils.getGlobalAddressResolver(asyncDocumentClient); GlobalEndpointManager globalEndpointManager = ReflectionUtils.getGlobalEndpointManager(asyncDocumentClient); - List readEndpoints = globalEndpointManager.getReadEndpoints(); + List readEndpoints = globalEndpointManager.getReadEndpoints(); Map endpointCacheByURIMap = globalAddressResolver.addressCacheByEndpoint; Map httpClientWrapperByRegionMap = new ConcurrentHashMap<>(); for (int i = 0; i < preferredRegions.size(); i++) { - URI readEndpoint = readEndpoints.get(i); + URI readEndpoint = readEndpoints.get(i).getGatewayRegionalEndpoint(); GlobalAddressResolver.EndpointCache endpointCache = endpointCacheByURIMap.get(readEndpoint); GatewayAddressCache gatewayAddressCache = endpointCache.addressCache; HttpClientUnderTestWrapper httpClientUnderTestWrapper = getHttpClientUnderTestWrapper(configs); @@ -347,7 +348,7 @@ public void metadataRetryPolicyTest( .withException(cosmosException) .build()); - int desiredInvocationCount = request.requestContext.locationEndpointToRoute == null ? 0 : 1; + int desiredInvocationCount = request.requestContext.regionalRoutingContextToRoute == null ? 0 : 1; if (request.isReadOnlyRequest()) { Mockito @@ -553,9 +554,13 @@ private static RxDocumentServiceRequest createRequest( assert request.requestContext != null; + URI locationEndpointToRoute + = URI.create("https://account-name-some-region.documents.azure.com:443"); + if (hasLocationEndpointToRoute) { - request.requestContext.locationEndpointToRoute - = URI.create("https://account-name-some-region.documents.azure.com:443"); + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointToRoute); + } else { + request.setEndpointOverride(locationEndpointToRoute); } if (isAddressRefresh) { diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/RntbdTransportClientTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/RntbdTransportClientTest.java index 7d25deb33662..97c36e75cece 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/RntbdTransportClientTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/RntbdTransportClientTest.java @@ -61,6 +61,7 @@ import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdUUID; import com.azure.cosmos.implementation.guava25.base.Strings; import com.azure.cosmos.implementation.guava25.collect.ImmutableMap; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import io.micrometer.core.instrument.Tag; import io.netty.buffer.ByteBuf; import io.netty.buffer.Unpooled; @@ -839,7 +840,9 @@ public void cancelRequestMono() throws InterruptedException, URISyntaxException, ConnectionPolicy connectionPolicy = new ConnectionPolicy(DirectConnectionConfig.getDefaultConfig()); RntbdTransportClient.Options options = new RntbdTransportClient.Options.Builder(connectionPolicy).build(); final SslContext sslContext = SslContextBuilder.forClient().build(); - request.requestContext.locationEndpointToRoute = locationToRoute; + + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationToRoute); + RntbdRequestArgs requestArgs = new RntbdRequestArgs(request, addressUri); RntbdRequestTimer requestTimer = new RntbdRequestTimer(5000, 5000); RntbdRequestRecord rntbdRequestRecord = new AsyncRntbdRequestRecord(requestArgs, requestTimer); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/query/DocumentProducerTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/query/DocumentProducerTest.java index 136e6df38af8..2bc5679eb6e4 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/query/DocumentProducerTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/query/DocumentProducerTest.java @@ -29,6 +29,7 @@ import com.azure.cosmos.implementation.query.orderbyquery.OrderbyRowComparer; import com.azure.cosmos.implementation.routing.PartitionKeyRangeIdentity; import com.azure.cosmos.implementation.routing.Range; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import com.azure.cosmos.models.FeedResponse; import com.fasterxml.jackson.databind.node.ObjectNode; import io.reactivex.subscribers.TestSubscriber; @@ -111,9 +112,9 @@ public Object[][] mergeParamProvider() { } private IRetryPolicyFactory mockDocumentClientIRetryPolicyFactory() { - URI url; + RegionalRoutingContext regionalRoutingContext; try { - url = new URI("http://localhost"); + regionalRoutingContext = new RegionalRoutingContext(new URI("http://localhost")); } catch (Exception e) { throw new IllegalStateException(e); } @@ -121,7 +122,7 @@ private IRetryPolicyFactory mockDocumentClientIRetryPolicyFactory() { GlobalEndpointManager globalEndpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(url).when(globalEndpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); + Mockito.doReturn(regionalRoutingContext).when(globalEndpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); doReturn(false).when(globalEndpointManager).isClosed(); return new RetryPolicy(mockDiagnosticsClientContext(), globalEndpointManager, ConnectionPolicy.getDefaultPolicy(), globalPartitionEndpointManager); } diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/routing/LocationCacheTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/routing/LocationCacheTest.java index 098c2ed24825..eec53a04e1f0 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/routing/LocationCacheTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/routing/LocationCacheTest.java @@ -35,7 +35,6 @@ import java.util.HashSet; import java.util.Iterator; import java.util.List; -import java.util.Locale; import java.util.Map; import java.util.Set; import java.util.concurrent.TimeUnit; @@ -480,9 +479,9 @@ public void validateAsync(boolean useMultipleWriteEndpoints, @Test(groups = "long") public void validateWriteEndpointOrderWithClientSideDisableMultipleWriteLocation() throws Exception { this.initialize(false, true, false); - assertThat(this.cache.getWriteEndpoints().get(0)).isEqualTo(LocationCacheTest.Location1Endpoint); - assertThat(this.cache.getWriteEndpoints().get(1)).isEqualTo(LocationCacheTest.Location2Endpoint); - assertThat(this.cache.getWriteEndpoints().get(2)).isEqualTo(LocationCacheTest.Location3Endpoint); + assertThat(this.cache.getWriteEndpoints().get(0)).isEqualTo(new RegionalRoutingContext(LocationCacheTest.Location1Endpoint)); + assertThat(this.cache.getWriteEndpoints().get(1)).isEqualTo(new RegionalRoutingContext(LocationCacheTest.Location2Endpoint)); + assertThat(this.cache.getWriteEndpoints().get(2)).isEqualTo(new RegionalRoutingContext(LocationCacheTest.Location3Endpoint)); } @Test(groups = "unit", dataProvider = "excludedRegionsTestConfigs") @@ -512,11 +511,11 @@ public void validateExcludedRegions( request.requestContext.setExcludeRegions(excludedRegionsOnRequest); if (request.isReadOnlyRequest()) { - List applicableReadEndpoints = cache.getApplicableReadEndpoints(request); + List applicableReadEndpoints = cache.getApplicableReadEndpoints(request); assertThat(applicableReadEndpoints.size()).isEqualTo(expectedApplicableEndpoints.size()); expectedApplicableEndpoints.forEach(endpoint -> assertThat(expectedApplicableEndpoints.contains(endpoint)).isTrue()); } else { - List applicableWriteEndpoints = cache.getApplicableWriteEndpoints(request); + List applicableWriteEndpoints = cache.getApplicableWriteEndpoints(request); assertThat(applicableWriteEndpoints.size()).isEqualTo(expectedApplicableEndpoints.size()); expectedApplicableEndpoints.forEach(endpoint -> assertThat(expectedApplicableEndpoints.contains(endpoint)).isTrue()); } @@ -532,8 +531,8 @@ public void validateEffectivePreferredRegions( boolean isDefaultEndpointAlsoRegionalEndpoint) { this.initialize(true, true, isPreferredLocationsListEmpty, isDefaultEndpointAlsoRegionalEndpoint); - List applicableReadEndpoints = cache.getApplicableReadEndpoints(request); - List applicableWriteEndpoints = cache.getApplicableWriteEndpoints(request); + List applicableReadEndpoints = cache.getApplicableReadEndpoints(request); + List applicableWriteEndpoints = cache.getApplicableWriteEndpoints(request); if (request.isReadOnlyRequest()) { assertThat(applicableReadEndpoints.size()).isEqualTo(expectedApplicableReadEndpoints.size()); @@ -667,8 +666,8 @@ private void validateLocationCacheAsync( endpointDiscoveryEnabled, isPreferredListEmpty); - UnmodifiableList currentWriteEndpoints = this.cache.getWriteEndpoints(); - UnmodifiableList currentReadEndpoints = this.cache.getReadEndpoints(); + UnmodifiableList currentWriteEndpoints = this.cache.getWriteEndpoints(); + UnmodifiableList currentReadEndpoints = this.cache.getReadEndpoints(); for (int i = 0; i < readLocationIndex; i++) { this.cache.markEndpointUnavailableForRead(createUrl(Iterables.get(this.databaseAccount.getReadableLocations(), i).getEndpoint())); this.endpointManager.markEndpointUnavailableForRead(createUrl(Iterables.get(this.databaseAccount.getReadableLocations(), i).getEndpoint()));; @@ -775,7 +774,7 @@ private void validateEndpointRefresh( boolean isMostPreferredLocationUnavailableForRead = isFirstReadEndpointUnavailable; boolean isMostPreferredLocationUnavailableForWrite = useMultipleWriteLocations ? false : isFirstWriteEndpointUnavailable; - if (this.preferredLocations.size() > 0 || isPreferredListEmpty) { + if (!this.preferredLocations.isEmpty() || isPreferredListEmpty) { String mostPreferredReadLocationName = (isPreferredListEmpty && endpointDiscoveryEnabled) ? preferredAvailableReadRegionsOrAccountLevelReadEndpoints.get(0) : this.preferredLocations.stream() .filter(location -> toStream(databaseAccount.getReadableLocations()) @@ -898,28 +897,28 @@ private void validateRequestEndpointResolution( // If current write endpoint is unavailable, write endpoints order doesn't change // ALL write requests flip-flop between current write and alternate write endpoint - UnmodifiableList writeEndpoints = this.cache.getWriteEndpoints(); + UnmodifiableList writeEndpoints = this.cache.getWriteEndpoints(); - assertThat(firstAvailableWriteEndpoint).isEqualTo(writeEndpoints.get(0)); - assertThat(secondAvailableWriteEndpoint).isEqualTo(this.resolveEndpointForWriteRequest(ResourceType.Document, true)); - assertThat(firstAvailableWriteEndpoint).isEqualTo(this.resolveEndpointForWriteRequest(ResourceType.Document, false)); + assertThat(new RegionalRoutingContext(firstAvailableWriteEndpoint)).isEqualTo(writeEndpoints.get(0)); + assertThat(new RegionalRoutingContext(secondAvailableWriteEndpoint)).isEqualTo(this.resolveEndpointForWriteRequest(ResourceType.Document, true)); + assertThat(new RegionalRoutingContext(firstAvailableWriteEndpoint)).isEqualTo(this.resolveEndpointForWriteRequest(ResourceType.Document, false)); // Writes to other resource types should be directed to first/second write getEndpoint - assertThat(firstWriteEnpoint).isEqualTo(this.resolveEndpointForWriteRequest(ResourceType.Database, false)); - assertThat(secondWriteEnpoint).isEqualTo(this.resolveEndpointForWriteRequest(ResourceType.Database, true)); + assertThat(new RegionalRoutingContext(firstWriteEnpoint)).isEqualTo(this.resolveEndpointForWriteRequest(ResourceType.Database, false)); + assertThat(new RegionalRoutingContext(secondWriteEnpoint)).isEqualTo(this.resolveEndpointForWriteRequest(ResourceType.Database, true)); // Reads should be directed to available read endpoints regardless of resource type - assertThat(firstAvailableReadEndpoint).isEqualTo(this.resolveEndpointForReadRequest(true)); - assertThat(firstAvailableReadEndpoint).isEqualTo(this.resolveEndpointForReadRequest(false)); + assertThat(new RegionalRoutingContext(firstAvailableReadEndpoint)).isEqualTo(this.resolveEndpointForReadRequest(true)); + assertThat(new RegionalRoutingContext(firstAvailableReadEndpoint)).isEqualTo(this.resolveEndpointForReadRequest(false)); } - private URI resolveEndpointForReadRequest(boolean masterResourceType) { + private RegionalRoutingContext resolveEndpointForReadRequest(boolean masterResourceType) { RxDocumentServiceRequest request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, masterResourceType ? ResourceType.Database : ResourceType.Document); return this.cache.resolveServiceEndpoint(request); } - private URI resolveEndpointForWriteRequest(ResourceType resourceType, boolean useAlternateWriteEndpoint) { + private RegionalRoutingContext resolveEndpointForWriteRequest(ResourceType resourceType, boolean useAlternateWriteEndpoint) { RxDocumentServiceRequest request = RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Create, resourceType); request.requestContext.routeToLocation(useAlternateWriteEndpoint ? 1 : 0, resourceType.isCollectionChild()); return this.cache.resolveServiceEndpoint(request); diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index f1eb701c7a67..db9bf25dd06e 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -2,7 +2,7 @@ ### 4.67.0 (2025-02-20) -#### Bugs Fixed +#### Other Changes * Block `ChangeFeedProcessor` from starting by throwing an `IllegalStateException` when the lease container contains leases with the same lease prefix but different `ChangeFeedMode` - [PR 43798](https://github.com/Azure/azure-sdk-for-java/pull/43798). ### 4.66.1 (2025-02-08) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosDiagnostics.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosDiagnostics.java index 79d11aafbc3a..4642652800f3 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosDiagnostics.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosDiagnostics.java @@ -10,6 +10,7 @@ import com.azure.cosmos.implementation.RxDocumentServiceRequest; import com.azure.cosmos.implementation.SerializationDiagnosticsContext; import com.azure.cosmos.implementation.guava25.collect.ImmutableList; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import com.azure.cosmos.util.Beta; import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.core.JsonProcessingException; @@ -330,7 +331,7 @@ String getFirstContactedRegion() { return this.clientSideRequestStatistics.getFirstContactedRegion(); } - URI getFirstContactedLocationEndpoint() { + RegionalRoutingContext getFirstContactedLocationEndpoint() { return this.clientSideRequestStatistics.getFirstContactedLocationEndpoint(); } @@ -478,7 +479,7 @@ public void setDiagnosticsContext(CosmosDiagnostics cosmosDiagnostics, CosmosDia } @Override - public URI getFirstContactedLocationEndpoint(CosmosDiagnostics cosmosDiagnostics) { + public RegionalRoutingContext getFirstContactedLocationEndpoint(CosmosDiagnostics cosmosDiagnostics) { if (cosmosDiagnostics == null) { return null; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientRetryPolicy.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientRetryPolicy.java index 16e78c75e9dd..f12f67ea7d8b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientRetryPolicy.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientRetryPolicy.java @@ -12,6 +12,7 @@ import com.azure.cosmos.implementation.circuitBreaker.GlobalPartitionEndpointManagerForCircuitBreaker; import com.azure.cosmos.implementation.directconnectivity.WebExceptionUtility; import com.azure.cosmos.implementation.faultinjection.FaultInjectionRequestContext; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Mono; @@ -21,6 +22,7 @@ import java.util.concurrent.atomic.AtomicInteger; import static com.azure.cosmos.implementation.HttpConstants.HttpHeaders.INTENDED_COLLECTION_RID_HEADER; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; /** * While this class is public, but it is not part of our published public APIs. @@ -45,7 +47,7 @@ public class ClientRetryPolicy extends DocumentClientRetryPolicy { private int staleContainerRetryCount; private boolean isReadRequest; private boolean canUseMultipleWriteLocations; - private URI locationEndpoint; + private RegionalRoutingContext regionalRoutingContext; private RetryContext retryContext; private CosmosDiagnostics cosmosDiagnostics; private AtomicInteger cnt = new AtomicInteger(0); @@ -86,9 +88,9 @@ public Mono shouldRetry(Exception e) { isReadRequest, canUseMultipleWriteLocations, e); - if (this.locationEndpoint == null) { + if (this.regionalRoutingContext == null) { // on before request is not invoked because Document Service Request creation failed. - logger.error("locationEndpoint is null because ClientRetryPolicy::onBeforeRequest(.) is not invoked, " + + logger.error("regionalRoutingContext is null because ClientRetryPolicy::onBeforeRequest(.) is not invoked, " + "probably request creation failed due to invalid options, serialization setting, etc."); return Mono.just(ShouldRetryResult.error(e)); } @@ -228,7 +230,7 @@ private ShouldRetryResult shouldRetryOnSessionNotAvailable(RxDocumentServiceRequ return ShouldRetryResult.noRetry(); } else { if (this.canUseMultipleWriteLocations) { - UnmodifiableList endpoints = + UnmodifiableList endpoints = this.isReadRequest ? this.globalEndpointManager.getApplicableReadEndpoints(request) : this.globalEndpointManager.getApplicableWriteEndpoints(request); @@ -304,7 +306,7 @@ private Mono shouldRetryOnGatewayTimeout() { boolean canPerformCrossRegionRetryOnGatewayReadTimeout = canRequestToGatewayBeSafelyRetriedOnReadTimeout(this.request); if (this.globalPartitionEndpointManagerForCircuitBreaker.isPartitionLevelCircuitBreakingApplicable(this.request)) { - this.globalPartitionEndpointManagerForCircuitBreaker.handleLocationExceptionForPartitionKeyRange(this.request, this.request.requestContext.locationEndpointToRoute); + this.globalPartitionEndpointManagerForCircuitBreaker.handleLocationExceptionForPartitionKeyRange(this.request, this.request.requestContext.regionalRoutingContextToRoute); } //if operation is data plane read, metadata read, or query plan it can be retried on a different endpoint. @@ -336,12 +338,15 @@ private Mono refreshLocation(boolean isReadRequest, boolean forceRefresh, this.failoverRetryCount++; // Mark the current read endpoint as unavailable + + URI gatewayRegionalEndpoint = this.regionalRoutingContext.getGatewayRegionalEndpoint(); + if (isReadRequest) { - logger.warn("marking the endpoint {} as unavailable for read",this.locationEndpoint); - this.globalEndpointManager.markEndpointUnavailableForRead(this.locationEndpoint); + logger.warn("marking the endpoint {} as unavailable for read", gatewayRegionalEndpoint); + this.globalEndpointManager.markEndpointUnavailableForRead(gatewayRegionalEndpoint); } else { - logger.warn("marking the endpoint {} as unavailable for write",this.locationEndpoint); - this.globalEndpointManager.markEndpointUnavailableForWrite(this.locationEndpoint); + logger.warn("marking the endpoint {} as unavailable for write", gatewayRegionalEndpoint); + this.globalEndpointManager.markEndpointUnavailableForWrite(gatewayRegionalEndpoint); } this.retryContext = new RetryContext(this.failoverRetryCount, usePreferredLocations); @@ -356,7 +361,7 @@ private Mono shouldRetryOnBackendServiceUnavailableAsync( if (this.globalPartitionEndpointManagerForCircuitBreaker.isPartitionLevelCircuitBreakingApplicable(this.request)) { this.globalPartitionEndpointManagerForCircuitBreaker - .handleLocationExceptionForPartitionKeyRange(this.request, this.request.requestContext.locationEndpointToRoute); + .handleLocationExceptionForPartitionKeyRange(this.request, this.request.requestContext.regionalRoutingContextToRoute); } // The request has failed with 503, SDK need to decide whether it is safe to retry for write operations @@ -401,7 +406,9 @@ private Mono shouldRetryOnBackendServiceUnavailableAsync( return Mono.just(ShouldRetryResult.noRetry()); } - logger.info("shouldRetryOnServiceUnavailable() Retrying. Received on endpoint {}, IsReadRequest = {}", this.locationEndpoint, isReadRequest); + URI gatewayLocationEndpoint = this.regionalRoutingContext.getGatewayRegionalEndpoint(); + + logger.info("shouldRetryOnServiceUnavailable() Retrying. Received on endpoint {}, IsReadRequest = {}", gatewayLocationEndpoint, isReadRequest); // Retrying on second PreferredLocations // RetryCount is used as zero-based index @@ -415,9 +422,10 @@ private Mono shouldRetryOnRequestTimeout( if (this.globalPartitionEndpointManagerForCircuitBreaker.isPartitionLevelCircuitBreakingApplicable(this.request)) { if (!isReadRequest && !nonIdempotentWriteRetriesEnabled) { + this.globalPartitionEndpointManagerForCircuitBreaker.handleLocationExceptionForPartitionKeyRange( - request, - request.requestContext.locationEndpointToRoute); + this.request, + this.request.requestContext.regionalRoutingContextToRoute); } } @@ -427,9 +435,10 @@ private Mono shouldRetryOnRequestTimeout( private Mono shouldRetryOnInternalServerError() { if (this.globalPartitionEndpointManagerForCircuitBreaker.isPartitionLevelCircuitBreakingApplicable(this.request)) { + this.globalPartitionEndpointManagerForCircuitBreaker.handleLocationExceptionForPartitionKeyRange( - request, - request.requestContext.locationEndpointToRoute); + this.request, + this.request.requestContext.regionalRoutingContextToRoute); } return Mono.just(ShouldRetryResult.NO_RETRY); @@ -450,6 +459,7 @@ public void onBeforeSendRequest(RxDocumentServiceRequest request) { } if (this.retryContext != null) { // set location-based routing directive based on request retry context + checkNotNull(request.requestContext, "Argument 'request.requestContext' cannot be null!"); request.requestContext.routeToLocation(this.retryContext.retryCount, this.retryContext.retryRequestOnPreferredLocations); } @@ -458,9 +468,10 @@ public void onBeforeSendRequest(RxDocumentServiceRequest request) { // Resolve the endpoint for the request and pin the resolution to the resolved endpoint // This enables marking the endpoint unavailability on endpoint failover/unreachability - this.locationEndpoint = this.globalEndpointManager.resolveServiceEndpoint(request); + this.regionalRoutingContext = this.globalEndpointManager.resolveServiceEndpoint(request); + if (request.requestContext != null) { - request.requestContext.routeToLocation(this.locationEndpoint); + request.requestContext.routeToLocation(this.regionalRoutingContext); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientSideRequestStatistics.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientSideRequestStatistics.java index 278b53bc9bdd..d287ec73483e 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientSideRequestStatistics.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientSideRequestStatistics.java @@ -8,6 +8,7 @@ import com.azure.cosmos.implementation.directconnectivity.StoreResponseDiagnostics; import com.azure.cosmos.implementation.directconnectivity.StoreResultDiagnostics; import com.azure.cosmos.implementation.faultinjection.FaultInjectionRequestContext; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.core.JsonGenerator; @@ -35,6 +36,9 @@ import java.util.concurrent.ConcurrentLinkedDeque; import java.util.stream.Collectors; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + @JsonSerialize(using = ClientSideRequestStatistics.ClientSideRequestStatisticsSerializer.class) public class ClientSideRequestStatistics { private static final int MAX_SUPPLEMENTAL_REQUESTS_FOR_TO_STRING = 10; @@ -163,7 +167,9 @@ public void recordResponse(RxDocumentServiceRequest request, StoreResultDiagnost this.requestPayloadSizeInBytes = 0; } + RegionalRoutingContext regionalRoutingContext = null; URI locationEndPoint = null; + if (request.requestContext != null) { this.approximateInsertionCountInBloomFilter = request.requestContext.getApproximateBloomFilterInsertionCount(); @@ -175,7 +181,7 @@ public void recordResponse(RxDocumentServiceRequest request, StoreResultDiagnost request.requestContext.getEndToEndOperationLatencyPolicyConfig().toString(); } - locationEndPoint = request.requestContext.locationEndpointToRoute; + regionalRoutingContext = request.requestContext.regionalRoutingContextToRoute; List excludedRegions = request.requestContext.getExcludeRegions(); if (excludedRegions != null && !excludedRegions.isEmpty()) { @@ -189,12 +195,12 @@ public void recordResponse(RxDocumentServiceRequest request, StoreResultDiagnost this.requestEndTimeUTC = responseTime; } - if (locationEndPoint != null) { + if (regionalRoutingContext != null) { storeResponseStatistics.regionName = - globalEndpointManager.getRegionName(locationEndPoint, request.getOperationType()); + globalEndpointManager.getRegionName(regionalRoutingContext.getGatewayRegionalEndpoint(), request.getOperationType()); this.regionsContacted.add(storeResponseStatistics.regionName); this.locationEndpointsContacted.add(locationEndPoint); - this.regionsContactedWithContext.add(new RegionWithContext(storeResponseStatistics.regionName, locationEndPoint)); + this.regionsContactedWithContext.add(new RegionWithContext(storeResponseStatistics.regionName, regionalRoutingContext)); } if (storeResponseStatistics.requestOperationType == OperationType.Head @@ -218,22 +224,29 @@ public void recordGatewayResponse( this.requestEndTimeUTC = responseTime; } - URI locationEndPoint = null; + RegionalRoutingContext regionalRoutingContext = null; + if (rxDocumentServiceRequest != null && rxDocumentServiceRequest.requestContext != null) { - locationEndPoint = rxDocumentServiceRequest.requestContext.locationEndpointToRoute; + + if (rxDocumentServiceRequest.requestContext.regionalRoutingContextToRoute != null) { + regionalRoutingContext = rxDocumentServiceRequest.requestContext.regionalRoutingContextToRoute; + } + this.approximateInsertionCountInBloomFilter = rxDocumentServiceRequest.requestContext.getApproximateBloomFilterInsertionCount(); this.keywordIdentifiers = rxDocumentServiceRequest.requestContext.getKeywordIdentifiers(); } + this.recordRetryContextEndTime(); - if (locationEndPoint != null) { + if (regionalRoutingContext != null) { - String regionName = globalEndpointManager.getRegionName(locationEndPoint, rxDocumentServiceRequest.getOperationType()); + URI locationEndpoint = regionalRoutingContext.getGatewayRegionalEndpoint(); + String regionName = globalEndpointManager.getRegionName(locationEndpoint, rxDocumentServiceRequest.getOperationType()); this.regionsContacted.add(regionName); - this.locationEndpointsContacted.add(locationEndPoint); + this.locationEndpointsContacted.add(locationEndpoint); - this.regionsContactedWithContext.add(new RegionWithContext(regionName, locationEndPoint)); + this.regionsContactedWithContext.add(new RegionWithContext(regionName, regionalRoutingContext)); } GatewayStatistics gatewayStatistics = new GatewayStatistics(); @@ -651,7 +664,7 @@ public String getFirstContactedRegion() { return this.regionsContactedWithContext.first().regionContacted; } - public URI getFirstContactedLocationEndpoint() { + public RegionalRoutingContext getFirstContactedLocationEndpoint() { if (this.regionsContactedWithContext == null || this.regionsContactedWithContext.isEmpty()) { return null; } @@ -1048,10 +1061,10 @@ public static CosmosDiagnosticsSystemUsageSnapshot fetchSystemInformation() { static class RegionWithContext implements Comparable { private final String regionContacted; - private final URI locationEndpointsContacted; + private final RegionalRoutingContext locationEndpointsContacted; private final long recordedTimestamp; - RegionWithContext(String regionContacted, URI locationEndpointsContacted) { + RegionWithContext(String regionContacted, RegionalRoutingContext locationEndpointsContacted) { this.regionContacted = regionContacted; this.locationEndpointsContacted = locationEndpointsContacted; this.recordedTimestamp = System.currentTimeMillis(); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java index d49074416874..cfed0edadb55 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java @@ -5,7 +5,6 @@ import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.implementation.circuitBreaker.PartitionLevelCircuitBreakerConfig; import com.azure.cosmos.implementation.directconnectivity.Protocol; -import com.fasterxml.jackson.annotation.JsonInclude; import io.netty.handler.ssl.ApplicationProtocolConfig; import io.netty.handler.ssl.ApplicationProtocolNames; import io.netty.handler.ssl.SslContext; @@ -16,10 +15,8 @@ import javax.net.ssl.SSLException; import java.net.URI; -import java.net.URISyntaxException; import java.time.Duration; import java.util.Locale; -import java.util.Objects; import static com.azure.cosmos.implementation.guava25.base.MoreObjects.firstNonNull; import static com.azure.cosmos.implementation.guava25.base.Strings.emptyToNull; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DocumentServiceRequestContext.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DocumentServiceRequestContext.java index 0fbc9ecf0c93..911309489584 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DocumentServiceRequestContext.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DocumentServiceRequestContext.java @@ -12,10 +12,9 @@ import com.azure.cosmos.implementation.directconnectivity.StoreResult; import com.azure.cosmos.implementation.directconnectivity.TimeoutHelper; import com.azure.cosmos.implementation.directconnectivity.Uri; -import com.azure.cosmos.implementation.guava25.collect.ImmutableSet; import com.azure.cosmos.implementation.routing.PartitionKeyInternal; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; -import java.net.URI; import java.util.List; import java.util.Map; import java.util.Set; @@ -39,7 +38,7 @@ public class DocumentServiceRequestContext implements Cloneable { public volatile Integer regionIndex; public volatile Boolean usePreferredLocations; public volatile Integer locationIndexToRoute; - public volatile URI locationEndpointToRoute; + public volatile RegionalRoutingContext regionalRoutingContextToRoute; public volatile boolean performedBackgroundAddressRefresh; public volatile boolean performLocalRefreshOnGoneException; public volatile List storeResponses; @@ -80,17 +79,16 @@ public DocumentServiceRequestContext() {} public void routeToLocation(int locationIndex, boolean usePreferredLocations) { this.locationIndexToRoute = locationIndex; this.usePreferredLocations = usePreferredLocations; - this.locationEndpointToRoute = null; + this.regionalRoutingContextToRoute = null; } /** * Sets location-based routing directive for GlobalEndpointManager to resolve * the request to given locationEndpoint. * - * @param locationEndpoint Location endpoint to which the request should be routed. */ - public void routeToLocation(URI locationEndpoint) { - this.locationEndpointToRoute = locationEndpoint; + public void routeToLocation(RegionalRoutingContext regionalRoutingContextToRoute) { + this.regionalRoutingContextToRoute = regionalRoutingContextToRoute; this.locationIndexToRoute = null; this.usePreferredLocations = null; } @@ -100,7 +98,7 @@ public void routeToLocation(URI locationEndpoint) { */ public void clearRouteToLocation() { this.locationIndexToRoute = null; - this.locationEndpointToRoute = null; + this.regionalRoutingContextToRoute = null; this.usePreferredLocations = null; } @@ -141,7 +139,7 @@ public DocumentServiceRequestContext clone() { context.regionIndex = this.regionIndex; context.usePreferredLocations = this.usePreferredLocations; context.locationIndexToRoute = this.locationIndexToRoute; - context.locationEndpointToRoute = this.locationEndpointToRoute; + context.regionalRoutingContextToRoute = this.regionalRoutingContextToRoute; context.performLocalRefreshOnGoneException = this.performLocalRefreshOnGoneException; context.effectivePartitionKey = this.effectivePartitionKey; context.performedBackgroundAddressRefresh = this.performedBackgroundAddressRefresh; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/GlobalEndpointManager.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/GlobalEndpointManager.java index 1b46cd560ee4..5396db24872f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/GlobalEndpointManager.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/GlobalEndpointManager.java @@ -6,6 +6,7 @@ import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList; import com.azure.cosmos.implementation.routing.LocationCache; import com.azure.cosmos.implementation.routing.LocationHelper; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; @@ -88,32 +89,32 @@ public void init() { startRefreshLocationTimerAsync(true).block(maxInitializationTime); } - public UnmodifiableList getReadEndpoints() { + public UnmodifiableList getReadEndpoints() { // readonly return this.locationCache.getReadEndpoints(); } - public UnmodifiableList getWriteEndpoints() { + public UnmodifiableList getWriteEndpoints() { //readonly return this.locationCache.getWriteEndpoints(); } - public UnmodifiableList getApplicableReadEndpoints(RxDocumentServiceRequest request) { + public UnmodifiableList getApplicableReadEndpoints(RxDocumentServiceRequest request) { // readonly return this.locationCache.getApplicableReadEndpoints(request); } - public UnmodifiableList getApplicableWriteEndpoints(RxDocumentServiceRequest request) { + public UnmodifiableList getApplicableWriteEndpoints(RxDocumentServiceRequest request) { //readonly return this.locationCache.getApplicableWriteEndpoints(request); } - public UnmodifiableList getApplicableReadEndpoints(List excludedRegions) { + public UnmodifiableList getApplicableReadEndpoints(List excludedRegions) { // readonly return this.locationCache.getApplicableReadEndpoints(excludedRegions, Collections.emptyList()); } - public UnmodifiableList getApplicableWriteEndpoints(List excludedRegions) { + public UnmodifiableList getApplicableWriteEndpoints(List excludedRegions) { //readonly return this.locationCache.getApplicableWriteEndpoints(excludedRegions, Collections.emptyList()); } @@ -146,13 +147,14 @@ public static Mono getDatabaseAccountFromAnyLocationsAsync( }); } - public URI resolveServiceEndpoint(RxDocumentServiceRequest request) { - URI serviceEndpoint = this.locationCache.resolveServiceEndpoint(request); + public RegionalRoutingContext resolveServiceEndpoint(RxDocumentServiceRequest request) { + RegionalRoutingContext serviceEndpoints = this.locationCache.resolveServiceEndpoint(request); if (request.faultInjectionRequestContext != null) { - request.faultInjectionRequestContext.setLocationEndpointToRoute(serviceEndpoint); + // TODO: integrate thin client into fault injection + request.faultInjectionRequestContext.setLocationEndpointToRoute(serviceEndpoints.getGatewayRegionalEndpoint()); } - return serviceEndpoint; + return serviceEndpoints; } public URI resolveFaultInjectionServiceEndpoint(String region, boolean writeOnly) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java index dfab7c69fdb6..7def5159a963 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java @@ -41,6 +41,7 @@ import com.azure.cosmos.implementation.faultinjection.IFaultInjectorProvider; import com.azure.cosmos.implementation.patch.PatchOperation; import com.azure.cosmos.implementation.routing.PartitionKeyInternal; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple; import com.azure.cosmos.models.CosmosBatch; import com.azure.cosmos.models.CosmosBatchOperationResult; @@ -83,7 +84,6 @@ import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; -import java.net.URI; import java.time.Duration; import java.util.Collection; import java.util.EnumSet; @@ -838,7 +838,7 @@ void recordAddressResolutionEnd( void setDiagnosticsContext(CosmosDiagnostics cosmosDiagnostics, CosmosDiagnosticsContext ctx); - URI getFirstContactedLocationEndpoint(CosmosDiagnostics cosmosDiagnostics); + RegionalRoutingContext getFirstContactedLocationEndpoint(CosmosDiagnostics cosmosDiagnostics); void mergeMetadataDiagnosticContext(CosmosDiagnostics cosmosDiagnostics, MetadataDiagnosticsContext otherMetadataDiagnosticsContext); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/MetadataRequestRetryPolicy.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/MetadataRequestRetryPolicy.java index 4e7bf835250e..d17aed4326b6 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/MetadataRequestRetryPolicy.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/MetadataRequestRetryPolicy.java @@ -6,6 +6,7 @@ import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.CosmosException; import com.azure.cosmos.implementation.directconnectivity.WebExceptionUtility; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Mono; @@ -69,9 +70,10 @@ public Mono shouldRetry(Exception e) { if (shouldMarkRegionAsUnavailable(cosmosException)) { - if (request.requestContext != null && request.requestContext.locationEndpointToRoute != null) { + if (request.requestContext != null && request.requestContext.regionalRoutingContextToRoute != null) { - URI locationEndpointToRoute = request.requestContext.locationEndpointToRoute; + RegionalRoutingContext regionalRoutingContext = request.requestContext.regionalRoutingContextToRoute; + URI locationEndpointToRoute = regionalRoutingContext.getGatewayRegionalEndpoint(); if (request.isReadOnlyRequest()) { logger.warn("Marking the endpoint : {} as unavailable for read.", locationEndpointToRoute); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RegionScopedSessionContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RegionScopedSessionContainer.java index daf7204d6f93..b93761935dc5 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RegionScopedSessionContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RegionScopedSessionContainer.java @@ -8,6 +8,7 @@ import com.azure.cosmos.implementation.apachecommons.math.util.Pair; import com.azure.cosmos.implementation.routing.PartitionKeyInternal; import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import com.azure.cosmos.models.PartitionKeyDefinition; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -387,7 +388,9 @@ private void addSessionTokenAndTryRecordEpkInBloomFilter(RxDocumentServiceReques String regionRoutedTo = null; if (request.requestContext != null) { - URI regionEndpointRoutedTo = request.requestContext.locationEndpointToRoute; + RegionalRoutingContext regionalRoutingContext = request.requestContext.regionalRoutingContextToRoute; + URI regionEndpointRoutedTo = regionalRoutingContext.getGatewayRegionalEndpoint(); + regionRoutedTo = this.globalEndpointManager.getRegionName(regionEndpointRoutedTo, request.getOperationType()); } @@ -525,7 +528,7 @@ private String extractFirstEffectivePreferredReadableRegion() { List regionNamesForRead = globalEndpointManager .getReadEndpoints() .stream() - .map(endpoint -> globalEndpointManager.getRegionName(endpoint, OperationType.Read)) + .map(consolidatedReadLocationEndpoints -> globalEndpointManager.getRegionName(consolidatedReadLocationEndpoints.getGatewayRegionalEndpoint(), OperationType.Read)) .collect(Collectors.toList()); checkNotNull(regionNamesForRead, "regionNamesForRead cannot be null!"); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ResourceThrottleRetryPolicy.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ResourceThrottleRetryPolicy.java index bca5929dc8c7..bea73e015285 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ResourceThrottleRetryPolicy.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ResourceThrottleRetryPolicy.java @@ -74,7 +74,7 @@ public Mono shouldRetry(Exception exception) { "Operation will NOT be retried - not a throttled request. Current attempt {}", this.currentAttemptCount, exception); - return Mono.just(ShouldRetryResult.noRetryOnNonRelatedException()); + return Mono.just(ShouldRetryResult.errorOnNonRelatedException(exception)); } if (!retryOnClientSideThrottledBatchRequests && diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index e30ee7916a1c..e4aac881c87a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -62,6 +62,7 @@ import com.azure.cosmos.implementation.routing.PartitionKeyRangeIdentity; import com.azure.cosmos.implementation.routing.Range; import com.azure.cosmos.implementation.routing.RegionNameToRegionIdMap; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import com.azure.cosmos.implementation.spark.OperationContext; import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple; import com.azure.cosmos.implementation.spark.OperationListener; @@ -146,7 +147,7 @@ public class RxDocumentClientImpl implements AsyncDocumentClient, IAuthorization private final static List EMPTY_REGION_LIST = Collections.emptyList(); - private final static List EMPTY_ENDPOINT_LIST = Collections.emptyList(); + private final static List EMPTY_ENDPOINT_LIST = Collections.emptyList(); private final static ImplementationBridgeHelpers.CosmosDiagnosticsHelper.CosmosDiagnosticsAccessor diagnosticsAccessor = @@ -6460,7 +6461,7 @@ private DiagnosticsClientContext getEffectiveClientContext(DiagnosticsClientCont * @param operationType - the operationT * @return the applicable endpoints ordered by preference list if any */ - private List getApplicableEndPoints(OperationType operationType, List excludedRegions) { + private List getApplicableEndPoints(OperationType operationType, List excludedRegions) { if (operationType.isReadOnlyOperation()) { return withoutNulls(this.globalEndpointManager.getApplicableReadEndpoints(excludedRegions)); } else if (operationType.isWriteOperation()) { @@ -6470,7 +6471,7 @@ private List getApplicableEndPoints(OperationType operationType, List withoutNulls(List orderedEffectiveEndpointsList) { + private static List withoutNulls(List orderedEffectiveEndpointsList) { if (orderedEffectiveEndpointsList == null) { return EMPTY_ENDPOINT_LIST; } @@ -6529,7 +6530,7 @@ private List getApplicableRegionsForSpeculation( return EMPTY_REGION_LIST; } - List endpoints = getApplicableEndPoints(operationType, excludedRegions); + List regionalRoutingContextList = getApplicableEndPoints(operationType, excludedRegions); HashSet normalizedExcludedRegions = new HashSet<>(); if (excludedRegions != null) { @@ -6537,8 +6538,8 @@ private List getApplicableRegionsForSpeculation( } List orderedRegionsForSpeculation = new ArrayList<>(); - endpoints.forEach(uri -> { - String regionName = this.globalEndpointManager.getRegionName(uri, operationType); + regionalRoutingContextList.forEach(consolidatedLocationEndpoints -> { + String regionName = this.globalEndpointManager.getRegionName(consolidatedLocationEndpoints.getGatewayRegionalEndpoint(), operationType); if (!normalizedExcludedRegions.contains(regionName.toLowerCase(Locale.ROOT))) { orderedRegionsForSpeculation.add(regionName); } @@ -6742,7 +6743,7 @@ private Mono executeFeedOperationWithAvailabilityStrategy( private void handleLocationCancellationExceptionForPartitionKeyRange(RxDocumentServiceRequest failedRequest) { - URI firstContactedLocationEndpoint = diagnosticsAccessor + RegionalRoutingContext firstContactedLocationEndpoint = diagnosticsAccessor .getFirstContactedLocationEndpoint(failedRequest.requestContext.cosmosDiagnostics); if (firstContactedLocationEndpoint != null) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java index 0ef03f72a061..94c6d4a1a224 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java @@ -12,9 +12,7 @@ import com.azure.cosmos.implementation.directconnectivity.GatewayServiceConfigurationReader; import com.azure.cosmos.implementation.directconnectivity.HttpUtils; import com.azure.cosmos.implementation.directconnectivity.RequestHelper; -import com.azure.cosmos.implementation.directconnectivity.ResourceOperation; import com.azure.cosmos.implementation.directconnectivity.StoreResponse; -import com.azure.cosmos.implementation.directconnectivity.Uri; import com.azure.cosmos.implementation.directconnectivity.WebExceptionUtility; import com.azure.cosmos.implementation.faultinjection.GatewayServerErrorInjector; import com.azure.cosmos.implementation.faultinjection.IFaultInjectorProvider; @@ -322,9 +320,9 @@ private URI getUri(RxDocumentServiceRequest request) throws URISyntaxException { if (rootUri == null) { if (request.getIsMedia()) { // For media read request, always use the write endpoint. - rootUri = this.globalEndpointManager.getWriteEndpoints().get(0); + rootUri = this.globalEndpointManager.getWriteEndpoints().get(0).getGatewayRegionalEndpoint(); } else { - rootUri = this.globalEndpointManager.resolveServiceEndpoint(request); + rootUri = this.globalEndpointManager.resolveServiceEndpoint(request).getGatewayRegionalEndpoint(); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ShouldRetryResult.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ShouldRetryResult.java index 5869680a2bc4..80eff8fcc626 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ShouldRetryResult.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ShouldRetryResult.java @@ -41,6 +41,10 @@ public static ShouldRetryResult noRetryOnNonRelatedException() { return new ShouldRetryResult(null, null, false, null, true); } + public static ShouldRetryResult errorOnNonRelatedException(Exception e) { + return new ShouldRetryResult(null, e, false, null, true); + } + public static ShouldRetryResult noRetry(Quadruple policyArg) { return new ShouldRetryResult( null, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/WebExceptionRetryPolicy.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/WebExceptionRetryPolicy.java index c60eca10f47f..ab1452cd2537 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/WebExceptionRetryPolicy.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/WebExceptionRetryPolicy.java @@ -7,6 +7,7 @@ import com.azure.cosmos.implementation.directconnectivity.WebExceptionUtility; import com.azure.cosmos.implementation.http.HttpTimeoutPolicy; import com.azure.cosmos.implementation.http.HttpTimeoutPolicyDefault; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Mono; @@ -14,6 +15,8 @@ import java.net.URI; import java.time.Duration; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; + public class WebExceptionRetryPolicy implements IRetryPolicy { private final static Logger logger = LoggerFactory.getLogger(WebExceptionRetryPolicy.class); @@ -23,7 +26,8 @@ public class WebExceptionRetryPolicy implements IRetryPolicy { private HttpTimeoutPolicy timeoutPolicy; private boolean isReadRequest; private int retryCount = 0; - private URI locationEndpoint; + private RegionalRoutingContext regionalRoutingContext; + private URI overriddenEndpoint; public WebExceptionRetryPolicy() { durationTimer.start(); @@ -37,12 +41,17 @@ public WebExceptionRetryPolicy(RetryContext retryContext) { @Override public Mono shouldRetry(Exception e) { + + checkArgument(this.overriddenEndpoint != null || + this.regionalRoutingContext != null, + "Both overriddenEndpoint and regionalRoutingContext cannot null!"); + if (this.isOutOfRetries()) { logger .warn( "WebExceptionRetryPolicy() No more retrying on endpoint {}, operationType = {}, count = {}, " + "isAddressRefresh = {}", - this.locationEndpoint, + this.regionalRoutingContext != null ? this.regionalRoutingContext.getGatewayRegionalEndpoint() : this.overriddenEndpoint, this.request.getOperationType(), this.retryCount, this.request.isAddressRefresh()); @@ -63,7 +72,8 @@ public Mono shouldRetry(Exception e) { .debug("WebExceptionRetryPolicy() Retrying on endpoint {}, operationType = {}, resourceType = {}, count = {}, " + "isAddressRefresh = {}, shouldForcedAddressRefresh = {}, " + "shouldForceCollectionRoutingMapRefresh = {}", - this.locationEndpoint, this.request.getOperationType(), this.request.getResourceType(), this.retryCount, + this.regionalRoutingContext != null ? this.regionalRoutingContext.getGatewayRegionalEndpoint() : this.overriddenEndpoint, + this.request.getOperationType(), this.request.getResourceType(), this.retryCount, this.request.isAddressRefresh(), this.request.shouldForceAddressRefresh(), this.request.forceCollectionRoutingMapRefresh); @@ -77,7 +87,7 @@ public Mono shouldRetry(Exception e) { .debug( "WebExceptionRetryPolicy() No retrying on un-retryable exceptions on endpoint {}, operationType = {}, resourceType = {}, count = {}, " + "isAddressRefresh = {}", - this.locationEndpoint, + this.regionalRoutingContext != null ? this.regionalRoutingContext.getGatewayRegionalEndpoint() : this.overriddenEndpoint, this.request.getOperationType(), this.request.getResourceType(), this.retryCount, @@ -101,7 +111,8 @@ public void onBeforeSendRequest(RxDocumentServiceRequest request) { // set the initial response timeout this.request.setResponseTimeout(timeoutPolicy.getTimeoutAndDelaysList().get(0).getResponseTimeout()); - this.locationEndpoint = request.requestContext.locationEndpointToRoute; + this.regionalRoutingContext = request.requestContext.regionalRoutingContextToRoute; + this.overriddenEndpoint = request.getEndpointOverride(); } private boolean isOutOfRetries() { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/circuitBreaker/GlobalPartitionEndpointManagerForCircuitBreaker.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/circuitBreaker/GlobalPartitionEndpointManagerForCircuitBreaker.java index 92fccede1052..84fcf5579de8 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/circuitBreaker/GlobalPartitionEndpointManagerForCircuitBreaker.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/circuitBreaker/GlobalPartitionEndpointManagerForCircuitBreaker.java @@ -17,6 +17,7 @@ import com.azure.cosmos.implementation.apachecommons.lang.tuple.Pair; import com.azure.cosmos.implementation.directconnectivity.GatewayAddressCache; import com.azure.cosmos.implementation.directconnectivity.GlobalAddressResolver; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; @@ -49,7 +50,7 @@ public class GlobalPartitionEndpointManagerForCircuitBreaker implements AutoClos private final LocationSpecificHealthContextTransitionHandler locationSpecificHealthContextTransitionHandler; private final ConsecutiveExceptionBasedCircuitBreaker consecutiveExceptionBasedCircuitBreaker; private final AtomicReference globalAddressResolverSnapshot; - private final ConcurrentHashMap locationToRegion; + private final ConcurrentHashMap locationToRegion; private final AtomicBoolean isClosed = new AtomicBoolean(false); private final Scheduler partitionRecoveryScheduler = Schedulers.newSingle("partition-availability-staleness-check"); @@ -72,7 +73,7 @@ public void init() { } } - public void handleLocationExceptionForPartitionKeyRange(RxDocumentServiceRequest request, URI failedLocation) { + public void handleLocationExceptionForPartitionKeyRange(RxDocumentServiceRequest request, RegionalRoutingContext failedLocation) { checkNotNull(request, "Argument 'request' cannot be null!"); checkNotNull(request.requestContext, "Argument 'request.requestContext' cannot be null!"); @@ -109,7 +110,7 @@ public void handleLocationExceptionForPartitionKeyRange(RxDocumentServiceRequest if (isFailureThresholdBreached.get()) { - UnmodifiableList applicableEndpoints = request.isReadOnlyRequest() ? + UnmodifiableList applicableEndpoints = request.isReadOnlyRequest() ? this.globalEndpointManager.getApplicableReadEndpoints(request.requestContext.getExcludeRegions()) : this.globalEndpointManager.getApplicableWriteEndpoints(request.requestContext.getExcludeRegions()); @@ -132,7 +133,7 @@ public void handleLocationExceptionForPartitionKeyRange(RxDocumentServiceRequest logger.warn("It is not possible to mark region {} as Unavailable for partition key range {}-{} and collection rid {} " + "as all regions will be Unavailable in that case, will remove health status tracking for this partition!", this.globalEndpointManager.getRegionName( - failedLocation, request.isReadOnlyRequest() ? OperationType.Read : OperationType.Create), + failedLocation.getGatewayRegionalEndpoint(), request.isReadOnlyRequest() ? OperationType.Read : OperationType.Create), resolvedPartitionKeyRangeForCircuitBreaker.getMinInclusive(), resolvedPartitionKeyRangeForCircuitBreaker.getMaxExclusive(), collectionResourceId); @@ -161,7 +162,7 @@ public void handleLocationSuccessForPartitionKeyRange(RxDocumentServiceRequest r String resourceId = request.getResourceId(); PartitionKeyRangeWrapper partitionKeyRangeWrapper = new PartitionKeyRangeWrapper(resolvedPartitionKeyRangeForCircuitBreaker, resourceId); - URI succeededLocation = request.requestContext.locationEndpointToRoute; + RegionalRoutingContext succeededLocation = request.requestContext.regionalRoutingContextToRoute; String collectionLink = getCollectionLink(request); @@ -173,7 +174,6 @@ public void handleLocationSuccessForPartitionKeyRange(RxDocumentServiceRequest r partitionKeyRangeToFailoverInfoAsVal.handleSuccess( partitionKeyRangeWrapper, - collectionLink, succeededLocation, request.isReadOnlyRequest()); @@ -195,15 +195,18 @@ public List getUnavailableRegionsForPartitionKeyRange(String collectionR List unavailableRegions = new ArrayList<>(); if (partitionLevelLocationUnavailabilityInfoSnapshot != null) { - Map locationEndpointToFailureMetricsForPartition = + Map locationEndpointToFailureMetricsForPartition = partitionLevelLocationUnavailabilityInfoSnapshot.locationEndpointToLocationSpecificContextForPartition; - for (Map.Entry pair : locationEndpointToFailureMetricsForPartition.entrySet()) { - URI location = pair.getKey(); + for (Map.Entry pair : locationEndpointToFailureMetricsForPartition.entrySet()) { + RegionalRoutingContext regionalRoutingContext = pair.getKey(); + + URI gatewayLocationEndpoint = regionalRoutingContext.getGatewayRegionalEndpoint(); + LocationSpecificHealthContext locationSpecificHealthContext = pair.getValue(); if (locationSpecificHealthContext.getLocationHealthStatus() == LocationHealthStatus.Unavailable) { - unavailableRegions.add(this.globalEndpointManager.getRegionName(location, operationType)); + unavailableRegions.add(this.globalEndpointManager.getRegionName(gatewayLocationEndpoint, operationType)); } } } @@ -226,11 +229,11 @@ private Flux updateStaleLocationInfo() { if (partitionLevelLocationUnavailabilityInfo != null) { - List>> locationToLocationSpecificHealthContextList = new ArrayList<>(); + List>> locationToLocationSpecificHealthContextList = new ArrayList<>(); - for (Map.Entry locationToLocationLevelMetrics : partitionLevelLocationUnavailabilityInfo.locationEndpointToLocationSpecificContextForPartition.entrySet()) { + for (Map.Entry locationToLocationLevelMetrics : partitionLevelLocationUnavailabilityInfo.locationEndpointToLocationSpecificContextForPartition.entrySet()) { - URI locationWithStaleUnavailabilityInfo = locationToLocationLevelMetrics.getKey(); + RegionalRoutingContext locationWithStaleUnavailabilityInfo = locationToLocationLevelMetrics.getKey(); LocationSpecificHealthContext locationSpecificHealthContext = locationToLocationLevelMetrics.getValue(); if (!locationSpecificHealthContext.isRegionAvailableToProcessRequests()) { @@ -257,7 +260,7 @@ private Flux updateStaleLocationInfo() { .flatMap(locationToLocationSpecificHealthContextPair -> { PartitionKeyRangeWrapper partitionKeyRangeWrapper = locationToLocationSpecificHealthContextPair.getLeft(); - URI locationWithStaleUnavailabilityInfo = locationToLocationSpecificHealthContextPair.getRight().getLeft(); + RegionalRoutingContext locationWithStaleUnavailabilityInfo = locationToLocationSpecificHealthContextPair.getRight().getLeft(); PartitionLevelLocationUnavailabilityInfo partitionLevelLocationUnavailabilityInfo = this.partitionKeyRangeToLocationSpecificUnavailabilityInfo.get(partitionKeyRangeWrapper); @@ -267,7 +270,7 @@ private Flux updateStaleLocationInfo() { if (globalAddressResolver != null) { - GatewayAddressCache gatewayAddressCache = globalAddressResolver.getGatewayAddressCache(locationWithStaleUnavailabilityInfo); + GatewayAddressCache gatewayAddressCache = globalAddressResolver.getGatewayAddressCache(locationWithStaleUnavailabilityInfo.getGatewayRegionalEndpoint()); if (gatewayAddressCache != null) { @@ -354,7 +357,7 @@ public boolean isPartitionLevelCircuitBreakingApplicable(RxDocumentServiceReques return false; } - UnmodifiableList applicableWriteEndpoints = globalEndpointManager.getApplicableWriteEndpoints(Collections.emptyList()); + UnmodifiableList applicableWriteEndpoints = globalEndpointManager.getApplicableWriteEndpoints(Collections.emptyList()); return applicableWriteEndpoints != null && applicableWriteEndpoints.size() > 1; } @@ -371,7 +374,7 @@ public void close() { private class PartitionLevelLocationUnavailabilityInfo { - private final ConcurrentHashMap locationEndpointToLocationSpecificContextForPartition; + private final ConcurrentHashMap locationEndpointToLocationSpecificContextForPartition; private final ConcurrentHashMap regionToLocationSpecificHealthContext; private final LocationSpecificHealthContextTransitionHandler locationSpecificHealthContextTransitionHandler; @@ -383,7 +386,7 @@ private PartitionLevelLocationUnavailabilityInfo() { private boolean handleException( PartitionKeyRangeWrapper partitionKeyRangeWrapper, - URI locationWithException, + RegionalRoutingContext locationWithException, boolean isReadOnlyRequest) { AtomicBoolean isExceptionThresholdBreached = new AtomicBoolean(false); @@ -417,7 +420,7 @@ private boolean handleException( locationAsKey, GlobalPartitionEndpointManagerForCircuitBreaker .this.globalEndpointManager - .getRegionName(locationAsKey, isReadOnlyRequest ? OperationType.Read : OperationType.Create)); + .getRegionName(locationAsKey.getGatewayRegionalEndpoint(), isReadOnlyRequest ? OperationType.Read : OperationType.Create)); } String region = GlobalPartitionEndpointManagerForCircuitBreaker.this.locationToRegion.get(locationAsKey); @@ -432,8 +435,7 @@ private boolean handleException( private void handleSuccess( PartitionKeyRangeWrapper partitionKeyRangeWrapper, - String collectionLink, - URI succeededLocation, + RegionalRoutingContext succeededLocation, boolean isReadOnlyRequest) { this.locationEndpointToLocationSpecificContextForPartition.compute(succeededLocation, (locationAsKey, locationSpecificContextAsVal) -> { @@ -467,7 +469,7 @@ private void handleSuccess( locationAsKey, GlobalPartitionEndpointManagerForCircuitBreaker .this.globalEndpointManager - .getRegionName(locationAsKey, isReadOnlyRequest ? OperationType.Read : OperationType.Create)); + .getRegionName(locationAsKey.getGatewayRegionalEndpoint(), isReadOnlyRequest ? OperationType.Read : OperationType.Create)); } String region = GlobalPartitionEndpointManagerForCircuitBreaker.this.locationToRegion.get(locationAsKey); @@ -477,9 +479,9 @@ private void handleSuccess( }); } - public boolean areLocationsAvailableForPartitionKeyRange(List availableLocationsAtAccountLevel) { + public boolean areLocationsAvailableForPartitionKeyRange(List availableLocationsAtAccountLevel) { - for (URI availableLocation : availableLocationsAtAccountLevel) { + for (RegionalRoutingContext availableLocation : availableLocationsAtAccountLevel) { if (!this.locationEndpointToLocationSpecificContextForPartition.containsKey(availableLocation)) { return true; } else { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/GlobalAddressResolver.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/GlobalAddressResolver.java index ecd0fa583d35..00905682b4d1 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/GlobalAddressResolver.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/GlobalAddressResolver.java @@ -25,6 +25,7 @@ import com.azure.cosmos.implementation.http.HttpClient; import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper; import com.azure.cosmos.implementation.routing.PartitionKeyRangeIdentity; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import com.azure.cosmos.models.CosmosContainerIdentity; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -91,11 +92,11 @@ public GlobalAddressResolver( this.addressCacheByEndpoint = new ConcurrentHashMap<>(); this.apiType = apiType; - for (URI endpoint : endpointManager.getWriteEndpoints()) { - this.getOrAddEndpoint(endpoint); + for (RegionalRoutingContext endpoint : endpointManager.getWriteEndpoints()) { + this.getOrAddEndpoint(endpoint.getGatewayRegionalEndpoint()); } - for (URI endpoint : endpointManager.getReadEndpoints()) { - this.getOrAddEndpoint(endpoint); + for (RegionalRoutingContext endpoint : endpointManager.getReadEndpoints()) { + this.getOrAddEndpoint(endpoint.getGatewayRegionalEndpoint()); } } @@ -136,7 +137,7 @@ public Flux submitOpenConnectionTasksAndInitCaches(CosmosContainerProactiv .getCosmosContainerIdentityAccessor() .getContainerLink(cosmosContainerIdentity); - if (valueHolder == null || valueHolder.v == null || valueHolder.v.size() == 0) { + if (valueHolder == null || valueHolder.v == null || valueHolder.v.isEmpty()) { logger.warn( "There is no pkRanges found for collection {}, no connections will be opened", collection.getResourceId()); @@ -154,8 +155,8 @@ public Flux submitOpenConnectionTasksAndInitCaches(CosmosContainerProactiv if (proactiveContainerInitConfig.getProactiveConnectionRegionsCount() > 0) { return Flux.fromIterable(this.endpointManager.getReadEndpoints().subList(0, proactiveContainerInitConfig.getProactiveConnectionRegionsCount())) .flatMap(readEndpoint -> { - if (this.addressCacheByEndpoint.containsKey(readEndpoint)) { - EndpointCache endpointCache = this.addressCacheByEndpoint.get(readEndpoint); + if (this.addressCacheByEndpoint.containsKey(readEndpoint.getGatewayRegionalEndpoint())) { + EndpointCache endpointCache = this.addressCacheByEndpoint.get(readEndpoint.getGatewayRegionalEndpoint()); return this.resolveAddressesPerCollection( endpointCache, containerLinkToPkrs.left, @@ -272,8 +273,8 @@ public void configureFaultInjectorProvider(IFaultInjectorProvider faultInjectorP } private IAddressResolver getAddressResolver(RxDocumentServiceRequest rxDocumentServiceRequest) { - URI endpoint = this.endpointManager.resolveServiceEndpoint(rxDocumentServiceRequest); - return this.getOrAddEndpoint(endpoint).addressResolver; + RegionalRoutingContext endpoint = this.endpointManager.resolveServiceEndpoint(rxDocumentServiceRequest); + return this.getOrAddEndpoint(endpoint.getGatewayRegionalEndpoint()).addressResolver; } private EndpointCache getOrAddEndpoint(URI endpoint) { @@ -299,15 +300,15 @@ private EndpointCache getOrAddEndpoint(URI endpoint) { }); if (this.addressCacheByEndpoint.size() > this.maxEndpoints) { - List allEndpoints = new ArrayList<>(this.endpointManager.getWriteEndpoints()); - allEndpoints.addAll(this.endpointManager.getReadEndpoints()); - Collections.reverse(allEndpoints); - LinkedList endpoints = new LinkedList<>(allEndpoints); + List allConsolidatedEndpoints = new ArrayList<>(this.endpointManager.getWriteEndpoints()); + allConsolidatedEndpoints.addAll(this.endpointManager.getReadEndpoints()); + Collections.reverse(allConsolidatedEndpoints); + LinkedList endpoints = new LinkedList<>(allConsolidatedEndpoints); while (this.addressCacheByEndpoint.size() > this.maxEndpoints) { - if (endpoints.size() > 0) { - URI dequeueEndpoint = endpoints.pop(); - if (this.addressCacheByEndpoint.get(dequeueEndpoint) != null) { - this.addressCacheByEndpoint.remove(dequeueEndpoint); + if (!endpoints.isEmpty()) { + RegionalRoutingContext dequeueEndpoint = endpoints.pop(); + if (this.addressCacheByEndpoint.get(dequeueEndpoint.getGatewayRegionalEndpoint()) != null) { + this.addressCacheByEndpoint.remove(dequeueEndpoint.getGatewayRegionalEndpoint()); } } else { break; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/RntbdTransportClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/RntbdTransportClient.java index 614115a6d63e..bb56c878620a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/RntbdTransportClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/RntbdTransportClient.java @@ -283,7 +283,7 @@ public Mono invokeStoreAsync(final Uri addressUri, final RxDocume config.minConnectionPoolSizePerEndpoint() : 1; final RntbdEndpoint endpoint = this.endpointProvider.createIfAbsent( - request.requestContext.locationEndpointToRoute, + request.requestContext.regionalRoutingContextToRoute.getGatewayRegionalEndpoint(), addressUri, this.proactiveOpenConnectionsProcessor, minConnectionPoolSizePerEndpoint, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdOpenConnectionsHandler.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdOpenConnectionsHandler.java index 8058658bbab3..baab52b212a1 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdOpenConnectionsHandler.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdOpenConnectionsHandler.java @@ -10,6 +10,7 @@ import com.azure.cosmos.implementation.RxDocumentServiceRequest; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.implementation.directconnectivity.Uri; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; @@ -76,7 +77,8 @@ public Flux openConnections(String collectionRid, List> nextPage(RxDocumentServiceRequest request) { } private void handleCancellationExceptionForPartitionKeyRange(RxDocumentServiceRequest failedRequest) { - URI firstContactedLocationEndpoint = diagnosticsAccessor.getFirstContactedLocationEndpoint(failedRequest.requestContext.cosmosDiagnostics); + RegionalRoutingContext firstContactedLocationEndpoint = diagnosticsAccessor.getFirstContactedLocationEndpoint(failedRequest.requestContext.cosmosDiagnostics); if (firstContactedLocationEndpoint != null) { this.globalPartitionEndpointManagerForCircuitBreaker.handleLocationExceptionForPartitionKeyRange(failedRequest, firstContactedLocationEndpoint); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/routing/LocationCache.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/routing/LocationCache.java index fa95e2abe121..ecab529cbd74 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/routing/LocationCache.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/routing/LocationCache.java @@ -42,11 +42,11 @@ public class LocationCache { private final static Logger logger = LoggerFactory.getLogger(LocationCache.class); private final boolean enableEndpointDiscovery; - private final URI defaultEndpoint; + private final RegionalRoutingContext defaultRegionalRoutingContext; private final boolean useMultipleWriteLocations; private final Object lockObject; private final Duration unavailableLocationsExpirationTime; - private final ConcurrentHashMap locationUnavailabilityInfoByEndpoint; + private final ConcurrentHashMap locationUnavailabilityInfoByEndpoint; private final ConnectionPolicy connectionPolicy; private DatabaseAccountLocationsInfo locationInfo; @@ -65,7 +65,7 @@ public LocationCache( ); this.locationInfo = new DatabaseAccountLocationsInfo(preferredLocations, defaultEndpoint); - this.defaultEndpoint = defaultEndpoint; + this.defaultRegionalRoutingContext = new RegionalRoutingContext(defaultEndpoint); this.enableEndpointDiscovery = connectionPolicy.isEndpointDiscoveryEnabled(); this.useMultipleWriteLocations = connectionPolicy.isMultipleWriteRegionsEnabled(); @@ -86,7 +86,7 @@ public LocationCache( * 2. Endpoint availability * @return */ - public UnmodifiableList getReadEndpoints() { + public UnmodifiableList getReadEndpoints() { if (this.locationUnavailabilityInfoByEndpoint.size() > 0 && unavailableLocationsExpirationTimePassed()) { this.updateLocationCache(); @@ -101,7 +101,7 @@ && unavailableLocationsExpirationTimePassed()) { * 2. Endpoint availability * @return */ - public UnmodifiableList getWriteEndpoints() { + public UnmodifiableList getWriteEndpoints() { if (this.locationUnavailabilityInfoByEndpoint.size() > 0 && unavailableLocationsExpirationTimePassed()) { this.updateLocationCache(); @@ -119,7 +119,11 @@ && unavailableLocationsExpirationTimePassed()) { * @return */ public List getAvailableReadEndpoints() { - return this.locationInfo.availableReadEndpointByLocation.values().stream().collect(Collectors.toList()); + return this.locationInfo.availableReadEndpointsByLocation.values().stream().map(consolidatedLocationEndpoints -> { + // TODO (nehrao1): Integrate thinclient endpoints into fault injection + // TODO (nehrao1): https://github.com/Azure/azure-sdk-for-java/issues/44429 + return consolidatedLocationEndpoints.getGatewayRegionalEndpoint(); + }).collect(Collectors.toList()); } /*** @@ -130,7 +134,11 @@ public List getAvailableReadEndpoints() { * @return */ public List getAvailableWriteEndpoints() { - return this.locationInfo.availableWriteEndpointByLocation.values().stream().collect(Collectors.toList()); + return this.locationInfo.availableWriteEndpointsByLocation.values().stream().map(consolidatedLocationEndpoints -> { + // TODO(nehrao1): Integrate thinclient endpoints into fault injection + // TODO(nehrao1): https://github.com/Azure/azure-sdk-for-java/issues/44429 + return consolidatedLocationEndpoints.getGatewayRegionalEndpoint(); + }).collect(Collectors.toList()); } public List getEffectivePreferredLocations() { @@ -186,43 +194,44 @@ void onLocationPreferenceChanged(UnmodifiableList preferredLocations) { * @param request Request for which getEndpoint is to be resolved * @return Resolved getEndpoint */ - public URI resolveServiceEndpoint(RxDocumentServiceRequest request) { + public RegionalRoutingContext resolveServiceEndpoint(RxDocumentServiceRequest request) { Objects.requireNonNull(request.requestContext, "RxDocumentServiceRequest.requestContext is required and cannot be null."); - if(request.requestContext.locationEndpointToRoute != null) { - return request.requestContext.locationEndpointToRoute; + + if (request.requestContext.regionalRoutingContextToRoute != null) { + return request.requestContext.regionalRoutingContextToRoute; } int locationIndex = Utils.getValueOrDefault(request.requestContext.locationIndexToRoute, 0); boolean usePreferredLocations = request.requestContext.usePreferredLocations != null ? request.requestContext.usePreferredLocations : true; - if(!usePreferredLocations || (request.getOperationType().isWriteOperation() && !this.canUseMultipleWriteLocations(request))) { + if (!usePreferredLocations || (request.getOperationType().isWriteOperation() && !this.canUseMultipleWriteLocations(request))) { // For non-document resource types in case of client can use multiple write locations // or when client cannot use multiple write locations, flip-flop between the // first and the second writable region in DatabaseAccount (for manual failover) DatabaseAccountLocationsInfo currentLocationInfo = this.locationInfo; - if(this.enableEndpointDiscovery && currentLocationInfo.availableWriteLocations.size() > 0) { + if (this.enableEndpointDiscovery && !currentLocationInfo.availableWriteLocations.isEmpty()) { locationIndex = Math.min(locationIndex%2, currentLocationInfo.availableWriteLocations.size()-1); String writeLocation = currentLocationInfo.availableWriteLocations.get(locationIndex); - return currentLocationInfo.availableWriteEndpointByLocation.get(writeLocation); + return currentLocationInfo.availableWriteEndpointsByLocation.get(writeLocation); } else { - return this.defaultEndpoint; + return this.defaultRegionalRoutingContext; } } else { - UnmodifiableList endpoints = + UnmodifiableList endpoints = request.getOperationType().isWriteOperation()? this.getApplicableWriteEndpoints(request) : this.getApplicableReadEndpoints(request); return endpoints.get(locationIndex % endpoints.size()); } } - public UnmodifiableList getApplicableWriteEndpoints(RxDocumentServiceRequest request) { + public UnmodifiableList getApplicableWriteEndpoints(RxDocumentServiceRequest request) { return this.getApplicableWriteEndpoints(request.requestContext.getExcludeRegions(), request.requestContext.getUnavailableRegionsForPartition()); } - public UnmodifiableList getApplicableWriteEndpoints(List excludedRegionsOnRequest, List unavailableRegionsForPartition) { + public UnmodifiableList getApplicableWriteEndpoints(List excludedRegionsOnRequest, List unavailableRegionsForPartition) { - UnmodifiableList writeEndpoints = this.getWriteEndpoints(); + UnmodifiableList writeEndpoints = this.getWriteEndpoints(); Supplier excludedRegionsSupplier = this.connectionPolicy.getExcludedRegionsSupplier(); List effectiveExcludedRegions = isExcludedRegionsSupplierConfigured(excludedRegionsSupplier) ? @@ -246,16 +255,16 @@ public UnmodifiableList getApplicableWriteEndpoints(List excludedRe return this.getApplicableEndpoints( writeEndpoints, this.locationInfo.regionNameByWriteEndpoint, - this.defaultEndpoint, + this.defaultRegionalRoutingContext, effectiveExcludedRegionsWithPartitionUnavailableRegions); } - public UnmodifiableList getApplicableReadEndpoints(RxDocumentServiceRequest request) { + public UnmodifiableList getApplicableReadEndpoints(RxDocumentServiceRequest request) { return this.getApplicableReadEndpoints(request.requestContext.getExcludeRegions(), request.requestContext.getUnavailableRegionsForPartition()); } - public UnmodifiableList getApplicableReadEndpoints(List excludedRegionsOnRequest, List unavailableRegionsForPartition) { - UnmodifiableList readEndpoints = this.getReadEndpoints(); + public UnmodifiableList getApplicableReadEndpoints(List excludedRegionsOnRequest, List unavailableRegionsForPartition) { + UnmodifiableList readEndpoints = this.getReadEndpoints(); Supplier excludedRegionsSupplier = this.connectionPolicy.getExcludedRegionsSupplier(); List effectiveExcludedRegions = isExcludedRegionsSupplierConfigured(excludedRegionsSupplier) ? @@ -283,14 +292,14 @@ public UnmodifiableList getApplicableReadEndpoints(List excludedReg effectiveExcludedRegionsWithPartitionUnavailableRegions); } - private UnmodifiableList getApplicableEndpoints( - UnmodifiableList endpoints, - UnmodifiableMap regionNameByEndpoint, - URI fallbackEndpoint, + private UnmodifiableList getApplicableEndpoints( + UnmodifiableList endpoints, + UnmodifiableMap regionNameByEndpoint, + RegionalRoutingContext fallbackRegionalRoutingContext, List excludeRegionList) { - List applicableEndpoints = new ArrayList<>(); - for (URI endpoint : endpoints) { + List applicableEndpoints = new ArrayList<>(); + for (RegionalRoutingContext endpoint : endpoints) { Utils.ValueHolder regionName = new Utils.ValueHolder<>(); if (Utils.tryGetValue(regionNameByEndpoint, endpoint, regionName)) { if (!excludeRegionList.stream().anyMatch(regionName.v::equalsIgnoreCase)) { @@ -300,7 +309,7 @@ private UnmodifiableList getApplicableEndpoints( } if (applicableEndpoints.isEmpty()) { - applicableEndpoints.add(fallbackEndpoint); + applicableEndpoints.add(fallbackRegionalRoutingContext); } return new UnmodifiableList<>(applicableEndpoints); @@ -314,22 +323,23 @@ private boolean isExcludeRegionsConfigured(List excludedRegionsOnRequest } public URI resolveFaultInjectionEndpoint(String region, boolean writeOnly) { - Utils.ValueHolder endpointValueHolder = new Utils.ValueHolder<>(); + Utils.ValueHolder endpointValueHolder = new Utils.ValueHolder<>(); if (writeOnly) { - Utils.tryGetValue(this.locationInfo.availableWriteEndpointByLocation, region, endpointValueHolder); + Utils.tryGetValue(this.locationInfo.availableWriteEndpointsByLocation, region, endpointValueHolder); } else { - Utils.tryGetValue(this.locationInfo.availableReadEndpointByLocation, region, endpointValueHolder); + Utils.tryGetValue(this.locationInfo.availableReadEndpointsByLocation, region, endpointValueHolder); } if (endpointValueHolder.v != null) { - return endpointValueHolder.v; + // TODO: Figure out how to integrate thinclient into fault injection + return endpointValueHolder.v.getGatewayRegionalEndpoint(); } throw new IllegalArgumentException("Can not find service endpoint for region " + region); } public URI getDefaultEndpoint() { - return this.defaultEndpoint; + return this.defaultRegionalRoutingContext.getGatewayRegionalEndpoint(); } public boolean shouldRefreshEndpoints(Utils.ValueHolder canRefreshInBackground) { @@ -345,7 +355,7 @@ public boolean shouldRefreshEndpoints(Utils.ValueHolder canRefreshInBac if (this.enableEndpointDiscovery) { boolean shouldRefresh = this.useMultipleWriteLocations && !this.enableMultipleWriteLocations; - List readLocationEndpoints = currentLocationInfo.readEndpoints; + List readLocationEndpoints = currentLocationInfo.readEndpoints; if (this.isEndpointUnavailable(readLocationEndpoints.get(0), OperationType.Read)) { // Since most preferred read endpoint is unavailable, we can only refresh in background if // we have an alternate read endpoint @@ -358,10 +368,10 @@ public boolean shouldRefreshEndpoints(Utils.ValueHolder canRefreshInBac } if (!Strings.isNullOrEmpty(mostPreferredLocation)) { - Utils.ValueHolder mostPreferredReadEndpointHolder = new Utils.ValueHolder<>(); + Utils.ValueHolder mostPreferredReadEndpointHolder = new Utils.ValueHolder<>(); logger.debug("getReadEndpoints [{}]", readLocationEndpoints); - if (Utils.tryGetValue(currentLocationInfo.availableReadEndpointByLocation, mostPreferredLocation, mostPreferredReadEndpointHolder)) { + if (Utils.tryGetValue(currentLocationInfo.availableReadEndpointsByLocation, mostPreferredLocation, mostPreferredReadEndpointHolder)) { logger.debug("most preferred is [{}], most preferred available is [{}]", mostPreferredLocation, mostPreferredReadEndpointHolder.v); if (!areEqual(mostPreferredReadEndpointHolder.v, readLocationEndpoints.get(0))) { @@ -382,8 +392,8 @@ public boolean shouldRefreshEndpoints(Utils.ValueHolder canRefreshInBac } } - Utils.ValueHolder mostPreferredWriteEndpointHolder = new Utils.ValueHolder<>(); - List writeLocationEndpoints = currentLocationInfo.writeEndpoints; + Utils.ValueHolder mostPreferredWriteEndpointHolder = new Utils.ValueHolder<>(); + List writeLocationEndpoints = currentLocationInfo.writeEndpoints; logger.debug("getWriteEndpoints [{}]", writeLocationEndpoints); if (!this.canUseMultipleWriteLocations()) { @@ -403,7 +413,7 @@ public boolean shouldRefreshEndpoints(Utils.ValueHolder canRefreshInBac return shouldRefresh; } } else if (!Strings.isNullOrEmpty(mostPreferredLocation)) { - if (Utils.tryGetValue(currentLocationInfo.availableWriteEndpointByLocation, mostPreferredLocation, mostPreferredWriteEndpointHolder)) { + if (Utils.tryGetValue(currentLocationInfo.availableWriteEndpointsByLocation, mostPreferredLocation, mostPreferredWriteEndpointHolder)) { shouldRefresh = ! areEqual(mostPreferredWriteEndpointHolder.v,writeLocationEndpoints.get(0)); if (shouldRefresh) { @@ -433,11 +443,11 @@ public boolean shouldRefreshEndpoints(Utils.ValueHolder canRefreshInBac public String getRegionName(URI locationEndpoint, com.azure.cosmos.implementation.OperationType operationType) { Utils.ValueHolder regionName = new Utils.ValueHolder<>(); if (operationType.isWriteOperation()) { - if (Utils.tryGetValue(this.locationInfo.regionNameByWriteEndpoint, locationEndpoint, regionName)) { + if (Utils.tryGetValue(this.locationInfo.regionNameByWriteEndpoint, new RegionalRoutingContext(locationEndpoint), regionName)) { return regionName.v; } } else { - if (Utils.tryGetValue(this.locationInfo.regionNameByReadEndpoint, locationEndpoint, regionName)) { + if (Utils.tryGetValue(this.locationInfo.regionNameByReadEndpoint, new RegionalRoutingContext(locationEndpoint), regionName)) { return regionName.v; } } @@ -446,15 +456,15 @@ public String getRegionName(URI locationEndpoint, com.azure.cosmos.implementatio return this.locationInfo.availableWriteLocations.get(0).toLowerCase(Locale.ROOT); } - private boolean areEqual(URI url1, URI url2) { + private boolean areEqual(RegionalRoutingContext url1, RegionalRoutingContext url2) { return url1.equals(url2); } private void clearStaleEndpointUnavailabilityInfo() { if (!this.locationUnavailabilityInfoByEndpoint.isEmpty()) { - List unavailableEndpoints = new ArrayList<>(this.locationUnavailabilityInfoByEndpoint.keySet()); + List unavailableEndpoints = new ArrayList<>(this.locationUnavailabilityInfoByEndpoint.keySet()); - for (URI unavailableEndpoint: unavailableEndpoints) { + for (RegionalRoutingContext unavailableEndpoint: unavailableEndpoints) { Utils.ValueHolder unavailabilityInfoHolder = new Utils.ValueHolder<>(); Utils.ValueHolder removedHolder = new Utils.ValueHolder<>(); @@ -473,7 +483,7 @@ private void clearStaleEndpointUnavailabilityInfo() { } } - private boolean isEndpointUnavailable(URI endpoint, OperationType expectedAvailableOperations) { + private boolean isEndpointUnavailable(RegionalRoutingContext endpoint, OperationType expectedAvailableOperations) { Utils.ValueHolder unavailabilityInfoHolder = new Utils.ValueHolder<>(); if (expectedAvailableOperations == OperationType.None @@ -494,10 +504,10 @@ private boolean isEndpointUnavailable(URI endpoint, OperationType expectedAvaila } } - private boolean anyEndpointsAvailable(List endpoints, OperationType expectedAvailableOperations) { + private boolean anyEndpointsAvailable(List endpoints, OperationType expectedAvailableOperations) { Utils.ValueHolder unavailabilityInfoHolder = new Utils.ValueHolder<>(); boolean anyEndpointsAvailable = false; - for (URI endpoint : endpoints) { + for (RegionalRoutingContext endpoint : endpoints) { if (!isEndpointUnavailable(endpoint, expectedAvailableOperations)) { anyEndpointsAvailable = true; break; @@ -511,10 +521,10 @@ private void markEndpointUnavailable( OperationType unavailableOperationType) { Instant currentTime = Instant.now(); LocationUnavailabilityInfo updatedInfo = this.locationUnavailabilityInfoByEndpoint.compute( - unavailableEndpoint, - new BiFunction() { + new RegionalRoutingContext(unavailableEndpoint), + new BiFunction() { @Override - public LocationUnavailabilityInfo apply(URI url, LocationUnavailabilityInfo info) { + public LocationUnavailabilityInfo apply(RegionalRoutingContext url, LocationUnavailabilityInfo info) { if (info == null) { // not already present, add return new LocationUnavailabilityInfo(currentTime, unavailableOperationType); @@ -524,7 +534,6 @@ public LocationUnavailabilityInfo apply(URI url, LocationUnavailabilityInfo info info.unavailableOperations = OperationType.combine(info.unavailableOperations, unavailableOperationType); return info; } - } }); @@ -542,8 +551,8 @@ private void updateLocationCache(){ } private void updateLocationCache( - Iterable writeLocations, - Iterable readLocations, + Iterable gatewayWriteLocations, + Iterable gatewayReadLocations, UnmodifiableList preferenceList, Boolean enableMultipleWriteLocations) { synchronized (this.lockObject) { @@ -561,31 +570,31 @@ private void updateLocationCache( this.clearStaleEndpointUnavailabilityInfo(); - if (readLocations != null) { - Utils.ValueHolder> out = Utils.ValueHolder.initialize(nextLocationInfo.availableReadLocations); - Utils.ValueHolder> outReadRegionMap = Utils.ValueHolder.initialize(nextLocationInfo.regionNameByReadEndpoint); - nextLocationInfo.availableReadEndpointByLocation = this.getEndpointByLocation(readLocations, out, outReadRegionMap); - nextLocationInfo.availableReadLocations = out.v; - nextLocationInfo.regionNameByReadEndpoint = outReadRegionMap.v; + if (gatewayReadLocations != null) { + Utils.ValueHolder> readValueHolder = Utils.ValueHolder.initialize(nextLocationInfo.availableReadLocations); + Utils.ValueHolder> readRegionMapValueHolder = Utils.ValueHolder.initialize(nextLocationInfo.regionNameByReadEndpoint); + nextLocationInfo.availableReadEndpointsByLocation = this.getEndpointsByLocation(gatewayReadLocations, readValueHolder, readRegionMapValueHolder); + nextLocationInfo.availableReadLocations = readValueHolder.v; + nextLocationInfo.regionNameByReadEndpoint = readRegionMapValueHolder.v; } - if (writeLocations != null) { - Utils.ValueHolder> out = Utils.ValueHolder.initialize(nextLocationInfo.availableWriteLocations); - Utils.ValueHolder> outWriteRegionMap = Utils.ValueHolder.initialize(nextLocationInfo.regionNameByWriteEndpoint); - nextLocationInfo.availableWriteEndpointByLocation = this.getEndpointByLocation(writeLocations, out, outWriteRegionMap); - nextLocationInfo.availableWriteLocations = out.v; + if (gatewayWriteLocations != null) { + Utils.ValueHolder> writeValueHolder = Utils.ValueHolder.initialize(nextLocationInfo.availableWriteLocations); + Utils.ValueHolder> outWriteRegionMap = Utils.ValueHolder.initialize(nextLocationInfo.regionNameByWriteEndpoint); + nextLocationInfo.availableWriteEndpointsByLocation = this.getEndpointsByLocation(gatewayWriteLocations, writeValueHolder, outWriteRegionMap); + nextLocationInfo.availableWriteLocations = writeValueHolder.v; nextLocationInfo.regionNameByWriteEndpoint = outWriteRegionMap.v; } - nextLocationInfo.writeEndpoints = this.getPreferredAvailableEndpoints(nextLocationInfo.availableWriteEndpointByLocation, nextLocationInfo.availableWriteLocations, OperationType.Write, this.defaultEndpoint); - nextLocationInfo.readEndpoints = this.getPreferredAvailableEndpoints(nextLocationInfo.availableReadEndpointByLocation, nextLocationInfo.availableReadLocations, OperationType.Read, nextLocationInfo.writeEndpoints.get(0)); + nextLocationInfo.writeEndpoints = this.getPreferredAvailableEndpoints(nextLocationInfo.availableWriteEndpointsByLocation, nextLocationInfo.availableWriteLocations, OperationType.Write, this.defaultRegionalRoutingContext); + nextLocationInfo.readEndpoints = this.getPreferredAvailableEndpoints(nextLocationInfo.availableReadEndpointsByLocation, nextLocationInfo.availableReadLocations, OperationType.Read, nextLocationInfo.writeEndpoints.get(0)); if (nextLocationInfo.preferredLocations == null || nextLocationInfo.preferredLocations.isEmpty()) { Utils.ValueHolder regionForDefaultEndpoint = new Utils.ValueHolder<>(); // only set effective preferred locations when default endpoint doesn't map to a regional endpoint - if (!Utils.tryGetValue(nextLocationInfo.regionNameByReadEndpoint, this.defaultEndpoint, regionForDefaultEndpoint)) { + if (!Utils.tryGetValue(nextLocationInfo.regionNameByReadEndpoint, this.defaultRegionalRoutingContext, regionForDefaultEndpoint)) { nextLocationInfo.effectivePreferredLocations = nextLocationInfo.availableReadLocations; } } @@ -598,16 +607,16 @@ private void updateLocationCache( } } - private UnmodifiableList getPreferredAvailableEndpoints(UnmodifiableMap endpointsByLocation, - UnmodifiableList orderedLocations, - OperationType expectedAvailableOperation, - URI fallbackEndpoint) { - List endpoints = new ArrayList<>(); + private UnmodifiableList getPreferredAvailableEndpoints(UnmodifiableMap endpointsByLocation, + UnmodifiableList orderedLocations, + OperationType expectedAvailableOperation, + RegionalRoutingContext fallbackRegionalRoutingContext) { + List endpoints = new ArrayList<>(); DatabaseAccountLocationsInfo currentLocationInfo = this.locationInfo; // if enableEndpointDiscovery is false, we always use the defaultEndpoint that user passed in during documentClient init if (this.enableEndpointDiscovery) { if (this.canUseMultipleWriteLocations() || expectedAvailableOperation.supports(OperationType.Read)) { - List unavailableEndpoints = new ArrayList<>(); + List unavailableEndpoints = new ArrayList<>(); // When client can not use multiple write locations, preferred locations list should only be used // determining read endpoints order. @@ -616,7 +625,7 @@ private UnmodifiableList getPreferredAvailableEndpoints(UnmodifiableMap endpoint = new Utils.ValueHolder<>(); + Utils.ValueHolder endpoint = new Utils.ValueHolder<>(); if (Utils.tryGetValue(endpointsByLocation, location, endpoint)) { if (this.isEndpointUnavailable(endpoint.v, expectedAvailableOperation)) { unavailableEndpoints.add(endpoint.v); @@ -627,12 +636,12 @@ private UnmodifiableList getPreferredAvailableEndpoints(UnmodifiableMap endpoint = Utils.ValueHolder.initialize(null); + Utils.ValueHolder endpoint = Utils.ValueHolder.initialize(null); if (Utils.tryGetValue(endpointsByLocation, location, endpoint)) { // if defaultEndpoint equals a regional endpoint then use // whatever the fallback endpoint is - if (this.defaultEndpoint.equals(endpoint.v)) { + if (this.defaultRegionalRoutingContext.getGatewayRegionalEndpoint().equals(endpoint.v.getGatewayRegionalEndpoint())) { endpoints = new ArrayList<>(); break; } @@ -647,14 +656,13 @@ private UnmodifiableList getPreferredAvailableEndpoints(UnmodifiableMap endpoint = Utils.ValueHolder.initialize(null); + Utils.ValueHolder endpoint = Utils.ValueHolder.initialize(null); if (!Strings.isNullOrEmpty(location) && // location is empty during manual failover Utils.tryGetValue(endpointsByLocation, location, endpoint)) { endpoints.add(endpoint.v); @@ -664,40 +672,61 @@ private UnmodifiableList getPreferredAvailableEndpoints(UnmodifiableMap(endpoints); + return new UnmodifiableList<>(endpoints); } + private void addEndpoints( + Iterable gatewayDbAccountLocations, + Map endpointsByLocation, + Map regionByEndpoint, + List parsedLocations) { + if (gatewayDbAccountLocations != null) { + for (DatabaseAccountLocation gatewayDbAccountLocation : gatewayDbAccountLocations) { + if (!Strings.isNullOrEmpty(gatewayDbAccountLocation.getName())) { + try { - private UnmodifiableMap getEndpointByLocation(Iterable locations, - Utils.ValueHolder> orderedLocations, - Utils.ValueHolder> regionMap) { - Map endpointsByLocation = new CaseInsensitiveMap<>(); - Map regionByEndpoint = new CaseInsensitiveMap<>(); - List parsedLocations = new ArrayList<>(); + String location = gatewayDbAccountLocation.getName().toLowerCase(Locale.ROOT); + URI endpoint = new URI(gatewayDbAccountLocation.getEndpoint().toLowerCase(Locale.ROOT)); + + RegionalRoutingContext regionalRoutingContext = new RegionalRoutingContext(endpoint); + + if (!endpointsByLocation.containsKey(location)) { + endpointsByLocation.put(location, regionalRoutingContext); + } - for (DatabaseAccountLocation location: locations) { - if (!Strings.isNullOrEmpty(location.getName())) { - try { - URI endpoint = new URI(location.getEndpoint().toLowerCase(Locale.ROOT)); - endpointsByLocation.put(location.getName().toLowerCase(Locale.ROOT), endpoint); - regionByEndpoint.put(endpoint, location.getName().toLowerCase(Locale.ROOT)); - parsedLocations.add(location.getName()); - } catch (Exception e) { - logger.warn("GetAvailableEndpointsByLocation() - skipping add for location = [{}] as it is location name is either empty or endpoint is malformed [{}]", - location.getName(), - location.getEndpoint()); + if (!regionByEndpoint.containsKey(regionalRoutingContext)) { + regionByEndpoint.put(regionalRoutingContext, location); + } + + parsedLocations.add(gatewayDbAccountLocation.getName()); + } catch (Exception e) { + + logger.warn("Skipping add for location = [{}] and endpoint = [{}] due to exception [{}]", + gatewayDbAccountLocation.getName(), + gatewayDbAccountLocation.getEndpoint(), + e.getMessage()); + } } } } + } + + private UnmodifiableMap getEndpointsByLocation(Iterable gatewayLocations, + Utils.ValueHolder> orderedLocations, + Utils.ValueHolder> regionMap) { + Map endpointsByLocation = new CaseInsensitiveMap<>(); + Map regionByEndpoint = new CaseInsensitiveMap<>(); + List parsedLocations = new ArrayList<>(); + + addEndpoints(gatewayLocations, endpointsByLocation, regionByEndpoint, parsedLocations); - orderedLocations.v = new UnmodifiableList(parsedLocations); - regionMap.v = (UnmodifiableMap) UnmodifiableMap.unmodifiableMap(regionByEndpoint); + orderedLocations.v = new UnmodifiableList<>(parsedLocations); + regionMap.v = (UnmodifiableMap) UnmodifiableMap.unmodifiableMap(regionByEndpoint); - return (UnmodifiableMap) UnmodifiableMap.unmodifiableMap(endpointsByLocation); + return (UnmodifiableMap) UnmodifiableMap.unmodifiableMap(endpointsByLocation); } public boolean canUseMultipleWriteLocations() { @@ -773,35 +802,35 @@ private static boolean isExcludedRegionsSupplierConfigured(Supplier writeEndpoints; + private UnmodifiableList readEndpoints; private UnmodifiableList preferredLocations; private UnmodifiableList effectivePreferredLocations; // lower-case region private UnmodifiableList availableWriteLocations; // lower-case region private UnmodifiableList availableReadLocations; - private UnmodifiableMap availableWriteEndpointByLocation; - private UnmodifiableMap availableReadEndpointByLocation; - private UnmodifiableMap regionNameByWriteEndpoint; - private UnmodifiableMap regionNameByReadEndpoint; - private UnmodifiableList writeEndpoints; - private UnmodifiableList readEndpoints; + private UnmodifiableMap availableWriteEndpointsByLocation; + private UnmodifiableMap availableReadEndpointsByLocation; + private UnmodifiableMap regionNameByWriteEndpoint; + private UnmodifiableMap regionNameByReadEndpoint; public DatabaseAccountLocationsInfo(List preferredLocations, URI defaultEndpoint) { this.preferredLocations = new UnmodifiableList<>(preferredLocations.stream().map(loc -> loc.toLowerCase(Locale.ROOT)).collect(Collectors.toList())); this.effectivePreferredLocations = new UnmodifiableList<>(Collections.emptyList()); - this.availableWriteEndpointByLocation - = (UnmodifiableMap) UnmodifiableMap.unmodifiableMap(new CaseInsensitiveMap<>()); - this.availableReadEndpointByLocation - = (UnmodifiableMap) UnmodifiableMap.unmodifiableMap(new CaseInsensitiveMap<>()); + this.availableWriteEndpointsByLocation + = (UnmodifiableMap) UnmodifiableMap.unmodifiableMap(new CaseInsensitiveMap<>()); + this.availableReadEndpointsByLocation + = (UnmodifiableMap) UnmodifiableMap.unmodifiableMap(new CaseInsensitiveMap<>()); this.regionNameByWriteEndpoint - = (UnmodifiableMap) UnmodifiableMap.unmodifiableMap(new CaseInsensitiveMap<>()); + = (UnmodifiableMap) UnmodifiableMap.unmodifiableMap(new CaseInsensitiveMap<>()); this.regionNameByReadEndpoint - = (UnmodifiableMap) UnmodifiableMap.unmodifiableMap(new CaseInsensitiveMap<>()); + = (UnmodifiableMap) UnmodifiableMap.unmodifiableMap(new CaseInsensitiveMap<>()); this.availableReadLocations = new UnmodifiableList<>(Collections.emptyList()); this.availableWriteLocations = new UnmodifiableList<>(Collections.emptyList()); - this.readEndpoints = new UnmodifiableList<>(Collections.singletonList(defaultEndpoint)); - this.writeEndpoints = new UnmodifiableList<>(Collections.singletonList(defaultEndpoint)); + this.readEndpoints = new UnmodifiableList<>(Collections.singletonList(new RegionalRoutingContext(defaultEndpoint))); + this.writeEndpoints = new UnmodifiableList<>(Collections.singletonList(new RegionalRoutingContext(defaultEndpoint))); } public DatabaseAccountLocationsInfo(DatabaseAccountLocationsInfo other) { @@ -809,10 +838,10 @@ public DatabaseAccountLocationsInfo(DatabaseAccountLocationsInfo other) { this.effectivePreferredLocations = other.effectivePreferredLocations; this.availableWriteLocations = other.availableWriteLocations; this.availableReadLocations = other.availableReadLocations; - this.availableWriteEndpointByLocation = other.availableWriteEndpointByLocation; + this.availableWriteEndpointsByLocation = other.availableWriteEndpointsByLocation; this.regionNameByWriteEndpoint = other.regionNameByWriteEndpoint; this.regionNameByReadEndpoint = other.regionNameByReadEndpoint; - this.availableReadEndpointByLocation = other.availableReadEndpointByLocation; + this.availableReadEndpointsByLocation = other.availableReadEndpointsByLocation; this.writeEndpoints = other.writeEndpoints; this.readEndpoints = other.readEndpoints; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/routing/RegionalRoutingContext.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/routing/RegionalRoutingContext.java new file mode 100644 index 000000000000..6a796bb9866a --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/routing/RegionalRoutingContext.java @@ -0,0 +1,38 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.routing; + +import java.net.URI; +import java.util.Objects; + +public class RegionalRoutingContext { + + // IMPORTANT: + // Please reevaluate overridden equals() implementation + // when adding additional properties to this class + private final URI gatewayRegionalEndpoint; + private final String gatewayRegionalEndpointAsString; + + public RegionalRoutingContext(URI gatewayRegionalEndpoint) { + this.gatewayRegionalEndpoint = gatewayRegionalEndpoint; + this.gatewayRegionalEndpointAsString = gatewayRegionalEndpoint.toString(); + } + + public URI getGatewayRegionalEndpoint() { + return this.gatewayRegionalEndpoint; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + RegionalRoutingContext that = (RegionalRoutingContext) o; + return this.gatewayRegionalEndpoint.equals(that.gatewayRegionalEndpoint); + } + + @Override + public int hashCode() { + return Objects.hash(this.gatewayRegionalEndpointAsString); + } +}