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