Add action to write android_binary and lib R.classes directly

For android_binary rules, we regenerate all of
the R.java of the libraries to get the final
resource IDs. Compiling all of them together can
be slow with the normal JavaBuilder, so add a
specialized class writer.

Example build with many R.java classes:
- R.java -> R.class via JavaBuilder: over 80s
- ErrorProne takes about 40% of that. So turning off
  ErrorProne would be projected to be 48s.
  Some of ErrorProne slowness is from static field
  checks (e.g., on Flag classes), which may look
  up the same Type over and over.

In comparison, if we write our own bytecode with ASM:
- ~16s total
- 4.7s to parse R.txt
- 4.8s to write class files
- 5.8s to copy and compress .jar

TODO: clean up SymbolLoader patching (upstream)

This only creates the action. We will need to
separately wire this up in blaze.

NOTE: This also makes the exising R.txt loading
used by live code multi-threaded, since that is
partly I/O-bound. Something to watch out for
(for flakiness, etc.) once this is submitted.

--
MOS_MIGRATED_REVID=125384467

1 files changed, 289 insertions, 88 deletions

diff --git a/src/tools/android/java/com/google/devtools/build/android/AndroidResourceProcessor.java b/src/tools/android/java/com/google/devtools/build/android/AndroidResourceProcessor.java
index 33aa697b8e..9bc1ae5d57 100644
--- a/src/tools/android/java/com/google/devtools/build/android/AndroidResourceProcessor.java
+++ b/src/tools/android/java/com/google/devtools/build/android/AndroidResourceProcessor.java
@@ -21,10 +21,16 @@ import com.google.common.base.Strings;
 import com.google.common.collect.ArrayListMultimap;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Multimap;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+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.FullRevisionConverter;
+import com.google.devtools.build.android.resources.RClassWriter;
 import com.google.devtools.common.options.Converters.ColonSeparatedOptionListConverter;
 import com.google.devtools.common.options.Converters.CommaSeparatedOptionListConverter;
 import com.google.devtools.common.options.Option;
@@ -63,6 +69,8 @@ import com.android.utils.StdLogger;
 import org.xml.sax.SAXException;
 
 import java.io.BufferedOutputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.Closeable;
 import java.io.File;
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -80,6 +88,12 @@ import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executors;
+import java.util.jar.Attributes;
+import java.util.jar.JarFile;
+import java.util.jar.Manifest;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import java.util.zip.CRC32;
@@ -217,9 +231,31 @@ public class AndroidResourceProcessor {
     }
   }
 
