diff --git a/instrumentation/akka/akka-actor-2.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaScheduleInstrumentation.java b/instrumentation/akka/akka-actor-2.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaScheduleInstrumentation.java index 7a2c6230dcd7..12378d9f5493 100644 --- a/instrumentation/akka/akka-actor-2.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaScheduleInstrumentation.java +++ b/instrumentation/akka/akka-actor-2.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaScheduleInstrumentation.java @@ -8,8 +8,6 @@ import static net.bytebuddy.matcher.ElementMatchers.named; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; -import io.opentelemetry.context.Context; -import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge; import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation; import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer; import net.bytebuddy.asm.Advice; @@ -46,8 +44,7 @@ public static class ScheduleAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static void enterSchedule( @Advice.Argument(value = 2, readOnly = false) Runnable runnable) { - Context context = Java8BytecodeBridge.currentContext(); - runnable = context.wrap(runnable); + runnable = AkkaSchedulerTaskWrapper.wrap(runnable); } } @@ -57,8 +54,7 @@ public static class ScheduleOnceAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static void enterScheduleOnce( @Advice.Argument(value = 1, readOnly = false) Runnable runnable) { - Context context = Java8BytecodeBridge.currentContext(); - runnable = context.wrap(runnable); + runnable = AkkaSchedulerTaskWrapper.wrap(runnable); } } } diff --git a/instrumentation/akka/akka-actor-2.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaSchedulerTaskWrapper.java b/instrumentation/akka/akka-actor-2.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaSchedulerTaskWrapper.java new file mode 100644 index 000000000000..e0a6c2906fd7 --- /dev/null +++ b/instrumentation/akka/akka-actor-2.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaSchedulerTaskWrapper.java @@ -0,0 +1,37 @@ +/* + * Copyright The OpenTelemetry Authors + * SPDX-License-Identifier: Apache-2.0 + */ + +package io.opentelemetry.javaagent.instrumentation.akkaactor; + +import io.opentelemetry.context.Context; + +public final class AkkaSchedulerTaskWrapper { + private static final Class RUN_ON_CLOSE_TASK_CLASS = getRunOnCloseTaskClass(); + + private static Class getRunOnCloseTaskClass() { + try { + return Class.forName("akka.actor.Scheduler$TaskRunOnClose"); + } catch (ClassNotFoundException e) { + return null; + } + } + + private static boolean isRunOnCloseTask(Runnable runnable) { + return RUN_ON_CLOSE_TASK_CLASS != null && RUN_ON_CLOSE_TASK_CLASS.isInstance(runnable); + } + + public static Runnable wrap(Runnable runnable) { + // https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/13066 + // Skip wrapping shutdown tasks. Shutdown process hangs when shutdown tasks are wrapped here. + if (isRunOnCloseTask(runnable)) { + return runnable; + } + + Context context = Context.current(); + return context.wrap(runnable); + } + + private AkkaSchedulerTaskWrapper() {} +} diff --git a/instrumentation/pekko/pekko-actor-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pekkoactor/v1_0/PekkoScheduleInstrumentation.java b/instrumentation/pekko/pekko-actor-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pekkoactor/v1_0/PekkoScheduleInstrumentation.java index 8da23bc58028..0f8dcff947f6 100644 --- a/instrumentation/pekko/pekko-actor-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pekkoactor/v1_0/PekkoScheduleInstrumentation.java +++ b/instrumentation/pekko/pekko-actor-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pekkoactor/v1_0/PekkoScheduleInstrumentation.java @@ -8,7 +8,6 @@ import static net.bytebuddy.matcher.ElementMatchers.named; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; -import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge; import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation; import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer; import net.bytebuddy.asm.Advice; @@ -46,7 +45,7 @@ public static class ScheduleAdvice { @Advice.AssignReturned.ToArguments(@ToArgument(2)) @Advice.OnMethodEnter(suppress = Throwable.class) public static Runnable enterSchedule(@Advice.Argument(value = 2) Runnable runnable) { - return Java8BytecodeBridge.currentContext().wrap(runnable); + return PekkoSchedulerTaskWrapper.wrap(runnable); } } @@ -56,7 +55,7 @@ public static class ScheduleOnceAdvice { @Advice.AssignReturned.ToArguments(@ToArgument(1)) @Advice.OnMethodEnter(suppress = Throwable.class) public static Runnable enterScheduleOnce(@Advice.Argument(value = 1) Runnable runnable) { - return Java8BytecodeBridge.currentContext().wrap(runnable); + return PekkoSchedulerTaskWrapper.wrap(runnable); } } } diff --git a/instrumentation/pekko/pekko-actor-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pekkoactor/v1_0/PekkoSchedulerTaskWrapper.java b/instrumentation/pekko/pekko-actor-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pekkoactor/v1_0/PekkoSchedulerTaskWrapper.java new file mode 100644 index 000000000000..25e219e4f985 --- /dev/null +++ b/instrumentation/pekko/pekko-actor-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pekkoactor/v1_0/PekkoSchedulerTaskWrapper.java @@ -0,0 +1,25 @@ +/* + * Copyright The OpenTelemetry Authors + * SPDX-License-Identifier: Apache-2.0 + */ + +package io.opentelemetry.javaagent.instrumentation.pekkoactor.v1_0; + +import io.opentelemetry.context.Context; +import org.apache.pekko.actor.Scheduler; + +public final class PekkoSchedulerTaskWrapper { + + public static Runnable wrap(Runnable runnable) { + // https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/13066 + // Skip wrapping shutdown tasks. Shutdown process hangs when shutdown tasks are wrapped here. + if (runnable instanceof Scheduler.TaskRunOnClose) { + return runnable; + } + + Context context = Context.current(); + return context.wrap(runnable); + } + + private PekkoSchedulerTaskWrapper() {} +}