#16612 optimize method RemoteWebDriver.downloadFile()

Instead of reading the whole file to a byte array, just save given InputStream directly to the file.

Now it can download large files (I tried 4GB) while consuming very low memory.

Author: asolntsev

java/src/org/openqa/selenium/grid/node/local/LocalNode.java

@@ -54,6 +54,7 @@
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.nio.file.attribute.BasicFileAttributes;
+import java.nio.file.attribute.FileTime;
 import java.time.Clock;
 import java.time.Duration;
 import java.time.Instant;
@@ -844,12 +845,15 @@ private HttpResponse getDownloadedFile(File downloadsDirectory, String fileName)
     BasicFileAttributes attributes = readAttributes(file.toPath(), BasicFileAttributes.class);
     return new HttpResponse()
         .setHeader("Content-Type", MediaType.OCTET_STREAM.toString())
-        .setHeader(
-            "Last-Modified",
-            HTTP_DATE_FORMAT.format(attributes.lastModifiedTime().toInstant().atZone(UTC)))
+        .setHeader("Content-Length", String.valueOf(attributes.size()))
+        .setHeader("Last-Modified", lastModifiedHeader(attributes.lastModifiedTime()))
         .setContent(Contents.file(file));
   }
 
+  private String lastModifiedHeader(FileTime fileTime) {
+    return HTTP_DATE_FORMAT.format(fileTime.toInstant().atZone(UTC));
+  }
+
   private File findDownloadedFile(File downloadsDirectory, String filename)
       throws WebDriverException {
     List<File> matchingFiles =

java/src/org/openqa/selenium/remote/http/Contents.java

@@ -89,6 +89,10 @@ public static Supplier file(final File file) {
     return new FileContentSupplier(file);
   }
 
+  public static Supplier fromStream(InputStream stream, long length) {
+    return new InputStreamContentSupplier(stream, length);
+  }
+
   public static Supplier bytes(byte[] bytes) {
     return new BytesContentSupplier(bytes);
   }

java/src/org/openqa/selenium/remote/http/InputStreamContentSupplier.java

@@ -0,0 +1,71 @@
+// Licensed to the Software Freedom Conservancy (SFC) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The SFC licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.openqa.selenium.remote.http;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.charset.Charset;
+import org.openqa.selenium.internal.Require;
+
+class InputStreamContentSupplier implements Contents.Supplier {
+
+  private static final int MAX_TEXT_RESPONSE_SIZE = 256 * 1024 * 1024;
+  private final InputStream stream;
+  private final long length;
+
+  InputStreamContentSupplier(InputStream stream, long length) {
+    this.stream = Require.nonNull("InputStream", stream);
+    this.length = length;
+  }
+
+  @Override
+  public InputStream get() {
+    return stream;
+  }
+
+  @Override
+  public long length() {
+    return length;
+  }
+
+  public void close() {
+    try {
+      stream.close();
+    } catch (IOException ignore) {
+    }
+  }
+
+  @Override
+  public String toString() {
+    return String.format("Contents.fromStream(%s bytes)", length);
+  }
+
+  @Override
+  public String contentAsString(Charset charset) {
+    if (length > MAX_TEXT_RESPONSE_SIZE) {
+      throw new UnsupportedOperationException("Cannot print out too large stream content");
+    }
+    try {
+      return new String(stream.readAllBytes(), UTF_8);
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+}

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

@@ -20,6 +20,7 @@
 import com.google.auto.service.AutoService;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.io.InputStream;
 import java.io.UncheckedIOException;
 import java.net.Authenticator;
 import java.net.PasswordAuthentication;
@@ -440,8 +441,7 @@ private HttpResponse execute0(HttpRequest req) throws UncheckedIOException {
 
     LOG.log(Level.FINE, "Executing request: {0}", req);
     long start = System.currentTimeMillis();
-
-    BodyHandler<byte[]> byteHandler = BodyHandlers.ofByteArray();
+    BodyHandler<InputStream> byteHandler = BodyHandlers.ofInputStream();
     try {
       HttpMethod method = req.getMethod();
       URI rawUri = messages.getRawUri(req);
@@ -456,7 +456,7 @@ private HttpResponse execute0(HttpRequest req) throws UncheckedIOException {
         }
 
         java.net.http.HttpRequest request = messages.createRequest(req, method, rawUri);
-        java.net.http.HttpResponse<byte[]> response = client.send(request, byteHandler);
+        java.net.http.HttpResponse<InputStream> response = client.send(request, byteHandler);
 
         switch (response.statusCode()) {
           case 303:

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

@@ -19,6 +19,7 @@
 
 import static java.nio.charset.StandardCharsets.UTF_8;
 
+import java.io.InputStream;
 import java.net.URI;
 import java.net.URLEncoder;
 import java.net.http.HttpRequest.BodyPublisher;
@@ -152,7 +153,7 @@ public URI getRawUri(HttpRequest req) {
     return URI.create(rawUrl);
   }
 
-  public HttpResponse createResponse(java.net.http.HttpResponse<byte[]> response) {
+  public HttpResponse createResponse(java.net.http.HttpResponse<InputStream> response) {
     HttpResponse res = new HttpResponse();
     res.setStatus(response.statusCode());
     response
@@ -163,10 +164,8 @@ public HttpResponse createResponse(java.net.http.HttpResponse<byte[]> response)
                 values.stream()
                     .filter(Objects::nonNull)
                     .forEach(value -> res.addHeader(name, value)));
-    byte[] responseBody = response.body();
-    if (responseBody != null) {
-      res.setContent(Contents.bytes(responseBody));
-    }
+    long length = response.headers().firstValueAsLong("Content-Length").orElse(-1);
+    res.setContent(Contents.fromStream(response.body(), length));
 
     return res;
   }

java/test/org/openqa/selenium/grid/router/ReverseProxyEndToEndTest.java

@@ -125,7 +125,6 @@ private static void waitUntilReady(Server<?> server, Duration duration) {
           .until(
               c -> {
                 HttpResponse response = c.execute(new HttpRequest(GET, "/status"));
-                System.out.println(Contents.string(response));
                 Map<String, Object> status = Values.get(response, MAP_TYPE);
                 return Boolean.TRUE.equals(
                     status != null && Boolean.parseBoolean(status.get("ready").toString()));