Cleanup and javadoc
MODIFY
dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/InstrumentationContext.java
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.
+ *
+ * <p>
+ *
+ * <p>Conceptually, this can be thought of as a two pass map look up.
+ *
+ * <p>For example: <em>RunnableState runnableState = get(runnableImpl, Runnable.class,
+ * RunnableState.class)</em> --> <em>RunnableState runnableState = (RunnableState)
+ * GlobalContextMap.get(Runnable.class).get(runnableImpl)</em>
+ *
+ * <p>
+ *
+ * <p>However, the implementation is actually provided by bytecode transformation for performance
+ * reasons.
+ *
+ * <p>
+ *
+ * <p>Context classes are weakly referenced and will be garbage collected when their corresponding
+ * user instance is collected.
+ *
+ * <p>
+ *
+ * <p>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 <K> user class
+ * @param <V> context class
+ * @return The context instance attached to userInstance.
+ */
public static <K, V> V get(Object userInstance, Class<K> userClass, Class<V> contextClass) {
- throw new RuntimeException("calls to this method should be rewritten");
+ throw new RuntimeException("calls to this method will be rewritten");
}
}

MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java
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<String> 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

MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/InstrumentationContextProvider.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/InstrumentationContextProvider.java
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<String, byte[]> 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);
- }
- }
}

MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/MapBackedProvider.java
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<Map<String, byte[]>> dynamicClasses = new AtomicReference<>(null);
-
- /** user-class-name -> dynamic-class-name */
- private final AtomicReference<Map<String, String>> dynamicClassNames =
- new AtomicReference<>(null);
static {
try {
@@ -49,6 +43,13 @@ public class MapBackedProvider implements InstrumentationContextProvider {
}
}
+ /** dynamic-class-name -> dynamic-class-bytes */
+ private final AtomicReference<Map<String, byte[]>> dynamicClasses = new AtomicReference<>(null);
+
+ /** user-class-name -> dynamic-class-name */
+ private final AtomicReference<Map<String, String>> 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);
}
}
}

Cleanup and javadoc