Merge pull request #265 from DataDog/tyler/jax-client

JAX-RS Client Tracing

dd-java-agent/instrumentation/jax-rs-annotations/src/main/java/datadog/trace/instrumentation/jaxrs/JaxRsAnnotationsInstrumentation.java

@@ -22,10 +22,10 @@ import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.asm.Advice;
 
 @AutoService(Instrumenter.class)
-public final class JaxRsInstrumentation extends Instrumenter.Configurable {
+public final class JaxRsAnnotationsInstrumentation extends Instrumenter.Configurable {
 
-  public JaxRsInstrumentation() {
-    super("jax-rs", "jaxrs");
+  public JaxRsAnnotationsInstrumentation() {
+    super("jax-rs", "jaxrs", "jax-rs-annotations");
   }
 
   @Override
@@ -54,11 +54,11 @@ public final class JaxRsInstrumentation extends Instrumenter.Configurable {
                             .or(named("javax.ws.rs.OPTIONS"))
                             .or(named("javax.ws.rs.POST"))
                             .or(named("javax.ws.rs.PUT"))),
-                    JaxRsAdvice.class.getName()))
+                    JaxRsAnnotationsAdvice.class.getName()))
         .asDecorator();
   }
 
-  public static class JaxRsAdvice {
+  public static class JaxRsAnnotationsAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void nameSpan(@Advice.This final Object obj, @Advice.Origin final Method method) {

dd-java-agent/instrumentation/jax-rs-client/jax-rs-client.gradle

@@ -0,0 +1,45 @@
+apply plugin: 'version-scan'
+
+versionScan {
+  group = "javax.ws.rs"
+  module = "jsr311-api"
+  versions = "(,)"
+}
+
+apply from: "${rootDir}/gradle/java.gradle"
+
+if (JavaVersion.current() != JavaVersion.VERSION_1_8) {
+  sourceSets {
+    test {
+      groovy {
+        // These classes use Ratpack which requires Java 8. (Currently also incompatible with Java 9.)
+        exclude '**/*.groovy'
+      }
+    }
+  }
+}
+
+dependencies {
+  compileOnly group: 'javax.ws.rs', name: 'javax.ws.rs-api', version: '2.1'
+  compileOnly group: 'javax.annotation', name: 'javax.annotation-api', version: '1.2'
+
+
+  compile deps.bytebuddy
+  compile deps.opentracing
+  compile deps.autoservice
+
+  compile project(':dd-java-agent:agent-tooling')
+
+  testCompile project(':dd-java-agent:testing')
+
+//  testCompile group: 'javax.ws.rs', name: 'javax.ws.rs-api', version: '2.1'
+
+  testCompile group: 'org.glassfish.jersey.core', name: 'jersey-client', version: '2.25.1'
+  testCompile group: 'org.apache.cxf', name: 'cxf-rt-rs-client', version: '3.2.2'
+  testCompile group: 'org.jboss.resteasy', name: 'resteasy-client', version: '3.5.0.Final'
+
+//  testCompile group: 'com.sun.jersey', name: 'jersey-core', version: '1.19.4'
+//  testCompile group: 'com.sun.jersey', name: 'jersey-servlet', version: '1.19.4'
+//  testCompile group: 'io.dropwizard', name: 'dropwizard-testing', version: '0.7.1'
+//  testCompile group: 'javax.xml.bind', name: 'jaxb-api', version: '2.2.3'
+}

dd-java-agent/instrumentation/jax-rs-client/src/main/java/datadog/trace/instrumentation/jaxrs/ClientTracingFeature.java

@@ -0,0 +1,15 @@
+package datadog.trace.instrumentation.jaxrs;
+
+import javax.ws.rs.core.Feature;
+import javax.ws.rs.core.FeatureContext;
+import lombok.extern.slf4j.Slf4j;
+
+@Slf4j
+public class ClientTracingFeature implements Feature {
+  @Override
+  public boolean configure(final FeatureContext context) {
+    context.register(new ClientTracingFilter());
+    log.debug("ClientTracingFilter registered");
+    return true;
+  }
+}

dd-java-agent/instrumentation/jax-rs-client/src/main/java/datadog/trace/instrumentation/jaxrs/ClientTracingFilter.java

@@ -0,0 +1,60 @@
+package datadog.trace.instrumentation.jaxrs;
+
+import datadog.trace.api.DDSpanTypes;
+import datadog.trace.api.DDTags;
+import io.opentracing.Span;
+import io.opentracing.propagation.Format;
+import io.opentracing.tag.Tags;
+import io.opentracing.util.GlobalTracer;
+import java.io.IOException;
+import javax.annotation.Priority;
+import javax.ws.rs.Priorities;
+import javax.ws.rs.client.ClientRequestContext;
+import javax.ws.rs.client.ClientRequestFilter;
+import javax.ws.rs.client.ClientResponseContext;
+import javax.ws.rs.client.ClientResponseFilter;
+import lombok.extern.slf4j.Slf4j;
+
+@Slf4j
+@Priority(Priorities.HEADER_DECORATOR)
+public class ClientTracingFilter implements ClientRequestFilter, ClientResponseFilter {
+  private static final String PROPERTY_NAME = ClientTracingFilter.class.getName() + ".span";
+
+  @Override
+  public void filter(final ClientRequestContext requestContext) throws IOException {
+
+    final Span span =
+        GlobalTracer.get()
+            .buildSpan("jax-rs.client.call")
+            .withTag(Tags.COMPONENT.getKey(), "jax-rs.client")
+            .withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CLIENT)
+            .withTag(Tags.HTTP_METHOD.getKey(), requestContext.getMethod())
+            .withTag(Tags.HTTP_URL.getKey(), requestContext.getUri().toString())
+            .withTag(DDTags.SPAN_TYPE, DDSpanTypes.HTTP_CLIENT)
+            .withTag(DDTags.RESOURCE_NAME, requestContext.getMethod() + " jax-rs.client.call")
+            .start();
+
+    log.debug("{} - client span started", span);
+
+    GlobalTracer.get()
+        .inject(
+            span.context(),
+            Format.Builtin.HTTP_HEADERS,
+            new InjectAdapter(requestContext.getHeaders()));
+    requestContext.setProperty(PROPERTY_NAME, span);
+  }
+
+  @Override
+  public void filter(
+      final ClientRequestContext requestContext, final ClientResponseContext responseContext)
+      throws IOException {
+    final Object spanObj = requestContext.getProperty(PROPERTY_NAME);
+    if (spanObj instanceof Span) {
+      final Span span = (Span) spanObj;
+      Tags.HTTP_STATUS.set(span, responseContext.getStatus());
+
+      span.finish();
+      log.debug("{} - client spanObj finished", spanObj);
+    }
+  }
+}

dd-java-agent/instrumentation/jax-rs-client/src/main/java/datadog/trace/instrumentation/jaxrs/InjectAdapter.java

@@ -0,0 +1,26 @@
+package datadog.trace.instrumentation.jaxrs;
+
+import io.opentracing.propagation.TextMap;
+import java.util.Iterator;
+import java.util.Map;
+import javax.ws.rs.core.MultivaluedMap;
+
+public final class InjectAdapter implements TextMap {
+  private final MultivaluedMap<String, Object> map;
+
+  public InjectAdapter(final MultivaluedMap<String, Object> map) {
+    this.map = map;
+  }
+
+  @Override
+  public Iterator<Map.Entry<String, String>> iterator() {
+    throw new UnsupportedOperationException(
+        "InjectAdapter should only be used with Tracer.inject()");
+  }
+
+  @Override
+  public void put(final String key, final String value) {
+    // Don't allow duplicates.
+    map.putSingle(key, value);
+  }
+}

dd-java-agent/instrumentation/jax-rs-client/src/main/java/datadog/trace/instrumentation/jaxrs/JaxRsClientInstrumentation.java

@@ -0,0 +1,51 @@
+package datadog.trace.instrumentation.jaxrs;
+
+import static net.bytebuddy.matcher.ElementMatchers.failSafe;
+import static net.bytebuddy.matcher.ElementMatchers.hasSuperType;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.DDAdvice;
+import datadog.trace.agent.tooling.HelperInjector;
+import datadog.trace.agent.tooling.Instrumenter;
+import javax.ws.rs.client.ClientBuilder;
+import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.asm.Advice;
+
+@AutoService(Instrumenter.class)
+public final class JaxRsClientInstrumentation extends Instrumenter.Configurable {
+
+  public JaxRsClientInstrumentation() {
+    super("jax-rs", "jaxrs", "jax-rs-client");
+  }
+
+  @Override
+  protected boolean defaultEnabled() {
+    return false;
+  }
+
+  @Override
+  protected AgentBuilder apply(final AgentBuilder agentBuilder) {
+    return agentBuilder
+        .type(failSafe(hasSuperType(named("javax.ws.rs.client.ClientBuilder"))))
+        .transform(
+            new HelperInjector(
+                "datadog.trace.instrumentation.jaxrs.ClientTracingFeature",
+                "datadog.trace.instrumentation.jaxrs.ClientTracingFilter"))
+        .transform(
+            DDAdvice.create()
+                .advice(
+                    named("build").and(returns(hasSuperType(named("javax.ws.rs.client.Client")))),
+                    ClientBuilderAdvice.class.getName()))
+        .asDecorator();
+  }
+
+  public static class ClientBuilderAdvice {
+
+    @Advice.OnMethodEnter
+    public static void registerFeature(@Advice.This final ClientBuilder builder) {
+      builder.register(ClientTracingFeature.class);
+    }
+  }
+}

dd-java-agent/instrumentation/jax-rs-client/src/test/groovy/JaxRsClientTest.groovy

@@ -0,0 +1,96 @@
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.api.DDSpanTypes
+import datadog.trace.api.DDTags
+import io.opentracing.tag.Tags
+import org.apache.cxf.jaxrs.client.spec.ClientBuilderImpl
+import org.glassfish.jersey.client.JerseyClientBuilder
+import org.jboss.resteasy.client.jaxrs.ResteasyClientBuilder
+import ratpack.http.Headers
+import spock.lang.Unroll
+
+import javax.ws.rs.client.AsyncInvoker
+import javax.ws.rs.client.Client
+import javax.ws.rs.client.Invocation
+import javax.ws.rs.client.WebTarget
+import javax.ws.rs.core.MediaType
+import javax.ws.rs.core.Response
+import java.util.concurrent.atomic.AtomicReference
+
+import static ratpack.groovy.test.embed.GroovyEmbeddedApp.ratpack
+
+//@Timeout(10)
+class JaxRsClientTest extends AgentTestRunner {
+  static {
+    System.setProperty("dd.integration.jax-rs.enabled", "true")
+  }
+
+  def receivedHeaders = new AtomicReference<Headers>()
+  def server = ratpack {
+    handlers {
+      all {
+        receivedHeaders.set(request.headers)
+        response.status(200).send("pong")
+      }
+    }
+  }
+
+  @Unroll
+  def "#lib request creates spans and sends headers"() {
+    setup:
+    Client client = builder.build()
+    WebTarget service = client.target("http://localhost:$server.address.port/ping")
+    Response response
+    if (async) {
+      AsyncInvoker request = service.request(MediaType.TEXT_PLAIN).async()
+      response = request.get().get()
+    } else {
+      Invocation.Builder request = service.request(MediaType.TEXT_PLAIN)
+      response = request.get()
+    }
+
+    expect:
+    response.readEntity(String) == "pong"
+
+    TEST_WRITER.size() == 1
+
+    def trace = TEST_WRITER.firstTrace()
+    trace.size() == 1
+
+    and:
+    def span = trace[0]
+
+    span.context().operationName == "jax-rs.client.call"
+    span.serviceName == "unnamed-java-app"
+    span.resourceName == "GET jax-rs.client.call"
+    span.type == "http"
+    !span.context().getErrorFlag()
+    span.context().parentId == 0
+
+
+    def tags = span.context().tags
+    tags[Tags.COMPONENT.key] == "jax-rs.client"
+    tags[Tags.SPAN_KIND.key] == Tags.SPAN_KIND_CLIENT
+    tags[Tags.HTTP_METHOD.key] == "GET"
+    tags[Tags.HTTP_STATUS.key] == 200
+    tags[Tags.HTTP_URL.key] == "http://localhost:$server.address.port/ping"
+    tags[DDTags.SPAN_TYPE] == DDSpanTypes.HTTP_CLIENT
+    tags[DDTags.THREAD_NAME] != null
+    tags[DDTags.THREAD_ID] != null
+    tags.size() == 8
+
+    receivedHeaders.get().get("x-datadog-trace-id") == "$span.traceId"
+    receivedHeaders.get().get("x-datadog-parent-id") == "$span.spanId"
+
+    cleanup:
+    server.close()
+
+    where:
+    builder                     | async | lib
+    new JerseyClientBuilder()   | false | "jersey"
+    new ClientBuilderImpl()     | false | "cxf"
+    new ResteasyClientBuilder() | false | "resteasy"
+    new JerseyClientBuilder()   | true  | "jersey async"
+    new ClientBuilderImpl()     | true  | "cxf async"
+    new ResteasyClientBuilder() | true  | "resteasy async"
+  }
+}

dd-trace-ot/src/test/groovy/datadog/opentracing/scopemanager/ScopeManagerTest.groovy

@@ -209,8 +209,10 @@ class ScopeManagerTest extends Specification {
 
     when:
     def newScope = continuation.activate()
+    def newContinuation = newScope.capture(true)
 
     then:
+    newScope instanceof ContinuableScope.Continuation.ClosingScope
     scopeManager.active() == newScope.wrappedScope
     newScope != childScope && newScope != parentScope
     newScope.span() == childSpan
@@ -220,6 +222,7 @@ class ScopeManagerTest extends Specification {
 
     when:
     newScope.close()
+    newContinuation.activate().close()
 
     then:
     scopeManager.active() == null
@@ -240,6 +243,7 @@ class ScopeManagerTest extends Specification {
     def newScope = continuation.activate()
 
     expect:
+    newScope instanceof ContinuableScope.Continuation.ClosingScope
     newScope != scope
     scopeManager.active() == newScope.wrappedScope
     spanFinished(span)

settings.gradle

@@ -14,7 +14,8 @@ include ':dd-java-agent:instrumentation:apache-httpclient-4.3'
 include ':dd-java-agent:instrumentation:aws-sdk'
 include ':dd-java-agent:instrumentation:classloaders'
 include ':dd-java-agent:instrumentation:datastax-cassandra-3.2'
-include ':dd-java-agent:instrumentation:jax-rs'
+include ':dd-java-agent:instrumentation:jax-rs-annotations'
+include ':dd-java-agent:instrumentation:jax-rs-client'
 include ':dd-java-agent:instrumentation:java-concurrent'
 include ':dd-java-agent:instrumentation:jboss-classloading'
 include ':dd-java-agent:instrumentation:jdbc'