diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/GlobalPartitionEndpointManagerForPPCBUnitTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/GlobalPartitionEndpointManagerForPPCBUnitTests.java index cc7d589836eb..5929d2d255be 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/GlobalPartitionEndpointManagerForPPCBUnitTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/GlobalPartitionEndpointManagerForPPCBUnitTests.java @@ -4,9 +4,12 @@ package com.azure.cosmos; import com.azure.cosmos.implementation.AvailabilityStrategyContext; +import com.azure.cosmos.implementation.ConnectionPolicy; import com.azure.cosmos.implementation.CrossRegionAvailabilityContextForRxDocumentServiceRequest; import com.azure.cosmos.implementation.GlobalEndpointManager; import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.implementation.IAuthorizationTokenProvider; +import com.azure.cosmos.implementation.OpenConnectionResponse; import com.azure.cosmos.implementation.OperationType; import com.azure.cosmos.implementation.PartitionKeyRange; import com.azure.cosmos.implementation.PartitionKeyRangeWrapper; @@ -15,6 +18,14 @@ import com.azure.cosmos.implementation.RxDocumentServiceRequest; import com.azure.cosmos.implementation.SerializationDiagnosticsContext; import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList; +import com.azure.cosmos.implementation.directconnectivity.Address; +import com.azure.cosmos.implementation.directconnectivity.GatewayAddressCache; +import com.azure.cosmos.implementation.directconnectivity.GlobalAddressResolver; +import com.azure.cosmos.implementation.directconnectivity.Protocol; +import com.azure.cosmos.implementation.directconnectivity.Uri; +import com.azure.cosmos.implementation.directconnectivity.rntbd.OpenConnectionTask; +import com.azure.cosmos.implementation.directconnectivity.rntbd.ProactiveOpenConnectionsProcessor; +import com.azure.cosmos.implementation.http.HttpClient; import com.azure.cosmos.implementation.perPartitionAutomaticFailover.PerPartitionAutomaticFailoverInfoHolder; import com.azure.cosmos.implementation.perPartitionCircuitBreaker.GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker; import com.azure.cosmos.implementation.perPartitionCircuitBreaker.LocationHealthStatus; @@ -22,6 +33,7 @@ import com.azure.cosmos.implementation.guava25.collect.ImmutableList; import com.azure.cosmos.implementation.perPartitionCircuitBreaker.PerPartitionCircuitBreakerInfoHolder; import com.azure.cosmos.implementation.routing.RegionalRoutingContext; +import io.netty.channel.ConnectTimeoutException; import org.apache.commons.lang3.tuple.Pair; import org.mockito.Mockito; import org.slf4j.Logger; @@ -29,17 +41,29 @@ import org.testng.annotations.BeforeClass; import org.testng.annotations.DataProvider; import org.testng.annotations.Test; +import reactor.core.Disposable; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; +import reactor.test.StepVerifier; +import reactor.test.scheduler.VirtualTimeScheduler; import java.lang.reflect.Field; +import java.lang.reflect.Method; 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.Map; +import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ScheduledFuture; import java.util.concurrent.ScheduledThreadPoolExecutor; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; import static com.azure.cosmos.implementation.TestUtils.mockDiagnosticsClientContext; @@ -54,6 +78,11 @@ public class GlobalPartitionEndpointManagerForPPCBUnitTests { private final static Pair LocationCentralUsEndpointToLocationPair = Pair.of(createUrl("https://contoso-central-us.documents.azure.com"), "centralus"); private static final boolean READ_OPERATION_TRUE = true; + private static final String PPCB_RECOVERY_CONFIG + = "{\"isPartitionLevelCircuitBreakerEnabled\":true," + + "\"circuitBreakerType\":\"CONSECUTIVE_EXCEPTION_COUNT_BASED\"," + + "\"consecutiveExceptionCountToleratedForReads\":10," + + "\"consecutiveExceptionCountToleratedForWrites\":5}"; private GlobalEndpointManager globalEndpointManagerMock; @@ -123,6 +152,15 @@ public Object[][] nullPartitionKeyRangeHandlingArgs() { }; } + @DataProvider(name = "addressCacheStates") + public Object[][] addressCacheStates() { + return new Object[][] { + { false, false }, + { true, false }, + { true, true } + }; + } + @Test(groups = {"unit"}, dataProvider = "partitionLevelCircuitBreakerConfigs") public void recordHealthyStatus(String partitionLevelCircuitBreakerConfigAsJsonString, boolean readOperationTrue) throws IllegalAccessException, NoSuchFieldException { @@ -1009,6 +1047,227 @@ public void validateHandlingOnNullPartitionKeyRange(boolean setResolvedPartition } } + @Test(groups = "unit", dataProvider = "addressCacheStates") + @SuppressWarnings("unchecked") + public void scheduledRecoveryHandlesMissingAndStaleAddressCacheEntries( + boolean populateStaleAddress, + boolean refreshedProbeFails) + throws Exception { + + String originalPpcbConfig = System.getProperty("COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG"); + + URI failedRegionEndpoint = createUrl("https://contoso-east-us.documents.azure.com"); + URI healthyRegionEndpoint = createUrl("https://contoso-west-us.documents.azure.com"); + RegionalRoutingContext failedRegion = new RegionalRoutingContext(failedRegionEndpoint); + List applicableRegions = Arrays.asList( + failedRegion, + new RegionalRoutingContext(healthyRegionEndpoint)); + String collectionRid = "collectionRid"; + String partitionKeyRangeId = "0"; + + GlobalEndpointManager globalEndpointManager = Mockito.mock(GlobalEndpointManager.class); + Mockito.when(globalEndpointManager.getApplicableReadRegionalRoutingContexts(Mockito.anyList())) + .thenReturn((UnmodifiableList) UnmodifiableList.unmodifiableList(applicableRegions)); + Mockito.when(globalEndpointManager.getRegionName(failedRegionEndpoint, OperationType.Read)) + .thenReturn("East US"); + + AtomicInteger addressResolutionCount = new AtomicInteger(); + List forceRefreshValues = new CopyOnWriteArrayList<>(); + Address staleAddress = createAddress("rntbd://stale:10250/", partitionKeyRangeId); + Address refreshedAddress = createAddress("rntbd://refreshed:10250/", partitionKeyRangeId); + AtomicInteger staleConnectionAttempts = new AtomicInteger(); + AtomicInteger refreshedConnectionAttempts = new AtomicInteger(); + + ProactiveOpenConnectionsProcessor openConnectionsProcessor + = Mockito.mock(ProactiveOpenConnectionsProcessor.class); + Mockito.when(openConnectionsProcessor.submitOpenConnectionTaskOutsideLoop( + Mockito.anyString(), Mockito.any(), Mockito.any(), Mockito.anyInt())) + .thenAnswer(invocation -> { + Uri uri = invocation.getArgument(2); + Throwable failure = null; + if (populateStaleAddress + && uri.getURIAsString().equals(staleAddress.getPhyicalUri()) + && staleConnectionAttempts.incrementAndGet() == 2) { + + failure = new ConnectTimeoutException("Cached replica address is stale"); + } else if (refreshedProbeFails + && uri.getURIAsString().equals(refreshedAddress.getPhyicalUri())) { + + refreshedConnectionAttempts.incrementAndGet(); + failure = new ConnectTimeoutException("Refreshed replica is unavailable"); + } + + return completedOpenConnectionTask(collectionRid, failedRegionEndpoint, uri, failure); + }); + + GatewayAddressCache gatewayAddressCache = new GatewayAddressCache( + mockDiagnosticsClientContext(), + failedRegionEndpoint, + Protocol.TCP, + Mockito.mock(IAuthorizationTokenProvider.class), + null, + Mockito.mock(HttpClient.class), + null, + globalEndpointManager, + ConnectionPolicy.getDefaultPolicy(), + openConnectionsProcessor, + null, + null) { + @Override + public Mono> getServerAddressesViaGatewayAsync( + RxDocumentServiceRequest request, + String requestedCollectionRid, + List partitionKeyRangeIds, + boolean forceRefresh) { + + forceRefreshValues.add(forceRefresh); + addressResolutionCount.incrementAndGet(); + return Mono.just(Collections.singletonList( + populateStaleAddress && !forceRefresh ? staleAddress : refreshedAddress)); + } + }; + + GlobalAddressResolver globalAddressResolver = Mockito.mock(GlobalAddressResolver.class); + Mockito.when(globalAddressResolver.getGatewayAddressCache(failedRegionEndpoint)) + .thenReturn(gatewayAddressCache); + + GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker ppcbManager = null; + try { + System.setProperty("COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG", PPCB_RECOVERY_CONFIG); + ppcbManager = new GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker(globalEndpointManager); + ppcbManager.setGlobalAddressResolver(globalAddressResolver); + assertThat(ppcbManager.getCircuitBreakerConfig().isPartitionLevelCircuitBreakerEnabled()).isTrue(); + if (populateStaleAddress) { + StepVerifier.create(gatewayAddressCache.submitOpenConnectionTasks( + new PartitionKeyRange(partitionKeyRangeId, "AA", "BB"), + collectionRid, + false)) + .expectNextCount(1) + .verifyComplete(); + } + + RxDocumentServiceRequest request = constructRxDocumentServiceRequestInstance( + OperationType.Read, + ResourceType.Document, + collectionRid, + partitionKeyRangeId, + collectionRid, + "AA", + "BB", + failedRegionEndpoint); + PartitionKeyRange partitionKeyRange = request.requestContext.resolvedPartitionKeyRange; + for (int i = 0; i < 10; i++) { + ppcbManager.handleLocationExceptionForPartitionKeyRange(request, failedRegion, false); + } + assertThat(ppcbManager.getUnavailableRegionsForPartitionKeyRange( + request, + collectionRid, + partitionKeyRange)).containsExactly("East US"); + backdateUnavailableSince(ppcbManager, partitionKeyRange, collectionRid, failedRegion); + + VirtualTimeScheduler virtualTimeScheduler = VirtualTimeScheduler.getOrSet(); + Disposable recoverySubscription = invokeRecoveryPublisher(ppcbManager).subscribe(); + try { + virtualTimeScheduler.advanceTimeBy(Duration.ofSeconds(61)); + } finally { + recoverySubscription.dispose(); + VirtualTimeScheduler.reset(); + } + + if (refreshedProbeFails) { + assertThat(ppcbManager.getUnavailableRegionsForPartitionKeyRange( + request, + collectionRid, + partitionKeyRange)).containsExactly("East US"); + assertThat(refreshedConnectionAttempts).hasValue(1); + } else { + assertThat(ppcbManager.getUnavailableRegionsForPartitionKeyRange( + request, + collectionRid, + partitionKeyRange)).isEmpty(); + } + + if (populateStaleAddress) { + assertThat(forceRefreshValues).containsExactly(false, true); + assertThat(addressResolutionCount).hasValue(2); + assertThat(staleConnectionAttempts).hasValue(2); + } else { + assertThat(forceRefreshValues).containsExactly(false); + assertThat(addressResolutionCount).hasValue(1); + } + } finally { + if (ppcbManager != null) { + ppcbManager.close(); + } + if (originalPpcbConfig == null) { + System.clearProperty("COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG"); + } else { + System.setProperty("COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG", originalPpcbConfig); + } + } + } + + private static Address createAddress(String physicalUri, String partitionKeyRangeId) { + return new Address( + "{\"isPrimary\":true," + + "\"protocol\":\"rntbd\"," + + "\"physcialUri\":\"" + physicalUri + "\"," + + "\"partitionKeyRangeId\":\"" + partitionKeyRangeId + "\"}"); + } + + private static OpenConnectionTask completedOpenConnectionTask( + String collectionRid, + URI serviceEndpoint, + Uri uri, + Throwable failure) { + + OpenConnectionTask task = new OpenConnectionTask(collectionRid, serviceEndpoint, uri, 1); + task.complete(new OpenConnectionResponse(uri, failure == null, failure, failure == null ? 1 : 0)); + return task; + } + + @SuppressWarnings("unchecked") + private static void backdateUnavailableSince( + GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker ppcbManager, + PartitionKeyRange partitionKeyRange, + String collectionRid, + RegionalRoutingContext failedRegion) throws Exception { + + Field partitionMapField = GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker.class + .getDeclaredField("partitionKeyRangeToLocationSpecificUnavailabilityInfo"); + partitionMapField.setAccessible(true); + Map partitionMap + = (Map) partitionMapField.get(ppcbManager); + Object partitionInfo = partitionMap.get(new PartitionKeyRangeWrapper(partitionKeyRange, collectionRid)); + + Field locationMapField = partitionInfo.getClass() + .getDeclaredField("locationEndpointToLocationSpecificContextForPartition"); + locationMapField.setAccessible(true); + Map locationMap + = (Map) locationMapField.get(partitionInfo); + + Field unavailableSinceField = LocationSpecificHealthContext.class.getDeclaredField("unavailableSince"); + unavailableSinceField.setAccessible(true); + LocationSpecificHealthContext context = locationMap.get(failedRegion); + // Virtual time advances the recovery scheduler but not the Instant-based unavailability duration. + Instant backdatedUnavailableSince = Instant.now().minus(Duration.ofMinutes(2)); + unavailableSinceField.set(context, backdatedUnavailableSince); + assertThat(context.getUnavailableSince()).isEqualTo(backdatedUnavailableSince); + } + + private static Flux invokeRecoveryPublisher( + GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker ppcbManager) { + + try { + Method updateStaleLocationInfo = GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker.class + .getDeclaredMethod("updateStaleLocationInfo"); + updateStaleLocationInfo.setAccessible(true); + return (Flux) updateStaleLocationInfo.invoke(ppcbManager); + } catch (ReflectiveOperationException exception) { + return Flux.error(exception); + } + } + private static void validateAllRegionsAreNotUnavailableAfterExceptionInLocation( GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker globalPartitionEndpointManagerForCircuitBreaker, RxDocumentServiceRequest request, diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PerPartitionCircuitBreakerE2ETests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PerPartitionCircuitBreakerE2ETests.java index 8e1ad9f6c37f..409a9b5d025d 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PerPartitionCircuitBreakerE2ETests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PerPartitionCircuitBreakerE2ETests.java @@ -3,6 +3,7 @@ package com.azure.cosmos; +import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.faultinjection.FaultInjectionTestBase; import com.azure.cosmos.implementation.ClientSideRequestStatistics; import com.azure.cosmos.implementation.Configs; @@ -5659,6 +5660,26 @@ private static double getEstimatedFailureCountSeenPerRegionPerPartitionKeyRange( return 0d; } + @SuppressWarnings("unchecked") + private static boolean hasUnavailableLocationForPartition( + PartitionKeyRangeWrapper partitionKeyRangeWrapper, + ConcurrentHashMap partitionKeyRangeToLocationSpecificUnavailabilityInfo, + Field locationEndpointToLocationSpecificContextForPartitionField) throws IllegalAccessException { + + Object partitionUnavailabilityInfo + = partitionKeyRangeToLocationSpecificUnavailabilityInfo.get(partitionKeyRangeWrapper); + if (partitionUnavailabilityInfo == null) { + return false; + } + + ConcurrentHashMap locationContexts + = (ConcurrentHashMap) + locationEndpointToLocationSpecificContextForPartitionField.get(partitionUnavailabilityInfo); + + return locationContexts.values().stream() + .anyMatch(context -> context.getLocationHealthStatus() == LocationHealthStatus.Unavailable); + } + private static FaultInjectionConnectionType evaluateFaultInjectionConnectionType(ConnectionMode connectionMode) { if (connectionMode == ConnectionMode.DIRECT) { @@ -5690,6 +5711,180 @@ public AccountLevelLocationContext( } } + @Test(groups = {"circuit-breaker-misc-direct"}, timeOut = 20 * TIMEOUT) + public void ppcbRecoveryResolvesAddressesAfterInitialAddressRefreshFailures() throws Exception { + if (this.readRegions == null || this.readRegions.size() <= 1) { + throw new SkipException("Test requires a multi-region account"); + } + + ConnectionPolicy connectionPolicy = ReflectionUtils.getConnectionPolicy(getClientBuilder()); + if (connectionPolicy.getConnectionMode() != ConnectionMode.DIRECT) { + throw new SkipException("Test only applicable to DIRECT mode"); + } + + if (!Boolean.FALSE.equals(Configs.isThinClientEnabled()) && Configs.isHttp2Enabled()) { + throw new SkipException("DIRECT mode is not supported with thin client"); + } + + String originalPpcbConfig = System.getProperty("COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG"); + TestObject testObject = TestObject.create(); + PartitionKey partitionKey = new PartitionKey(testObject.getId()); + try (CosmosAsyncClient bootstrapClient = getClientBuilder().buildAsyncClient()) { + bootstrapClient + .getDatabase(this.sharedAsyncDatabaseId) + .getContainer(this.sharedMultiPartitionAsyncContainerIdWhereIdIsPartitionKey) + .createItem(testObject, partitionKey, new CosmosItemRequestOptions()) + .block(); + } + + CosmosAsyncClient testClient = null; + FaultInjectionRule addressRefreshRule = null; + try { + System.setProperty( + "COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG", + "{\"isPartitionLevelCircuitBreakerEnabled\":true," + + "\"circuitBreakerType\":\"CONSECUTIVE_EXCEPTION_COUNT_BASED\"," + + "\"consecutiveExceptionCountToleratedForReads\":10," + + "\"consecutiveExceptionCountToleratedForWrites\":5}"); + testClient = getClientBuilder() + .preferredRegions(this.readRegions) + .buildAsyncClient(); + CosmosAsyncContainer container = testClient + .getDatabase(this.sharedAsyncDatabaseId) + .getContainer(this.sharedMultiPartitionAsyncContainerIdWhereIdIsPartitionKey); + + RxDocumentClientImpl documentClient + = (RxDocumentClientImpl) ReflectionUtils.getAsyncDocumentClient(testClient); + RxCollectionCache collectionCache = ReflectionUtils.getClientCollectionCache(documentClient); + RxPartitionKeyRangeCache partitionKeyRangeCache = ReflectionUtils.getPartitionKeyRangeCache(documentClient); + DocumentCollection documentCollection = collectionCache + .resolveByNameAsync(null, containerAccessor.getLinkWithoutTrailingSlash(container), null) + .block(); + List partitionKeyRanges = partitionKeyRangeCache + .tryGetOverlappingRangesAsync( + null, + documentCollection.getResourceId(), + new FeedRangePartitionKeyImpl(BridgeInternal.getPartitionKeyInternal(partitionKey)) + .getEffectiveRange(documentCollection.getPartitionKey()), + true, + null) + .block() + .v; + assertThat(partitionKeyRanges).hasSize(1); + PartitionKeyRangeWrapper partitionKeyRangeWrapper + = new PartitionKeyRangeWrapper(partitionKeyRanges.get(0), documentCollection.getResourceId()); + + GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker ppcbManager + = documentClient.getGlobalPartitionEndpointManagerForCircuitBreaker(); + assertThat(ppcbManager.getCircuitBreakerConfig().isPartitionLevelCircuitBreakerEnabled()).isTrue(); + Class partitionUnavailabilityInfoClass = getClassBySimpleName( + GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker.class.getDeclaredClasses(), + "PartitionLevelLocationUnavailabilityInfo"); + assertThat(partitionUnavailabilityInfoClass).isNotNull(); + + Field partitionUnavailabilityMapField + = GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker.class + .getDeclaredField("partitionKeyRangeToLocationSpecificUnavailabilityInfo"); + partitionUnavailabilityMapField.setAccessible(true); + ConcurrentHashMap partitionUnavailabilityMap + = (ConcurrentHashMap) partitionUnavailabilityMapField.get(ppcbManager); + + Field locationContextMapField = partitionUnavailabilityInfoClass + .getDeclaredField("locationEndpointToLocationSpecificContextForPartition"); + locationContextMapField.setAccessible(true); + + addressRefreshRule = new FaultInjectionRuleBuilder( + "ppcb-address-refresh-connection-delay-" + UUID.randomUUID()) + .condition(new FaultInjectionConditionBuilder() + .region(this.readRegions.get(0)) + .operationType(FaultInjectionOperationType.METADATA_REQUEST_ADDRESS_REFRESH) + .build()) + .result(FaultInjectionResultBuilders + .getResultBuilder(FaultInjectionServerErrorType.RESPONSE_DELAY) + .delay(Duration.ofSeconds(11)) + .times(3) + .build()) + .duration(Duration.ofMinutes(10)) + // Keep recovery probes faulted until the test has observed failover. + .hitLimit(60) + .build(); + CosmosFaultInjectionHelper.configureFaultInjectionRules( + container, + Collections.singletonList(addressRefreshRule)).block(); + + CosmosItemRequestOptions readOptions = new CosmosItemRequestOptions() + .setCosmosEndToEndOperationLatencyPolicyConfig(NO_END_TO_END_TIMEOUT); + CosmosDiagnostics lastDiagnostics = null; + for (int i = 0; i < 20 + && !hasUnavailableLocationForPartition( + partitionKeyRangeWrapper, + partitionUnavailabilityMap, + locationContextMapField); i++) { + + try { + CosmosItemResponse response = container + .readItem(testObject.getId(), partitionKey, readOptions, TestObject.class) + .block(); + lastDiagnostics = response.getDiagnostics(); + } catch (CosmosException exception) { + lastDiagnostics = exception.getDiagnostics(); + } + } + + assertThat(addressRefreshRule.getHitCount()).isGreaterThanOrEqualTo(30); + assertThat(hasUnavailableLocationForPartition( + partitionKeyRangeWrapper, + partitionUnavailabilityMap, + locationContextMapField)).isTrue(); + assertThat(lastDiagnostics).isNotNull(); + + CosmosItemResponse failedOverResponse = container + .readItem(testObject.getId(), partitionKey, readOptions, TestObject.class) + .block(); + assertContactedRegionsContain( + failedOverResponse.getDiagnostics().getDiagnosticsContext(), + getRegionNameForAssertion(this.readRegions.get(1)), + "PPCB should route the partition to the second preferred region"); + + addressRefreshRule.disable(); + long recoveryDeadline = System.nanoTime() + Duration.ofSeconds(120).toNanos(); + while (hasUnavailableLocationForPartition( + partitionKeyRangeWrapper, + partitionUnavailabilityMap, + locationContextMapField) && System.nanoTime() < recoveryDeadline) { + + Thread.sleep(Duration.ofSeconds(1).toMillis()); + } + + assertThat(hasUnavailableLocationForPartition( + partitionKeyRangeWrapper, + partitionUnavailabilityMap, + locationContextMapField)).isFalse(); + + CosmosItemResponse recoveredResponse = container + .readItem(testObject.getId(), partitionKey, readOptions, TestObject.class) + .block(); + assertContactedRegionCount( + recoveredResponse.getDiagnostics().getDiagnosticsContext(), + 1, + "Recovered partition should use one preferred region"); + assertContactedRegionsContain( + recoveredResponse.getDiagnostics().getDiagnosticsContext(), + getRegionNameForAssertion(this.readRegions.get(0)), + "PPCB should fail back to the first preferred region after recovery"); + } finally { + if (addressRefreshRule != null) { + addressRefreshRule.disable(); + } + safeClose(testClient); + if (originalPpcbConfig == null) { + System.clearProperty("COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG"); + } else { + System.setProperty("COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG", originalPpcbConfig); + } + } + } + @Test(groups = {"circuit-breaker-misc-direct"}, timeOut = 4 * TIMEOUT) public void nonCanonicalPreferredRegions_ppcbShouldStillRouteCorrectly() { diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/GatewayAddressCacheTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/GatewayAddressCacheTest.java index ab75fc4aa503..9f717d46c589 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/GatewayAddressCacheTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/GatewayAddressCacheTest.java @@ -17,7 +17,9 @@ import com.azure.cosmos.implementation.HttpClientUnderTestWrapper; import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.IAuthorizationTokenProvider; +import com.azure.cosmos.implementation.OpenConnectionResponse; import com.azure.cosmos.implementation.OperationType; +import com.azure.cosmos.implementation.PartitionKeyRange; import com.azure.cosmos.implementation.RequestOptions; import com.azure.cosmos.implementation.ResourceType; import com.azure.cosmos.implementation.RxDocumentClientImpl; @@ -32,6 +34,7 @@ import com.azure.cosmos.implementation.http.HttpClientConfig; import com.azure.cosmos.implementation.routing.PartitionKeyRangeIdentity; import com.azure.cosmos.models.PartitionKeyDefinition; +import io.netty.channel.ConnectTimeoutException; import org.assertj.core.api.AssertionsForClassTypes; import org.mockito.ArgumentCaptor; import org.mockito.ArgumentMatchers; @@ -55,11 +58,14 @@ import java.time.Instant; import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Set; import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; @@ -1606,6 +1612,292 @@ public static void validateSuccess(Mono> observable, assertThat(httpClient.capturedRequests.get(requestIndex).headers().value(HttpConstants.HttpHeaders.ACTIVITY_ID)).isEqualTo(addressResolutionActivityId); } + @Test(groups = { "direct" }, timeOut = TIMEOUT) + public void submitOpenConnectionTasksResolvesAddressesWhenCacheEntryIsMissing() throws Exception { + String collectionRid = "collectionRid"; + String partitionKeyRangeId = "0"; + URI serviceEndpoint = new URI("https://localhost"); + Address address = createAddress("rntbd://localhost:10250/", partitionKeyRangeId, true); + + AtomicInteger addressResolutionCount = new AtomicInteger(); + ProactiveOpenConnectionsProcessor processor = Mockito.mock(ProactiveOpenConnectionsProcessor.class); + Mockito.when(processor.submitOpenConnectionTaskOutsideLoop( + Mockito.anyString(), Mockito.any(), Mockito.any(), Mockito.anyInt())) + .thenReturn(completedOpenConnectionTask( + collectionRid, + serviceEndpoint, + new Uri(address.getPhyicalUri()), + null)); + + GatewayAddressCache cache = createGatewayAddressCache( + serviceEndpoint, + processor, + (request, requestedCollectionRid, partitionKeyRangeIds, forceRefresh) -> { + addressResolutionCount.incrementAndGet(); + assertThat(request.requestContext.regionalRoutingContextToRoute.getGatewayRegionalEndpoint()) + .isEqualTo(serviceEndpoint); + assertThat(request.faultInjectionRequestContext.getRegionalRoutingContextToRoute() + .getGatewayRegionalEndpoint()).isEqualTo(serviceEndpoint); + assertThat(requestedCollectionRid).isEqualTo(collectionRid); + assertThat(partitionKeyRangeIds).containsExactly(partitionKeyRangeId); + assertThat(forceRefresh).isFalse(); + return Collections.singletonList(address); + }); + + PartitionKeyRange partitionKeyRange = new PartitionKeyRange().setId(partitionKeyRangeId); + StepVerifier.create(cache.submitOpenConnectionTasks(partitionKeyRange, collectionRid, false)) + .expectNextCount(1) + .verifyComplete(); + StepVerifier.create(cache.submitOpenConnectionTasks(partitionKeyRange, collectionRid, false)) + .expectNextCount(1) + .verifyComplete(); + + assertThat(addressResolutionCount).hasValue(1); + Mockito.verify(processor, Mockito.times(2)) + .submitOpenConnectionTaskOutsideLoop( + Mockito.eq(collectionRid), + Mockito.eq(serviceEndpoint), + Mockito.argThat(uri -> uri.getURIAsString().equals(address.getPhyicalUri())), + Mockito.eq(1)); + } + + @Test(groups = { "direct" }, timeOut = TIMEOUT) + public void submitOpenConnectionTasksRefreshesAddressesAfterNetworkFailure() throws Exception { + String collectionRid = "collectionRid"; + String partitionKeyRangeId = "0"; + URI serviceEndpoint = new URI("https://localhost"); + Address stalePrimary = createAddress("rntbd://localhost:10250/", partitionKeyRangeId, true); + Address staleSecondary = createAddress("rntbd://localhost:10251/", partitionKeyRangeId, false); + Address refreshedPrimary = createAddress("rntbd://localhost:10252/", partitionKeyRangeId, true); + Address refreshedSecondary = createAddress("rntbd://localhost:10253/", partitionKeyRangeId, false); + ConnectTimeoutException staleAddressException = new ConnectTimeoutException("Connection timed out"); + + AtomicInteger addressResolutionCount = new AtomicInteger(); + List forceRefreshValues = new CopyOnWriteArrayList<>(); + Map connectionAttempts = new ConcurrentHashMap<>(); + ProactiveOpenConnectionsProcessor processor = Mockito.mock(ProactiveOpenConnectionsProcessor.class); + Mockito.when(processor.submitOpenConnectionTaskOutsideLoop( + Mockito.anyString(), Mockito.any(), Mockito.any(), Mockito.anyInt())) + .thenAnswer(invocation -> { + Uri uri = invocation.getArgument(2); + int attempt = connectionAttempts + .computeIfAbsent(uri.getURIAsString(), ignored -> new AtomicInteger()) + .incrementAndGet(); + Throwable exception = uri.getURIAsString().equals(stalePrimary.getPhyicalUri()) && attempt == 2 + ? staleAddressException + : null; + return completedOpenConnectionTask(collectionRid, serviceEndpoint, uri, exception); + }); + + GatewayAddressCache cache = createGatewayAddressCache( + serviceEndpoint, + processor, + (request, requestedCollectionRid, partitionKeyRangeIds, forceRefresh) -> { + assertThat(requestedCollectionRid).isEqualTo(collectionRid); + assertThat(partitionKeyRangeIds).containsExactly(partitionKeyRangeId); + forceRefreshValues.add(forceRefresh); + return addressResolutionCount.incrementAndGet() == 1 + ? Arrays.asList(stalePrimary, staleSecondary) + : Arrays.asList(refreshedPrimary, refreshedSecondary); + }); + + PartitionKeyRange partitionKeyRange = new PartitionKeyRange().setId(partitionKeyRangeId); + StepVerifier.create(cache.submitOpenConnectionTasks(partitionKeyRange, collectionRid, false)) + .expectNextCount(2) + .verifyComplete(); + StepVerifier.create(cache.submitOpenConnectionTasks(partitionKeyRange, collectionRid, false)) + .expectErrorMatches(throwable -> throwable == staleAddressException) + .verify(); + StepVerifier.create(cache.submitOpenConnectionTasks(partitionKeyRange, collectionRid, true)) + .expectNextCount(2) + .verifyComplete(); + + assertThat(addressResolutionCount).hasValue(2); + assertThat(forceRefreshValues).containsExactly(false, true); + assertThat(connectionAttempts.get(stalePrimary.getPhyicalUri())).hasValue(2); + assertThat(connectionAttempts.get(refreshedPrimary.getPhyicalUri())).hasValue(1); + assertThat(connectionAttempts.get(refreshedSecondary.getPhyicalUri())).hasValue(1); + } + + @Test(groups = { "direct" }, timeOut = TIMEOUT) + public void submitOpenConnectionTasksPropagatesFailureAfterRefreshedAddressFails() throws Exception { + String collectionRid = "collectionRid"; + String partitionKeyRangeId = "0"; + URI serviceEndpoint = new URI("https://localhost"); + Address staleAddress = createAddress("rntbd://localhost:10250/", partitionKeyRangeId, true); + Address refreshedAddress = createAddress("rntbd://localhost:10251/", partitionKeyRangeId, true); + ConnectTimeoutException connectionFailure = new ConnectTimeoutException("Connection timed out"); + + AtomicInteger addressResolutionCount = new AtomicInteger(); + AtomicInteger connectionAttemptCount = new AtomicInteger(); + ProactiveOpenConnectionsProcessor processor = Mockito.mock(ProactiveOpenConnectionsProcessor.class); + Mockito.when(processor.submitOpenConnectionTaskOutsideLoop( + Mockito.anyString(), Mockito.any(), Mockito.any(), Mockito.anyInt())) + .thenAnswer(invocation -> { + connectionAttemptCount.incrementAndGet(); + return completedOpenConnectionTask( + collectionRid, + serviceEndpoint, + invocation.getArgument(2), + connectionFailure); + }); + + GatewayAddressCache cache = createGatewayAddressCache( + serviceEndpoint, + processor, + (request, requestedCollectionRid, partitionKeyRangeIds, forceRefresh) -> + Collections.singletonList(addressResolutionCount.incrementAndGet() == 1 + ? staleAddress + : refreshedAddress)); + + StepVerifier.create(cache.submitOpenConnectionTasks( + new PartitionKeyRange().setId(partitionKeyRangeId), + collectionRid, + false)) + .expectErrorMatches(throwable -> throwable == connectionFailure) + .verify(); + StepVerifier.create(cache.submitOpenConnectionTasks( + new PartitionKeyRange().setId(partitionKeyRangeId), + collectionRid, + true)) + .expectErrorMatches(throwable -> throwable == connectionFailure) + .verify(); + + assertThat(addressResolutionCount).hasValue(2); + assertThat(connectionAttemptCount).hasValue(2); + } + + @DataProvider(name = "networkFailureResponseOrders") + public Object[][] networkFailureResponseOrders() { + return new Object[][] { + { true }, + { false } + }; + } + + @Test(groups = { "direct" }, dataProvider = "networkFailureResponseOrders", timeOut = TIMEOUT) + public void submitOpenConnectionTasksPrefersNetworkFailureAcrossReplicas(boolean networkFailureFirst) + throws Exception { + + String collectionRid = "collectionRid"; + String partitionKeyRangeId = "0"; + URI serviceEndpoint = new URI("https://localhost"); + Address networkFailureAddress = createAddress("rntbd://localhost:10250/", partitionKeyRangeId, true); + Address nonNetworkFailureAddress = createAddress("rntbd://localhost:10251/", partitionKeyRangeId, false); + ConnectTimeoutException networkFailure = new ConnectTimeoutException("Connection timed out"); + IllegalStateException nonNetworkFailure = new IllegalStateException("Context negotiation failed"); + OpenConnectionTask networkFailureTask = new OpenConnectionTask( + collectionRid, + serviceEndpoint, + new Uri(networkFailureAddress.getPhyicalUri()), + 1); + OpenConnectionTask nonNetworkFailureTask = new OpenConnectionTask( + collectionRid, + serviceEndpoint, + new Uri(nonNetworkFailureAddress.getPhyicalUri()), + 1); + + ProactiveOpenConnectionsProcessor processor = Mockito.mock(ProactiveOpenConnectionsProcessor.class); + Mockito.when(processor.submitOpenConnectionTaskOutsideLoop( + Mockito.anyString(), Mockito.any(), Mockito.any(), Mockito.anyInt())) + .thenAnswer(invocation -> ((Uri) invocation.getArgument(2)).getURIAsString() + .equals(networkFailureAddress.getPhyicalUri()) + ? networkFailureTask + : nonNetworkFailureTask); + + GatewayAddressCache cache = createGatewayAddressCache( + serviceEndpoint, + processor, + (request, requestedCollectionRid, partitionKeyRangeIds, forceRefresh) -> + Arrays.asList(networkFailureAddress, nonNetworkFailureAddress)); + + StepVerifier.create(cache.submitOpenConnectionTasks( + new PartitionKeyRange().setId(partitionKeyRangeId), + collectionRid, + false)) + .then(() -> { + OpenConnectionResponse networkFailureResponse = new OpenConnectionResponse( + networkFailureTask.getAddressUri(), false, networkFailure, 0); + OpenConnectionResponse nonNetworkFailureResponse = new OpenConnectionResponse( + nonNetworkFailureTask.getAddressUri(), false, nonNetworkFailure, 0); + if (networkFailureFirst) { + networkFailureTask.complete(networkFailureResponse); + } else { + nonNetworkFailureTask.complete(nonNetworkFailureResponse); + networkFailureTask.complete(networkFailureResponse); + } + }) + .expectErrorMatches(throwable -> throwable == networkFailure) + .verify(); + + if (networkFailureFirst) { + assertThat(nonNetworkFailureTask.isDone()).isFalse(); + } + } + + private static GatewayAddressCache createGatewayAddressCache( + URI serviceEndpoint, + ProactiveOpenConnectionsProcessor processor, + AddressResolver addressResolver) { + + return new GatewayAddressCache( + mockDiagnosticsClientContext(), + serviceEndpoint, + Protocol.TCP, + Mockito.mock(IAuthorizationTokenProvider.class), + null, + Mockito.mock(HttpClient.class), + null, + null, + ConnectionPolicy.getDefaultPolicy(), + processor, + null, + null) { + @Override + public Mono> getServerAddressesViaGatewayAsync( + RxDocumentServiceRequest request, + String collectionRid, + List partitionKeyRangeIds, + boolean forceRefresh) { + + return Mono.just(addressResolver.resolve( + request, + collectionRid, + partitionKeyRangeIds, + forceRefresh)); + } + }; + } + + private static Address createAddress(String physicalUri, String partitionKeyRangeId, boolean primary) { + Address address = new Address(); + address.setIsPrimary(primary); + address.setProtocol(Protocol.TCP.scheme()); + address.setPhysicalUri(physicalUri); + address.setPartitionKeyRangeId(partitionKeyRangeId); + return address; + } + + private static OpenConnectionTask completedOpenConnectionTask( + String collectionRid, + URI serviceEndpoint, + Uri uri, + Throwable exception) { + + OpenConnectionTask task = new OpenConnectionTask(collectionRid, serviceEndpoint, uri, 1); + task.complete(new OpenConnectionResponse(uri, exception == null, exception, exception == null ? 1 : 0)); + return task; + } + + @FunctionalInterface + private interface AddressResolver { + List
resolve( + RxDocumentServiceRequest request, + String collectionRid, + List partitionKeyRangeIds, + boolean forceRefresh); + } + @BeforeClass(groups = { "direct" }, timeOut = SETUP_TIMEOUT) public void before_GatewayAddressCacheTest() { client = clientBuilder().build(); diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 1afb0d0e805c..0c0115a04fb8 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -9,6 +9,7 @@ #### Breaking Changes #### Bugs Fixed +* Fixed Per-Partition Circuit Breaker failback getting stuck when partition recovery encounters missing or stale replica addresses. - See [PR 50182](https://github.com/Azure/azure-sdk-for-java/pull/50182). * Fixed an intermittent `IndexOutOfBoundsException` in cross-partition hybrid search queries caused by multiple subscriptions to the coalesced component query results. - See PR [49831](https://github.com/Azure/azure-sdk-for-java/issues/49831) * Fixed document requests failing when Gateway V2 is enabled with resource-token or permission-feed authentication by routing those requests through Compute Gateway. - See PR [50084](https://github.com/Azure/azure-sdk-for-java/pull/50084). * Unified request-level consistency override behavior across transports: invalid attempts to upgrade the request consistency level above the account default are now silently ignored instead of returning `BadRequest` in some gateway paths. - See PR [49606](https://github.com/Azure/azure-sdk-for-java/pull/49606). diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/GatewayAddressCache.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/GatewayAddressCache.java index 97773c6aae32..d1cdc9c1d745 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/GatewayAddressCache.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/GatewayAddressCache.java @@ -51,6 +51,7 @@ import com.azure.cosmos.implementation.http.HttpResponse; import com.azure.cosmos.implementation.http.HttpTimeoutPolicy; import com.azure.cosmos.implementation.routing.PartitionKeyRangeIdentity; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import io.netty.handler.codec.http.HttpMethod; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -1162,7 +1163,8 @@ public Mono submitOpenConnectionTask( public Flux submitOpenConnectionTasks( PartitionKeyRange partitionKeyRange, - String collectionRid) { + String collectionRid, + boolean forceRefresh) { if (this.proactiveOpenConnectionsProcessor == null) { return Flux.empty(); @@ -1173,14 +1175,77 @@ public Flux submitOpenConnectionTasks( PartitionKeyRangeIdentity partitionKeyRangeIdentity = new PartitionKeyRangeIdentity(collectionRid, partitionKeyRange.getId()); - return this.serverPartitionAddressCache.getAsync(partitionKeyRangeIdentity, cachedAddresses -> Mono.just(cachedAddresses), cachedAddresses -> true) - .flatMapMany(cachedAddresses -> Flux.fromArray(cachedAddresses)) + return this.serverPartitionAddressCache.getAsync( + partitionKeyRangeIdentity, + cachedAddresses -> cachedAddresses != null && !forceRefresh + ? Mono.just(cachedAddresses) + : this.getAddressesForRangeId( + this.createPartitionAddressRequest(collectionRid), + partitionKeyRangeIdentity, + forceRefresh, + cachedAddresses), + cachedAddresses -> forceRefresh) + .flatMapMany(cachedAddresses -> this.openConnections(collectionRid, cachedAddresses)) + .handle((response, sink) -> { + Throwable exception = response.getException(); + if (!response.isConnected() + && exception instanceof Exception + && WebExceptionUtility.isNetworkFailure((Exception) exception)) { + + // Fail on the first network exception so PPCB can refresh addresses without waiting for other probes. + sink.error(exception); + } else { + // Keep non-network failures until all probes finish in case a later probe reports a network failure. + sink.next(response); + } + }) + .collectList() + .flatMapMany(this::validateOpenConnectionResponses); + } + + private RxDocumentServiceRequest createPartitionAddressRequest(String collectionRid) { + RxDocumentServiceRequest request = RxDocumentServiceRequest.create( + this.clientContext, + OperationType.Read, + collectionRid, + ResourceType.DocumentCollection, + Collections.emptyMap()); + request.requestContext.regionalRoutingContextToRoute = new RegionalRoutingContext(this.serviceEndpoint); + request.faultInjectionRequestContext.setRegionalRoutingContextToRoute( + request.requestContext.regionalRoutingContextToRoute); + return request; + } + + private Flux openConnections( + String collectionRid, + AddressInformation[] addresses) { + + return Flux.fromArray(addresses) .flatMap(addressInformation -> Mono.fromFuture( this.proactiveOpenConnectionsProcessor.submitOpenConnectionTaskOutsideLoop( collectionRid, - this.addressEndpoint, + this.serviceEndpoint, addressInformation.getPhysicalUri(), - 1))); + 1), + true) + .onErrorResume(throwable -> Mono.just( + new OpenConnectionResponse(addressInformation.getPhysicalUri(), false, throwable, 0)))); + } + + private Flux validateOpenConnectionResponses( + List openConnectionResponses) { + + // No network exception short-circuited the probes, so surface the first remaining connection failure. + for (OpenConnectionResponse response : openConnectionResponses) { + if (!response.isConnected()) { + Throwable exception = response.getException(); + return Flux.error(exception != null + ? exception + : new IllegalStateException("Failed to open a connection without an exception.")); + } + } + + return Flux.fromIterable(openConnectionResponses); } private Mono> getServerAddressesViaGatewayWithRetry( diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker.java index cf5ffe7aa874..47da7d37155d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker.java @@ -20,6 +20,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.directconnectivity.WebExceptionUtility; import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -37,6 +38,7 @@ import java.util.Map; import java.util.PriorityQueue; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; @@ -343,8 +345,19 @@ private Flux updateStaleLocationInfo() { if (gatewayAddressCache != null) { return gatewayAddressCache - .submitOpenConnectionTasks(partitionKeyRangeWrapper.getPartitionKeyRange(), partitionKeyRangeWrapper.getCollectionResourceId()) - .timeout(Duration.ofSeconds(Configs.getConnectionEstablishmentTimeoutForPartitionRecoveryInSeconds())) + .submitOpenConnectionTasks( + partitionKeyRangeWrapper.getPartitionKeyRange(), + partitionKeyRangeWrapper.getCollectionResourceId(), + false) + .timeout(this.getPartitionRecoveryAttemptTimeout()) + .onErrorResume(throwable -> this.shouldForceRefreshAddresses(throwable) + ? gatewayAddressCache + .submitOpenConnectionTasks( + partitionKeyRangeWrapper.getPartitionKeyRange(), + partitionKeyRangeWrapper.getCollectionResourceId(), + true) + .timeout(this.getPartitionRecoveryAttemptTimeout()) + : Flux.error(throwable)) .doOnComplete(() -> { logger.debug("Partition health recovery query for partitionKeyRange : " + @@ -364,6 +377,7 @@ private Flux updateStaleLocationInfo() { false, true); } + return locationSpecificContextAsVal; }); }) @@ -401,6 +415,16 @@ private Flux updateStaleLocationInfo() { }); } + private Duration getPartitionRecoveryAttemptTimeout() { + return Duration.ofSeconds(Configs.getConnectionEstablishmentTimeoutForPartitionRecoveryInSeconds()); + } + + private boolean shouldForceRefreshAddresses(Throwable throwable) { + return throwable instanceof TimeoutException + || throwable instanceof Exception + && WebExceptionUtility.isNetworkFailure((Exception) throwable); + } + public boolean isPerPartitionLevelCircuitBreakingApplicable(RxDocumentServiceRequest request) { if (!this.consecutiveExceptionBasedCircuitBreaker.isPartitionLevelCircuitBreakerEnabled()) {