Support instrumentation of repackaged libraries #8153
There are no files selected for viewing
This file contains 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 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,30 @@ | ||
| package datadog.trace.agent.tooling; | ||
|
|
||
| import java.io.IOException; | ||
| import net.bytebuddy.dynamic.ClassFileLocator; | ||
|
|
||
| /** Locates and shades class-file resources from the advice class-loader. */ | ||
| public final class ShadedAdviceLocator implements ClassFileLocator { | ||
| private final ClassFileLocator adviceLocator; | ||
| private final AdviceShader adviceShader; | ||
|
|
||
| public ShadedAdviceLocator(ClassLoader adviceLoader, AdviceShader adviceShader) { | ||
| this.adviceLocator = ClassFileLocator.ForClassLoader.of(adviceLoader); | ||
| this.adviceShader = adviceShader; | ||
| } | ||
|
|
||
| @Override | ||
| public Resolution locate(String className) throws IOException { | ||
| final Resolution resolution = adviceLocator.locate(className); | ||
| if (resolution.isResolved()) { | ||
| return new Resolution.Explicit(adviceShader.shade(resolution.resolve())); | ||
| } else { | ||
| return resolution; | ||
| } | ||
| } | ||
|
|
||
| @Override | ||
| public void close() throws IOException { | ||
| adviceLocator.close(); | ||
| } | ||
| } |
This file contains 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,98 @@ | ||
| package datadog.trace.agent.tooling; | ||
|
|
||
| import datadog.trace.api.cache.DDCache; | ||
| import datadog.trace.api.cache.DDCaches; | ||
| import java.util.Map; | ||
| import net.bytebuddy.jar.asm.ClassReader; | ||
| import net.bytebuddy.jar.asm.ClassVisitor; | ||
| import net.bytebuddy.jar.asm.ClassWriter; | ||
| import net.bytebuddy.jar.asm.commons.ClassRemapper; | ||
| import net.bytebuddy.jar.asm.commons.Remapper; | ||
|
|
||
| /** Shades advice bytecode by applying relocations to all references. */ | ||
| public final class AdviceShader extends Remapper { | ||
| private final DDCache<String, String> cache = DDCaches.newFixedSizeCache(64); | ||
|
|
||
| /** Flattened sequence of old-prefix, new-prefix relocations. */ | ||
| private final String[] prefixes; | ||
|
|
||
| public static AdviceShader with(Map<String, String> relocations) { | ||
| return relocations != null ? new AdviceShader(relocations) : null; | ||
| } | ||
|
|
||
| AdviceShader(Map<String, String> relocations) { | ||
| // convert relocations to a flattened sequence: old-prefix, new-prefix, etc. | ||
| this.prefixes = new String[relocations.size() * 2]; | ||
| int i = 0; | ||
| for (Map.Entry<String, String> e : relocations.entrySet()) { | ||
| String oldPrefix = e.getKey(); | ||
| String newPrefix = e.getValue(); | ||
| if (oldPrefix.indexOf('.') > 0) { | ||
| // accept dotted prefixes, but store them in their internal form | ||
| this.prefixes[i++] = oldPrefix.replace('.', '/'); | ||
| this.prefixes[i++] = newPrefix.replace('.', '/'); | ||
| } else { | ||
| this.prefixes[i++] = oldPrefix; | ||
| this.prefixes[i++] = newPrefix; | ||
| } | ||
| } | ||
| } | ||
mcculls marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| /** Applies shading before calling the given {@link ClassVisitor}. */ | ||
| public ClassVisitor shade(ClassVisitor cv) { | ||
| return new ClassRemapper(cv, this); | ||
| } | ||
|
|
||
| /** Returns the result of shading the given bytecode. */ | ||
| public byte[] shade(byte[] bytecode) { | ||
| ClassReader cr = new ClassReader(bytecode); | ||
| ClassWriter cw = new ClassWriter(null, 0); | ||
| cr.accept(shade(cw), 0); | ||
| return cw.toByteArray(); | ||
| } | ||
|
|
||
| @Override | ||
| public String map(String internalName) { | ||
| if (internalName.startsWith("java/") | ||
| || internalName.startsWith("datadog/") | ||
| || internalName.startsWith("net/bytebuddy/")) { | ||
| return internalName; // never shade these references | ||
| } | ||
| return cache.computeIfAbsent(internalName, this::shade); | ||
| } | ||
|
|
||
| @Override | ||
| public Object mapValue(Object value) { | ||
| if (value instanceof String) { | ||
| String text = (String) value; | ||
| if (text.isEmpty()) { | ||
| return text; | ||
| } else if (text.indexOf('.') > 0) { | ||
| return shadeDottedName(text); | ||
|
There was a problem hiding this comment. Quick question: when do you expect having dotted names? There was a problem hiding this comment. Needed for #8157 - we want to relocate the context store string, which is in an inner advice class which we cannot use the The context store key here is a string literal containing a dotted class-name. (Also the reason the AWS instrumentation uses a string literal for the context store key is so it can access the context-store without needing a compile-time dependency, for reasons that come down to how the AWS client may be deployed) |
||
| } else { | ||
| return shade(text); | ||
| } | ||
| } else { | ||
| return super.mapValue(value); | ||
| } | ||
| } | ||
|
|
||
| private String shade(String internalName) { | ||
| for (int i = 0; i < prefixes.length; i += 2) { | ||
| if (internalName.startsWith(prefixes[i])) { | ||
| return prefixes[i + 1] + internalName.substring(prefixes[i].length()); | ||
| } | ||
| } | ||
| return internalName; | ||
| } | ||
|
|
||
| private String shadeDottedName(String name) { | ||
| String internalName = name.replace('.', '/'); | ||
| for (int i = 0; i < prefixes.length; i += 2) { | ||
| if (internalName.startsWith(prefixes[i])) { | ||
| return prefixes[i + 1].replace('/', '.') + name.substring(prefixes[i].length()); | ||
| } | ||
| } | ||
| return name; | ||
| } | ||
| } | ||
This file contains 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 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 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 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 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.
Can we just consider 'safe' to leave it unbounded not to forget to raise the limits if needed in a near future?
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 general recommendation would be to always use bounded collections.
What's the worst case according to you if the cache is too small? Loosing CPU cycle at startup or having unreclaimed memory? 🤷
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.
FYI this cache was tuned by running it over some sample advice and optimising for both space and reducing duplicate requests - as Bruce says if we ever have advice whose bytecode contains type references that overflow the cache (and repeat over the same references) then the worst case is a small amount of CPU needed for pattern matching (also this cache is only used when the main relocated matcher applies to a class-loader and byte-buddy resolves the advice for transformation)