multi-thread DexFileMerger

RELNOTES: speedup of incremental dexing tools

PiperOrigin-RevId: 164926895

3 files changed, 158 insertions, 74 deletions

diff --git a/src/tools/android/java/com/google/devtools/build/android/dexer/DexFileAggregator.java b/src/tools/android/java/com/google/devtools/build/android/dexer/DexFileAggregator.java
index bd7d572fe2..f90c5f78bd 100644
--- a/src/tools/android/java/com/google/devtools/build/android/dexer/DexFileAggregator.java
+++ b/src/tools/android/java/com/google/devtools/build/android/dexer/DexFileAggregator.java
@@ -13,6 +13,7 @@
 // limitations under the License.
 package com.google.devtools.build.android.dexer;
 
+import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkState;
 
 import com.android.dex.Dex;
@@ -26,13 +27,19 @@ import com.android.dx.merge.CollisionPolicy;
 import com.android.dx.merge.DexMerger;
 import com.google.auto.value.AutoValue;
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableList;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.ListeningExecutorService;
 import java.io.Closeable;
 import java.io.IOException;
 import java.nio.BufferOverflowException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashSet;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
 import java.util.zip.ZipEntry;
 
 /**
@@ -59,23 +66,28 @@ class DexFileAggregator implements Closeable {
   private final int wasteThresholdPerDex;
   private final MultidexStrategy multidex;
   private final DxContext context;
-  private DexFileArchive dest;
+  private final ListeningExecutorService executor;
+  private final DexFileArchive dest;
+
   private int nextDexFileIndex = 0;
+  private ListenableFuture<Void> lastWriter = Futures.<Void>immediateFuture(null);
 
   public DexFileAggregator(
       DxContext context,
       DexFileArchive dest,
+      ListeningExecutorService executor,
       MultidexStrategy multidex,
       int maxNumberOfIdxPerDex,
       int wasteThresholdPerDex) {
     this.context = context;
     this.dest = dest;
+    this.executor = executor;
     this.multidex = multidex;
     this.maxNumberOfIdxPerDex = maxNumberOfIdxPerDex;
     this.wasteThresholdPerDex = wasteThresholdPerDex;
   }
 
-  public DexFileAggregator add(Dex dexFile) throws IOException {
+  public DexFileAggregator add(Dex dexFile) {
     if (multidex.isMultidexAllowed()) {
       // To determine whether currentShard is "full" we track unique field and method signatures,
       // which predicts precisely the number of field and method indices.
@@ -114,13 +126,20 @@ class DexFileAggregator implements Closeable {
       if (!currentShard.isEmpty()) {
         rotateDexFile();
       }
+      // Wait for last shard to be written before closing underlying archive
+      lastWriter.get();
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+    } catch (ExecutionException e) {
+      Throwables.throwIfInstanceOf(e.getCause(), IOException.class);
+      Throwables.throwIfUnchecked(e.getCause());
+      throw new AssertionError("Unexpected execution exception", e);
     } finally {
       dest.close();
-      dest = null;
     }
   }
 
-  public void flush() throws IOException {
+  public void flush() {
     checkState(multidex.isMultidexAllowed());
     if (!currentShard.isEmpty()) {
       rotateDexFile();
@@ -131,16 +150,24 @@ class DexFileAggregator implements Closeable {
     return nextDexFileIndex;
   }
 
-  private void rotateDexFile() throws IOException {
+  private void rotateDexFile() {
     writeMergedFile(currentShard.toArray(/* apparently faster than pre-sized array */ new Dex[0]));
     currentShard.clear();
     fieldsInCurrentShard.clear();
     methodsInCurrentShard.clear();
   }
 
