Skip to content

Commit 630c930

Browse files
xinlian12Annie LiangCopilot
authored
perf: reduce HashMap/collection allocation overhead in gateway path (#48662)
* perf: reduce HashMap/collection allocation overhead in gateway path --------- Co-authored-by: Annie Liang <xinlian@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent ea8931d commit 630c930

11 files changed

Lines changed: 375 additions & 64 deletions

File tree

‎sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/HttpUtilsTest.java‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -32,11 +32,5 @@ public void verifyConversionOfHttpResponseHeadersToMap() {
3232
Entry<String, String> entry = resultHeadersSet.iterator().next();
3333
assertThat(entry.getKey()).isEqualTo(HttpConstants.HttpHeaders.OWNER_FULL_NAME);
3434
assertThat(entry.getValue()).isEqualTo(HttpUtils.urlDecode(OWNER_FULL_NAME_VALUE));
35-
36-
Map<String, String> resultHeaders = HttpUtils.unescape(httpResponseHeaders.toMap());
37-
assertThat(resultHeaders.size()).isEqualTo(1);
38-
entry = resultHeadersSet.iterator().next();
39-
assertThat(entry.getKey()).isEqualTo(HttpConstants.HttpHeaders.OWNER_FULL_NAME);
40-
assertThat(entry.getValue()).isEqualTo(HttpUtils.urlDecode(OWNER_FULL_NAME_VALUE));
4135
}
4236
}

‎sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/JsonNodeStorePayloadTests.java‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,9 @@
1111

1212
import java.util.HashMap;
1313

