Merge pull request #925 from DataDog/landerson/resourcename-on-annotation

Add resource name to the @Trace annotation

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

@@ -2,8 +2,10 @@ package datadog.trace.instrumentation.trace_annotation;
 
 import static datadog.trace.instrumentation.trace_annotation.TraceDecorator.DECORATE;
 
+import datadog.trace.api.DDTags;
 import datadog.trace.api.Trace;
 import io.opentracing.Scope;
+import io.opentracing.Tracer;
 import io.opentracing.util.GlobalTracer;
 import java.lang.reflect.Method;
 import net.bytebuddy.asm.Advice;
@@ -12,12 +14,20 @@ public class TraceAdvice {
 
   @Advice.OnMethodEnter(suppress = Throwable.class)
   public static Scope startSpan(@Advice.Origin final Method method) {
-    final Trace trace = method.getAnnotation(Trace.class);
+    final Trace traceAnnotation = method.getAnnotation(Trace.class);
-    String operationName = trace == null ? null : trace.operationName();
+    String operationName = traceAnnotation == null ? null : traceAnnotation.operationName();
     if (operationName == null || operationName.isEmpty()) {
       operationName = DECORATE.spanNameForMethod(method);
     }
-    return DECORATE.afterStart(GlobalTracer.get().buildSpan(operationName).startActive(true));
+
+    Tracer.SpanBuilder spanBuilder = GlobalTracer.get().buildSpan(operationName);
+
+    final String resourceName = traceAnnotation == null ? null : traceAnnotation.resourceName();
+    if (resourceName != null && !resourceName.isEmpty()) {
+      spanBuilder = spanBuilder.withTag(DDTags.RESOURCE_NAME, resourceName);
+    }
+
+    return DECORATE.afterStart(spanBuilder.startActive(true));
   }
 
   @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)

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

@@ -38,6 +38,77 @@ class TraceAnnotationsTest extends AgentTestRunner {
     }
   }
 
+  def "test simple case with only operation name set"() {
+    setup:
+    // Test single span in new trace
+    SayTracedHello.sayHA()
+
+    expect:
+    assertTraces(1) {
+      trace(0, 1) {
+        span(0) {
+          serviceName "test"
+          resourceName "SAY_HA"
+          operationName "SAY_HA"
+          spanType "DB"
+          parent()
+          errored false
+          tags {
+            "$Tags.COMPONENT.key" "trace"
+            defaultTags()
+          }
+        }
+      }
+    }
+  }
+
+  def "test simple case with only resource name set"() {
+    setup:
+    // Test single span in new trace
+    SayTracedHello.sayHelloOnlyResourceSet()
+
+    expect:
+    assertTraces(1) {
+      trace(0, 1) {
+        span(0) {
+          serviceName "test"
+          resourceName "WORLD"
+          operationName "SayTracedHello.sayHelloOnlyResourceSet"
+          parent()
+          errored false
+          tags {
+            "$Tags.COMPONENT.key" "trace"
+            defaultTags()
+          }
+        }
+      }
+    }
+  }
+
+  def "test simple case with both resource and operation name set"() {
+    setup:
+    // Test single span in new trace
+    SayTracedHello.sayHAWithResource()
+
+    expect:
+    assertTraces(1) {
+      trace(0, 1) {
+        span(0) {
+          serviceName "test"
+          resourceName "EARTH"
+          operationName "SAY_HA"
+          spanType "DB"
+          parent()
+          errored false
+          tags {
+            "$Tags.COMPONENT.key" "trace"
+            defaultTags()
+          }
+        }
+      }
+    }
+  }
+
   def "test complex case annotations"() {
     when:
     // Test new trace with 2 children spans
@@ -82,6 +153,94 @@ class TraceAnnotationsTest extends AgentTestRunner {
     }
   }
 
