#4056 remove oringin handler when removelast in netty (#4123)

* removelast * removelast * addAfter and removeLast * addAfter and removeLast * addAfter and removeLast
open-telemetry · Sep 17, 2021 · 3bc0e07 · 3bc0e07
1 parent dc4ddf7
commit 3bc0e07
Show file tree

Hide file tree

Showing 3 changed files with 136 additions and 4 deletions.
diff --git a/...y/javaagent/instrumentation/netty/common/AbstractNettyChannelPipelineInstrumentation.java b/...y/javaagent/instrumentation/netty/common/AbstractNettyChannelPipelineInstrumentation.java
@@ -11,6 +11,7 @@
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.returns;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import io.netty.channel.ChannelHandler;
 import io.netty.channel.ChannelPipeline;
@@ -47,11 +48,18 @@ public void transform(TypeTransformer transformer) {
     transformer.applyAdviceToMethod(
         isMethod().and(named("remove").or(named("replace"))).and(takesArgument(0, Class.class)),
         AbstractNettyChannelPipelineInstrumentation.class.getName() + "$RemoveByClassAdvice");
+    transformer.applyAdviceToMethod(
+        isMethod().and(named("removeFirst")).and(returns(named("io.netty.channel.ChannelHandler"))),
+        AbstractNettyChannelPipelineInstrumentation.class.getName() + "$RemoveFirstAdvice");
+    transformer.applyAdviceToMethod(
+        isMethod().and(named("removeLast")).and(returns(named("io.netty.channel.ChannelHandler"))),
+        AbstractNettyChannelPipelineInstrumentation.class.getName() + "$RemoveLastAdvice");
     transformer.applyAdviceToMethod(
         isMethod()
-            .and(named("removeFirst").or(named("removeLast")))
-            .and(returns(named("io.netty.channel.ChannelHandler"))),
-        AbstractNettyChannelPipelineInstrumentation.class.getName() + "$RemoveFirstLastAdvice");
+            .and(named("addAfter"))
+            .and(takesArgument(1, String.class))
+            .and(takesArguments(4)),
+        AbstractNettyChannelPipelineInstrumentation.class.getName() + "$AddAfterAdvice");
   }
 
   @SuppressWarnings("unused")
@@ -114,7 +122,23 @@ public static void removeHandler(
   }
 
   @SuppressWarnings("unused")
-  public static class RemoveFirstLastAdvice {
+  public static class RemoveFirstAdvice {
+
+    @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);
+      }
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class RemoveLastAdvice {
 
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void removeHandler(
@@ -125,6 +149,30 @@ public static void removeHandler(
       if (ourHandler != null) {
         pipeline.remove(ourHandler);
         contextStore.put(handler, null);
+      } else if (handler
+          .getClass()
+          .getName()
+          .startsWith("io.opentelemetry.javaagent.instrumentation.netty.")) {
+        pipeline.removeLast();
+      }
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class AddAfterAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void addAfterHandler(
+        @Advice.This ChannelPipeline pipeline,
+        @Advice.Argument(value = 1, readOnly = false) String name) {
+      ChannelHandler handler = pipeline.get(name);
+      if (handler != null) {
+        ContextStore<ChannelHandler, ChannelHandler> contextStore =
+            InstrumentationContext.get(ChannelHandler.class, ChannelHandler.class);
+        ChannelHandler ourHandler = contextStore.get(handler);
+        if (ourHandler != null) {
+          name = ourHandler.getClass().getName();
+        }
       }
     }
   }

diff --git a/instrumentation/netty/netty-4.0/javaagent/src/test/groovy/ChannelPipelineTest.groovy b/instrumentation/netty/netty-4.0/javaagent/src/test/groovy/ChannelPipelineTest.groovy
@@ -89,6 +89,48 @@ class ChannelPipelineTest extends AgentInstrumentationSpecification {
     "by name"     | { pipeline, oldName, oldHandler, newName, newHandler -> pipeline.replace(oldName, newName, newHandler) }
   }
 
+  // regression test for https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/4056
+  def "should addAfter and removeLast 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:
+    channelPipeline.addLast("http", httpHandler)
+
+    then: "add http and instrumentation handlers"
+    channelPipeline.size() == 2
+    channelPipeline.first() == httpHandler
+    channelPipeline.last().getClass() == HttpClientTracingHandler
+
+    when:
+    def noopHandler = new NoopChannelHandler()
+    channelPipeline.addAfter("http", "noop", noopHandler)
+
+    then: "instrumentation handler is between with http and noop"
+    channelPipeline.size() == 3
+    channelPipeline.first() == httpHandler
+    channelPipeline.last() == noopHandler
+
+    when:
+    channelPipeline.removeLast()
+
+    then: "http and instrumentation handlers will be remained"
+    channelPipeline.size() == 2
+    channelPipeline.first() == httpHandler
+    channelPipeline.last().getClass() == HttpClientTracingHandler
+
+    when:
+    channelPipeline.removeLast()
+
+    then: "there is no handler in pipeline"
+    channelPipeline.size() == 0
+  }
+
   private static class NoopChannelHandler extends ChannelHandlerAdapter {
   }
 }
diff --git a/instrumentation/netty/netty-4.1/javaagent/src/test/groovy/ChannelPipelineTest.groovy b/instrumentation/netty/netty-4.1/javaagent/src/test/groovy/ChannelPipelineTest.groovy
@@ -87,6 +87,48 @@ class ChannelPipelineTest extends AgentInstrumentationSpecification {
     "by name"     | { pipeline, oldName, oldHandler, newName, newHandler -> pipeline.replace(oldName, newName, newHandler) }
   }
 
+  // regression test for https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/4056
+  def "should addAfter and removeLast handler #desc"() {
+    setup:
+    def channel = new EmbeddedChannel()
+    def channelPipeline = new DefaultChannelPipeline(channel)
+    def httpHandler = new HttpClientCodec()
+
+    expect: "no handlers initially"
+    channelPipeline.size() == 0
+
+    when:
+    channelPipeline.addLast("http", httpHandler)
+
+    then: "add http and instrumentation handlers"
+    channelPipeline.size() == 2
+    channelPipeline.first() == httpHandler
+    channelPipeline.last().getClass() == HttpClientTracingHandler
+
+    when:
+    def noopHandler = new NoopChannelHandler()
+    channelPipeline.addAfter("http", "noop", noopHandler)
+
+    then: "instrumentation handler is between with http and noop"
+    channelPipeline.size() == 3
+    channelPipeline.first() == httpHandler
+    channelPipeline.last() == noopHandler
+
+    when:
+    channelPipeline.removeLast()
+
+    then: "http and instrumentation handlers will be remained"
+    channelPipeline.size() == 2
+    channelPipeline.first() == httpHandler
+    channelPipeline.last().getClass() == HttpClientTracingHandler
+
+    when:
+    channelPipeline.removeLast()
+
+    then: "there is no handler in pipeline"
+    channelPipeline.size() == 0
+  }
+
   private static class NoopChannelHandler extends ChannelHandlerAdapter {
   }
 }