Fix prematurely ending server span on undertow

docs/supported-libraries.md

@@ -84,6 +84,7 @@ These are the supported libraries and frameworks:
 | [Spymemcached](https://github.com/couchbase/spymemcached)                                                                             | 2.12+                          |
 | [Struts2](https://github.com/apache/struts)                                                                                           | 2.3+                           |
 | [Twilio](https://github.com/twilio/twilio-java)                                                                                       | 6.6+ (not including 8.x yet)   |
+| [Undertow](https://undertow.io/)                                                                                                      | 1.4+                           |
 | [Vert.x](https://vertx.io)                                                                                                            | 3.0+                           |
 | [Vert.x RxJava2](https://vertx.io/docs/vertx-rx/java2/)                                                                               | 3.5+                           |
 

instrumentation/undertow/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/undertow/EndSpanListener.java

@@ -8,7 +8,6 @@
 import static io.opentelemetry.javaagent.instrumentation.undertow.UndertowHttpServerTracer.tracer;
 
 import io.opentelemetry.context.Context;
-import io.undertow.server.DefaultResponseListener;
 import io.undertow.server.ExchangeCompletionListener;
 import io.undertow.server.HttpServerExchange;
 
@@ -21,12 +20,7 @@ public EndSpanListener(Context context) {
 
   @Override
   public void exchangeEvent(HttpServerExchange exchange, NextListener nextListener) {
-    Throwable throwable = exchange.getAttachment(DefaultResponseListener.EXCEPTION);
-    if (throwable != null) {
-      tracer().endExceptionally(context, throwable, exchange);
-    } else {
-      tracer().end(context, exchange);
-    }
+    tracer().exchangeCompleted(context, exchange);
     nextListener.proceed();
   }
 }

instrumentation/undertow/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/undertow/UndertowHttpServerTracer.java

@@ -11,9 +11,10 @@
 import io.opentelemetry.instrumentation.api.servlet.ServletSpanNaming;
 import io.opentelemetry.instrumentation.api.tracer.HttpServerTracer;
 import io.opentelemetry.javaagent.instrumentation.api.undertow.KeyHolder;
+import io.opentelemetry.javaagent.instrumentation.api.undertow.UndertowRequestContext;
+import io.undertow.server.DefaultResponseListener;
 import io.undertow.server.HttpServerExchange;
 import io.undertow.util.AttachmentKey;
-import java.lang.reflect.Method;
 import java.net.InetSocketAddress;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
@@ -31,17 +32,53 @@ protected String getInstrumentationName() {
     return "io.opentelemetry.javaagent.undertow";
   }
 
-  public Context startServerSpan(HttpServerExchange exchange, Method instrumentedMethod) {
+  public Context startServerSpan(HttpServerExchange exchange) {
     return startSpan(
         exchange, exchange, exchange, "HTTP " + exchange.getRequestMethod().toString());
   }
 
   @Override
   protected Context customizeContext(Context context, HttpServerExchange exchange) {
     context = ServletSpanNaming.init(context);
+    // span is ended when counter reaches 0, we start from 2 which accounts for the
+    // handler that started the span and exchange completion listener
+    context = UndertowRequestContext.init(context, 2);
     return AppServerBridge.init(context);
   }
 
+  public void handlerStarted(Context context) {
+    // request was dispatched to a new thread, handler on the original thread
+    // may exit before this one so we need to wait for this handler to complete
+    // before ending span
+    UndertowRequestContext.enter(context);
+  }
+
+  public void handlerCompleted(Context context, Throwable throwable, HttpServerExchange exchange) {
+    // end the span when this is the last handler to complete and exchange has
+    // been completed
+    if (UndertowRequestContext.exit(context)) {
+      endSpan(context, throwable, exchange);
+    }
+  }
+
+  public void exchangeCompleted(Context context, HttpServerExchange exchange) {
+    // after exchange is completed we can read response status
+    // if all handlers have completed we can end the span, if there are running
+    // handlers we'll end the span when last handler exits
+    if (UndertowRequestContext.exit(context)) {
+      Throwable throwable = exchange.getAttachment(DefaultResponseListener.EXCEPTION);
+      endSpan(context, throwable, exchange);
+    }
+  }
+
+  private void endSpan(Context context, Throwable throwable, HttpServerExchange exchange) {
+    if (throwable != null) {
+      tracer().endExceptionally(context, throwable, exchange);
+    } else {
+      tracer().end(context, exchange);
+    }
+  }
+
   @SuppressWarnings("unchecked")
   @Override
   public @Nullable Context getServerContext(HttpServerExchange exchange) {

instrumentation/undertow/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/undertow/UndertowInstrumentationModule.java

@@ -21,7 +21,6 @@
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
 import io.undertow.server.HttpServerExchange;
-import java.lang.reflect.Method;
 import java.util.List;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
@@ -33,7 +32,7 @@
 public class UndertowInstrumentationModule extends InstrumentationModule {
 
   public UndertowInstrumentationModule() {
-    super("undertow", "undertow-2.0");
+    super("undertow", "undertow-1.4");
   }
 
   @Override
@@ -65,33 +64,38 @@ public Map<? extends ElementMatcher<? super MethodDescription>, String> transfor
     public static class UndertowHandlerAdvice {
       @Advice.OnMethodEnter(suppress = Throwable.class)
       public static void onEnter(
-          @Advice.Origin Method method,
           @Advice.Argument(value = 0, readOnly = false) HttpServerExchange exchange,
+          @Advice.Local("otelContext") Context context,
           @Advice.Local("otelScope") Scope scope) {
         Context attachedContext = tracer().getServerContext(exchange);
         if (attachedContext != null) {
           if (!Java8BytecodeBridge.currentContext().equals(attachedContext)) {
+            // request processing is dispatched to another thread
             scope = attachedContext.makeCurrent();
+            context = attachedContext;
+            tracer().handlerStarted(attachedContext);
           }
           return;
         }
 
-        Context context = tracer().startServerSpan(exchange, method);
+        context = tracer().startServerSpan(exchange);
         scope = context.makeCurrent();
 
         exchange.addExchangeCompleteListener(new EndSpanListener(context));
       }
 
       @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
       public static void onExit(
-          @Advice.Argument(value = 0, readOnly = false) HttpServerExchange exchange,
+          @Advice.Argument(0) HttpServerExchange exchange,
+          @Advice.Thrown Throwable throwable,
+          @Advice.Local("otelContext") Context context,
           @Advice.Local("otelScope") Scope scope) {
         if (scope == null) {
           return;
         }
-
         scope.close();
-        // span is closed by EndSpanListener
+
+        tracer().handlerCompleted(context, throwable, exchange);
       }
     }
   }

instrumentation/undertow/javaagent/src/test/groovy/UndertowServerTest.groovy

@@ -8,13 +8,19 @@ import static io.opentelemetry.instrumentation.test.base.HttpServerTest.ServerEn
 import static io.opentelemetry.instrumentation.test.base.HttpServerTest.ServerEndpoint.QUERY_PARAM
 import static io.opentelemetry.instrumentation.test.base.HttpServerTest.ServerEndpoint.REDIRECT
 import static io.opentelemetry.instrumentation.test.base.HttpServerTest.ServerEndpoint.SUCCESS
+import static io.opentelemetry.instrumentation.test.utils.TraceUtils.basicSpan
+import static io.opentelemetry.instrumentation.test.utils.TraceUtils.runUnderTrace
 
+import io.opentelemetry.api.trace.Span
 import io.opentelemetry.instrumentation.test.AgentTestTrait
 import io.opentelemetry.instrumentation.test.base.HttpServerTest
 import io.undertow.Handlers
 import io.undertow.Undertow
 import io.undertow.util.Headers
 import io.undertow.util.StatusCodes
+import okhttp3.HttpUrl
+import okhttp3.Response
+import spock.lang.Unroll
 
 //TODO make test which mixes handlers and servlets
 class UndertowServerTest extends HttpServerTest<Undertow> implements AgentTestTrait {
@@ -52,6 +58,27 @@ class UndertowServerTest extends HttpServerTest<Undertow> implements AgentTestTr
             throw new Exception(EXCEPTION.body)
           }
         }
+        .addExactPath("sendResponse") { exchange ->
+          Span.current().addEvent("before-event")
+          runUnderTrace("sendResponse") {
+            exchange.setStatusCode(StatusCodes.OK)
+            exchange.getResponseSender().send("sendResponse")
+          }
+          // event is added only when server span has not been ended
+          // we need to make sure that sending response does not end server span
+          Span.current().addEvent("after-event")
+        }
+        .addExactPath("sendResponseWithException") { exchange ->
+          Span.current().addEvent("before-event")
+          runUnderTrace("sendResponseWithException") {
+            exchange.setStatusCode(StatusCodes.OK)
+            exchange.getResponseSender().send("sendResponseWithException")
+          }
+          // event is added only when server span has not been ended
+          // we need to make sure that sending response does not end server span
+          Span.current().addEvent("after-event")
+          throw new Exception("exception after sending response")
+        }
       ).build()
     server.start()
     return server
@@ -66,4 +93,65 @@ class UndertowServerTest extends HttpServerTest<Undertow> implements AgentTestTr
   String expectedServerSpanName(ServerEndpoint endpoint) {
     return "HTTP GET"
   }
+
+  @Unroll
+  def "test send response"() {
+    setup:
+    def url = HttpUrl.get(address.resolve("sendResponse")).newBuilder().build()
+    def request = request(url, "GET", null).build()
+    Response response = client.newCall(request).execute()
+
+    expect:
+    response.code() == 200
+    response.body().string().trim() == "sendResponse"
+
+    and:
+    assertTraces(1) {
+      trace(0, 2) {
+        it.span(0) {
+          hasNoParent()
+          name "HTTP GET"
+
+          event(0) {
+            eventName "before-event"
+          }
+          event(1) {
+            eventName "after-event"
+          }
+        }
+        basicSpan(it, 1, "sendResponse", span(0))
+      }
+    }
+  }
+
+  @Unroll
+  def "test send response with exception"() {
+    setup:
+    def url = HttpUrl.get(address.resolve("sendResponseWithException")).newBuilder().build()
+    def request = request(url, "GET", null).build()
+    Response response = client.newCall(request).execute()
+
+    expect:
+    response.code() == 200
+    response.body().string().trim() == "sendResponseWithException"
+
+    and:
+    assertTraces(1) {
+      trace(0, 2) {
+        it.span(0) {
+          hasNoParent()
+          name "HTTP GET"
+
+          event(0) {
+            eventName "before-event"
+          }
+          event(1) {
+            eventName "after-event"
+          }
+          errorEvent(Exception, "exception after sending response", 2)
+        }
+        basicSpan(it, 1, "sendResponseWithException", span(0))
+      }
+    }
+  }
 }

instrumentation/undertow/javaagent/undertow-javaagent.gradle

@@ -4,7 +4,8 @@ muzzle {
   pass {
     group = "io.undertow"
     module = 'undertow-core'
-    versions = "[2.0.0.Final,)"
+    versions = "[1.4.0.Final,)"
+    assertInverse = true
   }
 }
 

javaagent-api/src/main/java/io/opentelemetry/javaagent/instrumentation/api/undertow/UndertowRequestContext.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.api.undertow;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.ContextKey;
+import java.util.concurrent.atomic.AtomicInteger;
+
+/** Helper container for keeping track of request processing state in undertow. */
+public final class UndertowRequestContext {
+  private static final ContextKey<AtomicInteger> CONTEXT_KEY =
+      ContextKey.named("opentelemetry-undertow-context-key");
+
+  private UndertowRequestContext() {}
+
+  /**
+   * Attach to context.
+   *
+   * @param context server context
+   * @param initialValue initial value for counter
+   * @return new context
+   */
+  public static Context init(Context context, int initialValue) {
+    return context.with(CONTEXT_KEY, new AtomicInteger(initialValue));
+  }
+
+  /**
+   * Increment counter.
+   *
+   * @param context server context
+   */
+  public static void enter(Context context) {
+    context.get(CONTEXT_KEY).incrementAndGet();
+  }
+
+  /**
+   * Decrement counter.
+   *
+   * @param context server context
+   * @return true when counter has reached zero, false otherwise
+   */
+  public static boolean exit(Context context) {
+    return context.get(CONTEXT_KEY).decrementAndGet() == 0;
+  }
+}