Merge pull request #1652 from rztao/fix-remove-traffic-tag-issue-for-…

…dubbo-provider-interceptor-in-consumer-side Fix remove tag issue for DubboProviderInterceptor in consumer side
sermant-io · Oct 29, 2024 · b8f7c5a · b8f7c5a
2 parents dd4a1cd + a5be709
commit b8f7c5a
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);