Reduce runtime overhead of line-level per-test coverage

dd-java-agent/agent-ci-visibility/civisibility-instrumentation-test-fixtures/src/main/groovy/datadog/trace/civisibility/CiVisibilityInstrumentationTest.groovy

@@ -15,7 +15,6 @@ import datadog.trace.api.civisibility.config.LibraryCapability
 import datadog.trace.api.civisibility.config.TestFQN
 import datadog.trace.api.civisibility.config.TestIdentifier
 import datadog.trace.api.civisibility.config.TestMetadata
-import datadog.trace.api.civisibility.coverage.CoveragePerTestBridge
 import datadog.trace.api.civisibility.events.TestEventsHandler
 import datadog.trace.api.civisibility.telemetry.CiVisibilityMetricCollector
 import datadog.trace.api.civisibility.telemetry.tag.Provider
@@ -201,8 +200,6 @@ abstract class CiVisibilityInstrumentationTest extends InstrumentationSpecificat
     InstrumentationBridge.registerBuildEventsHandlerFactory {
       decorator -> new BuildEventsHandlerImpl<>(buildSystemSessionFactory, new JvmInfoFactoryImpl())
     }
-
-    CoveragePerTestBridge.registerCoverageStoreRegistry(coverageStoreFactory)
   }
 
   private static final class MockExecutionSettingsFactory implements ExecutionSettingsFactory {

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/CiVisibilitySystem.java

@@ -8,7 +8,6 @@
 import datadog.trace.api.civisibility.DDTestSuite;
 import datadog.trace.api.civisibility.InstrumentationBridge;
 import datadog.trace.api.civisibility.config.LibraryCapability;
-import datadog.trace.api.civisibility.coverage.CoveragePerTestBridge;
 import datadog.trace.api.civisibility.events.BuildEventsHandler;
 import datadog.trace.api.civisibility.events.TestEventsHandler;
 import datadog.trace.api.civisibility.telemetry.CiVisibilityMetricCollector;
@@ -118,7 +117,6 @@ public static void start(Instrumentation inst, SharedCommunicationObjects sco) {
       TestEventsHandlerFactory testEventsHandlerFactory =
           new TestEventsHandlerFactory(services, repoServices, coverageServices, executionSettings);
       InstrumentationBridge.registerTestEventsHandlerFactory(testEventsHandlerFactory);
-      CoveragePerTestBridge.registerCoverageStoreRegistry(coverageServices.coverageStoreFactory);
 
       AgentTracer.TracerAPI tracerAPI = AgentTracer.get();
       tracerAPI.addShutdownListener(testEventsHandlerFactory::shutdown);

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/coverage/SkippableAwareCoverageStoreFactory.java

@@ -31,9 +31,4 @@ public CoverageStore create(@Nullable TestIdentifier testIdentifier) {
       return delegate.create(testIdentifier);
     }
   }
-
-  @Override
-  public void setTotalProbeCount(String className, int totalProbeCount) {
-    delegate.setTotalProbeCount(className, totalProbeCount);
-  }
 }

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/coverage/file/FileCoverageStore.java

@@ -118,10 +118,5 @@ public CoverageStore create(@Nullable TestIdentifier testIdentifier) {
     private FileProbes createProbes(boolean isTestThread) {
       return new FileProbes(metrics, isTestThread);
     }
-
-    @Override
-    public void setTotalProbeCount(String className, int totalProbeCount) {
-      // no op
-    }
   }
 }

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/coverage/file/FileProbes.java

@@ -31,11 +31,6 @@ public class FileProbes implements CoverageProbes {
     nonCodeResources = isTestThread ? new HashMap<>() : new ConcurrentHashMap<>();
   }
 
