diff --git a/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md index 5045991f3542..acaee40c6cd9 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md @@ -8,6 +8,7 @@ #### Bugs Fixed * Improved partition planning performance for change feed with large number of feed ranges. - See [PR 49086](https://github.com/Azure/azure-sdk-for-java/pull/49086) +* Fixed `UnsupportedOperationException` when using `readManyByPartitionKeys` for empty pages. - See [PR 49311](https://github.com/Azure/azure-sdk-for-java/pull/49311) #### Other Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md index 777177708547..09b19a4119c5 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md @@ -8,6 +8,7 @@ #### Bugs Fixed * Improved partition planning performance for change feed with large number of feed ranges. - See [PR 49086](https://github.com/Azure/azure-sdk-for-java/pull/49086) +* Fixed `UnsupportedOperationException` when using `readManyByPartitionKeys` for empty pages. - See [PR 49311](https://github.com/Azure/azure-sdk-for-java/pull/49311) #### Other Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md index 1360614308a3..a36a8521c71b 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md @@ -8,6 +8,7 @@ #### Bugs Fixed * Improved partition planning performance for change feed with large number of feed ranges. - See [PR 49086](https://github.com/Azure/azure-sdk-for-java/pull/49086) +* Fixed `UnsupportedOperationException` when using `readManyByPartitionKeys` for empty pages. - See [PR 49311](https://github.com/Azure/azure-sdk-for-java/pull/49311) #### Other Changes diff --git a/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md index 6674b6f8bb74..5d46eb30dde5 100644 --- a/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md @@ -8,6 +8,7 @@ #### Bugs Fixed * Improved partition planning performance for change feed with large number of feed ranges. - See [PR 49086](https://github.com/Azure/azure-sdk-for-java/pull/49086) +* Fixed `UnsupportedOperationException` when using `readManyByPartitionKeys` for empty pages. - See [PR 49311](https://github.com/Azure/azure-sdk-for-java/pull/49311) #### Other Changes diff --git a/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md index 9b87a3fcf675..9b971d9786c0 100644 --- a/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md @@ -8,6 +8,7 @@ #### Bugs Fixed * Improved partition planning performance for change feed with large number of feed ranges. - See [PR 49086](https://github.com/Azure/azure-sdk-for-java/pull/49086) +* Fixed `UnsupportedOperationException` when using `readManyByPartitionKeys` for empty pages. - See [PR 49311](https://github.com/Azure/azure-sdk-for-java/pull/49311) #### Other Changes diff --git a/sdk/cosmos/azure-cosmos-spark_4-1_2-13/CHANGELOG.md b/sdk/cosmos/azure-cosmos-spark_4-1_2-13/CHANGELOG.md index 570aec149b2c..b3c35541850b 100644 --- a/sdk/cosmos/azure-cosmos-spark_4-1_2-13/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos-spark_4-1_2-13/CHANGELOG.md @@ -8,6 +8,7 @@ #### Bugs Fixed * Improved partition planning performance for change feed with large number of feed ranges. - See [PR 49086](https://github.com/Azure/azure-sdk-for-java/pull/49086) +* Fixed `UnsupportedOperationException` when using `readManyByPartitionKeys` for empty pages. - See [PR 49311](https://github.com/Azure/azure-sdk-for-java/pull/49311) #### Other Changes diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java index 86057a353021..8afec5905d11 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyContinuationTokenTest.java @@ -127,17 +127,20 @@ public void roundtrip_lastBatchNoContinuation() { @Test(groups = { "unit" }) public void setFeedResponseContinuationToken_handlesEmptyHeadersWithoutCopyingNormalCase() { - Map immutableEmptyHeaders = Collections.emptyMap(); + // Immutable inputs are normalized to a mutable map at FeedResponse construction + // time (so the field stays final). Clearing a continuation token on an empty + // header map is a no-op and must not throw. FeedResponse emptyResponse = ModelBridgeInternal.createFeedResponse( Collections.emptyList(), - immutableEmptyHeaders); + Collections.emptyMap()); ModelBridgeInternal.setFeedResponseContinuationToken(null, emptyResponse); assertThat(emptyResponse.getContinuationToken()).isNull(); - assertThat(emptyResponse.getResponseHeaders()).isSameAs(immutableEmptyHeaders); assertThat(emptyResponse.getResponseHeaders()).isEmpty(); + // Mutable header maps are passed through without copying, preserving the + // reference returned by getResponseHeaders(). Map normalHeaders = new HashMap<>(); normalHeaders.put(HttpConstants.HttpHeaders.ACTIVITY_ID, "test-activity-id"); FeedResponse normalResponse = ModelBridgeInternal.createFeedResponse( @@ -150,6 +153,30 @@ public void setFeedResponseContinuationToken_handlesEmptyHeadersWithoutCopyingNo assertThat(normalResponse.getResponseHeaders()).isSameAs(normalHeaders); } + /** + * Reproduces the customer-reported failure path: the parallel query pipeline's + * "artificial empty page" branch (ParallelDocumentQueryExecutionContext.headerResponse) + * emits a FeedResponse whose header map is {@code Utils.immutableMapOf(...)} - i.e. a + * non-empty {@code Collections.unmodifiableMap} wrapper. When such a page reaches the + * readManyByPartitionKeys stamping lambda, setFeedResponseContinuationToken attempts + * to put the composite token into the immutable map, throwing UnsupportedOperationException. + */ + @Test(groups = { "unit" }) + public void setFeedResponseContinuationToken_immutableNonEmptyHeaders_doesNotThrow() { + Map immutableSingleEntryHeaders = + Utils.immutableMapOf(HttpConstants.HttpHeaders.REQUEST_CHARGE, "1.23"); + FeedResponse response = ModelBridgeInternal.createFeedResponse( + Collections.emptyList(), + immutableSingleEntryHeaders); + + ModelBridgeInternal.setFeedResponseContinuationToken("composite-token", response); + + assertThat(response.getContinuationToken()).isEqualTo("composite-token"); + assertThat(response.getResponseHeaders()) + .containsEntry(HttpConstants.HttpHeaders.CONTINUATION, "composite-token") + .containsEntry(HttpConstants.HttpHeaders.REQUEST_CHARGE, "1.23"); + } + @Test(groups = { "unit" }) public void deserialize_malformedInput_throws() { // Either the base64 decoder or the JSON parsing layer rejects garbage; both raise diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/UtilsImmutableMapTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/UtilsImmutableMapTests.java new file mode 100644 index 000000000000..f0524146dbb9 --- /dev/null +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/UtilsImmutableMapTests.java @@ -0,0 +1,52 @@ +/* + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. + */ + +package com.azure.cosmos.implementation; + +import org.testng.annotations.Test; + +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Unit tests for {@link Utils#immutableMapOf(Object, Object)} and the matching detector + * {@link Utils#isImmutableMap(Map)}. These live next to the factory so that any future + * change to the factory's runtime class shape is caught by the same regression suite. + */ +public class UtilsImmutableMapTests { + + @Test(groups = { "unit" }) + public void immutableMapOf_isDetectedAsImmutable() { + Map m = Utils.immutableMapOf("k", "v"); + assertThat(Utils.isImmutableMap(m)).isTrue(); + } + + @Test(groups = { "unit" }) + public void emptyMap_isDetectedAsImmutable() { + assertThat(Utils.isImmutableMap(Collections.emptyMap())).isTrue(); + } + + @Test(groups = { "unit" }) + public void hashMap_isNotDetectedAsImmutable() { + assertThat(Utils.isImmutableMap(new HashMap<>())).isFalse(); + Map populated = new HashMap<>(); + populated.put("k", "v"); + assertThat(Utils.isImmutableMap(populated)).isFalse(); + } + + @Test(groups = { "unit" }) + public void linkedHashMap_isNotDetectedAsImmutable() { + assertThat(Utils.isImmutableMap(new LinkedHashMap<>())).isFalse(); + } + + @Test(groups = { "unit" }) + public void nullMap_isNotDetectedAsImmutable() { + assertThat(Utils.isImmutableMap(null)).isFalse(); + } +} diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 4d6c3781b00b..5180cd9a9b71 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -8,6 +8,7 @@ #### Breaking Changes #### Bugs Fixed +* Fixed `UnsupportedOperationException` when using `readManyByPartitionKeys` for empty pages. - See [PR 49311](https://github.com/Azure/azure-sdk-for-java/pull/49311) #### Other Changes * 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) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Utils.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Utils.java index 0b01e87d985b..1e57ff4bedca 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Utils.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Utils.java @@ -574,6 +574,33 @@ public static List immutableListOf() { return map; } + // Runtime classes produced by the immutable factory methods above. Captured once at + // class-init time so that callers can perform an O(1) reference-equality check to + // decide whether they need a defensive mutable copy, without resorting to + // exception-driven probing on the hot path. + private static final Class UNMODIFIABLE_MAP_CLASS = + Collections.unmodifiableMap(new HashMap<>()).getClass(); + private static final Class EMPTY_MAP_CLASS = Collections.emptyMap().getClass(); + + /** + * Returns {@code true} if {@code map} is one of the immutable map shapes produced by + * the factory methods in this class ({@link #immutableMapOf(Object, Object)}) or by + * {@link Collections#emptyMap()}. The check is a single reference comparison and is + * safe to call from hot paths. + *

