mirror of
https://github.com/google/nomulus
synced 2026-07-29 11:32:48 +00:00
Fix OOM in UploadBsaUnavailableDomains action (#2817)
* Fix OOM in UploadBsaUnavailableDomains action The action was using string concatenation to generate the upload content. This causes an OOM when string length exceeds 25MB on our current VM. This PR witches to streaming upload. Also added an HTTP upload test. * Fix OOM in UploadBsaUnavailableDomains action The action was using string concatenation to generate the upload content. This causes an OOM when string length exceeds 25MB on our current VM. This PR witches to streaming upload. Also added an HTTP upload test.
This commit is contained in:
Weimin Yu
@@ -20,13 +20,24 @@ import static google.registry.testing.DatabaseHelper.persistActiveDomain;
|
||||
import static google.registry.testing.DatabaseHelper.persistDeletedDomain;
|
||||
import static google.registry.testing.DatabaseHelper.persistReservedList;
|
||||
import static google.registry.testing.DatabaseHelper.persistResource;
|
||||
import static google.registry.testing.LogsSubject.assertAboutLogs;
|
||||
import static google.registry.util.DateTimeUtils.START_OF_TIME;
|
||||
import static google.registry.util.NetworkUtils.pickUnusedPort;
|
||||
import static java.nio.charset.StandardCharsets.UTF_8;
|
||||
import static java.util.concurrent.Executors.newSingleThreadExecutor;
|
||||
import static org.mockito.Mockito.times;
|
||||
import static org.mockito.Mockito.verify;
|
||||
|
||||
import com.google.cloud.storage.BlobId;
|
||||
import com.google.cloud.storage.contrib.nio.testing.LocalStorageHelper;
|
||||
import com.google.common.collect.ImmutableList;
|
||||
import com.google.common.collect.ImmutableMap;
|
||||
import com.google.common.flogger.FluentLogger;
|
||||
import com.google.common.hash.Hashing;
|
||||
import com.google.common.io.ByteStreams;
|
||||
import com.google.common.net.HostAndPort;
|
||||
import com.google.common.testing.TestLogHandler;
|
||||
import com.google.gson.Gson;
|
||||
import google.registry.bsa.api.BsaCredential;
|
||||
import google.registry.gcs.GcsUtils;
|
||||
import google.registry.model.tld.Tld;
|
||||
@@ -35,9 +46,25 @@ import google.registry.model.tld.label.ReservedList;
|
||||
import google.registry.persistence.transaction.JpaTestExtensions;
|
||||
import google.registry.persistence.transaction.JpaTestExtensions.JpaIntegrationTestExtension;
|
||||
import google.registry.request.UrlConnectionService;
|
||||
import google.registry.server.Route;
|
||||
import google.registry.server.TestServer;
|
||||
import google.registry.testing.FakeClock;
|
||||
import google.registry.testing.FakeResponse;
|
||||
import jakarta.servlet.ServletException;
|
||||
import jakarta.servlet.annotation.MultipartConfig;
|
||||
import jakarta.servlet.http.HttpServlet;
|
||||
import jakarta.servlet.http.HttpServletRequest;
|
||||
import jakarta.servlet.http.HttpServletResponse;
|
||||
import jakarta.servlet.http.Part;
|
||||
import java.io.IOException;
|
||||
import java.io.InputStream;
|
||||
import java.io.PrintWriter;
|
||||
import java.net.InetAddress;
|
||||
import java.util.Map;
|
||||
import java.util.Optional;
|
||||
import java.util.logging.Level;
|
||||
import java.util.logging.Logger;
|
||||
import java.util.zip.GZIPInputStream;
|
||||
import org.joda.time.DateTime;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
@@ -102,13 +129,112 @@ public class UploadBsaUnavailableDomainsActionTest {
|
||||
BlobId existingFile =
|
||||
BlobId.of(BUCKET, String.format("unavailable_domains_%s.txt", clock.nowUtc()));
|
||||
String blockList = new String(gcsUtils.readBytesFrom(existingFile), UTF_8);
|
||||
assertThat(blockList).isEqualTo("ace.tld\nflagrant.tld\nfoobar.tld\njimmy.tld\ntine.tld");
|
||||
assertThat(blockList).isEqualTo("ace.tld\nflagrant.tld\nfoobar.tld\njimmy.tld\ntine.tld\n");
|
||||
assertThat(blockList).doesNotContain("not-blocked.tld");
|
||||
|
||||
// This test currently fails in the upload-to-bsa step.
|
||||
verify(emailSender, times(1))
|
||||
.sendNotification("BSA daily upload completed with errors", "Please see logs for details.");
|
||||
}
|
||||
|
||||
// TODO(mcilwain): Add test of BSA API upload as well.
|
||||
@Test
|
||||
void uploadToBsaTest() throws Exception {
|
||||
TestLogHandler logHandler = new TestLogHandler();
|
||||
Logger loggerToIntercept =
|
||||
Logger.getLogger(UploadBsaUnavailableDomainsAction.class.getCanonicalName());
|
||||
loggerToIntercept.addHandler(logHandler);
|
||||
|
||||
persistActiveDomain("foobar.tld");
|
||||
persistActiveDomain("ace.tld");
|
||||
persistDeletedDomain("not-blocked.tld", clock.nowUtc().minusDays(1));
|
||||
|
||||
var testServer = startTestServer();
|
||||
action.apiUrl = testServer.getUrl("/upload").toURI().toString();
|
||||
try {
|
||||
action.run();
|
||||
} finally {
|
||||
testServer.stop();
|
||||
}
|
||||
String dataSent = "ace.tld\nflagrant.tld\nfoobar.tld\njimmy.tld\ntine.tld\n";
|
||||
String checkSum = Hashing.sha512().hashString(dataSent, UTF_8).toString();
|
||||
String expectedResponse =
|
||||
"Received response with code 200 from server: "
|
||||
+ String.format("Checksum: [%s]\n%s\n", checkSum, dataSent);
|
||||
assertAboutLogs().that(logHandler).hasLogAtLevelWithMessage(Level.INFO, expectedResponse);
|
||||
verify(emailSender, times(1)).sendNotification("BSA daily upload completed successfully", "");
|
||||
}
|
||||
|
||||
private TestServer startTestServer() throws Exception {
|
||||
TestServer testServer =
|
||||
new TestServer(
|
||||
HostAndPort.fromParts(InetAddress.getLocalHost().getHostAddress(), pickUnusedPort()),
|
||||
ImmutableMap.of(),
|
||||
ImmutableList.of(Route.route("/upload", Servelet.class)));
|
||||
testServer.start();
|
||||
newSingleThreadExecutor()
|
||||
.execute(
|
||||
() -> {
|
||||
try {
|
||||
while (true) {
|
||||
testServer.process();
|
||||
}
|
||||
} catch (InterruptedException e) {
|
||||
// Expected
|
||||
}
|
||||
});
|
||||
return testServer;
|
||||
}
|
||||
|
||||
@MultipartConfig(
|
||||
location = "", // Directory for storing uploaded files. Use default when blank
|
||||
maxFileSize = 10485760L, // 10MB
|
||||
maxRequestSize = 20971520L, // 20MB
|
||||
fileSizeThreshold = 1048576 // Save in memory if file size < 1MB
|
||||
)
|
||||
public static class Servelet extends HttpServlet {
|
||||
private static final FluentLogger logger = FluentLogger.forEnclosingClass();
|
||||
|
||||
@Override
|
||||
protected void doPost(HttpServletRequest req, HttpServletResponse resp)
|
||||
throws ServletException, IOException {
|
||||
String checkSum = null;
|
||||
String content = null;
|
||||
try {
|
||||
for (Part part : req.getParts()) {
|
||||
switch (part.getName()) {
|
||||
case "zone" -> checkSum = readChecksum(part);
|
||||
case "file" -> content = readGzipped(part);
|
||||
}
|
||||
}
|
||||
} catch (Exception e) {
|
||||
logger.atInfo().withCause(e).log("");
|
||||
}
|
||||
int status = checkSum == null || content == null ? 400 : 200;
|
||||
resp.setStatus(status);
|
||||
resp.setContentType("text/plain");
|
||||
try (PrintWriter writer = resp.getWriter()) {
|
||||
writer.printf("Checksum: [%s]\n%s\n", checkSum, content);
|
||||
}
|
||||
}
|
||||
|
||||
private String readChecksum(Part part) {
|
||||
try (InputStream is = part.getInputStream()) {
|
||||
return new Gson()
|
||||
.fromJson(new String(ByteStreams.toByteArray(is), UTF_8), Map.class)
|
||||
.getOrDefault("checkSum", "Not found")
|
||||
.toString();
|
||||
} catch (IOException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
}
|
||||
|
||||
private String readGzipped(Part part) {
|
||||
try (InputStream is = part.getInputStream();
|
||||
GZIPInputStream gis = new GZIPInputStream(is)) {
|
||||
return new String(ByteStreams.toByteArray(gis), UTF_8);
|
||||
} catch (IOException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -26,10 +26,16 @@ import com.google.common.net.HostAndPort;
|
||||
import com.google.common.util.concurrent.SimpleTimeLimiter;
|
||||
import google.registry.util.RegistryEnvironment;
|
||||
import google.registry.util.UrlChecker;
|
||||
import jakarta.servlet.MultipartConfigElement;
|
||||
import jakarta.servlet.annotation.MultipartConfig;
|
||||
import jakarta.servlet.http.HttpServlet;
|
||||
import java.io.IOException;
|
||||
import java.net.MalformedURLException;
|
||||
import java.net.URL;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.BlockingQueue;
|
||||
import java.util.concurrent.FutureTask;
|
||||
@@ -49,6 +55,8 @@ import org.eclipse.jetty.server.ServerConnector;
|
||||
* {@link #stop()} methods. However, a {@link #process()} method was added, which is used to process
|
||||
* requests made to servlets (not static files) in the calling thread.
|
||||
*
|
||||
* <p>A servlet that expects multi-part requests should be annotated with {@link MultipartConfig}.
|
||||
*
|
||||
* <p><b>Note:</b> This server is intended for development purposes. For the love all that is good,
|
||||
* do not make this public-facing.
|
||||
*
|
||||
@@ -70,6 +78,7 @@ public final class TestServer {
|
||||
private final HostAndPort urlAddress;
|
||||
private final Server server = new Server();
|
||||
private final BlockingQueue<FutureTask<Void>> requestQueue = new LinkedBlockingDeque<>();
|
||||
private List<Path> multiPartTmpDirs = new ArrayList<>();
|
||||
|
||||
/**
|
||||
* Creates a new instance, but does not begin serving.
|
||||
@@ -134,6 +143,13 @@ public final class TestServer {
|
||||
},
|
||||
SHUTDOWN_TIMEOUT_MS,
|
||||
TimeUnit.MILLISECONDS);
|
||||
for (var dir : multiPartTmpDirs) {
|
||||
try {
|
||||
Files.delete(dir);
|
||||
} catch (Exception e) {
|
||||
// Ignore
|
||||
}
|
||||
}
|
||||
} catch (Exception e) {
|
||||
throwIfUnchecked(e);
|
||||
throw new RuntimeException(e);
|
||||
@@ -161,7 +177,27 @@ public final class TestServer {
|
||||
StaticResourceServlet.configureServletHolder(holder, runfile.getKey(), runfile.getValue());
|
||||
}
|
||||
for (Route route : routes) {
|
||||
context.addServlet(wrapServlet(route.servletClass()), route.path());
|
||||
holder = context.addServlet(wrapServlet(route.servletClass()), route.path());
|
||||
MultipartConfig multipartConfig = route.servletClass().getAnnotation(MultipartConfig.class);
|
||||
if (multipartConfig != null) {
|
||||
try {
|
||||
var location = multipartConfig.location();
|
||||
if (location == null || location.isBlank()) {
|
||||
Path tmpDir = Files.createTempDirectory("TestServer_");
|
||||
multiPartTmpDirs.add(tmpDir);
|
||||
location = tmpDir.toString();
|
||||
}
|
||||
MultipartConfigElement multipartConfigElement =
|
||||
new MultipartConfigElement(
|
||||
location,
|
||||
multipartConfig.maxFileSize(),
|
||||
multipartConfig.maxRequestSize(),
|
||||
multipartConfig.fileSizeThreshold());
|
||||
holder.getRegistration().setMultipartConfig(multipartConfigElement);
|
||||
} catch (IOException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
holder = context.addServlet(DefaultServlet.class, "/*");
|
||||
holder.setInitParameter("aliases", "1");
|
||||
|
||||
Reference in New Issue
Block a user