End jedis span when operation actually ends (#5256)

* End jedis span when operation actually ends * address review comments
open-telemetry · Feb 1, 2022 · 8240a5f · 8240a5f
1 parent fbf0076
commit 8240a5f
Show file tree

Hide file tree

Showing 16 changed files with 330 additions and 20 deletions.
diff --git a/...tion-api/src/main/java/io/opentelemetry/instrumentation/api/db/RedisCommandSanitizer.java b/...tion-api/src/main/java/io/opentelemetry/instrumentation/api/db/RedisCommandSanitizer.java
@@ -52,21 +52,31 @@ public final class RedisCommandSanitizer {
     CommandSanitizer setMultiField = new MultiKeyValue(0);
 
     // Cluster
-    for (String command : asList("CLUSTER", "READONLY", "READWRITE")) {
+    for (String command : asList("CLUSTER", "FAILOVER", "READONLY", "READWRITE")) {
       sanitizers.put(command, KeepAllArgs.INSTANCE);
     }
 
     // Connection
     sanitizers.put("AUTH", DEFAULT);
     // HELLO can contain AUTH data
     sanitizers.put("HELLO", keepTwoArgs);
-    for (String command : asList("CLIENT", "ECHO", "PING", "QUIT", "SELECT")) {
+    for (String command : asList("CLIENT", "ECHO", "PING", "QUIT", "RESET", "SELECT")) {
       sanitizers.put(command, KeepAllArgs.INSTANCE);
     }
 
     // Geo
     for (String command :
-        asList("GEOADD", "GEODIST", "GEOHASH", "GEOPOS", "GEORADIUS", "GEORADIUSBYMEMBER")) {
+        asList(
+            "GEOADD",
+            "GEODIST",
+            "GEOHASH",
+            "GEOPOS",
+            "GEORADIUS",
+            "GEORADIUS_RO",
+            "GEORADIUSBYMEMBER",
+            "GEORADIUSBYMEMBER_RO",
+            "GEOSEARCH",
+            "GEOSEARCHSTORE")) {
       sanitizers.put(command, KeepAllArgs.INSTANCE);
     }
 
@@ -85,6 +95,7 @@ public final class RedisCommandSanitizer {
             "HKEYS",
             "HLEN",
             "HMGET",
+            "HRANDFIELD",
             "HSCAN",
             "HSTRLEN",
             "HVALS")) {
@@ -103,23 +114,27 @@ public final class RedisCommandSanitizer {
     sanitizers.put("RESTORE", keepTwoArgs);
     for (String command :
         asList(
+            "COPY",
             "DEL",
             "DUMP",
             "EXISTS",
             "EXPIRE",
             "EXPIREAT",
+            "EXPIRETIME",
             "KEYS",
             "MOVE",
             "OBJECT",
             "PERSIST",
             "PEXPIRE",
             "PEXPIREAT",
+            "PEXPIRETIME",
             "PTTL",
             "RANDOMKEY",
             "RENAME",
             "RENAMENX",
             "SCAN",
             "SORT",
+            "SORT_RO",
             "TOUCH",
             "TTL",
             "TYPE",
@@ -140,12 +155,14 @@ public final class RedisCommandSanitizer {
     for (String command :
         asList(
             "BLMOVE",
+            "BLMPOP",
             "BLPOP",
             "BRPOP",
             "BRPOPLPUSH",
             "LINDEX",
             "LLEN",
             "LMOVE",
+            "LMPOP",
             "LPOP",
             "LRANGE",
             "LTRIM",
@@ -157,13 +174,23 @@ public final class RedisCommandSanitizer {
     // Pub/Sub
     sanitizers.put("PUBLISH", keepOneArg);
     for (String command :
-        asList("PSUBSCRIBE", "PUBSUB", "PUNSUBSCRIBE", "SUBSCRIBE", "UNSUBSCRIBE")) {
+        asList(
+            "PSUBSCRIBE",
+            "PUBSUB",
+            "PUNSUBSCRIBE",
+            "SPUBLISH",
+            "SSUBSCRIBE",
+            "SUBSCRIBE",
+            "SUNSUBSCRIBE",
+            "UNSUBSCRIBE")) {
       sanitizers.put(command, KeepAllArgs.INSTANCE);
     }
 
     // Scripting
     sanitizers.put("EVAL", Eval.INSTANCE);
+    sanitizers.put("EVAL_RO", Eval.INSTANCE);
     sanitizers.put("EVALSHA", Eval.INSTANCE);
+    sanitizers.put("EVALSHA_RO", Eval.INSTANCE);
     sanitizers.put("SCRIPT", KeepAllArgs.INSTANCE);
 
     // Server
@@ -211,6 +238,7 @@ public final class RedisCommandSanitizer {
             "SDIFF",
             "SDIFFSTORE",
             "SINTER",
+            "SINTERCARD",
             "SINTERSTORE",
             "SMEMBERS",
             "SPOP",
@@ -239,14 +267,21 @@ public final class RedisCommandSanitizer {
     sanitizers.put("ZSCORE", keepOneArg);
     for (String command :
         asList(
+            "BZMPOP",
             "BZPOPMAX",
             "BZPOPMIN",
             "ZCARD",
+            "ZDIFF",
+            "ZDIFFSTORE",
             "ZINTER",
+            "ZINTERCARD",
             "ZINTERSTORE",
+            "ZMPOP",
             "ZPOPMAX",
             "ZPOPMIN",
+            "ZRANDMEMBER",
             "ZRANGE",
+            "ZRANGESTORE",
             "ZREMRANGEBYRANK",
             "ZREVRANGE",
             "ZSCAN",
@@ -260,6 +295,7 @@ public final class RedisCommandSanitizer {
     for (String command :
         asList(
             "XACK",
+            "XAUTOCLAIM",
             "XCLAIM",
             "XDEL",
             "XGROUP",
@@ -288,16 +324,20 @@ public final class RedisCommandSanitizer {
         asList(
             "BITCOUNT",
             "BITFIELD",
+            "BITFIELD_RO",
             "BITOP",
             "BITPOS",
             "DECR",
             "DECRBY",
             "GET",
             "GETBIT",
+            "GETDEL",
+            "GETEX",
             "GETRANGE",
             "INCR",
             "INCRBY",
             "INCRBYFLOAT",
+            "LCS",
             "MGET",
             "SETBIT",
             "STRALGO",
@@ -329,7 +369,7 @@ static String argToString(Object arg) {
       if (arg instanceof byte[]) {
         return new String((byte[]) arg, StandardCharsets.UTF_8);
       } else {
-        return arg.toString();
+        return String.valueOf(arg);
       }
     }
 

diff --git a/instrumentation/jedis/jedis-1.4/javaagent/build.gradle.kts b/instrumentation/jedis/jedis-1.4/javaagent/build.gradle.kts
@@ -17,6 +17,8 @@ dependencies {
   compileOnly("com.google.auto.value:auto-value-annotations")
   annotationProcessor("com.google.auto.value:auto-value")
 
+  implementation(project(":instrumentation:jedis:jedis-common:javaagent"))
+
   testInstrumentation(project(":instrumentation:jedis:jedis-3.0:javaagent"))
   testInstrumentation(project(":instrumentation:jedis:jedis-4.0:javaagent"))
 

diff --git a/...io/opentelemetry/javaagent/instrumentation/jedis/v1_4/JedisConnectionInstrumentation.java b/...io/opentelemetry/javaagent/instrumentation/jedis/v1_4/JedisConnectionInstrumentation.java
@@ -18,6 +18,7 @@
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.opentelemetry.javaagent.instrumentation.jedis.JedisRequestContext;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -32,8 +33,6 @@ public ElementMatcher<TypeDescription> typeMatcher() {
 
   @Override
   public void transform(TypeTransformer transformer) {
-    // FIXME: This instrumentation only incorporates sending the command, not processing the
-    // result.
     transformer.applyAdviceToMethod(
         isMethod()
             .and(named("sendCommand"))
@@ -80,7 +79,7 @@ public static void stopSpan(
       }
 
       scope.close();
-      instrumenter().end(context, request, null, throwable);
+      JedisRequestContext.endIfNotAttached(instrumenter(), context, request, throwable);
     }
   }
 
@@ -116,7 +115,7 @@ public static void stopSpan(
       }
 
       scope.close();
-      instrumenter().end(context, request, null, throwable);
+      JedisRequestContext.endIfNotAttached(instrumenter(), context, request, throwable);
     }
   }
 }
diff --git a/...main/java/io/opentelemetry/javaagent/instrumentation/jedis/v1_4/JedisInstrumentation.java b/...main/java/io/opentelemetry/javaagent/instrumentation/jedis/v1_4/JedisInstrumentation.java
@@ -0,0 +1,62 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.jedis.v1_4;
+
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.isStatic;
+import static net.bytebuddy.matcher.ElementMatchers.namedOneOf;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.opentelemetry.javaagent.instrumentation.jedis.JedisRequestContext;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class JedisInstrumentation implements TypeInstrumentation {
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return namedOneOf("redis.clients.jedis.Jedis", "redis.clients.jedis.BinaryJedis");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        isMethod()
+            .and(isPublic())
+            .and(not(isStatic()))
+            .and(
+                not(
+                    namedOneOf(
+                        "close",
+                        "setDataSource",
+                        "getDB",
+                        "isConnected",
+                        "connect",
+                        "resetState",
+                        "getClient",
+                        "disconnect"))),
+        this.getClass().getName() + "$JedisMethodAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class JedisMethodAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static JedisRequestContext<JedisRequest> onEnter() {
+      return JedisRequestContext.attach();
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(@Advice.Enter JedisRequestContext<JedisRequest> requestContext) {
+      if (requestContext != null) {
+        requestContext.detachAndEnd();
+      }
+    }
+  }
+}
diff --git a/...ava/io/opentelemetry/javaagent/instrumentation/jedis/v1_4/JedisInstrumentationModule.java b/...ava/io/opentelemetry/javaagent/instrumentation/jedis/v1_4/JedisInstrumentationModule.java
@@ -6,7 +6,7 @@
 package io.opentelemetry.javaagent.instrumentation.jedis.v1_4;
 
 import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
-import static java.util.Collections.singletonList;
+import static java.util.Arrays.asList;
 import static net.bytebuddy.matcher.ElementMatchers.not;
 
 import com.google.auto.service.AutoService;
@@ -30,6 +30,6 @@ public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
 
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
-    return singletonList(new JedisConnectionInstrumentation());
+    return asList(new JedisConnectionInstrumentation(), new JedisInstrumentation());
   }
 }
diff --git a/instrumentation/jedis/jedis-3.0/javaagent/build.gradle.kts b/instrumentation/jedis/jedis-3.0/javaagent/build.gradle.kts
@@ -18,6 +18,8 @@ dependencies {
   compileOnly("com.google.auto.value:auto-value-annotations")
   annotationProcessor("com.google.auto.value:auto-value")
 
+  implementation(project(":instrumentation:jedis:jedis-common:javaagent"))
+
   // ensures jedis-1.4 instrumentation does not load with jedis 3.0+ by failing
   // the tests in the event it does. The tests will end up with double spans
   testInstrumentation(project(":instrumentation:jedis:jedis-1.4:javaagent"))

diff --git a/...io/opentelemetry/javaagent/instrumentation/jedis/v3_0/JedisConnectionInstrumentation.java b/...io/opentelemetry/javaagent/instrumentation/jedis/v3_0/JedisConnectionInstrumentation.java
@@ -18,6 +18,7 @@
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.opentelemetry.javaagent.instrumentation.jedis.JedisRequestContext;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -39,7 +40,6 @@ public void transform(TypeTransformer transformer) {
             .and(takesArgument(0, named("redis.clients.jedis.commands.ProtocolCommand")))
             .and(takesArgument(1, is(byte[][].class))),
         this.getClass().getName() + "$SendCommandAdvice");
-    // FIXME: This instrumentation only incorporates sending the command, not processing the result.
   }
 
   @SuppressWarnings("unused")
@@ -74,7 +74,7 @@ public static void stopSpan(
       }
 
       scope.close();
-      instrumenter().end(context, request, null, throwable);
+      JedisRequestContext.endIfNotAttached(instrumenter(), context, request, throwable);
     }
   }
 }