Add stable session ID headers to telemetry requests

dd-java-agent/instrumentation/java/java-lang/java-lang-1.8/src/main/java/datadog/trace/instrumentation/java/lang/ProcessImplStartAdvice.java

@@ -1,13 +1,22 @@
 package datadog.trace.instrumentation.java.lang;
 
+import datadog.trace.api.Config;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.java.lang.ProcessImplInstrumentationHelpers;
+import java.util.Map;
 import net.bytebuddy.asm.Advice;
 
 class ProcessImplStartAdvice {
   @Advice.OnMethodEnter(suppress = Throwable.class)
-  public static AgentSpan beforeStart(@Advice.Argument(0) final String[] command) {
+  public static AgentSpan beforeStart(
+      @Advice.Argument(0) final String[] command,
+      @Advice.Argument(1) final Map<String, String> environment) {
+    String rootSessionId = Config.get().getRootSessionId();
+    if (rootSessionId != null && environment != null) {
+      environment.put("_DD_ROOT_JAVA_SESSION_ID", rootSessionId);
+    }
+
     if (!ProcessImplInstrumentationHelpers.ONLINE) {
       return null;
     }

dd-java-agent/instrumentation/java/java-lang/java-lang-1.8/src/test/groovy/datadog/trace/instrumentation/java/lang/ProcessImplInstrumentationSpecification.groovy

@@ -3,6 +3,7 @@ package datadog.trace.instrumentation.java.lang
 import datadog.trace.agent.test.InstrumentationSpecification
 import datadog.trace.agent.test.asserts.SpanAssert
 import datadog.trace.agent.test.utils.TraceUtils
+import datadog.trace.api.Config
 import datadog.trace.bootstrap.ActiveSubsystems
 
 import java.util.concurrent.TimeUnit
@@ -331,4 +332,18 @@ class ProcessImplInstrumentationSpecification extends InstrumentationSpecificati
     ['/does/not/exist/cmd', '--pass', 'abc', '--token=def'] | '["/does/not/exist/cmd","--pass","?","--token=?"]'
     ['/does/not/exist/md5', '-s', 'pony']                   | '["/does/not/exist/md5","?","?"]'
   }
+
+  void 'child process inherits root session ID'() {
+    when:
+    def command = ['/bin/sh', '-c', 'echo $_DD_ROOT_JAVA_SESSION_ID']
+    def builder = new ProcessBuilder(command)
+    builder.environment()
+    Process p = builder.start()
+    def output = p.inputStream.text.trim()
+    def terminated = p.waitFor(5, TimeUnit.SECONDS)
+
+    then:
+    terminated
+    output == Config.get().getRootSessionId()
+  }
 }

internal-api/src/main/java/datadog/trace/api/Config.java

@@ -792,6 +792,12 @@ public class Config {
    */
   static class RuntimeIdHolder {
     static final String runtimeId = RandomUtils.randomUUID().toString();
+    static final String rootSessionId = initRootSessionId();
+
+    private static String initRootSessionId() {
+      String inherited = ConfigHelper.env("_DD_ROOT_JAVA_SESSION_ID");
+      return inherited != null ? inherited : runtimeId;
+    }
   }
 
   static class HostNameHolder {
@@ -3123,6 +3129,10 @@ public String getRuntimeId() {
     return runtimeIdEnabled ? RuntimeIdHolder.runtimeId : "";
   }
 
+  public String getRootSessionId() {
+    return runtimeIdEnabled ? RuntimeIdHolder.rootSessionId : "";
+  }
+
   public Long getProcessId() {
     return PidHelper.getPidAsLong();
   }
@@ -5863,6 +5873,9 @@ public String toString() {
         + ", runtimeId='"
         + getRuntimeId()
         + '\''
+        + ", rootSessionId='"
+        + getRootSessionId()
+        + '\''
         + ", runtimeVersion='"
         + runtimeVersion
         + ", apiKey="

telemetry/src/main/java/datadog/telemetry/TelemetryRequest.java

@@ -80,6 +80,16 @@ public Request.Builder httpRequest() {
       builder.addHeader("DD-Telemetry-Debug-Enabled", "true");
     }
 
+    Config config = Config.get();
+    String sessionId = config.getRuntimeId();
+    if (sessionId != null && !sessionId.isEmpty()) {
+      builder.addHeader("DD-Session-ID", sessionId);
+    }
+    String rootSessionId = config.getRootSessionId();
+    if (rootSessionId != null && !rootSessionId.equals(sessionId)) {
+      builder.addHeader("DD-Root-Session-ID", rootSessionId);
+    }
+
     return builder;
   }
 

telemetry/src/test/groovy/datadog/telemetry/TestTelemetryRouter.groovy

@@ -7,6 +7,7 @@ import datadog.telemetry.api.DistributionSeries
 import datadog.telemetry.api.LogMessage
 import datadog.telemetry.api.Metric
 import datadog.telemetry.api.RequestType
+import datadog.trace.api.Config
 import datadog.trace.api.ConfigSetting
 import datadog.trace.api.telemetry.Endpoint
 import datadog.trace.api.telemetry.ProductChange
@@ -84,7 +85,8 @@ class TestTelemetryRouter extends TelemetryRouter {
         'DD-Client-Library-Language',
         'DD-Client-Library-Version',
         'DD-Telemetry-API-Version',
-        'DD-Telemetry-Request-Type'
+        'DD-Telemetry-Request-Type',
+        'DD-Session-ID'
       ])
       assert this.request.header('Content-Type') == 'application/json; charset=utf-8'
       assert this.request.header('Content-Length').toInteger() > 0
@@ -94,6 +96,17 @@ class TestTelemetryRouter extends TelemetryRouter {
       assert this.request.header('DD-Telemetry-Request-Type') == requestType.toString()
       def entityId = this.request.header('Datadog-Entity-ID')
       assert entityId == null || entityId.startsWith("in-") || entityId.startsWith("cin-")
+      def sessionId = this.request.header('DD-Session-ID')
+      assert sessionId =~ /[\da-f]{8}-([\da-f]{4}-){3}[\da-f]{12}/
+      assert sessionId == Config.get().getRuntimeId()
+      // DD-Root-Session-ID should only be present when inherited from a parent process
+      // (i.e., when rootSessionId != runtimeId). In normal test context, they're equal.
+      def rootSessionId = this.request.header('DD-Root-Session-ID')
+      if (Config.get().getRootSessionId() == Config.get().getRuntimeId()) {
+        assert rootSessionId == null
+      } else {
+        assert rootSessionId == Config.get().getRootSessionId()
+      }
       return this
     }