From dfe476aec966f2ffd6c7369c5c78a9f14f39fabf Mon Sep 17 00:00:00 2001 From: Eugene Orlovsky Date: Thu, 12 Dec 2024 14:36:08 +0100 Subject: [PATCH] feat: custom redis sampler that uses span name instead of db.statement (#621) --- .../javaagent/RedisSamplingConfigurer.java | 70 +++++++++++++------ .../RedisSamplingConfiguratorTest.java | 16 ++--- 2 files changed, 55 insertions(+), 31 deletions(-) diff --git a/custom/src/main/java/io/lumigo/javaagent/RedisSamplingConfigurer.java b/custom/src/main/java/io/lumigo/javaagent/RedisSamplingConfigurer.java index d5ba78108a..aab5d3b45c 100644 --- a/custom/src/main/java/io/lumigo/javaagent/RedisSamplingConfigurer.java +++ b/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 dbSystemKey = AttributeKey.stringKey("db.system"); - AttributeKey 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 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 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"; + } } } diff --git a/custom/src/test/java/io/lumigo/javaagent/RedisSamplingConfiguratorTest.java b/custom/src/test/java/io/lumigo/javaagent/RedisSamplingConfiguratorTest.java index 39101e7891..b9bb614402 100644 --- a/custom/src/test/java/io/lumigo/javaagent/RedisSamplingConfiguratorTest.java +++ b/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());