pekko: propagate context for schedule/scheduleOnce

class names

Author: pjfanning

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

@@ -23,6 +23,8 @@ public List<TypeInstrumentation> typeInstrumentations() {
     return asList(
         new PekkoDispatcherInstrumentation(),
         new PekkoActorCellInstrumentation(),
-        new PekkoDefaultSystemMessageQueueInstrumentation());
+        new PekkoDefaultSystemMessageQueueInstrumentation(),
+        new PekkoScheduleInstrumentation(),
+        new PekkoScheduleOnceInstrumentation());
   }
 }

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

@@ -0,0 +1,47 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.pekkoactor.v1_0;
+
+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;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class PekkoScheduleInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("org.apache.pekko.actor.Scheduler");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        named("schedule")
+            .and(takesArgument(0, named("scala.concurrent.duration.FiniteDuration")))
+            .and(takesArgument(1, named("scala.concurrent.duration.FiniteDuration")))
+            .and(takesArgument(2, named("java.lang.Runnable")))
+            .and(takesArgument(3, named("scala.concurrent.ExecutionContext"))),
+        PekkoScheduleInstrumentation.class.getName() + "$ScheduleAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class ScheduleAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void enterScheduleOnce(
+        @Advice.Argument(value = 2, readOnly = false) Runnable runnable) {
+      Context context = Java8BytecodeBridge.currentContext();
+      runnable = context.wrap(runnable);
+    }
+  }
+}

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

@@ -0,0 +1,46 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.pekkoactor.v1_0;
+
+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;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class PekkoScheduleOnceInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("org.apache.pekko.actor.Scheduler");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        named("scheduleOnce")
+            .and(takesArgument(0, named("scala.concurrent.duration.FiniteDuration")))
+            .and(takesArgument(1, named("java.lang.Runnable")))
+            .and(takesArgument(2, named("scala.concurrent.ExecutionContext"))),
+        PekkoScheduleOnceInstrumentation.class.getName() + "$ScheduleOnceAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  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);
+    }
+  }
+}