From 6cd8be67d52102b505b56aa5c3f621bc93d2e27d Mon Sep 17 00:00:00 2001
From: Andrew Kent <andrew.kent@datadoghq.com>
Date: Thu, 14 Jun 2018 20:16:16 -0700
Subject: [PATCH] Close scope for async servlet requests

---
 .../servlet3/FilterChain3Instrumentation.java |  1 +
 .../servlet3/HttpServlet3Instrumentation.java |  1 +
 .../src/test/groovy/JettyServlet3Test.groovy  | 25 +++++++++++++++++++
 .../trace/common/writer/ListWriter.java       |  3 ++-
 4 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/dd-java-agent/instrumentation/servlet-3/src/main/java/datadog/trace/instrumentation/servlet3/FilterChain3Instrumentation.java b/dd-java-agent/instrumentation/servlet-3/src/main/java/datadog/trace/instrumentation/servlet3/FilterChain3Instrumentation.java
index 802603ed6a..ed78a1b23d 100644
--- a/dd-java-agent/instrumentation/servlet-3/src/main/java/datadog/trace/instrumentation/servlet3/FilterChain3Instrumentation.java
+++ b/dd-java-agent/instrumentation/servlet-3/src/main/java/datadog/trace/instrumentation/servlet3/FilterChain3Instrumentation.java
@@ -128,6 +128,7 @@ public final class FilterChain3Instrumentation extends Instrumenter.Configurable
             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());
             scope.close();
diff --git a/dd-java-agent/instrumentation/servlet-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServlet3Instrumentation.java b/dd-java-agent/instrumentation/servlet-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServlet3Instrumentation.java
index 4aa78bdc04..229b728a81 100644
--- a/dd-java-agent/instrumentation/servlet-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServlet3Instrumentation.java
+++ b/dd-java-agent/instrumentation/servlet-3/src/main/java/datadog/trace/instrumentation/servlet3/HttpServlet3Instrumentation.java
@@ -120,6 +120,7 @@ public final class HttpServlet3Instrumentation extends Instrumenter.Configurable
           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());
           scope.close();
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 a581a229c6..ab3a373f1f 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
@@ -122,6 +122,31 @@ class JettyServlet3Test extends AgentTestRunner {
     "sync"  | "Hello Sync"
   }
 
+  def "servlet instrumentation clears state after async request"() {
+    setup:
+    def request = new Request.Builder()
+      .url("http://localhost:$PORT/async")
+      .get()
+      .build()
+    def numTraces = 5
+    for (int i = 0; i < numTraces; ++i) {
+      client.newCall(request).execute()
+    }
+
+    expect:
+    assertTraces(writer, numTraces) {
+      for (int i = 0; i < numTraces; ++i) {
+        trace(i, 1) {
+          span(0) {
+            serviceName "unnamed-java-app"
+            operationName "servlet.request"
+            resourceName "GET /async"
+          }
+        }
+      }
+    }
+  }
+
   def "test #path error servlet call"() {
     setup:
     def request = new Request.Builder()
diff --git a/dd-trace-ot/src/main/java/datadog/trace/common/writer/ListWriter.java b/dd-trace-ot/src/main/java/datadog/trace/common/writer/ListWriter.java
index 9ea0a37966..d60ac9bb66 100644
--- a/dd-trace-ot/src/main/java/datadog/trace/common/writer/ListWriter.java
+++ b/dd-trace-ot/src/main/java/datadog/trace/common/writer/ListWriter.java
@@ -39,7 +39,8 @@ public class ListWriter extends CopyOnWriteArrayList<List<DDSpan>> implements Wr
       latches.add(latch);
     }
     if (!latch.await(20, TimeUnit.SECONDS)) {
-      throw new TimeoutException("Timeout waiting for " + number + " trace(s).");
+      throw new TimeoutException(
+          "Timeout waiting for " + number + " trace(s). ListWriter.size() == " + size());
     }
   }