Fix remove tag issue for DubboProviderInterceptor in consumer side, i…

…t will cause previously traffic tag removed and the follow up invocation can not get traffic tag from thread local Signed-off-by: rztao <[email protected]>
provenceee · Nov 15, 2024 · ec2d183 · ec2d183
1 parent 87035c5
commit ec2d183
Show file tree

Hide file tree

Showing 6 changed files with 97 additions and 0 deletions.
diff --git a/...o/sermant/tag/transmission/alibabadubbo/interceptors/AlibabaDubboProviderInterceptor.java b/...o/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;
     }

diff --git a/...rmant/tag/transmission/alibabadubbo/interceptors/AlibabaDubboProviderInterceptorTest.java b/...rmant/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);

diff --git a/...o/sermant/tag/transmission/apachedubbov2/interceptors/ApacheDubboProviderInterceptor.java b/...o/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;
     }

diff --git a/...rmant/tag/transmission/apachedubbov2/interceptors/ApacheDubboProviderInterceptorTest.java b/...rmant/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);

diff --git a/...java/io/sermant/tag/transmission/dubbov3/interceptors/ApacheDubboProviderInterceptor.java b/...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;
     }

diff --git a/.../io/sermant/tag/transmission/dubbov3/interceptors/ApacheDubboProviderInterceptorTest.java b/.../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);