Cleanup and javadoc

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

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 = applyContextStoreTransform(agentBuilder);
-        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 = 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

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

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