Skip to content

Commit f5d0a60

Browse files
asolntsevasolntsev
committed
#16612 support downloading large files from Grid
I added a new Grid endpoint "/se/files/:name" which allows downloading the file directly, without encoding it to Base64 and adding to Json. This transformation kills the performance and causes OutOfMemory errors for large files (e.g. 256+ MB). NB! Be sure that `toString()` method of objects (HttpRequest, HttpResponse, Contents.Supplier) never returns too long string - it spam debug logs and can cause OOM during debugging.
1 parent a2b4269 commit f5d0a60

File tree

15 files changed

+243
-38
lines changed

15 files changed

+243
-38
lines changed

‎java/src/org/openqa/selenium/grid/data/Session.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,4 +144,9 @@ public boolean equals(Object that) {
144144
public int hashCode() {
145145
return Objects.hash(id, uri);
146146
}
147+
148+
@Override
149+
public String toString() {
150+
return String.format("Session(id:%s, url:%s, started at: %s)", id, uri, getStartTime());
151+
}
147152
}

‎java/src/org/openqa/selenium/grid/node/BUILD.bazel‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,5 +22,6 @@ java_library(
2222
"//java/src/org/openqa/selenium/remote",
2323
"//java/src/org/openqa/selenium/status",
2424
artifact("com.google.guava:guava"),
25+
artifact("org.jspecify:jspecify"),
2526
],
2627
)

