diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/InstrumentationContext.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/InstrumentationContext.java index 8bb5719ff0..ff3e57b432 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/InstrumentationContext.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/InstrumentationContext.java @@ -4,7 +4,40 @@ package datadog.trace.bootstrap; public class InstrumentationContext { private InstrumentationContext() {} + /** + * Fetch a context instance out of the context store. + * + *

+ * + *

Conceptually, this can be thought of as a two pass map look up. + * + *

For example: RunnableState runnableState = get(runnableImpl, Runnable.class, + * RunnableState.class) --> RunnableState runnableState = (RunnableState) + * GlobalContextMap.get(Runnable.class).get(runnableImpl) + * + *

+ * + *

However, the implementation is actually provided by bytecode transformation for performance + * reasons. + * + *

+ * + *

Context classes are weakly referenced and will be garbage collected when their corresponding + * user instance is collected. + * + *

+ * + *

Instrumenters making this call must define the user-context class relationship in + * datadog.trace.agent.tooling.Instrumenter.Default#contextStore. + * + * @param userInstance The instance to store context on. + * @param userClass The user class context is attached to. + * @param contextClass The context class attached to the user class. + * @param user class + * @param context class + * @return The context instance attached to userInstance. + */ public static V get(Object userInstance, Class userClass, Class contextClass) { - throw new RuntimeException("calls to this method should be rewritten"); + throw new RuntimeException("calls to this method will be rewritten"); } } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java index dfea6aedb2..b0a8d4defe 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java @@ -5,6 +5,7 @@ import static datadog.trace.agent.tooling.Utils.getConfigEnabled; import static net.bytebuddy.matcher.ElementMatchers.any; import datadog.trace.agent.tooling.context.InstrumentationContextProvider; +import datadog.trace.agent.tooling.context.MapBackedProvider; import datadog.trace.agent.tooling.muzzle.Reference; import datadog.trace.agent.tooling.muzzle.ReferenceMatcher; import java.security.ProtectionDomain; @@ -52,7 +53,6 @@ public interface Instrumenter { abstract class Default implements Instrumenter { private final Set instrumentationNames; private final String instrumentationPrimaryName; - // TODO: create a generic instrumentaton pipeline attacher? private final InstrumentationContextProvider contextProvider; protected final boolean enabled; @@ -78,7 +78,7 @@ public interface Instrumenter { } } enabled = anyEnabled; - contextProvider = InstrumentationContextProvider.Creator.contextProviderFor(this); + contextProvider = new MapBackedProvider(this); } @Override @@ -101,21 +101,7 @@ public interface Instrumenter { .and(new MuzzleMatcher()) .transform(DDTransformers.defaultTransformers()); agentBuilder = injectHelperClasses(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())); - } + agentBuilder = applyContextStoreTransform(agentBuilder); agentBuilder = applyInstrumentationTransformers(agentBuilder); return agentBuilder; } @@ -142,6 +128,26 @@ 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 diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/InstrumentationContextProvider.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/InstrumentationContextProvider.java index 90dddfd6b3..b3b3ddfc30 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/InstrumentationContextProvider.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/InstrumentationContextProvider.java @@ -1,6 +1,5 @@ package datadog.trace.agent.tooling.context; -import datadog.trace.agent.tooling.Instrumenter; import java.util.Map; import net.bytebuddy.agent.builder.AgentBuilder; import net.bytebuddy.asm.AsmVisitorWrapper; @@ -17,14 +16,4 @@ public interface InstrumentationContextProvider { Map dynamicClasses(); AgentBuilder additionalInstrumentation(AgentBuilder builder); - - // TODO: better place to put factory/creator - class Creator { - private Creator() {} - - public static InstrumentationContextProvider contextProviderFor( - Instrumenter.Default instrumenter) { - return new MapBackedProvider(instrumenter); - } - } } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/MapBackedProvider.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/MapBackedProvider.java index c912c647b6..a886f0b933 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/MapBackedProvider.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/MapBackedProvider.java @@ -31,12 +31,6 @@ import net.bytebuddy.pool.TypePool; public class MapBackedProvider implements InstrumentationContextProvider { private static final Method contextGetMethod; private static final Method mapGetMethod; - /** dynamic-class-name -> dynamic-class-bytes */ - private final AtomicReference> dynamicClasses = new AtomicReference<>(null); - - /** user-class-name -> dynamic-class-name */ - private final AtomicReference> dynamicClassNames = - new AtomicReference<>(null); static { try { @@ -49,6 +43,13 @@ public class MapBackedProvider implements InstrumentationContextProvider { } } + /** dynamic-class-name -> dynamic-class-bytes */ + private final AtomicReference> dynamicClasses = new AtomicReference<>(null); + + /** user-class-name -> dynamic-class-name */ + private final AtomicReference> dynamicClassNames = + new AtomicReference<>(null); + private final Instrumenter.Default instrumenter; public MapBackedProvider(Instrumenter.Default instrumenter) { @@ -258,7 +259,8 @@ public class MapBackedProvider implements InstrumentationContextProvider { contextInstance = contextClass.newInstance(); MAP.put(instance, contextInstance); } catch (Exception e) { - throw new RuntimeException(e); + throw new RuntimeException( + contextClass.getName() + " must define a public, no-arg constructor.", e); } } }