🍒 11111 - Preload scope classes to prevent virtual thread deadlock

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java

@@ -90,6 +90,7 @@ public int order() {
   }
 
   public List<Instrumenter> typeInstrumentations() {
+    preloadClasses();
     return singletonList(this);
   }
 
@@ -214,6 +215,28 @@ protected final boolean isShortcutMatchingEnabled(boolean defaultToShortcut) {
         .isIntegrationShortcutMatchingEnabled(singletonList(name()), defaultToShortcut);
   }
 
+  /**
+   * Force loading of classes that need to be instrumented, but are using during instrumentation.
+   */
+  @SuppressForbidden // allow this use of Class.forName()
+  protected void preloadClasses() {
+    String[] list = preloadClassNames();
+    if (list != null) {
+      for (String clazz : list) {
+        try {
+          Class.forName(clazz);
+        } catch (Throwable t) {
+          log.debug("Error force loading {} class", clazz);
+        }
+      }
+    }
+  }
+
+  /** Get classes to force load */
+  public String[] preloadClassNames() {
+    return null;
+  }
+
   /** Parent class for all tracing related instrumentations */
   public abstract static class Tracing extends InstrumenterModule {
     public Tracing(String instrumentationName, String... additionalNames) {
@@ -258,18 +281,11 @@ public final boolean isApplicable(Set<TargetSystem> enabledSystems) {
   }
 
   /** Parent class for all IAST related instrumentations */
-  @SuppressForbidden
   public abstract static class Iast extends InstrumenterModule {
     public Iast(String instrumentationName, String... additionalNames) {
       super(instrumentationName, additionalNames);
     }
 
-    @Override
-    public List<Instrumenter> typeInstrumentations() {
-      preloadClassNames();
-      return super.typeInstrumentations();
-    }
-
     @Override
     public final boolean isApplicable(Set<TargetSystem> enabledSystems) {
       if (enabledSystems.contains(TargetSystem.IAST)) {
@@ -282,27 +298,6 @@ public final boolean isApplicable(Set<TargetSystem> enabledSystems) {
       return cfg.getAppSecActivation() == ProductActivation.FULLY_ENABLED;
     }
 
-    /**
-     * Force loading of classes that need to be instrumented, but are using during instrumentation.
-     */
-    private void preloadClassNames() {
-      String[] list = getClassNamesToBePreloaded();
-      if (list != null) {
-        for (String clazz : list) {
-          try {
-            Class.forName(clazz);
-          } catch (Throwable t) {
-            log.debug("Error force loading {} class", clazz);
-          }
-        }
-      }
-    }
-
-    /** Get classes to force load* */
-    public String[] getClassNamesToBePreloaded() {
-      return null;
-    }
-
     @Override
     public Advice.PostProcessor.Factory postProcessor() {
       return IastPostProcessorFactory.INSTANCE;

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

@@ -20,7 +20,7 @@
 public class InputStreamInstrumentation extends InstrumenterModule.Iast
     implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
 
-  private static final String[] FORCE_LOADING = {"java.io.PushbackInputStream"};
+  private static final String[] PRELOAD_CLASS_NAMES = {"java.io.PushbackInputStream"};
 
   public InputStreamInstrumentation() {
     super("inputStream");
@@ -44,8 +44,8 @@ public void methodAdvice(MethodTransformer transformer) {
   }
 
   @Override
-  public String[] getClassNamesToBePreloaded() {
-    return FORCE_LOADING;
+  public String[] preloadClassNames() {
+    return PRELOAD_CLASS_NAMES;
   }
 
   public static class InputStreamAdvice {

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

@@ -45,8 +45,8 @@
  *   <li>{@code unmount()}: swaps the carrier thread's original context back, saving the virtual
  *       thread's (possibly modified) context for the next mount.
  *   <li>Steps 2-3 repeat on each park/unpark cycle, potentially on different carrier threads.
- *   <li>{@code afterDone()}: cancels the help continuation, releasing the context scope to be
- *       closed.
+ *   <li>{@code afterDone()} / {@code afterTerminate()} for early VirtualThread support: cancels the
+ *       help continuation, releasing the context scope to be closed.
  * </ol>
  *
  * <p>Instrumenting the internal {@code VirtualThread.runContinuation()} method does not work as the
@@ -66,10 +66,22 @@ public final class VirtualThreadInstrumentation extends InstrumenterModule.Conte
         Instrumenter.HasMethodAdvice,
         ExcludeFilterProvider {
 
+  // Preload classes used by Context.swap() to avoid class loading on the virtual thread mount path.
+  // DatadogClassLoader loads these from a JarFile using synchronized I/O, which pins
+  // virtual thread carrier threads and can deadlock the application.
+  private static final String[] PRELOAD_CLASS_NAMES = {
+    "datadog.trace.core.scopemanager.ScopeContext", "datadog.trace.core.scopemanager.ScopeStack"
+  };
+
   public VirtualThreadInstrumentation() {
     super("java-lang", "java-lang-21", "virtual-thread");
   }
 
+  @Override
+  public String[] preloadClassNames() {
+    return PRELOAD_CLASS_NAMES;
+  }
+
   @Override
   public String instrumentedType() {
     return VIRTUAL_THREAD_CLASS_NAME;
@@ -99,6 +111,9 @@ public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(
         isMethod().and(named("afterDone")).and(takesArguments(boolean.class)),
         getClass().getName() + "$AfterDone");
+    transformer.applyAdvice(
+        isMethod().and(named("afterTerminate")).and(takesArguments(boolean.class, boolean.class)),
+        getClass().getName() + "$AfterDone");
   }
 
   public static final class Construct {
@@ -141,7 +156,7 @@ public static void onUnmount(@Advice.This Object virtualThread) {
 
   public static final class AfterDone {
     @OnMethodEnter(suppress = Throwable.class)
-    public static void onTerminate(@Advice.This Object virtualThread) {
+    public static void onDone(@Advice.This Object virtualThread) {
       ContextStore<Object, VirtualThreadState> store =
           InstrumentationContext.get(VIRTUAL_THREAD_CLASS_NAME, VIRTUAL_THREAD_STATE_CLASS_NAME);
       VirtualThreadState state = store.remove(virtualThread);