JVM build time improvements via tracking unused deps or classes [Java part]

src/java_tools/buildjar/java/com/google/devtools/build/buildjar/JavaLibraryBuildRequest.java

@@ -159,6 +159,7 @@ public JavaLibraryBuildRequest(
     if (optionsParser.reduceClasspathMode() != OptionsParser.ReduceClasspathMode.NONE) {
       depsBuilder.setReduceClasspath();
     }
+    depsBuilder.setUsageTrackerMode(optionsParser.usageTrackerMode());
     if (optionsParser.getTargetLabel() != null) {
       depsBuilder.setTargetLabel(optionsParser.getTargetLabel());
     }

src/java_tools/buildjar/java/com/google/devtools/build/buildjar/OptionsParser.java

@@ -67,6 +67,8 @@ public enum ReduceClasspathMode {
    */
   private ReduceClasspathMode reduceClasspathMode = ReduceClasspathMode.NONE;
 
+  private boolean usageTrackerMode = false;
+
   private int fullClasspathLength = -1;
   private int reducedClasspathLength = -1;
 
@@ -156,6 +158,9 @@ private void processCommandlineArgs(Deque<String> argQueue) throws InvalidComman
         case "--reduce_classpath_mode":
           reduceClasspathMode = ReduceClasspathMode.valueOf(getArgument(argQueue, arg));
           break;
+        case "--experimental_track_class_usage":
+          usageTrackerMode = true;
+          break;
         case "--full_classpath_length":
           fullClasspathLength = Integer.parseInt(getArgument(argQueue, arg));
           break;
@@ -378,6 +383,10 @@ public ReduceClasspathMode reduceClasspathMode() {
     return reduceClasspathMode;
   }
 
+  public boolean usageTrackerMode() {
+    return usageTrackerMode;
+  }
+
   public int fullClasspathLength() {
     return fullClasspathLength;
   }

src/java_tools/buildjar/java/com/google/devtools/build/buildjar/javac/plugins/BUILD

@@ -36,10 +36,12 @@ java_library(
         ":plugins",
         "//src/java_tools/buildjar/java/com/google/devtools/build/buildjar:JarOwner",
         "//src/java_tools/buildjar/java/com/google/devtools/build/buildjar/javac/statistics",
+        "//src/main/java/com/google/devtools/build/lib/vfs",
         "//src/main/protobuf:deps_java_proto",
         "//third_party:auto_value",
         "//third_party:guava",
         "//third_party:tomcat_annotations_api",
+        "//third_party/protobuf:protobuf_java",
     ],
 )
 

src/java_tools/buildjar/java/com/google/devtools/build/buildjar/javac/plugins/dependency/DependencyModule.java

@@ -24,6 +24,7 @@
 import com.google.common.collect.Streams;
 import com.google.devtools.build.buildjar.JarOwner;
 import com.google.devtools.build.buildjar.javac.plugins.BlazeJavaCompilerPlugin;
+import com.google.devtools.build.lib.view.proto.Deps;
 import com.google.devtools.build.lib.view.proto.Deps.Dependencies;
 import com.google.devtools.build.lib.view.proto.Deps.Dependency;
 import com.google.devtools.build.lib.view.proto.Deps.Dependency.Kind;
@@ -76,12 +77,15 @@ public static enum StrictJavaDeps {
   private final FixTool fixDepsTool;
   private final ImmutableSet<Path> directJars;
   private final boolean strictClasspathMode;
+  private final boolean usageTrackerMode;
   private final Set<Path> depsArtifacts;
   private final String targetLabel;
   private final Path outputDepsProtoFile;
   private boolean hasMissingTargets;
   private final Map<Path, Dependency> explicitDependenciesMap;
   private final Map<Path, Dependency> implicitDependenciesMap;
+
+  private final Map<Path, Set<Deps.UsedClass>> usedClassesMap;
   private final ImmutableSet<Path> platformJars;
   Set<Path> requiredClasspath;
   private final FixMessage fixMessage;
@@ -93,6 +97,7 @@ public static enum StrictJavaDeps {
       FixTool fixDepsTool,
       ImmutableSet<Path> directJars,
       boolean strictClasspathMode,
+      boolean usageTrackerMode,
       Set<Path> depsArtifacts,
       ImmutableSet<Path> platformJars,
       String targetLabel,
@@ -103,11 +108,13 @@ public static enum StrictJavaDeps {
     this.fixDepsTool = fixDepsTool;
     this.directJars = directJars;
     this.strictClasspathMode = strictClasspathMode;
+    this.usageTrackerMode = usageTrackerMode;
     this.depsArtifacts = depsArtifacts;
     this.targetLabel = targetLabel;
     this.outputDepsProtoFile = outputDepsProtoFile;
     this.explicitDependenciesMap = new HashMap<>();
     this.implicitDependenciesMap = new HashMap<>();
+    this.usedClassesMap = new HashMap<>();
     this.platformJars = platformJars;
     this.fixMessage = fixMessage;
     this.exemptGenerators = exemptGenerators;
@@ -116,7 +123,7 @@ public static enum StrictJavaDeps {
 
   /** Returns a plugin to be enabled in the compiler. */
   public BlazeJavaCompilerPlugin getPlugin() {
-    return new StrictJavaDepsPlugin(this);
+    return new StrictJavaDepsPlugin(this, usageTrackerMode);
   }
 
   /**
@@ -162,11 +169,18 @@ Dependencies buildDependenciesProto(
     // Filter using the original classpath, to preserve ordering.
     for (Path entry : classpath) {
       if (explicitDependenciesMap.containsKey(entry)) {
-        deps.addDependency(explicitDependenciesMap.get(entry));
+        Deps.Dependency d = explicitDependenciesMap.get(entry).toBuilder()
+                .addAllUsedClasses(usedClassesMap.getOrDefault(entry, Set.of()))
+                .build();
+        deps.addDependency(d);
       } else if (implicitDependenciesMap.containsKey(entry)) {
-        deps.addDependency(implicitDependenciesMap.get(entry));
+        Deps.Dependency d = implicitDependenciesMap.get(entry).toBuilder()
+                .addAllUsedClasses(usedClassesMap.getOrDefault(entry, Set.of()))
+                .build();
+        deps.addDependency(d);
       }
     }
+
     return deps.build();
   }
 
@@ -195,6 +209,10 @@ public Map<Path, Dependency> getImplicitDependenciesMap() {
     return implicitDependenciesMap;
   }
 
+  public Map<Path, Set<Deps.UsedClass>> getUsedClassesMap() {
+    return usedClassesMap;
+  }
+
   /** Returns the jars in the platform classpath. */
   public ImmutableSet<Path> getPlatformJars() {
     return platformJars;
@@ -230,6 +248,11 @@ public boolean reduceClasspath() {
     return strictClasspathMode;
   }
 
+  /** Writes used classes information in deps file. */
+  public boolean usageTrackerMode() {
+    return usageTrackerMode;
+  }
+
   void setHasMissingTargets() {
     hasMissingTargets = true;
   }
@@ -337,6 +360,7 @@ public static class Builder {
     private String targetLabel;
     private Path outputDepsProtoFile;
     private boolean strictClasspathMode = false;
+    private boolean usageTrackerMode = false;
     private FixMessage fixMessage = new DefaultFixMessage();
     private final Set<String> exemptGenerators = new LinkedHashSet<>(SJD_EXEMPT_PROCESSORS);
 
@@ -370,6 +394,7 @@ public DependencyModule build() {
           fixDepsTool,
           directJars,
           strictClasspathMode,
+          usageTrackerMode,
           depsArtifacts,
           platformJars,
           targetLabel,
@@ -456,6 +481,16 @@ public Builder setReduceClasspath() {
       return this;
     }
 
+    /**
+     * Set action input usage tracking behavior. Used for build time optimization via compilation avoidance.
+     *
+     * @return this Builder instance
+     */
+    public Builder setUsageTrackerMode(boolean usageTrackerMode) {
+      this.usageTrackerMode = usageTrackerMode;
+      return this;
+    }
+
     /**
      * Set the message to display when a missing indirect dependency is found.
      *

src/java_tools/buildjar/java/com/google/devtools/build/buildjar/javac/plugins/dependency/StrictJavaDepsPlugin.java

@@ -19,12 +19,15 @@
 
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.io.ByteStreams;
 import com.google.devtools.build.buildjar.JarOwner;
 import com.google.devtools.build.buildjar.javac.plugins.BlazeJavaCompilerPlugin;
 import com.google.devtools.build.buildjar.javac.plugins.dependency.DependencyModule.StrictJavaDeps;
 import com.google.devtools.build.buildjar.javac.statistics.BlazeJavacStatistics;
+import com.google.devtools.build.lib.vfs.DigestHashFunction;
 import com.google.devtools.build.lib.view.proto.Deps;
 import com.google.devtools.build.lib.view.proto.Deps.Dependency;
+import com.google.protobuf.ByteString;
 import com.sun.tools.javac.code.Flags;
 import com.sun.tools.javac.code.Kinds;
 import com.sun.tools.javac.code.Symbol;
@@ -42,6 +45,7 @@
 import com.sun.tools.javac.util.Log.WriterKind;
 import com.sun.tools.javac.util.Name;
 import java.io.IOException;
+import java.io.InputStream;
 import java.io.PrintWriter;
 import java.io.UncheckedIOException;
 import java.lang.reflect.Method;
@@ -59,6 +63,7 @@
 import java.util.jar.Manifest;
 import javax.lang.model.element.AnnotationValue;
 import javax.lang.model.util.SimpleAnnotationValueVisitor8;
+import javax.tools.FileObject;
 import javax.tools.JavaFileObject;
 
 /**
@@ -76,6 +81,7 @@ public final class StrictJavaDepsPlugin extends BlazeJavaCompilerPlugin {
   private ImplicitDependencyExtractor implicitDependencyExtractor;
   private CheckingTreeScanner checkingTreeScanner;
   private final DependencyModule dependencyModule;
+  private final boolean usageTrackerMode;
 
   /** Marks seen compilation toplevels and their import sections */
   private final Set<JCTree.JCCompilationUnit> toplevels;
@@ -111,8 +117,9 @@ static SjdDiagnostic create(int pos, String message, JavaFileObject source) {
    * flagging that dependency. Also, we can check whether the direct dependencies were actually
    * necessary, i.e. if their associated jars were used at all for looking up class definitions.
    */
-  public StrictJavaDepsPlugin(DependencyModule dependencyModule) {
+  public StrictJavaDepsPlugin(DependencyModule dependencyModule, boolean usageTrackerMode) {
     this.dependencyModule = dependencyModule;
+    this.usageTrackerMode = usageTrackerMode;
     toplevels = new HashSet<>();
     trees = new HashSet<>();
     missingTargets = new HashSet<>();
@@ -135,7 +142,7 @@ public void init(
     if (checkingTreeScanner == null) {
       Set<Path> platformJars = dependencyModule.getPlatformJars();
       checkingTreeScanner =
-          new CheckingTreeScanner(dependencyModule, diagnostics, missingTargets, platformJars);
+          new CheckingTreeScanner(dependencyModule, diagnostics, missingTargets, platformJars, usageTrackerMode);
       context.put(CheckingTreeScanner.class, checkingTreeScanner);
     }
   }
@@ -230,6 +237,8 @@ private static class CheckingTreeScanner extends TreeScanner {
     /** Collect seen direct dependencies and their associated information */
     private final Map<Path, Deps.Dependency> directDependenciesMap;
 
+    private final Map<Path, Set<Deps.UsedClass>> usedClassesMap;
+
     /** We only emit one warning/error per class symbol */
     private final Set<ClassSymbol> seenClasses = new HashSet<>();
 
@@ -240,19 +249,25 @@ private static class CheckingTreeScanner extends TreeScanner {
     /** The set of jars on the compilation bootclasspath. */
     private final Set<Path> platformJars;
 
+    /** The action usage tracker mode. */
+    private boolean usageTrackerMode;
+
     /** The current source, for diagnostics. */
     private JavaFileObject source = null;
 
     public CheckingTreeScanner(
         DependencyModule dependencyModule,
         List<SjdDiagnostic> diagnostics,
         Set<JarOwner> missingTargets,
-        Set<Path> platformJars) {
+        Set<Path> platformJars,
+        boolean usageTrackerMode) {
       this.directJars = dependencyModule.directJars();
       this.diagnostics = diagnostics;
       this.missingTargets = missingTargets;
       this.directDependenciesMap = dependencyModule.getExplicitDependenciesMap();
+      this.usedClassesMap = dependencyModule.getUsedClassesMap();
       this.platformJars = platformJars;
+      this.usageTrackerMode = usageTrackerMode;
     }
 
     Set<ClassSymbol> getSeenClasses() {
@@ -270,6 +285,33 @@ private void checkTypeLiteral(JCTree node, Symbol sym) {
       // whether that jar was a direct dependency and error out otherwise.
       if (jarPath != null && seenClasses.add(sym.enclClass())) {
         collectExplicitDependency(jarPath, node, sym);
+
+        // Track used classes
+        if (usageTrackerMode) {
+          if (!usedClassesMap.containsKey(jarPath)) {
+            usedClassesMap.put(jarPath, new HashSet());
+          }
+          String internalPath = sym.enclClass().classfile.toString().split(":")[1];
+          Deps.UsedClass usedClass = Deps.UsedClass.newBuilder()
+                  .setFullyQualifiedName(sym.enclClass().fullname.toString())
+                  .setJarInternalPath(internalPath.substring(1, internalPath.length() - 1))
+                  .setHash(hashFile(sym.enclClass().classfile))
+                  .build();
+          usedClassesMap.get(jarPath).add(usedClass);
+        }
+      }
+    }
+
+    /*
+     * Generate Sha256 of input fileObject content.
+     */
+    private static ByteString hashFile(FileObject fileObject) {
+      try {
+        InputStream stream = fileObject.openInputStream();
+        byte[] targetArray = ByteStreams.toByteArray(stream);
+        return ByteString.copyFrom(DigestHashFunction.SHA256.getHashFunction().hashBytes(targetArray).asBytes());
+      } catch (IOException ex) {
+        throw new RuntimeException("Failure to compute hash for " + fileObject, ex);
       }
     }