diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxDocumentServiceRequestTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxDocumentServiceRequestTest.java index dc269aff7488..fd97351cb87a 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxDocumentServiceRequestTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxDocumentServiceRequestTest.java @@ -3,15 +3,24 @@ package com.azure.cosmos.implementation; +import com.azure.cosmos.implementation.http.HttpTransportSerializer; import org.apache.commons.collections4.map.HashedMap; import org.apache.commons.lang3.StringUtils; +import org.mockito.Mockito; import org.testng.annotations.DataProvider; import org.testng.annotations.Test; +import java.lang.reflect.Field; +import java.lang.reflect.Modifier; import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.Arrays; +import java.util.HashSet; import java.util.Map; +import java.util.Set; import java.util.UUID; +import java.util.stream.Collectors; import static com.azure.cosmos.implementation.TestUtils.mockDiagnosticsClientContext; import static org.assertj.core.api.Assertions.assertThat; @@ -19,6 +28,20 @@ public class RxDocumentServiceRequestTest { + // activityId is regenerated for the clone and isDisposed is reset by its constructor. All other fields are copied. + // Keeping the intentional exceptions in this inventory makes adding a field fail this test until clone() is audited. + private static final Set CLONE_ACCOUNTED_FIELDS = new HashSet<>(Arrays.asList( + "clientContext", "forcePartitionKeyRangeRefresh", "forceCollectionRoutingMapRefresh", "resourceId", + "resourceType", "headers", "continuation", "isMedia", "isNameBased", "operationType", "resourceAddress", + "forceNameCacheRefresh", "endpointOverride", "activityId", "originalSessionToken", "partitionKeyRangeIdentity", + "defaultReplicaIndex", "isAddressRefresh", "isForcedAddressRefresh", "requestContext", + "faultInjectionRequestContext", "partitionKeyInternal", "partitionKeyDefinition", "effectivePartitionKey", + "feedRange", "effectiveRange", "numberOfItemsInBatchRequest", "contentAsByteArray", "useGatewayMode", + "useThinClientMode", "isDisposed", "entityId", "isFeed", "authorizationTokenType", "properties", + "throughputControlGroupName", "intendedCollectionRidPassedIntoSDK", "isBarrierRequest", "responseTimeout", + "nonIdempotentWriteRetriesEnabled", "hasFeedRangeFilteringBeenApplied", + "isPerPartitionAutomaticFailoverEnabledAndWriteRequest", "httpTransportSerializer")); + private final static String DOCUMENT_DEFINITION = "{ " + "\"id\": \"%s\", " + "\"mypk\": \"%s\", " + "\"sgmts\": [[6519456, 1471916863], [2498434, 1455671440]]" + "}"; private static final String PARTITION_KEY_VALUE = "1"; @@ -441,6 +464,54 @@ public void cloneProducesIndependentHeaderMap() { assertThat(cloned.getHeaders().get("x-new-header")).isEqualTo("new-value"); } + @Test(groups = { "unit" }) + public void cloneCopiesRequestScopedFields() { + RxDocumentServiceRequest original = RxDocumentServiceRequest.createFromName( + mockDiagnosticsClientContext(), + OperationType.Query, + "/dbs/db/colls/coll/docs", + ResourceType.Document); + original.setAddressRefresh(true, true); + original.setEffectivePartitionKey("effective-partition-key"); + original.setNumberOfItemsInBatchRequest(42); + original.entityId = "entity-id"; + original.authorizationTokenType = AuthorizationTokenType.ResourceToken; + original.properties = new HashedMap<>(); + original.properties.put("excludedRegions", Arrays.asList("region1")); + original.throughputControlGroupName = "throughput-control-group"; + original.intendedCollectionRidPassedIntoSDK = true; + original.isBarrierRequest = true; + original.setResponseTimeout(Duration.ofSeconds(5)); + HttpTransportSerializer transportSerializer = Mockito.mock(HttpTransportSerializer.class); + HttpTransportSerializer defaultTransportSerializer = Mockito.mock(HttpTransportSerializer.class); + original.setHttpTransportSerializer(transportSerializer); + + RxDocumentServiceRequest cloned = original.clone(); + + assertThat(cloned.isAddressRefresh()).isTrue(); + assertThat(cloned.shouldForceAddressRefresh()).isTrue(); + assertThat(cloned.getEffectivePartitionKey()).isEqualTo("effective-partition-key"); + assertThat(cloned.getNumberOfItemsInBatchRequest()).isEqualTo(42); + assertThat(cloned.entityId).isEqualTo("entity-id"); + assertThat(cloned.authorizationTokenType).isEqualTo(AuthorizationTokenType.ResourceToken); + assertThat(cloned.properties).isEqualTo(original.properties).isNotSameAs(original.properties); + assertThat(cloned.throughputControlGroupName).isEqualTo("throughput-control-group"); + assertThat(cloned.intendedCollectionRidPassedIntoSDK).isTrue(); + assertThat(cloned.isBarrierRequest).isTrue(); + assertThat(cloned.getResponseTimeout()).isEqualTo(Duration.ofSeconds(5)); + assertThat(cloned.getEffectiveHttpTransportSerializer(defaultTransportSerializer)).isSameAs(transportSerializer); + } + + @Test(groups = { "unit" }) + public void cloneAccountsForEveryInstanceField() { + Set instanceFields = Arrays.stream(RxDocumentServiceRequest.class.getDeclaredFields()) + .filter(field -> !Modifier.isStatic(field.getModifiers())) + .map(Field::getName) + .collect(Collectors.toSet()); + + assertThat(CLONE_ACCOUNTED_FIELDS).isEqualTo(instanceFields); + } + @Test(groups = { "unit" }) public void createWithInvalidPath() { try { diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 9cbcee519b89..d2d4750c7b58 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -9,6 +9,7 @@ #### Breaking Changes #### Bugs Fixed +* 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). * 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). diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java index 1e2958040978..f3b5d6323165 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java @@ -1075,11 +1075,11 @@ public RxDocumentServiceRequest clone() { rxDocumentServiceRequest.setIsMedia(this.getIsMedia()); rxDocumentServiceRequest.setOriginalSessionToken(this.getOriginalSessionToken()); rxDocumentServiceRequest.setPartitionKeyRangeIdentity(this.getPartitionKeyRangeIdentity()); + rxDocumentServiceRequest.setAddressRefresh(this.isAddressRefresh(), this.shouldForceAddressRefresh()); rxDocumentServiceRequest.forceCollectionRoutingMapRefresh = this.forceCollectionRoutingMapRefresh; rxDocumentServiceRequest.forcePartitionKeyRangeRefresh = this.forcePartitionKeyRangeRefresh; rxDocumentServiceRequest.useGatewayMode = this.useGatewayMode; rxDocumentServiceRequest.useThinClientMode = this.useThinClientMode; - rxDocumentServiceRequest.requestContext = this.requestContext; rxDocumentServiceRequest.faultInjectionRequestContext = new FaultInjectionRequestContext(this.faultInjectionRequestContext); rxDocumentServiceRequest.nonIdempotentWriteRetriesEnabled = this.nonIdempotentWriteRetriesEnabled; rxDocumentServiceRequest.setResourceAddress(this.resourceAddress); @@ -1091,6 +1091,16 @@ public RxDocumentServiceRequest clone() { rxDocumentServiceRequest.hasFeedRangeFilteringBeenApplied = this.hasFeedRangeFilteringBeenApplied; rxDocumentServiceRequest.isPerPartitionAutomaticFailoverEnabledAndWriteRequest = this.isPerPartitionAutomaticFailoverEnabledAndWriteRequest; rxDocumentServiceRequest.partitionKeyDefinition = this.partitionKeyDefinition; + rxDocumentServiceRequest.effectivePartitionKey = this.effectivePartitionKey; + rxDocumentServiceRequest.numberOfItemsInBatchRequest = this.numberOfItemsInBatchRequest; + rxDocumentServiceRequest.entityId = this.entityId; + rxDocumentServiceRequest.authorizationTokenType = this.authorizationTokenType; + rxDocumentServiceRequest.properties = this.properties != null ? new HashMap<>(this.properties) : null; + rxDocumentServiceRequest.throughputControlGroupName = this.throughputControlGroupName; + rxDocumentServiceRequest.intendedCollectionRidPassedIntoSDK = this.intendedCollectionRidPassedIntoSDK; + rxDocumentServiceRequest.isBarrierRequest = this.isBarrierRequest; + rxDocumentServiceRequest.responseTimeout = this.responseTimeout; + rxDocumentServiceRequest.httpTransportSerializer.set(this.httpTransportSerializer.get()); return rxDocumentServiceRequest; }