Optimize per-test code coverage

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

@@ -3,6 +3,8 @@
 import datadog.communication.BackendApi;
 import datadog.trace.api.Config;
 import datadog.trace.api.civisibility.config.ModuleExecutionSettings;
+import datadog.trace.api.civisibility.coverage.CoverageStore;
+import datadog.trace.api.civisibility.coverage.NoOpCoverageStore;
 import datadog.trace.api.civisibility.telemetry.CiVisibilityMetricCollector;
 import datadog.trace.api.civisibility.telemetry.tag.Provider;
 import datadog.trace.api.git.GitInfoProvider;
@@ -18,6 +20,8 @@
 import datadog.trace.civisibility.config.JvmInfo;
 import datadog.trace.civisibility.config.ModuleExecutionSettingsFactory;
 import datadog.trace.civisibility.config.ModuleExecutionSettingsFactoryImpl;
+import datadog.trace.civisibility.coverage.file.FileCoverageStore;
+import datadog.trace.civisibility.coverage.line.LineCoverageStore;
 import datadog.trace.civisibility.git.tree.GitClient;
 import datadog.trace.civisibility.git.tree.GitDataApi;
 import datadog.trace.civisibility.git.tree.GitDataUploader;
@@ -53,6 +57,7 @@ public class CiVisibilityRepoServices {
   final RepoIndexProvider repoIndexProvider;
   final Codeowners codeowners;
   final SourcePathResolver sourcePathResolver;
+  final CoverageStore.Factory coverageStoreFactory;
   final ModuleExecutionSettingsFactory moduleExecutionSettingsFactory;
 
   CiVisibilityRepoServices(CiVisibilityServices services, Path path) {
@@ -75,6 +80,8 @@ public class CiVisibilityRepoServices {
     repoIndexProvider = services.repoIndexProviderFactory.create(repoRoot);
     codeowners = buildCodeowners(repoRoot);
     sourcePathResolver = buildSourcePathResolver(repoRoot, repoIndexProvider);
+    coverageStoreFactory =
+        buildCoverageStoreFactory(services.config, services.metricCollector, sourcePathResolver);
 
     if (ProcessHierarchyUtils.isChild()) {
       moduleExecutionSettingsFactory =
@@ -107,6 +114,17 @@ private static String getModuleName(Config config, Path path, CIInfo ciInfo) {
     return config.getServiceName();
   }
 
+  private static CoverageStore.Factory buildCoverageStoreFactory(
+      Config config, CiVisibilityMetricCollector metrics, SourcePathResolver sourcePathResolver) {
+    if (!config.isCiVisibilityCodeCoverageEnabled()) {
+      return new NoOpCoverageStore.Factory();
+    }
+    if (!config.isCiVisibilityCoverageSegmentsEnabled()) {
+      return new FileCoverageStore.Factory(metrics, sourcePathResolver);
+    }
+    return new LineCoverageStore.Factory(metrics, sourcePathResolver);
+  }
+
   private static ModuleExecutionSettingsFactory buildModuleExecutionSettingsFetcher(
       SignalClient.Factory signalClientFactory) {
     return (JvmInfo jvmInfo, String moduleName) -> {

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

@@ -10,10 +10,6 @@
 import datadog.trace.civisibility.config.CachingJvmInfoFactory;
 import datadog.trace.civisibility.config.JvmInfoFactory;
 import datadog.trace.civisibility.config.JvmInfoFactoryImpl;
-import datadog.trace.civisibility.coverage.CoverageProbeStoreFactory;
-import datadog.trace.civisibility.coverage.NoopCoverageProbeStore;
-import datadog.trace.civisibility.coverage.SegmentlessTestProbes;
-import datadog.trace.civisibility.coverage.TestProbes;
 import datadog.trace.civisibility.git.CILocalGitInfoBuilder;
 import datadog.trace.civisibility.git.CIProviderGitInfoBuilder;
 import datadog.trace.civisibility.git.GitClientGitInfoBuilder;
@@ -50,7 +46,6 @@ public class CiVisibilityServices {
   final GitClient.Factory gitClientFactory;
   final GitInfoProvider gitInfoProvider;
   final MethodLinesResolver methodLinesResolver;
-  final CoverageProbeStoreFactory coverageProbeStoreFactory;
   final RepoIndexProvider.Factory repoIndexProviderFactory;
   @Nullable final SignalClient.Factory signalClientFactory;
 
@@ -69,7 +64,6 @@ public class CiVisibilityServices {
     this.methodLinesResolver =
         new BestEffortMethodLinesResolver(
             new CompilerAidedMethodLinesResolver(), new ByteCodeMethodLinesResolver());
-    this.coverageProbeStoreFactory = buildTestProbesFactory(config, metricCollector);
 
     this.gitInfoProvider = gitInfoProvider;
     gitInfoProvider.registerGitInfoBuilder(new CIProviderGitInfoBuilder());
@@ -97,17 +91,6 @@ public class CiVisibilityServices {
     }
   }
 
-  private static CoverageProbeStoreFactory buildTestProbesFactory(
-      Config config, CiVisibilityMetricCollector metricCollector) {
-    if (!config.isCiVisibilityCodeCoverageEnabled()) {
-      return new NoopCoverageProbeStore.NoopCoverageProbeStoreFactory();
-    }
-    if (!config.isCiVisibilityCoverageSegmentsEnabled()) {
-      return new SegmentlessTestProbes.SegmentlessTestProbesFactory(metricCollector);
-    }
-    return new TestProbes.TestProbesFactory(metricCollector);
-  }
-
   CiVisibilityRepoServices repoServices(Path path) {
     return new CiVisibilityRepoServices(this, path);
   }

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

@@ -100,7 +100,7 @@ public static void start(Instrumentation inst, SharedCommunicationObjects sco) {
 
       InstrumentationBridge.registerTestEventsHandlerFactory(
           new TestEventsHandlerFactory(services, repoServices, executionSettings));
-      CoverageBridge.registerCoverageProbeStoreRegistry(services.coverageProbeStoreFactory);
+      CoverageBridge.registerCoverageStoreRegistry(repoServices.coverageStoreFactory);
     }
   }
 
@@ -218,7 +218,7 @@ private static BuildSystemSession.Factory buildSystemSessionFactory(
           repoServices.codeowners,
           services.methodLinesResolver,
           repoServices.moduleExecutionSettingsFactory,
-          services.coverageProbeStoreFactory,
+          repoServices.coverageStoreFactory,
           signalServer,
           repoServices.repoIndexProvider);
     };
@@ -243,7 +243,7 @@ private static TestFrameworkSession.Factory childTestFrameworkSessionFactory(
           repoServices.sourcePathResolver,
           repoServices.codeowners,
           services.methodLinesResolver,
-          services.coverageProbeStoreFactory,
+          repoServices.coverageStoreFactory,
           coverageDataSupplier,
           services.signalClientFactory,
           moduleExecutionSettings);
@@ -268,7 +268,7 @@ private static TestFrameworkSession.Factory headlessTestFrameworkEssionFactory(
           repoServices.sourcePathResolver,
           repoServices.codeowners,
           services.methodLinesResolver,
-          services.coverageProbeStoreFactory,
+          repoServices.coverageStoreFactory,
           moduleExecutionSettings);
     };
   }
@@ -288,7 +288,7 @@ private static CIVisibility.SessionFactory manualApiSessionFactory(
           repoServices.sourcePathResolver,
           repoServices.codeowners,
           services.methodLinesResolver,
-          services.coverageProbeStoreFactory);
+          repoServices.coverageStoreFactory);
     };
   }
 }

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

