Refactor netty test and fix instrumentation

DataDog · ValentinZakharov · May 13, 2025 · Mar 28, 2025 · Mar 28, 2025 · Mar 30, 2025
commit b2d45a70529fb4b1cc99c2513004364378116644
@@ -34,8 +34,11 @@ public final class AttributeKeys {
   public static final AttributeKey<Boolean> BLOCKED_RESPONSE_KEY =
       attributeKey("datadog.server.blocked_response");
 
-  public static final AttributeKey<HandlerContext.Sender> WEBSOCKET_HANDLER_CONTEXT =
-      attributeKey("datadog.server.websocket.handler_context");
+  public static final AttributeKey<HandlerContext.Sender> WEBSOCKET_SENDER_HANDLER_CONTEXT =
+      attributeKey("datadog.server.websocket.sender.handler_context");
+
+  public static final AttributeKey<HandlerContext.Receiver> WEBSOCKET_RECEIVER_HANDLER_CONTEXT =
+      attributeKey("datadog.server.websocket.receiver.handler_context");
 
   /**
    * Generate an attribute key or reuse the one existing in the global app map. This implementation

@@ -73,7 +73,7 @@ public void methodAdvice(MethodTransformer transformer) {
   }
 
   public static class OperationCompleteAdvice {
-    @Advice.OnMethodEnter
+    @Advice.OnMethodEnter(suppress = Throwable.class)
     public static AgentScope activateScope(@Advice.Argument(0) final ChannelFuture future) {
       /*
       Idea here is:

@@ -22,7 +22,6 @@
 import datadog.trace.instrumentation.netty41.server.HttpServerResponseTracingHandler;
 import datadog.trace.instrumentation.netty41.server.HttpServerTracingHandler;
 import datadog.trace.instrumentation.netty41.server.MaybeBlockResponseHandler;
-import datadog.trace.instrumentation.netty41.server.websocket.WebSocketProtocolHandshakeHandler;
 import datadog.trace.instrumentation.netty41.server.websocket.WebSocketServerTracingHandler;
 import io.netty.channel.ChannelHandler;
 import io.netty.channel.ChannelHandlerContext;
@@ -33,7 +32,6 @@
 import io.netty.handler.codec.http.HttpResponseDecoder;
 import io.netty.handler.codec.http.HttpResponseEncoder;
 import io.netty.handler.codec.http.HttpServerCodec;
-import io.netty.handler.codec.http.websocketx.WebSocketServerProtocolHandler;
 import io.netty.util.Attribute;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
@@ -80,7 +78,6 @@ public String[] helperClassNames() {
       packageName + ".server.HttpServerResponseTracingHandler",
       packageName + ".server.HttpServerTracingHandler",
       packageName + ".server.MaybeBlockResponseHandler",
-      packageName + ".server.websocket.WebSocketProtocolHandshakeHandler",
       packageName + ".server.websocket.WebSocketServerTracingHandler",
       packageName + ".server.websocket.WebSocketServerResponseTracingHandler",
       packageName + ".server.websocket.WebSocketServerRequestTracingHandler",
@@ -157,8 +154,6 @@ public static void addHandler(
         } else if (handler instanceof HttpResponseEncoder) {
           toAdd = HttpServerResponseTracingHandler.INSTANCE;
           toAdd2 = MaybeBlockResponseHandler.INSTANCE;
-        } else if (handler instanceof WebSocketServerProtocolHandler) {
-          toAdd = WebSocketProtocolHandshakeHandler.INSTANCE;
         } else
         // Client pipeline handlers
         if (handler instanceof HttpClientCodec) {

@@ -2,10 +2,12 @@
 
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.instrumentation.netty41.AttributeKeys.SPAN_ATTRIBUTE_KEY;
+import static datadog.trace.instrumentation.netty41.AttributeKeys.WEBSOCKET_SENDER_HANDLER_CONTEXT;
 import static datadog.trace.instrumentation.netty41.server.NettyHttpServerDecorator.DECORATE;
 
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext;
 import io.netty.channel.ChannelHandler;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelOutboundHandlerAdapter;
@@ -39,9 +41,16 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann
         ctx.channel().attr(SPAN_ATTRIBUTE_KEY).remove();
         throw throwable;
       }
+      final boolean isWebsocketUpgrade =
+          response.status() == HttpResponseStatus.SWITCHING_PROTOCOLS
+              && "websocket".equals(response.headers().get(HttpHeaderNames.UPGRADE));
+      if (isWebsocketUpgrade) {
+        ctx.channel()
+            .attr(WEBSOCKET_SENDER_HANDLER_CONTEXT)
+            .set(new HandlerContext.Sender(span, ctx.channel().id().asShortText()));
+      }
       if (response.status() != HttpResponseStatus.CONTINUE
-          && (response.status() != HttpResponseStatus.SWITCHING_PROTOCOLS
-              || "websocket".equals(response.headers().get(HttpHeaderNames.UPGRADE)))) {
+          && (response.status() != HttpResponseStatus.SWITCHING_PROTOCOLS || isWebsocketUpgrade)) {
         DECORATE.onResponse(span, response);
         DECORATE.beforeFinish(span);
         span.finish(); // Finish the span manually since finishSpanOnClose was false

diff --git a/...dog/trace/instrumentation/netty41/server/websocket/WebSocketProtocolHandshakeHandler.java b/...dog/trace/instrumentation/netty41/server/websocket/WebSocketProtocolHandshakeHandler.java
@@ -2,14 +2,17 @@
 
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
 import static datadog.trace.bootstrap.instrumentation.decorator.WebsocketDecorator.DECORATE;
-import static datadog.trace.instrumentation.netty41.AttributeKeys.WEBSOCKET_HANDLER_CONTEXT;
+import static datadog.trace.bootstrap.instrumentation.websocket.HandlersExtractor.MESSAGE_TYPE_TEXT;
+import static datadog.trace.instrumentation.netty41.AttributeKeys.WEBSOCKET_RECEIVER_HANDLER_CONTEXT;
+import static datadog.trace.instrumentation.netty41.AttributeKeys.WEBSOCKET_SENDER_HANDLER_CONTEXT;
 
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext;
 import io.netty.channel.*;
 import io.netty.handler.codec.http.websocketx.BinaryWebSocketFrame;
 import io.netty.handler.codec.http.websocketx.CloseWebSocketFrame;
+import io.netty.handler.codec.http.websocketx.ContinuationWebSocketFrame;
 import io.netty.handler.codec.http.websocketx.TextWebSocketFrame;
 import io.netty.handler.codec.http.websocketx.WebSocketFrame;
 
@@ -23,29 +26,32 @@ public void channelRead(ChannelHandlerContext ctx, Object frame) {
 
     if (frame instanceof WebSocketFrame) {
       Channel channel = ctx.channel();
-      final HandlerContext.Sender sessionState = channel.attr(WEBSOCKET_HANDLER_CONTEXT).get();
-      if (sessionState != null) {
-        final HandlerContext.Receiver handlerContext =
-            new HandlerContext.Receiver(
-                sessionState.getHandshakeSpan(), ctx.channel().id().asShortText());
-
+      HandlerContext.Receiver receiverContext =
+          channel.attr(WEBSOCKET_RECEIVER_HANDLER_CONTEXT).get();
+      if (receiverContext == null) {
+        HandlerContext.Sender sessionState = channel.attr(WEBSOCKET_SENDER_HANDLER_CONTEXT).get();
+        if (sessionState != null) {
+          receiverContext =
+              new HandlerContext.Receiver(
+                  sessionState.getHandshakeSpan(), ctx.channel().id().asShortText());
+          channel.attr(WEBSOCKET_RECEIVER_HANDLER_CONTEXT).set(receiverContext);
+        }
+      }
+      if (receiverContext != null) {
         if (frame instanceof TextWebSocketFrame) {
           // WebSocket Read Text Start
           TextWebSocketFrame textFrame = (TextWebSocketFrame) frame;
 
           final AgentSpan span =
               DECORATE.onReceiveFrameStart(
-                  handlerContext, textFrame.text(), textFrame.isFinalFragment());
-          if (span == null) {
+                  receiverContext, textFrame.text(), textFrame.isFinalFragment());
+          try (final AgentScope scope = activateSpan(span)) {
             ctx.fireChannelRead(textFrame);
-          } else {
-            try (final AgentScope scope = activateSpan(span)) {
-              ctx.fireChannelRead(textFrame);
-
-              // WebSocket Read Text Start
-              if (textFrame.isFinalFragment()) {
-                DECORATE.onFrameEnd(handlerContext);
-              }
+            // WebSocket Read Text Start
+          } finally {
+            if (textFrame.isFinalFragment()) {
+              channel.attr(WEBSOCKET_RECEIVER_HANDLER_CONTEXT).remove();
+              DECORATE.onFrameEnd(receiverContext);
             }
           }
           return;
@@ -56,17 +62,38 @@ public void channelRead(ChannelHandlerContext ctx, Object frame) {
           BinaryWebSocketFrame binaryFrame = (BinaryWebSocketFrame) frame;
           final AgentSpan span =
               DECORATE.onReceiveFrameStart(
-                  handlerContext, binaryFrame.content().nioBuffer(), binaryFrame.isFinalFragment());
-          if (span == null) {
+                  receiverContext,
+                  binaryFrame.content().nioBuffer(),
+                  binaryFrame.isFinalFragment());
+          try (final AgentScope scope = activateSpan(span)) {
             ctx.fireChannelRead(binaryFrame);
-          } else {
-            try (final AgentScope scope = activateSpan(span)) {
-              ctx.fireChannelRead(binaryFrame);
+          } finally {
+            // WebSocket Read Binary End
+            if (binaryFrame.isFinalFragment()) {
+              channel.attr(WEBSOCKET_RECEIVER_HANDLER_CONTEXT).remove();
+              DECORATE.onFrameEnd(receiverContext);
+            }
+          }
 
-              // WebSocket Read Binary End
-              if (binaryFrame.isFinalFragment()) {
-                DECORATE.onFrameEnd(handlerContext);
-              }
+          return;
+        }
+
+        if (frame instanceof ContinuationWebSocketFrame) {
+          ContinuationWebSocketFrame continuationWebSocketFrame =
+              (ContinuationWebSocketFrame) frame;
+          final AgentSpan span =
+              DECORATE.onReceiveFrameStart(
+                  receiverContext,
+                  MESSAGE_TYPE_TEXT.equals(receiverContext.getMessageType())
+                      ? continuationWebSocketFrame.text()
+                      : continuationWebSocketFrame.content().nioBuffer(),
+                  continuationWebSocketFrame.isFinalFragment());
+          try (final AgentScope scope = activateSpan(span)) {
+            ctx.fireChannelRead(continuationWebSocketFrame);
+          } finally {
+            if (continuationWebSocketFrame.isFinalFragment()) {
+              channel.attr(WEBSOCKET_RECEIVER_HANDLER_CONTEXT).remove();
+              DECORATE.onFrameEnd(receiverContext);
             }
           }
           return;
@@ -77,25 +104,21 @@ public void channelRead(ChannelHandlerContext ctx, Object frame) {
           CloseWebSocketFrame closeFrame = (CloseWebSocketFrame) frame;
           int statusCode = closeFrame.statusCode();
           String reasonText = closeFrame.reasonText();
-          channel.attr(WEBSOCKET_HANDLER_CONTEXT).remove();
+          channel.attr(WEBSOCKET_SENDER_HANDLER_CONTEXT).remove();
+          channel.attr(WEBSOCKET_RECEIVER_HANDLER_CONTEXT).remove();
           final AgentSpan span =
-              DECORATE.onSessionCloseReceived(handlerContext, reasonText, statusCode);
-          if (span == null) {
+              DECORATE.onSessionCloseReceived(receiverContext, reasonText, statusCode);
+          try (final AgentScope scope = activateSpan(span)) {
             ctx.fireChannelRead(closeFrame);
-          } else {
-            try (final AgentScope scope = activateSpan(span)) {
-              ctx.fireChannelRead(closeFrame);
-
-              if (closeFrame.isFinalFragment()) {
-                DECORATE.onFrameEnd(handlerContext);
-              }
+            if (closeFrame.isFinalFragment()) {
+              DECORATE.onFrameEnd(receiverContext);
             }
           }
           return;
         }
       }
     }
-    // can be other messages we do not handle like ping, pong or continuations
+    // can be other messages we do not handle like ping, pong
     ctx.fireChannelRead(frame);
   }
 }
@@ -4,14 +4,15 @@
 import static datadog.trace.bootstrap.instrumentation.decorator.WebsocketDecorator.DECORATE;
 import static datadog.trace.bootstrap.instrumentation.websocket.HandlersExtractor.MESSAGE_TYPE_BINARY;
 import static datadog.trace.bootstrap.instrumentation.websocket.HandlersExtractor.MESSAGE_TYPE_TEXT;
-import static datadog.trace.instrumentation.netty41.AttributeKeys.WEBSOCKET_HANDLER_CONTEXT;
+import static datadog.trace.instrumentation.netty41.AttributeKeys.WEBSOCKET_SENDER_HANDLER_CONTEXT;
 
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.websocket.HandlerContext;
 import io.netty.channel.*;
 import io.netty.handler.codec.http.websocketx.BinaryWebSocketFrame;
 import io.netty.handler.codec.http.websocketx.CloseWebSocketFrame;
+import io.netty.handler.codec.http.websocketx.ContinuationWebSocketFrame;
 import io.netty.handler.codec.http.websocketx.TextWebSocketFrame;
 import io.netty.handler.codec.http.websocketx.WebSocketFrame;
 
@@ -26,7 +27,7 @@ public void write(ChannelHandlerContext ctx, Object frame, ChannelPromise promis
 
     if (frame instanceof WebSocketFrame) {
       Channel channel = ctx.channel();
-      HandlerContext.Sender handlerContext = channel.attr(WEBSOCKET_HANDLER_CONTEXT).get();
+      HandlerContext.Sender handlerContext = channel.attr(WEBSOCKET_SENDER_HANDLER_CONTEXT).get();
       if (handlerContext != null) {
 
         if (frame instanceof TextWebSocketFrame) {
@@ -35,16 +36,12 @@ public void write(ChannelHandlerContext ctx, Object frame, ChannelPromise promis
           final AgentSpan span =
               DECORATE.onSendFrameStart(
                   handlerContext, MESSAGE_TYPE_TEXT, textFrame.text().length());
-          if (span == null) {
+          try (final AgentScope scope = activateSpan(span)) {
             ctx.write(frame, promise);
-          } else {
-            try (final AgentScope scope = activateSpan(span)) {
-              ctx.write(frame, promise);
-
-              // WebSocket Write Text End
-              if (textFrame.isFinalFragment()) {
-                DECORATE.onFrameEnd(handlerContext);
-              }
+          } finally {
+            // WebSocket Write Text End
+            if (textFrame.isFinalFragment()) {
+              DECORATE.onFrameEnd(handlerContext);
             }
           }
           return;
@@ -56,16 +53,33 @@ public void write(ChannelHandlerContext ctx, Object frame, ChannelPromise promis
           final AgentSpan span =
               DECORATE.onSendFrameStart(
                   handlerContext, MESSAGE_TYPE_BINARY, binaryFrame.content().readableBytes());
-          if (span == null) {
+          try (final AgentScope scope = activateSpan(span)) {
             ctx.write(frame, promise);
-          } else {
-            try (final AgentScope scope = activateSpan(span)) {
-              ctx.write(frame, promise);
+          } finally {
+            // WebSocket Write Binary End
+            if (binaryFrame.isFinalFragment()) {
+              DECORATE.onFrameEnd(handlerContext);
+            }
+          }
+          return;
+        }
 
-              // WebSocket Write Binary End
-              if (binaryFrame.isFinalFragment()) {
-                DECORATE.onFrameEnd(handlerContext);
-              }
+        if (frame instanceof ContinuationWebSocketFrame) {
+          ContinuationWebSocketFrame continuationWebSocketFrame =
+              (ContinuationWebSocketFrame) frame;
+          final AgentSpan span =
+              DECORATE.onSendFrameStart(
+                  handlerContext,
+                  handlerContext.getMessageType(),
+                  MESSAGE_TYPE_TEXT.equals(handlerContext.getMessageType())
+                      ? continuationWebSocketFrame.text().length()
+                      : continuationWebSocketFrame.content().readableBytes());
+          try (final AgentScope scope = activateSpan(span)) {
+            ctx.write(frame, promise);
+          } finally {
+            // WebSocket Write Binary End
+            if (continuationWebSocketFrame.isFinalFragment()) {
+              DECORATE.onFrameEnd(handlerContext);
             }
           }
           return;
@@ -76,24 +90,21 @@ public void write(ChannelHandlerContext ctx, Object frame, ChannelPromise promis
           CloseWebSocketFrame closeFrame = (CloseWebSocketFrame) frame;
           int statusCode = closeFrame.statusCode();
           String reasonText = closeFrame.reasonText();
-          channel.attr(WEBSOCKET_HANDLER_CONTEXT).remove();
+          channel.attr(WEBSOCKET_SENDER_HANDLER_CONTEXT).remove();
           final AgentSpan span =
               DECORATE.onSessionCloseIssued(handlerContext, reasonText, statusCode);
-          if (span == null) {
+          try (final AgentScope scope = activateSpan(span)) {
             ctx.write(frame, promise);
-          } else {
-            try (final AgentScope scope = activateSpan(span)) {
-              ctx.write(frame, promise);
-              if (closeFrame.isFinalFragment()) {
-                DECORATE.onFrameEnd(handlerContext);
-              }
+          } finally {
+            if (closeFrame.isFinalFragment()) {
+              DECORATE.onFrameEnd(handlerContext);
             }
           }
           return;
         }
       }
     }
-    // can be other messages we do not handle like ping, pong or continuations
+    // can be other messages we do not handle like ping, pong
     ctx.write(frame, promise);
   }
 }