Fix ending server span with servlet async request

instrumentation/servlet/servlet-3.0/testing/src/test/java/io/opentelemetry/javaagent/instrumentation/servlet/v3_0/TestServlet3.java

@@ -17,6 +17,7 @@
 import static io.opentelemetry.javaagent.instrumentation.servlet.v3_0.AbstractServlet3Test.HTML_PRINT_WRITER;
 import static io.opentelemetry.javaagent.instrumentation.servlet.v3_0.AbstractServlet3Test.HTML_SERVLET_OUTPUT_STREAM;
 
+import io.opentelemetry.instrumentation.testing.GlobalTraceUtil;
 import io.opentelemetry.instrumentation.testing.junit.http.ServerEndpoint;
 import java.io.IOException;
 import java.io.PrintWriter;
@@ -100,7 +101,12 @@ public static class Async extends HttpServlet {
     protected void service(HttpServletRequest req, HttpServletResponse resp) {
       ServerEndpoint endpoint = ServerEndpoint.forPath(req.getServletPath());
       CountDownLatch latch = new CountDownLatch(1);
-      AsyncContext context = req.startAsync();
+      boolean startAsyncInSpan =
+          SUCCESS.equals(endpoint) && "true".equals(req.getParameter("startAsyncInSpan"));
+      AsyncContext context =
+          startAsyncInSpan
+              ? GlobalTraceUtil.runWithSpan("startAsync", () -> req.startAsync())
+              : req.startAsync();
       if (endpoint.equals(EXCEPTION)) {
         context.setTimeout(5000);
       }

instrumentation/servlet/servlet-3.0/testing/src/test/java/io/opentelemetry/javaagent/instrumentation/servlet/v3_0/jetty/JettyServlet3AsyncTest.java

@@ -5,8 +5,16 @@
 
 package io.opentelemetry.javaagent.instrumentation.servlet.v3_0.jetty;
 
+import static io.opentelemetry.instrumentation.testing.junit.http.ServerEndpoint.SUCCESS;
+import static org.assertj.core.api.Assertions.assertThat;
+
+import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.javaagent.instrumentation.servlet.v3_0.TestServlet3;
+import io.opentelemetry.testing.internal.armeria.common.AggregatedHttpRequest;
+import io.opentelemetry.testing.internal.armeria.common.AggregatedHttpResponse;
+import io.opentelemetry.testing.internal.armeria.common.HttpMethod;
 import javax.servlet.Servlet;
+import org.junit.jupiter.api.Test;
 
 class JettyServlet3AsyncTest extends JettyServlet3Test {
 
@@ -24,4 +32,32 @@ public boolean errorEndpointUsesSendError() {
   public boolean isAsyncTest() {
     return true;
   }
+
+  @Test
+  void startAsyncInSpan() {
+    AggregatedHttpRequest request =
+        AggregatedHttpRequest.of(
+            HttpMethod.GET, resolveAddress(SUCCESS, "h1c://") + "?startAsyncInSpan=true");
+    AggregatedHttpResponse response = client.execute(request).aggregate().join();
+
+    assertThat(response.status().code()).isEqualTo(SUCCESS.getStatus());
+    assertThat(response.contentUtf8()).isEqualTo(SUCCESS.getBody());
+
+    testing()
+        .waitAndAssertTraces(
+            trace ->
+                trace.hasSpansSatisfyingExactly(
+                    span ->
+                        span.hasName("GET " + getContextPath() + "/success")
+                            .hasKind(SpanKind.SERVER)
+                            .hasNoParent(),
+                    span ->
+                        span.hasName("startAsync")
+                            .hasKind(SpanKind.INTERNAL)
+                            .hasParent(trace.getSpan(0)),
+                    span ->
+                        span.hasName("controller")
+                            .hasKind(SpanKind.INTERNAL)
+                            .hasParent(trace.getSpan(0))));
+  }
 }

instrumentation/servlet/servlet-3.0/testing/src/test/java/io/opentelemetry/javaagent/instrumentation/servlet/v3_0/tomcat/TomcatServlet3AsyncTest.java

@@ -5,8 +5,16 @@
 
 package io.opentelemetry.javaagent.instrumentation.servlet.v3_0.tomcat;
 
+import static io.opentelemetry.instrumentation.testing.junit.http.ServerEndpoint.SUCCESS;
+import static org.assertj.core.api.Assertions.assertThat;
+
+import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.javaagent.instrumentation.servlet.v3_0.TestServlet3;
+import io.opentelemetry.testing.internal.armeria.common.AggregatedHttpRequest;
+import io.opentelemetry.testing.internal.armeria.common.AggregatedHttpResponse;
+import io.opentelemetry.testing.internal.armeria.common.HttpMethod;
 import javax.servlet.Servlet;
+import org.junit.jupiter.api.Test;
 
 class TomcatServlet3AsyncTest extends TomcatServlet3Test {
 
@@ -19,4 +27,32 @@ public Class<? extends Servlet> servlet() {
   public boolean errorEndpointUsesSendError() {
     return false;
   }
+
+  @Test
+  void startAsyncInSpan() {
+    AggregatedHttpRequest request =
+        AggregatedHttpRequest.of(
+            HttpMethod.GET, resolveAddress(SUCCESS, "h1c://") + "?startAsyncInSpan=true");
+    AggregatedHttpResponse response = client.execute(request).aggregate().join();
+
+    assertThat(response.status().code()).isEqualTo(SUCCESS.getStatus());
+    assertThat(response.contentUtf8()).isEqualTo(SUCCESS.getBody());
+
+    testing()
+        .waitAndAssertTraces(
+            trace ->
+                trace.hasSpansSatisfyingExactly(
+                    span ->
+                        span.hasName("GET " + getContextPath() + "/success")
+                            .hasKind(SpanKind.SERVER)
+                            .hasNoParent(),
+                    span ->
+                        span.hasName("startAsync")
+                            .hasKind(SpanKind.INTERNAL)
+                            .hasParent(trace.getSpan(0)),
+                    span ->
+                        span.hasName("controller")
+                            .hasKind(SpanKind.INTERNAL)
+                            .hasParent(trace.getSpan(0))));
+  }
 }

instrumentation/servlet/servlet-5.0/jetty11-testing/src/test/java/io/opentelemetry/javaagent/instrumentation/servlet/v5_0/jetty/JettyServlet5AsyncTest.java

@@ -5,8 +5,16 @@
 
 package io.opentelemetry.javaagent.instrumentation.servlet.v5_0.jetty;
 
+import static io.opentelemetry.instrumentation.testing.junit.http.ServerEndpoint.SUCCESS;
+import static org.assertj.core.api.Assertions.assertThat;
+
+import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.javaagent.instrumentation.servlet.v5_0.TestServlet5;
+import io.opentelemetry.testing.internal.armeria.common.AggregatedHttpRequest;
+import io.opentelemetry.testing.internal.armeria.common.AggregatedHttpResponse;
+import io.opentelemetry.testing.internal.armeria.common.HttpMethod;
 import jakarta.servlet.Servlet;
+import org.junit.jupiter.api.Test;
 
 class JettyServlet5AsyncTest extends JettyServlet5Test {
 
@@ -19,4 +27,32 @@ public Class<? extends Servlet> servlet() {
   public boolean errorEndpointUsesSendError() {
     return false;
   }
+
+  @Test
+  void startAsyncInSpan() {
+    AggregatedHttpRequest request =
+        AggregatedHttpRequest.of(
+            HttpMethod.GET, resolveAddress(SUCCESS, "h1c://") + "?startAsyncInSpan=true");
+    AggregatedHttpResponse response = client.execute(request).aggregate().join();
+
+    assertThat(response.status().code()).isEqualTo(SUCCESS.getStatus());
+    assertThat(response.contentUtf8()).isEqualTo(SUCCESS.getBody());
+
+    testing()
+        .waitAndAssertTraces(
+            trace ->
+                trace.hasSpansSatisfyingExactly(
+                    span ->
+                        span.hasName("GET " + getContextPath() + "/success")
+                            .hasKind(SpanKind.SERVER)
+                            .hasNoParent(),
+                    span ->
+                        span.hasName("startAsync")
+                            .hasKind(SpanKind.INTERNAL)
+                            .hasParent(trace.getSpan(0)),
+                    span ->
+                        span.hasName("controller")
+                            .hasKind(SpanKind.INTERNAL)
+                            .hasParent(trace.getSpan(0))));
+  }
 }

instrumentation/servlet/servlet-5.0/jetty12-testing/src/test/java/io/opentelemetry/javaagent/instrumentation/servlet/v5_0/jetty12/Jetty12Servlet5AsyncTest.java

@@ -5,8 +5,16 @@
 
 package io.opentelemetry.javaagent.instrumentation.servlet.v5_0.jetty12;
 
+import static io.opentelemetry.instrumentation.testing.junit.http.ServerEndpoint.SUCCESS;
+import static org.assertj.core.api.Assertions.assertThat;
+
+import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.javaagent.instrumentation.servlet.v5_0.TestServlet5;
+import io.opentelemetry.testing.internal.armeria.common.AggregatedHttpRequest;
+import io.opentelemetry.testing.internal.armeria.common.AggregatedHttpResponse;
+import io.opentelemetry.testing.internal.armeria.common.HttpMethod;
 import jakarta.servlet.Servlet;
+import org.junit.jupiter.api.Test;
 
 class Jetty12Servlet5AsyncTest extends Jetty12Servlet5Test {
 
@@ -19,4 +27,32 @@ public Class<? extends Servlet> servlet() {
   public boolean errorEndpointUsesSendError() {
     return false;
   }
+
+  @Test
+  void startAsyncInSpan() {
+    AggregatedHttpRequest request =
+        AggregatedHttpRequest.of(
+            HttpMethod.GET, resolveAddress(SUCCESS, "h1c://") + "?startAsyncInSpan=true");
+    AggregatedHttpResponse response = client.execute(request).aggregate().join();
+
+    assertThat(response.status().code()).isEqualTo(SUCCESS.getStatus());
+    assertThat(response.contentUtf8()).isEqualTo(SUCCESS.getBody());
+
+    testing()
+        .waitAndAssertTraces(
+            trace ->
+                trace.hasSpansSatisfyingExactly(
+                    span ->
+                        span.hasName("GET " + getContextPath() + "/success")
+                            .hasKind(SpanKind.SERVER)
+                            .hasNoParent(),
+                    span ->
+                        span.hasName("startAsync")
+                            .hasKind(SpanKind.INTERNAL)
+                            .hasParent(trace.getSpan(0)),
+                    span ->
+                        span.hasName("controller")
+                            .hasKind(SpanKind.INTERNAL)
+                            .hasParent(trace.getSpan(0))));
+  }
 }

instrumentation/servlet/servlet-5.0/testing/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v5_0/TestServlet5.java

@@ -17,6 +17,7 @@
 import static io.opentelemetry.javaagent.instrumentation.servlet.v5_0.AbstractServlet5Test.HTML_PRINT_WRITER;
 import static io.opentelemetry.javaagent.instrumentation.servlet.v5_0.AbstractServlet5Test.HTML_SERVLET_OUTPUT_STREAM;
 
+import io.opentelemetry.instrumentation.testing.GlobalTraceUtil;
 import io.opentelemetry.instrumentation.testing.junit.http.ServerEndpoint;
 import jakarta.servlet.AsyncContext;
 import jakarta.servlet.RequestDispatcher;
@@ -100,7 +101,12 @@ public static class Async extends HttpServlet {
     protected void service(HttpServletRequest req, HttpServletResponse resp) {
       ServerEndpoint endpoint = ServerEndpoint.forPath(req.getServletPath());
       CountDownLatch latch = new CountDownLatch(1);
-      AsyncContext context = req.startAsync();
+      boolean startAsyncInSpan =
+          SUCCESS.equals(endpoint) && "true".equals(req.getParameter("startAsyncInSpan"));
+      AsyncContext context =
+          startAsyncInSpan
+              ? GlobalTraceUtil.runWithSpan("startAsync", () -> req.startAsync())
+              : req.startAsync();
       if (endpoint.equals(EXCEPTION)) {
         context.setTimeout(5000);
       }

instrumentation/servlet/servlet-5.0/tomcat-testing/src/test/java/io/opentelemetry/javaagent/instrumentation/servlet/v5_0/tomcat/TomcatServlet5AsyncTest.java

@@ -5,8 +5,16 @@
 
 package io.opentelemetry.javaagent.instrumentation.servlet.v5_0.tomcat;
 
+import static io.opentelemetry.instrumentation.testing.junit.http.ServerEndpoint.SUCCESS;
+import static org.assertj.core.api.Assertions.assertThat;
+
+import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.javaagent.instrumentation.servlet.v5_0.TestServlet5;
+import io.opentelemetry.testing.internal.armeria.common.AggregatedHttpRequest;
+import io.opentelemetry.testing.internal.armeria.common.AggregatedHttpResponse;
+import io.opentelemetry.testing.internal.armeria.common.HttpMethod;
 import jakarta.servlet.Servlet;
+import org.junit.jupiter.api.Test;
 
 class TomcatServlet5AsyncTest extends TomcatServlet5Test {
 
@@ -19,4 +27,32 @@ public Class<? extends Servlet> servlet() {
   public boolean errorEndpointUsesSendError() {
     return false;
   }
+
+  @Test
+  void startAsyncInSpan() {
+    AggregatedHttpRequest request =
+        AggregatedHttpRequest.of(
+            HttpMethod.GET, resolveAddress(SUCCESS, "h1c://") + "?startAsyncInSpan=true");
+    AggregatedHttpResponse response = client.execute(request).aggregate().join();
+
+    assertThat(response.status().code()).isEqualTo(SUCCESS.getStatus());
+    assertThat(response.contentUtf8()).isEqualTo(SUCCESS.getBody());
+
+    testing()
+        .waitAndAssertTraces(
+            trace ->
+                trace.hasSpansSatisfyingExactly(
+                    span ->
+                        span.hasName("GET " + getContextPath() + "/success")
+                            .hasKind(SpanKind.SERVER)
+                            .hasNoParent(),
+                    span ->
+                        span.hasName("startAsync")
+                            .hasKind(SpanKind.INTERNAL)
+                            .hasParent(trace.getSpan(0)),
+                    span ->
+                        span.hasName("controller")
+                            .hasKind(SpanKind.INTERNAL)
+                            .hasParent(trace.getSpan(0))));
+  }
 }

instrumentation/servlet/servlet-common/bootstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/servlet/ServletAsyncContext.java

@@ -19,6 +19,7 @@ public class ServletAsyncContext implements ImplicitContextKeyed {
   private boolean isAsyncListenerAttached;
   private Throwable throwable;
   private Object response;
+  private Context context;
 
   public static Context init(Context context) {
     if (context.get(CONTEXT_KEY) != null) {
@@ -61,13 +62,22 @@ public static Object getAsyncListenerResponse(@Nullable Context context) {
     return servletAsyncContext != null ? servletAsyncContext.response : null;
   }
 
-  public static void setAsyncListenerResponse(@Nullable Context context, Object response) {
+  public static void setAsyncListenerResponse(Context context, Object response) {
     ServletAsyncContext servletAsyncContext = get(context);
     if (servletAsyncContext != null) {
       servletAsyncContext.response = response;
+      servletAsyncContext.context = context;
     }
   }
 
+  public static Context getAsyncListenerContext(Context context) {
+    ServletAsyncContext servletAsyncContext = get(context);
+    if (servletAsyncContext != null) {
+      return servletAsyncContext.context;
+    }
+    return null;
+  }
+
   @Override
   public Context storeInContext(Context context) {
     return context.with(CONTEXT_KEY, this);

instrumentation/servlet/servlet-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/AsyncRequestCompletionListener.java

@@ -23,10 +23,14 @@ public AsyncRequestCompletionListener(
       Instrumenter<ServletRequestContext<REQUEST>, ServletResponseContext<RESPONSE>> instrumenter,
       ServletRequestContext<REQUEST> requestContext,
       Context context) {
+    // The context passed into this method may contain other spans besides the server span. To end
+    // the server span we get the context that set at the start of the request with
+    // ServletHelper#setAsyncListenerResponse that contains just the server span.
+    Context serverSpanContext = servletHelper.getAsyncListenerContext(context);
     this.servletHelper = servletHelper;
     this.instrumenter = instrumenter;
     this.requestContext = requestContext;
-    this.context = context;
+    this.context = serverSpanContext != null ? serverSpanContext : context;
   }
 
   @Override

instrumentation/servlet/servlet-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/ServletHelper.java

@@ -120,4 +120,8 @@ public void recordAsyncException(Context context, Throwable throwable) {
   public Throwable getAsyncException(Context context) {
     return ServletAsyncContext.getAsyncException(context);
   }
+
+  public Context getAsyncListenerContext(Context context) {
+    return ServletAsyncContext.getAsyncListenerContext(context);
+  }
 }