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

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

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

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(

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")

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();