Fix instrumentation being dropped when a JDK AOT cache is used

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

@@ -12,6 +12,7 @@
 import datadog.trace.agent.tooling.bytebuddy.matcher.DDElementMatchers;
 import datadog.trace.agent.tooling.bytebuddy.memoize.MemoizedMatchers;
 import datadog.trace.agent.tooling.bytebuddy.outline.TypePoolFacade;
+import datadog.trace.agent.tooling.bytebuddy.profiling.TaskWrapperRedefinitionStrategyListener;
 import datadog.trace.agent.tooling.usm.UsmExtractorImpl;
 import datadog.trace.agent.tooling.usm.UsmMessageFactoryImpl;
 import datadog.trace.api.InstrumenterConfig;
@@ -401,9 +402,12 @@ private static void temporaryOverride(String key, String value, BooleanSupplier
   private static AgentBuilder.RedefinitionStrategy.Listener redefinitionStrategyListener(
       final Set<InstrumenterModule.TargetSystem> enabledSystems) {
     if (enabledSystems.contains(InstrumenterModule.TargetSystem.IAST)) {
-      return TaintableRedefinitionStrategyListener.INSTANCE;
+      // both visitors change the structure of a class, so both need a chance to back off
+      return new AgentBuilder.RedefinitionStrategy.Listener.Compound(
+          TaintableRedefinitionStrategyListener.INSTANCE,
+          TaskWrapperRedefinitionStrategyListener.INSTANCE);
     } else {
-      return AgentBuilder.RedefinitionStrategy.Listener.NoOp.INSTANCE;
+      return TaskWrapperRedefinitionStrategyListener.INSTANCE;
     }
   }
 

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/profiling/TaskWrapperRedefinitionStrategyListener.java

@@ -0,0 +1,67 @@
+package datadog.trace.agent.tooling.bytebuddy.profiling;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import javax.annotation.Nonnull;
+import net.bytebuddy.agent.builder.AgentBuilder;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * {@link UnwrappingVisitor} redefines the structure of a class by adding an interface, meaning that
+ * it cannot be applied to already loaded classes.
+ *
+ * <p>This listener will disable the visitor to prevent a failure with the whole redefinition batch.
+ */
+public final class TaskWrapperRedefinitionStrategyListener
+    extends AgentBuilder.RedefinitionStrategy.Listener.Adapter {
+
+  private static final Logger LOGGER =
+      LoggerFactory.getLogger(TaskWrapperRedefinitionStrategyListener.class);
+  private static final boolean DEBUG = LOGGER.isDebugEnabled();
+
+  public static final TaskWrapperRedefinitionStrategyListener INSTANCE =
+      new TaskWrapperRedefinitionStrategyListener();
+
+  private TaskWrapperRedefinitionStrategyListener() {}
+
+  @Override
+  @Nonnull
+  public Iterable<? extends List<Class<?>>> onError(
+      final int index,
+      @Nonnull final List<Class<?>> batch,
+      @Nonnull final Throwable throwable,
+      @Nonnull final List<Class<?>> types) {
+    if (UnwrappingVisitor.ENABLED) {
+      // logged once: the flag is one-way, so this branch cannot be re-entered
+      LOGGER.info(
+          "Disabling task unwrapping for queueing time profiling: retransformation failed for a"
+              + " batch of {} classes because adding the TaskWrapper interface is a structural"
+              + " change the JVM rejects for already loaded classes (e.g. materialized from a"
+              + " JDK 24+ AOT cache). Retrying the batch without it. Queueing time is still"
+              + " recorded, but the task type is reported as the wrapper class. All other"
+              + " instrumentation is preserved.",
+          batch.size());
+      UnwrappingVisitor.ENABLED = false;
+      return Collections.singletonList(batch);
+    } else {
+      if (DEBUG) {
+        LOGGER.debug(
+            "Exception while retransforming after disabling the visitor in batch {}, task unwrapping is disabled",
+            index);
+      }
+      return Collections.emptyList();
+    }
+  }
+
+  @Override
+  public void onComplete(
+      final int amount, final List<Class<?>> types, final Map<List<Class<?>>, Throwable> failures) {
+    if (DEBUG) {
+      if (!UnwrappingVisitor.ENABLED) {
+        LOGGER.debug("Retransforming succeeded with a disabled visitor");
+      }
+    }
+  }
+}

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/profiling/UnwrappingVisitor.java

@@ -18,6 +18,21 @@
 
 public class UnwrappingVisitor implements AsmVisitorWrapper {
 
+  /**
+   * Adding the {@code TaskWrapper} interface changes the structure of a class, which the JVM
+   * rejects for classes that are already loaded ("attempted to change superclass or interfaces").
+   * That normally does not happen because the agent installs before those classes load, but a JDK
+   * 24+ AOT cache (JEP 483) materializes them up-front, so the whole retransform batch fails and no
+   * further instrumentation is installed.
+   *
+   * <p>When that happens {@link TaskWrapperRedefinitionStrategyListener} clears this flag and
+   * retries the batch, trading task unwrapping for a working install. Queueing time itself keeps
+   * working, since {@code TaskWrapper} is only read by {@code QueueTimeEvent.setTask()} and {@code
+   * TaskWrapper.getUnwrappedType} falls back to the object's own class; only the reported task type
+   * loses its resolution.
+   */
+  public static volatile boolean ENABLED = true;
+
   private final Map<String, String> classNameToDelegateFieldNames;
 
   public UnwrappingVisitor(String... classAndDelegateFieldNames) {
@@ -81,7 +96,7 @@ public void visit(
         String signature,
         String superName,
         String[] interfaces) {
-      if (interfaces == null || !Arrays.asList(interfaces).contains(TASK_WRAPPER)) {
+      if (ENABLED && (interfaces == null || !Arrays.asList(interfaces).contains(TASK_WRAPPER))) {
         interfaces = append(interfaces, TASK_WRAPPER);
         if (signature != null) {
           signature += 'L' + TASK_WRAPPER + ';';

dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/bytebuddy/profiling/UnwrappingVisitorTest.groovy

@@ -6,6 +6,7 @@ import java.util.concurrent.FutureTask
 import net.bytebuddy.jar.asm.ClassReader
 import net.bytebuddy.jar.asm.ClassVisitor
 import net.bytebuddy.jar.asm.ClassWriter
+import net.bytebuddy.jar.asm.MethodVisitor
 import net.bytebuddy.jar.asm.signature.SignatureReader
 import net.bytebuddy.jar.asm.signature.SignatureVisitor
 import org.apache.commons.io.IOUtils
@@ -57,6 +58,44 @@ class UnwrappingVisitorTest extends Specification {
     genericInterfaces.last() == 'datadog/trace/bootstrap/instrumentation/api/TaskWrapper'
   }
 
+  void 'test the visitor backs off when disabled'(){
+    setup:
+    def classFile = readClassBytes(FutureTask)
+    def classReader = new ClassReader(classFile)
+    def classWriter = new ClassWriter(classReader, 0)
+
+    def declaredInterfaces = []
+    def methodNames = []
+
+    def classVisitor = new ClassVisitor(ASM_API, classWriter) {
+        @Override
+        void visit(int version, int access, String name, String signature, String superName, String[] interfaces) {
+          declaredInterfaces = interfaces
+          super.visit(version, access, name, signature, superName, interfaces)
+        }
+
+        @Override
+        MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
+          methodNames += name
+          return super.visitMethod(access, name, descriptor, signature, exceptions)
+        }
+      }
+
+    def taskVisitor = new UnwrappingVisitor.ImplementTaskWrapperClassVisitor(
+      classVisitor, 'java.util.concurrent.FutureTask', 'callable')
+
+    when:
+    UnwrappingVisitor.ENABLED = false
+    classReader.accept(taskVisitor, 0)
+
+    then: 'neither the interface nor the unwrap method is added'
+    !declaredInterfaces.contains('datadog/trace/bootstrap/instrumentation/api/TaskWrapper')
+    !methodNames.contains('$$DD$$__unwrap')
+
+    cleanup:
+    UnwrappingVisitor.ENABLED = true
+  }
+
   static byte [] readClassBytes(Class<?> clazz){
     final String classResourceName = '/' + clazz.getName().replace('.', '/') + '.class'
     try (InputStream is = clazz.getResourceAsStream(classResourceName)) {