diff --git a/dd-java-agent/agent-bootstrap/build.gradle b/dd-java-agent/agent-bootstrap/build.gradle index 7a59e3d7aca..d626307ed64 100644 --- a/dd-java-agent/agent-bootstrap/build.gradle +++ b/dd-java-agent/agent-bootstrap/build.gradle @@ -5,12 +5,17 @@ plugins { id 'idea' } +apply from: "$rootDir/gradle/tries.gradle" + // The shadowJar of this project will be injected into the JVM's bootstrap classloader tasks.named("compileJava", JavaCompile) { configureCompiler(it, 8, JavaVersion.VERSION_1_8, "Need access to sun.* packages") + dependsOn 'generateClassNameTries' } +tasks.named("sourcesJar") { dependsOn 'generateClassNameTries' } + // FIXME: Improve test coverage. minimumBranchCoverage = 0.0 minimumInstructionCoverage = 0.0 diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/RunnableWrapper.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/RunnableWrapper.java index 5d32250ae4c..54570a00559 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/RunnableWrapper.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/RunnableWrapper.java @@ -1,11 +1,10 @@ package datadog.trace.bootstrap.instrumentation.java.concurrent; +import datadog.trace.bootstrap.FieldBackedContextAccessor; import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType; /** - * This is used to wrap lambda runnables so we can apply field-injection. RunnableWrapper can be - * transformed to add the necessary context-store fields, while lambdas currently cannot until the - * issue reported in https://github.com/raphw/byte-buddy/issues/558 is addressed. + * Wraps anonymous Runnable classes that were not field-injected. * *

