fix: corrects h2 handling and cleans up otel Context currency when executing Promises

dmarkwat · dmarkwat · commit bac29f72106e · 2026-04-30T08:25:07.000-07:00
diff --git a/instrumentation/finagle-http-23.11/javaagent/build.gradle.kts b/instrumentation/finagle-http-23.11/javaagent/build.gradle.kts
@@ -36,6 +36,7 @@ dependencies {
 
   implementation(project(":instrumentation:netty:netty-4.1:javaagent"))
   implementation(project(":instrumentation:netty:netty-4.1:library"))
+  implementation(project(":instrumentation:netty:netty-common-4.0:javaagent"))
   implementation(project(":instrumentation:netty:netty-common-4.0:library"))
 }
 
diff --git a/instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/GenStreamingServerDispatcherInstrumentation.java b/instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/GenStreamingServerDispatcherInstrumentation.java
@@ -48,8 +48,9 @@ public static Scope methodEnter(@Advice.Argument(0) Object req) {
       if (req instanceof Request) {
         // practically this will always be a Request, from HttpServerDispatcher
         Request request = (Request) req;
+        // if this is null, there's a bug in the instrumentation
         Context context = request.ctx().apply(Helpers.OTEL_CONTEXT_KEY);
-        return context != null && context != Context.root() ? context.makeCurrent() : null;
+        return context.makeCurrent();
       }
       return null;
     }
diff --git a/instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/Helpers.java b/instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/Helpers.java
@@ -13,6 +13,7 @@
 import com.twitter.finagle.http.Request;
 import com.twitter.finagle.http.Request$;
 import com.twitter.finagle.http.collection.RecordSchema;
+import com.twitter.finagle.http2.transport.common.Http2StreamMessageHandler;
 import io.netty.channel.Channel;
 import io.netty.channel.ChannelHandler;
 import io.netty.channel.ChannelHandlerContext;
@@ -26,16 +27,18 @@
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.util.VirtualField;
 import io.opentelemetry.instrumentation.netty.v4_1.internal.AttributeKeys;
-import io.opentelemetry.instrumentation.netty.v4_1.internal.ServerContext;
 import io.opentelemetry.instrumentation.netty.v4_1.internal.ServerContexts;
 import io.opentelemetry.instrumentation.netty.v4_1.internal.client.HttpClientTracingHandler;
 import io.opentelemetry.instrumentation.netty.v4_1.internal.server.HttpServerTracingHandler;
+import io.opentelemetry.javaagent.instrumentation.netty.common.v4_0.VirtualFieldHelper;
 import io.opentelemetry.javaagent.instrumentation.netty.v4_1.NettyServerSingletons;
 
 public class Helpers {
 
-  private static final VirtualField<ChannelHandler, ChannelHandler> CHANNEL_HANDLER =
-      VirtualField.find(ChannelHandler.class, ChannelHandler.class);
+  private static final VirtualField<FullHttpRequest, Context> FULL_HTTP_REQUEST_CONTEXT = VirtualField.find(
+      FullHttpRequest.class, Context.class);
+  private static final VirtualField<HttpRequest, Context> HTTP_REQUEST_CONTEXT = VirtualField.find(
+      HttpRequest.class, Context.class);
 
   public static final RecordSchema.Field<Context> OTEL_CONTEXT_KEY =
       Request$.MODULE$.Schema().newField();
@@ -76,7 +79,7 @@ protected void initChannel(C channel) throws Exception {
                 .pipeline()
                 .addAfter(codecCtx.name(), ourHandler.getClass().getName(), ourHandler);
             // attach this in this way to match up with how netty instrumentation expects things
-            CHANNEL_HANDLER.set(codecCtx.handler(), ourHandler);
+            VirtualFieldHelper.CHANNEL_HANDLER.set(codecCtx.handler(), ourHandler);
           }
         }
       }
