Clear context propagation virtual field (#12397)

open-telemetry · Oct 8, 2024 · 77be7c6 · 77be7c6
1 parent 2edd778
commit 77be7c6
Show file tree

Hide file tree

Showing 17 changed files with 119 additions and 37 deletions.
diff --git a/...try/javaagent/instrumentation/akkaactor/AkkaDefaultSystemMessageQueueInstrumentation.java b/...try/javaagent/instrumentation/akkaactor/AkkaDefaultSystemMessageQueueInstrumentation.java
@@ -58,8 +58,13 @@ public static PropagatedContext enter(@Advice.Argument(1) SystemMessage systemMe
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void exit(
-        @Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
-      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
+        @Advice.Argument(1) SystemMessage systemMessage,
+        @Advice.Enter PropagatedContext propagatedContext,
+        @Advice.Thrown Throwable throwable) {
+      VirtualField<SystemMessage, PropagatedContext> virtualField =
+          VirtualField.find(SystemMessage.class, PropagatedContext.class);
+      ExecutorAdviceHelper.cleanUpAfterSubmit(
+          propagatedContext, throwable, virtualField, systemMessage);
     }
   }
 }
diff --git a/...a/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaDispatcherInstrumentation.java b/...a/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaDispatcherInstrumentation.java
@@ -52,8 +52,12 @@ public static PropagatedContext enterDispatch(@Advice.Argument(1) Envelope envel
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void exitDispatch(
-        @Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
-      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
+        @Advice.Argument(1) Envelope envelope,
+        @Advice.Enter PropagatedContext propagatedContext,
+        @Advice.Thrown Throwable throwable) {
+      VirtualField<Envelope, PropagatedContext> virtualField =
+          VirtualField.find(Envelope.class, PropagatedContext.class);
+      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, envelope);
     }
   }
 }
diff --git a/...io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaForkJoinPoolInstrumentation.java b/...io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaForkJoinPoolInstrumentation.java
@@ -60,8 +60,12 @@ public static PropagatedContext enterJobSubmit(@Advice.Argument(0) ForkJoinTask<
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void exitJobSubmit(
-        @Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
-      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
+        @Advice.Argument(0) ForkJoinTask<?> task,
+        @Advice.Enter PropagatedContext propagatedContext,
+        @Advice.Thrown Throwable throwable) {
+      VirtualField<ForkJoinTask<?>, PropagatedContext> virtualField =
+          VirtualField.find(ForkJoinTask.class, PropagatedContext.class);
+      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, task);
     }
   }
 }