We also make this class final to stop instrumentations from extending it in their injected * helper classes, because if this class is loaded during helper injection then we can miss the @@ -25,9 +24,11 @@ public void run() { } public static Runnable wrapIfNeeded(final Runnable task) { - if (!(task instanceof RunnableWrapper) && !ExcludeFilter.exclude(ExcludeType.RUNNABLE, task)) { - // We wrap only lambdas' anonymous classes and if given object has not already been wrapped. - // Anonymous classes have '/' in class name which is not allowed in 'normal' classes. + // Field-injected tasks are already instrumented and must retain their identity. + if (!(task instanceof RunnableWrapper) + && !(task instanceof FieldBackedContextAccessor) + && !ExcludeFilter.exclude(ExcludeType.RUNNABLE, task)) { + // Hidden lambda class names contain '/'. final String className = task.getClass().getName(); if (className.indexOf('/', className.lastIndexOf('.')) > 0) { return new RunnableWrapper(task); diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/TPEHelper.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/TPEHelper.java index e95a6580849..9a81f06faeb 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/TPEHelper.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/TPEHelper.java @@ -8,6 +8,7 @@ import datadog.trace.api.InstrumenterConfig; import datadog.trace.api.Platform; import datadog.trace.bootstrap.ContextStore; +import datadog.trace.bootstrap.FieldBackedContextAccessor; import java.util.Set; import java.util.concurrent.ThreadPoolExecutor; @@ -22,7 +23,7 @@ public final class TPEHelper { // If legacy is enabled, we will try to propagate via wrapping, if not we will try to propagate // via storing the state in the existing field in the Runnable private static final boolean useWrapping; - // A ThreadPoolExecutor with one of these types will newer be propagated/wrapped + // A ThreadPoolExecutor with one of these types will never be propagated/wrapped private static final Set excludedClasses; // A ThreadLocal to store the Scope between beforeExecute and afterExecute if wrapping is not used private static final ThreadLocal threadLocalScope; @@ -30,10 +31,11 @@ public final class TPEHelper { private static final ClassValue WRAP = GenericClassValue.of( input -> { + if (FieldBackedContextAccessor.class.isAssignableFrom(input)) { + return false; + } String className = input.getName(); - // We should always wrap anonymous lambda classes since we can't inject fields into - // them, and they can never be anything more than a _pure_ Runnable. They have '/' in - // their class name which is not allowed in 'normal' classes. + // Wrap anonymous lambda classes that were not field-injected. return className.indexOf('/', className.lastIndexOf('.')) > 0; }); diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformer.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformer.java new file mode 100644 index 00000000000..daf9525a1dd --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformer.java @@ -0,0 +1,12 @@ +package datadog.trace.bootstrap.instrumentation.java.lang.invoke; + +/** Transforms a generated lambda class before it is defined. */ +public interface LambdaTransformer { + /** + * @param slashClassName internal (slash-separated) name of the generated lambda class + * @param targetClass the class declaring the lambda + * @param classBytes the freshly generated lambda class bytes + * @return the transformed bytes, or {@code null}/the original bytes if unchanged + */ + byte[] transform(String slashClassName, Class targetClass, byte[] classBytes); +} diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformerHelper.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformerHelper.java new file mode 100644 index 00000000000..5054e657589 --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformerHelper.java @@ -0,0 +1,64 @@ +package datadog.trace.bootstrap.instrumentation.java.lang.invoke; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Transforms eligible lambda bytes before definition, falling back to the original on failure. */ +public final class LambdaTransformerHelper { + private static final Logger log = LoggerFactory.getLogger(LambdaTransformerHelper.class); + + // Agent transformation may itself create lambdas. + private static final ThreadLocal TRANSFORMING = new ThreadLocal<>(); + + private LambdaTransformerHelper() {} + + /** + * @param classBytes the generated lambda class bytes + * @param lambdaClassName internal (slash-separated) name of the generated lambda class + * @param targetClass the class declaring the lambda + * @param interfaceClass the functional interface implemented by the lambda + * @return possibly transformed bytes; the original bytes on any failure + */ + public static byte[] transform( + byte[] classBytes, String lambdaClassName, Class targetClass, Class interfaceClass) { + try { + // Only exact allowlisted interfaces enter the transformer. + if (interfaceClass == null || LambdaInterfaceNameTrie.apply(interfaceClass.getName()) != 1) { + return classBytes; + } + LambdaTransformer transformer = LambdaTransformerHolder.get(); + if (transformer == null) { + log.debug("Lambda {} skipped: no transformer registered", lambdaClassName); + return classBytes; + } + if (targetClass == null) { + log.debug("Lambda {} skipped: no target class", lambdaClassName); + return classBytes; + } + // Skip lambdas declared by the agent itself to avoid self-instrumentation and recursion. + String targetName = targetClass.getName(); + if (targetName.startsWith("datadog.") || targetName.startsWith("net.bytebuddy.")) { + log.debug("Lambda {} skipped: declared by the agent", lambdaClassName); + return classBytes; + } + if (Boolean.TRUE.equals(TRANSFORMING.get())) { + log.debug("Lambda {} skipped: re-entrant transform", lambdaClassName); + return classBytes; + } + TRANSFORMING.set(Boolean.TRUE); + try { + byte[] result = transformer.transform(lambdaClassName, targetClass, classBytes); + if (result == null) { + log.debug("Lambda {} not transformed", lambdaClassName); + return classBytes; + } + return result; + } finally { + TRANSFORMING.set(Boolean.FALSE); + } + } catch (Throwable e) { + log.debug("Lambda {} skipped: {}", lambdaClassName, e.toString()); + return classBytes; + } + } +} diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformerHolder.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformerHolder.java new file mode 100644 index 00000000000..9c596a659e4 --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformerHolder.java @@ -0,0 +1,19 @@ +package datadog.trace.bootstrap.instrumentation.java.lang.invoke; + +/** + * Holds the {@link LambdaTransformer} registered by the agent installer. Lives on the bootstrap + * class path so it is reachable from instrumented {@code java.lang.invoke} code. + */ +public final class LambdaTransformerHolder { + private static volatile LambdaTransformer transformer; + + private LambdaTransformerHolder() {} + + public static void set(LambdaTransformer transformer) { + LambdaTransformerHolder.transformer = transformer; + } + + public static LambdaTransformer get() { + return transformer; + } +} diff --git a/dd-java-agent/agent-bootstrap/src/main/resources/datadog/trace/bootstrap/instrumentation/java/lang/invoke/lambda_interface_name.trie b/dd-java-agent/agent-bootstrap/src/main/resources/datadog/trace/bootstrap/instrumentation/java/lang/invoke/lambda_interface_name.trie new file mode 100644 index 00000000000..93487386424 --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/main/resources/datadog/trace/bootstrap/instrumentation/java/lang/invoke/lambda_interface_name.trie @@ -0,0 +1,7 @@ +# Generates 'LambdaInterfaceNameTrie.java' + +# Exact functional interfaces whose generated lambda classes should be sent through the agent's +# matching and transformation pipeline. Keep this list narrow: the lookup runs for every lambda +# linkage in the application. + +1 java.lang.Runnable 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 3a8c7065362..501f97866c6 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 @@ -5,6 +5,7 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.GlobalIgnoresMatcher.globalIgnoresMatcher; import static net.bytebuddy.matcher.ElementMatchers.isDefaultFinalizer; +import datadog.environment.JavaVirtualMachine; import datadog.environment.SystemProperties; import datadog.trace.agent.tooling.bytebuddy.SharedTypePools; import datadog.trace.agent.tooling.bytebuddy.iast.TaintableRedefinitionStrategyListener; @@ -19,6 +20,9 @@ import datadog.trace.api.telemetry.IntegrationsCollector; import datadog.trace.bootstrap.FieldBackedContextAccessor; import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter; +import datadog.trace.bootstrap.instrumentation.java.lang.invoke.LambdaTransformer; +import datadog.trace.bootstrap.instrumentation.java.lang.invoke.LambdaTransformerHelper; +import datadog.trace.bootstrap.instrumentation.java.lang.invoke.LambdaTransformerHolder; import datadog.trace.bootstrap.instrumentation.java.module.JpmsHelper; import datadog.trace.util.AgentTaskScheduler; import de.thetaphi.forbiddenapis.SuppressForbidden; @@ -35,6 +39,7 @@ import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.TimeUnit; import java.util.function.BooleanSupplier; +import java.util.function.Function; import net.bytebuddy.ByteBuddy; import net.bytebuddy.agent.builder.AgentBuilder; import net.bytebuddy.description.type.TypeDescription; @@ -147,7 +152,7 @@ public static ClassFileTransformer installBytebuddyAgent( agentBuilder = agentBuilder .disableClassFormatChanges() - .assureReadEdgeTo(inst, FieldBackedContextAccessor.class) + .assureReadEdgeTo(inst, FieldBackedContextAccessor.class, LambdaTransformerHelper.class) .with(AgentStrategies.transformerDecorator()) .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION) .with(AgentStrategies.rediscoveryStrategy()) @@ -253,12 +258,65 @@ public void applied(Iterable instrumentationNames) { InstrumenterState.resetDefaultState(); try { - return transformerBuilder.installOn(inst); + ClassFileTransformer classFileTransformer = transformerBuilder.installOn(inst); + registerLambdaTransformer(classFileTransformer); + return classFileTransformer; } finally { SharedTypePools.endInstall(); } } + /** Registers the installed class-file transformer for generated lambdas. */ + private static void registerLambdaTransformer(final ClassFileTransformer classFileTransformer) { + LambdaTransformer lambdaTransformer = newLambdaTransformer(classFileTransformer); + if (null != lambdaTransformer) { + LambdaTransformerHolder.set(lambdaTransformer); + } + } + + /** + * Java 9+ requires the module-aware transformer for injected read edges. Failure must disable + * lambda transformation rather than fall back to the module-less overload. + */ + @SuppressWarnings("unchecked") + private static LambdaTransformer newLambdaTransformer( + final ClassFileTransformer classFileTransformer) { + if (JavaVirtualMachine.isJavaVersionAtLeast(9)) { + try { + Function factory = + (Function) + Instrumenter.class + .getClassLoader() + .loadClass("datadog.trace.agent.tooling.bytebuddy.DDJava9LambdaTransformer") + .getField("FACTORY") + .get(null); + return factory.apply(classFileTransformer); + } catch (Throwable e) { + log.debug("Problem loading Java 9 lambda transformer, disabling lambda field-injection", e); + return null; + } + } + // Avoid invoking the instrumented metafactory while installing its transformer. + return new LambdaTransformer() { + @Override + public byte[] transform(String slashClassName, Class targetClass, byte[] classBytes) { + TypePoolFacade.beginLambdaTransform(); + try { + return classFileTransformer.transform( + targetClass.getClassLoader(), + slashClassName, + null, + targetClass.getProtectionDomain(), + classBytes); + } catch (Throwable ignored) { + return null; + } finally { + TypePoolFacade.endLambdaTransform(); + } + } + }; + } + /** Returns an iterable that combines the original sequence with any discovered extensions. */ private static Iterable withExtensions(Iterable initial) { String extensionsPath = InstrumenterConfig.get().getTraceExtensionsPath(); diff --git a/dd-java-agent/agent-installer/src/main/java11/datadog/trace/agent/tooling/bytebuddy/DDJava9LambdaTransformer.java b/dd-java-agent/agent-installer/src/main/java11/datadog/trace/agent/tooling/bytebuddy/DDJava9LambdaTransformer.java new file mode 100644 index 00000000000..1ffa6640693 --- /dev/null +++ b/dd-java-agent/agent-installer/src/main/java11/datadog/trace/agent/tooling/bytebuddy/DDJava9LambdaTransformer.java @@ -0,0 +1,43 @@ +package datadog.trace.agent.tooling.bytebuddy; + +import datadog.trace.agent.tooling.bytebuddy.outline.TypePoolFacade; +import datadog.trace.bootstrap.instrumentation.java.lang.invoke.LambdaTransformer; +import java.lang.instrument.ClassFileTransformer; +import java.util.function.Function; + +/** Routes generated lambdas through the module-aware Java 9+ transformer overload. */ +public final class DDJava9LambdaTransformer implements LambdaTransformer { + + /** Read reflectively by the agent installer, which cannot name {@link Module} itself. */ + public static final Function FACTORY = + new Function() { + @Override + public LambdaTransformer apply(ClassFileTransformer classFileTransformer) { + return new DDJava9LambdaTransformer(classFileTransformer); + } + }; + + private final ClassFileTransformer classFileTransformer; + + public DDJava9LambdaTransformer(ClassFileTransformer classFileTransformer) { + this.classFileTransformer = classFileTransformer; + } + + @Override + public byte[] transform(String slashClassName, Class targetClass, byte[] classBytes) { + TypePoolFacade.beginLambdaTransform(); + try { + return classFileTransformer.transform( + targetClass.getModule(), + targetClass.getClassLoader(), + slashClassName, + null, + targetClass.getProtectionDomain(), + classBytes); + } catch (Throwable ignored) { + return null; + } finally { + TypePoolFacade.endLambdaTransform(); + } + } +} diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/memoize/Memoizer.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/memoize/Memoizer.java index 6c8d01152ca..b29989fc259 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/memoize/Memoizer.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/memoize/Memoizer.java @@ -155,6 +155,8 @@ static final class MemoizingMatcher @Override protected boolean doMatch(TypeDescription target) { String targetName = target.getName(); + // Same-owner hidden lambdas share a symbolic name. Bypass these caches before supporting + // lambda interfaces with different matcher results. if (noMatchFilter.contains(targetName) || "java.lang.Object".equals(targetName) || target.isPrimitive()) { diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/TypeFactory.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/TypeFactory.java index 51650b3d2b2..6f044d89c27 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/TypeFactory.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/TypeFactory.java @@ -100,6 +100,8 @@ final class TypeFactory { boolean createOutlines = OUTLINING_ENABLED; + boolean transformingLambda; + ClassLoader originalClassLoader; ClassLoader currentClassLoader; @@ -159,6 +161,14 @@ void beginTransform(String name, byte[] bytecode) { } } + void beginLambdaTransform() { + transformingLambda = true; + } + + void endLambdaTransform() { + transformingLambda = false; + } + /** Once matching is complete we need full descriptions for the actual transformation. */ void enableFullDescriptions() { createOutlines = false; @@ -258,8 +268,9 @@ private TypeDescription lookupType( boolean isOutline = typeParser == outlineTypeParser; long fromTick = InstrumenterMetrics.tick(); - // existing type description from same classloader? - SharedTypeInfo sharedType = types.find(name); + // Same-owner lambdas share a symbolic name, so build their target from the supplied bytes. + SharedTypeInfo sharedType = + transformingLambda && name.equals(targetName) ? null : types.find(name); if (null != sharedType && (name.startsWith("java.") || sharedType.sameClassLoader(classLoaderId))) { InstrumenterMetrics.reuseTypeDescription(fromTick, isOutline); diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/TypePoolFacade.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/TypePoolFacade.java index b5b2e3b39a8..f4c08b8cc1c 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/TypePoolFacade.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/TypePoolFacade.java @@ -49,6 +49,14 @@ public static void beginTransform(String name, byte[] bytecode) { typeFactory.get().beginTransform(name, bytecode); } + public static void beginLambdaTransform() { + typeFactory.get().beginLambdaTransform(); + } + + public static void endLambdaTransform() { + typeFactory.get().endLambdaTransform(); + } + /** Switch to full descriptions, needed for the actual class transformation. */ public static void enableFullDescriptions() { typeFactory.get().enableFullDescriptions(); diff --git a/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie b/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie index 36cedaf6081..dc90bdf2c60 100644 --- a/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie +++ b/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie @@ -57,6 +57,8 @@ 0 java.lang.Runtime # allow context tracking for VirtualThread 0 java.lang.VirtualThread +# allow instrumenting the lambda metafactory to field-inject generated lambda classes +0 java.lang.invoke.InnerClassLambdaMetafactory 0 java.net.http.* 0 java.net.HttpURLConnection 0 java.net.InetAddress diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/bytebuddy/outline/TypeFactoryTest.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/bytebuddy/outline/TypeFactoryTest.java new file mode 100644 index 00000000000..e00f938ce3d --- /dev/null +++ b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/bytebuddy/outline/TypeFactoryTest.java @@ -0,0 +1,57 @@ +package datadog.trace.agent.tooling.bytebuddy.outline; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import java.util.concurrent.Callable; +import net.bytebuddy.ByteBuddy; +import net.bytebuddy.description.type.TypeDescription; +import org.junit.jupiter.api.Test; + +class TypeFactoryTest { + @Test + void reusesCachedDescriptionForRegularTransformationTarget() { + String name = getClass().getName() + "$RegularTarget"; + + assertEquals( + Runnable.class.getName(), resolveInterface(name, bytes(name, Runnable.class), false)); + assertEquals( + Runnable.class.getName(), resolveInterface(name, bytes(name, Callable.class), false)); + } + + @Test + void rebuildsLambdaTransformationTargetFromSuppliedBytes() { + String name = getClass().getName() + "$LambdaTarget"; + + assertEquals( + Runnable.class.getName(), resolveInterface(name, bytes(name, Runnable.class), false)); + assertEquals( + Callable.class.getName(), resolveInterface(name, bytes(name, Callable.class), true)); + } + + private static String resolveInterface(String name, byte[] bytecode, boolean lambda) { + TypeFactory typeFactory = TypeFactory.typeFactory.get(); + typeFactory.switchContext(TypeFactoryTest.class.getClassLoader()); + if (lambda) { + typeFactory.beginLambdaTransform(); + } + typeFactory.beginTransform(name, bytecode); + try { + TypeDescription type = TypeFactory.findType(name); + return type.getInterfaces().getOnly().asErasure().getName(); + } finally { + typeFactory.endTransform(); + if (lambda) { + typeFactory.endLambdaTransform(); + } + } + } + + private static byte[] bytes(String name, Class implementedInterface) { + return new ByteBuddy() + .subclass(Object.class) + .name(name) + .implement(implementedInterface) + .make() + .getBytes(); + } +} diff --git a/dd-java-agent/benchmark/build.gradle b/dd-java-agent/benchmark/build.gradle index 178eefd2def..bc963b78d8e 100644 --- a/dd-java-agent/benchmark/build.gradle +++ b/dd-java-agent/benchmark/build.gradle @@ -38,8 +38,16 @@ jmh { jmhVersion = libs.versions.jmh.get() } +// Copy the agent to a fixed, version-independent path so benchmarks that attach it can name it +// with a compile-time constant (JMH's @Fork annotation cannot read a system property). +def agentJarForBenchmarks = tasks.register('agentJarForBenchmarks', Copy) { + from project(':dd-java-agent').tasks.named('shadowJar') + into layout.buildDirectory.dir('agent') + rename { 'dd-java-agent.jar' } +} + tasks.named('jmh') { - dependsOn ':dd-java-agent:shadowJar' + dependsOn agentJarForBenchmarks } /* diff --git a/dd-java-agent/benchmark/src/jmh/java/lambdabench/LambdaExecutorBenchmark.java b/dd-java-agent/benchmark/src/jmh/java/lambdabench/LambdaExecutorBenchmark.java new file mode 100644 index 00000000000..c1a59f85c77 --- /dev/null +++ b/dd-java-agent/benchmark/src/jmh/java/lambdabench/LambdaExecutorBenchmark.java @@ -0,0 +1,161 @@ +package lambdabench; + +import datadog.trace.api.Trace; +import java.lang.invoke.CallSite; +import java.lang.invoke.LambdaConversionException; +import java.lang.invoke.LambdaMetafactory; +import java.lang.invoke.MethodHandle; +import java.lang.invoke.MethodHandles; +import java.lang.invoke.MethodType; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.function.Supplier; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; + +/** + * Compares lambda {@code Runnable} allocation, execution, and submission with no agent, wrapping, + * and field injection. + * + *

+ * + *

{@code runUntracedLambda} isolates the advice cost when no context was attached. With the GC + * profiler, {@code allocateCapturingRunnable} isolates the injected field's object size cost while + * {@code submitLambda} includes the wrapper allocation tradeoff. Run with: + * + * 

{@code
+ * ./gradlew :dd-java-agent:benchmark:jmh \
+ *   '-Pjmh.includes=LambdaExecutorBenchmark.*allocateCapturingRunnable' \
+ *   -Pjmh.profilers=gc
+ * }
+ */ +public abstract class LambdaExecutorBenchmark { + + // Must remain outside datadog.* because the helper skips agent-owned lambdas. + // Relative to the JMH working directory, which Gradle sets to this project's directory. + private static final String AGENT = "-javaagent:build/agent/dd-java-agent.jar"; + + private static final MethodHandles.Lookup LOOKUP = MethodHandles.lookup(); + private static final MethodHandle RUNNABLE_TARGET; + private static final MethodHandle SUPPLIER_TARGET; + + static { + try { + RUNNABLE_TARGET = + LOOKUP.findStatic( + LambdaExecutorBenchmark.class, "runTarget", MethodType.methodType(void.class)); + SUPPLIER_TARGET = + LOOKUP.findStatic( + LambdaExecutorBenchmark.class, "supplyTarget", MethodType.methodType(Object.class)); + } catch (NoSuchMethodException | IllegalAccessException e) { + throw new ExceptionInInitializerError(e); + } + } + + @State(Scope.Benchmark) + public static class ExecutorState { + ExecutorService pool; + + @Setup + public void setup() { + pool = Executors.newSingleThreadExecutor(); + } + + @TearDown + public void tearDown() { + pool.shutdownNow(); + } + } + + @State(Scope.Thread) + public static class CapturingLambdaState { + public void run() {} + } + + @State(Scope.Thread) + public static class DirectRunState { + Runnable runnable; + int executions; + + @Setup + public void setup() { + runnable = () -> executions++; + } + } + + /** Allocates an untraced capturing Runnable. */ + @Benchmark + public Runnable allocateCapturingRunnable(CapturingLambdaState state) { + return state::run; + } + + /** Executes an already-created lambda Runnable without an active trace. */ + @Benchmark + public int runUntracedLambda(DirectRunState state) { + state.runnable.run(); + return state.executions; + } + + /** Submit a lambda Runnable to the executor under an active trace, and wait for it to run. */ + @Benchmark + public void submitLambda(ExecutorState state) throws InterruptedException { + runUnderTrace(state.pool); + } + + /** Measures cold linkage of an eligible {@link Runnable} lambda class. */ + @Benchmark + public CallSite linkRunnableLambda() throws LambdaConversionException { + return LambdaMetafactory.metafactory( + LOOKUP, + "run", + MethodType.methodType(Runnable.class), + MethodType.methodType(void.class), + RUNNABLE_TARGET, + MethodType.methodType(void.class)); + } + + /** Measures cold linkage of a non-task lambda, which should bypass the agent transformer. */ + @Benchmark + public CallSite linkSupplierLambda() throws LambdaConversionException { + return LambdaMetafactory.metafactory( + LOOKUP, + "get", + MethodType.methodType(Supplier.class), + MethodType.methodType(Object.class), + SUPPLIER_TARGET, + MethodType.methodType(Object.class)); + } + + @Trace(operationName = "parent") + private void runUnderTrace(ExecutorService pool) throws InterruptedException { + CountDownLatch latch = new CountDownLatch(1); + pool.execute(latch::countDown); + latch.await(); + } + + private static void runTarget() {} + + private static Object supplyTarget() { + return null; + } + + @Fork + public static class NoAgent extends LambdaExecutorBenchmark {} + + @Fork(jvmArgsAppend = {AGENT, "-Ddd.trace.lambda.enabled=false"}) + public static class AgentLambdaOff extends LambdaExecutorBenchmark {} + + @Fork(jvmArgsAppend = AGENT) + public static class AgentLambdaOn extends LambdaExecutorBenchmark {} +} diff --git a/dd-java-agent/instrumentation/java/java-lambda/java-lambda-1.8/build.gradle b/dd-java-agent/instrumentation/java/java-lambda/java-lambda-1.8/build.gradle new file mode 100644 index 00000000000..a47efe00d4a --- /dev/null +++ b/dd-java-agent/instrumentation/java/java-lambda/java-lambda-1.8/build.gradle @@ -0,0 +1,20 @@ +plugins { + id 'dd-trace-java.module.instrumentation' +} + +muzzle { + pass { + coreJdk() + } +} + +tasks.named("compileJava") { + configureCompiler(it, 8) +} + +dependencies { + testImplementation project(':dd-java-agent:instrumentation:datadog:tracing:trace-annotation') + testImplementation 'org.clojure:clojure:1.11.1' + // Runs the existing Runnable propagation advice on injected lambdas. + testRuntimeOnly project(':dd-java-agent:instrumentation:java:java-concurrent:java-concurrent-1.8') +} diff --git a/dd-java-agent/instrumentation/java/java-lambda/java-lambda-1.8/src/main/java/datadog/trace/instrumentation/java/lang/invoke/LambdaMetafactoryInstrumentation.java b/dd-java-agent/instrumentation/java/java-lambda/java-lambda-1.8/src/main/java/datadog/trace/instrumentation/java/lang/invoke/LambdaMetafactoryInstrumentation.java new file mode 100644 index 00000000000..c63c9a98b07 --- /dev/null +++ b/dd-java-agent/instrumentation/java/java-lambda/java-lambda-1.8/src/main/java/datadog/trace/instrumentation/java/lang/invoke/LambdaMetafactoryInstrumentation.java @@ -0,0 +1,233 @@ +package datadog.trace.instrumentation.java.lang.invoke; + +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.agent.tooling.JavaModuleOpenProvider; +import datadog.trace.api.Platform; +import datadog.trace.bootstrap.instrumentation.java.lang.invoke.LambdaTransformerHelper; +import java.util.Collection; +import java.util.Collections; +import net.bytebuddy.asm.AsmVisitorWrapper; +import net.bytebuddy.description.field.FieldDescription; +import net.bytebuddy.description.field.FieldList; +import net.bytebuddy.description.method.MethodList; +import net.bytebuddy.description.type.TypeDefinition; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.implementation.Implementation; +import net.bytebuddy.jar.asm.ClassVisitor; +import net.bytebuddy.jar.asm.ClassWriter; +import net.bytebuddy.jar.asm.MethodVisitor; +import net.bytebuddy.jar.asm.Opcodes; +import net.bytebuddy.jar.asm.Type; +import net.bytebuddy.matcher.ElementMatcher; +import net.bytebuddy.pool.TypePool; +import net.bytebuddy.utility.OpenedClassReader; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Routes generated lambda bytes through the agent transformer before definition, allowing + * allowlisted interfaces such as {@link Runnable} to receive field injection and advice. + * + *

An ASM visitor is required because the transform call must be inserted immediately after the + * lambda bytes are generated, in the middle of the metafactory method. + */ +@AutoService(InstrumenterModule.class) +public final class LambdaMetafactoryInstrumentation extends InstrumenterModule.ContextTracking + implements Instrumenter.ForBootstrap, + Instrumenter.ForSingleType, + Instrumenter.HasTypeAdvice, + Instrumenter.WithTypeStructure, + JavaModuleOpenProvider { + + private static final Logger log = LoggerFactory.getLogger(LambdaMetafactoryInstrumentation.class); + + private static final String METAFACTORY = "java.lang.invoke.InnerClassLambdaMetafactory"; + + private static final String LAMBDA_CLASS_NAME_FIELD = "lambdaClassName"; + private static final String TARGET_CLASS_FIELD = "targetClass"; + private static final String INTERFACE_CLASS_FIELD = "interfaceClass"; + private static final String LEGACY_INTERFACE_CLASS_FIELD = "samBase"; + + public LambdaMetafactoryInstrumentation() { + super("lambda"); + } + + @Override + public boolean isEnabled() { + return super.isEnabled() && !Platform.isNativeImageBuilder(); + } + + @Override + public String instrumentedType() { + return METAFACTORY; + } + + @Override + public Collection triggerClasses() { + return Collections.singleton(METAFACTORY); + } + + /** Require every field read by the injected bytecode. */ + @Override + public ElementMatcher structureMatcher() { + return HasMetafactoryFields.INSTANCE; + } + + /** Public because this matcher is loaded across agent class-loader boundaries. */ + public static final class HasMetafactoryFields implements ElementMatcher { + public static final HasMetafactoryFields INSTANCE = new HasMetafactoryFields(); + + @Override + public boolean matches(TypeDescription target) { + return declaresField(target, LAMBDA_CLASS_NAME_FIELD, String.class.getName()) + && declaresField(target, TARGET_CLASS_FIELD, Class.class.getName()) + && interfaceClassField(target) != null; + } + + static String interfaceClassField(TypeDescription type) { + // JDK 8 and 11 use samBase; newer JDKs use interfaceClass. + if (declaresField(type, INTERFACE_CLASS_FIELD, Class.class.getName())) { + return INTERFACE_CLASS_FIELD; + } + if (declaresField(type, LEGACY_INTERFACE_CLASS_FIELD, Class.class.getName())) { + return LEGACY_INTERFACE_CLASS_FIELD; + } + return null; + } + + private static boolean declaresField(TypeDescription type, String name, String fieldType) { + for (TypeDefinition current = type; current != null; current = current.getSuperClass()) { + for (FieldDescription field : current.asErasure().getDeclaredFields()) { + if (name.equals(field.getName()) + && fieldType.equals(field.getType().asErasure().getName())) { + return true; + } + } + } + return false; + } + } + + @Override + public void typeAdvice(TypeTransformer transformer) { + transformer.applyAdvice(new MetafactoryVisitorWrapper()); + } + + public static final class MetafactoryVisitorWrapper implements AsmVisitorWrapper { + @Override + public int mergeWriter(int flags) { + return flags | ClassWriter.COMPUTE_MAXS; + } + + @Override + public int mergeReader(int flags) { + return flags; + } + + @Override + public ClassVisitor wrap( + TypeDescription instrumentedType, + ClassVisitor classVisitor, + Implementation.Context implementationContext, + TypePool typePool, + FieldList fields, + MethodList methods, + int writerFlags, + int readerFlags) { + return new MetafactoryClassVisitor( + classVisitor, + instrumentedType.getInternalName(), + HasMetafactoryFields.interfaceClassField(instrumentedType)); + } + } + + private static final class MetafactoryClassVisitor extends ClassVisitor { + private final String slashClassName; + private final String interfaceClassField; + private boolean injected; + + MetafactoryClassVisitor(ClassVisitor cv, String slashClassName, String interfaceClassField) { + super(OpenedClassReader.ASM_API, cv); + this.slashClassName = slashClassName; + this.interfaceClassField = interfaceClassField; + } + + @Override + public MethodVisitor visitMethod( + int access, String name, String descriptor, String signature, String[] exceptions) { + MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions); + // The byte-generation method changed in JDK 25. + if (("spinInnerClass".equals(name) || "generateInnerClass".equals(name)) + && "()Ljava/lang/Class;".equals(descriptor)) { + return new MetafactoryMethodVisitor(api, mv, slashClassName, interfaceClassField, this); + } + return mv; + } + + @Override + public void visitEnd() { + super.visitEnd(); + if (!injected) { + log.debug( + "No injection site found in {}; lambda field-injection is inactive.", slashClassName); + } + } + } + + private static final class MetafactoryMethodVisitor extends MethodVisitor { + private final String slashClassName; + private final String interfaceClassField; + private final MetafactoryClassVisitor declaringVisitor; + + MetafactoryMethodVisitor( + int api, + MethodVisitor mv, + String slashClassName, + String interfaceClassField, + MetafactoryClassVisitor declaringVisitor) { + super(api, mv); + this.slashClassName = slashClassName; + this.interfaceClassField = interfaceClassField; + this.declaringVisitor = declaringVisitor; + } + + @Override + public void visitMethodInsn( + int opcode, String owner, String name, String descriptor, boolean isInterface) { + super.visitMethodInsn(opcode, owner, name, descriptor, isInterface); + // Match repackaged JDK APIs while excluding unrelated byte-array producers. The generated + // byte[] remains on the operand stack after the original call. + if ((opcode == Opcodes.INVOKEVIRTUAL + && "toByteArray".equals(name) + && "()[B".equals(descriptor) + && owner.endsWith("/ClassWriter")) + || (opcode == Opcodes.INVOKEINTERFACE + && "build".equals(name) + && descriptor.endsWith(")[B") + && owner.endsWith("/ClassFile"))) { + // stack: ..., byte[] + super.visitVarInsn(Opcodes.ALOAD, 0); + super.visitFieldInsn( + Opcodes.GETFIELD, slashClassName, LAMBDA_CLASS_NAME_FIELD, "Ljava/lang/String;"); + super.visitVarInsn(Opcodes.ALOAD, 0); + // Resolves the defining class loader and module. + super.visitFieldInsn( + Opcodes.GETFIELD, slashClassName, TARGET_CLASS_FIELD, "Ljava/lang/Class;"); + super.visitVarInsn(Opcodes.ALOAD, 0); + // Allows the helper to reject non-allowlisted interfaces before full matching. + super.visitFieldInsn( + Opcodes.GETFIELD, slashClassName, interfaceClassField, "Ljava/lang/Class;"); + super.visitMethodInsn( + Opcodes.INVOKESTATIC, + Type.getInternalName(LambdaTransformerHelper.class), + "transform", + "([BLjava/lang/String;Ljava/lang/Class;Ljava/lang/Class;)[B", + false); + // stack: ..., transformed byte[] + declaringVisitor.injected = true; + } + } + } +} diff --git a/dd-java-agent/instrumentation/java/java-lambda/java-lambda-1.8/src/test/java/datadog/trace/instrumentation/java/lang/invoke/LambdaMetafactoryInstrumentationTest.java b/dd-java-agent/instrumentation/java/java-lambda/java-lambda-1.8/src/test/java/datadog/trace/instrumentation/java/lang/invoke/LambdaMetafactoryInstrumentationTest.java new file mode 100644 index 00000000000..7b173e0578a --- /dev/null +++ b/dd-java-agent/instrumentation/java/java-lambda/java-lambda-1.8/src/test/java/datadog/trace/instrumentation/java/lang/invoke/LambdaMetafactoryInstrumentationTest.java @@ -0,0 +1,350 @@ +package datadog.trace.instrumentation.java.lang.invoke; + +import static net.bytebuddy.jar.asm.Opcodes.ACC_PRIVATE; +import static net.bytebuddy.jar.asm.Opcodes.ACC_PUBLIC; +import static net.bytebuddy.jar.asm.Opcodes.ACONST_NULL; +import static net.bytebuddy.jar.asm.Opcodes.ARETURN; +import static net.bytebuddy.jar.asm.Opcodes.ASM7; +import static net.bytebuddy.jar.asm.Opcodes.INVOKEINTERFACE; +import static net.bytebuddy.jar.asm.Opcodes.INVOKESTATIC; +import static net.bytebuddy.jar.asm.Opcodes.INVOKEVIRTUAL; +import static net.bytebuddy.jar.asm.Opcodes.POP; +import static net.bytebuddy.jar.asm.Opcodes.V1_8; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.trace.bootstrap.instrumentation.java.lang.invoke.LambdaTransformerHelper; +import datadog.trace.bootstrap.instrumentation.java.lang.invoke.LambdaTransformerHolder; +import datadog.trace.instrumentation.java.lang.invoke.LambdaMetafactoryInstrumentation.MetafactoryVisitorWrapper; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Supplier; +import net.bytebuddy.description.type.TypeDescription; +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 org.junit.jupiter.api.Test; + +class LambdaMetafactoryInstrumentationTest { + + private static final String HELPER = + "datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformerHelper"; + + private static boolean injectsTransformCall( + String methodName, String methodDescriptor, ClassBody body) { + ClassWriter in = new ClassWriter(0); + in.visit(V1_8, ACC_PUBLIC, "Dummy", null, "java/lang/Object", null); + MethodVisitor mv = in.visitMethod(ACC_PRIVATE, methodName, methodDescriptor, null, null); + mv.visitCode(); + body.write(mv); + mv.visitMaxs(0, 0); + mv.visitEnd(); + in.visitEnd(); + + ClassWriter out = new ClassWriter(0); + ClassVisitor visitor = + new MetafactoryVisitorWrapper() + .wrap(realMetafactoryDescription(), out, null, null, null, null, 0, 0); + new ClassReader(in.toByteArray()).accept(visitor, 0); + + AtomicBoolean found = new AtomicBoolean(false); + new ClassReader(out.toByteArray()) + .accept( + new ClassVisitor(ASM7) { + @Override + public MethodVisitor visitMethod( + int access, String name, String desc, String sig, String[] ex) { + return new MethodVisitor(ASM7) { + @Override + public void visitMethodInsn( + int opcode, String owner, String name, String desc, boolean itf) { + if (opcode == INVOKESTATIC + && HELPER.equals(owner) + && "transform".equals(name) + && "([BLjava/lang/String;Ljava/lang/Class;Ljava/lang/Class;)[B" + .equals(desc)) { + found.set(true); + } + } + }; + } + }, + 0); + return found.get(); + } + + @Test + void injectsAfterToByteArrayInSpinInnerClass() { + assertTrue( + injectsTransformCall( + "spinInnerClass", + "()Ljava/lang/Class;", + mv -> { + mv.visitInsn(ACONST_NULL); + mv.visitMethodInsn( + INVOKEVIRTUAL, + "jdk/internal/org/objectweb/asm/ClassWriter", + "toByteArray", + "()[B", + false); + mv.visitInsn(POP); + mv.visitInsn(ACONST_NULL); + mv.visitInsn(ARETURN); + })); + } + + @Test + void injectsAfterToByteArrayInGenerateInnerClass() { + assertTrue( + injectsTransformCall( + "generateInnerClass", + "()Ljava/lang/Class;", + mv -> { + mv.visitInsn(ACONST_NULL); + mv.visitMethodInsn( + INVOKEVIRTUAL, + "jdk/internal/org/objectweb/asm/ClassWriter", + "toByteArray", + "()[B", + false); + mv.visitInsn(POP); + mv.visitInsn(ACONST_NULL); + mv.visitInsn(ARETURN); + })); + } + + @Test + void injectsAfterBuildOnClassFileApi() { + assertTrue( + injectsTransformCall( + "spinInnerClass", + "()Ljava/lang/Class;", + mv -> { + mv.visitInsn(ACONST_NULL); + mv.visitMethodInsn( + INVOKEINTERFACE, + "java/lang/classfile/ClassFile", + "build", + "(Ljava/lang/classfile/constantpool/ClassEntry;" + + "Ljava/lang/classfile/constantpool/ConstantPoolBuilder;" + + "Ljava/util/function/Consumer;)[B", + true); + mv.visitInsn(POP); + mv.visitInsn(ACONST_NULL); + mv.visitInsn(ARETURN); + })); + } + + @Test + void doesNotInjectInUnrelatedMethod() { + assertFalse( + injectsTransformCall( + "someOtherMethod", + "()Ljava/lang/Class;", + mv -> { + mv.visitInsn(ACONST_NULL); + mv.visitMethodInsn( + INVOKEVIRTUAL, + "jdk/internal/org/objectweb/asm/ClassWriter", + "toByteArray", + "()[B", + false); + mv.visitInsn(POP); + mv.visitInsn(ACONST_NULL); + mv.visitInsn(ARETURN); + })); + } + + @Test + void doesNotInjectOnUnrelatedToByteArrayOwner() { + assertFalse( + injectsTransformCall( + "spinInnerClass", + "()Ljava/lang/Class;", + mv -> { + mv.visitInsn(ACONST_NULL); + mv.visitMethodInsn( + INVOKEVIRTUAL, "java/io/ByteArrayOutputStream", "toByteArray", "()[B", false); + mv.visitInsn(POP); + mv.visitInsn(ACONST_NULL); + mv.visitInsn(ARETURN); + })); + } + + /** Verifies every field read by the injected bytecode, including inherited fields. */ + @Test + void structureMatcherAcceptsTheRealMetafactory() { + assertTrue( + new LambdaMetafactoryInstrumentation() + .structureMatcher() + .matches(realMetafactoryDescription())); + } + + @Test + void structureMatcherAcceptsCurrentInterfaceField() { + assertTrue( + new LambdaMetafactoryInstrumentation() + .structureMatcher() + .matches(TypeDescription.ForLoadedType.of(CurrentMetafactoryFields.class))); + } + + @Test + void structureMatcherAcceptsLegacyInterfaceField() { + assertTrue( + new LambdaMetafactoryInstrumentation() + .structureMatcher() + .matches(TypeDescription.ForLoadedType.of(LegacyMetafactoryFields.class))); + } + + @Test + void structureMatcherRejectsTypeWithoutTheFields() { + assertFalse( + new LambdaMetafactoryInstrumentation() + .structureMatcher() + .matches(TypeDescription.ForLoadedType.of(Object.class))); + } + + private static TypeDescription realMetafactoryDescription() { + try { + return TypeDescription.ForLoadedType.of( + Class.forName("java.lang.invoke.InnerClassLambdaMetafactory")); + } catch (ClassNotFoundException e) { + throw new AssertionError(e); + } + } + + private static final class CurrentMetafactoryFields { + private String lambdaClassName; + private Class targetClass; + private Class interfaceClass; + } + + private static final class LegacyMetafactoryFields { + private String lambdaClassName; + private Class targetClass; + private Class samBase; + } + + @Test + void nonRunnableLambdaBypassesTransformer() { + byte[] originalBytes = new byte[0]; + AtomicBoolean transformed = new AtomicBoolean(); + LambdaTransformerHolder.set( + (className, targetClass, classBytes) -> { + transformed.set(true); + return classBytes; + }); + try { + byte[] result = + LambdaTransformerHelper.transform( + originalBytes, "test/Lambda", Object.class, Supplier.class); + + assertSame(originalBytes, result); + assertFalse(transformed.get()); + } finally { + LambdaTransformerHolder.set(null); + } + } + + @Test + void exactRunnableInterfaceUsesTransformer() { + byte[] originalBytes = new byte[0]; + AtomicBoolean transformed = new AtomicBoolean(); + LambdaTransformerHolder.set( + (className, targetClass, classBytes) -> { + transformed.set(true); + return classBytes; + }); + try { + byte[] result = + LambdaTransformerHelper.transform( + originalBytes, "test/Lambda", Object.class, Runnable.class); + + assertSame(originalBytes, result); + assertTrue(transformed.get()); + } finally { + LambdaTransformerHolder.set(null); + } + } + + @Test + void runnableSubinterfaceBypassesTransformer() { + byte[] originalBytes = new byte[0]; + AtomicBoolean transformed = new AtomicBoolean(); + LambdaTransformerHolder.set( + (className, targetClass, classBytes) -> { + transformed.set(true); + return classBytes; + }); + try { + byte[] result = + LambdaTransformerHelper.transform( + originalBytes, "test/Lambda", Object.class, RunnableSubtype.class); + + assertSame(originalBytes, result); + assertFalse(transformed.get()); + } finally { + LambdaTransformerHolder.set(null); + } + } + + @Test + void transformerFailureFallsBackAndDoesNotPoisonNextLambda() { + byte[] originalBytes = new byte[0]; + byte[] transformedBytes = new byte[1]; + AtomicInteger calls = new AtomicInteger(); + LambdaTransformerHolder.set( + (className, targetClass, classBytes) -> { + if (calls.getAndIncrement() == 0) { + throw new IllegalStateException("expected test failure"); + } + return transformedBytes; + }); + try { + assertSame( + originalBytes, + LambdaTransformerHelper.transform( + originalBytes, "test/Lambda", Object.class, Runnable.class)); + assertSame( + transformedBytes, + LambdaTransformerHelper.transform( + originalBytes, "test/Lambda", Object.class, Runnable.class)); + assertEquals(2, calls.get()); + } finally { + LambdaTransformerHolder.set(null); + } + } + + @Test + void nullTransformFallsBackAndDoesNotPoisonNextLambda() { + byte[] originalBytes = new byte[0]; + byte[] transformedBytes = new byte[1]; + AtomicInteger calls = new AtomicInteger(); + LambdaTransformerHolder.set( + (className, targetClass, classBytes) -> + calls.getAndIncrement() == 0 ? null : transformedBytes); + try { + assertSame( + originalBytes, + LambdaTransformerHelper.transform( + originalBytes, "test/Lambda", Object.class, Runnable.class)); + assertSame( + transformedBytes, + LambdaTransformerHelper.transform( + originalBytes, "test/Lambda", Object.class, Runnable.class)); + assertEquals(2, calls.get()); + } finally { + LambdaTransformerHolder.set(null); + } + } + + private interface RunnableSubtype extends Runnable {} + + @FunctionalInterface + private interface ClassBody { + void write(MethodVisitor mv); + } +} diff --git a/dd-java-agent/instrumentation/java/java-lambda/java-lambda-1.8/src/test/java/testdog/trace/instrumentation/lambda/ClojureAFnIntegrationTest.java b/dd-java-agent/instrumentation/java/java-lambda/java-lambda-1.8/src/test/java/testdog/trace/instrumentation/lambda/ClojureAFnIntegrationTest.java new file mode 100644 index 00000000000..a52ac12a473 --- /dev/null +++ b/dd-java-agent/instrumentation/java/java-lambda/java-lambda-1.8/src/test/java/testdog/trace/instrumentation/lambda/ClojureAFnIntegrationTest.java @@ -0,0 +1,26 @@ +package testdog.trace.instrumentation.lambda; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import clojure.java.api.Clojure; +import clojure.lang.AFn; +import datadog.trace.agent.test.AbstractInstrumentationTest; +import datadog.trace.bootstrap.FieldBackedContextAccessor; +import datadog.trace.test.junit.utils.config.WithConfig; +import org.junit.jupiter.api.Test; + +@WithConfig(key = "trace.runnable.enabled", value = "false") +public class ClojureAFnIntegrationTest extends AbstractInstrumentationTest { + + @Test + void afnIsNotFieldInjected() { + // Runnable instrumentation can be disabled to avoid inflating every AFn; see + // https://github.com/DataDog/dd-trace-java/pull/2925. + Object function = Clojure.var("clojure.core", "eval").invoke(Clojure.read("(fn [] nil)")); + + assertTrue(function instanceof AFn); + assertTrue(function instanceof Runnable); + assertFalse(function instanceof FieldBackedContextAccessor); + } +} diff --git a/dd-java-agent/instrumentation/java/java-lambda/java-lambda-1.8/src/test/java/testdog/trace/instrumentation/lambda/LambdaMetafactoryDisabledForkedTest.java b/dd-java-agent/instrumentation/java/java-lambda/java-lambda-1.8/src/test/java/testdog/trace/instrumentation/lambda/LambdaMetafactoryDisabledForkedTest.java new file mode 100644 index 00000000000..7edd1873b90 --- /dev/null +++ b/dd-java-agent/instrumentation/java/java-lambda/java-lambda-1.8/src/test/java/testdog/trace/instrumentation/lambda/LambdaMetafactoryDisabledForkedTest.java @@ -0,0 +1,19 @@ +package testdog.trace.instrumentation.lambda; + +import static org.junit.jupiter.api.Assertions.assertFalse; + +import datadog.trace.agent.test.AbstractInstrumentationTest; +import datadog.trace.bootstrap.FieldBackedContextAccessor; +import datadog.trace.test.junit.utils.config.WithConfig; +import org.junit.jupiter.api.Test; + +@WithConfig(key = "trace.lambda.enabled", value = "false") +public class LambdaMetafactoryDisabledForkedTest extends AbstractInstrumentationTest { + + @Test + void runnableLambdaIsNotFieldInjected() { + Runnable lambda = () -> {}; + + assertFalse(lambda instanceof FieldBackedContextAccessor); + } +} diff --git a/dd-java-agent/instrumentation/java/java-lambda/java-lambda-1.8/src/test/java/testdog/trace/instrumentation/lambda/LambdaMetafactoryIntegrationTest.java b/dd-java-agent/instrumentation/java/java-lambda/java-lambda-1.8/src/test/java/testdog/trace/instrumentation/lambda/LambdaMetafactoryIntegrationTest.java new file mode 100644 index 00000000000..eeb1f805042 --- /dev/null +++ b/dd-java-agent/instrumentation/java/java-lambda/java-lambda-1.8/src/test/java/testdog/trace/instrumentation/lambda/LambdaMetafactoryIntegrationTest.java @@ -0,0 +1,91 @@ +package testdog.trace.instrumentation.lambda; + +import static datadog.trace.agent.test.assertions.SpanMatcher.span; +import static datadog.trace.agent.test.assertions.TraceMatcher.SORT_BY_START_TIME; +import static datadog.trace.agent.test.assertions.TraceMatcher.trace; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.trace.agent.test.AbstractInstrumentationTest; +import datadog.trace.api.Trace; +import datadog.trace.bootstrap.FieldBackedContextAccessor; +import datadog.trace.bootstrap.instrumentation.java.concurrent.RunnableWrapper; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Supplier; +import org.junit.jupiter.api.Test; + +/** Lambda integration tests outside the ignored {@code datadog.*} prefix. */ +public class LambdaMetafactoryIntegrationTest extends AbstractInstrumentationTest { + + @Test + void lambdaRunnableIsFieldInjectedNotWrapped() { + // Link after the agent is installed. + Runnable lambda = () -> {}; + + assertTrue( + lambda instanceof FieldBackedContextAccessor, + "lambda Runnable should be field-injected via the metafactory instrumentation"); + assertSame(lambda, RunnableWrapper.wrapIfNeeded(lambda)); + } + + @Test + void sameOwnerRunnableCaptureShapesAreAllFieldInjected() { + AtomicInteger counter = new AtomicInteger(); + int delta = 7; + Runnable[] lambdas = {() -> {}, counter::incrementAndGet, () -> counter.addAndGet(delta)}; + + for (Runnable lambda : lambdas) { + assertTrue( + lambda instanceof FieldBackedContextAccessor, + "every Runnable lambda shape should be field-injected"); + assertSame(lambda, RunnableWrapper.wrapIfNeeded(lambda)); + lambda.run(); + } + assertEquals(8, counter.get()); + } + + @Test + void nonRunnableLambdaIsNotTransformed() { + Supplier lambda = Object::new; + + assertFalse( + lambda instanceof FieldBackedContextAccessor, + "non-Runnable lambda should bypass the agent transformer"); + } + + @Test + void lambdaPropagatesContextAcrossExecutor() throws Exception { + ExecutorService pool = Executors.newSingleThreadExecutor(); + try { + CountDownLatch latch = new CountDownLatch(1); + submitUnderParent(pool, latch); + assertTrue(latch.await(10, TimeUnit.SECONDS), "child task did not run"); + + assertTraces( + trace( + SORT_BY_START_TIME, + span().root().operationName("parent"), + span().childOfPrevious().operationName("lambda-child"))); + } finally { + pool.shutdownNow(); + } + } + + @Trace(operationName = "parent") + void submitUnderParent(ExecutorService pool, CountDownLatch latch) { + pool.execute( + () -> { + child(); + latch.countDown(); + }); + } + + @Trace(operationName = "lambda-child") + void child() {} +} diff --git a/dd-smoke-tests/java9-modules/src/main/java11/datadog/smoketest/moduleapp/ModuleApplication.java b/dd-smoke-tests/java9-modules/src/main/java11/datadog/smoketest/moduleapp/ModuleApplication.java index b7e484b2d7b..41808269b44 100644 --- a/dd-smoke-tests/java9-modules/src/main/java11/datadog/smoketest/moduleapp/ModuleApplication.java +++ b/dd-smoke-tests/java9-modules/src/main/java11/datadog/smoketest/moduleapp/ModuleApplication.java @@ -1,7 +1,10 @@ package datadog.smoketest.moduleapp; +import testdog.moduleapp.LambdaTask; + public class ModuleApplication { public static void main(final String[] args) throws InterruptedException { + LambdaTask.runOnExecutor(); Thread.sleep(600); } } diff --git a/dd-smoke-tests/java9-modules/src/main/java11/testdog/moduleapp/LambdaTask.java b/dd-smoke-tests/java9-modules/src/main/java11/testdog/moduleapp/LambdaTask.java new file mode 100644 index 00000000000..9d8338bd1b8 --- /dev/null +++ b/dd-smoke-tests/java9-modules/src/main/java11/testdog/moduleapp/LambdaTask.java @@ -0,0 +1,46 @@ +package testdog.moduleapp; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; + +/** Links a Runnable lambda from a named module and outside the ignored {@code datadog.*} prefix. */ +public final class LambdaTask { + private static final String FIELD_BACKED_CONTEXT_ACCESSOR = + "datadog.trace.bootstrap.FieldBackedContextAccessor"; + + private LambdaTask() {} + + public static void runOnExecutor() throws InterruptedException { + final ExecutorService pool = Executors.newSingleThreadExecutor(); + try { + final CountDownLatch latch = new CountDownLatch(1); + final Runnable task = latch::countDown; + assertFieldInjection(task); + pool.execute(task); + if (!latch.await(10, TimeUnit.SECONDS)) { + throw new IllegalStateException("lambda task did not run"); + } + } finally { + pool.shutdownNow(); + } + } + + private static void assertFieldInjection(final Runnable task) { + final List interfaces = new ArrayList<>(); + for (final Class type : task.getClass().getInterfaces()) { + interfaces.add(type.getName()); + } + final boolean injected = interfaces.contains(FIELD_BACKED_CONTEXT_ACCESSOR); + if (!injected) { + throw new IllegalStateException( + "expected lambda field-injection; " + + task.getClass().getName() + + " implements " + + interfaces); + } + } +} diff --git a/metadata/agent-jar-checks.properties b/metadata/agent-jar-checks.properties index a7bde3e2525..c02458675a1 100644 --- a/metadata/agent-jar-checks.properties +++ b/metadata/agent-jar-checks.properties @@ -110,6 +110,7 @@ expected.integrations = IastInstrumentation,\ jwt,\ kafka,\ kotlin_coroutine,\ + lambda,\ lettuce,\ liberty,\ log4j,\ diff --git a/metadata/supported-configurations.json b/metadata/supported-configurations.json index 6c3fe354b68..c5d035c44a4 100644 --- a/metadata/supported-configurations.json +++ b/metadata/supported-configurations.json @@ -7972,6 +7972,14 @@ "aliases": ["DD_LEGACY_E2E_DURATION_ENABLED"] } ], + "DD_TRACE_LAMBDA_ENABLED": [ + { + "version": "A", + "type": "boolean", + "default": "true", + "aliases": ["DD_TRACE_INTEGRATION_LAMBDA_ENABLED", "DD_INTEGRATION_LAMBDA_ENABLED"] + } + ], "DD_TRACE_LETTUCE_4_ASYNC_ENABLED": [ { "version": "A", diff --git a/settings.gradle.kts b/settings.gradle.kts index 14e988af3fe..88316d62a6c 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -399,6 +399,7 @@ include( ":dd-java-agent:instrumentation:java:java-concurrent:java-concurrent-21.0", ":dd-java-agent:instrumentation:java:java-concurrent:java-concurrent-25.0", ":dd-java-agent:instrumentation:java:java-io-1.8", + ":dd-java-agent:instrumentation:java:java-lambda:java-lambda-1.8", ":dd-java-agent:instrumentation:java:java-lang:java-lang-1.8", ":dd-java-agent:instrumentation:java:java-lang:java-lang-11.0", ":dd-java-agent:instrumentation:java:java-lang:java-lang-15.0",