From 30d662953026976f0809df124f90234626c11413 Mon Sep 17 00:00:00 2001 From: Annie Liang <64233642+xinlian12@users.noreply.github.com> Date: Thu, 20 Aug 2026 20:29:52 -0700 Subject: [PATCH 01/14] Fix PPCB failback with missing or stale addresses (#50182) * Fix PPCB failback with missing or stale addresses --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...titionEndpointManagerForPPCBUnitTests.java | 260 +++++++++++++++ .../PerPartitionCircuitBreakerE2ETests.java | 302 ++++++++++++++++++ .../GatewayAddressCacheTest.java | 293 ++++++++++++++++- sdk/cosmos/azure-cosmos/CHANGELOG.md | 114 +++++++ .../GatewayAddressCache.java | 75 ++++- ...tManagerForPerPartitionCircuitBreaker.java | 28 +- 6 files changed, 1064 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 11d14758f352..db4f665a3f05 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,11 +18,21 @@ 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; import com.azure.cosmos.implementation.perPartitionCircuitBreaker.LocationSpecificHealthContext; import com.azure.cosmos.implementation.guava25.collect.ImmutableList; 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; @@ -27,17 +40,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; @@ -52,6 +77,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; @@ -121,6 +151,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 { @@ -1007,6 +1046,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 10c5aa400894..0f7580c80ed7 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.ConnectionPolicy; import com.azure.cosmos.implementation.DatabaseAccount; @@ -5175,6 +5176,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) { @@ -5205,4 +5226,285 @@ public AccountLevelLocationContext( this.regionNameToEndpoint = regionNameToEndpoint; } } + + @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() { + + if (this.writeRegions == null || this.writeRegions.size() <= 1) { + throw new SkipException("Test requires multi-region account"); + } + + // Build non-canonical preferred regions: "West US 3" → "westus3", "East US" → "eastus" + List nonCanonicalRegions = new ArrayList<>(); + for (String region : this.writeRegions) { + nonCanonicalRegions.add(region.toLowerCase(Locale.ROOT).replace(" ", "")); + } + + String firstRegionCanonicalLower = this.writeRegions.get(0).toLowerCase(Locale.ROOT); + String secondRegionCanonicalLower = this.writeRegions.get(1).toLowerCase(Locale.ROOT); + + System.setProperty( + "COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG", + "{\"isPartitionLevelCircuitBreakerEnabled\": true, " + + "\"circuitBreakerType\": \"CONSECUTIVE_EXCEPTION_COUNT_BASED\"," + + "\"consecutiveExceptionCountToleratedForReads\": 10," + + "\"consecutiveExceptionCountToleratedForWrites\": 5," + + "}"); + + CosmosClientBuilder clientBuilder = getClientBuilder() + .multipleWriteRegionsEnabled(true) + .preferredRegions(nonCanonicalRegions); + + ConnectionPolicy connectionPolicy = ReflectionUtils.getConnectionPolicy(clientBuilder); + 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"); + } + + CosmosAsyncClient asyncClient = null; + + try { + asyncClient = clientBuilder.buildAsyncClient(); + + CosmosAsyncContainer container = asyncClient + .getDatabase(this.sharedAsyncDatabaseId) + .getContainer(this.sharedMultiPartitionAsyncContainerIdWhereIdIsPartitionKey); + + // Bootstrap: create a test item + TestObject testObject = TestObject.create(); + container.createItem(testObject, new PartitionKey(testObject.getId()), new CosmosItemRequestOptions()).block(); + + // Step 1: Inject 503 (ServiceUnavailable) into the first preferred region for READ_ITEM + FaultInjectionCondition faultCondition = new FaultInjectionConditionBuilder() + .region(this.writeRegions.get(0)) + .operationType(FaultInjectionOperationType.READ_ITEM) + .build(); + + FaultInjectionServerErrorResult serverError = FaultInjectionResultBuilders + .getResultBuilder(FaultInjectionServerErrorType.SERVICE_UNAVAILABLE) + .build(); + + FaultInjectionRule faultRule = new FaultInjectionRuleBuilder("ppcb-non-canonical-region-test-" + UUID.randomUUID()) + .condition(faultCondition) + .result(serverError) + .hitLimit(15) + .build(); + + CosmosFaultInjectionHelper.configureFaultInjectionRules(container, Arrays.asList(faultRule)).block(); + + // Step 2: Issue reads until circuit breaker trips — expect failover to second region + boolean circuitBreakerTripped = false; + + for (int i = 0; i < 20; i++) { + CosmosItemRequestOptions readOptions = new CosmosItemRequestOptions(); + readOptions.setCosmosEndToEndOperationLatencyPolicyConfig(NO_END_TO_END_TIMEOUT); + + CosmosItemResponse readResponse = container + .readItem(testObject.getId(), new PartitionKey(testObject.getId()), readOptions, TestObject.class) + .block(); + + assertThat(readResponse).isNotNull(); + assertThat(readResponse.getStatusCode()).isEqualTo(200); + + CosmosDiagnosticsContext ctx = readResponse.getDiagnostics().getDiagnosticsContext(); + + // Once we see only the second region contacted, the circuit breaker has tripped + if (ctx.getContactedRegionNames().contains(secondRegionCanonicalLower) + && !ctx.getContactedRegionNames().contains(firstRegionCanonicalLower)) { + circuitBreakerTripped = true; + logger.info("Circuit breaker tripped at iteration {}, routing to second region: {}", i, secondRegionCanonicalLower); + break; + } + } + + assertThat(circuitBreakerTripped) + .as("PPCB should have tripped and routed reads to the second preferred region (%s) " + + "even though preferred regions were passed in non-canonical form (%s)", + secondRegionCanonicalLower, nonCanonicalRegions) + .isTrue(); + + } finally { + System.clearProperty("COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG"); + if (asyncClient != null) { + asyncClient.close(); + } + } + } } 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 8285ea915603..2247c2eea902 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 @@ -16,7 +16,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,7 +34,7 @@ import com.azure.cosmos.implementation.http.HttpClientConfig; import com.azure.cosmos.implementation.routing.PartitionKeyRangeIdentity; import com.azure.cosmos.models.PartitionKeyDefinition; -import io.reactivex.subscribers.TestSubscriber; +import io.netty.channel.ConnectTimeoutException; import org.assertj.core.api.AssertionsForClassTypes; import org.mockito.ArgumentCaptor; import org.mockito.ArgumentMatchers; @@ -55,10 +57,13 @@ 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.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; @@ -1594,6 +1599,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 c76189f5b3dd..0c0115a04fb8 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -1,5 +1,119 @@ ## Release History +### 4.82.0-beta.1 (Unreleased) + +#### Features Added +* Enabled Gateway V2 (thin-client) data-plane routing by default for `Cosmos(Async)Client` instances configured with `gatewayMode` and HTTP/2, gated by an HTTP/2 connectivity probe with automatic fallback to Gateway V1. - See [PR 49437](https://github.com/Azure/azure-sdk-for-java/pull/49437) +* Added support for QueryPlan and Execute Stored Procedure requests to be routed to Gateway V2. - See [PR 47759](https://github.com/Azure/azure-sdk-for-java/pull/47759) + +#### 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). +* 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). +* Fixed thin-client (Gateway V2) queries with a prefix (partial) hierarchical partition key returning co-located documents from other logical partitions. - See PR [49688](https://github.com/Azure/azure-sdk-for-java/pull/49688). +* Fixed hedged requests losing request-scoped routing, timeout, authorization, throughput-control, and metadata state when cloning the original request. - See [PR 50069](https://github.com/Azure/azure-sdk-for-java/pull/50069). + +#### Other Changes +* Reduced memory footprint of deserialized `PartitionKeyRange` instances by stripping unused fields in the `PartitionKeyRange(ObjectNode)` constructor - See PR [49513](https://github.com/Azure/azure-sdk-for-java/pull/49513). +* Added bounded retries for transient "collection routing map / partition key range metadata not available" responses (HTTP 404 with sub-status `0`, `1003`, or `1013`) that can briefly occur right after a container is (re)created, improving the robustness of data-plane operations against the post-creation metadata-propagation race. As part of this change, when the routing map remains unavailable after retries an operation now fails with a `CosmosException` (HTTP 404, sub-status `1024` / `INCORRECT_CONTAINER_RID`) instead of an internal `IllegalStateException`. - See [PR 49639](https://github.com/Azure/azure-sdk-for-java/pull/49639). +* Reduced memory footprint and redundant `/pkranges` reads when multiple `CosmosClient` / `CosmosAsyncClient` instances in the same JVM are configured with the same service endpoint. Disable with system property `COSMOS.SHARED_PARTITION_KEY_RANGE_CACHE_ENABLED=false` if needed. - See [PR 49560](https://github.com/Azure/azure-sdk-for-java/pull/49560). + +### 4.81.0 (2026-06-08) + +#### Features Added +* Added support for creating Global Secondary Index (GSI) containers via `CosmosContainerProperties.setGlobalSecondaryIndexDefinition()` / `getGlobalSecondaryIndexDefinition()`, the new `CosmosGlobalSecondaryIndexDefinition` model, and the `CosmosGlobalSecondaryIndexBuildStatus` enum returned by `getStatus()`. - See [PR 48480](https://github.com/Azure/azure-sdk-for-java/pull/48480) +* Promoted the Full Fidelity Change Feed (AllVersionsAndDeletes) APIs to GA - See [PR 49283](https://github.com/Azure/azure-sdk-for-java/pull/49283) +* Enabled `ReadConsistencyStrategy` for Gateway V1 (compute gateway) and Gateway V2 (thin client proxy). Previously only supported in Direct mode. - See [PR 48787](https://github.com/Azure/azure-sdk-for-java/pull/48787) + +#### Bugs Fixed +* Fixed region name normalization for preferred and excluded regions — non-canonical inputs (e.g., `"westus3"`, `"WEST US 3"`) are now mapped to the canonical form. Also fixed a case-sensitive exclude-region check in PPCB reevaluate logic. - See [PR 49090](https://github.com/Azure/azure-sdk-for-java/pull/49090) +* Fixed `UnsupportedOperationException` when using `readManyByPartitionKeys` for empty pages. - See [PR 49311](https://github.com/Azure/azure-sdk-for-java/pull/49311) +* Fixed silent drift in `CosmosChangeFeedRequestOptions` when resuming from a continuation token via `byPage(savedContinuation)`. Previously only `maxPrefetchPageCount` and `throughputControlGroupName` were inherited onto the rebuilt impl; `endLSN`, `customSerializer`, `excludeRegions`, `readConsistencyStrategy`, `completeAfterAllCurrentChangesRetrieved`, and other caller-supplied configuration were silently dropped. All non-token-encoded fields are now propagated. - See [PR 49276](https://github.com/Azure/azure-sdk-for-java/pull/49276) +* Fixed HTTP/2 PING keepalive handler (introduced in [PR 49095](https://github.com/Azure/azure-sdk-for-java/pull/49095)) so it observes child-stream HEADERS/DATA reads via `Http2PingCloseRewrapHandler.channelReadComplete`, preventing spurious PINGs (and spurious closes) on connections actively serving requests through `Http2MultiplexHandler`. + +#### Other Changes +* Added HTTP/2 PING keepalive (default ON) for Gateway service endpoints to detect silently-broken connections. - See [PR 49095](https://github.com/Azure/azure-sdk-for-java/pull/49095) +* Replaced per-client `Schedulers.newSingle()` schedulers in `GlobalEndpointManager` and `GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker` with shared `BoundedElastic` schedulers in `CosmosSchedulers` to prevent thread count from scaling linearly with client/tenant count. - See [PR 49062](https://github.com/Azure/azure-sdk-for-java/pull/49062) +* Promoted the `ReadConsistencyStrategy` and `Http2ConnectionConfig` related `@Beta` APIs to GA. - See [PR 49345](https://github.com/Azure/azure-sdk-for-java/pull/49345) +* Fixed a sporadic `NullPointerException` in `JsonSerializable.getWithMapping` triggered by concurrent first-time calls to `DatabaseAccount.getConsistencyPolicy()` and its sibling lazy getters (`getReplicationPolicy`, `getSystemReplicationPolicy`, `getQueryEngineConfiguration`). The fix makes `JsonSerializable.propertyBag` `final`, closing an unsafe-publication race in the lazy-initialisation pattern. - See [Issue 49256](https://github.com/Azure/azure-sdk-for-java/issues/49256) and [PR #49258](https://github.com/Azure/azure-sdk-for-java/pull/49258) +* Changed 449 (`Retry With`) retries in Gateway V1 and Gateway V2 to be consistently orchestrated client-side. - See [PR 49332](https://github.com/Azure/azure-sdk-for-java/pull/49332) +* Added client-side fast-fail validation for `ReadConsistencyStrategy.GLOBAL_STRONG`: requests that specify `GLOBAL_STRONG` against an account whose default consistency is not `STRONG` are now rejected client-side with a `BadRequestException` (HTTP 400). - See [PR 48787](https://github.com/Azure/azure-sdk-for-java/pull/48787) + +### 4.80.0 (2026-05-01) + +#### Features Added +* Added support for Query Advisor feature - See [48160](https://github.com/Azure/azure-sdk-for-java/pull/48160) +* Added `additionalHeaders` support to allow setting additional headers (e.g., `x-ms-cosmos-workload-id`) that are sent with every request. - See [PR 48128](https://github.com/Azure/azure-sdk-for-java/pull/48128) +* Added `IGNORE_UNKNOWN_RNTBD_TOKENS` SDK capability flag and propagated SDK supported capabilities to barrier requests, enabling N-Region Synchronous Commit to function correctly with backends that return new RNTBD response tokens. - See [PR 48965](https://github.com/Azure/azure-sdk-for-java/pull/48965) +* Added support for change feed with `startFrom` point-in-time on merged partitions by enabling the `CHANGE_FEED_WITH_START_TIME_POST_MERGE` SDK capability. - See [PR 48752](https://github.com/Azure/azure-sdk-for-java/pull/48752) +* Added new `readManyByPartitionKeys` API on `CosmosAsyncContainer` / `CosmosContainer` to bulk-query all documents matching a list of partition key values with better efficiency than issuing individual queries. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added `CosmosReadManyByPartitionKeysRequestOptions` - a dedicated request-options type for `readManyByPartitionKeys` that exposes `setContinuationToken(String)` for resuming previous invocations and `setMaxConcurrentBatchPrefetch(int)` to bound per-call prefetch parallelism. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Added `CosmosReadManyByPartitionKeysRequestOptions.setMaxBatchSize(Integer)` to set the max. number of partition keys used for a single batch. See [PR 48930](https://github.com/Azure/azure-sdk-for-java/pull/48930) +* Added `getCustomItemSerializer()` to `CosmosRequestContext` and `setCustomItemSerializer(CosmosItemSerializer)` to `CosmosRequestOptions` to allow overriding the custom item serializer via operation policies. - See [PR 48963](https://github.com/Azure/azure-sdk-for-java/pull/48963) + +#### Bugs Fixed +* Fixed `readMany` and `readAllItems` returning incorrect results on containers whose partition key path is nested (e.g. `/address/city`) due to malformed selector generation. - See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) +* Fixed an issue where the throughput control `throughputQueryMono` was always subscribed even when `targetThroughput` is used (not `targetThroughputThreshold`), causing unnecessary `throughputSettings/read` permission requirement for AAD principals. - See [PR 48800](https://github.com/Azure/azure-sdk-for-java/pull/48800) +* Fixed JVM `` deadlock when multiple threads concurrently trigger Cosmos SDK class loading for the first time. - See [PR 48689](https://github.com/Azure/azure-sdk-for-java/pull/48689) +* Fixed an issue where `CustomItemSerializer` was incorrectly applied to internal SDK query pipeline structures (e.g., `OrderByRowResult`, `Document`), causing deserialization failures in ORDER BY, GROUP BY, aggregate, DISTINCT, and hybrid search queries. - See [PR 48811](https://github.com/Azure/azure-sdk-for-java/pull/48811) +* Fixed an issue where `SqlParameter` ignored the configured `CustomItemSerializer`, always using the internal default serializer instead. - See [PR 48811](https://github.com/Azure/azure-sdk-for-java/pull/48811) +* Fixed a `ClientTelemetry` static initialization failure when IMDS access is disabled, preventing `NoClassDefFoundError` during Cosmos client creation in non-Azure environments. - See [PR 48888](https://github.com/Azure/azure-sdk-for-java/pull/48888) +* Fixed an issue where Netty could log "An exceptionCaught() event was fired, and it reached at the tail of the pipeline" on HTTP/2 connections when the server resets idle TCP connections by adding an exception handler on the HTTP/2 parent channel to handle these connection-level exceptions more appropriately. - See [PR 48890](https://github.com/Azure/azure-sdk-for-java/pull/48890) +* Fixed an issue where `CustomItemSerializer` configured on `CosmosClientBuilder` was not honored for response deserialization in `CosmosAsyncContainer.upsertItem` when no request-level serializer was set. - See [PR 48962](https://github.com/Azure/azure-sdk-for-java/pull/48962) + +### 4.79.1 (2026-04-06) + +#### Bugs Fixed +* Fixing an NPE caused due to boxed Boolean conversion. - See [PR 48656](https://github.com/Azure/azure-sdk-for-java/pull/48656/) + +### 4.79.0 (2026-03-27) + +#### Features Added +* Added support for N-Region synchronous commit feature - See [PR 47757](https://github.com/Azure/azure-sdk-for-java/pull/47757) +* Added support for Query Advisor feature - See [48160](https://github.com/Azure/azure-sdk-for-java/pull/48160) +* Added `CosmosFullTextScoreScope` enum and `setFullTextScoreScope()` on `CosmosQueryRequestOptions` for controlling BM25 statistics scope in hybrid search queries. Supports `LOCAL` (scoped to target partitions) and `GLOBAL` (default, all partitions) scopes. See [PR 48431](https://github.com/Azure/azure-sdk-for-java/pull/48431) + +#### Bugs Fixed +* Fixed Remote Code Execution (RCE) vulnerability (CWE-502) by replacing Java deserialization with JSON-based serialization in `CosmosClientMetadataCachesSnapshot`, `AsyncCache`, and `DocumentCollection`. The metadata cache snapshot now uses Jackson for serialization/deserialization, eliminating the entire class of Java deserialization attacks. - [PR 47971](https://github.com/Azure/azure-sdk-for-java/pull/47971) +* Fixed `NullPointerException` in `DocumentQueryExecutionContextFactory.tryCacheQueryPlan` when executing hybrid search queries with a partition key filter. See [PR 48431](https://github.com/Azure/azure-sdk-for-java/pull/48431) +* Fixed `ConcurrentModificationException` in hybrid search component query execution caused by concurrent access to shared mutable state. See [PR 48431](https://github.com/Azure/azure-sdk-for-java/pull/48431) +* Fixed availability strategy for Gateway V2 (thin client) by ensuring `RegionalRoutingContext` identity is based only on the immutable gateway endpoint. - See [PR 48432](https://github.com/Azure/azure-sdk-for-java/pull/48432) +* Fixed an issue where `replaceItem` bypassed the `customItemSerializer`, serialising POJOs with the SDK's internal `ObjectMapper` instead of the user-configured one. - See [PR 48529](https://github.com/Azure/azure-sdk-for-java/pull/48529) +* Fixed `ClassCastException` (`ArrayNode cannot be cast to ObjectNode`) when executing `SELECT VALUE ... GROUP BY` queries. See - [PR 48507](https://github.com/Azure/azure-sdk-for-java/pull/48507) + +#### Other Changes +* Promoted the following `@Beta` APIs to GA: `CosmosContainerProperties.getFullTextPolicy()`/`setFullTextPolicy()`, `IndexingPolicy.getCosmosFullTextIndexes()`/`setCosmosFullTextIndexes()`. - See [PR 48538](https://github.com/Azure/azure-sdk-for-java/pull/48538) +* Added `appendUserAgentSuffix` method to `AsyncDocumentClient` to allow downstream libraries to append to the user agent after client construction. - See [PR 48505](https://github.com/Azure/azure-sdk-for-java/pull/48505) +* Added aggressive HTTP timeout policies for document operations routed to Gateway V2. - [PR 47879](https://github.com/Azure/azure-sdk-for-java/pull/47879) +* Added a default connect timeout of 5s for Gateway V2 (thin client) data-plane endpoints. - See [PR 48174](https://github.com/Azure/azure-sdk-for-java/pull/48174) +* Added system property `COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_MS` and environment variable `COSMOS_CONNECTION_ACQUIRE_TIMEOUT_IN_MS` to allow overriding the gateway connection acquire timeout in milliseconds (default 45000ms). Minimum accepted value is 500ms. Replaces the previous `_IN_SECONDS` variants. - See [PR 48580](https://github.com/Azure/azure-sdk-for-java/pull/48580) +* Changed system property for thin client connection timeout from `COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS` to `COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_MS` (default 5000ms, minimum 500ms). - See [PR 48580](https://github.com/Azure/azure-sdk-for-java/pull/48580) + +### 4.78.0 (2026-02-10) + +#### Features Added +* Added shardKey support in `DedicatedGatewayRequestOptions` to allow specifying a shard key for dedicated gateway sharding support. - See [PR 47796](https://github.com/Azure/azure-sdk-for-java/pull/47796) + +#### Bugs Fixed +* Fixed an issue where `query plan` failed with `400` or query return empty result when `CosmosQueryRequestOptions` has partition key filter and partition key value contains non-ascii character. See [PR 47881](https://github.com/Azure/azure-sdk-for-java/pull/47881) +* Fixed an issue where operation failed with `400` when configured with pre-trigger or post-trigger with non-ascii character. Only impact for gateway mode. See [PR 47881](https://github.com/Azure/azure-sdk-for-java/pull/47881) + +#### Other Changes +* Added `x-ms-hub-region-processing-only` header to allow hub-region stickiness when 404 `READ SESSION NOT AVAILABLE` is hit for Single-Writer accounts. - [PR 47631](https://github.com/Azure/azure-sdk-for-java/pull/47631) + +### 4.77.0 (2026-01-26) + +#### Features Added +* Added `ChangeFeedProcessorOptions#setMaxLeasesToAcquirePerCycle(int)` to allow faster acquisition of unused/expired leases during scale-out and rolling deployments (default `0` preserves legacy behavior). - [47606](https://github.com/Azure/azure-sdk-for-java/pull/47606) +* Added the `QuantizerType` to the vectorIndexSpec: `product`/`spherical`. - [PR 47566](https://github.com/Azure/azure-sdk-for-java/pull/47566) + +#### Other Changes +* Remaps sub-status to 1003 for requests to child resources against non-existent container. - [PR 47604](https://github.com/Azure/azure-sdk-for-java/pull/47604) + ### 4.76.0 (2025-12-09) #### Bugs Fixed 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 e62d7b8c6ca4..fc4f7f0575fa 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; @@ -1145,7 +1146,8 @@ public Mono submitOpenConnectionTask( public Flux submitOpenConnectionTasks( PartitionKeyRange partitionKeyRange, - String collectionRid) { + String collectionRid, + boolean forceRefresh) { if (this.proactiveOpenConnectionsProcessor == null) { return Flux.empty(); @@ -1156,14 +1158,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 0213f2255144..3eacc111e522 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 @@ -19,6 +19,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; @@ -341,8 +343,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 : " + @@ -362,6 +375,7 @@ private Flux updateStaleLocationInfo() { false, true); } + return locationSpecificContextAsVal; }); }) @@ -399,6 +413,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 fcd5023779889ef288325d3f78eb2264fab2fdf9 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Fri, 21 Aug 2026 19:32:33 -0400 Subject: [PATCH 02/14] Fix PPCB diagnostics and failback observability. (#50158) * Improve PPCB failback diagnostics * Validate PPCB state in diagnostics E2E tests * Limit PPCB diagnostics assertion to data-plane requests * Log PPCB failback backlog progress * Add PPCB failback remaining meter * Clarify PPCB failback meter name * Reduce PPCB failback meter allocations * Track PPCB pending recoveries by collection * Correlate PPCB failback recovery diagnostics * Harden PPCB failback recovery tests * Fix PPCB all-region diagnostics assertion * Optimize PPCB diagnostics snapshots * Cache PPCB diagnostics snapshots * Align PPCB failback flow with main * Add benchmark fault injection support. * Harden PPCB failback telemetry Use a single injectable logger, prevent backlog telemetry failures from escaping the recovery flow, and align the pending failback metric name. * Scope PR to PPCB diagnostics Remove benchmark, fault-injection, metric, and recovery behavior changes. Retain immutable CosmosDiagnostics PPCB snapshots, lifecycle E2E assertions, and WARN logging for every failback failure. * Reduce PPCB diagnostics overhead Reuse immutable PPCB map references in response snapshots and shorten per-region diagnostic field names. * Refactoring * Simplify PPCB diagnostics snapshots Represent holder state with one volatile immutable-map reference and align the compact timestamp serialization test. * Ignore updates to empty PPCB diagnostics Make updates to the shared uninitialized diagnostics sentinel a no-op and verify it remains null-serializing. * Avoid copying PPCB diagnostics state Retain the live PPCB diagnostics map reference to avoid per-publication map and wrapper allocations, accepting weak consistency. * Log PPCB diagnostics lifecycle snapshots Emit one full CosmosDiagnostics JSON payload for failed, post-failover, and post-failback E2E phases for PR evidence. * Document PPCB diagnostics improvements Add the unreleased changelog entry for per-region PPCB snapshots and failback WARN logging. * Add PPCB failback outcome diagnostics Track the latest background failback attempt time, outcome, and failure reason per partition-region; validate lifecycle state in E2E and focused recovery tests. * Reset unavailable timestamp after PPCB failback Use the available-state sentinel when recovery moves a region to HealthyTentative and cover it in the scheduled recovery test. * Refine PPCB failback diagnostics Keep failback attempt metadata partition-scoped, retain only the latest full failure message per region, and clear retained messages when recovery backlog drains. --- ...titionEndpointManagerForPPCBUnitTests.java | 21 + .../PerPartitionCircuitBreakerE2ETests.java | 366 +++++++++++++++++- ...PartitionCircuitBreakerInfoHolderTest.java | 250 ++++++++++++ .../PpcbFailbackLoggingTest.java | 171 ++++++++ sdk/cosmos/azure-cosmos/CHANGELOG.md | 1 + .../ClientSideRequestStatistics.java | 42 +- ...nsecutiveExceptionBasedCircuitBreaker.java | 6 +- ...tManagerForPerPartitionCircuitBreaker.java | 296 ++++++++++++-- .../LocationSpecificHealthContext.java | 97 ++++- ...pecificHealthContextTransitionHandler.java | 46 ++- .../PerPartitionCircuitBreakerInfoHolder.java | 60 ++- 11 files changed, 1282 insertions(+), 74 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/PerPartitionCircuitBreakerInfoHolderTest.java create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/PpcbFailbackLoggingTest.java 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 db4f665a3f05..22fd70774881 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 @@ -32,6 +32,7 @@ import com.azure.cosmos.implementation.perPartitionCircuitBreaker.LocationSpecificHealthContext; import com.azure.cosmos.implementation.guava25.collect.ImmutableList; import com.azure.cosmos.implementation.routing.RegionalRoutingContext; +import com.fasterxml.jackson.databind.ObjectMapper; import io.netty.channel.ConnectTimeoutException; import org.apache.commons.lang3.tuple.Pair; import org.mockito.Mockito; @@ -1179,13 +1180,33 @@ public Mono> getServerAddressesViaGatewayAsync( collectionRid, partitionKeyRange)).containsExactly("East US"); assertThat(refreshedConnectionAttempts).hasValue(1); + String diagnostics = new ObjectMapper().writeValueAsString( + request.requestContext.getPerPartitionCircuitBreakerInfoHolder()); + assertThat(diagnostics) + .contains("\"outcome\":\"Failed\"") + .contains("\"stage\":\"OPEN_CONNECTION_TASK\"") + .contains("\"type\":\"io.netty.channel.ConnectTimeoutException\"") + .contains("\"latestFailbackMessageByRegion\":{") + .contains("\"East US\":\"Refreshed replica is unavailable\""); } else { assertThat(ppcbManager.getUnavailableRegionsForPartitionKeyRange( request, collectionRid, partitionKeyRange)).isEmpty(); + assertThat(request.requestContext.getPerPartitionCircuitBreakerInfoHolder() + .getPerPartitionCircuitBreakerInfoHolder() + .get("East US") + .getUnavailableSince()).isEqualTo(Instant.MAX); + assertThat(new ObjectMapper().writeValueAsString( + request.requestContext.getPerPartitionCircuitBreakerInfoHolder())) + .contains("\"outcome\":\"Succeeded\"") + .doesNotContain("\"failure\"", "\"latestFailbackMessageByRegion\""); } + assertThat(new ObjectMapper().writeValueAsString( + request.requestContext.getPerPartitionCircuitBreakerInfoHolder())) + .contains("\"lastAttemptedAt\":"); + if (populateStaleAddress) { assertThat(forceRefreshValues).containsExactly(false, true); assertThat(addressResolutionCount).hasValue(2); 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 0f7580c80ed7..99cef9422503 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 @@ -14,6 +14,7 @@ import com.azure.cosmos.implementation.ImplementationBridgeHelpers; import com.azure.cosmos.implementation.OperationType; import com.azure.cosmos.implementation.PartitionKeyRange; +import com.azure.cosmos.implementation.ResourceType; import com.azure.cosmos.implementation.RxDocumentClientImpl; import com.azure.cosmos.implementation.TestConfigurations; import com.azure.cosmos.implementation.Utils; @@ -3550,6 +3551,8 @@ private void execute( boolean hasReachedCircuitBreakingThreshold = false; int executionCountAfterCircuitBreakingThresholdBreached = 0; + boolean failbackExpected = false; + Set loggedPpcbDiagnosticsPhases = new HashSet<>(); List testObjects = operationInvocationParamsWrapper.testObjectsForDataPlaneOperationToWorkWith; PartitionKeyRangeWrapper partitionKeyRangeWrapper @@ -3563,7 +3566,12 @@ private void execute( validateNonEmptyList(operationInvocationParamsWrapper.itemIdentitiesForReadManyOperation); } - ResponseWrapper response = executeDataPlaneOperation.apply(operationInvocationParamsWrapper); + ResponseWrapper response = executeDataPlaneOperationWithTransient4041002Retry( + testId, + executeDataPlaneOperation, + operationInvocationParamsWrapper); + assertPpcbSnapshotsPopulated(response, PpcbDiagnosticsPhase.FAILURE, false); + logPpcbDiagnosticsOnce(response, PpcbDiagnosticsPhase.FAILURE, loggedPpcbDiagnosticsPhases); ConsecutiveExceptionBasedCircuitBreaker consecutiveExceptionBasedCircuitBreaker = globalPartitionEndpointManagerForPerPartitionCircuitBreaker.getConsecutiveExceptionBasedCircuitBreaker(); @@ -3589,6 +3597,14 @@ private void execute( if (executionCountAfterCircuitBreakingThresholdBreached > 1) { validateResponseInAbsenceOfFailures.accept(response); + failbackExpected |= assertPpcbSnapshotsPopulated( + response, + PpcbDiagnosticsPhase.POST_FAILOVER, + false); + logPpcbDiagnosticsOnce( + response, + PpcbDiagnosticsPhase.POST_FAILOVER, + loggedPpcbDiagnosticsPhases); } if (response.cosmosItemResponse != null) { @@ -3640,6 +3656,14 @@ private void execute( ResponseWrapper response = executeDataPlaneOperation.apply(operationInvocationParamsWrapper); validateResponseInAbsenceOfFailures.accept(response); + assertPpcbSnapshotsPopulated( + response, + PpcbDiagnosticsPhase.POST_FAILBACK, + failbackExpected); + logPpcbDiagnosticsOnce( + response, + PpcbDiagnosticsPhase.POST_FAILBACK, + loggedPpcbDiagnosticsPhases); if (response.cosmosItemResponse != null) { assertThat(response.cosmosItemResponse).isNotNull(); @@ -3677,6 +3701,334 @@ private void execute( } } + private static CosmosDiagnosticsContext getDiagnosticsContext(ResponseWrapper response) { + if (response.cosmosItemResponse != null) { + return response.cosmosItemResponse.getDiagnostics().getDiagnosticsContext(); + } else if (response.feedResponse != null) { + return response.feedResponse.getCosmosDiagnostics().getDiagnosticsContext(); + } else if (response.cosmosException != null) { + return response.cosmosException.getDiagnostics().getDiagnosticsContext(); + } else if (response.batchResponse != null) { + return response.batchResponse.getDiagnostics().getDiagnosticsContext(); + } + return null; + } + + private static void logPpcbDiagnosticsOnce( + ResponseWrapper response, + PpcbDiagnosticsPhase phase, + Set loggedPhases) { + + if (loggedPhases.add(phase)) { + CosmosDiagnosticsContext diagnosticsContext = getDiagnosticsContext(response); + if (diagnosticsContext != null) { + logger.info("PPCB CosmosDiagnostics [{}]: {}", phase.label, diagnosticsContext.toJson()); + } + } + } + + private static boolean assertPpcbSnapshotsPopulated( + ResponseWrapper response, + PpcbDiagnosticsPhase phase, + boolean failbackExpected) { + + CosmosDiagnosticsContext diagnosticsContext = getDiagnosticsContext(response); + assertThat(diagnosticsContext) + .as("Expected CosmosDiagnostics for %s", phase.label) + .isNotNull(); + assertThat(diagnosticsContext.getDiagnostics()) + .as("Expected diagnostics entries for %s", phase.label) + .isNotNull(); + + int applicableStatisticCount = 0; + List healthContexts = new ArrayList<>(); + for (CosmosDiagnostics cosmosDiagnostics : diagnosticsContext.getDiagnostics()) { + Collection statisticsCollection = + cosmosDiagnosticsAccessor.getClientSideRequestStatistics(cosmosDiagnostics); + if (statisticsCollection == null) { + continue; + } + + for (ClientSideRequestStatistics statistics : statisticsCollection) { + if (statistics == null) { + continue; + } + + for (ClientSideRequestStatistics.StoreResponseStatistics storeStatistics + : statistics.getResponseStatisticsList()) { + + if (isPpcbApplicableDataPlaneStatistic( + storeStatistics.getRequestResourceType(), + storeStatistics.getRequestOperationType())) { + + applicableStatisticCount++; + assertThat(storeStatistics.getPerPartitionCircuitBreakerInfoHolder()) + .as("Expected direct PPCB holder for %s", phase.label) + .isNotNull(); + Map stateByRegion + = storeStatistics.getPerPartitionCircuitBreakerInfoHolder() + .getPerPartitionCircuitBreakerInfoHolder(); + assertThat(stateByRegion) + .as("Expected populated direct PPCB snapshot for %s", phase.label) + .isNotNull(); + healthContexts.addAll(stateByRegion.values()); + } + } + + for (ClientSideRequestStatistics.GatewayStatistics gatewayStatistics + : statistics.getGatewayStatisticsList()) { + + if (isPpcbApplicableDataPlaneStatistic( + gatewayStatistics.getResourceType(), + gatewayStatistics.getOperationType())) { + + applicableStatisticCount++; + assertThat(gatewayStatistics.getPerPartitionCircuitBreakerInfoHolder()) + .as("Expected gateway PPCB holder for %s", phase.label) + .isNotNull(); + Map stateByRegion + = gatewayStatistics.getPerPartitionCircuitBreakerInfoHolder() + .getPerPartitionCircuitBreakerInfoHolder(); + assertThat(stateByRegion) + .as("Expected populated gateway PPCB snapshot for %s", phase.label) + .isNotNull(); + healthContexts.addAll(stateByRegion.values()); + } + } + } + } + + if (applicableStatisticCount == 0) { + assertThat(hasOnlyQueryPlanStatistics(diagnosticsContext)) + .as("Expected PPCB-applicable data-plane statistics or QueryPlan-only diagnostics for %s", phase.label) + .isTrue(); + } + + boolean unavailableRegionFound = false; + boolean successfulFailbackFound = false; + for (LocationSpecificHealthContext healthContext : healthContexts) { + if (healthContext.getLocationHealthStatus() == LocationHealthStatus.Unavailable) { + unavailableRegionFound = true; + if (phase == PpcbDiagnosticsPhase.POST_FAILOVER) { + assertThat(healthContext.getLastFailbackOutcome()) + .as("Failback must not have succeeded while the region remains unavailable") + .isNotEqualTo(LocationSpecificHealthContext.FailbackOutcome.Succeeded); + } + } + + if (healthContext.getLastFailbackOutcome() + == LocationSpecificHealthContext.FailbackOutcome.Succeeded) { + + successfulFailbackFound = true; + assertThat(healthContext.getLastFailbackAttemptTime()) + .as("Expected failback attempt timestamp after successful failback") + .isNotNull(); + assertThat(healthContext.getLocationHealthStatus()) + .as("Expected recovered region after successful failback") + .isIn(LocationHealthStatus.HealthyTentative, LocationHealthStatus.Healthy); + } + } + + if (phase == PpcbDiagnosticsPhase.POST_FAILBACK && failbackExpected) { + assertThat(successfulFailbackFound) + .as("Expected a successful failback outcome for a previously unavailable region") + .isTrue(); + } + + return unavailableRegionFound; + } + + private static boolean isPpcbApplicableDataPlaneStatistic( + ResourceType resourceType, + OperationType operationType) { + + return resourceType == ResourceType.Document && operationType != OperationType.QueryPlan; + } + + private static boolean hasOnlyQueryPlanStatistics(CosmosDiagnosticsContext diagnosticsContext) { + boolean queryPlanStatisticFound = false; + for (CosmosDiagnostics cosmosDiagnostics : diagnosticsContext.getDiagnostics()) { + Collection statisticsCollection = + cosmosDiagnosticsAccessor.getClientSideRequestStatistics(cosmosDiagnostics); + if (statisticsCollection == null) { + continue; + } + + for (ClientSideRequestStatistics statistics : statisticsCollection) { + if (statistics == null) { + continue; + } + + for (ClientSideRequestStatistics.StoreResponseStatistics storeStatistics + : statistics.getResponseStatisticsList()) { + + if (storeStatistics.getRequestResourceType() != ResourceType.Document) { + continue; + } + if (storeStatistics.getRequestOperationType() != OperationType.QueryPlan) { + return false; + } + queryPlanStatisticFound = true; + } + + for (ClientSideRequestStatistics.GatewayStatistics gatewayStatistics + : statistics.getGatewayStatisticsList()) { + + if (gatewayStatistics.getResourceType() != ResourceType.Document) { + continue; + } + if (gatewayStatistics.getOperationType() != OperationType.QueryPlan) { + return false; + } + queryPlanStatisticFound = true; + } + } + } + + return queryPlanStatisticFound; + } + + private ResponseWrapper executeDataPlaneOperationWithTransient4041002Retry( + String testId, + Function> executeDataPlaneOperation, + OperationInvocationParamsWrapper operationInvocationParamsWrapper) throws InterruptedException { + + long retryStartNanos = System.nanoTime(); + int retryAttempt = 0; + ResponseWrapper response = executeDataPlaneOperation.apply(operationInvocationParamsWrapper); + + while (hasNonFaultInjected404RetryableResponse(response)) { + Duration elapsed = Duration.ofNanos(System.nanoTime() - retryStartNanos); + if (elapsed.compareTo(TRANSIENT_404_1002_MAX_RETRY_DURATION) >= 0) { + logger.warn( + "Detected non-fault-injected retryable 404 in diagnostics for test {} for {}. " + + "Continuing with latest response so normal assertions can report diagnostics.", + testId, + elapsed); + return response; + } + + retryAttempt++; + logger.warn( + "Detected non-fault-injected retryable 404 in diagnostics for test {}. " + + "Waiting {} before retry attempt {}.", + testId, + TRANSIENT_404_1002_RETRY_DELAY, + retryAttempt); + Thread.sleep(TRANSIENT_404_1002_RETRY_DELAY.toMillis()); + response = executeDataPlaneOperation.apply(operationInvocationParamsWrapper); + } + + return response; + } + + private static boolean hasNonFaultInjected404RetryableResponse(ResponseWrapper response) { + if (!hasRetryableTerminal404(response)) { + return false; + } + + CosmosDiagnosticsContext diagnosticsContext = getDiagnosticsContext(response); + if (diagnosticsContext == null || diagnosticsContext.getDiagnostics() == null) { + return false; + } + + for (CosmosDiagnostics cosmosDiagnostics : diagnosticsContext.getDiagnostics()) { + Collection clientSideRequestStatisticsCollection = + cosmosDiagnosticsAccessor.getClientSideRequestStatistics(cosmosDiagnostics); + if (clientSideRequestStatisticsCollection == null) { + continue; + } + + for (ClientSideRequestStatistics clientSideRequestStatistics : clientSideRequestStatisticsCollection) { + if (clientSideRequestStatistics == null) { + continue; + } + + if (hasNonFaultInjected404RetryableGatewayResponse(clientSideRequestStatistics.getGatewayStatisticsList())) { + return true; + } + + if (hasNonFaultInjected404RetryableStoreResponse(clientSideRequestStatistics.getResponseStatisticsList()) + || hasNonFaultInjected404RetryableStoreResponse(clientSideRequestStatistics.getSupplementalResponseStatisticsList())) { + + return true; + } + } + } + + return false; + } + + private static boolean hasRetryableTerminal404(ResponseWrapper response) { + if (response == null) { + return false; + } + + if (response.cosmosException != null) { + return isRetryable404( + response.cosmosException.getStatusCode(), + response.cosmosException.getSubStatusCode()); + } + + return response.batchResponse != null + && isRetryable404( + response.batchResponse.getStatusCode(), + response.batchResponse.getSubStatusCode()); + } + + private static boolean hasNonFaultInjected404RetryableGatewayResponse( + List gatewayStatisticsList) { + + if (gatewayStatisticsList == null) { + return false; + } + + for (ClientSideRequestStatistics.GatewayStatistics gatewayStatistics : gatewayStatisticsList) { + if (gatewayStatistics != null + && isRetryable404(gatewayStatistics.getStatusCode(), gatewayStatistics.getSubStatusCode()) + && isNullOrEmpty(gatewayStatistics.getFaultInjectionRuleId())) { + + return true; + } + } + + return false; + } + + private static boolean hasNonFaultInjected404RetryableStoreResponse( + Collection storeResponseStatisticsCollection) { + + if (storeResponseStatisticsCollection == null) { + return false; + } + + for (ClientSideRequestStatistics.StoreResponseStatistics storeResponseStatistics : storeResponseStatisticsCollection) { + StoreResultDiagnostics storeResultDiagnostics = + storeResponseStatistics == null ? null : storeResponseStatistics.getStoreResult(); + StoreResponseDiagnostics storeResponseDiagnostics = + storeResultDiagnostics == null ? null : storeResultDiagnostics.getStoreResponseDiagnostics(); + + if (storeResponseDiagnostics != null + && isRetryable404(storeResponseDiagnostics.getStatusCode(), storeResponseDiagnostics.getSubStatusCode()) + && isNullOrEmpty(storeResponseDiagnostics.getFaultInjectionRuleId())) { + + return true; + } + } + + return false; + } + + private static boolean isRetryable404(int statusCode, int subStatusCode) { + return statusCode == HttpConstants.StatusCodes.NOTFOUND + && (subStatusCode == HttpConstants.SubStatusCodes.UNKNOWN + || subStatusCode == HttpConstants.SubStatusCodes.READ_SESSION_NOT_AVAILABLE); + } + + private static boolean isNullOrEmpty(String value) { + return value == null || value.isEmpty(); + } + private static int resolveTestObjectCountToBootstrapFrom(FaultInjectionOperationType faultInjectionOperationType, int opCount) { switch (faultInjectionOperationType) { case READ_ITEM: @@ -5211,6 +5563,18 @@ private enum QueryType { READ_MANY, READ_ALL } + private enum PpcbDiagnosticsPhase { + FAILURE("failed operation"), + POST_FAILOVER("post-failover operation"), + POST_FAILBACK("post-failback operation"); + + private final String label; + + PpcbDiagnosticsPhase(String label) { + this.label = label; + } + } + private static class AccountLevelLocationContext { private final List serviceOrderedReadableRegions; private final List serviceOrderedWriteableRegions; diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/PerPartitionCircuitBreakerInfoHolderTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/PerPartitionCircuitBreakerInfoHolderTest.java new file mode 100644 index 000000000000..87b7193d9ed4 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/PerPartitionCircuitBreakerInfoHolderTest.java @@ -0,0 +1,250 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.perPartitionCircuitBreaker; + +import com.azure.cosmos.implementation.ClientSideRequestStatistics; +import com.azure.cosmos.implementation.CrossRegionAvailabilityContextForRxDocumentServiceRequest; +import com.azure.cosmos.implementation.DiagnosticsClientContext; +import com.azure.cosmos.implementation.GlobalEndpointManager; +import com.azure.cosmos.implementation.OperationType; +import com.azure.cosmos.implementation.PartitionKeyRange; +import com.azure.cosmos.implementation.ResourceType; +import com.azure.cosmos.implementation.RxDocumentServiceRequest; +import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList; +import com.azure.cosmos.implementation.directconnectivity.StoreResponseDiagnostics; +import com.azure.cosmos.implementation.perPartitionAutomaticFailover.PerPartitionAutomaticFailoverInfoHolder; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.mockito.Mockito; +import org.testng.annotations.Test; + +import java.net.URI; +import java.time.Instant; +import java.util.Arrays; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.doReturn; + +public class PerPartitionCircuitBreakerInfoHolderTest { + + @Test(groups = {"unit"}) + public void storesStateReferenceWithoutCopying() { + LocationSpecificHealthContext healthContext = createHealthContext(LocationHealthStatus.Unavailable); + Map currentState = new LinkedHashMap<>(); + currentState.put("eastus", healthContext); + + PerPartitionCircuitBreakerInfoHolder holder = new PerPartitionCircuitBreakerInfoHolder(); + holder.setPerPartitionCircuitBreakerInfoHolder(currentState); + PerPartitionCircuitBreakerInfoHolder snapshot = holder.snapshot(); + + assertThat(holder.getPerPartitionCircuitBreakerInfoHolder()).isSameAs(currentState); + assertThat(snapshot.getPerPartitionCircuitBreakerInfoHolder()) + .isSameAs(currentState); + + currentState.clear(); + assertThat(snapshot.getPerPartitionCircuitBreakerInfoHolder()).isEmpty(); + } + + @Test(groups = {"unit"}) + public void uninitializedSnapshotIsSharedAndIgnoresUpdates() throws Exception { + PerPartitionCircuitBreakerInfoHolder holder = new PerPartitionCircuitBreakerInfoHolder(); + + assertThat(holder.snapshot()).isSameAs(PerPartitionCircuitBreakerInfoHolder.EMPTY); + PerPartitionCircuitBreakerInfoHolder.EMPTY + .setPerPartitionCircuitBreakerInfoHolder(Collections.emptyMap()); + assertThat(new ObjectMapper().writeValueAsString(PerPartitionCircuitBreakerInfoHolder.EMPTY)) + .isEqualTo("null"); + } + + @Test(groups = {"unit"}) + public void initializedEmptyStateIsSerialized() throws Exception { + PerPartitionCircuitBreakerInfoHolder holder = new PerPartitionCircuitBreakerInfoHolder(); + holder.setPerPartitionCircuitBreakerInfoHolder(Collections.emptyMap()); + + ObjectMapper objectMapper = new ObjectMapper(); + + assertThat(objectMapper.writeValueAsString(holder)) + .isEqualTo("{\"stateByRegion\":{}}"); + assertThat(objectMapper.writeValueAsString(PerPartitionCircuitBreakerInfoHolder.EMPTY)) + .isEqualTo("null"); + } + + @Test(groups = {"unit"}) + public void stateIsSerializedUsingCompactFieldNames() throws Exception { + LocationSpecificHealthContext healthContext = new LocationSpecificHealthContext.Builder() + .withLocationHealthStatus(LocationHealthStatus.Unavailable) + .withExceptionCountForReadForCircuitBreaking(1) + .withExceptionCountForWriteForCircuitBreaking(2) + .withSuccessCountForReadForRecovery(3) + .withSuccessCountForWriteForRecovery(4) + .withUnavailableSince(Instant.EPOCH) + .build(); + PerPartitionCircuitBreakerInfoHolder holder = new PerPartitionCircuitBreakerInfoHolder(); + holder.setPerPartitionCircuitBreakerInfoHolder(Collections.singletonMap("eastus", healthContext)); + + assertThat(new ObjectMapper().writeValueAsString(holder)) + .isEqualTo("{\"stateByRegion\":{\"eastus\":{\"st\":\"Unavailable\",\"rErr\":1,\"wErr\":2," + + "\"rOk\":3,\"wOk\":4,\"unavailableSince\":\"1970-01-01T00:00:00Z\"}}}"); + } + + @Test(groups = {"unit"}) + public void failedFailbackAttemptIsSerialized() throws Exception { + LocationSpecificHealthContext healthContext = createHealthContext(LocationHealthStatus.Unavailable) + .withFailbackAttempt( + Instant.EPOCH, + LocationSpecificHealthContext.FailbackOutcome.Failed, + "OPEN_CONNECTION_TASK", + new IllegalStateException("connection failed")); + PerPartitionCircuitBreakerInfoHolder holder = new PerPartitionCircuitBreakerInfoHolder(); + holder.setPerPartitionCircuitBreakerInfoHolder(Collections.singletonMap("eastus", healthContext)); + + assertThat(new ObjectMapper().writeValueAsString(holder)) + .contains("\"failback\":{\"lastAttemptedAt\":\"1970-01-01T00:00:00Z\",\"outcome\":\"Failed\"," + + "\"failure\":{\"stage\":\"OPEN_CONNECTION_TASK\"," + + "\"type\":\"java.lang.IllegalStateException\"}}") + .doesNotContain("connection failed"); + } + + @Test(groups = {"unit"}) + public void latestFailbackMessageIsSerializedByRegion() throws Exception { + PerPartitionCircuitBreakerInfoHolder holder = new PerPartitionCircuitBreakerInfoHolder(); + Map latestMessageByRegion = new LinkedHashMap<>(); + latestMessageByRegion.put("eastus", "first failure"); + latestMessageByRegion.put("westus", "second failure"); + holder.setPerPartitionCircuitBreakerInfoHolder( + Collections.emptyMap(), + latestMessageByRegion); + + assertThat(new ObjectMapper().writeValueAsString(holder)) + .contains("\"latestFailbackMessageByRegion\":{") + .contains("\"eastus\":\"first failure\"") + .contains("\"westus\":\"second failure\""); + } + + @Test(groups = {"unit"}) + public void nonFailedFailbackDoesNotRetainFailureStrings() throws Exception { + LocationSpecificHealthContext healthContext = createHealthContext(LocationHealthStatus.HealthyTentative) + .withFailbackAttempt( + Instant.EPOCH, + LocationSpecificHealthContext.FailbackOutcome.Succeeded, + "SHOULD_NOT_BE_RETAINED", + new IllegalStateException("should not be retained")); + PerPartitionCircuitBreakerInfoHolder holder = new PerPartitionCircuitBreakerInfoHolder(); + holder.setPerPartitionCircuitBreakerInfoHolder(Collections.singletonMap("eastus", healthContext)); + + String serialized = new ObjectMapper().writeValueAsString(holder); + assertThat(serialized).contains("\"outcome\":\"Succeeded\""); + assertThat(serialized).doesNotContain("SHOULD_NOT_BE_RETAINED", "should not be retained", "\"failure\""); + } + + @Test(groups = {"unit"}) + public void responseStatisticsRetainStateAtRecordTime() { + DiagnosticsClientContext diagnosticsClientContext = Mockito.mock(DiagnosticsClientContext.class); + PerPartitionCircuitBreakerInfoHolder holder = new PerPartitionCircuitBreakerInfoHolder(); + holder.setPerPartitionCircuitBreakerInfoHolder(Collections.singletonMap( + "eastus", + createHealthContext(LocationHealthStatus.Unavailable))); + RxDocumentServiceRequest request = createRequest(diagnosticsClientContext, holder); + + ClientSideRequestStatistics statistics = new ClientSideRequestStatistics(diagnosticsClientContext); + statistics.recordResponse(request, null, null); + holder.setPerPartitionCircuitBreakerInfoHolder(Collections.singletonMap( + "westus", + createHealthContext(LocationHealthStatus.Healthy))); + + PerPartitionCircuitBreakerInfoHolder recordedHolder = statistics.getResponseStatisticsList() + .iterator() + .next() + .getPerPartitionCircuitBreakerInfoHolder(); + assertThat(recordedHolder.getPerPartitionCircuitBreakerInfoHolder()).containsOnlyKeys("eastus"); + } + + @Test(groups = {"unit"}) + public void gatewayStatisticsRetainStateAtRecordTime() throws Exception { + DiagnosticsClientContext diagnosticsClientContext = Mockito.mock(DiagnosticsClientContext.class); + PerPartitionCircuitBreakerInfoHolder holder = new PerPartitionCircuitBreakerInfoHolder(); + holder.setPerPartitionCircuitBreakerInfoHolder(Collections.singletonMap( + "eastus", + createHealthContext(LocationHealthStatus.Unavailable))); + RxDocumentServiceRequest request = createRequest(diagnosticsClientContext, holder); + + ClientSideRequestStatistics statistics = new ClientSideRequestStatistics(diagnosticsClientContext); + statistics.recordGatewayResponse(request, Mockito.mock(StoreResponseDiagnostics.class), null); + holder.setPerPartitionCircuitBreakerInfoHolder(Collections.singletonMap( + "westus", + createHealthContext(LocationHealthStatus.Healthy))); + + PerPartitionCircuitBreakerInfoHolder recordedHolder = statistics.getGatewayStatisticsList() + .get(0) + .getPerPartitionCircuitBreakerInfoHolder(); + assertThat(recordedHolder.getPerPartitionCircuitBreakerInfoHolder()).containsOnlyKeys("eastus"); + assertThat(new ObjectMapper().writeValueAsString(statistics)) + .contains("\"ppcb\":{\"stateByRegion\":{\"eastus\":"); + } + + @Test(groups = {"unit"}) + public void routingLookupInitializesEmptyStateWhenNoCircuitExists() throws Exception { + DiagnosticsClientContext diagnosticsClientContext = Mockito.mock(DiagnosticsClientContext.class); + PerPartitionCircuitBreakerInfoHolder holder = new PerPartitionCircuitBreakerInfoHolder(); + RxDocumentServiceRequest request = createRequest(diagnosticsClientContext, holder); + request.setResourceId("collectionRid"); + PartitionKeyRange partitionKeyRange = new PartitionKeyRange("0", "AA", "BB"); + request.requestContext.resolvedPartitionKeyRange = partitionKeyRange; + request.requestContext.resolvedPartitionKeyRangeForCircuitBreaker = partitionKeyRange; + + RegionalRoutingContext eastUs = new RegionalRoutingContext(URI.create("https://eastus.documents.azure.com")); + RegionalRoutingContext westUs = new RegionalRoutingContext(URI.create("https://westus.documents.azure.com")); + GlobalEndpointManager globalEndpointManager = Mockito.mock(GlobalEndpointManager.class); + doReturn(false).when(globalEndpointManager).canUseMultipleWriteLocations(request); + doReturn(UnmodifiableList.unmodifiableList(Arrays.asList(eastUs, westUs))) + .when(globalEndpointManager) + .getApplicableReadRegionalRoutingContexts(Collections.emptyList()); + + GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker manager + = new GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker(globalEndpointManager); + manager.resetCircuitBreakerConfig(PartitionLevelCircuitBreakerConfig.fromJsonString( + "{\"isPartitionLevelCircuitBreakerEnabled\":true," + + "\"consecutiveExceptionCountToleratedForReads\":10," + + "\"consecutiveExceptionCountToleratedForWrites\":5}")); + + assertThat(manager.getUnavailableRegionsForPartitionKeyRange(request, "collectionRid", partitionKeyRange)) + .isEmpty(); + assertThat(holder.getPerPartitionCircuitBreakerInfoHolder()).isEmpty(); + + ClientSideRequestStatistics statistics = new ClientSideRequestStatistics(diagnosticsClientContext); + statistics.recordResponse(request, null, null); + assertThat(new ObjectMapper().writeValueAsString(statistics)) + .contains("\"ppcb\":{\"stateByRegion\":{}}"); + } + + private static RxDocumentServiceRequest createRequest( + DiagnosticsClientContext diagnosticsClientContext, + PerPartitionCircuitBreakerInfoHolder holder) { + + RxDocumentServiceRequest request = RxDocumentServiceRequest.create( + diagnosticsClientContext, + OperationType.Read, + ResourceType.Document); + request.requestContext.setCrossRegionAvailabilityContext( + new CrossRegionAvailabilityContextForRxDocumentServiceRequest( + null, + null, + null, + new AtomicBoolean(false), + holder, + new PerPartitionAutomaticFailoverInfoHolder())); + return request; + } + + private static LocationSpecificHealthContext createHealthContext(LocationHealthStatus healthStatus) { + return new LocationSpecificHealthContext.Builder() + .withLocationHealthStatus(healthStatus) + .withUnavailableSince(Instant.EPOCH) + .build(); + } +} \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/PpcbFailbackLoggingTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/PpcbFailbackLoggingTest.java new file mode 100644 index 000000000000..bb2c526000f9 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/PpcbFailbackLoggingTest.java @@ -0,0 +1,171 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.perPartitionCircuitBreaker; + +import com.azure.cosmos.implementation.GlobalEndpointManager; +import com.azure.cosmos.implementation.OperationType; +import com.azure.cosmos.implementation.PartitionKeyRange; +import com.azure.cosmos.implementation.PartitionKeyRangeWrapper; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; +import org.mockito.Mockito; +import org.slf4j.Logger; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; +import reactor.core.publisher.Flux; +import reactor.core.scheduler.Schedulers; + +import java.net.URI; +import java.time.Duration; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.entry; +import static org.mockito.ArgumentMatchers.contains; +import static org.mockito.ArgumentMatchers.same; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; + +public class PpcbFailbackLoggingTest { + + private static final PartitionKeyRangeWrapper PARTITION = new PartitionKeyRangeWrapper( + new PartitionKeyRange("0", "AA", "BB"), + "collectionRid"); + private static final RegionalRoutingContext REGION = new RegionalRoutingContext( + URI.create("https://contoso-east-us.documents.azure.com")); + private static final RegionalRoutingContext SECOND_REGION = new RegionalRoutingContext( + URI.create("https://contoso-west-us.documents.azure.com")); + + private GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker manager; + private Logger logger; + + @BeforeMethod(groups = {"unit"}) + public void setup() { + GlobalEndpointManager globalEndpointManager = Mockito.mock(GlobalEndpointManager.class); + doReturn("eastus").when(globalEndpointManager).getRegionName( + REGION.getGatewayRegionalEndpoint(), + OperationType.Read); + doReturn("westus").when(globalEndpointManager).getRegionName( + SECOND_REGION.getGatewayRegionalEndpoint(), + OperationType.Read); + this.logger = Mockito.mock(Logger.class); + this.manager = new GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker( + globalEndpointManager, + this.logger); + } + + @Test(groups = {"unit"}) + public void repeatedFailuresAreWarnedAndContainRecoveryIdentity() { + RuntimeException failure = new RuntimeException("connection failed"); + + for (int failureIndex = 0; failureIndex < 10; failureIndex++) { + this.manager.logFailbackFailure(PARTITION, REGION, "OPEN_CONNECTION_TASK", failure); + } + + String expectedFields = "PPCB failback failed: collectionResourceId=collectionRid, " + + "partitionKeyRangeId=0, region=eastus, stage=OPEN_CONNECTION_TASK, " + + "exceptionType=java.lang.RuntimeException, exceptionMessage=connection failed"; + verify(this.logger, times(10)).warn(contains(expectedFields), same(failure)); + assertThat(this.manager.getLatestFailbackMessageByRegion()) + .containsOnly(entry("eastus", "connection failed")); + } + + @Test(groups = {"unit"}) + public void changedFailureReasonIsWarned() { + RuntimeException firstFailure = new RuntimeException("first"); + IllegalStateException changedFailure = new IllegalStateException("changed"); + + this.manager.logFailbackFailure(PARTITION, REGION, "OPEN_CONNECTION_TASK", firstFailure); + this.manager.logFailbackFailure(PARTITION, REGION, "OPEN_CONNECTION_TASK", changedFailure); + + verify(this.logger).warn(contains("exceptionMessage=first"), same(firstFailure)); + verify(this.logger).warn( + contains("exceptionType=java.lang.IllegalStateException, exceptionMessage=changed"), + same(changedFailure)); + } + + @Test(groups = {"unit"}) + public void latestMessageIsRetainedPerRegion() { + this.manager.logFailbackFailure( + PARTITION, + REGION, + "OPEN_CONNECTION_TASK", + new RuntimeException("east-first")); + this.manager.logFailbackFailure( + PARTITION, + SECOND_REGION, + "RECOVERY_PIPELINE", + new RuntimeException("west-latest")); + this.manager.logFailbackFailure( + PARTITION, + REGION, + "OPEN_CONNECTION_TASK", + new RuntimeException("east-latest")); + + assertThat(this.manager.getLatestFailbackMessageByRegion()) + .containsOnly( + entry("eastus", "east-latest"), + entry("westus", "west-latest")); + } + + @Test(groups = {"unit"}) + public void differentStagesAreWarned() { + RuntimeException failure = new RuntimeException("failure"); + + this.manager.logFailbackFailure(PARTITION, REGION, "OPEN_CONNECTION_TASK", failure); + this.manager.logFailbackFailure(PARTITION, REGION, "RECOVERY_PIPELINE", failure); + + verify(this.logger).warn(contains("stage=OPEN_CONNECTION_TASK"), same(failure)); + verify(this.logger).warn(contains("stage=RECOVERY_PIPELINE"), same(failure)); + } + + @Test(groups = {"unit"}) + public void streamFailureWithoutPartitionIdentityIsStillLogged() { + RuntimeException failure = new RuntimeException("stream failed"); + + this.manager.logFailbackFailure(null, null, "RECOVERY_STREAM", failure); + + verify(this.logger).warn( + contains("collectionResourceId=, partitionKeyRangeId=, region=, stage=RECOVERY_STREAM"), + same(failure)); + assertThat(this.manager.getLatestFailbackMessageByRegion()).isEmpty(); + } + + @Test(groups = {"unit"}) + public void failuresForManyPartitionsAreWarned() { + for (int rangeId = 0; rangeId < 100; rangeId++) { + RuntimeException failure = new RuntimeException("failure-" + rangeId); + this.manager.logFailbackFailure( + new PartitionKeyRangeWrapper( + new PartitionKeyRange(String.valueOf(rangeId), "AA", "BB"), + "collectionRid"), + REGION, + "OPEN_CONNECTION_TASK", + failure); + } + + verify(this.logger, times(100)).warn( + contains("exceptionMessage=failure-"), + Mockito.any(RuntimeException.class)); + assertThat(this.manager.getLatestFailbackMessageByRegion()) + .containsOnly(entry("eastus", "failure-99")); + } + + @Test(groups = {"unit"}) + public void concurrentFailuresAreWarned() { + RuntimeException failure = new RuntimeException("failure"); + + Flux.range(0, 100) + .parallel(4) + .runOn(Schedulers.parallel()) + .doOnNext(ignored -> this.manager.logFailbackFailure( + PARTITION, + REGION, + "OPEN_CONNECTION_TASK", + failure)) + .sequential() + .blockLast(Duration.ofSeconds(5)); + + verify(this.logger, times(100)).warn(contains("exceptionMessage=failure"), same(failure)); + } +} \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 0c0115a04fb8..4a094db78ef4 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -18,6 +18,7 @@ * Fixed hedged requests losing request-scoped routing, timeout, authorization, throughput-control, and metadata state when cloning the original request. - See [PR 50069](https://github.com/Azure/azure-sdk-for-java/pull/50069). #### Other Changes +* Added per-region Per-Partition Circuit Breaker health and last failback outcome snapshots to `CosmosDiagnostics`, including structured failure reasons, and WARN logging for failback failures. - See [PR 50158](https://github.com/Azure/azure-sdk-for-java/pull/50158). * Reduced memory footprint of deserialized `PartitionKeyRange` instances by stripping unused fields in the `PartitionKeyRange(ObjectNode)` constructor - See PR [49513](https://github.com/Azure/azure-sdk-for-java/pull/49513). * Added bounded retries for transient "collection routing map / partition key range metadata not available" responses (HTTP 404 with sub-status `0`, `1003`, or `1013`) that can briefly occur right after a container is (re)created, improving the robustness of data-plane operations against the post-creation metadata-propagation race. As part of this change, when the routing map remains unavailable after retries an operation now fails with a `CosmosException` (HTTP 404, sub-status `1024` / `INCORRECT_CONTAINER_RID`) instead of an internal `IllegalStateException`. - See [PR 49639](https://github.com/Azure/azure-sdk-for-java/pull/49639). * Reduced memory footprint and redundant `/pkranges` reads when multiple `CosmosClient` / `CosmosAsyncClient` instances in the same JVM are configured with the same service endpoint. Disable with system property `COSMOS.SHARED_PARTITION_KEY_RANGE_CACHE_ENABLED=false` if needed. - See [PR 49560](https://github.com/Azure/azure-sdk-for-java/pull/49560). diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientSideRequestStatistics.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientSideRequestStatistics.java index fbfaf776edc8..50bd2bfce028 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientSideRequestStatistics.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientSideRequestStatistics.java @@ -12,6 +12,7 @@ import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; import com.fasterxml.jackson.core.JsonGenerator; import com.fasterxml.jackson.databind.SerializerProvider; import com.fasterxml.jackson.databind.annotation.JsonSerialize; @@ -171,8 +172,20 @@ public void recordResponse(RxDocumentServiceRequest request, StoreResultDiagnost this.approximateInsertionCountInBloomFilter = request.requestContext.getApproximateBloomFilterInsertionCount(); storeResponseStatistics.sessionTokenEvaluationResults = request.requestContext.getSessionTokenEvaluationResults(); - storeResponseStatistics.perPartitionCircuitBreakerInfoHolder = request.requestContext.getPerPartitionCircuitBreakerInfoHolder(); - storeResponseStatistics.perPartitionFailoverInfoHolder = request.requestContext.getPerPartitionFailoverContextHolder(); + storeResponseStatistics.perPartitionCircuitBreakerInfoHolder + = request.requestContext.getPerPartitionCircuitBreakerInfoHolder().snapshot(); + storeResponseStatistics.perPartitionAutomaticFailoverInfoHolder = request.requestContext.getPerPartitionFailoverContextHolder(); + + if (request.requestContext.getCrossRegionAvailabilityContext() != null) { + CrossRegionAvailabilityContextForRxDocumentServiceRequest crossRegionAvailabilityContextForRequest + = request.requestContext.getCrossRegionAvailabilityContext(); + + if (crossRegionAvailabilityContextForRequest.shouldAddHubRegionProcessingOnlyHeader()) { + storeResponseStatistics.isHubRegionProcessingOnly = "true"; + } else { + storeResponseStatistics.isHubRegionProcessingOnly = "false"; + } + } if (request.requestContext.getEndToEndOperationLatencyPolicyConfig() != null) { storeResponseStatistics.e2ePolicyCfg = @@ -254,8 +267,24 @@ public void recordGatewayResponse( if (rxDocumentServiceRequest.requestContext != null) { gatewayStatistics.sessionTokenEvaluationResults = rxDocumentServiceRequest.requestContext.getSessionTokenEvaluationResults(); - gatewayStatistics.perPartitionCircuitBreakerInfoHolder = rxDocumentServiceRequest.requestContext.getPerPartitionCircuitBreakerInfoHolder(); - gatewayStatistics.perPartitionFailoverInfoHolder = rxDocumentServiceRequest.requestContext.getPerPartitionFailoverContextHolder(); + gatewayStatistics.perPartitionCircuitBreakerInfoHolder + = rxDocumentServiceRequest.requestContext.getPerPartitionCircuitBreakerInfoHolder().snapshot(); + gatewayStatistics.perPartitionAutomaticFailoverInfoHolder = rxDocumentServiceRequest.requestContext.getPerPartitionFailoverContextHolder(); + gatewayStatistics.isHubRegionProcessingOnly = "false"; + + CrossRegionAvailabilityContextForRxDocumentServiceRequest crossRegionAvailabilityContextForRequest + = rxDocumentServiceRequest.requestContext.getCrossRegionAvailabilityContext(); + + if (crossRegionAvailabilityContextForRequest != null) { + if (crossRegionAvailabilityContextForRequest.shouldAddHubRegionProcessingOnlyHeader()) { + gatewayStatistics.isHubRegionProcessingOnly = "true"; + } + } + + if (rxDocumentServiceRequest.requestContext.getEndToEndOperationLatencyPolicyConfig() != null) { + gatewayStatistics.e2ePolicyCfg = + rxDocumentServiceRequest.requestContext.getEndToEndOperationLatencyPolicyConfig().toString(); + } } } gatewayStatistics.statusCode = storeResponseDiagnostics.getStatusCode(); @@ -698,6 +727,7 @@ public static class StoreResponseStatistics { private Set sessionTokenEvaluationResults; @JsonSerialize(using = PerPartitionCircuitBreakerInfoHolder.PerPartitionCircuitBreakerInfoHolderSerializer.class) + @JsonProperty("ppcb") private PerPartitionCircuitBreakerInfoHolder perPartitionCircuitBreakerInfoHolder; @JsonSerialize(using = PerPartitionFailoverInfoHolder.PerPartitionFailoverInfoHolderSerializer.class) @@ -1025,8 +1055,8 @@ public void serialize(GatewayStatistics gatewayStatistics, } this.writeNonEmptyStringSetField(jsonGenerator, "sessionTokenEvaluationResults", gatewayStatistics.getSessionTokenEvaluationResults()); - this.writeNonNullObjectField(jsonGenerator, "perPartitionCircuitBreakerInfoHolder", gatewayStatistics.getPerPartitionCircuitBreakerInfoHolder()); - this.writeNonNullObjectField(jsonGenerator, "perPartitionFailoverInfoHolder", gatewayStatistics.getPerPartitionFailoverInfoHolder()); + this.writeNonNullObjectField(jsonGenerator, "ppcb", gatewayStatistics.getPerPartitionCircuitBreakerInfoHolder()); + this.writeNonNullObjectField(jsonGenerator, "perPartitionAutomaticFailoverInfoHolder", gatewayStatistics.getPerPartitionFailoverInfoHolder()); this.writeNonNullStringField(jsonGenerator, "requestTCG", gatewayStatistics.getRequestThroughputControlGroupName()); this.writeNonNullStringField(jsonGenerator, "requestTCGConfig", gatewayStatistics.getRequestThroughputControlGroupConfig()); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/ConsecutiveExceptionBasedCircuitBreaker.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/ConsecutiveExceptionBasedCircuitBreaker.java index 6af0848583ab..d43a7a121889 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/ConsecutiveExceptionBasedCircuitBreaker.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/ConsecutiveExceptionBasedCircuitBreaker.java @@ -36,7 +36,7 @@ public LocationSpecificHealthContext handleException( exceptionCountAfterHandling++; int successCountAfterHandling = 0; - LocationSpecificHealthContext.Builder builder = new LocationSpecificHealthContext.Builder() + LocationSpecificHealthContext.Builder builder = new LocationSpecificHealthContext.Builder(locationSpecificHealthContext) .withUnavailableSince(locationSpecificHealthContext.getUnavailableSince()) .withLocationHealthStatus(locationSpecificHealthContext.getLocationHealthStatus()) .withExceptionThresholdBreached(locationSpecificHealthContext.isExceptionThresholdBreached()); @@ -97,7 +97,7 @@ public LocationSpecificHealthContext handleSuccess( exceptionCountAfterHandling = 0; - LocationSpecificHealthContext.Builder builder = new LocationSpecificHealthContext.Builder() + LocationSpecificHealthContext.Builder builder = new LocationSpecificHealthContext.Builder(locationSpecificHealthContext) .withUnavailableSince(locationSpecificHealthContext.getUnavailableSince()) .withLocationHealthStatus(locationSpecificHealthContext.getLocationHealthStatus()) .withExceptionThresholdBreached(locationSpecificHealthContext.isExceptionThresholdBreached()); @@ -124,7 +124,7 @@ public LocationSpecificHealthContext handleSuccess( successCountAfterHandling++; - builder = new LocationSpecificHealthContext.Builder() + builder = new LocationSpecificHealthContext.Builder(locationSpecificHealthContext) .withUnavailableSince(locationSpecificHealthContext.getUnavailableSince()) .withLocationHealthStatus(locationSpecificHealthContext.getLocationHealthStatus()) .withExceptionThresholdBreached(locationSpecificHealthContext.isExceptionThresholdBreached()); 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 3eacc111e522..8a9525c20970 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 @@ -34,6 +34,7 @@ import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.PriorityQueue; @@ -58,13 +59,22 @@ public class GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker impleme private final ConcurrentHashMap regionalRoutingContextToRegion; private final AtomicBoolean isClosed = new AtomicBoolean(false); private final AtomicBoolean isPartitionRecoveryTaskRunning = new AtomicBoolean(false); - private final Scheduler partitionRecoveryScheduler = Schedulers.newSingle( - "partition-availability-staleness-check", - true); + private final AtomicReference partitionRecoveryDisposable = new AtomicReference<>(); + private final Logger failbackLogger; + private final Object latestFailbackMessageByRegionLock = new Object(); + private volatile Map latestFailbackMessageByRegion = Collections.emptyMap(); public GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker(GlobalEndpointManager globalEndpointManager) { + this(globalEndpointManager, logger); + } + + GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker( + GlobalEndpointManager globalEndpointManager, + Logger failbackLogger) { + this.partitionKeyRangeToLocationSpecificUnavailabilityInfo = new ConcurrentHashMap<>(); this.globalEndpointManager = globalEndpointManager; + this.failbackLogger = checkNotNull(failbackLogger, "Argument 'failbackLogger' cannot be null!"); PartitionLevelCircuitBreakerConfig partitionLevelCircuitBreakerConfig = Configs.getPartitionLevelCircuitBreakerConfig(); this.consecutiveExceptionBasedCircuitBreaker = new ConsecutiveExceptionBasedCircuitBreaker(partitionLevelCircuitBreakerConfig); @@ -148,7 +158,7 @@ public void handleLocationExceptionForPartitionKeyRange( partitionLevelLocationUnavailabilityInfoAsVal.areLocationsAvailableForPartitionKeyRange(applicableRegionalRoutingContexts)); } - request.requestContext.setPerPartitionCircuitBreakerInfoHolder(partitionLevelLocationUnavailabilityInfoAsVal.regionToLocationSpecificHealthContext); + this.publishSnapshot(request, partitionLevelLocationUnavailabilityInfoAsVal); return partitionLevelLocationUnavailabilityInfoAsVal; }); @@ -209,7 +219,7 @@ public void handleLocationSuccessForPartitionKeyRange(RxDocumentServiceRequest r succeededRegionalRoutingContext, request.isReadOnlyRequest()); - request.requestContext.setPerPartitionCircuitBreakerInfoHolder(partitionKeyRangeToFailoverInfoAsVal.regionToLocationSpecificHealthContext); + this.publishSnapshot(request, partitionKeyRangeToFailoverInfoAsVal); return partitionKeyRangeToFailoverInfoAsVal; }); } catch (Exception e) { @@ -236,6 +246,7 @@ public List getUnavailableRegionsForPartitionKeyRange( this.partitionKeyRangeToLocationSpecificUnavailabilityInfo.get(partitionKeyRangeWrapper); List unavailableRegions = new ArrayList<>(); + this.publishSnapshot(request, partitionLevelLocationUnavailabilityInfoSnapshot); if (partitionLevelLocationUnavailabilityInfoSnapshot != null) { Map locationEndpointToFailureMetricsForPartition = @@ -278,10 +289,21 @@ public List getUnavailableRegionsForPartitionKeyRange( } } + private void publishSnapshot( + RxDocumentServiceRequest request, + PartitionLevelLocationUnavailabilityInfo info) { + + request.requestContext.getPerPartitionCircuitBreakerInfoHolder() + .setPerPartitionCircuitBreakerInfoHolder( + info == null ? Collections.emptyMap() : info.regionToLocationSpecificHealthContext, + this.latestFailbackMessageByRegion); + } + private Flux updateStaleLocationInfo() { return Mono.just(1) .delayElement(Duration.ofSeconds(Configs.getStalePartitionUnavailabilityRefreshIntervalInSeconds())) .repeat(() -> !this.isClosed.get()) + .doOnNext(ignore -> this.clearLatestFailbackMessagesIfNoBacklog()) .flatMap(ignore -> Flux.fromIterable(this.partitionKeyRangeToLocationSpecificUnavailabilityInfo.entrySet()), 1, 1) .flatMap(partitionKeyRangeWrapperToPartitionKeyRangeWrapperPair -> { @@ -320,19 +342,27 @@ private Flux updateStaleLocationInfo() { return Mono.empty(); } } catch (Exception e) { - logger.warn("An exception was thrown trying to recover an Unavailable partitionKeyRange!", e); + this.logFailbackFailure( + partitionKeyRangeWrapperToPartitionKeyRangeWrapperPair.getKey(), + null, + "SCAN_UNAVAILABLE_PARTITIONS", + e); return Flux.empty(); } }, 1, 1) .flatMap(locationToLocationSpecificHealthContextPair -> { - try { - PartitionKeyRangeWrapper partitionKeyRangeWrapper = locationToLocationSpecificHealthContextPair.getLeft(); - RegionalRoutingContext locationWithStaleUnavailabilityInfo = locationToLocationSpecificHealthContextPair.getRight().getLeft(); - - PartitionLevelLocationUnavailabilityInfo partitionLevelLocationUnavailabilityInfo = this.partitionKeyRangeToLocationSpecificUnavailabilityInfo.get(partitionKeyRangeWrapper); + PartitionKeyRangeWrapper partitionKeyRangeWrapper = locationToLocationSpecificHealthContextPair.getLeft(); + RegionalRoutingContext locationWithStaleUnavailabilityInfo = locationToLocationSpecificHealthContextPair.getRight().getLeft(); + PartitionLevelLocationUnavailabilityInfo partitionLevelLocationUnavailabilityInfo + = this.partitionKeyRangeToLocationSpecificUnavailabilityInfo.get(partitionKeyRangeWrapper); + Instant failbackAttemptTime = Instant.now(); + try { if (partitionLevelLocationUnavailabilityInfo != null) { + partitionLevelLocationUnavailabilityInfo.recordFailbackAttempt( + locationWithStaleUnavailabilityInfo, + failbackAttemptTime); GlobalAddressResolver globalAddressResolver = this.globalAddressResolverSnapshot.get(); @@ -364,51 +394,65 @@ private Flux updateStaleLocationInfo() { + partitionKeyRangeWrapper.getCollectionResourceId() + " has succeeded..."); - partitionLevelLocationUnavailabilityInfo.locationEndpointToLocationSpecificContextForPartition.compute(locationWithStaleUnavailabilityInfo, (locationWithStaleUnavailabilityInfoAsKey, locationSpecificContextAsVal) -> { - - if (locationSpecificContextAsVal != null) { - locationSpecificContextAsVal = GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker - .this.locationSpecificHealthContextTransitionHandler.handleSuccess( - locationSpecificContextAsVal, - partitionKeyRangeWrapper, - this.regionalRoutingContextToRegion.getOrDefault(locationWithStaleUnavailabilityInfoAsKey, StringUtils.EMPTY), - false, - true); - } - - return locationSpecificContextAsVal; - }); + partitionLevelLocationUnavailabilityInfo.recordFailbackSuccess( + partitionKeyRangeWrapper, + locationWithStaleUnavailabilityInfo, + failbackAttemptTime); }) .onErrorResume(throwable -> { - logger.debug("An exception was thrown trying to recover an Unavailable partition key range!", throwable); + partitionLevelLocationUnavailabilityInfo.recordFailbackFailure( + locationWithStaleUnavailabilityInfo, + failbackAttemptTime, + "OPEN_CONNECTION_TASK", + throwable); + this.logFailbackFailure( + partitionKeyRangeWrapper, + locationWithStaleUnavailabilityInfo, + "OPEN_CONNECTION_TASK", + throwable); return Mono.empty(); }); + } else { + IllegalStateException failure + = new IllegalStateException("GatewayAddressCache is not available."); + partitionLevelLocationUnavailabilityInfo.recordFailbackFailure( + locationWithStaleUnavailabilityInfo, + failbackAttemptTime, + "RESOLVE_GATEWAY_ADDRESS_CACHE", + failure); + this.logFailbackFailure( + partitionKeyRangeWrapper, + locationWithStaleUnavailabilityInfo, + "RESOLVE_GATEWAY_ADDRESS_CACHE", + failure); } } else { - partitionLevelLocationUnavailabilityInfo.locationEndpointToLocationSpecificContextForPartition.compute(locationWithStaleUnavailabilityInfo, (locationWithStaleUnavailabilityInfoAsKey, locationSpecificContextAsVal) -> { - - if (locationSpecificContextAsVal != null) { - locationSpecificContextAsVal = GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker - .this.locationSpecificHealthContextTransitionHandler.handleSuccess( - locationSpecificContextAsVal, - partitionKeyRangeWrapper, - this.regionalRoutingContextToRegion.getOrDefault(locationWithStaleUnavailabilityInfoAsKey, StringUtils.EMPTY), - false, - true); - } - return locationSpecificContextAsVal; - }); + partitionLevelLocationUnavailabilityInfo.recordFailbackSuccess( + partitionKeyRangeWrapper, + locationWithStaleUnavailabilityInfo, + failbackAttemptTime); } } } catch (Exception e) { - logger.debug("An exception was thrown trying to recover an Unavailable partition key range!", e); + if (partitionLevelLocationUnavailabilityInfo != null) { + partitionLevelLocationUnavailabilityInfo.recordFailbackFailure( + locationWithStaleUnavailabilityInfo, + failbackAttemptTime, + "RECOVERY_PIPELINE", + e); + } + this.logFailbackFailure( + partitionKeyRangeWrapper, + locationWithStaleUnavailabilityInfo, + "RECOVERY_PIPELINE", + e); return Flux.empty(); } return Flux.empty(); }, 1, 1) .onErrorResume(throwable -> { - logger.warn("An exception : was thrown trying to recover an Unavailable partitionKeyRange!, fail-back flow won't be executed!", throwable); + this.logFailbackFailure(null, null, "RECOVERY_STREAM", throwable); return Flux.empty(); }); } @@ -423,6 +467,93 @@ private boolean shouldForceRefreshAddresses(Throwable throwable) { && WebExceptionUtility.isNetworkFailure((Exception) throwable); } + void logFailbackFailure( + PartitionKeyRangeWrapper partitionKeyRangeWrapper, + RegionalRoutingContext regionalRoutingContext, + String stage, + Throwable throwable) { + + String collectionResourceId = partitionKeyRangeWrapper == null + ? StringUtils.EMPTY + : partitionKeyRangeWrapper.getCollectionResourceId(); + String partitionKeyRangeId = partitionKeyRangeWrapper == null + || partitionKeyRangeWrapper.getPartitionKeyRange() == null + ? StringUtils.EMPTY + : partitionKeyRangeWrapper.getPartitionKeyRange().getId(); + String exceptionType = throwable == null + ? StringUtils.EMPTY + : throwable.getClass().getName(); + String exceptionMessage = throwable == null || throwable.getMessage() == null + ? StringUtils.EMPTY + : throwable.getMessage(); + String region = this.resolveRegionName(regionalRoutingContext); + String message = "PPCB failback failed: collectionResourceId=" + + collectionResourceId + + ", partitionKeyRangeId=" + + partitionKeyRangeId + + ", region=" + + region + + ", stage=" + + stage + + ", exceptionType=" + + exceptionType + + ", exceptionMessage=" + + exceptionMessage; + + if (!StringUtils.isEmpty(region)) { + this.recordLatestFailbackMessage(region, exceptionMessage); + } + this.failbackLogger.warn(message, throwable); + } + + private void recordLatestFailbackMessage( + String region, + String failureMessage) { + + synchronized (this.latestFailbackMessageByRegionLock) { + Map updatedMessages = new LinkedHashMap<>(this.latestFailbackMessageByRegion); + updatedMessages.put(region, failureMessage); + this.latestFailbackMessageByRegion = Collections.unmodifiableMap(updatedMessages); + } + } + + private void clearLatestFailbackMessagesIfNoBacklog() { + synchronized (this.latestFailbackMessageByRegionLock) { + for (PartitionLevelLocationUnavailabilityInfo info + : this.partitionKeyRangeToLocationSpecificUnavailabilityInfo.values()) { + + for (LocationSpecificHealthContext healthContext + : info.locationEndpointToLocationSpecificContextForPartition.values()) { + + if (!healthContext.isRegionAvailableToProcessRequests()) { + return; + } + } + } + + this.latestFailbackMessageByRegion = Collections.emptyMap(); + } + } + + Map getLatestFailbackMessageByRegion() { + return this.latestFailbackMessageByRegion; + } + + private String resolveRegionName(RegionalRoutingContext regionalRoutingContext) { + if (regionalRoutingContext == null) { + return StringUtils.EMPTY; + } + + String region = this.regionalRoutingContextToRegion.get(regionalRoutingContext); + if (!StringUtils.isEmpty(region)) { + return region; + } + + return this.globalEndpointManager.getRegionName( + regionalRoutingContext.getGatewayRegionalEndpoint(), + OperationType.Read); + } + public boolean isPerPartitionLevelCircuitBreakingApplicable(RxDocumentServiceRequest request) { if (!this.consecutiveExceptionBasedCircuitBreaker.isPartitionLevelCircuitBreakerEnabled()) { @@ -582,6 +713,89 @@ private void handleSuccess( }); } + private void recordFailbackAttempt( + RegionalRoutingContext regionalRoutingContext, + Instant attemptTime) { + + this.updateFailbackDiagnostics( + regionalRoutingContext, + attemptTime, + LocationSpecificHealthContext.FailbackOutcome.Attempting, + null, + null); + } + + private void recordFailbackSuccess( + PartitionKeyRangeWrapper partitionKeyRangeWrapper, + RegionalRoutingContext regionalRoutingContext, + Instant attemptTime) { + + this.locationEndpointToLocationSpecificContextForPartition.computeIfPresent( + regionalRoutingContext, + (routingContext, healthContext) -> { + LocationSpecificHealthContext updatedContext + = this.locationSpecificHealthContextTransitionHandler.handleSuccess( + healthContext, + partitionKeyRangeWrapper, + GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker.this + .regionalRoutingContextToRegion.getOrDefault(routingContext, StringUtils.EMPTY), + false, + true) + .withFailbackAttempt( + attemptTime, + LocationSpecificHealthContext.FailbackOutcome.Succeeded, + null, + null); + this.updateRegionDiagnostics(routingContext, updatedContext); + return updatedContext; + }); + GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker.this + .clearLatestFailbackMessagesIfNoBacklog(); + } + + private void recordFailbackFailure( + RegionalRoutingContext regionalRoutingContext, + Instant attemptTime, + String failureStage, + Throwable failure) { + + this.updateFailbackDiagnostics( + regionalRoutingContext, + attemptTime, + LocationSpecificHealthContext.FailbackOutcome.Failed, + failureStage, + failure); + } + + private void updateFailbackDiagnostics( + RegionalRoutingContext regionalRoutingContext, + Instant attemptTime, + LocationSpecificHealthContext.FailbackOutcome outcome, + String failureStage, + Throwable failure) { + + this.locationEndpointToLocationSpecificContextForPartition.computeIfPresent( + regionalRoutingContext, + (routingContext, healthContext) -> { + LocationSpecificHealthContext updatedContext = healthContext.withFailbackAttempt( + attemptTime, + outcome, + failureStage, + failure); + this.updateRegionDiagnostics(routingContext, updatedContext); + return updatedContext; + }); + } + + private void updateRegionDiagnostics( + RegionalRoutingContext regionalRoutingContext, + LocationSpecificHealthContext healthContext) { + + String region = GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker.this + .regionalRoutingContextToRegion.getOrDefault(regionalRoutingContext, StringUtils.EMPTY); + this.regionToLocationSpecificHealthContext.put(region, healthContext); + } + public boolean areLocationsAvailableForPartitionKeyRange(List availableLocationsAtAccountLevel) { for (RegionalRoutingContext availableLocation : availableLocationsAtAccountLevel) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/LocationSpecificHealthContext.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/LocationSpecificHealthContext.java index 2031f4d3e270..46a196a57937 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/LocationSpecificHealthContext.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/LocationSpecificHealthContext.java @@ -24,6 +24,7 @@ public class LocationSpecificHealthContext implements Serializable { private final Instant unavailableSince; private final LocationHealthStatus locationHealthStatus; private final boolean isExceptionThresholdBreached; + private final FailbackDiagnostics failbackDiagnostics; LocationSpecificHealthContext( int successCountForWriteForRecovery, @@ -32,7 +33,8 @@ public class LocationSpecificHealthContext implements Serializable { int exceptionCountForReadForCircuitBreaking, Instant unavailableSince, LocationHealthStatus locationHealthStatus, - boolean isExceptionThresholdBreached) { + boolean isExceptionThresholdBreached, + FailbackDiagnostics failbackDiagnostics) { this.successCountForWriteForRecovery = successCountForWriteForRecovery; this.exceptionCountForWriteForCircuitBreaking = exceptionCountForWriteForCircuitBreaking; @@ -41,6 +43,7 @@ public class LocationSpecificHealthContext implements Serializable { this.unavailableSince = unavailableSince; this.locationHealthStatus = locationHealthStatus; this.isExceptionThresholdBreached = isExceptionThresholdBreached; + this.failbackDiagnostics = failbackDiagnostics; } public boolean isExceptionThresholdBreached() { @@ -77,6 +80,55 @@ public LocationHealthStatus getLocationHealthStatus() { return this.locationHealthStatus; } + public Instant getLastFailbackAttemptTime() { + return this.failbackDiagnostics == null ? null : this.failbackDiagnostics.lastAttemptedAt; + } + + public FailbackOutcome getLastFailbackOutcome() { + return this.failbackDiagnostics == null ? null : this.failbackDiagnostics.outcome; + } + + LocationSpecificHealthContext withFailbackAttempt( + Instant attemptTime, + FailbackOutcome outcome, + String failureStage, + Throwable failure) { + + boolean failed = outcome == FailbackOutcome.Failed; + return new Builder(this) + .withFailbackDiagnostics(new FailbackDiagnostics( + attemptTime, + outcome, + failed ? failureStage : null, + failed && failure != null ? failure.getClass().getName() : null)) + .build(); + } + + public enum FailbackOutcome { + Attempting, + Succeeded, + Failed + } + + private static class FailbackDiagnostics { + private final Instant lastAttemptedAt; + private final FailbackOutcome outcome; + private final String failureStage; + private final String failureType; + + private FailbackDiagnostics( + Instant lastAttemptedAt, + FailbackOutcome outcome, + String failureStage, + String failureType) { + + this.lastAttemptedAt = lastAttemptedAt; + this.outcome = outcome; + this.failureStage = failureStage; + this.failureType = failureType; + } + } + static class Builder { private int exceptionCountForWriteForCircuitBreaking; @@ -86,9 +138,21 @@ static class Builder { private Instant unavailableSince; private LocationHealthStatus locationHealthStatus; private boolean isExceptionThresholdBreached; + private FailbackDiagnostics failbackDiagnostics; public Builder() {} + Builder(LocationSpecificHealthContext source) { + this.exceptionCountForWriteForCircuitBreaking = source.exceptionCountForWriteForCircuitBreaking; + this.successCountForWriteForRecovery = source.successCountForWriteForRecovery; + this.exceptionCountForReadForCircuitBreaking = source.exceptionCountForReadForCircuitBreaking; + this.successCountForReadForRecovery = source.successCountForReadForRecovery; + this.unavailableSince = source.unavailableSince; + this.locationHealthStatus = source.locationHealthStatus; + this.isExceptionThresholdBreached = source.isExceptionThresholdBreached; + this.failbackDiagnostics = source.failbackDiagnostics; + } + public Builder withExceptionCountForWriteForCircuitBreaking(int exceptionCountForWriteForCircuitBreaking) { this.exceptionCountForWriteForCircuitBreaking = exceptionCountForWriteForCircuitBreaking; return this; @@ -124,6 +188,11 @@ public Builder withExceptionThresholdBreached(boolean exceptionThresholdBreached return this; } + Builder withFailbackDiagnostics(FailbackDiagnostics failbackDiagnostics) { + this.failbackDiagnostics = failbackDiagnostics; + return this; + } + public LocationSpecificHealthContext build() { return new LocationSpecificHealthContext( @@ -133,7 +202,8 @@ public LocationSpecificHealthContext build() { this.exceptionCountForReadForCircuitBreaking, this.unavailableSince, this.locationHealthStatus, - this.isExceptionThresholdBreached); + this.isExceptionThresholdBreached, + this.failbackDiagnostics); } } @@ -143,13 +213,26 @@ static class LocationSpecificHealthContextSerializer extends com.fasterxml.jacks public void serialize(LocationSpecificHealthContext value, JsonGenerator gen, SerializerProvider provider) throws IOException { gen.writeStartObject(); - gen.writeNumberField("exceptionCountForWriteForCircuitBreaking", value.exceptionCountForWriteForCircuitBreaking); - gen.writeNumberField("exceptionCountForReadForCircuitBreaking", value.exceptionCountForReadForCircuitBreaking); - gen.writeNumberField("successCountForWriteForRecovery", value.successCountForWriteForRecovery); - gen.writeNumberField("successCountForReadForRecovery", value.successCountForReadForRecovery); - gen.writePOJOField("locationHealthStatus", value.locationHealthStatus); + gen.writePOJOField("st", value.locationHealthStatus); + gen.writeNumberField("rErr", value.exceptionCountForReadForCircuitBreaking); + gen.writeNumberField("wErr", value.exceptionCountForWriteForCircuitBreaking); + gen.writeNumberField("rOk", value.successCountForReadForRecovery); + gen.writeNumberField("wOk", value.successCountForWriteForRecovery); gen.writeStringField("unavailableSince", toInstantString(value.unavailableSince)); + if (value.failbackDiagnostics != null) { + gen.writeObjectFieldStart("failback"); + gen.writeStringField("lastAttemptedAt", toInstantString(value.failbackDiagnostics.lastAttemptedAt)); + gen.writePOJOField("outcome", value.failbackDiagnostics.outcome); + if (value.failbackDiagnostics.outcome == FailbackOutcome.Failed) { + gen.writeObjectFieldStart("failure"); + gen.writeStringField("stage", value.failbackDiagnostics.failureStage); + gen.writeStringField("type", value.failbackDiagnostics.failureType); + gen.writeEndObject(); + } + gen.writeEndObject(); + } + gen.writeEndObject(); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/LocationSpecificHealthContextTransitionHandler.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/LocationSpecificHealthContextTransitionHandler.java index 174cf4eda822..c75c34684d84 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/LocationSpecificHealthContextTransitionHandler.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/LocationSpecificHealthContextTransitionHandler.java @@ -67,7 +67,10 @@ public LocationSpecificHealthContext handleSuccess( partitionKeyRangeWrapper.getCollectionResourceId() + " marked as Healthy from HealthyTentative for region : " + regionWithSuccess); - return this.transitionHealthStatus(LocationHealthStatus.Healthy, isReadOnlyRequest); + return this.transitionHealthStatus( + LocationHealthStatus.Healthy, + isReadOnlyRequest, + locationSpecificHealthContextInner); } else { return locationSpecificHealthContextInner; } @@ -83,11 +86,17 @@ public LocationSpecificHealthContext handleSuccess( partitionKeyRangeWrapper.getCollectionResourceId() + " marked as HealthyTentative from Unavailable for region :" + regionWithSuccess); - return this.transitionHealthStatus(LocationHealthStatus.HealthyTentative, isReadOnlyRequest); + return this.transitionHealthStatus( + LocationHealthStatus.HealthyTentative, + isReadOnlyRequest, + locationSpecificHealthContext); } } else { logger.debug("PartitionKeyRange " + partitionKeyRangeWrapper.getPartitionKeyRange() + " and collectionResourceId : " + partitionKeyRangeWrapper.getCollectionResourceId() + " marked as HealthyTentative from Unavailable for region : " + regionWithSuccess);; - return this.transitionHealthStatus(LocationHealthStatus.HealthyTentative, isReadOnlyRequest); + return this.transitionHealthStatus( + LocationHealthStatus.HealthyTentative, + isReadOnlyRequest, + locationSpecificHealthContext); } break; default: @@ -108,7 +117,10 @@ public LocationSpecificHealthContext handleException( switch (currentLocationHealthStatusSnapshot) { case Healthy: logger.debug("PartitionKeyRange " + partitionKeyRangeWrapper.getPartitionKeyRange() + " of collectionResourceId : " + partitionKeyRangeWrapper.getCollectionResourceId() + " marked as HealthyWithFailures from Healthy for region : " + regionWithException); - return this.transitionHealthStatus(LocationHealthStatus.HealthyWithFailures, isReadOnlyRequest); + return this.transitionHealthStatus( + LocationHealthStatus.HealthyWithFailures, + isReadOnlyRequest, + locationSpecificHealthContext); case HealthyWithFailures: if (!this.consecutiveExceptionBasedCircuitBreaker.shouldHealthStatusBeDowngraded(locationSpecificHealthContext, isReadOnlyRequest)) { @@ -138,7 +150,10 @@ public LocationSpecificHealthContext handleException( partitionKeyRangeWrapper.getCollectionResourceId() + " marked as Unavailable from HealthyWithFailures for region : " + regionWithException); - return this.transitionHealthStatus(LocationHealthStatus.Unavailable, isReadOnlyRequest); + return this.transitionHealthStatus( + LocationHealthStatus.Unavailable, + isReadOnlyRequest, + locationSpecificHealthContext); } case HealthyTentative: if (!this.consecutiveExceptionBasedCircuitBreaker.shouldHealthStatusBeDowngraded(locationSpecificHealthContext, isReadOnlyRequest)) { @@ -155,7 +170,10 @@ public LocationSpecificHealthContext handleException( partitionKeyRangeWrapper.getCollectionResourceId() + " marked as Unavailable from HealthyTentative for region : " + regionWithException); - return this.transitionHealthStatus(LocationHealthStatus.Unavailable, isReadOnlyRequest); + return this.transitionHealthStatus( + LocationHealthStatus.Unavailable, + isReadOnlyRequest, + locationSpecificHealthContext); } case Unavailable: return this.consecutiveExceptionBasedCircuitBreaker @@ -173,7 +191,19 @@ public LocationSpecificHealthContext transitionHealthStatus( LocationHealthStatus newStatus, boolean isReadOnlyRequest) { - LocationSpecificHealthContext.Builder builder = new LocationSpecificHealthContext.Builder() + return this.transitionHealthStatus(newStatus, isReadOnlyRequest, null); + } + + private LocationSpecificHealthContext transitionHealthStatus( + LocationHealthStatus newStatus, + boolean isReadOnlyRequest, + LocationSpecificHealthContext previousContext) { + + LocationSpecificHealthContext.Builder builder = previousContext == null + ? new LocationSpecificHealthContext.Builder() + : new LocationSpecificHealthContext.Builder(previousContext); + + builder .withSuccessCountForWriteForRecovery(0) .withExceptionCountForWriteForCircuitBreaking(0) .withSuccessCountForReadForRecovery(0) @@ -216,7 +246,7 @@ public LocationSpecificHealthContext transitionHealthStatus( case HealthyTentative: return builder - .withUnavailableSince(Instant.now()) + .withUnavailableSince(Instant.MAX) .withLocationHealthStatus(LocationHealthStatus.HealthyTentative) .withExceptionThresholdBreached(false) .build(); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/PerPartitionCircuitBreakerInfoHolder.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/PerPartitionCircuitBreakerInfoHolder.java index 90b1d804ffe0..46890bdd314f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/PerPartitionCircuitBreakerInfoHolder.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/PerPartitionCircuitBreakerInfoHolder.java @@ -3,24 +3,64 @@ package com.azure.cosmos.implementation.perPartitionCircuitBreaker; -import com.azure.cosmos.implementation.Utils; import com.fasterxml.jackson.core.JsonGenerator; import com.fasterxml.jackson.databind.SerializerProvider; +import com.fasterxml.jackson.databind.annotation.JsonSerialize; import java.io.IOException; import java.io.Serializable; +import java.util.Collections; import java.util.Map; +@JsonSerialize(using = PerPartitionCircuitBreakerInfoHolder.PerPartitionCircuitBreakerInfoHolderSerializer.class) public class PerPartitionCircuitBreakerInfoHolder implements Serializable { - private final Utils.ValueHolder> perPartitionCircuitBreakerInfoHolder = new Utils.ValueHolder>(); + public static final PerPartitionCircuitBreakerInfoHolder EMPTY = new PerPartitionCircuitBreakerInfoHolder(); - public synchronized void setPerPartitionCircuitBreakerInfoHolder(final Map locationSpecificHealthContext) { - this.perPartitionCircuitBreakerInfoHolder.v = locationSpecificHealthContext; + private volatile Map perPartitionCircuitBreakerInfoHolder; + private volatile Map latestFailbackMessageByRegion = Collections.emptyMap(); + + public PerPartitionCircuitBreakerInfoHolder() { + } + + private PerPartitionCircuitBreakerInfoHolder( + Map perPartitionCircuitBreakerInfoHolder, + Map latestFailbackMessageByRegion) { + + this.perPartitionCircuitBreakerInfoHolder = perPartitionCircuitBreakerInfoHolder; + this.latestFailbackMessageByRegion = latestFailbackMessageByRegion; } - public synchronized Map getPerPartitionCircuitBreakerInfoHolder() { - return perPartitionCircuitBreakerInfoHolder.v; + public void setPerPartitionCircuitBreakerInfoHolder(final Map locationSpecificHealthContext) { + this.setPerPartitionCircuitBreakerInfoHolder(locationSpecificHealthContext, this.latestFailbackMessageByRegion); + } + + void setPerPartitionCircuitBreakerInfoHolder( + Map locationSpecificHealthContext, + Map latestFailbackMessageByRegion) { + + if (this == EMPTY) { + return; + } + + this.perPartitionCircuitBreakerInfoHolder = locationSpecificHealthContext == null + ? Collections.emptyMap() + : locationSpecificHealthContext; + this.latestFailbackMessageByRegion = latestFailbackMessageByRegion == null + ? Collections.emptyMap() + : latestFailbackMessageByRegion; + } + + public Map getPerPartitionCircuitBreakerInfoHolder() { + return this.perPartitionCircuitBreakerInfoHolder; + } + + public PerPartitionCircuitBreakerInfoHolder snapshot() { + Map snapshot = this.perPartitionCircuitBreakerInfoHolder; + + return snapshot == null + ? EMPTY + : new PerPartitionCircuitBreakerInfoHolder(snapshot, this.latestFailbackMessageByRegion); } public static class PerPartitionCircuitBreakerInfoHolderSerializer extends com.fasterxml.jackson.databind.JsonSerializer { @@ -30,10 +70,14 @@ public void serialize(PerPartitionCircuitBreakerInfoHolder value, JsonGenerator Map locationToLocationSpecificHealthContext = value.getPerPartitionCircuitBreakerInfoHolder(); - if (locationToLocationSpecificHealthContext != null && !locationToLocationSpecificHealthContext.isEmpty()) { + if (locationToLocationSpecificHealthContext != null) { gen.writeStartObject(); - gen.writePOJOField("locSpecificHealthCtx", locationToLocationSpecificHealthContext); + gen.writePOJOField("stateByRegion", locationToLocationSpecificHealthContext); + + if (!value.latestFailbackMessageByRegion.isEmpty()) { + gen.writePOJOField("latestFailbackMessageByRegion", value.latestFailbackMessageByRegion); + } gen.writeEndObject(); } From 0342704328f8ac77e2ce1c6451306b9237d8c55b Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Mon, 6 Jul 2026 11:17:13 -0400 Subject: [PATCH 03/14] Fix partitionLevelCircuitBreakerCfg missing from CosmosDiagnostics clientCfgs The partitionLevelCircuitBreakerCfg field disappeared from the clientCfgs section of CosmosDiagnostics when a customer explicitly enabled Per-Partition Circuit Breaker (PPCB) client-side. Root cause: the diagnostics write was coupled to the Per-Partition Automatic Failover (PPAF) initialization path, so the field was only populated when the service mandated PPAF, not when PPCB was configured client-side. Fix: move the diagnostics write into initializePerPartitionCircuitBreaker(), which is invoked unconditionally at client init, so the field appears whenever the circuit breaker is configured client-side. Adds a CI-runnable regression test asserting all clientCfgs keys are present, including partitionLevelCircuitBreakerCfg. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../RxDocumentClientImplTest.java | 122 ++++++++++++++++++ .../implementation/RxDocumentClientImpl.java | 7 +- 2 files changed, 128 insertions(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxDocumentClientImplTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxDocumentClientImplTest.java index cbc9301142f4..242ada8ab223 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxDocumentClientImplTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxDocumentClientImplTest.java @@ -5,6 +5,7 @@ import com.azure.core.credential.AzureKeyCredential; import com.azure.core.http.ProxyOptions; import com.azure.cosmos.BridgeInternal; +import com.azure.cosmos.ConnectionMode; import com.azure.cosmos.ConsistencyLevel; import com.azure.cosmos.CosmosContainerProactiveInitConfig; import com.azure.cosmos.CosmosDiagnostics; @@ -34,6 +35,11 @@ import com.azure.cosmos.models.ModelBridgeInternal; import com.azure.cosmos.models.PartitionKey; import com.azure.cosmos.models.PartitionKeyDefinition; +import com.fasterxml.jackson.core.JsonFactory; +import com.fasterxml.jackson.core.JsonGenerator; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.SerializerProvider; +import com.fasterxml.jackson.databind.node.ObjectNode; import io.netty.buffer.ByteBufInputStream; import io.netty.buffer.Unpooled; import io.netty.handler.codec.http.HttpResponseStatus; @@ -46,6 +52,8 @@ import reactor.test.StepVerifier; import java.net.URI; +import java.io.StringWriter; +import java.lang.reflect.Method; import java.nio.charset.StandardCharsets; import java.time.Duration; import java.util.ArrayList; @@ -317,6 +325,120 @@ public void readMany() { } } + // Regression test for the "partitionLevelCircuitBreakerCfg" diagnostics field silently disappearing from the + // CosmosDiagnostics "clientCfgs" section. Prior to the fix, the field was only written on the PPAF + // (service-mandated) path, so a client that explicitly enabled Per-Partition Circuit Breaker never surfaced it. + // This test constructs a real RxDocumentClientImpl (exercising the actual constructor wiring), drives the + // private initializePerPartitionCircuitBreaker() init path, serializes the resulting DiagnosticsClientConfig, + // and asserts that every expected "clientCfgs" key is present (guarding against future serialization + // truncation as well as the specific regression). It also asserts the effective PPCB config string. + @Test(groups = {"unit"}) + public void diagnosticsClientConfigContainsAllClientCfgKeysIncludingPartitionLevelCircuitBreaker() throws Exception { + System.setProperty( + "COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG", + "{\"isPartitionLevelCircuitBreakerEnabled\": true, " + + "\"circuitBreakerType\": \"CONSECUTIVE_EXCEPTION_COUNT_BASED\"," + + "\"consecutiveExceptionCountToleratedForReads\": 10," + + "\"consecutiveExceptionCountToleratedForWrites\": 5,}"); + + Mockito.when(this.connectionPolicyMock.getIdleHttpConnectionTimeout()).thenReturn(Duration.ZERO); + Mockito.when(this.connectionPolicyMock.getMaxConnectionPoolSize()).thenReturn(1); + Mockito.when(this.connectionPolicyMock.getProxy()).thenReturn(null); + Mockito.when(this.connectionPolicyMock.getHttpNetworkRequestTimeout()).thenReturn(Duration.ZERO); + Mockito.when(this.connectionPolicyMock.getHttp2ConnectionConfig()).thenReturn(new Http2ConnectionConfig()); + // The serializer eagerly calls getConnectionMode().toString() for the very first "connectionMode" key; if this + // returns null (Mockito default), serialization would NPE and silently drop every subsequent key. + Mockito.when(this.connectionPolicyMock.getConnectionMode()).thenReturn(ConnectionMode.DIRECT); + + MockedStatic httpClientMock = Mockito.mockStatic(HttpClient.class); + httpClientMock + .when(() -> HttpClient.createFixed(Mockito.any(HttpClientConfig.class))) + .thenReturn(dummyHttpClient()); + + RxDocumentClientImpl rxDocumentClient = null; + + try { + rxDocumentClient = new RxDocumentClientImpl( + this.serviceEndpointMock, + this.masterKeyOrResourceTokenMock, + this.permissionFeedMock, + this.connectionPolicyMock, + this.consistencyLevelMock, + null, + this.configsMock, + this.cosmosAuthorizationTokenResolverMock, + this.azureKeyCredentialMock, + false, + false, + false, + this.metadataCachesSnapshotMock, + this.apiTypeMock, + this.cosmosClientTelemetryConfigMock, + this.clientCorrelationIdMock, + this.endToEndOperationLatencyPolicyConfig, + this.sessionRetryOptionsMock, + this.containerProactiveInitConfigMock, + this.defaultItemSerializer, + false + ); + + // Drive the exact wiring that regressed: explicit (client-side) Per-Partition Circuit Breaker + // initialization. The constructor does not invoke init() (which would require network), so invoke the + // private no-arg initializer reflectively. + Method initPpcb = RxDocumentClientImpl.class.getDeclaredMethod("initializePerPartitionCircuitBreaker"); + initPpcb.setAccessible(true); + initPpcb.invoke(rxDocumentClient); + + ObjectMapper objectMapper = new ObjectMapper(); + StringWriter jsonWriter = new StringWriter(); + JsonGenerator jsonGenerator = new JsonFactory().createGenerator(jsonWriter); + SerializerProvider serializerProvider = objectMapper.getSerializerProvider(); + DiagnosticsClientContext.DiagnosticsClientConfigSerializer.INSTANCE + .serialize(rxDocumentClient.getConfig(), jsonGenerator, serializerProvider); + jsonGenerator.flush(); + ObjectNode clientCfgs = (ObjectNode) objectMapper.readTree(jsonWriter.toString()); + + String serializedJson = clientCfgs.toString(); + + // Every key the serializer unconditionally writes, plus the (previously regressed) + // partitionLevelCircuitBreakerCfg which is present whenever PPCB is enabled. + String[] expectedKeys = new String[] { + "id", + "machineId", + "connectionMode", + "numberOfClients", + "isPpafEnabled", + "isFalseProgSessionTokenMergeEnabled", + "excrgns", + "clientEndpoints", + "connCfg", + "consistencyCfg", + "proactiveInitCfg", + "e2ePolicyCfg", + "sessionRetryCfg", + "partitionLevelCircuitBreakerCfg" + }; + + for (String expectedKey : expectedKeys) { + assertThat(clientCfgs.has(expectedKey)) + .withFailMessage("Expected clientCfgs key '%s' to be present. Serialized clientCfgs: %s", + expectedKey, serializedJson) + .isTrue(); + } + + assertThat(clientCfgs.get("partitionLevelCircuitBreakerCfg").asText()) + .withFailMessage("Unexpected partitionLevelCircuitBreakerCfg value. Serialized clientCfgs: %s", + serializedJson) + .isEqualTo("(cb: true, type: CONSECUTIVE_EXCEPTION_COUNT_BASED, rexcntt: 10, wexcntt: 5)"); + } finally { + System.clearProperty("COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG"); + if (rxDocumentClient != null) { + rxDocumentClient.close(); + } + httpClientMock.close(); + } + } + private static HttpClient dummyHttpClient() { return new HttpClient() { @Override diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 9b035b28dff6..4754dd82ac5a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -7937,7 +7937,6 @@ private synchronized void initializePerPartitionFailover(DatabaseAccount databas checkNotNull(this.globalPartitionEndpointManagerForPerPartitionAutomaticFailover, "Argument 'globalPartitionEndpointManagerForPerPartitionAutomaticFailover' cannot be null."); checkNotNull(this.globalPartitionEndpointManagerForPerPartitionCircuitBreaker, "Argument 'globalPartitionEndpointManagerForPerPartitionCircuitBreaker' cannot be null."); - this.diagnosticsClientConfig.withPartitionLevelCircuitBreakerConfig(this.globalPartitionEndpointManagerForPerPartitionCircuitBreaker.getCircuitBreakerConfig()); this.diagnosticsClientConfig.withIsPerPartitionAutomaticFailoverEnabled(this.globalPartitionEndpointManagerForPerPartitionAutomaticFailover.isPerPartitionAutomaticFailoverEnabled()); } @@ -7966,6 +7965,12 @@ private void initializePerPartitionCircuitBreaker() { this.globalPartitionEndpointManagerForPerPartitionCircuitBreaker.resetCircuitBreakerConfig(partitionLevelCircuitBreakerConfig); this.globalPartitionEndpointManagerForPerPartitionCircuitBreaker.init(); + + // Populate the circuit breaker config in the diagnostics client config here (rather than in + // initializePerPartitionFailover) so the "partitionLevelCircuitBreakerCfg" field appears in + // CosmosDiagnostics whenever the circuit breaker is configured client-side, not only when + // Per-Partition Automatic Failover is mandated by the service. + this.diagnosticsClientConfig.withPartitionLevelCircuitBreakerConfig(this.globalPartitionEndpointManagerForPerPartitionCircuitBreaker.getCircuitBreakerConfig()); } private void enableAvailabilityStrategyForReads() { From e846a7943b9d83d1fe54a6549059322fb9130e69 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Mon, 6 Jul 2026 17:33:34 -0400 Subject: [PATCH 04/14] Address Copilot review: use strictly valid JSON in PPCB test system property Remove trailing comma before closing brace in the COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG JSON so the value is strictly valid and clearer as a customer reference. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../azure/cosmos/implementation/RxDocumentClientImplTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxDocumentClientImplTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxDocumentClientImplTest.java index 242ada8ab223..d6d2eb430c1f 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxDocumentClientImplTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxDocumentClientImplTest.java @@ -339,7 +339,7 @@ public void diagnosticsClientConfigContainsAllClientCfgKeysIncludingPartitionLev "{\"isPartitionLevelCircuitBreakerEnabled\": true, " + "\"circuitBreakerType\": \"CONSECUTIVE_EXCEPTION_COUNT_BASED\"," + "\"consecutiveExceptionCountToleratedForReads\": 10," - + "\"consecutiveExceptionCountToleratedForWrites\": 5,}"); + + "\"consecutiveExceptionCountToleratedForWrites\": 5}"); Mockito.when(this.connectionPolicyMock.getIdleHttpConnectionTimeout()).thenReturn(Duration.ZERO); Mockito.when(this.connectionPolicyMock.getMaxConnectionPoolSize()).thenReturn(1); From 7fa1c57e65acf1fc6e7888467b302f3a37df4d11 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Tue, 7 Jul 2026 10:08:26 -0400 Subject: [PATCH 05/14] Add PPCB clientCfgs diagnostics validation to PerPartitionCircuitBreakerE2ETests Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../PerPartitionCircuitBreakerE2ETests.java | 71 +++++++++++++++++++ 1 file changed, 71 insertions(+) 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 99cef9422503..5d9d14f74dda 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 @@ -4592,6 +4592,77 @@ public void validateHandlingOnNullPartitionKeyRangeOnSmallE2ETimeout_allOps(Oper } } + /** + * Regression validation for the Per-Partition Circuit Breaker (PPCB) diagnostics fix (see PR 49734). + * + * When PPCB is explicitly enabled via the {@code COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG} + * system property, the {@code clientCfgs} section of the emitted {@link CosmosDiagnostics} must + * include the {@code partitionLevelCircuitBreakerCfg} field. A prior regression silently dropped + * this field. This test asserts that all the expected {@code clientCfgs} keys - including + * {@code partitionLevelCircuitBreakerCfg} - are present in the diagnostics of a real operation. + */ + @Test(groups = { "circuit-breaker-misc-direct" }, timeOut = TIMEOUT) + public void partitionLevelCircuitBreakerConfigIsPresentInClientCfgsDiagnostics() { + + System.setProperty( + "COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG", + "{\"isPartitionLevelCircuitBreakerEnabled\": true, " + + "\"circuitBreakerType\": \"CONSECUTIVE_EXCEPTION_COUNT_BASED\"," + + "\"consecutiveExceptionCountToleratedForReads\": 10," + + "\"consecutiveExceptionCountToleratedForWrites\": 5," + + "}"); + + try (CosmosAsyncClient client = getClientBuilder().buildAsyncClient()) { + + CosmosAsyncContainer container = client + .getDatabase(this.sharedAsyncDatabaseId) + .getContainer(this.sharedMultiPartitionAsyncContainerIdWhereIdIsPartitionKey); + + TestObject item = TestObject.create(); + + CosmosItemResponse createResponse = container + .createItem(item, new PartitionKey(item.getId()), new CosmosItemRequestOptions()) + .block(); + + assertThat(createResponse).isNotNull(); + + String diagnosticsString = createResponse.getDiagnostics().toString(); + + assertThat(diagnosticsString) + .as("clientCfgs section should be present in the CosmosDiagnostics") + .contains("\"clientCfgs\""); + + // All the clientCfgs keys unconditionally emitted by DiagnosticsClientConfigSerializer. + List expectedClientCfgsKeys = Arrays.asList( + "id", + "machineId", + "connectionMode", + "numberOfClients", + "isPpafEnabled", + "isFalseProgSessionTokenMergeEnabled", + "excrgns", + "clientEndpoints", + "connCfg", + "consistencyCfg", + "proactiveInitCfg", + "e2ePolicyCfg", + "sessionRetryCfg"); + + for (String expectedKey : expectedClientCfgsKeys) { + assertThat(diagnosticsString) + .as("clientCfgs key '%s' should be present in the CosmosDiagnostics", expectedKey) + .contains("\"" + expectedKey + "\""); + } + + // The regression fix: PPCB config must be present in clientCfgs when explicitly enabled. + assertThat(diagnosticsString) + .as("partitionLevelCircuitBreakerCfg should be present in clientCfgs when PPCB is enabled") + .contains("\"partitionLevelCircuitBreakerCfg\""); + } finally { + System.clearProperty("COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG"); + } + } + private static Function> resolveDataPlaneOperation(FaultInjectionOperationType faultInjectionOperationType) { switch (faultInjectionOperationType) { From 0125e45506e9f0d3e61a74825919245bf6863dec Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Tue, 7 Jul 2026 10:28:23 -0400 Subject: [PATCH 06/14] Run PPCB clientCfgs diagnostics E2E test across all TestNG groups Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../com/azure/cosmos/PerPartitionCircuitBreakerE2ETests.java | 2 +- 1 file changed, 1 insertion(+), 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 5d9d14f74dda..c468d79a57fe 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 @@ -4601,7 +4601,7 @@ public void validateHandlingOnNullPartitionKeyRangeOnSmallE2ETimeout_allOps(Oper * this field. This test asserts that all the expected {@code clientCfgs} keys - including * {@code partitionLevelCircuitBreakerCfg} - are present in the diagnostics of a real operation. */ - @Test(groups = { "circuit-breaker-misc-direct" }, timeOut = TIMEOUT) + @Test(groups = { "circuit-breaker-misc-gateway", "circuit-breaker-misc-direct", "circuit-breaker-read-all-read-many", "multi-region", "fi-thinclient-multi-master" }, timeOut = TIMEOUT) public void partitionLevelCircuitBreakerConfigIsPresentInClientCfgsDiagnostics() { System.setProperty( From b1b6862f620f2f5ba6f7779d0cfd4c461c2c5c4f Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Tue, 7 Jul 2026 10:39:51 -0400 Subject: [PATCH 07/14] Simplify PPCB E2E diagnostics test to verify unconditional clientCfgs keys The partitionLevelCircuitBreakerCfg field now appears in clientCfgs for every client regardless of PPCB configuration, so the E2E test no longer sets the COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG system property. It just builds a plain client and asserts all expected clientCfgs keys (including partitionLevelCircuitBreakerCfg) are present in CosmosDiagnostics. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../PerPartitionCircuitBreakerE2ETests.java | 32 ++++++------------- 1 file changed, 10 insertions(+), 22 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 c468d79a57fe..584f00595812 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 @@ -4595,23 +4595,16 @@ public void validateHandlingOnNullPartitionKeyRangeOnSmallE2ETimeout_allOps(Oper /** * Regression validation for the Per-Partition Circuit Breaker (PPCB) diagnostics fix (see PR 49734). * - * When PPCB is explicitly enabled via the {@code COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG} - * system property, the {@code clientCfgs} section of the emitted {@link CosmosDiagnostics} must - * include the {@code partitionLevelCircuitBreakerCfg} field. A prior regression silently dropped - * this field. This test asserts that all the expected {@code clientCfgs} keys - including - * {@code partitionLevelCircuitBreakerCfg} - are present in the diagnostics of a real operation. + * The {@code clientCfgs} section of the emitted {@link CosmosDiagnostics} must always include the + * {@code partitionLevelCircuitBreakerCfg} field for every client, regardless of whether PPCB is + * explicitly enabled. A prior regression silently dropped this field unless PPAF mandated it. This + * test asserts that all the expected {@code clientCfgs} keys - including + * {@code partitionLevelCircuitBreakerCfg} - are present in the diagnostics of a real operation + * without setting any PPCB configuration. */ @Test(groups = { "circuit-breaker-misc-gateway", "circuit-breaker-misc-direct", "circuit-breaker-read-all-read-many", "multi-region", "fi-thinclient-multi-master" }, timeOut = TIMEOUT) public void partitionLevelCircuitBreakerConfigIsPresentInClientCfgsDiagnostics() { - System.setProperty( - "COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG", - "{\"isPartitionLevelCircuitBreakerEnabled\": true, " - + "\"circuitBreakerType\": \"CONSECUTIVE_EXCEPTION_COUNT_BASED\"," - + "\"consecutiveExceptionCountToleratedForReads\": 10," - + "\"consecutiveExceptionCountToleratedForWrites\": 5," - + "}"); - try (CosmosAsyncClient client = getClientBuilder().buildAsyncClient()) { CosmosAsyncContainer container = client @@ -4632,7 +4625,8 @@ public void partitionLevelCircuitBreakerConfigIsPresentInClientCfgsDiagnostics() .as("clientCfgs section should be present in the CosmosDiagnostics") .contains("\"clientCfgs\""); - // All the clientCfgs keys unconditionally emitted by DiagnosticsClientConfigSerializer. + // All the clientCfgs keys unconditionally emitted by DiagnosticsClientConfigSerializer, + // including partitionLevelCircuitBreakerCfg (the field the regression previously dropped). List expectedClientCfgsKeys = Arrays.asList( "id", "machineId", @@ -4646,20 +4640,14 @@ public void partitionLevelCircuitBreakerConfigIsPresentInClientCfgsDiagnostics() "consistencyCfg", "proactiveInitCfg", "e2ePolicyCfg", - "sessionRetryCfg"); + "sessionRetryCfg", + "partitionLevelCircuitBreakerCfg"); for (String expectedKey : expectedClientCfgsKeys) { assertThat(diagnosticsString) .as("clientCfgs key '%s' should be present in the CosmosDiagnostics", expectedKey) .contains("\"" + expectedKey + "\""); } - - // The regression fix: PPCB config must be present in clientCfgs when explicitly enabled. - assertThat(diagnosticsString) - .as("partitionLevelCircuitBreakerCfg should be present in clientCfgs when PPCB is enabled") - .contains("\"partitionLevelCircuitBreakerCfg\""); - } finally { - System.clearProperty("COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG"); } } From ce34f6e97a36ee7ef5d1643d5cc2bff45610efd4 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Tue, 7 Jul 2026 11:03:10 -0400 Subject: [PATCH 08/14] Simplify PPCB unit test to assert clientCfgs key presence only Remove PPCB System.setProperty/clearProperty and the value-specific assertion; assert only that all clientCfgs keys are present, matching the E2E test simplification. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../RxDocumentClientImplTest.java | 29 ++++++------------- 1 file changed, 9 insertions(+), 20 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxDocumentClientImplTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxDocumentClientImplTest.java index d6d2eb430c1f..39a834ad2267 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxDocumentClientImplTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxDocumentClientImplTest.java @@ -327,20 +327,15 @@ public void readMany() { // Regression test for the "partitionLevelCircuitBreakerCfg" diagnostics field silently disappearing from the // CosmosDiagnostics "clientCfgs" section. Prior to the fix, the field was only written on the PPAF - // (service-mandated) path, so a client that explicitly enabled Per-Partition Circuit Breaker never surfaced it. - // This test constructs a real RxDocumentClientImpl (exercising the actual constructor wiring), drives the - // private initializePerPartitionCircuitBreaker() init path, serializes the resulting DiagnosticsClientConfig, - // and asserts that every expected "clientCfgs" key is present (guarding against future serialization - // truncation as well as the specific regression). It also asserts the effective PPCB config string. + // (service-mandated) path, so a client that did not have PPAF-mandated PPCB never surfaced it. The field must + // now be present for every client regardless of any PPCB configuration. This test constructs a real + // RxDocumentClientImpl (exercising the actual constructor wiring), drives the private + // initializePerPartitionCircuitBreaker() init path without setting any PPCB configuration, serializes the + // resulting DiagnosticsClientConfig, and asserts that every expected "clientCfgs" key - including + // partitionLevelCircuitBreakerCfg - is present (guarding against future serialization truncation as well as + // the specific regression). @Test(groups = {"unit"}) public void diagnosticsClientConfigContainsAllClientCfgKeysIncludingPartitionLevelCircuitBreaker() throws Exception { - System.setProperty( - "COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG", - "{\"isPartitionLevelCircuitBreakerEnabled\": true, " - + "\"circuitBreakerType\": \"CONSECUTIVE_EXCEPTION_COUNT_BASED\"," - + "\"consecutiveExceptionCountToleratedForReads\": 10," - + "\"consecutiveExceptionCountToleratedForWrites\": 5}"); - Mockito.when(this.connectionPolicyMock.getIdleHttpConnectionTimeout()).thenReturn(Duration.ZERO); Mockito.when(this.connectionPolicyMock.getMaxConnectionPoolSize()).thenReturn(1); Mockito.when(this.connectionPolicyMock.getProxy()).thenReturn(null); @@ -400,8 +395,8 @@ public void diagnosticsClientConfigContainsAllClientCfgKeysIncludingPartitionLev String serializedJson = clientCfgs.toString(); - // Every key the serializer unconditionally writes, plus the (previously regressed) - // partitionLevelCircuitBreakerCfg which is present whenever PPCB is enabled. + // Every key the serializer unconditionally writes, including the (previously regressed) + // partitionLevelCircuitBreakerCfg which must be present for every client regardless of PPCB config. String[] expectedKeys = new String[] { "id", "machineId", @@ -425,13 +420,7 @@ public void diagnosticsClientConfigContainsAllClientCfgKeysIncludingPartitionLev expectedKey, serializedJson) .isTrue(); } - - assertThat(clientCfgs.get("partitionLevelCircuitBreakerCfg").asText()) - .withFailMessage("Unexpected partitionLevelCircuitBreakerCfg value. Serialized clientCfgs: %s", - serializedJson) - .isEqualTo("(cb: true, type: CONSECUTIVE_EXCEPTION_COUNT_BASED, rexcntt: 10, wexcntt: 5)"); } finally { - System.clearProperty("COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG"); if (rxDocumentClient != null) { rxDocumentClient.close(); } From 1011ce309e58d5c4e82275ca1d2fb1d01df7a437 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Thu, 27 Aug 2026 07:50:31 -0400 Subject: [PATCH 09/14] Prepare azure-cosmos 4.76.1-hotfix backport Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 16ce0941-555c-4190-8c4d-96c705086350 --- sdk/cosmos/azure-cosmos-encryption/pom.xml | 2 +- sdk/cosmos/azure-cosmos-kafka-connect/pom.xml | 2 +- sdk/cosmos/azure-cosmos-spark_3/pom.xml | 2 +- sdk/cosmos/azure-cosmos-test/pom.xml | 2 +- sdk/cosmos/azure-cosmos-tests/pom.xml | 2 +- ...titionEndpointManagerForPPCBUnitTests.java | 281 -------- .../PerPartitionCircuitBreakerE2ETests.java | 668 +----------------- .../GatewayAddressCacheTest.java | 293 +------- ...PartitionCircuitBreakerInfoHolderTest.java | 27 +- sdk/cosmos/azure-cosmos/CHANGELOG.md | 110 +-- sdk/cosmos/azure-cosmos/pom.xml | 2 +- .../ClientSideRequestStatistics.java | 42 +- ...tManagerForPerPartitionCircuitBreaker.java | 9 +- .../PerPartitionCircuitBreakerInfoHolder.java | 2 + 14 files changed, 41 insertions(+), 1403 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-encryption/pom.xml b/sdk/cosmos/azure-cosmos-encryption/pom.xml index 2ec6e7beb33e..a6b4c32243c0 100644 --- a/sdk/cosmos/azure-cosmos-encryption/pom.xml +++ b/sdk/cosmos/azure-cosmos-encryption/pom.xml @@ -61,7 +61,7 @@ Licensed under the MIT License. com.azure azure-cosmos - 4.76.0 + 4.76.1-hotfix diff --git a/sdk/cosmos/azure-cosmos-kafka-connect/pom.xml b/sdk/cosmos/azure-cosmos-kafka-connect/pom.xml index 5590ebb9ff63..9a438ce3e5ba 100644 --- a/sdk/cosmos/azure-cosmos-kafka-connect/pom.xml +++ b/sdk/cosmos/azure-cosmos-kafka-connect/pom.xml @@ -92,7 +92,7 @@ Licensed under the MIT License. com.azure azure-cosmos - 4.76.0 + 4.76.1-hotfix + 4.76.1-hotfix org.slf4j diff --git a/sdk/cosmos/azure-cosmos-test/pom.xml b/sdk/cosmos/azure-cosmos-test/pom.xml index db566dd3bdc2..1d310ced4e51 100644 --- a/sdk/cosmos/azure-cosmos-test/pom.xml +++ b/sdk/cosmos/azure-cosmos-test/pom.xml @@ -59,7 +59,7 @@ Licensed under the MIT License. com.azure azure-cosmos - 4.76.0 + 4.76.1-hotfix diff --git a/sdk/cosmos/azure-cosmos-tests/pom.xml b/sdk/cosmos/azure-cosmos-tests/pom.xml index cf7a2c8cb5e3..5579b257c9c5 100644 --- a/sdk/cosmos/azure-cosmos-tests/pom.xml +++ b/sdk/cosmos/azure-cosmos-tests/pom.xml @@ -100,7 +100,7 @@ Licensed under the MIT License. com.azure azure-cosmos - 4.76.0 + 4.76.1-hotfix com.azure 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 22fd70774881..11d14758f352 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,12 +4,9 @@ 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; @@ -18,22 +15,11 @@ 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; import com.azure.cosmos.implementation.perPartitionCircuitBreaker.LocationSpecificHealthContext; import com.azure.cosmos.implementation.guava25.collect.ImmutableList; import com.azure.cosmos.implementation.routing.RegionalRoutingContext; -import com.fasterxml.jackson.databind.ObjectMapper; -import io.netty.channel.ConnectTimeoutException; import org.apache.commons.lang3.tuple.Pair; import org.mockito.Mockito; import org.slf4j.Logger; @@ -41,29 +27,17 @@ 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; @@ -78,11 +52,6 @@ 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; @@ -152,15 +121,6 @@ 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 { @@ -1047,247 +1007,6 @@ 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); - String diagnostics = new ObjectMapper().writeValueAsString( - request.requestContext.getPerPartitionCircuitBreakerInfoHolder()); - assertThat(diagnostics) - .contains("\"outcome\":\"Failed\"") - .contains("\"stage\":\"OPEN_CONNECTION_TASK\"") - .contains("\"type\":\"io.netty.channel.ConnectTimeoutException\"") - .contains("\"latestFailbackMessageByRegion\":{") - .contains("\"East US\":\"Refreshed replica is unavailable\""); - } else { - assertThat(ppcbManager.getUnavailableRegionsForPartitionKeyRange( - request, - collectionRid, - partitionKeyRange)).isEmpty(); - assertThat(request.requestContext.getPerPartitionCircuitBreakerInfoHolder() - .getPerPartitionCircuitBreakerInfoHolder() - .get("East US") - .getUnavailableSince()).isEqualTo(Instant.MAX); - assertThat(new ObjectMapper().writeValueAsString( - request.requestContext.getPerPartitionCircuitBreakerInfoHolder())) - .contains("\"outcome\":\"Succeeded\"") - .doesNotContain("\"failure\"", "\"latestFailbackMessageByRegion\""); - } - - assertThat(new ObjectMapper().writeValueAsString( - request.requestContext.getPerPartitionCircuitBreakerInfoHolder())) - .contains("\"lastAttemptedAt\":"); - - 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 584f00595812..4390b3f83c9c 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,7 +3,6 @@ package com.azure.cosmos; -import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.faultinjection.FaultInjectionTestBase; import com.azure.cosmos.implementation.ConnectionPolicy; import com.azure.cosmos.implementation.DatabaseAccount; @@ -14,7 +13,6 @@ import com.azure.cosmos.implementation.ImplementationBridgeHelpers; import com.azure.cosmos.implementation.OperationType; import com.azure.cosmos.implementation.PartitionKeyRange; -import com.azure.cosmos.implementation.ResourceType; import com.azure.cosmos.implementation.RxDocumentClientImpl; import com.azure.cosmos.implementation.TestConfigurations; import com.azure.cosmos.implementation.Utils; @@ -3551,8 +3549,6 @@ private void execute( boolean hasReachedCircuitBreakingThreshold = false; int executionCountAfterCircuitBreakingThresholdBreached = 0; - boolean failbackExpected = false; - Set loggedPpcbDiagnosticsPhases = new HashSet<>(); List testObjects = operationInvocationParamsWrapper.testObjectsForDataPlaneOperationToWorkWith; PartitionKeyRangeWrapper partitionKeyRangeWrapper @@ -3566,12 +3562,7 @@ private void execute( validateNonEmptyList(operationInvocationParamsWrapper.itemIdentitiesForReadManyOperation); } - ResponseWrapper response = executeDataPlaneOperationWithTransient4041002Retry( - testId, - executeDataPlaneOperation, - operationInvocationParamsWrapper); - assertPpcbSnapshotsPopulated(response, PpcbDiagnosticsPhase.FAILURE, false); - logPpcbDiagnosticsOnce(response, PpcbDiagnosticsPhase.FAILURE, loggedPpcbDiagnosticsPhases); + ResponseWrapper response = executeDataPlaneOperation.apply(operationInvocationParamsWrapper); ConsecutiveExceptionBasedCircuitBreaker consecutiveExceptionBasedCircuitBreaker = globalPartitionEndpointManagerForPerPartitionCircuitBreaker.getConsecutiveExceptionBasedCircuitBreaker(); @@ -3597,14 +3588,6 @@ private void execute( if (executionCountAfterCircuitBreakingThresholdBreached > 1) { validateResponseInAbsenceOfFailures.accept(response); - failbackExpected |= assertPpcbSnapshotsPopulated( - response, - PpcbDiagnosticsPhase.POST_FAILOVER, - false); - logPpcbDiagnosticsOnce( - response, - PpcbDiagnosticsPhase.POST_FAILOVER, - loggedPpcbDiagnosticsPhases); } if (response.cosmosItemResponse != null) { @@ -3656,14 +3639,6 @@ private void execute( ResponseWrapper response = executeDataPlaneOperation.apply(operationInvocationParamsWrapper); validateResponseInAbsenceOfFailures.accept(response); - assertPpcbSnapshotsPopulated( - response, - PpcbDiagnosticsPhase.POST_FAILBACK, - failbackExpected); - logPpcbDiagnosticsOnce( - response, - PpcbDiagnosticsPhase.POST_FAILBACK, - loggedPpcbDiagnosticsPhases); if (response.cosmosItemResponse != null) { assertThat(response.cosmosItemResponse).isNotNull(); @@ -3701,334 +3676,6 @@ private void execute( } } - private static CosmosDiagnosticsContext getDiagnosticsContext(ResponseWrapper response) { - if (response.cosmosItemResponse != null) { - return response.cosmosItemResponse.getDiagnostics().getDiagnosticsContext(); - } else if (response.feedResponse != null) { - return response.feedResponse.getCosmosDiagnostics().getDiagnosticsContext(); - } else if (response.cosmosException != null) { - return response.cosmosException.getDiagnostics().getDiagnosticsContext(); - } else if (response.batchResponse != null) { - return response.batchResponse.getDiagnostics().getDiagnosticsContext(); - } - return null; - } - - private static void logPpcbDiagnosticsOnce( - ResponseWrapper response, - PpcbDiagnosticsPhase phase, - Set loggedPhases) { - - if (loggedPhases.add(phase)) { - CosmosDiagnosticsContext diagnosticsContext = getDiagnosticsContext(response); - if (diagnosticsContext != null) { - logger.info("PPCB CosmosDiagnostics [{}]: {}", phase.label, diagnosticsContext.toJson()); - } - } - } - - private static boolean assertPpcbSnapshotsPopulated( - ResponseWrapper response, - PpcbDiagnosticsPhase phase, - boolean failbackExpected) { - - CosmosDiagnosticsContext diagnosticsContext = getDiagnosticsContext(response); - assertThat(diagnosticsContext) - .as("Expected CosmosDiagnostics for %s", phase.label) - .isNotNull(); - assertThat(diagnosticsContext.getDiagnostics()) - .as("Expected diagnostics entries for %s", phase.label) - .isNotNull(); - - int applicableStatisticCount = 0; - List healthContexts = new ArrayList<>(); - for (CosmosDiagnostics cosmosDiagnostics : diagnosticsContext.getDiagnostics()) { - Collection statisticsCollection = - cosmosDiagnosticsAccessor.getClientSideRequestStatistics(cosmosDiagnostics); - if (statisticsCollection == null) { - continue; - } - - for (ClientSideRequestStatistics statistics : statisticsCollection) { - if (statistics == null) { - continue; - } - - for (ClientSideRequestStatistics.StoreResponseStatistics storeStatistics - : statistics.getResponseStatisticsList()) { - - if (isPpcbApplicableDataPlaneStatistic( - storeStatistics.getRequestResourceType(), - storeStatistics.getRequestOperationType())) { - - applicableStatisticCount++; - assertThat(storeStatistics.getPerPartitionCircuitBreakerInfoHolder()) - .as("Expected direct PPCB holder for %s", phase.label) - .isNotNull(); - Map stateByRegion - = storeStatistics.getPerPartitionCircuitBreakerInfoHolder() - .getPerPartitionCircuitBreakerInfoHolder(); - assertThat(stateByRegion) - .as("Expected populated direct PPCB snapshot for %s", phase.label) - .isNotNull(); - healthContexts.addAll(stateByRegion.values()); - } - } - - for (ClientSideRequestStatistics.GatewayStatistics gatewayStatistics - : statistics.getGatewayStatisticsList()) { - - if (isPpcbApplicableDataPlaneStatistic( - gatewayStatistics.getResourceType(), - gatewayStatistics.getOperationType())) { - - applicableStatisticCount++; - assertThat(gatewayStatistics.getPerPartitionCircuitBreakerInfoHolder()) - .as("Expected gateway PPCB holder for %s", phase.label) - .isNotNull(); - Map stateByRegion - = gatewayStatistics.getPerPartitionCircuitBreakerInfoHolder() - .getPerPartitionCircuitBreakerInfoHolder(); - assertThat(stateByRegion) - .as("Expected populated gateway PPCB snapshot for %s", phase.label) - .isNotNull(); - healthContexts.addAll(stateByRegion.values()); - } - } - } - } - - if (applicableStatisticCount == 0) { - assertThat(hasOnlyQueryPlanStatistics(diagnosticsContext)) - .as("Expected PPCB-applicable data-plane statistics or QueryPlan-only diagnostics for %s", phase.label) - .isTrue(); - } - - boolean unavailableRegionFound = false; - boolean successfulFailbackFound = false; - for (LocationSpecificHealthContext healthContext : healthContexts) { - if (healthContext.getLocationHealthStatus() == LocationHealthStatus.Unavailable) { - unavailableRegionFound = true; - if (phase == PpcbDiagnosticsPhase.POST_FAILOVER) { - assertThat(healthContext.getLastFailbackOutcome()) - .as("Failback must not have succeeded while the region remains unavailable") - .isNotEqualTo(LocationSpecificHealthContext.FailbackOutcome.Succeeded); - } - } - - if (healthContext.getLastFailbackOutcome() - == LocationSpecificHealthContext.FailbackOutcome.Succeeded) { - - successfulFailbackFound = true; - assertThat(healthContext.getLastFailbackAttemptTime()) - .as("Expected failback attempt timestamp after successful failback") - .isNotNull(); - assertThat(healthContext.getLocationHealthStatus()) - .as("Expected recovered region after successful failback") - .isIn(LocationHealthStatus.HealthyTentative, LocationHealthStatus.Healthy); - } - } - - if (phase == PpcbDiagnosticsPhase.POST_FAILBACK && failbackExpected) { - assertThat(successfulFailbackFound) - .as("Expected a successful failback outcome for a previously unavailable region") - .isTrue(); - } - - return unavailableRegionFound; - } - - private static boolean isPpcbApplicableDataPlaneStatistic( - ResourceType resourceType, - OperationType operationType) { - - return resourceType == ResourceType.Document && operationType != OperationType.QueryPlan; - } - - private static boolean hasOnlyQueryPlanStatistics(CosmosDiagnosticsContext diagnosticsContext) { - boolean queryPlanStatisticFound = false; - for (CosmosDiagnostics cosmosDiagnostics : diagnosticsContext.getDiagnostics()) { - Collection statisticsCollection = - cosmosDiagnosticsAccessor.getClientSideRequestStatistics(cosmosDiagnostics); - if (statisticsCollection == null) { - continue; - } - - for (ClientSideRequestStatistics statistics : statisticsCollection) { - if (statistics == null) { - continue; - } - - for (ClientSideRequestStatistics.StoreResponseStatistics storeStatistics - : statistics.getResponseStatisticsList()) { - - if (storeStatistics.getRequestResourceType() != ResourceType.Document) { - continue; - } - if (storeStatistics.getRequestOperationType() != OperationType.QueryPlan) { - return false; - } - queryPlanStatisticFound = true; - } - - for (ClientSideRequestStatistics.GatewayStatistics gatewayStatistics - : statistics.getGatewayStatisticsList()) { - - if (gatewayStatistics.getResourceType() != ResourceType.Document) { - continue; - } - if (gatewayStatistics.getOperationType() != OperationType.QueryPlan) { - return false; - } - queryPlanStatisticFound = true; - } - } - } - - return queryPlanStatisticFound; - } - - private ResponseWrapper executeDataPlaneOperationWithTransient4041002Retry( - String testId, - Function> executeDataPlaneOperation, - OperationInvocationParamsWrapper operationInvocationParamsWrapper) throws InterruptedException { - - long retryStartNanos = System.nanoTime(); - int retryAttempt = 0; - ResponseWrapper response = executeDataPlaneOperation.apply(operationInvocationParamsWrapper); - - while (hasNonFaultInjected404RetryableResponse(response)) { - Duration elapsed = Duration.ofNanos(System.nanoTime() - retryStartNanos); - if (elapsed.compareTo(TRANSIENT_404_1002_MAX_RETRY_DURATION) >= 0) { - logger.warn( - "Detected non-fault-injected retryable 404 in diagnostics for test {} for {}. " - + "Continuing with latest response so normal assertions can report diagnostics.", - testId, - elapsed); - return response; - } - - retryAttempt++; - logger.warn( - "Detected non-fault-injected retryable 404 in diagnostics for test {}. " - + "Waiting {} before retry attempt {}.", - testId, - TRANSIENT_404_1002_RETRY_DELAY, - retryAttempt); - Thread.sleep(TRANSIENT_404_1002_RETRY_DELAY.toMillis()); - response = executeDataPlaneOperation.apply(operationInvocationParamsWrapper); - } - - return response; - } - - private static boolean hasNonFaultInjected404RetryableResponse(ResponseWrapper response) { - if (!hasRetryableTerminal404(response)) { - return false; - } - - CosmosDiagnosticsContext diagnosticsContext = getDiagnosticsContext(response); - if (diagnosticsContext == null || diagnosticsContext.getDiagnostics() == null) { - return false; - } - - for (CosmosDiagnostics cosmosDiagnostics : diagnosticsContext.getDiagnostics()) { - Collection clientSideRequestStatisticsCollection = - cosmosDiagnosticsAccessor.getClientSideRequestStatistics(cosmosDiagnostics); - if (clientSideRequestStatisticsCollection == null) { - continue; - } - - for (ClientSideRequestStatistics clientSideRequestStatistics : clientSideRequestStatisticsCollection) { - if (clientSideRequestStatistics == null) { - continue; - } - - if (hasNonFaultInjected404RetryableGatewayResponse(clientSideRequestStatistics.getGatewayStatisticsList())) { - return true; - } - - if (hasNonFaultInjected404RetryableStoreResponse(clientSideRequestStatistics.getResponseStatisticsList()) - || hasNonFaultInjected404RetryableStoreResponse(clientSideRequestStatistics.getSupplementalResponseStatisticsList())) { - - return true; - } - } - } - - return false; - } - - private static boolean hasRetryableTerminal404(ResponseWrapper response) { - if (response == null) { - return false; - } - - if (response.cosmosException != null) { - return isRetryable404( - response.cosmosException.getStatusCode(), - response.cosmosException.getSubStatusCode()); - } - - return response.batchResponse != null - && isRetryable404( - response.batchResponse.getStatusCode(), - response.batchResponse.getSubStatusCode()); - } - - private static boolean hasNonFaultInjected404RetryableGatewayResponse( - List gatewayStatisticsList) { - - if (gatewayStatisticsList == null) { - return false; - } - - for (ClientSideRequestStatistics.GatewayStatistics gatewayStatistics : gatewayStatisticsList) { - if (gatewayStatistics != null - && isRetryable404(gatewayStatistics.getStatusCode(), gatewayStatistics.getSubStatusCode()) - && isNullOrEmpty(gatewayStatistics.getFaultInjectionRuleId())) { - - return true; - } - } - - return false; - } - - private static boolean hasNonFaultInjected404RetryableStoreResponse( - Collection storeResponseStatisticsCollection) { - - if (storeResponseStatisticsCollection == null) { - return false; - } - - for (ClientSideRequestStatistics.StoreResponseStatistics storeResponseStatistics : storeResponseStatisticsCollection) { - StoreResultDiagnostics storeResultDiagnostics = - storeResponseStatistics == null ? null : storeResponseStatistics.getStoreResult(); - StoreResponseDiagnostics storeResponseDiagnostics = - storeResultDiagnostics == null ? null : storeResultDiagnostics.getStoreResponseDiagnostics(); - - if (storeResponseDiagnostics != null - && isRetryable404(storeResponseDiagnostics.getStatusCode(), storeResponseDiagnostics.getSubStatusCode()) - && isNullOrEmpty(storeResponseDiagnostics.getFaultInjectionRuleId())) { - - return true; - } - } - - return false; - } - - private static boolean isRetryable404(int statusCode, int subStatusCode) { - return statusCode == HttpConstants.StatusCodes.NOTFOUND - && (subStatusCode == HttpConstants.SubStatusCodes.UNKNOWN - || subStatusCode == HttpConstants.SubStatusCodes.READ_SESSION_NOT_AVAILABLE); - } - - private static boolean isNullOrEmpty(String value) { - return value == null || value.isEmpty(); - } - private static int resolveTestObjectCountToBootstrapFrom(FaultInjectionOperationType faultInjectionOperationType, int opCount) { switch (faultInjectionOperationType) { case READ_ITEM: @@ -5587,26 +5234,6 @@ 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) { @@ -5622,18 +5249,6 @@ private enum QueryType { READ_MANY, READ_ALL } - private enum PpcbDiagnosticsPhase { - FAILURE("failed operation"), - POST_FAILOVER("post-failover operation"), - POST_FAILBACK("post-failback operation"); - - private final String label; - - PpcbDiagnosticsPhase(String label) { - this.label = label; - } - } - private static class AccountLevelLocationContext { private final List serviceOrderedReadableRegions; private final List serviceOrderedWriteableRegions; @@ -5649,285 +5264,4 @@ public AccountLevelLocationContext( this.regionNameToEndpoint = regionNameToEndpoint; } } - - @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() { - - if (this.writeRegions == null || this.writeRegions.size() <= 1) { - throw new SkipException("Test requires multi-region account"); - } - - // Build non-canonical preferred regions: "West US 3" → "westus3", "East US" → "eastus" - List nonCanonicalRegions = new ArrayList<>(); - for (String region : this.writeRegions) { - nonCanonicalRegions.add(region.toLowerCase(Locale.ROOT).replace(" ", "")); - } - - String firstRegionCanonicalLower = this.writeRegions.get(0).toLowerCase(Locale.ROOT); - String secondRegionCanonicalLower = this.writeRegions.get(1).toLowerCase(Locale.ROOT); - - System.setProperty( - "COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG", - "{\"isPartitionLevelCircuitBreakerEnabled\": true, " - + "\"circuitBreakerType\": \"CONSECUTIVE_EXCEPTION_COUNT_BASED\"," - + "\"consecutiveExceptionCountToleratedForReads\": 10," - + "\"consecutiveExceptionCountToleratedForWrites\": 5," - + "}"); - - CosmosClientBuilder clientBuilder = getClientBuilder() - .multipleWriteRegionsEnabled(true) - .preferredRegions(nonCanonicalRegions); - - ConnectionPolicy connectionPolicy = ReflectionUtils.getConnectionPolicy(clientBuilder); - 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"); - } - - CosmosAsyncClient asyncClient = null; - - try { - asyncClient = clientBuilder.buildAsyncClient(); - - CosmosAsyncContainer container = asyncClient - .getDatabase(this.sharedAsyncDatabaseId) - .getContainer(this.sharedMultiPartitionAsyncContainerIdWhereIdIsPartitionKey); - - // Bootstrap: create a test item - TestObject testObject = TestObject.create(); - container.createItem(testObject, new PartitionKey(testObject.getId()), new CosmosItemRequestOptions()).block(); - - // Step 1: Inject 503 (ServiceUnavailable) into the first preferred region for READ_ITEM - FaultInjectionCondition faultCondition = new FaultInjectionConditionBuilder() - .region(this.writeRegions.get(0)) - .operationType(FaultInjectionOperationType.READ_ITEM) - .build(); - - FaultInjectionServerErrorResult serverError = FaultInjectionResultBuilders - .getResultBuilder(FaultInjectionServerErrorType.SERVICE_UNAVAILABLE) - .build(); - - FaultInjectionRule faultRule = new FaultInjectionRuleBuilder("ppcb-non-canonical-region-test-" + UUID.randomUUID()) - .condition(faultCondition) - .result(serverError) - .hitLimit(15) - .build(); - - CosmosFaultInjectionHelper.configureFaultInjectionRules(container, Arrays.asList(faultRule)).block(); - - // Step 2: Issue reads until circuit breaker trips — expect failover to second region - boolean circuitBreakerTripped = false; - - for (int i = 0; i < 20; i++) { - CosmosItemRequestOptions readOptions = new CosmosItemRequestOptions(); - readOptions.setCosmosEndToEndOperationLatencyPolicyConfig(NO_END_TO_END_TIMEOUT); - - CosmosItemResponse readResponse = container - .readItem(testObject.getId(), new PartitionKey(testObject.getId()), readOptions, TestObject.class) - .block(); - - assertThat(readResponse).isNotNull(); - assertThat(readResponse.getStatusCode()).isEqualTo(200); - - CosmosDiagnosticsContext ctx = readResponse.getDiagnostics().getDiagnosticsContext(); - - // Once we see only the second region contacted, the circuit breaker has tripped - if (ctx.getContactedRegionNames().contains(secondRegionCanonicalLower) - && !ctx.getContactedRegionNames().contains(firstRegionCanonicalLower)) { - circuitBreakerTripped = true; - logger.info("Circuit breaker tripped at iteration {}, routing to second region: {}", i, secondRegionCanonicalLower); - break; - } - } - - assertThat(circuitBreakerTripped) - .as("PPCB should have tripped and routed reads to the second preferred region (%s) " - + "even though preferred regions were passed in non-canonical form (%s)", - secondRegionCanonicalLower, nonCanonicalRegions) - .isTrue(); - - } finally { - System.clearProperty("COSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_CONFIG"); - if (asyncClient != null) { - asyncClient.close(); - } - } - } } 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 2247c2eea902..8285ea915603 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 @@ -16,9 +16,7 @@ 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; @@ -34,7 +32,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 io.reactivex.subscribers.TestSubscriber; import org.assertj.core.api.AssertionsForClassTypes; import org.mockito.ArgumentCaptor; import org.mockito.ArgumentMatchers; @@ -57,13 +55,10 @@ 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.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; @@ -1599,292 +1594,6 @@ 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-tests/src/test/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/PerPartitionCircuitBreakerInfoHolderTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/PerPartitionCircuitBreakerInfoHolderTest.java index 87b7193d9ed4..3ef69fa0c9f0 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/PerPartitionCircuitBreakerInfoHolderTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/PerPartitionCircuitBreakerInfoHolderTest.java @@ -4,7 +4,6 @@ package com.azure.cosmos.implementation.perPartitionCircuitBreaker; import com.azure.cosmos.implementation.ClientSideRequestStatistics; -import com.azure.cosmos.implementation.CrossRegionAvailabilityContextForRxDocumentServiceRequest; import com.azure.cosmos.implementation.DiagnosticsClientContext; import com.azure.cosmos.implementation.GlobalEndpointManager; import com.azure.cosmos.implementation.OperationType; @@ -13,7 +12,6 @@ import com.azure.cosmos.implementation.RxDocumentServiceRequest; import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList; import com.azure.cosmos.implementation.directconnectivity.StoreResponseDiagnostics; -import com.azure.cosmos.implementation.perPartitionAutomaticFailover.PerPartitionAutomaticFailoverInfoHolder; import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import com.fasterxml.jackson.databind.ObjectMapper; import org.mockito.Mockito; @@ -25,7 +23,6 @@ import java.util.Collections; import java.util.LinkedHashMap; import java.util.Map; -import java.util.concurrent.atomic.AtomicBoolean; import static org.assertj.core.api.Assertions.assertThat; import static org.mockito.Mockito.doReturn; @@ -153,7 +150,7 @@ public void responseStatisticsRetainStateAtRecordTime() { ClientSideRequestStatistics statistics = new ClientSideRequestStatistics(diagnosticsClientContext); statistics.recordResponse(request, null, null); - holder.setPerPartitionCircuitBreakerInfoHolder(Collections.singletonMap( + request.requestContext.setPerPartitionCircuitBreakerInfoHolder(Collections.singletonMap( "westus", createHealthContext(LocationHealthStatus.Healthy))); @@ -175,7 +172,7 @@ public void gatewayStatisticsRetainStateAtRecordTime() throws Exception { ClientSideRequestStatistics statistics = new ClientSideRequestStatistics(diagnosticsClientContext); statistics.recordGatewayResponse(request, Mockito.mock(StoreResponseDiagnostics.class), null); - holder.setPerPartitionCircuitBreakerInfoHolder(Collections.singletonMap( + request.requestContext.setPerPartitionCircuitBreakerInfoHolder(Collections.singletonMap( "westus", createHealthContext(LocationHealthStatus.Healthy))); @@ -190,8 +187,11 @@ public void gatewayStatisticsRetainStateAtRecordTime() throws Exception { @Test(groups = {"unit"}) public void routingLookupInitializesEmptyStateWhenNoCircuitExists() throws Exception { DiagnosticsClientContext diagnosticsClientContext = Mockito.mock(DiagnosticsClientContext.class); - PerPartitionCircuitBreakerInfoHolder holder = new PerPartitionCircuitBreakerInfoHolder(); - RxDocumentServiceRequest request = createRequest(diagnosticsClientContext, holder); + RxDocumentServiceRequest request = RxDocumentServiceRequest.create( + diagnosticsClientContext, + OperationType.Read, + ResourceType.Document); + assertThat(request.requestContext.getPerPartitionCircuitBreakerInfoHolder()).isNull(); request.setResourceId("collectionRid"); PartitionKeyRange partitionKeyRange = new PartitionKeyRange("0", "AA", "BB"); request.requestContext.resolvedPartitionKeyRange = partitionKeyRange; @@ -214,7 +214,8 @@ public void routingLookupInitializesEmptyStateWhenNoCircuitExists() throws Excep assertThat(manager.getUnavailableRegionsForPartitionKeyRange(request, "collectionRid", partitionKeyRange)) .isEmpty(); - assertThat(holder.getPerPartitionCircuitBreakerInfoHolder()).isEmpty(); + assertThat(request.requestContext.getPerPartitionCircuitBreakerInfoHolder() + .getPerPartitionCircuitBreakerInfoHolder()).isEmpty(); ClientSideRequestStatistics statistics = new ClientSideRequestStatistics(diagnosticsClientContext); statistics.recordResponse(request, null, null); @@ -230,14 +231,8 @@ private static RxDocumentServiceRequest createRequest( diagnosticsClientContext, OperationType.Read, ResourceType.Document); - request.requestContext.setCrossRegionAvailabilityContext( - new CrossRegionAvailabilityContextForRxDocumentServiceRequest( - null, - null, - null, - new AtomicBoolean(false), - holder, - new PerPartitionAutomaticFailoverInfoHolder())); + request.requestContext.setPerPartitionCircuitBreakerInfoHolder( + holder.getPerPartitionCircuitBreakerInfoHolder()); return request; } diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 4a094db78ef4..93e547540b54 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -1,119 +1,13 @@ ## Release History -### 4.82.0-beta.1 (Unreleased) - -#### Features Added -* Enabled Gateway V2 (thin-client) data-plane routing by default for `Cosmos(Async)Client` instances configured with `gatewayMode` and HTTP/2, gated by an HTTP/2 connectivity probe with automatic fallback to Gateway V1. - See [PR 49437](https://github.com/Azure/azure-sdk-for-java/pull/49437) -* Added support for QueryPlan and Execute Stored Procedure requests to be routed to Gateway V2. - See [PR 47759](https://github.com/Azure/azure-sdk-for-java/pull/47759) - -#### Breaking Changes +### 4.76.1-hotfix (Unreleased) #### 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). * 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). -* Fixed thin-client (Gateway V2) queries with a prefix (partial) hierarchical partition key returning co-located documents from other logical partitions. - See PR [49688](https://github.com/Azure/azure-sdk-for-java/pull/49688). -* Fixed hedged requests losing request-scoped routing, timeout, authorization, throughput-control, and metadata state when cloning the original request. - See [PR 50069](https://github.com/Azure/azure-sdk-for-java/pull/50069). +* 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). #### Other Changes * Added per-region Per-Partition Circuit Breaker health and last failback outcome snapshots to `CosmosDiagnostics`, including structured failure reasons, and WARN logging for failback failures. - See [PR 50158](https://github.com/Azure/azure-sdk-for-java/pull/50158). -* Reduced memory footprint of deserialized `PartitionKeyRange` instances by stripping unused fields in the `PartitionKeyRange(ObjectNode)` constructor - See PR [49513](https://github.com/Azure/azure-sdk-for-java/pull/49513). -* Added bounded retries for transient "collection routing map / partition key range metadata not available" responses (HTTP 404 with sub-status `0`, `1003`, or `1013`) that can briefly occur right after a container is (re)created, improving the robustness of data-plane operations against the post-creation metadata-propagation race. As part of this change, when the routing map remains unavailable after retries an operation now fails with a `CosmosException` (HTTP 404, sub-status `1024` / `INCORRECT_CONTAINER_RID`) instead of an internal `IllegalStateException`. - See [PR 49639](https://github.com/Azure/azure-sdk-for-java/pull/49639). -* Reduced memory footprint and redundant `/pkranges` reads when multiple `CosmosClient` / `CosmosAsyncClient` instances in the same JVM are configured with the same service endpoint. Disable with system property `COSMOS.SHARED_PARTITION_KEY_RANGE_CACHE_ENABLED=false` if needed. - See [PR 49560](https://github.com/Azure/azure-sdk-for-java/pull/49560). - -### 4.81.0 (2026-06-08) - -#### Features Added -* Added support for creating Global Secondary Index (GSI) containers via `CosmosContainerProperties.setGlobalSecondaryIndexDefinition()` / `getGlobalSecondaryIndexDefinition()`, the new `CosmosGlobalSecondaryIndexDefinition` model, and the `CosmosGlobalSecondaryIndexBuildStatus` enum returned by `getStatus()`. - See [PR 48480](https://github.com/Azure/azure-sdk-for-java/pull/48480) -* Promoted the Full Fidelity Change Feed (AllVersionsAndDeletes) APIs to GA - See [PR 49283](https://github.com/Azure/azure-sdk-for-java/pull/49283) -* Enabled `ReadConsistencyStrategy` for Gateway V1 (compute gateway) and Gateway V2 (thin client proxy). Previously only supported in Direct mode. - See [PR 48787](https://github.com/Azure/azure-sdk-for-java/pull/48787) - -#### Bugs Fixed -* Fixed region name normalization for preferred and excluded regions — non-canonical inputs (e.g., `"westus3"`, `"WEST US 3"`) are now mapped to the canonical form. Also fixed a case-sensitive exclude-region check in PPCB reevaluate logic. - See [PR 49090](https://github.com/Azure/azure-sdk-for-java/pull/49090) -* Fixed `UnsupportedOperationException` when using `readManyByPartitionKeys` for empty pages. - See [PR 49311](https://github.com/Azure/azure-sdk-for-java/pull/49311) -* Fixed silent drift in `CosmosChangeFeedRequestOptions` when resuming from a continuation token via `byPage(savedContinuation)`. Previously only `maxPrefetchPageCount` and `throughputControlGroupName` were inherited onto the rebuilt impl; `endLSN`, `customSerializer`, `excludeRegions`, `readConsistencyStrategy`, `completeAfterAllCurrentChangesRetrieved`, and other caller-supplied configuration were silently dropped. All non-token-encoded fields are now propagated. - See [PR 49276](https://github.com/Azure/azure-sdk-for-java/pull/49276) -* Fixed HTTP/2 PING keepalive handler (introduced in [PR 49095](https://github.com/Azure/azure-sdk-for-java/pull/49095)) so it observes child-stream HEADERS/DATA reads via `Http2PingCloseRewrapHandler.channelReadComplete`, preventing spurious PINGs (and spurious closes) on connections actively serving requests through `Http2MultiplexHandler`. - -#### Other Changes -* Added HTTP/2 PING keepalive (default ON) for Gateway service endpoints to detect silently-broken connections. - See [PR 49095](https://github.com/Azure/azure-sdk-for-java/pull/49095) -* Replaced per-client `Schedulers.newSingle()` schedulers in `GlobalEndpointManager` and `GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker` with shared `BoundedElastic` schedulers in `CosmosSchedulers` to prevent thread count from scaling linearly with client/tenant count. - See [PR 49062](https://github.com/Azure/azure-sdk-for-java/pull/49062) -* Promoted the `ReadConsistencyStrategy` and `Http2ConnectionConfig` related `@Beta` APIs to GA. - See [PR 49345](https://github.com/Azure/azure-sdk-for-java/pull/49345) -* Fixed a sporadic `NullPointerException` in `JsonSerializable.getWithMapping` triggered by concurrent first-time calls to `DatabaseAccount.getConsistencyPolicy()` and its sibling lazy getters (`getReplicationPolicy`, `getSystemReplicationPolicy`, `getQueryEngineConfiguration`). The fix makes `JsonSerializable.propertyBag` `final`, closing an unsafe-publication race in the lazy-initialisation pattern. - See [Issue 49256](https://github.com/Azure/azure-sdk-for-java/issues/49256) and [PR #49258](https://github.com/Azure/azure-sdk-for-java/pull/49258) -* Changed 449 (`Retry With`) retries in Gateway V1 and Gateway V2 to be consistently orchestrated client-side. - See [PR 49332](https://github.com/Azure/azure-sdk-for-java/pull/49332) -* Added client-side fast-fail validation for `ReadConsistencyStrategy.GLOBAL_STRONG`: requests that specify `GLOBAL_STRONG` against an account whose default consistency is not `STRONG` are now rejected client-side with a `BadRequestException` (HTTP 400). - See [PR 48787](https://github.com/Azure/azure-sdk-for-java/pull/48787) - -### 4.80.0 (2026-05-01) - -#### Features Added -* Added support for Query Advisor feature - See [48160](https://github.com/Azure/azure-sdk-for-java/pull/48160) -* Added `additionalHeaders` support to allow setting additional headers (e.g., `x-ms-cosmos-workload-id`) that are sent with every request. - See [PR 48128](https://github.com/Azure/azure-sdk-for-java/pull/48128) -* Added `IGNORE_UNKNOWN_RNTBD_TOKENS` SDK capability flag and propagated SDK supported capabilities to barrier requests, enabling N-Region Synchronous Commit to function correctly with backends that return new RNTBD response tokens. - See [PR 48965](https://github.com/Azure/azure-sdk-for-java/pull/48965) -* Added support for change feed with `startFrom` point-in-time on merged partitions by enabling the `CHANGE_FEED_WITH_START_TIME_POST_MERGE` SDK capability. - See [PR 48752](https://github.com/Azure/azure-sdk-for-java/pull/48752) -* Added new `readManyByPartitionKeys` API on `CosmosAsyncContainer` / `CosmosContainer` to bulk-query all documents matching a list of partition key values with better efficiency than issuing individual queries. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) -* Added `CosmosReadManyByPartitionKeysRequestOptions` - a dedicated request-options type for `readManyByPartitionKeys` that exposes `setContinuationToken(String)` for resuming previous invocations and `setMaxConcurrentBatchPrefetch(int)` to bound per-call prefetch parallelism. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) -* Added `CosmosReadManyByPartitionKeysRequestOptions.setMaxBatchSize(Integer)` to set the max. number of partition keys used for a single batch. See [PR 48930](https://github.com/Azure/azure-sdk-for-java/pull/48930) -* Added `getCustomItemSerializer()` to `CosmosRequestContext` and `setCustomItemSerializer(CosmosItemSerializer)` to `CosmosRequestOptions` to allow overriding the custom item serializer via operation policies. - See [PR 48963](https://github.com/Azure/azure-sdk-for-java/pull/48963) - -#### Bugs Fixed -* Fixed `readMany` and `readAllItems` returning incorrect results on containers whose partition key path is nested (e.g. `/address/city`) due to malformed selector generation. - See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801) -* Fixed an issue where the throughput control `throughputQueryMono` was always subscribed even when `targetThroughput` is used (not `targetThroughputThreshold`), causing unnecessary `throughputSettings/read` permission requirement for AAD principals. - See [PR 48800](https://github.com/Azure/azure-sdk-for-java/pull/48800) -* Fixed JVM `` deadlock when multiple threads concurrently trigger Cosmos SDK class loading for the first time. - See [PR 48689](https://github.com/Azure/azure-sdk-for-java/pull/48689) -* Fixed an issue where `CustomItemSerializer` was incorrectly applied to internal SDK query pipeline structures (e.g., `OrderByRowResult`, `Document`), causing deserialization failures in ORDER BY, GROUP BY, aggregate, DISTINCT, and hybrid search queries. - See [PR 48811](https://github.com/Azure/azure-sdk-for-java/pull/48811) -* Fixed an issue where `SqlParameter` ignored the configured `CustomItemSerializer`, always using the internal default serializer instead. - See [PR 48811](https://github.com/Azure/azure-sdk-for-java/pull/48811) -* Fixed a `ClientTelemetry` static initialization failure when IMDS access is disabled, preventing `NoClassDefFoundError` during Cosmos client creation in non-Azure environments. - See [PR 48888](https://github.com/Azure/azure-sdk-for-java/pull/48888) -* Fixed an issue where Netty could log "An exceptionCaught() event was fired, and it reached at the tail of the pipeline" on HTTP/2 connections when the server resets idle TCP connections by adding an exception handler on the HTTP/2 parent channel to handle these connection-level exceptions more appropriately. - See [PR 48890](https://github.com/Azure/azure-sdk-for-java/pull/48890) -* Fixed an issue where `CustomItemSerializer` configured on `CosmosClientBuilder` was not honored for response deserialization in `CosmosAsyncContainer.upsertItem` when no request-level serializer was set. - See [PR 48962](https://github.com/Azure/azure-sdk-for-java/pull/48962) - -### 4.79.1 (2026-04-06) - -#### Bugs Fixed -* Fixing an NPE caused due to boxed Boolean conversion. - See [PR 48656](https://github.com/Azure/azure-sdk-for-java/pull/48656/) - -### 4.79.0 (2026-03-27) - -#### Features Added -* Added support for N-Region synchronous commit feature - See [PR 47757](https://github.com/Azure/azure-sdk-for-java/pull/47757) -* Added support for Query Advisor feature - See [48160](https://github.com/Azure/azure-sdk-for-java/pull/48160) -* Added `CosmosFullTextScoreScope` enum and `setFullTextScoreScope()` on `CosmosQueryRequestOptions` for controlling BM25 statistics scope in hybrid search queries. Supports `LOCAL` (scoped to target partitions) and `GLOBAL` (default, all partitions) scopes. See [PR 48431](https://github.com/Azure/azure-sdk-for-java/pull/48431) - -#### Bugs Fixed -* Fixed Remote Code Execution (RCE) vulnerability (CWE-502) by replacing Java deserialization with JSON-based serialization in `CosmosClientMetadataCachesSnapshot`, `AsyncCache`, and `DocumentCollection`. The metadata cache snapshot now uses Jackson for serialization/deserialization, eliminating the entire class of Java deserialization attacks. - [PR 47971](https://github.com/Azure/azure-sdk-for-java/pull/47971) -* Fixed `NullPointerException` in `DocumentQueryExecutionContextFactory.tryCacheQueryPlan` when executing hybrid search queries with a partition key filter. See [PR 48431](https://github.com/Azure/azure-sdk-for-java/pull/48431) -* Fixed `ConcurrentModificationException` in hybrid search component query execution caused by concurrent access to shared mutable state. See [PR 48431](https://github.com/Azure/azure-sdk-for-java/pull/48431) -* Fixed availability strategy for Gateway V2 (thin client) by ensuring `RegionalRoutingContext` identity is based only on the immutable gateway endpoint. - See [PR 48432](https://github.com/Azure/azure-sdk-for-java/pull/48432) -* Fixed an issue where `replaceItem` bypassed the `customItemSerializer`, serialising POJOs with the SDK's internal `ObjectMapper` instead of the user-configured one. - See [PR 48529](https://github.com/Azure/azure-sdk-for-java/pull/48529) -* Fixed `ClassCastException` (`ArrayNode cannot be cast to ObjectNode`) when executing `SELECT VALUE ... GROUP BY` queries. See - [PR 48507](https://github.com/Azure/azure-sdk-for-java/pull/48507) - -#### Other Changes -* Promoted the following `@Beta` APIs to GA: `CosmosContainerProperties.getFullTextPolicy()`/`setFullTextPolicy()`, `IndexingPolicy.getCosmosFullTextIndexes()`/`setCosmosFullTextIndexes()`. - See [PR 48538](https://github.com/Azure/azure-sdk-for-java/pull/48538) -* Added `appendUserAgentSuffix` method to `AsyncDocumentClient` to allow downstream libraries to append to the user agent after client construction. - See [PR 48505](https://github.com/Azure/azure-sdk-for-java/pull/48505) -* Added aggressive HTTP timeout policies for document operations routed to Gateway V2. - [PR 47879](https://github.com/Azure/azure-sdk-for-java/pull/47879) -* Added a default connect timeout of 5s for Gateway V2 (thin client) data-plane endpoints. - See [PR 48174](https://github.com/Azure/azure-sdk-for-java/pull/48174) -* Added system property `COSMOS.CONNECTION_ACQUIRE_TIMEOUT_IN_MS` and environment variable `COSMOS_CONNECTION_ACQUIRE_TIMEOUT_IN_MS` to allow overriding the gateway connection acquire timeout in milliseconds (default 45000ms). Minimum accepted value is 500ms. Replaces the previous `_IN_SECONDS` variants. - See [PR 48580](https://github.com/Azure/azure-sdk-for-java/pull/48580) -* Changed system property for thin client connection timeout from `COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_SECONDS` to `COSMOS.THINCLIENT_CONNECTION_TIMEOUT_IN_MS` (default 5000ms, minimum 500ms). - See [PR 48580](https://github.com/Azure/azure-sdk-for-java/pull/48580) - -### 4.78.0 (2026-02-10) - -#### Features Added -* Added shardKey support in `DedicatedGatewayRequestOptions` to allow specifying a shard key for dedicated gateway sharding support. - See [PR 47796](https://github.com/Azure/azure-sdk-for-java/pull/47796) - -#### Bugs Fixed -* Fixed an issue where `query plan` failed with `400` or query return empty result when `CosmosQueryRequestOptions` has partition key filter and partition key value contains non-ascii character. See [PR 47881](https://github.com/Azure/azure-sdk-for-java/pull/47881) -* Fixed an issue where operation failed with `400` when configured with pre-trigger or post-trigger with non-ascii character. Only impact for gateway mode. See [PR 47881](https://github.com/Azure/azure-sdk-for-java/pull/47881) - -#### Other Changes -* Added `x-ms-hub-region-processing-only` header to allow hub-region stickiness when 404 `READ SESSION NOT AVAILABLE` is hit for Single-Writer accounts. - [PR 47631](https://github.com/Azure/azure-sdk-for-java/pull/47631) - -### 4.77.0 (2026-01-26) - -#### Features Added -* Added `ChangeFeedProcessorOptions#setMaxLeasesToAcquirePerCycle(int)` to allow faster acquisition of unused/expired leases during scale-out and rolling deployments (default `0` preserves legacy behavior). - [47606](https://github.com/Azure/azure-sdk-for-java/pull/47606) -* Added the `QuantizerType` to the vectorIndexSpec: `product`/`spherical`. - [PR 47566](https://github.com/Azure/azure-sdk-for-java/pull/47566) - -#### Other Changes -* Remaps sub-status to 1003 for requests to child resources against non-existent container. - [PR 47604](https://github.com/Azure/azure-sdk-for-java/pull/47604) ### 4.76.0 (2025-12-09) diff --git a/sdk/cosmos/azure-cosmos/pom.xml b/sdk/cosmos/azure-cosmos/pom.xml index 2fd9dc103215..483d62ca6b96 100644 --- a/sdk/cosmos/azure-cosmos/pom.xml +++ b/sdk/cosmos/azure-cosmos/pom.xml @@ -13,7 +13,7 @@ Licensed under the MIT License. com.azure azure-cosmos - 4.76.0 + 4.76.1-hotfix Microsoft Azure SDK for SQL API of Azure Cosmos DB Service This Package contains Microsoft Azure Cosmos SDK (with Reactive Extension Reactor support) for Azure Cosmos DB SQL API jar diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientSideRequestStatistics.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientSideRequestStatistics.java index 50bd2bfce028..4f49f9d6e4d4 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientSideRequestStatistics.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientSideRequestStatistics.java @@ -173,19 +173,8 @@ public void recordResponse(RxDocumentServiceRequest request, StoreResultDiagnost this.approximateInsertionCountInBloomFilter = request.requestContext.getApproximateBloomFilterInsertionCount(); storeResponseStatistics.sessionTokenEvaluationResults = request.requestContext.getSessionTokenEvaluationResults(); storeResponseStatistics.perPartitionCircuitBreakerInfoHolder - = request.requestContext.getPerPartitionCircuitBreakerInfoHolder().snapshot(); - storeResponseStatistics.perPartitionAutomaticFailoverInfoHolder = request.requestContext.getPerPartitionFailoverContextHolder(); - - if (request.requestContext.getCrossRegionAvailabilityContext() != null) { - CrossRegionAvailabilityContextForRxDocumentServiceRequest crossRegionAvailabilityContextForRequest - = request.requestContext.getCrossRegionAvailabilityContext(); - - if (crossRegionAvailabilityContextForRequest.shouldAddHubRegionProcessingOnlyHeader()) { - storeResponseStatistics.isHubRegionProcessingOnly = "true"; - } else { - storeResponseStatistics.isHubRegionProcessingOnly = "false"; - } - } + = snapshot(request.requestContext.getPerPartitionCircuitBreakerInfoHolder()); + storeResponseStatistics.perPartitionFailoverInfoHolder = request.requestContext.getPerPartitionFailoverContextHolder(); if (request.requestContext.getEndToEndOperationLatencyPolicyConfig() != null) { storeResponseStatistics.e2ePolicyCfg = @@ -268,23 +257,8 @@ public void recordGatewayResponse( if (rxDocumentServiceRequest.requestContext != null) { gatewayStatistics.sessionTokenEvaluationResults = rxDocumentServiceRequest.requestContext.getSessionTokenEvaluationResults(); gatewayStatistics.perPartitionCircuitBreakerInfoHolder - = rxDocumentServiceRequest.requestContext.getPerPartitionCircuitBreakerInfoHolder().snapshot(); - gatewayStatistics.perPartitionAutomaticFailoverInfoHolder = rxDocumentServiceRequest.requestContext.getPerPartitionFailoverContextHolder(); - gatewayStatistics.isHubRegionProcessingOnly = "false"; - - CrossRegionAvailabilityContextForRxDocumentServiceRequest crossRegionAvailabilityContextForRequest - = rxDocumentServiceRequest.requestContext.getCrossRegionAvailabilityContext(); - - if (crossRegionAvailabilityContextForRequest != null) { - if (crossRegionAvailabilityContextForRequest.shouldAddHubRegionProcessingOnlyHeader()) { - gatewayStatistics.isHubRegionProcessingOnly = "true"; - } - } - - if (rxDocumentServiceRequest.requestContext.getEndToEndOperationLatencyPolicyConfig() != null) { - gatewayStatistics.e2ePolicyCfg = - rxDocumentServiceRequest.requestContext.getEndToEndOperationLatencyPolicyConfig().toString(); - } + = snapshot(rxDocumentServiceRequest.requestContext.getPerPartitionCircuitBreakerInfoHolder()); + gatewayStatistics.perPartitionFailoverInfoHolder = rxDocumentServiceRequest.requestContext.getPerPartitionFailoverContextHolder(); } } gatewayStatistics.statusCode = storeResponseDiagnostics.getStatusCode(); @@ -309,6 +283,12 @@ public void recordGatewayResponse( } } + private static PerPartitionCircuitBreakerInfoHolder snapshot( + PerPartitionCircuitBreakerInfoHolder holder) { + + return holder == null ? PerPartitionCircuitBreakerInfoHolder.EMPTY : holder.snapshot(); + } + public int getRequestPayloadSizeInBytes() { return this.requestPayloadSizeInBytes; } @@ -1056,7 +1036,7 @@ public void serialize(GatewayStatistics gatewayStatistics, this.writeNonEmptyStringSetField(jsonGenerator, "sessionTokenEvaluationResults", gatewayStatistics.getSessionTokenEvaluationResults()); this.writeNonNullObjectField(jsonGenerator, "ppcb", gatewayStatistics.getPerPartitionCircuitBreakerInfoHolder()); - this.writeNonNullObjectField(jsonGenerator, "perPartitionAutomaticFailoverInfoHolder", gatewayStatistics.getPerPartitionFailoverInfoHolder()); + this.writeNonNullObjectField(jsonGenerator, "perPartitionFailoverInfoHolder", gatewayStatistics.getPerPartitionFailoverInfoHolder()); this.writeNonNullStringField(jsonGenerator, "requestTCG", gatewayStatistics.getRequestThroughputControlGroupName()); this.writeNonNullStringField(jsonGenerator, "requestTCGConfig", gatewayStatistics.getRequestThroughputControlGroupConfig()); 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 8a9525c20970..17fe0de7dd3c 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 @@ -59,7 +59,9 @@ public class GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker impleme private final ConcurrentHashMap regionalRoutingContextToRegion; private final AtomicBoolean isClosed = new AtomicBoolean(false); private final AtomicBoolean isPartitionRecoveryTaskRunning = new AtomicBoolean(false); - private final AtomicReference partitionRecoveryDisposable = new AtomicReference<>(); + private final Scheduler partitionRecoveryScheduler = Schedulers.newSingle( + "partition-availability-staleness-check", + true); private final Logger failbackLogger; private final Object latestFailbackMessageByRegionLock = new Object(); private volatile Map latestFailbackMessageByRegion = Collections.emptyMap(); @@ -293,9 +295,12 @@ private void publishSnapshot( RxDocumentServiceRequest request, PartitionLevelLocationUnavailabilityInfo info) { + Map stateByRegion + = info == null ? Collections.emptyMap() : info.regionToLocationSpecificHealthContext; + request.requestContext.setPerPartitionCircuitBreakerInfoHolder(stateByRegion); request.requestContext.getPerPartitionCircuitBreakerInfoHolder() .setPerPartitionCircuitBreakerInfoHolder( - info == null ? Collections.emptyMap() : info.regionToLocationSpecificHealthContext, + stateByRegion, this.latestFailbackMessageByRegion); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/PerPartitionCircuitBreakerInfoHolder.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/PerPartitionCircuitBreakerInfoHolder.java index 46890bdd314f..b5df443271d4 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/PerPartitionCircuitBreakerInfoHolder.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/perPartitionCircuitBreaker/PerPartitionCircuitBreakerInfoHolder.java @@ -80,6 +80,8 @@ public void serialize(PerPartitionCircuitBreakerInfoHolder value, JsonGenerator } gen.writeEndObject(); + } else { + gen.writeNull(); } } } From 5c0a5c5bbded27a07127dfd2c270f5613c2944ab Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Fri, 28 Aug 2026 15:04:02 -0400 Subject: [PATCH 10/14] Restore PPCB hotfix regression coverage --- eng/versioning/version_client.txt | 2 +- ...titionEndpointManagerForPPCBUnitTests.java | 278 ++++++++++++ .../PerPartitionCircuitBreakerE2ETests.java | 408 ++++++++++++++++++ 3 files changed, 687 insertions(+), 1 deletion(-) diff --git a/eng/versioning/version_client.txt b/eng/versioning/version_client.txt index 552488a5703b..50e555b1d43d 100644 --- a/eng/versioning/version_client.txt +++ b/eng/versioning/version_client.txt @@ -104,7 +104,7 @@ com.azure:azure-core-test;1.27.0-beta.13;1.27.0-beta.14 com.azure:azure-core-tracing-opentelemetry;1.0.0-beta.61;1.0.0-beta.62 com.azure:azure-core-tracing-opentelemetry-samples;1.0.0-beta.1;1.0.0-beta.1 com.azure:azure-core-version-tests;1.0.0-beta.1;1.0.0-beta.1 -com.azure:azure-cosmos;4.75.0;4.76.0 +com.azure:azure-cosmos;4.75.0;4.76.1-hotfix com.azure:azure-cosmos-benchmark;4.0.1-beta.1;4.0.1-beta.1 com.azure.cosmos.spark:azure-cosmos-spark_3;0.0.1-beta.1;0.0.1-beta.1 com.azure.cosmos.spark:azure-cosmos-spark_3-5;0.0.1-beta.1;0.0.1-beta.1 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 11d14758f352..4a6f3266f2b3 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,11 +18,21 @@ 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.perPartitionCircuitBreaker.GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker; import com.azure.cosmos.implementation.perPartitionCircuitBreaker.LocationHealthStatus; import com.azure.cosmos.implementation.perPartitionCircuitBreaker.LocationSpecificHealthContext; import com.azure.cosmos.implementation.guava25.collect.ImmutableList; +import com.azure.cosmos.implementation.http.HttpClient; import com.azure.cosmos.implementation.routing.RegionalRoutingContext; +import com.fasterxml.jackson.databind.ObjectMapper; +import io.netty.channel.ConnectTimeoutException; import org.apache.commons.lang3.tuple.Pair; import org.mockito.Mockito; import org.slf4j.Logger; @@ -27,17 +40,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.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; 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; @@ -52,6 +77,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; @@ -121,6 +151,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 { @@ -1007,6 +1046,245 @@ 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) { + @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); + String diagnostics = new ObjectMapper().writeValueAsString( + request.requestContext.getPerPartitionCircuitBreakerInfoHolder()); + assertThat(diagnostics) + .contains("\"outcome\":\"Failed\"") + .contains("\"stage\":\"OPEN_CONNECTION_TASK\"") + .contains("\"type\":\"io.netty.channel.ConnectTimeoutException\"") + .contains("\"latestFailbackMessageByRegion\":{") + .contains("\"East US\":\"Refreshed replica is unavailable\""); + } else { + assertThat(ppcbManager.getUnavailableRegionsForPartitionKeyRange( + request, + collectionRid, + partitionKeyRange)).isEmpty(); + assertThat(request.requestContext.getPerPartitionCircuitBreakerInfoHolder() + .getPerPartitionCircuitBreakerInfoHolder() + .get("East US") + .getUnavailableSince()).isEqualTo(Instant.MAX); + assertThat(new ObjectMapper().writeValueAsString( + request.requestContext.getPerPartitionCircuitBreakerInfoHolder())) + .contains("\"outcome\":\"Succeeded\"") + .doesNotContain("\"failure\"", "\"latestFailbackMessageByRegion\""); + } + + assertThat(new ObjectMapper().writeValueAsString( + request.requestContext.getPerPartitionCircuitBreakerInfoHolder())) + .contains("\"lastAttemptedAt\":"); + + 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); + 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 4390b3f83c9c..ccdc2937825c 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 @@ -4,6 +4,7 @@ package com.azure.cosmos; import com.azure.cosmos.faultinjection.FaultInjectionTestBase; +import com.azure.cosmos.implementation.ClientSideRequestStatistics; import com.azure.cosmos.implementation.ConnectionPolicy; import com.azure.cosmos.implementation.DatabaseAccount; import com.azure.cosmos.implementation.DatabaseAccountLocation; @@ -13,6 +14,7 @@ import com.azure.cosmos.implementation.ImplementationBridgeHelpers; import com.azure.cosmos.implementation.OperationType; import com.azure.cosmos.implementation.PartitionKeyRange; +import com.azure.cosmos.implementation.ResourceType; import com.azure.cosmos.implementation.RxDocumentClientImpl; import com.azure.cosmos.implementation.TestConfigurations; import com.azure.cosmos.implementation.Utils; @@ -67,6 +69,7 @@ import java.time.Duration; import java.util.ArrayList; import java.util.Arrays; +import java.util.Collection; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; @@ -88,6 +91,9 @@ public class PerPartitionCircuitBreakerE2ETests extends FaultInjectionTestBase { private static final ImplementationBridgeHelpers.CosmosAsyncContainerHelper.CosmosAsyncContainerAccessor containerAccessor = ImplementationBridgeHelpers.CosmosAsyncContainerHelper.getCosmosAsyncContainerAccessor(); + private static final ImplementationBridgeHelpers.CosmosDiagnosticsHelper.CosmosDiagnosticsAccessor cosmosDiagnosticsAccessor + = ImplementationBridgeHelpers.CosmosDiagnosticsHelper.getCosmosDiagnosticsAccessor(); + private List writeRegions; private List readRegions; @@ -246,6 +252,7 @@ public void beforeClass() { DatabaseAccount databaseAccount = globalEndpointManager.getLatestDatabaseAccount(); this.writeRegions = new ArrayList<>(this.getAccountLevelLocationContext(databaseAccount, true).serviceOrderedWriteableRegions); + this.readRegions = new ArrayList<>(this.getAccountLevelLocationContext(databaseAccount, false).serviceOrderedReadableRegions); CosmosAsyncDatabase sharedAsyncDatabase = getSharedCosmosDatabase(testClient); CosmosAsyncContainer sharedMultiPartitionCosmosContainerWithIdAsPartitionKey = getSharedMultiPartitionCosmosContainerWithIdAsPartitionKey(testClient); @@ -3549,6 +3556,8 @@ private void execute( boolean hasReachedCircuitBreakingThreshold = false; int executionCountAfterCircuitBreakingThresholdBreached = 0; + boolean failbackExpected = false; + Set loggedPpcbDiagnosticsPhases = new HashSet<>(); List testObjects = operationInvocationParamsWrapper.testObjectsForDataPlaneOperationToWorkWith; PartitionKeyRangeWrapper partitionKeyRangeWrapper @@ -3563,6 +3572,8 @@ private void execute( } ResponseWrapper response = executeDataPlaneOperation.apply(operationInvocationParamsWrapper); + assertPpcbSnapshotsPopulated(response, PpcbDiagnosticsPhase.FAILURE, false); + logPpcbDiagnosticsOnce(response, PpcbDiagnosticsPhase.FAILURE, loggedPpcbDiagnosticsPhases); ConsecutiveExceptionBasedCircuitBreaker consecutiveExceptionBasedCircuitBreaker = globalPartitionEndpointManagerForPerPartitionCircuitBreaker.getConsecutiveExceptionBasedCircuitBreaker(); @@ -3588,6 +3599,14 @@ private void execute( if (executionCountAfterCircuitBreakingThresholdBreached > 1) { validateResponseInAbsenceOfFailures.accept(response); + failbackExpected |= assertPpcbSnapshotsPopulated( + response, + PpcbDiagnosticsPhase.POST_FAILOVER, + false); + logPpcbDiagnosticsOnce( + response, + PpcbDiagnosticsPhase.POST_FAILOVER, + loggedPpcbDiagnosticsPhases); } if (response.cosmosItemResponse != null) { @@ -3639,6 +3658,14 @@ private void execute( ResponseWrapper response = executeDataPlaneOperation.apply(operationInvocationParamsWrapper); validateResponseInAbsenceOfFailures.accept(response); + assertPpcbSnapshotsPopulated( + response, + PpcbDiagnosticsPhase.POST_FAILBACK, + failbackExpected); + logPpcbDiagnosticsOnce( + response, + PpcbDiagnosticsPhase.POST_FAILBACK, + loggedPpcbDiagnosticsPhases); if (response.cosmosItemResponse != null) { assertThat(response.cosmosItemResponse).isNotNull(); @@ -3676,6 +3703,193 @@ private void execute( } } + private static CosmosDiagnosticsContext getDiagnosticsContext(ResponseWrapper response) { + if (response.cosmosItemResponse != null) { + return response.cosmosItemResponse.getDiagnostics().getDiagnosticsContext(); + } else if (response.feedResponse != null) { + return response.feedResponse.getCosmosDiagnostics().getDiagnosticsContext(); + } else if (response.cosmosException != null) { + return response.cosmosException.getDiagnostics().getDiagnosticsContext(); + } else if (response.batchResponse != null) { + return response.batchResponse.getDiagnostics().getDiagnosticsContext(); + } + return null; + } + + private static void logPpcbDiagnosticsOnce( + ResponseWrapper response, + PpcbDiagnosticsPhase phase, + Set loggedPhases) { + + if (loggedPhases.add(phase)) { + CosmosDiagnosticsContext diagnosticsContext = getDiagnosticsContext(response); + if (diagnosticsContext != null) { + logger.info("PPCB CosmosDiagnostics [{}]: {}", phase.label, diagnosticsContext.toJson()); + } + } + } + + private static boolean assertPpcbSnapshotsPopulated( + ResponseWrapper response, + PpcbDiagnosticsPhase phase, + boolean failbackExpected) { + + CosmosDiagnosticsContext diagnosticsContext = getDiagnosticsContext(response); + assertThat(diagnosticsContext) + .as("Expected CosmosDiagnostics for %s", phase.label) + .isNotNull(); + assertThat(diagnosticsContext.getDiagnostics()) + .as("Expected diagnostics entries for %s", phase.label) + .isNotNull(); + + int applicableStatisticCount = 0; + List healthContexts = new ArrayList<>(); + for (CosmosDiagnostics cosmosDiagnostics : diagnosticsContext.getDiagnostics()) { + Collection statisticsCollection + = cosmosDiagnosticsAccessor.getClientSideRequestStatistics(cosmosDiagnostics); + if (statisticsCollection == null) { + continue; + } + + for (ClientSideRequestStatistics statistics : statisticsCollection) { + if (statistics == null) { + continue; + } + + for (ClientSideRequestStatistics.StoreResponseStatistics storeStatistics + : statistics.getResponseStatisticsList()) { + + if (isPpcbApplicableDataPlaneStatistic( + storeStatistics.getRequestResourceType(), + storeStatistics.getRequestOperationType())) { + + applicableStatisticCount++; + assertThat(storeStatistics.getPerPartitionCircuitBreakerInfoHolder()) + .as("Expected direct PPCB holder for %s", phase.label) + .isNotNull(); + Map stateByRegion + = storeStatistics.getPerPartitionCircuitBreakerInfoHolder() + .getPerPartitionCircuitBreakerInfoHolder(); + assertThat(stateByRegion) + .as("Expected populated direct PPCB snapshot for %s", phase.label) + .isNotNull(); + healthContexts.addAll(stateByRegion.values()); + } + } + + for (ClientSideRequestStatistics.GatewayStatistics gatewayStatistics + : statistics.getGatewayStatisticsList()) { + + if (isPpcbApplicableDataPlaneStatistic( + gatewayStatistics.getResourceType(), + gatewayStatistics.getOperationType())) { + + applicableStatisticCount++; + assertThat(gatewayStatistics.getPerPartitionCircuitBreakerInfoHolder()) + .as("Expected gateway PPCB holder for %s", phase.label) + .isNotNull(); + Map stateByRegion + = gatewayStatistics.getPerPartitionCircuitBreakerInfoHolder() + .getPerPartitionCircuitBreakerInfoHolder(); + assertThat(stateByRegion) + .as("Expected populated gateway PPCB snapshot for %s", phase.label) + .isNotNull(); + healthContexts.addAll(stateByRegion.values()); + } + } + } + } + + if (applicableStatisticCount == 0) { + assertThat(hasOnlyQueryPlanStatistics(diagnosticsContext)) + .as("Expected PPCB-applicable data-plane statistics or QueryPlan-only diagnostics for %s", phase.label) + .isTrue(); + } + + boolean unavailableRegionFound = false; + boolean successfulFailbackFound = false; + for (LocationSpecificHealthContext healthContext : healthContexts) { + if (healthContext.getLocationHealthStatus() == LocationHealthStatus.Unavailable) { + unavailableRegionFound = true; + if (phase == PpcbDiagnosticsPhase.POST_FAILOVER) { + assertThat(healthContext.getLastFailbackOutcome()) + .as("Failback must not have succeeded while the region remains unavailable") + .isNotEqualTo(LocationSpecificHealthContext.FailbackOutcome.Succeeded); + } + } + + if (healthContext.getLastFailbackOutcome() + == LocationSpecificHealthContext.FailbackOutcome.Succeeded) { + + successfulFailbackFound = true; + assertThat(healthContext.getLastFailbackAttemptTime()) + .as("Expected failback attempt timestamp after successful failback") + .isNotNull(); + assertThat(healthContext.getLocationHealthStatus()) + .as("Expected recovered region after successful failback") + .isIn(LocationHealthStatus.HealthyTentative, LocationHealthStatus.Healthy); + } + } + + if (phase == PpcbDiagnosticsPhase.POST_FAILBACK && failbackExpected) { + assertThat(successfulFailbackFound) + .as("Expected a successful failback outcome for a previously unavailable region") + .isTrue(); + } + + return unavailableRegionFound; + } + + private static boolean isPpcbApplicableDataPlaneStatistic( + ResourceType resourceType, + OperationType operationType) { + + return resourceType == ResourceType.Document && operationType != OperationType.QueryPlan; + } + + private static boolean hasOnlyQueryPlanStatistics(CosmosDiagnosticsContext diagnosticsContext) { + boolean queryPlanStatisticFound = false; + for (CosmosDiagnostics cosmosDiagnostics : diagnosticsContext.getDiagnostics()) { + Collection statisticsCollection + = cosmosDiagnosticsAccessor.getClientSideRequestStatistics(cosmosDiagnostics); + if (statisticsCollection == null) { + continue; + } + + for (ClientSideRequestStatistics statistics : statisticsCollection) { + if (statistics == null) { + continue; + } + + for (ClientSideRequestStatistics.StoreResponseStatistics storeStatistics + : statistics.getResponseStatisticsList()) { + + if (storeStatistics.getRequestResourceType() != ResourceType.Document) { + continue; + } + if (storeStatistics.getRequestOperationType() != OperationType.QueryPlan) { + return false; + } + queryPlanStatisticFound = true; + } + + for (ClientSideRequestStatistics.GatewayStatistics gatewayStatistics + : statistics.getGatewayStatisticsList()) { + + if (gatewayStatistics.getResourceType() != ResourceType.Document) { + continue; + } + if (gatewayStatistics.getOperationType() != OperationType.QueryPlan) { + return false; + } + queryPlanStatisticFound = true; + } + } + } + + return queryPlanStatisticFound; + } + private static int resolveTestObjectCountToBootstrapFrom(FaultInjectionOperationType faultInjectionOperationType, int opCount) { switch (faultInjectionOperationType) { case READ_ITEM: @@ -5249,6 +5463,18 @@ private enum QueryType { READ_MANY, READ_ALL } + private enum PpcbDiagnosticsPhase { + FAILURE("failed operation"), + POST_FAILOVER("post-failover operation"), + POST_FAILBACK("post-failback operation"); + + private final String label; + + PpcbDiagnosticsPhase(String label) { + this.label = label; + } + } + private static class AccountLevelLocationContext { private final List serviceOrderedReadableRegions; private final List serviceOrderedWriteableRegions; @@ -5264,4 +5490,186 @@ public AccountLevelLocationContext( this.regionNameToEndpoint = regionNameToEndpoint; } } + + @Test(groups = {"circuit-breaker-misc-direct"}, timeOut = 20 * TIMEOUT) + @SuppressWarnings("unchecked") + 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"); + } + + 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(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(); + assertThat(failedOverResponse.getDiagnostics().getDiagnosticsContext().getContactedRegionNames()) + .contains(this.readRegions.get(1).toLowerCase(Locale.ROOT)); + + 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(); + assertThat(recoveredResponse.getDiagnostics().getDiagnosticsContext().getContactedRegionNames()) + .containsExactly(this.readRegions.get(0).toLowerCase(Locale.ROOT)); + } 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); + } + } + } + + @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); + } } From 409c12c754cda15f8d33ebaad9e50e8da61ab0f5 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Sat, 29 Aug 2026 06:22:52 -0400 Subject: [PATCH 11/14] Backport customer workflow tests from PR #49568 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 16ce0941-555c-4190-8c4d-96c705086350 --- sdk/cosmos/azure-cosmos-tests/pom.xml | 54 ++ .../com/azure/cosmos/rx/TestSuiteBase.java | 4 +- ...erWorkflowAvailabilityFaultMatrixTest.java | 199 +++++++ ...stomerWorkflowChangeFeedProcessorTest.java | 219 ++++++++ ...ustomerWorkflowDaoStyleOperationsTest.java | 145 +++++ .../CustomerWorkflowHighE2ETimeoutTest.java | 250 +++++++++ .../CustomerWorkflowLatestCommittedTest.java | 171 ++++++ ...kflowPartitionLevelCircuitBreakerTest.java | 128 +++++ .../CustomerWorkflowRequestOptionsTest.java | 157 ++++++ .../CustomerWorkflowSessionTokenTest.java | 92 ++++ ...rWorkflowSingleMasterAvailabilityTest.java | 291 ++++++++++ .../CustomerWorkflowStoredProcedureTest.java | 134 +++++ .../customer/CustomerWorkflowTestBase.java | 499 ++++++++++++++++++ .../fi-customer-workflows-testng.xml | 38 ++ .../fi-sm-customer-workflows-testng.xml | 38 ++ ...fi-customer-workflows-platform-matrix.json | 42 ++ ...sm-customer-workflows-platform-matrix.json | 41 ++ sdk/cosmos/tests.yml | 64 +++ 18 files changed, 2564 insertions(+), 2 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowAvailabilityFaultMatrixTest.java create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowChangeFeedProcessorTest.java create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowDaoStyleOperationsTest.java create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowHighE2ETimeoutTest.java create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowLatestCommittedTest.java create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowPartitionLevelCircuitBreakerTest.java create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowRequestOptionsTest.java create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowSessionTokenTest.java create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowSingleMasterAvailabilityTest.java create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowStoredProcedureTest.java create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowTestBase.java create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/resources/fi-customer-workflows-testng.xml create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/resources/fi-sm-customer-workflows-testng.xml create mode 100644 sdk/cosmos/live-fi-customer-workflows-platform-matrix.json create mode 100644 sdk/cosmos/live-fi-sm-customer-workflows-platform-matrix.json diff --git a/sdk/cosmos/azure-cosmos-tests/pom.xml b/sdk/cosmos/azure-cosmos-tests/pom.xml index 5579b257c9c5..8ddba3d83ba2 100644 --- a/sdk/cosmos/azure-cosmos-tests/pom.xml +++ b/sdk/cosmos/azure-cosmos-tests/pom.xml @@ -665,6 +665,60 @@ Licensed under the MIT License. + + + fi-customer-workflows + + fi-customer-workflows + + + + + org.apache.maven.plugins + maven-failsafe-plugin + 3.5.3 + + + src/test/resources/fi-customer-workflows-testng.xml + + + true + 1 + 256 + paranoid + + + + + + + + + fi-sm-customer-workflows + + fi-sm-customer-workflows + + + + + org.apache.maven.plugins + maven-failsafe-plugin + 3.5.3 + + + src/test/resources/fi-sm-customer-workflows-testng.xml + + + true + 1 + 256 + paranoid + + + + + + multi-region diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java index d4e06ca7407b..05f91f3218d9 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java @@ -205,7 +205,7 @@ public CosmosAsyncDatabase getDatabase(String id) { @BeforeSuite(groups = {"thinclient", "fast", "long", "direct", "multi-region", "multi-master", "flaky-multi-master", "emulator", "emulator-vnext", "split", "query", "cfp-split", "circuit-breaker-misc-gateway", "circuit-breaker-misc-direct", - "circuit-breaker-read-all-read-many", "fi-multi-master", "long-emulator", "fi-thinclient-multi-region", "fi-thinclient-multi-master", "multi-region-strong"}, timeOut = SUITE_SETUP_TIMEOUT) + "circuit-breaker-read-all-read-many", "fi-multi-master", "fi-customer-workflows", "fi-sm-customer-workflows", "long-emulator", "fi-thinclient-multi-region", "fi-thinclient-multi-master", "multi-region-strong"}, timeOut = SUITE_SETUP_TIMEOUT) public void beforeSuite() { logger.info("beforeSuite Started"); @@ -223,7 +223,7 @@ public void beforeSuite() { @AfterSuite(groups = {"thinclient", "fast", "long", "direct", "multi-region", "multi-master", "flaky-multi-master", "emulator", "split", "query", "cfp-split", "circuit-breaker-misc-gateway", "circuit-breaker-misc-direct", - "circuit-breaker-read-all-read-many", "fi-multi-master", "long-emulator", "fi-thinclient-multi-region", "fi-thinclient-multi-master", "multi-region-strong"}, timeOut = SUITE_SHUTDOWN_TIMEOUT) + "circuit-breaker-read-all-read-many", "fi-multi-master", "fi-customer-workflows", "fi-sm-customer-workflows", "long-emulator", "fi-thinclient-multi-region", "fi-thinclient-multi-master", "multi-region-strong"}, timeOut = SUITE_SHUTDOWN_TIMEOUT) public void afterSuite() { logger.info("afterSuite Started"); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowAvailabilityFaultMatrixTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowAvailabilityFaultMatrixTest.java new file mode 100644 index 000000000000..176715cd5a2a --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowAvailabilityFaultMatrixTest.java @@ -0,0 +1,199 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.workflows.customer; + +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.CosmosDiagnosticsContext; +import com.azure.cosmos.CosmosException; +import com.azure.cosmos.TestObject; +import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.models.CosmosItemIdentity; +import com.azure.cosmos.models.CosmosItemRequestOptions; +import com.azure.cosmos.models.CosmosItemResponse; +import com.azure.cosmos.models.CosmosPatchOperations; +import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.CosmosReadManyRequestOptions; +import com.azure.cosmos.models.FeedResponse; +import com.azure.cosmos.test.faultinjection.FaultInjectionOperationType; +import com.azure.cosmos.test.faultinjection.FaultInjectionRule; +import com.azure.cosmos.test.faultinjection.FaultInjectionServerErrorType; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Factory; +import org.testng.annotations.Test; + +import java.util.Collections; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +public class CustomerWorkflowAvailabilityFaultMatrixTest extends CustomerWorkflowTestBase { + + @Factory(dataProvider = "clientBuildersWithSessionConsistency") + public CustomerWorkflowAvailabilityFaultMatrixTest(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + } + + @BeforeClass(groups = {"fi-customer-workflows"}, timeOut = SETUP_TIMEOUT) + public void beforeClass() { + initializeSharedSinglePartitionContainer("Customer availability fault workflow tests"); + } + + @AfterClass(groups = {"fi-customer-workflows"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) + public void afterClass() { + closeClient(); + } + + @DataProvider(name = "availabilityFaultScenarios") + public Object[][] availabilityFaultScenarios() { + return new Object[][]{ + {"read", FaultInjectionOperationType.READ_ITEM, FaultInjectionServerErrorType.GONE}, + {"read", FaultInjectionOperationType.READ_ITEM, FaultInjectionServerErrorType.TIMEOUT}, + {"read", FaultInjectionOperationType.READ_ITEM, FaultInjectionServerErrorType.READ_SESSION_NOT_AVAILABLE}, + {"read", FaultInjectionOperationType.READ_ITEM, FaultInjectionServerErrorType.INTERNAL_SERVER_ERROR}, + {"query", FaultInjectionOperationType.QUERY_ITEM, FaultInjectionServerErrorType.SERVICE_UNAVAILABLE}, + {"query", FaultInjectionOperationType.QUERY_ITEM, FaultInjectionServerErrorType.GONE}, + {"query", FaultInjectionOperationType.QUERY_ITEM, FaultInjectionServerErrorType.TIMEOUT}, + {"query", FaultInjectionOperationType.QUERY_ITEM, FaultInjectionServerErrorType.INTERNAL_SERVER_ERROR}, + {"readMany", FaultInjectionOperationType.QUERY_ITEM, FaultInjectionServerErrorType.GONE}, + {"readMany", FaultInjectionOperationType.QUERY_ITEM, FaultInjectionServerErrorType.READ_SESSION_NOT_AVAILABLE}, + {"readMany", FaultInjectionOperationType.QUERY_ITEM, FaultInjectionServerErrorType.TOO_MANY_REQUEST}, + {"create", FaultInjectionOperationType.CREATE_ITEM, FaultInjectionServerErrorType.INTERNAL_SERVER_ERROR}, + {"create", FaultInjectionOperationType.CREATE_ITEM, FaultInjectionServerErrorType.TOO_MANY_REQUEST}, + {"create", FaultInjectionOperationType.CREATE_ITEM, FaultInjectionServerErrorType.TIMEOUT}, + {"create", FaultInjectionOperationType.CREATE_ITEM, FaultInjectionServerErrorType.RETRY_WITH}, + {"create", FaultInjectionOperationType.CREATE_ITEM, FaultInjectionServerErrorType.PARTITION_IS_MIGRATING}, + {"upsert", FaultInjectionOperationType.UPSERT_ITEM, FaultInjectionServerErrorType.SERVICE_UNAVAILABLE}, + {"upsert", FaultInjectionOperationType.UPSERT_ITEM, FaultInjectionServerErrorType.PARTITION_IS_MIGRATING}, + {"upsert", FaultInjectionOperationType.UPSERT_ITEM, FaultInjectionServerErrorType.TOO_MANY_REQUEST}, + {"replace", FaultInjectionOperationType.REPLACE_ITEM, FaultInjectionServerErrorType.GONE}, + {"replace", FaultInjectionOperationType.REPLACE_ITEM, FaultInjectionServerErrorType.TIMEOUT}, + {"replace", FaultInjectionOperationType.REPLACE_ITEM, FaultInjectionServerErrorType.SERVICE_UNAVAILABLE}, + {"delete", FaultInjectionOperationType.DELETE_ITEM, FaultInjectionServerErrorType.SERVICE_UNAVAILABLE}, + {"delete", FaultInjectionOperationType.DELETE_ITEM, FaultInjectionServerErrorType.GONE}, + {"delete", FaultInjectionOperationType.DELETE_ITEM, FaultInjectionServerErrorType.TIMEOUT}, + {"patch", FaultInjectionOperationType.PATCH_ITEM, FaultInjectionServerErrorType.INTERNAL_SERVER_ERROR}, + {"patch", FaultInjectionOperationType.PATCH_ITEM, FaultInjectionServerErrorType.SERVICE_UNAVAILABLE}, + {"patch", FaultInjectionOperationType.PATCH_ITEM, FaultInjectionServerErrorType.GONE} + }; + } + + @Test(groups = {"fi-customer-workflows"}, dataProvider = "availabilityFaultScenarios", timeOut = TIMEOUT) + public void representativeDirectMultiMasterFaultWorkflow( + String operation, + FaultInjectionOperationType faultInjectionOperationType, + FaultInjectionServerErrorType errorType) { + + skipIfNotDirectMode("Customer availability fault workflow (direct multi-master)"); + + TestObject item = TestObject.create(); + if (!"create".equals(operation)) { + this.container.createItem(item).block(); + registerForCleanup(item); + } + + List faultRules = "readMany".equals(operation) + ? configureReadManyServerErrorRules(this.container, errorType, this.writableRegions.get(0), 1) + : Collections.singletonList(configureServerErrorRule( + this.container, + faultInjectionOperationType, + errorType, + this.writableRegions.get(0), + currentFaultInjectionConnectionType(), + 1)); + + try { + CosmosDiagnosticsContext diagnosticsContext = executeOperation(operation, item); + + assertFaultInjectedOperation(diagnosticsContext, faultRules); + assertThat(diagnosticsContext.getDuration()).isNotNull(); + } finally { + faultRules.forEach(FaultInjectionRule::disable); + } + } + + private CosmosDiagnosticsContext executeOperation(String operation, TestObject item) { + try { + if ("read".equals(operation)) { + CosmosItemResponse response = this.container + .readItem(item.getId(), partitionKey(item), new CosmosItemRequestOptions(), TestObject.class) + .block(); + + return response.getDiagnostics().getDiagnosticsContext(); + } + + if ("query".equals(operation)) { + FeedResponse response = this.container + .queryItems( + String.format("SELECT * FROM c WHERE c.id = '%s'", item.getId()), + new CosmosQueryRequestOptions().setQueryName("AvailabilityFaultWorkflowQuery"), + TestObject.class) + .byPage() + .blockFirst(); + + return response.getCosmosDiagnostics().getDiagnosticsContext(); + } + + if ("readMany".equals(operation)) { + FeedResponse response = this.container + .readMany( + Collections.singletonList(new CosmosItemIdentity(partitionKey(item), item.getId())), + new CosmosReadManyRequestOptions(), + TestObject.class) + .block(); + + return response.getCosmosDiagnostics().getDiagnosticsContext(); + } + + if ("upsert".equals(operation)) { + item.setStringProp("fault-upsert-" + item.getStringProp()); + CosmosItemResponse response = this.container + .upsertItem(item, new CosmosItemRequestOptions().setContentResponseOnWriteEnabled(true)) + .block(); + + return response.getDiagnostics().getDiagnosticsContext(); + } + + if ("replace".equals(operation)) { + item.setStringProp("fault-replace-" + item.getStringProp()); + CosmosItemResponse response = this.container + .replaceItem(item, item.getId(), partitionKey(item), new CosmosItemRequestOptions()) + .block(); + + return response.getDiagnostics().getDiagnosticsContext(); + } + + if ("delete".equals(operation)) { + CosmosItemResponse response = this.container + .deleteItem(item.getId(), partitionKey(item), new CosmosItemRequestOptions()) + .block(); + + return response.getDiagnostics().getDiagnosticsContext(); + } + + if ("patch".equals(operation)) { + CosmosItemResponse response = this.container + .patchItem( + item.getId(), + partitionKey(item), + CosmosPatchOperations.create().set("/stringProp", "fault-patch-" + item.getStringProp()), + TestObject.class) + .block(); + + return response.getDiagnostics().getDiagnosticsContext(); + } + + CosmosItemResponse response = this.container + .createItem(item, new CosmosItemRequestOptions().setContentResponseOnWriteEnabled(true)) + .block(); + + registerForCleanup(item); + return response.getDiagnostics().getDiagnosticsContext(); + } catch (CosmosException error) { + CosmosDiagnosticsContext diagnosticsContext = error.getDiagnostics().getDiagnosticsContext(); + assertThat(error.getStatusCode()).isGreaterThanOrEqualTo(HttpConstants.StatusCodes.BADREQUEST); + return diagnosticsContext; + } + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowChangeFeedProcessorTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowChangeFeedProcessorTest.java new file mode 100644 index 000000000000..2d5b88cacd65 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowChangeFeedProcessorTest.java @@ -0,0 +1,219 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.workflows.customer; + +import com.azure.cosmos.ChangeFeedProcessor; +import com.azure.cosmos.ChangeFeedProcessorBuilder; +import com.azure.cosmos.CosmosAsyncContainer; +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.TestObject; +import com.azure.cosmos.models.ChangeFeedProcessorItem; +import com.azure.cosmos.models.ChangeFeedProcessorOptions; +import com.azure.cosmos.models.ChangeFeedProcessorState; +import com.azure.cosmos.test.faultinjection.FaultInjectionOperationType; +import com.azure.cosmos.test.faultinjection.FaultInjectionRule; +import com.fasterxml.jackson.databind.JsonNode; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Factory; +import org.testng.annotations.Test; + +import java.time.Duration; +import java.util.Collections; +import java.util.List; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; + +import static org.assertj.core.api.Assertions.assertThat; + +public class CustomerWorkflowChangeFeedProcessorTest extends CustomerWorkflowTestBase { + + @Factory(dataProvider = "clientBuildersWithDirectTcpSession") + public CustomerWorkflowChangeFeedProcessorTest(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + } + + @BeforeClass(groups = {"fi-customer-workflows"}, timeOut = SETUP_TIMEOUT) + public void beforeClass() { + initializeSharedSinglePartitionContainer("Customer change feed processor workflow tests"); + } + + @AfterClass(groups = {"fi-customer-workflows"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) + public void afterClass() { + closeClient(); + } + + @Test(groups = {"fi-customer-workflows"}, timeOut = 2 * TIMEOUT) + public void latestVersionProcessorRestartResumesFromLeasesWorkflow() throws InterruptedException { + CosmosAsyncContainer feedContainer = createTemporaryContainer("customer-cfp-feed", "/mypk"); + CosmosAsyncContainer leaseContainer = createTemporaryContainer("customer-cfp-lease", "/id"); + ChangeFeedProcessor processor = null; + FaultInjectionRule readFeedDelayRule = null; + + try { + Set expectedIds = Collections.newSetFromMap(new ConcurrentHashMap()); + Set receivedIds = Collections.newSetFromMap(new ConcurrentHashMap()); + CountDownLatch initialLatch = new CountDownLatch(2); + + createFeedItem(feedContainer, expectedIds, "cfp-initial-1"); + createFeedItem(feedContainer, expectedIds, "cfp-initial-2"); + + // Use a single, stable lease prefix so the second processor instance resumes from the persisted + // continuation instead of reprocessing from the beginning - this validates a genuine restart. + String leasePrefix = "resume"; + processor = createLatestVersionProcessor(feedContainer, leaseContainer, expectedIds, receivedIds, initialLatch, leasePrefix); + processor.start().block(); + ChangeFeedProcessor initialProcessor = processor; + + assertThat(processor.isStarted()).isTrue(); + assertThat(initialLatch.await(30, TimeUnit.SECONDS)).isTrue(); + assertThat(receivedIds).containsAll(expectedIds); + + awaitCondition( + () -> hasAcquiredLeases(initialProcessor), + Duration.ofSeconds(20), + "Change feed processor did not acquire leases."); + + processor.stop().block(); + assertThat(processor.isStarted()).isFalse(); + + CountDownLatch restartLatch = new CountDownLatch(1); + TestObject restartedItem = createFeedItem(feedContainer, expectedIds, "cfp-restart"); + readFeedDelayRule = configureResponseDelayRule(feedContainer, FaultInjectionOperationType.READ_FEED_ITEM, Duration.ofMillis(100), 1); + + processor = createLatestVersionProcessor(feedContainer, leaseContainer, expectedIds, receivedIds, restartLatch, leasePrefix); + processor.start().block(); + + assertThat(processor.isStarted()).isTrue(); + assertThat(restartLatch.await(30, TimeUnit.SECONDS)).isTrue(); + assertThat(receivedIds).contains(restartedItem.getId()); + + // getEstimatedLag() is not supported for a latest-version processor; query the per-lease state + // (which exposes the estimated lag) via the supported getCurrentState() API instead. + List currentState = processor.getCurrentState().block(); + assertThat(currentState).isNotNull().isNotEmpty(); + assertThat(currentState).allSatisfy(state -> assertThat(state.getEstimatedLag()).isGreaterThanOrEqualTo(0)); + } finally { + if (readFeedDelayRule != null) { + readFeedDelayRule.disable(); + } + if (processor != null && processor.isStarted()) { + processor.stop().block(); + } + deleteTemporaryContainer(feedContainer); + deleteTemporaryContainer(leaseContainer); + } + } + + @Test(groups = {"fi-customer-workflows"}, timeOut = 2 * TIMEOUT) + public void latestVersionProcessorWithNewLeasePrefixReprocessesFromBeginningWorkflow() throws InterruptedException { + CosmosAsyncContainer feedContainer = createTemporaryContainer("customer-cfp-feed", "/mypk"); + CosmosAsyncContainer leaseContainer = createTemporaryContainer("customer-cfp-lease", "/id"); + ChangeFeedProcessor processor = null; + + try { + Set expectedIds = Collections.newSetFromMap(new ConcurrentHashMap()); + Set initialReceivedIds = Collections.newSetFromMap(new ConcurrentHashMap()); + CountDownLatch initialLatch = new CountDownLatch(2); + + createFeedItem(feedContainer, expectedIds, "cfp-initial-1"); + createFeedItem(feedContainer, expectedIds, "cfp-initial-2"); + + processor = createLatestVersionProcessor(feedContainer, leaseContainer, expectedIds, initialReceivedIds, initialLatch, "initial"); + processor.start().block(); + ChangeFeedProcessor initialProcessor = processor; + + assertThat(processor.isStarted()).isTrue(); + assertThat(initialLatch.await(30, TimeUnit.SECONDS)).isTrue(); + assertThat(initialReceivedIds).containsAll(expectedIds); + + awaitCondition( + () -> hasAcquiredLeases(initialProcessor), + Duration.ofSeconds(20), + "Change feed processor did not acquire leases."); + + processor.stop().block(); + assertThat(processor.isStarted()).isFalse(); + + // A different lease prefix creates a fresh lease set, so a from-beginning processor reprocesses all + // existing items. A separate received-id set is required because the original set already contains them. + Set reprocessedIds = Collections.newSetFromMap(new ConcurrentHashMap()); + CountDownLatch reprocessLatch = new CountDownLatch(expectedIds.size()); + + processor = createLatestVersionProcessor(feedContainer, leaseContainer, expectedIds, reprocessedIds, reprocessLatch, "fresh"); + processor.start().block(); + + assertThat(processor.isStarted()).isTrue(); + assertThat(reprocessLatch.await(30, TimeUnit.SECONDS)).isTrue(); + assertThat(reprocessedIds).containsAll(expectedIds); + + // getEstimatedLag() is not supported for a latest-version processor; query the per-lease state + // (which exposes the estimated lag) via the supported getCurrentState() API instead. + List currentState = processor.getCurrentState().block(); + assertThat(currentState).isNotNull().isNotEmpty(); + assertThat(currentState).allSatisfy(state -> assertThat(state.getEstimatedLag()).isGreaterThanOrEqualTo(0)); + } finally { + if (processor != null && processor.isStarted()) { + processor.stop().block(); + } + deleteTemporaryContainer(feedContainer); + deleteTemporaryContainer(leaseContainer); + } + } + + private static boolean hasAcquiredLeases(ChangeFeedProcessor processor) { + List currentState = processor.getCurrentState().block(); + return currentState != null && !currentState.isEmpty(); + } + + private TestObject createFeedItem(CosmosAsyncContainer feedContainer, Set expectedIds, String partitionKey) { + TestObject item = TestObject.create(partitionKey + "-" + UUID.randomUUID()); + feedContainer.createItem(item).block(); + expectedIds.add(item.getId()); + return item; + } + + private ChangeFeedProcessor createLatestVersionProcessor( + CosmosAsyncContainer feedContainer, + CosmosAsyncContainer leaseContainer, + Set expectedIds, + Set receivedIds, + CountDownLatch latch, + String leasePrefix) { + + return new ChangeFeedProcessorBuilder() + .hostName("customer-workflow-" + leasePrefix + "-" + UUID.randomUUID()) + .feedContainer(feedContainer) + .leaseContainer(leaseContainer) + .handleLatestVersionChanges(items -> recordLatestVersionItems(items, expectedIds, receivedIds, latch)) + .options(new ChangeFeedProcessorOptions() + .setStartFromBeginning(true) + .setFeedPollDelay(Duration.ofMillis(500)) + .setLeaseAcquireInterval(Duration.ofSeconds(1)) + .setLeaseRenewInterval(Duration.ofSeconds(2)) + .setLeaseExpirationInterval(Duration.ofSeconds(6)) + .setMaxItemCount(10) + .setLeasePrefix("customer-" + leasePrefix)) + .buildChangeFeedProcessor(); + } + + private static void recordLatestVersionItems( + List items, + Set expectedIds, + Set receivedIds, + CountDownLatch latch) { + + for (ChangeFeedProcessorItem item : items) { + JsonNode current = item.getCurrent(); + if (current != null && current.has("id")) { + String id = current.get("id").asText(); + if (expectedIds.contains(id) && receivedIds.add(id)) { + latch.countDown(); + } + } + } + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowDaoStyleOperationsTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowDaoStyleOperationsTest.java new file mode 100644 index 000000000000..a375d6540879 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowDaoStyleOperationsTest.java @@ -0,0 +1,145 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.workflows.customer; + +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.CosmosItemSerializer; +import com.azure.cosmos.TestObject; +import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.models.CosmosBatch; +import com.azure.cosmos.models.CosmosBatchResponse; +import com.azure.cosmos.models.CosmosBulkExecutionOptions; +import com.azure.cosmos.models.CosmosBulkOperationResponse; +import com.azure.cosmos.models.CosmosBulkOperations; +import com.azure.cosmos.models.CosmosItemOperation; +import com.azure.cosmos.models.CosmosItemRequestOptions; +import com.azure.cosmos.models.CosmosItemResponse; +import com.azure.cosmos.models.CosmosPatchOperations; +import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.FeedResponse; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Factory; +import org.testng.annotations.Test; +import reactor.core.publisher.Flux; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.UUID; + +import static org.assertj.core.api.Assertions.assertThat; + +public class CustomerWorkflowDaoStyleOperationsTest extends CustomerWorkflowTestBase { + + @Factory(dataProvider = "clientBuildersWithDirectTcpSession") + public CustomerWorkflowDaoStyleOperationsTest(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + } + + @BeforeClass(groups = {"fi-customer-workflows"}, timeOut = SETUP_TIMEOUT) + public void beforeClass() { + initializeSharedSinglePartitionContainer("Customer DAO-style workflow tests", true); + } + + @AfterClass(groups = {"fi-customer-workflows"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) + public void afterClass() { + closeClient(); + } + + @Test(groups = {"fi-customer-workflows"}, timeOut = TIMEOUT) + public void crudReadAllPatchBatchAndBulkWorkflow() { + List excludedRegions = excludeFirstWritableRegion(); + TestObject item = TestObject.create(); + + CosmosItemRequestOptions createOptions = new CosmosItemRequestOptions() + .setKeywordIdentifiers(Collections.singleton("workflow-crud-create")) + .setExcludedRegions(excludedRegions) + .setCustomItemSerializer(CosmosItemSerializer.DEFAULT_SERIALIZER) + .setContentResponseOnWriteEnabled(true); + + CosmosItemResponse createResponse = this.container + .createItem(item, createOptions) + .block(); + + assertThat(createResponse).isNotNull(); + registerForCleanup(item); + assertThat(createResponse.getStatusCode()).isEqualTo(HttpConstants.StatusCodes.CREATED); + assertKeywordIdentifier(createResponse.getDiagnostics().getDiagnosticsContext(), "workflow-crud-create"); + assertDidNotContactExcludedRegions(createResponse.getDiagnostics().getDiagnosticsContext(), excludedRegions); + + CosmosItemResponse readResponse = this.container + .readItem(item.getId(), partitionKey(item), new CosmosItemRequestOptions().setExcludedRegions(excludedRegions), TestObject.class) + .block(); + + assertThat(readResponse).isNotNull(); + assertThat(readResponse.getItem()).isEqualTo(item); + + FeedResponse readAllResponse = this.container + .readAllItems( + partitionKey(item), + new CosmosQueryRequestOptions() + .setExcludedRegions(excludedRegions) + .setCustomItemSerializer(CosmosItemSerializer.DEFAULT_SERIALIZER), + TestObject.class) + .byPage() + .blockFirst(); + + assertThat(readAllResponse).isNotNull(); + assertThat(readAllResponse.getResults()).extracting(TestObject::getId).contains(item.getId()); + assertExcludedRegions(readAllResponse.getCosmosDiagnostics().getDiagnosticsContext(), excludedRegions); + + CosmosPatchOperations patchOperations = CosmosPatchOperations.create() + .set("/stringProp", "patched-" + item.getStringProp()); + + CosmosItemResponse patchResponse = this.container + .patchItem(item.getId(), partitionKey(item), patchOperations, TestObject.class) + .block(); + + assertThat(patchResponse).isNotNull(); + assertThat(patchResponse.getStatusCode()).isEqualTo(HttpConstants.StatusCodes.OK); + assertThat(patchResponse.getItem().getStringProp()).startsWith("patched-"); + + String batchPk = "batch-" + UUID.randomUUID(); + TestObject batchItem = TestObject.create(batchPk); + CosmosBatch batch = CosmosBatch.createCosmosBatch(partitionKey(batchItem)); + batch.createItemOperation(batchItem); + batch.readItemOperation(batchItem.getId()); + + CosmosBatchResponse batchResponse = this.container.executeCosmosBatch(batch).block(); + + assertThat(batchResponse).isNotNull(); + registerForCleanup(batchItem); + assertThat(batchResponse.isSuccessStatusCode()).isTrue(); + assertThat(batchResponse.size()).isEqualTo(2); + assertThat(batchResponse.getDiagnostics()).isNotNull(); + + TestObject bulkItem = TestObject.create(); + this.container.createItem(bulkItem).block(); + registerForCleanup(bulkItem); + CosmosPatchOperations bulkPatchOperations = CosmosPatchOperations.create() + .set("/stringProp", "bulk-patched-" + bulkItem.getStringProp()); + + List bulkOperations = new ArrayList<>(); + bulkOperations.add(CosmosBulkOperations.getReadItemOperation(bulkItem.getId(), partitionKey(bulkItem))); + bulkOperations.add(CosmosBulkOperations.getPatchItemOperation(bulkItem.getId(), partitionKey(bulkItem), bulkPatchOperations)); + + CosmosBulkExecutionOptions bulkExecutionOptions = new CosmosBulkExecutionOptions() + .setMaxMicroBatchSize(2) + .setExcludedRegions(excludedRegions) + .setKeywordIdentifiers(Collections.singleton("workflow-bulk")); + + List> bulkResponses = this.container + .executeBulkOperations(Flux.fromIterable(bulkOperations), bulkExecutionOptions) + .collectList() + .block(); + + assertThat(bulkResponses).isNotNull(); + assertThat(bulkResponses).hasSize(2); + assertThat(bulkResponses).allSatisfy(response -> { + assertThat(response.getException()).isNull(); + assertThat(response.getResponse().getStatusCode()).isIn(HttpConstants.StatusCodes.OK, HttpConstants.StatusCodes.CREATED); + assertThat(response.getResponse().getCosmosDiagnostics()).isNotNull(); + }); + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowHighE2ETimeoutTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowHighE2ETimeoutTest.java new file mode 100644 index 000000000000..8d2d2bd4d892 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowHighE2ETimeoutTest.java @@ -0,0 +1,250 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.workflows.customer; + +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.CosmosDiagnosticsContext; +import com.azure.cosmos.CosmosEndToEndOperationLatencyPolicyConfig; +import com.azure.cosmos.CosmosEndToEndOperationLatencyPolicyConfigBuilder; +import com.azure.cosmos.CosmosException; +import com.azure.cosmos.TestObject; +import com.azure.cosmos.ThresholdBasedAvailabilityStrategy; +import com.azure.cosmos.implementation.ImplementationBridgeHelpers; +import com.azure.cosmos.models.CosmosBatch; +import com.azure.cosmos.models.CosmosBatchRequestOptions; +import com.azure.cosmos.models.CosmosBatchResponse; +import com.azure.cosmos.models.CosmosItemIdentity; +import com.azure.cosmos.models.CosmosItemRequestOptions; +import com.azure.cosmos.models.CosmosItemResponse; +import com.azure.cosmos.models.CosmosPatchItemRequestOptions; +import com.azure.cosmos.models.CosmosPatchOperations; +import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.CosmosReadManyRequestOptions; +import com.azure.cosmos.models.FeedResponse; +import com.azure.cosmos.test.faultinjection.FaultInjectionOperationType; +import com.azure.cosmos.test.faultinjection.FaultInjectionRule; +import com.azure.cosmos.test.faultinjection.FaultInjectionServerErrorType; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Factory; +import org.testng.annotations.Test; +import org.testng.SkipException; + +import java.time.Duration; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +public class CustomerWorkflowHighE2ETimeoutTest extends CustomerWorkflowTestBase { + + @Factory(dataProvider = "clientBuildersWithDirectTcpSession") + public CustomerWorkflowHighE2ETimeoutTest(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + } + + @BeforeClass(groups = {"fi-customer-workflows"}, timeOut = SETUP_TIMEOUT) + public void beforeClass() { + initializeSharedSinglePartitionContainer("Customer high E2E timeout workflow tests"); + } + + @AfterClass(groups = {"fi-customer-workflows"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) + public void afterClass() { + closeClient(); + } + + @DataProvider(name = "timeoutWorkflowOperations") + public Object[][] timeoutWorkflowOperations() { + return new Object[][]{ + {"create", FaultInjectionOperationType.CREATE_ITEM}, + {"read", FaultInjectionOperationType.READ_ITEM}, + {"query", FaultInjectionOperationType.QUERY_ITEM}, + {"readMany", FaultInjectionOperationType.QUERY_ITEM}, + {"upsert", FaultInjectionOperationType.UPSERT_ITEM}, + {"batch", FaultInjectionOperationType.BATCH_ITEM}, + {"patch", FaultInjectionOperationType.PATCH_ITEM} + }; + } + + @Test(groups = {"fi-customer-workflows"}, dataProvider = "timeoutWorkflowOperations", timeOut = 2 * TIMEOUT) + public void responseDelayWithAvailabilityStrategyWorkflow(String operation, FaultInjectionOperationType faultInjectionOperationType) { + TestObject item = TestObject.create(); + if (!"create".equals(operation)) { + this.container.createItem(item).block(); + registerForCleanup(item); + } + + CosmosEndToEndOperationLatencyPolicyConfig e2ePolicy = new CosmosEndToEndOperationLatencyPolicyConfigBuilder(Duration.ofSeconds(4)) + .availabilityStrategy(new ThresholdBasedAvailabilityStrategy(Duration.ofMillis(100), Duration.ofMillis(200))) + .build(); + + // readMany resolves to a point read for a single item, so the QUERY_ITEM data-provider value alone would not + // exercise the fault - inject the delay for both the point-read and query operation types. + List delayRules = new ArrayList<>(); + if ("readMany".equals(operation)) { + delayRules.add(configureResponseDelayRule(this.container, FaultInjectionOperationType.READ_ITEM, Duration.ofMillis(1500), 1)); + delayRules.add(configureResponseDelayRule(this.container, FaultInjectionOperationType.QUERY_ITEM, Duration.ofMillis(1500), 1)); + } else { + delayRules.add(configureResponseDelayRule(this.container, faultInjectionOperationType, Duration.ofMillis(1500), 1)); + } + + try { + CosmosDiagnosticsContext diagnosticsContext = executeWithE2EPolicy(operation, item, e2ePolicy); + + assertFaultInjectedOperation(diagnosticsContext, delayRules); + assertThat(diagnosticsContext.getDuration()).isLessThan(Duration.ofSeconds(10)); + } finally { + delayRules.forEach(FaultInjectionRule::disable); + } + } + + @Test(groups = {"fi-customer-workflows"}, timeOut = 2 * TIMEOUT) + public void partitionMigratingFaultWithE2EPolicyWorkflow() { + TestObject item = TestObject.create(); + this.container.createItem(item).block(); + registerForCleanup(item); + + CosmosEndToEndOperationLatencyPolicyConfig e2ePolicy = new CosmosEndToEndOperationLatencyPolicyConfigBuilder(Duration.ofSeconds(4)) + .availabilityStrategy(new ThresholdBasedAvailabilityStrategy(Duration.ofMillis(100), Duration.ofMillis(200))) + .build(); + + FaultInjectionRule migratingRule = configureServerErrorRule( + this.container, + FaultInjectionOperationType.READ_ITEM, + FaultInjectionServerErrorType.PARTITION_IS_MIGRATING, + 1); + + try { + CosmosDiagnosticsContext diagnosticsContext = executeWithE2EPolicy("read", item, e2ePolicy); + + assertFaultInjectedOperation(diagnosticsContext, migratingRule); + assertThat(diagnosticsContext.getDuration()).isLessThan(Duration.ofSeconds(10)); + } finally { + migratingRule.disable(); + } + } + + private CosmosDiagnosticsContext executeWithE2EPolicy( + String operation, + TestObject item, + CosmosEndToEndOperationLatencyPolicyConfig e2ePolicy) { + + try { + if ("create".equals(operation)) { + TestObject createdItem = TestObject.create(); + CosmosItemRequestOptions options = new CosmosItemRequestOptions() + .setContentResponseOnWriteEnabled(true) + .setCosmosEndToEndOperationLatencyPolicyConfig(e2ePolicy); + + CosmosItemResponse response = this.container + .createItem(createdItem, options) + .block(); + + registerForCleanup(createdItem); + return response.getDiagnostics().getDiagnosticsContext(); + } + + if ("read".equals(operation)) { + CosmosItemRequestOptions options = new CosmosItemRequestOptions() + .setCosmosEndToEndOperationLatencyPolicyConfig(e2ePolicy); + + return this.container + .readItem(item.getId(), partitionKey(item), options, TestObject.class) + .block() + .getDiagnostics() + .getDiagnosticsContext(); + } + + if ("query".equals(operation)) { + CosmosQueryRequestOptions options = new CosmosQueryRequestOptions() + .setCosmosEndToEndOperationLatencyPolicyConfig(e2ePolicy) + .setQueryName("HighE2ETimeoutWorkflowQuery"); + + FeedResponse response = this.container + .queryItems(String.format("SELECT * FROM c WHERE c.id = '%s'", item.getId()), options, TestObject.class) + .byPage() + .blockFirst(); + + return response.getCosmosDiagnostics().getDiagnosticsContext(); + } + + if ("readMany".equals(operation)) { + CosmosReadManyRequestOptions options = new CosmosReadManyRequestOptions() + .setCosmosEndToEndOperationLatencyPolicyConfig(e2ePolicy); + + FeedResponse response = this.container + .readMany(Collections.singletonList(new CosmosItemIdentity(partitionKey(item), item.getId())), options, TestObject.class) + .block(); + + return response.getCosmosDiagnostics().getDiagnosticsContext(); + } + + if ("upsert".equals(operation)) { + item.setStringProp("timeout-upsert-" + item.getStringProp()); + CosmosItemRequestOptions options = new CosmosItemRequestOptions() + .setContentResponseOnWriteEnabled(true) + .setCosmosEndToEndOperationLatencyPolicyConfig(e2ePolicy); + + return this.container + .upsertItem(item, options) + .block() + .getDiagnostics() + .getDiagnosticsContext(); + } + + if ("batch".equals(operation)) { + TestObject batchItem = TestObject.create("timeout-batch"); + CosmosBatch batch = CosmosBatch.createCosmosBatch(partitionKey(batchItem)); + batch.createItemOperation(batchItem); + batch.readItemOperation(batchItem.getId()); + + CosmosBatchRequestOptions batchOptions = new CosmosBatchRequestOptions(); + setBatchEndToEndOperationLatencyPolicyConfig(batchOptions, e2ePolicy); + + CosmosBatchResponse response = this.container.executeCosmosBatch(batch, batchOptions).block(); + + registerForCleanup(batchItem); + return response.getDiagnostics().getDiagnosticsContext(); + } + + CosmosPatchItemRequestOptions options = new CosmosPatchItemRequestOptions(); + options.setCosmosEndToEndOperationLatencyPolicyConfig(e2ePolicy); + + CosmosItemResponse response = this.container + .patchItem( + item.getId(), + partitionKey(item), + CosmosPatchOperations.create().set("/stringProp", "timeout-patched-" + item.getStringProp()), + options, + TestObject.class) + .block(); + + return response.getDiagnostics().getDiagnosticsContext(); + } catch (CosmosException error) { + return error.getDiagnostics().getDiagnosticsContext(); + } + } + + private static void setBatchEndToEndOperationLatencyPolicyConfig( + CosmosBatchRequestOptions batchOptions, + CosmosEndToEndOperationLatencyPolicyConfig e2ePolicy) { + + Object accessor = ImplementationBridgeHelpers.CosmosBatchRequestOptionsHelper + .getCosmosBatchRequestOptionsAccessor(); + try { + Method setter = accessor.getClass().getMethod( + "setEndToEndOperationLatencyPolicyConfig", + CosmosBatchRequestOptions.class, + CosmosEndToEndOperationLatencyPolicyConfig.class); + setter.invoke(accessor, batchOptions, e2ePolicy); + } catch (NoSuchMethodException error) { + throw new SkipException("Batch end-to-end latency policy is unavailable in this historical SDK."); + } catch (IllegalAccessException | InvocationTargetException error) { + throw new AssertionError("Unable to configure the batch end-to-end latency policy.", error); + } + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowLatestCommittedTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowLatestCommittedTest.java new file mode 100644 index 000000000000..a5eeaca9822d --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowLatestCommittedTest.java @@ -0,0 +1,171 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.workflows.customer; + +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.CosmosDiagnosticsContext; +import com.azure.cosmos.CosmosException; +import com.azure.cosmos.ReadConsistencyStrategy; +import com.azure.cosmos.TestObject; +import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.models.CosmosChangeFeedRequestOptions; +import com.azure.cosmos.models.CosmosItemIdentity; +import com.azure.cosmos.models.CosmosItemRequestOptions; +import com.azure.cosmos.models.CosmosItemResponse; +import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.CosmosReadManyRequestOptions; +import com.azure.cosmos.models.FeedRange; +import com.azure.cosmos.models.FeedResponse; +import com.azure.cosmos.test.faultinjection.FaultInjectionOperationType; +import com.azure.cosmos.test.faultinjection.FaultInjectionRule; +import com.azure.cosmos.test.faultinjection.FaultInjectionServerErrorType; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Factory; +import org.testng.annotations.Test; + +import java.util.Collections; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +public class CustomerWorkflowLatestCommittedTest extends CustomerWorkflowTestBase { + + @Factory(dataProvider = "clientBuildersWithSessionConsistency") + public CustomerWorkflowLatestCommittedTest(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + } + + @BeforeClass(groups = {"fi-customer-workflows"}, timeOut = SETUP_TIMEOUT) + public void beforeClass() { + initializeSharedSinglePartitionContainer("Customer latest-committed workflow tests", true); + } + + @AfterClass(groups = {"fi-customer-workflows"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) + public void afterClass() { + closeClient(); + } + + @Test(groups = {"fi-customer-workflows"}, timeOut = TIMEOUT) + public void latestCommittedAndExcludedRegionsFlowAcrossReadOperations() { + List excludedRegions = excludeFirstWritableRegion(); + TestObject item = TestObject.create(); + + CosmosItemResponse createResponse = this.container + .createItem(item, new CosmosItemRequestOptions().setExcludedRegions(excludedRegions)) + .block(); + + assertThat(createResponse).isNotNull(); + registerForCleanup(item); + CosmosDiagnosticsContext createDiagnostics = createResponse.getDiagnostics().getDiagnosticsContext(); + assertThat(createResponse.getStatusCode()).isEqualTo(HttpConstants.StatusCodes.CREATED); + assertThat(createDiagnostics.getEffectiveReadConsistencyStrategy()).isEqualTo(ReadConsistencyStrategy.DEFAULT); + assertExcludedRegions(createDiagnostics, excludedRegions); + assertDidNotContactExcludedRegions(createDiagnostics, excludedRegions); + + CosmosItemRequestOptions readOptions = new CosmosItemRequestOptions() + .setExcludedRegions(excludedRegions) + .setKeywordIdentifiers(Collections.singleton("latest-committed-read")) + .setReadConsistencyStrategy(ReadConsistencyStrategy.LATEST_COMMITTED); + + CosmosItemResponse readResponse = this.container + .readItem(item.getId(), partitionKey(item), readOptions, TestObject.class) + .block(); + + assertThat(readResponse).isNotNull(); + CosmosDiagnosticsContext readDiagnostics = readResponse.getDiagnostics().getDiagnosticsContext(); + assertThat(readResponse.getStatusCode()).isEqualTo(HttpConstants.StatusCodes.OK); + assertThat(readDiagnostics.getEffectiveReadConsistencyStrategy()).isEqualTo(ReadConsistencyStrategy.LATEST_COMMITTED); + assertThat(readDiagnostics.getTotalRequestCharge()).isGreaterThan(0); + assertKeywordIdentifier(readDiagnostics, "latest-committed-read"); + assertExcludedRegions(readDiagnostics, excludedRegions); + assertDidNotContactExcludedRegions(readDiagnostics, excludedRegions); + + CosmosQueryRequestOptions queryOptions = new CosmosQueryRequestOptions() + .setExcludedRegions(excludedRegions) + .setReadConsistencyStrategy(ReadConsistencyStrategy.LATEST_COMMITTED) + .setQueryName("LatestCommittedCustomerWorkflowQuery"); + + FeedResponse queryResponse = this.container + .queryItems(String.format("SELECT * FROM c WHERE c.id = '%s'", item.getId()), queryOptions, TestObject.class) + .byPage() + .blockFirst(); + + assertThat(queryResponse).isNotNull(); + assertThat(queryResponse.getResults()).hasSize(1); + CosmosDiagnosticsContext queryDiagnostics = queryResponse.getCosmosDiagnostics().getDiagnosticsContext(); + assertThat(queryDiagnostics.getEffectiveReadConsistencyStrategy()).isEqualTo(ReadConsistencyStrategy.LATEST_COMMITTED); + assertExcludedRegions(queryDiagnostics, excludedRegions); + + CosmosReadManyRequestOptions readManyOptions = new CosmosReadManyRequestOptions() + .setExcludedRegions(excludedRegions) + .setReadConsistencyStrategy(ReadConsistencyStrategy.LATEST_COMMITTED); + + FeedResponse readManyResponse = this.container + .readMany(Collections.singletonList(new CosmosItemIdentity(partitionKey(item), item.getId())), readManyOptions, TestObject.class) + .block(); + + assertThat(readManyResponse).isNotNull(); + assertThat(readManyResponse.getResults()).hasSize(1); + CosmosDiagnosticsContext readManyDiagnostics = readManyResponse.getCosmosDiagnostics().getDiagnosticsContext(); + assertThat(readManyDiagnostics.getEffectiveReadConsistencyStrategy()).isEqualTo(ReadConsistencyStrategy.LATEST_COMMITTED); + assertExcludedRegions(readManyDiagnostics, excludedRegions); + assertDidNotContactExcludedRegions(readManyDiagnostics, excludedRegions); + + CosmosChangeFeedRequestOptions changeFeedOptions = CosmosChangeFeedRequestOptions + .createForProcessingFromBeginning(FeedRange.forLogicalPartition(partitionKey(item))) + .setReadConsistencyStrategy(ReadConsistencyStrategy.LATEST_COMMITTED) + .setExcludedRegions(excludedRegions); + + FeedResponse changeFeedResponse = this.container + .queryChangeFeed(changeFeedOptions, TestObject.class) + .byPage() + .blockFirst(); + + assertThat(changeFeedResponse) + .as("change feed query should return at least one page before reading diagnostics") + .isNotNull(); + CosmosDiagnosticsContext changeFeedDiagnostics = changeFeedResponse.getCosmosDiagnostics().getDiagnosticsContext(); + assertThat(changeFeedDiagnostics.getEffectiveReadConsistencyStrategy()).isEqualTo(ReadConsistencyStrategy.LATEST_COMMITTED); + assertExcludedRegions(changeFeedDiagnostics, excludedRegions); + } + + @Test(groups = {"fi-customer-workflows"}, timeOut = TIMEOUT) + public void latestCommittedReadWithRegionalLeaseNotFoundFault() { + TestObject item = TestObject.create(); + this.container.createItem(item).block(); + registerForCleanup(item); + + FaultInjectionRule leaseNotFoundRule = configureServerErrorRule( + this.container, + FaultInjectionOperationType.READ_ITEM, + FaultInjectionServerErrorType.LEASE_NOT_FOUND, + this.writableRegions.get(0), + currentFaultInjectionConnectionType(), + 1); + + try { + CosmosItemRequestOptions readOptions = new CosmosItemRequestOptions() + .setReadConsistencyStrategy(ReadConsistencyStrategy.LATEST_COMMITTED) + .setKeywordIdentifiers(Collections.singleton("latest-committed-fault-read")); + + CosmosDiagnosticsContext diagnosticsContext; + try { + CosmosItemResponse readResponse = this.container + .readItem(item.getId(), partitionKey(item), readOptions, TestObject.class) + .block(); + + assertThat(readResponse).isNotNull(); + diagnosticsContext = readResponse.getDiagnostics().getDiagnosticsContext(); + } catch (CosmosException error) { + diagnosticsContext = error.getDiagnostics().getDiagnosticsContext(); + } + + assertThat(diagnosticsContext).isNotNull(); + assertThat(diagnosticsContext.getEffectiveReadConsistencyStrategy()).isEqualTo(ReadConsistencyStrategy.LATEST_COMMITTED); + assertFaultInjectedOperation(diagnosticsContext, leaseNotFoundRule); + } finally { + leaseNotFoundRule.disable(); + } + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowPartitionLevelCircuitBreakerTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowPartitionLevelCircuitBreakerTest.java new file mode 100644 index 000000000000..8ffaf5571295 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowPartitionLevelCircuitBreakerTest.java @@ -0,0 +1,128 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.workflows.customer; + +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.CosmosDiagnosticsContext; +import com.azure.cosmos.CosmosEndToEndOperationLatencyPolicyConfig; +import com.azure.cosmos.CosmosEndToEndOperationLatencyPolicyConfigBuilder; +import com.azure.cosmos.CosmosException; +import com.azure.cosmos.TestObject; +import com.azure.cosmos.ThresholdBasedAvailabilityStrategy; +import com.azure.cosmos.models.CosmosItemRequestOptions; +import com.azure.cosmos.models.CosmosItemResponse; +import com.azure.cosmos.models.CosmosPatchItemRequestOptions; +import com.azure.cosmos.models.CosmosPatchOperations; +import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.FeedResponse; +import com.azure.cosmos.test.faultinjection.FaultInjectionOperationType; +import com.azure.cosmos.test.faultinjection.FaultInjectionRule; +import com.azure.cosmos.test.faultinjection.FaultInjectionServerErrorType; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Factory; +import org.testng.annotations.Test; + +import java.time.Duration; + +import static org.assertj.core.api.Assertions.assertThat; + +public class CustomerWorkflowPartitionLevelCircuitBreakerTest extends CustomerWorkflowTestBase { + + @Factory(dataProvider = "clientBuildersWithDirectTcpSession") + public CustomerWorkflowPartitionLevelCircuitBreakerTest(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + } + + @BeforeClass(groups = {"fi-customer-workflows"}, timeOut = SETUP_TIMEOUT) + public void beforeClass() { + initializeSharedSinglePartitionContainer("Customer PCLB workflow tests"); + } + + @AfterClass(groups = {"fi-customer-workflows"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) + public void afterClass() { + closeClient(); + } + + @Test(groups = {"fi-customer-workflows"}, timeOut = 2 * TIMEOUT) + public void pointOperationCircuitBreakerAndQueryPlanWorkflow() { + TestObject item = TestObject.create(); + this.container.createItem(item).block(); + registerForCleanup(item); + + CosmosEndToEndOperationLatencyPolicyConfig e2ePolicy = new CosmosEndToEndOperationLatencyPolicyConfigBuilder(Duration.ofSeconds(3)) + .availabilityStrategy(new ThresholdBasedAvailabilityStrategy(Duration.ofMillis(100), Duration.ofMillis(200))) + .build(); + + FaultInjectionRule readFaultRule = configureServerErrorRule( + this.container, + FaultInjectionOperationType.READ_ITEM, + FaultInjectionServerErrorType.SERVICE_UNAVAILABLE, + 1); + + try { + CosmosDiagnosticsContext readDiagnostics = readWithPolicy(item, e2ePolicy); + + assertFaultInjectedOperation(readDiagnostics, readFaultRule); + } finally { + readFaultRule.disable(); + } + + CosmosDiagnosticsContext queryDiagnostics = queryWithPolicy(item, e2ePolicy); + assertThat(queryDiagnostics).isNotNull(); + assertThat(queryDiagnostics.getStatusCode()).isBetween(200, 599); + assertThat(queryDiagnostics.getContactedRegionNames()).isNotNull(); + assertThat(queryDiagnostics.toJson()).contains("queryPlanDiagnosticsContext"); + + CosmosPatchItemRequestOptions patchOptions = new CosmosPatchItemRequestOptions(); + patchOptions.setCosmosEndToEndOperationLatencyPolicyConfig(e2ePolicy); + CosmosItemResponse patchResponse = this.container + .patchItem( + item.getId(), + partitionKey(item), + CosmosPatchOperations.create().set("/stringProp", "pclb-patched-" + item.getStringProp()), + patchOptions, + TestObject.class) + .block(); + + assertThat(patchResponse).isNotNull(); + assertThat(patchResponse.getDiagnostics()).isNotNull(); + } + + private CosmosDiagnosticsContext readWithPolicy(TestObject item, CosmosEndToEndOperationLatencyPolicyConfig e2ePolicy) { + try { + CosmosItemRequestOptions options = new CosmosItemRequestOptions() + .setCosmosEndToEndOperationLatencyPolicyConfig(e2ePolicy); + + return this.container + .readItem(item.getId(), partitionKey(item), options, TestObject.class) + .block() + .getDiagnostics() + .getDiagnosticsContext(); + } catch (CosmosException error) { + return error.getDiagnostics().getDiagnosticsContext(); + } + } + + private CosmosDiagnosticsContext queryWithPolicy(TestObject item, CosmosEndToEndOperationLatencyPolicyConfig e2ePolicy) { + try { + CosmosQueryRequestOptions queryOptions = new CosmosQueryRequestOptions() + .setCosmosEndToEndOperationLatencyPolicyConfig(e2ePolicy) + .setQueryName("PclbCustomerWorkflowQuery"); + + // ORDER BY forces the gateway query-plan round-trip so the queryPlanDiagnosticsContext is always present, + // independent of single-partition / ServiceInterop query-plan optimizations. + FeedResponse response = this.container + .queryItems( + String.format("SELECT * FROM c WHERE c.id = '%s' ORDER BY c.id", item.getId()), + queryOptions, + TestObject.class) + .byPage() + .blockFirst(); + + return response.getCosmosDiagnostics().getDiagnosticsContext(); + } catch (CosmosException error) { + return error.getDiagnostics().getDiagnosticsContext(); + } + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowRequestOptionsTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowRequestOptionsTest.java new file mode 100644 index 000000000000..f46b517813a6 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowRequestOptionsTest.java @@ -0,0 +1,157 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.workflows.customer; + +import com.azure.cosmos.ConsistencyLevel; +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.CosmosDiagnosticsContext; +import com.azure.cosmos.ReadConsistencyStrategy; +import com.azure.cosmos.TestObject; +import com.azure.cosmos.implementation.OverridableRequestOptions; +import com.azure.cosmos.models.CosmosItemIdentity; +import com.azure.cosmos.models.CosmosItemRequestOptions; +import com.azure.cosmos.models.CosmosItemResponse; +import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.CosmosReadManyRequestOptions; +import com.azure.cosmos.models.FeedResponse; +import com.azure.cosmos.models.PartitionKey; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Factory; +import org.testng.annotations.Test; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +public class CustomerWorkflowRequestOptionsTest extends CustomerWorkflowTestBase { + @Factory(dataProvider = "clientBuildersWithDirectTcpSession") + public CustomerWorkflowRequestOptionsTest(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + } + + @BeforeClass(groups = {"fi-customer-workflows"}, timeOut = SETUP_TIMEOUT) + public void beforeClass() { + initializeSharedSinglePartitionContainer("Customer workflow request option tests", true); + } + + @AfterClass(groups = {"fi-customer-workflows"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) + public void afterClass() { + closeClient(); + } + + @Test(groups = {"fi-customer-workflows"}, timeOut = TIMEOUT) + public void excludedRegionAndKeywordIdentifiersFlowAcrossOperations() { + String excludedRegion = this.writableRegions.get(0); + List excludedRegions = Collections.singletonList(excludedRegion); + TestObject item = TestObject.create(); + + CosmosItemRequestOptions createOptions = new CosmosItemRequestOptions() + .setKeywordIdentifiers(Collections.singleton("customer-create")) + .setContentResponseOnWriteEnabled(true) + .setExcludedRegions(excludedRegions); + + CosmosItemResponse createResponse = this.container + .createItem(item, createOptions) + .block(); + + assertThat(createResponse).isNotNull(); + registerForCleanup(item); + assertThat(createResponse.getStatusCode()).isEqualTo(201); + assertKeywordIdentifier(createResponse.getDiagnostics().getDiagnosticsContext(), "customer-create"); + assertExcludedRegions(createResponse.getDiagnostics().getDiagnosticsContext(), excludedRegions); + assertDidNotContactExcludedRegions(createResponse.getDiagnostics().getDiagnosticsContext(), excludedRegions); + + CosmosItemRequestOptions readOptions = new CosmosItemRequestOptions() + .setKeywordIdentifiers(Collections.singleton("customer-read")) + .setExcludedRegions(excludedRegions) + .setReadConsistencyStrategy(ReadConsistencyStrategy.LATEST_COMMITTED); + + CosmosItemResponse readResponse = this.container + .readItem(item.getId(), new PartitionKey(item.getMypk()), readOptions, TestObject.class) + .block(); + + assertThat(readResponse).isNotNull(); + CosmosDiagnosticsContext readDiagnostics = readResponse.getDiagnostics().getDiagnosticsContext(); + assertThat(readResponse.getStatusCode()).isEqualTo(200); + assertThat(readDiagnostics.getEffectiveReadConsistencyStrategy()).isEqualTo(ReadConsistencyStrategy.LATEST_COMMITTED); + assertKeywordIdentifier(readDiagnostics, "customer-read"); + assertExcludedRegions(readDiagnostics, excludedRegions); + assertDidNotContactExcludedRegions(readDiagnostics, excludedRegions); + + CosmosQueryRequestOptions queryOptions = new CosmosQueryRequestOptions() + .setKeywordIdentifiers(Collections.singleton("customer-query")) + .setExcludedRegions(excludedRegions) + .setConsistencyLevel(ConsistencyLevel.EVENTUAL) + .setQueryMetricsEnabled(true) + .setQueryName("CustomerWorkflowQuery"); + + String query = String.format("SELECT * FROM c WHERE c.id = '%s'", item.getId()); + FeedResponse queryResponse = this.container + .queryItems(query, queryOptions, TestObject.class) + .byPage() + .blockFirst(); + + assertThat(queryResponse).isNotNull(); + assertThat(queryResponse.getResults()).hasSize(1); + CosmosDiagnosticsContext queryDiagnostics = queryResponse.getCosmosDiagnostics().getDiagnosticsContext(); + assertKeywordIdentifier(queryDiagnostics, "customer-query"); + assertExcludedRegions(queryDiagnostics, excludedRegions); + OverridableRequestOptions queryRequestOptions = getRequestOptions(queryDiagnostics); + assertThat(queryRequestOptions.getConsistencyLevel()).isEqualTo(ConsistencyLevel.EVENTUAL); + assertThat(queryRequestOptions.isQueryMetricsEnabled()).isTrue(); + assertThat(queryRequestOptions.getQueryNameOrDefault(null)).isEqualTo("CustomerWorkflowQuery"); + + CosmosReadManyRequestOptions readManyOptions = new CosmosReadManyRequestOptions() + .setKeywordIdentifiers(Collections.singleton("customer-read-many")) + .setExcludedRegions(excludedRegions) + .setReadConsistencyStrategy(ReadConsistencyStrategy.LATEST_COMMITTED); + + FeedResponse readManyResponse = this.container + .readMany( + Arrays.asList(new CosmosItemIdentity(new PartitionKey(item.getMypk()), item.getId())), + readManyOptions, + TestObject.class) + .block(); + + assertThat(readManyResponse).isNotNull(); + assertThat(readManyResponse.getResults()).hasSize(1); + CosmosDiagnosticsContext readManyDiagnostics = readManyResponse.getCosmosDiagnostics().getDiagnosticsContext(); + assertThat(readManyDiagnostics.getEffectiveReadConsistencyStrategy()).isEqualTo(ReadConsistencyStrategy.LATEST_COMMITTED); + assertKeywordIdentifier(readManyDiagnostics, "customer-read-many"); + assertExcludedRegions(readManyDiagnostics, excludedRegions); + assertDidNotContactExcludedRegions(readManyDiagnostics, excludedRegions); + + item.setStringProp("updated-" + item.getStringProp()); + CosmosItemRequestOptions upsertOptions = new CosmosItemRequestOptions() + .setKeywordIdentifiers(Collections.singleton("customer-upsert")) + .setExcludedRegions(excludedRegions) + .setContentResponseOnWriteEnabled(true); + + CosmosItemResponse upsertResponse = this.container + .upsertItem(item, upsertOptions) + .block(); + + assertThat(upsertResponse).isNotNull(); + assertThat(upsertResponse.getStatusCode()).isEqualTo(200); + assertKeywordIdentifier(upsertResponse.getDiagnostics().getDiagnosticsContext(), "customer-upsert"); + assertExcludedRegions(upsertResponse.getDiagnostics().getDiagnosticsContext(), excludedRegions); + assertDidNotContactExcludedRegions(upsertResponse.getDiagnostics().getDiagnosticsContext(), excludedRegions); + + CosmosItemRequestOptions deleteOptions = new CosmosItemRequestOptions() + .setKeywordIdentifiers(Collections.singleton("customer-delete")) + .setExcludedRegions(excludedRegions); + + CosmosItemResponse deleteResponse = this.container + .deleteItem(item.getId(), new PartitionKey(item.getMypk()), deleteOptions) + .block(); + + assertThat(deleteResponse).isNotNull(); + assertThat(deleteResponse.getStatusCode()).isEqualTo(204); + assertKeywordIdentifier(deleteResponse.getDiagnostics().getDiagnosticsContext(), "customer-delete"); + assertExcludedRegions(deleteResponse.getDiagnostics().getDiagnosticsContext(), excludedRegions); + assertDidNotContactExcludedRegions(deleteResponse.getDiagnostics().getDiagnosticsContext(), excludedRegions); + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowSessionTokenTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowSessionTokenTest.java new file mode 100644 index 000000000000..2ad18a9c09a7 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowSessionTokenTest.java @@ -0,0 +1,92 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.workflows.customer; + +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.CosmosException; +import com.azure.cosmos.TestObject; +import com.azure.cosmos.implementation.ConsistencyTestsBase; +import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.implementation.ISessionToken; +import com.azure.cosmos.implementation.SessionTokenHelper; +import com.azure.cosmos.implementation.Utils; +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.models.CosmosItemIdentity; +import com.azure.cosmos.models.CosmosItemResponse; +import com.azure.cosmos.models.FeedResponse; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Factory; +import org.testng.annotations.Test; + +import java.util.ArrayList; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.fail; + +public class CustomerWorkflowSessionTokenTest extends CustomerWorkflowTestBase { + + @Factory(dataProvider = "clientBuildersWithDirectTcpSession") + public CustomerWorkflowSessionTokenTest(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + } + + @BeforeClass(groups = {"fi-customer-workflows"}, timeOut = SETUP_TIMEOUT) + public void beforeClass() { + initializeSharedSinglePartitionContainer("Customer session-token workflow tests"); + } + + @AfterClass(groups = {"fi-customer-workflows"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) + public void afterClass() { + closeClient(); + } + + @Test(groups = {"fi-customer-workflows"}, timeOut = TIMEOUT) + public void readManyWithAdvancedSessionTokenReturnsReadSessionNotAvailable() throws Exception { + List itemIdentities = new ArrayList<>(); + String lastSessionToken = null; + + for (int index = 0; index < 3; index++) { + TestObject item = TestObject.create("session-token-workflow"); + CosmosItemResponse createResponse = this.container.createItem(item).block(); + + assertThat(createResponse).isNotNull(); + registerForCleanup(item); + lastSessionToken = createResponse.getSessionToken(); + itemIdentities.add(new CosmosItemIdentity(partitionKey(item), item.getId())); + } + + FeedResponse validReadManyResponse = this.container + .readMany(itemIdentities, lastSessionToken, TestObject.class) + .block(); + + assertThat(validReadManyResponse).isNotNull(); + assertThat(validReadManyResponse.getResults()).hasSize(3); + + String advancedSessionToken = advanceSessionToken(lastSessionToken); + + try { + this.container + .readMany(itemIdentities, advancedSessionToken, TestObject.class) + .block(); + + fail("Should have hit read session not available error."); + } catch (Exception error) { + CosmosException cosmosException = Utils.as(error, CosmosException.class); + + assertThat(cosmosException).isNotNull(); + assertThat(cosmosException.getStatusCode()).isEqualTo(HttpConstants.StatusCodes.NOTFOUND); + assertThat(cosmosException.getSubStatusCode()).isEqualTo(HttpConstants.SubStatusCodes.READ_SESSION_NOT_AVAILABLE); + assertThat(cosmosException.getDiagnostics()).isNotNull(); + } + } + + private static String advanceSessionToken(String originalSessionToken) throws Exception { + String[] tokenParts = StringUtils.split(originalSessionToken, ":"); + ISessionToken sessionToken = SessionTokenHelper.parse(tokenParts[1]); + ISessionToken modifiedSessionToken = ConsistencyTestsBase.createSessionToken(sessionToken, sessionToken.getLSN() + 1000000); + + return tokenParts[0] + ":" + modifiedSessionToken.convertToString(); + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowSingleMasterAvailabilityTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowSingleMasterAvailabilityTest.java new file mode 100644 index 000000000000..d3249ae10940 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowSingleMasterAvailabilityTest.java @@ -0,0 +1,291 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.workflows.customer; + +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.CosmosDiagnosticsContext; +import com.azure.cosmos.CosmosEndToEndOperationLatencyPolicyConfig; +import com.azure.cosmos.CosmosEndToEndOperationLatencyPolicyConfigBuilder; +import com.azure.cosmos.CosmosException; +import com.azure.cosmos.ReadConsistencyStrategy; +import com.azure.cosmos.TestObject; +import com.azure.cosmos.ThresholdBasedAvailabilityStrategy; +import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.models.CosmosItemRequestOptions; +import com.azure.cosmos.models.CosmosItemResponse; +import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.test.faultinjection.FaultInjectionOperationType; +import com.azure.cosmos.test.faultinjection.FaultInjectionRule; +import com.azure.cosmos.test.faultinjection.FaultInjectionServerErrorType; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Factory; +import org.testng.annotations.Test; + +import java.time.Duration; +import java.util.List; +import java.util.Locale; +import java.util.Set; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; + +public class CustomerWorkflowSingleMasterAvailabilityTest extends CustomerWorkflowTestBase { + + @Factory(dataProvider = "clientBuildersWithSessionConsistency") + public CustomerWorkflowSingleMasterAvailabilityTest(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + } + + @BeforeClass(groups = {"fi-sm-customer-workflows"}, timeOut = SETUP_TIMEOUT) + public void beforeClass() { + initializeSharedSingleWriteMultiRegionContainer("Customer single-master workflow tests"); + } + + @AfterClass(groups = {"fi-sm-customer-workflows"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) + public void afterClass() { + closeClient(); + } + + @Test(groups = {"fi-sm-customer-workflows"}, timeOut = TIMEOUT) + public void excludedReadableRegionRoutesReadToRemainingReadableRegion() { + TestObject item = TestObject.create(); + this.container.createItem(item).block(); + registerForCleanup(item); + + List excludedRegions = excludeFirstReadableRegion(); + CosmosItemRequestOptions readOptions = new CosmosItemRequestOptions() + .setExcludedRegions(excludedRegions) + .setReadConsistencyStrategy(ReadConsistencyStrategy.LATEST_COMMITTED); + + // Excluding the preferred readable region forces the read onto the remaining readable region, which may + // lag behind the just-completed write. Retry until cross-region replication catches up before asserting. + CosmosItemResponse readResponse = readWithReplicationRetry(item, readOptions); + + assertThat(readResponse).isNotNull(); + CosmosDiagnosticsContext diagnosticsContext = readResponse.getDiagnostics().getDiagnosticsContext(); + assertThat(readResponse.getStatusCode()).isEqualTo(HttpConstants.StatusCodes.OK); + assertThat(diagnosticsContext.getEffectiveReadConsistencyStrategy()).isEqualTo(ReadConsistencyStrategy.LATEST_COMMITTED); + assertExcludedRegions(diagnosticsContext, excludedRegions); + assertDidNotContactExcludedRegions(diagnosticsContext, excludedRegions); + } + + @Test(groups = {"fi-sm-customer-workflows"}, timeOut = TIMEOUT) + public void readFaultInPreferredReadableRegionCanUseRemoteReadableRegion() { + TestObject item = TestObject.create(); + this.container.createItem(item).block(); + registerForCleanup(item); + + FaultInjectionRule readSessionNotAvailableRule = configureServerErrorRule( + this.container, + FaultInjectionOperationType.READ_ITEM, + FaultInjectionServerErrorType.READ_SESSION_NOT_AVAILABLE, + this.readableRegions.get(0), + currentFaultInjectionConnectionType(), + 1); + + CosmosEndToEndOperationLatencyPolicyConfig e2ePolicy = new CosmosEndToEndOperationLatencyPolicyConfigBuilder(Duration.ofSeconds(5)) + .availabilityStrategy(new ThresholdBasedAvailabilityStrategy(Duration.ofMillis(100), Duration.ofMillis(200))) + .build(); + + try { + CosmosItemRequestOptions readOptions = new CosmosItemRequestOptions() + .setCosmosEndToEndOperationLatencyPolicyConfig(e2ePolicy); + + CosmosDiagnosticsContext diagnosticsContext = readWithDiagnostics(item, readOptions); + + assertThat(diagnosticsContext).isNotNull(); + assertThat(readSessionNotAvailableRule.getHitCount()) + .as("the injected read-session-not-available fault should have been hit in the preferred readable region") + .isGreaterThanOrEqualTo(1); + assertThat(diagnosticsContext.getStatusCode()).isBetween(HttpConstants.StatusCodes.OK, 599); + assertThat(diagnosticsContext.getContactedRegionNames()).isNotNull(); + if (diagnosticsContext.getStatusCode() < HttpConstants.StatusCodes.BADREQUEST) { + assertThat(diagnosticsContext.getContactedRegionNames()).isNotEmpty(); + } else { + assertThat(diagnosticsContext.getStatusCode()).isEqualTo(HttpConstants.StatusCodes.NOTFOUND); + assertThat(diagnosticsContext.getSubStatusCode()).isEqualTo(HttpConstants.SubStatusCodes.READ_SESSION_NOT_AVAILABLE); + } + } finally { + readSessionNotAvailableRule.disable(); + } + } + + @Test(groups = {"fi-sm-customer-workflows"}, timeOut = TIMEOUT) + public void writeFaultStaysOnSingleWritableRegion() { + FaultInjectionRule partitionMigratingRule = configureServerErrorRule( + this.container, + FaultInjectionOperationType.CREATE_ITEM, + FaultInjectionServerErrorType.PARTITION_IS_MIGRATING, + this.writableRegions.get(0), + currentFaultInjectionConnectionType(), + 1); + + try { + CosmosEndToEndOperationLatencyPolicyConfig e2ePolicy = new CosmosEndToEndOperationLatencyPolicyConfigBuilder(Duration.ofSeconds(5)) + .availabilityStrategy(new ThresholdBasedAvailabilityStrategy(Duration.ofMillis(100), Duration.ofMillis(200))) + .build(); + CosmosItemRequestOptions createOptions = new CosmosItemRequestOptions() + .setContentResponseOnWriteEnabled(true) + .setCosmosEndToEndOperationLatencyPolicyConfig(e2ePolicy); + + CosmosDiagnosticsContext diagnosticsContext = createWithDiagnostics(TestObject.create(), createOptions); + + assertThat(diagnosticsContext).isNotNull(); + assertThat(partitionMigratingRule.getHitCount()) + .as("the injected write fault should have been hit in the single writable region") + .isGreaterThanOrEqualTo(1); + assertThat(diagnosticsContext.getStatusCode()).isBetween(HttpConstants.StatusCodes.OK, 599); + assertThat(diagnosticsContext.getContactedRegionNames()).isNotNull(); + + // A single-write account cannot hedge writes to another region, so even with an availability strategy + // configured the write must never be routed to a read-only region. + Set readOnlyRegions = this.readableRegions + .stream() + .map(region -> region.toLowerCase(Locale.ROOT)) + .filter(region -> !region.equals(this.writableRegions.get(0).toLowerCase(Locale.ROOT))) + .collect(Collectors.toSet()); + assertThat(diagnosticsContext.getContactedRegionNames()).doesNotContainAnyElementsOf(readOnlyRegions); + } finally { + partitionMigratingRule.disable(); + } + } + + @DataProvider(name = "singleWriteReadFaultScenarios") + public Object[][] singleWriteReadFaultScenarios() { + return new Object[][]{ + {FaultInjectionServerErrorType.GONE}, + {FaultInjectionServerErrorType.TIMEOUT}, + {FaultInjectionServerErrorType.READ_SESSION_NOT_AVAILABLE}, + {FaultInjectionServerErrorType.INTERNAL_SERVER_ERROR}, + {FaultInjectionServerErrorType.SERVICE_UNAVAILABLE} + }; + } + + @Test(groups = {"fi-sm-customer-workflows"}, dataProvider = "singleWriteReadFaultScenarios", timeOut = TIMEOUT) + public void singleWriteReadFaultMatrix(FaultInjectionServerErrorType errorType) { + skipIfFaultTypeUnsupportedOnGateway(errorType, "Customer single-master read fault matrix"); + + TestObject item = TestObject.create(); + this.container.createItem(item).block(); + registerForCleanup(item); + + FaultInjectionRule faultRule = configureServerErrorRule( + this.container, + FaultInjectionOperationType.READ_ITEM, + errorType, + this.readableRegions.get(0), + currentFaultInjectionConnectionType(), + 1); + + try { + CosmosItemRequestOptions readOptions = new CosmosItemRequestOptions() + .setCosmosEndToEndOperationLatencyPolicyConfig( + new CosmosEndToEndOperationLatencyPolicyConfigBuilder(Duration.ofSeconds(5)) + .availabilityStrategy(new ThresholdBasedAvailabilityStrategy(Duration.ofMillis(100), Duration.ofMillis(200))) + .build()); + + CosmosDiagnosticsContext diagnosticsContext = readWithDiagnostics(item, readOptions); + + assertFaultInjectedOperation(diagnosticsContext, faultRule); + } finally { + faultRule.disable(); + } + } + + @DataProvider(name = "singleWriteMutationFaultScenarios") + public Object[][] singleWriteMutationFaultScenarios() { + return new Object[][]{ + {FaultInjectionServerErrorType.PARTITION_IS_MIGRATING}, + {FaultInjectionServerErrorType.TIMEOUT}, + {FaultInjectionServerErrorType.TOO_MANY_REQUEST}, + {FaultInjectionServerErrorType.RETRY_WITH}, + {FaultInjectionServerErrorType.INTERNAL_SERVER_ERROR}, + {FaultInjectionServerErrorType.SERVICE_UNAVAILABLE} + }; + } + + @Test(groups = {"fi-sm-customer-workflows"}, dataProvider = "singleWriteMutationFaultScenarios", timeOut = TIMEOUT) + public void singleWriteCreateFaultMatrix(FaultInjectionServerErrorType errorType) { + FaultInjectionRule faultRule = configureServerErrorRule( + this.container, + FaultInjectionOperationType.CREATE_ITEM, + errorType, + this.writableRegions.get(0), + currentFaultInjectionConnectionType(), + 1); + + try { + CosmosItemRequestOptions createOptions = new CosmosItemRequestOptions() + .setContentResponseOnWriteEnabled(true) + .setCosmosEndToEndOperationLatencyPolicyConfig( + new CosmosEndToEndOperationLatencyPolicyConfigBuilder(Duration.ofSeconds(5)) + .availabilityStrategy(new ThresholdBasedAvailabilityStrategy(Duration.ofMillis(100), Duration.ofMillis(200))) + .build()); + + CosmosDiagnosticsContext diagnosticsContext = createWithDiagnostics(TestObject.create(), createOptions); + + // The availability strategy cannot hedge writes on a single-write account; the assertion below confirms + // the injected write fault was still exercised and produced a real HTTP outcome. + assertFaultInjectedOperation(diagnosticsContext, faultRule); + } finally { + faultRule.disable(); + } + } + + private CosmosDiagnosticsContext readWithDiagnostics(TestObject item, CosmosItemRequestOptions options) { + try { + return this.container + .readItem(item.getId(), partitionKey(item), options, TestObject.class) + .block() + .getDiagnostics() + .getDiagnosticsContext(); + } catch (CosmosException error) { + return error.getDiagnostics().getDiagnosticsContext(); + } + } + + private CosmosItemResponse readWithReplicationRetry(TestObject item, CosmosItemRequestOptions options) { + Duration deadline = Duration.ofSeconds(30); + long deadlineNanos = System.nanoTime() + deadline.toNanos(); + CosmosException lastNotFound = null; + + while (System.nanoTime() < deadlineNanos) { + try { + return this.container + .readItem(item.getId(), partitionKey(item), options, TestObject.class) + .block(); + } catch (CosmosException error) { + if (error.getStatusCode() != HttpConstants.StatusCodes.NOTFOUND) { + throw error; + } + // Item not yet replicated to the remaining readable region - wait and retry. + lastNotFound = error; + try { + Thread.sleep(500); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + throw new AssertionError("Interrupted while waiting for cross-region replication.", interrupted); + } + } + } + + throw new AssertionError("Item was not replicated to the remaining readable region within " + deadline, lastNotFound); + } + + private CosmosDiagnosticsContext createWithDiagnostics(TestObject item, CosmosItemRequestOptions options) { + try { + CosmosDiagnosticsContext diagnosticsContext = this.container + .createItem(item, new PartitionKey(item.getMypk()), options) + .block() + .getDiagnostics() + .getDiagnosticsContext(); + + registerForCleanup(item); + return diagnosticsContext; + } catch (CosmosException error) { + return error.getDiagnostics().getDiagnosticsContext(); + } + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowStoredProcedureTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowStoredProcedureTest.java new file mode 100644 index 000000000000..762ae05b0a72 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowStoredProcedureTest.java @@ -0,0 +1,134 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.workflows.customer; + +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.CosmosException; +import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.models.CosmosStoredProcedureProperties; +import com.azure.cosmos.models.CosmosStoredProcedureRequestOptions; +import com.azure.cosmos.models.CosmosStoredProcedureResponse; +import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.test.faultinjection.FaultInjectionOperationType; +import com.azure.cosmos.test.faultinjection.FaultInjectionRule; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Factory; +import org.testng.annotations.Test; + +import java.time.Duration; +import java.util.Collections; +import java.util.UUID; +import java.util.function.Supplier; + +import static org.assertj.core.api.Assertions.assertThat; + +public class CustomerWorkflowStoredProcedureTest extends CustomerWorkflowTestBase { + + @Factory(dataProvider = "clientBuildersWithDirectTcpSession") + public CustomerWorkflowStoredProcedureTest(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + } + + @BeforeClass(groups = {"fi-customer-workflows"}, timeOut = SETUP_TIMEOUT) + public void beforeClass() { + initializeSharedSinglePartitionContainer("Customer stored procedure workflow tests"); + } + + @AfterClass(groups = {"fi-customer-workflows"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) + public void afterClass() { + closeClient(); + } + + @Test(groups = {"fi-customer-workflows"}, timeOut = TIMEOUT) + public void storedProcedureCreateReadExecuteWithMetadataFaultRule() { + String storedProcedureId = "customer-sproc-" + UUID.randomUUID(); + CosmosStoredProcedureProperties storedProcedureProperties = new CosmosStoredProcedureProperties( + storedProcedureId, + "function(input) {" + + " var value = input || 'workflow';" + + " console.log('stored procedure workflow ' + value);" + + " getContext().getResponse().setBody('sproc-ok:' + value);" + + "}"); + + CosmosStoredProcedureResponse createResponse = this.container + .getScripts() + .createStoredProcedure(storedProcedureProperties) + .block(); + + assertThat(createResponse).isNotNull(); + assertThat(createResponse.getStatusCode()).isEqualTo(HttpConstants.StatusCodes.CREATED); + assertThat(createResponse.getDiagnostics()).isNotNull(); + + FaultInjectionRule metadataDelayRule = configureResponseDelayRule( + this.container, + FaultInjectionOperationType.METADATA_REQUEST_CONTAINER, + Duration.ofMillis(100), + 1); + + try { + CosmosStoredProcedureRequestOptions options = new CosmosStoredProcedureRequestOptions(); + options.setPartitionKey(new PartitionKey("sproc-workflow")); + options.setScriptLoggingEnabled(true); + + CosmosStoredProcedureResponse readResponse = withStoredProcedureReplicationRetry(() -> this.container + .getScripts() + .getStoredProcedure(storedProcedureId) + .read() + .block()); + + assertThat(readResponse).isNotNull(); + assertThat(readResponse.getProperties().getId()).isEqualTo(storedProcedureId); + + CosmosStoredProcedureResponse executeResponse = withStoredProcedureReplicationRetry(() -> this.container + .getScripts() + .getStoredProcedure(storedProcedureId) + .execute(Collections.singletonList("workflow"), options) + .block()); + + assertThat(executeResponse).isNotNull(); + assertThat(executeResponse.getStatusCode()).isEqualTo(HttpConstants.StatusCodes.OK); + assertThat(executeResponse.getResponseAsString()).contains("sproc-ok:workflow"); + assertThat(executeResponse.getScriptLog()).contains("stored procedure workflow workflow"); + assertThat(executeResponse.getDiagnostics()).isNotNull(); + } finally { + metadataDelayRule.disable(); + try { + this.container.getScripts().getStoredProcedure(storedProcedureId).delete().block(); + } catch (Exception error) { + // best-effort cleanup of the stored procedure created by this test + } + } + } + + /** + * Retries a stored-procedure operation while it returns 404. A stored procedure that was just created can + * be temporarily not found when the request is routed to a region the metadata has not yet replicated to + * (possible on a multi-write account, where stored-procedure metadata is not covered by session + * read-your-write the way document operations are). + */ + private CosmosStoredProcedureResponse withStoredProcedureReplicationRetry(Supplier operation) { + Duration deadline = Duration.ofSeconds(30); + long deadlineNanos = System.nanoTime() + deadline.toNanos(); + CosmosException lastNotFound = null; + + while (System.nanoTime() < deadlineNanos) { + try { + return operation.get(); + } catch (CosmosException error) { + if (error.getStatusCode() != HttpConstants.StatusCodes.NOTFOUND) { + throw error; + } + lastNotFound = error; + try { + Thread.sleep(500); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + throw new AssertionError("Interrupted while waiting for stored procedure replication.", interrupted); + } + } + } + + throw new AssertionError("Stored procedure was not available to read within " + deadline, lastNotFound); + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowTestBase.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowTestBase.java new file mode 100644 index 000000000000..5f35d2defc31 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowTestBase.java @@ -0,0 +1,499 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.workflows.customer; + +import com.azure.cosmos.CosmosAsyncClient; +import com.azure.cosmos.CosmosAsyncContainer; +import com.azure.cosmos.CosmosAsyncDatabase; +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.CosmosDiagnosticsContext; +import com.azure.cosmos.ConnectionMode; +import com.azure.cosmos.ConsistencyLevel; +import com.azure.cosmos.TestObject; +import com.azure.cosmos.implementation.AsyncDocumentClient; +import com.azure.cosmos.implementation.DatabaseAccount; +import com.azure.cosmos.implementation.DatabaseAccountLocation; +import com.azure.cosmos.implementation.GlobalEndpointManager; +import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.implementation.ImplementationBridgeHelpers; +import com.azure.cosmos.implementation.OverridableRequestOptions; +import com.azure.cosmos.implementation.RxDocumentClientImpl; +import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils; +import com.azure.cosmos.models.CosmosItemIdentity; +import com.azure.cosmos.models.CosmosItemRequestOptions; +import com.azure.cosmos.models.ThroughputProperties; +import com.azure.cosmos.rx.TestSuiteBase; +import com.azure.cosmos.test.faultinjection.CosmosFaultInjectionHelper; +import com.azure.cosmos.test.faultinjection.FaultInjectionCondition; +import com.azure.cosmos.test.faultinjection.FaultInjectionConditionBuilder; +import com.azure.cosmos.test.faultinjection.FaultInjectionConnectionType; +import com.azure.cosmos.test.faultinjection.FaultInjectionOperationType; +import com.azure.cosmos.test.faultinjection.FaultInjectionResultBuilders; +import com.azure.cosmos.test.faultinjection.FaultInjectionRule; +import com.azure.cosmos.test.faultinjection.FaultInjectionRuleBuilder; +import com.azure.cosmos.test.faultinjection.FaultInjectionServerErrorType; +import com.azure.cosmos.test.faultinjection.IFaultInjectionResult; +import org.testng.SkipException; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Collection; +import java.util.List; +import java.util.Locale; +import java.util.Set; +import java.util.UUID; +import java.util.function.BooleanSupplier; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; + +public abstract class CustomerWorkflowTestBase extends TestSuiteBase { + protected CosmosAsyncClient client; + protected CosmosAsyncContainer container; + protected List writableRegions; + protected List readableRegions; + private final List itemsToCleanup = Collections.synchronizedList(new ArrayList<>()); + + protected CustomerWorkflowTestBase(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + } + + protected final void initializeSharedSinglePartitionContainer(String scenarioName) { + initializeSharedSinglePartitionContainer(scenarioName, false); + } + + protected final void initializeSharedSinglePartitionContainer(String scenarioName, boolean forceSessionConsistency) { + if (forceSessionConsistency) { + skipIfAccountConsistencyWeakerThanSession(scenarioName); + } + + CosmosAsyncClient discoveryClient = null; + + try { + discoveryClient = getClientBuilder().buildAsyncClient(); + this.writableRegions = discoverWritableRegions(discoveryClient); + skipIfInsufficientRegions(this.writableRegions, scenarioName); + + CosmosClientBuilder clientBuilder = getClientBuilder() + .preferredRegions(this.writableRegions) + .multipleWriteRegionsEnabled(true) + .contentResponseOnWriteEnabled(true); + + if (forceSessionConsistency) { + // Read-your-write across an excluded write region is only deterministic with session (or + // stronger) consistency, so pin the client to session consistency for these scenarios. + clientBuilder.consistencyLevel(ConsistencyLevel.SESSION); + } + + this.client = clientBuilder.buildAsyncClient(); + this.container = getSharedSinglePartitionCosmosContainer(this.client); + waitForCollectionToBeAvailableToRead(this.container, /* probeClient */ null); + } finally { + safeClose(discoveryClient); + } + } + + protected final void closeClient() { + cleanupRegisteredItems(); + safeClose(this.client); + this.client = null; + this.container = null; + this.writableRegions = null; + this.readableRegions = null; + } + + protected final void initializeSharedSingleWriteMultiRegionContainer(String scenarioName) { + CosmosAsyncClient discoveryClient = null; + + try { + discoveryClient = getClientBuilder() + .multipleWriteRegionsEnabled(false) + .contentResponseOnWriteEnabled(true) + .buildAsyncClient(); + this.writableRegions = discoverWritableRegions(discoveryClient); + this.readableRegions = discoverReadableRegions(discoveryClient); + skipIfInsufficientReadableRegions(this.readableRegions, scenarioName); + skipIfNotSingleWriteRegion(this.writableRegions, scenarioName); + + this.client = getClientBuilder() + .preferredRegions(this.readableRegions) + .multipleWriteRegionsEnabled(false) + .contentResponseOnWriteEnabled(true) + .buildAsyncClient(); + this.container = getSharedSinglePartitionCosmosContainer(this.client); + waitForCollectionToBeAvailableToRead(this.container, /* probeClient */ null); + } finally { + safeClose(discoveryClient); + } + } + + /** + * Registers an item to be best-effort deleted from the shared container when the test class finishes, + * so the shared single-partition container does not accumulate items across runs. + */ + protected final void registerForCleanup(TestObject item) { + if (item != null) { + this.itemsToCleanup.add(new CosmosItemIdentity(partitionKey(item), item.getId())); + } + } + + private void cleanupRegisteredItems() { + CosmosAsyncContainer cleanupContainer = this.container; + List snapshot; + synchronized (this.itemsToCleanup) { + snapshot = new ArrayList<>(this.itemsToCleanup); + this.itemsToCleanup.clear(); + } + + if (cleanupContainer == null) { + return; + } + + for (CosmosItemIdentity identity : snapshot) { + try { + cleanupContainer + .deleteItem(identity.getId(), identity.getPartitionKey(), new CosmosItemRequestOptions()) + .block(); + } catch (Exception error) { + // best-effort cleanup - ignore (for example item already deleted by the test itself) + } + } + } + + protected final List excludeFirstWritableRegion() { + return Collections.singletonList(this.writableRegions.get(0)); + } + + protected final List excludeFirstReadableRegion() { + return Collections.singletonList(this.readableRegions.get(0)); + } + + protected static com.azure.cosmos.models.PartitionKey partitionKey(TestObject item) { + return new com.azure.cosmos.models.PartitionKey(item.getMypk()); + } + + protected final CosmosAsyncContainer createTemporaryContainer(String prefix, String partitionKeyPath) { + CosmosAsyncDatabase database = getSharedCosmosDatabase(this.client); + String containerId = prefix + "-" + UUID.randomUUID(); + + database + .createContainerIfNotExists(containerId, partitionKeyPath, ThroughputProperties.createManualThroughput(400)) + .block(); + + return database.getContainer(containerId); + } + + protected static void deleteTemporaryContainer(CosmosAsyncContainer container) { + safeDeleteCollection(container); + } + + protected static void waitForCollectionToBeAvailableToRead( + CosmosAsyncContainer container, + CosmosAsyncClient probeClient) { + + CosmosAsyncContainer probeContainer = probeClient == null + ? container + : probeClient.getDatabase(container.getDatabase().getId()).getContainer(container.getId()); + awaitCondition( + () -> { + try { + probeContainer.read().block(); + return true; + } catch (RuntimeException ignored) { + return false; + } + }, + Duration.ofMinutes(2), + "Container '" + container.getId() + "' was not available to read within 2 minutes."); + } + + protected static void awaitCondition(BooleanSupplier condition, Duration timeout, String failureMessage) { + long deadline = System.nanoTime() + timeout.toNanos(); + + while (System.nanoTime() < deadline) { + if (condition.getAsBoolean()) { + return; + } + + try { + Thread.sleep(250); + } catch (InterruptedException error) { + Thread.currentThread().interrupt(); + throw new AssertionError("Interrupted while waiting for condition: " + failureMessage, error); + } + } + + throw new AssertionError(failureMessage); + } + + protected final FaultInjectionRule configureServerErrorRule( + CosmosAsyncContainer targetContainer, + FaultInjectionOperationType operationType, + FaultInjectionServerErrorType errorType, + int hitLimit) { + + return configureServerErrorRule(targetContainer, operationType, errorType, this.writableRegions.get(0), hitLimit); + } + + protected final FaultInjectionRule configureServerErrorRule( + CosmosAsyncContainer targetContainer, + FaultInjectionOperationType operationType, + FaultInjectionServerErrorType errorType, + String region, + int hitLimit) { + + return configureServerErrorRule(targetContainer, operationType, errorType, region, currentFaultInjectionConnectionType(), hitLimit); + } + + protected final FaultInjectionRule configureServerErrorRule( + CosmosAsyncContainer targetContainer, + FaultInjectionOperationType operationType, + FaultInjectionServerErrorType errorType, + String region, + FaultInjectionConnectionType connectionType, + int hitLimit) { + + FaultInjectionConditionBuilder conditionBuilder = new FaultInjectionConditionBuilder() + .operationType(operationType) + .connectionType(connectionType); + + if (region != null) { + conditionBuilder.region(region); + } + + FaultInjectionRule rule = new FaultInjectionRuleBuilder("customer-workflow-" + errorType + "-" + UUID.randomUUID()) + .condition(conditionBuilder.build()) + .result(FaultInjectionResultBuilders.getResultBuilder(errorType).build()) + .duration(Duration.ofMinutes(5)) + .hitLimit(hitLimit) + .build(); + + CosmosFaultInjectionHelper.configureFaultInjectionRules(targetContainer, Collections.singletonList(rule)).block(); + return rule; + } + + protected final FaultInjectionConnectionType currentFaultInjectionConnectionType() { + if (getConnectionPolicy().getConnectionMode() == ConnectionMode.GATEWAY) { + return FaultInjectionConnectionType.GATEWAY; + } + + return FaultInjectionConnectionType.DIRECT; + } + + protected final FaultInjectionRule configureResponseDelayRule( + CosmosAsyncContainer targetContainer, + FaultInjectionOperationType operationType, + Duration delay, + int hitLimit) { + + FaultInjectionCondition condition = new FaultInjectionConditionBuilder() + .operationType(operationType) + .connectionType(currentFaultInjectionConnectionType()) + .build(); + + IFaultInjectionResult result = FaultInjectionResultBuilders + .getResultBuilder(FaultInjectionServerErrorType.RESPONSE_DELAY) + .delay(delay) + .times(hitLimit) + .build(); + + FaultInjectionRule rule = new FaultInjectionRuleBuilder("customer-workflow-response-delay-" + UUID.randomUUID()) + .condition(condition) + .result(result) + .duration(Duration.ofMinutes(5)) + .hitLimit(hitLimit) + .build(); + + CosmosFaultInjectionHelper.configureFaultInjectionRules(targetContainer, Collections.singletonList(rule)).block(); + return rule; + } + + protected static List discoverWritableRegions(CosmosAsyncClient client) { + DatabaseAccount databaseAccount = readDatabaseAccount(client); + + List writableRegions = new ArrayList<>(); + for (DatabaseAccountLocation accountLocation : databaseAccount.getWritableLocations()) { + writableRegions.add(accountLocation.getName()); + } + + return writableRegions; + } + + protected static List discoverReadableRegions(CosmosAsyncClient client) { + DatabaseAccount databaseAccount = readDatabaseAccount(client); + + List readableRegions = new ArrayList<>(); + for (DatabaseAccountLocation accountLocation : databaseAccount.getReadableLocations()) { + readableRegions.add(accountLocation.getName()); + } + + return readableRegions; + } + + private static DatabaseAccount readDatabaseAccount(CosmosAsyncClient client) { + AsyncDocumentClient asyncDocumentClient = ReflectionUtils.getAsyncDocumentClient(client); + RxDocumentClientImpl rxDocumentClient = (RxDocumentClientImpl) asyncDocumentClient; + GlobalEndpointManager globalEndpointManager = ReflectionUtils.getGlobalEndpointManager(rxDocumentClient); + + // The latest database account is populated during client initialization. Poll briefly to defend against + // an initialization race instead of forcing a synthetic database-account read (which is not routable in + // direct connection mode). + DatabaseAccount databaseAccount = globalEndpointManager.getLatestDatabaseAccount(); + long deadlineNanos = System.nanoTime() + Duration.ofSeconds(10).toNanos(); + while (databaseAccount == null && System.nanoTime() < deadlineNanos) { + try { + Thread.sleep(200); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + throw new AssertionError("Interrupted while waiting for the database account to be available.", interrupted); + } + databaseAccount = globalEndpointManager.getLatestDatabaseAccount(); + } + + assertThat(databaseAccount) + .as("database account must be available for region discovery") + .isNotNull(); + + return databaseAccount; + } + + protected static void skipIfInsufficientRegions(List regions, String scenarioName) { + if (regions == null || regions.size() < 2) { + throw new SkipException(scenarioName + " requires a live multi-region account."); + } + } + + protected static void skipIfInsufficientReadableRegions(List regions, String scenarioName) { + if (regions == null || regions.size() < 2) { + throw new SkipException(scenarioName + " requires a live multi-region single-write account."); + } + } + + protected static void skipIfNotSingleWriteRegion(List regions, String scenarioName) { + if (regions == null || regions.size() != 1) { + throw new SkipException(scenarioName + " requires exactly one write region."); + } + } + + protected static void skipIfAccountConsistencyWeakerThanSession(String scenarioName) { + if (accountConsistency == ConsistencyLevel.EVENTUAL || accountConsistency == ConsistencyLevel.CONSISTENT_PREFIX) { + throw new SkipException( + scenarioName + " requires an account with session or stronger default consistency for deterministic read-your-write."); + } + } + + protected final void skipIfNotDirectMode(String scenarioName) { + if (getConnectionPolicy().getConnectionMode() != ConnectionMode.DIRECT) { + throw new SkipException(scenarioName + " only applies to the direct connection mode client builder."); + } + } + + protected final void skipIfNotGatewayMode(String scenarioName) { + if (getConnectionPolicy().getConnectionMode() != ConnectionMode.GATEWAY) { + throw new SkipException(scenarioName + " only applies to the gateway connection mode client builder."); + } + } + + /** + * Skips fault-injection scenarios that cannot be injected for the gateway connection type. The gateway + * internally retries 410/0, so {@code GONE} and {@code STALED_ADDRESSES_SERVER_GONE} rules are rejected at + * configuration time for gateway-mode clients. + */ + protected final void skipIfFaultTypeUnsupportedOnGateway(FaultInjectionServerErrorType errorType, String scenarioName) { + if (currentFaultInjectionConnectionType() == FaultInjectionConnectionType.GATEWAY + && (errorType == FaultInjectionServerErrorType.GONE + || errorType == FaultInjectionServerErrorType.STALED_ADDRESSES_SERVER_GONE)) { + + throw new SkipException( + scenarioName + " cannot inject " + errorType + " for the gateway connection type."); + } + } + + /** + * Configures the same server-error fault for both the point-read ({@code READ_ITEM}) and query + * ({@code QUERY_ITEM}) operation types. {@code readMany} resolves to a point read for a single item in a + * partition and to a query for multiple items, so both rules are needed for the fault to reliably apply. + */ + protected final List configureReadManyServerErrorRules( + CosmosAsyncContainer targetContainer, + FaultInjectionServerErrorType errorType, + String region, + int hitLimit) { + + List rules = new ArrayList<>(); + rules.add(configureServerErrorRule( + targetContainer, FaultInjectionOperationType.READ_ITEM, errorType, region, currentFaultInjectionConnectionType(), hitLimit)); + rules.add(configureServerErrorRule( + targetContainer, FaultInjectionOperationType.QUERY_ITEM, errorType, region, currentFaultInjectionConnectionType(), hitLimit)); + return rules; + } + + /** + * Asserts that a fault-injected operation produced a real HTTP outcome and that at least one of the supplied + * fault rules was actually hit, so the scenario cannot silently pass without exercising the injected fault. + */ + protected static void assertFaultInjectedOperation( + CosmosDiagnosticsContext diagnosticsContext, + FaultInjectionRule... rules) { + + assertThat(diagnosticsContext).isNotNull(); + assertThat(diagnosticsContext.getStatusCode()).isBetween(HttpConstants.StatusCodes.OK, 599); + assertThat(diagnosticsContext.getContactedRegionNames()).isNotNull(); + + long totalHits = 0; + for (FaultInjectionRule rule : rules) { + totalHits += rule.getHitCount(); + } + + assertThat(totalHits) + .as("expected at least one injected fault to be hit") + .isGreaterThanOrEqualTo(1); + } + + protected static void assertFaultInjectedOperation( + CosmosDiagnosticsContext diagnosticsContext, + List rules) { + + assertFaultInjectedOperation(diagnosticsContext, rules.toArray(new FaultInjectionRule[0])); + } + + protected static OverridableRequestOptions getRequestOptions(CosmosDiagnosticsContext diagnosticsContext) { + assertThat(diagnosticsContext).isNotNull(); + return ImplementationBridgeHelpers + .CosmosDiagnosticsContextHelper + .getCosmosDiagnosticsContextAccessor() + .getRequestOptions(diagnosticsContext); + } + + protected static void assertKeywordIdentifier(CosmosDiagnosticsContext diagnosticsContext, String expectedKeywordIdentifier) { + OverridableRequestOptions requestOptions = getRequestOptions(diagnosticsContext); + + assertThat(requestOptions.getKeywordIdentifiers()) + .contains(expectedKeywordIdentifier); + } + + protected static void assertExcludedRegions( + CosmosDiagnosticsContext diagnosticsContext, + List expectedExcludedRegions) { + + OverridableRequestOptions requestOptions = getRequestOptions(diagnosticsContext); + + assertThat(requestOptions.getExcludedRegions()) + .containsExactlyElementsOf(expectedExcludedRegions); + } + + protected static void assertDidNotContactExcludedRegions( + CosmosDiagnosticsContext diagnosticsContext, + Collection excludedRegions) { + + Set contactedRegionNames = diagnosticsContext.getContactedRegionNames(); + Set normalizedExcludedRegions = excludedRegions + .stream() + .map(region -> region.toLowerCase(Locale.ROOT)) + .collect(Collectors.toSet()); + + assertThat(contactedRegionNames).isNotNull(); + assertThat(contactedRegionNames).doesNotContainAnyElementsOf(normalizedExcludedRegions); + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/resources/fi-customer-workflows-testng.xml b/sdk/cosmos/azure-cosmos-tests/src/test/resources/fi-customer-workflows-testng.xml new file mode 100644 index 000000000000..edfa8a57770f --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/resources/fi-customer-workflows-testng.xml @@ -0,0 +1,38 @@ + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/resources/fi-sm-customer-workflows-testng.xml b/sdk/cosmos/azure-cosmos-tests/src/test/resources/fi-sm-customer-workflows-testng.xml new file mode 100644 index 000000000000..976b8fbdc204 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/resources/fi-sm-customer-workflows-testng.xml @@ -0,0 +1,38 @@ + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/sdk/cosmos/live-fi-customer-workflows-platform-matrix.json b/sdk/cosmos/live-fi-customer-workflows-platform-matrix.json new file mode 100644 index 000000000000..c3b1e2b841fe --- /dev/null +++ b/sdk/cosmos/live-fi-customer-workflows-platform-matrix.json @@ -0,0 +1,42 @@ +{ + "displayNames": { + "-Pfi-customer-workflows": "FaultInjectionCustomerWorkflows", + "Session": "", + "ubuntu": "", + "@{ enableMultipleWriteLocations = $true; defaultConsistencyLevel = 'Session'; enableMultipleRegions = $true }": "" + }, + "include": [ + { + "DESIRED_CONSISTENCIES": "[\"Session\"]", + "ACCOUNT_CONSISTENCY": "Session", + "ArmConfig": { + "MultiMaster_MultiRegion_FI_CustomerWorkflows": { + "ArmTemplateParameters": "@{ enableMultipleWriteLocations = $true; defaultConsistencyLevel = 'Session'; enableMultipleRegions = $true }", + "PREFERRED_LOCATIONS": "[\"East US 2\"]" + } + }, + "PROTOCOLS": "[\"Tcp\"]", + "ProfileFlag": [ "-Pfi-customer-workflows" ], + "AdditionalArgs": "\"-DCOSMOS.PARTITION_LEVEL_CIRCUIT_BREAKER_DEFAULT_CONFIG_OPT_IN=TRUE\"", + "Agent": { + "ubuntu": { "OSVmImage": "env:LINUXVMIMAGE", "Pool": "env:LINUXPOOL" } + } + }, + { + "DESIRED_CONSISTENCIES": "[\"Session\"]", + "ACCOUNT_CONSISTENCY": "Session", + "ArmConfig": { + "MultiMaster_MultiRegion_FI_CustomerWorkflows_ThinClient_Http2": { + "ArmTemplateParameters": "@{ enableMultipleWriteLocations = $true; defaultConsistencyLevel = 'Session'; enableMultipleRegions = $true }", + "PREFERRED_LOCATIONS": "[]" + } + }, + "PROTOCOLS": "[\"Tcp\"]", + "ProfileFlag": [ "-Pfi-customer-workflows" ], + "AdditionalArgs": "-DCOSMOS.CLIENT_LEAK_DETECTION_ENABLED=true -DACCOUNT_HOST=$(thin-client-canary-multi-writer-session-endpoint) -DACCOUNT_KEY=$(thin-client-canary-multi-writer-session-key) -DCOSMOS.THINCLIENT_ENABLED=true -DCOSMOS.HTTP2_ENABLED=true", + "Agent": { + "ubuntu": { "OSVmImage": "env:LINUXVMIMAGE", "Pool": "env:LINUXPOOL" } + } + } + ] +} diff --git a/sdk/cosmos/live-fi-sm-customer-workflows-platform-matrix.json b/sdk/cosmos/live-fi-sm-customer-workflows-platform-matrix.json new file mode 100644 index 000000000000..cd56d0a830a5 --- /dev/null +++ b/sdk/cosmos/live-fi-sm-customer-workflows-platform-matrix.json @@ -0,0 +1,41 @@ +{ + "displayNames": { + "-Pfi-sm-customer-workflows": "FaultInjectionSingleMasterCustomerWorkflows", + "Session": "", + "ubuntu": "", + "@{ enableMultipleWriteLocations = $false; defaultConsistencyLevel = 'Session'; enableMultipleRegions = $true }": "" + }, + "include": [ + { + "DESIRED_CONSISTENCIES": "[\"Session\"]", + "ACCOUNT_CONSISTENCY": "Session", + "ArmConfig": { + "SingleMaster_MultiRegion_FI_CustomerWorkflows": { + "ArmTemplateParameters": "@{ enableMultipleWriteLocations = $false; defaultConsistencyLevel = 'Session'; enableMultipleRegions = $true }", + "PREFERRED_LOCATIONS": "[\"East US 2\"]" + } + }, + "PROTOCOLS": "[\"Tcp\"]", + "ProfileFlag": [ "-Pfi-sm-customer-workflows" ], + "Agent": { + "ubuntu": { "OSVmImage": "env:LINUXVMIMAGE", "Pool": "env:LINUXPOOL" } + } + }, + { + "DESIRED_CONSISTENCIES": "[\"Session\"]", + "ACCOUNT_CONSISTENCY": "Session", + "ArmConfig": { + "SingleMaster_MultiRegion_FI_CustomerWorkflows_ThinClient_Http2": { + "ArmTemplateParameters": "@{ enableMultipleWriteLocations = $false; defaultConsistencyLevel = 'Session'; enableMultipleRegions = $true }", + "PREFERRED_LOCATIONS": "[]" + } + }, + "PROTOCOLS": "[\"Tcp\"]", + "ProfileFlag": [ "-Pfi-sm-customer-workflows" ], + "AdditionalArgs": "-DCOSMOS.CLIENT_LEAK_DETECTION_ENABLED=true -DACCOUNT_HOST=$(thin-client-canary-multi-region-session-endpoint) -DACCOUNT_KEY=$(thin-client-canary-multi-region-session-key) -DCOSMOS.THINCLIENT_ENABLED=true -DCOSMOS.HTTP2_ENABLED=true", + "Agent": { + "ubuntu": { "OSVmImage": "env:LINUXVMIMAGE", "Pool": "env:LINUXPOOL" } + } + } + ] +} diff --git a/sdk/cosmos/tests.yml b/sdk/cosmos/tests.yml index 69d782fcc9a0..bec960b6c133 100644 --- a/sdk/cosmos/tests.yml +++ b/sdk/cosmos/tests.yml @@ -163,6 +163,70 @@ extends: - name: AdditionalArgs value: '-DCOSMOS.CLIENT_LEAK_DETECTION_ENABLED=true -DACCOUNT_HOST=$(thin-client-canary-multi-writer-session-endpoint) -DACCOUNT_KEY=$(thin-client-canary-multi-writer-session-key) -DCOSMOS.THINCLIENT_ENABLED=true' + - template: /eng/pipelines/templates/stages/archetype-sdk-tests-isolated.yml + parameters: + TestName: 'Cosmos_Live_Test_FaultInjectionCustomerWorkflows' + CloudConfig: + Public: + ServiceConnection: azure-sdk-tests-cosmos + MatrixConfigs: + - Name: Cosmos_live_test_fi_customer_workflows + Path: sdk/cosmos/live-fi-customer-workflows-platform-matrix.json + Selection: all + GenerateVMJobs: true + MatrixReplace: + - .*Version=1.2(1|5)/1.17 + ServiceDirectory: cosmos + Artifacts: + - name: azure-cosmos + groupId: com.azure + safeName: azurecosmos + AdditionalModules: + - name: azure-cosmos-tests + groupId: com.azure + - name: azure-cosmos-benchmark + groupId: com.azure + TimeoutInMinutes: 210 + MaxParallel: 20 + TestGoals: 'verify' + TestOptions: '$(ProfileFlag) $(AdditionalArgs) -DskipCompile=true -DskipTestCompile=true -DcreateSourcesJar=false' + TestResultsFiles: '**/junitreports/TEST-*.xml' + AdditionalVariables: + - name: AdditionalArgs + value: '-DCOSMOS.CLIENT_LEAK_DETECTION_ENABLED=true' + + - template: /eng/pipelines/templates/stages/archetype-sdk-tests-isolated.yml + parameters: + TestName: 'Cosmos_Live_Test_FaultInjectionSingleMasterCustomerWorkflows' + CloudConfig: + Public: + ServiceConnection: azure-sdk-tests-cosmos + MatrixConfigs: + - Name: Cosmos_live_test_fi_sm_customer_workflows + Path: sdk/cosmos/live-fi-sm-customer-workflows-platform-matrix.json + Selection: all + GenerateVMJobs: true + MatrixReplace: + - .*Version=1.2(1|5)/1.17 + ServiceDirectory: cosmos + Artifacts: + - name: azure-cosmos + groupId: com.azure + safeName: azurecosmos + AdditionalModules: + - name: azure-cosmos-tests + groupId: com.azure + - name: azure-cosmos-benchmark + groupId: com.azure + TimeoutInMinutes: 210 + MaxParallel: 20 + TestGoals: 'verify' + TestOptions: '$(ProfileFlag) $(AdditionalArgs) -DskipCompile=true -DskipTestCompile=true -DcreateSourcesJar=false' + TestResultsFiles: '**/junitreports/TEST-*.xml' + AdditionalVariables: + - name: AdditionalArgs + value: '-DCOSMOS.CLIENT_LEAK_DETECTION_ENABLED=true' + - template: /eng/pipelines/templates/stages/archetype-sdk-tests-isolated.yml parameters: TestName: 'Spring_Data_Cosmos_Integration' From 00aa877b9dda9f807c2851a0ff2a526a9ceb2fc8 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Sat, 29 Aug 2026 06:48:16 -0400 Subject: [PATCH 12/14] Use authenticated Maven mirror in pipelines Route CI and live-test Maven dependency resolution through the authenticated Azure SDK feed so network-isolated agents do not access Maven Central directly. --- eng/pipelines/templates/jobs/ci.tests.yml | 3 +++ eng/pipelines/templates/jobs/live.tests.yml | 3 +++ .../templates/steps/maven-authenticate.yml | 14 ++++++++++++++ eng/pipelines/templates/variables/globals.yml | 2 +- eng/settings.xml | 15 ++++++++++++--- 5 files changed, 33 insertions(+), 4 deletions(-) create mode 100644 eng/pipelines/templates/steps/maven-authenticate.yml diff --git a/eng/pipelines/templates/jobs/ci.tests.yml b/eng/pipelines/templates/jobs/ci.tests.yml index fdf148f93e98..82d12ff45fdf 100644 --- a/eng/pipelines/templates/jobs/ci.tests.yml +++ b/eng/pipelines/templates/jobs/ci.tests.yml @@ -127,6 +127,9 @@ jobs: - ${{ parameters.PreTestSteps }} + # Authenticate with Azure Artifacts + - template: /eng/pipelines/templates/steps/maven-authenticate.yml + - template: /eng/pipelines/templates/steps/run-and-validate-linting.yml parameters: JavaBuildVersion: $(JavaTestVersion) diff --git a/eng/pipelines/templates/jobs/live.tests.yml b/eng/pipelines/templates/jobs/live.tests.yml index 90575cda9782..5dd6d9d3c918 100644 --- a/eng/pipelines/templates/jobs/live.tests.yml +++ b/eng/pipelines/templates/jobs/live.tests.yml @@ -121,6 +121,9 @@ jobs: - ${{ parameters.PreSteps }} + # Authenticate with Azure Artifacts + - template: /eng/pipelines/templates/steps/maven-authenticate.yml + - template: /eng/pipelines/templates/steps/build-and-test.yml parameters: PreTestRunSteps: ${{ parameters.PreTestRunSteps }} diff --git a/eng/pipelines/templates/steps/maven-authenticate.yml b/eng/pipelines/templates/steps/maven-authenticate.yml new file mode 100644 index 000000000000..31e6a0a635b9 --- /dev/null +++ b/eng/pipelines/templates/steps/maven-authenticate.yml @@ -0,0 +1,14 @@ +steps: + # Copy mirror settings to default Maven location so all requests go through CFS + - pwsh: | + $m2Dir = if ($env:USERPROFILE) { "$env:USERPROFILE\.m2" } else { "$HOME/.m2" } + New-Item -ItemType Directory -Force -Path $m2Dir | Out-Null + Copy-Item -Path "$(Build.SourcesDirectory)/eng/settings.xml" -Destination "$m2Dir/settings.xml" + displayName: 'Setup Maven mirror settings' + + # Authenticate with Azure Artifacts feeds + # MavenAuthenticate adds entries to ~/.m2/settings.xml matching mirror id 'azure-sdk-for-java' + - task: MavenAuthenticate@0 + displayName: 'Maven Authenticate' + inputs: + artifactsFeeds: 'azure-sdk-for-java' \ No newline at end of file diff --git a/eng/pipelines/templates/variables/globals.yml b/eng/pipelines/templates/variables/globals.yml index 674aaebc4750..cbd64a9a98d1 100644 --- a/eng/pipelines/templates/variables/globals.yml +++ b/eng/pipelines/templates/variables/globals.yml @@ -26,7 +26,7 @@ variables: # See https://github.com/actions/virtual-environments/issues/1499 for more info about the wagon options # If reports about Maven dependency downloads become more common investigate re-introducing "-Dhttp.keepAlive=false -Dmaven.wagon.http.pool=false", or other iterations of the configurations. WagonOptions: '-Dmaven.wagon.httpconnectionManager.ttlSeconds=60 -Dmaven.wagon.http.pool=false' - DefaultOptions: '-Dmaven.repo.local=$(MAVEN_CACHE_FOLDER) --batch-mode --fail-at-end --settings eng/settings.xml $(WagonOptions)' + DefaultOptions: '-Dmaven.repo.local=$(MAVEN_CACHE_FOLDER) --batch-mode --fail-at-end $(WagonOptions)' LoggingOptions: '-Dorg.slf4j.simpleLogger.defaultLogLevel=$(MavenLogLevel) -Dorg.slf4j.simpleLogger.log.org.apache.maven.cli.transfer.Slf4jMavenTransferListener=warn' MemoryOptions: '-Xmx4096m' DefaultSkipOptions: '-Dgpg.skip -Dmaven.javadoc.skip=true -Dcodesnippet.skip=true -Dspotbugs.skip=true -Dcheckstyle.skip=true -Drevapi.skip=true -DtrimStackTrace=false -Dspotless.apply.skip=true -Dspotless.check.skip=true' diff --git a/eng/settings.xml b/eng/settings.xml index b1b7cb0d1d0d..8e65b655fd04 100644 --- a/eng/settings.xml +++ b/eng/settings.xml @@ -1,4 +1,13 @@ - + + + + azure-sdk-for-java + Azure Artifacts Maven Mirror + https://pkgs.dev.azure.com/azure-sdk/public/_packaging/azure-sdk-for-java/maven/v1 + external:*,!confluent,!repository.spring.milestone + + From 2cba37f166291b32936f2628f728d5df44102b9e Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Sat, 29 Aug 2026 18:17:42 -0400 Subject: [PATCH 13/14] Run LatestCommitted workflow in Direct mode Exclude the unsupported Gateway factory case while retaining Direct-mode LatestCommitted customer workflow coverage. --- .../workflows/customer/CustomerWorkflowLatestCommittedTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowLatestCommittedTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowLatestCommittedTest.java index a5eeaca9822d..ce345fb89ce7 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowLatestCommittedTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowLatestCommittedTest.java @@ -31,7 +31,7 @@ public class CustomerWorkflowLatestCommittedTest extends CustomerWorkflowTestBase { - @Factory(dataProvider = "clientBuildersWithSessionConsistency") + @Factory(dataProvider = "clientBuilderSolelyDirectWithSessionConsistency") public CustomerWorkflowLatestCommittedTest(CosmosClientBuilder clientBuilder) { super(clientBuilder); } From a607a78b058909bee4016acb6562a2585fadfb55 Mon Sep 17 00:00:00 2001 From: Abhijeet Mohanty Date: Sun, 30 Aug 2026 15:02:40 -0400 Subject: [PATCH 14/14] Fix hotfix test and Kafka pipeline compatibility Restrict unsupported customer workflow modes, initialize PPCB fixtures for the thin-client group, route Confluent dependencies through authenticated CFS, and authenticate Testcontainers image pulls through ACR. --- eng/settings.xml | 2 +- .../PerPartitionCircuitBreakerE2ETests.java | 8 ++--- ...ustomerWorkflowDaoStyleOperationsTest.java | 3 ++ .../CustomerWorkflowRequestOptionsTest.java | 3 ++ ...rWorkflowSingleMasterAvailabilityTest.java | 3 ++ .../customer/CustomerWorkflowTestBase.java | 7 +++++ sdk/cosmos/kafka.yml | 30 +++++++++++++++++++ 7 files changed, 51 insertions(+), 5 deletions(-) diff --git a/eng/settings.xml b/eng/settings.xml index 8e65b655fd04..06591ce96ef2 100644 --- a/eng/settings.xml +++ b/eng/settings.xml @@ -7,7 +7,7 @@ azure-sdk-for-java Azure Artifacts Maven Mirror https://pkgs.dev.azure.com/azure-sdk/public/_packaging/azure-sdk-for-java/maven/v1 - external:*,!confluent,!repository.spring.milestone + external:*,!repository.spring.milestone 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 ccdc2937825c..6411743f5127 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 @@ -243,7 +243,7 @@ public PerPartitionCircuitBreakerE2ETests(CosmosClientBuilder cosmosClientBuilde super(cosmosClientBuilder); } - @BeforeClass(groups = {"circuit-breaker-misc-gateway", "circuit-breaker-misc-direct", "circuit-breaker-read-all-read-many", "multi-region"}) + @BeforeClass(groups = {"circuit-breaker-misc-gateway", "circuit-breaker-misc-direct", "circuit-breaker-read-all-read-many", "multi-region", "fi-thinclient-multi-master"}) public void beforeClass() { try (CosmosAsyncClient testClient = getClientBuilder().buildAsyncClient()) { RxDocumentClientImpl documentClient = (RxDocumentClientImpl) ReflectionUtils.getAsyncDocumentClient(testClient); @@ -4773,18 +4773,18 @@ private String resolveContainerIdByFaultInjectionOperationType(FaultInjectionOpe } } - @BeforeMethod(groups = { "circuit-breaker-misc-gateway", "circuit-breaker-misc-direct", "circuit-breaker-read-all-read-many", "multi-region" }, timeOut = 2 * SETUP_TIMEOUT, alwaysRun = true) + @BeforeMethod(groups = { "circuit-breaker-misc-gateway", "circuit-breaker-misc-direct", "circuit-breaker-read-all-read-many", "multi-region", "fi-thinclient-multi-master" }, timeOut = 2 * SETUP_TIMEOUT, alwaysRun = true) public void beforeMethod() throws Exception { // add a cool off time CosmosNettyLeakDetectorFactory.resetIdentifiedLeaks(); } - @AfterMethod(groups = { "circuit-breaker-misc-gateway", "circuit-breaker-misc-direct", "circuit-breaker-read-all-read-many", "multi-region" }, timeOut = SETUP_TIMEOUT, alwaysRun = true) + @AfterMethod(groups = { "circuit-breaker-misc-gateway", "circuit-breaker-misc-direct", "circuit-breaker-read-all-read-many", "multi-region", "fi-thinclient-multi-master" }, timeOut = SETUP_TIMEOUT, alwaysRun = true) public void afterMethod() throws Exception { logger.info("captureNettyLeaks: {}", captureNettyLeaks()); } - @AfterClass(groups = { "circuit-breaker-misc-gateway", "circuit-breaker-misc-direct", "circuit-breaker-read-all-read-many", "multi-region" }) + @AfterClass(groups = { "circuit-breaker-misc-gateway", "circuit-breaker-misc-direct", "circuit-breaker-read-all-read-many", "multi-region", "fi-thinclient-multi-master" }) public void afterClass() { CosmosClientBuilder clientBuilder = new CosmosClientBuilder() .endpoint(TestConfigurations.HOST) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowDaoStyleOperationsTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowDaoStyleOperationsTest.java index a375d6540879..c9e51baa3e5a 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowDaoStyleOperationsTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowDaoStyleOperationsTest.java @@ -49,6 +49,9 @@ public void afterClass() { @Test(groups = {"fi-customer-workflows"}, timeOut = TIMEOUT) public void crudReadAllPatchBatchAndBulkWorkflow() { + // Thin Client excluded-region and availability-strategy routing was fixed in PR #48432 in azure-cosmos 4.79.0. + skipIfThinClient("Customer DAO-style workflow"); + List excludedRegions = excludeFirstWritableRegion(); TestObject item = TestObject.create(); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowRequestOptionsTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowRequestOptionsTest.java index f46b517813a6..4a6a7a147ee0 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowRequestOptionsTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowRequestOptionsTest.java @@ -44,6 +44,9 @@ public void afterClass() { @Test(groups = {"fi-customer-workflows"}, timeOut = TIMEOUT) public void excludedRegionAndKeywordIdentifiersFlowAcrossOperations() { + // Gateway ReadConsistencyStrategy support was added in PR #48787 after azure-cosmos 4.81.0. + skipIfNotDirectMode("Customer request options workflow"); + String excludedRegion = this.writableRegions.get(0); List excludedRegions = Collections.singletonList(excludedRegion); TestObject item = TestObject.create(); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowSingleMasterAvailabilityTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowSingleMasterAvailabilityTest.java index d3249ae10940..fa844db25ea6 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowSingleMasterAvailabilityTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowSingleMasterAvailabilityTest.java @@ -50,6 +50,9 @@ public void afterClass() { @Test(groups = {"fi-sm-customer-workflows"}, timeOut = TIMEOUT) public void excludedReadableRegionRoutesReadToRemainingReadableRegion() { + // Gateway ReadConsistencyStrategy support was added in PR #48787 after azure-cosmos 4.81.0. + skipIfNotDirectMode("Customer excluded readable region workflow"); + TestObject item = TestObject.create(); this.container.createItem(item).block(); registerForCleanup(item); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowTestBase.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowTestBase.java index 5f35d2defc31..8b6db40eb49a 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowTestBase.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/workflows/customer/CustomerWorkflowTestBase.java @@ -395,6 +395,13 @@ protected final void skipIfNotGatewayMode(String scenarioName) { } } + protected final void skipIfThinClient(String scenarioName) { + RxDocumentClientImpl rxDocumentClient = (RxDocumentClientImpl) ReflectionUtils.getAsyncDocumentClient(this.client); + if (rxDocumentClient.useThinClient()) { + throw new SkipException(scenarioName + " does not apply when Thin Client is selected."); + } + } + /** * Skips fault-injection scenarios that cannot be injected for the gateway connection type. The gateway * internally retries 410/0, so {@code GONE} and {@code STALED_ADDRESSES_SERVER_GONE} rules are rejected at diff --git a/sdk/cosmos/kafka.yml b/sdk/cosmos/kafka.yml index 771b39c6e56a..82321b977487 100644 --- a/sdk/cosmos/kafka.yml +++ b/sdk/cosmos/kafka.yml @@ -14,6 +14,7 @@ extends: COSMOS.CLIENT_TELEMETRY_ENDPOINT: $(cosmos-client-telemetry-endpoint) COSMOS.CLIENT_TELEMETRY_COSMOS_ACCOUNT: $(cosmos-client-telemetry-cosmos-account) COSMOS_ACR_NAME: $(kafka-mcr-name) + TESTCONTAINERS_HUB_IMAGE_NAME_PREFIX: $(kafka-acr-login-server)/ CloudConfig: Public: ServiceConnection: azure-sdk-tests-cosmos @@ -35,3 +36,32 @@ extends: AdditionalVariables: - name: AdditionalArgs value: '' + PreTestRunSteps: + - script: | + if ! command -v docker &>/dev/null; then + echo "Docker not found; please install Docker or ensure it's available on the agent." + exit 1 + fi + + # wait for docker daemon to be ready + for i in {1..30}; do + if docker info >/dev/null 2>&1; then + echo "Docker is running" + break + fi + echo "Waiting for Docker to start... ($i/30)" + sleep 2 + done + + if ! docker info >/dev/null 2>&1; then + echo "Docker failed to start" + exit 1 + fi + displayName: 'Ensure Docker is installed and running' + - script: | + printf '%s' "$(kafka-acr-sp-client-secret)" | docker login "$(kafka-acr-login-server)" --username "$(kafka-acr-sp-client-id)" --password-stdin + displayName: 'Login to ACR for Testcontainers image cache' + PostSteps: + - script: docker logout "$(kafka-acr-login-server)" + displayName: 'Logout from ACR' + condition: always()