feat: custom redis sampler that uses span name instead of db.statement

custom/src/main/java/io/lumigo/javaagent/RedisSamplingConfigurer.java

@@ -19,19 +19,22 @@
 
 import com.google.auto.service.AutoService;
 import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.trace.SpanKind;
-import io.opentelemetry.contrib.sampler.RuleBasedRoutingSampler;
-import io.opentelemetry.contrib.sampler.RuleBasedRoutingSamplerBuilder;
+import io.opentelemetry.context.Context;
 import io.opentelemetry.sdk.autoconfigure.spi.AutoConfigurationCustomizer;
 import io.opentelemetry.sdk.autoconfigure.spi.AutoConfigurationCustomizerProvider;
 import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
+import io.opentelemetry.sdk.trace.data.LinkData;
 import io.opentelemetry.sdk.trace.samplers.Sampler;
+import io.opentelemetry.sdk.trace.samplers.SamplingResult;
+import java.util.List;
 import java.util.logging.Logger;
+import java.util.regex.Pattern;
 
 @AutoService(AutoConfigurationCustomizerProvider.class)
 public class RedisSamplingConfigurer implements AutoConfigurationCustomizerProvider {
   private static final Logger LOGGER = Logger.getLogger(RedisSamplingConfigurer.class.getName());
-
   public static final String LUMIGO_REDUCED_REDIS_INSTRUMENTATION =
       "lumigo.reduced.redis.instrumentation";
 
@@ -42,10 +45,6 @@ public void customize(AutoConfigurationCustomizer autoConfiguration) {
 
   private static Sampler customizeRedisSpans(
       Sampler defaultSampler, ConfigProperties configProperties) {
-
-    RuleBasedRoutingSamplerBuilder samplerBuilder =
-        RuleBasedRoutingSampler.builder(SpanKind.CLIENT, defaultSampler);
-
     String reduceRedisInstrumentation =
         configProperties.getString(LUMIGO_REDUCED_REDIS_INSTRUMENTATION);
     boolean isReducedRedisInstrumentationEnabled;
@@ -65,26 +64,51 @@ private static Sampler customizeRedisSpans(
     }
 
     if (isReducedRedisInstrumentationEnabled) {
-
-      // Setting the environment variable `LUMIGO_REDUCED_REDIS_INSTRUMENTATION=false` will disable
-      // this optimization.
       LOGGER.finest(
-          "Lumigo reduces Redis instrumentation. The `db.statement` attribute (e.g., `INFO server`) is excluded by default. Set `LUMIGO_REDUCED_REDIS_INSTRUMENTATION=false` to disable this behavior.");
+          "Lumigo reduces Redis instrumentation. Redis spans are sampled based on span name using regex. Set `LUMIGO_REDUCED_REDIS_INSTRUMENTATION=false` to disable this behavior.");
+      return new RedisReduceInfoSpanSampler(defaultSampler);
+    }
 
-      // Define attribute keys
-      AttributeKey<String> dbSystemKey = AttributeKey.stringKey("db.system");
-      AttributeKey<String> dbStatementKey = AttributeKey.stringKey("db.statement");
+    // Use the default sampler if reduced instrumentation is disabled
+    return defaultSampler;
+  }
 
-      samplerBuilder.customize(
-          dbSystemKey,
-          "redis",
-          RuleBasedRoutingSampler.builder(SpanKind.CLIENT, defaultSampler)
-              // have to use regex to match the db.statement attribute, that can be "INFO server" or
-              // "server", depending on the Redis configuration
-              .drop(dbStatementKey, "(INFO\\s+)?server")
-              .build());
+  // Custom Sampler Implementation with Regex
+  static class RedisReduceInfoSpanSampler implements Sampler {
+    private static final AttributeKey<String> DB_SYSTEM_KEY = AttributeKey.stringKey("db.system");
+    private final Sampler delegateSampler;
+
+    // Regex pattern to match span names containing "INFO," (case insensitive)
+    private final Pattern spanNamePattern = Pattern.compile("INFO.*", Pattern.CASE_INSENSITIVE);
+
+    public RedisReduceInfoSpanSampler(Sampler delegateSampler) {
+      this.delegateSampler = delegateSampler;
     }
 
-    return samplerBuilder.build();
+    @Override
+    public SamplingResult shouldSample(
+        Context parentContext,
+        String traceId,
+        String spanName,
+        SpanKind spanKind,
+        Attributes attributes,
+        List<LinkData> parentLinks) {
+      // Check if the db.system attribute is "redis"
+      String dbSystem = attributes.get(DB_SYSTEM_KEY);
+      if ("redis".equalsIgnoreCase(dbSystem)) {
+        // Match the span name against the regex
+        if (spanNamePattern.matcher(spanName).matches()) {
+          return SamplingResult.drop();
+        }
+      }
+      // Fallback to the delegate sampler
+      return delegateSampler.shouldSample(
+          parentContext, traceId, spanName, spanKind, attributes, parentLinks);
+    }
+
+    @Override
+    public String getDescription() {
+      return "RedisReduceInfoSpanSampler";
+    }
   }
 }

custom/src/test/java/io/lumigo/javaagent/RedisSamplingConfiguratorTest.java

@@ -55,24 +55,24 @@ public void shouldDropByDefaultInfoServer() {
         sampler.shouldSample(
             Context.root(),
             IdGenerator.random().generateTraceId(),
-            "name",
+            "INFO",
             SpanKind.CLIENT,
             Attributes.of(
                 AttributeKey.stringKey("db.system"), "redis",
-                AttributeKey.stringKey("db.statement"), "INFO server"),
+                AttributeKey.stringKey("db.statement"), "server"),
             Collections.emptyList());
     Assertions.assertEquals(SamplingResult.drop(), infoServerResult);
 
-    // Test that the default behavior is to drop "server" commands
+    // Test that the default behavior is to drop other INFO commands
     SamplingResult serverResult =
         sampler.shouldSample(
             Context.root(),
             IdGenerator.random().generateTraceId(),
-            "name",
+            "INFO",
             SpanKind.CLIENT,
             Attributes.of(
                 AttributeKey.stringKey("db.system"), "redis",
-                AttributeKey.stringKey("db.statement"), "server"),
+                AttributeKey.stringKey("db.statement"), "other"),
             Collections.emptyList());
     Assertions.assertEquals(SamplingResult.drop(), serverResult);
   }
@@ -93,12 +93,12 @@ public void shouldNotDropIfFalse() {
     Attributes attributes =
         Attributes.of(
             AttributeKey.stringKey("db.system"), "redis",
-            AttributeKey.stringKey("db.statement"), "INFO server");
+            AttributeKey.stringKey("db.statement"), "server");
     SamplingResult result =
         sampler.shouldSample(
             Context.root(),
             IdGenerator.random().generateTraceId(),
-            "name",
+            "INFO",
             SpanKind.CLIENT,
             attributes,
             Collections.emptyList());
@@ -125,7 +125,7 @@ public void shouldNotDropIfNotInfo() {
         sampler.shouldSample(
             Context.root(),
             IdGenerator.random().generateTraceId(),
-            "name",
+            "Other",
             SpanKind.CLIENT,
             attributes,
             Collections.emptyList());