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..7682743b3a96 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.consolidatedRegionalEndpointToRoute = new LocationCache.ConsolidatedRegionalEndpoint(regionEndpoint, null); faultInjectionAddressRequest.setPartitionKeyRangeIdentity(new PartitionKeyRangeIdentity(pkRangeId)); if (isWriteOnly) { diff --git a/sdk/cosmos/azure-cosmos-tests/pom.xml b/sdk/cosmos/azure-cosmos-tests/pom.xml index a0c21994fed7..14c9ecf3c88e 100644 --- a/sdk/cosmos/azure-cosmos-tests/pom.xml +++ b/sdk/cosmos/azure-cosmos-tests/pom.xml @@ -222,18 +222,6 @@ Licensed under the MIT License. test 2.17.2 - - com.azure - azure-cosmos - 4.67.0-beta.1 - test - - - com.azure - azure-cosmos - 4.67.0-beta.1 - test - 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..51f7b9fb1915 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.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getKey(), null)); 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.ConsolidatedRegionalEndpoint(uriStringPair.getLeft(), null)) .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.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getKey(), null)); 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.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getKey(), null)); 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.ConsolidatedRegionalEndpoint(uriStringPair.getLeft(), null)) .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.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getKey(), null)); } 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.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getKey(), null)); 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.ConsolidatedRegionalEndpoint(uriStringPair.getLeft(), null)) .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.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getKey(), null)); } 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.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getKey(), null)); 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.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getKey(), null)); 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.ConsolidatedRegionalEndpoint(uriStringPair.getLeft(), null)) .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.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getKey(), null)); } 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.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getKey(), null)); 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.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getKey(), null)); 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.ConsolidatedRegionalEndpoint(uriStringPair.getLeft(), null)) .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.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getKey(), null)); } 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.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getKey(), null)); 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.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getKey(), null)); } - locationSpecificHealthContext = locationEndpointToLocationSpecificContextForPartition.get(LocationEastUs2EndpointToLocationPair.getKey()); + locationSpecificHealthContext = locationEndpointToLocationSpecificContextForPartition.get(new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getKey(), null)); 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.ConsolidatedRegionalEndpoint(uriStringPair.getLeft(), null)) .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.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getKey(), null)); globalPartitionEndpointManagerForCircuitBreaker - .handleLocationExceptionForPartitionKeyRange(request, LocationEastUsEndpointToLocationPair.getKey()); + .handleLocationExceptionForPartitionKeyRange(request, new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getKey(), null)); globalPartitionEndpointManagerForCircuitBreaker - .handleLocationExceptionForPartitionKeyRange(request, LocationCentralUsEndpointToLocationPair.getKey()); + .handleLocationExceptionForPartitionKeyRange(request, new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getKey(), null)); } 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.ConsolidatedRegionalEndpoint(uriStringPair.getLeft(), null)) .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.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getKey(), null)); } 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.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getKey(), null)); LocationSpecificHealthContext locationSpecificHealthContext2 - = locationEndpointToLocationSpecificContextForPartitionForColl2.get(LocationEastUs2EndpointToLocationPair.getKey()); + = locationEndpointToLocationSpecificContextForPartitionForColl2.get(new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getKey(), null)); 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.ConsolidatedRegionalEndpoint(uriStringPair.getLeft(), null)) .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.ConsolidatedRegionalEndpoint(locationWithFailure, null)); 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.consolidatedRegionalEndpointToRoute = new LocationCache.ConsolidatedRegionalEndpoint(locationEndpointToRoute, null); + 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..c1c5e50d832c 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..75ec2b6d4013 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.ConsolidatedRegionalEndpoint::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.ConsolidatedRegionalEndpoint::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.ConsolidatedRegionalEndpoint::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.ConsolidatedRegionalEndpoint::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..f8ed32b0a9ee 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.ConsolidatedRegionalEndpoint(new URI("http://localhost"), null)).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.ConsolidatedRegionalEndpoint(new URI("http://localhost"), null)).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.ConsolidatedRegionalEndpoint(new URI("http://localhost"), null)).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.ConsolidatedRegionalEndpoint(new URI("http://localhost"), null)).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.ConsolidatedRegionalEndpoint(new URI("http://localhost"), null)).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.ConsolidatedRegionalEndpoint(new URI("http://localhost"), null)).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.ConsolidatedRegionalEndpoint(new URI("http://localhost"), null)).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.ConsolidatedRegionalEndpoint(new URI("http://localhost"), null)).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.ConsolidatedRegionalEndpoint(new URI("http://localhost"), null)).when(endpointManager).resolveServiceEndpoint(Mockito.any(RxDocumentServiceRequest.class)); Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(Mockito.eq(null), Mockito.eq(false)); Mockito.doReturn(2).when(endpointManager).getPreferredLocationCount(); ClientRetryPolicy clientRetryPolicy = new ClientRetryPolicy(mockDiagnosticsClientContext(), endpointManager, true, retryOptions, null, globalPartitionEndpointManager); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java index 26c34a1c423c..8e480f3ed666 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ConfigsTests.java @@ -6,8 +6,6 @@ import com.azure.cosmos.implementation.clienttelemetry.MetricCategory; import com.azure.cosmos.implementation.clienttelemetry.TagName; import com.azure.cosmos.implementation.directconnectivity.Protocol; -import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils; -import io.netty.handler.ssl.SslContext; import org.testng.annotations.Test; import java.net.URI; @@ -170,12 +168,12 @@ public void http2MaxConcurrentStreams() { @Test(groups = { "unit" }) public void thinClientEnabledTest() { Configs config = new Configs(); - assertThat(config.getThinclientEnabled()).isFalse(); + assertThat(config.isThinClientEnabled()).isFalse(); System.clearProperty("COSMOS.THINCLIENT_ENABLED"); System.setProperty("COSMOS.THINCLIENT_ENABLED", "true"); try { - assertThat(config.getThinclientEnabled()).isTrue(); + assertThat(config.isThinClientEnabled()).isTrue(); } finally { System.clearProperty("COSMOS.THINCLIENT_ENABLED"); } 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..02a735a8f451 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.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getLeft(), null))); 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.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getLeft(), null))); 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.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getLeft(), null))); 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.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getLeft(), null))); 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.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getLeft(), null))); 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.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getLeft(), null))); 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.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getLeft(), null))); 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.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getLeft(), null))); 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.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getLeft(), null))); 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.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getLeft(), null))); 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.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getLeft(), null))); 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.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getLeft(), null))); 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.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getLeft(), null))); 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.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getLeft(), null))); 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.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getLeft(), null))); 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.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getLeft(), null))); 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.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getLeft(), null))); 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.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getLeft(), null))); 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.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getLeft(), null))); 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.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getLeft(), null))); 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.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getLeft(), null))); 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.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getLeft(), null))); 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.ConsolidatedRegionalEndpoint(LocationEastUsEndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationEastUs2EndpointToLocationPair.getLeft(), null), + new LocationCache.ConsolidatedRegionalEndpoint(LocationCentralUsEndpointToLocationPair.getLeft(), null))); 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.ConsolidatedRegionalEndpoint(uriToLocationMappings.getLeft(), null)) + .collect(Collectors.toList()); + + List consolidatedReadRegionalEndpoints = readEndpoints + .stream() + .map(readEndpoint -> new LocationCache.ConsolidatedRegionalEndpoint(readEndpoint, null)) .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.consolidatedRegionalEndpointToRoute = new LocationCache.ConsolidatedRegionalEndpoint(locationEndpointToRoute, null); 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..44f5896961d1 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.ConsolidatedRegionalEndpoint consolidatedLocationEndpointToRoute = new LocationCache.ConsolidatedRegionalEndpoint(locationEndToRoute, null); + + 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.ConsolidatedRegionalEndpoint consolidatedLocationEndpointToRoute = new LocationCache.ConsolidatedRegionalEndpoint(locationEndToRoute, null); + + 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.ConsolidatedRegionalEndpoint consolidatedLocationEndpointToRoute = new LocationCache.ConsolidatedRegionalEndpoint(locationEndToRoute, null); + + Mockito.when(endpointManager.resolveServiceEndpoint(Mockito.any())).thenReturn(consolidatedLocationEndpointToRoute); + Mockito.doReturn(Mono.empty()).when(endpointManager).refreshLocationAsync(eq(null), eq(false)); IRetryPolicyFactory retryPolicyFactory = new RetryPolicy(mockDiagnosticsClientContext(), endpointManager, ConnectionPolicy.getDefaultPolicy(), globalPartitionEndpointManager); RxClientCollectionCache rxClientCollectionCache = Mockito.mock(RxClientCollectionCache.class); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxGatewayStoreModelTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxGatewayStoreModelTest.java index bd26bb53a0b5..ed505f63fa53 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.ConsolidatedRegionalEndpoint(new URI("https://localhost"), null)) .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.ConsolidatedRegionalEndpoint(new URI("https://localhost"), null)) .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.ConsolidatedRegionalEndpoint(new URI("https://localhost"), null)) .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.ConsolidatedRegionalEndpoint(new URI("https://localhost"), null)) .when(globalEndpointManager).resolveServiceEndpoint(any()); HttpClient httpClient = Mockito.mock(HttpClient.class); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ThinClientStoreModelTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ThinClientStoreModelTest.java index 9b1b60c6fd7f..f53393b6ef90 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 @@ -2,6 +2,7 @@ import com.azure.cosmos.ConsistencyLevel; import com.azure.cosmos.implementation.http.HttpClient; +import com.azure.cosmos.implementation.routing.LocationCache; import io.netty.channel.ConnectTimeoutException; import org.mockito.Mockito; import org.testng.annotations.Test; @@ -29,7 +30,7 @@ public void testThinClientStoreModel() throws Exception { GlobalEndpointManager globalEndpointManager = Mockito.mock(GlobalEndpointManager.class); - Mockito.doReturn(new URI("https://localhost")) + Mockito.doReturn(new LocationCache.ConsolidatedRegionalEndpoint(new URI("https://localhost"), null)) .when(globalEndpointManager).resolveServiceEndpoint(any()); // mocking with HTTP/1.1 client, just using this test as basic store model validation. e2e request flow diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ThinClientTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ThinClientTest.java index 71b6a47c4de8..3889fa2bbe3c 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ThinClientTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ThinClientTest.java @@ -1,14 +1,18 @@ package com.azure.cosmos.implementation; +import com.azure.cosmos.ConsistencyLevel; import com.azure.cosmos.CosmosAsyncClient; import com.azure.cosmos.CosmosAsyncContainer; import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.GatewayConnectionConfig; import com.azure.cosmos.implementation.throughputControl.TestItem; import com.azure.cosmos.models.CosmosDatabaseProperties; +import com.azure.cosmos.models.CosmosDatabaseResponse; import com.azure.cosmos.models.PartitionKey; import com.azure.cosmos.util.CosmosPagedFlux; import com.fasterxml.jackson.databind.JsonNode; import org.testng.annotations.Test; +import reactor.core.publisher.Mono; public class ThinClientTest { @Test @@ -16,7 +20,6 @@ public void testThinclientHttp2() { try { //String thinclientEndpoint = "https://cdb-ms-stage-eastus2-fe2.eastus2.cloudapp.azure.com:10650"; String thinclientEndpoint = "https://chukangzhongstagesignoff.documents-staging.windows-ppe.net:443/"; - // TODO: figure out how to set up auth, using TestConfigurations key for now System.setProperty(Configs.THINCLIENT_ENABLED, "true"); System.setProperty(Configs.THINCLIENT_ENDPOINT, thinclientEndpoint); System.setProperty(Configs.HTTP2_ENABLED, "true"); @@ -25,15 +28,20 @@ public void testThinclientHttp2() { .key(TestConfigurations.MASTER_KEY) .endpoint(TestConfigurations.HOST) .gatewayMode() + .consistencyLevel(ConsistencyLevel.SESSION) .buildAsyncClient(); - //CosmosPagedFlux feedObservable = client.readAllDatabases(); - - CosmosAsyncContainer container = client.getDatabase("TestDatabase").getContainer("ChangeFeedTestContainer"); + CosmosAsyncContainer container = client.getDatabase("NehaTestDb").getContainer("NehaTestContainer"); TestItem testItem = TestItem.createNewItem(); System.out.println(testItem.getId()); container.createItem(testItem).block(); container.readItem(testItem.getId(), new PartitionKey(testItem.getId()), JsonNode.class).block(); + +/* CosmosAsyncContainer container = client.getDatabase("TestDatabase").getContainer("ChangeFeedTestContainer"); + TestItem testItem = TestItem.createNewItem(); + System.out.println(testItem.getId()); + container.createItem(testItem).block(); + container.readItem(testItem.getId(), new PartitionKey(testItem.getId()), JsonNode.class).block();*/ } finally { System.clearProperty(Configs.THINCLIENT_ENABLED); System.clearProperty(Configs.THINCLIENT_ENDPOINT); 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..f576126daca6 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.ConsolidatedRegionalEndpoint(new URI("http://localhost:"), null)).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.ConsolidatedRegionalEndpoint(new URI("http://localhost:"), null)).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.ConsolidatedRegionalEndpoint(new URI("http://localhost:"), null)).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.ConsolidatedRegionalEndpoint(new URI("http://localhost:"), null)).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.ConsolidatedRegionalEndpoint(new URI("http://localhost:"), null)).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..483adfe5e328 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 @@ -25,6 +25,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 +79,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.ConsolidatedRegionalEndpoint(urlforRead1, null)); + readEndPointList.add(new LocationCache.ConsolidatedRegionalEndpoint(urlforRead2, null)); + readEndPointList.add(new LocationCache.ConsolidatedRegionalEndpoint(urlforRead3, null)); + 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.ConsolidatedRegionalEndpoint(urlforWrite1, null)); + writeEndPointList.add(new LocationCache.ConsolidatedRegionalEndpoint(urlforWrite2, null)); + writeEndPointList.add(new LocationCache.ConsolidatedRegionalEndpoint(urlforWrite3, null)); + UnmodifiableList writeList = new UnmodifiableList<>(writeEndPointList); Mockito.when(endpointManager.getReadEndpoints()).thenReturn(readList); Mockito.when(endpointManager.getWriteEndpoints()).thenReturn(writeList); @@ -122,13 +123,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.ConsolidatedRegionalEndpoint testUrl = new LocationCache.ConsolidatedRegionalEndpoint(new URI("http://Test.com/"), null); 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 +157,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.ConsolidatedRegionalEndpoint(urlforRead1, null), new LocationCache.ConsolidatedRegionalEndpoint(urlforRead2, null)))); 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..85433c8cbb82 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..f2e08bb3c306 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.consolidatedRegionalEndpointToRoute = new LocationCache.ConsolidatedRegionalEndpoint(request.requestContext.locationEndpointToRoute, null); } 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..33ca6c7b089a 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.ConsolidatedRegionalEndpoint consolidatedRegionalEndpoint; try { - url = new URI("http://localhost"); + consolidatedRegionalEndpoint = new LocationCache.ConsolidatedRegionalEndpoint(new URI("http://localhost"), null); } 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(consolidatedRegionalEndpoint).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..4b162208e885 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; @@ -63,6 +65,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 +484,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.ConsolidatedRegionalEndpoint(LocationCacheTest.Location1Endpoint, null)); + assertThat(this.cache.getWriteEndpoints().get(1)).isEqualTo(new LocationCache.ConsolidatedRegionalEndpoint(LocationCacheTest.Location2Endpoint, null)); + assertThat(this.cache.getWriteEndpoints().get(2)).isEqualTo(new LocationCache.ConsolidatedRegionalEndpoint(LocationCacheTest.Location3Endpoint, null)); } @Test(groups = "unit", dataProvider = "excludedRegionsTestConfigs") @@ -512,11 +516,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 +536,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 +671,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 +779,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 +902,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.ConsolidatedRegionalEndpoint(firstAvailableWriteEndpoint, null)).isEqualTo(writeEndpoints.get(0)); + assertThat(new LocationCache.ConsolidatedRegionalEndpoint(secondAvailableWriteEndpoint, null)).isEqualTo(this.resolveEndpointForWriteRequest(ResourceType.Document, true)); + assertThat(new LocationCache.ConsolidatedRegionalEndpoint(firstAvailableWriteEndpoint, null)).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.ConsolidatedRegionalEndpoint(firstWriteEnpoint, null)).isEqualTo(this.resolveEndpointForWriteRequest(ResourceType.Database, false)); + assertThat(new LocationCache.ConsolidatedRegionalEndpoint(secondWriteEnpoint, null)).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.ConsolidatedRegionalEndpoint(firstAvailableReadEndpoint, null)).isEqualTo(this.resolveEndpointForReadRequest(true)); + assertThat(new LocationCache.ConsolidatedRegionalEndpoint(firstAvailableReadEndpoint, null)).isEqualTo(this.resolveEndpointForReadRequest(false)); } - private URI resolveEndpointForReadRequest(boolean masterResourceType) { + private LocationCache.ConsolidatedRegionalEndpoint 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.ConsolidatedRegionalEndpoint 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..0c565330ecf5 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.ConsolidatedRegionalEndpoint getFirstContactedLocationEndpoint() { return this.clientSideRequestStatistics.getFirstContactedLocationEndpoint(); } @@ -478,7 +479,7 @@ public void setDiagnosticsContext(CosmosDiagnostics cosmosDiagnostics, CosmosDia } @Override - public URI getFirstContactedLocationEndpoint(CosmosDiagnostics cosmosDiagnostics) { + public LocationCache.ConsolidatedRegionalEndpoint 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 1855fd253fa6..a281c8933f7a 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 @@ -19,6 +19,7 @@ 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,7 +47,8 @@ public class ClientRetryPolicy extends DocumentClientRetryPolicy { private int staleContainerRetryCount; private boolean isReadRequest; private boolean canUseMultipleWriteLocations; - private LocationCache.LocationEndpoints locationEndpoints; + private URI locationEndpoint; + private LocationCache.ConsolidatedRegionalEndpoint consolidatedRegionalEndpoint; private RetryContext retryContext; private CosmosDiagnostics cosmosDiagnostics; private AtomicInteger cnt = new AtomicInteger(0); @@ -87,8 +89,7 @@ public Mono shouldRetry(Exception e) { isReadRequest, canUseMultipleWriteLocations, e); - // TODO: verify this logic - if (this.locationEndpoints.gatewayEndpoint == null && this.locationEndpoints.thinClientEndpoint == null) { + if (this.locationEndpoint == null || this.consolidatedRegionalEndpoint == 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."); @@ -230,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); @@ -306,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.consolidatedRegionalEndpointToRoute); } //if operation is data plane read, metadata read, or query plan it can be retried on a different endpoint. @@ -337,19 +338,13 @@ private Mono shouldNotRetryOnEndpointFailureAsync(boolean isR private Mono refreshLocation(boolean isReadRequest, boolean forceRefresh, boolean usePreferredLocations) { this.failoverRetryCount++; - // Mark the current read endpoints as unavailable + // Mark the current read endpoint as unavailable if (isReadRequest) { - logger.warn("marking the endpoint {} as unavailable for read",this.locationEndpoints.gatewayEndpoint); - logger.warn("marking the endpoint {} as unavailable for read",this.locationEndpoints.thinClientEndpoint); - // adding an extra call here will cause updateLocationCache to be called twice, is this ok or should we consolidate - this.globalEndpointManager.markEndpointUnavailableForRead(this.locationEndpoints.gatewayEndpoint); - this.globalEndpointManager.markEndpointUnavailableForRead(this.locationEndpoints.thinClientEndpoint); + logger.warn("marking the endpoint {} as unavailable for read",this.locationEndpoint); + this.globalEndpointManager.markEndpointUnavailableForRead(this.consolidatedRegionalEndpoint.getGatewayLocationEndpoint()); } else { - logger.warn("marking the endpoint {} as unavailable for write",this.locationEndpoints.gatewayEndpoint); - logger.warn("marking the endpoint {} as unavailable for write",this.locationEndpoints.thinClientEndpoint); - // adding an extra call here will cause updateLocationCache to be called twice, is this ok or should we consolidate - this.globalEndpointManager.markEndpointUnavailableForWrite(this.locationEndpoints.gatewayEndpoint); - this.globalEndpointManager.markEndpointUnavailableForWrite(this.locationEndpoints.thinClientEndpoint); + logger.warn("marking the endpoint {} as unavailable for write",this.locationEndpoint); + this.globalEndpointManager.markEndpointUnavailableForWrite(this.consolidatedRegionalEndpoint.getGatewayLocationEndpoint()); } this.retryContext = new RetryContext(this.failoverRetryCount, usePreferredLocations); @@ -364,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.consolidatedRegionalEndpointToRoute); } // The request has failed with 503, SDK need to decide whether it is safe to retry for write operations @@ -409,10 +404,7 @@ private Mono shouldRetryOnBackendServiceUnavailableAsync( return Mono.just(ShouldRetryResult.noRetry()); } - logger.info("shouldRetryOnServiceUnavailable() Retrying. Received on endpoints {}, {}, IsReadRequest = {}", - this.locationEndpoints.gatewayEndpoint, - this.locationEndpoints.thinClientEndpoint, - isReadRequest); + logger.info("shouldRetryOnServiceUnavailable() Retrying. Received on endpoint {}, IsReadRequest = {}", this.locationEndpoint, isReadRequest); // Retrying on second PreferredLocations // RetryCount is used as zero-based index @@ -426,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.consolidatedRegionalEndpointToRoute); } } @@ -438,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.consolidatedRegionalEndpointToRoute); } return Mono.just(ShouldRetryResult.NO_RETRY); @@ -469,11 +463,17 @@ 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.locationEndpoints = this.globalEndpointManager.resolveServiceEndpoint(request); + this.consolidatedRegionalEndpoint = this.globalEndpointManager.resolveServiceEndpoint(request); + this.locationEndpoint = request.useThinProxy && this.consolidatedRegionalEndpoint.getThinClientLocationEndpoint() != null ? + this.consolidatedRegionalEndpoint.getThinClientLocationEndpoint() : + this.consolidatedRegionalEndpoint.getGatewayLocationEndpoint(); + + if (this.consolidatedRegionalEndpoint.getThinClientLocationEndpoint() == null) { + request.useThinProxy = false; + } + if (request.requestContext != null) { - // TODO: try both endpoints before we force cross-region retry - request.requestContext.locationEndpoints = this.locationEndpoints; - request.requestContext.routeToLocation(Configs.getThinclientEnabled() ? this.locationEndpoints.thinClientEndpoint : this.locationEndpoints.gatewayEndpoint); + request.requestContext.routeToLocation(this.locationEndpoint, this.consolidatedRegionalEndpoint); } } @@ -519,4 +519,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.consolidatedRegionalEndpointToRoute, "Argument 'request.requestContext.consolidatedRegionalEndpointToRoute' must not be null'"); + + return request.requestContext.consolidatedRegionalEndpointToRoute.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..72c05373f21a 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.ConsolidatedRegionalEndpoint consolidatedRegionalEndpoint = 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; + consolidatedRegionalEndpoint = request.requestContext.consolidatedRegionalEndpointToRoute; 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 (consolidatedRegionalEndpoint != null) { storeResponseStatistics.regionName = - globalEndpointManager.getRegionName(locationEndPoint, request.getOperationType()); + globalEndpointManager.getRegionName(consolidatedRegionalEndpoint.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, consolidatedRegionalEndpoint)); } if (storeResponseStatistics.requestOperationType == OperationType.Head @@ -218,9 +222,13 @@ public void recordGatewayResponse( this.requestEndTimeUTC = responseTime; } + LocationCache.ConsolidatedRegionalEndpoint consolidatedRegionalEndpoint = null; URI locationEndPoint = null; if (rxDocumentServiceRequest != null && rxDocumentServiceRequest.requestContext != null) { + + consolidatedRegionalEndpoint = rxDocumentServiceRequest.requestContext.consolidatedRegionalEndpointToRoute; 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, consolidatedRegionalEndpoint)); } GatewayStatistics gatewayStatistics = new GatewayStatistics(); @@ -651,7 +659,7 @@ public String getFirstContactedRegion() { return this.regionsContactedWithContext.first().regionContacted; } - public URI getFirstContactedLocationEndpoint() { + public LocationCache.ConsolidatedRegionalEndpoint 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.ConsolidatedRegionalEndpoint locationEndpointsContacted; private final long recordedTimestamp; - RegionWithContext(String regionContacted, URI locationEndpointsContacted) { + RegionWithContext(String regionContacted, LocationCache.ConsolidatedRegionalEndpoint 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 75606dd768e6..6c19e759920d 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 @@ -15,10 +15,8 @@ import javax.net.ssl.SSLException; import java.net.URI; -import java.net.URISyntaxException; import java.time.Duration; import java.util.Locale; -import java.util.Objects; import static com.azure.cosmos.implementation.guava25.base.MoreObjects.firstNonNull; import static com.azure.cosmos.implementation.guava25.base.Strings.emptyToNull; @@ -428,7 +426,7 @@ public URI getThinclientEndpoint() { return URI.create(DEFAULT_THINCLIENT_ENDPOINT); } - public static boolean getThinclientEnabled() { + public static boolean isThinClientEnabled() { String valueFromSystemProperty = System.getProperty(THINCLIENT_ENABLED); if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) { return Boolean.parseBoolean(valueFromSystemProperty); 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 6398bf637929..5f34f5653225 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,6 @@ 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; @@ -41,7 +40,7 @@ public class DocumentServiceRequestContext implements Cloneable { public volatile Boolean usePreferredLocations; public volatile Integer locationIndexToRoute; public volatile URI locationEndpointToRoute; - public volatile LocationCache.LocationEndpoints locationEndpoints; + public volatile LocationCache.ConsolidatedRegionalEndpoint consolidatedRegionalEndpointToRoute; public volatile boolean performedBackgroundAddressRefresh; public volatile boolean performLocalRefreshOnGoneException; public volatile List storeResponses; @@ -83,6 +82,7 @@ public void routeToLocation(int locationIndex, boolean usePreferredLocations) { this.locationIndexToRoute = locationIndex; this.usePreferredLocations = usePreferredLocations; this.locationEndpointToRoute = null; + this.consolidatedRegionalEndpointToRoute = null; } /** @@ -91,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.ConsolidatedRegionalEndpoint consolidatedRegionalEndpointToRoute) { this.locationEndpointToRoute = locationEndpoint; + this.consolidatedRegionalEndpointToRoute = consolidatedRegionalEndpointToRoute; this.locationIndexToRoute = null; this.usePreferredLocations = null; } @@ -103,6 +104,7 @@ public void routeToLocation(URI locationEndpoint) { public void clearRouteToLocation() { this.locationIndexToRoute = null; this.locationEndpointToRoute = null; + this.consolidatedRegionalEndpointToRoute = null; this.usePreferredLocations = null; } @@ -144,6 +146,7 @@ public DocumentServiceRequestContext clone() { context.usePreferredLocations = this.usePreferredLocations; context.locationIndexToRoute = this.locationIndexToRoute; context.locationEndpointToRoute = this.locationEndpointToRoute; + context.consolidatedRegionalEndpointToRoute = this.consolidatedRegionalEndpointToRoute; 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 67b5e84c6579..b0cd09f37c86 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,11 +146,11 @@ public static Mono getDatabaseAccountFromAnyLocationsAsync( }); } - public LocationCache.LocationEndpoints resolveServiceEndpoint(RxDocumentServiceRequest request) { - LocationCache.LocationEndpoints serviceEndpoints = this.locationCache.resolveServiceEndpoint(request); + public LocationCache.ConsolidatedRegionalEndpoint resolveServiceEndpoint(RxDocumentServiceRequest request) { + LocationCache.ConsolidatedRegionalEndpoint serviceEndpoints = this.locationCache.resolveServiceEndpoint(request); if (request.faultInjectionRequestContext != null) { // TODO: integrate thin client into fault injection - request.faultInjectionRequestContext.setLocationEndpointToRoute(serviceEndpoints.gatewayEndpoint); + request.faultInjectionRequestContext.setLocationEndpointToRoute(serviceEndpoints.getGatewayLocationEndpoint()); } 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 73f692adf149..1afa8c2a5bbc 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.ConsolidatedRegionalEndpoint 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..06b549096daa 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.ConsolidatedRegionalEndpoint consolidatedRegionalEndpoint = request.requestContext.consolidatedRegionalEndpointToRoute; + URI locationEndpointToRoute = consolidatedRegionalEndpoint.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 cb934f9f39df..5dbe4a054451 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 = @@ -490,10 +491,8 @@ private RxDocumentClientImpl(URI serviceEndpoint, this.connectionSharingAcrossClientsEnabled = connectionSharingAcrossClientsEnabled; this.configs = configs; this.masterKeyOrResourceToken = masterKeyOrResourceToken; - this.isThinClientEnabled = this.configs.getThinclientEnabled(); - // TODO: we have a default thinclient endpoint in configs, but it is not set to anything - need - // to figure out what default will be - this.serviceEndpoint = this.isThinClientEnabled ? this.configs.getThinclientEndpoint() : serviceEndpoint; + this.isThinClientEnabled = Configs.isThinClientEnabled(); + this.serviceEndpoint = serviceEndpoint; this.credential = credential; this.tokenCredential = tokenCredential; this.contentResponseOnWriteEnabled = contentResponseOnWriteEnabled; @@ -673,6 +672,14 @@ private void updateGatewayProxy() { (this.gatewayProxy).setSessionContainer(this.sessionContainer); } + private void updateThinProxy() { + (this.thinProxy).setGatewayServiceConfigurationReader(this.gatewayConfigurationReader); + (this.thinProxy).setCollectionCache(this.collectionCache); + (this.thinProxy).setPartitionKeyRangeCache(this.partitionKeyRangeCache); + (this.thinProxy).setUseMultipleWriteLocations(this.useMultipleWriteLocations); + (this.thinProxy).setSessionContainer(this.sessionContainer); + } + public void init(CosmosClientMetadataCachesSnapshot metadataCachesSnapshot, Function httpClientInterceptor) { try { @@ -722,6 +729,7 @@ public void init(CosmosClientMetadataCachesSnapshot metadataCachesSnapshot, Func collectionCache); updateGatewayProxy(); + updateThinProxy(); clientTelemetry = new ClientTelemetry( this, null, @@ -1885,7 +1893,9 @@ private Mono false for some reason, fix it + request.useThinProxy = Configs.isThinClientEnabled() && request.useGatewayMode ? true : false; + request.useThinProxy = true; if (operationType.isWriteOperation() && options != null && options.getNonIdempotentWriteRetriesEnabled() != null && options.getNonIdempotentWriteRetriesEnabled()) { request.setNonIdempotentWriteRetriesEnabled(true); @@ -5742,6 +5752,13 @@ public Flux getDatabaseAccountFromEndpoint(URI endpoint) { * @return RxStoreModel */ private RxStoreModel getStoreProxy(RxDocumentServiceRequest request) { + // TODO: request.useThinProxy is false here, fix it + + // TODO: remove, for testing + if (request.isMetadataRequest()) { + return this.gatewayProxy; + } + if (request.useThinProxy) { return this.thinProxy; } @@ -6486,7 +6503,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()) { @@ -6496,7 +6513,7 @@ private List getApplicableEndPoints(OperationType operationType, List withoutNulls(List orderedEffectiveEndpointsList) { + private static List withoutNulls(List orderedEffectiveEndpointsList) { if (orderedEffectiveEndpointsList == null) { return EMPTY_ENDPOINT_LIST; } @@ -6555,7 +6572,7 @@ private List getApplicableRegionsForSpeculation( return EMPTY_REGION_LIST; } - List endpoints = getApplicableEndPoints(operationType, excludedRegions); + List consolidatedRegionalEndpointList = getApplicableEndPoints(operationType, excludedRegions); HashSet normalizedExcludedRegions = new HashSet<>(); if (excludedRegions != null) { @@ -6563,8 +6580,8 @@ private List getApplicableRegionsForSpeculation( } List orderedRegionsForSpeculation = new ArrayList<>(); - endpoints.forEach(uri -> { - String regionName = this.globalEndpointManager.getRegionName(uri, operationType); + consolidatedRegionalEndpointList.forEach(consolidatedLocationEndpoints -> { + String regionName = this.globalEndpointManager.getRegionName(consolidatedLocationEndpoints.getGatewayLocationEndpoint(), operationType); if (!normalizedExcludedRegions.contains(regionName.toLowerCase(Locale.ROOT))) { orderedRegionsForSpeculation.add(regionName); } @@ -6768,7 +6785,7 @@ private Mono executeFeedOperationWithAvailabilityStrategy( private void handleLocationCancellationExceptionForPartitionKeyRange(RxDocumentServiceRequest failedRequest) { - URI firstContactedLocationEndpoint = diagnosticsAccessor + LocationCache.ConsolidatedRegionalEndpoint 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 28e0668977e6..1cc6f316388a 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 @@ -88,7 +88,7 @@ public class RxDocumentServiceRequest implements Cloneable { public String throughputControlGroupName; public volatile boolean intendedCollectionRidPassedIntoSDK = false; private volatile Duration responseTimeout; - public boolean useThinProxy = false; + 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 422653431cdd..f5820ad69c9d 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).gatewayEndpoint; + return this.globalEndpointManager.resolveServiceEndpoint(request).getGatewayLocationEndpoint(); } private URI getUri(RxDocumentServiceRequest request) throws URISyntaxException { @@ -326,7 +324,7 @@ private URI getUri(RxDocumentServiceRequest request) throws URISyntaxException { if (rootUri == null) { if (request.getIsMedia()) { // For media read request, always use the write endpoint. - rootUri = this.globalEndpointManager.getWriteEndpoints().get(0); + rootUri = this.globalEndpointManager.getWriteEndpoints().get(0).getGatewayLocationEndpoint(); } else { rootUri = getRootUri(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/ThinClientStoreModel.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java index 5d4998ed3659..a242d7d71db3 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ThinClientStoreModel.java @@ -8,6 +8,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.models.FeedRange; import io.netty.buffer.ByteBuf; import io.netty.buffer.Unpooled; import io.netty.handler.codec.http.HttpMethod; @@ -18,6 +19,7 @@ import java.util.HashMap; import java.util.Map; +import static com.azure.cosmos.implementation.directconnectivity.WFConstants.BackendHeaders.EFFECTIVE_PARTITION_KEY; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; /** @@ -77,7 +79,8 @@ protected Map getDefaultHeaders( @Override public URI getRootUri(RxDocumentServiceRequest request) { - return this.globalEndpointManager.resolveServiceEndpoint(request).thinClientEndpoint; + //var uri = this.globalEndpointManager.resolveServiceEndpoint(request).getThinClientLocationEndpoint(); + return URI.create("https://chukangzhongstagesignoff-eastus2.documents-staging.windows-ppe.net:10650/"); } @Override @@ -86,6 +89,12 @@ public HttpRequest wrapInHttpRequest(RxDocumentServiceRequest request, URI reque // todo - neharao1 - validate b/w name() v/s toString() request.setThinclientHeaders(request.getOperationType().name(), request.getResourceType().name()); + String epk = request.getPartitionKeyInternal().getEffectivePartitionKeyString(request.getPartitionKeyInternal(), request.getPartitionKeyDefinition()); + if (request.properties == null) { + request.properties = new HashMap<>(); + } + request.properties.put(EFFECTIVE_PARTITION_KEY, epk); + //request.getHeaders().put(EFFECTIVE_PARTITION_KEY, epk); // todo - neharao1: no concept of a replica / service endpoint that can be passed RntbdRequestArgs rntbdRequestArgs = new RntbdRequestArgs(request); @@ -106,8 +115,10 @@ public HttpRequest wrapInHttpRequest(RxDocumentServiceRequest request, URI reque return new HttpRequest( HttpMethod.POST, - requestUri, - requestUri.getPort(), + //requestUri, + URI.create("https://chukangzhongstagesignoff-eastus2.documents-staging.windows-ppe.net:10650/"), + //requestUri.getPort(), + 10650, headers, Flux.just(byteBuf.array())); } @@ -121,7 +132,6 @@ private HttpHeaders getHttpHeaders() { httpHeaders.set(header.getKey(), header.getValue()); } - // todo: add thin client resourcetype/operationtype headers return httpHeaders; } } 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..e4a74005508b 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.ConsolidatedRegionalEndpoint 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.ConsolidatedRegionalEndpoint succeededLocation = request.requestContext.consolidatedRegionalEndpointToRoute; 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.ConsolidatedRegionalEndpoint consolidatedRegionalEndpoint = pair.getKey(); + + URI gatewayLocationEndpoint = consolidatedRegionalEndpoint.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.ConsolidatedRegionalEndpoint 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.ConsolidatedRegionalEndpoint 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.ConsolidatedRegionalEndpoint 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.ConsolidatedRegionalEndpoint 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.ConsolidatedRegionalEndpoint 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..25094e6882d8 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.ConsolidatedRegionalEndpoint endpoint : endpointManager.getWriteEndpoints()) { + this.getOrAddEndpoint(endpoint.getGatewayLocationEndpoint()); } - for (URI endpoint : endpointManager.getReadEndpoints()) { - this.getOrAddEndpoint(endpoint); + for (LocationCache.ConsolidatedRegionalEndpoint 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.ConsolidatedRegionalEndpoint 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.ConsolidatedRegionalEndpoint 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..228841e30d93 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.consolidatedRegionalEndpointToRoute = new LocationCache.ConsolidatedRegionalEndpoint(serviceEndpoint, null); 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..b3d421dcfae5 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.ConsolidatedRegionalEndpoint consolidatedRegionalEndpointToRoute; /*** * This usually is called during retries. @@ -96,6 +99,7 @@ public String getFaultInjectionRuleId(long transportRequestId) { public void setLocationEndpointToRoute(URI locationEndpointToRoute) { this.locationEndpointToRoute = locationEndpointToRoute; + this.consolidatedRegionalEndpointToRoute = new LocationCache.ConsolidatedRegionalEndpoint(locationEndpointToRoute, null); } 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..d3a422bad013 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.ConsolidatedRegionalEndpoint 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 9892f3d5e3be..e5e6b7ba7ea9 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 @@ -20,12 +20,10 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import javax.xml.stream.Location; import java.net.URI; import java.time.Duration; import java.time.Instant; import java.util.ArrayList; -import java.util.Arrays; import java.util.Collections; import java.util.List; import java.util.Locale; @@ -44,11 +42,11 @@ public class LocationCache { private final static Logger logger = LoggerFactory.getLogger(LocationCache.class); private final boolean enableEndpointDiscovery; - private final LocationEndpoints defaultEndpoint; + private final URI defaultEndpoint; 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; @@ -67,7 +65,7 @@ public LocationCache( ); this.locationInfo = new DatabaseAccountLocationsInfo(preferredLocations, defaultEndpoint); - this.defaultEndpoint = new LocationEndpoints(defaultEndpoint); // assuming default endpoint as gateway + this.defaultEndpoint = defaultEndpoint; this.enableEndpointDiscovery = connectionPolicy.isEndpointDiscoveryEnabled(); this.useMultipleWriteLocations = connectionPolicy.isMultipleWriteRegionsEnabled(); @@ -88,7 +86,7 @@ public LocationCache( * 2. Endpoint availability * @return */ - public UnmodifiableList getReadEndpoints() { + public UnmodifiableList getReadEndpoints() { if (this.locationUnavailabilityInfoByEndpoint.size() > 0 && unavailableLocationsExpirationTimePassed()) { this.updateLocationCache(); @@ -103,7 +101,7 @@ && unavailableLocationsExpirationTimePassed()) { * 2. Endpoint availability * @return */ - public UnmodifiableList getWriteEndpoints() { + public UnmodifiableList getWriteEndpoints() { if (this.locationUnavailabilityInfoByEndpoint.size() > 0 && unavailableLocationsExpirationTimePassed()) { this.updateLocationCache(); @@ -120,9 +118,9 @@ && unavailableLocationsExpirationTimePassed()) { * @return */ public List getAvailableReadEndpoints() { - return this.locationInfo.availableReadEndpointsByLocation.values().stream().map(locationEndpoints -> { + return this.locationInfo.availableReadEndpointsByLocation.values().stream().map(consolidatedLocationEndpoints -> { // TODO: Integrate thinclient endpoints into fault injection - return locationEndpoints.gatewayEndpoint; + return consolidatedLocationEndpoints.gatewayLocationEndpoint; }).collect(Collectors.toList()); } @@ -134,9 +132,9 @@ public List getAvailableReadEndpoints() { * @return */ public List getAvailableWriteEndpoints() { - return this.locationInfo.availableWriteEndpointsByLocation.values().stream().map(locationEndpoints -> { + return this.locationInfo.availableWriteEndpointsByLocation.values().stream().map(consolidatedLocationEndpoints -> { // TODO: Integrate thinclient endpoints into fault injection - return locationEndpoints.gatewayEndpoint; + return consolidatedLocationEndpoints.gatewayLocationEndpoint; }).collect(Collectors.toList()); } @@ -195,47 +193,47 @@ void onLocationPreferenceChanged(UnmodifiableList preferredLocations) { * @param request Request for which getEndpoint is to be resolved * @return Resolved getEndpoint */ - public LocationEndpoints resolveServiceEndpoint(RxDocumentServiceRequest request) { + public ConsolidatedRegionalEndpoint resolveServiceEndpoint(RxDocumentServiceRequest request) { Objects.requireNonNull(request.requestContext, "RxDocumentServiceRequest.requestContext is required and cannot be null."); - if(request.requestContext.locationEndpointToRoute != null) { - return new LocationEndpoints(request.requestContext.locationEndpointToRoute); + + if (request.requestContext.locationEndpointToRoute != null) { + + Objects.requireNonNull(request.requestContext.consolidatedRegionalEndpointToRoute); + + return request.requestContext.consolidatedRegionalEndpointToRoute; } 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); - URI thinclientEndpoint = currentLocationInfo.availableWriteEndpointsByLocation.get(writeLocation).thinClientEndpoint; - URI gatewayEndpoint = currentLocationInfo.availableWriteEndpointsByLocation.get(writeLocation).gatewayEndpoint; - // Bubble up both endpoints so we can try both before we have to do a cross-region retry - return new LocationEndpoints(gatewayEndpoint, thinclientEndpoint); + return currentLocationInfo.availableWriteEndpointsByLocation.get(writeLocation); } else { - // assuming defaultEndpoint would be gateway endpoint here - return this.defaultEndpoint; + return new ConsolidatedRegionalEndpoint(this.defaultEndpoint, null); } } 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) ? @@ -263,12 +261,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) ? @@ -292,24 +290,20 @@ public UnmodifiableList getApplicableReadEndpoints(List getApplicableEndpoints( - UnmodifiableList endpoints, - UnmodifiableMap regionNameByEndpoint, - LocationEndpoints fallbackEndpoint, + private UnmodifiableList getApplicableEndpoints( + UnmodifiableList endpoints, + UnmodifiableMap regionNameByEndpoint, + URI fallbackEndpoint, List excludeRegionList) { - List applicableEndpoints = new ArrayList<>(); - for (LocationEndpoints endpoint : endpoints) { + List applicableEndpoints = new ArrayList<>(); + for (ConsolidatedRegionalEndpoint endpoint : endpoints) { Utils.ValueHolder regionName = new Utils.ValueHolder<>(); - // TODO: @jeet1995: validate this logic - URI endpointToTry = endpoint.gatewayEndpoint != null ? endpoint.gatewayEndpoint : fallbackEndpoint.gatewayEndpoint; - // TODO: @jeet1995: validate this logic - if (Utils.tryGetValue(regionNameByEndpoint, endpoint.gatewayEndpoint, regionName) || - Utils.tryGetValue(regionNameByEndpoint, endpoint.thinClientEndpoint, regionName)) { + if (Utils.tryGetValue(regionNameByEndpoint, endpoint, regionName)) { if (!excludeRegionList.stream().anyMatch(regionName.v::equalsIgnoreCase)) { applicableEndpoints.add(endpoint); } @@ -317,7 +311,7 @@ private UnmodifiableList getApplicableEndpoints( } if (applicableEndpoints.isEmpty()) { - applicableEndpoints.add(fallbackEndpoint); + applicableEndpoints.add(new ConsolidatedRegionalEndpoint(fallbackEndpoint, null)); } return new UnmodifiableList<>(applicableEndpoints); @@ -331,7 +325,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 { @@ -340,15 +334,14 @@ public URI resolveFaultInjectionEndpoint(String region, boolean writeOnly) { if (endpointValueHolder.v != null) { // TODO: Figure out how to integrate thinclient into fault injection - return endpointValueHolder.v.gatewayEndpoint; + return endpointValueHolder.v.gatewayLocationEndpoint; } throw new IllegalArgumentException("Can not find service endpoint for region " + region); } public URI getDefaultEndpoint() { - // assuming default endpoint as gateway, don't want to pull LocationEndpoints type up from here - return this.defaultEndpoint.gatewayEndpoint; + return this.defaultEndpoint; } public boolean shouldRefreshEndpoints(Utils.ValueHolder canRefreshInBackground) { @@ -364,11 +357,9 @@ public boolean shouldRefreshEndpoints(Utils.ValueHolder canRefreshInBac if (this.enableEndpointDiscovery) { boolean shouldRefresh = this.useMultipleWriteLocations && !this.enableMultipleWriteLocations; - List readLocationEndpoints = currentLocationInfo.readEndpoints; - // TODO: @jeet1995 validate this logic - if (this.isEndpointUnavailable(readLocationEndpoints.get(0).gatewayEndpoint, OperationType.Read) && - this.isEndpointUnavailable(readLocationEndpoints.get(0).thinClientEndpoint, OperationType.Read)) { - // Since most preferred read endpoints are unavailable, we can only refresh in background if + 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 canRefreshInBackground.v = anyEndpointsAvailable(readLocationEndpoints,OperationType.Read); logger.debug("shouldRefreshEndpoints = true, since the first read endpoint " + @@ -379,25 +370,17 @@ 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)) { - URI mostPreferredLocationGatewayEndpoint = mostPreferredReadEndpointHolder.v.gatewayEndpoint; - URI mostPreferredLocationThinclientEndpoint = mostPreferredReadEndpointHolder.v.thinClientEndpoint; logger.debug("most preferred is [{}], most preferred available is [{}]", mostPreferredLocation, mostPreferredReadEndpointHolder.v); - if (!areEqual(mostPreferredLocationThinclientEndpoint, readLocationEndpoints.get(0).thinClientEndpoint)) { - // For reads, we can always refresh in background as we can alternate to - // other available read endpoints - logger.debug("shouldRefreshEndpoints = true, most preferred location [{}]" + - " is not available for read.", mostPreferredLocationThinclientEndpoint); - return true; - } else if (!areEqual(mostPreferredLocationGatewayEndpoint, readLocationEndpoints.get(0).gatewayEndpoint)) { + if (!areEqual(mostPreferredReadEndpointHolder.v, readLocationEndpoints.get(0))) { // For reads, we can always refresh in background as we can alternate to // other available read endpoints logger.debug("shouldRefreshEndpoints = true, most preferred location [{}]" + - " is not available for read.", mostPreferredLocationGatewayEndpoint); + " is not available for read.", mostPreferredLocation); return true; } @@ -411,14 +394,12 @@ 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()) { - // TODO: @jeet1995 validate this logic - if (this.isEndpointUnavailable(writeLocationEndpoints.get(0).thinClientEndpoint, OperationType.Write) && - this.isEndpointUnavailable(writeLocationEndpoints.get(0).gatewayEndpoint, OperationType.Write)) { + if (this.isEndpointUnavailable(writeLocationEndpoints.get(0), OperationType.Write)) { // Since most preferred write endpoint is unavailable, we can only refresh in background if // we have an alternate write endpoint canRefreshInBackground.v = anyEndpointsAvailable(writeLocationEndpoints,OperationType.Write); @@ -435,27 +416,14 @@ public boolean shouldRefreshEndpoints(Utils.ValueHolder canRefreshInBac } } else if (!Strings.isNullOrEmpty(mostPreferredLocation)) { if (Utils.tryGetValue(currentLocationInfo.availableWriteEndpointsByLocation, mostPreferredLocation, mostPreferredWriteEndpointHolder)) { - URI mostPreferredLocationGatewayEndpoint = mostPreferredWriteEndpointHolder.v.gatewayEndpoint; - URI mostPreferredLocationThinclientEndpoint = mostPreferredWriteEndpointHolder.v.thinClientEndpoint; - - if (mostPreferredLocationThinclientEndpoint != null) { - shouldRefresh = ! areEqual(mostPreferredLocationThinclientEndpoint, writeLocationEndpoints.get(0).thinClientEndpoint); - if (shouldRefresh) { - logger.debug("shouldRefreshEndpoints: true, write endpoint [{}] is not the same as most preferred [{}]", - writeLocationEndpoints.get(0), mostPreferredLocationThinclientEndpoint); - } else { - logger.debug("shouldRefreshEndpoints: false, write endpoint [{}] is the same as most preferred [{}]", - writeLocationEndpoints.get(0), mostPreferredLocationThinclientEndpoint); - } - } else if (mostPreferredLocationGatewayEndpoint != null) { - shouldRefresh = ! areEqual(mostPreferredLocationGatewayEndpoint, writeLocationEndpoints.get(0).gatewayEndpoint); - if (shouldRefresh) { - logger.debug("shouldRefreshEndpoints: true, write endpoint [{}] is not the same as most preferred [{}]", - writeLocationEndpoints.get(0), mostPreferredLocationGatewayEndpoint); - } else { - logger.debug("shouldRefreshEndpoints: false, write endpoint [{}] is the same as most preferred [{}]", - writeLocationEndpoints.get(0), mostPreferredLocationGatewayEndpoint); - } + shouldRefresh = ! areEqual(mostPreferredWriteEndpointHolder.v,writeLocationEndpoints.get(0)); + + if (shouldRefresh) { + logger.debug("shouldRefreshEndpoints: true, write endpoint [{}] is not the same as most preferred [{}]", + writeLocationEndpoints.get(0), mostPreferredWriteEndpointHolder.v); + } else { + logger.debug("shouldRefreshEndpoints: false, write endpoint [{}] is the same as most preferred [{}]", + writeLocationEndpoints.get(0), mostPreferredWriteEndpointHolder.v); } return shouldRefresh; @@ -477,11 +445,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 ConsolidatedRegionalEndpoint(locationEndpoint, locationEndpoint), regionName)) { return regionName.v; } } else { - if (Utils.tryGetValue(this.locationInfo.regionNameByReadEndpoint, locationEndpoint, regionName)) { + if (Utils.tryGetValue(this.locationInfo.regionNameByReadEndpoint, new ConsolidatedRegionalEndpoint(locationEndpoint, locationEndpoint), regionName)) { return regionName.v; } } @@ -490,15 +458,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(ConsolidatedRegionalEndpoint url1, ConsolidatedRegionalEndpoint 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 (ConsolidatedRegionalEndpoint unavailableEndpoint: unavailableEndpoints) { Utils.ValueHolder unavailabilityInfoHolder = new Utils.ValueHolder<>(); Utils.ValueHolder removedHolder = new Utils.ValueHolder<>(); @@ -517,7 +485,7 @@ private void clearStaleEndpointUnavailabilityInfo() { } } - private boolean isEndpointUnavailable(URI endpoint, OperationType expectedAvailableOperations) { + private boolean isEndpointUnavailable(ConsolidatedRegionalEndpoint endpoint, OperationType expectedAvailableOperations) { Utils.ValueHolder unavailabilityInfoHolder = new Utils.ValueHolder<>(); if (expectedAvailableOperations == OperationType.None @@ -538,13 +506,11 @@ 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 (LocationEndpoints endpoint : endpoints) { - // TODO: @jeet1995 validate this logic - if (!isEndpointUnavailable(endpoint.gatewayEndpoint, expectedAvailableOperations) || - !isEndpointUnavailable(endpoint.thinClientEndpoint, expectedAvailableOperations)) { + for (ConsolidatedRegionalEndpoint endpoint : endpoints) { + if (!isEndpointUnavailable(endpoint, expectedAvailableOperations)) { anyEndpointsAvailable = true; break; } @@ -557,10 +523,10 @@ private void markEndpointUnavailable( OperationType unavailableOperationType) { Instant currentTime = Instant.now(); LocationUnavailabilityInfo updatedInfo = this.locationUnavailabilityInfoByEndpoint.compute( - unavailableEndpoint, - new BiFunction() { + new ConsolidatedRegionalEndpoint(unavailableEndpoint, null), + new BiFunction() { @Override - public LocationUnavailabilityInfo apply(URI url, LocationUnavailabilityInfo info) { + public LocationUnavailabilityInfo apply(ConsolidatedRegionalEndpoint url, LocationUnavailabilityInfo info) { if (info == null) { // not already present, add return new LocationUnavailabilityInfo(currentTime, unavailableOperationType); @@ -570,7 +536,6 @@ public LocationUnavailabilityInfo apply(URI url, LocationUnavailabilityInfo info info.unavailableOperations = OperationType.combine(info.unavailableOperations, unavailableOperationType); return info; } - } }); @@ -590,8 +555,8 @@ private void updateLocationCache(){ private void updateLocationCache( Iterable gatewayWriteLocations, Iterable gatewayReadLocations, - Iterable thinclientWriteLocations, - Iterable thinclientReadLocations, + Iterable thinClientWriteLocations, + Iterable thinClientReadLocations, UnmodifiableList preferenceList, Boolean enableMultipleWriteLocations) { synchronized (this.lockObject) { @@ -609,27 +574,31 @@ private void updateLocationCache( this.clearStaleEndpointUnavailabilityInfo(); - Utils.ValueHolder> readValueHolder = Utils.ValueHolder.initialize(nextLocationInfo.availableReadLocations); - Utils.ValueHolder> readRegionMapValueHolder = Utils.ValueHolder.initialize(nextLocationInfo.regionNameByReadEndpoint); - nextLocationInfo.availableReadEndpointsByLocation = this.getEndpointsByLocation(gatewayReadLocations, thinclientReadLocations, readValueHolder, readRegionMapValueHolder); - nextLocationInfo.availableReadLocations = readValueHolder.v; - nextLocationInfo.regionNameByReadEndpoint = readRegionMapValueHolder.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, thinClientReadLocations, readValueHolder, readRegionMapValueHolder); + nextLocationInfo.availableReadLocations = readValueHolder.v; + nextLocationInfo.regionNameByReadEndpoint = readRegionMapValueHolder.v; + } - Utils.ValueHolder> writeValueHolder = Utils.ValueHolder.initialize(nextLocationInfo.availableWriteLocations); - Utils.ValueHolder> outWriteRegionMap = Utils.ValueHolder.initialize(nextLocationInfo.regionNameByWriteEndpoint); - nextLocationInfo.availableWriteEndpointsByLocation = this.getEndpointsByLocation(gatewayWriteLocations, thinclientWriteLocations, writeValueHolder, outWriteRegionMap); - nextLocationInfo.availableWriteLocations = writeValueHolder.v; - nextLocationInfo.regionNameByWriteEndpoint = outWriteRegionMap.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, thinClientWriteLocations, 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)); + 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 ConsolidatedRegionalEndpoint(this.defaultEndpoint, null), regionForDefaultEndpoint)) { nextLocationInfo.effectivePreferredLocations = nextLocationInfo.availableReadLocations; } } @@ -642,16 +611,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. @@ -660,139 +629,134 @@ private UnmodifiableList getPreferredAvailableEndpoints(Unmod if (currentLocationInfo.preferredLocations != null && !currentLocationInfo.preferredLocations.isEmpty()) { for (String location: currentLocationInfo.preferredLocations) { - Utils.ValueHolder locationEndpoints = new Utils.ValueHolder<>(); - if (Utils.tryGetValue(endpointsByLocation, location, locationEndpoints)) { - URI thinclientEndpoint = locationEndpoints.v.thinClientEndpoint; - boolean addThinclientEndpoint = false; - if (thinclientEndpoint != null && this.isEndpointUnavailable(thinclientEndpoint, expectedAvailableOperation)) { - unavailableEndpoints.add(thinclientEndpoint); + Utils.ValueHolder endpoint = new Utils.ValueHolder<>(); + if (Utils.tryGetValue(endpointsByLocation, location, endpoint)) { + if (this.isEndpointUnavailable(endpoint.v, expectedAvailableOperation)) { + unavailableEndpoints.add(endpoint.v); } else { - addThinclientEndpoint = true; - } - URI gatewayEndpoint = locationEndpoints.v.gatewayEndpoint; - boolean addGatewayEndpoint = false; - if (gatewayEndpoint != null && this.isEndpointUnavailable(gatewayEndpoint, expectedAvailableOperation)) { - unavailableEndpoints.add(gatewayEndpoint); - } else { - addGatewayEndpoint = true; + endpoints.add(endpoint.v); } } } } else { for (String location : orderedLocations) { - Utils.ValueHolder locationEndpoints = Utils.ValueHolder.initialize(null); - if (Utils.tryGetValue(endpointsByLocation, location, locationEndpoints)) { - URI gatewayEndpoint = locationEndpoints.v.gatewayEndpoint; - URI thinClientEndpoint = locationEndpoints.v.thinClientEndpoint; + 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 (thinClientEndpoint != null && this.defaultEndpoint.equals(thinClientEndpoint)) { - endpoints = new ArrayList<>(); - break; - } - if (gatewayEndpoint != null && this.defaultEndpoint.equals(gatewayEndpoint)) { + if (this.defaultEndpoint.equals(endpoint.v.getGatewayLocationEndpoint())) { endpoints = new ArrayList<>(); break; } - LocationEndpoints endpointsToAdd = new LocationEndpoints(null); - if (gatewayEndpoint != null && this.isEndpointUnavailable(gatewayEndpoint, expectedAvailableOperation)) { - unavailableEndpoints.add(gatewayEndpoint); - } else { - endpointsToAdd.gatewayEndpoint = gatewayEndpoint; - } - if (thinClientEndpoint != null && this.isEndpointUnavailable(thinClientEndpoint, expectedAvailableOperation)) { - unavailableEndpoints.add(thinClientEndpoint); + if (this.isEndpointUnavailable(endpoint.v, expectedAvailableOperation)) { + unavailableEndpoints.add(endpoint.v); } else { - endpointsToAdd.thinClientEndpoint = thinClientEndpoint; - } - - if (endpointsToAdd.gatewayEndpoint != null || endpointsToAdd.thinClientEndpoint != null) { - endpoints.add(endpointsToAdd); + endpoints.add(endpoint.v); } } } } if (endpoints.isEmpty()) { - endpoints.add(new LocationEndpoints(fallbackEndpoint)); + endpoints.add(new ConsolidatedRegionalEndpoint(fallbackEndpoint, null)); } endpoints.addAll(unavailableEndpoints); } else { for (String location : orderedLocations) { - Utils.ValueHolder locationEndpoint = 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, locationEndpoint)) { - if (locationEndpoint.v.thinClientEndpoint != null) { - endpoints.add(locationEndpoint.v.thinClientEndpoint); - } - if (locationEndpoint.v.gatewayEndpoint != null) { - endpoints.add(locationEndpoint.v.gatewayEndpoint); - } + Utils.tryGetValue(endpointsByLocation, location, endpoint)) { + endpoints.add(endpoint.v); } } } } if (endpoints.isEmpty()) { - endpoints.add(fallbackEndpoint); + endpoints.add(new ConsolidatedRegionalEndpoint(fallbackEndpoint, null)); } - return new UnmodifiableList(endpoints); + return new UnmodifiableList<>(endpoints); } - private void addEndpoints(Iterable locations, - Map endpointsByLocation, - Map regionByEndpoint, - List parsedLocations, - boolean isThinclient) { - for (DatabaseAccountLocation dbAccountLocation: locations) { - if (!Strings.isNullOrEmpty(dbAccountLocation.getName())) { - try { - String location = dbAccountLocation.getName().toLowerCase(Locale.ROOT); - URI endpoint = new URI(dbAccountLocation.getEndpoint().toLowerCase(Locale.ROOT)); - if (isThinclient) { - // there should be a LocationEndpoints object with gateway endpoint at this point - endpointsByLocation.get(location).thinClientEndpoint = endpoint; - } else { + private void addEndpoints( + Iterable gatewayDbAccountLocations, + Iterable thinClientDbAccountLocations, + Map endpointsByLocation, + Map regionByEndpoint, + List parsedLocations) { + + if (gatewayDbAccountLocations != null) { + for (DatabaseAccountLocation gatewayDbAccountLocation : gatewayDbAccountLocations) { + if (!Strings.isNullOrEmpty(gatewayDbAccountLocation.getName())) { + try { + + String location = gatewayDbAccountLocation.getName().toLowerCase(Locale.ROOT); + URI endpoint = new URI(gatewayDbAccountLocation.getEndpoint().toLowerCase(Locale.ROOT)); + + ConsolidatedRegionalEndpoint consolidatedRegionalEndpoint = new ConsolidatedRegionalEndpoint(endpoint, null); + if (!endpointsByLocation.containsKey(location)) { - endpointsByLocation.put(location, new LocationEndpoints(endpoint)); + endpointsByLocation.put(location, consolidatedRegionalEndpoint); + } + + if (!regionByEndpoint.containsKey(consolidatedRegionalEndpoint)) { + regionByEndpoint.put(consolidatedRegionalEndpoint, 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()); } - regionByEndpoint.put(endpoint, dbAccountLocation.getName().toLowerCase(Locale.ROOT)); - parsedLocations.add(dbAccountLocation.getName()); - } catch (Exception e) { - logger.warn("GetAvailableEndpointsByLocation() - skipping add for location = [{}] as its location name is either empty or endpoint is malformed [{}]", - dbAccountLocation.getName(), - dbAccountLocation.getEndpoint()); } } } - } - private UnmodifiableMap getEndpointsByLocation(Iterable gatewayLocations, - Iterable thinclientLocations, - Utils.ValueHolder> orderedLocations, - Utils.ValueHolder> regionMap) { - Map endpointsByLocation = new CaseInsensitiveMap<>(); - Map regionByEndpoint = new CaseInsensitiveMap<>(); - List parsedLocations = new ArrayList<>(); + if (thinClientDbAccountLocations != null) { + for (DatabaseAccountLocation thinClientDbAccountLocation: thinClientDbAccountLocations) { + if (!Strings.isNullOrEmpty(thinClientDbAccountLocation.getName())) { + try { + String location = thinClientDbAccountLocation.getName().toLowerCase(Locale.ROOT); + URI endpoint = new URI(thinClientDbAccountLocation.getEndpoint().toLowerCase(Locale.ROOT)); - // order is important here, always check for gateway first to ensure we have a gateway endpoint - if (gatewayLocations != null) { - addEndpoints(gatewayLocations, endpointsByLocation, regionByEndpoint, parsedLocations, false); - } + ConsolidatedRegionalEndpoint consolidatedRegionalEndpoint = endpointsByLocation.get(location); + + if (consolidatedRegionalEndpoint == null) { + throw new IllegalStateException(String.format("Gateway location endpoint doesn't exist while thin client location endpoint exists for location %s", location)); + } - if (thinclientLocations != null) { - addEndpoints(thinclientLocations, endpointsByLocation, regionByEndpoint, parsedLocations, true); + consolidatedRegionalEndpoint.thinClientLocationEndpoint = endpoint; + + } catch (Exception e) { + logger.warn("GetAvailableEndpointsByLocation() - skipping add for location = [{}] as its location name is either empty or endpoint is malformed [{}]", + thinClientDbAccountLocation.getName(), + thinClientDbAccountLocation.getEndpoint()); + } + } + } } + } + + private UnmodifiableMap getEndpointsByLocation(Iterable gatewayLocations, + Iterable thinClientLocations, + Utils.ValueHolder> orderedLocations, + Utils.ValueHolder> regionMap) { + Map endpointsByLocation = new CaseInsensitiveMap<>(); + Map regionByEndpoint = new CaseInsensitiveMap<>(); + List parsedLocations = new ArrayList<>(); - orderedLocations.v = new UnmodifiableList(parsedLocations); - regionMap.v = (UnmodifiableMap) UnmodifiableMap.unmodifiableMap(regionByEndpoint); + addEndpoints(gatewayLocations, thinClientLocations, endpointsByLocation, regionByEndpoint, parsedLocations); - return (UnmodifiableMap) UnmodifiableMap.unmodifiableMap(endpointsByLocation); + orderedLocations.v = new UnmodifiableList<>(parsedLocations); + regionMap.v = (UnmodifiableMap) UnmodifiableMap.unmodifiableMap(regionByEndpoint); + + return (UnmodifiableMap) UnmodifiableMap.unmodifiableMap(endpointsByLocation); } public boolean canUseMultipleWriteLocations() { @@ -867,51 +831,67 @@ 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(defaultEndpoint)); - this.writeEndpoints = new UnmodifiableList<>(Collections.singletonList(defaultEndpoint)); + this.readEndpoints = new UnmodifiableList<>(Collections.singletonList(new ConsolidatedRegionalEndpoint(defaultEndpoint, null))); + this.writeEndpoints = new UnmodifiableList<>(Collections.singletonList(new ConsolidatedRegionalEndpoint(defaultEndpoint, null))); } public DatabaseAccountLocationsInfo(DatabaseAccountLocationsInfo other) {