Preserve old implementation when legacy manager is disabled

Author: amarziali

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/VirtualThreadState.java

@@ -1,44 +1,63 @@
 package datadog.trace.bootstrap.instrumentation.java.lang;
 
 import datadog.context.Context;
+import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope.Continuation;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 
 /**
- * The virtual thread's scope stack lives in a virtual-thread-aware {@code ThreadLocal}, so it is
- * seeded once on the first mount and then follows the thread across park/unpark and carrier
- * migration on its own. ddprof's profiler context is keyed by carrier thread instead, so it is
- * re-bound on mount and cleared on unmount when a carrier-bound profiling integration is active.
+ * Holds the context and continuation for a virtual thread.
+ *
+ * <p>With the legacy context manager, {@code swap()} rebuilds the whole scope stack on every
+ * mount/unmount, which is costly on the virtual-thread park/unpark hot path. There the context is
+ * seeded once on the first mount — it then follows the thread across park/unpark and carrier
+ * migration via its virtual-thread-aware {@code ThreadLocal} scope stack — and the profiler context
+ * (which is keyed by carrier thread) is re-applied on each subsequent mount and cleared on unmount.
+ *
+ * <p>With the new context manager {@code swap()} is cheap and drives the profiler through its
+ * context listener, so we simply swap in on mount and out on unmount.
  */
 public final class VirtualThreadState {
-  private Context seedContext;
-  // Keeps the enclosing trace alive until the virtual thread finishes.
+  private static final boolean LEGACY_CONTEXT_MANAGER =
+      InstrumenterConfig.get().isLegacyContextManagerEnabled();
+
+  private Context context;
   private final Continuation continuation;
+  private Context previousContext;
   private boolean seeded;
 
-  public VirtualThreadState(Context seedContext, Continuation continuation) {
-    this.seedContext = seedContext;
+  public VirtualThreadState(Context context, Continuation continuation) {
+    this.context = context;
     this.continuation = continuation;
   }
 
   public void onMount() {
-    if (!seeded) {
-      // First mount also applies the profiler context to the carrier.
-      seedContext.swap();
-      seeded = true;
-      seedContext = null;
+    if (LEGACY_CONTEXT_MANAGER) {
+      if (!seeded) {
+        // First mount also applies the profiler context to the carrier.
+        context.swap();
+        seeded = true;
+        context = null;
+      } else {
+        AgentTracer.get().getProfilingContext().setContext(Context.current());
+      }
     } else {
-      AgentTracer.get().rebindProfilingContextToCarrier();
+      previousContext = context.swap();
     }
   }
 
   public void onUnmount() {
-    AgentTracer.get().unbindProfilingContextFromCarrier();
+    if (LEGACY_CONTEXT_MANAGER) {
+      AgentTracer.get().getProfilingContext().clearContext();
+    } else if (previousContext != null) {
+      context = previousContext.swap();
+      previousContext = null;
+    }
   }
 
   public void onTerminate() {
-    if (this.continuation != null) {
-      this.continuation.cancel();
+    if (continuation != null) {
+      continuation.cancel();
     }
   }
 }

dd-java-agent/agent-profiling/profiling-ddprof/src/main/java/com/datadog/profiling/ddprof/DatadogProfilingIntegration.java

@@ -1,5 +1,6 @@
 package com.datadog.profiling.ddprof;
 
+import datadog.context.Context;
 import datadog.trace.api.EndpointTracker;
 import datadog.trace.api.Stateful;
 import datadog.trace.api.profiling.ProfilingContextAttribute;
@@ -79,10 +80,14 @@ public String name() {
   }
 
   @Override
-  public boolean isCarrierThreadBound() {
-    return true;
+  public void setContext(Context context) {
+    AgentSpan span = AgentSpan.fromContext(context);
+    if (span != null) {
+      contextManager.activate(span.spanContext());
+    }
   }
 
+  @Override
   public void clearContext() {
     DDPROF.clearSpanContext();
     DDPROF.clearContextValue(SPAN_NAME_INDEX);

dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/main/java/datadog/trace/instrumentation/java/lang/jdk21/VirtualThreadInstrumentation.java

@@ -40,12 +40,13 @@
  * <ol>
  *   <li>{@code init()}: captures the current {@link Context} and an {@link Continuation} to prevent
  *       the enclosing context scope from completing early.
- *   <li>{@code mount()}: on the first mount, seeds the virtual thread's scope stack with the
- *       captured context; on later mounts, re-applies the profiler context to the current carrier
- *       (no-op unless carrier-bound profiling is active).
- *   <li>{@code unmount()}: clears the carrier's profiler context (no-op unless carrier-bound
- *       profiling is active). The scope stack is NOT swapped out — it lives in the virtual thread's
- *       own thread-local and follows it across park/unpark and carrier migration.
+ *   <li>{@code mount()}: with the legacy context manager, seeds the virtual thread's scope stack
+ *       with the captured context on the first mount only (it then follows the thread via its
+ *       virtual-thread-aware thread-local) and re-applies the profiler context on later mounts;
+ *       with the new context manager, swaps the context in on every mount (cheap, and drives the
+ *       profiler through its context listener).
+ *   <li>{@code unmount()}: with the legacy context manager, clears the carrier's profiler context;
+ *       with the new context manager, swaps the carrier's context back.
  *   <li>Steps 2-3 repeat on each park/unpark cycle, potentially on different carrier threads.
  *   <li>{@code afterDone()} / {@code afterTerminate()} for early VirtualThread support: cancels the
  *       help continuation, releasing the context scope to be closed.

dd-java-agent/instrumentation/java/java-lang/java-lang-21.0/src/test/java/testdog/trace/instrumentation/java/lang/jdk21/VirtualThreadApiInstrumentationForkedTest.java

@@ -0,0 +1,11 @@
+package testdog.trace.instrumentation.java.lang.jdk21;
+
+import datadog.trace.junit.utils.config.WithConfig;
+
+/**
+ * Runs the {@link VirtualThreadApiInstrumentationTest} cases with the legacy context manager
+ * disabled, so {@code VirtualThreadState} takes the swap path instead of seed-once. Forked because
+ * the legacy-context-manager choice is captured once per JVM.
+ */
+@WithConfig(key = "legacy.context-manager.enabled", value = "false")
+class VirtualThreadApiInstrumentationForkedTest extends VirtualThreadApiInstrumentationTest {}

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

@@ -1438,16 +1438,6 @@ public void addScopeListener(final ScopeListener listener) {
     this.scopeManager.addScopeListener(listener);
   }
 
-  @Override
-  public void rebindProfilingContextToCarrier() {
-    scopeManager.rebindProfilingContextToCarrier();
-  }
-
-  @Override
-  public void unbindProfilingContextFromCarrier() {
-    scopeManager.unbindProfilingContextFromCarrier();
-  }
-
   @Override
   public SubscriptionService getSubscriptionService(RequestContextSlot slot) {
     return (SubscriptionService) instrumentationGateway.getCallbackProvider(slot);

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

@@ -177,23 +177,6 @@ public final void beforeActivated() {
     }
   }
 
-  /**
-   * Clears the profiler context for this scope's state without closing the scope itself. The scope
-   * state must be idempotent and reusable, as this may be followed by another {@link
-   * #beforeActivated()} on the same scope (e.g. across virtual-thread unmount/remount).
-   */
-  public final void clearProfilingContext() {
-    if (scopeState == Stateful.DEFAULT) {
-      return;
-    }
-    try {
-      scopeState.close();
-    } catch (Throwable e) {
-      ContinuableScopeManager.ratelimitedLog.warn(
-          "ScopeState {} threw exception in clearProfilingContext()", scopeState.getClass(), e);
-    }
-  }
-
   public final void afterActivated() {
     boolean hasScopeListeners = !scopeManager.scopeListeners.isEmpty();
     boolean hasExtendedListeners = !scopeManager.extendedScopeListeners.isEmpty();

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

@@ -62,7 +62,6 @@ public final class ContinuableScopeManager implements ContextManager {
   final HealthMetrics healthMetrics;
   private final ProfilingContextIntegration profilingContextIntegration;
   private final boolean profilingEnabled;
-  private final boolean profilerCarrierBound;
   private final boolean hasDepthLimit;
 
   /**
@@ -96,8 +95,6 @@ public ContinuableScopeManager(
     this.profilingContextIntegration = profilingContextIntegration;
     this.profilingEnabled =
         !(profilingContextIntegration instanceof ProfilingContextIntegration.NoOp);
-    this.profilerCarrierBound =
-        this.profilingEnabled && profilingContextIntegration.isCarrierThreadBound();
 
     ContextManager.register(this);
   }
@@ -373,34 +370,6 @@ ScopeStack scopeStack() {
     return this.tlsScopeStack.get();
   }
 
-  /**
-   * Re-applies the active scope's profiler context to the current carrier; no-op unless a
-   * carrier-bound profiling integration is active.
-   */
-  public void rebindProfilingContextToCarrier() {
-    if (!profilerCarrierBound) {
-      return;
-    }
-    final ContinuableScope active = scopeStack().active();
-    if (active != null) {
-      active.beforeActivated();
-    }
-  }
-
-  /**
-   * Clears the current carrier's profiler context; no-op unless a carrier-bound profiling
-   * integration is active.
-   */
-  public void unbindProfilingContextFromCarrier() {
-    if (!profilerCarrierBound) {
-      return;
-    }
-    final ContinuableScope active = scopeStack().active();
-    if (active != null) {
-      active.clearProfilingContext();
-    }
-  }
-
   @Override
   public Context current() {
     final ContinuableScope active = scopeStack().active();