‎java/src/org/openqa/selenium/grid/node/Node.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -171,6 +171,9 @@ protected Node(
171171
get("/session/{sessionId}/se/files")
172172
.to(params -> new DownloadFile(this, sessionIdFrom(params)))
173173
.with(spanDecorator("node.download_file")),
174+
get("/session/{sessionId}/se/files/{fileName}")
175+
.to(params -> new DownloadFile(this, sessionIdFrom(params)))
176+
.with(spanDecorator("node.download_file")),
174177
post("/session/{sessionId}/se/files")
175178
.to(params -> new DownloadFile(this, sessionIdFrom(params)))
176179
.with(spanDecorator("node.download_file")),

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

Lines changed: 73 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,18 @@
1919

2020
import static com.google.common.collect.ImmutableSet.toImmutableSet;
2121
import static java.nio.file.Files.readAttributes;
22+
import static java.time.ZoneOffset.UTC;
23+
import static java.time.format.DateTimeFormatter.RFC_1123_DATE_TIME;
24+
import static java.util.Arrays.asList;
25+
import static java.util.Locale.US;
26+
import static java.util.Objects.requireNonNullElseGet;
2227
import static org.openqa.selenium.HasDownloads.DownloadedFile;
2328
import static org.openqa.selenium.concurrent.ExecutorServices.shutdownGracefully;
2429
import static org.openqa.selenium.grid.data.Availability.DOWN;
2530
import static org.openqa.selenium.grid.data.Availability.DRAINING;
2631
import static org.openqa.selenium.grid.data.Availability.UP;
2732
import static org.openqa.selenium.grid.node.CapabilityResponseEncoder.getEncoder;
33+
import static org.openqa.selenium.net.Urls.urlDecode;
2834
import static org.openqa.selenium.remote.CapabilityType.ENABLE_DOWNLOADS;
2935
import static org.openqa.selenium.remote.HttpSessionId.getSessionId;
3036
import static org.openqa.selenium.remote.RemoteTags.CAPABILITIES;
@@ -39,6 +45,7 @@
3945
import com.github.benmanes.caffeine.cache.Ticker;
4046
import com.google.common.annotations.VisibleForTesting;
4147
import com.google.common.collect.ImmutableList;
48+
import com.google.common.net.MediaType;
4249
import java.io.Closeable;
4350
import java.io.File;
4451
import java.io.IOException;
@@ -50,6 +57,7 @@
5057
import java.time.Clock;
5158
import java.time.Duration;
5259
import java.time.Instant;
60+
import java.time.format.DateTimeFormatter;
5361
import java.util.Arrays;
5462
import java.util.HashMap;
5563
import java.util.List;
@@ -107,6 +115,7 @@
107115
import org.openqa.selenium.json.Json;
108116
import org.openqa.selenium.remote.Browser;
109117
import org.openqa.selenium.remote.SessionId;
118+
import org.openqa.selenium.remote.http.Contents;
110119
import org.openqa.selenium.remote.http.HttpMethod;
111120
import org.openqa.selenium.remote.http.HttpRequest;
112121
import org.openqa.selenium.remote.http.HttpResponse;
@@ -123,6 +132,7 @@ public class LocalNode extends Node implements Closeable {
123132

124133
private static final Json JSON = new Json();
125134
private static final Logger LOG = Logger.getLogger(LocalNode.class.getName());
135+
private static final DateTimeFormatter HTTP_DATE_FORMAT = RFC_1123_DATE_TIME.withLocale(US);
126136

127137
private final EventBus bus;
128138
private final URI externalUri;
@@ -742,9 +752,12 @@ public HttpResponse downloadFile(HttpRequest req, SessionId id) {
742752
Optional.ofNullable(tempFS.getBaseDir().listFiles()).orElse(new File[] {})[0];
743753

744754
try {
745-
if (req.getMethod().equals(HttpMethod.GET)) {
755+
if (req.getMethod().equals(HttpMethod.GET) && req.getUri().endsWith("/se/files")) {
746756
return listDownloadedFiles(downloadsDirectory);
747757
}
758+
if (req.getMethod().equals(HttpMethod.GET)) {
759+
return getDownloadedFile(downloadsDirectory, extractFileName(req));
760+
}
748761
if (req.getMethod().equals(HttpMethod.DELETE)) {
749762
return deleteDownloadedFile(downloadsDirectory);
750763
}
@@ -754,6 +767,19 @@ public HttpResponse downloadFile(HttpRequest req, SessionId id) {
754767
}
755768
}
756769

770+
private String extractFileName(HttpRequest req) {
771+
return extractFileName(req.getUri());
772+
}
773+
774+
String extractFileName(String uri) {
775+
String prefix = "/se/files/";
776+
int index = uri.lastIndexOf(prefix);
777+
if (index < 0) {
778+
throw new IllegalArgumentException("Unexpected URL for downloading a file: " + uri);
779+
}
780+
return urlDecode(uri.substring(index + prefix.length())).replace(' ', '+');
781+
}
782+
757783
/** User wants to list files that can be downloaded */
758784
private HttpResponse listDownloadedFiles(File downloadsDirectory) {
759785
File[] files = Optional.ofNullable(downloadsDirectory.listFiles()).orElse(new File[] {});
@@ -798,29 +824,60 @@ private HttpResponse getDownloadedFile(HttpRequest req, File downloadsDirectory)
798824
new WebDriverException(
799825
"Please specify file to download in payload as {\"name\":"
800826
+ " \"fileToDownload\"}"));
801-
File[] allFiles =
802-
Optional.ofNullable(downloadsDirectory.listFiles((dir, name) -> name.equals(filename)))
803-
.orElse(new File[] {});
804-
if (allFiles.length == 0) {
805-
throw new WebDriverException(
806-
String.format(
807-
"Cannot find file [%s] in directory %s.",
808-
filename, downloadsDirectory.getAbsolutePath()));
809-
}
810-
if (allFiles.length != 1) {
811-
throw new WebDriverException(
812-
String.format("Expected there to be only 1 file. There were: %s.", allFiles.length));
813-
}
814-
String content = Zip.zip(allFiles[0]);
827+
File file = findDownloadedFile(downloadsDirectory, filename);
828+
String content = Zip.zip(file);
815829
Map<String, Object> data =
816830
Map.of(
817831
"filename", filename,
818-
"file", getFileInfo(allFiles[0]),
832+
"file", getFileInfo(file),
819833
"contents", content);
820834
Map<String, Map<String, Object>> result = Map.of("value", data);
821835
return new HttpResponse().setContent(asJson(result));
822836
}
823837

838+
private HttpResponse getDownloadedFile(File downloadsDirectory, String fileName)
839+
throws IOException {
840+
if (fileName.isEmpty()) {
841+
throw new WebDriverException("Please specify file to download in URL");
842+
}
843+
File file = findDownloadedFile(downloadsDirectory, fileName);
844+
BasicFileAttributes attributes = readAttributes(file.toPath(), BasicFileAttributes.class);
845+
return new HttpResponse()
846+
.setHeader("Content-Type", MediaType.OCTET_STREAM.toString())
847+
.setHeader(
848+
"Last-Modified",
849+
HTTP_DATE_FORMAT.format(attributes.lastModifiedTime().toInstant().atZone(UTC)))
850+
.setContent(Contents.file(file));
851+
}
852+
853+
private File findDownloadedFile(File downloadsDirectory, String filename)
854+
throws WebDriverException {
855+
List<File> matchingFiles =
856+
asList(
857+
requireNonNullElseGet(
858+
downloadsDirectory.listFiles((dir, name) -> name.equals(filename)),
859+
() -> new File[0]));
860+
if (matchingFiles.isEmpty()) {
861+
List<File> files = downloadedFiles(downloadsDirectory);
862+
throw new WebDriverException(
863+
String.format(
864+
"Cannot find file [%s] in directory %s. Found %s files: %s.",
865+
filename, downloadsDirectory.getAbsolutePath(), files.size(), files));
866+
}
867+
if (matchingFiles.size() != 1) {
868+
throw new WebDriverException(
869+
String.format(
870+
"Expected there to be only 1 file. Found %s files: %s.",
871+
matchingFiles.size(), matchingFiles));
872+
}
873+
return matchingFiles.get(0);
874+
}
875+
876+
private static List<File> downloadedFiles(File downloadsDirectory) {
877+
File[] files = requireNonNullElseGet(downloadsDirectory.listFiles(), () -> new File[0]);
878+
return asList(files);
879+
}
880+
824881
private HttpResponse deleteDownloadedFile(File downloadsDirectory) {
825882
File[] files = Optional.ofNullable(downloadsDirectory.listFiles()).orElse(new File[] {});
826883
for (File file : files) {

‎java/src/org/openqa/selenium/net/Urls.java‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,14 +17,16 @@
1717

1818
package org.openqa.selenium.net;
1919

20+
import static java.nio.charset.StandardCharsets.UTF_8;
21+
2022
import java.io.IOException;
2123
import java.io.UncheckedIOException;
2224
import java.net.MalformedURLException;
2325
import java.net.URI;
2426
import java.net.URISyntaxException;
2527
import java.net.URL;
28+
import java.net.URLDecoder;
2629
import java.net.URLEncoder;
27-
import java.nio.charset.StandardCharsets;
2830
import org.jspecify.annotations.NullMarked;
2931
import org.openqa.selenium.internal.Require;
3032

@@ -43,7 +45,11 @@ private Urls() {
4345
* @see URLEncoder#encode(java.lang.String, java.lang.String)
4446
*/
4547
public static String urlEncode(String value) {
46-
return URLEncoder.encode(value, StandardCharsets.UTF_8);
48+
return URLEncoder.encode(value, UTF_8);
49+
}
50+
51+
public static String urlDecode(String encodedValue) {
52+
return URLDecoder.decode(encodedValue, UTF_8);
4753
}
4854

4955
public static URL fromUri(URI uri) {

‎java/src/org/openqa/selenium/netty/server/RequestConverter.java‎

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -37,9 +37,12 @@
3737
import io.netty.handler.codec.http.LastHttpContent;
3838
import io.netty.handler.codec.http.QueryStringDecoder;
3939
import io.netty.util.ReferenceCountUtil;
40+
41+
import java.io.ByteArrayOutputStream;
4042
import java.io.IOException;
4143
import java.io.InputStream;
4244
import java.io.UncheckedIOException;
45+
import java.nio.charset.Charset;
4346
import java.util.Arrays;
4447
import java.util.List;
4548
import java.util.logging.Logger;
@@ -56,7 +59,7 @@ class RequestConverter extends SimpleChannelInboundHandler<HttpObject> {
5659
private static final List<io.netty.handler.codec.http.HttpMethod> SUPPORTED_METHODS =
5760
Arrays.asList(DELETE, GET, POST, OPTIONS);
5861
private volatile FileBackedOutputStream buffer;
59-
private volatile int length;
62+
private volatile long length;
6063
private volatile HttpRequest request;
6164

6265
@Override
@@ -119,7 +122,7 @@ protected void channelRead0(ChannelHandlerContext ctx, HttpObject msg) throws Ex
119122

120123
if (buffer != null) {
121124
ByteSource source = buffer.asByteSource();
122-
int len = length;
125+
long len = length;
123126

124127
request.setContent(
125128
new Contents.Supplier() {
@@ -133,14 +136,30 @@ public InputStream get() {
133136
}
134137

135138
@Override
136-
public int length() {
139+
public long length() {
137140
return len;
138141
}
139142

140143
@Override
141144
public void close() throws IOException {
142145
buffer.reset();
143146
}
147+
148+
@Override
149+
public String toString() {
150+
return String.format("Last content for %s (%s bytes)", request.toString(), len);
151+
}
152+
153+
@Override
154+
public String contentAsString(Charset charset) {
155+
ByteArrayOutputStream out = new ByteArrayOutputStream();
156+
try {
157+
source.copyTo(out);
158+
} catch (IOException e) {
159+
throw new RuntimeException(e);
160+
}
161+
return out.toString(charset);
162+
}
144163
});
145164
} else {
146165
request.setContent(Contents.empty());

‎java/src/org/openqa/selenium/remote/DriverCommand.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -154,6 +154,7 @@ public interface DriverCommand {
154154
String RESET_COOLDOWN = "resetCooldown";
155155
String GET_DOWNLOADABLE_FILES = "getDownloadableFiles";
156156
String DOWNLOAD_FILE = "downloadFile";
157+
String GET_DOWNLOADED_FILE = "getDownloadedFile";
157158
String DELETE_DOWNLOADABLE_FILES = "deleteDownloadableFiles";
158159

159160
static CommandPayload NEW_SESSION(Capabilities capabilities) {

‎java/src/org/openqa/selenium/remote/RemoteWebDriver.java‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,12 @@
2323
import static org.openqa.selenium.HasDownloads.DownloadedFile;
2424
import static org.openqa.selenium.remote.CapabilityType.PLATFORM_NAME;
2525

26+
import java.io.BufferedInputStream;
2627
import java.io.IOException;
28+
import java.io.InputStream;
2729
import java.net.MalformedURLException;
2830
import java.net.URL;
31+
import java.nio.file.Files;
2932
import java.nio.file.Path;
3033
import java.time.Duration;
3134
import java.util.ArrayList;
@@ -80,14 +83,14 @@
8083
import org.openqa.selenium.interactions.Sequence;
8184
import org.openqa.selenium.internal.Debug;
8285
import org.openqa.selenium.internal.Require;
83-
import org.openqa.selenium.io.Zip;
8486
import org.openqa.selenium.logging.LocalLogs;
8587
import org.openqa.selenium.logging.LoggingHandler;
8688
import org.openqa.selenium.logging.Logs;
8789
import org.openqa.selenium.logging.NeedsLocalLogs;
8890
import org.openqa.selenium.print.PrintOptions;
8991
import org.openqa.selenium.remote.http.ClientConfig;
9092
import org.openqa.selenium.remote.http.ConnectionFailedException;
93+
import org.openqa.selenium.remote.http.Contents;
9194
import org.openqa.selenium.remote.http.HttpClient;
9295
import org.openqa.selenium.remote.service.DriverCommandExecutor;
9396
import org.openqa.selenium.remote.tracing.TracedHttpClient;
@@ -727,9 +730,12 @@ public List<DownloadedFile> getDownloadedFiles() {
727730
public void downloadFile(String fileName, Path targetLocation) throws IOException {
728731
requireDownloadsEnabled(capabilities);
729732

730-
Response response = execute(DriverCommand.DOWNLOAD_FILE, Map.of("name", fileName));
731-
String contents = ((Map<String, String>) response.getValue()).get("contents");
732-
Zip.unzip(contents, targetLocation.toFile());
733+
Response response = execute(DriverCommand.GET_DOWNLOADED_FILE, Map.of("name", fileName));
734+
735+
Contents.Supplier content = (Contents.Supplier) response.getValue();
736+
try (InputStream fileContent = content.get()) {
737+
Files.copy(new BufferedInputStream(fileContent), targetLocation.resolve(fileName));
738+
}
733739
}
734740

735741
/**

‎java/src/org/openqa/selenium/remote/codec/AbstractHttpCommandCodec.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
import static org.openqa.selenium.remote.DriverCommand.GET_CREDENTIALS;
4747
import static org.openqa.selenium.remote.DriverCommand.GET_CURRENT_URL;
4848
import static org.openqa.selenium.remote.DriverCommand.GET_DOWNLOADABLE_FILES;
49+
import static org.openqa.selenium.remote.DriverCommand.GET_DOWNLOADED_FILE;
4950
import static org.openqa.selenium.remote.DriverCommand.GET_ELEMENT_RECT;
5051
import static org.openqa.selenium.remote.DriverCommand.GET_ELEMENT_TAG_NAME;
5152
import static org.openqa.selenium.remote.DriverCommand.GET_ELEMENT_TEXT;
@@ -199,6 +200,7 @@ public AbstractHttpCommandCodec() {
199200

200201
defineCommand(GET_DOWNLOADABLE_FILES, get(sessionId + "/se/files"));
201202
defineCommand(DOWNLOAD_FILE, post(sessionId + "/se/files"));
203+
defineCommand(GET_DOWNLOADED_FILE, get(sessionId + "/se/files/:name"));
202204
defineCommand(DELETE_DOWNLOADABLE_FILES, delete(sessionId + "/se/files"));
203205
}
204206

0 commit comments

Comments
 (0)