Move ASM visiting method into context-store impl
This commit is contained in:
parent
cc27f1507e
commit
b89aa54a93
|
@ -18,7 +18,6 @@ import java.util.Set;
|
|||
import lombok.extern.slf4j.Slf4j;
|
||||
import net.bytebuddy.agent.builder.AgentBuilder;
|
||||
import net.bytebuddy.description.type.TypeDescription;
|
||||
import net.bytebuddy.dynamic.DynamicType;
|
||||
import net.bytebuddy.matcher.ElementMatcher;
|
||||
import net.bytebuddy.utility.JavaModule;
|
||||
|
||||
|
@ -101,8 +100,9 @@ public interface Instrumenter {
|
|||
.and(new MuzzleMatcher())
|
||||
.transform(DDTransformers.defaultTransformers());
|
||||
agentBuilder = injectHelperClasses(agentBuilder);
|
||||
agentBuilder = applyContextStoreTransform(agentBuilder);
|
||||
agentBuilder = contextProvider.instrumentationTransformer(agentBuilder);
|
||||
agentBuilder = applyInstrumentationTransformers(agentBuilder);
|
||||
agentBuilder = contextProvider.additionalInstrumentation(agentBuilder);
|
||||
return agentBuilder;
|
||||
}
|
||||
|
||||
|
@ -128,26 +128,6 @@ public interface Instrumenter {
|
|||
return agentBuilder;
|
||||
}
|
||||
|
||||
private AgentBuilder.Identified.Extendable applyContextStoreTransform(
|
||||
AgentBuilder.Identified.Extendable agentBuilder) {
|
||||
if (contextStore().size() > 0) {
|
||||
agentBuilder =
|
||||
agentBuilder.transform(
|
||||
new AgentBuilder.Transformer() {
|
||||
@Override
|
||||
public DynamicType.Builder<?> transform(
|
||||
DynamicType.Builder<?> builder,
|
||||
TypeDescription typeDescription,
|
||||
ClassLoader classLoader,
|
||||
JavaModule module) {
|
||||
return builder.visit(contextProvider.getInstrumentationVisitor());
|
||||
}
|
||||
});
|
||||
agentBuilder = agentBuilder.transform(new HelperInjector(contextProvider.dynamicClasses()));
|
||||
}
|
||||
return agentBuilder;
|
||||
}
|
||||
|
||||
/** Matches classes for which instrumentation is not muzzled. */
|
||||
private class MuzzleMatcher implements AgentBuilder.RawMatcher {
|
||||
@Override
|
||||
|
|
|
@ -1,20 +1,17 @@
|
|||
package datadog.trace.agent.tooling.context;
|
||||
|
||||
import java.util.Map;
|
||||
import net.bytebuddy.agent.builder.AgentBuilder;
|
||||
import net.bytebuddy.asm.AsmVisitorWrapper;
|
||||
|
||||
public interface InstrumentationContextProvider {
|
||||
|
||||
/** @return An AsmVisitorWrapper to run on the instrumentation's bytecode. */
|
||||
AsmVisitorWrapper getInstrumentationVisitor();
|
||||
|
||||
/**
|
||||
* @return A map of dynamic-class-name -> dynamic-class-bytes. These classes will be injected into
|
||||
* the runtime classloader.
|
||||
* Hook to provide an agent builder after advice is applied to target class. Used to implement
|
||||
* context-store lookup.
|
||||
*/
|
||||
Map<String, byte[]> dynamicClasses();
|
||||
AgentBuilder.Identified.Extendable instrumentationTransformer(
|
||||
AgentBuilder.Identified.Extendable builder);
|
||||
|
||||
/** Hook for the context impl to define additional instrumentation if needed. */
|
||||
AgentBuilder additionalInstrumentation(AgentBuilder builder);
|
||||
/** Hook to define additional instrumentation. Run at instrumentation advice is hooked up. */
|
||||
AgentBuilder.Identified.Extendable additionalInstrumentation(
|
||||
AgentBuilder.Identified.Extendable builder);
|
||||
}
|
||||
|
|
|
@ -1,5 +1,6 @@
|
|||
package datadog.trace.agent.tooling.context;
|
||||
|
||||
import datadog.trace.agent.tooling.HelperInjector;
|
||||
import datadog.trace.agent.tooling.Instrumenter;
|
||||
import datadog.trace.agent.tooling.Utils;
|
||||
import datadog.trace.bootstrap.InstrumentationContext;
|
||||
|
@ -18,6 +19,7 @@ import net.bytebuddy.description.field.FieldDescription;
|
|||
import net.bytebuddy.description.field.FieldList;
|
||||
import net.bytebuddy.description.method.MethodList;
|
||||
import net.bytebuddy.description.type.TypeDescription;
|
||||
import net.bytebuddy.dynamic.DynamicType;
|
||||
import net.bytebuddy.implementation.Implementation;
|
||||
import net.bytebuddy.jar.asm.ClassVisitor;
|
||||
import net.bytebuddy.jar.asm.Label;
|
||||
|
@ -25,6 +27,7 @@ import net.bytebuddy.jar.asm.MethodVisitor;
|
|||
import net.bytebuddy.jar.asm.Opcodes;
|
||||
import net.bytebuddy.jar.asm.Type;
|
||||
import net.bytebuddy.pool.TypePool;
|
||||
import net.bytebuddy.utility.JavaModule;
|
||||
|
||||
/**
|
||||
* InstrumentationContextProvider which stores context in a global map.
|
||||
|
@ -77,8 +80,27 @@ public class MapBackedProvider implements InstrumentationContextProvider {
|
|||
this.instrumenter = instrumenter;
|
||||
}
|
||||
|
||||
@Override
|
||||
public AsmVisitorWrapper getInstrumentationVisitor() {
|
||||
public AgentBuilder.Identified.Extendable instrumentationTransformer(
|
||||
AgentBuilder.Identified.Extendable builder) {
|
||||
if (instrumenter.contextStore().size() > 0) {
|
||||
builder =
|
||||
builder.transform(
|
||||
new AgentBuilder.Transformer() {
|
||||
@Override
|
||||
public DynamicType.Builder<?> transform(
|
||||
DynamicType.Builder<?> builder,
|
||||
TypeDescription typeDescription,
|
||||
ClassLoader classLoader,
|
||||
JavaModule module) {
|
||||
return builder.visit(getInstrumentationVisitor());
|
||||
}
|
||||
});
|
||||
builder = builder.transform(new HelperInjector(dynamicClasses()));
|
||||
}
|
||||
return builder;
|
||||
}
|
||||
|
||||
private AsmVisitorWrapper getInstrumentationVisitor() {
|
||||
return new AsmVisitorWrapper() {
|
||||
@Override
|
||||
public int mergeWriter(int flags) {
|
||||
|
@ -217,13 +239,13 @@ public class MapBackedProvider implements InstrumentationContextProvider {
|
|||
};
|
||||
}
|
||||
|
||||
@Override
|
||||
public Map<String, byte[]> dynamicClasses() {
|
||||
private Map<String, byte[]> dynamicClasses() {
|
||||
return createOrGetDynamicClasses();
|
||||
}
|
||||
|
||||
@Override
|
||||
public AgentBuilder additionalInstrumentation(AgentBuilder builder) {
|
||||
public AgentBuilder.Identified.Extendable additionalInstrumentation(
|
||||
AgentBuilder.Identified.Extendable builder) {
|
||||
return builder;
|
||||
}
|
||||
|
||||
|
|
Loading…
Reference in New Issue