From a5be709284d2ba00596beaf75691f86554a37c7e Mon Sep 17 00:00:00 2001 From: rztao Date: Thu, 24 Oct 2024 00:09:41 +0800 Subject: [PATCH] Fix remove tag issue for DubboProviderInterceptor in consumer side, it will cause previously traffic tag removed and the follow up invocation can not get traffic tag from thread local Signed-off-by: rztao --- .../AlibabaDubboProviderInterceptor.java | 7 +++++ .../AlibabaDubboProviderInterceptorTest.java | 26 +++++++++++++++++++ .../ApacheDubboProviderInterceptor.java | 6 +++++ .../ApacheDubboProviderInterceptorTest.java | 26 +++++++++++++++++++ .../ApacheDubboProviderInterceptor.java | 6 +++++ .../ApacheDubboProviderInterceptorTest.java | 26 +++++++++++++++++++ 6 files changed, 97 insertions(+) diff --git a/sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo2.6.x-plugin/src/main/java/io/sermant/tag/transmission/alibabadubbo/interceptors/AlibabaDubboProviderInterceptor.java b/sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo2.6.x-plugin/src/main/java/io/sermant/tag/transmission/alibabadubbo/interceptors/AlibabaDubboProviderInterceptor.java index 45c891f22a..11bcf94fb1 100644 --- a/sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo2.6.x-plugin/src/main/java/io/sermant/tag/transmission/alibabadubbo/interceptors/AlibabaDubboProviderInterceptor.java +++ b/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> extractTrafficTagFromCarrier(RpcInvocation i @Override public ExecuteContext onThrow(ExecuteContext context) { + if (isConsumer(context)) { + return context; + } TrafficUtils.removeTrafficTag(); return context; } diff --git a/sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo2.6.x-plugin/src/test/java/io/sermant/tag/transmission/alibabadubbo/interceptors/AlibabaDubboProviderInterceptorTest.java b/sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo2.6.x-plugin/src/test/java/io/sermant/tag/transmission/alibabadubbo/interceptors/AlibabaDubboProviderInterceptorTest.java index 5089474bff..61c907ba81 100644 --- a/sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo2.6.x-plugin/src/test/java/io/sermant/tag/transmission/alibabadubbo/interceptors/AlibabaDubboProviderInterceptorTest.java +++ b/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> 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> 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 headers, String side) { URL url = new URL("http", "127.0.0.1", 8080); url = url.addParameter("side", side); diff --git a/sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo2.7.x-plugin/src/main/java/io/sermant/tag/transmission/apachedubbov2/interceptors/ApacheDubboProviderInterceptor.java b/sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo2.7.x-plugin/src/main/java/io/sermant/tag/transmission/apachedubbov2/interceptors/ApacheDubboProviderInterceptor.java index d3121c700b..52df52e9cf 100644 --- a/sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo2.7.x-plugin/src/main/java/io/sermant/tag/transmission/apachedubbov2/interceptors/ApacheDubboProviderInterceptor.java +++ b/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; } diff --git a/sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo2.7.x-plugin/src/test/java/io/sermant/tag/transmission/apachedubbov2/interceptors/ApacheDubboProviderInterceptorTest.java b/sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo2.7.x-plugin/src/test/java/io/sermant/tag/transmission/apachedubbov2/interceptors/ApacheDubboProviderInterceptorTest.java index badb537fc2..28db05a657 100644 --- a/sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo2.7.x-plugin/src/test/java/io/sermant/tag/transmission/apachedubbov2/interceptors/ApacheDubboProviderInterceptorTest.java +++ b/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> 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> 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 headers, String side) { URL url = new URL("http", "127.0.0.1", 8080); url = url.addParameter("side", side); diff --git a/sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo3.x-plugin/src/main/java/io/sermant/tag/transmission/dubbov3/interceptors/ApacheDubboProviderInterceptor.java b/sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo3.x-plugin/src/main/java/io/sermant/tag/transmission/dubbov3/interceptors/ApacheDubboProviderInterceptor.java index f75e98311e..cad77d4e57 100644 --- a/sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo3.x-plugin/src/main/java/io/sermant/tag/transmission/dubbov3/interceptors/ApacheDubboProviderInterceptor.java +++ b/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; } diff --git a/sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo3.x-plugin/src/test/java/io/sermant/tag/transmission/dubbov3/interceptors/ApacheDubboProviderInterceptorTest.java b/sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo3.x-plugin/src/test/java/io/sermant/tag/transmission/dubbov3/interceptors/ApacheDubboProviderInterceptorTest.java index 41dee4e0a0..1f9a9f61c9 100644 --- a/sermant-plugins/sermant-tag-transmission/tag-transmission-dubbo3.x-plugin/src/test/java/io/sermant/tag/transmission/dubbov3/interceptors/ApacheDubboProviderInterceptorTest.java +++ b/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> 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> 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 headers, String side) { URL url = new URL("http", "127.0.0.1", 8080); url = url.addParameter("side", side);