diff --git a/...ap/src/main/java/io/opentelemetry/javaagent/bootstrap/executors/ExecutorAdviceHelper.java b/...ap/src/main/java/io/opentelemetry/javaagent/bootstrap/executors/ExecutorAdviceHelper.java
@@ -94,8 +94,11 @@ public static <T> PropagatedContext attachContextToTask(
    * Clean up {@code propagatedContext} in case of any submission errors. Call this method after the
    * submission method has exited.
    */
-  public static void cleanUpAfterSubmit(
-      @Nullable PropagatedContext propagatedContext, @Nullable Throwable throwable) {
+  public static <T> void cleanUpAfterSubmit(
+      @Nullable PropagatedContext propagatedContext,
+      @Nullable Throwable throwable,
+      VirtualField<T, PropagatedContext> virtualField,
+      T task) {
     if (propagatedContext != null && throwable != null) {
       /*
       Note: this may potentially clear somebody else's parent span if we didn't set it
@@ -106,6 +109,8 @@ public static void cleanUpAfterSubmit(
       exceptions.
        */
       propagatedContext.clear();
+      // setting the field to null removes it from the fallback map
+      virtualField.set(task, null);
     }
   }
 
@@ -119,6 +124,8 @@ public static <T> void cleanPropagatedContext(
     PropagatedContext propagatedContext = virtualField.get(task);
     if (propagatedContext != null) {
       propagatedContext.clear();
+      // setting the field to null removes it from the fallback map
+      virtualField.set(task, null);
     }
   }
 

diff --git a/...tstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/executors/TaskAdviceHelper.java b/...tstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/executors/TaskAdviceHelper.java
@@ -22,6 +22,8 @@ public static <T> Scope makePropagatedContextCurrent(
       VirtualField<T, PropagatedContext> virtualField, T task) {
     PropagatedContext propagatedContext = virtualField.get(task);
     if (propagatedContext != null) {
+      // setting the field to null removes it from the fallback map
+      virtualField.set(task, null);
       Context context = propagatedContext.getAndClear();
       if (context != null) {
         return context.makeCurrent();

diff --git a/...ava/io/opentelemetry/javaagent/instrumentation/executors/JavaExecutorInstrumentation.java b/...ava/io/opentelemetry/javaagent/instrumentation/executors/JavaExecutorInstrumentation.java
@@ -105,8 +105,12 @@ public static PropagatedContext enterJobSubmit(
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void exitJobSubmit(
-        @Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
-      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
+        @Advice.Argument(0) Runnable task,
+        @Advice.Enter PropagatedContext propagatedContext,
+        @Advice.Thrown Throwable throwable) {
+      VirtualField<Runnable, PropagatedContext> virtualField =
+          VirtualField.find(Runnable.class, PropagatedContext.class);
+      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, task);
     }
   }
 
@@ -126,8 +130,12 @@ public static PropagatedContext enterJobSubmit(@Advice.Argument(0) ForkJoinTask<
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void exitJobSubmit(
-        @Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
-      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
+        @Advice.Argument(0) ForkJoinTask<?> task,
+        @Advice.Enter PropagatedContext propagatedContext,
+        @Advice.Thrown Throwable throwable) {
+      VirtualField<ForkJoinTask<?>, PropagatedContext> virtualField =
+          VirtualField.find(ForkJoinTask.class, PropagatedContext.class);
+      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, task);
     }
   }
 
@@ -148,6 +156,7 @@ public static PropagatedContext enterJobSubmit(
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void exitJobSubmit(
+        @Advice.Argument(0) Runnable task,
         @Advice.Enter PropagatedContext propagatedContext,
         @Advice.Thrown Throwable throwable,
         @Advice.Return Future<?> future) {
@@ -156,7 +165,9 @@ public static void exitJobSubmit(
             VirtualField.find(Future.class, PropagatedContext.class);
         virtualField.set(future, propagatedContext);
       }
-      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
+      VirtualField<Runnable, PropagatedContext> virtualField =
+          VirtualField.find(Runnable.class, PropagatedContext.class);
+      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, task);
     }
   }
 
@@ -176,6 +187,7 @@ public static PropagatedContext enterJobSubmit(@Advice.Argument(0) Callable<?> t
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void exitJobSubmit(
+        @Advice.Argument(0) Callable<?> task,
         @Advice.Enter PropagatedContext propagatedContext,
         @Advice.Thrown Throwable throwable,
         @Advice.Return Future<?> future) {
@@ -184,7 +196,9 @@ public static void exitJobSubmit(
             VirtualField.find(Future.class, PropagatedContext.class);
         virtualField.set(future, propagatedContext);
       }
-      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
+      VirtualField<Callable<?>, PropagatedContext> virtualField =
+          VirtualField.find(Callable.class, PropagatedContext.class);
+      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, task);
     }
   }
 
@@ -230,7 +244,8 @@ public static void submitExit(
             VirtualField<Callable<?>, PropagatedContext> virtualField =
                 VirtualField.find(Callable.class, PropagatedContext.class);
             PropagatedContext propagatedContext = virtualField.get(task);
-            ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
+            ExecutorAdviceHelper.cleanUpAfterSubmit(
+                propagatedContext, throwable, virtualField, task);
           }
         }
       }

diff --git a/...io/opentelemetry/javaagent/instrumentation/executors/JavaForkJoinTaskInstrumentation.java b/...io/opentelemetry/javaagent/instrumentation/executors/JavaForkJoinTaskInstrumentation.java
@@ -117,8 +117,12 @@ public static PropagatedContext enterFork(@Advice.This ForkJoinTask<?> task) {
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void exitFork(
-        @Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
-      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
+        @Advice.This ForkJoinTask<?> task,
+        @Advice.Enter PropagatedContext propagatedContext,
+        @Advice.Thrown Throwable throwable) {
+      VirtualField<ForkJoinTask<?>, PropagatedContext> virtualField =
+          VirtualField.find(ForkJoinTask.class, PropagatedContext.class);
+      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, task);
     }
   }
 }
