Undertow: run response customizer on all ServerConnection implementat…

…ions (#11539)
open-telemetry · Jun 8, 2024 · cc1c042 · cc1c042
1 parent ab4432c
commit cc1c042
Showing 1 changed file with 24 additions and 5 deletions.
diff --git a/...opentelemetry/javaagent/instrumentation/undertow/HttpServerConnectionInstrumentation.java b/...opentelemetry/javaagent/instrumentation/undertow/HttpServerConnectionInstrumentation.java
@@ -5,26 +5,33 @@
 
 package io.opentelemetry.javaagent.instrumentation.undertow;
 
+import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.extendsClass;
+import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
 import static io.opentelemetry.javaagent.instrumentation.undertow.UndertowSingletons.helper;
 import static net.bytebuddy.matcher.ElementMatchers.named;
-import static net.bytebuddy.matcher.ElementMatchers.namedOneOf;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
 import io.opentelemetry.context.Context;
+import io.opentelemetry.javaagent.bootstrap.CallDepth;
 import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseCustomizerHolder;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import io.undertow.server.HttpServerExchange;
+import io.undertow.server.ServerConnection;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
 public class HttpServerConnectionInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<ClassLoader> classLoaderOptimization() {
+    return hasClassesNamed("io.undertow.server.ServerConnection");
+  }
+
   @Override
   public ElementMatcher<TypeDescription> typeMatcher() {
-    return namedOneOf(
-        "io.undertow.server.protocol.http.HttpServerConnection",
-        "io.undertow.server.protocol.http2.Http2ServerConnection");
+    return extendsClass(named("io.undertow.server.ServerConnection"));
   }
 
   @Override
@@ -39,10 +46,22 @@ public void transform(TypeTransformer transformer) {
   public static class ResponseAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void onEnter(@Advice.Argument(0) HttpServerExchange exchange) {
+    public static void onEnter(
+        @Advice.Argument(0) HttpServerExchange exchange,
+        @Advice.Local("otelCallDepth") CallDepth callDepth) {
+      callDepth = CallDepth.forClass(ServerConnection.class);
+      if (callDepth.getAndIncrement() > 0) {
+        return;
+      }
+
       Context context = helper().getServerContext(exchange);
       HttpServerResponseCustomizerHolder.getCustomizer()
           .customize(context, exchange, UndertowHttpResponseMutator.INSTANCE);
     }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(@Advice.Local("otelCallDepth") CallDepth callDepth) {
+      callDepth.decrementAndGet();
+    }
   }
 }