From 233d241e07ac2b47aa7c1a93b05da74cbfb39f9a Mon Sep 17 00:00:00 2001
From: Tyler Benson <tyler@newrelic.com>
Date: Wed, 23 Aug 2017 16:01:42 -0700
Subject: [PATCH] Mark servlet trace filter as async supported

---
 .../trace/agent/integration/JettyServletHelper.java   | 11 ++++++-----
 .../trace/agent/integration/TomcatServletHelper.java  |  9 +++++----
 .../dropwizard/resources/SimpleCrudResource.java      |  9 +++++++++
 3 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/dd-java-agent/src/main/java/com/datadoghq/trace/agent/integration/JettyServletHelper.java b/dd-java-agent/src/main/java/com/datadoghq/trace/agent/integration/JettyServletHelper.java
index 318d69fa9a..e24c8e6297 100644
--- a/dd-java-agent/src/main/java/com/datadoghq/trace/agent/integration/JettyServletHelper.java
+++ b/dd-java-agent/src/main/java/com/datadoghq/trace/agent/integration/JettyServletHelper.java
@@ -3,6 +3,7 @@ package com.datadoghq.trace.agent.integration;
 import io.opentracing.contrib.web.servlet.filter.TracingFilter;
 import java.util.EnumSet;
 import javax.servlet.Filter;
+import javax.servlet.FilterRegistration;
 import org.eclipse.jetty.servlet.ServletContextHandler;
 import org.jboss.byteman.rule.Rule;
 
@@ -25,11 +26,11 @@ public class JettyServletHelper extends DDAgentTracingHelper<ServletContextHandl
     String[] patterns = {"/*"};
 
     Filter filter = new TracingFilter(tracer);
-    contextHandler
-        .getServletContext()
-        .addFilter("tracingFilter", filter)
-        .addMappingForUrlPatterns(
-            EnumSet.allOf(javax.servlet.DispatcherType.class), true, patterns);
+    FilterRegistration.Dynamic registration =
+        contextHandler.getServletContext().addFilter("tracingFilter", filter);
+    registration.setAsyncSupported(true);
+    registration.addMappingForUrlPatterns(
+        EnumSet.allOf(javax.servlet.DispatcherType.class), true, patterns);
 
     setState(contextHandler.getServletContext(), 1);
     return contextHandler;
diff --git a/dd-java-agent/src/main/java/com/datadoghq/trace/agent/integration/TomcatServletHelper.java b/dd-java-agent/src/main/java/com/datadoghq/trace/agent/integration/TomcatServletHelper.java
index 2508aae843..19c0451539 100644
--- a/dd-java-agent/src/main/java/com/datadoghq/trace/agent/integration/TomcatServletHelper.java
+++ b/dd-java-agent/src/main/java/com/datadoghq/trace/agent/integration/TomcatServletHelper.java
@@ -3,6 +3,7 @@ package com.datadoghq.trace.agent.integration;
 import io.opentracing.contrib.web.servlet.filter.TracingFilter;
 import java.util.EnumSet;
 import javax.servlet.Filter;
+import javax.servlet.FilterRegistration;
 import org.apache.catalina.core.ApplicationContext;
 import org.jboss.byteman.rule.Rule;
 
@@ -25,10 +26,10 @@ public class TomcatServletHelper extends DDAgentTracingHelper<ApplicationContext
     String[] patterns = {"/*"};
 
     Filter filter = new TracingFilter(tracer);
-    contextHandler
-        .addFilter("tracingFilter", filter)
-        .addMappingForUrlPatterns(
-            EnumSet.allOf(javax.servlet.DispatcherType.class), true, patterns);
+    FilterRegistration.Dynamic registration = contextHandler.addFilter("tracingFilter", filter);
+    registration.setAsyncSupported(true);
+    registration.addMappingForUrlPatterns(
+        EnumSet.allOf(javax.servlet.DispatcherType.class), true, patterns);
 
     return contextHandler;
   }
diff --git a/dd-trace-examples/dropwizard-mongo-client/src/main/java/com/datadoghq/example/dropwizard/resources/SimpleCrudResource.java b/dd-trace-examples/dropwizard-mongo-client/src/main/java/com/datadoghq/example/dropwizard/resources/SimpleCrudResource.java
index 27855bd8d2..7473e6e1c9 100644
--- a/dd-trace-examples/dropwizard-mongo-client/src/main/java/com/datadoghq/example/dropwizard/resources/SimpleCrudResource.java
+++ b/dd-trace-examples/dropwizard-mongo-client/src/main/java/com/datadoghq/example/dropwizard/resources/SimpleCrudResource.java
@@ -14,6 +14,8 @@ import javax.ws.rs.GET;
 import javax.ws.rs.Path;
 import javax.ws.rs.Produces;
 import javax.ws.rs.QueryParam;
+import javax.ws.rs.container.AsyncResponse;
+import javax.ws.rs.container.Suspended;
 import javax.ws.rs.core.MediaType;
 import org.bson.Document;
 
@@ -96,6 +98,13 @@ public class SimpleCrudResource {
     return books;
   }
 
+  @GET
+  @Path("/async")
+  public void async(@Suspended AsyncResponse response) {
+    // not actually async, but useful for testing that codepath.
+    response.resume("Returned from async");
+  }
+
   /**
    * The beforeDB is traced using the annotation @Trace with a custom operationName and a custom
    * tag.