fixed the traffic label is lost in the HikariDataSource scenario

Signed-off-by: chengyouling <[email protected]>
sermant-io · Jan 2, 2025 · 284bf61 · 284bf61
1 parent 18c63df
commit 284bf61
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/...ansmit-plugin/src/main/java/io/sermant/router/transmit/wrapper/AbstractThreadWrapper.java b/...ansmit-plugin/src/main/java/io/sermant/router/transmit/wrapper/AbstractThreadWrapper.java
@@ -45,6 +45,8 @@ public abstract class AbstractThreadWrapper<T> {
 
     private final boolean cannotTransmit;
 
+    private final long threadId;
+
     /**
      * Constructor
      *
@@ -66,6 +68,7 @@ public AbstractThreadWrapper(Runnable runnable, Callable<T> callable, RequestTag
             this.requestData = requestData;
         }
         this.cannotTransmit = cannotTransmit;
+        this.threadId = Thread.currentThread().getId();
     }
 
     /**
@@ -117,6 +120,9 @@ private void before(Object obj) {
     }
 
     private void after() {
+        if (threadId == Thread.currentThread().getId()) {
+            return;
+        }
         ThreadLocalUtils.removeRequestTag();
         ThreadLocalUtils.removeRequestData();
     }

diff --git a/.../src/main/java/io/sermant/tag/transmission/crossthread/wrapper/AbstractThreadWrapper.java b/.../src/main/java/io/sermant/tag/transmission/crossthread/wrapper/AbstractThreadWrapper.java
@@ -51,6 +51,8 @@ public abstract class AbstractThreadWrapper<T> {
 
     private final boolean cannotTransmit;
 
+    private final long threadId;
+
     /**
      * constructor
      *
@@ -73,6 +75,7 @@ public AbstractThreadWrapper(Runnable runnable, Callable<T> callable, TrafficMes
         }
         this.cannotTransmit = cannotTransmit;
         this.executorName = executorName;
+        this.threadId = Thread.currentThread().getId();
     }
 
     /**
@@ -133,6 +136,9 @@ protected void before(Object obj) {
      * Post method executed by thread object
      */
     protected void after() {
+        if (threadId == Thread.currentThread().getId()) {
+            return;
+        }
         TrafficUtils.removeTrafficTag();
         TrafficUtils.removeTrafficData();
     }