Improve Akka instrumentation

instrumentation/akka-actor-2.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaActorCellInstrumentation.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.akkaactor;
+
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import akka.dispatch.Envelope;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
+import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
+import io.opentelemetry.javaagent.instrumentation.api.concurrent.AdviceUtils;
+import io.opentelemetry.javaagent.instrumentation.api.concurrent.State;
+import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
+import java.util.Collections;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class AkkaActorCellInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<? super TypeDescription> typeMatcher() {
+    return named("akka.actor.ActorCell");
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return Collections.singletonMap(
+        named("invoke").and(takesArgument(0, named("akka.dispatch.Envelope"))),
+        AkkaActorCellInstrumentation.class.getName() + "$InvokeStateAdvice");
+  }
+
+  public static class InvokeStateAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static Scope enter(@Advice.Argument(0) Envelope envelope) {
+      ContextStore<Envelope, State> contextStore =
+          InstrumentationContext.get(Envelope.class, State.class);
+      return AdviceUtils.startTaskScope(contextStore, envelope);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void exit(@Advice.Enter Scope scope) {
+      if (scope != null) {
+        scope.close();
+      }
+    }
+  }
+}

instrumentation/akka-actor-2.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaActorInstrumentationModule.java

@@ -25,7 +25,11 @@ public AkkaActorInstrumentationModule() {
 
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
-    return asList(new AkkaForkJoinPoolInstrumentation(), new AkkaForkJoinTaskInstrumentation());
+    return asList(
+        new AkkaForkJoinPoolInstrumentation(),
+        new AkkaForkJoinTaskInstrumentation(),
+        new AkkaDispatcherInstrumentation(),
+        new AkkaActorCellInstrumentation());
   }
 
   @Override
@@ -34,6 +38,7 @@ public Map<String, String> contextStore() {
     map.put(Runnable.class.getName(), State.class.getName());
     map.put(Callable.class.getName(), State.class.getName());
     map.put(AkkaForkJoinTaskInstrumentation.TASK_CLASS_NAME, State.class.getName());
+    map.put("akka.dispatch.Envelope", State.class.getName());
     return Collections.unmodifiableMap(map);
   }
 

instrumentation/akka-actor-2.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkaactor/AkkaDispatcherInstrumentation.java

@@ -0,0 +1,54 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.akkaactor;
+
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import akka.dispatch.Envelope;
+import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
+import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
+import io.opentelemetry.javaagent.instrumentation.api.concurrent.ExecutorInstrumentationUtils;
+import io.opentelemetry.javaagent.instrumentation.api.concurrent.State;
+import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class AkkaDispatcherInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<? super TypeDescription> typeMatcher() {
+    return named("akka.dispatch.Dispatcher");
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+        named("dispatch")
+            .and(takesArgument(0, named("akka.actor.ActorCell")))
+            .and(takesArgument(1, named("akka.dispatch.Envelope"))),
+        AkkaDispatcherInstrumentation.class.getName() + "$DispatcherStateAdvice");
+  }
+
+  public static class DispatcherStateAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static State enterDispatch(@Advice.Argument(1) Envelope envelope) {
+      if (ExecutorInstrumentationUtils.shouldAttachStateToTask(envelope)) {
+        ContextStore<Envelope, State> contextStore =
+            InstrumentationContext.get(Envelope.class, State.class);
+        return ExecutorInstrumentationUtils.setupState(
+            contextStore, envelope, Java8BytecodeBridge.currentContext());
+      }
+      return null;
+    }
+  }
+}

instrumentation/akka-actor-2.5/javaagent/src/test/groovy/AkkaActorTest.groovy

@@ -7,44 +7,39 @@ import io.opentelemetry.instrumentation.test.AgentInstrumentationSpecification
 
 class AkkaActorTest extends AgentInstrumentationSpecification {
 
-  // TODO this test doesn't really depend on otel.instrumentation.akka-actor.enabled=true
-  //  but setting this property here is needed when running both this test
-  //  and AkkaExecutorInstrumentationTest in the run, otherwise get
-  //  "class redefinition failed: attempted to change superclass or interfaces"
-  //  on whichever test runs second
-  //  (related question is what's the purpose of this test if it doesn't depend on any of the
-  //  instrumentation in this module, is it just to verify that the instrumentation doesn't break
-  //  this scenario?)
-  static {
-    System.setProperty("otel.instrumentation.akka-actor.enabled", "true")
-  }
-
-  def "akka #testMethod"() {
+  def "akka #testMethod #count"() {
     setup:
     AkkaActors akkaTester = new AkkaActors()
-    akkaTester."$testMethod"()
+    count.times {
+      akkaTester."$testMethod"()
+    }
 
     expect:
-    assertTraces(1) {
-      trace(0, 2) {
-        span(0) {
-          name "parent"
-          attributes {
+    assertTraces(count) {
+      count.times {
+        trace(it, 2) {
+          span(0) {
+            name "parent"
+            attributes {
+            }
           }
-        }
-        span(1) {
-          name "$expectedGreeting, Akka"
-          childOf span(0)
-          attributes {
+          span(1) {
+            name "$expectedGreeting, Akka"
+            childOf span(0)
+            attributes {
+            }
           }
         }
       }
     }
 
     where:
-    testMethod     | expectedGreeting
-    "basicTell"    | "Howdy"
-    "basicAsk"     | "Howdy"
-    "basicForward" | "Hello"
+    testMethod     | expectedGreeting | count
+    "basicTell"    | "Howdy"          | 1
+    "basicAsk"     | "Howdy"          | 1
+    "basicForward" | "Hello"          | 1
+    "basicTell"    | "Howdy"          | 150
+    "basicAsk"     | "Howdy"          | 150
+    "basicForward" | "Hello"          | 150
   }
 }