feat: improvements to line level granularity in per-test coverage

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

@@ -7,17 +7,34 @@ public class ExecutionDataAdapter {
   private final String className;
   // Unbounded data structure that only exists within a single test span
   private final boolean[] probeActivations;
+  private boolean[] jacocoArray;
 
   public ExecutionDataAdapter(long classId, String className, int totalProbeCount) {
     this.classId = classId;
     this.className = className;
     this.probeActivations = new boolean[totalProbeCount];
   }
 
+  long getClassId() {
+    return classId;
+  }
+
   public String getClassName() {
     return className;
   }
 
+  boolean[] getProbeActivations() {
+    return probeActivations;
+  }
+
+  void setJacocoArray(boolean[] jacocoArray) {
+    this.jacocoArray = jacocoArray;
+  }
+
+  boolean[] getJacocoArray() {
+    return jacocoArray;
+  }
+
   void record(int probeId) {
     probeActivations[probeId] = true;
   }

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

@@ -12,6 +12,8 @@
 import datadog.trace.civisibility.coverage.ConcurrentCoverageStore;
 import datadog.trace.civisibility.source.SourcePathResolver;
 import datadog.trace.civisibility.source.Utils;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
 import java.io.InputStream;
 import java.util.ArrayList;
 import java.util.BitSet;
@@ -25,6 +27,7 @@
 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;
@@ -39,14 +42,17 @@ public class LineCoverageStore extends ConcurrentCoverageStore<LineProbes> {
 
   private final CiVisibilityMetricCollector metrics;
   private final SourcePathResolver sourcePathResolver;
+  private final ConcurrentHashMap<Long, BitSet[]> probeLineMappingCache;
 
   private LineCoverageStore(
       Function<Boolean, LineProbes> probesFactory,
       CiVisibilityMetricCollector metrics,
-      SourcePathResolver sourcePathResolver) {
+      SourcePathResolver sourcePathResolver,
+      ConcurrentHashMap<Long, BitSet[]> probeLineMappingCache) {
     super(probesFactory);
     this.metrics = metrics;
     this.sourcePathResolver = sourcePathResolver;
+    this.probeLineMappingCache = probeLineMappingCache;
   }
 
   @Nullable
@@ -63,6 +69,10 @@ protected TestReport report(
       combinedNonCodeResources.addAll(probe.getNonCodeResources());
     }
 
+    // Copy per-test probe data back into JaCoCo's shared $jacocoData arrays so that
+    // JaCoCo's aggregate coverage reports remain accurate.
+    copyProbeDataToJacoco(combinedExecutionData);
+
     if (combinedExecutionData.isEmpty() && combinedNonCodeResources.isEmpty()) {
       return null;
     }
@@ -83,16 +93,33 @@ protected TestReport report(
       }
       String sourcePath = sourcePaths.iterator().next();
 
-      try (InputStream is = Utils.getClassStream(clazz)) {
+      try {
+        long classId = executionDataAdapter.getClassId();
+        boolean[] probeActivations = executionDataAdapter.getProbeActivations();
+
+        BitSet[] probeToLines =
+            probeLineMappingCache.computeIfAbsent(
+                classId,
+                id -> {
+                  try (InputStream is = Utils.getClassStream(clazz)) {
+                    if (is == null) {
+                      return new BitSet[0];
+                    }
+                    byte[] classBytes = readAllBytes(is);
+                    return buildProbeLineMapping(
+                        id, className, probeActivations.length, classBytes);
+                  } catch (Exception ex) {
+                    return new BitSet[0];
+                  }
+                });
+
         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);
+        for (int i = 0; i < probeActivations.length && i < probeToLines.length; i++) {
+          if (probeActivations[i]) {
+            coveredLines.or(probeToLines[i]);
+          }
+        }
 
       } catch (Exception exception) {
         log.debug(
@@ -132,9 +159,53 @@ protected TestReport report(
     return report;
   }
 
+  private static byte[] readAllBytes(InputStream is) throws IOException {
+    ByteArrayOutputStream buffer = new ByteArrayOutputStream();
+    byte[] data = new byte[4096];
+    int len;
+    while ((len = is.read(data, 0, data.length)) != -1) {
+      buffer.write(data, 0, len);
+    }
+    return buffer.toByteArray();
+  }
+
+  private static BitSet[] buildProbeLineMapping(
+      long classId, String className, int probeCount, byte[] classBytes) {
+    BitSet[] mapping = new BitSet[probeCount];
+    for (int i = 0; i < probeCount; i++) {
+      boolean[] singleProbe = new boolean[probeCount];
+      singleProbe[i] = true;
+      ExecutionDataStore store = new ExecutionDataStore();
+      store.put(new ExecutionData(classId, className, singleProbe));
+      BitSet lines = new BitSet();
+      Analyzer analyzer = new Analyzer(store, new SourceAnalyzer(lines));
+      try {
+        analyzer.analyzeClass(classBytes, null);
+      } catch (Exception e) {
+        // Analysis failure — empty BitSet is safe
+      }
+      mapping[i] = lines;
+    }
+    return mapping;
+  }
+
+  private static void copyProbeDataToJacoco(Map<Class<?>, ExecutionDataAdapter> executionData) {
+    for (ExecutionDataAdapter adapter : executionData.values()) {
+      boolean[] probeActivations = adapter.getProbeActivations();
+      boolean[] jacocoData = adapter.getJacocoArray();
+      if (jacocoData != null) {
+        for (int i = 0; i < probeActivations.length && i < jacocoData.length; i++) {
+          jacocoData[i] |= probeActivations[i];
+        }
+      }
+    }
+  }
+
   public static final class Factory implements CoverageStore.Factory {
 
     private final Map<String, Integer> probeCounts = new ConcurrentHashMap<>();
+    private final ConcurrentHashMap<Long, BitSet[]> probeLineMappingCache =
+        new ConcurrentHashMap<>();
 
     private final CiVisibilityMetricCollector metrics;
     private final SourcePathResolver sourcePathResolver;
@@ -146,7 +217,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, probeLineMappingCache);
     }
 
     private LineProbes createProbes(boolean isTestThread) {

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

@@ -58,6 +58,27 @@ public void record(Class<?> clazz, long classId, int probeId) {
     }
   }
 
+  @Override
+  public boolean[] resolveProbeArray(Class<?> clazz, long classId, boolean[] jacocoArray) {
+    try {
+      if (lastCoveredClass != clazz) {
+        lastCoveredExecutionData =
+            executionData.computeIfAbsent(
+                lastCoveredClass = clazz,
+                k -> {
+                  ExecutionDataAdapter adapter =
+                      new ExecutionDataAdapter(classId, k.getName(), jacocoArray.length);
+                  adapter.setJacocoArray(jacocoArray);
+                  return adapter;
+                });
+      }
+      return lastCoveredExecutionData.getProbeActivations();
+    } catch (Exception e) {
+      metrics.add(CiVisibilityCountMetric.CODE_COVERAGE_ERRORS, 1, CoverageErrorType.RECORD);
+      return null;
+    }
+  }
+
   @Override
   public void recordNonCodeResource(String absolutePath) {
     nonCodeResources.put(absolutePath, absolutePath);

dd-java-agent/instrumentation/jacoco-0.8.9/src/main/java/datadog/trace/instrumentation/jacoco/MethodVisitorWrapper.java

@@ -18,6 +18,7 @@ public class MethodVisitorWrapper {
   private static final MethodHandle visitInsnHandle;
   private static final MethodHandle visitIntInsnHandle;
   private static final MethodHandle visitLdcInsnHandle;
+  private static final MethodHandle visitVarInsnHandle;
   private static final MethodHandle getTypeHandle;
 
   static {
@@ -45,6 +46,8 @@ public class MethodVisitorWrapper {
         accessMethod(lookup, shadedMethodVisitorClass, "visitIntInsn", int.class, int.class);
     visitLdcInsnHandle =
         accessMethod(lookup, shadedMethodVisitorClass, "visitLdcInsn", Object.class);
+    visitVarInsnHandle =
+        accessMethod(lookup, shadedMethodVisitorClass, "visitVarInsn", int.class, int.class);
 
     Class<?> shadedTypeClass = getJacocoClass(jacocoClassLoader, jacocoPackageName, ".asm.Type");
     getTypeHandle = accessMethod(lookup, shadedTypeClass, "getType", String.class);
@@ -119,6 +122,10 @@ public void push(int value) throws Throwable {
     }
   }
 
+  public void visitVarInsn(int opcode, int var) throws Throwable {
+    visitVarInsnHandle.invoke(mv, opcode, var);
+  }
+
   public void pushClass(String className) throws Throwable {
     Object clazz = getTypeHandle.invoke('L' + className + ';');
     visitLdcInsnHandle.invoke(mv, clazz);

dd-java-agent/instrumentation/jacoco-0.8.9/src/main/java/datadog/trace/instrumentation/jacoco/ProbeInserterInstrumentation.java

@@ -93,7 +93,13 @@ private ElementMatcher<FieldDescription> methodVisitor() {
                         declaresMethod(
                             named("visitLdcInsn")
                                 .and(takesArguments(1))
-                                .and(takesArgument(0, Object.class))))));
+                                .and(takesArgument(0, Object.class))))
+                    .and(
+                        declaresMethod(
+                            named("visitVarInsn")
+                                .and(takesArguments(2))
+                                .and(takesArgument(0, int.class))
+                                .and(takesArgument(1, int.class))))));
   }
 
   @Override
