Merge pull request #1304 from DataDog/landerson/httprequest-attributes

Add traceId and spanId as request attributes

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

@@ -13,6 +13,8 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.api.CorrelationIdentifier;
+import datadog.trace.api.GlobalTracer;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan.Context;
@@ -81,6 +83,8 @@ public class GrizzlyHttpHandlerInstrumentation extends Instrumenter.Default {
       scope.setAsyncPropagation(true);
 
       request.setAttribute(DD_SPAN_ATTRIBUTE, span);
+      request.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());
+      request.setAttribute(CorrelationIdentifier.getSpanIdKey(), GlobalTracer.get().getSpanId());
       request.addAfterServiceListener(SpanClosingListener.LISTENER);
 
       return scope;

dd-java-agent/instrumentation/jetty-8/src/main/java/datadog/trace/instrumentation/jetty8/JettyHandlerAdvice.java

@@ -7,7 +7,9 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
 import static datadog.trace.instrumentation.jetty8.HttpServletRequestExtractAdapter.GETTER;
 import static datadog.trace.instrumentation.jetty8.JettyDecorator.DECORATE;
 
+import datadog.trace.api.CorrelationIdentifier;
 import datadog.trace.api.DDTags;
+import datadog.trace.api.GlobalTracer;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
@@ -41,6 +43,8 @@ public class JettyHandlerAdvice {
     final AgentScope scope = activateSpan(span, false);
     scope.setAsyncPropagation(true);
     req.setAttribute(DD_SPAN_ATTRIBUTE, span);
+    req.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());
+    req.setAttribute(CorrelationIdentifier.getSpanIdKey(), GlobalTracer.get().getSpanId());
     return scope;
   }
 

dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java

@@ -7,7 +7,9 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
 import static datadog.trace.instrumentation.servlet2.HttpServletRequestExtractAdapter.GETTER;
 import static datadog.trace.instrumentation.servlet2.Servlet2Decorator.DECORATE;
 
+import datadog.trace.api.CorrelationIdentifier;
 import datadog.trace.api.DDTags;
+import datadog.trace.api.GlobalTracer;
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -56,10 +58,15 @@ public class Servlet2Advice {
     DECORATE.onConnection(span, httpServletRequest);
     DECORATE.onRequest(span, httpServletRequest);
 
-    httpServletRequest.setAttribute(DD_SPAN_ATTRIBUTE, span);
-
     final AgentScope scope = activateSpan(span, true);
     scope.setAsyncPropagation(true);
+
+    httpServletRequest.setAttribute(DD_SPAN_ATTRIBUTE, span);
+    httpServletRequest.setAttribute(
+        CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());
+    httpServletRequest.setAttribute(
+        CorrelationIdentifier.getSpanIdKey(), GlobalTracer.get().getSpanId());
+
     return scope;
   }
 

dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java

@@ -7,7 +7,9 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
 import static datadog.trace.instrumentation.servlet3.HttpServletRequestExtractAdapter.GETTER;
 import static datadog.trace.instrumentation.servlet3.Servlet3Decorator.DECORATE;
 
+import datadog.trace.api.CorrelationIdentifier;
 import datadog.trace.api.DDTags;
+import datadog.trace.api.GlobalTracer;
 import datadog.trace.bootstrap.InstrumentationContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
@@ -51,10 +53,15 @@ public class Servlet3Advice {
     DECORATE.onConnection(span, httpServletRequest);
     DECORATE.onRequest(span, httpServletRequest);
 
-    httpServletRequest.setAttribute(DD_SPAN_ATTRIBUTE, span);
-
     final AgentScope scope = activateSpan(span, false);
     scope.setAsyncPropagation(true);
+
+    httpServletRequest.setAttribute(DD_SPAN_ATTRIBUTE, span);
+    httpServletRequest.setAttribute(
+        CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId());
+    httpServletRequest.setAttribute(
+        CorrelationIdentifier.getSpanIdKey(), GlobalTracer.get().getSpanId());
+
     return scope;
   }
 

dd-java-agent/instrumentation/servlet/request-3/src/test/groovy/TomcatServlet3Test.groovy

@@ -1,11 +1,13 @@
 import com.google.common.io.Files
 import datadog.opentracing.DDSpan
 import datadog.trace.agent.test.asserts.ListWriterAssert
+import datadog.trace.api.CorrelationIdentifier
 import datadog.trace.api.DDSpanTypes
 import datadog.trace.api.DDTags
 import datadog.trace.bootstrap.instrumentation.api.Tags
 import groovy.transform.stc.ClosureParams
 import groovy.transform.stc.SimpleType
+import org.apache.catalina.AccessLog
 import org.apache.catalina.Context
 import org.apache.catalina.connector.Request
 import org.apache.catalina.connector.Response
