From 82ee01cadf61cc8f72022e8a24054ca4904e88cc Mon Sep 17 00:00:00 2001 From: Laplie Anderson Date: Fri, 12 Jul 2019 15:41:14 -0400 Subject: [PATCH] Implement instrumentation for async requests --- .../GoogleHttpClientInstrumentation.java | 113 ++++++++++++++---- .../googlehttpclient/RequestState.java | 10 ++ .../AbstractGoogleHttpClientTest.groovy | 69 +++++++++++ .../groovy/GoogleHttpClientAsyncTest.groovy | 9 ++ .../test/groovy/GoogleHttpClientTest.groovy | 64 +--------- 5 files changed, 182 insertions(+), 83 deletions(-) create mode 100644 dd-java-agent/instrumentation/google-http-client/src/main/java/datadog/trace/instrumentation/googlehttpclient/RequestState.java create mode 100644 dd-java-agent/instrumentation/google-http-client/src/test/groovy/AbstractGoogleHttpClientTest.groovy create mode 100644 dd-java-agent/instrumentation/google-http-client/src/test/groovy/GoogleHttpClientAsyncTest.groovy diff --git a/dd-java-agent/instrumentation/google-http-client/src/main/java/datadog/trace/instrumentation/googlehttpclient/GoogleHttpClientInstrumentation.java b/dd-java-agent/instrumentation/google-http-client/src/main/java/datadog/trace/instrumentation/googlehttpclient/GoogleHttpClientInstrumentation.java index dba8cae9da..d587ab0aba 100644 --- a/dd-java-agent/instrumentation/google-http-client/src/main/java/datadog/trace/instrumentation/googlehttpclient/GoogleHttpClientInstrumentation.java +++ b/dd-java-agent/instrumentation/google-http-client/src/main/java/datadog/trace/instrumentation/googlehttpclient/GoogleHttpClientInstrumentation.java @@ -5,13 +5,15 @@ import static java.util.Collections.singletonMap; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.isPublic; import static net.bytebuddy.matcher.ElementMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.api.client.http.HttpRequest; import com.google.api.client.http.HttpResponse; import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; -import datadog.trace.bootstrap.CallDepthThreadLocalMap; +import datadog.trace.bootstrap.ContextStore; +import datadog.trace.bootstrap.InstrumentationContext; import io.opentracing.Scope; import io.opentracing.Span; import io.opentracing.log.Fields; @@ -19,6 +21,8 @@ import io.opentracing.propagation.Format; import io.opentracing.propagation.TextMap; import io.opentracing.tag.Tags; import io.opentracing.util.GlobalTracer; +import java.util.Collections; +import java.util.HashMap; import java.util.Iterator; import java.util.Map; import net.bytebuddy.asm.Advice; @@ -38,6 +42,12 @@ public class GoogleHttpClientInstrumentation extends Instrumenter.Default { return named("com.google.api.client.http.HttpRequest"); } + @Override + public Map contextStore() { + return Collections.singletonMap( + "com.google.api.client.http.HttpRequest", RequestState.class.getName()); + } + @Override public String[] helperClassNames() { return new String[] { @@ -45,45 +55,71 @@ public class GoogleHttpClientInstrumentation extends Instrumenter.Default { "datadog.trace.agent.decorator.ClientDecorator", "datadog.trace.agent.decorator.HttpClientDecorator", packageName + ".GoogleHttpClientDecorator", + packageName + ".RequestState", getClass().getName() + "$GoogleHttpClientAdvice", + getClass().getName() + "$GoogleHttpClientAsyncAdvice", getClass().getName() + "$HeadersInjectAdapter" }; } @Override public Map, String> transformers() { - return singletonMap( - isMethod().and(isPublic()).and(named("execute").and(takesArguments(0))), + final Map, String> transformers = new HashMap<>(); + transformers.put( + isMethod().and(isPublic()).and(named("execute")).and(takesArguments(0)), GoogleHttpClientAdvice.class.getName()); + + transformers.put( + isMethod() + .and(isPublic()) + .and(named("executeAsync")) + .and(takesArguments(1)) + .and(takesArgument(0, (named("java.util.concurrent.Executor")))), + GoogleHttpClientAsyncAdvice.class.getName()); + + return transformers; } public static class GoogleHttpClientAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static Scope methodEnter(@Advice.This final HttpRequest request) { - final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(HttpRequest.class); - if (callDepth > 0) { - return null; + public static void methodEnter(@Advice.This final HttpRequest request) { + + final ContextStore contextStore = + InstrumentationContext.get(HttpRequest.class, RequestState.class); + + RequestState state = contextStore.get(request); + + if (state == null) { + state = new RequestState(GlobalTracer.get().buildSpan("http.request").start()); + contextStore.put(request, state); } - final Span span = GlobalTracer.get().buildSpan("http.request").start(); - final Scope scope = GlobalTracer.get().scopeManager().activate(span, false); - DECORATE.afterStart(span); - DECORATE.onRequest(span, request); - GlobalTracer.get() - .inject(span.context(), Format.Builtin.HTTP_HEADERS, new HeadersInjectAdapter(request)); - return scope; + final Span span = state.getSpan(); + + try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, false)) { + DECORATE.afterStart(span); + DECORATE.onRequest(span, request); + GlobalTracer.get() + .inject(span.context(), Format.Builtin.HTTP_HEADERS, new HeadersInjectAdapter(request)); + } } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void methodExit( - @Advice.Enter final Scope scope, + @Advice.This final HttpRequest request, @Advice.Return final HttpResponse response, @Advice.Thrown final Throwable throwable) { - if (scope != null) { - try { - final Span span = scope.span(); + final ContextStore contextStore = + InstrumentationContext.get(HttpRequest.class, RequestState.class); + final RequestState state = contextStore.get(request); + + if (state != null) { + final Span span = state.getSpan(); + + try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, false)) { DECORATE.onResponse(span, response); + DECORATE.onError(span, throwable); // If HttpRequest.setThrowExceptionOnExecuteError is set to false, there are no exceptions // for a failed request. Thus, check the response code @@ -91,13 +127,46 @@ public class GoogleHttpClientInstrumentation extends Instrumenter.Default { Tags.ERROR.set(span, true); span.log(singletonMap(Fields.MESSAGE, response.getStatusMessage())); } - DECORATE.onError(span, throwable); DECORATE.beforeFinish(span); span.finish(); - } finally { - scope.close(); - CallDepthThreadLocalMap.reset(HttpRequest.class); + } + } + } + } + + public static class GoogleHttpClientAsyncAdvice { + + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void methodEnter(@Advice.This final HttpRequest request) { + final Span span = GlobalTracer.get().buildSpan("http.request").start(); + + final ContextStore contextStore = + InstrumentationContext.get(HttpRequest.class, RequestState.class); + + final RequestState state = new RequestState(span); + contextStore.put(request, state); + } + + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) + public static void methodExit( + @Advice.This final HttpRequest request, @Advice.Thrown final Throwable throwable) { + + if (throwable != null) { + + final ContextStore contextStore = + InstrumentationContext.get(HttpRequest.class, RequestState.class); + final RequestState state = contextStore.get(request); + + if (state != null) { + final Span span = state.getSpan(); + + try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, false)) { + DECORATE.onError(span, throwable); + + DECORATE.beforeFinish(span); + span.finish(); + } } } } diff --git a/dd-java-agent/instrumentation/google-http-client/src/main/java/datadog/trace/instrumentation/googlehttpclient/RequestState.java b/dd-java-agent/instrumentation/google-http-client/src/main/java/datadog/trace/instrumentation/googlehttpclient/RequestState.java new file mode 100644 index 0000000000..8e52df3b72 --- /dev/null +++ b/dd-java-agent/instrumentation/google-http-client/src/main/java/datadog/trace/instrumentation/googlehttpclient/RequestState.java @@ -0,0 +1,10 @@ +package datadog.trace.instrumentation.googlehttpclient; + +import io.opentracing.Span; +import lombok.Data; +import lombok.NonNull; + +@Data +public class RequestState { + @NonNull public Span span; +} diff --git a/dd-java-agent/instrumentation/google-http-client/src/test/groovy/AbstractGoogleHttpClientTest.groovy b/dd-java-agent/instrumentation/google-http-client/src/test/groovy/AbstractGoogleHttpClientTest.groovy new file mode 100644 index 0000000000..2dce2bdb81 --- /dev/null +++ b/dd-java-agent/instrumentation/google-http-client/src/test/groovy/AbstractGoogleHttpClientTest.groovy @@ -0,0 +1,69 @@ +import com.google.api.client.http.GenericUrl +import com.google.api.client.http.HttpRequest +import com.google.api.client.http.HttpResponse +import com.google.api.client.http.javanet.NetHttpTransport +import datadog.trace.agent.test.base.HttpClientTest +import datadog.trace.api.DDSpanTypes +import datadog.trace.instrumentation.googlehttpclient.GoogleHttpClientDecorator +import spock.lang.Shared + +abstract class AbstractGoogleHttpClientTest extends HttpClientTest { + + @Shared + def requestFactory = new NetHttpTransport().createRequestFactory(); + + @Override + int doRequest(String method, URI uri, Map headers, Closure callback) { + doRequest(method, uri, headers, callback, false); + } + + int doRequest(String method, URI uri, Map headers, Closure callback, boolean throwExceptionOnError) { + GenericUrl genericUrl = new GenericUrl(uri) + + HttpRequest request = requestFactory.buildRequest(method, genericUrl, null) + request.getHeaders().putAll(headers) + request.setThrowExceptionOnExecuteError(throwExceptionOnError) + + HttpResponse response = executeRequest(request); + callback?.call() + + return response.getStatusCode() + } + + abstract HttpResponse executeRequest(HttpRequest request); + + @Override + GoogleHttpClientDecorator decorator() { + return GoogleHttpClientDecorator.DECORATE + } + + @Override + boolean testRedirects() { + // Circular redirects don't throw an exception with Google Http Client + return false + } + + def "error traces when exception is not thrown"() { + given: + def uri = server.address.resolve("/error") + + when: + def status = doRequest(method, uri) + + then: + status == 500 + assertTraces(2) { + server.distributedRequestTrace(it, 0, trace(1).last()) + trace(1, size(1)) { + span(0) { + resourceName "$method $uri.path" + spanType DDSpanTypes.HTTP_CLIENT + errored true + } + } + } + + where: + method = "GET" + } +} diff --git a/dd-java-agent/instrumentation/google-http-client/src/test/groovy/GoogleHttpClientAsyncTest.groovy b/dd-java-agent/instrumentation/google-http-client/src/test/groovy/GoogleHttpClientAsyncTest.groovy new file mode 100644 index 0000000000..934e31dc54 --- /dev/null +++ b/dd-java-agent/instrumentation/google-http-client/src/test/groovy/GoogleHttpClientAsyncTest.groovy @@ -0,0 +1,9 @@ +import com.google.api.client.http.HttpRequest +import com.google.api.client.http.HttpResponse + +class GoogleHttpClientAsyncTest extends AbstractGoogleHttpClientTest { + @Override + HttpResponse executeRequest(HttpRequest request) { + return request.executeAsync().get() + } +} diff --git a/dd-java-agent/instrumentation/google-http-client/src/test/groovy/GoogleHttpClientTest.groovy b/dd-java-agent/instrumentation/google-http-client/src/test/groovy/GoogleHttpClientTest.groovy index 944b72dfaf..70272d34e7 100644 --- a/dd-java-agent/instrumentation/google-http-client/src/test/groovy/GoogleHttpClientTest.groovy +++ b/dd-java-agent/instrumentation/google-http-client/src/test/groovy/GoogleHttpClientTest.groovy @@ -1,67 +1,9 @@ -import com.google.api.client.http.GenericUrl import com.google.api.client.http.HttpRequest import com.google.api.client.http.HttpResponse -import com.google.api.client.http.javanet.NetHttpTransport -import datadog.trace.agent.test.base.HttpClientTest -import datadog.trace.api.DDSpanTypes -import datadog.trace.instrumentation.googlehttpclient.GoogleHttpClientDecorator -import spock.lang.Shared - -class GoogleHttpClientTest extends HttpClientTest { - - @Shared - def requestFactory = new NetHttpTransport().createRequestFactory(); +class GoogleHttpClientTest extends AbstractGoogleHttpClientTest { @Override - int doRequest(String method, URI uri, Map headers, Closure callback) { - doRequest(method, uri, headers, callback, false); - } - - int doRequest(String method, URI uri, Map headers, Closure callback, boolean throwExceptionOnError) { - GenericUrl genericUrl = new GenericUrl(uri) - - HttpRequest request = requestFactory.buildRequest(method, genericUrl, null) - request.getHeaders().putAll(headers) - request.setThrowExceptionOnExecuteError(throwExceptionOnError) - - HttpResponse response = request.execute() - callback?.call() - - return response.getStatusCode() - } - - @Override - GoogleHttpClientDecorator decorator() { - return GoogleHttpClientDecorator.DECORATE - } - - @Override - boolean testRedirects() { - // Circular redirects don't throw an exception with Google Http Client - return false - } - - def "error traces when exception is not thrown"() { - given: - def uri = server.address.resolve("/error") - - when: - def status = doRequest(method, uri) - - then: - status == 500 - assertTraces(2) { - server.distributedRequestTrace(it, 0, trace(1).last()) - trace(1, size(1)) { - span(0) { - resourceName "$method $uri.path" - spanType DDSpanTypes.HTTP_CLIENT - errored true - } - } - } - - where: - method = "GET" + HttpResponse executeRequest(HttpRequest request) { + return request.execute(); } }