+  def "test complex case with resource name at top level"() {
+    when:
+    // Test new trace with 2 children spans
+    SayTracedHello.sayHELLOsayHAWithResource()
+
+    then:
+    assertTraces(1) {
+      trace(0, 3) {
+        span(0) {
+          resourceName "WORLD"
+          operationName "NEW_TRACE"
+          parent()
+          errored false
+          tags {
+            "$Tags.COMPONENT.key" "trace"
+            defaultTags()
+          }
+        }
+        span(1) {
+          resourceName "SAY_HA"
+          operationName "SAY_HA"
+          spanType "DB"
+          childOf span(0)
+          errored false
+          tags {
+            "$Tags.COMPONENT.key" "trace"
+            defaultTags()
+          }
+        }
+        span(2) {
+          serviceName "test"
+          resourceName "SayTracedHello.sayHello"
+          operationName "SayTracedHello.sayHello"
+          childOf span(0)
+          errored false
+          tags {
+            "$Tags.COMPONENT.key" "trace"
+            defaultTags()
+          }
+        }
+      }
+    }
+  }
+
+  def "test complex case with resource name at various levels"() {
+    when:
+    // Test new trace with 2 children spans
+    SayTracedHello.sayHELLOsayHAMixedResourceChildren()
+
+    then:
+    assertTraces(1) {
+      trace(0, 3) {
+        span(0) {
+          resourceName "WORLD"
+          operationName "NEW_TRACE"
+          parent()
+          errored false
+          tags {
+            "$Tags.COMPONENT.key" "trace"
+            defaultTags()
+          }
+        }
+        span(1) {
+          resourceName "EARTH"
+          operationName "SAY_HA"
+          spanType "DB"
+          childOf span(0)
+          errored false
+          tags {
+            "$Tags.COMPONENT.key" "trace"
+            defaultTags()
+          }
+        }
+        span(2) {
+          serviceName "test"
+          resourceName "SayTracedHello.sayHello"
+          operationName "SayTracedHello.sayHello"
+          childOf span(0)
+          errored false
+          tags {
+            "$Tags.COMPONENT.key" "trace"
+            defaultTags()
+          }
+        }
+      }
+    }
+  }
+
   def "test exception exit"() {
     setup:
 
@@ -111,6 +270,35 @@ class TraceAnnotationsTest extends AgentTestRunner {
     }
   }
 
+  def "test exception exit with resource name"() {
+    setup:
+
+    TEST_TRACER.addDecorator(new ErrorFlag())
+
+    Throwable error = null
+    try {
+      SayTracedHello.sayERRORWithResource()
+    } catch (final Throwable ex) {
+      error = ex
+    }
+
+    expect:
+    assertTraces(1) {
+      trace(0, 1) {
+        span(0) {
+          resourceName "WORLD"
+          operationName "ERROR"
+          errored true
+          tags {
+            "$Tags.COMPONENT.key" "trace"
+            errorTags(error.class)
+            defaultTags()
+          }
+        }
+      }
+    }
+  }
+
   def "test annonymous class annotations"() {
     setup:
     // Test anonymous classes with package.

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

@@ -15,6 +15,13 @@ public class SayTracedHello {
     return "hello!";
   }
 
+  @Trace(resourceName = "WORLD")
+  public static String sayHelloOnlyResourceSet() {
+    new StringTag(DDTags.SERVICE_NAME)
+        .set(GlobalTracer.get().scopeManager().active().span(), "test");
+    return "hello!";
+  }
+
   @Trace(operationName = "SAY_HA")
   public static String sayHA() {
     new StringTag(DDTags.SERVICE_NAME)
@@ -23,6 +30,14 @@ public class SayTracedHello {
     return "HA!!";
   }
 
+  @Trace(operationName = "SAY_HA", resourceName = "EARTH")
+  public static String sayHAWithResource() {
+    new StringTag(DDTags.SERVICE_NAME)
+        .set(GlobalTracer.get().scopeManager().active().span(), "test");
+    new StringTag(DDTags.SPAN_TYPE).set(GlobalTracer.get().scopeManager().active().span(), "DB");
+    return "HA EARTH!!";
+  }
+
   @Trace(operationName = "NEW_TRACE")
   public static String sayHELLOsayHA() {
     new StringTag(DDTags.SERVICE_NAME)
@@ -30,11 +45,30 @@ public class SayTracedHello {
     return sayHello() + sayHA();
   }
 
+  @Trace(operationName = "NEW_TRACE", resourceName = "WORLD")
+  public static String sayHELLOsayHAWithResource() {
+    new StringTag(DDTags.SERVICE_NAME)
+        .set(GlobalTracer.get().scopeManager().active().span(), "test2");
+    return sayHello() + sayHA();
+  }
+
+  @Trace(operationName = "NEW_TRACE", resourceName = "WORLD")
+  public static String sayHELLOsayHAMixedResourceChildren() {
+    new StringTag(DDTags.SERVICE_NAME)
+        .set(GlobalTracer.get().scopeManager().active().span(), "test2");
+    return sayHello() + sayHAWithResource();
+  }
+
   @Trace(operationName = "ERROR")
   public static String sayERROR() {
     throw new RuntimeException();
   }
 
+  @Trace(operationName = "ERROR", resourceName = "WORLD")
+  public static String sayERRORWithResource() {
+    throw new RuntimeException();
+  }
+
   public static String fromCallable() throws Exception {
     return new Callable<String>() {
       @com.newrelic.api.agent.Trace

dd-trace-api/src/main/java/datadog/trace/api/Trace.java

@@ -13,4 +13,7 @@ public @interface Trace {
 
   /** The operation name to set. By default it takes the method's name */
   String operationName() default "";
+
+  /** The resource name. By default it uses the same value as the operation name */
+  String resourceName() default "";
 }