Improve rediscala instrumentation (#10301)

open-telemetry · Feb 16, 2024 · 4b235b9 · 4b235b9
1 parent e39d735
commit 4b235b9
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 13 deletions.
diff --git a/instrumentation/rediscala-1.8/javaagent/build.gradle.kts b/instrumentation/rediscala-1.8/javaagent/build.gradle.kts
@@ -44,12 +44,19 @@ muzzle {
     versions.set("[1.9.0,)")
     assertInverse.set(true)
   }
+
+  pass {
+    group.set("io.github.rediscala")
+    module.set("rediscala_2.13")
+    versions.set("[1.10.0,)")
+    assertInverse.set(true)
+  }
 }
 
 dependencies {
   library("com.github.etaty:rediscala_2.11:1.8.0")
 
-  latestDepTestLibrary("com.github.etaty:rediscala_2.13:+")
+  latestDepTestLibrary("io.github.rediscala:rediscala_2.13:+")
 }
 
 tasks {

diff --git a/...ain/java/io/opentelemetry/javaagent/instrumentation/rediscala/RequestInstrumentation.java b/...ain/java/io/opentelemetry/javaagent/instrumentation/rediscala/RequestInstrumentation.java
@@ -23,7 +23,11 @@
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
+import redis.ActorRequest;
+import redis.BufferedRequest;
 import redis.RedisCommand;
+import redis.Request;
+import redis.RoundRobinPoolRequest;
 import scala.concurrent.ExecutionContext;
 import scala.concurrent.Future;
 
@@ -74,18 +78,29 @@ public static void onEnter(
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void onExit(
+        @Advice.This Object action,
         @Advice.Argument(0) RedisCommand<?, ?> cmd,
         @Advice.Local("otelContext") Context context,
         @Advice.Local("otelScope") Scope scope,
         @Advice.Thrown Throwable throwable,
-        @Advice.FieldValue("executionContext") ExecutionContext ctx,
         @Advice.Return(readOnly = false) Future<Object> responseFuture) {
 
       if (scope == null) {
         return;
       }
       scope.close();
 
+      ExecutionContext ctx = null;
+      if (action instanceof ActorRequest) {
+        ctx = ((ActorRequest) action).executionContext();
+      } else if (action instanceof Request) {
+        ctx = ((Request) action).executionContext();
+      } else if (action instanceof BufferedRequest) {
+        ctx = ((BufferedRequest) action).executionContext();
+      } else if (action instanceof RoundRobinPoolRequest) {
+        ctx = ((RoundRobinPoolRequest) action).executionContext();
+      }
+
       if (throwable != null) {
         instrumenter().end(context, cmd, null, throwable);
       } else {

diff --git a/instrumentation/rediscala-1.8/javaagent/src/test/groovy/RediscalaClientTest.groovy b/instrumentation/rediscala-1.8/javaagent/src/test/groovy/RediscalaClientTest.groovy
@@ -3,7 +3,6 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import akka.actor.ActorSystem
 import io.opentelemetry.instrumentation.test.AgentInstrumentationSpecification
 import io.opentelemetry.semconv.SemanticAttributes
 import org.testcontainers.containers.GenericContainer
@@ -26,23 +25,44 @@ class RediscalaClientTest extends AgentInstrumentationSpecification {
   int port
 
   @Shared
-  ActorSystem system
+  def system
 
   @Shared
   RedisClient redisClient
 
   def setupSpec() {
     redisServer.start()
     port = redisServer.getMappedPort(6379)
-    system = ActorSystem.create()
-    redisClient = new RedisClient("localhost",
-      port,
-      Option.apply(null),
-      Option.apply(null),
-      "RedisClient",
-      Option.apply(null),
-      system,
-      new RedisDispatcher("rediscala.rediscala-client-worker-dispatcher"))
+    // latest has separate artifacts for akka an pekko, currently latestDepTestLibrary picks the
+    // pekko one
+    try {
+      def clazz = Class.forName("akka.actor.ActorSystem")
+      system = clazz.getMethod("create").invoke(null)
+    } catch (ClassNotFoundException exception) {
+      def clazz = Class.forName("org.apache.pekko.actor.ActorSystem")
+      system = clazz.getMethod("create").invoke(null)
+    }
+    // latest RedisClient constructor takes username as argument
+    if (RedisClient.metaClass.getMetaMethod("username") != null) {
+      redisClient = new RedisClient("localhost",
+        port,
+        Option.apply(null),
+        Option.apply(null),
+        Option.apply(null),
+        "RedisClient",
+        Option.apply(null),
+        system,
+        new RedisDispatcher("rediscala.rediscala-client-worker-dispatcher"))
+    } else {
+      redisClient = new RedisClient("localhost",
+        port,
+        Option.apply(null),
+        Option.apply(null),
+        "RedisClient",
+        Option.apply(null),
+        system,
+        new RedisDispatcher("rediscala.rediscala-client-worker-dispatcher"))
+    }
   }
 
   def cleanupSpec() {