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,14 @@ | ||
| 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 | ||
| * @param interfaceClassName the functional interface implemented by the lambda | ||
| * @return the transformed bytes, or {@code null}/the original bytes if unchanged | ||
| */ | ||
| byte[] transform( | ||
| String slashClassName, Class<?> targetClass, byte[] classBytes, String interfaceClassName); | ||
| } |
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,69 @@ | ||
| 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) { | ||
| return classBytes; | ||
| } | ||
| String interfaceName = interfaceClass.getName(); | ||
| if (LambdaInterfaceNameTrie.apply(interfaceName) != 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); | ||
|
Contributor
There was a problem hiding this comment. The |
||
| return classBytes; | ||
| } | ||
| TRANSFORMING.set(Boolean.TRUE); | ||
| try { | ||
| byte[] result = | ||
| transformer.transform(lambdaClassName, targetClass, classBytes, interfaceName); | ||
| 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,84 @@ public void applied(Iterable<String> instrumentationNames) { | |
|
|
||
| InstrumenterState.resetDefaultState(); | ||
| try { | ||
| ClassFileTransformer classFileTransformer = transformerBuilder.installOn(inst); | ||
| registerLambdaTransformer(classFileTransformer, transformerBuilder.lambdaInterfaces()); | ||
|
Contributor
There was a problem hiding this comment. Narrow startup race: |
||
| return classFileTransformer; | ||
| } finally { | ||
| SharedTypePools.endInstall(); | ||
| } | ||
| } | ||
|
|
||
| /** Registers the installed class-file transformer for generated lambdas. */ | ||
| private static void registerLambdaTransformer( | ||
| final ClassFileTransformer classFileTransformer, final String[] lambdaInterfaces) { | ||
| LambdaTransformer transformer = | ||
| lambdaInterfaces.length == 0 ? null : newLambdaTransformer(classFileTransformer); | ||
| LambdaTransformerHolder.set(filterLambdaTransformer(transformer, lambdaInterfaces)); | ||
| } | ||
|
|
||
| static LambdaTransformer filterLambdaTransformer( | ||
| final LambdaTransformer transformer, final String[] lambdaInterfaces) { | ||
| if (transformer == null) { | ||
| return null; | ||
| } | ||
| return (className, targetClass, classBytes, interfaceName) -> { | ||
| for (String enabledInterface : lambdaInterfaces) { | ||
|
Contributor
There was a problem hiding this comment.
|
||
| if (enabledInterface.equals(interfaceName)) { | ||
| return transformer.transform(className, targetClass, classBytes, interfaceName); | ||
| } | ||
| } | ||
| return null; | ||
| }; | ||
| } | ||
|
|
||
| /** | ||
| * 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, | ||
| String interfaceClassName) { | ||
| TypePoolFacade.beginLambdaTransform(interfaceClassName); | ||
| 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
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.