Skip to content

Commit 77be7c6

Browse files
authored
Clear context propagation virtual field (#12397)
1 parent 2edd778 commit 77be7c6

File tree

17 files changed

+119
-37
lines changed

17 files changed

+119
-37
lines changed

instrumentation/akka/akka-actor-2.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaDefaultSystemMessageQueueInstrumentation.java

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -58,8 +58,13 @@ public static PropagatedContext enter(@Advice.Argument(1) SystemMessage systemMe
5858

5959
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
6060
public static void exit(
61-
@Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
62-
ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
61+
@Advice.Argument(1) SystemMessage systemMessage,
62+
@Advice.Enter PropagatedContext propagatedContext,
63+
@Advice.Thrown Throwable throwable) {
64+
VirtualField<SystemMessage, PropagatedContext> virtualField =
65+
VirtualField.find(SystemMessage.class, PropagatedContext.class);
66+
ExecutorAdviceHelper.cleanUpAfterSubmit(
67+
propagatedContext, throwable, virtualField, systemMessage);
6368
}
6469
}
6570
}

instrumentation/akka/akka-actor-2.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaDispatcherInstrumentation.java

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -52,8 +52,12 @@ public static PropagatedContext enterDispatch(@Advice.Argument(1) Envelope envel
5252

5353
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
5454
public static void exitDispatch(
55-
@Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
56-
ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
55+
@Advice.Argument(1) Envelope envelope,
56+
@Advice.Enter PropagatedContext propagatedContext,
57+
@Advice.Thrown Throwable throwable) {
58+
VirtualField<Envelope, PropagatedContext> virtualField =
59+
VirtualField.find(Envelope.class, PropagatedContext.class);
60+
ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, envelope);
5761
}
5862
}
5963
}

instrumentation/akka/akka-actor-fork-join-2.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaForkJoinPoolInstrumentation.java

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,8 +60,12 @@ public static PropagatedContext enterJobSubmit(@Advice.Argument(0) ForkJoinTask<
6060

6161
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
6262
public static void exitJobSubmit(
63-
@Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
64-
ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
63+
@Advice.Argument(0) ForkJoinTask<?> task,
64+
@Advice.Enter PropagatedContext propagatedContext,
65+
@Advice.Thrown Throwable throwable) {
66+
VirtualField<ForkJoinTask<?>, PropagatedContext> virtualField =
67+
VirtualField.find(ForkJoinTask.class, PropagatedContext.class);
68+
ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, task);
6569
}
6670
}
6771
}

instrumentation/executors/bootstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/executors/ExecutorAdviceHelper.java

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -94,8 +94,11 @@ public static <T> PropagatedContext attachContextToTask(
9494
* Clean up {@code propagatedContext} in case of any submission errors. Call this method after the
9595
* submission method has exited.
9696
*/
97-
public static void cleanUpAfterSubmit(
98-
@Nullable PropagatedContext propagatedContext, @Nullable Throwable throwable) {
97+
public static <T> void cleanUpAfterSubmit(
98+
@Nullable PropagatedContext propagatedContext,
99+
@Nullable Throwable throwable,
100+
VirtualField<T, PropagatedContext> virtualField,
101+
T task) {
99102
if (propagatedContext != null && throwable != null) {
100103
/*
101104
Note: this may potentially clear somebody else's parent span if we didn't set it
@@ -106,6 +109,8 @@ public static void cleanUpAfterSubmit(
106109
exceptions.
107110
*/
108111
propagatedContext.clear();
112+
// setting the field to null removes it from the fallback map
113+
virtualField.set(task, null);
109114
}
110115
}
111116

@@ -119,6 +124,8 @@ public static <T> void cleanPropagatedContext(
119124
PropagatedContext propagatedContext = virtualField.get(task);
120125
if (propagatedContext != null) {
121126
propagatedContext.clear();
127+
// setting the field to null removes it from the fallback map
128+
virtualField.set(task, null);
122129
}
123130
}
124131

instrumentation/executors/bootstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/executors/TaskAdviceHelper.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@ public static <T> Scope makePropagatedContextCurrent(
2222
VirtualField<T, PropagatedContext> virtualField, T task) {
2323
PropagatedContext propagatedContext = virtualField.get(task);
2424
if (propagatedContext != null) {
25+
// setting the field to null removes it from the fallback map
26+
virtualField.set(task, null);
2527
Context context = propagatedContext.getAndClear();
2628
if (context != null) {
2729
return context.makeCurrent();

instrumentation/executors/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/executors/JavaExecutorInstrumentation.java

Lines changed: 22 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -105,8 +105,12 @@ public static PropagatedContext enterJobSubmit(
105105

106106
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
107107
public static void exitJobSubmit(
108-
@Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
109-
ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
108+
@Advice.Argument(0) Runnable task,
109+
@Advice.Enter PropagatedContext propagatedContext,
110+
@Advice.Thrown Throwable throwable) {
111+
VirtualField<Runnable, PropagatedContext> virtualField =
112+
VirtualField.find(Runnable.class, PropagatedContext.class);
113+
ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, task);
110114
}
111115
}
112116

@@ -126,8 +130,12 @@ public static PropagatedContext enterJobSubmit(@Advice.Argument(0) ForkJoinTask<
126130

127131
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
128132
public static void exitJobSubmit(
129-
@Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
130-
ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
133+
@Advice.Argument(0) ForkJoinTask<?> task,
134+
@Advice.Enter PropagatedContext propagatedContext,
135+
@Advice.Thrown Throwable throwable) {
136+
VirtualField<ForkJoinTask<?>, PropagatedContext> virtualField =
137+
VirtualField.find(ForkJoinTask.class, PropagatedContext.class);
138+
ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, task);
131139
}
132140
}
133141

@@ -148,6 +156,7 @@ public static PropagatedContext enterJobSubmit(
148156

149157
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
150158
public static void exitJobSubmit(
159+
@Advice.Argument(0) Runnable task,
151160
@Advice.Enter PropagatedContext propagatedContext,
152161
@Advice.Thrown Throwable throwable,
153162
@Advice.Return Future<?> future) {
@@ -156,7 +165,9 @@ public static void exitJobSubmit(
156165
VirtualField.find(Future.class, PropagatedContext.class);
157166
virtualField.set(future, propagatedContext);
158167
}
159-
ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
168+
VirtualField<Runnable, PropagatedContext> virtualField =
169+
VirtualField.find(Runnable.class, PropagatedContext.class);
170+
ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, task);
160171
}
161172
}
162173

@@ -176,6 +187,7 @@ public static PropagatedContext enterJobSubmit(@Advice.Argument(0) Callable<?> t
176187

177188
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
178189
public static void exitJobSubmit(
190+
@Advice.Argument(0) Callable<?> task,
179191
@Advice.Enter PropagatedContext propagatedContext,
180192
@Advice.Thrown Throwable throwable,
181193
@Advice.Return Future<?> future) {
@@ -184,7 +196,9 @@ public static void exitJobSubmit(
184196
VirtualField.find(Future.class, PropagatedContext.class);
185197
virtualField.set(future, propagatedContext);
186198
}
187-
ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
199+
VirtualField<Callable<?>, PropagatedContext> virtualField =
200+
VirtualField.find(Callable.class, PropagatedContext.class);
201+
ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, task);
188202
}
189203
}
190204

@@ -230,7 +244,8 @@ public static void submitExit(
230244
VirtualField<Callable<?>, PropagatedContext> virtualField =
231245
VirtualField.find(Callable.class, PropagatedContext.class);
232246
PropagatedContext propagatedContext = virtualField.get(task);
233-
ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
247+
ExecutorAdviceHelper.cleanUpAfterSubmit(
248+
propagatedContext, throwable, virtualField, task);
234249
}
235250
}
236251
}