diff --git a/...opentelemetry/javaagent/instrumentation/executors/StructuredTaskScopeInstrumentation.java b/...opentelemetry/javaagent/instrumentation/executors/StructuredTaskScopeInstrumentation.java
@@ -50,8 +50,12 @@ public static PropagatedContext enterCallableFork(@Advice.Argument(0) Callable<?
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void exitCallableFork(
-        @Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
-      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
+        @Advice.Argument(0) Callable<?> task,
+        @Advice.Enter PropagatedContext propagatedContext,
+        @Advice.Thrown Throwable throwable) {
+      VirtualField<Callable<?>, PropagatedContext> virtualField =
+          VirtualField.find(Callable.class, PropagatedContext.class);
+      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, task);
     }
   }
 }
diff --git a/...telemetry/javaagent/instrumentation/guava/v10_0/GuavaListenableFutureInstrumentation.java b/...telemetry/javaagent/instrumentation/guava/v10_0/GuavaListenableFutureInstrumentation.java
@@ -61,8 +61,12 @@ public static PropagatedContext addListenerEnter(@Advice.Argument(0) Runnable ta
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void addListenerExit(
-        @Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
-      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
+        @Advice.Argument(0) Runnable task,
+        @Advice.Enter PropagatedContext propagatedContext,
+        @Advice.Thrown Throwable throwable) {
+      VirtualField<Runnable, PropagatedContext> virtualField =
+          VirtualField.find(Runnable.class, PropagatedContext.class);
+      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, task);
     }
   }
 }
diff --git a/...ntelemetry/javaagent/instrumentation/jetty/v8_0/JettyQueuedThreadPoolInstrumentation.java b/...ntelemetry/javaagent/instrumentation/jetty/v8_0/JettyQueuedThreadPoolInstrumentation.java
@@ -50,8 +50,12 @@ public static PropagatedContext enterJobSubmit(@Advice.Argument(0) Runnable task
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void exitJobSubmit(
-        @Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
-      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
+        @Advice.Argument(0) Runnable task,
+        @Advice.Enter PropagatedContext propagatedContext,
+        @Advice.Thrown Throwable throwable) {
+      VirtualField<Runnable, PropagatedContext> virtualField =
+          VirtualField.find(Runnable.class, PropagatedContext.class);
+      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, task);
     }
   }
 }
diff --git a/...ava/io/opentelemetry/javaagent/instrumentation/okhttp/v2_2/DispatcherInstrumentation.java b/...ava/io/opentelemetry/javaagent/instrumentation/okhttp/v2_2/DispatcherInstrumentation.java
@@ -48,8 +48,12 @@ public static PropagatedContext onEnter(@Advice.Argument(0) Runnable call) {
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void onExit(
-        @Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
-      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
+        @Advice.Argument(0) Runnable call,
+        @Advice.Enter PropagatedContext propagatedContext,
+        @Advice.Thrown Throwable throwable) {
+      VirtualField<Runnable, PropagatedContext> virtualField =
+          VirtualField.find(Runnable.class, PropagatedContext.class);
+      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, call);
     }
   }
 }
diff --git a/...opentelemetry/javaagent/instrumentation/okhttp/v3_0/OkHttp3DispatcherInstrumentation.java b/...opentelemetry/javaagent/instrumentation/okhttp/v3_0/OkHttp3DispatcherInstrumentation.java
@@ -51,8 +51,12 @@ public static PropagatedContext onEnter(@Advice.Argument(0) Runnable call) {
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void onExit(
-        @Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
-      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
+        @Advice.Argument(0) Runnable call,
+        @Advice.Enter PropagatedContext propagatedContext,
+        @Advice.Thrown Throwable throwable) {
+      VirtualField<Runnable, PropagatedContext> virtualField =
+          VirtualField.find(Runnable.class, PropagatedContext.class);
+      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, call);
     }
   }
 }