@@ -112,29 +118,17 @@ public ElementMatcher<TypeDescription> hierarchyMatcher() {
   @Override
   public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(
-        isMethod().and(named("visitMaxs")).and(takesArguments(2)).and(takesArgument(0, int.class)),
-        getClass().getName() + "$VisitMaxsAdvice");
-    transformer.applyAdvice(
-        isMethod()
-            .and(named("insertProbe"))
-            .and(takesArguments(1))
-            .and(takesArgument(0, int.class)),
-        getClass().getName() + "$InsertProbeAdvice");
+        isMethod().and(named("visitCode")).and(takesArguments(0)),
+        getClass().getName() + "$VisitCodeAdvice");
   }
 
-  public static class VisitMaxsAdvice {
-    @Advice.OnMethodEnter(suppress = Throwable.class)
-    static void enter(@Advice.Argument(value = 0, readOnly = false) int maxStack) {
-      maxStack = maxStack + 2;
-    }
-  }
-
-  public static class InsertProbeAdvice {
+  public static class VisitCodeAdvice {
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     static void exit(
         @Advice.FieldValue(value = "mv") final Object mv,
         @Advice.FieldValue(value = "arrayStrategy") final Object arrayStrategy,
-        @Advice.Argument(0) final int id)
+        @Advice.FieldValue(value = "variable") final int variable,
+        @Advice.FieldValue(value = "accessorStackSize", readOnly = false) int accessorStackSize)
         throws Throwable {
       Field classNameField = arrayStrategy.getClass().getDeclaredField("className");
       classNameField.setAccessible(true);
@@ -163,20 +157,28 @@ static void exit(
 
       Field classIdField = arrayStrategy.getClass().getDeclaredField("classId");
       classIdField.setAccessible(true);
-      Long classId = classIdField.getLong(arrayStrategy);
+      long classId = classIdField.getLong(arrayStrategy);
 
       MethodVisitorWrapper methodVisitor = MethodVisitorWrapper.wrap(mv);
 
+      // ALOAD variable — load JaCoCo's shared boolean[] array
+      methodVisitor.visitVarInsn(Opcodes.ALOAD, variable);
+      // LDC ClassName.class — push Class reference
       methodVisitor.pushClass(className);
+      // LDC classId — push long classId
       methodVisitor.visitLdcInsn(classId);
-      methodVisitor.push(id);
-
+      // INVOKESTATIC resolveProbeArray(boolean[], Class, long) → boolean[]
       methodVisitor.visitMethodInsn(
           Opcodes.INVOKESTATIC,
           "datadog/trace/api/civisibility/coverage/CoveragePerTestBridge",
-          "recordCoverage",
-          "(Ljava/lang/Class;JI)V",
+          "resolveProbeArray",
+          "([ZLjava/lang/Class;J)[Z",
           false);
+      // ASTORE variable — replace local variable with per-test array
+      methodVisitor.visitVarInsn(Opcodes.ASTORE, variable);
+
+      // Ensure enough stack space for our bytecodes (1 ref + 1 ref + 1 long = 4 slots)
+      accessorStackSize = Math.max(accessorStackSize, 4);
     }
   }
 }

