Add tracing for other annotations and config

dd-java-agent/dd-java-agent.gradle

@@ -44,7 +44,10 @@ def includeShadowJar(subproject, jarname) {
         relocate 'datadog.trace.common', 'datadog.trace.agent.common'
         relocate 'datadog.opentracing', 'datadog.trace.agent.ot'
 
-        relocate 'io.opentracing.contrib', 'datadog.trace.agent.opentracing.contrib'
+        relocate('io.opentracing.contrib', 'datadog.trace.agent.opentracing.contrib') {
+          // Don't want to change the annotation we're looking for.
+          exclude 'io.opentracing.contrib.dropwizard.Trace'
+        }
 
         relocate 'org.yaml', 'datadog.trace.agent.deps.yaml'
         relocate 'org.msgpack', 'datadog.trace.agent.deps.msgpack'

dd-java-agent/instrumentation/trace-annotation/src/main/java/datadog/trace/instrumentation/trace_annotation/TraceAdvice.java

@@ -0,0 +1,48 @@
+package datadog.trace.instrumentation.trace_annotation;
+
+import static io.opentracing.log.Fields.ERROR_OBJECT;
+
+import datadog.trace.api.Trace;
+import io.opentracing.Scope;
+import io.opentracing.Span;
+import io.opentracing.tag.Tags;
+import io.opentracing.util.GlobalTracer;
+import java.lang.reflect.Method;
+import java.util.Collections;
+import net.bytebuddy.asm.Advice;
+
+public class TraceAdvice {
+
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  public static Scope startSpan(@Advice.Origin final Method method) {
+    final Trace trace = method.getAnnotation(Trace.class);
+    String operationName = trace == null ? null : trace.operationName();
+    if (operationName == null || operationName.isEmpty()) {
+      final Class<?> declaringClass = method.getDeclaringClass();
+      String className = declaringClass.getSimpleName();
+      if (className.isEmpty()) {
+        className = declaringClass.getName();
+        if (declaringClass.getPackage() != null) {
+          final String pkgName = declaringClass.getPackage().getName();
+          if (!pkgName.isEmpty()) {
+            className = declaringClass.getName().replace(pkgName, "").substring(1);
+          }
+        }
+      }
+      operationName = className + "." + method.getName();
+    }
+
+    return 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));
+    }
+    scope.close();
+  }
+}

dd-java-agent/instrumentation/trace-annotation/src/main/java/datadog/trace/instrumentation/trace_annotation/TraceAnnotationInstrumentation.java

@@ -1,75 +1,97 @@
 package datadog.trace.instrumentation.trace_annotation;
 
-import static io.opentracing.log.Fields.ERROR_OBJECT;
+import static datadog.trace.instrumentation.trace_annotation.TraceConfigInstrumentation.PACKAGE_CLASS_NAME_REGEX;
 import static net.bytebuddy.matcher.ElementMatchers.declaresMethod;
 import static net.bytebuddy.matcher.ElementMatchers.failSafe;
 import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
+import static net.bytebuddy.matcher.ElementMatchers.is;
 import static net.bytebuddy.matcher.ElementMatchers.isAnnotatedWith;
+import static net.bytebuddy.matcher.ElementMatchers.named;
 
 import com.google.auto.service.AutoService;
+import com.google.common.collect.Sets;
 import datadog.trace.agent.tooling.DDAdvice;
 import datadog.trace.agent.tooling.DDTransformers;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.api.Trace;
-import io.opentracing.Scope;
-import io.opentracing.Span;
-import io.opentracing.tag.Tags;
-import io.opentracing.util.GlobalTracer;
-import java.lang.reflect.Method;
 import java.util.Collections;
+import java.util.Set;
+import lombok.extern.slf4j.Slf4j;
 import net.bytebuddy.agent.builder.AgentBuilder;
-import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.NamedElement;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
 
