Skip to content

Commit ace9cf2

Browse files
Annie LiangCopilot
Annie Liang
and
Copilot
committed
perf: reduce HashMap/collection allocation overhead in gateway path
Eliminate per-response intermediate HashMap allocation by adding a new StoreResponse constructor that accepts HttpHeaders directly. Header names and values are populated into String[] arrays without materializing an intermediate Map. The JsonNodeStorePayload is updated to accept header arrays and only builds a Map lazily on error paths (extremely rare). Pre-size HashMaps throughout the hot path to avoid resize/rehash: - HttpHeaders request construction: sized to defaultHeaders + request headers - StoreResponse.replicaStatusList: pre-sized to 4 - StoreResponse.withRemappedStatusCode: pre-sized to header count - RxDocumentServiceRequest fallback maps: pre-sized to 32 Fix HttpUtils.asMap() double-allocation by iterating HttpHeaders directly instead of calling toMap() which creates an intermediate HashMap. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent 156e48e commit ace9cf2

6 files changed

Lines changed: 148 additions & 14 deletions

File tree

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -183,7 +183,7 @@ private RxDocumentServiceRequest(DiagnosticsClientContext clientContext,
183183
this.forceNameCacheRefresh = false;
184184
this.resourceType = resourceType;
185185
this.contentAsByteArray = toByteArray(byteBuffer);
186-
this.headers = headers != null ? headers : new HashMap<>();
186+
this.headers = headers != null ? headers : new HashMap<>(32);
187187
this.activityId = UUIDs.nonBlockingRandomUUID();
188188
this.isFeed = false;
189189
this.isNameBased = isNameBased;
@@ -217,7 +217,7 @@ private RxDocumentServiceRequest(DiagnosticsClientContext clientContext,
217217
this.operationType = operationType;
218218
this.resourceType = resourceType;
219219
this.requestContext.sessionToken = null;
220-
this.headers = headers != null ? headers : new HashMap<>();
220+
this.headers = headers != null ? headers : new HashMap<>(32);
221221
this.activityId = UUIDs.nonBlockingRandomUUID();
222222
this.isFeed = false;
223223

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -238,7 +238,7 @@ public StoreResponse unwrapToStoreResponse(
238238
return new StoreResponse(
239239
endpoint,
240240
statusCode,
241-
HttpUtils.unescape(headers.toLowerCaseMap()),
241+
headers,
242242
new ByteBufInputStream(retainedContent, true),
243243
size);
244244
} else {
@@ -248,7 +248,7 @@ public StoreResponse unwrapToStoreResponse(
248248
return new StoreResponse(
249249
endpoint,
250250
statusCode,
251-
HttpUtils.unescape(headers.toLowerCaseMap()),
251+
headers,
252252
null,
253253
0);
254254
}
@@ -343,7 +343,7 @@ private Mono<RxDocumentServiceResponse> performRequestInternalCore(RxDocumentSer
343343
}
344344

345345
private HttpHeaders getHttpRequestHeaders(Map<String, String> headers) {
346-
HttpHeaders httpHeaders = new HttpHeaders(this.defaultHeaders.size());
346+
HttpHeaders httpHeaders = new HttpHeaders(this.defaultHeaders.size() + (headers != null ? headers.size() : 0));
347347
// Add default headers.
348348
for (Entry<String, String> entry : this.defaultHeaders.entrySet()) {
349349
if (!headers.containsKey(entry.getKey())) {

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/HttpUtils.java‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import com.azure.cosmos.implementation.HttpConstants;
88
import com.azure.cosmos.implementation.Strings;
99
import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
10+
import com.azure.cosmos.implementation.http.HttpHeader;
1011
import com.azure.cosmos.implementation.http.HttpHeaders;
1112
import org.slf4j.Logger;
1213
import org.slf4j.LoggerFactory;
@@ -51,14 +52,14 @@ public static String urlDecode(String url) {
5152

5253
public static Map<String, String> asMap(HttpHeaders headers) {
5354
if (headers == null) {
54-
return new HashMap<>();
55+
return new HashMap<>(4);
5556
}
5657
HashMap<String, String> map = new HashMap<>(headers.size());
57-
for (Entry<String, String> entry : headers.toMap().entrySet()) {
58-
if (entry.getKey().equals(HttpConstants.HttpHeaders.OWNER_FULL_NAME)) {
59-
map.put(entry.getKey(), HttpUtils.urlDecode(entry.getValue()));
58+
for (HttpHeader header : headers) {
59+
if (header.name().equals(HttpConstants.HttpHeaders.OWNER_FULL_NAME)) {
60+
map.put(header.name(), HttpUtils.urlDecode(header.value()));
6061
} else {
61-
map.put(entry.getKey(), entry.getValue());
62+
map.put(header.name(), header.value());
6263
}
6364
}
6465
return map;

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/JsonNodeStorePayload.java‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import java.nio.charset.CodingErrorAction;
1919
import java.nio.charset.StandardCharsets;
2020
import java.util.Base64;
21+
import java.util.HashMap;
2122
import java.util.Map;
2223

2324
public class JsonNodeStorePayload implements StorePayload<JsonNode> {
@@ -36,6 +37,25 @@ public JsonNodeStorePayload(ByteBufInputStream bufferStream, int readableBytes,
3637
}
3738
}
3839

40+
/**
41+
* Creates a JsonNodeStorePayload using pre-populated header arrays instead of a Map.
42+
* The Map is constructed lazily only if needed for error reporting.
43+
*/
44+
public JsonNodeStorePayload(
45+
ByteBufInputStream bufferStream,
46+
int readableBytes,
47+
String[] headerNames,
48+
String[] headerValues) {
49+
50+
if (readableBytes > 0) {
51+
this.responsePayloadSize = readableBytes;
52+
this.jsonValue = fromJsonWithArrayHeaders(bufferStream, readableBytes, headerNames, headerValues);
53+
} else {
54+
this.responsePayloadSize = 0;
55+
this.jsonValue = null;
56+
}
57+
}
58+
3959
private static JsonNode fromJson(ByteBufInputStream bufferStream, int readableBytes, Map<String, String> responseHeaders) {
4060
byte[] bytes = new byte[readableBytes];
4161
try {
@@ -67,6 +87,51 @@ private static JsonNode fromJson(ByteBufInputStream bufferStream, int readableBy
6787
}
6888
}
6989

90+
private static JsonNode fromJsonWithArrayHeaders(
91+
ByteBufInputStream bufferStream,
92+
int readableBytes,
93+
String[] headerNames,
94+
String[] headerValues) {
95+
96+
byte[] bytes = new byte[readableBytes];
97+
try {
98+
bufferStream.read(bytes);
99+
return Utils.getSimpleObjectMapper().readTree(bytes);
100+
} catch (IOException e) {
101+
// Build Map lazily only on error (extremely rare in production)
102+
Map<String, String> responseHeaders = buildHeaderMap(headerNames, headerValues);
103+
if (fallbackCharsetDecoder != null) {
104+
logger.warn("Unable to parse JSON, fallback to use customized charset decoder.", e);
105+
return fromJsonWithFallbackCharsetDecoder(bytes, responseHeaders);
106+
} else {
107+
String baseErrorMessage = "Failed to parse JSON document. No fallback charset decoder configured.";
108+
109+
if (Configs.isNonParseableDocumentLoggingEnabled()) {
110+
String documentSample = Base64.getEncoder().encodeToString(bytes);
111+
logger.error(baseErrorMessage + " " + "Document in Base64 format: [" + documentSample + "]", e);
112+
} else {
113+
logger.error(baseErrorMessage);
114+
}
115+
116+
IllegalStateException innerException = new IllegalStateException("Unable to parse JSON.", e);
117+
118+
throw Utils.createCosmosException(
119+
HttpConstants.StatusCodes.BADREQUEST,
120+
HttpConstants.SubStatusCodes.FAILED_TO_PARSE_SERVER_RESPONSE,
121+
innerException,
122+
responseHeaders);
123+
}
124+
}
125+
}
126+
127+
private static Map<String, String> buildHeaderMap(String[] headerNames, String[] headerValues) {
128+
Map<String, String> map = new HashMap<>(headerNames.length * 4 / 3 + 1);
129+
for (int i = 0; i < headerNames.length; i++) {
130+
map.put(headerNames[i], headerValues[i]);
131+
}
132+
return map;
133+
}
134+
70135
private static JsonNode fromJsonWithFallbackCharsetDecoder(byte[] bytes, Map<String, String> responseHeaders) {
71136
try {
72137
String sanitizedJson = fallbackCharsetDecoder.decode(ByteBuffer.wrap(bytes)).toString();

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/StoreResponse.java‎

Lines changed: 55 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdChannelAcquisitionTimeline;
1010
import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdChannelStatistics;
1111
import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdEndpointStatistics;
12+
import com.azure.cosmos.implementation.http.HttpHeaders;
1213
import com.fasterxml.jackson.databind.JsonNode;
1314
import io.netty.buffer.ByteBufInputStream;
1415
import io.netty.util.IllegalReferenceCountException;
@@ -68,7 +69,7 @@ public StoreResponse(
6869
}
6970

7071
this.status = status;
71-
replicaStatusList = new HashMap<>();
72+
replicaStatusList = new HashMap<>(4);
7273
if (contentStream != null) {
7374
try {
7475
this.responsePayload = new JsonNodeStorePayload(contentStream, responsePayloadLength, headerMap);
@@ -87,6 +88,57 @@ public StoreResponse(
8788
}
8889
}
8990

91+
/**
92+
* Creates a StoreResponse directly from HttpHeaders, avoiding intermediate HashMap allocation.
93+
* Header names are stored as lowercase keys (matching HttpHeaders internal representation).
94+
* The OWNER_FULL_NAME header value is URL-decoded inline.
95+
*/
96+
public StoreResponse(
97+
String endpoint,
98+
int status,
99+
HttpHeaders httpHeaders,
100+
ByteBufInputStream contentStream,
101+
int responsePayloadLength) {
102+
103+
checkArgument((contentStream == null) == (responsePayloadLength == 0),
104+
"Parameter 'contentStream' must be consistent with 'responsePayloadLength'.");
105+
requestTimeline = RequestTimeline.empty();
106+
107+
int headerCount = httpHeaders.size();
108+
responseHeaderNames = new String[headerCount];
109+
responseHeaderValues = new String[headerCount];
110+
this.endpoint = endpoint != null ? endpoint : "";
111+
112+
httpHeaders.populateLowerCaseHeaders(responseHeaderNames, responseHeaderValues);
113+
114+
// URL-decode OWNER_FULL_NAME header value inline (replaces HttpUtils.unescape)
115+
for (int i = 0; i < headerCount; i++) {
116+
if (HttpConstants.HttpHeaders.OWNER_FULL_NAME.equals(responseHeaderNames[i])) {
117+
responseHeaderValues[i] = HttpUtils.urlDecode(responseHeaderValues[i]);
118+
break;
119+
}
120+
}
121+
122+
this.status = status;
123+
replicaStatusList = new HashMap<>(4);
124+
if (contentStream != null) {
125+
try {
126+
this.responsePayload = new JsonNodeStorePayload(
127+
contentStream, responsePayloadLength, responseHeaderNames, responseHeaderValues);
128+
} finally {
129+
try {
130+
contentStream.close();
131+
} catch (Throwable e) {
132+
if (!(e instanceof IllegalReferenceCountException)) {
133+
logger.warn("Failed to close content stream. This may cause a Netty ByteBuf leak.", e);
134+
}
135+
}
136+
}
137+
} else {
138+
this.responsePayload = null;
139+
}
140+
}
141+
90142
private StoreResponse(
91143
String endpoint,
92144
int status,
@@ -108,7 +160,7 @@ private StoreResponse(
108160
}
109161

110162
this.status = status;
111-
replicaStatusList = new HashMap<>();
163+
replicaStatusList = new HashMap<>(4);
112164
this.responsePayload = responsePayload;
113165
}
114166

@@ -310,7 +362,7 @@ public void setFaultInjectionRuleEvaluationResults(List<String> results) {
310362

311363
public StoreResponse withRemappedStatusCode(int newStatusCode, double additionalRequestCharge) {
312364

313-
Map<String, String> headers = new HashMap<>();
365+
Map<String, String> headers = new HashMap<>(this.responseHeaderNames.length * 4 / 3 + 1);
314366
for (int i = 0; i < this.responseHeaderNames.length; i++) {
315367
String headerName = this.responseHeaderNames[i];
316368
if (headerName.equalsIgnoreCase(HttpConstants.HttpHeaders.REQUEST_CHARGE)) {

‎sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/http/HttpHeaders.java‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ public class HttpHeaders implements Iterable<HttpHeader>, JsonSerializable {
2424
* Create an empty HttpHeaders instance.
2525
*/
2626
public HttpHeaders() {
27-
this.headers = new HashMap<>();
27+
this.headers = new HashMap<>(16);
2828
}
2929

3030
/**
@@ -130,6 +130,22 @@ public Map<String, String> toLowerCaseMap() {
130130
return result;
131131
}
132132

133+
/**
134+
* Populates the provided arrays with lowercased header names and their values
135+
* directly from the internal map, avoiding intermediate HashMap allocation.
136+
*
137+
* @param names array to populate with lowercased header names (must be at least size() long)
138+
* @param values array to populate with header values (must be at least size() long)
139+
*/
140+
public void populateLowerCaseHeaders(String[] names, String[] values) {
141+
int i = 0;
142+
for (Map.Entry<String, HttpHeader> entry : headers.entrySet()) {
143+
names[i] = entry.getKey();
144+
values[i] = entry.getValue().value();
145+
i++;
146+
}
147+
}
148+
133149
@Override
134150
public Iterator<HttpHeader> iterator() {
135151
return headers.values().iterator();

0 commit comments

Comments
 (0)