Merge pull request #582 from DataDog/tyler/servlet-dispatch

Fix servlet async dispatch
DataDog · Nov 16, 2018 · e888b17 · e888b17
2 parents 7b6c25b + e4a1240
commit e888b17
Show file tree

Hide file tree

Showing 13 changed files with 703 additions and 395 deletions.
diff --git a/...src/main/java/datadog/trace/instrumentation/servlet3/AbstractServlet3Instrumentation.java b/...src/main/java/datadog/trace/instrumentation/servlet3/AbstractServlet3Instrumentation.java
diff --git a/...t-3/src/main/java/datadog/trace/instrumentation/servlet3/AsyncContextInstrumentation.java b/...t-3/src/main/java/datadog/trace/instrumentation/servlet3/AsyncContextInstrumentation.java
@@ -0,0 +1,97 @@
+package datadog.trace.instrumentation.servlet3;
+
+import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
+import static datadog.trace.instrumentation.servlet3.Servlet3Advice.SERVLET_SPAN;
+import static net.bytebuddy.matcher.ElementMatchers.isInterface;
+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.not;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.bootstrap.CallDepthThreadLocalMap;
+import io.opentracing.Span;
+import io.opentracing.propagation.Format;
+import io.opentracing.util.GlobalTracer;
+import java.util.Collections;
+import java.util.Map;
+import javax.servlet.AsyncContext;
+import javax.servlet.ServletRequest;
+import javax.servlet.http.HttpServletRequest;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(Instrumenter.class)
+public final class AsyncContextInstrumentation extends Instrumenter.Default {
+
+  public AsyncContextInstrumentation() {
+    super("servlet", "servlet-3");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {"datadog.trace.instrumentation.servlet3.HttpServletRequestInjectAdapter"};
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return not(isInterface()).and(safeHasSuperType(named("javax.servlet.AsyncContext")));
+  }
+
+  @Override
+  public Map<? extends ElementMatcher, String> transformers() {
+    return Collections.singletonMap(
+        isMethod().and(isPublic()).and(named("dispatch")), DispatchAdvice.class.getName());
+  }
+
+  /**
+   * When a request is dispatched, we want to close out the existing request and replace the
+   * propagation info in the headers with the closed span.
+   */
+  public static class DispatchAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static boolean enter(
+        @Advice.This final AsyncContext context, @Advice.AllArguments final Object[] args) {
+      final int depth = CallDepthThreadLocalMap.incrementCallDepth(AsyncContext.class);
+      if (depth > 0) {
+        return false;
+      }
+
+      final ServletRequest request = context.getRequest();
+      final Object spanAttr = request.getAttribute(SERVLET_SPAN);
+      if (spanAttr instanceof Span) {
+        request.removeAttribute(SERVLET_SPAN);
+        final Span span = (Span) spanAttr;
+        // Override propagation headers by injecting attributes with new values.
+        if (request instanceof HttpServletRequest) {
+          GlobalTracer.get()
+              .inject(
+                  span.context(),
+                  Format.Builtin.TEXT_MAP,
+                  new HttpServletRequestInjectAdapter((HttpServletRequest) request));
+        }
+        final String path;
+        if (args.length == 1 && args[0] instanceof String) {
+          path = (String) args[0];
+        } else if (args.length == 2 && args[1] instanceof String) {
+          path = (String) args[1];
+        } else {
+          path = "true";
+        }
+        span.setTag("servlet.dispatch", path);
+        span.finish();
+      }
+      return true;
+    }
+
+    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+    public static void exit(@Advice.Enter final boolean topLevel) {
+      if (topLevel) {
+        CallDepthThreadLocalMap.reset(AsyncContext.class);
+      }
+    }
+  }
+}
diff --git a/...t-3/src/main/java/datadog/trace/instrumentation/servlet3/FilterChain3Instrumentation.java b/...t-3/src/main/java/datadog/trace/instrumentation/servlet3/FilterChain3Instrumentation.java
@@ -15,7 +15,19 @@
 import net.bytebuddy.matcher.ElementMatcher;
 
 @AutoService(Instrumenter.class)
