ExecutionContextWrapper

Author: pjfanning

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

@@ -0,0 +1,30 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.pekkoactor.v1_0;
+
+import scala.concurrent.ExecutionContext;
+
+import io.opentelemetry.context.Context;
+
+public class ExecutionContextWrapper implements ExecutionContext {
+  private final ExecutionContext executionContext;
+  private final Context otelContext;
+
+  public ExecutionContextWrapper(ExecutionContext executionContext, Context otelContext) {
+    this.executionContext = executionContext;
+    this.otelContext = otelContext;
+  }
+
+  @Override
+  public void execute(Runnable runnable) {
+    executionContext.execute(otelContext.wrap(runnable));
+  }
+
+  @Override
+  public void reportFailure(Throwable cause) {
+    executionContext.reportFailure(cause);
+  }
+}

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

@@ -16,6 +16,8 @@
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
+import scala.concurrent.ExecutionContext;
+
 public class PekkoScheduleInstrumentation implements TypeInstrumentation {
 
   @Override
@@ -39,9 +41,9 @@ public static class ScheduleAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void enterSchedule(
-        @Advice.Argument(value = 2, readOnly = false) Runnable runnable) {
+        @Advice.Argument(value = 3, readOnly = false) ExecutionContext executionContext) {
       Context context = Java8BytecodeBridge.currentContext();
-      runnable = context.wrap(runnable);
+      executionContext = new ExecutionContextWrapper(executionContext, context);
     }
   }
 }

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

@@ -16,6 +16,8 @@
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
+import scala.concurrent.ExecutionContext;
+
 public class PekkoScheduleOnceInstrumentation implements TypeInstrumentation {
 
   @Override
@@ -38,9 +40,9 @@ public static class ScheduleOnceAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void enterScheduleOnce(
-        @Advice.Argument(value = 1, readOnly = false) Runnable runnable) {
+        @Advice.Argument(value = 2, readOnly = false) ExecutionContext executionContext) {
       Context context = Java8BytecodeBridge.currentContext();
-      runnable = context.wrap(runnable);
+      executionContext = new ExecutionContextWrapper(executionContext, context);
     }
   }
 }