fix: Lettuce redis client, does not apply RedisReduceSpan Sampler

lumigo-io · Dec 19, 2024 · d1a3609 · d1a3609
1 parent dfe476a
commit d1a3609
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 11 deletions.
diff --git a/custom/src/main/java/io/lumigo/javaagent/RedisSamplingConfigurer.java b/custom/src/main/java/io/lumigo/javaagent/RedisSamplingConfigurer.java
@@ -28,6 +28,7 @@
 import io.opentelemetry.sdk.trace.data.LinkData;
 import io.opentelemetry.sdk.trace.samplers.Sampler;
 import io.opentelemetry.sdk.trace.samplers.SamplingResult;
+import io.opentelemetry.semconv.SemanticAttributes;
 import java.util.List;
 import java.util.logging.Logger;
 import java.util.regex.Pattern;
@@ -79,7 +80,7 @@ static class RedisReduceInfoSpanSampler implements Sampler {
     private final Sampler delegateSampler;
 
     // Regex pattern to match span names containing "INFO," (case insensitive)
-    private final Pattern spanNamePattern = Pattern.compile("INFO.*", Pattern.CASE_INSENSITIVE);
+    private final Pattern infoCommandPattern = Pattern.compile("INFO.*", Pattern.CASE_INSENSITIVE);
 
     public RedisReduceInfoSpanSampler(Sampler delegateSampler) {
       this.delegateSampler = delegateSampler;
@@ -97,7 +98,14 @@ public SamplingResult shouldSample(
       String dbSystem = attributes.get(DB_SYSTEM_KEY);
       if ("redis".equalsIgnoreCase(dbSystem)) {
         // Match the span name against the regex
-        if (spanNamePattern.matcher(spanName).matches()) {
+        if (infoCommandPattern.matcher(spanName).matches()) {
+          LOGGER.finest("Dropping Redis INFO span because of span name: " + spanName);
+          return SamplingResult.drop();
+        }
+        // Math the span attribute db.statement against the regex
+        String dbStatement = attributes.get(SemanticAttributes.DB_STATEMENT);
+        if (dbStatement != null && infoCommandPattern.matcher(dbStatement).matches()) {
+          LOGGER.finest("Dropping Redis INFO span because of db.statement: " + dbStatement);
           return SamplingResult.drop();
         }
       }

diff --git a/custom/src/test/java/io/lumigo/javaagent/RedisSamplingConfiguratorTest.java b/custom/src/test/java/io/lumigo/javaagent/RedisSamplingConfiguratorTest.java
@@ -39,7 +39,7 @@ public void tearDown() {
   }
 
   @Test
-  public void shouldDropByDefaultInfoServer() {
+  public void shouldDropByInfoCommand() {
     AutoConfiguredOpenTelemetrySdkBuilder builder = AutoConfiguredOpenTelemetrySdk.builder();
 
     new RedisSamplingConfigurer().customize(builder);
@@ -50,8 +50,9 @@ public void shouldDropByDefaultInfoServer() {
     SdkTracerProvider tracerProvider = sdk.getOpenTelemetrySdk().getSdkTracerProvider();
     Sampler sampler = tracerProvider.getSampler();
 
-    // Test that the default behavior is to drop "INFO server" commands
-    SamplingResult infoServerResult =
+    // Test that the default behavior is to drop "INFO" command in the span name
+    Assertions.assertEquals(
+        SamplingResult.drop(),
         sampler.shouldSample(
             Context.root(),
             IdGenerator.random().generateTraceId(),
@@ -60,11 +61,11 @@ public void shouldDropByDefaultInfoServer() {
             Attributes.of(
                 AttributeKey.stringKey("db.system"), "redis",
                 AttributeKey.stringKey("db.statement"), "server"),
-            Collections.emptyList());
-    Assertions.assertEquals(SamplingResult.drop(), infoServerResult);
+            Collections.emptyList()));
 
     // Test that the default behavior is to drop other INFO commands
-    SamplingResult serverResult =
+    Assertions.assertEquals(
+        SamplingResult.drop(),
         sampler.shouldSample(
             Context.root(),
             IdGenerator.random().generateTraceId(),
@@ -73,8 +74,20 @@ public void shouldDropByDefaultInfoServer() {
             Attributes.of(
                 AttributeKey.stringKey("db.system"), "redis",
                 AttributeKey.stringKey("db.statement"), "other"),
-            Collections.emptyList());
-    Assertions.assertEquals(SamplingResult.drop(), serverResult);
+            Collections.emptyList()));
+
+    // Test that the default behavior is to drop other INFO commands
+    Assertions.assertEquals(
+        SamplingResult.drop(),
+        sampler.shouldSample(
+            Context.root(),
+            IdGenerator.random().generateTraceId(),
+            "redis",
+            SpanKind.CLIENT,
+            Attributes.of(
+                AttributeKey.stringKey("db.system"), "redis",
+                AttributeKey.stringKey("db.statement"), "INFO"),
+            Collections.emptyList()));
   }
 
   @Test

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
@@ -215,13 +215,18 @@ private void fillEndpoint(OpenTelemetryEndpoint endpoint) {
         @CanIgnoreReturnValue
         @SuppressWarnings("UnusedMethod")
         public synchronized Tracer.Span start(RedisCommand<?, ?, ?> command) {
+          // Here we would like to set db.statement before starting the span, because Sampler may use it.
+          // At that point db.statement having only command name is fine, command args will be added later.
+            String commandName = command.getType().name();
+            this.spanBuilder.setAttribute(SemanticAttributes.DB_STATEMENT, commandName);
+
             start();
 
             Span span = this.span;
             if (span == null) {
                 throw new IllegalStateException("Span started but null, this is a programming error.");
             }
-            span.updateName(command.getType().name());
+            span.updateName(commandName);
 
             if (command.getArgs() != null) {
                 argsList = OtelCommandArgsUtil.getCommandArgs(command.getArgs());