-public final class FilterChain3Instrumentation extends AbstractServlet3Instrumentation {
+public final class FilterChain3Instrumentation extends Instrumenter.Default {
+  public FilterChain3Instrumentation() {
+    super("servlet", "servlet-3");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      "datadog.trace.instrumentation.servlet3.HttpServletRequestExtractAdapter",
+      "datadog.trace.instrumentation.servlet3.HttpServletRequestExtractAdapter$MultivaluedMapFlatIterator",
+      "datadog.trace.instrumentation.servlet3.TagSettingAsyncListener"
+    };
+  }
 
   @Override
   public ElementMatcher<TypeDescription> typeMatcher() {

diff --git a/...t-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServlet3Instrumentation.java b/...t-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServlet3Instrumentation.java
@@ -15,7 +15,19 @@
 import net.bytebuddy.matcher.ElementMatcher;
 
 @AutoService(Instrumenter.class)
-public final class HttpServlet3Instrumentation extends AbstractServlet3Instrumentation {
+public final class HttpServlet3Instrumentation extends Instrumenter.Default {
+  public HttpServlet3Instrumentation() {
+    super("servlet", "servlet-3");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      "datadog.trace.instrumentation.servlet3.HttpServletRequestExtractAdapter",
+      "datadog.trace.instrumentation.servlet3.HttpServletRequestExtractAdapter$MultivaluedMapFlatIterator",
+      "datadog.trace.instrumentation.servlet3.TagSettingAsyncListener"
+    };
+  }
 
   @Override
   public ElementMatcher<TypeDescription> typeMatcher() {

diff --git a/...rc/main/java/datadog/trace/instrumentation/servlet3/HttpServletRequestExtractAdapter.java b/...rc/main/java/datadog/trace/instrumentation/servlet3/HttpServletRequestExtractAdapter.java
@@ -3,6 +3,7 @@
 import io.opentracing.propagation.TextMap;
 import java.util.AbstractMap;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.Enumeration;
 import java.util.HashMap;
 import java.util.Iterator;
@@ -52,6 +53,20 @@ protected Map<String, List<String>> servletHeadersToMultiMap(
       headersResult.put(headerName, valuesList);
     }
 
+    /*
+     * Read from the attributes and override the headers.
+     * This is used by HttpServletRequestInjectAdapter when a request is async-dispatched.
+     */
+    final Enumeration<String> attributeNamesIt = httpServletRequest.getAttributeNames();
+    while (attributeNamesIt.hasMoreElements()) {
+      final String attributeName = attributeNamesIt.nextElement();
+
+      final Object valuesIt = httpServletRequest.getAttribute(attributeName);
+      if (valuesIt instanceof String) {
+        headersResult.put(attributeName, Collections.singletonList((String) valuesIt));
+      }
+    }
+
     return headersResult;
   }
 

diff --git a/...src/main/java/datadog/trace/instrumentation/servlet3/HttpServletRequestInjectAdapter.java b/...src/main/java/datadog/trace/instrumentation/servlet3/HttpServletRequestInjectAdapter.java
@@ -0,0 +1,27 @@
+package datadog.trace.instrumentation.servlet3;
+
+import io.opentracing.propagation.TextMap;
+import java.util.Iterator;
+import java.util.Map;
+import javax.servlet.http.HttpServletRequest;
+
+/** Inject into request attributes since the request headers can't be modified. */
+public class HttpServletRequestInjectAdapter implements TextMap {
+
+  private final HttpServletRequest httpServletRequest;
+
+  public HttpServletRequestInjectAdapter(final HttpServletRequest httpServletRequest) {
+    this.httpServletRequest = httpServletRequest;
+  }
+
+  @Override
+  public Iterator<Map.Entry<String, String>> iterator() {
+    throw new UnsupportedOperationException(
+        "This class should be used only with Tracer.extract()!");
+  }
+
+  @Override
+  public void put(final String key, final String value) {
+    httpServletRequest.setAttribute(key, value);
+  }
+}
diff --git a/...tation/servlet-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java b/...tation/servlet-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java
@@ -21,11 +21,13 @@
 import net.bytebuddy.asm.Advice;
 
 public class Servlet3Advice {
+  public static final String SERVLET_SPAN = "datadog.servlet.span";
 
   @Advice.OnMethodEnter(suppress = Throwable.class)
   public static Scope startSpan(
       @Advice.This final Object servlet, @Advice.Argument(0) final ServletRequest req) {
-    if (GlobalTracer.get().activeSpan() != null || !(req instanceof HttpServletRequest)) {
+    final Object spanAttr = req.getAttribute(SERVLET_SPAN);
+    if (!(req instanceof HttpServletRequest) || spanAttr != null) {
       // Tracing might already be applied by the FilterChain.  If so ignore this.
       return null;
     }
@@ -53,6 +55,8 @@ public static Scope startSpan(
     if (scope instanceof TraceScope) {
       ((TraceScope) scope).setAsyncPropagation(true);
     }
+
+    req.setAttribute(SERVLET_SPAN, scope.span());
     return scope;
   }
 
@@ -63,13 +67,11 @@ public static void stopSpan(
       @Advice.Enter final Scope scope,
       @Advice.Thrown final Throwable throwable) {
     // Set user.principal regardless of who created this span.
-    final Span currentSpan = GlobalTracer.get().activeSpan();
-    if (currentSpan != null) {
-      if (request instanceof HttpServletRequest) {
-        final Principal principal = ((HttpServletRequest) request).getUserPrincipal();
-        if (principal != null) {
-          currentSpan.setTag(DDTags.USER_NAME, principal.getName());
-        }
+    final Object spanAttr = request.getAttribute(SERVLET_SPAN);
+    if (spanAttr instanceof Span && request instanceof HttpServletRequest) {
+      final Principal principal = ((HttpServletRequest) request).getUserPrincipal();
+      if (principal != null) {
+        ((Span) spanAttr).setTag(DDTags.USER_NAME, principal.getName());
       }
     }
 
@@ -90,19 +92,23 @@ public static void stopSpan(
             ((TraceScope) scope).setAsyncPropagation(false);
           }
           scope.close();
+          req.removeAttribute(SERVLET_SPAN);
           span.finish(); // Finish the span manually since finishSpanOnClose was false
-        } else if (req.isAsyncStarted()) {
-          final AtomicBoolean activated = new AtomicBoolean(false);
-          // what if async is already finished? This would not be called
-          req.getAsyncContext().addListener(new TagSettingAsyncListener(activated, span));
-          scope.close();
         } else {
-          Tags.HTTP_STATUS.set(span, resp.getStatus());
-          if (scope instanceof TraceScope) {
-            ((TraceScope) scope).setAsyncPropagation(false);
+          final AtomicBoolean activated = new AtomicBoolean(false);
+          if (req.isAsyncStarted()) {
+            req.getAsyncContext().addListener(new TagSettingAsyncListener(activated, span));
+          }
+          // Check again in case the request finished before adding the listener.
+          if (!req.isAsyncStarted() && activated.compareAndSet(false, true)) {
+            Tags.HTTP_STATUS.set(span, resp.getStatus());
+            if (scope instanceof TraceScope) {
+              ((TraceScope) scope).setAsyncPropagation(false);
+            }
+            req.removeAttribute(SERVLET_SPAN);
+            span.finish(); // Finish the span manually since finishSpanOnClose was false
           }
           scope.close();
-          span.finish(); // Finish the span manually since finishSpanOnClose was false
         }
       }
     }

diff --git a/...rvlet-3/src/main/java/datadog/trace/instrumentation/servlet3/TagSettingAsyncListener.java b/...rvlet-3/src/main/java/datadog/trace/instrumentation/servlet3/TagSettingAsyncListener.java
@@ -10,6 +10,7 @@
 import java.io.IOException;
 import java.util.Collections;
 import java.util.concurrent.atomic.AtomicBoolean;
+import javax.servlet.AsyncContext;
 import javax.servlet.AsyncEvent;
 import javax.servlet.AsyncListener;
 import javax.servlet.http.HttpServletResponse;
@@ -72,6 +73,12 @@ public void onError(final AsyncEvent event) throws IOException {
     }
   }
 
+  /** Re-attach listener for dispatch. */
   @Override
-  public void onStartAsync(final AsyncEvent event) throws IOException {}
+  public void onStartAsync(final AsyncEvent event) {
+    final AsyncContext eventAsyncContext = event.getAsyncContext();
+    if (eventAsyncContext != null) {
+      eventAsyncContext.addListener(this, event.getSuppliedRequest(), event.getSuppliedResponse());
+    }
+  }
 }