Add support for jax-rs AsyncResponse

dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/AkkaForkJoinTaskInstrumentation.java

@@ -1,6 +1,7 @@
 package datadog.trace.instrumentation.java.concurrent;
 
 import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
+import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.isAbstract;
 import static net.bytebuddy.matcher.ElementMatchers.isInterface;
 import static net.bytebuddy.matcher.ElementMatchers.named;
@@ -64,11 +65,9 @@ public final class AkkaForkJoinTaskInstrumentation extends Instrumenter.Default
 
   @Override
   public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
-    final Map<ElementMatcher<? super MethodDescription>, String> transformers = new HashMap<>();
-    transformers.put(
+    return singletonMap(
         named("exec").and(takesArguments(0)).and(not(isAbstract())),
         ForkJoinTaskAdvice.class.getName());
-    return transformers;
   }
 
   public static class ForkJoinTaskAdvice {

dd-java-agent/instrumentation/jax-rs-annotations-1/jax-rs-annotations-1.gradle

@@ -4,6 +4,7 @@ muzzle {
     module = "jsr311-api"
     versions = "[0.5,)"
   }
+  // Muzzle doesn't detect the classLoaderMatcher, so we can't assert fail for v2 api.
 }
 
 apply from: "${rootDir}/gradle/java.gradle"
@@ -11,8 +12,6 @@ apply from: "${rootDir}/gradle/java.gradle"
 dependencies {
   compileOnly group: 'javax.ws.rs', name: 'jsr311-api', version: '1.1.1'
 
-  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-annotations-1/src/main/java/datadog/trace/instrumentation/jaxrs1/JaxRsAnnotationsDecorator.java

@@ -1,4 +1,4 @@
-package datadog.trace.instrumentation.jaxrs;
+package datadog.trace.instrumentation.jaxrs1;
 
 import static datadog.trace.bootstrap.WeakMap.Provider.newWeakMap;
 

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

@@ -1,11 +1,13 @@
-package datadog.trace.instrumentation.jaxrs;
+package datadog.trace.instrumentation.jaxrs1;
 
 import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
-import static datadog.trace.instrumentation.jaxrs.JaxRsAnnotationsDecorator.DECORATE;
+import static datadog.trace.agent.tooling.ClassLoaderMatcher.classLoaderHasClasses;
+import static datadog.trace.instrumentation.jaxrs1.JaxRsAnnotationsDecorator.DECORATE;
 import static java.util.Collections.singletonMap;
 import static net.bytebuddy.matcher.ElementMatchers.declaresMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isAnnotatedWith;
 import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.not;
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
@@ -29,6 +31,12 @@ public final class JaxRsAnnotationsInstrumentation extends Instrumenter.Default
     super("jax-rs", "jaxrs", "jax-rs-annotations");
   }
 
