From 7173494eafdc0ddaaf5da4f29beb10adb888d8c3 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Fri, 29 May 2026 14:39:43 +0000 Subject: [PATCH 1/4] Fixes UnsupportedOperationException when using readManyByPartitionKey for empty pages --- ...nyByPartitionKeyContinuationTokenTest.java | 24 +++++++++ .../UtilsImmutableMapTests.java | 52 +++++++++++++++++++ .../azure/cosmos/implementation/Utils.java | 27 ++++++++++ .../com/azure/cosmos/models/FeedResponse.java | 27 +++++----- 4 files changed, 117 insertions(+), 13 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/UtilsImmutableMapTests.java 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..7c31a8ea2565 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 @@ -150,6 +150,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/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..a62720d7970e 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; @@ -44,7 +45,7 @@ private static ImplementationBridgeHelpers.CosmosDiagnosticsHelper.CosmosDiagnos private static final Pattern DELIMITER_CHARS_PATTERN = Pattern.compile(Constants.Quota.DELIMITER_CHARS); private final List results; - private final Map header; + private Map header; private final HashMap usageHeaders; private final HashMap quotaHeaders; private final boolean useEtagAsContinuation; @@ -449,19 +450,19 @@ void setContinuationToken(String continuationToken) { } private void setContinuationTokenInternal(String headerName, String continuationToken) { - if (!Strings.isNullOrWhiteSpace(continuationToken)) { + boolean setting = !Strings.isNullOrWhiteSpace(continuationToken); + boolean clearing = !setting && !this.header.isEmpty() && this.header.containsKey(headerName); + if (!setting && !clearing) { + return; + } + + if (Utils.isImmutableMap(this.header)) { + this.header = new HashMap<>(this.header); + } + + if (setting) { 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" + } else { this.header.remove(headerName); } } From a442db365fc0fd4df7d5050b6998d56ea2e83f1f Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Fri, 29 May 2026 14:49:43 +0000 Subject: [PATCH 2/4] Updated changelogs --- sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md | 1 + sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md | 1 + sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md | 1 + sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md | 1 + sdk/cosmos/azure-cosmos-spark_4-0_2-13/CHANGELOG.md | 1 + sdk/cosmos/azure-cosmos-spark_4-1_2-13/CHANGELOG.md | 1 + sdk/cosmos/azure-cosmos/CHANGELOG.md | 1 + 7 files changed, 7 insertions(+) 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/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 6a5ba9e04600..e1efa993fa73 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) From 364d93156e846e440ff56b28968224f17a183770 Mon Sep 17 00:00:00 2001 From: Annie Liang Date: Fri, 29 May 2026 10:22:19 -0700 Subject: [PATCH 3/4] Update spark.yml JarStorageAccountName for ephemeral tenant rotation 202605 Update JarStorageAccountName from oltpsparkcijarstore0326 to oltpsparkcijarstore0529 to point to the new storage account provisioned in the current ephemeral tenant rotation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- sdk/cosmos/spark.yml | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/sdk/cosmos/spark.yml b/sdk/cosmos/spark.yml index 91e8bfea69c6..05e5bd58232b 100644 --- a/sdk/cosmos/spark.yml +++ b/sdk/cosmos/spark.yml @@ -28,7 +28,7 @@ stages: JarStorageAccountKey: $(spark-databricks-cosmos-spn-clientIdCert) JarReadOnlySasUri: $(spark-databricks-cosmos-spn-clientCertBase64) JarReadOnlySasUriIndex: 0 - JarStorageAccountName: 'oltpsparkcijarstore0326' + JarStorageAccountName: 'oltpsparkcijarstore0529' JarName: 'azure-cosmos-spark_3-5_2-12-latest-ci-candidate.jar' - template: /sdk/cosmos/spark.databricks.yml parameters: @@ -50,7 +50,7 @@ stages: JarStorageAccountKey: $(spark-databricks-cosmos-spn-clientIdCert) JarReadOnlySasUri: $(spark-databricks-cosmos-spn-clientCertBase64) JarReadOnlySasUriIndex: 0 - JarStorageAccountName: 'oltpsparkcijarstore0326' + JarStorageAccountName: 'oltpsparkcijarstore0529' JarName: 'azure-cosmos-spark_3-5_2-12-latest-ci-candidate.jar' - template: /sdk/cosmos/spark.databricks.yml parameters: @@ -73,7 +73,7 @@ stages: JarStorageAccountKey: $(spark-databricks-cosmos-spn-clientIdCert) JarReadOnlySasUri: $(spark-databricks-cosmos-spn-clientCertBase64) JarReadOnlySasUriIndex: 0 - JarStorageAccountName: 'oltpsparkcijarstore0326' + JarStorageAccountName: 'oltpsparkcijarstore0529' JarName: 'azure-cosmos-spark_3-5_2-12-latest-ci-candidate.jar' - template: /sdk/cosmos/spark.databricks.yml parameters: @@ -96,7 +96,7 @@ stages: JarStorageAccountKey: $(spark-databricks-cosmos-spn-clientIdCert) JarReadOnlySasUri: $(spark-databricks-cosmos-spn-clientCertBase64) JarReadOnlySasUriIndex: 0 - JarStorageAccountName: 'oltpsparkcijarstore0326' + JarStorageAccountName: 'oltpsparkcijarstore0529' JarName: 'azure-cosmos-spark_3-5_2-12-latest-ci-candidate.jar' - template: /sdk/cosmos/spark.databricks.yml parameters: @@ -119,7 +119,7 @@ stages: JarStorageAccountKey: $(spark-databricks-cosmos-spn-clientIdCert) JarReadOnlySasUri: $(spark-databricks-cosmos-spn-clientCertBase64) JarReadOnlySasUriIndex: 1 - JarStorageAccountName: 'oltpsparkcijarstore0326' + JarStorageAccountName: 'oltpsparkcijarstore0529' JarName: 'azure-cosmos-spark_3-5_2-13-latest-ci-candidate.jar' - template: /sdk/cosmos/spark.databricks.yml parameters: @@ -142,7 +142,7 @@ stages: JarStorageAccountKey: $(spark-databricks-cosmos-spn-clientIdCert) JarReadOnlySasUri: $(spark-databricks-cosmos-spn-clientCertBase64) JarReadOnlySasUriIndex: 2 - JarStorageAccountName: 'oltpsparkcijarstore0326' + JarStorageAccountName: 'oltpsparkcijarstore0529' JarName: 'azure-cosmos-spark_4-0_2-13-latest-ci-candidate.jar' - template: /sdk/cosmos/spark.databricks.yml parameters: @@ -165,5 +165,5 @@ stages: JarStorageAccountKey: $(spark-databricks-cosmos-spn-clientIdCert) JarReadOnlySasUri: $(spark-databricks-cosmos-spn-clientCertBase64) JarReadOnlySasUriIndex: 3 - JarStorageAccountName: 'oltpsparkcijarstore0326' + JarStorageAccountName: 'oltpsparkcijarstore0529' JarName: 'azure-cosmos-spark_4-1_2-13-latest-ci-candidate.jar' From 671d37d7ec0bee8ffefb839a90ddd8b9d1d0d1d2 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Sun, 31 May 2026 23:38:12 +0000 Subject: [PATCH 4/4] Keeping FeedResponse.header final --- ...nyByPartitionKeyContinuationTokenTest.java | 9 ++-- .../com/azure/cosmos/models/FeedResponse.java | 42 ++++++++----------- 2 files changed, 23 insertions(+), 28 deletions(-) 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 7c31a8ea2565..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( 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 a62720d7970e..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 @@ -45,7 +45,7 @@ private static ImplementationBridgeHelpers.CosmosDiagnosticsHelper.CosmosDiagnos private static final Pattern DELIMITER_CHARS_PATTERN = Pattern.compile(Constants.Quota.DELIMITER_CHARS); private final List results; - private Map header; + private final Map header; private final HashMap usageHeaders; private final HashMap quotaHeaders; private final boolean useEtagAsContinuation; @@ -56,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) { @@ -450,19 +452,9 @@ void setContinuationToken(String continuationToken) { } private void setContinuationTokenInternal(String headerName, String continuationToken) { - boolean setting = !Strings.isNullOrWhiteSpace(continuationToken); - boolean clearing = !setting && !this.header.isEmpty() && this.header.containsKey(headerName); - if (!setting && !clearing) { - return; - } - - if (Utils.isImmutableMap(this.header)) { - this.header = new HashMap<>(this.header); - } - - if (setting) { + if (!Strings.isNullOrWhiteSpace(continuationToken)) { this.header.put(headerName, continuationToken); - } else { + } else if (!this.header.isEmpty() && this.header.containsKey(headerName)) { this.header.remove(headerName); } }