Propagate module context from build system process to child JVM processes (#7710)

Author: nikita-tkachenko-datadog

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

@@ -31,7 +31,6 @@
 import datadog.trace.civisibility.source.SourcePathResolver;
 import datadog.trace.civisibility.source.index.RepoIndexProvider;
 import datadog.trace.civisibility.source.index.RepoIndexSourcePathResolver;
-import datadog.trace.civisibility.utils.ProcessHierarchyUtils;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.Map;
@@ -76,11 +75,12 @@ public class CiVisibilityRepoServices {
     codeowners = buildCodeowners(repoRoot);
     sourcePathResolver = buildSourcePathResolver(repoRoot, repoIndexProvider);
 
-    if (ProcessHierarchyUtils.isChild()) {
+    if (services.processHierarchy.isChild()) {
       executionSettingsFactory = buildExecutionSettingsFetcher(services.signalClientFactory);
     } else {
       executionSettingsFactory =
           buildExecutionSettingsFactory(
+              services.processHierarchy,
               services.config,
               services.metricCollector,
               services.backendApi,
@@ -122,6 +122,7 @@ private static ExecutionSettingsFactory buildExecutionSettingsFetcher(
   }
 
   private static ExecutionSettingsFactory buildExecutionSettingsFactory(
+      ProcessHierarchy processHierarchy,
       Config config,
       CiVisibilityMetricCollector metricCollector,
       BackendApi backendApi,
@@ -138,7 +139,7 @@ private static ExecutionSettingsFactory buildExecutionSettingsFactory(
 
     ExecutionSettingsFactoryImpl factory =
         new ExecutionSettingsFactoryImpl(config, configurationApi, gitDataUploader, repoRoot);
-    if (ProcessHierarchyUtils.isHeadless()) {
+    if (processHierarchy.isHeadless()) {
       return factory;
     } else {
       return new MultiModuleExecutionSettingsFactory(config, factory);

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

@@ -28,7 +28,6 @@
 import datadog.trace.civisibility.source.CompilerAidedMethodLinesResolver;
 import datadog.trace.civisibility.source.MethodLinesResolver;
 import datadog.trace.civisibility.source.index.*;
-import datadog.trace.civisibility.utils.ProcessHierarchyUtils;
 import java.lang.reflect.Type;
 import java.net.InetSocketAddress;
 import java.nio.file.FileSystem;
@@ -56,6 +55,7 @@ public class CiVisibilityServices {
 
   static final String DD_ENV_VARS_PROVIDER_KEY_HEADER = "DD-Env-Vars-Provider-Key";
 
+  final ProcessHierarchy processHierarchy;
   final Config config;
   final CiVisibilityMetricCollector metricCollector;
   final BackendApi backendApi;
@@ -72,6 +72,7 @@ public class CiVisibilityServices {
       CiVisibilityMetricCollector metricCollector,
       SharedCommunicationObjects sco,
       GitInfoProvider gitInfoProvider) {
+    this.processHierarchy = new ProcessHierarchy();
     this.config = config;
     this.metricCollector = metricCollector;
     this.backendApi =
@@ -91,8 +92,8 @@ public class CiVisibilityServices {
         new CILocalGitInfoBuilder(gitClientFactory, GIT_FOLDER_NAME));
     gitInfoProvider.registerGitInfoBuilder(new GitClientGitInfoBuilder(config, gitClientFactory));
 
-    if (ProcessHierarchyUtils.isChild()) {
-      InetSocketAddress signalServerAddress = ProcessHierarchyUtils.getSignalServerAddress();
+    if (processHierarchy.isChild()) {
+      InetSocketAddress signalServerAddress = processHierarchy.getSignalServerAddress();
       this.signalClientFactory = new SignalClient.Factory(signalServerAddress, config);
 
       RepoIndexProvider indexFetcher = new RepoIndexFetcher(signalClientFactory);

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

@@ -33,7 +33,6 @@
 import datadog.trace.civisibility.telemetry.CiVisibilityMetricCollectorImpl;
 import datadog.trace.civisibility.test.ExecutionStrategy;
 import datadog.trace.civisibility.utils.ConcurrentHashMapContextStore;
-import datadog.trace.civisibility.utils.ProcessHierarchyUtils;
 import datadog.trace.util.throwable.FatalAgentMisconfigurationError;
 import java.lang.instrument.Instrumentation;
 import java.nio.file.Path;
@@ -79,7 +78,7 @@ public static void start(Instrumentation inst, SharedCommunicationObjects sco) {
     InstrumentationBridge.registerBuildEventsHandlerFactory(buildEventsHandlerFactory(services));
     CIVisibility.registerSessionFactory(manualApiSessionFactory(services));
 
-    if (ProcessHierarchyUtils.isChild() || ProcessHierarchyUtils.isHeadless()) {
+    if (services.processHierarchy.isChild() || services.processHierarchy.isHeadless()) {
       CiVisibilityRepoServices repoServices = services.repoServices(getCurrentPath());
 
       ExecutionSettings executionSettings =
@@ -142,7 +141,7 @@ private TestEventsHandlerFactory(
         ExecutionSettings executionSettings) {
       this.services = services;
       this.repoServices = repoServices;
-      if (ProcessHierarchyUtils.isChild()) {
+      if (services.processHierarchy.isChild()) {
         sessionFactory =
             childTestFrameworkSessionFactory(
                 services, repoServices, coverageServices, executionSettings);
@@ -222,19 +221,16 @@ private static TestFrameworkSession.Factory childTestFrameworkSessionFactory(
       CiVisibilityCoverageServices.Child coverageServices,
       ExecutionSettings executionSettings) {
     return (String projectName, String component, Long startTime) -> {
-      long parentProcessSessionId = ProcessHierarchyUtils.getParentSessionId();
-      long parentProcessModuleId = ProcessHierarchyUtils.getParentModuleId();
-
       String sessionName = services.config.getCiVisibilitySessionName();
       String testCommand = services.config.getCiVisibilityTestCommand();
       TestDecorator testDecorator =
           new TestDecoratorImpl(component, sessionName, testCommand, repoServices.ciTags);
 
       ExecutionStrategy executionStrategy =
           new ExecutionStrategy(services.config, executionSettings);
+
       return new ProxyTestSession(
-          parentProcessSessionId,
-          parentProcessModuleId,
+          services.processHierarchy.parentProcessModuleContext,
           services.config,
           services.metricCollector,
           testDecorator,

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

@@ -0,0 +1,91 @@
+package datadog.trace.civisibility;
+
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
+
+import datadog.trace.api.config.CiVisibilityConfig;
+import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.util.Strings;
+import java.net.InetSocketAddress;
+import java.util.Properties;
+import javax.annotation.Nullable;
+
+public class ProcessHierarchy {
+
+  private static final class SystemPropertiesPropagationGetter
+      implements AgentPropagation.ContextVisitor<Properties> {
+    static final AgentPropagation.ContextVisitor<Properties> INSTANCE =
+        new SystemPropertiesPropagationGetter();
+
+    private SystemPropertiesPropagationGetter() {}
+
+    @Override
+    public void forEachKey(Properties carrier, AgentPropagation.KeyClassifier classifier) {
+      for (String propertyName : carrier.stringPropertyNames()) {
+        if (!classifier.accept(propertyName, carrier.getProperty(propertyName))) {
+          return;
+        }
+      }
+    }
+  }
+
+  @Nullable public final AgentSpan.Context.Extracted parentProcessModuleContext;
+
+  ProcessHierarchy() {
+    parentProcessModuleContext =
+        propagate().extract(System.getProperties(), SystemPropertiesPropagationGetter.INSTANCE);
+  }
+
+  /**
+   * Module span context is propagated from the parent process if it runs with the tracer attached.
+   * If module span context is note there, either we are in the build system, or we are in the tests
+   * JVM and the build system is not instrumented.
+   */
+  public boolean isChild() {
+    return parentProcessModuleContext != null;
+  }
+
+  /**
+   * Determines if current process runs in "headless mode", i.e. has no instrumented parent and is
+   * not one of the supported build system processes.
+   */
+  public boolean isHeadless() {
+    return !isChild() && !isParent();
+  }
+
+  private boolean isParent() {
+    return isMavenParent() || isGradleDaemon();
+  }
+
+  private boolean isMavenParent() {
+    return System.getProperty("maven.home") != null
+        && System.getProperty("classworlds.conf") != null;
+  }
+
+  private boolean isGradleDaemon() {
+    return ClassLoader.getSystemClassLoader()
+                .getResource("org/gradle/launcher/daemon/bootstrap/GradleDaemon.class")
+            != null
+        // double-check this is not a Gradle Worker
+        && System.getProperties().getProperty("org.gradle.internal.worker.tmpdir") == null;
+  }
+
+  @Nullable
+  public InetSocketAddress getSignalServerAddress() {
+    // System.getProperty is used rather than Config,
+    // because system variables can be set after config was initialized
+    String host =
+        System.getProperty(
+            Strings.propertyNameToSystemPropertyName(
+                CiVisibilityConfig.CIVISIBILITY_SIGNAL_SERVER_HOST));
+    String port =
+        System.getProperty(
+            Strings.propertyNameToSystemPropertyName(
+                CiVisibilityConfig.CIVISIBILITY_SIGNAL_SERVER_PORT));
+    if (host != null && port != null) {
+      return new InetSocketAddress(host, Integer.parseInt(port));
+    } else {
+      return null;
+    }
+  }
+}

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

@@ -1,5 +1,6 @@
 package datadog.trace.civisibility.coverage;
 
+import datadog.trace.api.DDTraceId;
 import datadog.trace.api.civisibility.coverage.CoverageProbes;
 import datadog.trace.api.civisibility.coverage.CoverageStore;
 import datadog.trace.api.civisibility.coverage.TestReport;
@@ -34,14 +35,14 @@ private T create(Thread thread) {
   }
 
   @Override
-  public boolean report(Long testSessionId, Long testSuiteId, long testSpanId) {
+  public boolean report(DDTraceId 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);
+      DDTraceId testSessionId, Long testSuiteId, long testSpanId, Collection<T> probes);
 
   @Nullable
   @Override

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

@@ -1,5 +1,6 @@
 package datadog.trace.civisibility.coverage.file;
 
+import datadog.trace.api.DDTraceId;
 import datadog.trace.api.civisibility.config.TestIdentifier;
 import datadog.trace.api.civisibility.coverage.CoverageStore;
 import datadog.trace.api.civisibility.coverage.TestReport;
@@ -45,7 +46,7 @@ private FileCoverageStore(
   @Nullable
   @Override
   protected TestReport report(
-      Long testSessionId, Long testSuiteId, long testSpanId, Collection<FileProbes> probes) {
+      DDTraceId testSessionId, Long testSuiteId, long testSpanId, Collection<FileProbes> probes) {
     try {
       Set<Class<?>> combinedClasses = Collections.newSetFromMap(new IdentityHashMap<>());
       Collection<String> combinedNonCodeResources = new HashSet<>();

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

@@ -1,5 +1,6 @@
 package datadog.trace.civisibility.coverage.line;
 
+import datadog.trace.api.DDTraceId;
 import datadog.trace.api.civisibility.config.TestIdentifier;
 import datadog.trace.api.civisibility.coverage.CoverageStore;
 import datadog.trace.api.civisibility.coverage.TestReport;
@@ -51,7 +52,7 @@ private LineCoverageStore(
   @Nullable
   @Override
   protected TestReport report(
-      Long testSessionId, Long testSuiteId, long testSpanId, Collection<LineProbes> probes) {
+      DDTraceId testSessionId, Long testSuiteId, long testSpanId, Collection<LineProbes> probes) {
     try {
       Map<Class<?>, ExecutionDataAdapter> combinedExecutionData = new IdentityHashMap<>();
       Collection<String> combinedNonCodeResources = new HashSet<>();

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/coverage/percentage/child/ChildProcessCoverageReporter.java

@@ -1,5 +1,6 @@
 package datadog.trace.civisibility.coverage.percentage.child;
 
+import datadog.trace.api.DDTraceId;
 import datadog.trace.civisibility.ipc.ModuleSignal;
 import javax.annotation.Nullable;
 
@@ -12,5 +13,5 @@
  */
 public interface ChildProcessCoverageReporter {
   @Nullable
-  ModuleSignal createCoverageSignal(long sessionId, long moduleId);
+  ModuleSignal createCoverageSignal(DDTraceId sessionId, long moduleId);
 }

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/coverage/percentage/child/JacocoChildProcessCoverageReporter.java

@@ -1,5 +1,6 @@
 package datadog.trace.civisibility.coverage.percentage.child;
 
+import datadog.trace.api.DDTraceId;
 import datadog.trace.civisibility.ipc.ModuleCoverageDataJacoco;
 import datadog.trace.civisibility.ipc.ModuleSignal;
 import java.util.function.Supplier;
@@ -15,7 +16,7 @@ public JacocoChildProcessCoverageReporter(Supplier<byte[]> coverageDataSupplier)
 
   @Nullable
   @Override
-  public ModuleSignal createCoverageSignal(long sessionId, long moduleId) {
+  public ModuleSignal createCoverageSignal(DDTraceId sessionId, long moduleId) {
     byte[] coverageData = coverageDataSupplier.get();
     if (coverageData != null) {
       return new ModuleCoverageDataJacoco(sessionId, moduleId, coverageData);

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

@@ -19,7 +19,6 @@
 public abstract class AbstractTestModule {
 
   protected final AgentSpan span;
-  protected final long sessionId;
   protected final String moduleName;
   protected final Config config;
   protected final CiVisibilityMetricCollector metricCollector;
@@ -31,7 +30,6 @@ public abstract class AbstractTestModule {
 
   public AbstractTestModule(
       AgentSpan.Context sessionSpanContext,
-      long sessionId,
       String moduleName,
       @Nullable Long startTime,
       InstrumentationType instrumentationType,
@@ -42,7 +40,6 @@ public AbstractTestModule(
       Codeowners codeowners,
       MethodLinesResolver methodLinesResolver,
       Consumer<AgentSpan> onSpanFinish) {
-    this.sessionId = sessionId;
     this.moduleName = moduleName;
     this.config = config;
     this.metricCollector = metricCollector;
@@ -65,7 +62,7 @@ public AbstractTestModule(
     span.setTag(Tags.TEST_MODULE, moduleName);
 
     span.setTag(Tags.TEST_MODULE_ID, span.getSpanId());
-    span.setTag(Tags.TEST_SESSION_ID, sessionId);
+    span.setTag(Tags.TEST_SESSION_ID, span.getTraceId());
 
     // setting status to skip initially,
     // as we do not know in advance whether the module will have any children