From 7e80816851f0efa3c5ea4494e7ca2087812ea7cb Mon Sep 17 00:00:00 2001 From: Neha Rao Date: Tue, 25 Feb 2025 10:50:05 -0800 Subject: [PATCH 01/17] Add RegionalEndpoints --- .../FaultInjectionRuleProcessor.java | 2 + ...EndpointManagerForCircuitBreakerTests.java | 146 ++++----- .../PartitionLevelCircuitBreakerTests.java | 6 +- .../ProactiveConnectionManagementTest.java | 29 +- .../implementation/ClientRetryPolicyTest.java | 19 +- .../cosmos/implementation/ConfigsTests.java | 2 - .../RegionScopedSessionContainerTest.java | 154 ++++++++-- ...eCollectionAwareClientRetryPolicyTest.java | 27 +- .../RxGatewayStoreModelTest.java | 9 +- .../WebExceptionRetryPolicyTest.java | 11 +- .../GlobalAddressResolverTest.java | 31 +- .../GlobalEndPointManagerTest.java | 26 +- .../MetadataRequestRetryPolicyTests.java | 6 +- .../query/DocumentProducerTest.java | 7 +- .../routing/LocationCacheTest.java | 45 +-- .../com/azure/cosmos/CosmosDiagnostics.java | 5 +- .../implementation/ClientRetryPolicy.java | 40 ++- .../ClientSideRequestStatistics.java | 22 +- .../azure/cosmos/implementation/Configs.java | 1 + .../DocumentServiceRequestContext.java | 9 +- .../implementation/GlobalEndpointManager.java | 19 +- .../ImplementationBridgeHelpers.java | 4 +- .../MetadataRequestRetryPolicy.java | 4 +- .../RegionScopedSessionContainer.java | 2 +- .../ResourceThrottleRetryPolicy.java | 2 +- .../implementation/RxDocumentClientImpl.java | 15 +- .../RxDocumentServiceRequest.java | 1 + .../implementation/RxGatewayStoreModel.java | 6 +- .../implementation/ShouldRetryResult.java | 4 + ...itionEndpointManagerForCircuitBreaker.java | 50 ++-- .../GlobalAddressResolver.java | 35 +-- .../rntbd/RntbdOpenConnectionsHandler.java | 2 + .../FaultInjectionRequestContext.java | 4 + .../cosmos/implementation/query/Fetcher.java | 4 +- .../implementation/routing/LocationCache.java | 283 +++++++++++------- 35 files changed, 641 insertions(+), 391 deletions(-) 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..654607d6fb66 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 @@ -23,6 +23,7 @@ import com.azure.cosmos.implementation.directconnectivity.AddressSelector; import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdUtils; import com.azure.cosmos.implementation.feedranges.FeedRangeInternal; +import com.azure.cosmos.implementation.routing.LocationCache; import com.azure.cosmos.implementation.routing.PartitionKeyRangeIdentity; import com.azure.cosmos.test.faultinjection.FaultInjectionCondition; import com.azure.cosmos.test.faultinjection.FaultInjectionConnectionErrorResult; @@ -447,6 +448,7 @@ private Mono> resolvePhysicalAddresses( null); faultInjectionAddressRequest.requestContext.locationEndpointToRoute = regionEndpoint; + faultInjectionAddressRequest.requestContext.regionalEndpointsToRoute = new LocationCache.RegionalEndpoints(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..c7d12f277f1c 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.LocationCache; 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 LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getKey())); } - locationSpecificHealthContext = locationEndpointToLocationSpecificContextForPartition.get(LocationEastUs2EndpointToLocationPair.getKey()); + locationSpecificHealthContext = locationEndpointToLocationSpecificContextForPartition.get(new LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getKey())); globalPartitionEndpointManagerForCircuitBreaker - .handleLocationExceptionForPartitionKeyRange(request, LocationEastUsEndpointToLocationPair.getKey()); + .handleLocationExceptionForPartitionKeyRange(request, new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getKey())); globalPartitionEndpointManagerForCircuitBreaker - .handleLocationExceptionForPartitionKeyRange(request, LocationCentralUsEndpointToLocationPair.getKey()); + .handleLocationExceptionForPartitionKeyRange(request, new LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getKey())); LocationSpecificHealthContext locationSpecificHealthContext2 - = locationEndpointToLocationSpecificContextForPartitionForColl2.get(LocationEastUs2EndpointToLocationPair.getKey()); + = locationEndpointToLocationSpecificContextForPartitionForColl2.get(new LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(locationWithFailure)); List unavailableRegions = globalPartitionEndpointManagerForCircuitBreaker.getUnavailableRegionsForPartitionKeyRange(collectionResourceId, partitionKeyRange, request.getOperationType()); @@ -915,6 +915,8 @@ private RxDocumentServiceRequest constructRxDocumentServiceRequestInstance( request.requestContext.resolvedPartitionKeyRange = new PartitionKeyRange(partitionKeyRangeId, minInclusive, maxExclusive); request.requestContext.resolvedPartitionKeyRangeForCircuitBreaker = request.requestContext.resolvedPartitionKeyRange; request.requestContext.locationEndpointToRoute = locationEndpointToRoute; + request.requestContext.regionalEndpointsToRoute = new LocationCache.RegionalEndpoints(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..b6de3bcb21c9 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.LocationCache; 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..9b6e428bf47a 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 @@ -27,6 +27,7 @@ import com.azure.cosmos.implementation.directconnectivity.RntbdTransportClient; import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdEndpoint; import com.azure.cosmos.implementation.routing.CollectionRoutingMap; +import com.azure.cosmos.implementation.routing.LocationCache; import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper; import com.azure.cosmos.implementation.routing.PartitionKeyRangeIdentity; import com.azure.cosmos.models.CosmosContainerIdentity; @@ -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(LocationCache.RegionalEndpoints::getGatewayLocationEndpoint).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(LocationCache.RegionalEndpoints::getGatewayLocationEndpoint) + .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(LocationCache.RegionalEndpoints::getGatewayLocationEndpoint) + .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(LocationCache.RegionalEndpoints::getGatewayLocationEndpoint) + .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 2e4ae2c9f956..3a452ada1fbe 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.LocationCache; import io.netty.handler.timeout.ReadTimeoutException; import io.reactivex.subscribers.TestSubscriber; import org.mockito.Mockito; @@ -65,7 +66,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 LocationCache.RegionalEndpoints(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); @@ -106,7 +107,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 LocationCache.RegionalEndpoints(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( @@ -149,7 +150,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 LocationCache.RegionalEndpoints(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); @@ -197,7 +198,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 LocationCache.RegionalEndpoints(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); @@ -232,7 +233,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 LocationCache.RegionalEndpoints(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); @@ -292,7 +293,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 LocationCache.RegionalEndpoints(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); @@ -325,7 +326,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 LocationCache.RegionalEndpoints(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); @@ -361,7 +362,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 LocationCache.RegionalEndpoints(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); @@ -395,7 +396,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 LocationCache.RegionalEndpoints(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 f071799b08cc..6d21a79bcf18 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.util.EnumSet; 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..fe06e0ad6002 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.LocationCache; 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 LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(LocationEastUsEndpointToLocationPair.getLeft()), Mockito.any())).thenReturn(regionContacted); @@ -434,7 +439,11 @@ 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 LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(endpointContacted), Mockito.any())).thenReturn(regionContacted); @@ -497,7 +506,11 @@ 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 LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); @@ -550,12 +563,15 @@ 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 LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); @@ -583,7 +599,11 @@ 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 LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); @@ -625,7 +645,11 @@ 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 LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); @@ -669,7 +693,11 @@ 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 LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); @@ -709,7 +737,11 @@ 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 LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); @@ -754,7 +786,11 @@ 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 LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); @@ -795,7 +831,11 @@ 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 LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); @@ -834,7 +874,11 @@ 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 LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); @@ -876,7 +920,11 @@ 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 LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); @@ -935,7 +983,11 @@ 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 LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); @@ -995,7 +1047,11 @@ 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 LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); @@ -1050,7 +1106,11 @@ 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 LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); @@ -1088,7 +1148,11 @@ 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 LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); @@ -1131,7 +1195,11 @@ 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 LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); @@ -1181,7 +1249,11 @@ 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 LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(LocationEastUsEndpointToLocationPair.getLeft()), Mockito.any())).thenReturn(regionContacted); @@ -1245,7 +1317,11 @@ 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 LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); @@ -1289,7 +1365,11 @@ 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 LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); @@ -1332,7 +1412,11 @@ 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 LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); @@ -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 LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); @@ -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 LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), + new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); Mockito.when(globalEndpointManagerMock.getReadEndpoints()).thenReturn(endpoints); Mockito.when(globalEndpointManagerMock.getRegionName(Mockito.eq(locationEndpointContacted), Mockito.any())).thenReturn(regionContacted); @@ -1567,9 +1659,14 @@ public void resolvePartitionLocalSessionToken( GlobalEndpointManager globalEndpointManagerMock = null; RegionScopedSessionContainer sessionContainer = null; - List writableURIs = writableURIToLocationMappings + List consolidatedWriteRegionalEndpoints = writableURIToLocationMappings .stream() - .map(uriToLocationMappings -> uriToLocationMappings.getLeft()) + .map(uriToLocationMappings -> new LocationCache.RegionalEndpoints(uriToLocationMappings.getLeft())) + .collect(Collectors.toList()); + + List consolidatedReadRegionalEndpoints = readEndpoints + .stream() + .map(readEndpoint -> new LocationCache.RegionalEndpoints(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<>(consolidatedReadRegionalEndpoints); Mockito .when(globalEndpointManagerMock.getReadEndpoints()) @@ -1592,7 +1689,7 @@ public void resolvePartitionLocalSessionToken( Mockito .when(globalEndpointManagerMock.getApplicableWriteEndpoints(Mockito.anyList())) - .thenReturn(new UnmodifiableList<>(writableURIs)); + .thenReturn(new UnmodifiableList<>(consolidatedWriteRegionalEndpoints)); Mockito .when(globalEndpointManagerMock.canUseMultipleWriteLocations(Mockito.any())) @@ -1677,6 +1774,7 @@ private static RxDocumentServiceRequest createRequestEntity(OperationType operat resourceType); request.requestContext.locationEndpointToRoute = locationEndpointToRoute; + request.requestContext.regionalEndpointsToRoute = new LocationCache.RegionalEndpoints(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..05c402a7c336 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.LocationCache; 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"); + LocationCache.RegionalEndpoints consolidatedLocationEndpointToRoute = new LocationCache.RegionalEndpoints(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"); + LocationCache.RegionalEndpoints consolidatedLocationEndpointToRoute = new LocationCache.RegionalEndpoints(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"); + LocationCache.RegionalEndpoints consolidatedLocationEndpointToRoute = new LocationCache.RegionalEndpoints(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 0ea494c8dbc6..34059324da4e 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,6 +11,7 @@ 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.LocationCache; import io.netty.channel.ConnectTimeoutException; import io.netty.handler.codec.http.HttpMethod; import io.netty.handler.timeout.ReadTimeoutException; @@ -82,7 +83,7 @@ public void readTimeout() throws Exception { GlobalEndpointManager globalEndpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new URI("https://localhost")) + Mockito.doReturn(new LocationCache.RegionalEndpoints(new URI("https://localhost"))) .when(globalEndpointManager).resolveServiceEndpoint(any()); HttpClient httpClient = Mockito.mock(HttpClient.class); Mockito.doReturn(Mono.error(ReadTimeoutException.INSTANCE)) @@ -125,7 +126,7 @@ 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")) + Mockito.doReturn(new LocationCache.RegionalEndpoints(new URI("https://localhost"))) .when(globalEndpointManager).resolveServiceEndpoint(any()); HttpClient httpClient = Mockito.mock(HttpClient.class); Mockito.doReturn(Mono.error(new SocketException("Dummy SocketException"))) @@ -179,7 +180,7 @@ public void applySessionToken( GlobalEndpointManager globalEndpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new URI("https://localhost")) + Mockito.doReturn(new LocationCache.RegionalEndpoints(new URI("https://localhost"))) .when(globalEndpointManager).resolveServiceEndpoint(any()); HttpClient httpClient = Mockito.mock(HttpClient.class); @@ -250,7 +251,7 @@ public void validateApiType() throws Exception { GlobalEndpointManager globalEndpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new URI("https://localhost")) + Mockito.doReturn(new LocationCache.RegionalEndpoints(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/WebExceptionRetryPolicyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/WebExceptionRetryPolicyTest.java index 0a72963fe795..9481d4263f0b 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.LocationCache; import io.netty.handler.timeout.ReadTimeoutException; import io.reactivex.subscribers.TestSubscriber; import org.mockito.Mockito; @@ -37,7 +38,7 @@ 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)); + Mockito.doReturn(new LocationCache.RegionalEndpoints(new URI("http://localhost:"))).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(Mockito.eq(null), Mockito.eq(true)); RetryContext retryContext = new RetryContext(); @@ -89,7 +90,7 @@ 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)); + Mockito.doReturn(new LocationCache.RegionalEndpoints(new URI("http://localhost:"))).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(Mockito.eq(null), Mockito.eq(true)); RetryContext retryContext = new RetryContext(); @@ -143,7 +144,7 @@ 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)); + Mockito.doReturn(new LocationCache.RegionalEndpoints(new URI("http://localhost:"))).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(Mockito.eq(null), Mockito.eq(true)); RetryContext retryContext = new RetryContext(); @@ -195,7 +196,7 @@ 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)); + Mockito.doReturn(new LocationCache.RegionalEndpoints(new URI("http://localhost:"))).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(Mockito.eq(null), Mockito.eq(true)); @@ -236,7 +237,7 @@ 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)); + Mockito.doReturn(new LocationCache.RegionalEndpoints(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(); 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..6a060e845f0e 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; @@ -25,6 +24,7 @@ import com.azure.cosmos.implementation.caches.RxCollectionCache; import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; import com.azure.cosmos.implementation.http.HttpClient; +import com.azure.cosmos.implementation.routing.LocationCache; import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper; import com.azure.cosmos.implementation.routing.PartitionKeyRangeIdentity; import com.azure.cosmos.models.CosmosContainerIdentity; @@ -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 LocationCache.RegionalEndpoints(urlforRead1)); + readEndPointList.add(new LocationCache.RegionalEndpoints(urlforRead2)); + readEndPointList.add(new LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(urlforWrite1)); + writeEndPointList.add(new LocationCache.RegionalEndpoints(urlforWrite2)); + writeEndPointList.add(new LocationCache.RegionalEndpoints(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/"); + LocationCache.RegionalEndpoints testUrl = new LocationCache.RegionalEndpoints(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.getGatewayLocationEndpoint())).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 LocationCache.RegionalEndpoints(urlforRead1), new LocationCache.RegionalEndpoints(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..918d7d671280 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 @@ -99,7 +99,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 +155,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 +195,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 +261,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 +290,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 +353,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..2bbff8cfa59e 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.LocationCache; 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).getGatewayLocationEndpoint(); GlobalAddressResolver.EndpointCache endpointCache = endpointCacheByURIMap.get(readEndpoint); GatewayAddressCache gatewayAddressCache = endpointCache.addressCache; HttpClientUnderTestWrapper httpClientUnderTestWrapper = getHttpClientUnderTestWrapper(configs); @@ -556,6 +557,7 @@ private static RxDocumentServiceRequest createRequest( if (hasLocationEndpointToRoute) { request.requestContext.locationEndpointToRoute = URI.create("https://account-name-some-region.documents.azure.com:443"); + request.requestContext.regionalEndpointsToRoute = new LocationCache.RegionalEndpoints(request.requestContext.locationEndpointToRoute); } if (isAddressRefresh) { 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 eb05b05b5df9..e526dd4dbc9c 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 @@ -28,6 +28,7 @@ import com.azure.cosmos.implementation.guava25.collect.LinkedListMultimap; import com.azure.cosmos.implementation.query.orderbyquery.OrderByRowResult; import com.azure.cosmos.implementation.query.orderbyquery.OrderbyRowComparer; +import com.azure.cosmos.implementation.routing.LocationCache; import com.azure.cosmos.implementation.routing.PartitionKeyRangeIdentity; import com.azure.cosmos.implementation.routing.Range; import com.azure.cosmos.models.FeedResponse; @@ -112,9 +113,9 @@ public Object[][] mergeParamProvider() { } private IRetryPolicyFactory mockDocumentClientIRetryPolicyFactory() { - URI url; + LocationCache.RegionalEndpoints regionalEndpoints; try { - url = new URI("http://localhost"); + regionalEndpoints = new LocationCache.RegionalEndpoints(new URI("http://localhost")); } catch (Exception e) { throw new IllegalStateException(e); } @@ -122,7 +123,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(regionalEndpoints).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..313f6b005062 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 @@ -21,6 +21,8 @@ import com.azure.cosmos.models.ModelBridgeUtils; import com.azure.cosmos.implementation.guava25.collect.ImmutableList; import com.azure.cosmos.implementation.guava25.collect.Iterables; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import org.testng.annotations.AfterClass; import org.testng.annotations.DataProvider; import org.testng.annotations.Test; @@ -35,7 +37,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; @@ -63,6 +64,8 @@ public class LocationCacheTest { private static HashMap EndpointByLocation = new HashMap<>(); + private final static Logger logger = LoggerFactory.getLogger(LocationCache.class); + static { EndpointByLocation.put("location1", LocationCacheTest.Location1Endpoint); EndpointByLocation.put("location2", LocationCacheTest.Location2Endpoint); @@ -480,9 +483,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 LocationCache.RegionalEndpoints(LocationCacheTest.Location1Endpoint)); + assertThat(this.cache.getWriteEndpoints().get(1)).isEqualTo(new LocationCache.RegionalEndpoints(LocationCacheTest.Location2Endpoint)); + assertThat(this.cache.getWriteEndpoints().get(2)).isEqualTo(new LocationCache.RegionalEndpoints(LocationCacheTest.Location3Endpoint)); } @Test(groups = "unit", dataProvider = "excludedRegionsTestConfigs") @@ -512,11 +515,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 +535,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 +670,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 +778,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 +901,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 LocationCache.RegionalEndpoints(firstAvailableWriteEndpoint)).isEqualTo(writeEndpoints.get(0)); + assertThat(new LocationCache.RegionalEndpoints(secondAvailableWriteEndpoint)).isEqualTo(this.resolveEndpointForWriteRequest(ResourceType.Document, true)); + assertThat(new LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(firstWriteEnpoint)).isEqualTo(this.resolveEndpointForWriteRequest(ResourceType.Database, false)); + assertThat(new LocationCache.RegionalEndpoints(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 LocationCache.RegionalEndpoints(firstAvailableReadEndpoint)).isEqualTo(this.resolveEndpointForReadRequest(true)); + assertThat(new LocationCache.RegionalEndpoints(firstAvailableReadEndpoint)).isEqualTo(this.resolveEndpointForReadRequest(false)); } - private URI resolveEndpointForReadRequest(boolean masterResourceType) { + private LocationCache.RegionalEndpoints 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 LocationCache.RegionalEndpoints 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/src/main/java/com/azure/cosmos/CosmosDiagnostics.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosDiagnostics.java index 79d11aafbc3a..661b4f11c194 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.LocationCache; 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() { + LocationCache.RegionalEndpoints getFirstContactedLocationEndpoint() { return this.clientSideRequestStatistics.getFirstContactedLocationEndpoint(); } @@ -478,7 +479,7 @@ public void setDiagnosticsContext(CosmosDiagnostics cosmosDiagnostics, CosmosDia } @Override - public URI getFirstContactedLocationEndpoint(CosmosDiagnostics cosmosDiagnostics) { + public LocationCache.RegionalEndpoints 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 329c9f1a8f05..b4f2d81eda81 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,12 +12,14 @@ 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.LocationCache; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Mono; import java.net.URI; import java.time.Duration; +import java.util.Objects; import java.util.concurrent.atomic.AtomicInteger; import static com.azure.cosmos.implementation.HttpConstants.HttpHeaders.INTENDED_COLLECTION_RID_HEADER; @@ -46,6 +48,7 @@ public class ClientRetryPolicy extends DocumentClientRetryPolicy { private boolean isReadRequest; private boolean canUseMultipleWriteLocations; private URI locationEndpoint; + private LocationCache.RegionalEndpoints regionalEndpoints; private RetryContext retryContext; private CosmosDiagnostics cosmosDiagnostics; private AtomicInteger cnt = new AtomicInteger(0); @@ -86,7 +89,7 @@ public Mono shouldRetry(Exception e) { isReadRequest, canUseMultipleWriteLocations, e); - if (this.locationEndpoint == null) { + if (this.locationEndpoint == null || this.regionalEndpoints == null) { // on before request is not invoked because Document Service Request creation failed. logger.error("locationEndpoint is null because ClientRetryPolicy::onBeforeRequest(.) is not invoked, " + "probably request creation failed due to invalid options, serialization setting, etc."); @@ -228,7 +231,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 +307,7 @@ private Mono shouldRetryOnGatewayTimeout() { boolean canFailoverOnTimeout = canGatewayRequestFailoverOnTimeout(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.regionalEndpointsToRoute); } //if operation is data plane read, metadata read, or query plan it can be retried on a different endpoint. @@ -338,10 +341,10 @@ private Mono refreshLocation(boolean isReadRequest, boolean forceRefresh, // Mark the current read endpoint as unavailable if (isReadRequest) { logger.warn("marking the endpoint {} as unavailable for read",this.locationEndpoint); - this.globalEndpointManager.markEndpointUnavailableForRead(this.locationEndpoint); + this.globalEndpointManager.markEndpointUnavailableForRead(this.regionalEndpoints.getGatewayLocationEndpoint()); } else { logger.warn("marking the endpoint {} as unavailable for write",this.locationEndpoint); - this.globalEndpointManager.markEndpointUnavailableForWrite(this.locationEndpoint); + this.globalEndpointManager.markEndpointUnavailableForWrite(this.regionalEndpoints.getGatewayLocationEndpoint()); } this.retryContext = new RetryContext(this.failoverRetryCount, usePreferredLocations); @@ -356,7 +359,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.regionalEndpointsToRoute); } // The request has failed with 503, SDK need to decide whether it is safe to retry for write operations @@ -415,9 +418,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.regionalEndpointsToRoute); } } @@ -427,9 +431,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.regionalEndpointsToRoute); } return Mono.just(ShouldRetryResult.NO_RETRY); @@ -458,9 +463,11 @@ 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.regionalEndpoints = this.globalEndpointManager.resolveServiceEndpoint(request); + this.locationEndpoint = this.regionalEndpoints.getGatewayLocationEndpoint(); + if (request.requestContext != null) { - request.requestContext.routeToLocation(this.locationEndpoint); + request.requestContext.routeToLocation(this.locationEndpoint, this.regionalEndpoints); } } @@ -506,4 +513,13 @@ public RetryContext(int retryCount, this.retryRequestOnPreferredLocations = retryRequestOnPreferredLocations; } } + + private URI getGatewayLocationEndpoint(RxDocumentServiceRequest request) { + + Objects.requireNonNull(request, "Argument 'request' must not be null'"); + Objects.requireNonNull(request.requestContext, "Argument 'request.requestContext' must not be null'"); + Objects.requireNonNull(request.requestContext.regionalEndpointsToRoute, "Argument 'request.requestContext.consolidatedRegionalEndpointToRoute' must not be null'"); + + return request.requestContext.regionalEndpointsToRoute.getGatewayLocationEndpoint(); + } } 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..4fd366ce9b4c 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.LocationCache; import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.core.JsonGenerator; @@ -163,7 +164,9 @@ public void recordResponse(RxDocumentServiceRequest request, StoreResultDiagnost this.requestPayloadSizeInBytes = 0; } + LocationCache.RegionalEndpoints regionalEndpoints = null; URI locationEndPoint = null; + if (request.requestContext != null) { this.approximateInsertionCountInBloomFilter = request.requestContext.getApproximateBloomFilterInsertionCount(); @@ -176,6 +179,7 @@ public void recordResponse(RxDocumentServiceRequest request, StoreResultDiagnost } locationEndPoint = request.requestContext.locationEndpointToRoute; + regionalEndpoints = request.requestContext.regionalEndpointsToRoute; List excludedRegions = request.requestContext.getExcludeRegions(); if (excludedRegions != null && !excludedRegions.isEmpty()) { @@ -189,12 +193,12 @@ public void recordResponse(RxDocumentServiceRequest request, StoreResultDiagnost this.requestEndTimeUTC = responseTime; } - if (locationEndPoint != null) { + if (regionalEndpoints != null) { storeResponseStatistics.regionName = - globalEndpointManager.getRegionName(locationEndPoint, request.getOperationType()); + globalEndpointManager.getRegionName(regionalEndpoints.getGatewayLocationEndpoint(), 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, regionalEndpoints)); } if (storeResponseStatistics.requestOperationType == OperationType.Head @@ -218,9 +222,13 @@ public void recordGatewayResponse( this.requestEndTimeUTC = responseTime; } + LocationCache.RegionalEndpoints regionalEndpoints = null; URI locationEndPoint = null; if (rxDocumentServiceRequest != null && rxDocumentServiceRequest.requestContext != null) { + + regionalEndpoints = rxDocumentServiceRequest.requestContext.regionalEndpointsToRoute; locationEndPoint = rxDocumentServiceRequest.requestContext.locationEndpointToRoute; + this.approximateInsertionCountInBloomFilter = rxDocumentServiceRequest.requestContext.getApproximateBloomFilterInsertionCount(); this.keywordIdentifiers = rxDocumentServiceRequest.requestContext.getKeywordIdentifiers(); } @@ -233,7 +241,7 @@ public void recordGatewayResponse( this.regionsContacted.add(regionName); this.locationEndpointsContacted.add(locationEndPoint); - this.regionsContactedWithContext.add(new RegionWithContext(regionName, locationEndPoint)); + this.regionsContactedWithContext.add(new RegionWithContext(regionName, regionalEndpoints)); } GatewayStatistics gatewayStatistics = new GatewayStatistics(); @@ -651,7 +659,7 @@ public String getFirstContactedRegion() { return this.regionsContactedWithContext.first().regionContacted; } - public URI getFirstContactedLocationEndpoint() { + public LocationCache.RegionalEndpoints getFirstContactedLocationEndpoint() { if (this.regionsContactedWithContext == null || this.regionsContactedWithContext.isEmpty()) { return null; } @@ -1048,10 +1056,10 @@ public static CosmosDiagnosticsSystemUsageSnapshot fetchSystemInformation() { static class RegionWithContext implements Comparable { private final String regionContacted; - private final URI locationEndpointsContacted; + private final LocationCache.RegionalEndpoints locationEndpointsContacted; private final long recordedTimestamp; - RegionWithContext(String regionContacted, URI locationEndpointsContacted) { + RegionWithContext(String regionContacted, LocationCache.RegionalEndpoints 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 7f07616919e0..69cd23ea0314 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 @@ -14,6 +14,7 @@ import org.slf4j.LoggerFactory; import javax.net.ssl.SSLException; +import java.net.URI; import java.time.Duration; import java.util.Locale; 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..0623fcf3bb68 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,7 +12,7 @@ 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.LocationCache; import com.azure.cosmos.implementation.routing.PartitionKeyInternal; import java.net.URI; @@ -40,6 +40,7 @@ public class DocumentServiceRequestContext implements Cloneable { public volatile Boolean usePreferredLocations; public volatile Integer locationIndexToRoute; public volatile URI locationEndpointToRoute; + public volatile LocationCache.RegionalEndpoints regionalEndpointsToRoute; public volatile boolean performedBackgroundAddressRefresh; public volatile boolean performLocalRefreshOnGoneException; public volatile List storeResponses; @@ -81,6 +82,7 @@ public void routeToLocation(int locationIndex, boolean usePreferredLocations) { this.locationIndexToRoute = locationIndex; this.usePreferredLocations = usePreferredLocations; this.locationEndpointToRoute = null; + this.regionalEndpointsToRoute = null; } /** @@ -89,8 +91,9 @@ public void routeToLocation(int locationIndex, boolean usePreferredLocations) { * * @param locationEndpoint Location endpoint to which the request should be routed. */ - public void routeToLocation(URI locationEndpoint) { + public void routeToLocation(URI locationEndpoint, LocationCache.RegionalEndpoints regionalEndpointsToRoute) { this.locationEndpointToRoute = locationEndpoint; + this.regionalEndpointsToRoute = regionalEndpointsToRoute; this.locationIndexToRoute = null; this.usePreferredLocations = null; } @@ -101,6 +104,7 @@ public void routeToLocation(URI locationEndpoint) { public void clearRouteToLocation() { this.locationIndexToRoute = null; this.locationEndpointToRoute = null; + this.regionalEndpointsToRoute = null; this.usePreferredLocations = null; } @@ -142,6 +146,7 @@ public DocumentServiceRequestContext clone() { context.usePreferredLocations = this.usePreferredLocations; context.locationIndexToRoute = this.locationIndexToRoute; context.locationEndpointToRoute = this.locationEndpointToRoute; + context.regionalEndpointsToRoute = this.regionalEndpointsToRoute; 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..e71eab9db96f 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 @@ -88,32 +88,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,10 +146,11 @@ public static Mono getDatabaseAccountFromAnyLocationsAsync( }); } - public URI resolveServiceEndpoint(RxDocumentServiceRequest request) { - URI serviceEndpoint = this.locationCache.resolveServiceEndpoint(request); + public LocationCache.RegionalEndpoints resolveServiceEndpoint(RxDocumentServiceRequest request) { + LocationCache.RegionalEndpoints serviceEndpoints = this.locationCache.resolveServiceEndpoint(request); if (request.faultInjectionRequestContext != null) { - request.faultInjectionRequestContext.setLocationEndpointToRoute(serviceEndpoint); + // TODO: integrate thin client into fault injection + request.faultInjectionRequestContext.setLocationEndpointToRoute(serviceEndpoints.getGatewayLocationEndpoint()); } return serviceEndpoint; 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 73f692adf149..a3697151adbc 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 @@ -40,6 +40,7 @@ import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdChannelStatistics; import com.azure.cosmos.implementation.faultinjection.IFaultInjectorProvider; import com.azure.cosmos.implementation.patch.PatchOperation; +import com.azure.cosmos.implementation.routing.LocationCache; import com.azure.cosmos.implementation.routing.PartitionKeyInternal; import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple; import com.azure.cosmos.models.CosmosBatch; @@ -82,7 +83,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; @@ -837,7 +837,7 @@ void recordAddressResolutionEnd( void setDiagnosticsContext(CosmosDiagnostics cosmosDiagnostics, CosmosDiagnosticsContext ctx); - URI getFirstContactedLocationEndpoint(CosmosDiagnostics cosmosDiagnostics); + LocationCache.RegionalEndpoints 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..f11853aeca51 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.LocationCache; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Mono; @@ -71,7 +72,8 @@ public Mono shouldRetry(Exception e) { if (request.requestContext != null && request.requestContext.locationEndpointToRoute != null) { - URI locationEndpointToRoute = request.requestContext.locationEndpointToRoute; + LocationCache.RegionalEndpoints regionalEndpoints = request.requestContext.regionalEndpointsToRoute; + URI locationEndpointToRoute = regionalEndpoints.getGatewayLocationEndpoint(); 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..18d9237674cd 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 @@ -525,7 +525,7 @@ private String extractFirstEffectivePreferredReadableRegion() { List regionNamesForRead = globalEndpointManager .getReadEndpoints() .stream() - .map(endpoint -> globalEndpointManager.getRegionName(endpoint, OperationType.Read)) + .map(consolidatedReadLocationEndpoints -> globalEndpointManager.getRegionName(consolidatedReadLocationEndpoints.getGatewayLocationEndpoint(), 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..22c68665ce2c 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.noRetryOnNonRelatedException(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 451bd1991227..fcb1f4057c22 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 @@ -56,6 +56,7 @@ import com.azure.cosmos.implementation.query.PipelinedQueryExecutionContextBase; import com.azure.cosmos.implementation.query.QueryInfo; import com.azure.cosmos.implementation.routing.CollectionRoutingMap; +import com.azure.cosmos.implementation.routing.LocationCache; import com.azure.cosmos.implementation.routing.PartitionKeyAndResourceTokenPair; import com.azure.cosmos.implementation.routing.PartitionKeyInternal; import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper; @@ -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 = @@ -6451,7 +6452,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()) { @@ -6461,7 +6462,7 @@ private List getApplicableEndPoints(OperationType operationType, List withoutNulls(List orderedEffectiveEndpointsList) { + private static List withoutNulls(List orderedEffectiveEndpointsList) { if (orderedEffectiveEndpointsList == null) { return EMPTY_ENDPOINT_LIST; } @@ -6520,7 +6521,7 @@ private List getApplicableRegionsForSpeculation( return EMPTY_REGION_LIST; } - List endpoints = getApplicableEndPoints(operationType, excludedRegions); + List regionalEndpointsList = getApplicableEndPoints(operationType, excludedRegions); HashSet normalizedExcludedRegions = new HashSet<>(); if (excludedRegions != null) { @@ -6528,8 +6529,8 @@ private List getApplicableRegionsForSpeculation( } List orderedRegionsForSpeculation = new ArrayList<>(); - endpoints.forEach(uri -> { - String regionName = this.globalEndpointManager.getRegionName(uri, operationType); + regionalEndpointsList.forEach(consolidatedLocationEndpoints -> { + String regionName = this.globalEndpointManager.getRegionName(consolidatedLocationEndpoints.getGatewayLocationEndpoint(), operationType); if (!normalizedExcludedRegions.contains(regionName.toLowerCase(Locale.ROOT))) { orderedRegionsForSpeculation.add(regionName); } @@ -6733,7 +6734,7 @@ private Mono executeFeedOperationWithAvailabilityStrategy( private void handleLocationCancellationExceptionForPartitionKeyRange(RxDocumentServiceRequest failedRequest) { - URI firstContactedLocationEndpoint = diagnosticsAccessor + LocationCache.RegionalEndpoints firstContactedLocationEndpoint = diagnosticsAccessor .getFirstContactedLocationEndpoint(failedRequest.requestContext.cosmosDiagnostics); if (firstContactedLocationEndpoint != null) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java index a0b7daf4b351..1d9d38c7d30d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java @@ -84,6 +84,7 @@ public class RxDocumentServiceRequest implements Cloneable { public String throughputControlGroupName; public volatile boolean intendedCollectionRidPassedIntoSDK = false; private volatile Duration responseTimeout; + public volatile boolean useThinProxy = false; private volatile boolean nonIdempotentWriteRetriesEnabled = false; 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 4ff506b71165..15a6fab9561a 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 @@ -298,12 +298,16 @@ private HttpHeaders getHttpRequestHeaders(Map headers) { return httpHeaders; } + public URI getRootUri(RxDocumentServiceRequest request) { + return this.globalEndpointManager.resolveServiceEndpoint(request).getGatewayLocationEndpoint(); + } + private URI getUri(RxDocumentServiceRequest request) throws URISyntaxException { URI rootUri = request.getEndpointOverride(); 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).getGatewayLocationEndpoint(); } else { rootUri = this.globalEndpointManager.resolveServiceEndpoint(request); } 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..56a0ed924c44 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 noRetryOnNonRelatedException(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/circuitBreaker/GlobalPartitionEndpointManagerForCircuitBreaker.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/circuitBreaker/GlobalPartitionEndpointManagerForCircuitBreaker.java index 92fccede1052..21b2c9b0ac21 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.LocationCache; 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, LocationCache.RegionalEndpoints 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.getGatewayLocationEndpoint(), 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; + LocationCache.RegionalEndpoints succeededLocation = request.requestContext.regionalEndpointsToRoute; 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()) { + LocationCache.RegionalEndpoints regionalEndpoints = pair.getKey(); + + URI gatewayLocationEndpoint = regionalEndpoints.getGatewayLocationEndpoint(); + 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(); + LocationCache.RegionalEndpoints 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(); + LocationCache.RegionalEndpoints 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.getGatewayLocationEndpoint()); 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, + LocationCache.RegionalEndpoints 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.getGatewayLocationEndpoint(), 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, + LocationCache.RegionalEndpoints 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.getGatewayLocationEndpoint(), 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 (LocationCache.RegionalEndpoints 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..61986dd5dec4 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 @@ -23,6 +23,7 @@ import com.azure.cosmos.implementation.faultinjection.GatewayServerErrorInjector; import com.azure.cosmos.implementation.faultinjection.IFaultInjectorProvider; import com.azure.cosmos.implementation.http.HttpClient; +import com.azure.cosmos.implementation.routing.LocationCache; import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper; import com.azure.cosmos.implementation.routing.PartitionKeyRangeIdentity; import com.azure.cosmos.models.CosmosContainerIdentity; @@ -91,11 +92,11 @@ public GlobalAddressResolver( this.addressCacheByEndpoint = new ConcurrentHashMap<>(); this.apiType = apiType; - for (URI endpoint : endpointManager.getWriteEndpoints()) { - this.getOrAddEndpoint(endpoint); + for (LocationCache.RegionalEndpoints endpoint : endpointManager.getWriteEndpoints()) { + this.getOrAddEndpoint(endpoint.getGatewayLocationEndpoint()); } - for (URI endpoint : endpointManager.getReadEndpoints()) { - this.getOrAddEndpoint(endpoint); + for (LocationCache.RegionalEndpoints endpoint : endpointManager.getReadEndpoints()) { + this.getOrAddEndpoint(endpoint.getGatewayLocationEndpoint()); } } @@ -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.getGatewayLocationEndpoint())) { + EndpointCache endpointCache = this.addressCacheByEndpoint.get(readEndpoint.getGatewayLocationEndpoint()); 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; + LocationCache.RegionalEndpoints endpoint = this.endpointManager.resolveServiceEndpoint(rxDocumentServiceRequest); + return this.getOrAddEndpoint(endpoint.getGatewayLocationEndpoint()).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()) { + LocationCache.RegionalEndpoints dequeueEndpoint = endpoints.pop(); + if (this.addressCacheByEndpoint.get(dequeueEndpoint.getGatewayLocationEndpoint()) != null) { + this.addressCacheByEndpoint.remove(dequeueEndpoint.getGatewayLocationEndpoint()); } } else { break; 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..526b62ac51c2 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.LocationCache; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; @@ -77,6 +78,7 @@ private RxDocumentServiceRequest getOpenConnectionRequest(String collectionRid, RxDocumentServiceRequest openConnectionRequest = RxDocumentServiceRequest.create(null, OperationType.Create, ResourceType.Connection); openConnectionRequest.requestContext.locationEndpointToRoute = serviceEndpoint; + openConnectionRequest.requestContext.regionalEndpointsToRoute = new LocationCache.RegionalEndpoints(serviceEndpoint); openConnectionRequest.requestContext.resolvedCollectionRid = collectionRid; openConnectionRequest.faultInjectionRequestContext.setLocationEndpointToRoute(serviceEndpoint); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/faultinjection/FaultInjectionRequestContext.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/faultinjection/FaultInjectionRequestContext.java index c8dfcfa655d7..e713b2102d18 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/faultinjection/FaultInjectionRequestContext.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/faultinjection/FaultInjectionRequestContext.java @@ -3,6 +3,8 @@ package com.azure.cosmos.implementation.faultinjection; +import com.azure.cosmos.implementation.routing.LocationCache; + import java.net.URI; import java.util.ArrayList; import java.util.List; @@ -23,6 +25,7 @@ public class FaultInjectionRequestContext { private final AtomicBoolean addressForceRefreshed; private volatile URI locationEndpointToRoute; + private volatile LocationCache.RegionalEndpoints regionalEndpointsToRoute; /*** * This usually is called during retries. @@ -96,6 +99,7 @@ public String getFaultInjectionRuleId(long transportRequestId) { public void setLocationEndpointToRoute(URI locationEndpointToRoute) { this.locationEndpointToRoute = locationEndpointToRoute; + this.regionalEndpointsToRoute = new LocationCache.RegionalEndpoints(locationEndpointToRoute); } public URI getLocationEndpointToRoute() { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/Fetcher.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/Fetcher.java index dad3ff190abb..4e986dfeb7c1 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/Fetcher.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/Fetcher.java @@ -10,6 +10,7 @@ import com.azure.cosmos.implementation.circuitBreaker.GlobalPartitionEndpointManagerForCircuitBreaker; import com.azure.cosmos.implementation.ImplementationBridgeHelpers; import com.azure.cosmos.implementation.RxDocumentServiceRequest; +import com.azure.cosmos.implementation.routing.LocationCache; import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple; import com.azure.cosmos.models.FeedResponse; import com.azure.cosmos.models.ModelBridgeInternal; @@ -18,7 +19,6 @@ import reactor.core.publisher.Mono; import reactor.core.publisher.SignalType; -import java.net.URI; import java.util.List; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; @@ -228,7 +228,7 @@ private Mono> nextPage(RxDocumentServiceRequest request) { } private void handleCancellationExceptionForPartitionKeyRange(RxDocumentServiceRequest failedRequest) { - URI firstContactedLocationEndpoint = diagnosticsAccessor.getFirstContactedLocationEndpoint(failedRequest.requestContext.cosmosDiagnostics); + LocationCache.RegionalEndpoints 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..ead5c5037c8b 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 @@ -46,7 +46,7 @@ public class LocationCache { 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; @@ -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,10 @@ && unavailableLocationsExpirationTimePassed()) { * @return */ public List getAvailableReadEndpoints() { - return this.locationInfo.availableReadEndpointByLocation.values().stream().collect(Collectors.toList()); + return this.locationInfo.availableReadEndpointsByLocation.values().stream().map(consolidatedLocationEndpoints -> { + // TODO: Integrate thinclient endpoints into fault injection + return consolidatedLocationEndpoints.gatewayLocationEndpoint; + }).collect(Collectors.toList()); } /*** @@ -130,7 +133,10 @@ 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: Integrate thinclient endpoints into fault injection + return consolidatedLocationEndpoints.gatewayLocationEndpoint; + }).collect(Collectors.toList()); } public List getEffectivePreferredLocations() { @@ -186,43 +192,47 @@ void onLocationPreferenceChanged(UnmodifiableList preferredLocations) { * @param request Request for which getEndpoint is to be resolved * @return Resolved getEndpoint */ - public URI resolveServiceEndpoint(RxDocumentServiceRequest request) { + public RegionalEndpoints 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.locationEndpointToRoute != null) { + + Objects.requireNonNull(request.requestContext.regionalEndpointsToRoute); + + return request.requestContext.regionalEndpointsToRoute; } 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 new RegionalEndpoints(this.defaultEndpoint); } } 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) ? @@ -250,12 +260,12 @@ public UnmodifiableList getApplicableWriteEndpoints(List excludedRe 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) ? @@ -279,18 +289,18 @@ public UnmodifiableList getApplicableReadEndpoints(List excludedReg return this.getApplicableEndpoints( readEndpoints, this.locationInfo.regionNameByReadEndpoint, - this.locationInfo.writeEndpoints.get(0), // match the fallback region used in getPreferredAvailableEndpoints + this.locationInfo.writeEndpoints.get(0).getGatewayLocationEndpoint(), // match the fallback region used in getPreferredAvailableEndpoints effectiveExcludedRegionsWithPartitionUnavailableRegions); } - private UnmodifiableList getApplicableEndpoints( - UnmodifiableList endpoints, - UnmodifiableMap regionNameByEndpoint, + private UnmodifiableList getApplicableEndpoints( + UnmodifiableList endpoints, + UnmodifiableMap regionNameByEndpoint, URI fallbackEndpoint, List excludeRegionList) { - List applicableEndpoints = new ArrayList<>(); - for (URI endpoint : endpoints) { + List applicableEndpoints = new ArrayList<>(); + for (RegionalEndpoints endpoint : endpoints) { Utils.ValueHolder regionName = new Utils.ValueHolder<>(); if (Utils.tryGetValue(regionNameByEndpoint, endpoint, regionName)) { if (!excludeRegionList.stream().anyMatch(regionName.v::equalsIgnoreCase)) { @@ -300,7 +310,7 @@ private UnmodifiableList getApplicableEndpoints( } if (applicableEndpoints.isEmpty()) { - applicableEndpoints.add(fallbackEndpoint); + applicableEndpoints.add(new RegionalEndpoints(fallbackEndpoint)); } return new UnmodifiableList<>(applicableEndpoints); @@ -314,15 +324,16 @@ 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.gatewayLocationEndpoint; } throw new IllegalArgumentException("Can not find service endpoint for region " + region); @@ -345,7 +356,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 +369,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 +393,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 +414,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 +444,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 RegionalEndpoints(locationEndpoint), regionName)) { return regionName.v; } } else { - if (Utils.tryGetValue(this.locationInfo.regionNameByReadEndpoint, locationEndpoint, regionName)) { + if (Utils.tryGetValue(this.locationInfo.regionNameByReadEndpoint, new RegionalEndpoints(locationEndpoint), regionName)) { return regionName.v; } } @@ -446,15 +457,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(RegionalEndpoints url1, RegionalEndpoints 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 (RegionalEndpoints unavailableEndpoint: unavailableEndpoints) { Utils.ValueHolder unavailabilityInfoHolder = new Utils.ValueHolder<>(); Utils.ValueHolder removedHolder = new Utils.ValueHolder<>(); @@ -473,7 +484,7 @@ private void clearStaleEndpointUnavailabilityInfo() { } } - private boolean isEndpointUnavailable(URI endpoint, OperationType expectedAvailableOperations) { + private boolean isEndpointUnavailable(RegionalEndpoints endpoint, OperationType expectedAvailableOperations) { Utils.ValueHolder unavailabilityInfoHolder = new Utils.ValueHolder<>(); if (expectedAvailableOperations == OperationType.None @@ -494,10 +505,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 (RegionalEndpoints endpoint : endpoints) { if (!isEndpointUnavailable(endpoint, expectedAvailableOperations)) { anyEndpointsAvailable = true; break; @@ -511,10 +522,10 @@ private void markEndpointUnavailable( OperationType unavailableOperationType) { Instant currentTime = Instant.now(); LocationUnavailabilityInfo updatedInfo = this.locationUnavailabilityInfoByEndpoint.compute( - unavailableEndpoint, - new BiFunction() { + new RegionalEndpoints(unavailableEndpoint), + new BiFunction() { @Override - public LocationUnavailabilityInfo apply(URI url, LocationUnavailabilityInfo info) { + public LocationUnavailabilityInfo apply(RegionalEndpoints url, LocationUnavailabilityInfo info) { if (info == null) { // not already present, add return new LocationUnavailabilityInfo(currentTime, unavailableOperationType); @@ -524,7 +535,6 @@ public LocationUnavailabilityInfo apply(URI url, LocationUnavailabilityInfo info info.unavailableOperations = OperationType.combine(info.unavailableOperations, unavailableOperationType); return info; } - } }); @@ -542,8 +552,8 @@ private void updateLocationCache(){ } private void updateLocationCache( - Iterable writeLocations, - Iterable readLocations, + Iterable gatewayWriteLocations, + Iterable gatewayReadLocations, UnmodifiableList preferenceList, Boolean enableMultipleWriteLocations) { synchronized (this.lockObject) { @@ -561,31 +571,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.defaultEndpoint); + nextLocationInfo.readEndpoints = this.getPreferredAvailableEndpoints(nextLocationInfo.availableReadEndpointsByLocation, nextLocationInfo.availableReadLocations, OperationType.Read, nextLocationInfo.writeEndpoints.get(0).getGatewayLocationEndpoint()); 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, new RegionalEndpoints(this.defaultEndpoint), regionForDefaultEndpoint)) { nextLocationInfo.effectivePreferredLocations = nextLocationInfo.availableReadLocations; } } @@ -598,16 +608,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, + URI fallbackEndpoint) { + 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 +626,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 +637,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.defaultEndpoint.equals(endpoint.v.getGatewayLocationEndpoint())) { endpoints = new ArrayList<>(); break; } @@ -647,14 +657,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 +673,59 @@ 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)); + + RegionalEndpoints regionalEndpoints = new RegionalEndpoints(endpoint); + + if (!endpointsByLocation.containsKey(location)) { + endpointsByLocation.put(location, regionalEndpoints); + } - 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(regionalEndpoints)) { + regionByEndpoint.put(regionalEndpoints, location); + } + + parsedLocations.add(gatewayDbAccountLocation.getName()); + } catch (Exception e) { + logger.warn("GetAvailableEndpointsByLocation() - skipping add for location = [{}] as its location name is either empty or endpoint is malformed [{}]", + gatewayDbAccountLocation.getName(), + gatewayDbAccountLocation.getEndpoint()); + } } } } + } + + 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() { @@ -772,36 +800,61 @@ 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 RegionalEndpoints(defaultEndpoint))); + this.writeEndpoints = new UnmodifiableList<>(Collections.singletonList(new RegionalEndpoints(defaultEndpoint))); } public DatabaseAccountLocationsInfo(DatabaseAccountLocationsInfo other) { @@ -809,10 +862,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; } From 9b42df8f441589cff43afd35bca3611dea5a1d0c Mon Sep 17 00:00:00 2001 From: Neha Rao Date: Tue, 25 Feb 2025 11:33:12 -0800 Subject: [PATCH 02/17] fix variable name in resolveServiceEndpoint --- .../com/azure/cosmos/implementation/GlobalEndpointManager.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 e71eab9db96f..97bae5286c3e 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 @@ -153,7 +153,7 @@ public LocationCache.RegionalEndpoints resolveServiceEndpoint(RxDocumentServiceR request.faultInjectionRequestContext.setLocationEndpointToRoute(serviceEndpoints.getGatewayLocationEndpoint()); } - return serviceEndpoint; + return serviceEndpoints; } public URI resolveFaultInjectionServiceEndpoint(String region, boolean writeOnly) { From e65898396d7aaef588dbafa71a89ec92a4106734 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Wed, 26 Feb 2025 12:06:09 -0500 Subject: [PATCH 03/17] Fixing CI pipeline. --- .../com/azure/cosmos/implementation/RxGatewayStoreModel.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 21f25130ac46..d0ee340fc70a 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 @@ -328,7 +328,7 @@ private URI getUri(RxDocumentServiceRequest request) throws URISyntaxException { // For media read request, always use the write endpoint. rootUri = this.globalEndpointManager.getWriteEndpoints().get(0).getGatewayLocationEndpoint(); } else { - rootUri = this.globalEndpointManager.resolveServiceEndpoint(request); + rootUri = this.globalEndpointManager.resolveServiceEndpoint(request).getGatewayLocationEndpoint(); } } From b8611f7fbbf35e0d27df0857dfd521652749afc3 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Wed, 26 Feb 2025 13:01:14 -0500 Subject: [PATCH 04/17] Fixing CI pipeline. --- .../implementation/ThinClientStoreModelTest.java | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) 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 From 8db98bb153a3f13ac9151384592cd9b834cc8356 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Wed, 26 Feb 2025 13:01:30 -0500 Subject: [PATCH 05/17] Fixing CI pipeline. --- .../routing/RegionalRoutingContext.java | 32 +++++++++++++++++++ 1 file changed, 32 insertions(+) create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/routing/RegionalRoutingContext.java 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..1d8b44d59e3a --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/routing/RegionalRoutingContext.java @@ -0,0 +1,32 @@ +// 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 { + private final URI gatewayRegionalEndpoint; + + public RegionalRoutingContext(URI gatewayRegionalEndpoint) { + this.gatewayRegionalEndpoint = gatewayRegionalEndpoint; + } + + 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.gatewayRegionalEndpoint); + } +} From 85ddc94ab5126eb56a270376b0a1f3fbd45c9d46 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Wed, 26 Feb 2025 14:28:04 -0500 Subject: [PATCH 06/17] Fixing CI pipeline. --- .../FaultInjectionRuleProcessor.java | 5 +- ...EndpointManagerForCircuitBreakerTests.java | 148 ++++---- .../PartitionLevelCircuitBreakerTests.java | 6 +- .../ProactiveConnectionManagementTest.java | 18 +- .../implementation/ClientRetryPolicyTest.java | 20 +- .../RegionScopedSessionContainerTest.java | 335 +++++++++--------- ...eCollectionAwareClientRetryPolicyTest.java | 8 +- .../RxGatewayStoreModelTest.java | 11 +- .../WebExceptionRetryPolicyTest.java | 12 +- .../GlobalAddressResolverTest.java | 28 +- .../GlobalEndPointManagerTest.java | 21 +- .../MetadataRequestRetryPolicyTests.java | 12 +- .../RntbdTransportClientTest.java | 5 +- .../query/DocumentProducerTest.java | 8 +- .../routing/LocationCacheTest.java | 38 +- .../com/azure/cosmos/CosmosDiagnostics.java | 6 +- .../implementation/ClientRetryPolicy.java | 51 ++- .../ClientSideRequestStatistics.java | 29 +- .../DocumentServiceRequestContext.java | 24 +- .../implementation/GlobalEndpointManager.java | 19 +- .../ImplementationBridgeHelpers.java | 4 +- .../MetadataRequestRetryPolicy.java | 8 +- .../RegionScopedSessionContainer.java | 7 +- .../implementation/RxDocumentClientImpl.java | 16 +- .../RxDocumentServiceRequest.java | 1 - .../implementation/RxGatewayStoreModel.java | 8 +- .../WebExceptionRetryPolicy.java | 11 +- ...itionEndpointManagerForCircuitBreaker.java | 46 +-- .../GlobalAddressResolver.java | 28 +- .../RntbdTransportClient.java | 2 +- .../rntbd/RntbdOpenConnectionsHandler.java | 6 +- .../FaultInjectionRequestContext.java | 6 +- .../cosmos/implementation/query/Fetcher.java | 4 +- .../implementation/routing/LocationCache.java | 190 +++++----- 34 files changed, 554 insertions(+), 587 deletions(-) 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 654607d6fb66..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 @@ -23,8 +23,8 @@ import com.azure.cosmos.implementation.directconnectivity.AddressSelector; import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdUtils; import com.azure.cosmos.implementation.feedranges.FeedRangeInternal; -import com.azure.cosmos.implementation.routing.LocationCache; 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; @@ -447,8 +447,7 @@ private Mono> resolvePhysicalAddresses( ResourceType.Document, null); - faultInjectionAddressRequest.requestContext.locationEndpointToRoute = regionEndpoint; - faultInjectionAddressRequest.requestContext.regionalEndpointsToRoute = new LocationCache.RegionalEndpoints(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 c7d12f277f1c..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 @@ -16,7 +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.LocationCache; +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(new LocationCache.RegionalEndpoints(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(uriStringPair -> new LocationCache.RegionalEndpoints(uriStringPair.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, new LocationCache.RegionalEndpoints(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(new LocationCache.RegionalEndpoints(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(uriStringPair -> new LocationCache.RegionalEndpoints(uriStringPair.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, new LocationCache.RegionalEndpoints(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(new LocationCache.RegionalEndpoints(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(uriStringPair -> new LocationCache.RegionalEndpoints(uriStringPair.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, new LocationCache.RegionalEndpoints(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(new LocationCache.RegionalEndpoints(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(new LocationCache.RegionalEndpoints(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(uriStringPair -> new LocationCache.RegionalEndpoints(uriStringPair.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, new LocationCache.RegionalEndpoints(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(new LocationCache.RegionalEndpoints(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(new LocationCache.RegionalEndpoints(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(uriStringPair -> new LocationCache.RegionalEndpoints(uriStringPair.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, new LocationCache.RegionalEndpoints(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(new LocationCache.RegionalEndpoints(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, new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getKey())); + .handleLocationExceptionForPartitionKeyRange(request, new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getKey())); } - locationSpecificHealthContext = locationEndpointToLocationSpecificContextForPartition.get(new LocationCache.RegionalEndpoints(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(uriStringPair -> new LocationCache.RegionalEndpoints(uriStringPair.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, new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getKey())); + .handleLocationExceptionForPartitionKeyRange(request, new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getKey())); globalPartitionEndpointManagerForCircuitBreaker - .handleLocationExceptionForPartitionKeyRange(request, new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getKey())); + .handleLocationExceptionForPartitionKeyRange(request, new RegionalRoutingContext(LocationEastUsEndpointToLocationPair.getKey())); globalPartitionEndpointManagerForCircuitBreaker - .handleLocationExceptionForPartitionKeyRange(request, new LocationCache.RegionalEndpoints(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(uriStringPair -> new LocationCache.RegionalEndpoints(uriStringPair.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, new LocationCache.RegionalEndpoints(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(new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getKey())); + = locationEndpointToLocationSpecificContextForPartitionForColl1.get(new RegionalRoutingContext(LocationEastUs2EndpointToLocationPair.getKey())); LocationSpecificHealthContext locationSpecificHealthContext2 - = locationEndpointToLocationSpecificContextForPartitionForColl2.get(new LocationCache.RegionalEndpoints(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(uriStringPair -> new LocationCache.RegionalEndpoints(uriStringPair.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, new LocationCache.RegionalEndpoints(locationWithFailure)); + globalPartitionEndpointManagerForCircuitBreaker.handleLocationExceptionForPartitionKeyRange(request, new RegionalRoutingContext(locationWithFailure)); List unavailableRegions = globalPartitionEndpointManagerForCircuitBreaker.getUnavailableRegionsForPartitionKeyRange(collectionResourceId, partitionKeyRange, request.getOperationType()); @@ -914,8 +914,8 @@ private RxDocumentServiceRequest constructRxDocumentServiceRequestInstance( request.requestContext.resolvedPartitionKeyRange = new PartitionKeyRange(partitionKeyRangeId, minInclusive, maxExclusive); request.requestContext.resolvedPartitionKeyRangeForCircuitBreaker = request.requestContext.resolvedPartitionKeyRange; - request.requestContext.locationEndpointToRoute = locationEndpointToRoute; - request.requestContext.regionalEndpointsToRoute = new LocationCache.RegionalEndpoints(locationEndpointToRoute); + + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(locationEndpointToRoute); request.requestContext.setExcludeRegions(Collections.emptyList()); request.requestContext.setPointOperationContext( 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 b6de3bcb21c9..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,7 +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.LocationCache; +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; @@ -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 9b6e428bf47a..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 @@ -27,9 +27,9 @@ import com.azure.cosmos.implementation.directconnectivity.RntbdTransportClient; import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdEndpoint; import com.azure.cosmos.implementation.routing.CollectionRoutingMap; -import com.azure.cosmos.implementation.routing.LocationCache; 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; @@ -181,13 +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())) - .stream().map(LocationCache.RegionalEndpoints::getGatewayLocationEndpoint).collect(Collectors.toList()); + .stream().map(RegionalRoutingContext::getGatewayRegionalEndpoint).collect(Collectors.toList()); Mono asyncContainerMono = Mono.just(cosmosAsyncContainer); @@ -345,13 +345,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())) .stream() - .map(LocationCache.RegionalEndpoints::getGatewayLocationEndpoint) + .map(RegionalRoutingContext::getGatewayRegionalEndpoint) .collect(Collectors.toList()); Flux asyncContainerFlux = Flux.fromIterable(asyncContainers); @@ -495,12 +495,12 @@ 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())) .stream() - .map(LocationCache.RegionalEndpoints::getGatewayLocationEndpoint) + .map(RegionalRoutingContext::getGatewayRegionalEndpoint) .collect(Collectors.toList());; Flux asyncContainerFlux = Flux.fromIterable(asyncContainers); @@ -666,12 +666,12 @@ 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())) .stream() - .map(LocationCache.RegionalEndpoints::getGatewayLocationEndpoint) + .map(RegionalRoutingContext::getGatewayRegionalEndpoint) .collect(Collectors.toList());; Flux asyncContainerFlux = Flux.fromIterable(asyncContainers); 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 a7f8e94d30e6..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,7 +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.LocationCache; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import io.netty.handler.timeout.ReadTimeoutException; import io.reactivex.subscribers.TestSubscriber; import org.mockito.Mockito; @@ -67,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 LocationCache.RegionalEndpoints(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); @@ -108,7 +108,7 @@ public void shouldRetryOnGatewayTimeout( GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new LocationCache.RegionalEndpoints(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( @@ -151,7 +151,7 @@ public void tcpNetworkFailureOnRead() throws Exception { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new LocationCache.RegionalEndpoints(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); @@ -199,7 +199,7 @@ public void networkFailureOnWrite() throws Exception { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new LocationCache.RegionalEndpoints(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); @@ -234,7 +234,7 @@ public void tcpNetworkFailureOnWrite( GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new LocationCache.RegionalEndpoints(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); @@ -294,7 +294,7 @@ public void networkFailureOnUpsert() throws Exception { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new LocationCache.RegionalEndpoints(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); @@ -327,7 +327,7 @@ public void tcpNetworkFailureOnUpsert() throws Exception { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new LocationCache.RegionalEndpoints(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); @@ -363,7 +363,7 @@ public void networkFailureOnDelete() throws Exception { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new LocationCache.RegionalEndpoints(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); @@ -397,7 +397,7 @@ public void tcpNetworkFailureOnDelete() throws Exception { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new LocationCache.RegionalEndpoints(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/RegionScopedSessionContainerTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RegionScopedSessionContainerTest.java index fe06e0ad6002..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,7 +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.LocationCache; +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; @@ -375,11 +375,11 @@ public void sessionContainer() throws Exception { int numPartitionKeyRangeIds = 5; String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); - UnmodifiableList endpoints = new UnmodifiableList<>( + UnmodifiableList endpoints = new UnmodifiableList<>( ImmutableList.of( - new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); + 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); @@ -396,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, @@ -409,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); @@ -420,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); @@ -439,18 +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<>( + UnmodifiableList endpoints = new UnmodifiableList<>( ImmutableList.of( - new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); + 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; @@ -483,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(); @@ -506,18 +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<>( + UnmodifiableList endpoints = new UnmodifiableList<>( ImmutableList.of( - new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); + 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; @@ -539,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(); @@ -567,24 +567,24 @@ public void resolveGlobalSessionTokenReturnsEmptyStringOnCacheMiss() { GlobalEndpointManager globalEndpointManagerMock = Mockito.mock(GlobalEndpointManager.class); RegionScopedSessionContainer sessionContainer = new RegionScopedSessionContainer("127.0.0.1", false, globalEndpointManagerMock); - UnmodifiableList endpoints = new UnmodifiableList<>( + UnmodifiableList endpoints = new UnmodifiableList<>( ImmutableList.of( - new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); + 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)); } @@ -599,17 +599,17 @@ public void resolveGlobalSessionTokenReturnsTokenMapUsingName() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>( + UnmodifiableList endpoints = new UnmodifiableList<>( ImmutableList.of( - new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); + 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"; @@ -617,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"; @@ -645,17 +645,17 @@ public void resolveGlobalSessionTokenReturnsTokenMapUsingResourceId() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>( + UnmodifiableList endpoints = new UnmodifiableList<>( ImmutableList.of( - new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); + 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"; @@ -663,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"; @@ -672,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); @@ -693,17 +693,17 @@ public void resolveLocalSessionTokenReturnsTokenMapUsingName() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>( + UnmodifiableList endpoints = new UnmodifiableList<>( ImmutableList.of( - new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); + 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"; @@ -711,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"; @@ -737,17 +737,17 @@ public void resolveLocalSessionTokenReturnsTokenMapUsingResourceId() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>( + UnmodifiableList endpoints = new UnmodifiableList<>( ImmutableList.of( - new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); + 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"; @@ -755,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"; @@ -764,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(); @@ -786,17 +786,17 @@ public void resolveLocalSessionTokenReturnsNullOnPartitionMiss() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>( + UnmodifiableList endpoints = new UnmodifiableList<>( ImmutableList.of( - new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); + 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"; @@ -804,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"; @@ -813,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"); @@ -831,17 +831,17 @@ public void resolveLocalSessionTokenReturnsNullOnCollectionMiss() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>( + UnmodifiableList endpoints = new UnmodifiableList<>( ImmutableList.of( - new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); + 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"; @@ -849,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"; @@ -874,17 +874,17 @@ public void resolvePartitionLocalSessionTokenReturnsTokenOnParentMatch() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>( + UnmodifiableList endpoints = new UnmodifiableList<>( ImmutableList.of( - new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); + 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"; @@ -892,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"; @@ -920,17 +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<>( + UnmodifiableList endpoints = new UnmodifiableList<>( ImmutableList.of( - new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); + 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)); @@ -938,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); @@ -946,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); @@ -956,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(); @@ -964,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(); @@ -983,17 +983,17 @@ public void clearTokenByResourceIdRemovesToken() { String unparsedSessionToken = "range_0:1#100#1=20#2=5#3=30"; - UnmodifiableList endpoints = new UnmodifiableList<>( + UnmodifiableList endpoints = new UnmodifiableList<>( ImmutableList.of( - new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); + 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)); @@ -1001,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(); @@ -1010,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(); @@ -1021,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(); @@ -1029,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(); @@ -1047,17 +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<>( + UnmodifiableList endpoints = new UnmodifiableList<>( ImmutableList.of( - new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); + 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)); @@ -1065,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"; @@ -1073,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)); @@ -1106,18 +1106,18 @@ public void setSessionTokenSetsTokenWhenRequestIsntNameBased() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>( + UnmodifiableList endpoints = new UnmodifiableList<>( ImmutableList.of( - new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); + 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); @@ -1125,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(); @@ -1148,18 +1148,18 @@ public void setSessionTokenGivesPriorityToOwnerFullNameOverResourceAddress() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>( + UnmodifiableList endpoints = new UnmodifiableList<>( ImmutableList.of( - new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); + 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); @@ -1169,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(); @@ -1195,18 +1195,18 @@ public void setSessionTokenIgnoresOwnerIdWhenRequestIsntNameBased() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>( + UnmodifiableList endpoints = new UnmodifiableList<>( ImmutableList.of( - new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); + 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); @@ -1220,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(); @@ -1229,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(); @@ -1249,18 +1249,18 @@ public void setSessionTokenGivesPriorityToOwnerIdOverResourceIdWhenRequestIsName String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>( + UnmodifiableList endpoints = new UnmodifiableList<>( ImmutableList.of( - new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); + 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); @@ -1273,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); @@ -1317,18 +1317,18 @@ public void setSessionTokenDoesntOverwriteHigherLSN() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>( + UnmodifiableList endpoints = new UnmodifiableList<>( ImmutableList.of( - new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); + 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"; @@ -1338,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"; @@ -1347,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"); @@ -1365,18 +1365,18 @@ public void setSessionTokenOverwriteLowerLSN() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>( + UnmodifiableList endpoints = new UnmodifiableList<>( ImmutableList.of( - new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); + 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); @@ -1385,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"; @@ -1395,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"); @@ -1412,17 +1412,17 @@ public void setSessionTokenDoesNothingOnEmptySessionTokenHeader() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>( + UnmodifiableList endpoints = new UnmodifiableList<>( ImmutableList.of( - new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); + 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"; @@ -1430,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(",")); @@ -1438,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(",")); @@ -1508,11 +1508,11 @@ public void useParentSessionTokenAfterSplit() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>( + UnmodifiableList endpoints = new UnmodifiableList<>( ImmutableList.of( - new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); + 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); @@ -1523,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, @@ -1535,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( @@ -1560,11 +1560,11 @@ public void useParentSessionTokenAfterMerge() { String regionContacted = LocationEastUsEndpointToLocationPair.getRight(); URI locationEndpointContacted = LocationEastUsEndpointToLocationPair.getLeft(); - UnmodifiableList endpoints = new UnmodifiableList<>( + UnmodifiableList endpoints = new UnmodifiableList<>( ImmutableList.of( - new LocationCache.RegionalEndpoints(LocationEastUsEndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationEastUs2EndpointToLocationPair.getLeft()), - new LocationCache.RegionalEndpoints(LocationCentralUsEndpointToLocationPair.getLeft()))); + 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); @@ -1589,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, @@ -1607,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, @@ -1619,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( @@ -1659,14 +1659,14 @@ public void resolvePartitionLocalSessionToken( GlobalEndpointManager globalEndpointManagerMock = null; RegionScopedSessionContainer sessionContainer = null; - List consolidatedWriteRegionalEndpoints = writableURIToLocationMappings + List consolidatedWriteRegionalEndpointRoutingContexts = writableURIToLocationMappings .stream() - .map(uriToLocationMappings -> new LocationCache.RegionalEndpoints(uriToLocationMappings.getLeft())) + .map(uriToLocationMappings -> new RegionalRoutingContext(uriToLocationMappings.getLeft())) .collect(Collectors.toList()); - List consolidatedReadRegionalEndpoints = readEndpoints + List consolidatedReadRegionalEndpointRoutingContexts = readEndpoints .stream() - .map(readEndpoint -> new LocationCache.RegionalEndpoints(readEndpoint)) + .map(readEndpoint -> new RegionalRoutingContext(readEndpoint)) .collect(Collectors.toList()); DatabaseAccount databaseAccount = ModelBridgeUtils.createDatabaseAccount( @@ -1681,7 +1681,7 @@ public void resolvePartitionLocalSessionToken( .when(globalEndpointManagerMock.getLatestDatabaseAccount()) .thenReturn(databaseAccount); - UnmodifiableList readEndpointsInUnmodifiableList = new UnmodifiableList<>(consolidatedReadRegionalEndpoints); + UnmodifiableList readEndpointsInUnmodifiableList = new UnmodifiableList<>(consolidatedReadRegionalEndpointRoutingContexts); Mockito .when(globalEndpointManagerMock.getReadEndpoints()) @@ -1689,7 +1689,7 @@ public void resolvePartitionLocalSessionToken( Mockito .when(globalEndpointManagerMock.getApplicableWriteEndpoints(Mockito.anyList())) - .thenReturn(new UnmodifiableList<>(consolidatedWriteRegionalEndpoints)); + .thenReturn(new UnmodifiableList<>(consolidatedWriteRegionalEndpointRoutingContexts)); Mockito .when(globalEndpointManagerMock.canUseMultipleWriteLocations(Mockito.any())) @@ -1710,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) { @@ -1761,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); } @@ -1773,8 +1773,7 @@ private static RxDocumentServiceRequest createRequestEntity(OperationType operat operationType, resourceType); - request.requestContext.locationEndpointToRoute = locationEndpointToRoute; - request.requestContext.regionalEndpointsToRoute = new LocationCache.RegionalEndpoints(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 05c402a7c336..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,7 +6,7 @@ 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.LocationCache; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import io.netty.handler.timeout.ReadTimeoutException; import org.mockito.Mockito; import org.testng.annotations.Test; @@ -61,7 +61,7 @@ public void shouldRetryWithNotFoundStatusCode() throws URISyntaxException { Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(eq(null), eq(false)); URI locationEndToRoute = new URI("https://location1.documents.com"); - LocationCache.RegionalEndpoints consolidatedLocationEndpointToRoute = new LocationCache.RegionalEndpoints(locationEndToRoute); + RegionalRoutingContext consolidatedLocationEndpointToRoute = new RegionalRoutingContext(locationEndToRoute); Mockito.when(endpointManager.resolveServiceEndpoint(Mockito.any())).thenReturn(consolidatedLocationEndpointToRoute); @@ -95,7 +95,7 @@ public void shouldRetryWithNotFoundStatusCodeAndReadSessionNotAvailableSubStatus Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(eq(null), eq(false)); URI locationEndToRoute = new URI("https://location1.documents.com"); - LocationCache.RegionalEndpoints consolidatedLocationEndpointToRoute = new LocationCache.RegionalEndpoints(locationEndToRoute); + RegionalRoutingContext consolidatedLocationEndpointToRoute = new RegionalRoutingContext(locationEndToRoute); Mockito.when(endpointManager.resolveServiceEndpoint(Mockito.any())).thenReturn(consolidatedLocationEndpointToRoute); @@ -138,7 +138,7 @@ public void shouldRetryWithGenericException() throws URISyntaxException { GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); URI locationEndToRoute = new URI("https://location1.documents.com"); - LocationCache.RegionalEndpoints consolidatedLocationEndpointToRoute = new LocationCache.RegionalEndpoints(locationEndToRoute); + RegionalRoutingContext consolidatedLocationEndpointToRoute = new RegionalRoutingContext(locationEndToRoute); Mockito.when(endpointManager.resolveServiceEndpoint(Mockito.any())).thenReturn(consolidatedLocationEndpointToRoute); 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 b4d10a2ad5c9..3c92fc5ea062 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,9 +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.LocationCache; +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; @@ -83,7 +82,7 @@ public void readTimeout() throws Exception { GlobalEndpointManager globalEndpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new LocationCache.RegionalEndpoints(new URI("https://localhost"))) + Mockito.doReturn(new RegionalRoutingContext(new URI("https://localhost"))) .when(globalEndpointManager).resolveServiceEndpoint(any()); HttpClient httpClient = Mockito.mock(HttpClient.class); Mockito.doReturn(Mono.error(ReadTimeoutException.INSTANCE)) @@ -126,7 +125,7 @@ public void serviceUnavailable() throws Exception { UserAgentContainer userAgentContainer = new UserAgentContainer(); GlobalEndpointManager globalEndpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new LocationCache.RegionalEndpoints(new URI("https://localhost"))) + Mockito.doReturn(new RegionalRoutingContext(new URI("https://localhost"))) .when(globalEndpointManager).resolveServiceEndpoint(any()); HttpClient httpClient = Mockito.mock(HttpClient.class); Mockito.doReturn(Mono.error(new SocketException("Dummy SocketException"))) @@ -180,7 +179,7 @@ public void applySessionToken( GlobalEndpointManager globalEndpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new LocationCache.RegionalEndpoints(new URI("https://localhost"))) + Mockito.doReturn(new RegionalRoutingContext(new URI("https://localhost"))) .when(globalEndpointManager).resolveServiceEndpoint(any()); HttpClient httpClient = Mockito.mock(HttpClient.class); @@ -251,7 +250,7 @@ public void validateApiType() throws Exception { GlobalEndpointManager globalEndpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(new LocationCache.RegionalEndpoints(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/WebExceptionRetryPolicyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/WebExceptionRetryPolicyTest.java index 9481d4263f0b..79c7fdf1bbd7 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,7 +5,7 @@ import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.CosmosException; -import com.azure.cosmos.implementation.routing.LocationCache; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import io.netty.handler.timeout.ReadTimeoutException; import io.reactivex.subscribers.TestSubscriber; import org.mockito.Mockito; @@ -38,7 +38,7 @@ public static Object[][] operationTypeProvider() { @Test(groups = {"unit"}) public void shouldRetryOnTimeoutForReadOperations() throws Exception { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); - Mockito.doReturn(new LocationCache.RegionalEndpoints(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)); RetryContext retryContext = new RetryContext(); @@ -90,7 +90,7 @@ public void shouldRetryOnTimeoutForReadOperations() throws Exception { @Test(groups = {"unit"}) public void shouldRetryOnTimeoutForMetaDataReadOperations() throws Exception { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); - Mockito.doReturn(new LocationCache.RegionalEndpoints(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)); RetryContext retryContext = new RetryContext(); @@ -144,7 +144,7 @@ public void shouldRetryOnTimeoutForMetaDataReadOperations() throws Exception { @Test(groups = {"unit"}) public void shouldRetryOnTimeoutForQueryPlanOperations() throws Exception { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); - Mockito.doReturn(new LocationCache.RegionalEndpoints(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)); RetryContext retryContext = new RetryContext(); @@ -196,7 +196,7 @@ public void shouldRetryOnTimeoutForQueryPlanOperations() throws Exception { @Test(groups = "unit") public void shouldNotRetryOnTimeoutForWriteOperations() throws Exception { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); - Mockito.doReturn(new LocationCache.RegionalEndpoints(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)); @@ -237,7 +237,7 @@ 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 LocationCache.RegionalEndpoints(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(); 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 6a060e845f0e..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 @@ -24,9 +24,9 @@ import com.azure.cosmos.implementation.caches.RxCollectionCache; import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; import com.azure.cosmos.implementation.http.HttpClient; -import com.azure.cosmos.implementation.routing.LocationCache; 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(new LocationCache.RegionalEndpoints(urlforRead1)); - readEndPointList.add(new LocationCache.RegionalEndpoints(urlforRead2)); - readEndPointList.add(new LocationCache.RegionalEndpoints(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(new LocationCache.RegionalEndpoints(urlforWrite1)); - writeEndPointList.add(new LocationCache.RegionalEndpoints(urlforWrite2)); - writeEndPointList.add(new LocationCache.RegionalEndpoints(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(); - LocationCache.RegionalEndpoints testUrl = new LocationCache.RegionalEndpoints(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.getGatewayLocationEndpoint())).isTrue();//New endpoint will be added in addressCacheByEndpoint + assertThat(urlsBeforeResolve.contains(testUrl.getGatewayRegionalEndpoint())).isTrue();//New endpoint will be added in addressCacheByEndpoint } @Test(groups = "unit") @@ -158,7 +158,7 @@ public void submitOpenConnectionTasksAndInitCaches() { Mockito .when(endpointManager.getReadEndpoints()) .thenReturn(new UnmodifiableList<>( - Arrays.asList(new LocationCache.RegionalEndpoints(urlforRead1), new LocationCache.RegionalEndpoints(urlforRead2)))); + 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 918d7d671280..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,7 +291,7 @@ 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); @@ -302,11 +303,11 @@ private Map getAvailableWriteEndpointBy 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); @@ -316,7 +317,7 @@ private Map getAvailableReadEndpointByL 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 2bbff8cfa59e..ab1aa9562033 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,7 +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.LocationCache; +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; @@ -235,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).getGatewayLocationEndpoint(); + URI readEndpoint = readEndpoints.get(i).getGatewayRegionalEndpoint(); GlobalAddressResolver.EndpointCache endpointCache = endpointCacheByURIMap.get(readEndpoint); GatewayAddressCache gatewayAddressCache = endpointCache.addressCache; HttpClientUnderTestWrapper httpClientUnderTestWrapper = getHttpClientUnderTestWrapper(configs); @@ -348,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 @@ -555,9 +555,9 @@ private static RxDocumentServiceRequest createRequest( assert request.requestContext != null; if (hasLocationEndpointToRoute) { - request.requestContext.locationEndpointToRoute + URI locationEndpointToRoute = URI.create("https://account-name-some-region.documents.azure.com:443"); - request.requestContext.regionalEndpointsToRoute = new LocationCache.RegionalEndpoints(request.requestContext.locationEndpointToRoute); + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(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 a7f27e26176f..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 @@ -27,9 +27,9 @@ import com.azure.cosmos.implementation.guava25.collect.LinkedListMultimap; import com.azure.cosmos.implementation.query.orderbyquery.OrderByRowResult; import com.azure.cosmos.implementation.query.orderbyquery.OrderbyRowComparer; -import com.azure.cosmos.implementation.routing.LocationCache; 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; @@ -112,9 +112,9 @@ public Object[][] mergeParamProvider() { } private IRetryPolicyFactory mockDocumentClientIRetryPolicyFactory() { - LocationCache.RegionalEndpoints regionalEndpoints; + RegionalRoutingContext regionalRoutingContext; try { - regionalEndpoints = new LocationCache.RegionalEndpoints(new URI("http://localhost")); + regionalRoutingContext = new RegionalRoutingContext(new URI("http://localhost")); } catch (Exception e) { throw new IllegalStateException(e); } @@ -122,7 +122,7 @@ private IRetryPolicyFactory mockDocumentClientIRetryPolicyFactory() { GlobalEndpointManager globalEndpointManager = Mockito.mock(GlobalEndpointManager.class); GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); - Mockito.doReturn(regionalEndpoints).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 313f6b005062..0e5c78b4a8b8 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 @@ -483,9 +483,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(new LocationCache.RegionalEndpoints(LocationCacheTest.Location1Endpoint)); - assertThat(this.cache.getWriteEndpoints().get(1)).isEqualTo(new LocationCache.RegionalEndpoints(LocationCacheTest.Location2Endpoint)); - assertThat(this.cache.getWriteEndpoints().get(2)).isEqualTo(new LocationCache.RegionalEndpoints(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") @@ -515,11 +515,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()); } @@ -535,8 +535,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()); @@ -670,8 +670,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()));; @@ -901,28 +901,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(new LocationCache.RegionalEndpoints(firstAvailableWriteEndpoint)).isEqualTo(writeEndpoints.get(0)); - assertThat(new LocationCache.RegionalEndpoints(secondAvailableWriteEndpoint)).isEqualTo(this.resolveEndpointForWriteRequest(ResourceType.Document, true)); - assertThat(new LocationCache.RegionalEndpoints(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(new LocationCache.RegionalEndpoints(firstWriteEnpoint)).isEqualTo(this.resolveEndpointForWriteRequest(ResourceType.Database, false)); - assertThat(new LocationCache.RegionalEndpoints(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(new LocationCache.RegionalEndpoints(firstAvailableReadEndpoint)).isEqualTo(this.resolveEndpointForReadRequest(true)); - assertThat(new LocationCache.RegionalEndpoints(firstAvailableReadEndpoint)).isEqualTo(this.resolveEndpointForReadRequest(false)); + assertThat(new RegionalRoutingContext(firstAvailableReadEndpoint)).isEqualTo(this.resolveEndpointForReadRequest(true)); + assertThat(new RegionalRoutingContext(firstAvailableReadEndpoint)).isEqualTo(this.resolveEndpointForReadRequest(false)); } - private LocationCache.RegionalEndpoints 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 LocationCache.RegionalEndpoints 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/src/main/java/com/azure/cosmos/CosmosDiagnostics.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosDiagnostics.java index 661b4f11c194..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,7 +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.LocationCache; +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; @@ -331,7 +331,7 @@ String getFirstContactedRegion() { return this.clientSideRequestStatistics.getFirstContactedRegion(); } - LocationCache.RegionalEndpoints getFirstContactedLocationEndpoint() { + RegionalRoutingContext getFirstContactedLocationEndpoint() { return this.clientSideRequestStatistics.getFirstContactedLocationEndpoint(); } @@ -479,7 +479,7 @@ public void setDiagnosticsContext(CosmosDiagnostics cosmosDiagnostics, CosmosDia } @Override - public LocationCache.RegionalEndpoints 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 a3238d075521..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,17 +12,17 @@ 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.LocationCache; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Mono; import java.net.URI; import java.time.Duration; -import java.util.Objects; 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. @@ -47,8 +47,7 @@ public class ClientRetryPolicy extends DocumentClientRetryPolicy { private int staleContainerRetryCount; private boolean isReadRequest; private boolean canUseMultipleWriteLocations; - private URI locationEndpoint; - private LocationCache.RegionalEndpoints regionalEndpoints; + private RegionalRoutingContext regionalRoutingContext; private RetryContext retryContext; private CosmosDiagnostics cosmosDiagnostics; private AtomicInteger cnt = new AtomicInteger(0); @@ -89,9 +88,9 @@ public Mono shouldRetry(Exception e) { isReadRequest, canUseMultipleWriteLocations, e); - if (this.locationEndpoint == null || this.regionalEndpoints == 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)); } @@ -231,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); @@ -307,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.regionalEndpointsToRoute); + 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. @@ -339,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.regionalEndpoints.getGatewayLocationEndpoint()); + 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.regionalEndpoints.getGatewayLocationEndpoint()); + logger.warn("marking the endpoint {} as unavailable for write", gatewayRegionalEndpoint); + this.globalEndpointManager.markEndpointUnavailableForWrite(gatewayRegionalEndpoint); } this.retryContext = new RetryContext(this.failoverRetryCount, usePreferredLocations); @@ -359,7 +361,7 @@ private Mono shouldRetryOnBackendServiceUnavailableAsync( if (this.globalPartitionEndpointManagerForCircuitBreaker.isPartitionLevelCircuitBreakingApplicable(this.request)) { this.globalPartitionEndpointManagerForCircuitBreaker - .handleLocationExceptionForPartitionKeyRange(this.request, this.request.requestContext.regionalEndpointsToRoute); + .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 @@ -404,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 @@ -421,7 +425,7 @@ private Mono shouldRetryOnRequestTimeout( this.globalPartitionEndpointManagerForCircuitBreaker.handleLocationExceptionForPartitionKeyRange( this.request, - this.request.requestContext.regionalEndpointsToRoute); + this.request.requestContext.regionalRoutingContextToRoute); } } @@ -434,7 +438,7 @@ private Mono shouldRetryOnInternalServerError() { this.globalPartitionEndpointManagerForCircuitBreaker.handleLocationExceptionForPartitionKeyRange( this.request, - this.request.requestContext.regionalEndpointsToRoute); + this.request.requestContext.regionalRoutingContextToRoute); } return Mono.just(ShouldRetryResult.NO_RETRY); @@ -455,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); } @@ -463,11 +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.regionalEndpoints = this.globalEndpointManager.resolveServiceEndpoint(request); - this.locationEndpoint = this.regionalEndpoints.getGatewayLocationEndpoint(); + this.regionalRoutingContext = this.globalEndpointManager.resolveServiceEndpoint(request); if (request.requestContext != null) { - request.requestContext.routeToLocation(this.locationEndpoint, this.regionalEndpoints); + request.requestContext.routeToLocation(this.regionalRoutingContext); } } @@ -513,13 +517,4 @@ public RetryContext(int retryCount, this.retryRequestOnPreferredLocations = retryRequestOnPreferredLocations; } } - - private URI getGatewayLocationEndpoint(RxDocumentServiceRequest request) { - - Objects.requireNonNull(request, "Argument 'request' must not be null'"); - Objects.requireNonNull(request.requestContext, "Argument 'request.requestContext' must not be null'"); - Objects.requireNonNull(request.requestContext.regionalEndpointsToRoute, "Argument 'request.requestContext.consolidatedRegionalEndpointToRoute' must not be null'"); - - return request.requestContext.regionalEndpointsToRoute.getGatewayLocationEndpoint(); - } } 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 4fd366ce9b4c..ec945a0e4c6c 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,7 +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.LocationCache; +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; @@ -164,7 +164,7 @@ public void recordResponse(RxDocumentServiceRequest request, StoreResultDiagnost this.requestPayloadSizeInBytes = 0; } - LocationCache.RegionalEndpoints regionalEndpoints = null; + RegionalRoutingContext regionalRoutingContext = null; URI locationEndPoint = null; if (request.requestContext != null) { @@ -178,8 +178,7 @@ public void recordResponse(RxDocumentServiceRequest request, StoreResultDiagnost request.requestContext.getEndToEndOperationLatencyPolicyConfig().toString(); } - locationEndPoint = request.requestContext.locationEndpointToRoute; - regionalEndpoints = request.requestContext.regionalEndpointsToRoute; + regionalRoutingContext = request.requestContext.regionalRoutingContextToRoute; List excludedRegions = request.requestContext.getExcludeRegions(); if (excludedRegions != null && !excludedRegions.isEmpty()) { @@ -193,12 +192,12 @@ public void recordResponse(RxDocumentServiceRequest request, StoreResultDiagnost this.requestEndTimeUTC = responseTime; } - if (regionalEndpoints != null) { + if (regionalRoutingContext != null) { storeResponseStatistics.regionName = - globalEndpointManager.getRegionName(regionalEndpoints.getGatewayLocationEndpoint(), request.getOperationType()); + globalEndpointManager.getRegionName(regionalRoutingContext.getGatewayRegionalEndpoint(), request.getOperationType()); this.regionsContacted.add(storeResponseStatistics.regionName); this.locationEndpointsContacted.add(locationEndPoint); - this.regionsContactedWithContext.add(new RegionWithContext(storeResponseStatistics.regionName, regionalEndpoints)); + this.regionsContactedWithContext.add(new RegionWithContext(storeResponseStatistics.regionName, regionalRoutingContext)); } if (storeResponseStatistics.requestOperationType == OperationType.Head @@ -222,16 +221,18 @@ public void recordGatewayResponse( this.requestEndTimeUTC = responseTime; } - LocationCache.RegionalEndpoints regionalEndpoints = null; + RegionalRoutingContext regionalRoutingContext = null; URI locationEndPoint = null; + if (rxDocumentServiceRequest != null && rxDocumentServiceRequest.requestContext != null) { - regionalEndpoints = rxDocumentServiceRequest.requestContext.regionalEndpointsToRoute; - locationEndPoint = rxDocumentServiceRequest.requestContext.locationEndpointToRoute; + regionalRoutingContext = rxDocumentServiceRequest.requestContext.regionalRoutingContextToRoute; + locationEndPoint = regionalRoutingContext.getGatewayRegionalEndpoint(); this.approximateInsertionCountInBloomFilter = rxDocumentServiceRequest.requestContext.getApproximateBloomFilterInsertionCount(); this.keywordIdentifiers = rxDocumentServiceRequest.requestContext.getKeywordIdentifiers(); } + this.recordRetryContextEndTime(); if (locationEndPoint != null) { @@ -241,7 +242,7 @@ public void recordGatewayResponse( this.regionsContacted.add(regionName); this.locationEndpointsContacted.add(locationEndPoint); - this.regionsContactedWithContext.add(new RegionWithContext(regionName, regionalEndpoints)); + this.regionsContactedWithContext.add(new RegionWithContext(regionName, regionalRoutingContext)); } GatewayStatistics gatewayStatistics = new GatewayStatistics(); @@ -659,7 +660,7 @@ public String getFirstContactedRegion() { return this.regionsContactedWithContext.first().regionContacted; } - public LocationCache.RegionalEndpoints getFirstContactedLocationEndpoint() { + public RegionalRoutingContext getFirstContactedLocationEndpoint() { if (this.regionsContactedWithContext == null || this.regionsContactedWithContext.isEmpty()) { return null; } @@ -1056,10 +1057,10 @@ public static CosmosDiagnosticsSystemUsageSnapshot fetchSystemInformation() { static class RegionWithContext implements Comparable { private final String regionContacted; - private final LocationCache.RegionalEndpoints locationEndpointsContacted; + private final RegionalRoutingContext locationEndpointsContacted; private final long recordedTimestamp; - RegionWithContext(String regionContacted, LocationCache.RegionalEndpoints 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/DocumentServiceRequestContext.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DocumentServiceRequestContext.java index 0623fcf3bb68..298600afb0af 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.routing.LocationCache; 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,8 +38,8 @@ public class DocumentServiceRequestContext implements Cloneable { public volatile Integer regionIndex; public volatile Boolean usePreferredLocations; public volatile Integer locationIndexToRoute; - public volatile URI locationEndpointToRoute; - public volatile LocationCache.RegionalEndpoints regionalEndpointsToRoute; +// public volatile URI locationEndpointToRoute; + public volatile RegionalRoutingContext regionalRoutingContextToRoute; public volatile boolean performedBackgroundAddressRefresh; public volatile boolean performLocalRefreshOnGoneException; public volatile List storeResponses; @@ -81,19 +80,17 @@ public DocumentServiceRequestContext() {} public void routeToLocation(int locationIndex, boolean usePreferredLocations) { this.locationIndexToRoute = locationIndex; this.usePreferredLocations = usePreferredLocations; - this.locationEndpointToRoute = null; - this.regionalEndpointsToRoute = null; +// 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, LocationCache.RegionalEndpoints regionalEndpointsToRoute) { - this.locationEndpointToRoute = locationEndpoint; - this.regionalEndpointsToRoute = regionalEndpointsToRoute; + public void routeToLocation(RegionalRoutingContext regionalRoutingContextToRoute) { + this.regionalRoutingContextToRoute = regionalRoutingContextToRoute; this.locationIndexToRoute = null; this.usePreferredLocations = null; } @@ -103,8 +100,7 @@ public void routeToLocation(URI locationEndpoint, LocationCache.RegionalEndpoint */ public void clearRouteToLocation() { this.locationIndexToRoute = null; - this.locationEndpointToRoute = null; - this.regionalEndpointsToRoute = null; + this.regionalRoutingContextToRoute = null; this.usePreferredLocations = null; } @@ -145,8 +141,8 @@ public DocumentServiceRequestContext clone() { context.regionIndex = this.regionIndex; context.usePreferredLocations = this.usePreferredLocations; context.locationIndexToRoute = this.locationIndexToRoute; - context.locationEndpointToRoute = this.locationEndpointToRoute; - context.regionalEndpointsToRoute = this.regionalEndpointsToRoute; +// 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 97bae5286c3e..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,11 +147,11 @@ public static Mono getDatabaseAccountFromAnyLocationsAsync( }); } - public LocationCache.RegionalEndpoints resolveServiceEndpoint(RxDocumentServiceRequest request) { - LocationCache.RegionalEndpoints serviceEndpoints = this.locationCache.resolveServiceEndpoint(request); + public RegionalRoutingContext resolveServiceEndpoint(RxDocumentServiceRequest request) { + RegionalRoutingContext serviceEndpoints = this.locationCache.resolveServiceEndpoint(request); if (request.faultInjectionRequestContext != null) { // TODO: integrate thin client into fault injection - request.faultInjectionRequestContext.setLocationEndpointToRoute(serviceEndpoints.getGatewayLocationEndpoint()); + request.faultInjectionRequestContext.setLocationEndpointToRoute(serviceEndpoints.getGatewayRegionalEndpoint()); } return serviceEndpoints; 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 04a2b88167ee..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 @@ -40,8 +40,8 @@ import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdChannelStatistics; import com.azure.cosmos.implementation.faultinjection.IFaultInjectorProvider; import com.azure.cosmos.implementation.patch.PatchOperation; -import com.azure.cosmos.implementation.routing.LocationCache; 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; @@ -838,7 +838,7 @@ void recordAddressResolutionEnd( void setDiagnosticsContext(CosmosDiagnostics cosmosDiagnostics, CosmosDiagnosticsContext ctx); - LocationCache.RegionalEndpoints 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 f11853aeca51..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,7 +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.LocationCache; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Mono; @@ -70,10 +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) { - LocationCache.RegionalEndpoints regionalEndpoints = request.requestContext.regionalEndpointsToRoute; - URI locationEndpointToRoute = regionalEndpoints.getGatewayLocationEndpoint(); + 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 18d9237674cd..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(consolidatedReadLocationEndpoints -> globalEndpointManager.getRegionName(consolidatedReadLocationEndpoints.getGatewayLocationEndpoint(), 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/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index b96f386d9da5..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 @@ -56,13 +56,13 @@ import com.azure.cosmos.implementation.query.PipelinedQueryExecutionContextBase; import com.azure.cosmos.implementation.query.QueryInfo; import com.azure.cosmos.implementation.routing.CollectionRoutingMap; -import com.azure.cosmos.implementation.routing.LocationCache; import com.azure.cosmos.implementation.routing.PartitionKeyAndResourceTokenPair; import com.azure.cosmos.implementation.routing.PartitionKeyInternal; import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper; 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; @@ -147,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 = @@ -6461,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()) { @@ -6471,7 +6471,7 @@ private List getApplicableEndPoints(OperationTy return EMPTY_ENDPOINT_LIST; } - private static List withoutNulls(List orderedEffectiveEndpointsList) { + private static List withoutNulls(List orderedEffectiveEndpointsList) { if (orderedEffectiveEndpointsList == null) { return EMPTY_ENDPOINT_LIST; } @@ -6530,7 +6530,7 @@ private List getApplicableRegionsForSpeculation( return EMPTY_REGION_LIST; } - List regionalEndpointsList = getApplicableEndPoints(operationType, excludedRegions); + List regionalRoutingContextList = getApplicableEndPoints(operationType, excludedRegions); HashSet normalizedExcludedRegions = new HashSet<>(); if (excludedRegions != null) { @@ -6538,8 +6538,8 @@ private List getApplicableRegionsForSpeculation( } List orderedRegionsForSpeculation = new ArrayList<>(); - regionalEndpointsList.forEach(consolidatedLocationEndpoints -> { - String regionName = this.globalEndpointManager.getRegionName(consolidatedLocationEndpoints.getGatewayLocationEndpoint(), operationType); + regionalRoutingContextList.forEach(consolidatedLocationEndpoints -> { + String regionName = this.globalEndpointManager.getRegionName(consolidatedLocationEndpoints.getGatewayRegionalEndpoint(), operationType); if (!normalizedExcludedRegions.contains(regionName.toLowerCase(Locale.ROOT))) { orderedRegionsForSpeculation.add(regionName); } @@ -6743,7 +6743,7 @@ private Mono executeFeedOperationWithAvailabilityStrategy( private void handleLocationCancellationExceptionForPartitionKeyRange(RxDocumentServiceRequest failedRequest) { - LocationCache.RegionalEndpoints 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/RxDocumentServiceRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java index ea22a95ab03a..1f9f2cf2aaa3 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java @@ -86,7 +86,6 @@ public class RxDocumentServiceRequest implements Cloneable { public String throughputControlGroupName; public volatile boolean intendedCollectionRidPassedIntoSDK = false; private volatile Duration responseTimeout; - public volatile boolean useThinProxy = false; private volatile boolean nonIdempotentWriteRetriesEnabled = false; 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 d0ee340fc70a..2d1d43e8c692 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; @@ -318,7 +316,7 @@ private HttpHeaders getHttpRequestHeaders(Map headers) { } public URI getRootUri(RxDocumentServiceRequest request) { - return this.globalEndpointManager.resolveServiceEndpoint(request).getGatewayLocationEndpoint(); + return this.globalEndpointManager.resolveServiceEndpoint(request).getGatewayRegionalEndpoint(); } private URI getUri(RxDocumentServiceRequest request) throws URISyntaxException { @@ -326,9 +324,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).getGatewayLocationEndpoint(); + rootUri = this.globalEndpointManager.getWriteEndpoints().get(0).getGatewayRegionalEndpoint(); } else { - rootUri = this.globalEndpointManager.resolveServiceEndpoint(request).getGatewayLocationEndpoint(); + rootUri = this.globalEndpointManager.resolveServiceEndpoint(request).getGatewayRegionalEndpoint(); } } 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..a9114156499d 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; @@ -23,7 +24,7 @@ public class WebExceptionRetryPolicy implements IRetryPolicy { private HttpTimeoutPolicy timeoutPolicy; private boolean isReadRequest; private int retryCount = 0; - private URI locationEndpoint; + private RegionalRoutingContext regionalRoutingContext; public WebExceptionRetryPolicy() { durationTimer.start(); @@ -42,7 +43,7 @@ public Mono shouldRetry(Exception e) { .warn( "WebExceptionRetryPolicy() No more retrying on endpoint {}, operationType = {}, count = {}, " + "isAddressRefresh = {}", - this.locationEndpoint, + this.regionalRoutingContext.getGatewayRegionalEndpoint(), this.request.getOperationType(), this.retryCount, this.request.isAddressRefresh()); @@ -63,7 +64,7 @@ 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.getGatewayRegionalEndpoint(), this.request.getOperationType(), this.request.getResourceType(), this.retryCount, this.request.isAddressRefresh(), this.request.shouldForceAddressRefresh(), this.request.forceCollectionRoutingMapRefresh); @@ -77,7 +78,7 @@ public Mono shouldRetry(Exception e) { .debug( "WebExceptionRetryPolicy() No retrying on un-retryable exceptions on endpoint {}, operationType = {}, resourceType = {}, count = {}, " + "isAddressRefresh = {}", - this.locationEndpoint, + this.regionalRoutingContext.getGatewayRegionalEndpoint(), this.request.getOperationType(), this.request.getResourceType(), this.retryCount, @@ -101,7 +102,7 @@ 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; } 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 21b2c9b0ac21..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,7 +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.LocationCache; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; @@ -50,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"); @@ -73,7 +73,7 @@ public void init() { } } - public void handleLocationExceptionForPartitionKeyRange(RxDocumentServiceRequest request, LocationCache.RegionalEndpoints failedLocation) { + public void handleLocationExceptionForPartitionKeyRange(RxDocumentServiceRequest request, RegionalRoutingContext failedLocation) { checkNotNull(request, "Argument 'request' cannot be null!"); checkNotNull(request.requestContext, "Argument 'request.requestContext' cannot be null!"); @@ -110,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()); @@ -133,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.getGatewayLocationEndpoint(), request.isReadOnlyRequest() ? OperationType.Read : OperationType.Create), + failedLocation.getGatewayRegionalEndpoint(), request.isReadOnlyRequest() ? OperationType.Read : OperationType.Create), resolvedPartitionKeyRangeForCircuitBreaker.getMinInclusive(), resolvedPartitionKeyRangeForCircuitBreaker.getMaxExclusive(), collectionResourceId); @@ -162,7 +162,7 @@ public void handleLocationSuccessForPartitionKeyRange(RxDocumentServiceRequest r String resourceId = request.getResourceId(); PartitionKeyRangeWrapper partitionKeyRangeWrapper = new PartitionKeyRangeWrapper(resolvedPartitionKeyRangeForCircuitBreaker, resourceId); - LocationCache.RegionalEndpoints succeededLocation = request.requestContext.regionalEndpointsToRoute; + RegionalRoutingContext succeededLocation = request.requestContext.regionalRoutingContextToRoute; String collectionLink = getCollectionLink(request); @@ -195,13 +195,13 @@ public List getUnavailableRegionsForPartitionKeyRange(String collectionR List unavailableRegions = new ArrayList<>(); if (partitionLevelLocationUnavailabilityInfoSnapshot != null) { - Map locationEndpointToFailureMetricsForPartition = + Map locationEndpointToFailureMetricsForPartition = partitionLevelLocationUnavailabilityInfoSnapshot.locationEndpointToLocationSpecificContextForPartition; - for (Map.Entry pair : locationEndpointToFailureMetricsForPartition.entrySet()) { - LocationCache.RegionalEndpoints regionalEndpoints = pair.getKey(); + for (Map.Entry pair : locationEndpointToFailureMetricsForPartition.entrySet()) { + RegionalRoutingContext regionalRoutingContext = pair.getKey(); - URI gatewayLocationEndpoint = regionalEndpoints.getGatewayLocationEndpoint(); + URI gatewayLocationEndpoint = regionalRoutingContext.getGatewayRegionalEndpoint(); LocationSpecificHealthContext locationSpecificHealthContext = pair.getValue(); @@ -229,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()) { - LocationCache.RegionalEndpoints locationWithStaleUnavailabilityInfo = locationToLocationLevelMetrics.getKey(); + RegionalRoutingContext locationWithStaleUnavailabilityInfo = locationToLocationLevelMetrics.getKey(); LocationSpecificHealthContext locationSpecificHealthContext = locationToLocationLevelMetrics.getValue(); if (!locationSpecificHealthContext.isRegionAvailableToProcessRequests()) { @@ -260,7 +260,7 @@ private Flux updateStaleLocationInfo() { .flatMap(locationToLocationSpecificHealthContextPair -> { PartitionKeyRangeWrapper partitionKeyRangeWrapper = locationToLocationSpecificHealthContextPair.getLeft(); - LocationCache.RegionalEndpoints locationWithStaleUnavailabilityInfo = locationToLocationSpecificHealthContextPair.getRight().getLeft(); + RegionalRoutingContext locationWithStaleUnavailabilityInfo = locationToLocationSpecificHealthContextPair.getRight().getLeft(); PartitionLevelLocationUnavailabilityInfo partitionLevelLocationUnavailabilityInfo = this.partitionKeyRangeToLocationSpecificUnavailabilityInfo.get(partitionKeyRangeWrapper); @@ -270,7 +270,7 @@ private Flux updateStaleLocationInfo() { if (globalAddressResolver != null) { - GatewayAddressCache gatewayAddressCache = globalAddressResolver.getGatewayAddressCache(locationWithStaleUnavailabilityInfo.getGatewayLocationEndpoint()); + GatewayAddressCache gatewayAddressCache = globalAddressResolver.getGatewayAddressCache(locationWithStaleUnavailabilityInfo.getGatewayRegionalEndpoint()); if (gatewayAddressCache != null) { @@ -357,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; } @@ -374,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; @@ -386,7 +386,7 @@ private PartitionLevelLocationUnavailabilityInfo() { private boolean handleException( PartitionKeyRangeWrapper partitionKeyRangeWrapper, - LocationCache.RegionalEndpoints locationWithException, + RegionalRoutingContext locationWithException, boolean isReadOnlyRequest) { AtomicBoolean isExceptionThresholdBreached = new AtomicBoolean(false); @@ -420,7 +420,7 @@ private boolean handleException( locationAsKey, GlobalPartitionEndpointManagerForCircuitBreaker .this.globalEndpointManager - .getRegionName(locationAsKey.getGatewayLocationEndpoint(), isReadOnlyRequest ? OperationType.Read : OperationType.Create)); + .getRegionName(locationAsKey.getGatewayRegionalEndpoint(), isReadOnlyRequest ? OperationType.Read : OperationType.Create)); } String region = GlobalPartitionEndpointManagerForCircuitBreaker.this.locationToRegion.get(locationAsKey); @@ -435,7 +435,7 @@ private boolean handleException( private void handleSuccess( PartitionKeyRangeWrapper partitionKeyRangeWrapper, - LocationCache.RegionalEndpoints succeededLocation, + RegionalRoutingContext succeededLocation, boolean isReadOnlyRequest) { this.locationEndpointToLocationSpecificContextForPartition.compute(succeededLocation, (locationAsKey, locationSpecificContextAsVal) -> { @@ -469,7 +469,7 @@ private void handleSuccess( locationAsKey, GlobalPartitionEndpointManagerForCircuitBreaker .this.globalEndpointManager - .getRegionName(locationAsKey.getGatewayLocationEndpoint(), isReadOnlyRequest ? OperationType.Read : OperationType.Create)); + .getRegionName(locationAsKey.getGatewayRegionalEndpoint(), isReadOnlyRequest ? OperationType.Read : OperationType.Create)); } String region = GlobalPartitionEndpointManagerForCircuitBreaker.this.locationToRegion.get(locationAsKey); @@ -479,9 +479,9 @@ private void handleSuccess( }); } - public boolean areLocationsAvailableForPartitionKeyRange(List availableLocationsAtAccountLevel) { + public boolean areLocationsAvailableForPartitionKeyRange(List availableLocationsAtAccountLevel) { - for (LocationCache.RegionalEndpoints 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 61986dd5dec4..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 @@ -23,9 +23,9 @@ import com.azure.cosmos.implementation.faultinjection.GatewayServerErrorInjector; import com.azure.cosmos.implementation.faultinjection.IFaultInjectorProvider; import com.azure.cosmos.implementation.http.HttpClient; -import com.azure.cosmos.implementation.routing.LocationCache; 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; @@ -92,11 +92,11 @@ public GlobalAddressResolver( this.addressCacheByEndpoint = new ConcurrentHashMap<>(); this.apiType = apiType; - for (LocationCache.RegionalEndpoints endpoint : endpointManager.getWriteEndpoints()) { - this.getOrAddEndpoint(endpoint.getGatewayLocationEndpoint()); + for (RegionalRoutingContext endpoint : endpointManager.getWriteEndpoints()) { + this.getOrAddEndpoint(endpoint.getGatewayRegionalEndpoint()); } - for (LocationCache.RegionalEndpoints endpoint : endpointManager.getReadEndpoints()) { - this.getOrAddEndpoint(endpoint.getGatewayLocationEndpoint()); + for (RegionalRoutingContext endpoint : endpointManager.getReadEndpoints()) { + this.getOrAddEndpoint(endpoint.getGatewayRegionalEndpoint()); } } @@ -155,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.getGatewayLocationEndpoint())) { - EndpointCache endpointCache = this.addressCacheByEndpoint.get(readEndpoint.getGatewayLocationEndpoint()); + if (this.addressCacheByEndpoint.containsKey(readEndpoint.getGatewayRegionalEndpoint())) { + EndpointCache endpointCache = this.addressCacheByEndpoint.get(readEndpoint.getGatewayRegionalEndpoint()); return this.resolveAddressesPerCollection( endpointCache, containerLinkToPkrs.left, @@ -273,8 +273,8 @@ public void configureFaultInjectorProvider(IFaultInjectorProvider faultInjectorP } private IAddressResolver getAddressResolver(RxDocumentServiceRequest rxDocumentServiceRequest) { - LocationCache.RegionalEndpoints endpoint = this.endpointManager.resolveServiceEndpoint(rxDocumentServiceRequest); - return this.getOrAddEndpoint(endpoint.getGatewayLocationEndpoint()).addressResolver; + RegionalRoutingContext endpoint = this.endpointManager.resolveServiceEndpoint(rxDocumentServiceRequest); + return this.getOrAddEndpoint(endpoint.getGatewayRegionalEndpoint()).addressResolver; } private EndpointCache getOrAddEndpoint(URI endpoint) { @@ -300,15 +300,15 @@ private EndpointCache getOrAddEndpoint(URI endpoint) { }); if (this.addressCacheByEndpoint.size() > this.maxEndpoints) { - List allConsolidatedEndpoints = new ArrayList<>(this.endpointManager.getWriteEndpoints()); + List allConsolidatedEndpoints = new ArrayList<>(this.endpointManager.getWriteEndpoints()); allConsolidatedEndpoints.addAll(this.endpointManager.getReadEndpoints()); Collections.reverse(allConsolidatedEndpoints); - LinkedList endpoints = new LinkedList<>(allConsolidatedEndpoints); + LinkedList endpoints = new LinkedList<>(allConsolidatedEndpoints); while (this.addressCacheByEndpoint.size() > this.maxEndpoints) { if (!endpoints.isEmpty()) { - LocationCache.RegionalEndpoints dequeueEndpoint = endpoints.pop(); - if (this.addressCacheByEndpoint.get(dequeueEndpoint.getGatewayLocationEndpoint()) != null) { - this.addressCacheByEndpoint.remove(dequeueEndpoint.getGatewayLocationEndpoint()); + 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 526b62ac51c2..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,7 +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.LocationCache; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; @@ -77,8 +77,8 @@ public Flux openConnections(String collectionRid, List> nextPage(RxDocumentServiceRequest request) { } private void handleCancellationExceptionForPartitionKeyRange(RxDocumentServiceRequest failedRequest) { - LocationCache.RegionalEndpoints 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 ead5c5037c8b..085968524743 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 @@ -46,7 +46,7 @@ public class LocationCache { 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; @@ -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(); @@ -121,7 +121,7 @@ && unavailableLocationsExpirationTimePassed()) { public List getAvailableReadEndpoints() { return this.locationInfo.availableReadEndpointsByLocation.values().stream().map(consolidatedLocationEndpoints -> { // TODO: Integrate thinclient endpoints into fault injection - return consolidatedLocationEndpoints.gatewayLocationEndpoint; + return consolidatedLocationEndpoints.getGatewayRegionalEndpoint(); }).collect(Collectors.toList()); } @@ -135,7 +135,7 @@ public List getAvailableReadEndpoints() { public List getAvailableWriteEndpoints() { return this.locationInfo.availableWriteEndpointsByLocation.values().stream().map(consolidatedLocationEndpoints -> { // TODO: Integrate thinclient endpoints into fault injection - return consolidatedLocationEndpoints.gatewayLocationEndpoint; + return consolidatedLocationEndpoints.getGatewayRegionalEndpoint(); }).collect(Collectors.toList()); } @@ -192,15 +192,12 @@ void onLocationPreferenceChanged(UnmodifiableList preferredLocations) { * @param request Request for which getEndpoint is to be resolved * @return Resolved getEndpoint */ - public RegionalEndpoints 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) { - - Objects.requireNonNull(request.requestContext.regionalEndpointsToRoute); - - return request.requestContext.regionalEndpointsToRoute; + if (request.requestContext.regionalRoutingContextToRoute != null) { + return request.requestContext.regionalRoutingContextToRoute; } int locationIndex = Utils.getValueOrDefault(request.requestContext.locationIndexToRoute, 0); @@ -217,22 +214,22 @@ public RegionalEndpoints resolveServiceEndpoint(RxDocumentServiceRequest request String writeLocation = currentLocationInfo.availableWriteLocations.get(locationIndex); return currentLocationInfo.availableWriteEndpointsByLocation.get(writeLocation); } else { - return new RegionalEndpoints(this.defaultEndpoint); + return new RegionalRoutingContext(this.defaultEndpoint); } } 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) ? @@ -260,12 +257,12 @@ public UnmodifiableList getApplicableWriteEndpoints(List 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) ? @@ -289,18 +286,18 @@ public UnmodifiableList getApplicableReadEndpoints(List getApplicableEndpoints( - UnmodifiableList endpoints, - UnmodifiableMap regionNameByEndpoint, + private UnmodifiableList getApplicableEndpoints( + UnmodifiableList endpoints, + UnmodifiableMap regionNameByEndpoint, URI fallbackEndpoint, List excludeRegionList) { - List applicableEndpoints = new ArrayList<>(); - for (RegionalEndpoints 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)) { @@ -310,7 +307,7 @@ private UnmodifiableList getApplicableEndpoints( } if (applicableEndpoints.isEmpty()) { - applicableEndpoints.add(new RegionalEndpoints(fallbackEndpoint)); + applicableEndpoints.add(new RegionalRoutingContext(fallbackEndpoint)); } return new UnmodifiableList<>(applicableEndpoints); @@ -324,7 +321,7 @@ 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.availableWriteEndpointsByLocation, region, endpointValueHolder); } else { @@ -333,7 +330,7 @@ public URI resolveFaultInjectionEndpoint(String region, boolean writeOnly) { if (endpointValueHolder.v != null) { // TODO: Figure out how to integrate thinclient into fault injection - return endpointValueHolder.v.gatewayLocationEndpoint; + return endpointValueHolder.v.getGatewayRegionalEndpoint(); } throw new IllegalArgumentException("Can not find service endpoint for region " + region); @@ -356,7 +353,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 @@ -369,7 +366,7 @@ 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.availableReadEndpointsByLocation, mostPreferredLocation, mostPreferredReadEndpointHolder)) { @@ -393,8 +390,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()) { @@ -444,11 +441,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, new RegionalEndpoints(locationEndpoint), regionName)) { + if (Utils.tryGetValue(this.locationInfo.regionNameByWriteEndpoint, new RegionalRoutingContext(locationEndpoint), regionName)) { return regionName.v; } } else { - if (Utils.tryGetValue(this.locationInfo.regionNameByReadEndpoint, new RegionalEndpoints(locationEndpoint), regionName)) { + if (Utils.tryGetValue(this.locationInfo.regionNameByReadEndpoint, new RegionalRoutingContext(locationEndpoint), regionName)) { return regionName.v; } } @@ -457,15 +454,15 @@ public String getRegionName(URI locationEndpoint, com.azure.cosmos.implementatio return this.locationInfo.availableWriteLocations.get(0).toLowerCase(Locale.ROOT); } - private boolean areEqual(RegionalEndpoints url1, RegionalEndpoints 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 (RegionalEndpoints unavailableEndpoint: unavailableEndpoints) { + for (RegionalRoutingContext unavailableEndpoint: unavailableEndpoints) { Utils.ValueHolder unavailabilityInfoHolder = new Utils.ValueHolder<>(); Utils.ValueHolder removedHolder = new Utils.ValueHolder<>(); @@ -484,7 +481,7 @@ private void clearStaleEndpointUnavailabilityInfo() { } } - private boolean isEndpointUnavailable(RegionalEndpoints endpoint, OperationType expectedAvailableOperations) { + private boolean isEndpointUnavailable(RegionalRoutingContext endpoint, OperationType expectedAvailableOperations) { Utils.ValueHolder unavailabilityInfoHolder = new Utils.ValueHolder<>(); if (expectedAvailableOperations == OperationType.None @@ -505,10 +502,10 @@ private boolean isEndpointUnavailable(RegionalEndpoints endpoint, OperationType } } - private boolean anyEndpointsAvailable(List endpoints, OperationType expectedAvailableOperations) { + private boolean anyEndpointsAvailable(List endpoints, OperationType expectedAvailableOperations) { Utils.ValueHolder unavailabilityInfoHolder = new Utils.ValueHolder<>(); boolean anyEndpointsAvailable = false; - for (RegionalEndpoints endpoint : endpoints) { + for (RegionalRoutingContext endpoint : endpoints) { if (!isEndpointUnavailable(endpoint, expectedAvailableOperations)) { anyEndpointsAvailable = true; break; @@ -522,10 +519,10 @@ private void markEndpointUnavailable( OperationType unavailableOperationType) { Instant currentTime = Instant.now(); LocationUnavailabilityInfo updatedInfo = this.locationUnavailabilityInfoByEndpoint.compute( - new RegionalEndpoints(unavailableEndpoint), - new BiFunction() { + new RegionalRoutingContext(unavailableEndpoint), + new BiFunction() { @Override - public LocationUnavailabilityInfo apply(RegionalEndpoints url, LocationUnavailabilityInfo info) { + public LocationUnavailabilityInfo apply(RegionalRoutingContext url, LocationUnavailabilityInfo info) { if (info == null) { // not already present, add return new LocationUnavailabilityInfo(currentTime, unavailableOperationType); @@ -573,7 +570,7 @@ private void updateLocationCache( if (gatewayReadLocations != null) { Utils.ValueHolder> readValueHolder = Utils.ValueHolder.initialize(nextLocationInfo.availableReadLocations); - Utils.ValueHolder> readRegionMapValueHolder = Utils.ValueHolder.initialize(nextLocationInfo.regionNameByReadEndpoint); + Utils.ValueHolder> readRegionMapValueHolder = Utils.ValueHolder.initialize(nextLocationInfo.regionNameByReadEndpoint); nextLocationInfo.availableReadEndpointsByLocation = this.getEndpointsByLocation(gatewayReadLocations, readValueHolder, readRegionMapValueHolder); nextLocationInfo.availableReadLocations = readValueHolder.v; nextLocationInfo.regionNameByReadEndpoint = readRegionMapValueHolder.v; @@ -581,21 +578,21 @@ private void updateLocationCache( if (gatewayWriteLocations != null) { Utils.ValueHolder> writeValueHolder = Utils.ValueHolder.initialize(nextLocationInfo.availableWriteLocations); - Utils.ValueHolder> outWriteRegionMap = Utils.ValueHolder.initialize(nextLocationInfo.regionNameByWriteEndpoint); + 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.availableWriteEndpointsByLocation, nextLocationInfo.availableWriteLocations, OperationType.Write, this.defaultEndpoint); - nextLocationInfo.readEndpoints = this.getPreferredAvailableEndpoints(nextLocationInfo.availableReadEndpointsByLocation, nextLocationInfo.availableReadLocations, OperationType.Read, nextLocationInfo.writeEndpoints.get(0).getGatewayLocationEndpoint()); + nextLocationInfo.readEndpoints = this.getPreferredAvailableEndpoints(nextLocationInfo.availableReadEndpointsByLocation, nextLocationInfo.availableReadLocations, OperationType.Read, nextLocationInfo.writeEndpoints.get(0).getGatewayRegionalEndpoint()); 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, new RegionalEndpoints(this.defaultEndpoint), regionForDefaultEndpoint)) { + if (!Utils.tryGetValue(nextLocationInfo.regionNameByReadEndpoint, new RegionalRoutingContext(this.defaultEndpoint), regionForDefaultEndpoint)) { nextLocationInfo.effectivePreferredLocations = nextLocationInfo.availableReadLocations; } } @@ -608,16 +605,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, + URI fallbackEndpoint) { + 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. @@ -626,7 +623,7 @@ private UnmodifiableList getPreferredAvailableEndpoints(Unmod if (currentLocationInfo.preferredLocations != null && !currentLocationInfo.preferredLocations.isEmpty()) { for (String location: currentLocationInfo.preferredLocations) { - Utils.ValueHolder 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); @@ -637,12 +634,12 @@ private UnmodifiableList getPreferredAvailableEndpoints(Unmod } } else { for (String location : orderedLocations) { - Utils.ValueHolder 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.getGatewayLocationEndpoint())) { + if (this.defaultEndpoint.equals(endpoint.v.getGatewayRegionalEndpoint())) { endpoints = new ArrayList<>(); break; } @@ -657,13 +654,13 @@ private UnmodifiableList getPreferredAvailableEndpoints(Unmod } if (endpoints.isEmpty()) { - endpoints.add(new RegionalEndpoints(fallbackEndpoint)); + endpoints.add(new RegionalRoutingContext(fallbackEndpoint)); } endpoints.addAll(unavailableEndpoints); } else { for (String location : orderedLocations) { - Utils.ValueHolder 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); @@ -673,15 +670,15 @@ private UnmodifiableList getPreferredAvailableEndpoints(Unmod } if (endpoints.isEmpty()) { - endpoints.add(new RegionalEndpoints(fallbackEndpoint)); + endpoints.add(new RegionalRoutingContext(fallbackEndpoint)); } return new UnmodifiableList<>(endpoints); } private void addEndpoints( Iterable gatewayDbAccountLocations, - Map endpointsByLocation, - Map regionByEndpoint, + Map endpointsByLocation, + Map regionByEndpoint, List parsedLocations) { if (gatewayDbAccountLocations != null) { @@ -692,14 +689,14 @@ private void addEndpoints( String location = gatewayDbAccountLocation.getName().toLowerCase(Locale.ROOT); URI endpoint = new URI(gatewayDbAccountLocation.getEndpoint().toLowerCase(Locale.ROOT)); - RegionalEndpoints regionalEndpoints = new RegionalEndpoints(endpoint); + RegionalRoutingContext regionalRoutingContext = new RegionalRoutingContext(endpoint); if (!endpointsByLocation.containsKey(location)) { - endpointsByLocation.put(location, regionalEndpoints); + endpointsByLocation.put(location, regionalRoutingContext); } - if (!regionByEndpoint.containsKey(regionalEndpoints)) { - regionByEndpoint.put(regionalEndpoints, location); + if (!regionByEndpoint.containsKey(regionalRoutingContext)) { + regionByEndpoint.put(regionalRoutingContext, location); } parsedLocations.add(gatewayDbAccountLocation.getName()); @@ -713,19 +710,19 @@ private void addEndpoints( } } - private UnmodifiableMap getEndpointsByLocation(Iterable gatewayLocations, - Utils.ValueHolder> orderedLocations, - Utils.ValueHolder> regionMap) { - Map endpointsByLocation = new CaseInsensitiveMap<>(); - Map regionByEndpoint = new CaseInsensitiveMap<>(); + 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); + regionMap.v = (UnmodifiableMap) UnmodifiableMap.unmodifiableMap(regionByEndpoint); - return (UnmodifiableMap) UnmodifiableMap.unmodifiableMap(endpointsByLocation); + return (UnmodifiableMap) UnmodifiableMap.unmodifiableMap(endpointsByLocation); } public boolean canUseMultipleWriteLocations() { @@ -800,61 +797,36 @@ private static boolean isExcludedRegionsSupplierConfigured(Supplier writeEndpoints; - private UnmodifiableList readEndpoints; + private UnmodifiableList writeEndpoints; + private UnmodifiableList readEndpoints; private UnmodifiableList preferredLocations; private UnmodifiableList effectivePreferredLocations; // lower-case region private UnmodifiableList availableWriteLocations; // lower-case region private UnmodifiableList availableReadLocations; - private UnmodifiableMap availableWriteEndpointsByLocation; - private UnmodifiableMap availableReadEndpointsByLocation; - private UnmodifiableMap regionNameByWriteEndpoint; - private UnmodifiableMap regionNameByReadEndpoint; + 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.availableWriteEndpointsByLocation - = (UnmodifiableMap) UnmodifiableMap.unmodifiableMap(new CaseInsensitiveMap<>()); + = (UnmodifiableMap) UnmodifiableMap.unmodifiableMap(new CaseInsensitiveMap<>()); this.availableReadEndpointsByLocation - = (UnmodifiableMap) UnmodifiableMap.unmodifiableMap(new CaseInsensitiveMap<>()); + = (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(new RegionalEndpoints(defaultEndpoint))); - this.writeEndpoints = new UnmodifiableList<>(Collections.singletonList(new RegionalEndpoints(defaultEndpoint))); + this.readEndpoints = new UnmodifiableList<>(Collections.singletonList(new RegionalRoutingContext(defaultEndpoint))); + this.writeEndpoints = new UnmodifiableList<>(Collections.singletonList(new RegionalRoutingContext(defaultEndpoint))); } public DatabaseAccountLocationsInfo(DatabaseAccountLocationsInfo other) { From e2f1b7eba8f60c9a292eb09af18434a3480b7a8b Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Wed, 26 Feb 2025 19:05:51 -0500 Subject: [PATCH 07/17] Fixing CI pipeline. --- .../azure/cosmos/implementation/RxGatewayStoreModelTest.java | 2 -- .../cosmos/implementation/ClientSideRequestStatistics.java | 5 +++++ .../main/java/com/azure/cosmos/implementation/Configs.java | 2 -- 3 files changed, 5 insertions(+), 4 deletions(-) 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 3c92fc5ea062..f474023bb79c 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 @@ -177,7 +177,6 @@ public void applySessionToken( Mockito.doReturn(sdkGlobalSessionToken).when(sessionContainer).resolveGlobalSessionToken(any()); GlobalEndpointManager globalEndpointManager = Mockito.mock(GlobalEndpointManager.class); - GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManager = Mockito.mock(GlobalPartitionEndpointManagerForCircuitBreaker.class); Mockito.doReturn(new RegionalRoutingContext(new URI("https://localhost"))) .when(globalEndpointManager).resolveServiceEndpoint(any()); @@ -248,7 +247,6 @@ 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 RegionalRoutingContext(new URI("https://localhost"))) .when(globalEndpointManager).resolveServiceEndpoint(any()); 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 ec945a0e4c6c..34326be06ec3 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 @@ -36,6 +36,8 @@ import java.util.concurrent.ConcurrentLinkedDeque; import java.util.stream.Collectors; +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; @@ -227,6 +229,9 @@ public void recordGatewayResponse( if (rxDocumentServiceRequest != null && rxDocumentServiceRequest.requestContext != null) { regionalRoutingContext = rxDocumentServiceRequest.requestContext.regionalRoutingContextToRoute; + + checkNotNull(regionalRoutingContext, "Argument 'regionalRoutingContext' cannot be null!"); + locationEndPoint = regionalRoutingContext.getGatewayRegionalEndpoint(); this.approximateInsertionCountInBloomFilter = rxDocumentServiceRequest.requestContext.getApproximateBloomFilterInsertionCount(); 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 6f409d333319..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; @@ -18,7 +17,6 @@ import java.net.URI; 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; From 33a7b92884d35574ffdaed07fc29f73b20708f64 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Wed, 26 Feb 2025 21:05:09 -0500 Subject: [PATCH 08/17] Fixing CI pipeline. --- .../RxGatewayStoreModelTest.java | 22 ++++++++++++++----- .../ConnectionStateListenerTest.java | 7 +++++- .../MetadataRequestRetryPolicyTests.java | 7 ++++-- .../routing/LocationCacheTest.java | 2 -- .../ClientSideRequestStatistics.java | 17 +++++++------- .../DocumentServiceRequestContext.java | 3 --- .../WebExceptionRetryPolicy.java | 16 +++++++++++--- 7 files changed, 49 insertions(+), 25 deletions(-) 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 f474023bb79c..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 @@ -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 RegionalRoutingContext(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 RegionalRoutingContext(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,6 +180,8 @@ public void applySessionToken( Mockito.doReturn(sdkGlobalSessionToken).when(sessionContainer).resolveGlobalSessionToken(any()); GlobalEndpointManager globalEndpointManager = Mockito.mock(GlobalEndpointManager.class); + URI locationEndpointToRoute = new URI("https://localhost"); + RegionalRoutingContext regionalRoutingContext = new RegionalRoutingContext(locationEndpointToRoute); Mockito.doReturn(new RegionalRoutingContext(new URI("https://localhost"))) .when(globalEndpointManager).resolveServiceEndpoint(any()); @@ -206,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); } 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/MetadataRequestRetryPolicyTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/MetadataRequestRetryPolicyTests.java index ab1aa9562033..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 @@ -554,10 +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) { - URI 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/routing/LocationCacheTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/routing/LocationCacheTest.java index 0e5c78b4a8b8..dd34c0e4676e 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 @@ -64,8 +64,6 @@ public class LocationCacheTest { private static HashMap EndpointByLocation = new HashMap<>(); - private final static Logger logger = LoggerFactory.getLogger(LocationCache.class); - static { EndpointByLocation.put("location1", LocationCacheTest.Location1Endpoint); EndpointByLocation.put("location2", LocationCacheTest.Location2Endpoint); 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 34326be06ec3..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 @@ -36,6 +36,7 @@ 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) @@ -224,15 +225,12 @@ public void recordGatewayResponse( } RegionalRoutingContext regionalRoutingContext = null; - URI locationEndPoint = null; if (rxDocumentServiceRequest != null && rxDocumentServiceRequest.requestContext != null) { - regionalRoutingContext = rxDocumentServiceRequest.requestContext.regionalRoutingContextToRoute; - - checkNotNull(regionalRoutingContext, "Argument 'regionalRoutingContext' cannot be null!"); - - locationEndPoint = regionalRoutingContext.getGatewayRegionalEndpoint(); + if (rxDocumentServiceRequest.requestContext.regionalRoutingContextToRoute != null) { + regionalRoutingContext = rxDocumentServiceRequest.requestContext.regionalRoutingContextToRoute; + } this.approximateInsertionCountInBloomFilter = rxDocumentServiceRequest.requestContext.getApproximateBloomFilterInsertionCount(); this.keywordIdentifiers = rxDocumentServiceRequest.requestContext.getKeywordIdentifiers(); @@ -240,12 +238,13 @@ public void recordGatewayResponse( 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, regionalRoutingContext)); } 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 298600afb0af..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 @@ -38,7 +38,6 @@ 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; @@ -80,7 +79,6 @@ public DocumentServiceRequestContext() {} public void routeToLocation(int locationIndex, boolean usePreferredLocations) { this.locationIndexToRoute = locationIndex; this.usePreferredLocations = usePreferredLocations; -// this.locationEndpointToRoute = null; this.regionalRoutingContextToRoute = null; } @@ -141,7 +139,6 @@ 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; 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 a9114156499d..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 @@ -15,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); @@ -25,6 +27,7 @@ public class WebExceptionRetryPolicy implements IRetryPolicy { private boolean isReadRequest; private int retryCount = 0; private RegionalRoutingContext regionalRoutingContext; + private URI overriddenEndpoint; public WebExceptionRetryPolicy() { durationTimer.start(); @@ -38,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.regionalRoutingContext.getGatewayRegionalEndpoint(), + this.regionalRoutingContext != null ? this.regionalRoutingContext.getGatewayRegionalEndpoint() : this.overriddenEndpoint, this.request.getOperationType(), this.retryCount, this.request.isAddressRefresh()); @@ -64,7 +72,8 @@ public Mono shouldRetry(Exception e) { .debug("WebExceptionRetryPolicy() Retrying on endpoint {}, operationType = {}, resourceType = {}, count = {}, " + "isAddressRefresh = {}, shouldForcedAddressRefresh = {}, " + "shouldForceCollectionRoutingMapRefresh = {}", - this.regionalRoutingContext.getGatewayRegionalEndpoint(), 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); @@ -78,7 +87,7 @@ public Mono shouldRetry(Exception e) { .debug( "WebExceptionRetryPolicy() No retrying on un-retryable exceptions on endpoint {}, operationType = {}, resourceType = {}, count = {}, " + "isAddressRefresh = {}", - this.regionalRoutingContext.getGatewayRegionalEndpoint(), + this.regionalRoutingContext != null ? this.regionalRoutingContext.getGatewayRegionalEndpoint() : this.overriddenEndpoint, this.request.getOperationType(), this.request.getResourceType(), this.retryCount, @@ -103,6 +112,7 @@ public void onBeforeSendRequest(RxDocumentServiceRequest request) { // set the initial response timeout this.request.setResponseTimeout(timeoutPolicy.getTimeoutAndDelaysList().get(0).getResponseTimeout()); this.regionalRoutingContext = request.requestContext.regionalRoutingContextToRoute; + this.overriddenEndpoint = request.getEndpointOverride(); } private boolean isOutOfRetries() { From f332219194b9719397bc2fd1a847baeae1c21fe7 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Wed, 26 Feb 2025 21:48:35 -0500 Subject: [PATCH 09/17] Fixing CI pipeline. --- .../WebExceptionRetryPolicyTest.java | 31 ++++++++++++++++--- 1 file changed, 26 insertions(+), 5 deletions(-) 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 79c7fdf1bbd7..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 @@ -38,7 +38,10 @@ public static Object[][] operationTypeProvider() { @Test(groups = {"unit"}) public void shouldRetryOnTimeoutForReadOperations() throws Exception { GlobalEndpointManager endpointManager = Mockito.mock(GlobalEndpointManager.class); - Mockito.doReturn(new RegionalRoutingContext(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(); @@ -53,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); @@ -90,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 RegionalRoutingContext(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(); @@ -105,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); @@ -144,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 RegionalRoutingContext(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(); @@ -159,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); @@ -196,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 RegionalRoutingContext(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)); @@ -210,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); @@ -223,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); @@ -237,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 RegionalRoutingContext(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(); @@ -251,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); From 2169c9413c73eebef4dafc8cbd774e47587dd6f5 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Wed, 26 Feb 2025 22:06:02 -0500 Subject: [PATCH 10/17] Addressing review comments. --- .../cosmos/implementation/ResourceThrottleRetryPolicy.java | 2 +- .../cosmos/implementation/routing/RegionalRoutingContext.java | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) 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 22c68665ce2c..aa2f89f11d8f 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(exception)); + return Mono.just(ShouldRetryResult.noRetryOnNonRelatedException(/*exception*/)); } if (!retryOnClientSideThrottledBatchRequests && 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 index 1d8b44d59e3a..f4154921d7d6 100644 --- 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 @@ -7,6 +7,10 @@ import java.util.Objects; public class RegionalRoutingContext { + + // IMPORTANT: + // Please reevaluate overridden equals() implementation when adding additional properties + // to this class private final URI gatewayRegionalEndpoint; public RegionalRoutingContext(URI gatewayRegionalEndpoint) { From 0d401999158ae4fad4babc475f2bdc3916d3e811 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Wed, 26 Feb 2025 22:06:19 -0500 Subject: [PATCH 11/17] Addressing review comments. --- .../cosmos/implementation/routing/RegionalRoutingContext.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 index f4154921d7d6..b2955dc47d54 100644 --- 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 @@ -9,8 +9,8 @@ public class RegionalRoutingContext { // IMPORTANT: - // Please reevaluate overridden equals() implementation when adding additional properties - // to this class + // Please reevaluate overridden equals() implementation + // when adding additional properties to this class private final URI gatewayRegionalEndpoint; public RegionalRoutingContext(URI gatewayRegionalEndpoint) { From 4b9a6fe626dfb3e661d654f01ef83a8bffe1dfed Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Wed, 26 Feb 2025 23:21:22 -0500 Subject: [PATCH 12/17] Modify `ResourceThrottleRetryPolicy`. --- .../cosmos/implementation/ResourceThrottleRetryPolicy.java | 2 +- .../java/com/azure/cosmos/implementation/ShouldRetryResult.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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 aa2f89f11d8f..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(/*exception*/)); + return Mono.just(ShouldRetryResult.errorOnNonRelatedException(exception)); } if (!retryOnClientSideThrottledBatchRequests && 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 56a0ed924c44..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,7 +41,7 @@ public static ShouldRetryResult noRetryOnNonRelatedException() { return new ShouldRetryResult(null, null, false, null, true); } - public static ShouldRetryResult noRetryOnNonRelatedException(Exception e) { + public static ShouldRetryResult errorOnNonRelatedException(Exception e) { return new ShouldRetryResult(null, e, false, null, true); } From 14853c02350dc5b80fdfa5b33efe0d7cbfca5e9c Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Wed, 26 Feb 2025 23:29:50 -0500 Subject: [PATCH 13/17] Pinning relevant GitHub issues. --- .../azure/cosmos/implementation/routing/LocationCache.java | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) 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 085968524743..0441eea8d4d4 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 @@ -120,7 +120,8 @@ && unavailableLocationsExpirationTimePassed()) { */ public List getAvailableReadEndpoints() { return this.locationInfo.availableReadEndpointsByLocation.values().stream().map(consolidatedLocationEndpoints -> { - // TODO: Integrate thinclient endpoints into fault injection + // 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()); } @@ -134,7 +135,8 @@ public List getAvailableReadEndpoints() { */ public List getAvailableWriteEndpoints() { return this.locationInfo.availableWriteEndpointsByLocation.values().stream().map(consolidatedLocationEndpoints -> { - // TODO: Integrate thinclient endpoints into fault injection + // 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()); } From 531ec327a75214b08a8befb26cd0fc8ff6454b4f Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Thu, 27 Feb 2025 17:57:38 -0500 Subject: [PATCH 14/17] Addressing review comments. --- .../routing/LocationCacheTest.java | 2 -- .../implementation/RxGatewayStoreModel.java | 4 --- .../implementation/routing/LocationCache.java | 30 +++++++++---------- 3 files changed, 15 insertions(+), 21 deletions(-) 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 dd34c0e4676e..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 @@ -21,8 +21,6 @@ import com.azure.cosmos.models.ModelBridgeUtils; import com.azure.cosmos.implementation.guava25.collect.ImmutableList; import com.azure.cosmos.implementation.guava25.collect.Iterables; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; import org.testng.annotations.AfterClass; import org.testng.annotations.DataProvider; import org.testng.annotations.Test; 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 2d1d43e8c692..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 @@ -315,10 +315,6 @@ private HttpHeaders getHttpRequestHeaders(Map headers) { return httpHeaders; } - public URI getRootUri(RxDocumentServiceRequest request) { - return this.globalEndpointManager.resolveServiceEndpoint(request).getGatewayRegionalEndpoint(); - } - private URI getUri(RxDocumentServiceRequest request) throws URISyntaxException { URI rootUri = request.getEndpointOverride(); if (rootUri == null) { 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 0441eea8d4d4..f0050434fe42 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,7 +42,7 @@ 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; @@ -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(); @@ -216,7 +216,7 @@ public RegionalRoutingContext resolveServiceEndpoint(RxDocumentServiceRequest re String writeLocation = currentLocationInfo.availableWriteLocations.get(locationIndex); return currentLocationInfo.availableWriteEndpointsByLocation.get(writeLocation); } else { - return new RegionalRoutingContext(this.defaultEndpoint); + return this.defaultRegionalRoutingContext; } } else { UnmodifiableList endpoints = @@ -255,7 +255,7 @@ public UnmodifiableList getApplicableWriteEndpoints(List return this.getApplicableEndpoints( writeEndpoints, this.locationInfo.regionNameByWriteEndpoint, - this.defaultEndpoint, + this.defaultRegionalRoutingContext, effectiveExcludedRegionsWithPartitionUnavailableRegions); } @@ -288,14 +288,14 @@ public UnmodifiableList getApplicableReadEndpoints(List< return this.getApplicableEndpoints( readEndpoints, this.locationInfo.regionNameByReadEndpoint, - this.locationInfo.writeEndpoints.get(0).getGatewayRegionalEndpoint(), // match the fallback region used in getPreferredAvailableEndpoints + this.locationInfo.writeEndpoints.get(0), // match the fallback region used in getPreferredAvailableEndpoints effectiveExcludedRegionsWithPartitionUnavailableRegions); } private UnmodifiableList getApplicableEndpoints( UnmodifiableList endpoints, UnmodifiableMap regionNameByEndpoint, - URI fallbackEndpoint, + RegionalRoutingContext fallbackRegionalRoutingContext, List excludeRegionList) { List applicableEndpoints = new ArrayList<>(); @@ -309,7 +309,7 @@ private UnmodifiableList getApplicableEndpoints( } if (applicableEndpoints.isEmpty()) { - applicableEndpoints.add(new RegionalRoutingContext(fallbackEndpoint)); + applicableEndpoints.add(fallbackRegionalRoutingContext); } return new UnmodifiableList<>(applicableEndpoints); @@ -339,7 +339,7 @@ public URI resolveFaultInjectionEndpoint(String region, boolean writeOnly) { } public URI getDefaultEndpoint() { - return this.defaultEndpoint; + return this.defaultRegionalRoutingContext.getGatewayRegionalEndpoint(); } public boolean shouldRefreshEndpoints(Utils.ValueHolder canRefreshInBackground) { @@ -586,15 +586,15 @@ private void updateLocationCache( nextLocationInfo.regionNameByWriteEndpoint = outWriteRegionMap.v; } - nextLocationInfo.writeEndpoints = this.getPreferredAvailableEndpoints(nextLocationInfo.availableWriteEndpointsByLocation, nextLocationInfo.availableWriteLocations, OperationType.Write, this.defaultEndpoint); - nextLocationInfo.readEndpoints = this.getPreferredAvailableEndpoints(nextLocationInfo.availableReadEndpointsByLocation, nextLocationInfo.availableReadLocations, OperationType.Read, nextLocationInfo.writeEndpoints.get(0).getGatewayRegionalEndpoint()); + 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, new RegionalRoutingContext(this.defaultEndpoint), regionForDefaultEndpoint)) { + if (!Utils.tryGetValue(nextLocationInfo.regionNameByReadEndpoint, this.defaultRegionalRoutingContext, regionForDefaultEndpoint)) { nextLocationInfo.effectivePreferredLocations = nextLocationInfo.availableReadLocations; } } @@ -610,7 +610,7 @@ private void updateLocationCache( private UnmodifiableList getPreferredAvailableEndpoints(UnmodifiableMap endpointsByLocation, UnmodifiableList orderedLocations, OperationType expectedAvailableOperation, - URI fallbackEndpoint) { + 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 @@ -641,7 +641,7 @@ private UnmodifiableList getPreferredAvailableEndpoints( // if defaultEndpoint equals a regional endpoint then use // whatever the fallback endpoint is - if (this.defaultEndpoint.equals(endpoint.v.getGatewayRegionalEndpoint())) { + if (this.defaultRegionalRoutingContext.getGatewayRegionalEndpoint().equals(endpoint.v.getGatewayRegionalEndpoint())) { endpoints = new ArrayList<>(); break; } @@ -656,7 +656,7 @@ private UnmodifiableList getPreferredAvailableEndpoints( } if (endpoints.isEmpty()) { - endpoints.add(new RegionalRoutingContext(fallbackEndpoint)); + endpoints.add(fallbackRegionalRoutingContext); } endpoints.addAll(unavailableEndpoints); @@ -672,7 +672,7 @@ private UnmodifiableList getPreferredAvailableEndpoints( } if (endpoints.isEmpty()) { - endpoints.add(new RegionalRoutingContext(fallbackEndpoint)); + endpoints.add(fallbackRegionalRoutingContext); } return new UnmodifiableList<>(endpoints); } From 65199068f3bd96224764ac9cc0cc20d35bcd6b8d Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Sun, 2 Mar 2025 21:42:44 -0500 Subject: [PATCH 15/17] Addressing review comments. --- .../cosmos/implementation/routing/LocationCache.java | 8 ++++++-- .../implementation/routing/RegionalRoutingContext.java | 4 +++- 2 files changed, 9 insertions(+), 3 deletions(-) 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 f0050434fe42..a036e9b93bd9 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 @@ -703,9 +703,13 @@ private void addEndpoints( parsedLocations.add(gatewayDbAccountLocation.getName()); } catch (Exception e) { - logger.warn("GetAvailableEndpointsByLocation() - skipping add for location = [{}] as its location name is either empty or endpoint is malformed [{}]", + + String region = "GetAvailableEndpointsByLocation() - skipping add for location = [{}] as its location name is either empty or endpoint is malformed [{}]"; + + logger.warn("Skipping add for location = [{}] and endpoint = [{}] due to exception [{}]", gatewayDbAccountLocation.getName(), - gatewayDbAccountLocation.getEndpoint()); + gatewayDbAccountLocation.getEndpoint(), + e.getMessage()); } } } 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 index b2955dc47d54..6a796bb9866a 100644 --- 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 @@ -12,9 +12,11 @@ public class RegionalRoutingContext { // 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() { @@ -31,6 +33,6 @@ public boolean equals(Object o) { @Override public int hashCode() { - return Objects.hash(this.gatewayRegionalEndpoint); + return Objects.hash(this.gatewayRegionalEndpointAsString); } } From eca64cf5ac230405784f49a5afe64c8bbf11d113 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Sun, 2 Mar 2025 21:49:26 -0500 Subject: [PATCH 16/17] Addressing review comments. --- .../com/azure/cosmos/implementation/routing/LocationCache.java | 2 -- 1 file changed, 2 deletions(-) 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 a036e9b93bd9..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 @@ -704,8 +704,6 @@ private void addEndpoints( parsedLocations.add(gatewayDbAccountLocation.getName()); } catch (Exception e) { - String region = "GetAvailableEndpointsByLocation() - skipping add for location = [{}] as its location name is either empty or endpoint is malformed [{}]"; - logger.warn("Skipping add for location = [{}] and endpoint = [{}] due to exception [{}]", gatewayDbAccountLocation.getName(), gatewayDbAccountLocation.getEndpoint(), From 8eb9f465c98e3e0f91af8c96573ea4f7a035dcad Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Sun, 2 Mar 2025 21:50:20 -0500 Subject: [PATCH 17/17] Addressing review comments. --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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)