Merge pull request #143 from DataDog/tyler/trace-annotation

 Use byte buddy for @trace instead of byteman

Author: tylerbenson

…q/agent/TraceAnnotationsManagerTest.java …ion/annotation/TraceAnnotationsTest.javadd-java-agent-ittests/src/test/java/com/datadoghq/agent/TraceAnnotationsManagerTest.java renamed to dd-java-agent-ittests/src/test/java/com/datadoghq/agent/instrumentation/annotation/TraceAnnotationsTest.java dd-java-agent-ittests/src/test/java/com/datadoghq/agent/TraceAnnotationsManagerTest.java renamed to dd-java-agent-ittests/src/test/java/com/datadoghq/agent/instrumentation/annotation/TraceAnnotationsTest.java

@@ -1,4 +1,4 @@
-package com.datadoghq.agent;
+package com.datadoghq.agent.instrumentation.annotation;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -14,7 +14,7 @@
 import org.junit.Before;
 import org.junit.Test;
 
-public class TraceAnnotationsManagerTest {
+public class TraceAnnotationsTest {
 
   private final ListWriter writer = new ListWriter();
   private final DDTracer tracer = new DDTracer(writer);
@@ -39,7 +39,8 @@ public void testSimpleCaseAnnotations() {
     SayTracedHello.sayHello();
 
     assertThat(writer.firstTrace().size()).isEqualTo(1);
-    assertThat(writer.firstTrace().get(0).getOperationName()).isEqualTo("SAY_HELLO");
+    assertThat(writer.firstTrace().get(0).getOperationName())
+        .isEqualTo("com.datadoghq.agent.test.SayTracedHello.sayHello");
     assertThat(writer.firstTrace().get(0).getServiceName()).isEqualTo("test");
   }
 
@@ -56,7 +57,8 @@ public void testComplexCaseAnnotations() {
     assertThat(writer.firstTrace().get(0).context().getParentId()).isEqualTo(0);
     assertThat(writer.firstTrace().get(0).getServiceName()).isEqualTo("test2");
 
-    assertThat(writer.firstTrace().get(1).getOperationName()).isEqualTo("SAY_HELLO");
+    assertThat(writer.firstTrace().get(1).getOperationName())
+        .isEqualTo("com.datadoghq.agent.test.SayTracedHello.sayHello");
     assertThat(writer.firstTrace().get(1).getServiceName()).isEqualTo("test");
     assertThat(writer.firstTrace().get(1).getParentId()).isEqualTo(parentId);
 
@@ -82,7 +84,7 @@ public void testExceptionExit() {
 
     final DDBaseSpan<?> span = writer.firstTrace().get(0);
     assertThat(span.getOperationName()).isEqualTo("ERROR");
-    assertThat(span.getTags().get("error")).isEqualTo("true");
+    assertThat(span.getTags().get("error")).isEqualTo(true);
     assertThat(span.getTags().get("error.msg")).isEqualTo(error.getMessage());
     assertThat(span.getTags().get("error.type")).isEqualTo(error.getClass().getName());
     assertThat(span.getTags().get("error.stack")).isEqualTo(errorString.toString());

dd-java-agent-ittests/src/test/java/com/datadoghq/agent/test/SayTracedHello.java

@@ -7,7 +7,7 @@
 
 public class SayTracedHello {
 
-  @Trace(operationName = "SAY_HELLO")
+  @Trace
   public static String sayHello() {
     new StringTag(DDTags.SERVICE_NAME).set(GlobalTracer.get().activeSpan(), "test");
     return "hello!";

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

@@ -23,7 +23,6 @@ dependencies {
   compile group: 'net.bytebuddy', name: 'byte-buddy', version: '1.7.6'
   compile group: 'org.jboss.byteman', name: 'byteman', version: '3.0.10'
 
-  compile group: 'org.reflections', name: 'reflections', version: '0.9.11'
   compile group: 'com.google.auto.service', name: 'auto-service', version: '1.0-rc3'
   compile deps.slf4j
 
@@ -87,9 +86,7 @@ shadowJar {
     relocate 'org.msgpack', 'dd.deps.org.msgpack'
     relocate 'com.fasterxml', 'dd.deps.com.fasterxml'
 
-    relocate 'javassist', 'dd.deps.javassist'
     relocate 'net.bytebuddy', 'dd.deps.net.bytebuddy'
-    relocate 'org.reflections', 'dd.deps.org.reflections'
     relocate('org.jboss.byteman', 'dd.deps.org.jboss.byteman') {
       // Renaming these causes a verify error in the tests.
       exclude 'org.jboss.byteman.rule.*'

dd-java-agent/src/main/java/com/datadoghq/agent/AgentRulesManager.java

@@ -32,12 +32,10 @@ public class AgentRulesManager {
   protected final Retransformer transformer;
   protected final TracingAgentConfig agentTracerConfig;
   protected final InstrumentationRulesManager instrumentationRulesManager;
-  protected final TraceAnnotationsManager traceAnnotationsManager;
 
   public AgentRulesManager(final Retransformer trans, final TracingAgentConfig config) {
     transformer = trans;
     agentTracerConfig = config;
-    traceAnnotationsManager = new TraceAnnotationsManager(trans, config);
     instrumentationRulesManager = new InstrumentationRulesManager(trans, config, this);
   }
 
@@ -62,7 +60,6 @@ public static void initialize(final Retransformer trans) {
     INSTANCE = manager;
 
     manager.loadRules(INITIALIZER_RULES, ClassLoader.getSystemClassLoader());
-    manager.traceAnnotationsManager.initialize();
   }
 
   /**

dd-java-agent/src/main/java/com/datadoghq/agent/TraceAnnotationsManager.java

@@ -0,0 +1,57 @@
+package com.datadoghq.agent.instrumentation.annotation;
+
+import static net.bytebuddy.matcher.ElementMatchers.declaresMethod;
+import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
+import static net.bytebuddy.matcher.ElementMatchers.isAnnotatedWith;
+
+import com.datadoghq.agent.instrumentation.Instrumenter;
+import com.datadoghq.trace.Trace;
+import com.google.auto.service.AutoService;
+import io.opentracing.ActiveSpan;
+import io.opentracing.tag.Tags;
+import io.opentracing.util.GlobalTracer;
+import java.lang.reflect.Method;
+import java.sql.PreparedStatement;
+import java.util.Collections;
+import java.util.Map;
+import java.util.WeakHashMap;
+import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.asm.Advice;
+
+@AutoService(Instrumenter.class)
+public final class TraceAnnotationInstrumentation implements Instrumenter {
+  public static final Map<PreparedStatement, String> preparedStatements = new WeakHashMap<>();
+
+  @Override
+  public AgentBuilder instrument(final AgentBuilder agentBuilder) {
+    return agentBuilder
+        .type(hasSuperType(declaresMethod(isAnnotatedWith(Trace.class))))
+        .transform(
+            new AgentBuilder.Transformer.ForAdvice()
+                .advice(isAnnotatedWith(Trace.class), TraceAdvice.class.getName()));
+  }
+
+  public static class TraceAdvice {
+
+    @Advice.OnMethodEnter
+    public static ActiveSpan startSpan(@Advice.Origin final Method method) {
+      final Trace trace = method.getAnnotation(Trace.class);
+      String operationName = trace.operationName();
+      if (operationName == null || operationName.isEmpty()) {
+        operationName = method.getDeclaringClass().getName() + "." + method.getName();
+      }
+
+      return GlobalTracer.get().buildSpan(operationName).startActive();
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class)
+    public static void stopSpan(
+        @Advice.Enter final ActiveSpan activeSpan, @Advice.Thrown final Throwable throwable) {
+      if (throwable != null) {
+        Tags.ERROR.set(activeSpan, true);
+        activeSpan.log(Collections.singletonMap("error.object", throwable));
+      }
+      activeSpan.deactivate();
+    }
+  }
+}