Fix memleak in the Netty ChannelPipeline instrumentation

instrumentation/netty/netty-4-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/netty/common/AbstractNettyChannelPipelineInstrumentation.java

@@ -9,6 +9,7 @@
 import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
 import io.netty.channel.ChannelHandler;
@@ -37,22 +38,24 @@ public ElementMatcher<TypeDescription> typeMatcher() {
   public void transform(TypeTransformer transformer) {
     transformer.applyAdviceToMethod(
         isMethod()
-            .and(named("remove"))
+            .and(named("remove").or(named("replace")))
             .and(takesArgument(0, named("io.netty.channel.ChannelHandler"))),
-        AbstractNettyChannelPipelineInstrumentation.class.getName()
-            + "$ChannelPipelineRemoveAdvice");
+        AbstractNettyChannelPipelineInstrumentation.class.getName() + "$RemoveAdvice");
     transformer.applyAdviceToMethod(
-        isMethod().and(named("remove")).and(takesArgument(0, String.class)),
-        AbstractNettyChannelPipelineInstrumentation.class.getName()
-            + "$ChannelPipelineRemoveByNameAdvice");
+        isMethod().and(named("remove").or(named("replace"))).and(takesArgument(0, String.class)),
+        AbstractNettyChannelPipelineInstrumentation.class.getName() + "$RemoveByNameAdvice");
     transformer.applyAdviceToMethod(
-        isMethod().and(named("remove")).and(takesArgument(0, Class.class)),
-        AbstractNettyChannelPipelineInstrumentation.class.getName()
-            + "$ChannelPipelineRemoveByClassAdvice");
+        isMethod().and(named("remove").or(named("replace"))).and(takesArgument(0, Class.class)),
+        AbstractNettyChannelPipelineInstrumentation.class.getName() + "$RemoveByClassAdvice");
+    transformer.applyAdviceToMethod(
+        isMethod()
+            .and(named("removeFirst").or(named("removeLast")))
+            .and(returns(named("io.netty.channel.ChannelHandler"))),
+        AbstractNettyChannelPipelineInstrumentation.class.getName() + "$RemoveFirstLastAdvice");
   }
 
   @SuppressWarnings("unused")
