[DNM] JFR Context Integration PoC

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/Agent.java

@@ -847,18 +847,20 @@ public void withTracer(TracerAPI tracer) {
    * on JFR.
    */
   private static ProfilingContextIntegration createProfilingContextIntegration() {
+    ProfilingContextIntegration ddprof = ProfilingContextIntegration.NoOp.INSTANCE;
     if (Config.get().isProfilingEnabled() && !Platform.isWindows()) {
       try {
-        return (ProfilingContextIntegration)
-            AGENT_CLASSLOADER
-                .loadClass("com.datadog.profiling.ddprof.DatadogProfilingIntegration")
-                .getDeclaredConstructor()
-                .newInstance();
+        ddprof =
+            (ProfilingContextIntegration)
+                AGENT_CLASSLOADER
+                    .loadClass("com.datadog.profiling.ddprof.DatadogProfilingIntegration")
+                    .getDeclaredConstructor()
+                    .newInstance();
       } catch (Throwable t) {
         log.debug("Profiling context labeling not available. {}", t.getMessage());
       }
     }
-    return ProfilingContextIntegration.NoOp.INSTANCE;
+    return ddprof;
   }
 
   private static void startProfilingAgent(final boolean isStartingFirst) {

dd-java-agent/agent-profiling/build.gradle

@@ -16,6 +16,7 @@ dependencies {
   api project(':dd-java-agent:agent-profiling:profiling-controller-ddprof')
   api project(':dd-java-agent:agent-profiling:profiling-controller-openjdk')
   api project(':dd-java-agent:agent-profiling:profiling-controller-oracle')
+  api project(':dd-java-agent:agent-profiling:profiling-jfr-context')
 }
 
 Project parent_project = project

dd-java-agent/agent-profiling/profiling-jfr-context/build.gradle

@@ -0,0 +1,20 @@
+apply from: "$rootDir/gradle/java.gradle"
+
+tasks.withType(JavaCompile) {
+  options.fork = true
+  options.forkOptions.executable = project.getProperties()['jfr_context.javac']
+}
+
+dependencies {
+  api deps.slf4j
+  api project(':internal-api')
+  api project(':dd-trace-core')
+
+  testImplementation deps.junit5
+  testImplementation deps.mockito
+  testImplementation group: 'org.hamcrest', name: 'hamcrest', version: '2.1'
+}
+
+forbiddenApisMain {
+  failOnMissingClasses = false
+}

dd-java-agent/agent-profiling/profiling-jfr-context/src/main/java/com/datadog/profiling/jfr/context/ContextIntegration.java

@@ -0,0 +1,23 @@
+package com.datadog.profiling.jfr.context;
+
+import datadog.trace.bootstrap.instrumentation.api.ProfilerContext;
+import datadog.trace.core.DDSpanContext;
+import jdk.jfr.ContextAccess;
+
+public final class ContextIntegration {
+
+  public static void initialize() {
+    DDSpanContext.profilingContextAccess = new ProfilerContext.Access() {
+      private final ContextAccess access = ContextAccess.forType(DDSpanContext.class);
+      @Override
+      public void set(ProfilerContext ctx) {
+        access.set(ctx);
+      }
+
+      @Override
+      public void unset() {
+        access.unset();
+      }
+    };
+  }
+}

dd-java-agent/agent-profiling/profiling-uploader/src/main/java/com/datadog/profiling/uploader/ProfileUploader.java

@@ -270,15 +270,15 @@ public void upload(
           public void onResponse(final Call call, final Response response) throws IOException {
             if (handled.compareAndSet(false, true)) {
               handleResponse(call, response, data, onCompletion);
-              latch.countDown();
+              //              latch.countDown();
             }
           }
 
           @Override
           public void onFailure(final Call call, final IOException e) {
             if (handled.compareAndSet(false, true)) {
               handleFailure(call, e, data, onCompletion);
-              latch.countDown();
+              //              latch.countDown();
             }
           }
         });

dd-java-agent/agent-profiling/src/main/java/com/datadog/profiling/agent/ProfilingAgent.java

@@ -112,6 +112,17 @@ public static synchronized void run(final boolean isStartingFirst, ClassLoader a
         return;
       }
 
+      // initialize the JFR context; this needs to be done before JFR is initialized
+      try {
+        ProfilingAgent.class
+            .getClassLoader()
+            .loadClass("com.datadog.profiling.jfr.context.ContextIntegration")
+            .getMethod("initialize")
+            .invoke(null);
+      } catch (Throwable t) {
+        log.debug("Profiling context labeling via JFR not available. {}", t.getMessage());
+      }
+
       try {
         final Controller controller = ControllerFactory.createController(configProvider);
 

dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java

@@ -35,6 +35,8 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
+
+import jdk.jfr.Name;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -46,6 +48,7 @@
  * across Span boundaries and (2) any Datadog fields that are needed to identify or contextualize
  * the associated Span instance
  */
+@Name("dd-context")
 public class DDSpanContext
     implements AgentSpan.Context, RequestContext, TraceSegment, ProfilerContext {
   private static final Logger log = LoggerFactory.getLogger(DDSpanContext.class);
@@ -137,6 +140,8 @@ public class DDSpanContext
   private boolean injectBaggageAsTags;
   private volatile int encodedOperationName;
 
+  public static ProfilerContext.Access profilingContextAccess = Access.NOOP;
+
   public DDSpanContext(
       final DDTraceId traceId,
       final long spanId,
@@ -366,11 +371,13 @@ public long getParentId() {
     return parentId;
   }
 
+  @Name("spanid")
   @Override
   public long getSpanId() {
     return spanId;
   }
 
+  @Name("rootspanid")
   @Override
   public long getRootSpanId() {
     return getRootSpanContextOrThis().spanId;
@@ -417,6 +424,7 @@ private boolean isResourceNameSet() {
     return resourceName != null && resourceName.length() != 0;
   }
 
+  @Name("operation_name")
   public CharSequence getOperationName() {
     return operationName;
   }

dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeStack.java

@@ -4,6 +4,8 @@
 import datadog.trace.bootstrap.instrumentation.api.ProfilerContext;
 import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
 import datadog.trace.bootstrap.instrumentation.api.ScopeSource;
+import datadog.trace.core.DDSpanContext;
+
 import java.util.ArrayDeque;
 
 /**
@@ -112,6 +114,9 @@ void clear() {
 
   private void onTopChanged(ContinuableScope top) {
     AgentSpan.Context context = top.span.context();
+    if (context instanceof DDSpanContext) {
+      DDSpanContext.profilingContextAccess.set((DDSpanContext)context);
+    }
     if (context instanceof ProfilerContext) {
       try {
         profilingContextIntegration.setContext((ProfilerContext) context);
@@ -133,6 +138,7 @@ private void onBecomeNonEmpty() {
   /** Notifies profiler that this thread no longer has a context */
   private void onBecomeEmpty() {
     try {
+      DDSpanContext.profilingContextAccess.unset();
       profilingContextIntegration.clearContext();
       profilingContextIntegration.onDetach();
     } catch (Throwable e) {

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ProfilerContext.java

@@ -1,6 +1,25 @@
 package datadog.trace.bootstrap.instrumentation.api;
 
+import datadog.trace.api.profiling.ProfilingContext;
+
 public interface ProfilerContext {
+  interface Access {
+    Access NOOP = new Access() {
+      @Override
+      public void set(ProfilerContext ctx) {}
+
+      @Override
+      public void unset() {}
+
+      @Override
+      public String toString() {
+        return "NOOP";
+      }
+    };
+
+    void set(ProfilerContext ctx);
+    void unset();
+  }
 
   long getSpanId();
 

settings.gradle

@@ -54,6 +54,7 @@ include ':dd-java-agent:agent-profiling:profiling-controller-jfr'
 include ':dd-java-agent:agent-profiling:profiling-controller-ddprof'
 include ':dd-java-agent:agent-profiling:profiling-controller-openjdk'
 include ':dd-java-agent:agent-profiling:profiling-controller-oracle'
+include ':dd-java-agent:agent-profiling:profiling-jfr-context'
 include ':dd-java-agent:agent-profiling:profiling-testing'
 include ':dd-java-agent:agent-profiling:profiling-uploader'
 include ':dd-java-agent:agent-profiling:profiling-utils'