Memoize per-test class analysis in a module-wide cache

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

@@ -1,7 +1,5 @@
 package datadog.trace.civisibility.coverage.line;
 
-import org.jacoco.core.data.ExecutionData;
-
 public class ExecutionDataAdapter {
   private final long classId;
   private final String className;
@@ -18,6 +16,14 @@ public String getClassName() {
     return className;
   }
 
+  long getClassId() {
+    return classId;
+  }
+
+  boolean[] getProbeActivations() {
+    return probeActivations;
+  }
+
   void record(int probeId) {
     probeActivations[probeId] = true;
   }
@@ -28,8 +34,4 @@ ExecutionDataAdapter merge(ExecutionDataAdapter other) {
     }
     return this;
   }
-
-  ExecutionData toExecutionData() {
-    return new ExecutionData(classId, className, probeActivations);
-  }
 }

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

@@ -22,9 +22,11 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Function;
 import javax.annotation.Nullable;
 import org.jacoco.core.analysis.Analyzer;
+import org.jacoco.core.data.ExecutionData;
 import org.jacoco.core.data.ExecutionDataStore;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -37,16 +39,41 @@ public class LineCoverageStore extends ConcurrentCoverageStore<LineProbes> {
 
   private static final Logger log = LoggerFactory.getLogger(LineCoverageStore.class);
 
+  /**
+   * Upper bound on the approximate memory retained by the analysis cache. Coverage stays correct
+   * beyond it (analysis just isn't cached), this only guards memory for pathologically large
+   * suites. Bounding by bytes (rather than entry count) is what keeps a class with many probes,
+   * covered by many distinct probe sets, from retaining large arrays for the whole module lifetime.
+   */
+  private static final long MAX_ANALYSIS_CACHE_BYTES = 64L * 1024 * 1024;
+
+  /**
+   * Approximate fixed cost of one cache entry beyond its variable bit data: the {@link
+   * AnalysisCacheKey} and both {@link BitSet} objects (with their {@code long[]} + array headers)
+   * plus the {@code ConcurrentHashMap} node. Deliberately generous so small entries aren't
+   * undercounted and the byte bound stays a real ceiling.
+   */
+  private static final int APPROX_ENTRY_OVERHEAD_BYTES = 160;
+
   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;
+  // Approximate bytes retained by analysisCache, so the cache is bounded by size, not entry count.
+  private final AtomicLong analysisCacheBytes;
 
   private LineCoverageStore(
       Function<Boolean, LineProbes> probesFactory,
       CiVisibilityMetricCollector metrics,
-      SourcePathResolver sourcePathResolver) {
+      SourcePathResolver sourcePathResolver,
+      Map<AnalysisCacheKey, BitSet> analysisCache,
+      AtomicLong analysisCacheBytes) {
     super(probesFactory);
     this.metrics = metrics;
     this.sourcePathResolver = sourcePathResolver;
+    this.analysisCache = analysisCache;
+    this.analysisCacheBytes = analysisCacheBytes;
   }
 
   @Nullable
@@ -83,24 +110,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,9 +144,110 @@ protected TestReport report(
     return report;
   }
 
+  /**
+   * 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) {
+    long classId = executionDataAdapter.getClassId();
+    // Snapshot the activations once and use the same snapshot for both the cache key and the
+    // analysis. The per-test array is mutable and a propagated/background thread may record a late
+    // probe while report() runs; sharing one snapshot ensures the cached covered lines always match
+    // the key's probe set, so a late activation can't poison the entry for later tests.
+    boolean[] probes = executionDataAdapter.getProbeActivations().clone();
+    AnalysisCacheKey key = new AnalysisCacheKey(classId, probes);
+    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(new ExecutionData(classId, executionDataAdapter.getClassName(), probes));
+      Analyzer analyzer = new Analyzer(store, new SourceAnalyzer(coveredLines));
+      analyzer.analyzeClass(is, null);
+
+      // Reserve the entry's weight before inserting so concurrent inserts near the limit can't
+      // collectively overshoot the bound; release the reservation if we exceed it or another thread
+      // cached the class first.
+      long entryBytes =
+          APPROX_ENTRY_OVERHEAD_BYTES + key.packedBytes() + (coveredLines.size() >>> 3);
+      if (analysisCacheBytes.addAndGet(entryBytes) <= MAX_ANALYSIS_CACHE_BYTES) {
+        if (analysisCache.putIfAbsent(key, coveredLines) != null) {
+          analysisCacheBytes.addAndGet(-entryBytes);
+        }
+      } else {
+        analysisCacheBytes.addAndGet(-entryBytes);
+      }
+      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;
+    }
+  }
+
+  /**
+   * Cache key identifying a class (by Jacoco class id) covered by a specific set of probes. The
+   * probe activations are bit-packed into a {@link BitSet} rather than retaining the test's full
+   * {@code boolean[]} (1 byte/element), so a cached key uses ~8x less memory. Equality is exact:
+   * two keys match iff the same class was covered by the same set of probe ids.
+   */
+  static final class AnalysisCacheKey {
+    private final long classId;
+    private final BitSet probes;
+    private final int hash;
+
+    AnalysisCacheKey(long classId, boolean[] probeActivations) {
+      this.classId = classId;
+      BitSet bits = new BitSet(probeActivations.length);
+      for (int i = 0; i < probeActivations.length; i++) {
+        if (probeActivations[i]) {
+          bits.set(i);
+        }
+      }
+      this.probes = bits;
+      this.hash = 31 * Long.hashCode(classId) + bits.hashCode();
+    }
+
+    /** Bytes of the packed probe bits (the variable part of the retained key). */
+    int packedBytes() {
+      return probes.size() >>> 3;
+    }
+
+    @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 && probes.equals(other.probes);
+    }
+
+    @Override
+    public int hashCode() {
+      return hash;
+    }
+  }
+
   public static final class Factory implements CoverageStore.Factory {
 
     private final Map<String, Integer> probeCounts = new ConcurrentHashMap<>();
+    private final Map<AnalysisCacheKey, BitSet> analysisCache = new ConcurrentHashMap<>();
+    private final AtomicLong analysisCacheBytes = new AtomicLong();
 
     private final CiVisibilityMetricCollector metrics;
     private final SourcePathResolver sourcePathResolver;
@@ -146,7 +259,8 @@ 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, analysisCacheBytes);
     }
 
     private LineProbes createProbes(boolean isTestThread) {

dd-java-agent/agent-ci-visibility/src/test/java/datadog/trace/civisibility/coverage/line/LineCoverageStoreTest.java

@@ -0,0 +1,38 @@
+package datadog.trace.civisibility.coverage.line;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+
+import datadog.trace.civisibility.coverage.line.LineCoverageStore.AnalysisCacheKey;
+import org.junit.jupiter.api.Test;
+
+class LineCoverageStoreTest {
+
+  @Test
+  void cacheKeyReusesAnalysisForSameClassAndProbes() {
+    AnalysisCacheKey key = new AnalysisCacheKey(1L, new boolean[] {true, false, true, false});
+    AnalysisCacheKey same = new AnalysisCacheKey(1L, new boolean[] {true, false, true, false});
+    // identical class id + probe set must collide so the analysis is reused
+    assertEquals(key, same);
+    assertEquals(key.hashCode(), same.hashCode());
+  }
+
+  @Test
+  void cacheKeyDistinguishesClassesAndProbeSets() {
+    AnalysisCacheKey key = new AnalysisCacheKey(1L, new boolean[] {true, false, true});
+    // a different class or a different probe set must NOT hit the same cache entry
+    assertNotEquals(key, new AnalysisCacheKey(2L, new boolean[] {true, false, true}));
+    assertNotEquals(key, new AnalysisCacheKey(1L, new boolean[] {true, true, true}));
+  }
+
+  @Test
+  void cacheKeyIgnoresTrailingUnsetProbes() {
+    // The key bit-packs the activated probe set; trailing probes that never fire don't change
+    // coverage, so padding differences must not create distinct entries.
+    AnalysisCacheKey shortKey = new AnalysisCacheKey(1L, new boolean[] {true, false, true});
+    AnalysisCacheKey padded =
+        new AnalysisCacheKey(1L, new boolean[] {true, false, true, false, false});
+    assertEquals(shortKey, padded);
+    assertEquals(shortKey.hashCode(), padded.hashCode());
+  }
+}