Support W3C baggage propagation across various async instrumentation

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

@@ -234,6 +234,8 @@ public static void start(
 
     createAgentClassloader(agentJarURL);
 
+    AgentTracer.maybeInstallLegacyContextManager();
+
     if (Platform.isNativeImageBuilder()) {
       // these default services are not used during native-image builds
       remoteConfigEnabled = false;
@@ -338,8 +340,6 @@ public static void start(
       StaticEventLogger.end("crashtracking");
     }
 
-    AgentTracer.maybeInstallLegacyContextManager();
-
     startDatadogAgent(initTelemetry, inst);
 
     final EnumSet<Library> libraries = detectLibraries(log);

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/AdviceUtils.java

@@ -1,8 +1,8 @@
 package datadog.trace.bootstrap.instrumentation.java.concurrent;
 
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.isAsyncPropagationEnabled;
 
+import datadog.context.Context;
 import datadog.context.ContextContinuation;
 import datadog.context.ContextScope;
 import datadog.trace.bootstrap.ContextStore;
@@ -51,15 +51,24 @@ public static <T> void cancelTask(ContextStore<T, State> contextStore, final T t
     }
   }
 
+  public static boolean shouldCapture(Context context) {
+    if (context == Context.root()) {
+      return false;
+    }
+    AgentSpan span = AgentSpan.fromContext(context);
+    // propagate contexts with no span or a valid span, when flag is on
+    return (span == null || span.isValid()) && isAsyncPropagationEnabled();
+  }
+
   public static <T> void capture(ContextStore<T, State> contextStore, T task) {
-    AgentSpan span = activeSpan();
-    if (span != null && span.isValid() && isAsyncPropagationEnabled()) {
+    Context context = Context.current();
+    if (shouldCapture(context)) {
       State state = contextStore.get(task);
       if (null == state) {
         state = State.FACTORY.create();
         contextStore.put(task, state);
       }
-      state.captureAndSetContinuation(span);
+      state.captureAndSetContinuation(context);
     }
   }
 }

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ConcurrentState.java

@@ -1,13 +1,12 @@
 package datadog.trace.bootstrap.instrumentation.java.concurrent;
 
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.captureSpan;
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.isAsyncPropagationEnabled;
+import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.shouldCapture;
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.ContinuationClaim.CLAIMED;
 
+import datadog.context.Context;
 import datadog.context.ContextContinuation;
 import datadog.context.ContextScope;
 import datadog.trace.bootstrap.ContextStore;
-import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -33,16 +32,19 @@ public final class ConcurrentState {
   private ConcurrentState() {}
 
   public static <K> ConcurrentState captureContinuation(
-      ContextStore<K, ConcurrentState> contextStore, K key, AgentSpan span) {
-    if (span == null || !span.isValid() || !isAsyncPropagationEnabled()) {
+      ContextStore<K, ConcurrentState> contextStore, K key, Context context) {
+    if (shouldCapture(context)) {
+      final ConcurrentState state = contextStore.putIfAbsent(key, FACTORY);
+      if (!state.captureAndSetContinuation(context) && log.isDebugEnabled()) {
+        log.debug(
+            "continuation was already set for {} in context {}, no continuation captured.",
+            key,
+            context);
+      }
+      return state;
+    } else {
       return null;
     }
-    final ConcurrentState state = contextStore.putIfAbsent(key, FACTORY);
-    if (!state.captureAndSetContinuation(span) && log.isDebugEnabled()) {
-      log.debug(
-          "continuation was already set for {} in span {}, no continuation captured.", key, span);
-    }
-    return state;
   }
 
   public static <K> ContextScope activateAndContinueContinuation(
@@ -83,10 +85,10 @@ public static <K> void cancelAndClearContinuation(
     state.cancelAndClearContinuation();
   }
 
-  private boolean captureAndSetContinuation(final AgentSpan span) {
+  private boolean captureAndSetContinuation(final Context context) {
     if (CONTINUATION.compareAndSet(this, null, CLAIMED)) {
       // lazy write is guaranteed to be seen by getAndSet
-      CONTINUATION.lazySet(this, captureSpan(span).hold());
+      CONTINUATION.lazySet(this, context.capture().hold());
       return true;
     }
     return false;

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ExecutorInstrumentationUtils.java

@@ -1,10 +1,10 @@
 package datadog.trace.bootstrap.instrumentation.java.concurrent;
 
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.isAsyncPropagationEnabled;
+import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.shouldCapture;
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType;
 
+import datadog.context.Context;
 import datadog.trace.bootstrap.ContextStore;
-import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -17,38 +17,37 @@ public final class ExecutorInstrumentationUtils {
    * Checks if given task should get state attached.
    *
    * @param task task object
-   * @param span active span
+   * @param context active context
    * @return true iff given task object should be wrapped
    */
-  public static boolean shouldAttachStateToTask(final Object task, final AgentSpan span) {
+  public static boolean shouldAttachStateToTask(final Object task, final Context context) {
     if (task == null) {
       return false;
     }
     if (ExcludeFilter.exclude(ExcludeType.EXECUTOR, task)) {
       return false;
     }
-    return span != null && span.isValid() && isAsyncPropagationEnabled();
+    return shouldCapture(context);
   }
 
   /**
-   * Create task state given current span.
+   * Create task state given current context.
    *
    * @param contextStore context storage
    * @param task task instance
-   * @param span current span
+   * @param context current context
    * @param <T> task class type
    * @return new state
    */
   public static <T> State setupState(
-      final ContextStore<T, State> contextStore, final T task, final AgentSpan span) {
-
+      final ContextStore<T, State> contextStore, final T task, final Context context) {
     final State state = contextStore.putIfAbsent(task, State.FACTORY);
-
-    if (!state.captureAndSetContinuation(span)) {
+    if (!state.captureAndSetContinuation(context)) {
       log.debug(
-          "continuation was already set for {} in span {}, no continuation captured.", task, span);
+          "continuation was already set for {} in context {}, no continuation captured.",
+          task,
+          context);
     }
-
     return state;
   }
 

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/State.java

@@ -1,13 +1,13 @@
 package datadog.trace.bootstrap.instrumentation.java.concurrent;
 
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.captureSpan;
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.ContinuationClaim.CLAIMED;
 
+import datadog.context.Context;
 import datadog.context.ContextContinuation;
 import datadog.trace.api.profiling.Timing;
 import datadog.trace.bootstrap.ContextStore;
-import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
+import javax.annotation.Nullable;
 
 public final class State {
 
@@ -26,15 +26,15 @@ public final class State {
 
   private State() {}
 
-  public boolean captureAndSetContinuation(final AgentSpan span) {
+  public boolean captureAndSetContinuation(final Context context) {
     if (CONTINUATION.compareAndSet(this, null, CLAIMED)) {
       // it's a real pain to do this twice, and this can actually
       // happen systematically - WITHOUT RACES - because of broken
       // instrumentation, e.g. SetExecuteRunnableStateAdvice
       // "double instruments" calls to ScheduledExecutorService.submit/schedule
       //
       // lazy write is guaranteed to be seen by getAndSet
-      CONTINUATION.lazySet(this, captureSpan(span));
+      CONTINUATION.lazySet(this, context.capture());
       return true;
     }
     return false;
@@ -58,14 +58,15 @@ public void closeContinuation() {
     }
   }
 
-  public AgentSpan getSpan() {
+  public Context getContext() {
     ContextContinuation continuation = CONTINUATION.get(this);
     if (null == continuation || CLAIMED == continuation) {
-      return null;
+      return Context.root();
     }
-    return AgentSpan.fromContext(continuation.context());
+    return continuation.context();
   }
 
+  @Nullable
   public ContextContinuation getAndResetContinuation() {
     ContextContinuation continuation = CONTINUATION.get(this);
     if (null == continuation || CLAIMED == continuation) {

dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/nativeimage/TracerActivation.java

@@ -8,7 +8,6 @@
 import datadog.trace.agent.tooling.MeterInstaller;
 import datadog.trace.agent.tooling.ProfilerInstaller;
 import datadog.trace.agent.tooling.TracerInstaller;
-import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -19,7 +18,6 @@ public final class TracerActivation {
 
   public static void activate() {
     try {
-      AgentTracer.maybeInstallLegacyContextManager();
       // Initialize meter
       MeterInstaller.installMeter();
       // Initialize tracer

dd-java-agent/instrumentation/graal/graal-native-image-20.0/src/main/java/datadog/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java

@@ -69,6 +69,8 @@ public static void onEnter(@Advice.Argument(value = 0, readOnly = false) String[
               + "com.datadoghq.profiler.BufferWriter:run_time,"
               + "com.datadoghq.profiler.BufferWriter8:run_time,"
               + "com.datadoghq.profiler.BufferWriter9:run_time,"
+              + "datadog.context.ContextKey:build_time,"
+              + "datadog.context.ContextProviders$ProvidedManager:build_time,"
               + "datadog.environment.JavaVirtualMachine:rerun,"
               + "datadog.environment.OperatingSystem:rerun,"
               + "datadog.environment.OperatingSystem$Architecture:rerun,"
@@ -134,6 +136,8 @@ public static void onEnter(@Advice.Argument(value = 0, readOnly = false) String[
               + "datadog.trace.bootstrap.benchmark.StaticEventLogger:build_time,"
               + "datadog.trace.bootstrap.blocking.BlockingExceptionHandler:build_time,"
               + "datadog.trace.bootstrap.InstrumentationErrors:build_time,"
+              + "datadog.trace.bootstrap.instrumentation.api.AgentTracer$LegacyContextManager:build_time,"
+              + "datadog.trace.bootstrap.instrumentation.api.InternalContextKeys:build_time,"
               + "datadog.trace.bootstrap.instrumentation.java.concurrent.AsyncResultExtensions$CompletableAsyncResultExtension:build_time,"
               + "datadog.trace.bootstrap.instrumentation.java.concurrent.AsyncResultExtensions$1:build_time,"
               + "datadog.trace.bootstrap.instrumentation.java.concurrent.AsyncResultExtensions:build_time,"

dd-java-agent/instrumentation/guava-10.0/src/main/java/datadog/trace/instrumentation/guava10/ListenableFutureInstrumentation.java

@@ -1,17 +1,18 @@
 package datadog.trace.instrumentation.guava10;
 
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
+import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.currentContext;
+import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.rootContext;
 import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import com.google.auto.service.AutoService;
 import com.google.common.util.concurrent.AbstractFuture;
+import datadog.context.Context;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.bootstrap.InstrumentationContext;
-import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.ExecutorInstrumentationUtils;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.RunnableWrapper;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
@@ -55,16 +56,16 @@ public static class AddListenerAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static State addListenerEnter(
         @Advice.Argument(value = 0, readOnly = false) Runnable task) {
-      final AgentSpan span = activeSpan();
-      if (null != span) {
+      final Context context = currentContext();
+      if (context != rootContext()) {
         final Runnable newTask = RunnableWrapper.wrapIfNeeded(task);
         // It is important to check potentially wrapped task if we can instrument task in this
         // executor. Some executors do not support wrapped tasks.
-        if (ExecutorInstrumentationUtils.shouldAttachStateToTask(newTask, span)) {
+        if (ExecutorInstrumentationUtils.shouldAttachStateToTask(newTask, context)) {
           task = newTask;
           final ContextStore<Runnable, State> contextStore =
               InstrumentationContext.get(Runnable.class, State.class);
-          return ExecutorInstrumentationUtils.setupState(contextStore, newTask, span);
+          return ExecutorInstrumentationUtils.setupState(contextStore, newTask, context);
         }
       }
       return null;

dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/concurrent/executor/JavaExecutorInstrumentation.java

@@ -1,13 +1,14 @@
 package datadog.trace.instrumentation.java.concurrent.executor;
 
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
+import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.currentContext;
+import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.rootContext;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
+import datadog.context.Context;
 import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.bootstrap.InstrumentationContext;
-import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.ExecutorInstrumentationUtils;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.RunnableWrapper;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
@@ -33,16 +34,16 @@ public static State enterJobSubmit(
       // there are cased like ScheduledExecutorService.submit (which we instrument)
       // which calls ScheduledExecutorService.schedule (which we also instrument)
       // where all of this could be dodged the second time
-      final AgentSpan span = activeSpan();
-      if (null != span) {
+      final Context context = currentContext();
+      if (context != rootContext()) {
         final Runnable newTask = RunnableWrapper.wrapIfNeeded(task);
         // It is important to check potentially wrapped task if we can instrument task in this
         // executor. Some executors do not support wrapped tasks.
-        if (ExecutorInstrumentationUtils.shouldAttachStateToTask(newTask, span)) {
+        if (ExecutorInstrumentationUtils.shouldAttachStateToTask(newTask, context)) {
           task = newTask;
           final ContextStore<Runnable, State> contextStore =
               InstrumentationContext.get(Runnable.class, State.class);
-          return ExecutorInstrumentationUtils.setupState(contextStore, newTask, span);
+          return ExecutorInstrumentationUtils.setupState(contextStore, newTask, context);
         }
       }
       return null;

dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/main/java/datadog/trace/instrumentation/java/concurrent/runnable/ConsumerTaskInstrumentation.java

@@ -1,19 +1,20 @@
 package datadog.trace.instrumentation.java.concurrent.runnable;
 
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
-import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
+import static datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge.currentContext;
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.endTaskScope;
+import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.shouldCapture;
 import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.startTaskScope;
 import static datadog.trace.instrumentation.java.concurrent.ConcurrentInstrumentationNames.EXECUTOR_INSTRUMENTATION_NAME;
 import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 
 import com.google.auto.service.AutoService;
+import datadog.context.Context;
 import datadog.context.ContextScope;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.bootstrap.InstrumentationContext;
-import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
 import java.util.Map;
 import java.util.concurrent.ForkJoinTask;
@@ -51,10 +52,10 @@ public void methodAdvice(MethodTransformer transformer) {
   public static class Construct {
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void construct(@Advice.This ForkJoinTask<?> task) {
-      AgentSpan span = activeSpan();
-      if (null != span) {
+      Context context = currentContext();
+      if (shouldCapture(context)) {
         State state = State.FACTORY.create();
-        state.captureAndSetContinuation(span);
+        state.captureAndSetContinuation(context);
         InstrumentationContext.get(ForkJoinTask.class, State.class).put(task, state);
       }
     }