@@ -13,10 +15,14 @@ import org.apache.catalina.core.ApplicationFilterChain
 import org.apache.catalina.core.StandardHost
 import org.apache.catalina.startup.Tomcat
 import org.apache.catalina.valves.ErrorReportValve
+import org.apache.catalina.valves.ValveBase
 import org.apache.tomcat.JarScanFilter
 import org.apache.tomcat.JarScanType
+import spock.lang.Shared
+import spock.lang.Unroll
 
 import javax.servlet.Servlet
+import javax.servlet.ServletException
 
 import static datadog.trace.agent.test.asserts.TraceAssert.assertTrace
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.AUTH_REQUIRED
@@ -28,8 +34,12 @@ import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.REDIRE
 import static datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint.SUCCESS
 import static datadog.trace.agent.test.utils.TraceUtils.basicSpan
 
+@Unroll
 abstract class TomcatServlet3Test extends AbstractServlet3Test<Tomcat, Context> {
 
+  @Shared
+  def accessLogValue = new TestAccessLogValve()
+
   @Override
   Tomcat startServer(int port) {
     def tomcatServer = new Tomcat()
@@ -59,12 +69,17 @@ abstract class TomcatServlet3Test extends AbstractServlet3Test<Tomcat, Context>
     setupServlets(servletContext)
 
     (tomcatServer.host as StandardHost).errorReportValveClass = ErrorHandlerValve.name
+    (tomcatServer.host as StandardHost).getPipeline().addValve(accessLogValue)
 
     tomcatServer.start()
 
     return tomcatServer
   }
 
+  def setup() {
+    accessLogValue.loggedIds.clear()
+  }
+
   @Override
   void stopServer(Tomcat server) {
     server.stop()
@@ -83,6 +98,68 @@ abstract class TomcatServlet3Test extends AbstractServlet3Test<Tomcat, Context>
     servletContext.addServletMappingDecoded(path, name)
   }
 
+
+  def "access log has ids for #count requests"() {
+    given:
+    def request = request(SUCCESS, method, body).build()
+
+    when:
+    List<okhttp3.Response> responses = (1..count).collect {
+      return client.newCall(request).execute()
+    }
+
+    then:
+    responses.each { response ->
+      assert response.code() == SUCCESS.status
+      assert response.body().string() == SUCCESS.body
+    }
+
+    and:
+    cleanAndAssertTraces(count) {
+      (1..count).eachWithIndex { val, i ->
+        trace(i, 2) {
+          serverSpan(it, 0)
+          controllerSpan(it, 1, span(0))
+        }
+
+        def (String traceId, String spanId) = accessLogValue.loggedIds[i]
+        assert trace(i).get(0).traceId.toString() == traceId
+        assert trace(i).get(0).spanId.toString() == spanId
+      }
+    }
+
+    where:
+    method = "GET"
+    body = null
+    count << [1, 4] // make multiple requests.
+  }
+
+  def "access log has ids for error request"() {
+    setup:
+    def request = request(ERROR, method, body).build()
+    def response = client.newCall(request).execute()
+
+    expect:
+    response.code() == ERROR.status
+    response.body().string() == ERROR.body
+
+    and:
+    cleanAndAssertTraces(1) {
+      trace(0, 2) {
+        serverSpan(it, 0, null, null, method, ERROR)
+        controllerSpan(it, 1, span(0))
+      }
+
+      def (String traceId, String spanId) = accessLogValue.loggedIds[0]
+      assert trace(0).get(0).traceId.toString() == traceId
+      assert trace(0).get(0).spanId.toString() == spanId
+    }
+
+    where:
+    method = "GET"
+    body = null
+  }
+
   // FIXME: Add authentication tests back in...
 //  private setupAuthentication(Tomcat server, Context servletContext) {
 //    // Login Config
@@ -131,6 +208,33 @@ class ErrorHandlerValve extends ErrorReportValve {
   }
 }
 
+class TestAccessLogValve extends ValveBase implements AccessLog {
+  List<Tuple2<String, String>> loggedIds = []
+
+  TestAccessLogValve() {
+    super(true)
+  }
+
+  void log(Request request, Response response, long time) {
+    loggedIds.add(new Tuple2(request.getAttribute(CorrelationIdentifier.traceIdKey),
+      request.getAttribute(CorrelationIdentifier.spanIdKey)))
+  }
+
+  @Override
+  void setRequestAttributesEnabled(boolean requestAttributesEnabled) {
+  }
+
+  @Override
+  boolean getRequestAttributesEnabled() {
+    return false
+  }
+
+  @Override
+  void invoke(Request request, Response response) throws IOException, ServletException {
+    getNext().invoke(request, response)
+  }
+}
+
 class TomcatServlet3TestSync extends TomcatServlet3Test {
 
   @Override