From d598515d09dc3e7ab60feae8150aeeaf2d07efb6 Mon Sep 17 00:00:00 2001 From: Andrew Kent Date: Thu, 19 Jul 2018 15:52:17 -0700 Subject: [PATCH] Handle cases where connect() is called first --- .../HttpUrlConnectionInstrumentation.java | 68 +++++++++++++++---- .../test/groovy/HttpUrlConnectionTest.groovy | 20 +----- 2 files changed, 56 insertions(+), 32 deletions(-) diff --git a/dd-java-agent/instrumentation/http-url-connection/src/main/java/datadog/trace/instrumentation/http_url_connection/HttpUrlConnectionInstrumentation.java b/dd-java-agent/instrumentation/http-url-connection/src/main/java/datadog/trace/instrumentation/http_url_connection/HttpUrlConnectionInstrumentation.java index ba50831711..bc01ee69aa 100644 --- a/dd-java-agent/instrumentation/http-url-connection/src/main/java/datadog/trace/instrumentation/http_url_connection/HttpUrlConnectionInstrumentation.java +++ b/dd-java-agent/instrumentation/http-url-connection/src/main/java/datadog/trace/instrumentation/http_url_connection/HttpUrlConnectionInstrumentation.java @@ -24,6 +24,7 @@ import java.net.URL; import java.util.Collections; import java.util.HashMap; import java.util.Map; +import java.util.WeakHashMap; import javax.net.ssl.HttpsURLConnection; import net.bytebuddy.asm.Advice; import net.bytebuddy.matcher.ElementMatcher; @@ -50,7 +51,8 @@ public class HttpUrlConnectionInstrumentation extends Instrumenter.Default { @Override public String[] helperClassNames() { return new String[] { - "datadog.trace.instrumentation.http_url_connection.MessageHeadersInjectAdapter" + "datadog.trace.instrumentation.http_url_connection.MessageHeadersInjectAdapter", + HttpUrlConnectionInstrumentation.class.getName() + "$HttpURLState" }; } @@ -75,19 +77,32 @@ public class HttpUrlConnectionInstrumentation extends Instrumenter.Default { @Advice.OnMethodEnter(suppress = Throwable.class) public static Scope startSpan( @Advice.This final HttpURLConnection thiz, - @Advice.FieldValue("connected") final boolean connected) { + @Advice.FieldValue("connected") final boolean connected, + @Advice.Origin("#m") final String methodName) { - // This works with keep-alive because the "connected" flag is not initialized until the socket cache is checked. - // This allows us to use the connected flag to know if this instance has done any io. - if (connected) { - return null; + final HttpURLState state = HttpURLState.get(thiz); + if ("connect".equals(methodName)) { + if (connected) { + return null; + } + // We get here in cases where connect() is called first. + // We need to inject headers now because after connected=true no more headers can be added. + + // In total there will be two spans: + // - one for the connect() which does propagation + // - one for the input or output stream (presumably called after connect()) + } else { + if (state.hasDoneIO()) { + return null; + } + state.setHasDoneIO(true); } // AgentWriter uses HttpURLConnection to report to the trace-agent. We don't want to trace those requests. // Check after the connected test above because getRequestProperty will throw an exception if already connected. final boolean isTraceRequest = Thread.currentThread().getName().equals("dd-agent-writer") - || thiz.getRequestProperty("Datadog-Meta-Lang") != null; + || (!connected && thiz.getRequestProperty("Datadog-Meta-Lang") != null); if (isTraceRequest) { return null; } @@ -97,8 +112,6 @@ public class HttpUrlConnectionInstrumentation extends Instrumenter.Default { return null; } - final String protocol = thiz.getURL().getProtocol(); - final Tracer tracer = GlobalTracer.get(); final Scope scope = tracer @@ -140,15 +153,40 @@ public class HttpUrlConnectionInstrumentation extends Instrumenter.Default { if (throwable != null) { Tags.ERROR.set(span, true); span.log(Collections.singletonMap(ERROR_OBJECT, throwable)); - } else { - if (responseCode > 0) { - // responseCode field cache is sometimes not populated. - // We can't call getResponseCode() due to some unwanted side-effects (e.g. breaks getOutputStream). - Tags.HTTP_STATUS.set(span, responseCode); - } + } else if (responseCode > 0) { + // responseCode field cache is sometimes not populated. + // We can't call getResponseCode() due to some unwanted side-effects (e.g. breaks getOutputStream). + Tags.HTTP_STATUS.set(span, responseCode); } scope.close(); CallDepthThreadLocalMap.reset(HttpURLConnection.class); } } + + public static class HttpURLState { + private static final Map STATE_MAP = + Collections.synchronizedMap(new WeakHashMap()); + + public static HttpURLState get(HttpURLConnection connection) { + HttpURLState state = STATE_MAP.get(connection); + if (state == null) { + // not thread-safe, but neither is HttpURLConnection + state = new HttpURLState(); + STATE_MAP.put(connection, state); + } + return state; + } + + public boolean hasDoneIO = false; + + private HttpURLState() {} + + public boolean hasDoneIO() { + return hasDoneIO; + } + + public void setHasDoneIO(boolean value) { + this.hasDoneIO = value; + } + } } diff --git a/dd-java-agent/instrumentation/http-url-connection/src/test/groovy/HttpUrlConnectionTest.groovy b/dd-java-agent/instrumentation/http-url-connection/src/test/groovy/HttpUrlConnectionTest.groovy index 26c229cda1..8876b4fca7 100644 --- a/dd-java-agent/instrumentation/http-url-connection/src/test/groovy/HttpUrlConnectionTest.groovy +++ b/dd-java-agent/instrumentation/http-url-connection/src/test/groovy/HttpUrlConnectionTest.groovy @@ -295,7 +295,6 @@ class HttpUrlConnectionTest extends AgentTestRunner { "$DDTags.SPAN_TYPE" DDSpanTypes.HTTP_CLIENT "$Tags.HTTP_URL.key" "$server.address" "$Tags.HTTP_METHOD.key" "POST" - "$Tags.HTTP_STATUS.key" STATUS "$Tags.PEER_HOSTNAME.key" "localhost" "$Tags.PEER_PORT.key" server.address.port defaultTags() @@ -339,33 +338,21 @@ class HttpUrlConnectionTest extends AgentTestRunner { RestTemplate restTemplate = new RestTemplate() String res = restTemplate.getForObject(server.address.toString(), String) assert res == RESPONSE - String res2 = restTemplate.getForObject(server.address.toString(), String) - assert res2 == RESPONSE } expect: - assertTraces(TEST_WRITER, 3) { + assertTraces(TEST_WRITER, 2) { trace(0, 1) { span(0) { operationName "test-http-server" - childOf(TEST_WRITER[2][2]) + childOf(TEST_WRITER[1][2]) errored false tags { defaultTags() } } } - trace(1, 1) { - span(0) { - operationName "test-http-server" - childOf(TEST_WRITER[2][1]) - errored false - tags { - defaultTags() - } - } - } - trace(2, 3) { + trace(1, 3) { span(0) { operationName "someTrace" parent() @@ -400,7 +387,6 @@ class HttpUrlConnectionTest extends AgentTestRunner { "$DDTags.SPAN_TYPE" DDSpanTypes.HTTP_CLIENT "$Tags.HTTP_URL.key" "$server.address" "$Tags.HTTP_METHOD.key" "GET" - "$Tags.HTTP_STATUS.key" STATUS "$Tags.PEER_HOSTNAME.key" "localhost" "$Tags.PEER_PORT.key" server.address.port defaultTags()