Prevent OOM caused by loading large files into memory.

Previously when building Android APK zip archives, the source files are read entirely into memory as byte arrays. So it cannot support files larger than 2 GB and will crash with OOM when allocating the arrays.

The change adds the support for an arbitrary InputStream of unknown size which will work for large files.

To minimize the impact to the build speed, when a file is small enough, the whole file is still read into memory to avoid repeated disk I/O. Large files are read in a streaming fashion to limit their memory consumption at the cost of more disk I/O.

Currently the boundary between small and large size is set to 512 MB empirically. The Java byte arrays can support up to 2 GB, but loading such large files into memory can also cause a high memory pressure for the JVM, so I didn't set it to 2 GB.

PiperOrigin-RevId: 819382752
Change-Id: Id0b96172e1674b4a3556a61538c957fa599c5c79

Author: A Googler

src/tools/java/com/google/devtools/build/android/AndroidResourceOutputs.java

@@ -26,6 +26,7 @@
 import java.io.Closeable;
 import java.io.File;
 import java.io.IOException;
+import java.io.InputStream;
 import java.nio.ByteBuffer;
 import java.nio.file.FileVisitResult;
 import java.nio.file.Files;
@@ -58,6 +59,9 @@ static class ZipBuilder implements Closeable {
     // see http://www.info-zip.org/FAQ.html#limits
     private static final long MINIMUM_TIMESTAMP_INCREMENT = 2000L;
 
+    // The maximum size of an entry that will be read entirely into memory.
+    private static final long MAX_IN_MEMORY_ENTRY_SIZE = 512 * 1024 * 1024;  // 512 MB.
+
     /**
      * Normalized timestamp for zip entries We use the system's default timezone and locale and
      * additionally avoid using the DOS epoch to ensure Java's zip implementation does not add the
@@ -69,6 +73,13 @@ static class ZipBuilder implements Closeable {
     private static final long DEFAULT_TIMESTAMP =
         new GregorianCalendar(1980, Calendar.FEBRUARY, 01, 0, 0).getTimeInMillis();
 
+
+    /** Supplier for input streams. */
+    protected interface InputStreamSupplier {
+      /** Returns a new input stream for reading the content. */
+      InputStream get() throws IOException;
+    }
+
     private final ZipOutputStream zip;
 
     private ZipBuilder(ZipOutputStream zip) {
@@ -131,6 +142,80 @@ protected void addEntry(
       zip.closeEntry();
     }
 
+    /**
+     * Adds a new entry to the zip archive. Use this method when the content has an unknown size and
+     * may be too large to fit in memory.
+     *
+     * @param rawName The name of the entry, which may include path separators.
+     * @param contentSupplier A supplier for creating input streams to read the content. The
+     *     supplier may be called for more than once.
+     * @param storageMethod The storage method to use for the entry.
+     * @param comment The comment to add to the entry.
+     * @throws IOException if an I/O error occurs.
+     */
+    protected void addEntry(
+        String rawName,
+        InputStreamSupplier contentSupplier,
+        int storageMethod,
+        @Nullable String comment)
+        throws IOException {
+      // Fix the path for windows.
+      String relativeName = rawName.replace('\\', '/');
+      // Make sure the zip entry is not absolute.
+      Preconditions.checkArgument(
+          !relativeName.startsWith("/"), "Cannot add absolute resources %s", relativeName);
+      ZipEntry entry = new ZipEntry(relativeName);
+      entry.setMethod(storageMethod);
+      entry.setTime(normalizeTime(relativeName));
+      if (!Strings.isNullOrEmpty(comment)) {
+        entry.setComment(comment);
+      }
+
+      long size = 0;
+      CRC32 crc32 = new CRC32();
+      try (InputStream content = contentSupplier.get()) {
+        byte[] buf = new byte[8192];
+        int read;
+        while ((read = content.read(buf)) != -1) {
+          crc32.update(buf, 0, read);
+          size += read;
+        }
+      }
+      entry.setSize(size);
+      entry.setCrc(crc32.getValue());
+      zip.putNextEntry(entry);
+
+      // InputStream doesn't support reset(), so to safely read the stream twice, we need to create
+      // the stream twice.
+      try (InputStream content = contentSupplier.get()) {
+        content.transferTo(zip);
+      }
+      zip.closeEntry();
+    }
+
+    /**
+     * Adds a new entry to the zip archive.
+     *
+     * @param rawName The name of the entry, which may include path separators.
+     * @param sourcePath The path to the source file.
+     * @param storageMethod The storage method to use for the entry.
+     * @throws IOException if an I/O error occurs.
+     */
+    protected void addEntry(String rawName, Path sourcePath, int storageMethod, String comment)
+        throws IOException {
+      // If the file is known to be small, it's more efficient to read it entirely into memory.
+      // Otherwise, do streaming processing to avoid using too much memory.
+      if (Files.isRegularFile(sourcePath) && Files.size(sourcePath) <= MAX_IN_MEMORY_ENTRY_SIZE) {
+        addEntry(rawName, Files.readAllBytes(sourcePath), storageMethod, comment);
+      } else {
+        addEntry(rawName, () -> Files.newInputStream(sourcePath), storageMethod, comment);
+      }
+    }
+
+    protected void addEntry(String rawName, Path sourcePath, int storageMethod) throws IOException {
+      addEntry(rawName, sourcePath, storageMethod, /* comment= */ null);
+    }
+
     @Override
     public void close() throws IOException {
       zip.close();
@@ -310,6 +395,11 @@ protected void addEntry(String entry, byte[] content) throws IOException {
       zipBuilder.addEntry(entry, content, storageMethod);
     }
 
+    protected void addEntry(Path file) throws IOException {
+      Preconditions.checkArgument(file.startsWith(root), "%s does not start with %s", file, root);
+      zipBuilder.addEntry(directoryPrefix + root.relativize(file), file, storageMethod);
+    }
+
     public void setCompress(boolean compress) {
       storageMethod = compress ? ZipEntry.DEFLATED : ZipEntry.STORED;
     }
@@ -333,8 +423,7 @@ void writeEntries() throws IOException {
 
     protected void writeEntry(Path file) throws IOException {
       if (!Files.isDirectory(file)) {
-        byte[] content = Files.readAllBytes(file);
-        addEntry(file, content);
+        addEntry(file);
       }
     }
   }

src/tools/java/com/google/devtools/build/android/ResourcesZip.java

@@ -182,19 +182,19 @@ public void writeTo(Path output, boolean compress) throws IOException {
       }
       try {
         if (ids != null) {
-          zip.addEntry("ids.txt", Files.readAllBytes(ids), ZipEntry.STORED);
+          zip.addEntry("ids.txt", ids, ZipEntry.STORED);
         }
 
         if (proto != null && Files.exists(proto)) {
-          zip.addEntry("apk.pb", Files.readAllBytes(proto), ZipEntry.STORED);
+          zip.addEntry("apk.pb", proto, ZipEntry.STORED);
         }
 
         if (attributes != null && Files.exists(attributes)) {
-          zip.addEntry("tools.attributes.pb", Files.readAllBytes(attributes), ZipEntry.STORED);
+          zip.addEntry("tools.attributes.pb", attributes, ZipEntry.STORED);
         }
 
         if (packages != null && Files.exists(packages)) {
-          zip.addEntry("packages.txt", Files.readAllBytes(packages), ZipEntry.STORED);
+          zip.addEntry("packages.txt", packages, ZipEntry.STORED);
         }
 
       } catch (IOException e) {
@@ -295,7 +295,7 @@ ShrunkProtoApk writeReportTo(Path reportOut) throws IOException {
     @CanIgnoreReturnValue
     ShrunkProtoApk writeResourcesToZip(Path resourcesZip) throws IOException {
       try (final ZipBuilder zip = ZipBuilder.createFor(resourcesZip)) {
-        zip.addEntry("apk.pb", Files.readAllBytes(apk.asApkPath()), ZipEntry.STORED);
+        zip.addEntry("apk.pb", apk.asApkPath(), ZipEntry.STORED);
       }
       return this;
     }

src/tools/javatests/com/google/devtools/build/android/AndroidResourceOutputsTest.java

@@ -21,6 +21,7 @@
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.zip.CRC32;
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipInputStream;
 import java.util.zip.ZipOutputStream;
@@ -66,4 +67,36 @@ public void testZipEntryPaths() throws Exception {
     }
     assertThat(entries).containsExactly("some/prefix/data1.txt", "some/prefix/bar/data2.txt");
   }
+
+  @Test
+  public void zipBuilder_addEntry_contentFactory() throws Exception {
+    Path output = tmp.resolve("actual.zip");
+    Path fileToAdd = tmp.resolve("foo/data1.txt");
+    Files.createDirectories(fileToAdd.getParent());
+    byte[] content = "hello world".getBytes(Charset.defaultCharset());
+    Files.write(fileToAdd, content);
+
+    try (ZipOutputStream zout = new ZipOutputStream(Files.newOutputStream(output))) {
+      AndroidResourceOutputs.ZipBuilder zipBuilder = AndroidResourceOutputs.ZipBuilder.wrap(zout);
+      zipBuilder.addEntry(
+          "a/prefix/data1.txt",
+          () -> Files.newInputStream(fileToAdd),
+          ZipEntry.STORED,
+          /* comment= */ null);
+    }
+
+    try (ZipInputStream zin = new ZipInputStream(Files.newInputStream(output))) {
+      ZipEntry entry = zin.getNextEntry();
+      assertThat(entry).isNotNull();
+      assertThat(entry.getName()).isEqualTo("a/prefix/data1.txt");
+      assertThat(entry.getMethod()).isEqualTo(ZipEntry.STORED);
+      byte[] zippedContent = zin.readAllBytes();
+      assertThat(zippedContent).isEqualTo(content);
+      CRC32 crc32 = new CRC32();
+      crc32.update(content);
+      assertThat(entry.getCrc()).isEqualTo(crc32.getValue());
+      assertThat(entry.getSize()).isEqualTo(content.length);
+      assertThat(zin.getNextEntry()).isNull();
+    }
+  }
 }