Merge pull request #1031 from DataDog/landerson/trace-annotation-async

Enable Async Propagation for @Trace Annotatation instrumentation

dd-java-agent/instrumentation/netty-4.1/netty-4.1.gradle

@@ -38,6 +38,7 @@ dependencies {
   compileOnly group: 'io.netty', name: 'netty-codec-http', version: '4.1.0.Final'
 
   testCompile project(':dd-java-agent:instrumentation:java-concurrent')
+  testCompile project(':dd-java-agent:instrumentation:trace-annotation')
   testCompile group: 'io.netty', name: 'netty-codec-http', version: '4.1.0.Final'
   testCompile group: 'org.asynchttpclient', name: 'async-http-client', version: '2.1.0'
 

dd-java-agent/instrumentation/netty-4.1/src/test/groovy/Netty41ClientTest.groovy

@@ -1,4 +1,6 @@
+import datadog.opentracing.DDSpan
 import datadog.trace.agent.test.base.HttpClientTest
+import datadog.trace.api.Trace
 import datadog.trace.instrumentation.netty41.client.HttpClientTracingHandler
 import datadog.trace.instrumentation.netty41.client.NettyHttpClientDecorator
 import io.netty.channel.AbstractChannel
@@ -174,11 +176,54 @@ class Netty41ClientTest extends HttpClientTest<NettyHttpClientDecorator> {
     null != channel.pipeline().remove(HttpClientTracingHandler.getName())
   }
 
+  def "request with trace annotated method"() {
+    given:
+    def annotatedClass = new AnnotatedClass()
+
+    when:
+    def status = runUnderTrace("parent") {
+      annotatedClass.makeRequestUnderTrace(method)
+    }
+
+    then:
+    status == 200
+    assertTraces(2) {
+      server.distributedRequestTrace(it, 0, trace(1).last())
+      trace(1, size(3)) {
+        basicSpan(it, 0, "parent")
+        span(1) {
+          childOf((DDSpan) span(0))
+          serviceName "unnamed-java-app"
+          operationName "trace.annotation"
+          resourceName "AnnotatedClass.makeRequestUnderTrace"
+          errored false
+          tags {
+            defaultTags()
+            "$Tags.COMPONENT.key" "trace"
+          }
+        }
+        clientSpan(it, 2, span(1), method)
+      }
+    }
+
+    where:
+    method << BODY_METHODS
+  }
+
+  class AnnotatedClass {
+    @Trace
+    int makeRequestUnderTrace(String method) {
+      return doRequest(method, server.address.resolve("/success"))
+    }
+  }
+
   class SimpleHandler implements ChannelHandler {
     @Override
     void handlerAdded(ChannelHandlerContext ctx) throws Exception {}
+
     @Override
     void handlerRemoved(ChannelHandlerContext ctx) throws Exception {}
+
     @Override
     void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {}
   }
@@ -186,8 +231,10 @@ class Netty41ClientTest extends HttpClientTest<NettyHttpClientDecorator> {
   class OtherSimpleHandler implements ChannelHandler {
     @Override
     void handlerAdded(ChannelHandlerContext ctx) throws Exception {}
+
     @Override
     void handlerRemoved(ChannelHandlerContext ctx) throws Exception {}
+
     @Override
     void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {}
   }

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

@@ -4,6 +4,7 @@ import static datadog.trace.instrumentation.trace_annotation.TraceDecorator.DECO
 
 import datadog.trace.api.DDTags;
 import datadog.trace.api.Trace;
+import datadog.trace.context.TraceScope;
 import io.opentracing.Scope;
 import io.opentracing.Tracer;
 import io.opentracing.util.GlobalTracer;
@@ -29,7 +30,13 @@ public class TraceAdvice {
     }
     spanBuilder = spanBuilder.withTag(DDTags.RESOURCE_NAME, resourceName);
 
-    return DECORATE.afterStart(spanBuilder.startActive(true));
+    final Scope scope = DECORATE.afterStart(spanBuilder.startActive(true));
+
+    if (scope instanceof TraceScope) {
+      ((TraceScope) scope).setAsyncPropagation(true);
+    }
+
+    return scope;
   }
 
   @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)