From 7b788546a32d9cd434fa08e1dc88a29a596e64a4 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 4ca7e76232..9676e5eba2 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 86c60facf6..e10a1878e8 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 58a026e144..e4a4e58b39 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 abf59766d2..121c00bfe3 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 e5758f70a9..39336a7166 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);