🍒 11111 - Preload scope classes to prevent virtual thread deadlock (#11131)

* feat(agent): Move forced class preload into instrumenter module

(cherry picked from commit e95e91e)

* fix(virtual-thread): Preload scope classes to prevent virtual thread deadlock

Preload ScopeContext and ScopeStack classes in ContinuableScopeManager constructor to avoid class loading on virtual thread mount path. DatadogClassLoader's synchronized I/O from JarFile can pin carrier threads and deadlock the application.

(cherry picked from commit c521b50)

* fix(virtual-thread): Add support for early VirtualThread afterTerminate method

Support both afterDone() and afterTerminate() cleanup methods to handle different JDK 21 virtual thread implementations. Both methods cancel the help continuation and release the context scope when the virtual thread completes.

(cherry picked from commit d223f79)

---------

Co-authored-by: Brice Dutheil <brice.dutheil@gmail.com>

Author: PerfectSlayer

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);