Context propagation for undertow async dispatch
ADD
instrumentation/undertow-1.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/undertow/HttpServerExchangeInstrumentation.java
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.undertow;
+
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.field.VirtualField;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
+import io.opentelemetry.javaagent.instrumentation.api.concurrent.ExecutorAdviceHelper;
+import io.opentelemetry.javaagent.instrumentation.api.concurrent.PropagatedContext;
+import java.util.concurrent.Executor;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class HttpServerExchangeInstrumentation implements TypeInstrumentation {
+
+ @Override
+ public ElementMatcher<TypeDescription> typeMatcher() {
+ return named("io.undertow.server.HttpServerExchange");
+ }
+
+ @Override
+ public void transform(TypeTransformer transformer) {
+ transformer.applyAdviceToMethod(
+ named("dispatch").and(takesArguments(Executor.class, Runnable.class)),
+ this.getClass().getName() + "$DispatchAdvice");
+ }
+
+ @SuppressWarnings("unused")
+ public static class DispatchAdvice {
+
+ @Advice.OnMethodEnter(suppress = Throwable.class)
+ public static PropagatedContext enterJobSubmit(
+ @Advice.Argument(value = 1, readOnly = false) Runnable task) {
+ Context context = Java8BytecodeBridge.currentContext();
+ if (ExecutorAdviceHelper.shouldPropagateContext(context, task)) {
+ VirtualField<Runnable, PropagatedContext> virtualField =
+ VirtualField.find(Runnable.class, PropagatedContext.class);
+ return ExecutorAdviceHelper.attachContextToTask(context, virtualField, task);
+ }
+ return null;
+ }
+
+ @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+ public static void exitJobSubmit(
+ @Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
+ ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
+ }
+ }
+}

MODIFY
instrumentation/undertow-1.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/undertow/UndertowInstrumentationModule.java
instrumentation/undertow-1.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/undertow/UndertowInstrumentationModule.java
package io.opentelemetry.javaagent.instrumentation.undertow;
-import static java.util.Collections.singletonList;
+import static java.util.Arrays.asList;
import com.google.auto.service.AutoService;
import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
@@ -21,6 +21,6 @@ public class UndertowInstrumentationModule extends InstrumentationModule {
@Override
public List<TypeInstrumentation> typeInstrumentations() {
- return singletonList(new HandlerInstrumentation());
+ return asList(new HandlerInstrumentation(), new HttpServerExchangeInstrumentation());
}
}

ADD
instrumentation/undertow-1.4/javaagent/src/test/groovy/UndertowServerDispatchTest.groovy
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import io.opentelemetry.api.common.AttributeKey
+import io.opentelemetry.instrumentation.test.AgentTestTrait
+import io.opentelemetry.instrumentation.test.base.HttpServerTest
+import io.opentelemetry.semconv.trace.attributes.SemanticAttributes
+import io.undertow.Handlers
+import io.undertow.Undertow
+import io.undertow.util.Headers
+import io.undertow.util.HttpString
+import io.undertow.util.StatusCodes
+
+import static io.opentelemetry.instrumentation.test.base.HttpServerTest.ServerEndpoint.CAPTURE_HEADERS
+import static io.opentelemetry.instrumentation.test.base.HttpServerTest.ServerEndpoint.ERROR
+import static io.opentelemetry.instrumentation.test.base.HttpServerTest.ServerEndpoint.INDEXED_CHILD
+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
+
+class UndertowServerDispatchTest extends HttpServerTest<Undertow> implements AgentTestTrait {
+
+ @Override
+ boolean verifyServerSpanEndTime() {
+ return false
+ }
+
+ @Override
+ boolean testException() {
+ // throwing exception from dispatched task just makes the request time out
+ return false
+ }
+
+ @Override
+ Undertow startServer(int port) {
+ Undertow server = Undertow.builder()
+ .addHttpListener(port, "localhost")
+ .setHandler(Handlers.path()
+ .addExactPath(SUCCESS.rawPath()) { exchange ->
+ exchange.dispatch {
+ controller(SUCCESS) {
+ exchange.getResponseSender().send(SUCCESS.body)
+ }
+ }
+ }
+ .addExactPath(QUERY_PARAM.rawPath()) { exchange ->
+ exchange.dispatch {
+ controller(QUERY_PARAM) {
+ exchange.getResponseSender().send(exchange.getQueryString())
+ }
+ }
+ }
+ .addExactPath(REDIRECT.rawPath()) { exchange ->
+ exchange.dispatch {
+ controller(REDIRECT) {
+ exchange.setStatusCode(StatusCodes.FOUND)
+ exchange.getResponseHeaders().put(Headers.LOCATION, REDIRECT.body)
+ exchange.endExchange()
+ }
+ }
+ }
+ .addExactPath(CAPTURE_HEADERS.rawPath()) { exchange ->
+ exchange.dispatch {
+ controller(CAPTURE_HEADERS) {
+ exchange.setStatusCode(StatusCodes.OK)
+ exchange.getResponseHeaders().put(new HttpString("X-Test-Response"), exchange.getRequestHeaders().getFirst("X-Test-Request"))
+ exchange.getResponseSender().send(CAPTURE_HEADERS.body)
+ }
+ }
+ }
+ .addExactPath(ERROR.rawPath()) { exchange ->
+ exchange.dispatch {
+ controller(ERROR) {
+ exchange.setStatusCode(ERROR.status)
+ exchange.getResponseSender().send(ERROR.body)
+ }
+ }
+ }
+ .addExactPath(INDEXED_CHILD.rawPath()) { exchange ->
+ exchange.dispatch {
+ controller(INDEXED_CHILD) {
+ INDEXED_CHILD.collectSpanAttributes { name -> exchange.getQueryParameters().get(name).peekFirst() }
+ exchange.getResponseSender().send(INDEXED_CHILD.body)
+ }
+ }
+ }
+ ).build()
+ server.start()
+ return server
+ }
+
+ @Override
+ void stopServer(Undertow undertow) {
+ undertow.stop()
+ }
+
+ @Override
+ String expectedServerSpanName(ServerEndpoint endpoint) {
+ return "HTTP GET"
+ }
+
+ @Override
+ List<AttributeKey<?>> extraAttributes() {
+ [
+ SemanticAttributes.HTTP_RESPONSE_CONTENT_LENGTH,
+ SemanticAttributes.NET_PEER_NAME,
+ SemanticAttributes.NET_TRANSPORT
+ ]
+ }
+}

Context propagation for undertow async dispatch (#4950)