internal-api/src/main/java/datadog/trace/api/civisibility/coverage/CoveragePerTestBridge.java

@@ -66,6 +66,13 @@ private TotalProbeCount(String className, int count) {
     }
   }
 
+  /* This method is referenced by name in bytecode added in jacoco instrumentation module (see datadog.trace.instrumentation.jacoco.ProbeInserterInstrumentation.VisitCodeAdvice) */
+  public static boolean[] resolveProbeArray(boolean[] jacocoArray, Class<?> clazz, long classId) {
+    CoverageProbes probes = getCurrentCoverageProbes();
+    boolean[] result = probes.resolveProbeArray(clazz, classId, jacocoArray);
+    return result != null ? result : jacocoArray;
+  }
+
   /* This method is referenced by name in bytecode added in jacoco instrumentation module (see datadog.trace.instrumentation.jacoco.ProbeInserterInstrumentation.InsertProbeAdvice) */
   public static void recordCoverage(Class<?> clazz, long classId, int probeId) {
     getCurrentCoverageProbes().record(clazz, classId, probeId);

internal-api/src/main/java/datadog/trace/api/civisibility/coverage/CoverageProbes.java

@@ -6,4 +6,8 @@ public interface CoverageProbes {
   void record(Class<?> clazz, long classId, int probeId);
 
   void recordNonCodeResource(String absolutePath);
+
+  default boolean[] resolveProbeArray(Class<?> clazz, long classId, boolean[] jacocoArray) {
+    return null;
+  }
 }