Merge pull request #595 from DataDog/mar-kolya/servlet-instrumentation-fixes

Servlet3 instrumentation race condition fixes

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

@@ -47,8 +47,9 @@ public final class AsyncContextInstrumentation extends Instrumenter.Default {
   }
 
   /**
-   * When a request is dispatched, we want to close out the existing request and replace the
+   * When a request is dispatched, we want new request to have propagation headers from its parent
-   * propagation info in the headers with the closed span.
+   * request. The parent request's span is later closed by {@code
+   * TagSettingAsyncListener#onStartAsync}
    */
   public static class DispatchAdvice {
 
@@ -65,7 +66,8 @@ public final class AsyncContextInstrumentation extends Instrumenter.Default {
       if (spanAttr instanceof Span) {
         request.removeAttribute(SERVLET_SPAN);
         final Span span = (Span) spanAttr;
-        // Override propagation headers by injecting attributes with new values.
+        // Override propagation headers by injecting attributes from the current span
+        // into the new request
         if (request instanceof HttpServletRequest) {
           GlobalTracer.get()
               .inject(
@@ -82,7 +84,6 @@ public final class AsyncContextInstrumentation extends Instrumenter.Default {
           path = "true";
         }
         span.setTag("servlet.dispatch", path);
-        span.finish();
       }
       return true;
     }

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

@@ -88,10 +88,6 @@ public class Servlet3Advice {
           }
           Tags.ERROR.set(span, Boolean.TRUE);
           span.log(Collections.singletonMap(ERROR_OBJECT, throwable));
-          if (scope instanceof TraceScope) {
-            ((TraceScope) scope).setAsyncPropagation(false);
-          }
-          scope.close();
           req.removeAttribute(SERVLET_SPAN);
           span.finish(); // Finish the span manually since finishSpanOnClose was false
         } else {
@@ -102,14 +98,11 @@ public class Servlet3Advice {
           // 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();
         }
+        scope.close();
       }
     }
   }

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

@@ -2,7 +2,6 @@ package datadog.trace.instrumentation.servlet3;
 
 import static io.opentracing.log.Fields.ERROR_OBJECT;
 
-import datadog.trace.context.TraceScope;
 import io.opentracing.Scope;
 import io.opentracing.Span;
 import io.opentracing.tag.Tags;
@@ -10,7 +9,6 @@ import io.opentracing.util.GlobalTracer;
 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;
@@ -29,11 +27,6 @@ public class TagSettingAsyncListener implements AsyncListener {
     if (activated.compareAndSet(false, true)) {
       try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, true)) {
         Tags.HTTP_STATUS.set(span, ((HttpServletResponse) event.getSuppliedResponse()).getStatus());
-
-        if (scope instanceof TraceScope) {
-          // This doesn't do anything because we're in a new scope, but just to be safe...
-          ((TraceScope) scope).setAsyncPropagation(false);
-        }
       }
     }
   }
@@ -44,11 +37,6 @@ public class TagSettingAsyncListener implements AsyncListener {
       try (final Scope scope = GlobalTracer.get().scopeManager().activate(span, true)) {
         Tags.ERROR.set(span, Boolean.TRUE);
         span.setTag("timeout", event.getAsyncContext().getTimeout());
-
-        if (scope instanceof TraceScope) {
-          // This doesn't do anything because we're in a new scope, but just to be safe...
-          ((TraceScope) scope).setAsyncPropagation(false);
-        }
       }
     }
   }
@@ -64,21 +52,13 @@ public class TagSettingAsyncListener implements AsyncListener {
         }
         Tags.ERROR.set(span, Boolean.TRUE);
         span.log(Collections.singletonMap(ERROR_OBJECT, event.getThrowable()));
-
-        if (scope instanceof TraceScope) {
-          // This doesn't do anything because we're in a new scope, but just to be safe...
-          ((TraceScope) scope).setAsyncPropagation(false);
-        }
       }
     }
   }
 
-  /** Re-attach listener for dispatch. */
+  /** Finish current span on dispatch. New listener will be attached by Servlet3Advice */
   @Override
-  public void onStartAsync(final AsyncEvent event) {
+  public void onStartAsync(final AsyncEvent event) throws IOException {
-    final AsyncContext eventAsyncContext = event.getAsyncContext();
+    onComplete(event);
-    if (eventAsyncContext != null) {
-      eventAsyncContext.addListener(this, event.getSuppliedRequest(), event.getSuppliedResponse());
-    }
   }
 }

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

@@ -118,7 +118,7 @@ abstract class AbstractServlet3Test<CONTEXT> extends AgentTestRunner {
     "fake"       | "Hello FakeAsync"     | false | "FakeAsync"     | true
   }
 
