Apply CR changes

open-telemetry · Apr 29, 2024 · b47cfa4 · b47cfa4
1 parent a1e7307
commit b47cfa4
Showing 1 changed file with 12 additions and 3 deletions.
diff --git a/...opentelemetry/javaagent/instrumentation/executors/StructuredTaskScopeInstrumentation.java b/...opentelemetry/javaagent/instrumentation/executors/StructuredTaskScopeInstrumentation.java
@@ -40,9 +40,18 @@ public static class ForkCallableAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static PropagatedContext enterCallableFork(@Advice.Argument(0) Callable<?> task) {
       Context context = Java8BytecodeBridge.currentContext();
-      VirtualField<Callable<?>, PropagatedContext> virtualField =
-          VirtualField.find(Callable.class, PropagatedContext.class);
-      return ExecutorAdviceHelper.attachContextToTask(context, virtualField, task);
+      if (ExecutorAdviceHelper.shouldPropagateContext(context, task)) {
+        VirtualField<Callable<?>, PropagatedContext> virtualField =
+            VirtualField.find(Callable.class, PropagatedContext.class);
+        return ExecutorAdviceHelper.attachContextToTask(context, virtualField, task);
+      }
+      return null;
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void exitCallableFork(
+        @Advice.Enter PropagatedContext propagatedContext, @Advice.Thrown Throwable throwable) {
+      ExecutorAdviceHelper.cleanUpAfterSubmit(propagatedContext, throwable);
     }
   }
 }