Fix akka shutdown hanging (#13073)

Author: laurit

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);
     }
   }
 }

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() {}
+}

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);
     }
   }
 }

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() {}
+}