@@ -0,0 +1,49 @@
+package datadog.trace.civisibility.coverage;
+
+import datadog.trace.api.civisibility.coverage.CoverageProbes;
+import datadog.trace.api.civisibility.coverage.CoverageStore;
+import datadog.trace.api.civisibility.coverage.TestReport;
+import java.util.Collection;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Supplier;
+import org.jetbrains.annotations.Nullable;
+
+/** A store that keeps track of coverage probes allocated for multiple threads. */
+public abstract class ConcurrentCoverageStore<T extends CoverageProbes> implements CoverageStore {
+
+  private final Supplier<T> probesFactory;
+  private final Map<Thread, T> probes;
+
+  private volatile TestReport report;
+
+  protected ConcurrentCoverageStore(Supplier<T> probesFactory) {
+    this.probesFactory = probesFactory;
+    this.probes = new ConcurrentHashMap<>();
+  }
+
+  @Override
+  public CoverageProbes getProbes() {
+    return probes.computeIfAbsent(Thread.currentThread(), this::create);
+  }
+
+  private T create(Thread thread) {
+    return probesFactory.get();
+  }
+
+  @Override
+  public boolean report(Long testSessionId, Long testSuiteId, long testSpanId) {
+    report = report(testSessionId, testSuiteId, testSpanId, probes.values());
+    return report != null && report.isNotEmpty();
+  }
+
+  @Nullable
+  protected abstract TestReport report(
+      Long testSessionId, Long testSuiteId, long testSpanId, Collection<T> probes);
+
+  @Nullable
+  @Override
+  public TestReport getReport() {
+    return report;
+  }
+}