Extract query arguments without regex on lettuce 6 (#8932)

open-telemetry · Jul 12, 2023 · 70247f0 · 70247f0
1 parent 57cfddc
commit 70247f0
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 4 deletions.
diff --git a/...io/opentelemetry/javaagent/instrumentation/lettuce/v5_1/LettuceInstrumentationModule.java b/...io/opentelemetry/javaagent/instrumentation/lettuce/v5_1/LettuceInstrumentationModule.java
@@ -26,6 +26,11 @@ public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
     return hasClassesNamed("io.lettuce.core.tracing.Tracing");
   }
 
+  @Override
+  public boolean isHelperClass(String className) {
+    return className.startsWith("io.lettuce.core.protocol.OtelCommandArgsUtil");
+  }
+
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
     return asList(

diff --git a/...ttuce/lettuce-5.1/library/src/main/java/io/lettuce/core/protocol/OtelCommandArgsUtil.java b/...ttuce/lettuce-5.1/library/src/main/java/io/lettuce/core/protocol/OtelCommandArgsUtil.java
@@ -0,0 +1,40 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.lettuce.core.protocol;
+
+import io.lettuce.core.protocol.CommandArgs.KeyArgument;
+import io.lettuce.core.protocol.CommandArgs.SingularArgument;
+import io.lettuce.core.protocol.CommandArgs.ValueArgument;
+import io.opentelemetry.instrumentation.lettuce.common.LettuceArgSplitter;
+import java.util.ArrayList;
+import java.util.List;
+
+// Helper class for accessing package private fields in CommandArgs and its inner classes.
+// https://github.com/lettuce-io/lettuce-core/blob/main/src/main/java/io/lettuce/core/protocol/CommandArgs.java
+public final class OtelCommandArgsUtil {
+
+  /**
+   * Extract argument {@link List} from {@link CommandArgs} so that we wouldn't need to parse them
+   * from command {@link String} with {@link LettuceArgSplitter#splitArgs}.
+   */
+  public static List<String> getCommandArgs(CommandArgs<?, ?> commandArgs) {
+    List<String> result = new ArrayList<>();
+    for (SingularArgument argument : commandArgs.singularArguments) {
+      String value = argument.toString();
+      if (argument instanceof KeyArgument && value.startsWith("key<") && value.endsWith(">")) {
+        value = value.substring("key<".length(), value.length() - 1);
+      } else if (argument instanceof ValueArgument
+          && value.startsWith("value<")
+          && value.endsWith(">")) {
+        value = value.substring("value<".length(), value.length() - 1);
+      }
+      result.add(value);
+    }
+    return result;
+  }
+
+  private OtelCommandArgsUtil() {}
+}
diff --git a/...ary/src/main/java/io/opentelemetry/instrumentation/lettuce/v5_1/OpenTelemetryTracing.java b/...ary/src/main/java/io/opentelemetry/instrumentation/lettuce/v5_1/OpenTelemetryTracing.java
@@ -10,6 +10,7 @@
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import io.lettuce.core.output.CommandOutput;
 import io.lettuce.core.protocol.CompleteableCommand;
+import io.lettuce.core.protocol.OtelCommandArgsUtil;
 import io.lettuce.core.protocol.RedisCommand;
 import io.lettuce.core.tracing.TraceContext;
 import io.lettuce.core.tracing.TraceContextProvider;
@@ -170,7 +171,8 @@ private static class OpenTelemetrySpan extends Tracer.Span {
     @Nullable private List<Object> events;
     @Nullable private Throwable error;
     @Nullable private Span span;
-    @Nullable private String args;
+    @Nullable private List<String> argsList;
+    @Nullable private String argsString;
 
     OpenTelemetrySpan(Context context, SpanBuilder spanBuilder, RedisCommandSanitizer sanitizer) {
       this.context = context;
@@ -224,7 +226,7 @@ public synchronized Tracer.Span start(RedisCommand<?, ?, ?> command) {
       span.updateName(command.getType().name());
 
       if (command.getArgs() != null) {
-        args = command.getArgs().toCommandString();
+        argsList = OtelCommandArgsUtil.getCommandArgs(command.getArgs());
       }
 
       if (command instanceof CompleteableCommand) {
@@ -294,7 +296,7 @@ public synchronized Tracer.Span annotate(String value) {
     @CanIgnoreReturnValue
     public synchronized Tracer.Span tag(String key, String value) {
       if (key.equals("redis.args")) {
-        args = value;
+        argsString = value;
         return this;
       }
       if (span != null) {
@@ -325,7 +327,8 @@ public synchronized void finish() {
 
     private void finish(Span span) {
       if (name != null) {
-        String statement = sanitizer.sanitize(name, splitArgs(args));
+        String statement =
+            sanitizer.sanitize(name, argsList != null ? argsList : splitArgs(argsString));
         span.setAttribute(SemanticAttributes.DB_STATEMENT, statement);
       }
       span.end();