From d6b903665e480daff325d7d24b462507f9cbcf8a Mon Sep 17 00:00:00 2001
From: Tyler Benson <tyler.benson@datadoghq.com>
Date: Fri, 26 Jul 2019 16:30:23 -0700
Subject: [PATCH] Copy servlet listener over to new context on dispatch.

This allows it to finish the span with the correct response.
---
 .../servlet3/TagSettingAsyncListener.java     |  4 +-
 .../src/test/groovy/JettyServlet3Test.groovy  | 37 +++++++++----------
 .../src/test/groovy/TomcatServlet3Test.groovy | 37 +++++++++----------
 3 files changed, 38 insertions(+), 40 deletions(-)

diff --git a/dd-java-agent/instrumentation/servlet-3/src/main/java/datadog/trace/instrumentation/servlet3/TagSettingAsyncListener.java b/dd-java-agent/instrumentation/servlet-3/src/main/java/datadog/trace/instrumentation/servlet3/TagSettingAsyncListener.java
index 98517702b5..6d40f819b4 100644
--- a/dd-java-agent/instrumentation/servlet-3/src/main/java/datadog/trace/instrumentation/servlet3/TagSettingAsyncListener.java
+++ b/dd-java-agent/instrumentation/servlet-3/src/main/java/datadog/trace/instrumentation/servlet3/TagSettingAsyncListener.java
@@ -58,9 +58,9 @@ public class TagSettingAsyncListener implements AsyncListener {
     }
   }
 
-  /** Finish current span on dispatch. New listener will be attached by Servlet3Advice */
+  /** Transfer the listener over to the new context. */
   @Override
   public void onStartAsync(final AsyncEvent event) throws IOException {
-    onComplete(event);
+    event.getAsyncContext().addListener(this);
   }
 }
diff --git a/dd-java-agent/instrumentation/servlet-3/src/test/groovy/JettyServlet3Test.groovy b/dd-java-agent/instrumentation/servlet-3/src/test/groovy/JettyServlet3Test.groovy
index 199179740b..4133ea0edc 100644
--- a/dd-java-agent/instrumentation/servlet-3/src/test/groovy/JettyServlet3Test.groovy
+++ b/dd-java-agent/instrumentation/servlet-3/src/test/groovy/JettyServlet3Test.groovy
@@ -137,25 +137,24 @@ class JettyServlet3TestDispatchImmediate extends JettyDispatchTest {
   }
 }
 
-// TODO: Behavior in this test is pretty inconsistent with expectations. Fix and reenable.
-//class JettyServlet3TestDispatchAsync extends JettyDispatchTest {
-//  @Override
-//  Class<Servlet> servlet() {
-//    TestServlet3.Async
-//  }
-//
-//  @Override
-//  protected void setupServlets(ServletContextHandler context) {
-//    super.setupServlets(context)
-//
-//    addServlet(context, "/dispatch" + SUCCESS.path, TestServlet3.DispatchAsync)
-//    addServlet(context, "/dispatch" + ERROR.path, TestServlet3.DispatchAsync)
-//    addServlet(context, "/dispatch" + EXCEPTION.path, TestServlet3.DispatchAsync)
-//    addServlet(context, "/dispatch" + REDIRECT.path, TestServlet3.DispatchAsync)
-//    addServlet(context, "/dispatch" + AUTH_REQUIRED.path, TestServlet3.DispatchAsync)
-//    addServlet(context, "/dispatch/recursive", TestServlet3.DispatchRecursive)
-//  }
-//}
+class JettyServlet3TestDispatchAsync extends JettyDispatchTest {
+  @Override
+  Class<Servlet> servlet() {
+    TestServlet3.Async
+  }
+
+  @Override
+  protected void setupServlets(ServletContextHandler context) {
+    super.setupServlets(context)
+
+    addServlet(context, "/dispatch" + SUCCESS.path, TestServlet3.DispatchAsync)
+    addServlet(context, "/dispatch" + ERROR.path, TestServlet3.DispatchAsync)
+    addServlet(context, "/dispatch" + EXCEPTION.path, TestServlet3.DispatchAsync)
+    addServlet(context, "/dispatch" + REDIRECT.path, TestServlet3.DispatchAsync)
+    addServlet(context, "/dispatch" + AUTH_REQUIRED.path, TestServlet3.DispatchAsync)
+    addServlet(context, "/dispatch/recursive", TestServlet3.DispatchRecursive)
+  }
+}
 
 abstract class JettyDispatchTest extends JettyServlet3Test {
   @Override
diff --git a/dd-java-agent/instrumentation/servlet-3/src/test/groovy/TomcatServlet3Test.groovy b/dd-java-agent/instrumentation/servlet-3/src/test/groovy/TomcatServlet3Test.groovy
index 5b91cd5126..9ededd2a1e 100644
--- a/dd-java-agent/instrumentation/servlet-3/src/test/groovy/TomcatServlet3Test.groovy
+++ b/dd-java-agent/instrumentation/servlet-3/src/test/groovy/TomcatServlet3Test.groovy
@@ -178,25 +178,24 @@ class TomcatServlet3TestDispatchImmediate extends TomcatDispatchTest {
   }
 }
 
-// TODO: Behavior in this test is pretty inconsistent with expectations. Fix and reenable.
-//class TomcatServlet3TestDispatchAsync extends TomcatDispatchTest {
-//  @Override
-//  Class<Servlet> servlet() {
-//    TestServlet3.Async
-//  }
-//
-//  @Override
-//  protected void setupServlets(Context context) {
-//    super.setupServlets(context)
-//
-//    addServlet(context, "/dispatch" + SUCCESS.path, TestServlet3.DispatchAsync)
-//    addServlet(context, "/dispatch" + ERROR.path, TestServlet3.DispatchAsync)
-//    addServlet(context, "/dispatch" + EXCEPTION.path, TestServlet3.DispatchAsync)
-//    addServlet(context, "/dispatch" + REDIRECT.path, TestServlet3.DispatchAsync)
-//    addServlet(context, "/dispatch" + AUTH_REQUIRED.path, TestServlet3.DispatchAsync)
-//    addServlet(context, "/dispatch/recursive", TestServlet3.DispatchRecursive)
-//  }
-//}
+class TomcatServlet3TestDispatchAsync extends TomcatDispatchTest {
+  @Override
+  Class<Servlet> servlet() {
+    TestServlet3.Async
+  }
+
+  @Override
+  protected void setupServlets(Context context) {
+    super.setupServlets(context)
+
+    addServlet(context, "/dispatch" + SUCCESS.path, TestServlet3.DispatchAsync)
+    addServlet(context, "/dispatch" + ERROR.path, TestServlet3.DispatchAsync)
+    addServlet(context, "/dispatch" + EXCEPTION.path, TestServlet3.DispatchAsync)
+    addServlet(context, "/dispatch" + REDIRECT.path, TestServlet3.DispatchAsync)
+    addServlet(context, "/dispatch" + AUTH_REQUIRED.path, TestServlet3.DispatchAsync)
+    addServlet(context, "/dispatch/recursive", TestServlet3.DispatchRecursive)
+  }
+}
 
 abstract class TomcatDispatchTest extends TomcatServlet3Test {
   @Override