Instrument LambdaMetafactory and preserve Runnable lambda identity during context propagation #12346
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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); | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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<Boolean> 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) { | ||
|
amarziali marked this conversation as resolved.
Outdated
|
||
| 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); | ||
|
Contributor
There was a problem hiding this comment. The |
||
| 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; | ||
| } | ||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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; | ||
| } | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 | ||
|
Contributor
There was a problem hiding this comment. This allowlist is duplicated between this trie and each |
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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, LambdaTransformerHelper.class) | ||
| .with(AgentStrategies.transformerDecorator()) | ||
| .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION) | ||
| .with(AgentStrategies.rediscoveryStrategy()) | ||
|
|
@@ -253,12 +258,65 @@ public void applied(Iterable<String> instrumentationNames) { | |
|
|
||
| InstrumenterState.resetDefaultState(); | ||
| try { | ||
| 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)) { | ||
|
Contributor
There was a problem hiding this comment. This JDK9+ reflective-loading logic duplicates the existing idiom used elsewhere in this class/ |
||
| try { | ||
| Function<ClassFileTransformer, LambdaTransformer> factory = | ||
| (Function<ClassFileTransformer, LambdaTransformer>) | ||
| 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(), | ||
|
There was a problem hiding this comment.
Affected applications keep the old wrapper and do not preserve Runnable identity. Assertion details
Was this helpful? React 👍 or 👎
Contributor
Author
There was a problem hiding this comment. Accessing the target class’s protection domain can be denied. The code safely catches this and falls back. SecurityManager usage is increasingly rare, and IMHO changing the protection domain handling would be a risky change. I would steer not to fix it |
||
| 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<InstrumenterModule> withExtensions(Iterable<InstrumenterModule> initial) { | ||
| String extensionsPath = InstrumenterConfig.get().getTraceExtensionsPath(); | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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<ClassFileTransformer, LambdaTransformer> FACTORY = | ||
| new Function<ClassFileTransformer, LambdaTransformer>() { | ||
| @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(); | ||
| } | ||
| } | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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(); | ||
|
|
||
| // Same-owner lambdas share a symbolic name, so build their target from the supplied bytes. | ||
| SharedTypeInfo<TypeDescription> sharedType = | ||
| transformingLambda && name.equals(targetName) ? null : types.find(name); | ||
|
There was a problem hiding this comment.
The agent can apply wrong instrumentation to the normal class. Assertion details
Was this helpful? React 👍 or 👎
Contributor
Author
There was a problem hiding this comment. it's a good catch. I initially wanted to try not to touch this part as well but it seems something that can happen. A dedicate lambda matcher and a cache correctness fix have been done in 1bc6bbc . the perf test looks untouched |
||
| if (null != sharedType | ||
| && (name.startsWith("java.") || sharedType.sameClassLoader(classLoaderId))) { | ||
| InstrumenterMetrics.reuseTypeDescription(fromTick, isOutline); | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
instanceof FieldBackedContextAccessoris used here as a proxy for "already has propagation fields," but that marker isn't scoped to a specific context store. If a future secondForLambdainstrumenter targetsRunnablewith a different context store, this check would treat its lambdas as already handled and skip wrapping them too, silently dropping propagation for that store with no fallback.