From 3a1d3315253ff1e21574f8db548aeefc65a45901 Mon Sep 17 00:00:00 2001 From: Laplie Anderson Date: Wed, 30 Oct 2019 19:22:10 -0400 Subject: [PATCH] WSClient instrumentation --- .../wsclient/AsyncHandlerWrapper.java | 80 +++++++++ .../wsclient/HeadersInjectAdapter.java | 14 ++ .../wsclient/WSClientDecorator.java | 46 +++++ .../wsclient/WSClientInstrumentation.java | 91 ++++++++++ .../src/test/groovy/WSClientTest.groovy | 74 ++++++++ .../wsclient-1/wsclient-1.gradle | 51 ++++++ .../wsclient/AsyncHandlerWrapper.java | 159 ++++++++++++++++++ .../wsclient/HeadersInjectAdapter.java | 14 ++ .../wsclient/WSClientDecorator.java | 46 +++++ .../wsclient/WSClientInstrumentation.java | 91 ++++++++++ .../src/test/groovy/WSClientTest.groovy | 74 ++++++++ .../wsclient-2/wsclient-2.gradle | 74 ++++++++ settings.gradle | 2 + 13 files changed, 816 insertions(+) create mode 100644 dd-java-agent/instrumentation/wsclient-1/src/main/java/datadog/trace/instrumentation/wsclient/AsyncHandlerWrapper.java create mode 100644 dd-java-agent/instrumentation/wsclient-1/src/main/java/datadog/trace/instrumentation/wsclient/HeadersInjectAdapter.java create mode 100644 dd-java-agent/instrumentation/wsclient-1/src/main/java/datadog/trace/instrumentation/wsclient/WSClientDecorator.java create mode 100644 dd-java-agent/instrumentation/wsclient-1/src/main/java/datadog/trace/instrumentation/wsclient/WSClientInstrumentation.java create mode 100644 dd-java-agent/instrumentation/wsclient-1/src/test/groovy/WSClientTest.groovy create mode 100644 dd-java-agent/instrumentation/wsclient-1/wsclient-1.gradle create mode 100644 dd-java-agent/instrumentation/wsclient-2/src/main/java/datadog/trace/instrumentation/wsclient/AsyncHandlerWrapper.java create mode 100644 dd-java-agent/instrumentation/wsclient-2/src/main/java/datadog/trace/instrumentation/wsclient/HeadersInjectAdapter.java create mode 100644 dd-java-agent/instrumentation/wsclient-2/src/main/java/datadog/trace/instrumentation/wsclient/WSClientDecorator.java create mode 100644 dd-java-agent/instrumentation/wsclient-2/src/main/java/datadog/trace/instrumentation/wsclient/WSClientInstrumentation.java create mode 100644 dd-java-agent/instrumentation/wsclient-2/src/test/groovy/WSClientTest.groovy create mode 100644 dd-java-agent/instrumentation/wsclient-2/wsclient-2.gradle diff --git a/dd-java-agent/instrumentation/wsclient-1/src/main/java/datadog/trace/instrumentation/wsclient/AsyncHandlerWrapper.java b/dd-java-agent/instrumentation/wsclient-1/src/main/java/datadog/trace/instrumentation/wsclient/AsyncHandlerWrapper.java new file mode 100644 index 0000000000..014aa0b60f --- /dev/null +++ b/dd-java-agent/instrumentation/wsclient-1/src/main/java/datadog/trace/instrumentation/wsclient/AsyncHandlerWrapper.java @@ -0,0 +1,80 @@ +package datadog.trace.instrumentation.wsclient; + +import static datadog.trace.instrumentation.api.AgentTracer.propagate; +import static datadog.trace.instrumentation.wsclient.WSClientDecorator.DECORATE; + +import datadog.trace.context.TraceScope; +import datadog.trace.instrumentation.api.AgentSpan; +import play.shaded.ahc.org.asynchttpclient.AsyncHandler; +import play.shaded.ahc.org.asynchttpclient.HttpResponseBodyPart; +import play.shaded.ahc.org.asynchttpclient.HttpResponseHeaders; +import play.shaded.ahc.org.asynchttpclient.HttpResponseStatus; +import play.shaded.ahc.org.asynchttpclient.Response; + +public class AsyncHandlerWrapper implements AsyncHandler { + private final AsyncHandler delegate; + private final AgentSpan span; + private final TraceScope.Continuation continuation; + + private final Response.ResponseBuilder builder = new Response.ResponseBuilder(); + + public AsyncHandlerWrapper(final AsyncHandler delegate, final AgentSpan span) { + this.delegate = delegate; + this.span = span; + continuation = propagate().capture(); + } + + @Override + public State onBodyPartReceived(final HttpResponseBodyPart content) throws Exception { + builder.accumulate(content); + return delegate.onBodyPartReceived(content); + } + + @Override + public State onStatusReceived(final HttpResponseStatus status) throws Exception { + builder.reset(); + builder.accumulate(status); + return delegate.onStatusReceived(status); + } + + @Override + public State onHeadersReceived(final HttpResponseHeaders httpHeaders) throws Exception { + builder.accumulate(httpHeaders); + return delegate.onHeadersReceived(httpHeaders); + } + + @Override + public Object onCompleted() throws Exception { + final Response response = builder.build(); + if (response != null) { + DECORATE.onResponse(span, response); + } + DECORATE.beforeFinish(span); + span.finish(); + + if (continuation != null) { + try (final TraceScope scope = continuation.activate()) { + scope.setAsyncPropagation(true); + return delegate.onCompleted(); + } + } else { + return delegate.onCompleted(); + } + } + + @Override + public void onThrowable(final Throwable throwable) { + DECORATE.onError(span, throwable); + DECORATE.beforeFinish(span); + span.finish(); + + if (continuation != null) { + try (final TraceScope scope = continuation.activate()) { + scope.setAsyncPropagation(true); + delegate.onThrowable(throwable); + } + } else { + delegate.onThrowable(throwable); + } + } +} diff --git a/dd-java-agent/instrumentation/wsclient-1/src/main/java/datadog/trace/instrumentation/wsclient/HeadersInjectAdapter.java b/dd-java-agent/instrumentation/wsclient-1/src/main/java/datadog/trace/instrumentation/wsclient/HeadersInjectAdapter.java new file mode 100644 index 0000000000..50258719df --- /dev/null +++ b/dd-java-agent/instrumentation/wsclient-1/src/main/java/datadog/trace/instrumentation/wsclient/HeadersInjectAdapter.java @@ -0,0 +1,14 @@ +package datadog.trace.instrumentation.wsclient; + +import datadog.trace.instrumentation.api.AgentPropagation; +import play.shaded.ahc.org.asynchttpclient.Request; + +public class HeadersInjectAdapter implements AgentPropagation.Setter { + + public static final HeadersInjectAdapter SETTER = new HeadersInjectAdapter(); + + @Override + public void set(final Request carrier, final String key, final String value) { + carrier.getHeaders().add(key, value); + } +} diff --git a/dd-java-agent/instrumentation/wsclient-1/src/main/java/datadog/trace/instrumentation/wsclient/WSClientDecorator.java b/dd-java-agent/instrumentation/wsclient-1/src/main/java/datadog/trace/instrumentation/wsclient/WSClientDecorator.java new file mode 100644 index 0000000000..33714f849f --- /dev/null +++ b/dd-java-agent/instrumentation/wsclient-1/src/main/java/datadog/trace/instrumentation/wsclient/WSClientDecorator.java @@ -0,0 +1,46 @@ +package datadog.trace.instrumentation.wsclient; + +import datadog.trace.agent.decorator.HttpClientDecorator; +import java.net.URI; +import java.net.URISyntaxException; +import play.shaded.ahc.org.asynchttpclient.Request; +import play.shaded.ahc.org.asynchttpclient.Response; + +public class WSClientDecorator extends HttpClientDecorator { + public static final WSClientDecorator DECORATE = new WSClientDecorator(); + + @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 String hostname(final Request request) { + return request.getUri().getHost(); + } + + @Override + protected Integer port(final Request request) { + return request.getUri().getPort(); + } + + @Override + protected Integer status(final Response response) { + return response.getStatusCode(); + } + + @Override + protected String[] instrumentationNames() { + return new String[] {"wsclient"}; + } + + @Override + protected String component() { + return "wsclient"; + } +} diff --git a/dd-java-agent/instrumentation/wsclient-1/src/main/java/datadog/trace/instrumentation/wsclient/WSClientInstrumentation.java b/dd-java-agent/instrumentation/wsclient-1/src/main/java/datadog/trace/instrumentation/wsclient/WSClientInstrumentation.java new file mode 100644 index 0000000000..9f9caa57b8 --- /dev/null +++ b/dd-java-agent/instrumentation/wsclient-1/src/main/java/datadog/trace/instrumentation/wsclient/WSClientInstrumentation.java @@ -0,0 +1,91 @@ +package datadog.trace.instrumentation.wsclient; + +import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType; +import static datadog.trace.instrumentation.api.AgentTracer.propagate; +import static datadog.trace.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.instrumentation.wsclient.HeadersInjectAdapter.SETTER; +import static datadog.trace.instrumentation.wsclient.WSClientDecorator.DECORATE; +import static java.util.Collections.singletonMap; +import static net.bytebuddy.matcher.ElementMatchers.isMethod; +import static net.bytebuddy.matcher.ElementMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.not; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; + +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.instrumentation.api.AgentSpan; +import java.util.Map; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.method.MethodDescription; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; +import play.shaded.ahc.org.asynchttpclient.AsyncHandler; +import play.shaded.ahc.org.asynchttpclient.Request; + +@AutoService(Instrumenter.class) +public class WSClientInstrumentation extends Instrumenter.Default { + public WSClientInstrumentation() { + super("wsclient"); + } + + @Override + public ElementMatcher typeMatcher() { + // CachingAsyncHttpClient rejects overrides to AsyncHandler + // It also delegates to another AsyncHttpClient + return safeHasSuperType(named("play.shaded.ahc.org.asynchttpclient.AsyncHttpClient")) + .and(not(named("play.api.libs.ws.ahc.cache.CachingAsyncHttpClient"))); + } + + @Override + public Map, String> transformers() { + return singletonMap( + isMethod() + .and(named("execute")) + .and(takesArguments(2)) + .and(takesArgument(0, named("play.shaded.ahc.org.asynchttpclient.Request"))) + .and(takesArgument(1, named("play.shaded.ahc.org.asynchttpclient.AsyncHandler"))), + WSClientAdvice.class.getName()); + } + + @Override + public String[] helperClassNames() { + return new String[] { + "datadog.trace.agent.decorator.BaseDecorator", + "datadog.trace.agent.decorator.ClientDecorator", + "datadog.trace.agent.decorator.HttpClientDecorator", + packageName + ".WSClientDecorator", + packageName + ".HeadersInjectAdapter", + packageName + ".AsyncHandlerWrapper" + }; + } + + public static class WSClientAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static AgentSpan methodEnter( + @Advice.Argument(0) final Request request, + @Advice.Argument(value = 1, readOnly = false) AsyncHandler asyncHandler) { + + final AgentSpan span = startSpan("wsclient.request"); + + DECORATE.afterStart(span); + DECORATE.onRequest(span, request); + propagate().inject(span, request, SETTER); + + asyncHandler = new AsyncHandlerWrapper(asyncHandler, span); + + return span; + } + + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) + public static void methodExit( + @Advice.Enter final AgentSpan clientSpan, @Advice.Thrown final Throwable throwable) { + + if (throwable != null) { + DECORATE.onError(clientSpan, throwable); + DECORATE.beforeFinish(clientSpan); + clientSpan.finish(); + } + } + } +} diff --git a/dd-java-agent/instrumentation/wsclient-1/src/test/groovy/WSClientTest.groovy b/dd-java-agent/instrumentation/wsclient-1/src/test/groovy/WSClientTest.groovy new file mode 100644 index 0000000000..8a83473db4 --- /dev/null +++ b/dd-java-agent/instrumentation/wsclient-1/src/test/groovy/WSClientTest.groovy @@ -0,0 +1,74 @@ +import akka.actor.ActorSystem +import akka.stream.ActorMaterializer +import akka.stream.ActorMaterializerSettings +import datadog.trace.agent.test.base.HttpClientTest +import datadog.trace.instrumentation.wsclient.WSClientDecorator +import play.libs.ws.StandaloneWSClient +import play.libs.ws.StandaloneWSRequest +import play.libs.ws.StandaloneWSResponse +import play.libs.ws.ahc.StandaloneAhcWSClient +import play.shaded.ahc.org.asynchttpclient.AsyncHttpClient +import play.shaded.ahc.org.asynchttpclient.AsyncHttpClientConfig +import play.shaded.ahc.org.asynchttpclient.DefaultAsyncHttpClient +import play.shaded.ahc.org.asynchttpclient.DefaultAsyncHttpClientConfig +import spock.lang.Shared + +import java.util.concurrent.TimeUnit + +class WSClientTest extends HttpClientTest { + @Shared + ActorSystem system + + @Shared + StandaloneWSClient wsClient + + def setupSpec() { + String name = "wsclient" + system = ActorSystem.create(name) + ActorMaterializerSettings settings = ActorMaterializerSettings.create(system) + ActorMaterializer materializer = ActorMaterializer.create(settings, system, name) + + AsyncHttpClientConfig asyncHttpClientConfig = + new DefaultAsyncHttpClientConfig.Builder() + .setMaxRequestRetry(0) + .setShutdownQuietPeriod(0) + .setShutdownTimeout(0) + .build() + + AsyncHttpClient asyncHttpClient = new DefaultAsyncHttpClient(asyncHttpClientConfig) + + wsClient = new StandaloneAhcWSClient(asyncHttpClient, materializer) + } + + def cleanupSpec() { + wsClient?.close() + system?.terminate() + } + + @Override + int doRequest(String method, URI uri, Map headers, Closure callback) { + StandaloneWSRequest wsRequest = wsClient.url(uri.toURL().toString()).setFollowRedirects(true) + + headers.entrySet().each { entry -> wsRequest.addHeader(entry.getKey(), entry.getValue()) } + StandaloneWSResponse wsResponse = wsRequest.execute(method) + .whenComplete({ response, throwable -> + callback?.call() + }).toCompletableFuture().get(5, TimeUnit.SECONDS) + + return wsResponse.getStatus() + } + + @Override + WSClientDecorator decorator() { + return WSClientDecorator.DECORATE + } + + String expectedOperationName() { + return "wsclient.request" + } + + @Override + boolean testCircularRedirects() { + return false + } +} diff --git a/dd-java-agent/instrumentation/wsclient-1/wsclient-1.gradle b/dd-java-agent/instrumentation/wsclient-1/wsclient-1.gradle new file mode 100644 index 0000000000..7bcbbe5619 --- /dev/null +++ b/dd-java-agent/instrumentation/wsclient-1/wsclient-1.gradle @@ -0,0 +1,51 @@ +// Set properties before any plugins get loaded +ext { + minJavaVersionForTests = JavaVersion.VERSION_1_8 +} + +apply from: "${rootDir}/gradle/java.gradle" + +apply plugin: 'org.unbroken-dome.test-sets' + +testSets { + latestDepTest { + dirName = 'test' + } +} + +muzzle { + pass { + group = 'com.typesafe.play' + module = 'play-ahc-ws-standalone_2.11' + versions = '[1.0.0,2.0.0)' + assertInverse = true + } + pass { + group = 'com.typesafe.play' + module = 'play-ahc-ws-standalone_2.12' + versions = '[1.0.0,2.0.0)' + assertInverse = true + } + fail { + group = 'com.typesafe.play' + module = 'play-ahc-ws-standalone_2.13' + versions = '[,]' + } +} + +def scalaVersion = '2.12' + +dependencies { + compileOnly group: 'com.typesafe.play', name: "play-ahc-ws-standalone_$scalaVersion", version: '1.0.2' + + testCompile project(':dd-java-agent:instrumentation:java-concurrent') + + // These are to ensure cross compatibility + testCompile project(':dd-java-agent:instrumentation:netty-4.0') + testCompile project(':dd-java-agent:instrumentation:netty-4.1') + testCompile project(':dd-java-agent:instrumentation:akka-http-10.0') + + testCompile group: 'com.typesafe.play', name: "play-ahc-ws-standalone_$scalaVersion", version: '1.0.2' + + latestDepTestCompile group: 'com.typesafe.play', name: "play-ahc-ws_$scalaVersion", version: '1.+' +} diff --git a/dd-java-agent/instrumentation/wsclient-2/src/main/java/datadog/trace/instrumentation/wsclient/AsyncHandlerWrapper.java b/dd-java-agent/instrumentation/wsclient-2/src/main/java/datadog/trace/instrumentation/wsclient/AsyncHandlerWrapper.java new file mode 100644 index 0000000000..b9d862d103 --- /dev/null +++ b/dd-java-agent/instrumentation/wsclient-2/src/main/java/datadog/trace/instrumentation/wsclient/AsyncHandlerWrapper.java @@ -0,0 +1,159 @@ +package datadog.trace.instrumentation.wsclient; + +import static datadog.trace.instrumentation.api.AgentTracer.propagate; +import static datadog.trace.instrumentation.wsclient.WSClientDecorator.DECORATE; + +import datadog.trace.context.TraceScope; +import datadog.trace.instrumentation.api.AgentSpan; +import java.net.InetSocketAddress; +import java.util.List; +import play.shaded.ahc.io.netty.channel.Channel; +import play.shaded.ahc.io.netty.handler.codec.http.HttpHeaders; +import play.shaded.ahc.org.asynchttpclient.AsyncHandler; +import play.shaded.ahc.org.asynchttpclient.HttpResponseBodyPart; +import play.shaded.ahc.org.asynchttpclient.HttpResponseStatus; +import play.shaded.ahc.org.asynchttpclient.Response; +import play.shaded.ahc.org.asynchttpclient.netty.request.NettyRequest; + +public class AsyncHandlerWrapper implements AsyncHandler { + private final AsyncHandler delegate; + private final AgentSpan span; + private final TraceScope.Continuation continuation; + + private final Response.ResponseBuilder builder = new Response.ResponseBuilder(); + + public AsyncHandlerWrapper(final AsyncHandler delegate, final AgentSpan span) { + this.delegate = delegate; + this.span = span; + continuation = propagate().capture(); + } + + @Override + public State onBodyPartReceived(final HttpResponseBodyPart content) throws Exception { + builder.accumulate(content); + return delegate.onBodyPartReceived(content); + } + + @Override + public State onStatusReceived(final HttpResponseStatus status) throws Exception { + builder.reset(); + builder.accumulate(status); + return delegate.onStatusReceived(status); + } + + @Override + public State onHeadersReceived(final HttpHeaders httpHeaders) throws Exception { + builder.accumulate(httpHeaders); + return delegate.onHeadersReceived(httpHeaders); + } + + @Override + public Object onCompleted() throws Exception { + final Response response = builder.build(); + if (response != null) { + DECORATE.onResponse(span, response); + } + DECORATE.beforeFinish(span); + span.finish(); + + if (continuation != null) { + try (final TraceScope scope = continuation.activate()) { + scope.setAsyncPropagation(true); + return delegate.onCompleted(); + } + } else { + return delegate.onCompleted(); + } + } + + @Override + public void onThrowable(final Throwable throwable) { + DECORATE.onError(span, throwable); + DECORATE.beforeFinish(span); + span.finish(); + + if (continuation != null) { + try (final TraceScope scope = continuation.activate()) { + scope.setAsyncPropagation(true); + delegate.onThrowable(throwable); + } + } else { + delegate.onThrowable(throwable); + } + } + + @Override + public State onTrailingHeadersReceived(final HttpHeaders headers) throws Exception { + return delegate.onTrailingHeadersReceived(headers); + } + + @Override + public void onHostnameResolutionAttempt(final String name) { + delegate.onHostnameResolutionAttempt(name); + } + + @Override + public void onHostnameResolutionSuccess(final String name, final List list) { + delegate.onHostnameResolutionSuccess(name, list); + } + + @Override + public void onHostnameResolutionFailure(final String name, final Throwable cause) { + delegate.onHostnameResolutionFailure(name, cause); + } + + @Override + public void onTcpConnectAttempt(final InetSocketAddress remoteAddress) { + delegate.onTcpConnectAttempt(remoteAddress); + } + + @Override + public void onTcpConnectSuccess(final InetSocketAddress remoteAddress, final Channel connection) { + delegate.onTcpConnectSuccess(remoteAddress, connection); + } + + @Override + public void onTcpConnectFailure(final InetSocketAddress remoteAddress, final Throwable cause) { + delegate.onTcpConnectFailure(remoteAddress, cause); + } + + @Override + public void onTlsHandshakeAttempt() { + delegate.onTlsHandshakeAttempt(); + } + + @Override + public void onTlsHandshakeSuccess() { + delegate.onTlsHandshakeSuccess(); + } + + @Override + public void onTlsHandshakeFailure(final Throwable cause) { + delegate.onTlsHandshakeFailure(cause); + } + + @Override + public void onConnectionPoolAttempt() { + delegate.onConnectionPoolAttempt(); + } + + @Override + public void onConnectionPooled(final Channel connection) { + delegate.onConnectionPooled(connection); + } + + @Override + public void onConnectionOffer(final Channel connection) { + delegate.onConnectionOffer(connection); + } + + @Override + public void onRequestSend(final NettyRequest request) { + delegate.onRequestSend(request); + } + + @Override + public void onRetry() { + delegate.onRetry(); + } +} diff --git a/dd-java-agent/instrumentation/wsclient-2/src/main/java/datadog/trace/instrumentation/wsclient/HeadersInjectAdapter.java b/dd-java-agent/instrumentation/wsclient-2/src/main/java/datadog/trace/instrumentation/wsclient/HeadersInjectAdapter.java new file mode 100644 index 0000000000..50258719df --- /dev/null +++ b/dd-java-agent/instrumentation/wsclient-2/src/main/java/datadog/trace/instrumentation/wsclient/HeadersInjectAdapter.java @@ -0,0 +1,14 @@ +package datadog.trace.instrumentation.wsclient; + +import datadog.trace.instrumentation.api.AgentPropagation; +import play.shaded.ahc.org.asynchttpclient.Request; + +public class HeadersInjectAdapter implements AgentPropagation.Setter { + + public static final HeadersInjectAdapter SETTER = new HeadersInjectAdapter(); + + @Override + public void set(final Request carrier, final String key, final String value) { + carrier.getHeaders().add(key, value); + } +} diff --git a/dd-java-agent/instrumentation/wsclient-2/src/main/java/datadog/trace/instrumentation/wsclient/WSClientDecorator.java b/dd-java-agent/instrumentation/wsclient-2/src/main/java/datadog/trace/instrumentation/wsclient/WSClientDecorator.java new file mode 100644 index 0000000000..33714f849f --- /dev/null +++ b/dd-java-agent/instrumentation/wsclient-2/src/main/java/datadog/trace/instrumentation/wsclient/WSClientDecorator.java @@ -0,0 +1,46 @@ +package datadog.trace.instrumentation.wsclient; + +import datadog.trace.agent.decorator.HttpClientDecorator; +import java.net.URI; +import java.net.URISyntaxException; +import play.shaded.ahc.org.asynchttpclient.Request; +import play.shaded.ahc.org.asynchttpclient.Response; + +public class WSClientDecorator extends HttpClientDecorator { + public static final WSClientDecorator DECORATE = new WSClientDecorator(); + + @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 String hostname(final Request request) { + return request.getUri().getHost(); + } + + @Override + protected Integer port(final Request request) { + return request.getUri().getPort(); + } + + @Override + protected Integer status(final Response response) { + return response.getStatusCode(); + } + + @Override + protected String[] instrumentationNames() { + return new String[] {"wsclient"}; + } + + @Override + protected String component() { + return "wsclient"; + } +} diff --git a/dd-java-agent/instrumentation/wsclient-2/src/main/java/datadog/trace/instrumentation/wsclient/WSClientInstrumentation.java b/dd-java-agent/instrumentation/wsclient-2/src/main/java/datadog/trace/instrumentation/wsclient/WSClientInstrumentation.java new file mode 100644 index 0000000000..9f9caa57b8 --- /dev/null +++ b/dd-java-agent/instrumentation/wsclient-2/src/main/java/datadog/trace/instrumentation/wsclient/WSClientInstrumentation.java @@ -0,0 +1,91 @@ +package datadog.trace.instrumentation.wsclient; + +import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType; +import static datadog.trace.instrumentation.api.AgentTracer.propagate; +import static datadog.trace.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.instrumentation.wsclient.HeadersInjectAdapter.SETTER; +import static datadog.trace.instrumentation.wsclient.WSClientDecorator.DECORATE; +import static java.util.Collections.singletonMap; +import static net.bytebuddy.matcher.ElementMatchers.isMethod; +import static net.bytebuddy.matcher.ElementMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.not; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; + +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.instrumentation.api.AgentSpan; +import java.util.Map; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.method.MethodDescription; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; +import play.shaded.ahc.org.asynchttpclient.AsyncHandler; +import play.shaded.ahc.org.asynchttpclient.Request; + +@AutoService(Instrumenter.class) +public class WSClientInstrumentation extends Instrumenter.Default { + public WSClientInstrumentation() { + super("wsclient"); + } + + @Override + public ElementMatcher typeMatcher() { + // CachingAsyncHttpClient rejects overrides to AsyncHandler + // It also delegates to another AsyncHttpClient + return safeHasSuperType(named("play.shaded.ahc.org.asynchttpclient.AsyncHttpClient")) + .and(not(named("play.api.libs.ws.ahc.cache.CachingAsyncHttpClient"))); + } + + @Override + public Map, String> transformers() { + return singletonMap( + isMethod() + .and(named("execute")) + .and(takesArguments(2)) + .and(takesArgument(0, named("play.shaded.ahc.org.asynchttpclient.Request"))) + .and(takesArgument(1, named("play.shaded.ahc.org.asynchttpclient.AsyncHandler"))), + WSClientAdvice.class.getName()); + } + + @Override + public String[] helperClassNames() { + return new String[] { + "datadog.trace.agent.decorator.BaseDecorator", + "datadog.trace.agent.decorator.ClientDecorator", + "datadog.trace.agent.decorator.HttpClientDecorator", + packageName + ".WSClientDecorator", + packageName + ".HeadersInjectAdapter", + packageName + ".AsyncHandlerWrapper" + }; + } + + public static class WSClientAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static AgentSpan methodEnter( + @Advice.Argument(0) final Request request, + @Advice.Argument(value = 1, readOnly = false) AsyncHandler asyncHandler) { + + final AgentSpan span = startSpan("wsclient.request"); + + DECORATE.afterStart(span); + DECORATE.onRequest(span, request); + propagate().inject(span, request, SETTER); + + asyncHandler = new AsyncHandlerWrapper(asyncHandler, span); + + return span; + } + + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) + public static void methodExit( + @Advice.Enter final AgentSpan clientSpan, @Advice.Thrown final Throwable throwable) { + + if (throwable != null) { + DECORATE.onError(clientSpan, throwable); + DECORATE.beforeFinish(clientSpan); + clientSpan.finish(); + } + } + } +} diff --git a/dd-java-agent/instrumentation/wsclient-2/src/test/groovy/WSClientTest.groovy b/dd-java-agent/instrumentation/wsclient-2/src/test/groovy/WSClientTest.groovy new file mode 100644 index 0000000000..8a83473db4 --- /dev/null +++ b/dd-java-agent/instrumentation/wsclient-2/src/test/groovy/WSClientTest.groovy @@ -0,0 +1,74 @@ +import akka.actor.ActorSystem +import akka.stream.ActorMaterializer +import akka.stream.ActorMaterializerSettings +import datadog.trace.agent.test.base.HttpClientTest +import datadog.trace.instrumentation.wsclient.WSClientDecorator +import play.libs.ws.StandaloneWSClient +import play.libs.ws.StandaloneWSRequest +import play.libs.ws.StandaloneWSResponse +import play.libs.ws.ahc.StandaloneAhcWSClient +import play.shaded.ahc.org.asynchttpclient.AsyncHttpClient +import play.shaded.ahc.org.asynchttpclient.AsyncHttpClientConfig +import play.shaded.ahc.org.asynchttpclient.DefaultAsyncHttpClient +import play.shaded.ahc.org.asynchttpclient.DefaultAsyncHttpClientConfig +import spock.lang.Shared + +import java.util.concurrent.TimeUnit + +class WSClientTest extends HttpClientTest { + @Shared + ActorSystem system + + @Shared + StandaloneWSClient wsClient + + def setupSpec() { + String name = "wsclient" + system = ActorSystem.create(name) + ActorMaterializerSettings settings = ActorMaterializerSettings.create(system) + ActorMaterializer materializer = ActorMaterializer.create(settings, system, name) + + AsyncHttpClientConfig asyncHttpClientConfig = + new DefaultAsyncHttpClientConfig.Builder() + .setMaxRequestRetry(0) + .setShutdownQuietPeriod(0) + .setShutdownTimeout(0) + .build() + + AsyncHttpClient asyncHttpClient = new DefaultAsyncHttpClient(asyncHttpClientConfig) + + wsClient = new StandaloneAhcWSClient(asyncHttpClient, materializer) + } + + def cleanupSpec() { + wsClient?.close() + system?.terminate() + } + + @Override + int doRequest(String method, URI uri, Map headers, Closure callback) { + StandaloneWSRequest wsRequest = wsClient.url(uri.toURL().toString()).setFollowRedirects(true) + + headers.entrySet().each { entry -> wsRequest.addHeader(entry.getKey(), entry.getValue()) } + StandaloneWSResponse wsResponse = wsRequest.execute(method) + .whenComplete({ response, throwable -> + callback?.call() + }).toCompletableFuture().get(5, TimeUnit.SECONDS) + + return wsResponse.getStatus() + } + + @Override + WSClientDecorator decorator() { + return WSClientDecorator.DECORATE + } + + String expectedOperationName() { + return "wsclient.request" + } + + @Override + boolean testCircularRedirects() { + return false + } +} diff --git a/dd-java-agent/instrumentation/wsclient-2/wsclient-2.gradle b/dd-java-agent/instrumentation/wsclient-2/wsclient-2.gradle new file mode 100644 index 0000000000..2cfa7eed17 --- /dev/null +++ b/dd-java-agent/instrumentation/wsclient-2/wsclient-2.gradle @@ -0,0 +1,74 @@ +// Set properties before any plugins get loaded +ext { + minJavaVersionForTests = JavaVersion.VERSION_1_8 +} + +apply from: "${rootDir}/gradle/java.gradle" + +apply plugin: 'org.unbroken-dome.test-sets' + +testSets { + latestDepTest { + dirName = 'test' + } +} + +muzzle { + // 2.0.5 was a bad release + + fail { + group = 'com.typesafe.play' + module = 'play-ahc-ws-standalone_2.11' + versions = '[,2.0.0)' + } + pass { + group = 'com.typesafe.play' + module = 'play-ahc-ws-standalone_2.11' + versions = '[2.0.0,2.0.4]' + } + pass { + group = 'com.typesafe.play' + module = 'play-ahc-ws-standalone_2.11' + versions = '[2.0.6,]' + } + + fail { + group = 'com.typesafe.play' + module = 'play-ahc-ws-standalone_2.12' + versions = '[,2.0.0)' + } + pass { + group = 'com.typesafe.play' + module = 'play-ahc-ws-standalone_2.12' + versions = '[2.0.0,2.0.4]' + } + pass { + group = 'com.typesafe.play' + module = 'play-ahc-ws-standalone_2.12' + versions = '[2.0.6,]' + } + + // No Scala 2.13 versions below 2.0.6 exist + pass { + group = 'com.typesafe.play' + module = 'play-ahc-ws-standalone_2.13' + versions = '[2.0.6,]' + } +} + +def scalaVersion = '2.12' + +dependencies { + compileOnly group: 'com.typesafe.play', name: "play-ahc-ws-standalone_$scalaVersion", version: '2.0.0' + + testCompile project(':dd-java-agent:instrumentation:java-concurrent') + + // These are to ensure cross compatibility + testCompile project(':dd-java-agent:instrumentation:netty-4.0') + testCompile project(':dd-java-agent:instrumentation:netty-4.1') + testCompile project(':dd-java-agent:instrumentation:akka-http-10.0') + + testCompile group: 'com.typesafe.play', name: "play-ahc-ws-standalone_$scalaVersion", version: '2.0.0' + + latestDepTestCompile group: 'com.typesafe.play', name: "play-ahc-ws_$scalaVersion", version: '+' +} diff --git a/settings.gradle b/settings.gradle index 8f66ec0426..c221c7847d 100644 --- a/settings.gradle +++ b/settings.gradle @@ -99,6 +99,8 @@ include ':dd-java-agent:instrumentation:tomcat-classloading' include ':dd-java-agent:instrumentation:trace-annotation' include ':dd-java-agent:instrumentation:twilio' include ':dd-java-agent:instrumentation:vertx' +include ':dd-java-agent:instrumentation:wsclient-1' +include ':dd-java-agent:instrumentation:wsclient-2' // benchmark include ':dd-java-agent:benchmark'