Netty4.1: remove our handler when original handler is removed

instrumentation/netty/netty-4.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/netty/v4_1/NettyChannelPipelineInstrumentation.java

@@ -28,6 +28,8 @@
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import io.opentelemetry.javaagent.instrumentation.api.CallDepthThreadLocalMap;
+import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
+import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
 import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.instrumentation.netty.v4_1.client.HttpClientRequestTracingHandler;
 import io.opentelemetry.javaagent.instrumentation.netty.v4_1.client.HttpClientResponseTracingHandler;
@@ -59,15 +61,27 @@ public void transform(TypeTransformer transformer) {
             .and(takesArgument(1, String.class))
             .and(takesArgument(2, named("io.netty.channel.ChannelHandler"))),
         NettyChannelPipelineInstrumentation.class.getName() + "$ChannelPipelineAddAdvice");
+    transformer.applyAdviceToMethod(
+        isMethod()
+            .and(named("remove"))
+            .and(takesArgument(0, named("io.netty.channel.ChannelHandler"))),
+        NettyChannelPipelineInstrumentation.class.getName() + "$ChannelPipelineRemoveAdvice");
+    transformer.applyAdviceToMethod(
+        isMethod().and(named("remove")).and(takesArgument(0, String.class)),
+        NettyChannelPipelineInstrumentation.class.getName() + "$ChannelPipelineRemoveByNameAdvice");
+    transformer.applyAdviceToMethod(
+        isMethod().and(named("remove")).and(takesArgument(0, Class.class)),
+        NettyChannelPipelineInstrumentation.class.getName()
+            + "$ChannelPipelineRemoveByClassAdvice");
     transformer.applyAdviceToMethod(
         isMethod().and(named("connect")).and(returns(named("io.netty.channel.ChannelFuture"))),
         NettyChannelPipelineInstrumentation.class.getName() + "$ChannelPipelineConnectAdvice");
   }
 
   /**
    * When certain handlers are added to the pipeline, we want to add our corresponding tracing
-   * handlers. If those handlers are later removed, we may want to remove our handlers. That is not
-   * currently implemented.
+   * handlers. If those handlers are later removed, we also remove our handlers. Support for
+   * replacing handlers and removeFirst/removeLast is currently not implemented.
    */
   public static class ChannelPipelineAddAdvice {
     @Advice.OnMethodEnter
@@ -106,39 +120,84 @@ public static void addHandler(
         name = context.name();
       }
 
-      try {
-        // Server pipeline handlers
-        if (handler instanceof HttpServerCodec) {
-          pipeline.addAfter(
-              name, HttpServerTracingHandler.class.getName(), new HttpServerTracingHandler());
-        } else if (handler instanceof HttpRequestDecoder) {
-          pipeline.addAfter(
-              name,
-              HttpServerRequestTracingHandler.class.getName(),
-              new HttpServerRequestTracingHandler());
-        } else if (handler instanceof HttpResponseEncoder) {
-          pipeline.addAfter(
-              name,
-              HttpServerResponseTracingHandler.class.getName(),
-              new HttpServerResponseTracingHandler());
-        } else
+      ChannelHandler ourHandler = null;
+      // Server pipeline handlers
+      if (handler instanceof HttpServerCodec) {
+        ourHandler = new HttpServerTracingHandler();
+      } else if (handler instanceof HttpRequestDecoder) {
+        ourHandler = new HttpServerRequestTracingHandler();
+      } else if (handler instanceof HttpResponseEncoder) {
+        ourHandler = new HttpServerResponseTracingHandler();
         // Client pipeline handlers
-        if (handler instanceof HttpClientCodec) {
-          pipeline.addAfter(
-              name, HttpClientTracingHandler.class.getName(), new HttpClientTracingHandler());
-        } else if (handler instanceof HttpRequestEncoder) {
-          pipeline.addAfter(
-              name,
-              HttpClientRequestTracingHandler.class.getName(),
-              new HttpClientRequestTracingHandler());
-        } else if (handler instanceof HttpResponseDecoder) {
-          pipeline.addAfter(
-              name,
-              HttpClientResponseTracingHandler.class.getName(),
-              new HttpClientResponseTracingHandler());
+      } else if (handler instanceof HttpClientCodec) {
+        ourHandler = new HttpClientTracingHandler();
+      } else if (handler instanceof HttpRequestEncoder) {
+        ourHandler = new HttpClientRequestTracingHandler();
+      } else if (handler instanceof HttpResponseDecoder) {
+        ourHandler = new HttpClientResponseTracingHandler();
+      }
+
+      if (ourHandler != null) {
+        try {
+          pipeline.addAfter(name, ourHandler.getClass().getName(), ourHandler);
+          InstrumentationContext.get(ChannelHandler.class, ChannelHandler.class)
+              .putIfAbsent(handler, ourHandler);
+        } catch (IllegalArgumentException e) {
+          // Prevented adding duplicate handlers.
         }
-      } catch (IllegalArgumentException e) {
-        // Prevented adding duplicate handlers.
+      }
+    }
+  }
+
+  public static class ChannelPipelineRemoveAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void removeHandler(
+        @Advice.This ChannelPipeline pipeline, @Advice.Argument(0) ChannelHandler handler) {
+      ContextStore<ChannelHandler, ChannelHandler> contextStore =
+          InstrumentationContext.get(ChannelHandler.class, ChannelHandler.class);
+      ChannelHandler ourHandler = contextStore.get(handler);
+      if (ourHandler != null) {
+        pipeline.remove(ourHandler);
+        contextStore.put(handler, null);
+      }
+    }
+  }
+
+  public static class ChannelPipelineRemoveByNameAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void removeHandler(
+        @Advice.This ChannelPipeline pipeline, @Advice.Argument(0) String name) {
+      ChannelHandler handler = pipeline.get(name);
+      if (handler == null) {
+        return;
+      }
+
+      ContextStore<ChannelHandler, ChannelHandler> contextStore =
+          InstrumentationContext.get(ChannelHandler.class, ChannelHandler.class);
+      ChannelHandler ourHandler = contextStore.get(handler);
+      if (ourHandler != null) {
+        pipeline.remove(ourHandler);
+        contextStore.put(handler, null);
+      }
+    }
+  }
+
+  public static class ChannelPipelineRemoveByClassAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void removeHandler(
+        @Advice.This ChannelPipeline pipeline,
+        @Advice.Argument(0) Class<ChannelHandler> handlerClass) {
+      ChannelHandler handler = pipeline.get(handlerClass);
+      if (handler == null) {
+        return;
+      }
+
+      ContextStore<ChannelHandler, ChannelHandler> contextStore =
+          InstrumentationContext.get(ChannelHandler.class, ChannelHandler.class);
+      ChannelHandler ourHandler = contextStore.get(handler);
+      if (ourHandler != null) {
+        pipeline.remove(ourHandler);
+        contextStore.put(handler, null);
       }
     }
   }