-  def "test dispatch #path with depth #depth"() {
+  def "test dispatch #path with depth #depth, distributed tracing: #distributedTracing"() {
     setup:
     def requestBuilder = new Request.Builder()
       .url("http://localhost:$port/$context/dispatch/$path?depth=$depth")
@@ -131,7 +131,6 @@ abstract class AbstractServlet3Test<CONTEXT> extends AgentTestRunner {
 
     expect:
     response.body().string().trim() == "Hello $origin"
-
     assertTraces(2 + depth) {
       for (int i = 0; i < depth; i++) {
         trace(i, 1) {
@@ -166,7 +165,31 @@ abstract class AbstractServlet3Test<CONTEXT> extends AgentTestRunner {
           }
         }
       }
+      // In case of recursive requests or sync request the most 'bottom' span is closed before its parent
       trace(depth, 1) {
+        span(0) {
+          serviceName context
+          operationName "servlet.request"
+          resourceName "GET /$context/$path"
+          spanType DDSpanTypes.WEB_SERVLET
+          errored false
+          childOf TEST_WRITER[depth + 1][0]
+          tags {
+            "http.url" "http://localhost:$port/$context/$path"
+            "http.method" "GET"
+            "span.kind" "server"
+            "component" "java-web-servlet"
+            "span.origin.type" {
+              it == "TestServlet3\$$origin" || it == "TestServlet3\$DispatchRecursive" || it == ApplicationFilterChain.name
+            }
+            "span.type" DDSpanTypes.WEB_SERVLET
+            "http.status_code" 200
+            "servlet.context" "/$context"
+            defaultTags(true)
+          }
+        }
+      }
+      trace(depth + 1, 1) {
         span(0) {
           if (depth > 0) {
             childOf TEST_WRITER[depth - 1][0]
@@ -197,14 +220,72 @@ abstract class AbstractServlet3Test<CONTEXT> extends AgentTestRunner {
           }
         }
       }
-      trace(depth + 1, 1) {
+    }
+
+    where:
+    path        | distributedTracing | depth
+    "sync"      | true               | 0
+    "sync"      | false              | 0
+    "recursive" | true               | 0
+    "recursive" | false              | 0
+    "recursive" | true               | 1
+    "recursive" | false              | 1
+    "recursive" | true               | 20
+    "recursive" | false              | 20
+
+    origin = path.capitalize()
+  }
+
+  def "test dispatch async #path with depth #depth, distributed tracing: #distributedTracing"() {
+    setup:
+    def requestBuilder = new Request.Builder()
+      .url("http://localhost:$port/$context/dispatch/$path")
+      .get()
+    if (distributedTracing) {
+      requestBuilder.header("x-datadog-trace-id", "123")
+      requestBuilder.header("x-datadog-parent-id", "456")
+    }
+    def response = client.newCall(requestBuilder.build()).execute()
+
+    expect:
+    response.body().string().trim() == "Hello $origin"
+    assertTraces(2) {
+      // Async requests have their parent span closed before child span
+      trace(0, 1) {
+        span(0) {
+          if (distributedTracing) {
+            traceId "123"
+            parentId "456"
+          } else {
+            parent()
+          }
+          serviceName context
+          operationName "servlet.request"
+          resourceName "GET /$context/dispatch/$path"
+          spanType DDSpanTypes.WEB_SERVLET
+          errored false
+          tags {
+            "http.url" "http://localhost:$port/$context/dispatch/$path"
+            "http.method" "GET"
+            "span.kind" "server"
+            "component" "java-web-servlet"
+            "span.origin.type" { it == "TestServlet3\$Dispatch$origin" || it == ApplicationFilterChain.name }
+            "span.type" DDSpanTypes.WEB_SERVLET
+            "http.status_code" 200
+            "servlet.context" "/$context"
+            "servlet.dispatch" "/$path"
+            defaultTags(distributedTracing)
+          }
+        }
+      }
+      trace(1, 1) {
         span(0) {
           serviceName context
           operationName "servlet.request"
           resourceName "GET /$context/$path"
           spanType DDSpanTypes.WEB_SERVLET
           errored false
-          childOf TEST_WRITER[depth][0]
+          childOf TEST_WRITER[0][0]
           tags {
             "http.url" "http://localhost:$port/$context/$path"
             "http.method" "GET"
@@ -223,17 +304,9 @@ abstract class AbstractServlet3Test<CONTEXT> extends AgentTestRunner {
     }
 
     where:
-    path        | distributedTracing | depth
+    path    | distributedTracing
-    "sync"      | true               | 0
+    "async" | true
-    "sync"      | false              | 0
+    "async" | false
-    "async"     | true               | 0
-    "async"     | false              | 0
-    "recursive" | true               | 0
-    "recursive" | false              | 0
-    "recursive" | true               | 1
-    "recursive" | false              | 1
-    "recursive" | true               | 20
-    "recursive" | false              | 20
 
     origin = path.capitalize()
   }