+@Slf4j
 @AutoService(Instrumenter.class)
 public final class TraceAnnotationInstrumentation extends Instrumenter.Configurable {
+  private static final String CONFIG_NAME = "dd.trace.annotations";
+
+  static final String CONFIG_FORMAT =
+      "(?:\\s*"
+          + PACKAGE_CLASS_NAME_REGEX
+          + "\\s*;)*\\s*"
+          + PACKAGE_CLASS_NAME_REGEX
+          + "\\s*;?\\s*";
+
+  private static final String[] DEFAULT_ANNOTATIONS =
+      new String[] {
+        "com.newrelic.api.agent.Trace",
+        "kamon.annotation.Trace",
+        "com.tracelytics.api.ext.LogMethod",
+        "io.opentracing.contrib.dropwizard.Trace",
+        "org.springframework.cloud.sleuth.annotation.NewSpan"
+      };
+
+  private final Set<String> additionalTraceAnnotations;
 
   public TraceAnnotationInstrumentation() {
     super("trace", "trace-annotation");
+
+    final String configString = getPropOrEnv(CONFIG_NAME);
+    if (configString == null) {
+      additionalTraceAnnotations =
+          Collections.unmodifiableSet(Sets.<String>newHashSet(DEFAULT_ANNOTATIONS));
+    } else if (configString.trim().isEmpty()) {
+      additionalTraceAnnotations = Collections.emptySet();
+    } else if (!configString.matches(CONFIG_FORMAT)) {
+      log.warn(
+          "Invalid trace annotations config '{}'. Must match 'package.Annotation$Name;*'.",
+          configString);
+      additionalTraceAnnotations = Collections.emptySet();
+    } else {
+      final Set<String> annotations = Sets.newHashSet();
+      final String[] annotationClasses = configString.split(";", -1);
+      for (final String annotationClass : annotationClasses) {
+        if (!annotationClass.trim().isEmpty()) {
+          annotations.add(annotationClass.trim());
+        }
+      }
+      additionalTraceAnnotations = Collections.unmodifiableSet(annotations);
+    }
   }
 
   @Override
   public AgentBuilder apply(final AgentBuilder agentBuilder) {
+    ElementMatcher.Junction<NamedElement> methodTraceMatcher =
+        is(new TypeDescription.ForLoadedType(Trace.class));
+    for (final String annotationName : additionalTraceAnnotations) {
+      methodTraceMatcher = methodTraceMatcher.or(named(annotationName));
+    }
     return agentBuilder
-        .type(failSafe(hasSuperType(declaresMethod(isAnnotatedWith(Trace.class)))))
+        .type(failSafe(hasSuperType(declaresMethod(isAnnotatedWith(methodTraceMatcher)))))
         .transform(DDTransformers.defaultTransformers())
         .transform(
-            DDAdvice.create().advice(isAnnotatedWith(Trace.class), TraceAdvice.class.getName()))
+            DDAdvice.create()
+                .advice(isAnnotatedWith(methodTraceMatcher), TraceAdvice.class.getName()))
         .asDecorator();
   }
 
-  public static class TraceAdvice {
-
-    @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static Scope startSpan(@Advice.Origin final Method method) {
-      final Trace trace = method.getAnnotation(Trace.class);
-      String operationName = trace == null ? null : trace.operationName();
-      if (operationName == null || operationName.isEmpty()) {
-        final Class<?> declaringClass = method.getDeclaringClass();
-        String className = declaringClass.getSimpleName();
-        if (className.isEmpty()) {
-          className = declaringClass.getName();
-          if (declaringClass.getPackage() != null) {
-            final String pkgName = declaringClass.getPackage().getName();
-            if (!pkgName.isEmpty()) {
-              className = declaringClass.getName().replace(pkgName, "").substring(1);
-            }
-          }
-        }
-        operationName = className + "." + method.getName();
-      }
-
-      return GlobalTracer.get().buildSpan(operationName).startActive(true);
-    }
+  private String getPropOrEnv(final String name) {
+    return System.getProperty(name, System.getenv(propToEnvName(name)));
+  }
 
-    @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));
-      }
-      scope.close();
-    }
+  static String propToEnvName(final String name) {
+    return name.toUpperCase().replace(".", "_");
   }
 }

dd-java-agent/instrumentation/trace-annotation/src/main/java/datadog/trace/instrumentation/trace_annotation/TraceConfigInstrumentation.java