14+
import static com.azure.cosmos.implementation.Utils.getUTF8BytesOrNull;
15+
import static org.assertj.core.api.Assertions.assertThat;
16+
1417
public class JsonNodeStorePayloadTests {
1518
@Test(groups = {"unit"})
1619
@Ignore("fallbackCharsetDecoder will only be initialized during the first time when JsonNodeStorePayload loaded," +
@@ -46,4 +49,47 @@ private static byte[] hexStringToByteArray(String hex) {
4649

4750
return data;
4851
}
52+
53+
@Test(groups = {"unit"})
54+
public void arrayHeaderConstructorParsesValidJson() {
55+
String jsonContent = "{\"id\":\"test\",\"name\":\"value\"}";
56+
String[] headerNames = {"content-type", "x-request-id"};
57+
String[] headerValues = {"application/json", "req-123"};
58+
59+
ByteBuf buffer = getUTF8BytesOrNull(jsonContent);
60+
JsonNodeStorePayload payload = new JsonNodeStorePayload(
61+
new ByteBufInputStream(buffer, true), buffer.readableBytes(), headerNames, headerValues);
62+
63+
assertThat(payload.getPayload()).isNotNull();
64+
assertThat(payload.getPayload().get("id").asText()).isEqualTo("test");
65+
assertThat(payload.getPayload().get("name").asText()).isEqualTo("value");
66+
assertThat(payload.getResponsePayloadSize()).isEqualTo(jsonContent.getBytes().length);
67+
}
68+
69+
@Test(groups = {"unit"})
70+
public void arrayHeaderConstructorWithEmptyStreamReturnsNull() {
71+
String[] headerNames = {"content-type"};
72+
String[] headerValues = {"application/json"};
73+
74+
ByteBuf buffer = Unpooled.EMPTY_BUFFER;
75+
JsonNodeStorePayload payload = new JsonNodeStorePayload(
76+
new ByteBufInputStream(buffer), 0, headerNames, headerValues);
77+
78+
assertThat(payload.getPayload()).isNull();
79+
assertThat(payload.getResponsePayloadSize()).isEqualTo(0);
80+
}
81+
82+
@Test(groups = {"unit"})
83+
public void mapConstructorParsesValidJson() {
84+
String jsonContent = "{\"id\":\"test\"}";
85+
HashMap<String, String> headers = new HashMap<>();
86+
headers.put("content-type", "application/json");
87+
88+
ByteBuf buffer = getUTF8BytesOrNull(jsonContent);
89+
JsonNodeStorePayload payload = new JsonNodeStorePayload(
90+
new ByteBufInputStream(buffer, true), buffer.readableBytes(), headers);
91+
92+
assertThat(payload.getPayload()).isNotNull();
93+
assertThat(payload.getPayload().get("id").asText()).isEqualTo("test");
94+
}
4995
}

‎sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/StoreResponseTest.java‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@
33

44
package com.azure.cosmos.implementation.directconnectivity;
55

6+
import com.azure.cosmos.implementation.HttpConstants;
7+
import com.azure.cosmos.implementation.http.HttpHeaders;
68
import io.netty.buffer.ByteBuf;
79
import io.netty.buffer.ByteBufInputStream;
810
import org.testng.annotations.Test;
@@ -47,4 +49,84 @@ public void headerNamesAreCaseInsensitive() {
4749
assertThat(sp.getHeaderValue("kEy2")).isEqualTo("value2");
4850
assertThat(sp.getHeaderValue("KEY3")).isEqualTo("value3");
4951
}
52+
53+
@Test(groups = { "unit" })
54+
public void httpHeadersConstructorProducesSameResultAsMapConstructor() {
55+
String jsonContent = "{\"id\":\"test\"}";
56+
HashMap<String, String> headerMap = new HashMap<>();
57+
headerMap.put("key1", "value1");
58+
headerMap.put("Content-Type", "application/json");
59+
headerMap.put("X-Custom-Header", "customValue");
60+
61+
HttpHeaders httpHeaders = new HttpHeaders();
62+
httpHeaders.set("key1", "value1");
63+
httpHeaders.set("Content-Type", "application/json");
64+
httpHeaders.set("X-Custom-Header", "customValue");
65+
66+
ByteBuf buffer1 = getUTF8BytesOrNull(jsonContent);
67+
StoreResponse fromMap = new StoreResponse(
68+
"endpoint1", 200, headerMap, new ByteBufInputStream(buffer1, true), buffer1.readableBytes());
69+
70+
ByteBuf buffer2 = getUTF8BytesOrNull(jsonContent);
71+
StoreResponse fromHttpHeaders = new StoreResponse(
72+
"endpoint1", 200, httpHeaders, new ByteBufInputStream(buffer2, true), buffer2.readableBytes());
73+
74+
assertThat(fromHttpHeaders.getStatus()).isEqualTo(fromMap.getStatus());
75+
assertThat(fromHttpHeaders.getEndpoint()).isEqualTo(fromMap.getEndpoint());
76+
77+
// Verify all headers are accessible with case-insensitive lookup
78+
assertThat(fromHttpHeaders.getHeaderValue("key1")).isEqualTo("value1");
79+
assertThat(fromHttpHeaders.getHeaderValue("content-type")).isEqualTo("application/json");
80+
assertThat(fromHttpHeaders.getHeaderValue("x-custom-header")).isEqualTo("customValue");
81+
82+
// HttpHeaders constructor stores lowercase names
83+
String[] headerNames = fromHttpHeaders.getResponseHeaderNames();
84+
for (String name : headerNames) {
85+
assertThat(name).isEqualTo(name.toLowerCase());
86+
}
87+
}
88+
89+
@Test(groups = { "unit" })
90+
public void httpHeadersConstructorWithNullEndpoint() {
91+
String jsonContent = "{\"id\":\"test\"}";
92+
HttpHeaders httpHeaders = new HttpHeaders();
93+
httpHeaders.set("key1", "value1");
94+
95+
ByteBuf buffer = getUTF8BytesOrNull(jsonContent);
96+
StoreResponse sp = new StoreResponse(
97+
null, 200, httpHeaders, new ByteBufInputStream(buffer, true), buffer.readableBytes());
98+
99+
assertThat(sp.getEndpoint()).isEqualTo("");
100+
assertThat(sp.getHeaderValue("key1")).isEqualTo("value1");
101+
}
102+
103+
@Test(groups = { "unit" })
104+
public void httpHeadersConstructorWithNoContent() {
105+
HttpHeaders httpHeaders = new HttpHeaders();
106+
httpHeaders.set("key1", "value1");
107+
108+
StoreResponse sp = new StoreResponse("endpoint", 204, httpHeaders, null, 0);
109+
110+
assertThat(sp.getStatus()).isEqualTo(204);
111+
assertThat(sp.getResponseBodyAsJson()).isNull();
112+
assertThat(sp.getHeaderValue("key1")).isEqualTo("value1");
113+
}
114+
115+
@Test(groups = { "unit" })
116+
public void httpHeadersConstructorDecodesOwnerFullName() {
117+
// OWNER_FULL_NAME value with URL-encoded segments (e.g. spaces encoded as %20)
118+
String encodedOwner = "dbs%2FmyDb%2Fcolls%2Fmy%20Collection";
119+
String expectedDecoded = "dbs/myDb/colls/my Collection";
120+
121+
HttpHeaders httpHeaders = new HttpHeaders();
122+
httpHeaders.set(HttpConstants.HttpHeaders.OWNER_FULL_NAME, encodedOwner);
123+
httpHeaders.set("X-Other", "plain");
124+
125+
StoreResponse sp = new StoreResponse("endpoint", 200, httpHeaders, null, 0);
126+
127+
// The encoded OWNER_FULL_NAME should be URL-decoded when accessed via getHeaderValue
128+
assertThat(sp.getHeaderValue(HttpConstants.HttpHeaders.OWNER_FULL_NAME)).isEqualTo(expectedDecoded);
129+
// Other headers are left as-is
130+
assertThat(sp.getHeaderValue("X-Other")).isEqualTo("plain");
131+
}
50132
}

