diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/MetadataHedgingIntegrationTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/MetadataHedgingIntegrationTests.java new file mode 100644 index 000000000000..c3e80fb62c2b --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/MetadataHedgingIntegrationTests.java @@ -0,0 +1,232 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.faultinjection; + +import com.azure.cosmos.BridgeInternal; +import com.azure.cosmos.CosmosAsyncClient; +import com.azure.cosmos.CosmosAsyncContainer; +import com.azure.cosmos.CosmosClientBuilder; +import com.azure.cosmos.CosmosDiagnostics; +import com.azure.cosmos.CosmosException; +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.test.faultinjection.CosmosFaultInjectionHelper; +import com.azure.cosmos.test.faultinjection.FaultInjectionConditionBuilder; +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 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.ArrayList; +import java.util.Arrays; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; + +import static org.assertj.core.api.AssertionsForClassTypes.assertThat; + +/** + * Multi-region integration tests for cross-region metadata hedging (port of .NET PR #5999). + *

+ * The PartitionKeyRange routing-map read is forced to re-read (via a {@code PARTITION_IS_SPLITTING} fault on a + * data operation) while a primary-region {@code RESPONSE_DELAY} is injected on the metadata read. With hedging + * enabled, the SDK issues a hedge to a second region -- surfaced as a {@code METADATA_HEDGE} diagnostics datum and + * the secondary region being contacted; with hedging disabled, no hedge is issued. + *

+ * Requires a multi-region account (>= 2 read regions) via {@code COSMOSDB_MULTI_REGION} / {@code ACCOUNT_HOST} + + * {@code ACCOUNT_KEY}. + */ +public class MetadataHedgingIntegrationTests extends FaultInjectionTestBase { + + private static final String METADATA_HEDGING_ENABLED_PROPERTY = "COSMOS.metadataHedgingEnabled"; + private static final String METADATA_HEDGING_THRESHOLD_PROPERTY = "COSMOS.metadataHedgingThresholdInMs"; + private static final String METADATA_HEDGE_DATUM = "METADATA_HEDGE"; + // Above the (lowered-for-test) hedge threshold, below the gateway request timeout. + private static final Duration PRIMARY_DELAY = Duration.ofSeconds(4); + + private CosmosAsyncClient client; + private CosmosAsyncContainer cosmosAsyncContainer; + private List writePreferredLocations; + + @Factory(dataProvider = "simpleClientBuildersWithJustDirectTcp") + public MetadataHedgingIntegrationTests(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + this.subscriberValidationTimeout = TIMEOUT; + } + + @BeforeClass(groups = {"multi-region"}, timeOut = TIMEOUT) + public void beforeClass() { + // Lower the threshold so a hedge fires quickly once the primary is delayed. + System.setProperty(METADATA_HEDGING_THRESHOLD_PROPERTY, "500"); + + this.client = getClientBuilder().buildAsyncClient(); + AsyncDocumentClient asyncDocumentClient = BridgeInternal.getContextClient(this.client); + GlobalEndpointManager globalEndpointManager = asyncDocumentClient.getGlobalEndpointManager(); + DatabaseAccount databaseAccount = globalEndpointManager.getLatestDatabaseAccount(); + + AccountLevelLocationContext writeable = getAccountLevelLocationContext(databaseAccount, true); + assertThat(writeable.serviceOrderedWriteableRegions.size()) + .as("Metadata hedging tests require a multi-region (multi-master) account with >= 2 regions") + .isGreaterThanOrEqualTo(2); + + this.writePreferredLocations = writeable.serviceOrderedWriteableRegions; + this.cosmosAsyncContainer = getSharedMultiPartitionCosmosContainerWithIdAsPartitionKey(this.client); + } + + @AfterClass(groups = {"multi-region"}, timeOut = TIMEOUT, alwaysRun = true) + public void afterClass() { + System.clearProperty(METADATA_HEDGING_ENABLED_PROPERTY); + System.clearProperty(METADATA_HEDGING_THRESHOLD_PROPERTY); + safeClose(this.client); + } + + @Test(groups = {"multi-region"}, timeOut = 4 * TIMEOUT) + public void partitionKeyRangeHedged_whenEnabledAndPrimarySlow() { + runPartitionKeyRangeScenario(true); + } + + @Test(groups = {"multi-region"}, timeOut = 4 * TIMEOUT) + public void partitionKeyRangeNotHedged_whenDisabled() { + runPartitionKeyRangeScenario(false); + } + + private void runPartitionKeyRangeScenario(boolean hedgingEnabled) { + System.setProperty(METADATA_HEDGING_ENABLED_PROPERTY, Boolean.toString(hedgingEnabled)); + CosmosAsyncClient testClient = getClientBuilder() + .contentResponseOnWriteEnabled(true) + .preferredRegions(this.writePreferredLocations) + .buildAsyncClient(); + + CosmosAsyncContainer container = testClient + .getDatabase(this.cosmosAsyncContainer.getDatabase().getId()) + .getContainer(this.cosmosAsyncContainer.getId()); + + // Delay the PartitionKeyRange metadata read in the primary region. + FaultInjectionRule pkRangesDelayRule = + new FaultInjectionRuleBuilder("pkrange-delay-" + UUID.randomUUID()) + .condition( + new FaultInjectionConditionBuilder() + .region(this.writePreferredLocations.get(0)) + .operationType(FaultInjectionOperationType.METADATA_REQUEST_PARTITION_KEY_RANGES) + .build()) + .result( + FaultInjectionResultBuilders + .getResultBuilder(FaultInjectionServerErrorType.RESPONSE_DELAY) + .delay(PRIMARY_DELAY) + .times(1) + .build()) + .duration(Duration.ofMinutes(5)) + .build(); + + // Use a partition split on a write to force a routing-map (PartitionKeyRange) refresh AFTER the rule is active. + FaultInjectionRule splitRule = + new FaultInjectionRuleBuilder("split-" + UUID.randomUUID()) + .condition( + new FaultInjectionConditionBuilder() + .region(this.writePreferredLocations.get(0)) + .operationType(FaultInjectionOperationType.CREATE_ITEM) + .build()) + .result( + FaultInjectionResultBuilders + .getResultBuilder(FaultInjectionServerErrorType.PARTITION_IS_SPLITTING) + .times(1) + .build()) + .duration(Duration.ofMinutes(5)) + .build(); + + try { + // Populate the collection + PartitionKeyRange caches first. + for (int i = 0; i < 5; i++) { + container.createItem(TestObject.create()).block(); + } + + CosmosFaultInjectionHelper + .configureFaultInjectionRules(container, Arrays.asList(pkRangesDelayRule, splitRule)) + .block(); + + // This write hits PARTITION_IS_SPLITTING -> forces a PartitionKeyRange refresh -> hits the primary delay. + CosmosDiagnostics diagnostics = + container.createItem(TestObject.create()).block().getDiagnostics(); + + assertHedging(diagnostics, hedgingEnabled); + } catch (CosmosException e) { + throw new AssertionError("createItem should have succeeded. " + e.getDiagnostics(), e); + } finally { + pkRangesDelayRule.disable(); + splitRule.disable(); + System.clearProperty(METADATA_HEDGING_ENABLED_PROPERTY); + safeClose(testClient); + } + } + + private void assertHedging(CosmosDiagnostics diagnostics, boolean hedgingEnabled) { + assertThat(diagnostics).isNotNull(); + String json = diagnostics.toString(); + + if (hedgingEnabled) { + assertThat(json.contains(METADATA_HEDGE_DATUM)) + .as("expected a Metadata Hedge datum when hedging is enabled and the primary is slow; diagnostics: %s", json) + .isTrue(); + assertThat(json.contains("\"hedgeFired\":true")) + .as("expected the hedge to have fired; diagnostics: %s", json) + .isTrue(); + } else { + assertThat(json.contains(METADATA_HEDGE_DATUM)) + .as("no Metadata Hedge datum expected when hedging is disabled; diagnostics: %s", json) + .isFalse(); + } + } + + private static AccountLevelLocationContext getAccountLevelLocationContext(DatabaseAccount databaseAccount, boolean writeOnly) { + Iterator locationIterator = + writeOnly ? databaseAccount.getWritableLocations().iterator() : databaseAccount.getReadableLocations().iterator(); + + List serviceOrderedReadableRegions = new ArrayList<>(); + List serviceOrderedWriteableRegions = new ArrayList<>(); + Map regionMap = new ConcurrentHashMap<>(); + + while (locationIterator.hasNext()) { + DatabaseAccountLocation accountLocation = locationIterator.next(); + regionMap.put(accountLocation.getName(), accountLocation.getEndpoint()); + + if (writeOnly) { + serviceOrderedWriteableRegions.add(accountLocation.getName()); + } else { + serviceOrderedReadableRegions.add(accountLocation.getName()); + } + } + + return new AccountLevelLocationContext( + serviceOrderedReadableRegions, + serviceOrderedWriteableRegions, + regionMap); + } + + private static class AccountLevelLocationContext { + private final List serviceOrderedReadableRegions; + private final List serviceOrderedWriteableRegions; + private final Map regionNameToEndpoint; + + AccountLevelLocationContext( + List serviceOrderedReadableRegions, + List serviceOrderedWriteableRegions, + Map regionNameToEndpoint) { + + this.serviceOrderedReadableRegions = serviceOrderedReadableRegions; + this.serviceOrderedWriteableRegions = serviceOrderedWriteableRegions; + this.regionNameToEndpoint = regionNameToEndpoint; + } + } +} diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/caches/MetadataHedgingStrategyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/caches/MetadataHedgingStrategyTest.java new file mode 100644 index 000000000000..1e17d25e3568 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/caches/MetadataHedgingStrategyTest.java @@ -0,0 +1,171 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.implementation.caches; + +import com.azure.cosmos.implementation.HttpConstants; +import org.testng.annotations.Test; +import reactor.core.publisher.Mono; + +import java.time.Duration; +import java.util.function.Function; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; + +/** + * Unit tests for the primary-authoritative metadata hedging core. These exercise the pure + * regional-failure classifier and the Reactor race semantics (no live client / store model), + * mirroring the key cells of the .NET MetadataHedgingStrategyTests. + */ +public class MetadataHedgingStrategyTest { + + private static final Duration BLOCK = Duration.ofSeconds(15); + + // Sentinel branch errors + a classifier so the race can be tested without CosmosExceptions. + private static final class RegionalError extends RuntimeException { + RegionalError(String m) { super(m); } + } + private static final class DefinitiveError extends RuntimeException { + DefinitiveError(String m) { super(m); } + } + private static final Function CLASSIFIER = + e -> (e instanceof DefinitiveError) + ? MetadataHedgingStrategy.BranchOutcome.DEFINITIVE + : MetadataHedgingStrategy.BranchOutcome.REGIONAL_FAILURE; + + private final MetadataHedgingStrategy strategy = new MetadataHedgingStrategy(); + + private MetadataHedgingStrategy.HedgeResult run(Mono primary, Mono hedge, Duration threshold) { + return strategy.executeHedged(primary, "primaryRegion", hedge, "hedgeRegion", threshold, CLASSIFIER) + .block(BLOCK); + } + + // ---- classifier (pure) ---- + + @Test(groups = "unit") + public void isRegionalFailure_regionalSet() { + assertTrue(MetadataHedgingStrategy.isRegionalFailure(HttpConstants.StatusCodes.SERVICE_UNAVAILABLE, 0)); + assertTrue(MetadataHedgingStrategy.isRegionalFailure(HttpConstants.StatusCodes.INTERNAL_SERVER_ERROR, 0)); + assertTrue(MetadataHedgingStrategy.isRegionalFailure(HttpConstants.StatusCodes.REQUEST_TIMEOUT, 0)); + assertTrue(MetadataHedgingStrategy.isRegionalFailure( + HttpConstants.StatusCodes.GONE, HttpConstants.SubStatusCodes.LEASE_NOT_FOUND)); + assertTrue(MetadataHedgingStrategy.isRegionalFailure( + HttpConstants.StatusCodes.FORBIDDEN, HttpConstants.SubStatusCodes.DATABASE_ACCOUNT_NOTFOUND)); + assertTrue(MetadataHedgingStrategy.isRegionalFailure(200, HttpConstants.SubStatusCodes.GATEWAY_ENDPOINT_UNAVAILABLE)); + assertTrue(MetadataHedgingStrategy.isRegionalFailure(200, HttpConstants.SubStatusCodes.GATEWAY_ENDPOINT_READ_TIMEOUT)); + } + + @Test(groups = "unit") + public void isRegionalFailure_definitiveSet() { + assertFalse(MetadataHedgingStrategy.isRegionalFailure(HttpConstants.StatusCodes.NOTFOUND, 0)); // 404 + assertFalse(MetadataHedgingStrategy.isRegionalFailure(409, 0)); // conflict + assertFalse(MetadataHedgingStrategy.isRegionalFailure(412, 0)); // precondition + assertFalse(MetadataHedgingStrategy.isRegionalFailure(401, 0)); // unauthorized + assertFalse(MetadataHedgingStrategy.isRegionalFailure(HttpConstants.StatusCodes.FORBIDDEN, 0)); // plain 403 + assertFalse(MetadataHedgingStrategy.isRegionalFailure(HttpConstants.StatusCodes.GONE, 0)); // plain 410 (split) + assertFalse(MetadataHedgingStrategy.isRegionalFailure(HttpConstants.StatusCodes.TOO_MANY_REQUESTS, 0)); + } + + @Test(groups = "unit") + public void classifyThrowable_nonCosmosIsRegional() { + assertEquals(MetadataHedgingStrategy.classifyThrowable(new RuntimeException("boom")), + MetadataHedgingStrategy.BranchOutcome.REGIONAL_FAILURE); + } + + // ---- race semantics ---- + + @Test(groups = "unit") + public void primaryFast_noHedge() { + MetadataHedgingStrategy.HedgeResult r = + run(Mono.just("P"), Mono.just("H"), Duration.ofMillis(300)); + assertFalse(r.isError()); + assertEquals(r.getValue(), "P"); + assertFalse(r.isHedgeFired()); + assertFalse(r.isHedgeWon()); + } + + @Test(groups = "unit") + public void primarySlow_hedgeWins() { + MetadataHedgingStrategy.HedgeResult r = run( + Mono.just("P").delayElement(Duration.ofMillis(800)), + Mono.just("H"), + Duration.ofMillis(100)); + assertFalse(r.isError()); + assertEquals(r.getValue(), "H"); + assertTrue(r.isHedgeFired()); + assertTrue(r.isHedgeWon()); + assertEquals(r.getWinningRegion(), "hedgeRegion"); + } + + @Test(groups = "unit") + public void primaryFastDefinitiveError_noHedge() { + MetadataHedgingStrategy.HedgeResult r = run( + Mono.error(new DefinitiveError("404")), + Mono.just("H"), + Duration.ofMillis(300)); + assertTrue(r.isError()); + assertTrue(r.getError() instanceof DefinitiveError); + assertFalse(r.isHedgeFired()); + assertFalse(r.isHedgeWon()); + } + + @Test(groups = "unit") + public void primaryRegionalError_hedgeWins() { + MetadataHedgingStrategy.HedgeResult r = run( + Mono.error(new RegionalError("503")), + Mono.just("H"), + Duration.ofMillis(100)); + assertFalse(r.isError()); + assertEquals(r.getValue(), "H"); + assertTrue(r.isHedgeFired()); + assertTrue(r.isHedgeWon()); + } + + @Test(groups = "unit") + public void primaryFastRegionalFailure_hedgeFiresBeforeThreshold() { + // Large threshold: if the hedge only fired on the threshold timer this would take ~5s. Because + // the primary fails regionally immediately, the hedge must fire right away (mirroring .NET + // "primary hit a regional failure -> fire one hedge") and win well under the threshold. + Duration largeThreshold = Duration.ofSeconds(5); + long startNanos = System.nanoTime(); + MetadataHedgingStrategy.HedgeResult r = run( + Mono.error(new RegionalError("503-fast")), + Mono.just("H").delayElement(Duration.ofMillis(50)), + largeThreshold); + Duration elapsed = Duration.ofNanos(System.nanoTime() - startNanos); + assertFalse(r.isError()); + assertEquals(r.getValue(), "H"); + assertTrue(r.isHedgeFired()); + assertTrue(r.isHedgeWon()); + assertTrue(elapsed.compareTo(Duration.ofSeconds(2)) < 0, + "hedge should fire immediately on a fast primary regional failure, not wait the full " + + "threshold; elapsed=" + elapsed); + } + + @Test(groups = "unit") + public void primaryRegional_hedgeAlsoFails_primaryReturned() { + MetadataHedgingStrategy.HedgeResult r = run( + Mono.error(new RegionalError("primary-503")), + Mono.error(new RegionalError("hedge-401")), + Duration.ofMillis(100)); + assertTrue(r.isError()); + assertEquals(r.getError().getMessage(), "primary-503"); + assertTrue(r.isHedgeFired()); + assertFalse(r.isHedgeWon()); + } + + @Test(groups = "unit") + public void slowPrimaryDefinitive_beatsHedge_primaryAuthoritative() { + // Threshold elapses (100ms) so the hedge fires; the primary then settles Definitive at 250ms, + // before the hedge succeeds at ~100+400=500ms. The definitive primary must win. + MetadataHedgingStrategy.HedgeResult r = run( + Mono.error(new DefinitiveError("409")).delaySubscription(Duration.ofMillis(250)), + Mono.just("H").delayElement(Duration.ofMillis(400)), + Duration.ofMillis(100)); + assertTrue(r.isError()); + assertTrue(r.getError() instanceof DefinitiveError); + assertTrue(r.isHedgeFired()); + assertFalse(r.isHedgeWon()); + } +} diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index fe07a4d47941..7438e157eedc 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -3,8 +3,7 @@ ### 4.82.0-beta.1 (Unreleased) #### Features Added - -#### Breaking Changes +* Added **cross-region metadata hedging** for the two hot-path metadata cache reads (Collection Read and PartitionKeyRange ReadFeed). When the primary region is slow on one of these reads (or returns a regional failure such as `503`/`500`/`410`/network error), the SDK now issues a single hedged request to another region -- after a configurable threshold (default 1.5s) when the primary is slow, or immediately when the primary hits a regional failure -- and returns the first acceptable response, trimming latency tails on cold-start cache population and steady-state refresh. The primary region always remains authoritative, so a slow or misconfigured secondary can never change the operation result. Opt-in follows the account's PPAF (Per-Partition Automatic Failover) state and can be force-enabled or disabled via the `AZURE_COSMOS_METADATA_HEDGING_ENABLED` / `COSMOS_METADATA_HEDGING_ENABLED` environment variable or the `COSMOS.metadataHedgingEnabled` system property. Port of .NET [azure-cosmos-dotnet-v3 PR 5999](https://github.com/Azure/azure-cosmos-dotnet-v3/pull/5999). #### Bugs Fixed * Unified request-level consistency override behavior across transports: invalid attempts to upgrade the request consistency level above the account default are now silently ignored instead of returning `BadRequest` in some gateway paths. - See PR [49606](https://github.com/Azure/azure-sdk-for-java/pull/49606). diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java index cf1a812e10f2..78efa555725c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java @@ -55,6 +55,20 @@ public class Configs { private static final String THINCLIENT_ENABLED = "COSMOS.THINCLIENT_ENABLED"; private static final String THINCLIENT_ENABLED_VARIABLE = "COSMOS_THINCLIENT_ENABLED"; + // Metadata cross-region hedging (port of .NET azure-cosmos-dotnet-v3 PR #5999). + // Tri-state opt-in: + // - explicit override via the COSMOS.metadataHedgingEnabled system property, or the + // COSMOS_METADATA_HEDGING_ENABLED / AZURE_COSMOS_METADATA_HEDGING_ENABLED (.NET parity) env vars + // - when unset, the caller follows the account's PPAF (Per-Partition Automatic Failover) state, + // resolved once at client construction. + private static final String METADATA_HEDGING_ENABLED = "COSMOS.metadataHedgingEnabled"; + private static final String METADATA_HEDGING_ENABLED_VARIABLE = "COSMOS_METADATA_HEDGING_ENABLED"; + private static final String METADATA_HEDGING_ENABLED_AZURE_VARIABLE = "AZURE_COSMOS_METADATA_HEDGING_ENABLED"; + // Fixed configurable threshold (Q4): fire the single hedge after this delay when the primary is slow. + private static final int DEFAULT_METADATA_HEDGING_THRESHOLD_IN_MS = 1_500; + private static final String METADATA_HEDGING_THRESHOLD_IN_MS = "COSMOS.metadataHedgingThresholdInMs"; + private static final String METADATA_HEDGING_THRESHOLD_IN_MS_VARIABLE = "COSMOS_METADATA_HEDGING_THRESHOLD_IN_MS"; + private static final boolean DEFAULT_NETTY_HTTP_CLIENT_METRICS_ENABLED = false; private static final String NETTY_HTTP_CLIENT_METRICS_ENABLED = "COSMOS.NETTY_HTTP_CLIENT_METRICS_ENABLED"; private static final String NETTY_HTTP_CLIENT_METRICS_ENABLED_VARIABLE = "COSMOS_NETTY_HTTP_CLIENT_METRICS_ENABLED"; @@ -585,6 +599,54 @@ public static boolean isThinClientEnabled() { return DEFAULT_THINCLIENT_ENABLED; } + /** + * Tri-state opt-in for metadata cross-region hedging. + * + * @return {@code TRUE}/{@code FALSE} when explicitly overridden via the + * {@code COSMOS.metadataHedgingEnabled} system property or the + * {@code COSMOS_METADATA_HEDGING_ENABLED} / {@code AZURE_COSMOS_METADATA_HEDGING_ENABLED} + * environment variables; {@code null} when unset (the caller should then follow the + * account's PPAF state). + */ + public static Boolean getMetadataHedgingEnabledOverride() { + String valueFromSystemProperty = System.getProperty(METADATA_HEDGING_ENABLED); + if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) { + return Boolean.parseBoolean(valueFromSystemProperty); + } + + String valueFromEnvVariable = System.getenv(METADATA_HEDGING_ENABLED_VARIABLE); + if (valueFromEnvVariable != null && !valueFromEnvVariable.isEmpty()) { + return Boolean.parseBoolean(valueFromEnvVariable); + } + + String valueFromAzureEnvVariable = System.getenv(METADATA_HEDGING_ENABLED_AZURE_VARIABLE); + if (valueFromAzureEnvVariable != null && !valueFromAzureEnvVariable.isEmpty()) { + return Boolean.parseBoolean(valueFromAzureEnvVariable); + } + + return null; + } + + /** + * The fixed, configurable threshold after which a slow primary-region metadata read triggers a + * single cross-region hedge. Defaults to 1.5s; overridable via the + * {@code COSMOS.metadataHedgingThresholdInMs} system property or + * {@code COSMOS_METADATA_HEDGING_THRESHOLD_IN_MS} environment variable. + */ + public static Duration getMetadataHedgingThreshold() { + String valueFromSystemProperty = System.getProperty(METADATA_HEDGING_THRESHOLD_IN_MS); + if (valueFromSystemProperty != null && !valueFromSystemProperty.isEmpty()) { + return Duration.ofMillis(Long.parseLong(valueFromSystemProperty.trim())); + } + + String valueFromEnvVariable = System.getenv(METADATA_HEDGING_THRESHOLD_IN_MS_VARIABLE); + if (valueFromEnvVariable != null && !valueFromEnvVariable.isEmpty()) { + return Duration.ofMillis(Long.parseLong(valueFromEnvVariable.trim())); + } + + return Duration.ofMillis(DEFAULT_METADATA_HEDGING_THRESHOLD_IN_MS); + } + public static boolean isNettyHttpClientMetricsEnabled() { return Boolean.parseBoolean( System.getProperty(NETTY_HTTP_CLIENT_METRICS_ENABLED, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/MetadataDiagnosticsContext.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/MetadataDiagnosticsContext.java index 1dccc29b574c..44ac9f82bc09 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/MetadataDiagnosticsContext.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/MetadataDiagnosticsContext.java @@ -69,6 +69,28 @@ public ContainerLookupMetadataDiagnostics( } } + @JsonSerialize(using = MetaDataDiagnosticSerializer.class) + public static class MetadataHedgeDiagnostics extends MetadataDiagnostics { + public volatile String activityId; + public volatile boolean hedgeFired; + public volatile boolean hedgeWon; + public volatile String winningRegion; + + public MetadataHedgeDiagnostics( + Instant startTimeUTC, + Instant endTimeUTC, + String activityId, + boolean hedgeFired, + boolean hedgeWon, + String winningRegion) { + super(startTimeUTC, endTimeUTC, MetadataType.METADATA_HEDGE); + this.activityId = activityId; + this.hedgeFired = hedgeFired; + this.hedgeWon = hedgeWon; + this.winningRegion = winningRegion; + } + } + static class MetaDataDiagnosticSerializer extends StdSerializer { private static final long serialVersionUID = -6585518025594634820L; @@ -95,6 +117,14 @@ public void serialize(MetadataDiagnostics metaDataDiagnostic, JsonGenerator json jsonGenerator.writeStringField("collectionRid", ((ContainerLookupMetadataDiagnostics)metaDataDiagnostic).collectionRid); } + if (metaDataDiagnostic instanceof MetadataHedgeDiagnostics) { + MetadataHedgeDiagnostics hedge = (MetadataHedgeDiagnostics) metaDataDiagnostic; + jsonGenerator.writeStringField("activityId", hedge.activityId); + jsonGenerator.writeBooleanField("hedgeFired", hedge.hedgeFired); + jsonGenerator.writeBooleanField("hedgeWon", hedge.hedgeWon); + jsonGenerator.writeStringField("winningRegion", hedge.winningRegion); + } + jsonGenerator.writeEndObject(); } } @@ -103,7 +133,8 @@ public enum MetadataType { CONTAINER_LOOK_UP, PARTITION_KEY_RANGE_LOOK_UP, SERVER_ADDRESS_LOOKUP, - MASTER_ADDRESS_LOOK_UP + MASTER_ADDRESS_LOOK_UP, + METADATA_HEDGE } } 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 aa043d9df487..72555c5550f9 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 @@ -31,6 +31,7 @@ import com.azure.cosmos.implementation.batch.SinglePartitionKeyServerBatchRequest; import com.azure.cosmos.implementation.caches.AsyncCache; import com.azure.cosmos.implementation.caches.RxClientCollectionCache; +import com.azure.cosmos.implementation.caches.MetadataHedgingStrategy; import com.azure.cosmos.implementation.caches.RxCollectionCache; import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; import com.azure.cosmos.implementation.clienttelemetry.ClientTelemetry; @@ -944,6 +945,12 @@ public void init(CosmosClientMetadataCachesSnapshot metadataCachesSnapshot, Func DatabaseAccount databaseAccountSnapshot = this.initializeGatewayConfigurationReader(); this.resetSessionContainerIfNeeded(databaseAccountSnapshot); + MetadataHedgingStrategy metadataHedgingStrategy = new MetadataHedgingStrategy( + this.globalEndpointManager, + this.globalPartitionEndpointManagerForPerPartitionAutomaticFailover, + Configs.getMetadataHedgingEnabledOverride(), + Configs.getMetadataHedgingThreshold()); + if (metadataCachesSnapshot != null) { AsyncCache nameCache = metadataCachesSnapshot.getCollectionInfoByNameCache(); AsyncCache idCache = metadataCachesSnapshot.getCollectionInfoByIdCache(); @@ -954,7 +961,8 @@ public void init(CosmosClientMetadataCachesSnapshot metadataCachesSnapshot, Func this, this.retryPolicy, nameCache, - idCache + idCache, + metadataHedgingStrategy ); } else { // Cache data could not be deserialized (e.g., old format); fall back to fresh fetch @@ -962,19 +970,22 @@ public void init(CosmosClientMetadataCachesSnapshot metadataCachesSnapshot, Func this.sessionContainer, this.gatewayProxy, this, - this.retryPolicy); + this.retryPolicy, + metadataHedgingStrategy); } } else { this.collectionCache = new RxClientCollectionCache(this, this.sessionContainer, this.gatewayProxy, this, - this.retryPolicy); + this.retryPolicy, + metadataHedgingStrategy); } this.resetSessionTokenRetryPolicy = new ResetSessionTokenRetryPolicyFactory(this.sessionContainer, this.collectionCache, this.retryPolicy); this.partitionKeyRangeCache = new RxPartitionKeyRangeCache(RxDocumentClientImpl.this, - collectionCache); + collectionCache, + metadataHedgingStrategy); updateGatewayProxy(); updateThinProxy(); @@ -7285,6 +7296,15 @@ private Flux> nonDocumentReadFeedInternal( RxDocumentServiceRequest request = RxDocumentServiceRequest.create(this, OperationType.ReadFeed, resourceType, resourceLink, requestHeaders, nonNullOptions); + + // Honor per-request excluded regions on non-document ReadFeed reads so a region-pinned read + // (e.g. a PartitionKeyRange metadata-hedging branch) is actually routed to its intended + // region. Only applied when excluded regions are explicitly set, so all other callers are + // unaffected. + List readFeedExcludedRegions = nonNullOptions.getExcludedRegions(); + if (readFeedExcludedRegions != null && !readFeedExcludedRegions.isEmpty()) { + request.requestContext.setExcludeRegions(readFeedExcludedRegions); + } return request; }; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/caches/MetadataHedgingStrategy.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/caches/MetadataHedgingStrategy.java new file mode 100644 index 000000000000..d85685465636 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/caches/MetadataHedgingStrategy.java @@ -0,0 +1,396 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.implementation.caches; + +import com.azure.cosmos.CosmosException; +import com.azure.cosmos.implementation.GlobalEndpointManager; +import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.implementation.RxDocumentServiceRequest; +import com.azure.cosmos.implementation.RxDocumentServiceResponse; +import com.azure.cosmos.implementation.perPartitionAutomaticFailover.GlobalPartitionEndpointManagerForPerPartitionAutomaticFailover; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Consumer; +import java.util.function.Function; + +/** + * Bounded, primary-authoritative cross-region hedging for the two hot-path metadata cache reads + * (Collection Read and PartitionKeyRange ReadFeed). This is the Java port of the .NET metadata hedging + * strategy (Azure/azure-cosmos-dotnet-v3 PR #5999). + *

+ * While this class is public it is NOT part of the published public API; it is used internally by + * the SDK caches only. + *

+ * The PartitionKeyRange ReadFeed read is realized as region-pinned full routing-map racing: each + * branch reads all of its pages from a single pinned region (see {@code RxPartitionKeyRangeCache}), rather + * than a literal first-page-only race as in the .NET reference. This preserves continuation integrity + * (each branch's page chain stays within one region) while achieving the same latency win. + *

+ * Design invariant -- the primary is authoritative: a hedge can only win by being faster on a + * regional failure or latency; it can never override a definitive answer the primary has + * already produced (e.g. a 404 for a deleted collection, a 409, a 412). See {@code internal-spec.md}. + *

+ * The core race ({@link #executeHedged}) is intentionally kept generic over the payload type so it can + * be unit-tested in isolation, mirroring the proven data-plane hedging pattern + * ({@code RxDocumentClientImpl.executePointOperationWithAvailabilityStrategy}: a delayed hedge plus a + * first-acceptable combinator). Each branch is expected to be pinned to a single region with inner + * cross-region retry suppressed by the caller so that "one hedge" and "primary authoritative" hold. + */ +public class MetadataHedgingStrategy { + + private static final Duration DEFAULT_THRESHOLD = Duration.ofMillis(1500); + + // Per-client collaborators (null in the unit-test / no-arg construction path). + private final GlobalEndpointManager globalEndpointManager; + private final GlobalPartitionEndpointManagerForPerPartitionAutomaticFailover perPartitionAutomaticFailoverManager; + // Tri-state opt-in: TRUE/FALSE force on/off; null -> follow the account PPAF state. + private final Boolean enabledOverride; + private final Duration threshold; + + /** Test / core-only constructor: the generic race ({@link #executeHedged}) can be used without a client. */ + public MetadataHedgingStrategy() { + this(null, null, null, DEFAULT_THRESHOLD); + } + + public MetadataHedgingStrategy( + GlobalEndpointManager globalEndpointManager, + GlobalPartitionEndpointManagerForPerPartitionAutomaticFailover perPartitionAutomaticFailoverManager, + Boolean enabledOverride, + Duration threshold) { + this.globalEndpointManager = globalEndpointManager; + this.perPartitionAutomaticFailoverManager = perPartitionAutomaticFailoverManager; + this.enabledOverride = enabledOverride; + this.threshold = threshold == null ? DEFAULT_THRESHOLD : threshold; + } + + /** + * The effective per-request opt-in: an explicit override wins; otherwise follow the account's live PPAF state + * (there is no preview/GA package split in Java, so the .NET "preview=on" arm collapses to this). + */ + public boolean isEffectivelyEnabled() { + if (this.enabledOverride != null) { + return this.enabledOverride; + } + return this.perPartitionAutomaticFailoverManager != null + && this.perPartitionAutomaticFailoverManager.isPerPartitionAutomaticFailoverEnabled(); + } + + public Duration getThreshold() { + return this.threshold; + } + + /** + * The ordered applicable read regions for a request (honoring excluded + per-partition-unavailable regions). + * Returns an empty list when there is no endpoint manager (test path) or fewer than two regions. + */ + public List getApplicableReadRegions(RxDocumentServiceRequest request) { + if (this.globalEndpointManager == null) { + return new ArrayList<>(); + } + List ctxs = + new ArrayList<>(this.globalEndpointManager.getApplicableReadRegionalRoutingContexts(request)); + return ctxs; + } + + public String getRegionName(RegionalRoutingContext ctx, RxDocumentServiceRequest request) { + if (this.globalEndpointManager == null || ctx == null || ctx.getGatewayRegionalEndpoint() == null) { + return null; + } + return this.globalEndpointManager.getRegionName(ctx.getGatewayRegionalEndpoint(), request.getOperationType()); + } + + /** + * Exclude every applicable region except {@code target}, so a branch is pinned to a single region with inner + * cross-region retry suppressed. Mirrors {@code RxDocumentClientImpl.getEffectiveExcludedRegionsForHedging}. + */ + public static List excludedRegionsForTarget(List baseExcluded, List applicable, String target) { + List excluded = new ArrayList<>(); + if (baseExcluded != null) { + excluded.addAll(baseExcluded); + } + for (String r : applicable) { + if (r != null && !r.equals(target) && !excluded.contains(r)) { + excluded.add(r); + } + } + return excluded; + } + + /** + * Run a single-request metadata read (Collection Read) with hedging when eligible. Each branch is a + * region-pinned clone (via {@code routeToLocation} + excluded-regions); otherwise the original send is returned + * unchanged. + */ + public Mono executeMetadataRead( + RxDocumentServiceRequest request, + Function> sendFunc, + Function classifier, + Consumer> telemetrySink) { + + if (!isEffectivelyEnabled() || this.globalEndpointManager == null) { + return sendFunc.apply(request); + } + + List ctxs = getApplicableReadRegions(request); + if (ctxs.size() < 2) { + return sendFunc.apply(request); + } + + RegionalRoutingContext primaryCtx = ctxs.get(0); + RegionalRoutingContext hedgeCtx = ctxs.get(1); + String primaryRegion = getRegionName(primaryCtx, request); + String hedgeRegion = getRegionName(hedgeCtx, request); + + List allRegions = new ArrayList<>(); + for (RegionalRoutingContext ctx : ctxs) { + String name = getRegionName(ctx, request); + if (name != null && !allRegions.contains(name)) { + allRegions.add(name); + } + } + List baseExcluded = request.requestContext == null ? null : request.requestContext.getExcludeRegions(); + + RxDocumentServiceRequest primaryReq = request.clone(); + primaryReq.requestContext.routeToLocation(primaryCtx); + primaryReq.requestContext.setExcludeRegions(excludedRegionsForTarget(baseExcluded, allRegions, primaryRegion)); + + RxDocumentServiceRequest hedgeReq = request.clone(); + hedgeReq.requestContext.routeToLocation(hedgeCtx); + hedgeReq.requestContext.setExcludeRegions(excludedRegionsForTarget(baseExcluded, allRegions, hedgeRegion)); + + Mono primaryMono = Mono.defer(() -> sendFunc.apply(primaryReq)); + Mono hedgeMono = Mono.defer(() -> sendFunc.apply(hedgeReq)); + + return executeHedged(primaryMono, primaryRegion, hedgeMono, hedgeRegion, this.threshold, classifier) + .flatMap(result -> { + if (telemetrySink != null && result.isHedgeFired()) { + telemetrySink.accept(result); + } + return result.unwrap(); + }); + } + + /** + * Classification of a single branch's completed outcome. Only a {@link #REGIONAL_FAILURE} is worth hedging; + * a {@link #DEFINITIVE} error is authoritative and returned verbatim. + */ + public enum BranchOutcome { + SUCCESS, + REGIONAL_FAILURE, + DEFINITIVE + } + + private enum Origin { PRIMARY, HEDGE } + + /** + * The outcome of a hedged execution: the winning value or error plus the hedge telemetry + * ({@code HedgeFired}/{@code HedgeWon}/{@code winningRegion}). + */ + public static final class HedgeResult { + private final T value; + private final Throwable error; + private final boolean hedgeFired; + private final boolean hedgeWon; + private final String winningRegion; + + private HedgeResult(T value, Throwable error, boolean hedgeFired, boolean hedgeWon, String winningRegion) { + this.value = value; + this.error = error; + this.hedgeFired = hedgeFired; + this.hedgeWon = hedgeWon; + this.winningRegion = winningRegion; + } + + public boolean isError() { + return this.error != null; + } + + public T getValue() { + return this.value; + } + + public Throwable getError() { + return this.error; + } + + public boolean isHedgeFired() { + return this.hedgeFired; + } + + public boolean isHedgeWon() { + return this.hedgeWon; + } + + public String getWinningRegion() { + return this.winningRegion; + } + + /** Unwrap into the caller's Mono: propagate the error verbatim, or emit the value. */ + public Mono unwrap() { + return this.isError() ? Mono.error(this.error) : Mono.just(this.value); + } + } + + // An internal, never-erroring materialization of a single branch attempt. + private static final class Attempt { + private final Origin origin; + private final String region; + private final T value; + private final Throwable error; + private final BranchOutcome outcome; + + private Attempt(Origin origin, String region, T value, Throwable error, BranchOutcome outcome) { + this.origin = origin; + this.region = region; + this.value = value; + this.error = error; + this.outcome = outcome; + } + + private boolean isWinner() { + // Primary is authoritative on any non-regional outcome (success or definitive error). + // The hedge only wins by producing a good answer. + return this.origin == Origin.PRIMARY + ? this.outcome != BranchOutcome.REGIONAL_FAILURE + : this.outcome == BranchOutcome.SUCCESS; + } + } + + /** + * Run a primary-authoritative hedged execution. + *

+ * Semantics (mirrors the .NET flow): + *

+ * + * @param primary the primary-region attempt (must be pinned to a single region by the caller) + * @param primaryRegion the primary region name (for telemetry) + * @param hedge the hedge-region attempt (pinned to a different single region) + * @param hedgeRegion the hedge region name (for telemetry) + * @param threshold the delay after which the hedge is fired if the primary has not settled + * @param classifier maps a branch error to a {@link BranchOutcome} + */ + public Mono> executeHedged( + Mono primary, + String primaryRegion, + Mono hedge, + String hedgeRegion, + Duration threshold, + Function classifier) { + + AtomicBoolean hedgeStarted = new AtomicBoolean(false); + + // Materialize the primary once so the race, the early-hedge gate and the fallback all share a + // single subscription. + Mono> primaryAttempt = primary + .map(v -> new Attempt<>(Origin.PRIMARY, primaryRegion, v, null, BranchOutcome.SUCCESS)) + .onErrorResume(e -> Mono.just( + new Attempt(Origin.PRIMARY, primaryRegion, null, e, classifier.apply(e)))) + .cache(); + + // Fire the single hedge as soon as EITHER the threshold elapses (the primary is slow) OR the + // primary settles with a regional failure before the threshold -- mirroring the .NET flow + // ("primary is slow, or hit a regional failure -> fire one hedge"). A non-regional primary + // outcome (success or definitive error) never triggers an early hedge; in that case the primary + // has already won the race below and this gate is cancelled before the hedge subscribes. + Mono hedgeGate = Flux.merge( + Mono.delay(threshold).thenReturn(Boolean.TRUE), + primaryAttempt.flatMap(attempt -> attempt.outcome == BranchOutcome.REGIONAL_FAILURE + ? Mono.just(Boolean.TRUE) + : Mono.never())) + .next(); + + Mono> hedgeAttempt = hedgeGate.then( + hedge + .doOnSubscribe(s -> hedgeStarted.set(true)) + .map(v -> new Attempt<>(Origin.HEDGE, hedgeRegion, v, null, BranchOutcome.SUCCESS)) + .onErrorResume(e -> Mono.just( + new Attempt(Origin.HEDGE, hedgeRegion, null, e, classifier.apply(e))))); + + return Flux.merge(primaryAttempt.flux(), hedgeAttempt.flux()) + .filter(Attempt::isWinner) + .next() + // No winner at all -> return the primary's actual (regional) outcome, verbatim. + .switchIfEmpty(primaryAttempt) + .map(winner -> { + boolean hedgeWon = winner.origin == Origin.HEDGE; + return new HedgeResult<>( + winner.value, + winner.error, + hedgeStarted.get(), + hedgeWon, + winner.region); + }); + } + + /** + * Classify a completed branch throwable. Only a genuine regional failure of the region is + * worth hedging; everything else (including auth failures and definitive errors such as 404 / 409 / + * 412) is authoritative. This single classifier is shared by both the Collection Read and the + * PartitionKeyRange ReadFeed paths, matching the .NET reference (a 404 is always a definitive answer, + * never regional, so a lagging secondary can never override a primary that returned 404). + */ + public static BranchOutcome classifyThrowable(Throwable throwable) { + CosmosException cosmosException = asCosmosException(throwable); + if (cosmosException == null) { + // A non-Cosmos error (e.g. a transport error not yet mapped) is treated as regional so a + // healthy region can still answer. + return BranchOutcome.REGIONAL_FAILURE; + } + return isRegionalFailure(cosmosException.getStatusCode(), cosmosException.getSubStatusCode()) + ? BranchOutcome.REGIONAL_FAILURE + : BranchOutcome.DEFINITIVE; + } + + /** + * The shared regional-failure classifier. A regional failure means "the region is at fault" + * (worth hedging), as opposed to a definitive, request-level answer. + */ + public static boolean isRegionalFailure(int statusCode, int subStatusCode) { + // Transport / gateway-endpoint failures (network failure mapped by RxGatewayStoreModel). + if (subStatusCode == HttpConstants.SubStatusCodes.GATEWAY_ENDPOINT_UNAVAILABLE + || subStatusCode == HttpConstants.SubStatusCodes.GATEWAY_ENDPOINT_READ_TIMEOUT) { + return true; + } + + switch (statusCode) { + case HttpConstants.StatusCodes.SERVICE_UNAVAILABLE: // 503 + case HttpConstants.StatusCodes.INTERNAL_SERVER_ERROR: // 500 + case HttpConstants.StatusCodes.REQUEST_TIMEOUT: // 408 + return true; + case HttpConstants.StatusCodes.GONE: // 410 -- only lease-not-found is regional + return subStatusCode == HttpConstants.SubStatusCodes.LEASE_NOT_FOUND; + case HttpConstants.StatusCodes.FORBIDDEN: // 403 -- only database-account-not-found is regional + return subStatusCode == HttpConstants.SubStatusCodes.DATABASE_ACCOUNT_NOTFOUND; + default: + // 401, plain 403, 404 (any sub-status), 409, 412, 429, etc. are NOT regional -- + // authoritative/definitive, so a definitive primary answer is never overridden by a hedge. + return false; + } + } + + private static CosmosException asCosmosException(Throwable throwable) { + Throwable current = throwable; + while (current != null) { + if (current instanceof CosmosException) { + return (CosmosException) current; + } + current = current.getCause(); + } + return null; + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/caches/RxClientCollectionCache.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/caches/RxClientCollectionCache.java index 3d3ca54b2619..8cc41ef2dd75 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/caches/RxClientCollectionCache.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/caches/RxClientCollectionCache.java @@ -29,6 +29,7 @@ import java.time.Instant; import java.util.HashMap; import java.util.Map; +import java.util.function.Function; /** * Caches collection information. @@ -42,31 +43,56 @@ public class RxClientCollectionCache extends RxCollectionCache { private final IAuthorizationTokenProvider tokenProvider; private final IRetryPolicyFactory retryPolicy; private final ISessionContainer sessionContainer; + private final MetadataHedgingStrategy metadataHedgingStrategy; public RxClientCollectionCache(DiagnosticsClientContext diagnosticsClientContext, ISessionContainer sessionContainer, RxStoreModel storeModel, IAuthorizationTokenProvider tokenProvider, IRetryPolicyFactory retryPolicy, - AsyncCache collectionInfoByNameCache, AsyncCache collectionInfoByIdCache) { + AsyncCache collectionInfoByNameCache, + AsyncCache collectionInfoByIdCache, + MetadataHedgingStrategy metadataHedgingStrategy) { super(collectionInfoByNameCache, collectionInfoByIdCache); this.diagnosticsClientContext = diagnosticsClientContext; this.storeModel = storeModel; this.tokenProvider = tokenProvider; this.retryPolicy = retryPolicy; this.sessionContainer = sessionContainer; + this.metadataHedgingStrategy = metadataHedgingStrategy; } public RxClientCollectionCache(DiagnosticsClientContext diagnosticsClientContext, ISessionContainer sessionContainer, RxStoreModel storeModel, IAuthorizationTokenProvider tokenProvider, - IRetryPolicyFactory retryPolicy) { + IRetryPolicyFactory retryPolicy, + AsyncCache collectionInfoByNameCache, AsyncCache collectionInfoByIdCache) { + this(diagnosticsClientContext, sessionContainer, storeModel, tokenProvider, retryPolicy, + collectionInfoByNameCache, collectionInfoByIdCache, null); + } + + public RxClientCollectionCache(DiagnosticsClientContext diagnosticsClientContext, + ISessionContainer sessionContainer, + RxStoreModel storeModel, + IAuthorizationTokenProvider tokenProvider, + IRetryPolicyFactory retryPolicy, + MetadataHedgingStrategy metadataHedgingStrategy) { this.diagnosticsClientContext = diagnosticsClientContext; this.storeModel = storeModel; this.tokenProvider = tokenProvider; this.retryPolicy = retryPolicy; this.sessionContainer = sessionContainer; + this.metadataHedgingStrategy = metadataHedgingStrategy; + } + + public RxClientCollectionCache(DiagnosticsClientContext diagnosticsClientContext, + ISessionContainer sessionContainer, + RxStoreModel storeModel, + IAuthorizationTokenProvider tokenProvider, + IRetryPolicyFactory retryPolicy) { + this(diagnosticsClientContext, sessionContainer, storeModel, tokenProvider, retryPolicy, + (MetadataHedgingStrategy) null); } protected Mono getByRidAsync(MetadataDiagnosticsContext metaDataDiagnosticsContext, String collectionRid, Map properties) { @@ -120,13 +146,49 @@ private Mono readCollectionAsync(MetadataDiagnosticsContext } Instant addressCallStartTime = Instant.now(); + boolean isAadToken = tokenProvider.getAuthorizationTokenType() == AuthorizationTokenType.AadToken; + Function> sendFunc = req -> { + if (isAadToken) { + return tokenProvider + .populateAuthorizationHeader(req) + .flatMap(serviceRequest -> this.storeModel.processMessage(serviceRequest)); + } + return this.storeModel.processMessage(req); + }; + Mono responseObs; - if (tokenProvider.getAuthorizationTokenType() != AuthorizationTokenType.AadToken) { - responseObs = this.storeModel.processMessage(request); + if (this.metadataHedgingStrategy != null) { + responseObs = this.metadataHedgingStrategy.executeMetadataRead( + request, + sendFunc, + MetadataHedgingStrategy::classifyThrowable, + hedgeResult -> { + if (metaDataDiagnosticsContext != null) { + // Report the activityId that was actually sent on the wire by the winning branch: + // the hedging strategy clones the request per branch and each clone gets its own + // activityId, so the original request's activityId is never sent. The winning + // branch's response echoes its request activityId; fall back to the original + // request id only when the primary lost with no response (definitive error path). + String winningActivityId = request.getActivityId().toString(); + if (hedgeResult.getValue() != null && hedgeResult.getValue().getResponseHeaders() != null) { + String responseActivityId = + hedgeResult.getValue().getResponseHeaders().get(HttpConstants.HttpHeaders.ACTIVITY_ID); + if (responseActivityId != null) { + winningActivityId = responseActivityId; + } + } + metaDataDiagnosticsContext.addMetaDataDiagnostic( + new MetadataDiagnosticsContext.MetadataHedgeDiagnostics( + addressCallStartTime, + Instant.now(), + winningActivityId, + hedgeResult.isHedgeFired(), + hedgeResult.isHedgeWon(), + hedgeResult.getWinningRegion())); + } + }); } else { - responseObs = tokenProvider - .populateAuthorizationHeader(request) - .flatMap(serviceRequest -> this.storeModel.processMessage(serviceRequest)); + responseObs = sendFunc.apply(request); } return responseObs.map(response -> { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/caches/RxPartitionKeyRangeCache.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/caches/RxPartitionKeyRangeCache.java index b869489c7d75..926094568ad6 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/caches/RxPartitionKeyRangeCache.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/caches/RxPartitionKeyRangeCache.java @@ -23,6 +23,7 @@ import com.azure.cosmos.implementation.routing.IServerIdentity; import com.azure.cosmos.implementation.routing.InMemoryCollectionRoutingMap; import com.azure.cosmos.implementation.routing.Range; +import com.azure.cosmos.implementation.routing.RegionalRoutingContext; import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.FeedResponse; import com.azure.cosmos.models.ModelBridgeInternal; @@ -51,12 +52,20 @@ public class RxPartitionKeyRangeCache implements IPartitionKeyRangeCache { private final RxDocumentClientImpl client; private final RxCollectionCache collectionCache; private final DiagnosticsClientContext clientContext; + private final MetadataHedgingStrategy metadataHedgingStrategy; public RxPartitionKeyRangeCache(RxDocumentClientImpl client, RxCollectionCache collectionCache) { + this(client, collectionCache, null); + } + + public RxPartitionKeyRangeCache(RxDocumentClientImpl client, + RxCollectionCache collectionCache, + MetadataHedgingStrategy metadataHedgingStrategy) { this.routingMapCache = new AsyncCacheNonBlocking<>(); this.client = client; this.collectionCache = collectionCache; this.clientContext = client; + this.metadataHedgingStrategy = metadataHedgingStrategy; } /* (non-Javadoc) @@ -254,39 +263,126 @@ private Mono getRoutingMapForCollectionAsync( collectionRid, previousChangeFeedIfNoneMatch); - Flux> rangesObs = - getPartitionKeyRange( - metaDataDiagnosticsContext, - collectionRid, - previousChangeFeedIfNoneMatch, - properties); - - AtomicReference continuationToken = new AtomicReference<>(previousChangeFeedIfNoneMatch); - List ranges = new ArrayList<>(); - - return rangesObs - .doOnNext(response -> { - ranges.addAll(response.getResults()); - continuationToken.set(response.getContinuationToken()); - }) - .collectList() - .flatMap(allResults -> { - CollectionRoutingMap updatedMap = - updateRoutingMap(collectionRid, previousRoutingMap, ranges, continuationToken.get()); - - return Mono.just(updatedMap); - }) - .doFinally(signal -> { - if (metaDataDiagnosticsContext != null) { - Instant pkRangesCallEndTime = Instant.now(); - MetadataDiagnosticsContext.MetadataDiagnostics metaDataDiagnostic = - new MetadataDiagnosticsContext.MetadataDiagnostics( - pkRangesCallStartTime, - pkRangesCallEndTime, - MetadataDiagnosticsContext.MetadataType.PARTITION_KEY_RANGE_LOOK_UP); - metaDataDiagnosticsContext.addMetaDataDiagnostic(metaDataDiagnostic); + Mono result = + readRoutingMapWithOptionalHedging( + metaDataDiagnosticsContext, collectionRid, previousRoutingMap, properties, previousChangeFeedIfNoneMatch); + + return result.doFinally(signal -> { + if (metaDataDiagnosticsContext != null) { + Instant pkRangesCallEndTime = Instant.now(); + MetadataDiagnosticsContext.MetadataDiagnostics metaDataDiagnostic = + new MetadataDiagnosticsContext.MetadataDiagnostics( + pkRangesCallStartTime, + pkRangesCallEndTime, + MetadataDiagnosticsContext.MetadataType.PARTITION_KEY_RANGE_LOOK_UP); + metaDataDiagnosticsContext.addMetaDataDiagnostic(metaDataDiagnostic); + } + }); + } + + // When metadata hedging is eligible (opt-in/PPAF + >= 2 regions), race two region-pinned full routing-map reads. + // Each branch reads all its pages from a single region (via excluded-regions pinning), so continuation integrity + // is preserved and the winning region serves a consistent routing map. Otherwise the original single read is used. + private Mono readRoutingMapWithOptionalHedging( + MetadataDiagnosticsContext metaDataDiagnosticsContext, + String collectionRid, + CollectionRoutingMap previousRoutingMap, + Map properties, + String previousChangeFeedIfNoneMatch) { + + if (this.metadataHedgingStrategy != null && this.metadataHedgingStrategy.isEffectivelyEnabled()) { + RxDocumentServiceRequest probe = RxDocumentServiceRequest.create(this.clientContext, + OperationType.ReadFeed, collectionRid, ResourceType.PartitionKeyRange, null); + probe.requestContext.resolvedCollectionRid = collectionRid; + probe.setResourceId(collectionRid); + + List ctxs = this.metadataHedgingStrategy.getApplicableReadRegions(probe); + if (ctxs.size() >= 2) { + String primaryRegion = this.metadataHedgingStrategy.getRegionName(ctxs.get(0), probe); + String hedgeRegion = this.metadataHedgingStrategy.getRegionName(ctxs.get(1), probe); + List allRegions = new ArrayList<>(); + for (RegionalRoutingContext ctx : ctxs) { + String name = this.metadataHedgingStrategy.getRegionName(ctx, probe); + if (name != null && !allRegions.contains(name)) { + allRegions.add(name); + } } - }); + List primaryExcluded = + MetadataHedgingStrategy.excludedRegionsForTarget(null, allRegions, primaryRegion); + List hedgeExcluded = + MetadataHedgingStrategy.excludedRegionsForTarget(null, allRegions, hedgeRegion); + + // Capture the real wire activityId of each branch's routing-map read (from the winning + // page's FeedResponse). The per-branch requests are cloned deep inside the paginator and + // get their own activityId, so we report the winning branch's actual sent activityId + // rather than a throwaway probe id that never reaches the wire. + AtomicReference primaryActivityId = new AtomicReference<>(); + AtomicReference hedgeActivityId = new AtomicReference<>(); + + Mono primaryMono = readRoutingMapOnce( + metaDataDiagnosticsContext, collectionRid, previousRoutingMap, properties, + previousChangeFeedIfNoneMatch, primaryExcluded, primaryActivityId); + Mono hedgeMono = readRoutingMapOnce( + metaDataDiagnosticsContext, collectionRid, previousRoutingMap, properties, + previousChangeFeedIfNoneMatch, hedgeExcluded, hedgeActivityId); + + Instant hedgeStart = Instant.now(); + return this.metadataHedgingStrategy.executeHedged( + primaryMono, primaryRegion, hedgeMono, hedgeRegion, + this.metadataHedgingStrategy.getThreshold(), + MetadataHedgingStrategy::classifyThrowable) + .flatMap(hedgeResult -> { + if (metaDataDiagnosticsContext != null && hedgeResult.isHedgeFired()) { + String winningActivityId = + hedgeResult.isHedgeWon() ? hedgeActivityId.get() : primaryActivityId.get(); + metaDataDiagnosticsContext.addMetaDataDiagnostic( + new MetadataDiagnosticsContext.MetadataHedgeDiagnostics( + hedgeStart, Instant.now(), winningActivityId, + hedgeResult.isHedgeFired(), hedgeResult.isHedgeWon(), + hedgeResult.getWinningRegion())); + } + return hedgeResult.unwrap(); + }); + } + } + + return readRoutingMapOnce( + metaDataDiagnosticsContext, collectionRid, previousRoutingMap, properties, + previousChangeFeedIfNoneMatch, null, null); + } + + private Mono readRoutingMapOnce( + MetadataDiagnosticsContext metaDataDiagnosticsContext, + String collectionRid, + CollectionRoutingMap previousRoutingMap, + Map properties, + String previousChangeFeedIfNoneMatch, + List excludedRegions, + AtomicReference activityIdSink) { + + return Mono.defer(() -> { + AtomicReference continuationToken = new AtomicReference<>(previousChangeFeedIfNoneMatch); + List ranges = new ArrayList<>(); + + return getPartitionKeyRange( + metaDataDiagnosticsContext, collectionRid, previousChangeFeedIfNoneMatch, properties, excludedRegions) + .doOnNext(response -> { + ranges.addAll(response.getResults()); + continuationToken.set(response.getContinuationToken()); + if (activityIdSink != null && response.getActivityId() != null) { + // Record this branch's wire activityId (first page that reports one is sufficient + // to correlate the branch with server-side traces). + activityIdSink.compareAndSet(null, response.getActivityId()); + } + }) + .collectList() + .flatMap(allResults -> { + CollectionRoutingMap updatedMap = + updateRoutingMap(collectionRid, previousRoutingMap, ranges, continuationToken.get()); + + return Mono.just(updatedMap); + }); + }); } private CollectionRoutingMap updateRoutingMap( @@ -327,7 +423,8 @@ private Flux> getPartitionKeyRange( MetadataDiagnosticsContext metaDataDiagnosticsContext, String collectionRid, String previousRoutingMapContinuationToken, - Map properties) { + Map properties, + List excludedRegions) { RxDocumentServiceRequest request = RxDocumentServiceRequest.create(this.clientContext, OperationType.ReadFeed, collectionRid, @@ -351,6 +448,12 @@ private Flux> getPartitionKeyRange( ModelBridgeInternal.setQueryRequestOptionsProperties(cosmosQueryRequestOptions, properties); } + // Pin this branch to a single region (exclude the others) so the whole routing-map read -- all pages -- + // is served from one region, keeping the continuation chain consistent. + if (excludedRegions != null && !excludedRegions.isEmpty()) { + cosmosQueryRequestOptions.setExcludedRegions(excludedRegions); + } + return client.readPartitionKeyRanges(coll.getSelfLink(), cosmosQueryRequestOptions); }); }