@@ -0,0 +1,110 @@
+package datadog.trace.instrumentation.trace_annotation;
+
+import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+import com.google.auto.service.AutoService;
+import com.google.common.collect.Maps;
+import com.google.common.collect.Sets;
+import datadog.trace.agent.tooling.DDAdvice;
+import datadog.trace.agent.tooling.Instrumenter;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Set;
+import lombok.extern.slf4j.Slf4j;
+import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@Slf4j
+@AutoService(Instrumenter.class)
+public class TraceConfigInstrumentation extends Instrumenter.Configurable {
+  private static final String CONFIG_NAME = "dd.trace.methods";
+
+  static final String PACKAGE_CLASS_NAME_REGEX = "[\\w.\\$]+";
+  private static final String METHOD_LIST_REGEX = "\\s*(?:\\w+\\s*,)*\\s*(?:\\w+\\s*,?)\\s*";
+  private static final String CONFIG_FORMAT =
+      "(?:\\s*"
+          + PACKAGE_CLASS_NAME_REGEX
+          + "\\["
+          + METHOD_LIST_REGEX
+          + "\\]\\s*;)*\\s*"
+          + PACKAGE_CLASS_NAME_REGEX
+          + "\\["
+          + METHOD_LIST_REGEX
+          + "\\]\\s*;?\\s*";
+
+  private final Map<String, Set<String>> classMethodsToTrace;
+
+  public TraceConfigInstrumentation() {
+    super("trace", "trace-config");
+
+    final String configString = getPropOrEnv(CONFIG_NAME);
+    if (configString == null || configString.trim().isEmpty()) {
+      classMethodsToTrace = Collections.emptyMap();
+
+    } else if (!configString.matches(CONFIG_FORMAT)) {
+      log.warn(
+          "Invalid trace method config '{}'. Must match 'package.Class$Name[method1,method2];*'.",
+          configString);
+      classMethodsToTrace = Collections.emptyMap();
+
+    } else {
+      final Map<String, Set<String>> toTrace = Maps.newHashMap();
+      final String[] classMethods = configString.split(";");
+      for (final String classMethod : classMethods) {
+        final String[] splitClassMethod = classMethod.split("\\[");
+        final String className = splitClassMethod[0];
+        final String method = splitClassMethod[1].trim();
+        final String methodNames = method.substring(0, method.length() - 1);
+        final String[] splitMethodNames = methodNames.split(",");
+        final Set<String> trimmedMethodNames =
+            Sets.newHashSetWithExpectedSize(splitMethodNames.length);
+        for (final String methodName : splitMethodNames) {
+          final String trimmedMethodName = methodName.trim();
+          if (!trimmedMethodName.isEmpty()) {
+            trimmedMethodNames.add(trimmedMethodName);
+          }
+        }
+        if (!trimmedMethodNames.isEmpty()) {
+          toTrace.put(className.trim(), trimmedMethodNames);
+        }
+      }
+      classMethodsToTrace = Collections.unmodifiableMap(toTrace);
+    }
+  }
+
+  @Override
+  public AgentBuilder apply(final AgentBuilder agentBuilder) {
+    if (classMethodsToTrace.isEmpty()) {
+      return agentBuilder;
+    }
+    AgentBuilder builder = agentBuilder;
+
+    for (final Map.Entry<String, Set<String>> entry : classMethodsToTrace.entrySet()) {
+
+      ElementMatcher.Junction<MethodDescription> methodMatchers = null;
+      for (final String methodName : entry.getValue()) {
+        if (methodMatchers == null) {
+          methodMatchers = named(methodName);
+        } else {
+          methodMatchers = methodMatchers.or(named(methodName));
+        }
+      }
+      builder =
+          builder
+              .type(hasSuperType(named(entry.getKey())))
+              .transform(DDAdvice.create().advice(methodMatchers, TraceAdvice.class.getName()))
+              .asDecorator();
+    }
+    return builder;
+  }
+
+  private String getPropOrEnv(final String name) {
+    return System.getProperty(name, System.getenv(propToEnvName(name)));
+  }
+
+  static String propToEnvName(final String name) {
+    return name.toUpperCase().replace(".", "_");
+  }
+}

dd-java-agent/instrumentation/trace-annotation/src/test/groovy/ConfiguredTraceAnnotationsTest.groovy

@@ -0,0 +1,80 @@
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.instrumentation.trace_annotation.TraceAnnotationInstrumentation
+import dd.test.trace.annotation.SayTracedHello
+import spock.lang.Unroll
+
+import java.util.concurrent.Callable
+
+import static datadog.trace.agent.test.ListWriterAssert.assertTraces
+import static datadog.trace.agent.test.TestUtils.withSystemProperty
+import static datadog.trace.instrumentation.trace_annotation.TraceAnnotationInstrumentation.DEFAULT_ANNOTATIONS
+
+class ConfiguredTraceAnnotationsTest extends AgentTestRunner {
+
+  static {
+    // nr annotation not included here, so should be disabled.
+    System.setProperty("dd.trace.annotations", "package.Class\$Name;${OuterClass.InterestingMethod.name}")
+  }
+
+  def specCleanup() {
+    System.clearProperty("dd.trace.annotations")
+  }
+
+  def "test disabled nr annotation"() {
+    setup:
+    SayTracedHello.fromCallable()
+
+    expect:
+    TEST_WRITER == []
+  }
+
+  def "test custom annotation based trace"() {
+    expect:
+    new AnnotationTracedCallable().call() == "Hello!"
+
+    when:
+    TEST_WRITER.waitForTraces(1)
+
+    then:
+    assertTraces(TEST_WRITER, 1) {
+      trace(0, 1) {
+        span(0) {
+          resourceName "AnnotationTracedCallable.call"
+          operationName "AnnotationTracedCallable.call"
+        }
+      }
+    }
+  }
+
+  @Unroll
+  def "test configuration #value"() {
+    setup:
+    def config = null
+    withSystemProperty("dd.trace.annotations", value) {
+      def instrumentation = new TraceAnnotationInstrumentation()
+      config = instrumentation.additionalTraceAnnotations
+    }
+    expect:
+    config == expected.toSet()
+
+    where:
+    value                               | expected
+    null                                | DEFAULT_ANNOTATIONS.toList()
+    " "                                 | []
+    "some.Invalid[]"                    | []
+    "some.package.ClassName "           | ["some.package.ClassName"]
+    " some.package.Class\$Name"         | ["some.package.Class\$Name"]
+    "  ClassName  "                     | ["ClassName"]
+    "ClassName"                         | ["ClassName"]
+    "Class\$1;Class\$2;"                | ["Class\$1", "Class\$2"]
+    "Duplicate ;Duplicate ;Duplicate; " | ["Duplicate"]
+  }
+
+  class AnnotationTracedCallable implements Callable<String> {
+    @OuterClass.InterestingMethod
+    @Override
+    String call() throws Exception {
+      return "Hello!"
+    }
+  }
+}