From be0d92c2f795dc7b151ad73adc56b911cb854b59 Mon Sep 17 00:00:00 2001 From: Annie Liang Date: Thu, 20 Aug 2026 09:36:04 -0700 Subject: [PATCH 1/7] Fix PPCB failback with missing or stale addresses Resolve missing and stale partition addresses during PPCB recovery, retry forced address refreshes after network failures, and add cache-level, manager-level, and fault-injection coverage. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...titionEndpointManagerForPPCBUnitTests.java | 259 ++++++++++++++++ .../PerPartitionCircuitBreakerE2ETests.java | 181 +++++++++++ .../GatewayAddressCacheTest.java | 293 ++++++++++++++++++ sdk/cosmos/azure-cosmos/CHANGELOG.md | 1 + .../azure/cosmos/implementation/Configs.java | 3 +- .../GatewayAddressCache.java | 72 ++++- ...tManagerForPerPartitionCircuitBreaker.java | 28 +- 7 files changed, 829 insertions(+), 8 deletions(-) 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..ebd1180aabc7 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,166 @@ 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"); + } + + 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)) + .hitLimit(30) + .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()).isEqualTo(30); + assertThat(hasUnavailableLocationForPartition( + partitionKeyRangeWrapper, + partitionUnavailabilityMap, + locationContextMapField)).isTrue(); + assertThat(lastDiagnostics).isNotNull(); + assertContactedRegionsContain( + lastDiagnostics.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..fa7ac8e85946 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,293 @@ 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(staleSecondary.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 4105b1622f13..7ed3b0f07801 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 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). * Fixed `partitionLevelCircuitBreakerCfg` missing from the `clientCfgs` section of `CosmosDiagnostics` when Per-Partition Circuit Breaker is explicitly enabled. - See PR [49734](https://github.com/Azure/azure-sdk-for-java/pull/49734). 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 686efabe8188..30d619805045 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 @@ -359,7 +359,8 @@ public class Configs { // For partition-level circuit breaker, in order to recover a partition in a region, the SDK when configured // in the direct connectivity mode, establishes connections to replicas to attempt to recover a region // Below sets a time limit on how long these connection establishments be attempted for - private static final int DEFAULT_CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS = 10; + // Covers one address probe pass with the default proactive connection retry policy. + private static final int DEFAULT_CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS = 20; private static final String CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS = "COSMOS.CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS"; private static final boolean DEFAULT_SHOULD_LOG_INCORRECTLY_MAPPED_SESSION_TOKEN = true; 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..b475de3814f4 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,74 @@ 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)) { + + sink.error(exception); + } else { + 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) { + + 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()) { From e6fdb059a08bb1ffeefc2f023ca76a30e9e7d880 Mon Sep 17 00:00:00 2001 From: Annie Liang Date: Thu, 20 Aug 2026 09:53:02 -0700 Subject: [PATCH 2/7] Clarify PPCB probe exception handling Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../implementation/directconnectivity/GatewayAddressCache.java | 3 +++ 1 file changed, 3 insertions(+) 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 b475de3814f4..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 @@ -1192,8 +1192,10 @@ public Flux submitOpenConnectionTasks( && 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); } }) @@ -1233,6 +1235,7 @@ private Flux openConnections( 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(); From a89271eb58051ac79a7c4b4c54d5bee88e48710e Mon Sep 17 00:00:00 2001 From: Annie Liang Date: Thu, 20 Aug 2026 11:01:57 -0700 Subject: [PATCH 3/7] Fix PPCB failover test assertion race Verify failover using a request issued after PPCB marks the original region unavailable instead of relying on the threshold-triggering request to be rerouted. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../azure/cosmos/PerPartitionCircuitBreakerE2ETests.java | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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 ebd1180aabc7..ec41312ca4a4 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 @@ -5827,8 +5827,12 @@ public void ppcbRecoveryResolvesAddressesAfterInitialAddressRefreshFailures() th partitionUnavailabilityMap, locationContextMapField)).isTrue(); assertThat(lastDiagnostics).isNotNull(); + + CosmosItemResponse failedOverResponse = container + .readItem(testObject.getId(), partitionKey, readOptions, TestObject.class) + .block(); assertContactedRegionsContain( - lastDiagnostics.getDiagnosticsContext(), + failedOverResponse.getDiagnostics().getDiagnosticsContext(), getRegionNameForAssertion(this.readRegions.get(1)), "PPCB should route the partition to the second preferred region"); From 13ed4ec4caa1b87660915a6a42d71122bd0dcc4d Mon Sep 17 00:00:00 2001 From: Annie Liang Date: Thu, 20 Aug 2026 11:09:51 -0700 Subject: [PATCH 4/7] Address PPCB recovery review feedback Skip unsupported gateway and thin-client test rows and preserve the existing configurable recovery-timeout minimum while using the new 20-second default. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../PerPartitionCircuitBreakerE2ETests.java | 9 +++++++++ .../cosmos/implementation/ConfigsTests.java | 17 +++++++++++++++++ .../azure/cosmos/implementation/Configs.java | 9 +++++++-- 3 files changed, 33 insertions(+), 2 deletions(-) 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 ec41312ca4a4..04d28ff94d9e 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 @@ -5717,6 +5717,15 @@ public void ppcbRecoveryResolvesAddressesAfterInitialAddressRefreshFailures() th 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()); 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 f8d0a5c0292b..e82e7db3b87a 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 @@ -155,6 +155,23 @@ public void http2MinConnectionPoolSize() { } } + @Test(groups = { "unit" }) + public void connectionEstablishmentTimeoutForPartitionRecovery() { + String propertyName = "COSMOS.CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS"; + System.clearProperty(propertyName); + assertThat(Configs.getConnectionEstablishmentTimeoutForPartitionRecoveryInSeconds()).isEqualTo(20); + + System.setProperty(propertyName, "15"); + try { + assertThat(Configs.getConnectionEstablishmentTimeoutForPartitionRecoveryInSeconds()).isEqualTo(15); + + System.setProperty(propertyName, "5"); + assertThat(Configs.getConnectionEstablishmentTimeoutForPartitionRecoveryInSeconds()).isEqualTo(10); + } finally { + System.clearProperty(propertyName); + } + } + @Test(groups = { "unit" }) public void http2MaxConcurrentStreams() { assertThat(Configs.getHttp2MaxConcurrentStreams()).isEqualTo(30); 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 30d619805045..7da9a0937273 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 @@ -361,6 +361,7 @@ public class Configs { // Below sets a time limit on how long these connection establishments be attempted for // Covers one address probe pass with the default proactive connection retry policy. private static final int DEFAULT_CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS = 20; + private static final int MIN_CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS = 10; private static final String CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS = "COSMOS.CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS"; private static final boolean DEFAULT_SHOULD_LOG_INCORRECTLY_MAPPED_SESSION_TOKEN = true; @@ -1489,13 +1490,17 @@ public static int getConnectionEstablishmentTimeoutForPartitionRecoveryInSeconds String valueFromSystemProperty = System.getProperty(CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS); if (StringUtils.isNotEmpty(valueFromSystemProperty)) { - return Math.max(Integer.parseInt(valueFromSystemProperty), DEFAULT_CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS); + return Math.max( + Integer.parseInt(valueFromSystemProperty), + MIN_CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS); } String valueFromEnvVariable = System.getenv(CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS); if (StringUtils.isNotEmpty(valueFromEnvVariable)) { - return Math.max(Integer.parseInt(valueFromEnvVariable), DEFAULT_CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS); + return Math.max( + Integer.parseInt(valueFromEnvVariable), + MIN_CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS); } return DEFAULT_CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS; From 57979545f7a361cc0eb477c0146c4358bc6a08e6 Mon Sep 17 00:00:00 2001 From: Annie Liang Date: Thu, 20 Aug 2026 11:11:05 -0700 Subject: [PATCH 5/7] Keep PPCB recovery timeout at ten seconds Revert the recovery timeout configuration change and retain the existing ten-second default and minimum. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cosmos/implementation/ConfigsTests.java | 17 ----------------- .../azure/cosmos/implementation/Configs.java | 12 +++--------- 2 files changed, 3 insertions(+), 26 deletions(-) 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 e82e7db3b87a..f8d0a5c0292b 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 @@ -155,23 +155,6 @@ public void http2MinConnectionPoolSize() { } } - @Test(groups = { "unit" }) - public void connectionEstablishmentTimeoutForPartitionRecovery() { - String propertyName = "COSMOS.CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS"; - System.clearProperty(propertyName); - assertThat(Configs.getConnectionEstablishmentTimeoutForPartitionRecoveryInSeconds()).isEqualTo(20); - - System.setProperty(propertyName, "15"); - try { - assertThat(Configs.getConnectionEstablishmentTimeoutForPartitionRecoveryInSeconds()).isEqualTo(15); - - System.setProperty(propertyName, "5"); - assertThat(Configs.getConnectionEstablishmentTimeoutForPartitionRecoveryInSeconds()).isEqualTo(10); - } finally { - System.clearProperty(propertyName); - } - } - @Test(groups = { "unit" }) public void http2MaxConcurrentStreams() { assertThat(Configs.getHttp2MaxConcurrentStreams()).isEqualTo(30); 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 7da9a0937273..686efabe8188 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 @@ -359,9 +359,7 @@ public class Configs { // For partition-level circuit breaker, in order to recover a partition in a region, the SDK when configured // in the direct connectivity mode, establishes connections to replicas to attempt to recover a region // Below sets a time limit on how long these connection establishments be attempted for - // Covers one address probe pass with the default proactive connection retry policy. - private static final int DEFAULT_CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS = 20; - private static final int MIN_CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS = 10; + private static final int DEFAULT_CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS = 10; private static final String CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS = "COSMOS.CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS"; private static final boolean DEFAULT_SHOULD_LOG_INCORRECTLY_MAPPED_SESSION_TOKEN = true; @@ -1490,17 +1488,13 @@ public static int getConnectionEstablishmentTimeoutForPartitionRecoveryInSeconds String valueFromSystemProperty = System.getProperty(CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS); if (StringUtils.isNotEmpty(valueFromSystemProperty)) { - return Math.max( - Integer.parseInt(valueFromSystemProperty), - MIN_CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS); + return Math.max(Integer.parseInt(valueFromSystemProperty), DEFAULT_CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS); } String valueFromEnvVariable = System.getenv(CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS); if (StringUtils.isNotEmpty(valueFromEnvVariable)) { - return Math.max( - Integer.parseInt(valueFromEnvVariable), - MIN_CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS); + return Math.max(Integer.parseInt(valueFromEnvVariable), DEFAULT_CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS); } return DEFAULT_CONNECTION_ESTABLISHMENT_TIMEOUT_FOR_PARTITION_RECOVERY_IN_SECONDS; From c47c73ffa4a5c1df4c5c57ffec678124e125a925 Mon Sep 17 00:00:00 2001 From: Annie Liang Date: Thu, 20 Aug 2026 11:58:05 -0700 Subject: [PATCH 6/7] Keep PPCB recovery fault active through failover Avoid exhausting the address-refresh fault exactly as PPCB opens, which allowed the background recovery sweep to clear the unavailable state before the test could observe failover. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../com/azure/cosmos/PerPartitionCircuitBreakerE2ETests.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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 04d28ff94d9e..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 @@ -5805,7 +5805,8 @@ public void ppcbRecoveryResolvesAddressesAfterInitialAddressRefreshFailures() th .times(3) .build()) .duration(Duration.ofMinutes(10)) - .hitLimit(30) + // Keep recovery probes faulted until the test has observed failover. + .hitLimit(60) .build(); CosmosFaultInjectionHelper.configureFaultInjectionRules( container, @@ -5830,7 +5831,7 @@ public void ppcbRecoveryResolvesAddressesAfterInitialAddressRefreshFailures() th } } - assertThat(addressRefreshRule.getHitCount()).isEqualTo(30); + assertThat(addressRefreshRule.getHitCount()).isGreaterThanOrEqualTo(30); assertThat(hasUnavailableLocationForPartition( partitionKeyRangeWrapper, partitionUnavailabilityMap, From dab5ef1a8a2edfff13229a723aa51a8ceb210e9a Mon Sep 17 00:00:00 2001 From: Annie Liang Date: Thu, 20 Aug 2026 14:14:34 -0700 Subject: [PATCH 7/7] Fix fail-fast cache test assertion Do not require every stale replica probe to run after the first network failure terminates the pass. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../directconnectivity/GatewayAddressCacheTest.java | 1 - 1 file changed, 1 deletion(-) 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 fa7ac8e85946..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 @@ -1715,7 +1715,6 @@ public void submitOpenConnectionTasksRefreshesAddressesAfterNetworkFailure() thr assertThat(addressResolutionCount).hasValue(2); assertThat(forceRefreshValues).containsExactly(false, true); assertThat(connectionAttempts.get(stalePrimary.getPhyicalUri())).hasValue(2); - assertThat(connectionAttempts.get(staleSecondary.getPhyicalUri())).hasValue(2); assertThat(connectionAttempts.get(refreshedPrimary.getPhyicalUri())).hasValue(1); assertThat(connectionAttempts.get(refreshedSecondary.getPhyicalUri())).hasValue(1); }