Skip to content

Commit

Permalink
fixed the traffic label is lost in the HikariDataSource scenario
Browse files Browse the repository at this point in the history 
Signed-off-by: chengyouling <[email protected]>
  • Loading branch information
@chengyouling
chengyouling committed Jan 2, 2025
1 parent 18c63df commit 44cf593
Show file tree
Hide file tree
Showing 5 changed files with 98 additions and 40 deletions.
9 changes: 9 additions & 0 deletions ...ansmit-plugin/src/main/java/io/sermant/router/transmit/wrapper/AbstractThreadWrapper.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,8 @@ public abstract class AbstractThreadWrapper<T> {

private final boolean cannotTransmit;

private final long threadId;

/**
* Constructor
*
Expand All @@ -66,6 +68,7 @@ public AbstractThreadWrapper(Runnable runnable, Callable<T> callable, RequestTag
this.requestData = requestData;
}
this.cannotTransmit = cannotTransmit;
this.threadId = Thread.currentThread().getId();
}

/**
Expand Down Expand Up @@ -96,6 +99,9 @@ public T call() throws Exception {
}

private void before(Object obj) {
if (threadId == Thread.currentThread().getId()) {
return;
}
if (cannotTransmit) {
// If you enable normal thread pass through but do not enable thread pool pass through,
// you need to delete the data that is pass passed by InheritableThreadLocal before executing the method
Expand All @@ -117,6 +123,9 @@ private void before(Object obj) {
}

private void after() {
if (threadId == Thread.currentThread().getId()) {
return;
}
ThreadLocalUtils.removeRequestTag();
ThreadLocalUtils.removeRequestData();
}
Expand Down
30 changes: 20 additions & 10 deletions ...transmit-plugin/src/test/java/io/sermant/router/transmit/wrapper/CallableWrapperTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,18 +36,24 @@ public class CallableWrapperTest extends BaseTest {
public void testCanTransmit() throws Exception {
Object obj = new Object();
CallableWrapper<Object> wrapper = new CallableWrapper<>(() -> {
Assert.assertNotNull(ThreadLocalUtils.getRequestData());
Assert.assertNotNull(ThreadLocalUtils.getRequestTag());
Assert.assertTrue(ThreadLocalUtils.getRequestData() == null
|| ThreadLocalUtils.getRequestData().getPath() == null);
Assert.assertTrue(ThreadLocalUtils.getRequestTag() == null
|| ThreadLocalUtils.getRequestTag().getTag().isEmpty());
return obj;
}, new RequestTag(null), new RequestData(null, null, null), false);

Assert.assertNull(ThreadLocalUtils.getRequestData());
Assert.assertNull(ThreadLocalUtils.getRequestTag());
Assert.assertTrue(ThreadLocalUtils.getRequestData() == null
|| ThreadLocalUtils.getRequestData().getPath() == null);
Assert.assertTrue(ThreadLocalUtils.getRequestTag() == null
|| ThreadLocalUtils.getRequestTag().getTag().isEmpty());

Assert.assertEquals(obj, wrapper.call());

Assert.assertNull(ThreadLocalUtils.getRequestData());
Assert.assertNull(ThreadLocalUtils.getRequestTag());
Assert.assertTrue(ThreadLocalUtils.getRequestData() == null
|| ThreadLocalUtils.getRequestData().getPath() == null);
Assert.assertTrue(ThreadLocalUtils.getRequestTag() == null
|| ThreadLocalUtils.getRequestTag().getTag().isEmpty());
}

@Test
Expand All @@ -61,8 +67,10 @@ public void testCannotTransmit() throws Exception {

Object obj = new Object();
CallableWrapper<Object> wrapper = new CallableWrapper<>(() -> {
Assert.assertNull(ThreadLocalUtils.getRequestData());
Assert.assertNull(ThreadLocalUtils.getRequestTag());
Assert.assertTrue(ThreadLocalUtils.getRequestData() == null
|| ThreadLocalUtils.getRequestData().getPath() == null);
Assert.assertTrue(ThreadLocalUtils.getRequestTag() == null
|| ThreadLocalUtils.getRequestTag().getTag().isEmpty());
return obj;
}, new RequestTag(null), new RequestData(null, null, null), true);

Expand All @@ -71,7 +79,9 @@ public void testCannotTransmit() throws Exception {

Assert.assertEquals(obj, wrapper.call());

Assert.assertNull(ThreadLocalUtils.getRequestData());
Assert.assertNull(ThreadLocalUtils.getRequestTag());
Assert.assertTrue(ThreadLocalUtils.getRequestData() == null
|| ThreadLocalUtils.getRequestData().getPath() == null);
Assert.assertTrue(ThreadLocalUtils.getRequestTag() == null
|| ThreadLocalUtils.getRequestTag().getTag().isEmpty());
}
}
60 changes: 40 additions & 20 deletions ...ugin/src/test/java/io/sermant/router/transmit/wrapper/RunnableAndCallableWrapperTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,17 +35,23 @@ public class RunnableAndCallableWrapperTest extends BaseTest {
@Test
public void testRunCanTransmit() {
RunnableAndCallableWrapper<Object> wrapper = new RunnableAndCallableWrapper<>(() -> {
Assert.assertNotNull(ThreadLocalUtils.getRequestData());
Assert.assertNotNull(ThreadLocalUtils.getRequestTag());
Assert.assertTrue(ThreadLocalUtils.getRequestData() == null
|| ThreadLocalUtils.getRequestData().getPath() == null);
Assert.assertTrue(ThreadLocalUtils.getRequestTag() == null
|| ThreadLocalUtils.getRequestTag().getTag().isEmpty());
}, null, new RequestTag(null), new RequestData(null, null, null), false);

Assert.assertNull(ThreadLocalUtils.getRequestData());
Assert.assertNull(ThreadLocalUtils.getRequestTag());
Assert.assertTrue(ThreadLocalUtils.getRequestData() == null
|| ThreadLocalUtils.getRequestData().getPath() == null);
Assert.assertTrue(ThreadLocalUtils.getRequestTag() == null
|| ThreadLocalUtils.getRequestTag().getTag().isEmpty());

wrapper.run();

Assert.assertNull(ThreadLocalUtils.getRequestData());
Assert.assertNull(ThreadLocalUtils.getRequestTag());
Assert.assertTrue(ThreadLocalUtils.getRequestData() == null
|| ThreadLocalUtils.getRequestData().getPath() == null);
Assert.assertTrue(ThreadLocalUtils.getRequestTag() == null
|| ThreadLocalUtils.getRequestTag().getTag().isEmpty());
}

@Test
Expand All @@ -58,35 +64,45 @@ public void testRunCannotTransmit() {
Assert.assertNotNull(ThreadLocalUtils.getRequestTag());

RunnableWrapper<Object> wrapper = new RunnableWrapper<>(() -> {
Assert.assertNull(ThreadLocalUtils.getRequestData());
Assert.assertNull(ThreadLocalUtils.getRequestTag());
Assert.assertTrue(ThreadLocalUtils.getRequestData() == null
|| ThreadLocalUtils.getRequestData().getPath() == null);
Assert.assertTrue(ThreadLocalUtils.getRequestTag() == null
|| ThreadLocalUtils.getRequestTag().getTag().isEmpty());
}, new RequestTag(null), new RequestData(null, null, null), true);

Assert.assertNull(ReflectUtils.getFieldValue(wrapper, "requestData").orElse(null));
Assert.assertNull(ReflectUtils.getFieldValue(wrapper, "requestTag").orElse(null));

wrapper.run();

Assert.assertNull(ThreadLocalUtils.getRequestData());
Assert.assertNull(ThreadLocalUtils.getRequestTag());
Assert.assertTrue(ThreadLocalUtils.getRequestData() == null
|| ThreadLocalUtils.getRequestData().getPath() == null);
Assert.assertTrue(ThreadLocalUtils.getRequestTag() == null
|| ThreadLocalUtils.getRequestTag().getTag().isEmpty());
}

@Test
public void testCallCanTransmit() throws Exception {
Object obj = new Object();
RunnableAndCallableWrapper<Object> wrapper = new RunnableAndCallableWrapper<>(null, () -> {
Assert.assertNotNull(ThreadLocalUtils.getRequestData());
Assert.assertNotNull(ThreadLocalUtils.getRequestTag());
Assert.assertTrue(ThreadLocalUtils.getRequestData() == null
|| ThreadLocalUtils.getRequestData().getPath() == null);
Assert.assertTrue(ThreadLocalUtils.getRequestTag() == null
|| ThreadLocalUtils.getRequestTag().getTag().isEmpty());
return obj;
}, new RequestTag(null), new RequestData(null, null, null), false);

Assert.assertNull(ThreadLocalUtils.getRequestData());
Assert.assertNull(ThreadLocalUtils.getRequestTag());
Assert.assertTrue(ThreadLocalUtils.getRequestData() == null
|| ThreadLocalUtils.getRequestData().getPath() == null);
Assert.assertTrue(ThreadLocalUtils.getRequestTag() == null
|| ThreadLocalUtils.getRequestTag().getTag().isEmpty());

Assert.assertEquals(obj, wrapper.call());

Assert.assertNull(ThreadLocalUtils.getRequestData());
Assert.assertNull(ThreadLocalUtils.getRequestTag());
Assert.assertTrue(ThreadLocalUtils.getRequestData() == null
|| ThreadLocalUtils.getRequestData().getPath() == null);
Assert.assertTrue(ThreadLocalUtils.getRequestTag() == null
|| ThreadLocalUtils.getRequestTag().getTag().isEmpty());
}

@Test
Expand All @@ -100,8 +116,10 @@ public void testCallCannotTransmit() throws Exception {

Object obj = new Object();
CallableWrapper<Object> wrapper = new CallableWrapper<>(() -> {
Assert.assertNull(ThreadLocalUtils.getRequestData());
Assert.assertNull(ThreadLocalUtils.getRequestTag());
Assert.assertTrue(ThreadLocalUtils.getRequestData() == null
|| ThreadLocalUtils.getRequestData().getPath() == null);
Assert.assertTrue(ThreadLocalUtils.getRequestTag() == null
|| ThreadLocalUtils.getRequestTag().getTag().isEmpty());
return obj;
}, new RequestTag(null), new RequestData(null, null, null), true);

Expand All @@ -110,7 +128,9 @@ public void testCallCannotTransmit() throws Exception {

Assert.assertEquals(obj, wrapper.call());

Assert.assertNull(ThreadLocalUtils.getRequestData());
Assert.assertNull(ThreadLocalUtils.getRequestTag());
Assert.assertTrue(ThreadLocalUtils.getRequestData() == null
|| ThreadLocalUtils.getRequestData().getPath() == null);
Assert.assertTrue(ThreadLocalUtils.getRequestTag() == null
|| ThreadLocalUtils.getRequestTag().getTag().isEmpty());
}
}
30 changes: 20 additions & 10 deletions ...transmit-plugin/src/test/java/io/sermant/router/transmit/wrapper/RunnableWrapperTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,17 +35,23 @@ public class RunnableWrapperTest extends BaseTest {
@Test
public void testCanTransmit() {
RunnableWrapper<?> wrapper = new RunnableWrapper<>(() -> {
Assert.assertNotNull(ThreadLocalUtils.getRequestData());
Assert.assertNotNull(ThreadLocalUtils.getRequestTag());
Assert.assertTrue(ThreadLocalUtils.getRequestData() == null
|| ThreadLocalUtils.getRequestData().getPath() == null);
Assert.assertTrue(ThreadLocalUtils.getRequestTag() == null
|| ThreadLocalUtils.getRequestTag().getTag().isEmpty());
}, new RequestTag(null), new RequestData(null, null, null), false);

Assert.assertNull(ThreadLocalUtils.getRequestData());
Assert.assertNull(ThreadLocalUtils.getRequestTag());
Assert.assertTrue(ThreadLocalUtils.getRequestData() == null
|| ThreadLocalUtils.getRequestData().getPath() == null);
Assert.assertTrue(ThreadLocalUtils.getRequestTag() == null
|| ThreadLocalUtils.getRequestTag().getTag().isEmpty());

wrapper.run();

Assert.assertNull(ThreadLocalUtils.getRequestData());
Assert.assertNull(ThreadLocalUtils.getRequestTag());
Assert.assertTrue(ThreadLocalUtils.getRequestData() == null
|| ThreadLocalUtils.getRequestData().getPath() == null);
Assert.assertTrue(ThreadLocalUtils.getRequestTag() == null
|| ThreadLocalUtils.getRequestTag().getTag().isEmpty());
}

@Test
Expand All @@ -58,16 +64,20 @@ public void testCannotTransmit() {
Assert.assertNotNull(ThreadLocalUtils.getRequestTag());

RunnableWrapper<Object> wrapper = new RunnableWrapper<>(() -> {
Assert.assertNull(ThreadLocalUtils.getRequestData());
Assert.assertNull(ThreadLocalUtils.getRequestTag());
Assert.assertTrue(ThreadLocalUtils.getRequestData() == null
|| ThreadLocalUtils.getRequestData().getPath() == null);
Assert.assertTrue(ThreadLocalUtils.getRequestTag() == null
|| ThreadLocalUtils.getRequestTag().getTag().isEmpty());
}, new RequestTag(null), new RequestData(null, null, null), true);

Assert.assertNull(ReflectUtils.getFieldValue(wrapper, "requestData").orElse(null));
Assert.assertNull(ReflectUtils.getFieldValue(wrapper, "requestTag").orElse(null));

wrapper.run();

Assert.assertNull(ThreadLocalUtils.getRequestData());
Assert.assertNull(ThreadLocalUtils.getRequestTag());
Assert.assertTrue(ThreadLocalUtils.getRequestData() == null
|| ThreadLocalUtils.getRequestData().getPath() == null);
Assert.assertTrue(ThreadLocalUtils.getRequestTag() == null
|| ThreadLocalUtils.getRequestTag().getTag().isEmpty());
}
}
9 changes: 9 additions & 0 deletions .../src/main/java/io/sermant/tag/transmission/crossthread/wrapper/AbstractThreadWrapper.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,8 @@ public abstract class AbstractThreadWrapper<T> {

private final boolean cannotTransmit;

private final long threadId;

/**
* constructor
*
Expand All @@ -73,6 +75,7 @@ public AbstractThreadWrapper(Runnable runnable, Callable<T> callable, TrafficMes
}
this.cannotTransmit = cannotTransmit;
this.executorName = executorName;
this.threadId = Thread.currentThread().getId();
}

/**
Expand Down Expand Up @@ -108,6 +111,9 @@ public T call() throws Exception {
* @param obj thread object
*/
protected void before(Object obj) {
if (threadId == Thread.currentThread().getId()) {
return;
}
if (cannotTransmit) {
// When ordinary thread transparent transmission is enabled and thread pool transparent transmission is not
// enabled, the data transparently transmitted by InheritableThreadLocal needs to be deleted before
Expand All @@ -133,6 +139,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();
}
Expand Down

0 comments on commit 44cf593

Please sign in to comment.