@@ -116,7 +119,7 @@ protected void initChannel(C channel) throws Exception {
                 .pipeline()
                 .addAfter(codecCtx.name(), ourHandler.getClass().getName(), ourHandler);
             // attach this in this way to match up with how netty instrumentation expects things
-            CHANNEL_HANDLER.set(codecCtx.handler(), ourHandler);
+            VirtualFieldHelper.CHANNEL_HANDLER.set(codecCtx.handler(), ourHandler);
           }
         }
       }
@@ -127,9 +130,15 @@ protected void initChannel(C channel) throws Exception {
   Part 1/3 of bridging the otel Context from netty to finagle (for h2).
    */
   public static void mutateHandlerPipeline(Channel ch) {
-    ChannelHandler codec = ch.pipeline().get(Netty4HttpPackageHelpers.getHttpCodecName());
-
-    if (codec instanceof HttpServerCodec) {
+    ChannelHandler h1Handler = ch.pipeline().get(Netty4HttpPackageHelpers.getHttpCodecName());
+    Http2StreamMessageHandler h2Handler = ch.pipeline().get(Http2StreamMessageHandler.class);
+
+    // h1 server handler || h2 server handler;
+    // private class on a semi-private type -- not bothering to extract that any other way
+    if (
+        h1Handler instanceof HttpServerCodec || (h2Handler != null && h2Handler.getClass().getName()
+            .equals(
+                "com.twitter.finagle.http2.transport.common.Http2StreamMessageHandler$ServerHttp2StreamMessageHandler"))) {
       // ensure we capture the server context and assign it to the outgoing request before offering
       // to the AsyncQueue;
       // not applicable to clients
@@ -145,26 +154,12 @@ public static void mutateHandlerPipeline(Channel ch) {
                  */
                 @Override
                 public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
-                  ServerContexts serverContexts = ServerContexts.get(ctx.channel());
-                  if (serverContexts == null) {
-                    super.channelRead(ctx, msg);
-                    return;
-                  }
-
-                  ServerContext serverContext = serverContexts.peekLast();
-
                   // type switch courtesy of
                   // com.twitter.finagle.netty4.http.Netty4ServerStreamTransport.read()
                   if (msg instanceof FullHttpRequest) {
-                    VirtualField.find(FullHttpRequest.class, Context.class)
-                        .set(
-                            (FullHttpRequest) msg,
-                            serverContext != null ? serverContext.context() : null);
+                    FULL_HTTP_REQUEST_CONTEXT.set((FullHttpRequest) msg, Context.current());
                   } else if (msg instanceof HttpRequest) {
-                    VirtualField.find(HttpRequest.class, Context.class)
-                        .set(
-                            (HttpRequest) msg,
-                            serverContext != null ? serverContext.context() : null);
+                    HTTP_REQUEST_CONTEXT.set((HttpRequest) msg, Context.current());
                   } else {
                     throw new IllegalArgumentException("unexpected request type: " + msg);
                   }
@@ -182,9 +177,9 @@ public static void chainContextToFinagle(Object msg, Request request) {
     Context context;
     // type switch courtesy of com.twitter.finagle.netty4.http.Netty4ServerStreamTransport.read()
     if (msg instanceof FullHttpRequest) {
-      context = VirtualField.find(FullHttpRequest.class, Context.class).get((FullHttpRequest) msg);
+      context = FULL_HTTP_REQUEST_CONTEXT.get((FullHttpRequest) msg);
     } else if (msg instanceof HttpRequest) {
-      context = VirtualField.find(HttpRequest.class, Context.class).get((HttpRequest) msg);
+      context = HTTP_REQUEST_CONTEXT.get((HttpRequest) msg);
     } else {
       // shouldn't practically reach here
       throw new IllegalArgumentException("unexpected request type: " + msg);
diff --git a/instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/LocalSchedulerActivationInstrumentation.java b/instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/LocalSchedulerActivationInstrumentation.java
@@ -46,10 +46,8 @@ public static Runnable wrap(@Advice.Argument(0) Runnable task) {
         return null;
       }
 
+      // always set it: you never know what might be polluting the thread local context at the time
       Context context = Java8BytecodeBridge.currentContext();
-      if (context == Context.root()) {
-        return task;
-      }
       return ContextPropagatingRunnable.propagateContext(task, context);
     }
   }
diff --git a/instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/PromiseInterruptibleInstrumentation.java b/instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/PromiseInterruptibleInstrumentation.java
@@ -45,9 +45,6 @@ public static PartialFunction<Throwable, BoxedUnit> onEnter(
         return handler;
       }
       Context context = Context.current();
-      if (context == Context.root()) {
-        return handler;
-      }
       return new TwitterUtilCoreHelpers.InterruptibleWithContext(context, handler);
     }
   }
diff --git a/instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/PromiseKInstrumentation.java b/instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/PromiseKInstrumentation.java
@@ -48,28 +48,22 @@ public static class TrapContextAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class, inline = false)
     public static void onExit(@Advice.This Promise.K thiz) {
       Context current = Context.current();
-      if (current != Context.root()) {
-        TwitterUtilCoreHelpers.PROMISE_K_CONTEXT_FIELD.set(thiz, current);
-      }
+      TwitterUtilCoreHelpers.PROMISE_K_CONTEXT_FIELD.set(thiz, current);
     }
   }
 
   @SuppressWarnings("unused")
   public static class ApplyAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class, inline = false)
     public static Scope onApplyEnter(@Advice.This Promise.K thiz) {
+      // if this is null, there's a bug in the instrumentation
       Context savedContext = TwitterUtilCoreHelpers.PROMISE_K_CONTEXT_FIELD.get(thiz);
-      if (savedContext != null) {
-        return savedContext.makeCurrent();
-      }
-      return null;
+      return savedContext.makeCurrent();
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class, inline = false)
     public static void onApplyExit(@Advice.Enter Scope scope) {
-      if (scope != null) {
-        scope.close();
-      }
+      scope.close();
     }
   }
 }
