Change Test Runner to work with the new Bazel coverage implementation.

The new implementation doesn't use one metadata jar for the whole build anymore, but wraps all the uninstrumented classes in each of the build jars, among with a txt file that contains the paths of the files to be instrumented. The paths will be used to output the entire relative filepath to the partial coverage report.

Instead of one metadata jar containing all uninstrumented classes on the runtime classpath, the coverage runner will retrieve all the build jars on the runtime classpath.

Because the build jars contain now both classes and uninstrumented classes, the coverage runner will not analyze the entire given jar, but only the classes in the given jars that have the right suffix.

PiperOrigin-RevId: 161951370

3 files changed, 158 insertions, 21 deletions

diff --git a/src/java_tools/junitrunner/java/com/google/testing/coverage/BUILD b/src/java_tools/junitrunner/java/com/google/testing/coverage/BUILD
index dae20c6ff8..8741b0cf70 100644
--- a/src/java_tools/junitrunner/java/com/google/testing/coverage/BUILD
+++ b/src/java_tools/junitrunner/java/com/google/testing/coverage/BUILD
@@ -52,6 +52,7 @@ java_binary(
     ],
     deps = [
         ":bitfield",
+        "//third_party:guava",
         "//third_party/java/jacoco:blaze-agent",
         "//third_party/java/jacoco:core",
         "//third_party/java/jacoco:report",
diff --git a/src/java_tools/junitrunner/java/com/google/testing/coverage/JacocoCoverageRunner.java b/src/java_tools/junitrunner/java/com/google/testing/coverage/JacocoCoverageRunner.java
index 3ef86a006f..5d13aaf9ea 100644
--- a/src/java_tools/junitrunner/java/com/google/testing/coverage/JacocoCoverageRunner.java
+++ b/src/java_tools/junitrunner/java/com/google/testing/coverage/JacocoCoverageRunner.java
@@ -14,21 +14,33 @@
 
 package com.google.testing.coverage;
 
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Function;
+import com.google.common.base.Splitter;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Iterables;
+import java.io.BufferedReader;
 import java.io.ByteArrayInputStream;
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.InputStreamReader;
 import java.io.Reader;
 import java.lang.reflect.Method;
 import java.net.URL;
-import java.util.ArrayList;
 import java.util.Enumeration;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.TreeMap;
 import java.util.jar.Attributes;
+import java.util.jar.JarEntry;
+import java.util.jar.JarFile;
 import java.util.jar.JarInputStream;
 import java.util.jar.Manifest;
 import org.jacoco.agent.rt.IAgent;
@@ -50,24 +62,30 @@ import org.jacoco.report.ISourceFileLocator;
  * http://www.eclemma.org/jacoco/trunk/doc/examples/java/ReportGenerator.java
  *
  * <p>The following environment variables are expected:
- * JAVA_COVERAGE_FILE - specifies final location of the generated lcov file.
- * JACOCO_METADATA_JAR - specifies jar containing uninstrumented classes to be analyzed.
+ * <ul>
+ * <li>JAVA_COVERAGE_FILE - specifies final location of the generated lcov file.</li>
+ * <li>JACOCO_METADATA_JAR - specifies jar containing uninstrumented classes to be analyzed.</li>
+ * </ul>
  */
 public class JacocoCoverageRunner {
 
-  private final List<File> classesJars;
+  private final ImmutableList<File> classesJars;
   private final InputStream executionData;
   private final File reportFile;
+  private final boolean isNewCoverageImplementation;
   private ExecFileLoader execFileLoader;
 
   public JacocoCoverageRunner(InputStream jacocoExec, String reportPath, File... metadataJars) {
+    this(false, jacocoExec, reportPath, metadataJars);
+  }
+
+  private JacocoCoverageRunner(boolean isNewCoverageImplementation,
+      InputStream jacocoExec, String reportPath, File... metadataJars) {
     executionData = jacocoExec;
     reportFile = new File(reportPath);
 
-    classesJars = new ArrayList<>();
-    for (File metadataJar : metadataJars) {
-      classesJars.add(metadataJar);
-    }
+    this.classesJars = ImmutableList.copyOf(metadataJars);
+    this.isNewCoverageImplementation = isNewCoverageImplementation;
   }
 
   public void create() throws IOException {
@@ -85,10 +103,11 @@ public class JacocoCoverageRunner {
     createReport(bundleCoverage, branchDetails);
   }
 
-  private void createReport(
+  @VisibleForTesting
+  void createReport(
       final IBundleCoverage bundleCoverage, final Map<String, BranchCoverageDetail> branchDetails)
       throws IOException {
-    JacocoLCOVFormatter formatter = new JacocoLCOVFormatter();
+    JacocoLCOVFormatter formatter = new JacocoLCOVFormatter(createPathsSet());
     final IReportVisitor visitor = formatter.createVisitor(reportFile, branchDetails);
 
     // Initialize the report with all of the execution and session information. At this point the
@@ -122,11 +141,16 @@ public class JacocoCoverageRunner {
     visitor.visitEnd();
   }
 
-  private IBundleCoverage analyzeStructure() throws IOException {
+  @VisibleForTesting
+  IBundleCoverage analyzeStructure() throws IOException {
     final CoverageBuilder coverageBuilder = new CoverageBuilder();
     final Analyzer analyzer = new Analyzer(execFileLoader.getExecutionDataStore(), coverageBuilder);
     for (File classesJar : classesJars) {
-      analyzer.analyzeAll(classesJar);
+      if (isNewCoverageImplementation) {
+        analyzeUninstrumentedClassesFromJar(analyzer, classesJar);
+      } else {
+        analyzer.analyzeAll(classesJar);
+      }
     }
 
     // TODO(bazel-team): Find out where the name of the bundle can pop out in the report.
@@ -140,12 +164,81 @@ public class JacocoCoverageRunner {
 
     Map<String, BranchCoverageDetail> result = new TreeMap<>();
     for (File classesJar : classesJars) {
-      analyzer.analyzeAll(classesJar);
+      if (isNewCoverageImplementation) {
+        analyzeUninstrumentedClassesFromJar(analyzer, classesJar);
+      } else {
+        analyzer.analyzeAll(classesJar);
+      }
       result.putAll(analyzer.getBranchDetails());
     }
     return result;
   }
 
+  /**
+   * Analyzes all uninstrumented class files found in the given jar.
+   *
+   * <p>The uninstrumented classes are named using the .class.uninstrumented suffix.
+   */
+  private void analyzeUninstrumentedClassesFromJar(Analyzer analyzer, File jar) throws IOException {
+    JarFile jarFile = new JarFile(jar);
+    JarInputStream jarInputStream = new JarInputStream(new FileInputStream(jar));
+    for (JarEntry jarEntry = jarInputStream.getNextJarEntry();
+        jarEntry != null;
+        jarEntry = jarInputStream.getNextJarEntry()) {
+      String jarEntryName = jarEntry.getName();
+      if (jarEntryName.endsWith(".class.uninstrumented")) {
+        analyzer.analyzeAll(jarFile.getInputStream(jarEntry), jarEntryName);
+      }
+    }
+  }
+
+  /**
+   * Creates a {@link Set} containing the paths of the covered Java files.
+   *
+   * <p>The paths are retrieved from a txt file that is found inside each jar containing
+   * uninstrumented classes. Each line of the txt file represents a path to be added to the set.
+   *
+   * <p>This set is needed by {@link JacocoLCOVFormatter} in order to output the correct path for
+   * each covered class.
+   */
+  @VisibleForTesting
+  ImmutableSet<String> createPathsSet() throws IOException {
+    if (!isNewCoverageImplementation) {
+      return ImmutableSet.<String>of();
+    }
+    ImmutableSet.Builder<String> execPathsSetBuilder = ImmutableSet.builder();
+    for (File classJar : classesJars) {
+      addEntriesToExecPathsSet(classJar, execPathsSetBuilder);
+    }
+    return execPathsSetBuilder.build();
+  }
+
+  /**
+   * Adds to the given {@link Set} the paths found in a txt file inside the given jar.
+   *
+   * <p>If a jar contains uninstrumented classes it will also contain a txt file with the paths of
+   * each of these classes, one on each line.
+   */
+  @VisibleForTesting
+  static void addEntriesToExecPathsSet(
+      File jar, ImmutableSet.Builder<String> execPathsSetBuilder) throws IOException {
+    JarFile jarFile = new JarFile(jar);
+    JarInputStream jarInputStream = new JarInputStream(new FileInputStream(jar));
+    for (JarEntry jarEntry = jarInputStream.getNextJarEntry();
+        jarEntry != null;
+        jarEntry = jarInputStream.getNextJarEntry()) {
+      String jarEntryName = jarEntry.getName();
+      if (jarEntryName.endsWith("-paths-for-coverage.txt")) {
+        BufferedReader bufferedReader =
+            new BufferedReader(new InputStreamReader(jarFile.getInputStream(jarEntry), UTF_8));
+        String line;
+        while ((line = bufferedReader.readLine()) != null) {
+          execPathsSetBuilder.add(line);
+        }
+      }
+    }
+  }
+
   private static String getMainClass(String metadataJar) throws Exception {
     if (metadataJar != null) {
       // Blaze guarantees that JACOCO_METADATA_JAR has a proper manifest with a Main-Class entry.
@@ -195,6 +288,10 @@ public class JacocoCoverageRunner {
 
   public static void main(String[] args) throws Exception {
     final String metadataJar = System.getenv("JACOCO_METADATA_JAR");
+    String newMetadataJarsString = System.getenv("JACOCO_METADATA_JARS");
+    final List<String> newMetadataJars =
+        newMetadataJarsString == null ? null : Splitter.on(':').splitToList(newMetadataJarsString);
+
     final String coverageReportBase = System.getenv("JAVA_COVERAGE_FILE");
 
     // Disable Jacoco's default output mechanism, which runs as a shutdown hook. We generate the
@@ -256,6 +353,16 @@ public class JacocoCoverageRunner {
                     // warning if this happens. It's too late at this point.
                     new JacocoCoverageRunner(dataInputStream, coverageReport, new File(metadataJar))
                         .create();
+                  } else if (newMetadataJars != null){
+                    File[] metadataJars = Iterables.toArray(
+                        Iterables.transform(newMetadataJars, new Function<String, File>() {
+                          @Override
+                          public File apply(String input) {
+                            return new File(input);
+                          }
+                        }), File.class);
+                    new JacocoCoverageRunner(true, dataInputStream, coverageReport, metadataJars)
+                        .create();
                   }
                 } catch (IOException e) {
                   e.printStackTrace();
@@ -270,9 +377,9 @@ public class JacocoCoverageRunner {
     // the subprocess to match all JVM flags, runtime classpath, bootclasspath, etc is doable.
     // We'd share the same limitation if the system under test uses shutdown hooks internally, as
     // there's no way to collect coverage data on that code.
-    String mainClass = getMainClass(metadataJar);
-    Method main =
-        Class.forName(mainClass).getMethod("main", new Class[] {java.lang.String[].class});
+    String mainClass =
+        newMetadataJars == null ? getMainClass(metadataJar) : System.getenv("JACOCO_MAIN_CLASS");
+    Method main = Class.forName(mainClass).getMethod("main", String[].class);
     main.invoke(null, new Object[] {args});
   }
 }
diff --git a/src/java_tools/junitrunner/java/com/google/testing/coverage/JacocoLCOVFormatter.java b/src/java_tools/junitrunner/java/com/google/testing/coverage/JacocoLCOVFormatter.java
index 8c991a11d0..20f32e16cf 100644
--- a/src/java_tools/junitrunner/java/com/google/testing/coverage/JacocoLCOVFormatter.java
+++ b/src/java_tools/junitrunner/java/com/google/testing/coverage/JacocoLCOVFormatter.java
@@ -11,13 +11,18 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
-
 package com.google.testing.coverage;
 
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static java.nio.file.StandardOpenOption.APPEND;
+import static java.nio.file.StandardOpenOption.CREATE;
+
+import com.google.common.collect.ImmutableSet;
 import java.io.File;
-import java.io.FileWriter;
 import java.io.IOException;
 import java.io.PrintWriter;
+import java.io.Writer;
+import java.nio.file.Files;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
@@ -41,6 +46,19 @@ import org.jacoco.report.ISourceFileLocator;
  */
 public class JacocoLCOVFormatter {
 
+  // Exec paths of the uninstrumented files that are being analyzed. This is helpful for files in
+  // jars passed through java_import or some custom rule where blaze doesn't have enough context to
+  // compute the right paths, but relies on these pre-computed exec paths.
+  private final ImmutableSet<String> execPathsOfUninstrumentedFiles;
+
+  public JacocoLCOVFormatter(ImmutableSet<String> execPathsOfUninstrumentedFiles) {
+    this.execPathsOfUninstrumentedFiles = execPathsOfUninstrumentedFiles;
+  }
+
+  public JacocoLCOVFormatter() {
+    this.execPathsOfUninstrumentedFiles = ImmutableSet.of();
+  }
+
   public IReportVisitor createVisitor(
       final File output, final Map<String, BranchCoverageDetail> branchCoverageDetail) {
     return new IReportVisitor() {
@@ -48,13 +66,22 @@ public class JacocoLCOVFormatter {
       private Map<String, Map<String, IClassCoverage>> sourceToClassCoverage = new TreeMap<>();
       private Map<String, ISourceFileCoverage> sourceToFileCoverage = new TreeMap<>();
 
+      private String getExecPathForEntryName(String fileName) {
+        for (String execPath : execPathsOfUninstrumentedFiles) {
+          if (execPath.endsWith("/" + fileName)) {
+            return execPath;
+          }
+        }
+        return fileName;
+      }
+
       @Override
       public void visitInfo(List<SessionInfo> sessionInfos, Collection<ExecutionData> executionData)
           throws IOException {}
 
       @Override
       public void visitEnd() throws IOException {
-        try (FileWriter fileWriter = new FileWriter(output, true);
+        try (Writer fileWriter = Files.newBufferedWriter(output.toPath(), UTF_8, CREATE, APPEND);
             PrintWriter printWriter = new PrintWriter(fileWriter)) {
           for (String sourceFile : sourceToClassCoverage.keySet()) {
             processSourceFile(printWriter, sourceFile);
@@ -73,7 +100,8 @@ public class JacocoLCOVFormatter {
         // information and process everything at the end.
         for (IPackageCoverage pkgCoverage : bundle.getPackages()) {
           for (IClassCoverage clsCoverage : pkgCoverage.getClasses()) {
-            String fileName = clsCoverage.getPackageName() + "/" + clsCoverage.getSourceFileName();
+            String fileName = getExecPathForEntryName(
+                clsCoverage.getPackageName() + "/" + clsCoverage.getSourceFileName());
             if (!sourceToClassCoverage.containsKey(fileName)) {
               sourceToClassCoverage.put(fileName, new TreeMap<String, IClassCoverage>());
             }
@@ -81,7 +109,8 @@ public class JacocoLCOVFormatter {
           }
           for (ISourceFileCoverage srcCoverage : pkgCoverage.getSourceFiles()) {
             sourceToFileCoverage.put(
-                srcCoverage.getPackageName() + "/" + srcCoverage.getName(), srcCoverage);
+                getExecPathForEntryName(srcCoverage.getPackageName() + "/" + srcCoverage.getName()),
+                srcCoverage);
           }
         }
       }