Grizzly-http and grizzly-client instrumentation

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/InstrumentationContext.java

@@ -13,6 +13,8 @@ private InstrumentationContext() {}
    * <p>However, the implementation is actually provided by bytecode transformation for performance
    * reasons.
    *
+   * <p>This method must only be called within an Advice class.
+   *
    * @param keyClass The key class context is attached to.
    * @param contextClass The context class attached to the user class.
    * @param <K> key class

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/api/Pair.java

@@ -0,0 +1,32 @@
+package datadog.trace.bootstrap.instrumentation.api;
+
+public final class Pair<T, U> {
+
+  public static <T, U> Pair<T, U> of(T left, U right) {
+    return new Pair<>(left, right);
+  }
+
+  private final T left;
+  private final U right;
+
+  Pair(T left, U right) {
+    this.left = left;
+    this.right = right;
+  }
+
+  public T getLeft() {
+    return left;
+  }
+
+  public U getRight() {
+    return right;
+  }
+
+  public boolean hasLeft() {
+    return null != left;
+  }
+
+  public boolean hasRight() {
+    return null != right;
+  }
+}

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java

@@ -13,6 +13,7 @@
 @Slf4j
 public abstract class HttpServerDecorator<REQUEST, CONNECTION, RESPONSE> extends ServerDecorator {
   public static final String DD_SPAN_ATTRIBUTE = "datadog.span";
+  public static final String DD_RESPONSE_ATTRIBUTE = "datadog.response";
 
   // Source: https://www.regextester.com/22
   private static final Pattern VALID_IPV4_ADDRESS =

dd-java-agent/instrumentation/grizzly-2/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyHttpHandlerInstrumentation.java

@@ -3,6 +3,7 @@
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.bootstrap.instrumentation.api.DDSpanNames.GRIZZLY_REQUEST;
 import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE;
 import static datadog.trace.instrumentation.grizzly.GrizzlyDecorator.DECORATE;
 import static datadog.trace.instrumentation.grizzly.GrizzlyRequestExtractAdapter.GETTER;
@@ -71,7 +72,7 @@ public static AgentScope methodEnter(@Advice.Argument(0) final Request request)
       }
 
       final Context parentContext = propagate().extract(request, GETTER);
-      final AgentSpan span = startSpan("grizzly.request", parentContext);
+      final AgentSpan span = startSpan(GRIZZLY_REQUEST, parentContext);
       DECORATE.afterStart(span);
       DECORATE.onConnection(span, request);
       DECORATE.onRequest(span, request);

dd-java-agent/instrumentation/grizzly-2/src/test/groovy/GrizzlyTest.groovy

@@ -1,4 +1,5 @@
 import datadog.trace.agent.test.base.HttpServerTest
+import datadog.trace.bootstrap.instrumentation.api.DDSpanNames
 import datadog.trace.instrumentation.grizzly.GrizzlyDecorator
 import org.glassfish.grizzly.http.server.HttpServer
 import org.glassfish.jersey.grizzly2.httpserver.GrizzlyHttpServerFactory