diff --git a/instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/TwitterUtilCoreHelpers.java b/instrumentation/finagle-http-23.11/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/TwitterUtilCoreHelpers.java
@@ -49,21 +49,17 @@ public BoxedUnit apply(Throwable x) {
   }
 
   public static <T, O> Function1<T, O> wrap(Context context, Function1<T, O> fn) {
-    if (context == Context.root()) {
-      return fn;
-    }
     return (t) -> {
+      // always set it: you never know what might be polluting the thread local context at the time
       try (Scope ignored = context.makeCurrent()) {
         return fn.apply(t);
       }
     };
   }
 
   public static <T> Function0<T> wrap(Context context, Function0<T> fn) {
-    if (context == Context.root()) {
-      return fn;
-    }
     return () -> {
+      // always set it: you never know what might be polluting the thread local context at the time
       try (Scope ignored = context.makeCurrent()) {
         return fn.apply();
       }
diff --git a/instrumentation/finagle-http-23.11/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/ServerH2Test.java b/instrumentation/finagle-http-23.11/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/finaglehttp/v23_11/ServerH2Test.java
@@ -18,6 +18,7 @@
 import com.twitter.finagle.http2.param.PriorKnowledge;
 import com.twitter.util.Await;
 import com.twitter.util.Duration;
+import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.instrumentation.netty.v4_1.internal.ProtocolSpecificEvent;
 import io.opentelemetry.instrumentation.testing.internal.AutoCleanupExtension;
 import io.opentelemetry.instrumentation.testing.junit.http.ServerEndpoint;
@@ -55,6 +56,9 @@ private static void assertSwitchingProtocolsEvent(EventDataAssert eventDataAsser
             equalTo(ProtocolSpecificEvent.SWITCHING_PROTOCOLS_TO_KEY, singletonList("h2c")));
   }
 
+  /*
+  Bonus is that this implicitly tests both the server & client h2 upgrades.
+   */
   @Test
   void h2ProtocolUpgrade() throws Exception {
     URI uri = URI.create("http://localhost:" + port + SUCCESS.getPath());
@@ -66,7 +70,7 @@ void h2ProtocolUpgrade() throws Exception {
     cleanup.deferCleanup(client::close);
 
     Response response =
-        Await.result(
+        testing.runWithSpan("h2-upgrade-client", () -> Await.result(
             client.apply(
                 Utils.buildRequest(
                     "GET",
@@ -76,7 +80,7 @@ void h2ProtocolUpgrade() throws Exception {
                         TEST_USER_AGENT,
                         HttpHeaderNames.X_FORWARDED_FOR.toString(),
                         TEST_CLIENT_IP))),
-            Duration.fromSeconds(20));
+            Duration.fromSeconds(20)));
 
     assertThat(response.status().code()).isEqualTo(SUCCESS.getStatus());
     assertThat(response.contentString()).isEqualTo(SUCCESS.getBody());
@@ -87,19 +91,23 @@ void h2ProtocolUpgrade() throws Exception {
     testing.waitAndAssertTraces(
         trace -> {
           List<Consumer<SpanDataAssert>> spanAssertions = new ArrayList<>();
+          // client initiation
           spanAssertions.add(
-              s -> s.hasEventsSatisfyingExactly(ServerH2Test::assertSwitchingProtocolsEvent));
+              s -> s.hasName("h2-upgrade-client").hasNoParent().hasKind(SpanKind.INTERNAL));
+          // actual client netty span (including upgrade event)
           spanAssertions.add(
-              span -> {
-                assertServerSpan(span, method, endpoint, endpoint.getStatus());
-                span.hasEventsSatisfyingExactly(ServerH2Test::assertSwitchingProtocolsEvent);
-              });
-
-          int parentIndex = 1;
+              s -> s.hasKind(SpanKind.CLIENT).hasName(method).hasParent(trace.getSpan(0))
+                  .hasEventsSatisfyingExactly(ServerH2Test::assertSwitchingProtocolsEvent));
+          // server netty span (including upgrade event)
+          spanAssertions.add(
+              span -> assertServerSpan(span, method, endpoint, endpoint.getStatus())
+                  .hasParent(trace.getSpan(1))
+                  .hasEventsSatisfyingExactly(ServerH2Test::assertSwitchingProtocolsEvent));
+          // server controller span
           spanAssertions.add(
               span -> {
                 assertControllerSpan(span, null);
-                span.hasParent(trace.getSpan(parentIndex));
+                span.hasParent(trace.getSpan(2));
               });
 
           trace.hasSpansSatisfyingExactly(spanAssertions);

Original file line number	Diff line number	Diff line change
`@@ -36,6 +36,7 @@ dependencies {`
`36`	`36`
`37`	`37`	`implementation(project(":instrumentation:netty:netty-4.1:javaagent"))`
`38`	`38`	`implementation(project(":instrumentation:netty:netty-4.1:library"))`
	`39`	`+ implementation(project(":instrumentation:netty:netty-common-4.0:javaagent"))`
`39`	`40`	`implementation(project(":instrumentation:netty:netty-common-4.0:library"))`
`40`	`41`	`}`
`41`	`42`
Original file line number	Diff line number	Diff line change
`@@ -48,8 +48,9 @@ public static Scope methodEnter(@Advice.Argument(0) Object req) {`
`48`	`48`	`if (req instanceof Request) {`
`49`	`49`	`// practically this will always be a Request, from HttpServerDispatcher`
`50`	`50`	`Request request = (Request) req;`
	`51`	`+ // if this is null, there's a bug in the instrumentation`
`51`	`52`	`Context context = request.ctx().apply(Helpers.OTEL_CONTEXT_KEY);`
`52`		`- return context != null && context != Context.root() ? context.makeCurrent() : null;`
	`53`	`+ return context.makeCurrent();`
`53`	`54`	`}`
`54`	`55`	`return null;`
`55`	`56`	`}`
Original file line number	Diff line number	Diff line change
`@@ -46,10 +46,8 @@ public static Runnable wrap(@Advice.Argument(0) Runnable task) {`
`46`	`46`	`return null;`
`47`	`47`	`}`
`48`	`48`
	`49`	`+ // always set it: you never know what might be polluting the thread local context at the time`
`49`	`50`	`Context context = Java8BytecodeBridge.currentContext();`
`50`		`- if (context == Context.root()) {`
`51`		`- return task;`
`52`		`- }`
`53`	`51`	`return ContextPropagatingRunnable.propagateContext(task, context);`
`54`	`52`	`}`
`55`	`53`	`}`
Original file line number	Diff line number	Diff line change
`@@ -45,9 +45,6 @@ public static PartialFunction<Throwable, BoxedUnit> onEnter(`
`45`	`45`	`return handler;`
`46`	`46`	`}`
`47`	`47`	`Context context = Context.current();`
`48`		`- if (context == Context.root()) {`
`49`		`- return handler;`
`50`		`- }`
`51`	`48`	`return new TwitterUtilCoreHelpers.InterruptibleWithContext(context, handler);`
`52`	`49`	`}`
`53`	`50`	`}`
Original file line number	Diff line number	Diff line change
`@@ -48,28 +48,22 @@ public static class TrapContextAdvice {`
`48`	`48`	`@Advice.OnMethodExit(suppress = Throwable.class, inline = false)`
`49`	`49`	`public static void onExit(@Advice.This Promise.K thiz) {`
`50`	`50`	`Context current = Context.current();`
`51`		`- if (current != Context.root()) {`
`52`		`- TwitterUtilCoreHelpers.PROMISE_K_CONTEXT_FIELD.set(thiz, current);`
`53`		`- }`
	`51`	`+ TwitterUtilCoreHelpers.PROMISE_K_CONTEXT_FIELD.set(thiz, current);`
`54`	`52`	`}`
`55`	`53`	`}`
`56`	`54`
`57`	`55`	`@SuppressWarnings("unused")`
`58`	`56`	`public static class ApplyAdvice {`
`59`	`57`	`@Advice.OnMethodEnter(suppress = Throwable.class, inline = false)`
`60`	`58`	`public static Scope onApplyEnter(@Advice.This Promise.K thiz) {`
	`59`	`+ // if this is null, there's a bug in the instrumentation`
`61`	`60`	`Context savedContext = TwitterUtilCoreHelpers.PROMISE_K_CONTEXT_FIELD.get(thiz);`
`62`		`- if (savedContext != null) {`
`63`		`- return savedContext.makeCurrent();`
`64`		`- }`
`65`		`- return null;`
	`61`	`+ return savedContext.makeCurrent();`
`66`	`62`	`}`
`67`	`63`
`68`	`64`	`@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class, inline = false)`
`69`	`65`	`public static void onApplyExit(@Advice.Enter Scope scope) {`
`70`		`- if (scope != null) {`
`71`		`- scope.close();`
`72`		`- }`
	`66`	`+ scope.close();`
`73`	`67`	`}`
`74`	`68`	`}`
`75`	`69`	`}`
Original file line number	Diff line number	Diff line change
`@@ -49,21 +49,17 @@ public BoxedUnit apply(Throwable x) {`
`49`	`49`	`}`
`50`	`50`
`51`	`51`	`public static <T, O> Function1<T, O> wrap(Context context, Function1<T, O> fn) {`
`52`		`- if (context == Context.root()) {`
`53`		`- return fn;`
`54`		`- }`
`55`	`52`	`return (t) -> {`
	`53`	`+ // always set it: you never know what might be polluting the thread local context at the time`
`56`	`54`	`try (Scope ignored = context.makeCurrent()) {`
`57`	`55`	`return fn.apply(t);`
`58`	`56`	`}`
`59`	`57`	`};`
`60`	`58`	`}`
`61`	`59`
`62`	`60`	`public static <T> Function0<T> wrap(Context context, Function0<T> fn) {`
`63`		`- if (context == Context.root()) {`
`64`		`- return fn;`
`65`		`- }`
`66`	`61`	`return () -> {`
	`62`	`+ // always set it: you never know what might be polluting the thread local context at the time`
`67`	`63`	`try (Scope ignored = context.makeCurrent()) {`
`68`	`64`	`return fn.apply();`
`69`	`65`	`}`