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
| 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,65 @@ | ||
| 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 { | ||
| if (interfaceClass == null) { | ||
| return classBytes; | ||
| } | ||
| String interfaceName = interfaceClass.getName(); | ||
| 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, interfaceName); | ||
| if (result == null) { | ||
| log.debug("Lambda {} not transformed", lambdaClassName); | ||
| return classBytes; | ||
| } | ||
| return result; | ||
| } finally { | ||
| TRANSFORMING.remove(); | ||
| } | ||
| } 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,11 @@ | ||
| # Reserved lambda interface manifest (inactive) | ||
|
|
||
| # Production lambda transformation is intentionally not enabled for any functional interface. | ||
| # This file is not read at build time or runtime. Runtime selection is driven exclusively by | ||
| # enabled Instrumenter.ForLambda implementations and also requires trace.lambda.enabled. | ||
|
|
||
| # If a validated manifest is introduced later, entries will use the ClassNameTrie format: | ||
| # | ||
| # 1 java.lang.Runnable | ||
|
|
||
| # Tests register Runnable from TestRunnableLambdaInstrumentation; this file has no effect on them. |
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; | ||
|
|
@@ -163,6 +168,15 @@ public static ClassFileTransformer installBytebuddyAgent( | |
| // .with(AgentBuilder.LambdaInstrumentationStrategy.ENABLED) | ||
| .ignore(globalIgnoresMatcher(skipAdditionalLibraryMatcher)); | ||
|
|
||
| boolean lambdaTransformationEnabled = | ||
| !Platform.isNativeImageBuilder() | ||
| && InstrumenterConfig.get() | ||
| .isIntegrationEnabled(Collections.singleton("lambda"), false); | ||
| if (lambdaTransformationEnabled) { | ||
| // The injected metafactory call needs java.base to read the bootstrap helper's module. | ||
| agentBuilder = agentBuilder.assureReadEdgeTo(inst, LambdaTransformerHelper.class); | ||
| } | ||
|
|
||
| if (DEBUG) { | ||
| agentBuilder = | ||
| agentBuilder | ||
|
|
@@ -253,12 +267,86 @@ public void applied(Iterable<String> instrumentationNames) { | |
|
|
||
| InstrumenterState.resetDefaultState(); | ||
| try { | ||
| ClassFileTransformer classFileTransformer = transformerBuilder.installOn(inst); | ||
| if (lambdaTransformationEnabled) { | ||
| registerLambdaTransformer(classFileTransformer, transformerBuilder.lambdaInterfaces()); | ||
| } | ||
| 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 transformation", 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
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.
The
TRANSFORMINGre-entrancy guard drops field-injection for any lambda whose class definition is triggered as a side effect of processing another lambda's transform on the same thread (e.g. a lambda captured while defining/loading another lambda's supporting classes). That lambda permanently falls back toRunnableWrapperwith only this debug log as a trace.