Migrate Jax-rs annotation instrumentation to Decorator
MODIFY
dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMap.java
dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMap.java
@@ -16,6 +16,8 @@ public interface WeakMap<K, V> {
void put(K key, V value);
+ void putIfAbsent(K key, V value);
+
@Slf4j
class Provider {
private static final AtomicReference<Supplier> provider =
@@ -81,6 +83,19 @@ public interface WeakMap<K, V> {
map.put(key, value);
}
+ @Override
+ public void putIfAbsent(final K key, final V value) {
+ // We can't use putIfAbsent since it was added in 1.8.
+ // As a result, we must use double check locking.
+ if (!map.containsKey(key)) {
+ synchronized (this) {
+ if (!map.containsKey(key)) {
+ map.put(key, value);
+ }
+ }
+ }
+ }
+
@Override
public String toString() {
return map.toString();

MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/WeakMapSuppliers.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/WeakMapSuppliers.java
@@ -170,6 +170,11 @@ class WeakMapSuppliers {
public void put(final K key, final V value) {
map.put(key, value);
}
+
+ @Override
+ public void putIfAbsent(final K key, final V value) {
+ map.putIfAbsent(key, value);
+ }
}
static class Inline implements WeakMap.Supplier {

ADD
dd-java-agent/instrumentation/jax-rs-annotations/src/main/java/datadog/trace/instrumentation/jaxrs/JaxRsAnnotationsDecorator.java
+package datadog.trace.instrumentation.jaxrs;
+
+import static datadog.trace.bootstrap.WeakMap.Provider.newWeakMap;
+
+import datadog.trace.agent.decorator.BaseDecorator;
+import datadog.trace.api.DDTags;
+import datadog.trace.bootstrap.WeakMap;
+import io.opentracing.Scope;
+import io.opentracing.Span;
+import io.opentracing.tag.Tags;
+import java.lang.annotation.Annotation;
+import java.lang.reflect.Method;
+import java.util.LinkedList;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import javax.ws.rs.HttpMethod;
+import javax.ws.rs.Path;
+
+public class JaxRsAnnotationsDecorator extends BaseDecorator {
+ public static JaxRsAnnotationsDecorator DECORATE = new JaxRsAnnotationsDecorator();
+
+ private final WeakMap<Class, Map<String, String>> resourceNames = newWeakMap();
+
+ @Override
+ protected String[] instrumentationNames() {
+ return new String[0];
+ }
+
+ @Override
+ protected String spanType() {
+ return null;
+ }
+
+ @Override
+ protected String component() {
+ return "jax-rs-controller";
+ }
+
+ public void updateParent(final Scope scope, final Method method) {
+ if (scope == null) {
+ return;
+ }
+ final Span span = scope.span();
+ Tags.COMPONENT.set(span, "jax-rs");
+
+ Class<?> target = method.getDeclaringClass();
+ Map<String, String> classMap = resourceNames.get(target);
+
+ if (classMap == null) {
+ resourceNames.putIfAbsent(target, new ConcurrentHashMap<String, String>());
+ classMap = resourceNames.get(target);
+ }
+
+ final String methodName = method.toString();
+ String resourceName = classMap.get(methodName);
+ if (resourceName == null) {
+ final LinkedList<Path> paths = new LinkedList<>();
+ while (target != Object.class) {
+ final Path annotation = target.getAnnotation(Path.class);
+ if (annotation != null) {
+ paths.push(annotation);
+ }
+ target = target.getSuperclass();
+ }
+ final Path methodPath = method.getAnnotation(Path.class);
+ if (methodPath != null) {
+ paths.add(methodPath);
+ }
+ String httpMethod = null;
+ for (final Annotation ann : method.getDeclaredAnnotations()) {
+ if (ann.annotationType().getAnnotation(HttpMethod.class) != null) {
+ httpMethod = ann.annotationType().getSimpleName();
+ }
+ }
+
+ final StringBuilder resourceNameBuilder = new StringBuilder();
+ if (httpMethod != null) {
+ resourceNameBuilder.append(httpMethod);
+ resourceNameBuilder.append(" ");
+ }
+ Path last = null;
+ for (final Path path : paths) {
+ if (path.value().startsWith("/") || (last != null && last.value().endsWith("/"))) {
+ resourceNameBuilder.append(path.value());
+ } else {
+ resourceNameBuilder.append("/");
+ resourceNameBuilder.append(path.value());
+ }
+ last = path;
+ }
+ resourceName = resourceNameBuilder.toString().trim();
+ classMap.put(methodName, resourceName);
+ }
+
+ if (!resourceName.isEmpty()) {
+ span.setTag(DDTags.RESOURCE_NAME, resourceName);
+ }
+ }
+}

MODIFY
dd-java-agent/instrumentation/jax-rs-annotations/src/main/java/datadog/trace/instrumentation/jaxrs/JaxRsAnnotationsInstrumentation.java
dd-java-agent/instrumentation/jax-rs-annotations/src/main/java/datadog/trace/instrumentation/jaxrs/JaxRsAnnotationsInstrumentation.java
package datadog.trace.instrumentation.jaxrs;
import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
-import static io.opentracing.log.Fields.ERROR_OBJECT;
+import static datadog.trace.instrumentation.jaxrs.JaxRsAnnotationsDecorator.DECORATE;
import static java.util.Collections.singletonMap;
import static net.bytebuddy.matcher.ElementMatchers.declaresMethod;
import static net.bytebuddy.matcher.ElementMatchers.isAnnotatedWith;
@@ -9,18 +9,10 @@ import static net.bytebuddy.matcher.ElementMatchers.named;
import com.google.auto.service.AutoService;
import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.api.DDTags;
import io.opentracing.Scope;
-import io.opentracing.Span;
-import io.opentracing.tag.Tags;
import io.opentracing.util.GlobalTracer;
-import java.lang.annotation.Annotation;
import java.lang.reflect.Method;
-import java.util.Collections;
-import java.util.LinkedList;
import java.util.Map;
-import javax.ws.rs.HttpMethod;
-import javax.ws.rs.Path;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.description.type.TypeDescription;
@@ -58,83 +50,20 @@ public final class JaxRsAnnotationsInstrumentation extends Instrumenter.Default
@Advice.OnMethodEnter(suppress = Throwable.class)
public static Scope nameSpan(@Advice.Origin final Method method) {
-
- // TODO: do we need caching for this?
- final LinkedList<Path> paths = new LinkedList<>();
- Class<?> target = method.getDeclaringClass();
- while (target != Object.class) {
- final Path annotation = target.getAnnotation(Path.class);
- if (annotation != null) {
- paths.push(annotation);
- }
- target = target.getSuperclass();
- }
- final Path methodPath = method.getAnnotation(Path.class);
- if (methodPath != null) {
- paths.add(methodPath);
- }
- String httpMethod = null;
- for (final Annotation ann : method.getDeclaredAnnotations()) {
- if (ann.annotationType().getAnnotation(HttpMethod.class) != null) {
- httpMethod = ann.annotationType().getSimpleName();
- }
- }
-
- final StringBuilder resourceNameBuilder = new StringBuilder();
- if (httpMethod != null) {
- resourceNameBuilder.append(httpMethod);
- resourceNameBuilder.append(" ");
- }
- Path last = null;
- for (final Path path : paths) {
- if (path.value().startsWith("/") || (last != null && last.value().endsWith("/"))) {
- resourceNameBuilder.append(path.value());
- } else {
- resourceNameBuilder.append("/");
- resourceNameBuilder.append(path.value());
- }
- last = path;
- }
- final String resourceName = resourceNameBuilder.toString().trim();
-
+ // Rename the parent span according to the path represented by these annotations.
final Scope scope = GlobalTracer.get().scopeManager().active();
- if (scope != null && !resourceName.isEmpty()) {
- scope.span().setTag(DDTags.RESOURCE_NAME, resourceName);
- Tags.COMPONENT.set(scope.span(), "jax-rs");
- }
+ DECORATE.updateParent(scope, method);
// Now create a span representing the method execution.
-
- final Class<?> clazz = method.getDeclaringClass();
- final String methodName = method.getName();
-
- String className = clazz.getSimpleName();
- if (className.isEmpty()) {
- className = clazz.getName();
- if (clazz.getPackage() != null) {
- final String pkgName = clazz.getPackage().getName();
- if (!pkgName.isEmpty()) {
- className = clazz.getName().replace(pkgName, "").substring(1);
- }
- }
- }
-
- final String operationName = className + "." + methodName;
-
- return GlobalTracer.get()
- .buildSpan(operationName)
- .withTag(Tags.COMPONENT.getKey(), "jax-rs-controller")
- .startActive(true);
+ final String operationName = DECORATE.spanNameForMethod(method);
+ return DECORATE.afterStart(GlobalTracer.get().buildSpan(operationName).startActive(true));
}
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void stopSpan(
@Advice.Enter final Scope scope, @Advice.Thrown final Throwable throwable) {
- if (throwable != null) {
- final Span span = scope.span();
- Tags.ERROR.set(span, true);
- span.log(Collections.singletonMap(ERROR_OBJECT, throwable));
- }
+ DECORATE.onError(scope, throwable);
+ DECORATE.beforeFinish(scope);
scope.close();
}
}

MODIFY
dd-java-agent/instrumentation/jax-rs-annotations/src/test/groovy/JaxRsAnnotationsInstrumentationTest.groovy
dd-java-agent/instrumentation/jax-rs-annotations/src/test/groovy/JaxRsAnnotationsInstrumentationTest.groovy
@@ -10,11 +10,13 @@ import javax.ws.rs.PUT
import javax.ws.rs.Path
import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
+import static datadog.trace.instrumentation.jaxrs.JaxRsAnnotationsDecorator.DECORATE
class JaxRsAnnotationsInstrumentationTest extends AgentTestRunner {
def "span named '#name' from annotations on class"() {
setup:
+ def startingCacheSize = DECORATE.resourceNames.size()
runUnderTrace("test") {
obj.call()
}
@@ -42,6 +44,18 @@ class JaxRsAnnotationsInstrumentationTest extends AgentTestRunner {
}
}
}
+ DECORATE.resourceNames.size() == startingCacheSize + 1
+ DECORATE.resourceNames.get(obj.class).size() == 1
+
+ when: "multiple calls to the same method"
+ runUnderTrace("test") {
+ (1..10).each {
+ obj.call()
+ }
+ }
+ then: "doesn't increase the cache size"
+ DECORATE.resourceNames.size() == startingCacheSize + 1
+ DECORATE.resourceNames.get(obj.class).size() == 1
where:
name | obj

Migrate Jax-rs annotation instrumentation to Decorator