-  @Override
-  public void record(Class<?> clazz, long classId, int probeId) {
-    record(clazz);
-  }
-
   @Override
   public void record(Class<?> clazz) {
     try {

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/coverage/line/ExecutionDataAdapter.java

@@ -5,21 +5,28 @@
 public class ExecutionDataAdapter {
   private final long classId;
   private final String className;
-  // Unbounded data structure that only exists within a single test span
+  // Jacoco's shared probe array for the class, used to back-fill aggregate coverage at report time
+  private final boolean[] jacocoProbes;
+  // Per-test probe array that Jacoco's instrumentation writes into while a test is running
   private final boolean[] probeActivations;
 
-  public ExecutionDataAdapter(long classId, String className, int totalProbeCount) {
+  public ExecutionDataAdapter(long classId, String className, boolean[] jacocoProbes) {
     this.classId = classId;
     this.className = className;
-    this.probeActivations = new boolean[totalProbeCount];
+    this.jacocoProbes = jacocoProbes;
+    this.probeActivations = new boolean[jacocoProbes.length];
   }
 
   public String getClassName() {
     return className;
   }
 
-  void record(int probeId) {
-    probeActivations[probeId] = true;
+  long getClassId() {
+    return classId;
+  }
+
+  boolean[] getProbeActivations() {
+    return probeActivations;
   }
 
   ExecutionDataAdapter merge(ExecutionDataAdapter other) {
@@ -29,6 +36,31 @@ ExecutionDataAdapter merge(ExecutionDataAdapter other) {
     return this;
   }
 
+  /**
+   * Folds the per-test coverage back into Jacoco's shared probe array. Jacoco's aggregate coverage
+   * (used for total module/session coverage percentage and report uploads) no longer sees probes
+   * recorded into the per-test array directly, so they are OR-ed back here at report time. The
+   * write is monotonic (bits are only ever set), so concurrent back-fills from multiple tests are
+   * safe.
+   *
+   * <p>The per-test array is allocated when a method of the class is entered (so the probe array
+   * can be swapped in), which can happen even if no probe ends up firing (e.g. the method throws
+   * before reaching its first probe). Returning whether any probe was actually covered lets the
+   * caller skip such classes and avoid emitting empty coverage entries.
+   *
+   * @return {@code true} if at least one probe was covered by the test
+   */
+  boolean mergeIntoJacocoProbes() {
+    boolean covered = false;
+    for (int i = 0; i < probeActivations.length; i++) {
+      if (probeActivations[i]) {
+        jacocoProbes[i] = true;
+        covered = true;
+      }
+    }
+    return covered;
+  }
+
   ExecutionData toExecutionData() {
     return new ExecutionData(classId, className, probeActivations);
   }

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/coverage/line/LineCoverageStore.java

@@ -14,6 +14,7 @@
 import datadog.trace.civisibility.source.Utils;
 import java.io.InputStream;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.BitSet;
 import java.util.Collection;
 import java.util.HashMap;
@@ -37,16 +38,27 @@ public class LineCoverageStore extends ConcurrentCoverageStore<LineProbes> {
 
   private static final Logger log = LoggerFactory.getLogger(LineCoverageStore.class);
 
+  /**
+   * Upper bound on the number of cached class analyses. Coverage stays correct beyond it (analysis
+   * just isn't cached), this only guards memory for pathologically large suites.
+   */
+  private static final int MAX_ANALYSIS_CACHE_ENTRIES = 50_000;
+
   private final CiVisibilityMetricCollector metrics;
   private final SourcePathResolver sourcePathResolver;
+  // Module-wide cache: (class id + probe set) -> covered lines, shared across tests so a class
+  // covered identically by many tests is parsed by Jacoco's Analyzer only once.
+  private final Map<AnalysisCacheKey, BitSet> analysisCache;
 
   private LineCoverageStore(
       Function<Boolean, LineProbes> probesFactory,
       CiVisibilityMetricCollector metrics,
-      SourcePathResolver sourcePathResolver) {
+      SourcePathResolver sourcePathResolver,
+      Map<AnalysisCacheKey, BitSet> analysisCache) {
     super(probesFactory);
     this.metrics = metrics;
     this.sourcePathResolver = sourcePathResolver;
+    this.analysisCache = analysisCache;
   }
 
   @Nullable
@@ -70,6 +82,12 @@ protected TestReport report(
     Map<String, BitSet> coveredLinesBySourcePath = new HashMap<>();
     for (Map.Entry<Class<?>, ExecutionDataAdapter> e : combinedExecutionData.entrySet()) {
       ExecutionDataAdapter executionDataAdapter = e.getValue();
+      // Back-fill Jacoco's aggregate coverage (total coverage percentage and report uploads). Skip
+      // classes with no covered probes: the per-test array is allocated on method entry, so a
+      // method that throws before its first probe fires would otherwise yield an empty entry.
+      if (!executionDataAdapter.mergeIntoJacocoProbes()) {
+        continue;
+      }
       String className = executionDataAdapter.getClassName();
 
       Class<?> clazz = e.getKey();
@@ -83,24 +101,9 @@ protected TestReport report(
       }
       String sourcePath = sourcePaths.iterator().next();
 
-      try (InputStream is = Utils.getClassStream(clazz)) {
-        BitSet coveredLines =
-            coveredLinesBySourcePath.computeIfAbsent(sourcePath, key -> new BitSet());
-        ExecutionDataStore store = new ExecutionDataStore();
-        store.put(executionDataAdapter.toExecutionData());
-
-        // TODO optimize this part to avoid parsing
-        //  the same class multiple times for different test cases
-        Analyzer analyzer = new Analyzer(store, new SourceAnalyzer(coveredLines));
-        analyzer.analyzeClass(is, null);
-
-      } catch (Exception exception) {
-        log.debug(
-            "Skipping coverage reporting for {} ({}) because of error",
-            className,
-            sourcePath,
-            exception);
-        metrics.add(CiVisibilityCountMetric.CODE_COVERAGE_ERRORS, 1);
+      BitSet coveredLines = analyzeClass(clazz, executionDataAdapter);
+      if (coveredLines != null) {
+        coveredLinesBySourcePath.computeIfAbsent(sourcePath, key -> new BitSet()).or(coveredLines);
       }
     }
 
@@ -132,12 +135,81 @@ protected TestReport report(
     return report;
   }
 
-  public static final class Factory implements CoverageStore.Factory {
+  /**
+   * Resolves the covered lines for a class given a test's probe activations. Parsing the class with
+   * Jacoco's {@link Analyzer} is the dominant cost of reporting, and the result depends only on the
+   * class bytecode and the probe set, so it is memoized: the same class covered identically by
+   * different tests is parsed once.
+   *
+   * @return the covered lines, or {@code null} if the class could not be analyzed
+   */
+  @Nullable
+  private BitSet analyzeClass(Class<?> clazz, ExecutionDataAdapter executionDataAdapter) {
+    AnalysisCacheKey key =
+        new AnalysisCacheKey(
+            executionDataAdapter.getClassId(), executionDataAdapter.getProbeActivations());
+    BitSet cached = analysisCache.get(key);
+    if (cached != null) {
+      return cached;
+    }
+
+    try (InputStream is = Utils.getClassStream(clazz)) {
+      BitSet coveredLines = new BitSet();
+      ExecutionDataStore store = new ExecutionDataStore();
+      store.put(executionDataAdapter.toExecutionData());
+      Analyzer analyzer = new Analyzer(store, new SourceAnalyzer(coveredLines));
+      analyzer.analyzeClass(is, null);
+
+      if (analysisCache.size() < MAX_ANALYSIS_CACHE_ENTRIES) {
+        analysisCache.putIfAbsent(key, coveredLines);
+      }
+      return coveredLines;
+
+    } catch (Exception exception) {
+      log.debug(
+          "Skipping coverage reporting for {} because of error",
+          executionDataAdapter.getClassName(),
+          exception);
+      metrics.add(CiVisibilityCountMetric.CODE_COVERAGE_ERRORS, 1);
+      return null;
+    }
+  }
 
-    private final Map<String, Integer> probeCounts = new ConcurrentHashMap<>();
+  /** Cache key identifying a class (by Jacoco class id) covered by a specific set of probes. */
+  static final class AnalysisCacheKey {
+    private final long classId;
+    private final boolean[] probes;
+    private final int hash;
+
+    AnalysisCacheKey(long classId, boolean[] probes) {
+      this.classId = classId;
+      this.probes = probes;
+      this.hash = 31 * Long.hashCode(classId) + Arrays.hashCode(probes);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) {
+        return true;
+      }
+      if (!(o instanceof AnalysisCacheKey)) {
+        return false;
+      }
+      AnalysisCacheKey other = (AnalysisCacheKey) o;
+      return classId == other.classId && hash == other.hash && Arrays.equals(probes, other.probes);
+    }
+
+    @Override
+    public int hashCode() {
+      return hash;
+    }
+  }
+
+  public static final class Factory implements CoverageStore.Factory {
 
     private final CiVisibilityMetricCollector metrics;
     private final SourcePathResolver sourcePathResolver;
+    private final Map<AnalysisCacheKey, BitSet> analysisCache = new ConcurrentHashMap<>();
 
     public Factory(CiVisibilityMetricCollector metrics, SourcePathResolver sourcePathResolver) {
       this.metrics = metrics;
@@ -146,16 +218,11 @@ public Factory(CiVisibilityMetricCollector metrics, SourcePathResolver sourcePat
 
     @Override
     public CoverageStore create(@Nullable TestIdentifier testIdentifier) {
-      return new LineCoverageStore(this::createProbes, metrics, sourcePathResolver);
+      return new LineCoverageStore(this::createProbes, metrics, sourcePathResolver, analysisCache);
     }
 
     private LineProbes createProbes(boolean isTestThread) {
-      return new LineProbes(metrics, probeCounts, isTestThread);
-    }
-
-    @Override
-    public void setTotalProbeCount(String className, int totalProbeCount) {
-      probeCounts.put(className.replace('/', '.'), totalProbeCount);
+      return new LineProbes(metrics, isTestThread);
     }
   }
 }

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/coverage/line/LineProbes.java

@@ -19,18 +19,15 @@
 public class LineProbes implements CoverageProbes {
 
   private final CiVisibilityMetricCollector metrics;
-  private final Map<String, Integer> probeCounts;
 
   private final Map<Class<?>, ExecutionDataAdapter> executionData;
   private final Map<String, String> nonCodeResources;
 
   private Class<?> lastCoveredClass;
   private ExecutionDataAdapter lastCoveredExecutionData;
 
-  LineProbes(
-      CiVisibilityMetricCollector metrics, Map<String, Integer> probeCounts, boolean isTestThread) {
+  LineProbes(CiVisibilityMetricCollector metrics, boolean isTestThread) {
     this.metrics = metrics;
-    this.probeCounts = probeCounts;
     executionData = isTestThread ? new IdentityHashMap<>() : new ConcurrentHashMap<>();
     nonCodeResources = isTestThread ? new HashMap<>() : new ConcurrentHashMap<>();
   }
@@ -41,20 +38,21 @@ public void record(Class<?> clazz) {
   }
 
   @Override
-  public void record(Class<?> clazz, long classId, int probeId) {
+  public boolean[] resolveProbeArray(Class<?> clazz, long classId, boolean[] jacocoProbes) {
     try {
       if (lastCoveredClass != clazz) {
-        // optimization to avoid map lookup if activating several probes for same class in a row
+        // optimization to avoid map lookup if resolving the array for the same class in a row
         lastCoveredExecutionData =
             executionData.computeIfAbsent(
                 lastCoveredClass = clazz,
-                k -> new ExecutionDataAdapter(classId, k.getName(), probeCounts.get(k.getName())));
+                k -> new ExecutionDataAdapter(classId, k.getName(), jacocoProbes));
       }
-      lastCoveredExecutionData.record(probeId);
+      return lastCoveredExecutionData.getProbeActivations();
 
     } catch (Exception e) {
       metrics.add(CiVisibilityCountMetric.CODE_COVERAGE_ERRORS, 1, CoverageErrorType.RECORD);
-      throw e;
+      // fall back to Jacoco's shared array so coverage is still recorded
+      return jacocoProbes;
     }
   }