instrumentation/netty/netty-4.1/javaagent/src/test/groovy/ChannelPipelineTest.groovy

@@ -0,0 +1,42 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import io.netty.channel.DefaultChannelPipeline
+import io.netty.channel.embedded.EmbeddedChannel
+import io.netty.handler.codec.http.HttpClientCodec
+import io.opentelemetry.instrumentation.test.AgentInstrumentationSpecification
+import io.opentelemetry.javaagent.instrumentation.netty.v4_1.client.HttpClientTracingHandler
+
+class ChannelPipelineTest extends AgentInstrumentationSpecification {
+  // regression test for https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/1373
+  def "test remove our handler #testName"() {
+    setup:
+    def channel = new EmbeddedChannel()
+    def channelPipeline = new DefaultChannelPipeline(channel)
+    def handler = new HttpClientCodec()
+
+    when:
+    // no handlers
+    channelPipeline.first() == null
+
+    then:
+    // add handler
+    channelPipeline.addLast("http", handler)
+    channelPipeline.first() == handler
+    // our handler was also added
+    channelPipeline.last().getClass() == HttpClientTracingHandler
+
+    and:
+    removeMethod.call(channelPipeline, handler)
+    // removing handler also removes our handler
+    channelPipeline.first() == null
+
+    where:
+    testName      | removeMethod
+    "by instance" | { pipeline, h -> pipeline.remove(h) }
+    "by class"    | { pipeline, h -> pipeline.remove(h.getClass()) }
+    "by name"     | { pipeline, h -> pipeline.remove("http") }
+  }
+}

instrumentation/ratpack-1.4/javaagent/src/test/groovy/client/RatpackHttpClientTest.groovy

@@ -32,6 +32,10 @@ class RatpackHttpClientTest extends HttpClientTest<Void> implements AgentTestTra
     if (HttpClientSpec.metaClass.getMetaMethod("execController") != null) {
       it.execController(exec.getController())
     }
+    configureClient(it)
+  }
+
+  void configureClient(HttpClientSpec spec) {
   }
 
   @Override

instrumentation/ratpack-1.4/javaagent/src/test/groovy/client/RatpackPooledHttpClientTest.groovy

@@ -0,0 +1,16 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package client
+
+import ratpack.http.client.HttpClientSpec
+
+class RatpackPooledHttpClientTest extends RatpackHttpClientTest {
+
+  @Override
+  void configureClient(HttpClientSpec spec) {
+    spec.poolSize(5)
+  }
+}