Skip to content

Commit f82362e

Browse files
authored
fix: CallDepth thread safety (#559)
1 parent 0c269dd commit f82362e

File tree

3 files changed

+18
-2
lines changed

3 files changed

+18
-2
lines changed

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

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
package io.arex.agent.bootstrap.ctx;
22

3+
import io.arex.agent.bootstrap.internal.CallDepth;
4+
35
import java.util.HashMap;
46
import java.util.Iterator;
57
import java.util.Map;
@@ -45,7 +47,11 @@ public void superRemove() {
4547
}
4648

4749
public T copyValue() {
48-
return get();
50+
T value = get();
51+
if (value instanceof CallDepth) {
52+
return (T) ((CallDepth) value).copy();
53+
}
54+
return value;
4955
}
5056

5157
private static final ThreadLocal<WeakHashMap<ArexThreadLocal<Object>, ?>> holder =
@@ -76,7 +82,7 @@ public static class Transmitter {
7682

7783
public static Object capture() {
7884
HashMap<ArexThreadLocal<Object>, Object> values = captureValues();
79-
return values == null ? null : new Snapshot(captureValues());
85+
return values == null ? null : new Snapshot(values);
8086
}
8187

8288
private static HashMap<ArexThreadLocal<Object>, Object> captureValues() {

arex-agent-bootstrap/src/main/java/io/arex/agent/bootstrap/internal/CallDepth.java

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,4 +39,10 @@ public int getAndIncrement() {
3939
public int decrementAndGet() {
4040
return --this.depth;
4141
}
42+
43+
public CallDepth copy() {
44+
CallDepth copy = new CallDepth();
45+
copy.depth = this.depth;
46+
return copy;
47+
}
4248
}

arex-instrumentation/internal/arex-executors/src/main/java/io/arex/inst/executors/ForkJoinTaskConstructorInstrumentation.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package io.arex.inst.executors;
22

3+
import io.arex.agent.bootstrap.TraceContextManager;
34
import io.arex.agent.bootstrap.ctx.ArexThreadLocal;
45
import io.arex.agent.bootstrap.internal.Cache;
56
import io.arex.inst.extension.MethodInstrumentation;
@@ -31,6 +32,9 @@ public List<MethodInstrumentation> methodAdvices() {
3132
public static class ConstructorAdvice {
3233
@Advice.OnMethodExit(suppress = Throwable.class)
3334
public static void onExit(@Advice.This Object task) {
35+
if (TraceContextManager.get() == null) {
36+
return;
37+
}
3438
final Object captured = ArexThreadLocal.Transmitter.capture();
3539
if (captured != null) {
3640
Cache.CAPTURED_CACHE.put(task, captured);

0 commit comments

Comments
 (0)