From 14d1f4c2e52b72ba5f03ff45e8ce106b831f3ae7 Mon Sep 17 00:00:00 2001
From: Lauri Tulmin <ltulmin@splunk.com>
Date: Tue, 23 Jan 2024 21:47:12 +0200
Subject: [PATCH] Improve rediscala instrumentation (#10301)

---
 .../rediscala-1.8/javaagent/build.gradle.kts  |  9 +++-
 .../rediscala/RequestInstrumentation.java     | 17 +++++++-
 .../test/groovy/RediscalaClientTest.groovy    | 42 ++++++++++++++-----
 3 files changed, 55 insertions(+), 13 deletions(-)

diff --git a/instrumentation/rediscala-1.8/javaagent/build.gradle.kts b/instrumentation/rediscala-1.8/javaagent/build.gradle.kts
index bd73085bfe..8442cb3dbf 100644
--- 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/instrumentation/rediscala-1.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rediscala/RequestInstrumentation.java b/instrumentation/rediscala-1.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rediscala/RequestInstrumentation.java
index 018cd0cf85..0daf60166a 100644
--- a/instrumentation/rediscala-1.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rediscala/RequestInstrumentation.java
+++ b/instrumentation/rediscala-1.8/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/rediscala/RequestInstrumentation.java
@@ -23,7 +23,11 @@ import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 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,11 +78,11 @@ public class RequestInstrumentation implements TypeInstrumentation {
 
     @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) {
@@ -86,6 +90,17 @@ public class RequestInstrumentation implements TypeInstrumentation {
       }
       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
index 7e8048ddd8..6e4e9b98e0 100644
--- 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,7 +25,7 @@ class RediscalaClientTest extends AgentInstrumentationSpecification {
   int port
 
   @Shared
-  ActorSystem system
+  def system
 
   @Shared
   RedisClient redisClient
@@ -34,15 +33,36 @@ class RediscalaClientTest extends AgentInstrumentationSpecification {
   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() {