‎sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/http/HttpHeadersTests.java‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,11 @@
55

66
import org.testng.annotations.Test;
77

8+
import java.util.Locale;
89
import java.util.Map;
910

1011
import static org.assertj.core.api.Assertions.assertThat;
12+
import static org.assertj.core.api.Assertions.assertThatThrownBy;
1113

1214
public class HttpHeadersTests {
1315

@@ -26,4 +28,78 @@ public void caseInsensitiveToMap() {
2628
assertThat(caseSensitiveMap.get(headerName.toLowerCase())).isNull();
2729
assertThat(caseSensitiveMap.get(headerName)).isEqualTo(headerValue);
2830
}
31+
32+
@Test(groups = "unit")
33+
public void populateLowerCaseHeadersProducesLowercaseNames() {
34+
HttpHeaders headers = new HttpHeaders();
35+
headers.set("Content-Type", "application/json");
36+
headers.set("X-Ms-Request-Id", "abc-123");
37+
headers.set("ETag", "\"v1\"");
38+
39+
String[] names = new String[headers.size()];
40+
String[] values = new String[headers.size()];
41+
headers.populateLowerCaseHeaders(names, values);
42+
43+
// All names should be lowercase
44+
for (String name : names) {
45+
assertThat(name).isEqualTo(name.toLowerCase(Locale.ROOT));
46+
}
47+
48+
// Verify values are present (order depends on HashMap iteration, so use containment)
49+
Map<String, String> resultMap = new java.util.HashMap<>();
50+
for (int i = 0; i < names.length; i++) {
51+
resultMap.put(names[i], values[i]);
52+
}
53+
54+
assertThat(resultMap).containsEntry("content-type", "application/json");
55+
assertThat(resultMap).containsEntry("x-ms-request-id", "abc-123");
56+
assertThat(resultMap).containsEntry("etag", "\"v1\"");
57+
}
58+
59+
@Test(groups = "unit")
60+
public void populateLowerCaseHeadersWithEmptyHeaders() {
61+
HttpHeaders headers = new HttpHeaders();
62+
63+
String[] names = new String[0];
64+
String[] values = new String[0];
65+
headers.populateLowerCaseHeaders(names, values);
66+
67+
assertThat(names).isEmpty();
68+
assertThat(values).isEmpty();
69+
}
70+
71+
@Test(groups = "unit")
72+
public void populateLowerCaseHeadersRejectsNullNames() {
73+
HttpHeaders headers = new HttpHeaders();
74+
headers.set("Key", "value");
75+
76+
assertThatThrownBy(() -> headers.populateLowerCaseHeaders(null, new String[1]))
77+
.isInstanceOf(IllegalArgumentException.class)
78+
.hasMessageContaining("names");
79+
}
80+
81+
@Test(groups = "unit")
82+
public void populateLowerCaseHeadersRejectsNullValues() {
83+
HttpHeaders headers = new HttpHeaders();
84+
headers.set("Key", "value");
85+
86+
assertThatThrownBy(() -> headers.populateLowerCaseHeaders(new String[1], null))
87+
.isInstanceOf(IllegalArgumentException.class)
88+
.hasMessageContaining("values");
89+
}
90+
91+
@Test(groups = "unit")
92+
public void populateLowerCaseHeadersRejectsTooSmallArrays() {
93+
HttpHeaders headers = new HttpHeaders();
94+
headers.set("A", "1");
95+
headers.set("B", "2");
96+
97+
assertThatThrownBy(() -> headers.populateLowerCaseHeaders(new String[1], new String[2]))
98+
.isInstanceOf(IllegalArgumentException.class)
99+
.hasMessageContaining("names");
100+
101+
assertThatThrownBy(() -> headers.populateLowerCaseHeaders(new String[2], new String[1]))
102+
.isInstanceOf(IllegalArgumentException.class)
103+
.hasMessageContaining("values");
104+
}
29105
}

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -191,7 +191,9 @@ private RxDocumentServiceRequest(DiagnosticsClientContext clientContext,
191191
this.forceNameCacheRefresh = false;
192192
this.resourceType = resourceType;
193193
this.contentAsByteArray = toByteArray(byteBuffer);
194-
this.headers = headers != null ? headers : new HashMap<>();
194+
// Pre-size to 32 (threshold 24 at 0.75 load factor) to accommodate typical request
195+
// headers (auth, content-type, consistency, session-token, partition-key, etc.) without resize.
196+
this.headers = headers != null ? headers : new HashMap<>(32);
195197
this.activityId = UUIDs.nonBlockingRandomUUID();
196198
this.isFeed = false;
197199
this.isNameBased = isNameBased;
@@ -225,7 +227,9 @@ private RxDocumentServiceRequest(DiagnosticsClientContext clientContext,
225227
this.operationType = operationType;
226228
this.resourceType = resourceType;
227229
this.requestContext.sessionToken = null;
228-
this.headers = headers != null ? headers : new HashMap<>();
230+
// Pre-size to 32 (threshold 24 at 0.75 load factor) to accommodate typical request
231+
// headers (auth, content-type, consistency, session-token, partition-key, etc.) without resize.
232+
this.headers = headers != null ? headers : new HashMap<>(32);
229233
this.activityId = UUIDs.nonBlockingRandomUUID();
230234
this.isFeed = false;
231235

‎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
@@ -242,7 +242,7 @@ public StoreResponse unwrapToStoreResponse(
242242
return new StoreResponse(
243243
endpoint,
244244
statusCode,
245-
HttpUtils.unescape(headers.toLowerCaseMap()),
245+
headers,
246246
new ByteBufInputStream(retainedContent, true),
247247
size);
248248
} else {
@@ -252,7 +252,7 @@ public StoreResponse unwrapToStoreResponse(
252252
return new StoreResponse(
253253
endpoint,
254254
statusCode,
255-
HttpUtils.unescape(headers.toLowerCaseMap()),
255+
headers,
256256
null,
257257
0);
258258
}
@@ -347,7 +347,7 @@ private Mono<RxDocumentServiceResponse> performRequestInternalCore(RxDocumentSer
347347
}
348348

349349
private HttpHeaders getHttpRequestHeaders(Map<String, String> headers) {
350-
HttpHeaders httpHeaders = new HttpHeaders(this.defaultHeaders.size());
350+
HttpHeaders httpHeaders = new HttpHeaders(HttpUtils.mapCapacityForSize(this.defaultHeaders.size() + headers.size()));
351351
// Add default headers.
352352
for (Entry<String, String> entry : this.defaultHeaders.entrySet()) {
353353
if (!headers.containsKey(entry.getKey())) {

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

Lines changed: 19 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -7,20 +7,16 @@
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;
1314

1415
import java.io.UnsupportedEncodingException;
1516
import java.net.URLDecoder;
1617
import java.net.URLEncoder;
17-
import java.util.AbstractMap;
18-
import java.util.ArrayList;
1918
import java.util.HashMap;
20-
import java.util.List;
2119
import java.util.Map;
22-
import java.util.Map.Entry;
23-
import java.util.Set;
2420
import java.util.regex.Pattern;
2521

2622
public class HttpUtils {
@@ -29,6 +25,18 @@ public class HttpUtils {
2925

3026
private static final Pattern PLUS_SYMBOL_ESCAPE_PATTERN = Pattern.compile(UrlEncodingInfo.PLUS_SYMBOL_ESCAPED);
3127

28+
/**
29+
* Returns the initial capacity for a HashMap that will hold {@code expectedSize} entries
30+
* without resizing, accounting for the default load factor of 0.75.
31+
*/
32+
public static int mapCapacityForSize(int expectedSize) {
33+
if (expectedSize <= 0) {
34+
return 1;
35+
}
36+
long capacity = (long) expectedSize * 4 / 3 + 1;
37+
return capacity >= Integer.MAX_VALUE ? Integer.MAX_VALUE : (int) capacity;
38+
}
39+
3240
public static String urlEncode(String url) {
3341
try {
3442
return PLUS_SYMBOL_ESCAPE_PATTERN.matcher(URLEncoder.encode(url, UrlEncodingInfo.UTF_8))
@@ -51,14 +59,14 @@ public static String urlDecode(String url) {
5159

5260
public static Map<String, String> asMap(HttpHeaders headers) {
5361
if (headers == null) {
54-
return new HashMap<>();
62+
return new HashMap<>(4);
5563
}
56-
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()));
64+
HashMap<String, String> map = new HashMap<>(mapCapacityForSize(headers.size()));
65+
for (HttpHeader header : headers) {
66+
if (header.name().equals(HttpConstants.HttpHeaders.OWNER_FULL_NAME)) {
67+
map.put(header.name(), HttpUtils.urlDecode(header.value()));
6068
} else {
61-
map.put(entry.getKey(), entry.getValue());
69+
map.put(header.name(), header.value());
6270
}
6371
}
6472
return map;
@@ -78,24 +86,4 @@ public static String getDateHeader(Map<String, String> headerValues) {
7886

7987
return date != null ? date : StringUtils.EMPTY;
8088
}
81-
82-
public static List<Entry<String, String>> unescape(Set<Entry<String, String>> headers) {
83-
List<Entry<String, String>> result = new ArrayList<>(headers.size());
84-
for (Entry<String, String> entry : headers) {
85-
if (entry.getKey().equals(HttpConstants.HttpHeaders.OWNER_FULL_NAME)) {
86-
String unescapedUrl = HttpUtils.urlDecode(entry.getValue());
87-
entry = new AbstractMap.SimpleEntry<>(entry.getKey(), unescapedUrl);
88-
}
89-
result.add(entry);
90-
}
91-
return result;
92-
}
93-
94-
public static Map<String, String> unescape(Map<String, String> headers) {
95-
if (headers != null) {
96-
headers.computeIfPresent(HttpConstants.HttpHeaders.OWNER_FULL_NAME,
97-
(ownerKey, ownerValue) -> HttpUtils.urlDecode(ownerValue));
98-
}
99-
return headers;
100-
}
10189
}

0 commit comments

Comments
 (0)