diff --git a/...aagent/instrumentation/pekkoactor/v1_0/PekkoDefaultSystemMessageQueueInstrumentation.java b/...aagent/instrumentation/pekkoactor/v1_0/PekkoDefaultSystemMessageQueueInstrumentation.java
@@ -58,8 +58,13 @@ public static PropagatedContext enter(@Advice.Argument(1) SystemMessage systemMe
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void exit(
-        @Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
-      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
+        @Advice.Argument(1) SystemMessage systemMessage,
+        @Advice.Enter PropagatedContext propagatedContext,
+        @Advice.Thrown Throwable throwable) {
+      VirtualField<SystemMessage, PropagatedContext> virtualField =
+          VirtualField.find(SystemMessage.class, PropagatedContext.class);
+      ExecutorAdviceHelper.cleanUpAfterSubmit(
+          propagatedContext, throwable, virtualField, systemMessage);
     }
   }
 }
diff --git a/...entelemetry/javaagent/instrumentation/pekkoactor/v1_0/PekkoDispatcherInstrumentation.java b/...entelemetry/javaagent/instrumentation/pekkoactor/v1_0/PekkoDispatcherInstrumentation.java
@@ -52,8 +52,12 @@ public static PropagatedContext enterDispatch(@Advice.Argument(1) Envelope envel
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void exitDispatch(
-        @Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
-      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
+        @Advice.Argument(1) Envelope envelope,
+        @Advice.Enter PropagatedContext propagatedContext,
+        @Advice.Thrown Throwable throwable) {
+      VirtualField<Envelope, PropagatedContext> virtualField =
+          VirtualField.find(Envelope.class, PropagatedContext.class);
+      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, envelope);
     }
   }
 }
diff --git a/...ntelemetry/javaagent/instrumentation/scalaexecutors/ScalaForkJoinPoolInstrumentation.java b/...ntelemetry/javaagent/instrumentation/scalaexecutors/ScalaForkJoinPoolInstrumentation.java
@@ -53,8 +53,12 @@ public static PropagatedContext enterJobSubmit(@Advice.Argument(0) ForkJoinTask<
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void exitJobSubmit(
-        @Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
-      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
+        @Advice.Argument(0) ForkJoinTask<?> task,
+        @Advice.Enter PropagatedContext propagatedContext,
+        @Advice.Thrown Throwable throwable) {
+      VirtualField<ForkJoinTask<?>, PropagatedContext> virtualField =
+          VirtualField.find(ForkJoinTask.class, PropagatedContext.class);
+      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, task);
     }
   }
 }
diff --git a/...ry/javaagent/instrumentation/spring/core/v2_0/SimpleAsyncTaskExecutorInstrumentation.java b/...ry/javaagent/instrumentation/spring/core/v2_0/SimpleAsyncTaskExecutorInstrumentation.java
@@ -56,8 +56,12 @@ public static PropagatedContext enterJobSubmit(@Advice.Argument(0) Runnable task
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void exitJobSubmit(
-        @Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
-      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
+        @Advice.Argument(0) Runnable task,
+        @Advice.Enter PropagatedContext propagatedContext,
+        @Advice.Thrown Throwable throwable) {
+      VirtualField<Runnable, PropagatedContext> virtualField =
+          VirtualField.find(Runnable.class, PropagatedContext.class);
+      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, task);
     }
   }
 }
diff --git a/...o/opentelemetry/javaagent/instrumentation/undertow/HttpServerExchangeInstrumentation.java b/...o/opentelemetry/javaagent/instrumentation/undertow/HttpServerExchangeInstrumentation.java
@@ -50,8 +50,12 @@ public static PropagatedContext enterJobSubmit(@Advice.Argument(1) Runnable task
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void exitJobSubmit(
-        @Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
-      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
+        @Advice.Argument(1) Runnable task,
+        @Advice.Enter PropagatedContext propagatedContext,
+        @Advice.Thrown Throwable throwable) {
+      VirtualField<Runnable, PropagatedContext> virtualField =
+          VirtualField.find(Runnable.class, PropagatedContext.class);
+      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable, virtualField, task);
     }
   }
 }