Move ASM visiting method into context-store impl

This commit is contained in:
Andrew Kent 2018-10-26 14:06:57 -07:00
parent cc27f1507e
commit b89aa54a93
3 changed files with 36 additions and 37 deletions

View File

@ -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

View File

@ -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);
}

View File

@ -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;
}