Skip to content

Commit e0ade5c

Browse files
authored
[java] JDK Http client - avoid chunking without buffering to memory (#11198)
avoid chunking without buffering to memory Reduce the memory consumption especially while creating a new session.
1 parent f569d7f commit e0ade5c

File tree

2 files changed

+28
-5
lines changed

2 files changed

+28
-5
lines changed

java/src/org/openqa/selenium/remote/ProtocolHandshake.java

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -97,8 +97,7 @@ public Either<SessionNotCreatedException, Result> createSession(HttpHandler clie
9797
Writer writer = new OutputStreamWriter(counter, UTF_8)) {
9898
payload.writeTo(writer);
9999

100-
try (InputStream rawIn = os.asByteSource().openBufferedStream();
101-
BufferedInputStream contentStream = new BufferedInputStream(rawIn)) {
100+
try (InputStream contentStream = os.asByteSource().openBufferedStream()) {
102101
return createSession(client, contentStream, counter.getCount());
103102
}
104103
} finally {
@@ -113,6 +112,9 @@ private Either<SessionNotCreatedException, Result> createSession(HttpHandler cli
113112
HttpResponse response;
114113
long start = System.currentTimeMillis();
115114

115+
// Setting the CONTENT_LENGTH will allow a http client implementation not to read the data in
116+
// memory. Usually the payload is small and buffering it to memory is okay, except for a new
117+
// session e.g. with profiles.
116118
request.setHeader(CONTENT_LENGTH, String.valueOf(size));
117119
request.setHeader(CONTENT_TYPE, JSON_UTF_8);
118120
request.setContent(() -> newSessionBlob);

java/src/org/openqa/selenium/remote/http/jdk/JdkHttpMessages.java

Lines changed: 24 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import java.io.ByteArrayInputStream;
2727
import java.net.URI;
2828
import java.net.URLEncoder;
29+
import java.net.http.HttpRequest.BodyPublisher;
2930
import java.net.http.HttpRequest.BodyPublishers;
3031
import java.util.List;
3132
import java.util.Objects;
@@ -72,12 +73,11 @@ public java.net.http.HttpRequest createRequest(HttpRequest req) {
7273
break;
7374

7475
case POST:
75-
// Copy the content into a byte array to avoid reading the content inputstream multiple times.
76-
builder = builder.POST(BodyPublishers.ofByteArray(Contents.bytes(req.getContent())));
76+
builder = builder.POST(notChunkingBodyPublisher(req));
7777
break;
7878

7979
case PUT:
80-
builder = builder.PUT(BodyPublishers.ofByteArray(Contents.bytes(req.getContent())));
80+
builder = builder.PUT(notChunkingBodyPublisher(req));
8181
break;
8282

8383
default:
@@ -103,6 +103,27 @@ public java.net.http.HttpRequest createRequest(HttpRequest req) {
103103
return builder.build();
104104
}
105105

106+
/**
107+
* Some drivers do not support chunked transport, we ensure the http client is not using chunked
108+
* transport. This is done by using a BodyPublisher with a known size, in best case without
109+
* wasting memory by buffering the request.
110+
*
111+
* @return a BodyPublisher with a known size
112+
*/
113+
private BodyPublisher notChunkingBodyPublisher(HttpRequest req) {
114+
String length = req.getHeader("content-length");
115+
116+
if (length == null) {
117+
// read the data into a byte array to know the length
118+
return BodyPublishers.ofByteArray(Contents.bytes(req.getContent()));
119+
}
120+
121+
// we know the length of the request and use it
122+
BodyPublisher chunking = BodyPublishers.ofInputStream(req.getContent());
123+
124+
return BodyPublishers.fromPublisher(chunking, Long.parseLong(length));
125+
}
126+
106127
private String getRawUrl(URI baseUrl, String uri) {
107128
String rawUrl;
108129
if (uri.startsWith("ws://") || uri.startsWith("wss://") ||

0 commit comments

Comments
 (0)