Javadoc for map-backed context store. Type info on context store api
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
@@ -37,7 +37,7 @@ public class InstrumentationContext {
* @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) {
+ public static <K, V> V get(K userInstance, Class<K> userClass, Class<V> contextClass) {
throw new RuntimeException("calls to this method will be rewritten");
}
}

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
@@ -10,10 +10,11 @@ public interface InstrumentationContextProvider {
AsmVisitorWrapper getInstrumentationVisitor();
/**
- * @return A list of classes in byte-array format. These classes will be injected into the runtime
- * classloader.
+ * @return A map of dynamic-class-name -> dynamic-class-bytes. These classes will be injected into
+ * the runtime classloader.
*/
Map<String, byte[]> dynamicClasses();
+ /** Hook for the context impl to define additional instrumentation if needed. */
AgentBuilder additionalInstrumentation(AgentBuilder builder);
}

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
@@ -26,7 +26,28 @@ import net.bytebuddy.jar.asm.Opcodes;
import net.bytebuddy.jar.asm.Type;
import net.bytebuddy.pool.TypePool;
-/** InstrumentationContextProvider which stores context in a global map. */
+/**
+ * InstrumentationContextProvider which stores context in a global map.
+ *
+ * <p>This is accomplished by
+ *
+ * <ol>
+ * <li>Injecting a Dynamic Class to store a static map
+ * <li>Rewritting calls to the context-store to access the map on the dynamic class
+ * </ol>
+ *
+ * Storing the map on a dynamic class and doing bytecode rewrites allows for a 1-pass lookup.
+ * Without bytecode transformations a 2-pass lookup would be required.
+ *
+ * <p>Example:<br>
+ * <em>InstrumentationContext.get(runnableInstance, Runnable.class, RunnableState.class)")</em><br>
+ * is rewritten to:<br>
+ * <em>RunnableInstrumentation$ContextStore$RunnableState12345.getOrCreate(runnableInstance,
+ * Runnable.class, RunnableState.class)</em>
+ *
+ * <p>Map lookup implementation defined in template class: {@link MapHolder#getOrCreate(Object,
+ * Class, Class)}
+ */
@Slf4j
public class MapBackedProvider implements InstrumentationContextProvider {
private static final Method contextGetMethod;
@@ -144,11 +165,15 @@ public class MapBackedProvider implements InstrumentationContextProvider {
}
}
+ /** Tracking the most recently used opcodes to assert proper api usage. */
private void pushOpcode(final int opcode) {
System.arraycopy(insnStack, 0, insnStack, 1, insnStack.length - 1);
insnStack[0] = opcode;
}
+ /**
+ * Tracking the most recently pushed objects on the stack to assert proper api usage.
+ */
private void pushStack(Object o) {
System.arraycopy(stack, 0, stack, 1, stack.length - 1);
stack[0] = o;
@@ -242,6 +267,12 @@ public class MapBackedProvider implements InstrumentationContextProvider {
private static final class MapHolder {
public static final WeakMap MAP = WeakMap.Provider.newWeakMap();
+ /**
+ * Fetch a context class out of the backing map. Create and return a new context class if none
+ * currently exists.
+ *
+ * <p>This method is thread safe.
+ */
public static Object getOrCreate(Object instance, Class userClass, Class contextClass) {
if (!userClass.isAssignableFrom(instance.getClass())) {
throw new RuntimeException(

MODIFY
dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/context/MapBackedProviderTest.groovy
dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/context/MapBackedProviderTest.groovy
@@ -120,13 +120,6 @@ class MapBackedProviderTest extends Specification {
thrown RuntimeException
}
- def "context store fails if runtime types are incorrect" () {
- when:
- ClassToRemap.mapIncorrectObject()
- then:
- thrown RuntimeException
- }
-
static class TestInstrumenter extends Instrumenter.Default {
TestInstrumenter() {
super("test")

MODIFY
dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/test/context/ClassToRemap.java
dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/test/context/ClassToRemap.java
@@ -25,15 +25,6 @@ public class ClassToRemap {
return ++state.anInt;
}
- /**
- * Instance passed to the context api does not extend the user class. This will throw an
- * exception.
- */
- public static int mapIncorrectObject() {
- State state = InstrumentationContext.get(new Object(), Runnable.class, State.class);
- return state.anInt;
- }
-
public static class State {
public int anInt = 0;
public Object anObject = new Object();

Javadoc for map-backed context store. Type info on context store api