Fix tomcat async spans (#4339)

* Add test

* Fix tomcat async spans

* Preserve existing test controller behavior

* Comments

instrumentation/dropwizard-testing/src/test/groovy/DropwizardAsyncTest.groovy

@@ -35,6 +35,12 @@ class DropwizardAsyncTest extends DropwizardTest {
     AsyncServiceResource
   }
 
+  @Override
+  boolean verifyServerSpanEndTime() {
+    // server spans are ended inside of the JAX-RS controller spans
+    return false
+  }
+
   static class AsyncTestApp extends Application<Configuration> {
     @Override
     void initialize(Bootstrap<Configuration> bootstrap) {

instrumentation/grizzly-2.0/javaagent/src/test/groovy/GrizzlyAsyncTest.groovy

@@ -38,6 +38,12 @@ class GrizzlyAsyncTest extends GrizzlyTest {
     false
   }
 
+  @Override
+  boolean verifyServerSpanEndTime() {
+    // server spans are ended inside of the JAX-RS controller spans
+    return false
+  }
+
   @Path("/")
   static class AsyncServiceResource {
     private ExecutorService executor = Executors.newSingleThreadExecutor()

instrumentation/grizzly-2.0/javaagent/src/test/groovy/GrizzlyFilterchainServerTest.groovy

@@ -70,6 +70,12 @@ class GrizzlyFilterchainServerTest extends HttpServerTest<HttpServer> implements
     false
   }
 
+  @Override
+  boolean verifyServerSpanEndTime() {
+    // server spans are ended inside of the controller spans
+    return false
+  }
+
   void setUpTransport(FilterChain filterChain) {
     TCPNIOTransportBuilder transportBuilder = TCPNIOTransportBuilder.newInstance()
       .setOptimizedForMultiplexing(true)

instrumentation/play/play-2.4/javaagent/src/test/groovy/server/PlayServerTest.groovy

@@ -86,6 +86,12 @@ class PlayServerTest extends HttpServerTest<Server> implements AgentTestTrait {
     return true
   }
 
+  @Override
+  boolean verifyServerSpanEndTime() {
+    // server spans are ended inside of the controller spans
+    return false
+  }
+
   @Override
   void handlerSpan(TraceAssert trace, int index, Object parent, String method = "GET", ServerEndpoint endpoint = SUCCESS) {
     trace.span(index) {

instrumentation/ratpack-1.4/testing/src/main/groovy/io/opentelemetry/instrumentation/ratpack/server/AbstractRatpackHttpServerTest.groovy

@@ -137,6 +137,12 @@ abstract class AbstractRatpackHttpServerTest extends HttpServerTest<RatpackServe
     true
   }
 
+  @Override
+  boolean verifyServerSpanEndTime() {
+    // server spans are ended inside of the controller spans
+    return false
+  }
+
   @Override
   void handlerSpan(TraceAssert trace, int index, Object parent, String method = "GET", ServerEndpoint endpoint = SUCCESS) {
     trace.span(index) {

instrumentation/tomcat/tomcat-7.0/javaagent/src/test/groovy/io/opentelemetry/javaagent/instrumentation/tomcat/v7_0/AsyncServlet.groovy

@@ -11,7 +11,6 @@ import io.opentelemetry.instrumentation.test.base.HttpServerTest
 import javax.servlet.annotation.WebServlet
 import javax.servlet.http.HttpServletRequest
 import javax.servlet.http.HttpServletResponse
-import java.util.concurrent.CountDownLatch
 
 import static io.opentelemetry.instrumentation.test.base.HttpServerTest.ServerEndpoint.CAPTURE_HEADERS
 import static io.opentelemetry.instrumentation.test.base.HttpServerTest.ServerEndpoint.ERROR
@@ -26,7 +25,6 @@ class AsyncServlet extends AbstractHttpServlet {
   @Override
   protected void service(HttpServletRequest req, HttpServletResponse resp) {
     HttpServerTest.ServerEndpoint endpoint = HttpServerTest.ServerEndpoint.forPath(req.servletPath)
-    def latch = new CountDownLatch(1)
     def context = req.startAsync()
     context.start {
       try {
@@ -36,44 +34,37 @@ class AsyncServlet extends AbstractHttpServlet {
             case SUCCESS:
               resp.status = endpoint.status
               resp.writer.print(endpoint.body)
-              context.complete()
               break
             case INDEXED_CHILD:
               endpoint.collectSpanAttributes { req.getParameter(it) }
               resp.status = endpoint.status
-              context.complete()
               break
             case QUERY_PARAM:
               resp.status = endpoint.status
               resp.writer.print(req.queryString)
-              context.complete()
               break
             case REDIRECT:
               resp.sendRedirect(endpoint.body)
-              context.complete()
               break
             case CAPTURE_HEADERS:
               resp.setHeader("X-Test-Response", req.getHeader("X-Test-Request"))
               resp.status = endpoint.status
               resp.writer.print(endpoint.body)
-              context.complete()
               break
             case ERROR:
               resp.status = endpoint.status
               resp.writer.print(endpoint.body)
-              context.complete()
               break
             case EXCEPTION:
               resp.status = endpoint.status
               resp.writer.print(endpoint.body)
-              context.complete()
               throw new Exception(endpoint.body)
           }
         }
       } finally {
-        latch.countDown()
+        // complete at the end so the server span will end after the controller span
+        context.complete()
       }
     }
-    latch.await()
   }
 }

instrumentation/tomcat/tomcat-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/tomcat/common/TomcatHelper.java

@@ -9,6 +9,7 @@ import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import io.opentelemetry.instrumentation.api.servlet.AppServerBridge;
+import io.opentelemetry.instrumentation.api.tracer.HttpServerTracer;
 import io.opentelemetry.javaagent.instrumentation.servlet.ServletHelper;
 import org.apache.coyote.Request;
 import org.apache.coyote.Response;
@@ -32,7 +33,9 @@ public class TomcatHelper<REQUEST, RESPONSE> {
   }
 
   public Context start(Context parentContext, Request request) {
-    return instrumenter.start(parentContext, request);
+    Context context = instrumenter.start(parentContext, request);
+    request.setAttribute(HttpServerTracer.CONTEXT_ATTRIBUTE, context);
+    return context;
   }
 
   public void end(

instrumentation/vertx-reactive-3.5/javaagent/src/latestDepTest/groovy/server/VertxRxHttpServerTest.groovy

@@ -84,6 +84,12 @@ class VertxRxHttpServerTest extends HttpServerTest<Vertx> implements AgentTestTr
     return true
   }
 
+  @Override
+  boolean verifyServerSpanEndTime() {
+    // server spans are ended inside of the controller spans
+    return false
+  }
+
   protected Class<AbstractVerticle> verticle() {
     return VertxReactiveWebServer
   }

instrumentation/vertx-reactive-3.5/javaagent/src/version35Test/groovy/server/VertxRxHttpServerTest.groovy

@@ -84,6 +84,12 @@ class VertxRxHttpServerTest extends HttpServerTest<Vertx> implements AgentTestTr
     return true
   }
 
+  @Override
+  boolean verifyServerSpanEndTime() {
+    // server spans are ended inside of the controller spans
+    return false
+  }
+
   protected Class<AbstractVerticle> verticle() {
     return VertxReactiveWebServer
   }

instrumentation/vertx-web-3.0/testing/src/main/groovy/server/AbstractVertxHttpServerTest.groovy

@@ -63,6 +63,12 @@ abstract class AbstractVertxHttpServerTest extends HttpServerTest<Vertx> impleme
     false
   }
 
+  @Override
+  boolean verifyServerSpanEndTime() {
+    // server spans are ended inside of the controller spans
+    return false
+  }
+
   @Override
   String expectedServerSpanName(ServerEndpoint endpoint) {
     switch (endpoint) {

testing-common/src/main/groovy/io/opentelemetry/instrumentation/test/base/HttpServerTest.groovy

@@ -135,6 +135,10 @@ abstract class HttpServerTest<SERVER> extends InstrumentationSpecification imple
     false
   }
 
+  boolean verifyServerSpanEndTime() {
+    return true
+  }
+
   List<AttributeKey<?>> extraAttributes() {
     []
   }
@@ -506,6 +510,11 @@ abstract class HttpServerTest<SERVER> extends InstrumentationSpecification imple
       (0..size - 1).each {
         trace(it, spanCount) {
           def spanIndex = 0
+          if (verifyServerSpanEndTime() && spanCount > 1) {
+            (1..spanCount - 1).each { index ->
+              assert it.span(0).endEpochNanos - it.span(index).endEpochNanos >= 0
+            }
+          }
           serverSpan(it, spanIndex++, traceID, parentID, method, response?.content()?.length(), endpoint)
           if (hasHandlerSpan(endpoint)) {
             handlerSpan(it, spanIndex++, span(0), method, endpoint)