+ * Note: this is intentionally narrow - it only recognizes the wrappers the Cosmos + * pipeline actually emits. It does not attempt to recognize every possible JDK or + * third-party immutable map (e.g. {@code Map.of(...)}, Guava {@code ImmutableMap}). + * Add new sentinels here if a new immutable producer is introduced. + */ + public static boolean isImmutableMap(Map map) { + if (map == null) { + return false; + } + Class clazz = map.getClass(); + return clazz == UNMODIFIABLE_MAP_CLASS || clazz == EMPTY_MAP_CLASS; + } + public static V firstOrDefault(List list) { return list.size() > 0? list.get(0) : null ; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java index da8764c95f85..52792049d325 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java @@ -16,6 +16,7 @@ import com.azure.cosmos.implementation.QueryMetricsConstants; import com.azure.cosmos.implementation.RxDocumentServiceResponse; import com.azure.cosmos.implementation.Strings; +import com.azure.cosmos.implementation.Utils; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.implementation.query.queryadvisor.QueryAdvice; import com.azure.cosmos.implementation.query.QueryInfo; @@ -55,19 +56,21 @@ private static ImplementationBridgeHelpers.CosmosDiagnosticsHelper.CosmosDiagnos private QueryInfo queryInfo; private QueryInfo.QueryPlanDiagnosticsContext queryPlanDiagnosticsContext; - // All header maps are produced by the SDK's own query pipeline. Non-null maps - // are always mutable (HashMap or ConcurrentHashMap) - the SDK intentionally - // allows callers to add/modify headers on FeedResponse. The only known - // exception is empty-page responses where the query pipeline may pass null. - // We do NOT clone non-null maps here to avoid unnecessary allocations on every - // FeedResponse construction - the wider blast radius of cloning (every query, - // change feed, readMany response) is not justified by the narrow null case. - // If a future code path introduces an immutable non-null header map, the - // setContinuationTokenInternal method will fail fast with - // UnsupportedOperationException, and the fix should be to make the upstream - // pipeline emit a mutable map rather than adding defensive cloning here. + // The header map stored on FeedResponse must be mutable: downstream stages + // (e.g. the readManyByPartitionKeys stamping lambda) may add or replace + // headers in place. Normalize at construction time so the field can stay + // final and getResponseHeaders() consistently returns the same instance. + // Mutable inputs are passed through without copying; null and the known + // immutable shapes produced by Utils.immutableMapOf / Collections.emptyMap + // are replaced with a fresh HashMap (preserving entries). private static Map ensureMutableHeadersMap(Map headers) { - return headers == null ? new HashMap<>() : headers; + if (headers == null) { + return new HashMap<>(); + } + if (Utils.isImmutableMap(headers)) { + return new HashMap<>(headers); + } + return headers; } FeedResponse(List results, Map headers) { @@ -452,16 +455,6 @@ private void setContinuationTokenInternal(String headerName, String continuation if (!Strings.isNullOrWhiteSpace(continuationToken)) { this.header.put(headerName, continuationToken); } else if (!this.header.isEmpty() && this.header.containsKey(headerName)) { - // The query API returns unmodifiable header collections for empty - // responses (no documents returned - when only header set is request charge) - // the protection here to check for existence of the header before attempting - // to remove it would not be robust enough against unknown headers - // but since we only ever call our own query pipeline - // avoiding cloning in all cases and gating on continuation header - // existence is a reasonable trade-off - test coverage exists that uncovered - // the problem - so, this acts as regression test as well - // --> the test coverage is in ItemsPartitionReaderWithReadManyByPartitionKeyITest - // it should "return empty results for non-existent partition keys" this.header.remove(headerName); } }