diff --git a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java index 321b726d63b..72ebdeea455 100644 --- a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java +++ b/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 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; } } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/profiling/TaskWrapperRedefinitionStrategyListener.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/profiling/TaskWrapperRedefinitionStrategyListener.java new file mode 100644 index 00000000000..8835cccdbf6 --- /dev/null +++ b/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. + * + *

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>> onError( + final int index, + @Nonnull final List> batch, + @Nonnull final Throwable throwable, + @Nonnull final List> 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> types, final Map>, Throwable> failures) { + if (DEBUG) { + if (!UnwrappingVisitor.ENABLED) { + LOGGER.debug("Retransforming succeeded with a disabled visitor"); + } + } + } +} diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/profiling/UnwrappingVisitor.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/profiling/UnwrappingVisitor.java index 1c8e07ecaf4..97b20b6d90f 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/profiling/UnwrappingVisitor.java +++ b/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. + * + *

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 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 + ';'; diff --git a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/bytebuddy/profiling/UnwrappingVisitorTest.groovy b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/bytebuddy/profiling/UnwrappingVisitorTest.groovy index 2b46c89a6cc..2866ed52278 100644 --- a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/bytebuddy/profiling/UnwrappingVisitorTest.groovy +++ b/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)) {