+  // this is required to make sure instrumentation won't apply to jax-rs 2
+  @Override
+  public ElementMatcher<ClassLoader> classLoaderMatcher() {
+    return not(classLoaderHasClasses("javax.ws.rs.container.AsyncResponse"));
+  }
+
   @Override
   public ElementMatcher<TypeDescription> typeMatcher() {
     return safeHasSuperType(

dd-java-agent/instrumentation/jax-rs-annotations-1/src/test/groovy/JaxRsAnnotationsInstrumentationTest.groovy

@@ -1,6 +1,5 @@
 import datadog.trace.agent.test.AgentTestRunner
 import io.opentracing.tag.Tags
-
 import javax.ws.rs.DELETE
 import javax.ws.rs.GET
 import javax.ws.rs.HEAD
@@ -10,7 +9,7 @@ import javax.ws.rs.PUT
 import javax.ws.rs.Path
 
 import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
-import static datadog.trace.instrumentation.jaxrs.JaxRsAnnotationsDecorator.DECORATE
+import static datadog.trace.instrumentation.jaxrs1.JaxRsAnnotationsDecorator.DECORATE
 
 class JaxRsAnnotationsInstrumentationTest extends AgentTestRunner {
 
@@ -19,7 +18,8 @@ class JaxRsAnnotationsInstrumentationTest extends AgentTestRunner {
     new Jax() {
       @POST
       @Path("/a")
-      void call() {}
+      void call() {
+      }
     }.call()
 
     expect:
@@ -86,39 +86,47 @@ class JaxRsAnnotationsInstrumentationTest extends AgentTestRunner {
     name                        | obj
     "/a"                        | new Jax() {
       @Path("/a")
-      void call() {}
+      void call() {
+      }
     }
     "GET /b"                    | new Jax() {
       @GET
       @Path("/b")
-      void call() {}
+      void call() {
+      }
     }
     "POST /c"                   | new InterfaceWithPath() {
       @POST
       @Path("/c")
-      void call() {}
+      void call() {
+      }
     }
     "HEAD"                      | new InterfaceWithPath() {
       @HEAD
-      void call() {}
+      void call() {
+      }
     }
     "POST /abstract/d"          | new AbstractClassWithPath() {
       @POST
       @Path("/d")
-      void call() {}
+      void call() {
+      }
     }
     "PUT /abstract"             | new AbstractClassWithPath() {
       @PUT
-      void call() {}
+      void call() {
+      }
     }
     "OPTIONS /abstract/child/e" | new ChildClassWithPath() {
       @OPTIONS
       @Path("/e")
-      void call() {}
+      void call() {
+      }
     }
     "DELETE /abstract/child"    | new ChildClassWithPath() {
       @DELETE
-      void call() {}
+      void call() {
+      }
     }
     "POST /abstract/child/call" | new ChildClassWithPath()
 
@@ -147,16 +155,20 @@ class JaxRsAnnotationsInstrumentationTest extends AgentTestRunner {
     where:
     obj | _
     new Jax() {
-      void call() {}
+      void call() {
+      }
     }   | _
     new InterfaceWithPath() {
-      void call() {}
+      void call() {
+      }
     }   | _
     new AbstractClassWithPath() {
-      void call() {}
+      void call() {
+      }
     }   | _
     new ChildClassWithPath() {
-      void call() {}
+      void call() {
+      }
     }   | _
   }
 
@@ -180,7 +192,8 @@ class JaxRsAnnotationsInstrumentationTest extends AgentTestRunner {
   class ChildClassWithPath extends AbstractClassWithPath {
     @Path("call")
     @POST
-    void call() {}
+    void call() {
+    }
   }
 
   def getName(Class clazz) {

dd-java-agent/instrumentation/jax-rs-annotations-1/src/test/groovy/JerseyTest.groovy

@@ -7,7 +7,6 @@ import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
 
 class JerseyTest extends AgentTestRunner {
 
-  // FIXME: migrate test.
   @Shared
   @ClassRule
   ResourceTestRule resources = ResourceTestRule.builder().addResource(new TestResource()).build()

dd-java-agent/instrumentation/jax-rs-annotations-1/src/test/groovy/JettyTestInstrumentation.java

@@ -0,0 +1,27 @@
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.test.base.HttpServerTestAdvice;
+import datadog.trace.agent.tooling.Instrumenter;
+import net.bytebuddy.agent.builder.AgentBuilder;
+
+@AutoService(Instrumenter.class)
+public class JettyTestInstrumentation implements Instrumenter {
+
+  @Override
+  public AgentBuilder instrument(final AgentBuilder agentBuilder) {
+    return agentBuilder
+        // Jetty 8
+        .type(named("org.eclipse.jetty.server.AbstractHttpConnection"))
+        .transform(
+            new AgentBuilder.Transformer.ForAdvice()
+                .advice(
+                    named("headerComplete"),
+                    HttpServerTestAdvice.ServerEntryAdvice.class.getName()))
+        // Jetty 9
+        .type(named("org.eclipse.jetty.server.HttpChannel"))
+        .transform(
+            new AgentBuilder.Transformer.ForAdvice()
+                .advice(named("handle"), HttpServerTestAdvice.ServerEntryAdvice.class.getName()));
+  }
+}

dd-java-agent/instrumentation/jax-rs-annotations-2/jax-rs-annotations-2.gradle

@@ -0,0 +1,37 @@
+muzzle {
+  fail {
+    group = "javax.ws.rs"
+    module = "jsr311-api"
+    versions = "[,]"
+  }
+  pass {
+    group = "javax.ws.rs"
+    module = "javax.ws.rs-api"
+    versions = "[,]"
+  }
+}
+
+apply from: "${rootDir}/gradle/java.gradle"
+
+//apply plugin: 'org.unbroken-dome.test-sets'
+//
+//testSets {
+//  latestDepTest {
+//    dirName = 'test'
+//  }
+//}
+
+dependencies {
+  compileOnly group: 'javax.ws.rs', name: 'javax.ws.rs-api', version: '2.0'
+
+  testCompile project(':dd-java-agent:instrumentation:java-concurrent')
+  testCompile project(':dd-java-agent:instrumentation:servlet-3')
+
+  // First version with DropwizardTestSupport:
+  testCompile group: 'io.dropwizard', name: 'dropwizard-testing', version: '0.8.0'
+  testCompile group: 'javax.xml.bind', name: 'jaxb-api', version: '2.2.3'
+  testCompile group: 'com.fasterxml.jackson.module', name: 'jackson-module-afterburner', version: '2.9.10'
+
+  // Anything 1.0+ fails with a java.lang.NoClassDefFoundError: org/eclipse/jetty/server/RequestLog
+//  latestDepTestCompile group: 'io.dropwizard', name: 'dropwizard-testing', version: '1.+'
+}

dd-java-agent/instrumentation/jax-rs-annotations-2/src/main/java/datadog/trace/instrumentation/jaxrs2/JaxRsAnnotationsDecorator.java

@@ -0,0 +1,143 @@
+package datadog.trace.instrumentation.jaxrs2;
+
+import static datadog.trace.bootstrap.WeakMap.Provider.newWeakMap;
+
+import datadog.trace.agent.decorator.BaseDecorator;
+import datadog.trace.api.DDSpanTypes;
+import datadog.trace.api.DDTags;
+import datadog.trace.bootstrap.WeakMap;
+import io.opentracing.Scope;
+import io.opentracing.Span;
+import io.opentracing.tag.Tags;
+import java.lang.annotation.Annotation;
+import java.lang.reflect.Method;
+import java.util.LinkedList;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import javax.ws.rs.HttpMethod;
+import javax.ws.rs.Path;
+
+public class JaxRsAnnotationsDecorator extends BaseDecorator {
+  public static JaxRsAnnotationsDecorator DECORATE = new JaxRsAnnotationsDecorator();
+
+  private final WeakMap<Class, Map<Method, String>> resourceNames = newWeakMap();
+
+  @Override
+  protected String[] instrumentationNames() {
+    return new String[0];
+  }
+
+  @Override
+  protected String spanType() {
+    return null;
+  }
+
+  @Override
+  protected String component() {
+    return "jax-rs-controller";
+  }
+
+  public void onControllerStart(final Scope scope, final Scope parent, final Method method) {
+    final String resourceName = getPathResourceName(method);
+    updateParent(parent, resourceName);
+
+    final Span span = scope.span();
+    span.setTag(DDTags.SPAN_TYPE, DDSpanTypes.HTTP_SERVER);
+
+    // When jax-rs is the root, we want to name using the path, otherwise use the class/method.
+    final boolean isRootScope = parent == null;
+    if (isRootScope && !resourceName.isEmpty()) {
+      span.setTag(DDTags.RESOURCE_NAME, resourceName);
+    } else {
+      span.setTag(DDTags.RESOURCE_NAME, DECORATE.spanNameForMethod(method));
+    }
+  }
+
+  private void updateParent(final Scope scope, final String resourceName) {
+    if (scope == null) {
+      return;
+    }
+    final Span span = scope.span();
+    Tags.COMPONENT.set(span, "jax-rs");
+
+    if (!resourceName.isEmpty()) {
+      span.setTag(DDTags.RESOURCE_NAME, resourceName);
+    }
+  }
+
+  /**
+   * Returns the resource name given a JaxRS annotated method. Results are cached so this method can
+   * be called multiple times without significantly impacting performance.
+   *
+   * @return The result can be an empty string but will never be {@code null}.
+   */
+  private String getPathResourceName(final Method method) {
+    final Class<?> target = method.getDeclaringClass();
+    Map<Method, String> classMap = resourceNames.get(target);
+
+    if (classMap == null) {
+      resourceNames.putIfAbsent(target, new ConcurrentHashMap<Method, String>());
+      classMap = resourceNames.get(target);
+      // classMap should not be null at this point because we have a
+      // strong reference to target and don't manually clear the map.
+    }
+
+    String resourceName = classMap.get(method);
+    if (resourceName == null) {
+      final String httpMethod = locateHttpMethod(method);
+      final LinkedList<Path> paths = gatherPaths(method);
+      resourceName = buildResourceName(httpMethod, paths);
+      classMap.put(method, resourceName);
+    }
+
+    return resourceName;
+  }
+
+  private String locateHttpMethod(final Method method) {
+    String httpMethod = null;
+    for (final Annotation ann : method.getDeclaredAnnotations()) {
+      if (ann.annotationType().getAnnotation(HttpMethod.class) != null) {
+        httpMethod = ann.annotationType().getSimpleName();
+      }
+    }
+    return httpMethod;
+  }
+
+  private LinkedList<Path> gatherPaths(final Method method) {
+    Class<?> target = method.getDeclaringClass();
+    final LinkedList<Path> paths = new LinkedList<>();
+    while (target != Object.class) {
+      final Path annotation = target.getAnnotation(Path.class);
+      if (annotation != null) {
+        paths.push(annotation);
+      }
+      target = target.getSuperclass();
+    }
+    final Path methodPath = method.getAnnotation(Path.class);
+    if (methodPath != null) {
+      paths.add(methodPath);
+    }
+    return paths;
+  }
+
+  private String buildResourceName(final String httpMethod, final LinkedList<Path> paths) {
+    final String resourceName;
+    final StringBuilder resourceNameBuilder = new StringBuilder();
+    if (httpMethod != null) {
+      resourceNameBuilder.append(httpMethod);
+      resourceNameBuilder.append(" ");
+    }
+    Path last = null;
+    for (final Path path : paths) {
+      if (path.value().startsWith("/") || (last != null && last.value().endsWith("/"))) {
+        resourceNameBuilder.append(path.value());
+      } else {
+        resourceNameBuilder.append("/");
+        resourceNameBuilder.append(path.value());
+      }
+      last = path;
+    }
+    resourceName = resourceNameBuilder.toString().trim();
+    return resourceName;
+  }
+}

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

@@ -0,0 +1,116 @@
+package datadog.trace.instrumentation.jaxrs2;
+
+import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.declaresMethod;
+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.Instrumenter;
+import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.context.TraceScope;
+import io.opentracing.Scope;
+import io.opentracing.Span;
+import io.opentracing.Tracer;
+import io.opentracing.util.GlobalTracer;
+import java.lang.reflect.Method;
+import java.util.Collections;
+import java.util.Map;
+import javax.ws.rs.container.AsyncResponse;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(Instrumenter.class)
+public final class JaxRsAnnotationsInstrumentation extends Instrumenter.Default {
+
+  private static final String JAX_ENDPOINT_OPERATION_NAME = "jax-rs.request";
+
+  public JaxRsAnnotationsInstrumentation() {
+    super("jax-rs", "jaxrs", "jax-rs-annotations");
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    return Collections.singletonMap("javax.ws.rs.container.AsyncResponse", Span.class.getName());
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return safeHasSuperType(
+        isAnnotatedWith(named("javax.ws.rs.Path"))
+            .or(safeHasSuperType(declaresMethod(isAnnotatedWith(named("javax.ws.rs.Path"))))));
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      "datadog.trace.agent.decorator.BaseDecorator", packageName + ".JaxRsAnnotationsDecorator",
+    };
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+        isAnnotatedWith(
+            named("javax.ws.rs.Path")
+                .or(named("javax.ws.rs.DELETE"))
+                .or(named("javax.ws.rs.GET"))
+                .or(named("javax.ws.rs.HEAD"))
+                .or(named("javax.ws.rs.OPTIONS"))
+                .or(named("javax.ws.rs.POST"))
+                .or(named("javax.ws.rs.PUT"))),
+        JaxRsAnnotationsAdvice.class.getName());
+  }
+
+  public static class JaxRsAnnotationsAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static Scope nameSpan(@Advice.Origin final Method method) {
+      final Tracer tracer = GlobalTracer.get();
+      // Rename the parent span according to the path represented by these annotations.
+      final Scope parent = tracer.scopeManager().active();
+      final Scope scope = tracer.buildSpan(JAX_ENDPOINT_OPERATION_NAME).startActive(false);
+
+      if (scope instanceof TraceScope) {
+        ((TraceScope) scope).setAsyncPropagation(true);
+      }
+
+      JaxRsAnnotationsDecorator.DECORATE.onControllerStart(scope, parent, method);
+      return JaxRsAnnotationsDecorator.DECORATE.afterStart(scope);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void stopSpan(
+        @Advice.Enter final Scope scope,
+        @Advice.Thrown final Throwable throwable,
+        @Advice.AllArguments final Object[] args) {
+      final Span span = scope.span();
+      if (throwable != null) {
+        JaxRsAnnotationsDecorator.DECORATE.onError(span, throwable);
+        JaxRsAnnotationsDecorator.DECORATE.beforeFinish(span);
+        span.finish();
+        scope.close();
+        return;
+      }
+
+      AsyncResponse asyncResponse = null;
+      for (final Object arg : args) {
+        if (arg instanceof AsyncResponse) {
+          asyncResponse = (AsyncResponse) arg;
+          break;
+        }
+      }
+      if (asyncResponse != null && asyncResponse.isSuspended()) {
+        InstrumentationContext.get(AsyncResponse.class, Span.class).put(asyncResponse, span);
+      } else {
+        System.out.println("FINISHING: " + asyncResponse);
+        JaxRsAnnotationsDecorator.DECORATE.beforeFinish(span);
+        span.finish();
+      }
+      scope.close();
+    }
+  }
+}

dd-java-agent/instrumentation/jax-rs-annotations-2/src/main/java/datadog/trace/instrumentation/jaxrs2/JaxRsAsyncResponseInstrumentation.java

@@ -0,0 +1,114 @@
+package datadog.trace.instrumentation.jaxrs2;
+
+import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
+import static datadog.trace.instrumentation.jaxrs2.JaxRsAnnotationsDecorator.DECORATE;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.InstrumentationContext;
+import io.opentracing.Span;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import javax.ws.rs.container.AsyncResponse;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(Instrumenter.class)
+public final class JaxRsAsyncResponseInstrumentation extends Instrumenter.Default {
+
+  public JaxRsAsyncResponseInstrumentation() {
+    super("jax-rs", "jaxrs", "jax-rs-annotations");
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    return Collections.singletonMap("javax.ws.rs.container.AsyncResponse", Span.class.getName());
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return safeHasSuperType(named("javax.ws.rs.container.AsyncResponse"));
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      "datadog.trace.agent.decorator.BaseDecorator", packageName + ".JaxRsAnnotationsDecorator",
+    };
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    final Map<ElementMatcher<? super MethodDescription>, String> transformers = new HashMap<>();
+    transformers.put(
+        named("resume").and(takesArgument(0, Object.class)).and(isPublic()),
+        AsyncResponseAdvice.class.getName());
+    transformers.put(
+        named("resume").and(takesArgument(0, Throwable.class)).and(isPublic()),
+        AsyncResponseThrowableAdvice.class.getName());
+    transformers.put(named("cancel"), AsyncResponseCancelAdvice.class.getName());
+    return transformers;
+  }
+
+  public static class AsyncResponseAdvice {
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void stopSpan(@Advice.This final AsyncResponse asyncResponse) {
+
+      final ContextStore<AsyncResponse, Span> contextStore =
+          InstrumentationContext.get(AsyncResponse.class, Span.class);
+
+      final Span span = contextStore.get(asyncResponse);
+      if (span != null) {
+        contextStore.put(asyncResponse, null);
+        DECORATE.beforeFinish(span);
+        span.finish();
+      }
+    }
+  }
+
+  public static class AsyncResponseThrowableAdvice {
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void stopSpan(
+        @Advice.This final AsyncResponse asyncResponse,
+        @Advice.Argument(0) final Throwable throwable) {
+
+      final ContextStore<AsyncResponse, Span> contextStore =
+          InstrumentationContext.get(AsyncResponse.class, Span.class);
+
+      final Span span = contextStore.get(asyncResponse);
+      if (span != null) {
+        contextStore.put(asyncResponse, null);
+        DECORATE.onError(span, throwable);
+        DECORATE.beforeFinish(span);
+        span.finish();
+      }
+    }
+  }
+
+  public static class AsyncResponseCancelAdvice {
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void stopSpan(@Advice.This final AsyncResponse asyncResponse) {
+
+      final ContextStore<AsyncResponse, Span> contextStore =
+          InstrumentationContext.get(AsyncResponse.class, Span.class);
+
+      final Span span = contextStore.get(asyncResponse);
+      if (span != null) {
+        contextStore.put(asyncResponse, null);
+        span.setTag("canceled", true);
+        DECORATE.beforeFinish(span);
+        span.finish();
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/jax-rs-annotations-2/src/test/groovy/DropwizardAsyncTest.groovy

@@ -0,0 +1,91 @@
+import io.dropwizard.Application
+import io.dropwizard.Configuration
+import io.dropwizard.setup.Bootstrap
+import io.dropwizard.setup.Environment
+import javax.ws.rs.GET
+import javax.ws.rs.Path
+import javax.ws.rs.container.AsyncResponse
+import javax.ws.rs.container.Suspended
+import javax.ws.rs.core.Response
+
+import java.util.concurrent.Executors
+
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.ERROR
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.EXCEPTION
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.REDIRECT
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SUCCESS
+
+class DropwizardAsyncTest extends DropwizardTest {
+
+  Class testApp() {
+    AsyncTestApp
+  }
+
+  Class testResource() {
+    AsyncServiceResource
+  }
+
+  static class AsyncTestApp extends Application<Configuration> {
+    @Override
+    void initialize(Bootstrap<Configuration> bootstrap) {
+    }
+
+    @Override
+    void run(Configuration configuration, Environment environment) {
+      environment.jersey().register(AsyncServiceResource)
+    }
+  }
+
+  @Override
+  // Return the handler span's name
+  String reorderHandlerSpan() {
+    "jax-rs.request"
+  }
+
+  @Path("/")
+  static class AsyncServiceResource {
+    final executor = Executors.newSingleThreadExecutor()
+
+    @GET
+    @Path("success")
+    void success(@Suspended final AsyncResponse asyncResponse) {
+      executor.execute {
+        controller(SUCCESS) {
+          asyncResponse.resume(Response.status(SUCCESS.status).entity(SUCCESS.body).build())
+        }
+      }
+    }
+
+    @GET
+    @Path("redirect")
+    void redirect(@Suspended final AsyncResponse asyncResponse) {
+      executor.execute {
+        controller(REDIRECT) {
+          asyncResponse.resume(Response.status(REDIRECT.status).location(new URI(REDIRECT.body)).build())
+        }
+      }
+    }
+
+    @GET
+    @Path("error-status")
+    void error(@Suspended final AsyncResponse asyncResponse) {
+      executor.execute {
+        controller(ERROR) {
+          asyncResponse.resume(Response.status(ERROR.status).entity(ERROR.body).build())
+        }
+      }
+    }
+
+    @GET
+    @Path("exception")
+    void exception(@Suspended final AsyncResponse asyncResponse) {
+      executor.execute {
+        controller(EXCEPTION) {
+          def ex = new Exception(EXCEPTION.body)
+          asyncResponse.resume(ex)
+          throw ex
+        }
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/jax-rs-annotations-2/src/test/groovy/DropwizardTest.groovy

@@ -0,0 +1,179 @@
+import datadog.opentracing.DDSpan
+import datadog.trace.agent.test.asserts.TraceAssert
+import datadog.trace.agent.test.base.HttpServerTest
+import datadog.trace.api.DDSpanTypes
+import datadog.trace.instrumentation.jaxrs2.JaxRsAnnotationsDecorator
+import datadog.trace.instrumentation.servlet3.Servlet3Decorator
+import io.dropwizard.Application
+import io.dropwizard.Configuration
+import io.dropwizard.setup.Bootstrap
+import io.dropwizard.setup.Environment
+import io.dropwizard.testing.ConfigOverride
+import io.dropwizard.testing.DropwizardTestSupport
+import io.opentracing.tag.Tags
+import javax.ws.rs.GET
+import javax.ws.rs.Path
+import javax.ws.rs.core.Response
+import org.eclipse.jetty.servlet.ServletHandler
+
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.ERROR
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.EXCEPTION
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.REDIRECT
+import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SUCCESS
+
+class DropwizardTest extends HttpServerTest<DropwizardTestSupport, Servlet3Decorator> {
+  
+  @Override
+  DropwizardTestSupport startServer(int port) {
+    def testSupport = new DropwizardTestSupport(testApp(),
+      null,
+      ConfigOverride.config("server.applicationConnectors[0].port", "$port"))
+    testSupport.before()
+    return testSupport
+  }
+
+  Class testApp() {
+    TestApp
+  }
+
+  Class testResource() {
+    ServiceResource
+  }
+
+  @Override
+  void stopServer(DropwizardTestSupport testSupport) {
+    testSupport.after()
+  }
+
+  @Override
+  Servlet3Decorator decorator() {
+    return new Servlet3Decorator() {
+      @Override
+      protected String component() {
+        return "jax-rs"
+      }
+    }
+  }
+
+  @Override
+  String expectedOperationName() {
+    return "servlet.request"
+  }
+
+  @Override
+  boolean hasHandlerSpan() {
+    true
+  }
+
+  @Override
+  boolean testNotFound() {
+    false
+  }
+
+  boolean testExceptionBody() {
+    false
+  }
+
+  @Override
+  void handlerSpan(TraceAssert trace, int index, Object parent, ServerEndpoint endpoint = SUCCESS) {
+    trace.span(index) {
+      serviceName expectedServiceName()
+      operationName "jax-rs.request"
+      resourceName "${testResource().simpleName}.${endpoint.name().toLowerCase()}"
+      spanType DDSpanTypes.HTTP_SERVER
+      errored endpoint == EXCEPTION
+      childOf(parent as DDSpan)
+      tags {
+        "$Tags.COMPONENT.key" JaxRsAnnotationsDecorator.DECORATE.component()
+        defaultTags()
+        if (endpoint == EXCEPTION) {
+          errorTags(Exception, EXCEPTION.body)
+        }
+      }
+    }
+  }
+
+  @Override
+  void serverSpan(TraceAssert trace, int index, String traceID = null, String parentID = null, String method = "GET", ServerEndpoint endpoint = SUCCESS) {
+    trace.span(index) {
+      serviceName expectedServiceName()
+      operationName expectedOperationName()
+      resourceName endpoint.status == 404 ? "404" : "$method ${endpoint.resolve(address).path}"
+      spanType DDSpanTypes.HTTP_SERVER
+      errored endpoint.errored
+      if (parentID != null) {
+        traceId traceID
+        parentId parentID
+      } else {
+        parent()
+      }
+      tags {
+        "span.origin.type" ServletHandler.CachedChain.name
+
+        defaultTags(true)
+        "$Tags.COMPONENT.key" serverDecorator.component()
+        if (endpoint.errored) {
+          "$Tags.ERROR.key" endpoint.errored
+          "error.msg" { it == null || it == EXCEPTION.body }
+          "error.type" { it == null || it == Exception.name }
+          "error.stack" { it == null || it instanceof String }
+        }
+        "$Tags.HTTP_STATUS.key" endpoint.status
+        "$Tags.HTTP_URL.key" "${endpoint.resolve(address)}"
+        "$Tags.PEER_HOSTNAME.key" { it == "localhost" || it == "127.0.0.1" }
+        "$Tags.PEER_PORT.key" Integer
+        "$Tags.PEER_HOST_IPV4.key" { it == null || it == "127.0.0.1" } // Optional
+        "$Tags.HTTP_METHOD.key" method
+        "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_SERVER
+      }
+    }
+  }
+
+  static class TestApp extends Application<Configuration> {
+    @Override
+    void initialize(Bootstrap<Configuration> bootstrap) {
+    }
+
+    @Override
+    void run(Configuration configuration, Environment environment) {
+      environment.jersey().register(ServiceResource)
+    }
+  }
+
+  @Path("/")
+  static class ServiceResource {
+
+    @GET
+    @Path("success")
+    Response success() {
+      controller(SUCCESS) {
+        Response.status(SUCCESS.status).entity(SUCCESS.body).build()
+      }
+    }
+
+    @GET
+    @Path("redirect")
+    Response redirect() {
+      controller(REDIRECT) {
+        Response.status(REDIRECT.status).location(new URI(REDIRECT.body)).build()
+      }
+    }
+
+    @GET
+    @Path("error-status")
+    Response error() {
+      controller(ERROR) {
+        Response.status(ERROR.status).entity(ERROR.body).build()
+      }
+    }
+
+    @GET
+    @Path("exception")
+    Response exception() {
+      controller(EXCEPTION) {
+        throw new Exception(EXCEPTION.body)
+      }
+      return null
+    }
+  }
+}

dd-java-agent/instrumentation/jax-rs-annotations-2/src/test/groovy/JaxRsAnnotationsInstrumentationTest.groovy

@@ -0,0 +1,212 @@
+import datadog.trace.agent.test.AgentTestRunner
+import io.opentracing.tag.Tags
+import javax.ws.rs.DELETE
+import javax.ws.rs.GET
+import javax.ws.rs.HEAD
+import javax.ws.rs.OPTIONS
+import javax.ws.rs.POST
+import javax.ws.rs.PUT
+import javax.ws.rs.Path
+
+import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
+import static datadog.trace.instrumentation.jaxrs2.JaxRsAnnotationsDecorator.DECORATE
+
+class JaxRsAnnotationsInstrumentationTest extends AgentTestRunner {
+
+  def "instrumentation can be used as root span and resource is set to METHOD PATH"() {
+    setup:
+    new Jax() {
+      @POST
+      @Path("/a")
+      void call() {
+      }
+    }.call()
+
+    expect:
+    assertTraces(1) {
+      trace(0, 1) {
+        span(0) {
+          operationName "jax-rs.request"
+          resourceName "POST /a"
+          spanType "web"
+          tags {
+            "$Tags.COMPONENT.key" "jax-rs-controller"
+            defaultTags()
+          }
+        }
+      }
+    }
+  }
+
+  def "span named '#name' from annotations on class when is not root span"() {
+    setup:
+    def startingCacheSize = DECORATE.resourceNames.size()
+    runUnderTrace("test") {
+      obj.call()
+    }
+
+    expect:
+    assertTraces(1) {
+      trace(0, 2) {
+        span(0) {
+          operationName "test"
+          resourceName name
+          parent()
+          tags {
+            "$Tags.COMPONENT.key" "jax-rs"
+            defaultTags()
+          }
+        }
+        span(1) {
+          operationName "jax-rs.request"
+          resourceName "${className}.call"
+          spanType "web"
+          childOf span(0)
+          tags {
+            "$Tags.COMPONENT.key" "jax-rs-controller"
+            defaultTags()
+          }
+        }
+      }
+    }
+    DECORATE.resourceNames.size() == startingCacheSize + 1
+    DECORATE.resourceNames.get(obj.class).size() == 1
+
+    when: "multiple calls to the same method"
+    runUnderTrace("test") {
+      (1..10).each {
+        obj.call()
+      }
+    }
+    then: "doesn't increase the cache size"
+    DECORATE.resourceNames.size() == startingCacheSize + 1
+    DECORATE.resourceNames.get(obj.class).size() == 1
+
+    where:
+    name                        | obj
+    "/a"                        | new Jax() {
+      @Path("/a")
+      void call() {
+      }
+    }
+    "GET /b"                    | new Jax() {
+      @GET
+      @Path("/b")
+      void call() {
+      }
+    }
+    "POST /c"                   | new InterfaceWithPath() {
+      @POST
+      @Path("/c")
+      void call() {
+      }
+    }
+    "HEAD"                      | new InterfaceWithPath() {
+      @HEAD
+      void call() {
+      }
+    }
+    "POST /abstract/d"          | new AbstractClassWithPath() {
+      @POST
+      @Path("/d")
+      void call() {
+      }
+    }
+    "PUT /abstract"             | new AbstractClassWithPath() {
+      @PUT
+      void call() {
+      }
+    }
+    "OPTIONS /abstract/child/e" | new ChildClassWithPath() {
+      @OPTIONS
+      @Path("/e")
+      void call() {
+      }
+    }
+    "DELETE /abstract/child"    | new ChildClassWithPath() {
+      @DELETE
+      void call() {
+      }
+    }
+    "POST /abstract/child/call" | new ChildClassWithPath()
+
+    className = getName(obj.class)
+  }
+
+  def "no annotations has no effect"() {
+    setup:
+    runUnderTrace("test") {
+      obj.call()
+    }
+
+    expect:
+    assertTraces(1) {
+      trace(0, 1) {
+        span(0) {
+          operationName "test"
+          resourceName "test"
+          tags {
+            defaultTags()
+          }
+        }
+      }
+    }
+
+    where:
+    obj | _
+    new Jax() {
+      void call() {
+      }
+    }   | _
+    new InterfaceWithPath() {
+      void call() {
+      }
+    }   | _
+    new AbstractClassWithPath() {
+      void call() {
+      }
+    }   | _
+    new ChildClassWithPath() {
+      void call() {
+      }
+    }   | _
+  }
+
+  interface Jax {
+    void call()
+  }
+
+  @Path("/interface")
+  interface InterfaceWithPath extends Jax {
+    @GET
+    void call()
+  }
+
+  @Path("/abstract")
+  abstract class AbstractClassWithPath implements Jax {
+    @PUT
+    abstract void call()
+  }
+
+  @Path("child")
+  class ChildClassWithPath extends AbstractClassWithPath {
+    @Path("call")
+    @POST
+    void call() {
+    }
+  }
+
+  def getName(Class clazz) {
+    String className = clazz.getSimpleName()
+    if (className.isEmpty()) {
+      className = clazz.getName()
+      if (clazz.getPackage() != null) {
+        final String pkgName = clazz.getPackage().getName()
+        if (!pkgName.isEmpty()) {
+          className = clazz.getName().replace(pkgName, "").substring(1)
+        }
+      }
+    }
+    return className
+  }
+}

dd-java-agent/instrumentation/jax-rs-annotations-2/src/test/groovy/JettyTestInstrumentation.java

@@ -0,0 +1,27 @@
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.test.base.HttpServerTestAdvice;
+import datadog.trace.agent.tooling.Instrumenter;
+import net.bytebuddy.agent.builder.AgentBuilder;
+
+@AutoService(Instrumenter.class)
+public class JettyTestInstrumentation implements Instrumenter {
+
+  @Override
+  public AgentBuilder instrument(final AgentBuilder agentBuilder) {
+    return agentBuilder
+        // Jetty 8
+        .type(named("org.eclipse.jetty.server.AbstractHttpConnection"))
+        .transform(
+            new AgentBuilder.Transformer.ForAdvice()
+                .advice(
+                    named("headerComplete"),
+                    HttpServerTestAdvice.ServerEntryAdvice.class.getName()))
+        // Jetty 9
+        .type(named("org.eclipse.jetty.server.HttpChannel"))
+        .transform(
+            new AgentBuilder.Transformer.ForAdvice()
+                .advice(named("handle"), HttpServerTestAdvice.ServerEntryAdvice.class.getName()));
+  }
+}

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

@@ -449,7 +449,7 @@ abstract class HttpServerTest<SERVER, DECORATOR extends HttpServerDecorator> ext
 //          "$DDTags.HTTP_QUERY" uri.query
 //          "$DDTags.HTTP_FRAGMENT" { it == null || it == uri.fragment } // Optional
 //        }
-        "$Tags.PEER_HOSTNAME.key" "localhost"
+        "$Tags.PEER_HOSTNAME.key" { it == "localhost" || it == "127.0.0.1" }
         "$Tags.PEER_PORT.key" Integer
         "$Tags.PEER_HOST_IPV4.key" { it == null || it == "127.0.0.1" } // Optional
         "$Tags.HTTP_METHOD.key" method

settings.gradle

@@ -49,7 +49,8 @@ include ':dd-java-agent:instrumentation:hibernate:core-4.0'
 include ':dd-java-agent:instrumentation:hibernate:core-4.3'
 include ':dd-java-agent:instrumentation:http-url-connection'
 include ':dd-java-agent:instrumentation:hystrix-1.4'
-include ':dd-java-agent:instrumentation:jax-rs-annotations'
+include ':dd-java-agent:instrumentation:jax-rs-annotations-1'
+include ':dd-java-agent:instrumentation:jax-rs-annotations-2'
 include ':dd-java-agent:instrumentation:jax-rs-client-1.9'
 include ':dd-java-agent:instrumentation:jax-rs-client-2.0'
 include ':dd-java-agent:instrumentation:jax-rs-client-2.0:connection-error-handling-jersey'