-  public static class ChannelPipelineRemoveAdvice {
+  public static class RemoveAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void removeHandler(
@@ -68,7 +71,7 @@ public static void removeHandler(
   }
 
   @SuppressWarnings("unused")
-  public static class ChannelPipelineRemoveByNameAdvice {
+  public static class RemoveByNameAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void removeHandler(
@@ -89,7 +92,7 @@ public static void removeHandler(
   }
 
   @SuppressWarnings("unused")
-  public static class ChannelPipelineRemoveByClassAdvice {
+  public static class RemoveByClassAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void removeHandler(
@@ -109,4 +112,20 @@ public static void removeHandler(
       }
     }
   }
+
+  @SuppressWarnings("unused")
+  public static class RemoveFirstLastAdvice {
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void removeHandler(
+        @Advice.This ChannelPipeline pipeline, @Advice.Return 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);
+      }
+    }
+  }
 }

instrumentation/netty/netty-4.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/netty/v4_0/NettyChannelPipelineInstrumentation.java

@@ -39,7 +39,7 @@ public void transform(TypeTransformer transformer) {
 
     transformer.applyAdviceToMethod(
         isMethod()
-            .and(nameStartsWith("add"))
+            .and(nameStartsWith("add").or(named("replace")))
             .and(takesArgument(2, named("io.netty.channel.ChannelHandler"))),
         NettyChannelPipelineInstrumentation.class.getName() + "$ChannelPipelineAddAdvice");
   }

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

@@ -3,21 +3,22 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import io.netty.channel.ChannelHandler
-import io.netty.channel.ChannelHandlerContext
+import io.netty.channel.ChannelHandlerAdapter
 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_0.client.HttpClientTracingHandler
 import spock.lang.Unroll
 
+@Unroll
 class ChannelPipelineTest extends AgentInstrumentationSpecification {
+
   // regression test for https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/1373
-  @Unroll
+  // and https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/4040
   def "test remove our handler #testName"() {
     setup:
-    def channel = new EmbeddedChannel(new EmptyChannelHandler())
+    def channel = new EmbeddedChannel(new NoopChannelHandler())
     def channelPipeline = new DefaultChannelPipeline(channel)
     def handler = new HttpClientCodec()
 
@@ -44,19 +45,50 @@ class ChannelPipelineTest extends AgentInstrumentationSpecification {
     "by instance" | { pipeline, h -> pipeline.remove(h) }
     "by class"    | { pipeline, h -> pipeline.remove(h.getClass()) }
     "by name"     | { pipeline, h -> pipeline.remove("http") }
+    "first"       | { pipeline, h -> pipeline.removeFirst() }
   }
 
-  private static class EmptyChannelHandler implements ChannelHandler {
-    @Override
-    void handlerAdded(ChannelHandlerContext ctx) throws Exception {
-    }
+  // regression test for https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/4040
+  def "should replace handler #desc"() {
+    setup:
+    def channel = new EmbeddedChannel(new NoopChannelHandler())
+    def channelPipeline = new DefaultChannelPipeline(channel)
+    def httpHandler = new HttpClientCodec()
+
+    expect: "no handlers initially"
+    channelPipeline.size() == 0
+
+    when:
+    def noopHandler = new NoopChannelHandler()
+    channelPipeline.addFirst("test", noopHandler)
+
+    then: "only the noop handler"
+    channelPipeline.size() == 1
+    channelPipeline.first() == noopHandler
+
+    when:
+    replaceMethod(channelPipeline, "test", noopHandler, "http", httpHandler)
+
+    then: "noop handler was removed; http and instrumentation handlers were added"
+    channelPipeline.size() == 2
+    channelPipeline.first() == httpHandler
+    channelPipeline.last().getClass() == HttpClientTracingHandler
+
+    when:
+    def anotherNoopHandler = new NoopChannelHandler()
+    replaceMethod(channelPipeline, "http", httpHandler, "test", anotherNoopHandler)
+
+    then: "http and instrumentation handlers were removed; noop handler was added"
+    channelPipeline.size() == 1
+    channelPipeline.first() == anotherNoopHandler
 
-    @Override
-    void handlerRemoved(ChannelHandlerContext ctx) throws Exception {
-    }
+    where:
+    desc          | replaceMethod
+    "by instance" | { pipeline, oldName, oldHandler, newName, newHandler -> pipeline.replace(oldHandler, newName, newHandler) }
+    "by class"    | { pipeline, oldName, oldHandler, newName, newHandler -> pipeline.replace(oldHandler.getClass(), newName, newHandler) }
+    "by name"     | { pipeline, oldName, oldHandler, newName, newHandler -> pipeline.replace(oldName, newName, newHandler) }
+  }
 
-    @Override
-    void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
-    }
+  private static class NoopChannelHandler extends ChannelHandlerAdapter {
   }
 }

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

@@ -40,7 +40,7 @@ public void transform(TypeTransformer transformer) {
 
     transformer.applyAdviceToMethod(
         isMethod()
-            .and(nameStartsWith("add"))
+            .and(nameStartsWith("add").or(named("replace")))
             .and(takesArgument(1, String.class))
             .and(takesArgument(2, named("io.netty.channel.ChannelHandler"))),
         NettyChannelPipelineInstrumentation.class.getName() + "$ChannelPipelineAddAdvice");

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

@@ -3,16 +3,19 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import io.netty.channel.ChannelHandlerAdapter
 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
 import spock.lang.Unroll
 
+@Unroll
 class ChannelPipelineTest extends AgentInstrumentationSpecification {
+
   // regression test for https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/1373
-  @Unroll
+  // and https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/4040
   def "test remove our handler #testName"() {
     setup:
     def channel = new EmbeddedChannel()
@@ -40,5 +43,50 @@ class ChannelPipelineTest extends AgentInstrumentationSpecification {
     "by instance" | { pipeline, h -> pipeline.remove(h) }
     "by class"    | { pipeline, h -> pipeline.remove(h.getClass()) }
     "by name"     | { pipeline, h -> pipeline.remove("http") }
+    "first"       | { pipeline, h -> pipeline.removeFirst() }
+  }
+
+  // regression test for https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/4040
+  def "should replace handler #desc"() {
+    setup:
+    def channel = new EmbeddedChannel()
+    def channelPipeline = new DefaultChannelPipeline(channel)
+    def httpHandler = new HttpClientCodec()
+
+    expect: "no handlers initially"
+    channelPipeline.size() == 0
+
+    when:
+    def noopHandler = new NoopChannelHandler()
+    channelPipeline.addFirst("test", noopHandler)
+
+    then: "only the noop handler"
+    channelPipeline.size() == 1
+    channelPipeline.first() == noopHandler
+
+    when:
+    replaceMethod(channelPipeline, "test", noopHandler, "http", httpHandler)
+
+    then: "noop handler was removed; http and instrumentation handlers were added"
+    channelPipeline.size() == 2
+    channelPipeline.first() == httpHandler
+    channelPipeline.last().getClass() == HttpClientTracingHandler
+
+    when:
+    def anotherNoopHandler = new NoopChannelHandler()
+    replaceMethod(channelPipeline, "http", httpHandler, "test", anotherNoopHandler)
+
+    then: "http and instrumentation handlers were removed; noop handler was added"
+    channelPipeline.size() == 1
+    channelPipeline.first() == anotherNoopHandler
+
+    where:
+    desc          | replaceMethod
+    "by instance" | { pipeline, oldName, oldHandler, newName, newHandler -> pipeline.replace(oldHandler, newName, newHandler) }
+    "by class"    | { pipeline, oldName, oldHandler, newName, newHandler -> pipeline.replace(oldHandler.getClass(), newName, newHandler) }
+    "by name"     | { pipeline, oldName, oldHandler, newName, newHandler -> pipeline.replace(oldName, newName, newHandler) }
+  }
+
+  private static class NoopChannelHandler extends ChannelHandlerAdapter {
   }
 }