@@ -43,7 +44,7 @@ class GrizzlyTest extends HttpServerTest<HttpServer> {
 
   @Override
   String expectedOperationName() {
-    return "grizzly.request"
+    return DDSpanNames.GRIZZLY_REQUEST
   }
 
   static class SimpleExceptionMapper implements ExceptionMapper<Throwable> {

dd-java-agent/instrumentation/grizzly-client-1.9/grizzly-client-1.9.gradle

@@ -0,0 +1,36 @@
+ext {
+  minJavaVersionForTests = JavaVersion.VERSION_1_8
+}
+
+muzzle {
+  pass {
+    group = "org.glassfish.grizzly"
+    module = "grizzly-http-client"
+    versions = "[1.9,1.16]"
+    assertInverse = true
+  }
+  pass {
+    group = "com.ning"
+    module = "async-http-client"
+    versions = "[1.9.0,)"
+    assertInverse = true
+  }
+}
+
+apply from: "${rootDir}/gradle/java.gradle"
+
+apply plugin: 'org.unbroken-dome.test-sets'
+
+testSets {
+  latestDepTest {
+    dirName = 'test'
+  }
+}
+
+dependencies {
+  compileOnly group: 'org.glassfish.grizzly', name: 'grizzly-http-client', version: '1.9'
+  // for some reason, the tests don't *load* until 1.12, but muzzles works as far back as 1.9
+  testCompile group: 'org.glassfish.grizzly', name: 'grizzly-http-client', version: '1.12'
+
+  latestDepTestCompile group: 'org.glassfish.grizzly', name: 'grizzly-http-client', version: '1.16'
+}

dd-java-agent/instrumentation/grizzly-client-1.9/src/main/java/datadog/trace/instrumentation/grizzly/client/ClientDecorator.java

@@ -0,0 +1,39 @@
+package datadog.trace.instrumentation.grizzly.client;
+
+import static datadog.trace.bootstrap.instrumentation.api.DDComponents.GRIZZLY_HTTP_ASYNC_CLIENT;
+
+import com.ning.http.client.Request;
+import com.ning.http.client.Response;
+import datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator;
+import java.net.URI;
+import java.net.URISyntaxException;
+
+public class ClientDecorator extends HttpClientDecorator<Request, Response> {
+
+  public static final ClientDecorator DECORATE = new ClientDecorator();
+
+  @Override
+  protected String[] instrumentationNames() {
+    return new String[] {"grizzly-client", "ning"};
+  }
+
+  @Override
+  protected String component() {
+    return GRIZZLY_HTTP_ASYNC_CLIENT;
+  }
+
+  @Override
+  protected String method(final Request request) {
+    return request.getMethod();
+  }
+
+  @Override
+  protected URI url(final Request request) throws URISyntaxException {
+    return request.getUri().toJavaNetURI();
+  }
+
+  @Override
+  protected Integer status(final Response response) {
+    return response.getStatusCode();
+  }
+}

dd-java-agent/instrumentation/grizzly-client-1.9/src/main/java/datadog/trace/instrumentation/grizzly/client/ClientRequestAdvice.java

@@ -0,0 +1,40 @@
+package datadog.trace.instrumentation.grizzly.client;
+
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.propagate;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.bootstrap.instrumentation.api.DDSpanNames.HTTP_REQUEST;
+import static datadog.trace.instrumentation.grizzly.client.ClientDecorator.DECORATE;
+import static datadog.trace.instrumentation.grizzly.client.InjectAdapter.SETTER;
+
+import com.ning.http.client.AsyncHandler;
+import com.ning.http.client.Request;
+import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.Pair;
+import net.bytebuddy.asm.Advice;
+
+public class ClientRequestAdvice {
+
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  public static AgentScope onEnter(
+      @Advice.Argument(0) final Request request,
+      @Advice.Argument(1) final AsyncHandler<?> handler) {
+    AgentSpan parentSpan = activeSpan();
+    AgentSpan span = startSpan(HTTP_REQUEST);
+    DECORATE.afterStart(span);
+    DECORATE.onRequest(span, request);
+    propagate().inject(span, request, SETTER);
+    InstrumentationContext.get(AsyncHandler.class, Pair.class)
+        .put(handler, Pair.of(parentSpan, span));
+    return activateSpan(span);
+  }
+
+  @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+  public static void onExit(@Advice.Enter final AgentScope scope) {
+    // span closed in ClientResponseAdvice, scope only created for profiler's benefit
+    scope.close();
+  }
+}

dd-java-agent/instrumentation/grizzly-client-1.9/src/main/java/datadog/trace/instrumentation/grizzly/client/ClientRequestInstrumentation.java

@@ -0,0 +1,52 @@
+package datadog.trace.instrumentation.grizzly.client;
+
+import static java.util.Collections.singletonMap;
+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.instrumentation.api.Pair;
+import java.util.Map;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(Instrumenter.class)
+public final class ClientRequestInstrumentation extends Instrumenter.Default {
+
+  public ClientRequestInstrumentation() {
+    super("grizzly-client", "ning");
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    return singletonMap("com.ning.http.client.AsyncHandler", Pair.class.getName());
+  }
+
+  @Override
+  public ElementMatcher<? super TypeDescription> typeMatcher() {
+    return named("com.ning.http.client.AsyncHttpClient");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {packageName + ".ClientDecorator", packageName + ".InjectAdapter"};
+  }
+
+  @Override
+  protected boolean defaultEnabled() {
+    return false;
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+        named("executeRequest")
+            .and(takesArgument(0, named("com.ning.http.client.Request")))
+            .and(takesArgument(1, named("com.ning.http.client.AsyncHandler")))
+            .and(isPublic()),
+        packageName + ".ClientRequestAdvice");
+  }
+}

dd-java-agent/instrumentation/grizzly-client-1.9/src/main/java/datadog/trace/instrumentation/grizzly/client/ClientResponseAdvice.java

@@ -0,0 +1,43 @@
+package datadog.trace.instrumentation.grizzly.client;
+
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.instrumentation.grizzly.client.ClientDecorator.DECORATE;
+
+import com.ning.http.client.AsyncCompletionHandler;
+import com.ning.http.client.AsyncHandler;
+import com.ning.http.client.Response;
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.Pair;
+import net.bytebuddy.asm.Advice;
+
+@SuppressWarnings({"unchecked", "rawtypes"})
+public class ClientResponseAdvice {
+
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  public static AgentScope onEnter(
+      @Advice.This final AsyncCompletionHandler<?> handler,
+      @Advice.Argument(0) final Response response) {
+    ContextStore<AsyncHandler, Pair> contextStore =
+        InstrumentationContext.get(AsyncHandler.class, Pair.class);
+    Pair<AgentSpan, AgentSpan> spanWithParent = contextStore.get(handler);
+    if (null != spanWithParent) {
+      contextStore.put(handler, null);
+    }
+    if (spanWithParent.hasRight()) {
+      DECORATE.onResponse(spanWithParent.getRight(), response);
+      DECORATE.beforeFinish(spanWithParent.getRight());
+      spanWithParent.getRight().finish();
+    }
+    return spanWithParent.hasLeft() ? activateSpan(spanWithParent.getLeft()) : null;
+  }
+
+  @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+  public static void onExit(@Advice.Enter final AgentScope scope) {
+    if (null != scope) {
+      scope.close();
+    }
+  }
+}

dd-java-agent/instrumentation/grizzly-client-1.9/src/main/java/datadog/trace/instrumentation/grizzly/client/ClientResponseInstrumentation.java

@@ -0,0 +1,58 @@
+package datadog.trace.instrumentation.grizzly.client;
+
+import static datadog.trace.agent.tooling.ClassLoaderMatcher.hasClassesNamed;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.hasSuperClass;
+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.instrumentation.api.Pair;
+import java.util.Map;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(Instrumenter.class)
+public final class ClientResponseInstrumentation extends Instrumenter.Default {
+
+  public ClientResponseInstrumentation() {
+    super("grizzly-client", "ning");
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    return singletonMap("com.ning.http.client.AsyncHandler", Pair.class.getName());
+  }
+
+  @Override
+  public ElementMatcher<ClassLoader> classLoaderMatcher() {
+    return hasClassesNamed("com.ning.http.client.AsyncCompletionHandler");
+  }
+
+  @Override
+  protected boolean defaultEnabled() {
+    return false;
+  }
+
+  @Override
+  public ElementMatcher<? super TypeDescription> typeMatcher() {
+    return hasSuperClass(named("com.ning.http.client.AsyncCompletionHandler"));
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {packageName + ".ClientDecorator"};
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+        named("onCompleted")
+            .and(takesArgument(0, named("com.ning.http.client.Response")))
+            .and(isPublic()),
+        packageName + ".ClientResponseAdvice");
+  }
+}

dd-java-agent/instrumentation/grizzly-client-1.9/src/main/java/datadog/trace/instrumentation/grizzly/client/InjectAdapter.java

@@ -0,0 +1,14 @@
+package datadog.trace.instrumentation.grizzly.client;
+
+import com.ning.http.client.Request;
+import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
+
+public class InjectAdapter implements AgentPropagation.Setter<Request> {
+
+  public static final InjectAdapter SETTER = new InjectAdapter();
+
+  @Override
+  public void set(final Request carrier, final String key, final String value) {
+    carrier.getHeaders().replaceWith(key, value);
+  }
+}