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

Author: daniel-mohedano

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

@@ -18,6 +18,10 @@ public String getClassName() {
     return className;
   }
 
+  boolean[] getProbeActivations() {
+    return probeActivations;
+  }
+
   void record(int probeId) {
     probeActivations[probeId] = true;
   }

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

@@ -13,6 +13,7 @@
 import datadog.trace.civisibility.source.SourcePathResolver;
 import datadog.trace.civisibility.source.Utils;
 import java.io.InputStream;
+import java.lang.reflect.Field;
 import java.util.ArrayList;
 import java.util.BitSet;
 import java.util.Collection;
@@ -63,6 +64,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;
     }
@@ -132,6 +137,26 @@ protected TestReport report(
     return report;
   }
 
+  private static void copyProbeDataToJacoco(Map<Class<?>, ExecutionDataAdapter> executionData) {
+    for (Map.Entry<Class<?>, ExecutionDataAdapter> e : executionData.entrySet()) {
+      Class<?> clazz = e.getKey();
+      boolean[] probeActivations = e.getValue().getProbeActivations();
+      try {
+        Field jacocoDataField = clazz.getDeclaredField("$jacocoData");
+        jacocoDataField.setAccessible(true);
+        boolean[] jacocoData = (boolean[]) jacocoDataField.get(null);
+        if (jacocoData != null) {
+          for (int i = 0; i < probeActivations.length && i < jacocoData.length; i++) {
+            jacocoData[i] |= probeActivations[i];
+          }
+        }
+      } catch (Exception ex) {
+        // Class may not have $jacocoData field (e.g. not instrumented by JaCoCo agent),
+        // or field may not be accessible — silently ignore
+      }
+    }
+  }
+
   public static final class Factory implements CoverageStore.Factory {
 
     private final Map<String, Integer> probeCounts = new ConcurrentHashMap<>();

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

@@ -58,6 +58,22 @@ public void record(Class<?> clazz, long classId, int probeId) {
     }
   }
 
+  @Override
+  public boolean[] resolveProbeArray(Class<?> clazz, long classId, int probeCount) {
+    try {
+      if (lastCoveredClass != clazz) {
+        lastCoveredExecutionData =
+            executionData.computeIfAbsent(
+                lastCoveredClass = clazz,
+                k -> new ExecutionDataAdapter(classId, k.getName(), probeCount));
+      }
+      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.length);
+    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, int probeCount) {
+    return null;
+  }
 }