First iteration on a better context store api
Separate context storage from actual fetching/putting
ADD
dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/ContextStore.java
+package datadog.trace.bootstrap;
+
+/**
+ * Interface to represent context storage for instrumentations.
+ *
+ * <p>Context instances are weakly referenced and will be garbage collected when their corresponding
+ * key instance is collected.
+ *
+ * @param <K> key type to do context lookups
+ * @param <C> context type
+ */
+public interface ContextStore<K, C> {
+
+ /**
+ * Factory interface to create context instances
+ *
+ * @param <C> context type
+ */
+ interface Factory<C> {
+
+ /** @return new context instance */
+ C create();
+ }
+
+ /**
+ * Get context given the key
+ *
+ * @param key the key to looup
+ * @return context object
+ */
+ C get(K key);
+
+ /**
+ * Put new context instance for given key
+ *
+ * @param key key to use
+ * @param context context instance to save
+ */
+ void put(K key, C context);
+
+ /**
+ * Put new context instance if key is absent
+ *
+ * @param key key to use
+ * @param context new context instance to put
+ * @return old instance if it was present, or new instance
+ */
+ C putIfAbsent(K key, C context);
+
+ /**
+ * Put new context instance if key is absent. Uses context factory to avoid creating objects if
+ * not needed.
+ *
+ * @param key key to use
+ * @param contextFactory factory instance to produce new context object
+ * @return old instance if it was present, or new instance
+ */
+ C putIfAbsent(K key, Factory<C> contextFactory);
+}

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
@@ -5,39 +5,22 @@ public class InstrumentationContext {
private InstrumentationContext() {}
/**
- * Fetch a context instance out of the context store.
+ * Find a {@link ContextStore} instance for given key class and context class.
*
- * <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>Conceptually this can be thought of as a map lookup to fetch a second level map given
+ * keyClass.
*
* <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 keyClass The key 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.
+ * @param <K> key class
+ * @param <C> context class
+ * @return The instance of context store for given arguments.
*/
- public static <K, V> V get(K userInstance, Class<K> userClass, Class<V> contextClass) {
- throw new RuntimeException("calls to this method will be rewritten");
+ public static <K, C> ContextStore<K, C> get(
+ final Class<K> keyClass, final Class<C> contextClass) {
+ throw new RuntimeException("Calls to this method will be rewritten by MapBackedProvider");
}
}

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
@@ -5,6 +5,7 @@ import static datadog.trace.agent.tooling.ClassLoaderMatcher.BOOTSTRAP_CLASSLOAD
import datadog.trace.agent.tooling.HelperInjector;
import datadog.trace.agent.tooling.Instrumenter;
import datadog.trace.agent.tooling.Utils;
+import datadog.trace.bootstrap.ContextStore;
import datadog.trace.bootstrap.InstrumentationContext;
import datadog.trace.bootstrap.WeakMap;
import java.lang.reflect.Method;
@@ -37,8 +38,8 @@ import net.bytebuddy.utility.JavaModule;
* <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
+ * <li>Injecting a Dynamic Class created from {@link MapHolder} to store a static map
+ * <li>Rewritting calls to the context-store to access the specific dynamic {@link MapHolder}
* </ol>
*
* Storing the map on a dynamic class and doing bytecode rewrites allows for a 1-pass lookup.
@@ -47,11 +48,8 @@ import net.bytebuddy.utility.JavaModule;
* <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)}
+ * <em>RunnableInstrumentation$ContextStore$RunnableState12345.getMapHolder(runnableRunnable.class,
+ * RunnableState.class)</em>
*/
@Slf4j
public class MapBackedProvider implements InstrumentationContextProvider {
@@ -60,11 +58,9 @@ public class MapBackedProvider implements InstrumentationContextProvider {
static {
try {
- contextGetMethod =
- InstrumentationContext.class.getMethod("get", Object.class, Class.class, Class.class);
- mapGetMethod =
- MapHolder.class.getMethod("getOrCreate", Object.class, Class.class, Class.class);
- } catch (Exception e) {
+ contextGetMethod = InstrumentationContext.class.getMethod("get", Class.class, Class.class);
+ mapGetMethod = MapHolder.class.getMethod("getMapHolder", Class.class, Class.class);
+ } catch (final Exception e) {
throw new IllegalStateException(e);
}
}
@@ -72,16 +68,17 @@ 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 */
+ /** key-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) {
+ public MapBackedProvider(final Instrumenter.Default instrumenter) {
this.instrumenter = instrumenter;
}
+ @Override
public AgentBuilder.Identified.Extendable instrumentationTransformer(
AgentBuilder.Identified.Extendable builder) {
if (instrumenter.contextStore().size() > 0) {
@@ -90,10 +87,10 @@ public class MapBackedProvider implements InstrumentationContextProvider {
new AgentBuilder.Transformer() {
@Override
public DynamicType.Builder<?> transform(
- DynamicType.Builder<?> builder,
- TypeDescription typeDescription,
- ClassLoader classLoader,
- JavaModule module) {
+ final DynamicType.Builder<?> builder,
+ final TypeDescription typeDescription,
+ final ClassLoader classLoader,
+ final JavaModule module) {
return builder.visit(getInstrumentationVisitor());
}
});
@@ -104,10 +101,10 @@ public class MapBackedProvider implements InstrumentationContextProvider {
@Override
public DynamicType.Builder<?> transform(
- DynamicType.Builder<?> builder,
- TypeDescription typeDescription,
- ClassLoader classLoader,
- JavaModule module) {
+ final DynamicType.Builder<?> builder,
+ final TypeDescription typeDescription,
+ final ClassLoader classLoader,
+ final JavaModule module) {
return injector.transform(
builder,
typeDescription,
@@ -123,41 +120,45 @@ public class MapBackedProvider implements InstrumentationContextProvider {
private AsmVisitorWrapper getInstrumentationVisitor() {
return new AsmVisitorWrapper() {
@Override
- public int mergeWriter(int flags) {
+ public int mergeWriter(final int flags) {
return flags | ClassWriter.COMPUTE_MAXS;
}
@Override
- public int mergeReader(int flags) {
+ public int mergeReader(final int flags) {
return flags;
}
@Override
public ClassVisitor wrap(
final TypeDescription instrumentedType,
- ClassVisitor classVisitor,
- Implementation.Context implementationContext,
- TypePool typePool,
- FieldList<FieldDescription.InDefinedShape> fields,
- MethodList<?> methods,
- int writerFlags,
- int readerFlags) {
+ final ClassVisitor classVisitor,
+ final Implementation.Context implementationContext,
+ final TypePool typePool,
+ final FieldList<FieldDescription.InDefinedShape> fields,
+ final MethodList<?> methods,
+ final int writerFlags,
+ final int readerFlags) {
generateMapHolderClasses();
return new ClassVisitor(Opcodes.ASM7, classVisitor) {
@Override
public void visit(
- int version,
- int access,
- String name,
- String signature,
- String superName,
- String[] interfaces) {
+ final int version,
+ final int access,
+ final String name,
+ final String signature,
+ final String superName,
+ final String[] interfaces) {
super.visit(version, access, name, signature, superName, interfaces);
}
@Override
public MethodVisitor visitMethod(
- int access, String name, String descriptor, String signature, String[] exceptions) {
+ final int access,
+ final String name,
+ final String descriptor,
+ final String signature,
+ final String[] exceptions) {
final MethodVisitor mv =
super.visitMethod(access, name, descriptor, signature, exceptions);
return new MethodVisitor(Opcodes.ASM7, mv) {
@@ -168,7 +169,11 @@ public class MapBackedProvider implements InstrumentationContextProvider {
@Override
public void visitMethodInsn(
- int opcode, String owner, String name, String descriptor, boolean isInterface) {
+ final int opcode,
+ final String owner,
+ final String name,
+ final String descriptor,
+ final boolean isInterface) {
pushOpcode(opcode);
if (Utils.getInternalName(contextGetMethod.getDeclaringClass().getName())
.equals(owner)
@@ -180,17 +185,28 @@ public class MapBackedProvider implements InstrumentationContextProvider {
&& insnStack[2] == Opcodes.LDC)
&& (stack[0] instanceof Type && stack[1] instanceof Type)) {
final String contextClassName = ((Type) stack[0]).getClassName();
- final String userClassName = ((Type) stack[1]).getClassName();
- final String mapHolderClass = dynamicClassNames.get().get(userClassName);
+ final String keyClassName = ((Type) stack[1]).getClassName();
+ final String mapHolderClass = dynamicClassNames.get().get(keyClassName);
log.debug(
"Rewriting context-store map fetch for instrumenter {}: {} -> {}",
instrumenter.getClass().getName(),
- userClassName,
+ keyClassName,
contextClassName);
- if (mapHolderClass != null
- && contextClassName.equals(
- instrumenter.contextStore().get(userClassName))) {
- // stack: contextClass | userClass | instance
+ if (mapHolderClass == null) {
+ throw new IllegalStateException(
+ String.format(
+ "Incorrect Context Api Usage detected. Cannot find map holder class for %s. Was that class defined in contextStore for instrumentation %s?",
+ keyClassName, instrumenter.getClass().getName()));
+ }
+ if (!contextClassName.equals(instrumenter.contextStore().get(keyClassName))) {
+ throw new IllegalStateException(
+ String.format(
+ "Incorrect Context Api Usage detected. Incorrect context class %s, expected %s for instrumentation %s",
+ contextClassName,
+ instrumenter.contextStore().get(keyClassName),
+ instrumenter.getClass().getName()));
+ }
+ // stack: contextClass | keyClass
mv.visitMethodInsn(
Opcodes.INVOKESTATIC,
Utils.getInternalName(mapHolderClass),
@@ -199,9 +215,8 @@ public class MapBackedProvider implements InstrumentationContextProvider {
false);
return;
}
- }
throw new IllegalStateException(
- "Incorrect Context Api Usage detected. User and context class must be class-literals. Example of correct usage: InstrumentationContext.get(runnableInstance, Runnable.class, RunnableState.class)");
+ "Incorrect Context Api Usage detected. Key and context class must be class-literals. Example of correct usage: InstrumentationContext.get(Runnable.class, RunnableContext.class)");
} else {
super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
}
@@ -216,38 +231,38 @@ public class MapBackedProvider implements InstrumentationContextProvider {
/**
* Tracking the most recently pushed objects on the stack to assert proper api usage.
*/
- private void pushStack(Object o) {
+ private void pushStack(final Object o) {
System.arraycopy(stack, 0, stack, 1, stack.length - 1);
stack[0] = o;
}
@Override
- public void visitInsn(int opcode) {
+ public void visitInsn(final int opcode) {
pushOpcode(opcode);
super.visitInsn(opcode);
}
@Override
- public void visitJumpInsn(int opcode, Label label) {
+ public void visitJumpInsn(final int opcode, final Label label) {
pushOpcode(opcode);
super.visitJumpInsn(opcode, label);
}
@Override
- public void visitIntInsn(int opcode, int operand) {
+ public void visitIntInsn(final int opcode, final int operand) {
pushOpcode(opcode);
super.visitIntInsn(opcode, operand);
}
@Override
- public void visitVarInsn(int opcode, int var) {
+ public void visitVarInsn(final int opcode, final int var) {
pushOpcode(opcode);
pushStack(var);
super.visitVarInsn(opcode, var);
}
@Override
- public void visitLdcInsn(Object value) {
+ public void visitLdcInsn(final Object value) {
pushOpcode(Opcodes.LDC);
pushStack(value);
super.visitLdcInsn(value);
@@ -265,7 +280,7 @@ public class MapBackedProvider implements InstrumentationContextProvider {
@Override
public AgentBuilder.Identified.Extendable additionalInstrumentation(
- AgentBuilder.Identified.Extendable builder) {
+ final AgentBuilder.Identified.Extendable builder) {
return builder;
}
@@ -281,14 +296,14 @@ public class MapBackedProvider implements InstrumentationContextProvider {
final Map<String, byte[]> dynamicClasses = new HashMap<>(instrumenter.contextStore().size());
final Map<String, String> dynamicClassNames =
new HashMap<>(instrumenter.contextStore().size());
- for (final String userClassName : instrumenter.contextStore().keySet()) {
+ for (final String keyClassName : instrumenter.contextStore().keySet()) {
final String dynamicClassName =
instrumenter.getClass().getName()
+ "$ContextStore"
- + userClassName.replaceAll(".*([^\\.]+)$", "\\1")
+ + keyClassName.replaceAll(".*([^\\.]+)$", "\\1")
+ UUID.randomUUID().toString().replace('-', '_');
- dynamicClassNames.put(userClassName, dynamicClassName);
+ dynamicClassNames.put(keyClassName, dynamicClassName);
dynamicClasses.put(dynamicClassName, makeMapHolderClass(dynamicClassName));
}
this.dynamicClassNames.compareAndSet(null, dynamicClassNames);
@@ -296,7 +311,7 @@ public class MapBackedProvider implements InstrumentationContextProvider {
}
}
- private byte[] makeMapHolderClass(String className) {
+ private byte[] makeMapHolderClass(final String className) {
return new ByteBuddy()
.rebase(MapHolder.class)
.modifiers(Opcodes.ACC_PUBLIC | Opcodes.ACC_FINAL)
@@ -306,41 +321,62 @@ public class MapBackedProvider implements InstrumentationContextProvider {
}
/** Template class used to generate the class holding the global map. */
- private static final class MapHolder {
- public static final WeakMap MAP = WeakMap.Provider.newWeakMap();
+ private static final class MapHolder implements ContextStore<Object, Object> {
+ private static final MapHolder INSTANCE = new MapHolder(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(
- "Illegal context lookup. "
- + instance.getClass().getName()
- + " cannot be cast to "
- + userClass.getName());
- }
- Object contextInstance = MAP.get(instance);
- if (null == contextInstance) {
- synchronized (instance) {
- contextInstance = MAP.get(instance);
- if (null == contextInstance) {
- try {
- contextInstance = contextClass.newInstance();
- MAP.put(instance, contextInstance);
- } catch (Exception e) {
- throw new RuntimeException(
- contextClass.getName() + " must define a public, no-arg constructor.", e);
+ private final WeakMap map;
+
+ private MapHolder(final WeakMap map) {
+ this.map = map;
+ }
+
+ @Override
+ public Object get(final Object key) {
+ return map.get(key);
+ }
+
+ @Override
+ public Object putIfAbsent(final Object key, final Object context) {
+ Object existingContext = map.get(key);
+ if (null != existingContext) {
+ return existingContext;
+ }
+ synchronized (map) {
+ existingContext = map.get(key);
+ if (null != existingContext) {
+ return existingContext;
+ }
+ map.put(key, context);
+ return context;
}
}
+
+ @Override
+ public Object putIfAbsent(final Object key, final ContextStore.Factory<Object> contextFactory) {
+ Object existingContext = map.get(key);
+ if (null != existingContext) {
+ return existingContext;
}
+ synchronized (map) {
+ existingContext = map.get(key);
+ if (null != existingContext) {
+ return existingContext;
+ }
+ final Object context = contextFactory.create();
+ map.put(key, context);
+ return context;
}
- return contextInstance;
}
- private MapHolder() {}
+ @Override
+ public void put(final Object key, final Object context) {
+ map.put(key, context);
+ }
+
+ public static ContextStore getMapHolder(final Class keyClass, final Class contextClass) {
+ // We do not actually check the keyClass here - but that should be fine since compiler would
+ // check things for us.
+ return INSTANCE;
+ }
}
}

MODIFY
dd-java-agent/instrumentation/http-url-connection/src/main/java/datadog/trace/instrumentation/http_url_connection/HttpUrlConnectionInstrumentation.java
dd-java-agent/instrumentation/http-url-connection/src/main/java/datadog/trace/instrumentation/http_url_connection/HttpUrlConnectionInstrumentation.java
@@ -12,6 +12,7 @@ import datadog.trace.agent.tooling.Instrumenter;
import datadog.trace.api.DDSpanTypes;
import datadog.trace.api.DDTags;
import datadog.trace.bootstrap.CallDepthThreadLocalMap;
+import datadog.trace.bootstrap.ContextStore;
import datadog.trace.bootstrap.InstrumentationContext;
import io.opentracing.Scope;
import io.opentracing.Span;
@@ -48,9 +49,13 @@ public class HttpUrlConnectionInstrumentation extends Instrumenter.Default {
@Override
public String[] helperClassNames() {
- return new String[] {HttpUrlConnectionInstrumentation.class.getName() + "$HttpURLState"};
+ return new String[] {
+ HttpUrlConnectionInstrumentation.class.getName() + "$HttpURLState",
+ HttpUrlConnectionInstrumentation.class.getName() + "$HttpURLState$1"
+ };
}
+ @Override
public Map<String, String> contextStore() {
return Collections.singletonMap(
"java.net.HttpURLConnection", getClass().getName() + "$HttpURLState");
@@ -72,9 +77,9 @@ public class HttpUrlConnectionInstrumentation extends Instrumenter.Default {
@Advice.This final HttpURLConnection thiz,
@Advice.FieldValue("connected") final boolean connected,
@Advice.Origin("#m") final String methodName) {
-
- final HttpURLState state =
- InstrumentationContext.get(thiz, HttpURLConnection.class, HttpURLState.class);
+ final ContextStore<HttpURLConnection, HttpURLState> contextStore =
+ InstrumentationContext.get(HttpURLConnection.class, HttpURLState.class);
+ final HttpURLState state = contextStore.putIfAbsent(thiz, HttpURLState.FACTORY);
String operationName = "http.request";
@@ -183,6 +188,15 @@ public class HttpUrlConnectionInstrumentation extends Instrumenter.Default {
}
public static class HttpURLState {
+
+ public static final ContextStore.Factory<HttpURLState> FACTORY =
+ new ContextStore.Factory<HttpURLState>() {
+ @Override
+ public HttpURLState create() {
+ return new HttpURLState();
+ }
+ };
+
public boolean calledOutputStream = false;
public boolean calledInputStream = false;
}

MODIFY
dd-java-agent/testing/src/test/groovy/context/MapBackedProviderTest.groovy
dd-java-agent/testing/src/test/groovy/context/MapBackedProviderTest.groovy
@@ -6,10 +6,15 @@ import net.bytebuddy.utility.JavaModule
import java.lang.ref.WeakReference
+import static context.ContextTestInstrumentation.IncorrectCallUsageKeyClass
+import static context.ContextTestInstrumentation.IncorrectContextClassUsageKeyClass
+import static context.ContextTestInstrumentation.IncorrectKeyClassUsageKeyClass
+import static context.ContextTestInstrumentation.KeyClass
+
class MapBackedProviderTest extends AgentTestRunner {
def setupSpec() {
- assert new UserClass1().isInstrumented()
+ assert new KeyClass().isInstrumented()
}
@Override
@@ -19,14 +24,16 @@ class MapBackedProviderTest extends AgentTestRunner {
final JavaModule module,
final boolean loaded,
final Throwable throwable) {
- // UserClass2 asserts on incorrect api usage. Error expected.
- return !(typeName.equals(UserClass2.getName()) && throwable.getMessage().startsWith("Incorrect Context Api Usage detected."))
+ // Incorrect* classes assert on incorrect api usage. Error expected.
+ return !(typeName.startsWith(ContextTestInstrumentation.getName() + "$Incorrect") && throwable.getMessage().startsWith("Incorrect Context Api Usage detected."))
}
def "correct api usage stores state in map"() {
+ setup:
+ KeyClass instance1 = new KeyClass()
+ KeyClass instance2 = new KeyClass()
+
when:
- UserClass1 instance1 = new UserClass1()
- UserClass1 instance2 = new UserClass1()
instance1.incrementContextCount()
then:
@@ -34,11 +41,22 @@ class MapBackedProviderTest extends AgentTestRunner {
instance2.incrementContextCount() == 1
}
- def "backing map should not create strong refs to user instances"() {
+ def "get/put test"() {
+ setup:
+ KeyClass instance1 = new KeyClass()
+
+ when:
+ instance1.putContextCount(10)
+
+ then:
+ instance1.getContextCount() == 10
+ }
+
+ def "backing map should not create strong refs to key class instances"() {
when:
- UserClass1 instance = new UserClass1()
+ KeyClass instance = new KeyClass()
final int count = instance.incrementContextCount()
- WeakReference<UserClass1> instanceRef = new WeakReference(instance)
+ WeakReference<KeyClass> instanceRef = new WeakReference(instance)
instance = null
TestUtils.awaitGC(instanceRef)
@@ -47,8 +65,18 @@ class MapBackedProviderTest extends AgentTestRunner {
count == 1
}
- def "incorrect api usage fails at class load time"() {
+ def "incorrect key class usage fails at class load time"() {
+ expect:
+ !new IncorrectKeyClassUsageKeyClass().isInstrumented()
+ }
+
+ def "incorrect context class usage fails at class load time"() {
+ expect:
+ !new IncorrectContextClassUsageKeyClass().isInstrumented()
+ }
+
+ def "incorrect call usage fails at class load time"() {
expect:
- !new UserClass2().isInstrumented()
+ !new IncorrectCallUsageKeyClass().isInstrumented()
}
}

MODIFY
dd-java-agent/testing/src/test/java/context/ContextTestInstrumentation.java
dd-java-agent/testing/src/test/java/context/ContextTestInstrumentation.java
package context;
+import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
import static net.bytebuddy.matcher.ElementMatchers.named;
import com.google.auto.service.AutoService;
import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.bootstrap.ContextStore;
import datadog.trace.bootstrap.InstrumentationContext;
import java.util.Collections;
import java.util.HashMap;
@@ -20,59 +22,181 @@ public class ContextTestInstrumentation extends Instrumenter.Default {
@Override
public ElementMatcher<? super TypeDescription> typeMatcher() {
- return named(getClass().getPackage().getName() + ".UserClass1")
- .or(named(getClass().getPackage().getName() + ".UserClass2"));
+ return nameStartsWith(getClass().getName() + "$");
}
@Override
public Map<? extends ElementMatcher, String> transformers() {
final Map<ElementMatcher, String> transformers = new HashMap<>(2);
transformers.put(named("isInstrumented"), MarkInstrumentedAdvice.class.getName());
- transformers.put(named("incrementContextCount"), CorrectContextApiUsageAdvice.class.getName());
transformers.put(
- named("incrementContextCountCountBroken"), IncorrectContextApiUsageAdvice.class.getName());
+ named("incrementContextCount"), StoreAndIncrementApiUsageAdvice.class.getName());
+ transformers.put(named("getContextCount"), GetApiUsageAdvice.class.getName());
+ transformers.put(named("putContextCount"), PutApiUsageAdvice.class.getName());
+ transformers.put(
+ named("incorrectKeyClassUsage"), IncorrectKeyClassContextApiUsageAdvice.class.getName());
+ transformers.put(
+ named("incorrectContextClassUsage"),
+ IncorrectContextClassContextApiUsageAdvice.class.getName());
+ transformers.put(
+ named("incorrectCallUsage"), IncorrectCallContextApiUsageAdvice.class.getName());
return transformers;
}
@Override
public String[] helperClassNames() {
- return new String[] {getClass().getName() + "$UserClass1State"};
+ return new String[] {getClass().getName() + "$Context", getClass().getName() + "$Context$1"};
}
@Override
public Map<String, String> contextStore() {
return Collections.singletonMap(
- getClass().getPackage().getName() + ".UserClass1",
- getClass().getName() + "$UserClass1State");
+ getClass().getName() + "$KeyClass", getClass().getName() + "$Context");
}
public static class MarkInstrumentedAdvice {
@Advice.OnMethodExit
- public static void markInstrumented(@Advice.Return(readOnly = false) boolean isInstrumented) {
+ public static void methodExit(@Advice.Return(readOnly = false) boolean isInstrumented) {
isInstrumented = true;
}
}
- public static class CorrectContextApiUsageAdvice {
+ public static class StoreAndIncrementApiUsageAdvice {
+ @Advice.OnMethodExit
+ public static void methodExit(
+ @Advice.This final KeyClass thiz, @Advice.Return(readOnly = false) int contextCount) {
+ final ContextStore<KeyClass, Context> contextStore =
+ InstrumentationContext.get(KeyClass.class, Context.class);
+ final Context context = contextStore.putIfAbsent(thiz, new Context());
+ contextCount = ++context.count;
+ }
+ }
+
+ public static class StoreAndIncrementWithFactoryApiUsageAdvice {
+ @Advice.OnMethodExit
+ public static void methodExit(
+ @Advice.This final KeyClass thiz, @Advice.Return(readOnly = false) int contextCount) {
+ final ContextStore<KeyClass, Context> contextStore =
+ InstrumentationContext.get(KeyClass.class, Context.class);
+ final Context context = contextStore.putIfAbsent(thiz, Context.FACTORY);
+ contextCount = ++context.count;
+ }
+ }
+
+ public static class GetApiUsageAdvice {
+ @Advice.OnMethodExit
+ public static void methodExit(
+ @Advice.This final KeyClass thiz, @Advice.Return(readOnly = false) int contextCount) {
+ final ContextStore<KeyClass, Context> contextStore =
+ InstrumentationContext.get(KeyClass.class, Context.class);
+ contextCount = contextStore.get(thiz).count;
+ }
+ }
+
+ public static class PutApiUsageAdvice {
@Advice.OnMethodExit
- public static void storeAndIncrement(
- @Advice.This Object thiz, @Advice.Return(readOnly = false) int contextCount) {
- UserClass1State state =
- InstrumentationContext.get((UserClass1) thiz, UserClass1.class, UserClass1State.class);
- contextCount = ++state.count;
+ public static void methodExit(
+ @Advice.This final KeyClass thiz, @Advice.Argument(0) final int value) {
+ final ContextStore<KeyClass, Context> contextStore =
+ InstrumentationContext.get(KeyClass.class, Context.class);
+ final Context context = new Context();
+ context.count = value;
+ contextStore.put(thiz, context);
}
}
- public static class IncorrectContextApiUsageAdvice {
+ public static class IncorrectKeyClassContextApiUsageAdvice {
@Advice.OnMethodExit
- public static void storeAndIncrement(
- @Advice.This Object thiz, @Advice.Return(readOnly = false) int contextCount) {
- UserClass1State state = InstrumentationContext.get(thiz, Object.class, UserClass1State.class);
- contextCount = ++state.count;
+ public static void methodExit() {
+ InstrumentationContext.get(Object.class, Context.class);
+ }
+ }
+
+ public static class IncorrectContextClassContextApiUsageAdvice {
+ @Advice.OnMethodExit
+ public static void methodExit() {
+ InstrumentationContext.get(KeyClass.class, Object.class);
+ }
}
+
+ public static class IncorrectCallContextApiUsageAdvice {
+ @Advice.OnMethodExit
+ public static void methodExit() {
+ // Our instrumentation doesn't handle variables being passed to InstrumentationContext.get,
+ // so we make sure that this actually fails instrumentation.
+ final Class clazz = null;
+ InstrumentationContext.get(clazz, Object.class);
+ }
+ }
+
+ public static class Context {
+ public static final ContextStore.Factory<Context> FACTORY =
+ new ContextStore.Factory<Context>() {
+ @Override
+ public Context create() {
+ return new Context();
}
+ };
- public static class UserClass1State {
int count = 0;
}
+
+ public static class KeyClass {
+ public boolean isInstrumented() {
+ // implementation replaced with test instrumentation
+ return false;
+ }
+
+ public int incrementContextCount() {
+ // implementation replaced with test instrumentation
+ return -1;
+ }
+
+ public int incrementContextCountWithFactory() {
+ // implementation replaced with test instrumentation
+ return -1;
+ }
+
+ public int getContextCount() {
+ // implementation replaced with test instrumentation
+ return -1;
+ }
+
+ public void putContextCount(final int value) {
+ // implementation replaced with test instrumentation
+ }
+ }
+
+ public static class IncorrectKeyClassUsageKeyClass {
+ public boolean isInstrumented() {
+ return false;
+ }
+
+ public int incorrectKeyClassUsage() {
+ // instrumentation will not apply to this class because advice incorrectly uses context api
+ return -1;
+ }
+ }
+
+ public static class IncorrectContextClassUsageKeyClass {
+ public boolean isInstrumented() {
+ return false;
+ }
+
+ public int incorrectContextClassUsage() {
+ // instrumentation will not apply to this class because advice incorrectly uses context api
+ return -1;
+ }
+ }
+
+ public static class IncorrectCallUsageKeyClass {
+ public boolean isInstrumented() {
+ return false;
+ }
+
+ public int incorrectCallUsage() {
+ // instrumentation will not apply to this class because advice incorrectly uses context api
+ return -1;
+ }
+ }
}

DELETE
dd-java-agent/testing/src/test/java/context/UserClass1.java
-package context;
-
-public class UserClass1 {
- public boolean isInstrumented() {
- // implementation replaced with test instrumentation
- return false;
- }
-
- public int incrementContextCount() {
- // implementation replaced with test instrumentation
- return -1;
- }
-}

DELETE
dd-java-agent/testing/src/test/java/context/UserClass2.java
-package context;
-
-public class UserClass2 {
- public boolean isInstrumented() {
- return false;
- }
-
- public int incrementContextCountCountBroken() {
- // instrumentation will not apply to this class because advice incorrectly uses context api
- return -1;
- }
-}

First iteration on a better context store api

Separate context storage from actual fetching/putting