-  private void writeMergedFile(Dex... dexes) throws IOException {
-    Dex merged = merge(dexes);
-    dest.addFile(nextArchiveEntry(), merged);
+  private void writeMergedFile(Dex... dexes) {
+    checkArgument(0 < dexes.length);
+    checkState(multidex.isMultidexAllowed() || nextDexFileIndex == 0);
+    String filename = getDexFileName(nextDexFileIndex++);
+    ListenableFuture<Dex> merged =
+        dexes.length == 1
+            ? Futures.immediateFuture(dexes[0])
+            : executor.submit(new RunDexMerger(dexes));
+    lastWriter =
+        Futures.whenAllSucceed(lastWriter, merged)
+            .call(new WriteFile(filename, merged, dest), executor);
   }
 
   private Dex merge(Dex... dexes) throws IOException {
@@ -155,6 +182,9 @@ class DexFileAggregator implements Closeable {
           dexMerger.setCompactWasteThreshold(wasteThresholdPerDex);
           return dexMerger.merge();
         } catch (BufferOverflowException e) {
+          if (dexes.length <= 2) {
+            throw e;
+          }
           // Bug in dx can cause this for ~1500 or more classes
           Dex[] left = Arrays.copyOf(dexes, dexes.length / 2);
           Dex[] right = Arrays.copyOfRange(dexes, left.length, dexes.length);
@@ -169,19 +199,16 @@ class DexFileAggregator implements Closeable {
     }
   }
 
-  private ZipEntry nextArchiveEntry() {
-    checkState(multidex.isMultidexAllowed() || nextDexFileIndex == 0);
-    ZipEntry result = new ZipEntry(getDexFileName(nextDexFileIndex++));
-    result.setTime(0L); // Use simple stable timestamps for deterministic output
-    return result;
-  }
-
   // More or less copied from from com.android.dx.command.dexer.Main
   @VisibleForTesting
   static String getDexFileName(int i) {
     return i == 0 ? DexFormat.DEX_IN_JAR_NAME : DEX_PREFIX + (i + 1) + DEX_EXTENSION;
   }
 
+  private static String typeName(Dex dex, int typeIndex) {
+    return dex.typeNames().get(typeIndex);
+  }
+
   @AutoValue
   abstract static class FieldDescriptor {
     static FieldDescriptor fromDex(Dex dex, int fieldIndex) {
@@ -220,7 +247,58 @@ class DexFileAggregator implements Closeable {
     abstract String returnType();
   }
 
-  private static String typeName(Dex dex, int typeIndex) {
-    return dex.typeNames().get(typeIndex);
+  private class RunDexMerger implements Callable<Dex> {
+
+    private final Dex[] dexes;
+
+    public RunDexMerger(Dex... dexes) {
+      checkArgument(dexes.length >= 2, "Only got %s dex files to merge", dexes.length);
+      this.dexes = dexes;
+    }
+
+    @Override
+    public Dex call() throws IOException {
+      try {
+        return merge(dexes);
+      } catch (Throwable t) {
+        // Print out exceptions so they don't get swallowed completely
+        t.printStackTrace();
+        Throwables.throwIfInstanceOf(t, IOException.class);
+        Throwables.throwIfUnchecked(t);
+        throw new AssertionError(t);  // shouldn't get here
+      }
+    }
+  }
+
+  private static class WriteFile implements Callable<Void> {
+
+    private final ListenableFuture<Dex> dex;
+    private final String filename;
+    private final DexFileArchive dest;
+
+    public WriteFile(String filename, ListenableFuture<Dex> dex, DexFileArchive dest) {
+      this.filename = filename;
+      this.dex = dex;
+      this.dest = dest;
+    }
+
+    @Override
+    public Void call() throws Exception {
+      try {
+        checkState(dex.isDone());
+        ZipEntry entry = new ZipEntry(filename);
+        entry.setTime(0L); // Use simple stable timestamps for deterministic output
+        dest.addFile(entry, dex.get());
+        return null;
+      } catch (Exception e) {
+        // Print out exceptions so they don't get swallowed completely
+        e.printStackTrace();
+        throw e;
+      } catch (Throwable t) {
+        t.printStackTrace();
+        Throwables.throwIfUnchecked(t);
+        throw new AssertionError(t);  // shouldn't get here
+      }
+    }
   }
 }
diff --git a/src/tools/android/java/com/google/devtools/build/android/dexer/DexFileArchive.java b/src/tools/android/java/com/google/devtools/build/android/dexer/DexFileArchive.java
index cfc06885bf..9a36dfc6c1 100644
--- a/src/tools/android/java/com/google/devtools/build/android/dexer/DexFileArchive.java
+++ b/src/tools/android/java/com/google/devtools/build/android/dexer/DexFileArchive.java
@@ -13,11 +13,12 @@
 // limitations under the License.
 package com.google.devtools.build.android.dexer;
 
+import static com.google.common.base.Preconditions.checkState;
+
 import com.android.dex.Dex;
-import com.google.common.io.ByteStreams;
 import java.io.Closeable;
 import java.io.IOException;
-import java.io.InputStream;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipOutputStream;
 
@@ -25,37 +26,37 @@ import java.util.zip.ZipOutputStream;
  * Wrapper around a {@link ZipOutputStream} to simplify writing archives with {@code .dex} files.
  * Adding files generally requires a {@link ZipEntry} in order to control timestamps.
  */
+// TODO(kmb): Remove this class and inline into DexFileAggregator
 class DexFileArchive implements Closeable {
 
   private final ZipOutputStream out;
 
-  public DexFileArchive(ZipOutputStream out) {
-    this.out = out;
-  }
-
   /**
-   * Copies the content of the given {@link InputStream} into an entry with the given details.
+   * Used to ensure writes from different threads are sequenced, which {@link DexFileAggregator}
+   * ensures by making the writer futures wait on each oter.
    */
-  public DexFileArchive copy(ZipEntry entry, InputStream in) throws IOException {
-    out.putNextEntry(entry);
-    ByteStreams.copy(in, out);
-    out.closeEntry();
-    return this;
+  private final AtomicReference<ZipEntry> inUse = new AtomicReference<>(null);
+
+  public DexFileArchive(ZipOutputStream out) {
+    this.out = out;
   }
 
   /**
    * Adds a {@code .dex} file with the given details.
    */
   public DexFileArchive addFile(ZipEntry entry, Dex dex) throws IOException {
+    checkState(inUse.compareAndSet(null, entry), "Already in use");
     entry.setSize(dex.getLength());
     out.putNextEntry(entry);
     dex.writeTo(out);
     out.closeEntry();
+    checkState(inUse.compareAndSet(entry, null), "Swooped in: ", inUse.get());
     return this;
   }
 
   @Override
   public void close() throws IOException {
+    checkState(inUse.get() == null, "Still in use: ", inUse.get());
     out.close();
   }
 }
diff --git a/src/tools/android/java/com/google/devtools/build/android/dexer/DexFileMerger.java b/src/tools/android/java/com/google/devtools/build/android/dexer/DexFileMerger.java
index d81befc374..1ecbe00a83 100644
--- a/src/tools/android/java/com/google/devtools/build/android/dexer/DexFileMerger.java
+++ b/src/tools/android/java/com/google/devtools/build/android/dexer/DexFileMerger.java
@@ -16,6 +16,7 @@ package com.google.devtools.build.android.dexer;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkState;
 import static java.nio.charset.StandardCharsets.UTF_8;
+import static java.util.concurrent.TimeUnit.SECONDS;
 
 import com.android.dex.Dex;
 import com.android.dex.DexFormat;
@@ -24,9 +25,12 @@ import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Predicate;
 import com.google.common.base.Predicates;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Iterators;
 import com.google.common.collect.Lists;
 import com.google.common.io.ByteStreams;
+import com.google.common.util.concurrent.ListeningExecutorService;
+import com.google.common.util.concurrent.MoreExecutors;
 import com.google.devtools.build.android.Converters.ExistingPathConverter;
 import com.google.devtools.build.android.Converters.PathConverter;
 import com.google.devtools.common.options.EnumConverter;
@@ -44,7 +48,7 @@ import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Comparator;
-import java.util.Iterator;
+import java.util.concurrent.Executors;
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipFile;
 import java.util.zip.ZipOutputStream;
@@ -170,7 +174,10 @@ class DexFileMerger {
 
   @VisibleForTesting
   static void buildMergedDexFiles(Options options) throws IOException {
-    if (!options.multidexMode.isMultidexAllowed()) {
+    ListeningExecutorService executor;
+    if (options.multidexMode.isMultidexAllowed()) {
+      executor = createThreadPool();
+    } else {
       checkArgument(
           options.mainDexListFile == null,
           "--main-dex-list is only supported with multidex enabled, but mode is: %s",
@@ -179,61 +186,68 @@ class DexFileMerger {
           !options.minimalMainDex,
           "--minimal-main-dex is only supported with multidex enabled, but mode is: %s",
           options.multidexMode);
+      // We'll only ever merge and write one dex file, so multi-threading is pointless.
+      executor = MoreExecutors.newDirectExecutorService();
     }
+
     ImmutableSet<String> classesInMainDex = options.mainDexListFile != null
         ? ImmutableSet.copyOf(Files.readAllLines(options.mainDexListFile, UTF_8))
         : null;
     PrintStream originalStdOut = System.out;
     try (ZipFile zip = new ZipFile(options.inputArchive.toFile());
-        DexFileAggregator out = createDexFileAggregator(options)) {
-      checkForUnprocessedClasses(zip);
+        DexFileAggregator out = createDexFileAggregator(options, executor)) {
+      ArrayList<ZipEntry> dexFiles = filesToProcess(zip);
+
       if (!options.verbose) {
         // com.android.dx.merge.DexMerger prints status information to System.out that we silence
         // here unless it was explicitly requested.  (It also prints debug info to DxContext.out,
         // which we populate accordingly below.)
         System.setOut(Dexing.nullout);
       }
-
       if (classesInMainDex == null) {
-        processDexFiles(zip, out, Predicates.<ZipEntry>alwaysTrue());
+        processDexFiles(zip, dexFiles, out);
       } else {
         // To honor --main_dex_list make two passes:
         // 1. process only the classes listed in the given file
         // 2. process the remaining files
-        Predicate<ZipEntry> classFileFilter = ZipEntryPredicates.classFileFilter(classesInMainDex);
-        processDexFiles(zip, out, classFileFilter);
+        Predicate<ZipEntry> mainDexFilter = ZipEntryPredicates.classFileFilter(classesInMainDex);
+        processDexFiles(zip, Iterables.filter(dexFiles, mainDexFilter), out);
         // Fail if main_dex_list is too big, following dx's example
         checkState(out.getDexFilesWritten() == 0, "Too many classes listed in main dex list file "
             + "%s, main dex capacity exceeded", options.mainDexListFile);
         if (options.minimalMainDex) {
           out.flush(); // Start new .dex file if requested
         }
-        processDexFiles(zip, out, Predicates.not(classFileFilter));
+        processDexFiles(zip, Iterables.filter(dexFiles, Predicates.not(mainDexFilter)), out);
       }
     } finally {
+      // Kill threads in the pool so we don't hang
+      MoreExecutors.shutdownAndAwaitTermination(executor, 1, SECONDS);
       System.setOut(originalStdOut);
     }
-    // Use input's timestamp for output file so the output file is stable.
-    Files.setLastModifiedTime(options.outputArchive,
-        Files.getLastModifiedTime(options.inputArchive));
+  }
+
+  /**
+   * Returns all .dex and .class files in the given zip.  .class files are unexpected but we'll
+   * deal with them later.
+   */
+  private static ArrayList<ZipEntry> filesToProcess(ZipFile zip) {
+    ArrayList<ZipEntry> result = Lists.newArrayList(
+        Iterators.filter(
+            Iterators.forEnumeration(zip.entries()),
+            Predicates.and(
+                Predicates.not(ZipEntryPredicates.isDirectory()),
+                ZipEntryPredicates.suffixes(".dex", ".class"))));
+    Collections.sort(result, ZipEntryComparator.LIKE_DX);
+    return result;
   }
 
   private static void processDexFiles(
-      ZipFile zip, DexFileAggregator out, Predicate<ZipEntry> extraFilter) throws IOException {
-    @SuppressWarnings("unchecked") // Predicates.and uses varargs parameter with generics
-    ArrayList<? extends ZipEntry> filesToProcess =
-        Lists.newArrayList(
-            Iterators.filter(
-                Iterators.forEnumeration(zip.entries()),
-                Predicates.and(
-                    Predicates.not(ZipEntryPredicates.isDirectory()),
-                    ZipEntryPredicates.suffixes(".dex"),
-                    extraFilter)));
-    Collections.sort(filesToProcess, ZipEntryComparator.LIKE_DX);
+      ZipFile zip, Iterable<ZipEntry> filesToProcess, DexFileAggregator out) throws IOException {
     for (ZipEntry entry : filesToProcess) {
       String filename = entry.getName();
       try (InputStream content = zip.getInputStream(entry)) {
-        checkState(filename.endsWith(".dex"), "Shouldn't get here: %s", filename);
+        checkState(filename.endsWith(".dex"), "Input shouldn't contain .class files: %s", filename);
         // We don't want to use the Dex(InputStream) constructor because it closes the stream,
         // which will break the for loop, and it has its own bespoke way of reading the file into
         // a byte buffer before effectively calling Dex(byte[]) anyway.
@@ -242,37 +256,28 @@ class DexFileMerger {
     }
   }
 
-  private static void checkForUnprocessedClasses(ZipFile zip) {
-    Iterator<? extends ZipEntry> classes =
-        Iterators.filter(
-            Iterators.forEnumeration(zip.entries()),
-            Predicates.and(
-                Predicates.not(ZipEntryPredicates.isDirectory()),
-                ZipEntryPredicates.suffixes(".class")));
-    if (classes.hasNext()) {
-      // Hitting this error indicates Jar files not covered by incremental dexing (b/34949364).
-      // Bazel should prevent this error but if you do get this exception, you can use DexBuilder
-      // to convert offending classes first. In Bazel that typically means using java_import or to
-      // make sure Bazel rules use DexBuilder on implicit dependencies.
-      throw new IllegalArgumentException(
-          zip.getName()
-              + " should only contain .dex files but found the following .class files: "
-              + Iterators.toString(classes));
-    }
-  }
-
-  private static DexFileAggregator createDexFileAggregator(Options options) throws IOException {
+  private static DexFileAggregator createDexFileAggregator(
+      Options options, ListeningExecutorService executor) throws IOException {
     return new DexFileAggregator(
         new DxContext(options.verbose ? System.out : ByteStreams.nullOutputStream(), System.err),
         new DexFileArchive(
             new ZipOutputStream(
                 new BufferedOutputStream(Files.newOutputStream(options.outputArchive)))),
+        executor,
         options.multidexMode,
         options.maxNumberOfIdxPerDex,
         options.wasteThresholdPerDex);
   }
 
   /**
+   * Creates an unbounded thread pool executor, which is appropriate here since the number of tasks
+   * we will add to the thread pool is at most dozens and some of them perform I/O (ie, may block).
+   */
+  private static ListeningExecutorService createThreadPool() {
+    return MoreExecutors.listeningDecorator(Executors.newCachedThreadPool());
+  }
+
+  /**
    * Sorts java class names such that outer classes preceed their inner
    * classes and "package-info" preceeds all other classes in its package.
    *
@@ -309,7 +314,7 @@ class DexFileMerger {
 
     @Override
     // Copied from com.android.dx.cf.direct.ClassPathOpener
-    public int compare (ZipEntry a, ZipEntry b) {
+    public int compare(ZipEntry a, ZipEntry b) {
       return compareClassNames(a.getName(), b.getName());
     }
   }