Add some comments and some other minor CR tweaks
MODIFY
dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/CallableWrapper.java
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);

MODIFY
dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/RunnableWrapper.java
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);

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

MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/FieldBackedProvider.java
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
- // classloader where context class may be unavailable
+ // injected onto Bootstrap 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
- // classloader where context class may be unavailable
+ // onto Bootstrap 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)

MODIFY
dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/ExecutorInstrumentation.java
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("java.util.concurrent.Callable", State.class.getName());
- map.put("java.util.concurrent.Future", State.class.getName());
+ map.put(Runnable.class.getName(), State.class.getName());
+ map.put(Callable.class.getName(), State.class.getName());
+ map.put(Future.class.getName(), State.class.getName());
return Collections.unmodifiableMap(map);
}

MODIFY
dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/FutureInstrumentation.java
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);
}

RENAME
dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/RunnableInstrumentation.java
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("java.util.concurrent.Callable", State.class.getName());
+ map.put(Runnable.class.getName(), State.class.getName());
+ map.put(Callable.class.getName(), State.class.getName());
return Collections.unmodifiableMap(map);
}

MODIFY
dd-java-agent/testing/src/test/groovy/context/FieldBackedProviderTest.groovy
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 = instance.incrementContextCount()
- WeakReference<KeyClass> instanceRef = new WeakReference(instance)
- instance = null
+ final int count = keyValue.get().incrementContextCount()
+ WeakReference<KeyClass> instanceRef = new WeakReference(keyValue.get())
+ keyValue.set(null)
TestUtils.awaitGC(instanceRef)
then:
instanceRef.get() == null
count == 1
- }
-
- // can't use spock 'where' in this test because that creates strong references
- def "backing map should not create strong refs to untransformable key class instances"() {
- when:
- UntransformableKeyClass instance = new UntransformableKeyClass()
- final int count = instance.incrementContextCount()
- WeakReference<KeyClass> instanceRef = new WeakReference(instance)
- instance = null
- TestUtils.awaitGC(instanceRef)
- then:
- instanceRef.get() == null
- count == 1
+ where:
+ keyValue | _
+ new AtomicReference(new KeyClass()) | _
+ new AtomicReference(new UntransformableKeyClass()) | _
}
def "context classes are retransform safe"() {

Add some comments and some other minor CR tweaks