Added WebSocket tracing check

Author: ValentinZakharov

dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/ChannelPipelineAdviceUtil.java

@@ -1,5 +1,6 @@
 package datadog.trace.instrumentation.netty38;
 
+import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import datadog.trace.bootstrap.ContextStore;
 import datadog.trace.instrumentation.netty38.client.HttpClientRequestTracingHandler;
@@ -53,25 +54,31 @@ public static void wrapHandler(
         pipeline.addLast(
             MaybeBlockResponseHandler.class.getName(), new MaybeBlockResponseHandler(contextStore));
       } else if (handler instanceof WebSocketServerProtocolHandler) {
-        if (pipeline.get(HttpServerTracingHandler.class) != null) {
-          addHandlerAfter(
-              pipeline,
-              "datadog.trace.instrumentation.netty38.server.HttpServerTracingHandler",
-              new WebSocketServerTracingHandler(contextStore));
+        if (InstrumenterConfig.get().isWebsocketTracingEnabled()) {
+          if (pipeline.get(HttpServerTracingHandler.class) != null) {
+            addHandlerAfter(
+                pipeline,
+                "datadog.trace.instrumentation.netty38.server.HttpServerTracingHandler",
+                new WebSocketServerTracingHandler(contextStore));
+          }
         }
       } else if (handler instanceof WebSocket13FrameEncoder) {
-        if (pipeline.get(HttpServerRequestTracingHandler.class) != null) {
-          addHandlerAfter(
-              pipeline,
-              "datadog.trace.instrumentation.netty38.server.HttpServerRequestTracingHandler",
-              new WebSocketServerRequestTracingHandler(contextStore));
+        if (InstrumenterConfig.get().isWebsocketTracingEnabled()) {
+          if (pipeline.get(HttpServerRequestTracingHandler.class) != null) {
+            addHandlerAfter(
+                pipeline,
+                "datadog.trace.instrumentation.netty38.server.HttpServerRequestTracingHandler",
+                new WebSocketServerRequestTracingHandler(contextStore));
+          }
         }
       } else if (handler instanceof WebSocket13FrameDecoder) {
-        if (pipeline.get(HttpServerResponseTracingHandler.class) != null) {
-          addHandlerAfter(
-              pipeline,
-              "datadog.trace.instrumentation.netty38.server.HttpServerResponseTracingHandler",
-              new WebSocketServerResponseTracingHandler(contextStore));
+        if (InstrumenterConfig.get().isWebsocketTracingEnabled()) {
+          if (pipeline.get(HttpServerResponseTracingHandler.class) != null) {
+            addHandlerAfter(
+                pipeline,
+                "datadog.trace.instrumentation.netty38.server.HttpServerResponseTracingHandler",
+                new WebSocketServerResponseTracingHandler(contextStore));
+          }
         }
       } else
       // Client pipeline handlers

dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/NettyChannelPipelineInstrumentation.java

@@ -13,6 +13,7 @@
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.instrumentation.netty40.client.HttpClientRequestTracingHandler;
@@ -160,23 +161,25 @@ public static void addHandler(
               HttpServerResponseTracingHandler.INSTANCE,
               MaybeBlockResponseHandler.INSTANCE);
         } else if (handler instanceof WebSocketServerProtocolHandler) {
-          if (pipeline.get(HttpServerTracingHandler.class) != null) {
-            NettyPipelineHelper.addHandlerAfter(
-                pipeline,
-                "datadog.trace.instrumentation.netty40.server.HttpServerTracingHandler",
-                new WebSocketServerTracingHandler());
-          }
-          if (pipeline.get(HttpRequestDecoder.class) != null) {
-            NettyPipelineHelper.addHandlerAfter(
-                pipeline,
-                "datadog.trace.instrumentation.netty40.server.HttpServerRequestTracingHandler",
-                WebSocketServerRequestTracingHandler.INSTANCE);
-          }
-          if (pipeline.get(HttpResponseEncoder.class) != null) {
-            NettyPipelineHelper.addHandlerAfter(
-                pipeline,
-                "datadog.trace.instrumentation.netty40.server.HttpServerResponseTracingHandler",
-                WebSocketServerResponseTracingHandler.INSTANCE);
+          if (InstrumenterConfig.get().isWebsocketTracingEnabled()) {
+            if (pipeline.get(HttpServerTracingHandler.class) != null) {
+              NettyPipelineHelper.addHandlerAfter(
+                  pipeline,
+                  "datadog.trace.instrumentation.netty40.server.HttpServerTracingHandler",
+                  new WebSocketServerTracingHandler());
+            }
+            if (pipeline.get(HttpRequestDecoder.class) != null) {
+              NettyPipelineHelper.addHandlerAfter(
+                  pipeline,
+                  "datadog.trace.instrumentation.netty40.server.HttpServerRequestTracingHandler",
+                  WebSocketServerRequestTracingHandler.INSTANCE);
+            }
+            if (pipeline.get(HttpResponseEncoder.class) != null) {
+              NettyPipelineHelper.addHandlerAfter(
+                  pipeline,
+                  "datadog.trace.instrumentation.netty40.server.HttpServerResponseTracingHandler",
+                  WebSocketServerResponseTracingHandler.INSTANCE);
+            }
           }
         } else
         // Client pipeline handlers

dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/NettyChannelPipelineInstrumentation.java

@@ -13,6 +13,7 @@
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.bootstrap.CallDepthThreadLocalMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.instrumentation.netty41.client.HttpClientRequestTracingHandler;
@@ -161,21 +162,23 @@ public static void addHandler(
               HttpServerResponseTracingHandler.INSTANCE,
               MaybeBlockResponseHandler.INSTANCE);
         } else if (handler instanceof WebSocketServerProtocolHandler) {
-          if (pipeline.get(HttpServerTracingHandler.class) != null) {
-            NettyPipelineHelper.addHandlerAfter(
-                pipeline, "HttpServerTracingHandler#0", new WebSocketServerTracingHandler());
-          }
-          if (pipeline.get(HttpServerRequestTracingHandler.class) != null) {
-            NettyPipelineHelper.addHandlerAfter(
-                pipeline,
-                "HttpServerRequestTracingHandler#0",
-                WebSocketServerRequestTracingHandler.INSTANCE);
-          }
-          if (pipeline.get(HttpServerResponseTracingHandler.class) != null) {
-            NettyPipelineHelper.addHandlerAfter(
-                pipeline,
-                "HttpServerResponseTracingHandler#0",
-                WebSocketServerResponseTracingHandler.INSTANCE);
+          if (InstrumenterConfig.get().isWebsocketTracingEnabled()) {
+            if (pipeline.get(HttpServerTracingHandler.class) != null) {
+              NettyPipelineHelper.addHandlerAfter(
+                  pipeline, "HttpServerTracingHandler#0", new WebSocketServerTracingHandler());
+            }
+            if (pipeline.get(HttpServerRequestTracingHandler.class) != null) {
+              NettyPipelineHelper.addHandlerAfter(
+                  pipeline,
+                  "HttpServerRequestTracingHandler#0",
+                  WebSocketServerRequestTracingHandler.INSTANCE);
+            }
+            if (pipeline.get(HttpServerResponseTracingHandler.class) != null) {
+              NettyPipelineHelper.addHandlerAfter(
+                  pipeline,
+                  "HttpServerResponseTracingHandler#0",
+                  WebSocketServerResponseTracingHandler.INSTANCE);
+            }
           }
         }
         // Client pipeline handlers