diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/HttpUtilsTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/HttpUtilsTest.java index 209e216391b1..e1b9d71f74fd 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/HttpUtilsTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/HttpUtilsTest.java @@ -32,11 +32,5 @@ public void verifyConversionOfHttpResponseHeadersToMap() { Entry entry = resultHeadersSet.iterator().next(); assertThat(entry.getKey()).isEqualTo(HttpConstants.HttpHeaders.OWNER_FULL_NAME); assertThat(entry.getValue()).isEqualTo(HttpUtils.urlDecode(OWNER_FULL_NAME_VALUE)); - - Map resultHeaders = HttpUtils.unescape(httpResponseHeaders.toMap()); - assertThat(resultHeaders.size()).isEqualTo(1); - entry = resultHeadersSet.iterator().next(); - assertThat(entry.getKey()).isEqualTo(HttpConstants.HttpHeaders.OWNER_FULL_NAME); - assertThat(entry.getValue()).isEqualTo(HttpUtils.urlDecode(OWNER_FULL_NAME_VALUE)); } } diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/JsonNodeStorePayloadTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/JsonNodeStorePayloadTests.java index 16a15157118d..691fa0b72186 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/JsonNodeStorePayloadTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/JsonNodeStorePayloadTests.java @@ -11,6 +11,9 @@ import java.util.HashMap; +import static com.azure.cosmos.implementation.Utils.getUTF8BytesOrNull; +import static org.assertj.core.api.Assertions.assertThat; + public class JsonNodeStorePayloadTests { @Test(groups = {"unit"}) @Ignore("fallbackCharsetDecoder will only be initialized during the first time when JsonNodeStorePayload loaded," + @@ -46,4 +49,47 @@ private static byte[] hexStringToByteArray(String hex) { return data; } + + @Test(groups = {"unit"}) + public void arrayHeaderConstructorParsesValidJson() { + String jsonContent = "{\"id\":\"test\",\"name\":\"value\"}"; + String[] headerNames = {"content-type", "x-request-id"}; + String[] headerValues = {"application/json", "req-123"}; + + ByteBuf buffer = getUTF8BytesOrNull(jsonContent); + JsonNodeStorePayload payload = new JsonNodeStorePayload( + new ByteBufInputStream(buffer, true), buffer.readableBytes(), headerNames, headerValues); + + assertThat(payload.getPayload()).isNotNull(); + assertThat(payload.getPayload().get("id").asText()).isEqualTo("test"); + assertThat(payload.getPayload().get("name").asText()).isEqualTo("value"); + assertThat(payload.getResponsePayloadSize()).isEqualTo(jsonContent.getBytes().length); + } + + @Test(groups = {"unit"}) + public void arrayHeaderConstructorWithEmptyStreamReturnsNull() { + String[] headerNames = {"content-type"}; + String[] headerValues = {"application/json"}; + + ByteBuf buffer = Unpooled.EMPTY_BUFFER; + JsonNodeStorePayload payload = new JsonNodeStorePayload( + new ByteBufInputStream(buffer), 0, headerNames, headerValues); + + assertThat(payload.getPayload()).isNull(); + assertThat(payload.getResponsePayloadSize()).isEqualTo(0); + } + + @Test(groups = {"unit"}) + public void mapConstructorParsesValidJson() { + String jsonContent = "{\"id\":\"test\"}"; + HashMap headers = new HashMap<>(); + headers.put("content-type", "application/json"); + + ByteBuf buffer = getUTF8BytesOrNull(jsonContent); + JsonNodeStorePayload payload = new JsonNodeStorePayload( + new ByteBufInputStream(buffer, true), buffer.readableBytes(), headers); + + assertThat(payload.getPayload()).isNotNull(); + assertThat(payload.getPayload().get("id").asText()).isEqualTo("test"); + } } diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/StoreResponseTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/StoreResponseTest.java index 1e6c6cc147f8..4d823accfbdd 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/StoreResponseTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/StoreResponseTest.java @@ -3,6 +3,8 @@ package com.azure.cosmos.implementation.directconnectivity; +import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.implementation.http.HttpHeaders; import io.netty.buffer.ByteBuf; import io.netty.buffer.ByteBufInputStream; import org.testng.annotations.Test; @@ -47,4 +49,84 @@ public void headerNamesAreCaseInsensitive() { assertThat(sp.getHeaderValue("kEy2")).isEqualTo("value2"); assertThat(sp.getHeaderValue("KEY3")).isEqualTo("value3"); } + + @Test(groups = { "unit" }) + public void httpHeadersConstructorProducesSameResultAsMapConstructor() { + String jsonContent = "{\"id\":\"test\"}"; + HashMap headerMap = new HashMap<>(); + headerMap.put("key1", "value1"); + headerMap.put("Content-Type", "application/json"); + headerMap.put("X-Custom-Header", "customValue"); + + HttpHeaders httpHeaders = new HttpHeaders(); + httpHeaders.set("key1", "value1"); + httpHeaders.set("Content-Type", "application/json"); + httpHeaders.set("X-Custom-Header", "customValue"); + + ByteBuf buffer1 = getUTF8BytesOrNull(jsonContent); + StoreResponse fromMap = new StoreResponse( + "endpoint1", 200, headerMap, new ByteBufInputStream(buffer1, true), buffer1.readableBytes()); + + ByteBuf buffer2 = getUTF8BytesOrNull(jsonContent); + StoreResponse fromHttpHeaders = new StoreResponse( + "endpoint1", 200, httpHeaders, new ByteBufInputStream(buffer2, true), buffer2.readableBytes()); + + assertThat(fromHttpHeaders.getStatus()).isEqualTo(fromMap.getStatus()); + assertThat(fromHttpHeaders.getEndpoint()).isEqualTo(fromMap.getEndpoint()); + + // Verify all headers are accessible with case-insensitive lookup + assertThat(fromHttpHeaders.getHeaderValue("key1")).isEqualTo("value1"); + assertThat(fromHttpHeaders.getHeaderValue("content-type")).isEqualTo("application/json"); + assertThat(fromHttpHeaders.getHeaderValue("x-custom-header")).isEqualTo("customValue"); + + // HttpHeaders constructor stores lowercase names + String[] headerNames = fromHttpHeaders.getResponseHeaderNames(); + for (String name : headerNames) { + assertThat(name).isEqualTo(name.toLowerCase()); + } + } + + @Test(groups = { "unit" }) + public void httpHeadersConstructorWithNullEndpoint() { + String jsonContent = "{\"id\":\"test\"}"; + HttpHeaders httpHeaders = new HttpHeaders(); + httpHeaders.set("key1", "value1"); + + ByteBuf buffer = getUTF8BytesOrNull(jsonContent); + StoreResponse sp = new StoreResponse( + null, 200, httpHeaders, new ByteBufInputStream(buffer, true), buffer.readableBytes()); + + assertThat(sp.getEndpoint()).isEqualTo(""); + assertThat(sp.getHeaderValue("key1")).isEqualTo("value1"); + } + + @Test(groups = { "unit" }) + public void httpHeadersConstructorWithNoContent() { + HttpHeaders httpHeaders = new HttpHeaders(); + httpHeaders.set("key1", "value1"); + + StoreResponse sp = new StoreResponse("endpoint", 204, httpHeaders, null, 0); + + assertThat(sp.getStatus()).isEqualTo(204); + assertThat(sp.getResponseBodyAsJson()).isNull(); + assertThat(sp.getHeaderValue("key1")).isEqualTo("value1"); + } + + @Test(groups = { "unit" }) + public void httpHeadersConstructorDecodesOwnerFullName() { + // OWNER_FULL_NAME value with URL-encoded segments (e.g. spaces encoded as %20) + String encodedOwner = "dbs%2FmyDb%2Fcolls%2Fmy%20Collection"; + String expectedDecoded = "dbs/myDb/colls/my Collection"; + + HttpHeaders httpHeaders = new HttpHeaders(); + httpHeaders.set(HttpConstants.HttpHeaders.OWNER_FULL_NAME, encodedOwner); + httpHeaders.set("X-Other", "plain"); + + StoreResponse sp = new StoreResponse("endpoint", 200, httpHeaders, null, 0); + + // The encoded OWNER_FULL_NAME should be URL-decoded when accessed via getHeaderValue + assertThat(sp.getHeaderValue(HttpConstants.HttpHeaders.OWNER_FULL_NAME)).isEqualTo(expectedDecoded); + // Other headers are left as-is + assertThat(sp.getHeaderValue("X-Other")).isEqualTo("plain"); + } } diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/http/HttpHeadersTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/http/HttpHeadersTests.java index e59cf731e40d..013d0af00ff9 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/http/HttpHeadersTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/http/HttpHeadersTests.java @@ -5,9 +5,11 @@ import org.testng.annotations.Test; +import java.util.Locale; import java.util.Map; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; public class HttpHeadersTests { @@ -26,4 +28,78 @@ public void caseInsensitiveToMap() { assertThat(caseSensitiveMap.get(headerName.toLowerCase())).isNull(); assertThat(caseSensitiveMap.get(headerName)).isEqualTo(headerValue); } + + @Test(groups = "unit") + public void populateLowerCaseHeadersProducesLowercaseNames() { + HttpHeaders headers = new HttpHeaders(); + headers.set("Content-Type", "application/json"); + headers.set("X-Ms-Request-Id", "abc-123"); + headers.set("ETag", "\"v1\""); + + String[] names = new String[headers.size()]; + String[] values = new String[headers.size()]; + headers.populateLowerCaseHeaders(names, values); + + // All names should be lowercase + for (String name : names) { + assertThat(name).isEqualTo(name.toLowerCase(Locale.ROOT)); + } + + // Verify values are present (order depends on HashMap iteration, so use containment) + Map resultMap = new java.util.HashMap<>(); + for (int i = 0; i < names.length; i++) { + resultMap.put(names[i], values[i]); + } + + assertThat(resultMap).containsEntry("content-type", "application/json"); + assertThat(resultMap).containsEntry("x-ms-request-id", "abc-123"); + assertThat(resultMap).containsEntry("etag", "\"v1\""); + } + + @Test(groups = "unit") + public void populateLowerCaseHeadersWithEmptyHeaders() { + HttpHeaders headers = new HttpHeaders(); + + String[] names = new String[0]; + String[] values = new String[0]; + headers.populateLowerCaseHeaders(names, values); + + assertThat(names).isEmpty(); + assertThat(values).isEmpty(); + } + + @Test(groups = "unit") + public void populateLowerCaseHeadersRejectsNullNames() { + HttpHeaders headers = new HttpHeaders(); + headers.set("Key", "value"); + + assertThatThrownBy(() -> headers.populateLowerCaseHeaders(null, new String[1])) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("names"); + } + + @Test(groups = "unit") + public void populateLowerCaseHeadersRejectsNullValues() { + HttpHeaders headers = new HttpHeaders(); + headers.set("Key", "value"); + + assertThatThrownBy(() -> headers.populateLowerCaseHeaders(new String[1], null)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("values"); + } + + @Test(groups = "unit") + public void populateLowerCaseHeadersRejectsTooSmallArrays() { + HttpHeaders headers = new HttpHeaders(); + headers.set("A", "1"); + headers.set("B", "2"); + + assertThatThrownBy(() -> headers.populateLowerCaseHeaders(new String[1], new String[2])) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("names"); + + assertThatThrownBy(() -> headers.populateLowerCaseHeaders(new String[2], new String[1])) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("values"); + } } 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 58e838d482b8..dee68a1512b1 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 @@ -191,7 +191,9 @@ private RxDocumentServiceRequest(DiagnosticsClientContext clientContext, this.forceNameCacheRefresh = false; this.resourceType = resourceType; this.contentAsByteArray = toByteArray(byteBuffer); - this.headers = headers != null ? headers : new HashMap<>(); + // Pre-size to 32 (threshold 24 at 0.75 load factor) to accommodate typical request + // headers (auth, content-type, consistency, session-token, partition-key, etc.) without resize. + this.headers = headers != null ? headers : new HashMap<>(32); this.activityId = UUIDs.nonBlockingRandomUUID(); this.isFeed = false; this.isNameBased = isNameBased; @@ -225,7 +227,9 @@ private RxDocumentServiceRequest(DiagnosticsClientContext clientContext, this.operationType = operationType; this.resourceType = resourceType; this.requestContext.sessionToken = null; - this.headers = headers != null ? headers : new HashMap<>(); + // Pre-size to 32 (threshold 24 at 0.75 load factor) to accommodate typical request + // headers (auth, content-type, consistency, session-token, partition-key, etc.) without resize. + this.headers = headers != null ? headers : new HashMap<>(32); this.activityId = UUIDs.nonBlockingRandomUUID(); this.isFeed = false; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java index 81a553f01547..9106ab8a1abf 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java @@ -242,7 +242,7 @@ public StoreResponse unwrapToStoreResponse( return new StoreResponse( endpoint, statusCode, - HttpUtils.unescape(headers.toLowerCaseMap()), + headers, new ByteBufInputStream(retainedContent, true), size); } else { @@ -252,7 +252,7 @@ public StoreResponse unwrapToStoreResponse( return new StoreResponse( endpoint, statusCode, - HttpUtils.unescape(headers.toLowerCaseMap()), + headers, null, 0); } @@ -347,7 +347,7 @@ private Mono performRequestInternalCore(RxDocumentSer } private HttpHeaders getHttpRequestHeaders(Map headers) { - HttpHeaders httpHeaders = new HttpHeaders(this.defaultHeaders.size()); + HttpHeaders httpHeaders = new HttpHeaders(HttpUtils.mapCapacityForSize(this.defaultHeaders.size() + headers.size())); // Add default headers. for (Entry entry : this.defaultHeaders.entrySet()) { if (!headers.containsKey(entry.getKey())) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/HttpUtils.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/HttpUtils.java index 2e5d27ac03dd..3ec8a48d6da3 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/HttpUtils.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/HttpUtils.java @@ -7,6 +7,7 @@ import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.Strings; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.implementation.http.HttpHeader; import com.azure.cosmos.implementation.http.HttpHeaders; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -14,13 +15,8 @@ import java.io.UnsupportedEncodingException; import java.net.URLDecoder; import java.net.URLEncoder; -import java.util.AbstractMap; -import java.util.ArrayList; import java.util.HashMap; -import java.util.List; import java.util.Map; -import java.util.Map.Entry; -import java.util.Set; import java.util.regex.Pattern; public class HttpUtils { @@ -29,6 +25,18 @@ public class HttpUtils { private static final Pattern PLUS_SYMBOL_ESCAPE_PATTERN = Pattern.compile(UrlEncodingInfo.PLUS_SYMBOL_ESCAPED); + /** + * Returns the initial capacity for a HashMap that will hold {@code expectedSize} entries + * without resizing, accounting for the default load factor of 0.75. + */ + public static int mapCapacityForSize(int expectedSize) { + if (expectedSize <= 0) { + return 1; + } + long capacity = (long) expectedSize * 4 / 3 + 1; + return capacity >= Integer.MAX_VALUE ? Integer.MAX_VALUE : (int) capacity; + } + public static String urlEncode(String url) { try { return PLUS_SYMBOL_ESCAPE_PATTERN.matcher(URLEncoder.encode(url, UrlEncodingInfo.UTF_8)) @@ -51,14 +59,14 @@ public static String urlDecode(String url) { public static Map asMap(HttpHeaders headers) { if (headers == null) { - return new HashMap<>(); + return new HashMap<>(4); } - HashMap map = new HashMap<>(headers.size()); - for (Entry entry : headers.toMap().entrySet()) { - if (entry.getKey().equals(HttpConstants.HttpHeaders.OWNER_FULL_NAME)) { - map.put(entry.getKey(), HttpUtils.urlDecode(entry.getValue())); + HashMap map = new HashMap<>(mapCapacityForSize(headers.size())); + for (HttpHeader header : headers) { + if (header.name().equals(HttpConstants.HttpHeaders.OWNER_FULL_NAME)) { + map.put(header.name(), HttpUtils.urlDecode(header.value())); } else { - map.put(entry.getKey(), entry.getValue()); + map.put(header.name(), header.value()); } } return map; @@ -78,24 +86,4 @@ public static String getDateHeader(Map headerValues) { return date != null ? date : StringUtils.EMPTY; } - - public static List> unescape(Set> headers) { - List> result = new ArrayList<>(headers.size()); - for (Entry entry : headers) { - if (entry.getKey().equals(HttpConstants.HttpHeaders.OWNER_FULL_NAME)) { - String unescapedUrl = HttpUtils.urlDecode(entry.getValue()); - entry = new AbstractMap.SimpleEntry<>(entry.getKey(), unescapedUrl); - } - result.add(entry); - } - return result; - } - - public static Map unescape(Map headers) { - if (headers != null) { - headers.computeIfPresent(HttpConstants.HttpHeaders.OWNER_FULL_NAME, - (ownerKey, ownerValue) -> HttpUtils.urlDecode(ownerValue)); - } - return headers; - } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/JsonNodeStorePayload.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/JsonNodeStorePayload.java index bbf642ba667a..ae81cd22c1c3 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/JsonNodeStorePayload.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/JsonNodeStorePayload.java @@ -18,7 +18,9 @@ import java.nio.charset.CodingErrorAction; import java.nio.charset.StandardCharsets; import java.util.Base64; +import java.util.HashMap; import java.util.Map; +import java.util.function.Supplier; public class JsonNodeStorePayload implements StorePayload { private static final Logger logger = LoggerFactory.getLogger(JsonNodeStorePayload.class); @@ -29,24 +31,47 @@ public class JsonNodeStorePayload implements StorePayload { public JsonNodeStorePayload(ByteBufInputStream bufferStream, int readableBytes, Map responseHeaders) { if (readableBytes > 0) { this.responsePayloadSize = readableBytes; - this.jsonValue = fromJson(bufferStream, readableBytes, responseHeaders); + this.jsonValue = parseJson(bufferStream, readableBytes, () -> responseHeaders); } else { this.responsePayloadSize = 0; this.jsonValue = null; } } - private static JsonNode fromJson(ByteBufInputStream bufferStream, int readableBytes, Map responseHeaders) { + /** + * Creates a JsonNodeStorePayload using pre-populated header arrays instead of a Map. + * The Map is constructed lazily only if needed for error reporting. + */ + public JsonNodeStorePayload( + ByteBufInputStream bufferStream, + int readableBytes, + String[] headerNames, + String[] headerValues) { + + if (readableBytes > 0) { + this.responsePayloadSize = readableBytes; + this.jsonValue = parseJson(bufferStream, readableBytes, () -> buildHeaderMap(headerNames, headerValues)); + } else { + this.responsePayloadSize = 0; + this.jsonValue = null; + } + } + + private static JsonNode parseJson( + ByteBufInputStream bufferStream, + int readableBytes, + Supplier> headersSupplier) { + byte[] bytes = new byte[readableBytes]; try { bufferStream.read(bytes); return Utils.getSimpleObjectMapper().readTree(bytes); } catch (IOException e) { + Map responseHeaders = headersSupplier.get(); if (fallbackCharsetDecoder != null) { logger.warn("Unable to parse JSON, fallback to use customized charset decoder.", e); return fromJsonWithFallbackCharsetDecoder(bytes, responseHeaders); } else { - String baseErrorMessage = "Failed to parse JSON document. No fallback charset decoder configured."; if (Configs.isNonParseableDocumentLoggingEnabled()) { @@ -67,6 +92,14 @@ private static JsonNode fromJson(ByteBufInputStream bufferStream, int readableBy } } + private static Map buildHeaderMap(String[] headerNames, String[] headerValues) { + Map map = new HashMap<>(HttpUtils.mapCapacityForSize(headerNames.length)); + for (int i = 0; i < headerNames.length; i++) { + map.put(headerNames[i], headerValues[i]); + } + return map; + } + private static JsonNode fromJsonWithFallbackCharsetDecoder(byte[] bytes, Map responseHeaders) { try { String sanitizedJson = fallbackCharsetDecoder.decode(ByteBuffer.wrap(bytes)).toString(); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/ResponseUtils.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/ResponseUtils.java index bd5bf896eac4..fb00077053cd 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/ResponseUtils.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/ResponseUtils.java @@ -27,7 +27,7 @@ static Mono toStoreResponse(HttpResponse httpClientResponse, Stri return new StoreResponse( endpoint, httpClientResponse.statusCode(), - HttpUtils.unescape(httpResponseHeaders.toMap()), + httpResponseHeaders, null, 0); } @@ -35,7 +35,7 @@ static Mono toStoreResponse(HttpResponse httpClientResponse, Stri return new StoreResponse( endpoint, httpClientResponse.statusCode(), - HttpUtils.unescape(httpResponseHeaders.toMap()), + httpResponseHeaders, new ByteBufInputStream(byteBufContent, true), size); }); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResponse.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResponse.java index 10c7f40e9ae3..afce8c917428 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResponse.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResponse.java @@ -9,6 +9,7 @@ import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdChannelAcquisitionTimeline; import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdChannelStatistics; import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdEndpointStatistics; +import com.azure.cosmos.implementation.http.HttpHeaders; import com.fasterxml.jackson.databind.JsonNode; import io.netty.buffer.ByteBufInputStream; import io.netty.util.IllegalReferenceCountException; @@ -29,6 +30,11 @@ */ public class StoreResponse { private static final Logger logger = LoggerFactory.getLogger(StoreResponse.class.getSimpleName()); + + // Initial capacity for the replica-status map. Chosen to avoid resizing in + // the common case where only a handful of replica status entries are tracked. + private static final int REPLICA_STATUS_MAP_INITIAL_CAPACITY = 6; + final private int status; final private String[] responseHeaderNames; final private String[] responseHeaderValues; @@ -68,23 +74,48 @@ public StoreResponse( } this.status = status; - replicaStatusList = new HashMap<>(); - if (contentStream != null) { - try { - this.responsePayload = new JsonNodeStorePayload(contentStream, responsePayloadLength, headerMap); - } finally { - try { - contentStream.close(); - } catch (Throwable e) { - if (!(e instanceof IllegalReferenceCountException)) { - // Log as warning instead of debug to make ByteBuf leak issues more visible - logger.warn("Failed to close content stream. This may cause a Netty ByteBuf leak.", e); - } - } + replicaStatusList = new HashMap<>(REPLICA_STATUS_MAP_INITIAL_CAPACITY); + this.responsePayload = parseResponsePayload( + contentStream, responsePayloadLength, responseHeaderNames, responseHeaderValues); + } + + /** + * Creates a StoreResponse directly from HttpHeaders, avoiding intermediate HashMap allocation. + * Header names are stored as lowercase keys (matching HttpHeaders internal representation). + * The OWNER_FULL_NAME header value is URL-decoded inline (equivalent to HttpUtils.unescape). + */ + public StoreResponse( + String endpoint, + int status, + HttpHeaders httpHeaders, + ByteBufInputStream contentStream, + int responsePayloadLength) { + + checkArgument((contentStream == null) == (responsePayloadLength == 0), + "Parameter 'contentStream' must be consistent with 'responsePayloadLength'."); + requestTimeline = RequestTimeline.empty(); + + int headerCount = httpHeaders.size(); + responseHeaderNames = new String[headerCount]; + responseHeaderValues = new String[headerCount]; + this.endpoint = endpoint != null ? endpoint : ""; + + httpHeaders.populateLowerCaseHeaders(responseHeaderNames, responseHeaderValues); + + // URL-decode OWNER_FULL_NAME header value inline (replaces HttpUtils.unescape). + // This is kept separate from populateLowerCaseHeaders because HttpHeaders is a + // general-purpose HTTP class and should not contain Cosmos-specific URL-decoding logic. + for (int i = 0; i < headerCount; i++) { + if (HttpConstants.HttpHeaders.OWNER_FULL_NAME.equalsIgnoreCase(responseHeaderNames[i])) { + responseHeaderValues[i] = HttpUtils.urlDecode(responseHeaderValues[i]); + break; } - } else { - this.responsePayload = null; } + + this.status = status; + replicaStatusList = new HashMap<>(REPLICA_STATUS_MAP_INITIAL_CAPACITY); + this.responsePayload = parseResponsePayload( + contentStream, responsePayloadLength, responseHeaderNames, responseHeaderValues); } private StoreResponse( @@ -108,10 +139,32 @@ private StoreResponse( } this.status = status; - replicaStatusList = new HashMap<>(); + replicaStatusList = new HashMap<>(REPLICA_STATUS_MAP_INITIAL_CAPACITY); this.responsePayload = responsePayload; } + private static JsonNodeStorePayload parseResponsePayload( + ByteBufInputStream contentStream, + int responsePayloadLength, + String[] headerNames, + String[] headerValues) { + + if (contentStream == null) { + return null; + } + try { + return new JsonNodeStorePayload(contentStream, responsePayloadLength, headerNames, headerValues); + } finally { + try { + contentStream.close(); + } catch (Throwable e) { + if (!(e instanceof IllegalReferenceCountException)) { + logger.warn("Failed to close content stream. This may cause a Netty ByteBuf leak.", e); + } + } + } + } + public int getStatus() { return status; } @@ -310,7 +363,7 @@ public void setFaultInjectionRuleEvaluationResults(List results) { public StoreResponse withRemappedStatusCode(int newStatusCode, double additionalRequestCharge) { - Map headers = new HashMap<>(); + Map headers = new HashMap<>(HttpUtils.mapCapacityForSize(this.responseHeaderNames.length)); for (int i = 0; i < this.responseHeaderNames.length; i++) { String headerName = this.responseHeaderNames[i]; if (headerName.equalsIgnoreCase(HttpConstants.HttpHeaders.REQUEST_CHARGE)) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/http/HttpHeaders.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/http/HttpHeaders.java index 7090bd453a51..40bdafb5a807 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/http/HttpHeaders.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/http/HttpHeaders.java @@ -130,6 +130,41 @@ public Map toLowerCaseMap() { return result; } + /** + * Populates the provided arrays with lowercased header names and their values + * directly from the internal map, avoiding intermediate HashMap allocation. + * + *

Keys are guaranteed lowercase because {@link #set(String, String)} stores them + * via {@code name.toLowerCase(Locale.ROOT)} as the map key.

+ * + * @param names array to populate with lowercased header names (must be at least size() long) + * @param values array to populate with header values (must be at least size() long) + */ + public void populateLowerCaseHeaders(String[] names, String[] values) { + if (names == null) { + throw new IllegalArgumentException("Parameter 'names' must not be null."); + } + if (values == null) { + throw new IllegalArgumentException("Parameter 'values' must not be null."); + } + int headerCount = headers.size(); + if (names.length < headerCount) { + throw new IllegalArgumentException( + "Parameter 'names' must have length >= size(). Required: " + headerCount + ", actual: " + names.length); + } + if (values.length < headerCount) { + throw new IllegalArgumentException( + "Parameter 'values' must have length >= size(). Required: " + headerCount + ", actual: " + values.length); + } + + int i = 0; + for (Map.Entry entry : headers.entrySet()) { + names[i] = entry.getKey(); + values[i] = entry.getValue().value(); + i++; + } + } + @Override public Iterator iterator() { return headers.values().iterator();