Merge pull request #1661 from hanbingleixue/2.1.x

Fixed an issue where DubboProviderInterceptor remove tag on the consumer side. As a result, the traffic tag is removed and cannot be obtained from thread local

Author: lilai23

sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo2.6.x-plugin/src/main/java/io/sermant/tag/transmission/alibabadubbo/interceptors/AlibabaDubboProviderInterceptor.java

@@ -84,6 +84,10 @@ protected ExecuteContext doBefore(ExecuteContext context) {
 
     @Override
     protected ExecuteContext doAfter(ExecuteContext context) {
+        if (isConsumer(context)) {
+            return context;
+        }
+
         TrafficUtils.removeTrafficTag();
         return context;
     }
@@ -123,6 +127,9 @@ protected Map<String, List<String>> extractTrafficTagFromCarrier(RpcInvocation i
 
     @Override
     public ExecuteContext onThrow(ExecuteContext context) {
+        if (isConsumer(context)) {
+            return context;
+        }
         TrafficUtils.removeTrafficTag();
         return context;
     }

sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo2.6.x-plugin/src/test/java/io/sermant/tag/transmission/alibabadubbo/interceptors/AlibabaDubboProviderInterceptorTest.java

@@ -93,6 +93,32 @@ public void testAlibabaDubboProvider() {
         interceptor.after(returnContext);
     }
 
+    @Test
+    public void testConsumerSideRemoveTrafficTag() {
+        // If interceptor is invoked in consumer side, it should not remove traffic tag in after method.
+        Map<String, List<String>> expectTag = buildExpectTrafficTag("id", "name");
+        TrafficUtils.setTrafficTag(new TrafficTag(expectTag));
+
+        ExecuteContext context = buildContext(new RpcInvocation(), new HashMap<>(), "consumer");
+        interceptor.before(context);
+        Assert.assertEquals(TrafficUtils.getTrafficTag().getTag(), expectTag);
+        interceptor.after(context);
+        Assert.assertEquals(TrafficUtils.getTrafficTag().getTag(), expectTag);
+    }
+
+    @Test
+    public void testProviderSideRemoveTrafficTag() {
+        // If interceptor is invoked in provider side, it should remove traffic tag in after method.
+        Map<String, List<String>> expectTag = buildExpectTrafficTag("id", "name");
+        TrafficUtils.setTrafficTag(new TrafficTag(expectTag));
+
+        ExecuteContext context = buildContext(new RpcInvocation(), new HashMap<>(), "provider");
+        interceptor.before(context);
+        Assert.assertEquals(TrafficUtils.getTrafficTag().getTag(), expectTag);
+        interceptor.after(context);
+        Assert.assertNull(TrafficUtils.getTrafficTag());
+    }
+
     private ExecuteContext buildContext(RpcInvocation rpcInvocation, Map<String, String> headers, String side) {
         URL url = new URL("http", "127.0.0.1", 8080);
         url = url.addParameter("side", side);

sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo2.7.x-plugin/src/main/java/io/sermant/tag/transmission/apachedubbov2/interceptors/ApacheDubboProviderInterceptor.java

@@ -83,12 +83,18 @@ protected ExecuteContext doBefore(ExecuteContext context) {
 
     @Override
     protected ExecuteContext doAfter(ExecuteContext context) {
+        if (isConsumer(context)) {
+            return context;
+        }
         TrafficUtils.removeTrafficTag();
         return context;
     }
 
     @Override
     public ExecuteContext onThrow(ExecuteContext context) {
+        if (isConsumer(context)) {
+            return context;
+        }
         TrafficUtils.removeTrafficTag();
         return context;
     }

sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo2.7.x-plugin/src/test/java/io/sermant/tag/transmission/apachedubbov2/interceptors/ApacheDubboProviderInterceptorTest.java

@@ -92,6 +92,32 @@ public void testApacheDubboProvider() {
         interceptor.after(returnContext);
     }
 
+    @Test
+    public void testConsumerSideRemoveTrafficTag() {
+        // If interceptor is invoked in consumer side, it should not remove traffic tag in after method.
+        Map<String, List<String>> expectTag = buildExpectTrafficTag("id", "name");
+        TrafficUtils.setTrafficTag(new TrafficTag(expectTag));
+
+        ExecuteContext context = buildContext(new RpcInvocation(), new HashMap<>(), "consumer");
+        interceptor.before(context);
+        Assert.assertEquals(TrafficUtils.getTrafficTag().getTag(), expectTag);
+        interceptor.after(context);
+        Assert.assertEquals(TrafficUtils.getTrafficTag().getTag(), expectTag);
+    }
+
+    @Test
+    public void testProviderSideRemoveTrafficTag() {
+        // If interceptor is invoked in provider side, it should remove traffic tag in after method.
+        Map<String, List<String>> expectTag = buildExpectTrafficTag("id", "name");
+        TrafficUtils.setTrafficTag(new TrafficTag(expectTag));
+
+        ExecuteContext context = buildContext(new RpcInvocation(), new HashMap<>(), "provider");
+        interceptor.before(context);
+        Assert.assertEquals(TrafficUtils.getTrafficTag().getTag(), expectTag);
+        interceptor.after(context);
+        Assert.assertNull(TrafficUtils.getTrafficTag());
+    }
+
     private ExecuteContext buildContext(RpcInvocation rpcInvocation, Map<String, String> headers, String side) {
         URL url = new URL("http", "127.0.0.1", 8080);
         url = url.addParameter("side", side);

sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo3.x-plugin/src/main/java/io/sermant/tag/transmission/dubbov3/interceptors/ApacheDubboProviderInterceptor.java

@@ -83,12 +83,18 @@ protected ExecuteContext doBefore(ExecuteContext context) {
 
     @Override
     protected ExecuteContext doAfter(ExecuteContext context) {
+        if (isConsumer(context)) {
+            return context;
+        }
         TrafficUtils.removeTrafficTag();
         return context;
     }
 
     @Override
     public ExecuteContext onThrow(ExecuteContext context) {
+        if (isConsumer(context)) {
+            return context;
+        }
         TrafficUtils.removeTrafficTag();
         return context;
     }

sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo3.x-plugin/src/test/java/io/sermant/tag/transmission/dubbov3/interceptors/ApacheDubboProviderInterceptorTest.java

@@ -92,6 +92,32 @@ public void testApacheDubboProvider() {
         interceptor.after(returnContext);
     }
 
+    @Test
+    public void testConsumerSideRemoveTrafficTag() {
+        // If interceptor is invoked in consumer side, it should not remove traffic tag in after method.
+        Map<String, List<String>> expectTag = buildExpectTrafficTag("id", "name");
+        TrafficUtils.setTrafficTag(new TrafficTag(expectTag));
+
+        ExecuteContext context = buildContext(new RpcInvocation(), new HashMap<>(), "consumer");
+        interceptor.before(context);
+        Assert.assertEquals(TrafficUtils.getTrafficTag().getTag(), expectTag);
+        interceptor.after(context);
+        Assert.assertEquals(TrafficUtils.getTrafficTag().getTag(), expectTag);
+    }
+
+    @Test
+    public void testProviderSideRemoveTrafficTag() {
+        // If interceptor is invoked in provider side, it should remove traffic tag in after method.
+        Map<String, List<String>> expectTag = buildExpectTrafficTag("id", "name");
+        TrafficUtils.setTrafficTag(new TrafficTag(expectTag));
+
+        ExecuteContext context = buildContext(new RpcInvocation(), new HashMap<>(), "provider");
+        interceptor.before(context);
+        Assert.assertEquals(TrafficUtils.getTrafficTag().getTag(), expectTag);
+        interceptor.after(context);
+        Assert.assertNull(TrafficUtils.getTrafficTag());
+    }
+
     private ExecuteContext buildContext(RpcInvocation rpcInvocation, Map<String, String> headers, String side) {
         URL url = new URL("http", "127.0.0.1", 8080);
         url = url.addParameter("side", side);