instrumentation/executors/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/executors/JavaForkJoinTaskInstrumentation.java

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -117,8 +117,12 @@ public static PropagatedContext enterFork(@Advice.This ForkJoinTask<?> task) {
117117

118118
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
119119
public static void exitFork(
120-
@Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
121-
ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
120+
@Advice.This ForkJoinTask<?> task,
121+
@Advice.Enter PropagatedContext propagatedContext,
122+
@Advice.Thrown Throwable throwable) {
123+
VirtualField<ForkJoinTask<?>, PropagatedContext> virtualField =
124+
VirtualField.find(ForkJoinTask.class, PropagatedContext.class);
125+
ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, task);
122126
}
123127
}
124128
}

instrumentation/executors/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/executors/StructuredTaskScopeInstrumentation.java

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,8 +50,12 @@ public static PropagatedContext enterCallableFork(@Advice.Argument(0) Callable<?
5050

5151
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
5252
public static void exitCallableFork(
53-
@Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
54-
ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
53+
@Advice.Argument(0) Callable<?> task,
54+
@Advice.Enter PropagatedContext propagatedContext,
55+
@Advice.Thrown Throwable throwable) {
56+
VirtualField<Callable<?>, PropagatedContext> virtualField =
57+
VirtualField.find(Callable.class, PropagatedContext.class);
58+
ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, task);
5559
}
5660
}
5761
}

instrumentation/guava-10.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/guava/v10_0/GuavaListenableFutureInstrumentation.java

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -61,8 +61,12 @@ public static PropagatedContext addListenerEnter(@Advice.Argument(0) Runnable ta
6161

6262
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
6363
public static void addListenerExit(
64-
@Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
65-
ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
64+
@Advice.Argument(0) Runnable task,
65+
@Advice.Enter PropagatedContext propagatedContext,
66+
@Advice.Thrown Throwable throwable) {
67+
VirtualField<Runnable, PropagatedContext> virtualField =
68+
VirtualField.find(Runnable.class, PropagatedContext.class);
69+
ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, task);
6670
}
6771
}
6872
}

instrumentation/jetty/jetty-8.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jetty/v8_0/JettyQueuedThreadPoolInstrumentation.java

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,8 +50,12 @@ public static PropagatedContext enterJobSubmit(@Advice.Argument(0) Runnable task
5050

5151
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
5252
public static void exitJobSubmit(
53-
@Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
54-
ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
53+
@Advice.Argument(0) Runnable task,
54+
@Advice.Enter PropagatedContext propagatedContext,
55+
@Advice.Thrown Throwable throwable) {
56+
VirtualField<Runnable, PropagatedContext> virtualField =
57+
VirtualField.find(Runnable.class, PropagatedContext.class);
58+
ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, task);
5559
}
5660
}
5761
}

0 commit comments

Comments
 (0)