Add some comments and some other minor CR tweaks

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/CallableWrapper.java

@@ -22,6 +22,8 @@ public final class CallableWrapper implements Callable {
   }
 
   public static Callable<?> wrapIfNeeded(final Callable<?> task) {
+    // We wrap only lambdas' anonymous classes and if given object has not already been wrapped.
+    // Anonymous classes have '/' in class name which is not allowed in 'normal' classes.
     if (task.getClass().getName().contains("/") && (!(task instanceof CallableWrapper))) {
       log.debug("Wrapping callable task {}", task);
       return new CallableWrapper(task);

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/RunnableWrapper.java

@@ -21,6 +21,8 @@ public final class RunnableWrapper implements Runnable {
   }
 
   public static Runnable wrapIfNeeded(final Runnable task) {
+    // We wrap only lambdas' anonymous classes and if given object has not already been wrapped.
+    // Anonymous classes have '/' in class name which is not allowed in 'normal' classes.
     if (task.getClass().getName().contains("/") && (!(task instanceof RunnableWrapper))) {
       log.debug("Wrapping runnable task {}", task);
       return new RunnableWrapper(task);

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Utils.java

@@ -92,7 +92,7 @@ public class Utils {
 
   /**
    * Convert class name to a format that can be used as part of inner class name by replacing all
-   * ','s with '$'s.
+   * '.'s with '$'s.
    *
    * @param className class named to be converted
    * @return convertd name

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/FieldBackedProvider.java

@@ -64,7 +64,7 @@ import net.bytebuddy.utility.JavaModule;
  * <p>Example:<br>
  * <em>InstrumentationContext.get(Runnable.class, RunnableState.class)")</em><br>
  * is rewritten to:<br>
- * <em>RunnableInstrumentation$ContextStore$Runnable$RunnableState12345.getContextStore(runnableRunnable.class,
+ * <em>FieldBackedProvider$ContextStore$Runnable$RunnableState12345.getContextStore(runnableRunnable.class,
  * RunnableState.class)</em>
  */
 @Slf4j
@@ -88,6 +88,7 @@ public class FieldBackedProvider implements InstrumentationContextProvider {
   }
 
   private final Instrumenter.Default instrumenter;
+  private final ByteBuddy byteBuddy;
 
   /** fields-accessor-interface-name -> fields-accessor-interface-dynamic-type */
   private final Map<String, DynamicType.Unloaded<?>> fieldAccessorInterfaces;
@@ -99,6 +100,7 @@ public class FieldBackedProvider implements InstrumentationContextProvider {
 
   public FieldBackedProvider(final Instrumenter.Default instrumenter) {
     this.instrumenter = instrumenter;
+    byteBuddy = new ByteBuddy();
     fieldAccessorInterfaces = generateFieldAccessorInterfaces();
     contextStoreImplementations = generateContextStoreImplementationClasses();
     fieldInjectionEnabled = Config.get().isRuntimeContextFieldInjection();
@@ -108,7 +110,12 @@ public class FieldBackedProvider implements InstrumentationContextProvider {
   public AgentBuilder.Identified.Extendable instrumentationTransformer(
       AgentBuilder.Identified.Extendable builder) {
     if (instrumenter.contextStore().size() > 0) {
-      builder = builder.transform(getTransformerForASMVisitor(getInstrumentationVisitor()));
+      /*
+      Install transformer that rewrites accesses to context store with specialized bytecode that invokes appropriate
+      storage implementation.
+       */
+      builder =
+          builder.transform(getTransformerForASMVisitor(getContextStoreReadsRewritingVisitor()));
 
       /*
       We inject into bootstrap classloader because field accessor interfaces are needed by
@@ -149,7 +156,7 @@ public class FieldBackedProvider implements InstrumentationContextProvider {
     };
   }
 
-  private AsmVisitorWrapper getInstrumentationVisitor() {
+  private AsmVisitorWrapper getContextStoreReadsRewritingVisitor() {
     return new AsmVisitorWrapper() {
       @Override
       public int mergeWriter(final int flags) {
@@ -211,6 +218,12 @@ public class FieldBackedProvider implements InstrumentationContextProvider {
                     && CONTEXT_GET_METHOD.getName().equals(name)
                     && Type.getMethodDescriptor(CONTEXT_GET_METHOD).equals(descriptor)) {
                   log.debug("Found context-store access in {}", instrumenter.getClass().getName());
+                  /*
+                  The idea here is that the rest if this method visitor collects last three instructions in `insnStack`
+                  variable. Once we get here we check if those last three instructions constitute call that looks like
+                  `InstrumentationContext.get(K.class, V.class)`. If it does the inside of this if rewrites it to call
+                  dynamically injected context store implementation instead.
+                   */
                   if ((insnStack[0] == Opcodes.INVOKESTATIC
                           && insnStack[1] == Opcodes.LDC
                           && insnStack[2] == Opcodes.LDC)
@@ -312,6 +325,10 @@ public class FieldBackedProvider implements InstrumentationContextProvider {
 
     if (fieldInjectionEnabled) {
       for (final Map.Entry<String, String> entry : instrumenter.contextStore().entrySet()) {
+        /*
+        For each context store defined in a current instrumentation we create an agent builder
+        that injects necessary fields.
+         */
         builder =
             builder
                 .type(
@@ -375,8 +392,7 @@ public class FieldBackedProvider implements InstrumentationContextProvider {
           final int readerFlags) {
         return new ClassVisitor(Opcodes.ASM7, classVisitor) {
           // We are using Object class name instead of contextClassName here because this gets
-          // injected onto Bootstrap
+          // injected onto Bootstrap classloader where context class may be unavailable
-          // classloader where context class may be unavailable
           private final TypeDescription contextType =
               new TypeDescription.ForLoadedType(Object.class);
           private final String fieldName = getContextFieldName(keyClassName);
@@ -502,7 +518,7 @@ public class FieldBackedProvider implements InstrumentationContextProvider {
    */
   private DynamicType.Unloaded<?> makeContextStoreImplementationClass(
       final String keyClassName, final String contextClassName) {
-    return new ByteBuddy()
+    return byteBuddy
         .rebase(ContextStoreImplementationTemplate.class)
         .modifiers(Visibility.PUBLIC, TypeManifestation.FINAL)
         .name(getContextStoreImplementationClassName(keyClassName, contextClassName))
@@ -859,10 +875,9 @@ public class FieldBackedProvider implements InstrumentationContextProvider {
   private DynamicType.Unloaded<?> makeFieldAccessorInterface(
       final String keyClassName, final String contextClassName) {
     // We are using Object class name instead of contextClassName here because this gets injected
-    // onto Bootstrap
+    // onto Bootstrap classloader where context class may be unavailable
-    // classloader where context class may be unavailable
     final TypeDescription contextType = new TypeDescription.ForLoadedType(Object.class);
-    return new ByteBuddy()
+    return byteBuddy
         .makeInterface()
         .name(getContextAccessorInterfaceName(keyClassName, contextClassName))
         .defineMethod(getContextGetterName(keyClassName), contextType, Visibility.PUBLIC)

dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/ExecutorInstrumentation.java

@@ -135,9 +135,9 @@ public final class ExecutorInstrumentation extends Instrumenter.Default {
   @Override
   public Map<String, String> contextStore() {
     final Map<String, String> map = new HashMap<>();
-    map.put("java.lang.Runnable", State.class.getName());
+    map.put(Runnable.class.getName(), State.class.getName());
-    map.put("java.util.concurrent.Callable", State.class.getName());
+    map.put(Callable.class.getName(), State.class.getName());
-    map.put("java.util.concurrent.Future", State.class.getName());
+    map.put(Future.class.getName(), State.class.getName());
     return Collections.unmodifiableMap(map);
   }
 

dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/FutureInstrumentation.java

@@ -93,7 +93,7 @@ public final class FutureInstrumentation extends Instrumenter.Default {
   @Override
   public Map<String, String> contextStore() {
     final Map<String, String> map = new HashMap<>();
-    map.put("java.util.concurrent.Future", State.class.getName());
+    map.put(Future.class.getName(), State.class.getName());
     return Collections.unmodifiableMap(map);
   }
 

dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/RunnableCallableInstrumentation.java

@@ -23,9 +23,9 @@ import net.bytebuddy.matcher.ElementMatcher;
 
 @Slf4j
 @AutoService(Instrumenter.class)
-public final class RunnableInstrumentation extends Instrumenter.Default {
+public final class RunnableCallableInstrumentation extends Instrumenter.Default {
 
-  public RunnableInstrumentation() {
+  public RunnableCallableInstrumentation() {
     super(ExecutorInstrumentation.EXEC_NAME);
   }
 
@@ -38,15 +38,15 @@ public final class RunnableInstrumentation extends Instrumenter.Default {
   @Override
   public String[] helperClassNames() {
     return new String[] {
-      RunnableInstrumentation.class.getName() + "$RunnableUtils",
+      RunnableCallableInstrumentation.class.getName() + "$RunnableUtils",
     };
   }
 
   @Override
   public Map<String, String> contextStore() {
     final Map<String, String> map = new HashMap<>();
-    map.put("java.lang.Runnable", State.class.getName());
+    map.put(Runnable.class.getName(), State.class.getName());
-    map.put("java.util.concurrent.Callable", State.class.getName());
+    map.put(Callable.class.getName(), State.class.getName());
     return Collections.unmodifiableMap(map);
   }
 

dd-java-agent/testing/src/test/groovy/context/FieldBackedProviderTest.groovy

@@ -8,6 +8,7 @@ import net.bytebuddy.utility.JavaModule
 import java.lang.instrument.ClassDefinition
 import java.lang.ref.WeakReference
 import java.lang.reflect.Field
+import java.util.concurrent.atomic.AtomicReference
 
 import static context.ContextTestInstrumentation.IncorrectCallUsageKeyClass
 import static context.ContextTestInstrumentation.IncorrectContextClassUsageKeyClass
@@ -93,31 +94,21 @@ class FieldBackedProviderTest extends AgentTestRunner {
     new UntransformableKeyClass() | _
   }
 
-  def "backing map should not create strong refs to key class instances"() {
+  def "backing map should not create strong refs to key class instances #keyValue.get().getClass().getName()"() {
     when:
-    KeyClass instance = new KeyClass()
+    final int count = keyValue.get().incrementContextCount()
-    final int count = instance.incrementContextCount()
+    WeakReference<KeyClass> instanceRef = new WeakReference(keyValue.get())
-    WeakReference<KeyClass> instanceRef = new WeakReference(instance)
+    keyValue.set(null)
-    instance = null
     TestUtils.awaitGC(instanceRef)
 
     then:
     instanceRef.get() == null
     count == 1
-  }
 
-  // can't use spock 'where' in this test because that creates strong references
+    where:
-  def "backing map should not create strong refs to untransformable key class instances"() {
+    keyValue                                           | _
-    when:
+    new AtomicReference(new KeyClass())                | _
-    UntransformableKeyClass instance = new UntransformableKeyClass()
+    new AtomicReference(new UntransformableKeyClass()) | _
-    final int count = instance.incrementContextCount()
-    WeakReference<KeyClass> instanceRef = new WeakReference(instance)
-    instance = null
-    TestUtils.awaitGC(instanceRef)
-
-    then:
-    instanceRef.get() == null
-    count == 1
   }
 
   def "context classes are retransform safe"() {