Skip to content

Commit 4645aac

Browse files
lucas-myxmengqcc
lucas-myx
and
mengqcc
authored
feat: optimize trace context transmit (#171)
* enhancemeng: Optimizes the performance of ForkJoinTask by utilizing caching, not in the final version * feat: remove ForkJoinPool transform * feat: add UT * feat: correct UT * feat: restore InstrumentationInstaller ignore * feat: restore InstrumentationInstaller ignore format * feat: remove unused import --------- Co-authored-by: mengqcc <[email protected]>
1 parent cab7e95 commit 4645aac

File tree

20 files changed

+353
-71
lines changed

20 files changed

+353
-71
lines changed

‎arex-agent-bootstrap/src/main/java/io/arex/agent/bootstrap/ctx/ArexThreadLocal.java

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
* 3. transmittable-thread-local not work in nio.reactor.Worker（@see AbstractMultiworkerIOReactor）
1313
* 4. Change from InheritableThreadLocal to ThreadLocal，avoid collect unexpected data
1414
*/
15-
public class ArexThreadLocal<T> extends InheritableThreadLocal<T> {
15+
public class ArexThreadLocal<T> extends ThreadLocal<T> {
1616

1717
public ArexThreadLocal() {
1818
}
@@ -40,15 +40,6 @@ public final void remove() {
4040
super.remove();
4141
}
4242

43-
@Override
44-
protected T childValue(T parentValue) {
45-
// no need to copy arex-record-id, TraceTransmitter will capture and replay
46-
if (parentValue instanceof String && ((String) parentValue).startsWith("AREX-")) {
47-
return null;
48-
}
49-
return super.childValue(parentValue);
50-
}
51-
5243
public void superRemove() {
5344
super.remove();
5445
}
@@ -84,28 +75,34 @@ private void removeThisFromHolder() {
8475
public static class Transmitter {
8576

8677
public static Object capture() {
87-
return new Snapshot(captureValues());
78+
HashMap<ArexThreadLocal<Object>, Object> values = captureValues();
79+
return values == null ? null : new Snapshot(captureValues());
8880
}
8981

9082
private static HashMap<ArexThreadLocal<Object>, Object> captureValues() {
91-
HashMap<ArexThreadLocal<Object>, Object> values = new HashMap<>();
83+
HashMap<ArexThreadLocal<Object>, Object> values = null;
9284
for (ArexThreadLocal<Object> threadLocal : holder.get().keySet()) {
85+
if (null == values) {
86+
values = new HashMap<>();
87+
}
9388
values.put(threadLocal, threadLocal.copyValue());
9489
}
9590
return values;
9691
}
9792

9893
public static Object replay(Object captured) {
94+
if (captured == null) {
95+
return null;
96+
}
97+
9998
final Snapshot capturedSnapshot = (Snapshot) captured;
10099
return new Snapshot(replayValues(capturedSnapshot.values));
101100
}
102101

103102
private static HashMap<ArexThreadLocal<Object>, Object> replayValues(HashMap<ArexThreadLocal<Object>, Object> captured) {
104103
HashMap<ArexThreadLocal<Object>, Object> backup = new HashMap<>();
105-
106104
for (final Iterator<ArexThreadLocal<Object>> iterator = holder.get().keySet().iterator(); iterator.hasNext(); ) {
107105
ArexThreadLocal<Object> threadLocal = iterator.next();
108-
109106
backup.put(threadLocal, threadLocal.get());
110107
if (!captured.containsKey(threadLocal)) {
111108
iterator.remove();
@@ -118,14 +115,17 @@ private static HashMap<ArexThreadLocal<Object>, Object> replayValues(HashMap<Are
118115
}
119116

120117
public static void restore(Object backup) {
118+
if (backup == null) {
119+
return;
120+
}
121+
121122
final Snapshot backupSnapshot = (Snapshot) backup;
122123
restoreValues(backupSnapshot.values);
123124
}
124125

125126
private static void restoreValues(HashMap<ArexThreadLocal<Object>, Object> backup) {
126127
for (final Iterator<ArexThreadLocal<Object>> iterator = holder.get().keySet().iterator(); iterator.hasNext(); ) {
127128
ArexThreadLocal<Object> threadLocal = iterator.next();
128-
129129
if (!backup.containsKey(threadLocal)) {
130130
iterator.remove();
131131
threadLocal.superRemove();

‎arex-agent-bootstrap/src/main/java/io/arex/agent/bootstrap/ctx/CallableWrapper.java

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import io.arex.agent.bootstrap.TraceContextManager;
44

55
import java.util.concurrent.Callable;
6+
import java.util.concurrent.ForkJoinTask;
67

78

89
public class CallableWrapper<V> implements Callable<V> {
@@ -47,8 +48,8 @@ public static <T> Callable<T> get(Callable<T> callable) {
4748
return callable;
4849
}
4950

50-
if (callable instanceof CallableWrapper) {
51-
return (CallableWrapper<T>) callable;
51+
if (callable instanceof CallableWrapper || callable instanceof ForkJoinTask) {
52+
return callable;
5253
}
5354
return new CallableWrapper<T>(callable);
5455
}

‎arex-agent-bootstrap/src/main/java/io/arex/agent/bootstrap/ctx/RunnableWrapper.java

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22

33
import io.arex.agent.bootstrap.TraceContextManager;
44

5+
import java.util.concurrent.ForkJoinTask;
6+
57

68
public class RunnableWrapper implements Runnable {
79
private final Runnable runnable;
@@ -43,7 +45,7 @@ public static Runnable get(Runnable runnable) {
4345
return runnable;
4446
}
4547

46-
if (runnable instanceof RunnableWrapper) {
48+
if (runnable instanceof RunnableWrapper || runnable instanceof ForkJoinTask) {
4749
return runnable;
4850
}
4951
return new RunnableWrapper(runnable);

‎arex-agent-bootstrap/src/main/java/io/arex/agent/bootstrap/ctx/TraceTransmitter.java

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,12 @@ public class TraceTransmitter implements AutoCloseable {
88
private AtomicReference<Object> backupRef;
99

1010
private TraceTransmitter() {
11-
this.captureRef = new AtomicReference<>(ArexThreadLocal.Transmitter.capture());
11+
this.captureRef = null;
12+
this.backupRef = null;
13+
}
14+
15+
private TraceTransmitter(Object capture) {
16+
this.captureRef = new AtomicReference<>(capture);
1217
this.backupRef = new AtomicReference<>();
1318
}
1419

@@ -29,6 +34,23 @@ public void close() {
2934
}
3035

3136
public static TraceTransmitter create() {
32-
return new TraceTransmitter();
37+
Object capture = ArexThreadLocal.Transmitter.capture();
38+
if (capture == null) {
39+
return DoNothingTransmitter.INSTANCE;
40+
}
41+
return new TraceTransmitter(capture);
42+
}
43+
44+
static class DoNothingTransmitter extends TraceTransmitter {
45+
static final TraceTransmitter INSTANCE = new DoNothingTransmitter();
46+
47+
@Override
48+
public TraceTransmitter transmit() {
49+
return this;
50+
}
51+
52+
@Override
53+
public void close() {
54+
}
3355
}
3456
}

‎arex-agent-bootstrap/src/test/java/io/arex/agent/bootstrap/ctx/ArexThreadLocalTest.java

Lines changed: 38 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,47 @@
11
package io.arex.agent.bootstrap.ctx;
22

3-
import org.junit.jupiter.params.ParameterizedTest;
4-
import org.junit.jupiter.params.provider.CsvSource;
3+
import org.junit.jupiter.api.AfterAll;
4+
import org.junit.jupiter.api.BeforeAll;
5+
import org.junit.jupiter.api.Test;
6+
7+
import java.util.HashMap;
58

69
import static org.junit.jupiter.api.Assertions.*;
710

811
class ArexThreadLocalTest {
912

10-
static final ArexThreadLocal<String> TRACE_CONTEXT = new ArexThreadLocal<>();
11-
12-
@ParameterizedTest
13-
@CsvSource({
14-
"AREX-mock, null",
15-
"mock, mock"
16-
})
17-
void childValue(String request, String expect) {
18-
String result = String.valueOf(TRACE_CONTEXT.childValue(request));
19-
System.out.println(result);
20-
assertEquals(expect, result);
13+
static ArexThreadLocal<String> target = null;
14+
15+
static HashMap<ArexThreadLocal<Object>, Object> snapshotMap = null;
16+
17+
@BeforeAll
18+
static void setUp() {
19+
target = new ArexThreadLocal<>();
20+
snapshotMap = new HashMap<>();
21+
}
22+
23+
@AfterAll
24+
static void tearDown() {
25+
target = null;
26+
snapshotMap = null;
27+
}
28+
29+
@Test
30+
void capture() {
31+
target.set("mock");
32+
assertNotNull(ArexThreadLocal.Transmitter.capture());
33+
}
34+
35+
@Test
36+
void replay() {
37+
assertNull(ArexThreadLocal.Transmitter.replay(null));
38+
assertNotNull(ArexThreadLocal.Transmitter.replay(new ArexThreadLocal.Transmitter.Snapshot(snapshotMap)));
39+
}
40+
41+
@Test
42+
void restore() {
43+
assertDoesNotThrow(() -> ArexThreadLocal.Transmitter.restore(null));
44+
assertDoesNotThrow(() -> ArexThreadLocal.Transmitter.restore(
45+
new ArexThreadLocal.Transmitter.Snapshot(snapshotMap)));
2146
}
2247
}

‎arex-agent-bootstrap/src/test/java/io/arex/agent/bootstrap/ctx/CallableWrapperTest.java

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
package io.arex.agent.bootstrap.ctx;
2+
3+
import io.arex.agent.bootstrap.TraceContextManager;
4+
import org.junit.jupiter.api.Test;
5+
6+
import java.util.concurrent.Callable;
7+
import java.util.concurrent.ForkJoinTask;
8+
import java.util.concurrent.RunnableFuture;
9+
10+
import static org.junit.jupiter.api.Assertions.*;
11+
12+
class CallableWrapperTest {
13+
14+
@Test
15+
void get() {
16+
assertNull(CallableWrapper.get(null));
17+
TraceContextManager.set("mock");
18+
assertNotNull(CallableWrapper.get(new CallableTest<>()));
19+
assertNotNull(CallableWrapper.get(() -> "mock"));
20+
TraceContextManager.remove();
21+
}
22+
23+
static class CallableTest<T> extends ForkJoinTask<T> implements Callable<T> {
24+
public final T getRawResult() { return null; }
25+
public final void setRawResult(T v) {}
26+
public final boolean exec() { return true; }
27+
public final T call() { return null; }
28+
}
29+
}

‎arex-agent-bootstrap/src/test/java/io/arex/agent/bootstrap/ctx/RunnableWrapperTest.java

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
package io.arex.agent.bootstrap.ctx;
2+
3+
import io.arex.agent.bootstrap.TraceContextManager;
4+
import org.junit.jupiter.api.Test;
5+
6+
import java.util.concurrent.ForkJoinTask;
7+
import java.util.concurrent.RunnableFuture;
8+
9+
import static org.junit.jupiter.api.Assertions.*;
10+
11+
class RunnableWrapperTest {
12+
13+
@Test
14+
void get() {
15+
assertNull(RunnableWrapper.get(null));
16+
TraceContextManager.set("mock");
17+
assertNotNull(RunnableWrapper.get(new RunnableTest<>()));
18+
assertNotNull(RunnableWrapper.get(() -> {}));
19+
TraceContextManager.remove();
20+
}
21+
22+
static class RunnableTest<T> extends ForkJoinTask<T> implements RunnableFuture<T> {
23+
public final T getRawResult() { return null; }
24+
public final void setRawResult(T v) {}
25+
public final boolean exec() { return true; }
26+
public final void run() {}
27+
}
28+
}

‎arex-agent-bootstrap/src/test/java/io/arex/agent/bootstrap/ctx/TraceTransmitterTest.java

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
package io.arex.agent.bootstrap.ctx;
2+
3+
import org.junit.jupiter.api.Test;
4+
import org.mockito.Mockito;
5+
6+
import static org.junit.jupiter.api.Assertions.*;
7+
8+
class TraceTransmitterTest {
9+
10+
@Test
11+
void create() {
12+
assertNotNull(TraceTransmitter.create());
13+
Mockito.mockStatic(ArexThreadLocal.Transmitter.class);
14+
Mockito.when(ArexThreadLocal.Transmitter.capture()).thenReturn("mock");
15+
TraceTransmitter transmitter = TraceTransmitter.create();
16+
assertNotNull(transmitter);
17+
transmitter.transmit();
18+
transmitter.close();
19+
}
20+
21+
@Test
22+
void doNothingTransmit() {
23+
TraceTransmitter.DoNothingTransmitter.INSTANCE.transmit();
24+
}
25+
26+
@Test
27+
void doNothingClose() {
28+
TraceTransmitter.DoNothingTransmitter.INSTANCE.close();
29+
}
30+
}

‎arex-instrumentation-api/src/main/java/io/arex/inst/extension/matcher/IgnoredTypesMatcher.java

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
package io.arex.inst.extension.matcher;
2+
3+
import net.bytebuddy.description.type.TypeDescription;
4+
import net.bytebuddy.matcher.ElementMatcher;
5+
6+
import java.util.Arrays;
7+
import java.util.List;
8+
9+
public class IgnoredTypesMatcher extends ElementMatcher.Junction.AbstractBase<TypeDescription> {
10+
11+
private static final List<String> IGNORED_TYPES = Arrays.asList(
12+
"net.bytebuddy.",
13+
"io.arex.",
14+
"sun.reflect.",
15+
"com.intellij.",
16+
"shaded.");
17+
18+
@Override
19+
public boolean matches(TypeDescription target) {
20+
String name = target.getActualName();
21+
for (String ignoredType : IGNORED_TYPES) {
22+
if (name.startsWith(ignoredType)) {
23+
return true;
24+
}
25+
}
26+
return false;
27+
}
28+
}

‎arex-instrumentation-api/src/test/java/io/arex/inst/extension/matcher/IgnoredTypesMatcherTest.java

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
package io.arex.inst.extension.matcher;
2+
3+
import net.bytebuddy.description.type.TypeDescription;
4+
import org.junit.jupiter.api.Test;
5+
6+
import static org.junit.jupiter.api.Assertions.*;
7+
8+
class IgnoredTypesMatcherTest {
9+
10+
@Test
11+
void matches() {
12+
IgnoredTypesMatcher matcher = new IgnoredTypesMatcher();
13+
assertTrue(matcher.matches(new TypeDescription.ForLoadedType(IgnoredTypesMatcher.class)));
14+
assertFalse(matcher.matches(new TypeDescription.ForLoadedType(String.class)));
15+
}
16+
}

0 commit comments

Comments
 (0)