+  /** Shutdowns and verifies that no tasks are running in the executor service. */
+  private static final class ExecutorServiceCloser implements Closeable {
+    private final ListeningExecutorService executorService;
+    private ExecutorServiceCloser(ListeningExecutorService executorService) {
+      this.executorService = executorService;
+    }
+
+    @Override
+    public void close() throws IOException {
+      List<Runnable> unfinishedTasks = executorService.shutdownNow();
+      if (!unfinishedTasks.isEmpty()) {
+        throw new IOException(
+            "Shutting down the executor with unfinished tasks:" + unfinishedTasks);
+      }
+    }
+
+    public static Closeable createWith(ListeningExecutorService executorService) {
+      return new ExecutorServiceCloser(executorService);
+    }
+  }
+
   private static final ImmutableMap<SystemProperty, String> SYSTEM_PROPERTY_NAMES = Maps.toMap(
       Arrays.asList(SystemProperty.values()), new Function<SystemProperty, String>() {
-        @Override public String apply(SystemProperty property) {
+        @Override
+        public String apply(SystemProperty property) {
           if (property == SystemProperty.PACKAGE) {
             return "applicationId";
           } else {
@@ -237,6 +273,7 @@ public class AndroidResourceProcessor {
 
   /**
    * Copies the R.txt to the expected place.
+   *
    * @param generatedSourceRoot The path to the generated R.txt.
    * @param rOutput The Path to write the R.txt.
    * @param staticIds Boolean that indicates if the ids should be set to 0x1 for caching purposes.
@@ -273,8 +310,9 @@ public class AndroidResourceProcessor {
       Files.createDirectories(srcJar.getParent());
       try (final ZipOutputStream zip = new ZipOutputStream(
           new BufferedOutputStream(Files.newOutputStream(srcJar)))) {
-        Files.walkFileTree(generatedSourcesRoot,
-            new SymbolFileSrcJarBuildingVisitor(zip, generatedSourcesRoot, staticIds));
+        SymbolFileSrcJarBuildingVisitor visitor =
+            new SymbolFileSrcJarBuildingVisitor(zip, generatedSourcesRoot, staticIds);
+        Files.walkFileTree(generatedSourcesRoot, visitor);
       }
       // Set to the epoch for caching purposes.
       Files.setLastModifiedTime(srcJar, FileTime.fromMillis(0L));
@@ -284,6 +322,25 @@ public class AndroidResourceProcessor {
   }
 
   /**
+   * Creates a zip archive from all found R.class (and inner class) files.
+   */
+  public void createClassJar(Path generatedClassesRoot, Path classJar) {
+    try {
+      Files.createDirectories(classJar.getParent());
+      try (final ZipOutputStream zip = new ZipOutputStream(
+          new BufferedOutputStream(Files.newOutputStream(classJar)))) {
+        ClassJarBuildingVisitor visitor = new ClassJarBuildingVisitor(zip, generatedClassesRoot);
+        Files.walkFileTree(generatedClassesRoot, visitor);
+        visitor.writeManifestContent();
+      }
+      // Set to the epoch for caching purposes.
+      Files.setLastModifiedTime(classJar, FileTime.fromMillis(0L));
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  /**
    * Copies the AndroidManifest.xml to the specified output location.
    *
    * @param androidData The MergedAndroidData which contains the manifest to be written to
@@ -344,12 +401,6 @@ public class AndroidResourceProcessor {
       Path mainDexProguardOut,
       Path publicResourcesOut)
       throws IOException, InterruptedException, LoggedErrorException {
-    List<SymbolFileProvider> libraries = new ArrayList<>();
-    for (DependencyAndroidData dataDep : dependencyData) {
-      SymbolFileProvider library = dataDep.asSymbolFileProvider();
-      libraries.add(library);
-    }
-
     Path androidManifest = primaryData.getManifest();
     Path resourceDir = primaryData.getResourceDir();
     Path assetsDir = primaryData.getAssetDir();
@@ -406,9 +457,10 @@ public class AndroidResourceProcessor {
 
     // The R needs to be created for each library in the dependencies,
     // but only if the current project is not a library.
-    writeDependencyPackageRs(variantType, customPackageForR, libraries, androidManifest.toFile(),
-        sourceOut);
-
+    if (sourceOut != null && variantType != VariantConfiguration.Type.LIBRARY) {
+      writeDependencyPackageRJavaFiles(
+          dependencyData, customPackageForR, androidManifest, sourceOut);
+    }
     // Reset the output date stamps.
     if (proguardOut != null) {
       Files.setLastModifiedTime(proguardOut, FileTime.fromMillis(0L));
@@ -424,61 +476,157 @@ public class AndroidResourceProcessor {
     }
   }
 
-  private void writeDependencyPackageRs(VariantConfiguration.Type variantType,
-      String customPackageForR, List<SymbolFileProvider> libraries, File androidManifest,
-      Path sourceOut) throws IOException {
-    if (sourceOut != null && variantType != VariantConfiguration.Type.LIBRARY
-        && !libraries.isEmpty()) {
-      SymbolLoader fullSymbolValues = null;
+  /** Task to parse java package from AndroidManifest.xml */
+  private static final class PackageParsingTask implements Callable<String> {
 
-      String appPackageName = customPackageForR;
-      if (appPackageName == null) {
-        appPackageName = VariantConfiguration.getManifestPackage(androidManifest);
-      }
+    private final File manifest;
 
-      // List of all the symbol loaders per package names.
-      Multimap<String, SymbolLoader> libMap = ArrayListMultimap.create();
+    PackageParsingTask(File manifest) {
+      this.manifest = manifest;
+    }
 
-      for (SymbolFileProvider lib : libraries) {
-        String packageName = VariantConfiguration.getManifestPackage(lib.getManifest());
+    @Override
+    public String call() throws Exception {
+      return VariantConfiguration.getManifestPackage(manifest);
+    }
+  }
+
+  /** Task to load and parse R.txt symbols */
+  private static final class SymbolLoadingTask implements Callable<Object> {
+
+    private final SymbolLoader symbolLoader;
+
+    SymbolLoadingTask(SymbolLoader symbolLoader) {
+      this.symbolLoader = symbolLoader;
+    }
 
+    @Override
+    public Object call() throws Exception {
+      symbolLoader.load();
+      return null;
+    }
+  }
+
+  public SymbolLoader loadResourceSymbolTable(
+      List<SymbolFileProvider> libraries,
+      String appPackageName,
+      Path primaryRTxt,
+      Multimap<String, SymbolLoader> libMap) throws IOException {
+    // The reported availableProcessors may be higher than the actual resources
+    // (on a shared system). On the other hand, a lot of the work is I/O, so it's not completely
+    // CPU bound. As a compromise, divide by 2 the reported availableProcessors.
+    int numThreads = Math.max(1, Runtime.getRuntime().availableProcessors() / 2);
+    ListeningExecutorService executorService = MoreExecutors.listeningDecorator(
+        Executors.newFixedThreadPool(numThreads));
+    try (Closeable closeable = ExecutorServiceCloser.createWith(executorService)) {
+      // Load the package names from the manifest files.
+      Map<SymbolFileProvider, ListenableFuture<String>> packageJobs = new HashMap<>();
+      for (final SymbolFileProvider lib : libraries) {
+        packageJobs.put(lib, executorService.submit(new PackageParsingTask(lib.getManifest())));
+      }
+      Map<SymbolFileProvider, String> packageNames = new HashMap<>();
+      try {
+        for (Map.Entry<SymbolFileProvider, ListenableFuture<String>> entry : packageJobs
+            .entrySet()) {
+          packageNames.put(entry.getKey(), entry.getValue().get());
+        }
+      } catch (InterruptedException | ExecutionException e) {
+        throw new IOException("Failed to load package name: ", e);
+      }
+      // Associate the packages with symbol files.
+      // Eagerly load fullSymbolValues, in case we only have resources defined for the binary
+      // and not for any of the libraries.
+      SymbolLoader fullSymbolValues = new SymbolLoader(primaryRTxt.toFile(), stdLogger);
+      for (SymbolFileProvider lib : libraries) {
+        String packageName = packageNames.get(lib);
         // If the library package matches the app package skip -- the R class will contain
         // all the possible resources so it will not need to generate a new R.
         if (appPackageName.equals(packageName)) {
           continue;
         }
-
         File rFile = lib.getSymbolFile();
         // If the library has no resource, this file won't exist.
         if (rFile.isFile()) {
-          // Load the full values if that's not already been done.
-          // Doing it lazily allow us to support the case where there's no
-          // resources anywhere.
-          if (fullSymbolValues == null) {
-            fullSymbolValues = new SymbolLoader(sourceOut.resolve("R.txt").toFile(), stdLogger);
-            fullSymbolValues.load();
-          }
-
           SymbolLoader libSymbols = new SymbolLoader(rFile, stdLogger);
-          libSymbols.load();
-
-          // store these symbols by associating them with the package name.
           libMap.put(packageName, libSymbols);
         }
       }
+      // Now load the symbol files in parallel.
+      List<ListenableFuture<?>> loadJobs = new ArrayList<>();
+      for (final SymbolLoader loader : Iterables.concat(
+          libMap.values(), ImmutableList.of(fullSymbolValues))) {
+        loadJobs.add(executorService.submit(new SymbolLoadingTask(loader)));
+      }
+      try {
+        Futures.allAsList(loadJobs).get();
+      } catch (InterruptedException | ExecutionException e) {
+        throw new IOException("Failed to load SymbolFile: ", e);
+      }
+      return fullSymbolValues;
+    }
+  }
 
-      // Loop on all the package name, merge all the symbols to write, and write.
-      for (String packageName : libMap.keySet()) {
-        Collection<SymbolLoader> symbols = libMap.get(packageName);
-        SymbolWriter writer = new SymbolWriter(sourceOut.toString(), packageName, fullSymbolValues);
-        for (SymbolLoader symbolLoader : symbols) {
-          writer.addSymbolsToWrite(symbolLoader);
-        }
-        writer.write();
+  private void writeDependencyPackageRJavaFiles(
+      List<DependencyAndroidData> dependencyData,
+      String customPackageForR,
+      Path androidManifest,
+      Path sourceOut) throws IOException {
+    List<SymbolFileProvider> libraries = new ArrayList<>();
+    for (DependencyAndroidData dataDep : dependencyData) {
+      SymbolFileProvider library = dataDep.asSymbolFileProvider();
+      libraries.add(library);
+    }
+    String appPackageName = customPackageForR;
+    if (appPackageName == null) {
+      appPackageName = VariantConfiguration.getManifestPackage(androidManifest.toFile());
+    }
+    Multimap<String, SymbolLoader> libSymbolMap = ArrayListMultimap.create();
+    Path primaryRTxt = sourceOut != null ? sourceOut.resolve("R.txt") : null;
+    if (primaryRTxt != null && !libraries.isEmpty()) {
+      SymbolLoader fullSymbolValues = loadResourceSymbolTable(libraries,
+          appPackageName, primaryRTxt, libSymbolMap);
+      writePackageRJavaFiles(libSymbolMap, fullSymbolValues, sourceOut);
+    }
+  }
+
+  public void writePackageRJavaFiles(
+      Multimap<String, SymbolLoader> libMap,
+      SymbolLoader fullSymbolValues,
+      Path sourceOut) throws IOException {
+    // Loop on all the package name, merge all the symbols to write, and write.
+    for (String packageName : libMap.keySet()) {
+      Collection<SymbolLoader> symbols = libMap.get(packageName);
+      SymbolWriter writer = new SymbolWriter(sourceOut.toString(), packageName, fullSymbolValues);
+      for (SymbolLoader symbolLoader : symbols) {
+        writer.addSymbolsToWrite(symbolLoader);
       }
+      writer.write();
     }
   }
 
+  public void writePackageRClasses(
+      Multimap<String, SymbolLoader> libMap,
+      SymbolLoader fullSymbolValues,
+      String appPackageName,
+      Path classesOut,
+      boolean finalFields) throws IOException {
+    for (String packageName : libMap.keySet()) {
+      Collection<SymbolLoader> symbols = libMap.get(packageName);
+      RClassWriter classWriter =
+          new RClassWriter(classesOut.toFile(), packageName, fullSymbolValues, finalFields);
+      for (SymbolLoader symbolLoader : symbols) {
+        classWriter.addSymbolsToWrite(symbolLoader);
+      }
+      classWriter.write();
+    }
+    // Unlike the R.java generation, we also write the app's R.class file so that the class
+    // jar file can be complete (aapt doesn't generate it for us).
+    RClassWriter classWriter =
+        new RClassWriter(classesOut.toFile(), appPackageName, fullSymbolValues, finalFields);
+    classWriter.addSymbolsToWrite(fullSymbolValues);
+    classWriter.write();
+  }
+
   public MergedAndroidData processManifest(
       VariantConfiguration.Type variantType,
       String customPackageForR,
@@ -769,17 +917,79 @@ public class AndroidResourceProcessor {
     assetSets.add(mainAssets);
   }
 
-  @Nullable private Path prepareOutputPath(@Nullable Path out) throws IOException {
+  @Nullable
+  private Path prepareOutputPath(@Nullable Path out) throws IOException {
     if (out == null) {
       return null;
     }
     return Files.createDirectories(out);
   }
 
+  private static class ZipBuilderVisitor extends SimpleFileVisitor<Path> {
+
+    // The earliest date representable in a zip file, 1-1-1980 (the DOS epoch).
+    private static final long ZIP_EPOCH = 315561600000L;
+    // ZIP timestamps have a resolution of 2 seconds.
+    // see http://www.info-zip.org/FAQ.html#limits
+    private static final long MINIMUM_TIMESTAMP_INCREMENT = 2000L;
+
+    private final ZipOutputStream zip;
+    protected final Path root;
+    private final String directoryPrefix;
+    private int storageMethod = ZipEntry.STORED;
+
+    ZipBuilderVisitor(ZipOutputStream zip, Path root, String directory) {
+      this.zip = zip;
+      this.root = root;
+      this.directoryPrefix = directory;
+    }
+
+    public void setCompress(boolean compress) {
+      storageMethod = compress ? ZipEntry.DEFLATED : ZipEntry.STORED;
+    }
+
+    /**
+     * Normalize timestamps for deterministic builds. Stamp .class files to be a bit newer
+     * than .java files. See:
+     * {@link com.google.devtools.build.buildjar.jarhelper.JarHelper#normalizedTimestamp(String)}
+     */
+    protected long normalizeTime(String filename) {
+      if (filename.endsWith(".class")) {
+        return ZIP_EPOCH + MINIMUM_TIMESTAMP_INCREMENT;
+      } else {
+        return ZIP_EPOCH;
+      }
+    }
+
+    protected void addEntry(Path file, byte[] content) throws IOException {
+      String prefix = directoryPrefix != null ? (directoryPrefix + "/") : "";
+      String relativeName = root.relativize(file).toString();
+      ZipEntry entry = new ZipEntry(prefix + relativeName);
+      entry.setMethod(storageMethod);
+      entry.setTime(normalizeTime(relativeName));
+      entry.setSize(content.length);
+      CRC32 crc32 = new CRC32();
+      crc32.update(content);
+      entry.setCrc(crc32.getValue());
+
+      zip.putNextEntry(entry);
+      zip.write(content);
+      zip.closeEntry();
+    }
+
+    @Override
+    public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
+      byte[] content = Files.readAllBytes(file);
+      addEntry(file, content);
+      return FileVisitResult.CONTINUE;
+    }
+  }
+
   /**
    * A FileVisitor that will add all R.java files to be stored in a zip archive.
    */
-  private static final class SymbolFileSrcJarBuildingVisitor extends SimpleFileVisitor<Path> {
+  private static final class SymbolFileSrcJarBuildingVisitor extends ZipBuilderVisitor {
+
     static final Pattern PACKAGE_PATTERN = Pattern.compile(
         "\\s*package ([a-zA-Z_$][a-zA-Z\\d_$]*(?:\\.[a-zA-Z_$][a-zA-Z\\d_$]*)*)");
     static final Pattern ID_PATTERN = Pattern.compile(
@@ -787,15 +997,10 @@ public class AndroidResourceProcessor {
     static final Pattern INNER_CLASS = Pattern.compile("public static class ([a-z_]*) \\{(.*?)\\}",
         Pattern.DOTALL);
 
-    // The earliest date representable in a zip file, 1-1-1980.
-    private static final long ZIP_EPOCH = 315561600000L;
-    private final ZipOutputStream zip;
-    private final Path root;
     private final boolean staticIds;
 
     private SymbolFileSrcJarBuildingVisitor(ZipOutputStream zip, Path root, boolean staticIds) {
-      this.zip = zip;
-      this.root = root;
+      super(zip, root, null);
       this.staticIds = staticIds;
     }
 
@@ -832,52 +1037,48 @@ public class AndroidResourceProcessor {
           content = replaceIdsWithStaticIds(UTF_8.decode(
               ByteBuffer.wrap(content)).toString()).getBytes(UTF_8);
         }
-        ZipEntry entry = new ZipEntry(root.relativize(file).toString());
-
-        entry.setMethod(ZipEntry.STORED);
-        entry.setTime(ZIP_EPOCH);
-        entry.setSize(content.length);
-        CRC32 crc32 = new CRC32();
-        crc32.update(content);
-        entry.setCrc(crc32.getValue());
-        zip.putNextEntry(entry);
-        zip.write(content);
-        zip.closeEntry();
+        addEntry(file, content);
       }
       return FileVisitResult.CONTINUE;
     }
   }
 
-  private static final class ZipBuilderVisitor extends SimpleFileVisitor<Path> {
-    // The earliest date representable in a zip file, 1-1-1980.
-    private static final long ZIP_EPOCH = 315561600000L;
-    private final ZipOutputStream zip;
-    private final Path root;
-    private final String directory;
+  /**
+   * A FileVisitor that will add all R class files to be stored in a zip archive.
+   */
+  private static final class ClassJarBuildingVisitor extends ZipBuilderVisitor {
 
-    public ZipBuilderVisitor(ZipOutputStream zip, Path root, String directory) {
-      this.zip = zip;
-      this.root = root;
-      this.directory = directory;
+    ClassJarBuildingVisitor(ZipOutputStream zip, Path root) {
+      super(zip, root, null);
     }
 
     @Override
     public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
-      byte[] content = Files.readAllBytes(file);
-
-      CRC32 crc32 = new CRC32();
-      crc32.update(content);
+      Path filename = file.getFileName();
+      String name = filename.toString();
+      if (name.endsWith(".class")) {
+        byte[] content = Files.readAllBytes(file);
+        addEntry(file, content);
+      }
+      return FileVisitResult.CONTINUE;
+    }
 
-      ZipEntry entry = new ZipEntry(directory + "/" + root.relativize(file));
-      entry.setMethod(ZipEntry.STORED);
-      entry.setTime(ZIP_EPOCH);
-      entry.setSize(content.length);
-      entry.setCrc(crc32.getValue());
+    private byte[] manifestContent() throws IOException {
+      Manifest manifest = new Manifest();
+      Attributes attributes = manifest.getMainAttributes();
+      attributes.put(Attributes.Name.MANIFEST_VERSION, "1.0");
+      Attributes.Name createdBy = new Attributes.Name("Created-By");
+      if (attributes.getValue(createdBy) == null) {
+        attributes.put(createdBy, "bazel");
+      }
+      ByteArrayOutputStream out = new ByteArrayOutputStream();
+      manifest.write(out);
+      return out.toByteArray();
+    }
 
-      zip.putNextEntry(entry);
-      zip.write(content);
-      zip.closeEntry();
-      return FileVisitResult.CONTINUE;
+    void writeManifestContent() throws IOException {
+      addEntry(root.resolve(JarFile.MANIFEST_NAME), manifestContent());
     }
   }
+
 }