Clear context propagation virtual field (#12397)

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

@@ -58,8 +58,13 @@ public class AkkaDefaultSystemMessageQueueInstrumentation implements TypeInstrum
 
     @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);
     }
   }
 }

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

@@ -52,8 +52,12 @@ public class AkkaDispatcherInstrumentation implements TypeInstrumentation {
 
     @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);
     }
   }
 }

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

@@ -60,8 +60,12 @@ public class AkkaForkJoinPoolInstrumentation implements TypeInstrumentation {
 
     @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);
     }
   }
 }

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

@@ -94,8 +94,11 @@ public final class ExecutorAdviceHelper {
    * 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 final class ExecutorAdviceHelper {
       exceptions.
        */
       propagatedContext.clear();
+      // setting the field to null removes it from the fallback map
+      virtualField.set(task, null);
     }
   }
 
@@ -119,6 +124,8 @@ public final class ExecutorAdviceHelper {
     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);
     }
   }
 

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

@@ -22,6 +22,8 @@ public final class TaskAdviceHelper {
       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();

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

@@ -105,8 +105,12 @@ public class JavaExecutorInstrumentation implements TypeInstrumentation {
 
     @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 class JavaExecutorInstrumentation implements TypeInstrumentation {
 
     @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 class JavaExecutorInstrumentation implements TypeInstrumentation {
 
     @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 class JavaExecutorInstrumentation implements TypeInstrumentation {
             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 class JavaExecutorInstrumentation implements TypeInstrumentation {
 
     @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 class JavaExecutorInstrumentation implements TypeInstrumentation {
             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 class JavaExecutorInstrumentation implements TypeInstrumentation {
             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);
           }
         }
       }

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

@@ -117,8 +117,12 @@ public class JavaForkJoinTaskInstrumentation implements TypeInstrumentation {
 
     @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);
     }
   }
 }

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

@@ -50,8 +50,12 @@ public class StructuredTaskScopeInstrumentation implements TypeInstrumentation {
 
     @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);
     }
   }
 }

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

@@ -61,8 +61,12 @@ public class GuavaListenableFutureInstrumentation implements TypeInstrumentation
 
     @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);
     }
   }
 }

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

@@ -50,8 +50,12 @@ public class JettyQueuedThreadPoolInstrumentation implements TypeInstrumentation
 
     @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);
     }
   }
 }

instrumentation/okhttp/okhttp-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/okhttp/v2_2/DispatcherInstrumentation.java

@@ -48,8 +48,12 @@ public class DispatcherInstrumentation implements TypeInstrumentation {
 
     @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);
     }
   }
 }

instrumentation/okhttp/okhttp-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/okhttp/v3_0/OkHttp3DispatcherInstrumentation.java

@@ -51,8 +51,12 @@ public class OkHttp3DispatcherInstrumentation implements TypeInstrumentation {
 
     @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);
     }
   }
 }

instrumentation/pekko/pekko-actor-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pekkoactor/v1_0/PekkoDefaultSystemMessageQueueInstrumentation.java

@@ -58,8 +58,13 @@ public class PekkoDefaultSystemMessageQueueInstrumentation implements TypeInstru
 
     @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);
     }
   }
 }

instrumentation/pekko/pekko-actor-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pekkoactor/v1_0/PekkoDispatcherInstrumentation.java

@@ -52,8 +52,12 @@ public class PekkoDispatcherInstrumentation implements TypeInstrumentation {
 
     @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);
     }
   }
 }

instrumentation/scala-fork-join-2.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/scalaexecutors/ScalaForkJoinPoolInstrumentation.java

@@ -53,8 +53,12 @@ public class ScalaForkJoinPoolInstrumentation implements TypeInstrumentation {
 
     @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);
     }
   }
 }

instrumentation/spring/spring-core-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/core/v2_0/SimpleAsyncTaskExecutorInstrumentation.java

@@ -56,8 +56,12 @@ public class SimpleAsyncTaskExecutorInstrumentation implements TypeInstrumentati
 
     @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);
     }
   }
 }

instrumentation/undertow-1.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/undertow/HttpServerExchangeInstrumentation.java

@@ -50,8 +50,12 @@ public class HttpServerExchangeInstrumentation implements TypeInstrumentation {
 
     @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);
     }
   }
 }