Add tracing for other annotations and config

Allows for specifying additional traced methods via env var or sys prop 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,51 @@
 package datadog.trace.instrumentation.trace_annotation;
 
-import static io.opentracing.log.Fields.ERROR_OBJECT;
 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 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 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;
 
 @AutoService(Instrumenter.class)
 public final class TraceAnnotationInstrumentation extends Instrumenter.Configurable {
 
+  private static final String[] ADDITIONAL_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"
+      };
+
   public TraceAnnotationInstrumentation() {
     super("trace", "trace-annotation");
   }
 
   @Override
   public AgentBuilder apply(final AgentBuilder agentBuilder) {
+    ElementMatcher.Junction<NamedElement> methodTraceMatcher =
+        is(new TypeDescription.ForLoadedType(Trace.class));
+    for (final String annotationName : ADDITIONAL_ANNOTATIONS) {
+      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);
-    }
-
-    @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/TraceConfigInstrumentation.java

@@ -0,0 +1,87 @@
+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";
+
+  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(
+        "(?:([\\w.\\$]+)\\[((?:\\w+,)*(?:\\w+,?))\\];)*([\\w.\\$]+)\\[((?:\\w+,)*(?:\\w+,?))\\];?")) {
+      log.warn(
+          "Invalid 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 methodNames =
+            splitClassMethod[1].substring(0, splitClassMethod[1].length() - 1);
+        final String[] splitMethodNames = methodNames.split(",");
+        toTrace.put(className, Sets.newHashSet(splitMethodNames));
+      }
+      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/TraceAnnotationsTest.groovy

@@ -1,12 +1,12 @@
-import datadog.opentracing.DDSpan
 import datadog.opentracing.decorators.ErrorFlag
 import datadog.trace.agent.test.AgentTestRunner
 import datadog.trace.api.Trace
 import dd.test.trace.annotation.SayTracedHello
-import org.assertj.core.api.Assertions
 
 import java.util.concurrent.Callable
 
+import static datadog.trace.agent.test.ListWriterAssert.assertTraces
+
 class TraceAnnotationsTest extends AgentTestRunner {
 
   def "test simple case annotations"() {
@@ -14,34 +14,63 @@ class TraceAnnotationsTest extends AgentTestRunner {
     // Test single span in new trace
     SayTracedHello.sayHello()
 
-    Assertions.assertThat(TEST_WRITER.firstTrace().size()).isEqualTo(1)
+    expect:
-    Assertions.assertThat(TEST_WRITER.firstTrace().get(0).getOperationName())
+    assertTraces(TEST_WRITER, 1) {
-      .isEqualTo("SayTracedHello.sayHello")
+      trace(0, 1) {
-    Assertions.assertThat(TEST_WRITER.firstTrace().get(0).getServiceName()).isEqualTo("test")
+        span(0) {
+          serviceName "test"
+          resourceName "SayTracedHello.sayHello"
+          operationName "SayTracedHello.sayHello"
+          parent()
+          errored false
+          tags {
+            defaultTags()
+          }
+        }
+      }
+    }
   }
 
   def "test complex case annotations"() {
-    setup:
+    when:
-
     // Test new trace with 2 children spans
     SayTracedHello.sayHELLOsayHA()
-    Assertions.assertThat(TEST_WRITER.firstTrace().size()).isEqualTo(3)
-    final long parentId = TEST_WRITER.firstTrace().get(0).context().getSpanId()
 
-    Assertions.assertThat(TEST_WRITER.firstTrace().get(0).getOperationName()).isEqualTo("NEW_TRACE")
+    then:
-    Assertions.assertThat(TEST_WRITER.firstTrace().get(0).getParentId())
+    assertTraces(TEST_WRITER, 1) {
-      .isEqualTo(0) // ROOT / no parent
+      trace(0, 3) {
-    Assertions.assertThat(TEST_WRITER.firstTrace().get(0).context().getParentId()).isEqualTo(0)
+        span(0) {
-    Assertions.assertThat(TEST_WRITER.firstTrace().get(0).getServiceName()).isEqualTo("test2")
+          resourceName "NEW_TRACE"
-
+          operationName "NEW_TRACE"
-    Assertions.assertThat(TEST_WRITER.firstTrace().get(1).getOperationName()).isEqualTo("SAY_HA")
+          parent()
-    Assertions.assertThat(TEST_WRITER.firstTrace().get(1).getParentId()).isEqualTo(parentId)
+          errored false
-    Assertions.assertThat(TEST_WRITER.firstTrace().get(1).context().getSpanType()).isEqualTo("DB")
+          tags {
-
+            defaultTags()
-    Assertions.assertThat(TEST_WRITER.firstTrace().get(2).getOperationName())
+          }
-      .isEqualTo("SayTracedHello.sayHello")
+        }
-    Assertions.assertThat(TEST_WRITER.firstTrace().get(2).getServiceName()).isEqualTo("test")
+        span(1) {
-    Assertions.assertThat(TEST_WRITER.firstTrace().get(2).getParentId()).isEqualTo(parentId)
+          resourceName "SAY_HA"
+          operationName "SAY_HA"
+          spanType "DB"
+          childOf span(0)
+          errored false
+          tags {
+            "span.type" "DB"
+            defaultTags()
+          }
+        }
+        span(2) {
+          serviceName "test"
+          resourceName "SayTracedHello.sayHello"
+          operationName "SayTracedHello.sayHello"
+          childOf span(0)
+          errored false
+          tags {
+            defaultTags()
+          }
+        }
+      }
+    }
   }
 
   def "test exception exit"() {
@@ -56,16 +85,20 @@ class TraceAnnotationsTest extends AgentTestRunner {
       error = ex
     }
 
-    final StringWriter errorString = new StringWriter()
+    expect:
-    error.printStackTrace(new PrintWriter(errorString))
+    assertTraces(TEST_WRITER, 1) {
-
+      trace(0, 1) {
-    final DDSpan span = TEST_WRITER.firstTrace().get(0)
+        span(0) {
-    Assertions.assertThat(span.getOperationName()).isEqualTo("ERROR")
+          resourceName "ERROR"
-    Assertions.assertThat(span.getTags().get("error")).isEqualTo(true)
+          operationName "ERROR"
-    Assertions.assertThat(span.getTags().get("error.msg")).isEqualTo(error.getMessage())
+          errored true
-    Assertions.assertThat(span.getTags().get("error.type")).isEqualTo(error.getClass().getName())
+          tags {
-    Assertions.assertThat(span.getTags().get("error.stack")).isEqualTo(errorString.toString())
+            errorTags(error.class)
-    Assertions.assertThat(span.getError()).isEqualTo(1)
+            defaultTags()
+          }
+        }
+      }
+    }
   }
 
   def "test annonymous class annotations"() {
@@ -73,10 +106,15 @@ class TraceAnnotationsTest extends AgentTestRunner {
     // Test anonymous classes with package.
     SayTracedHello.fromCallable()
 
-    Assertions.assertThat(TEST_WRITER.size()).isEqualTo(1)
+    expect:
-    Assertions.assertThat(TEST_WRITER.firstTrace().size()).isEqualTo(1)
+    assertTraces(TEST_WRITER, 1) {
-    Assertions.assertThat(TEST_WRITER.firstTrace().get(0).getOperationName())
+      trace(0, 1) {
-      .isEqualTo("SayTracedHello\$1.call")
+        span(0) {
+          resourceName "SayTracedHello\$1.call"
+          operationName "SayTracedHello\$1.call"
+        }
+      }
+    }
 
     when:
     // Test anonymous classes with no package.
@@ -90,9 +128,48 @@ class TraceAnnotationsTest extends AgentTestRunner {
     TEST_WRITER.waitForTraces(2)
 
     then:
-    Assertions.assertThat(TEST_WRITER.size()).isEqualTo(2)
+    assertTraces(TEST_WRITER, 2) {
-    Assertions.assertThat(TEST_WRITER.get(1).size()).isEqualTo(1)
+      trace(0, 1) {
-    Assertions.assertThat(TEST_WRITER.get(1).get(0).getOperationName())
+        span(0) {
-      .isEqualTo("TraceAnnotationsTest\$1.call")
+          resourceName "SayTracedHello\$1.call"
+          operationName "SayTracedHello\$1.call"
+        }
+        trace(1, 1) {
+          span(0) {
+            resourceName "TraceAnnotationsTest\$1.call"
+            operationName "TraceAnnotationsTest\$1.call"
+          }
+        }
+      }
+    }
+  }
+
+  def "test configuration based trace"() {
+    expect:
+    new ConfigTracedCallable().call() == "Hello!"
+
+    when:
+    TEST_WRITER.waitForTraces(1)
+
+    then:
+    assertTraces(TEST_WRITER, 1) {
+      trace(0, 1) {
+        span(0) {
+          resourceName "ConfigTracedCallable.call"
+          operationName "ConfigTracedCallable.call"
+        }
+      }
+    }
+  }
+
+  static {
+    System.setProperty("dd.trace.methods", "package.ClassName[method1,method2];${ConfigTracedCallable.name}[call]")
+  }
+
+  class ConfigTracedCallable implements Callable<String> {
+    @Override
+    String call() throws Exception {
+      return "Hello!"
+    }
   }
 }

dd-java-agent/instrumentation/trace-annotation/src/test/java/dd/test/trace/annotation/SayTracedHello.java

@@ -37,7 +37,7 @@ public class SayTracedHello {
 
   public static String fromCallable() throws Exception {
     return new Callable<String>() {
-      @Trace
+      @com.newrelic.api.agent.Trace
       @Override
       public String call() throws Exception {
         return "Howdy!";

dd-java-agent/instrumentation/trace-annotation/trace-annotation.gradle

@@ -10,4 +10,5 @@ dependencies {
   implementation deps.autoservice
 
   testCompile project(':dd-java-agent:testing')
+  testCompile group: 'com.newrelic.agent.java', name: 'newrelic-api', version: '+'
 }

dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/TagsAssert.groovy

@@ -43,6 +43,7 @@ class TagsAssert {
     if (args.length > 1) {
       throw new IllegalArgumentException(args)
     }
+    assertedTags.add(name)
     assert tags[name] == args[0]
   }