lettuce redis client sync and async api instrumentation with sync unit tests complete
ADD
dd-java-agent/instrumentation/lettuce-5/lettuce-5.gradle
+apply from: "${rootDir}/gradle/java.gradle"
+
+sourceCompatibility = 1.8
+targetCompatibility = 1.8
+
+dependencies {
+ compileOnly group: 'io.lettuce', name: 'lettuce-core', version: '5.0.4.RELEASE'
+
+ compile project(':dd-java-agent:agent-tooling')
+
+ compile deps.bytebuddy
+ compile deps.opentracing
+ annotationProcessor deps.autoservice
+ implementation deps.autoservice
+
+ testCompile project(':dd-java-agent:testing')
+
+ testCompile group: 'com.github.kstyrc', name: 'embedded-redis', version: '0.6'
+ testCompile group: 'io.lettuce', name: 'lettuce-core', version: '5.0.4.RELEASE'
+}

ADD
dd-java-agent/instrumentation/lettuce-5/src/main/java/datadog/trace/instrumentation/lettuce/RedisAsyncCommandsInstrumentation.java
+package datadog.trace.instrumentation.lettuce;
+
+import static net.bytebuddy.matcher.ElementMatchers.*;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.DDAdvice;
+import datadog.trace.agent.tooling.DDTransformers;
+import datadog.trace.agent.tooling.HelperInjector;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.api.DDTags;
+import io.lettuce.core.protocol.AsyncCommand;
+import io.lettuce.core.protocol.RedisCommand;
+import io.opentracing.Scope;
+import io.opentracing.Span;
+import io.opentracing.tag.Tags;
+import io.opentracing.util.GlobalTracer;
+import java.util.Collections;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.asm.Advice;
+
+@AutoService(Instrumenter.class)
+public class RedisAsyncCommandsInstrumentation extends Instrumenter.Configurable {
+
+ private static final HelperInjector REDIS_ASYNC_COMMANDS_HELPERS =
+ new HelperInjector(
+ RedisAsyncCommandsInstrumentation.class.getName() + "$RedisAsyncConsumer",
+ RedisAsyncCommandsInstrumentation.class.getName() + "$RedisAsyncExceptionFunction");
+ private static final String SERVICE_NAME = "redis";
+ private static final String COMPONENT_NAME = SERVICE_NAME + "-client";
+
+ public RedisAsyncCommandsInstrumentation() {
+ super(SERVICE_NAME);
+ }
+
+ @Override
+ protected AgentBuilder apply(AgentBuilder agentBuilder) {
+ return agentBuilder
+ .type(named("io.lettuce.core.AbstractRedisAsyncCommands"))
+ .transform(REDIS_ASYNC_COMMANDS_HELPERS)
+ .transform(DDTransformers.defaultTransformers())
+ .transform(
+ DDAdvice.create()
+ .advice(
+ isMethod()
+ .and(named("dispatch"))
+ .and(takesArgument(0, named("io.lettuce.core.protocol.RedisCommand"))),
+ RedisAsyncCommandsAdvice.class.getName()))
+ .asDecorator();
+ }
+
+ public static class RedisAsyncCommandsAdvice {
+
+ @Advice.OnMethodEnter(suppress = Throwable.class)
+ public static Scope startSpan(@Advice.Argument(0) final RedisCommand command) {
+ final Scope scope = GlobalTracer.get().buildSpan(SERVICE_NAME + ".query").startActive(false);
+
+ final Span span = scope.span();
+ Tags.DB_TYPE.set(span, SERVICE_NAME);
+ Tags.SPAN_KIND.set(span, Tags.SPAN_KIND_CLIENT);
+ Tags.COMPONENT.set(span, COMPONENT_NAME);
+
+ String commandName = "Unknown Redis Command";
+ String commandArgs = null;
+ if (command != null) {
+ // get the redis command name (i.e. GET, SET, HMSET, etc)
+ if (command.getType() != null) {
+ commandName = command.getType().name();
+ }
+ // get the arguments passed into the redis command
+ if (command.getArgs() != null) {
+ commandArgs = command.getArgs().toCommandString();
+ }
+ }
+
+ span.setTag(DDTags.RESOURCE_NAME, commandName);
+ span.setTag("db.redis.command.args", commandArgs);
+ span.setTag(DDTags.SERVICE_NAME, SERVICE_NAME);
+ span.setTag(DDTags.SPAN_TYPE, SERVICE_NAME);
+
+ return scope;
+ }
+
+ @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+ public static void stopSpan(
+ @Advice.Enter final Scope scope,
+ @Advice.Thrown final Throwable throwable,
+ @Advice.Return(readOnly = false) final AsyncCommand<?, ?, ?> asyncCommand) {
+
+ if (throwable != null) {
+ final Span span = scope.span();
+ Tags.ERROR.set(span, true);
+ span.log(Collections.singletonMap("error.object", throwable));
+ scope.close();
+ return;
+ }
+
+ // use lettuce's callback interface to ensure no logic is lost
+ asyncCommand.onComplete(new RedisAsyncConsumer<>(scope.span()));
+ // exception callbacks are not part of lettuce's api, so use CompletableFuture's, in the future if
+ // lettuce does add some sort of exception callback handler, we should use that instead.
+ asyncCommand.exceptionally(new RedisAsyncExceptionFunction<>(scope.span()));
+ scope.close();
+ }
+ }
+
+ /**
+ * Callback class to close the span on an error in the RedisFuture returned by the lettuce async
+ * API
+ *
+ * @param <Throwable> the error
+ * @param <R> the object to return as a result of this callback, should be null as we are only
+ * closing the span and logging the error
+ */
+ public static class RedisAsyncExceptionFunction<Throwable, R extends Object>
+ implements Function<Throwable, R> {
+ private final Span span;
+
+ public RedisAsyncExceptionFunction(Span span) {
+ this.span = span;
+ }
+
+ @Override
+ public R apply(Throwable throwable) {
+ Tags.ERROR.set(this.span, true);
+ this.span.log(Collections.singletonMap("error.object", throwable));
+ this.span.finish();
+ return null;
+ }
+ }
+
+ /**
+ * Callback class to close the span on a normal execution of the RedisFuture returned by the
+ * lettice async API
+ *
+ * @param <T> the object (async result) from lettuce async API computation
+ */
+ public static class RedisAsyncConsumer<T> implements Consumer<T> {
+ private final Span span;
+
+ public RedisAsyncConsumer(Span span) {
+ this.span = span;
+ }
+
+ @Override
+ public void accept(Object o) {
+ this.span.finish();
+ }
+ }
+}

ADD
dd-java-agent/instrumentation/lettuce-5/src/main/java/datadog/trace/instrumentation/lettuce/RedisClientInstrumentation.java
+package datadog.trace.instrumentation.lettuce;
+
+import static net.bytebuddy.matcher.ElementMatchers.*;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.DDAdvice;
+import datadog.trace.agent.tooling.DDTransformers;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.api.DDTags;
+import io.lettuce.core.RedisURI;
+import io.opentracing.Scope;
+import io.opentracing.Span;
+import io.opentracing.tag.Tags;
+import io.opentracing.util.GlobalTracer;
+import java.util.Collections;
+import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.asm.Advice;
+
+@AutoService(Instrumenter.class)
+public final class RedisClientInstrumentation extends Instrumenter.Configurable {
+
+ public static final String SERVICE_NAME = "redis";
+ public static final String COMPONENT_NAME = SERVICE_NAME + "-client";
+ public static final String RESOURCE_NAME_PREFIX = "CONNECT:";
+ public static final String REDIS_URL_TAGNAME = "db.redis.url";
+ public static final String REDIS_DB_INDEX_TAG_NAME = "db.redis.dbIndex";
+
+ public RedisClientInstrumentation() {
+ super(SERVICE_NAME);
+ }
+
+ @Override
+ public AgentBuilder apply(final AgentBuilder agentBuilder) {
+ return agentBuilder
+ .type(
+ named("io.lettuce.core.RedisClient")
+ .and(hasSuperType(named("io.lettuce.core.AbstractRedisClient"))))
+ .transform(DDTransformers.defaultTransformers())
+ .transform(
+ DDAdvice.create()
+ .advice(
+ isMethod()
+ .and(isPrivate())
+ .and(returns(hasSuperType(named("io.lettuce.core.api.StatefulConnection"))))
+ .and(nameStartsWith("connect"))
+ .and(takesArgument(1, named("io.lettuce.core.RedisURI"))),
+ StatefulRedisConnectionAdvice.class.getName()))
+ .asDecorator();
+ }
+
+ public static class StatefulRedisConnectionAdvice {
+
+ @Advice.OnMethodEnter(suppress = Throwable.class)
+ public static Scope startSpan(@Advice.Argument(1) final RedisURI redisURI) {
+
+ final Scope scope = GlobalTracer.get().buildSpan(SERVICE_NAME + ".query").startActive(true);
+
+ final Span span = scope.span();
+ Tags.DB_TYPE.set(span, SERVICE_NAME);
+ Tags.SPAN_KIND.set(span, Tags.SPAN_KIND_CLIENT);
+ Tags.COMPONENT.set(span, COMPONENT_NAME);
+
+ final String url =
+ redisURI.getHost() + ":" + redisURI.getPort() + "/" + redisURI.getDatabase();
+ span.setTag(REDIS_URL_TAGNAME, url);
+ span.setTag(REDIS_DB_INDEX_TAG_NAME, redisURI.getDatabase());
+ span.setTag(DDTags.RESOURCE_NAME, RESOURCE_NAME_PREFIX + url);
+ span.setTag(DDTags.SERVICE_NAME, SERVICE_NAME);
+ span.setTag(DDTags.SPAN_TYPE, SERVICE_NAME);
+
+ return scope;
+ }
+
+ @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+ public static void stopSpan(
+ @Advice.Enter final Scope scope, @Advice.Thrown final Throwable throwable) {
+ if (throwable != null) {
+ final Span span = scope.span();
+ Tags.ERROR.set(span, true);
+ span.log(Collections.singletonMap("error.object", throwable));
+ }
+ scope.close();
+ }
+ }
+}

ADD
dd-java-agent/instrumentation/lettuce-5/src/test/groovy/LettuceSyncClientTest.groovy
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.instrumentation.lettuce.RedisAsyncCommandsInstrumentation
+import datadog.trace.instrumentation.lettuce.RedisClientInstrumentation
+import io.lettuce.core.RedisClient
+import io.lettuce.core.RedisConnectionException
+import io.lettuce.core.api.StatefulConnection
+import io.lettuce.core.api.sync.RedisCommands
+import redis.embedded.RedisServer
+import spock.lang.Shared
+
+class LettuceSyncClientTest extends AgentTestRunner {
+
+ @Shared
+ public static final String HOST = "127.0.0.1"
+ public static final int PORT = 6399
+ public static final int DB_INDEX = 0
+ @Shared
+ public static final String DB_ADDR = HOST + ":" + PORT + "/" + DB_INDEX
+ @Shared
+ public static final String DB_ADDR_NON_EXISTENT = HOST + ":" + 9999 + "/" + DB_INDEX
+ @Shared
+ public static final String DB_URI_NON_EXISTENT = "redis://" + DB_ADDR_NON_EXISTENT
+ public static final String EMBEDDED_DB_URI = "redis://" + DB_ADDR
+
+ @Shared
+ RedisServer redisServer = RedisServer.builder()
+ // bind to localhost to avoid firewall popup
+ .setting("bind " + HOST)
+ // set max memory to avoid problems in CI
+ .setting("maxmemory 128M")
+ .port(PORT).build()
+
+ @Shared
+ RedisClient redisClient = RedisClient.create(EMBEDDED_DB_URI)
+
+ @Shared
+ RedisCommands<String, ?> syncCommands = null
+
+ @Shared
+ Map<String, String> testHashMap = [
+ firstname: "John",
+ lastname: "Doe",
+ age: "53"
+ ]
+
+ def setupSpec() {
+ println "Using redis: $redisServer.args"
+ redisServer.start()
+ StatefulConnection connection = redisClient.connect()
+ syncCommands = connection.sync()
+ }
+
+ def cleanupSpec() {
+ redisServer.stop()
+ }
+
+ def setup() {
+ TEST_WRITER.start()
+ }
+
+ def "connect"() {
+ setup:
+ RedisClient testConnectionClient = RedisClient.create(EMBEDDED_DB_URI)
+ testConnectionClient.connect()
+
+ expect:
+ TEST_WRITER.size() == 1
+
+ def trace = TEST_WRITER.firstTrace()
+ trace.size() == 1
+ def span = trace[0]
+ span.getServiceName() == RedisClientInstrumentation.SERVICE_NAME
+ span.getOperationName() == "redis.query"
+ span.getType() == RedisClientInstrumentation.SERVICE_NAME
+ span.getResourceName() == RedisClientInstrumentation.RESOURCE_NAME_PREFIX + DB_ADDR
+ !span.context().getErrorFlag()
+
+ def tags = span.context().tags
+ tags[RedisClientInstrumentation.REDIS_URL_TAGNAME] == DB_ADDR
+ tags[RedisClientInstrumentation.REDIS_DB_INDEX_TAG_NAME] == 0
+ tags["span.kind"] == "client"
+ tags["span.type"] == RedisClientInstrumentation.SERVICE_NAME
+ tags["db.type"] == RedisClientInstrumentation.SERVICE_NAME
+ tags["component"] == RedisClientInstrumentation.COMPONENT_NAME
+ tags["thread.name"] != null
+ tags["thread.id"] != null
+ }
+
+ def "connect exception"() {
+ setup:
+ RedisClient testConnectionClient = RedisClient.create(DB_URI_NON_EXISTENT)
+ try {
+ testConnectionClient.connect()
+ } catch (RedisConnectionException rce) { }
+
+ expect:
+ TEST_WRITER.size() == 1
+
+ def trace = TEST_WRITER.firstTrace()
+ trace.size() == 1
+ def span = trace[0]
+ span.getServiceName() == RedisClientInstrumentation.SERVICE_NAME
+ span.getOperationName() == "redis.query"
+ span.getType() == RedisClientInstrumentation.SERVICE_NAME
+ span.getResourceName() == RedisClientInstrumentation.RESOURCE_NAME_PREFIX + DB_ADDR_NON_EXISTENT
+ span.context().getErrorFlag()
+
+ def tags = span.context().tags
+ tags[RedisClientInstrumentation.REDIS_URL_TAGNAME] == DB_ADDR_NON_EXISTENT
+ tags[RedisClientInstrumentation.REDIS_DB_INDEX_TAG_NAME] == 0
+ tags["span.kind"] == "client"
+ tags["span.type"] == RedisClientInstrumentation.SERVICE_NAME
+ tags["db.type"] == RedisClientInstrumentation.SERVICE_NAME
+ tags["component"] == RedisClientInstrumentation.COMPONENT_NAME
+ tags["thread.name"] != null
+ tags["thread.id"] != null
+ }
+
+ def "set command"() {
+ setup:
+ String res = syncCommands.set("TESTKEY", "TESTVAL")
+ TEST_WRITER.waitForTraces(1)
+
+ expect:
+ res == "OK"
+ TEST_WRITER.size() == 1
+
+
+ def trace = TEST_WRITER.firstTrace()
+ trace.size() == 1
+ def span = trace[0]
+ span.getServiceName() == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ span.getOperationName() == "redis.query"
+ span.getType() == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ span.getResourceName() == "SET"
+ !span.context().getErrorFlag()
+
+ def tags = span.context().tags
+ tags["span.kind"] == "client"
+ tags["span.type"] == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ tags["db.type"] == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ tags["db.redis.command.args"] == "key<TESTKEY> value<TESTVAL>"
+ tags["component"] == RedisAsyncCommandsInstrumentation.COMPONENT_NAME
+ tags["thread.name"] != null
+ tags["thread.id"] != null
+ }
+
+ def "get command"() {
+ setup:
+ String res = syncCommands.get("TESTKEY")
+ TEST_WRITER.waitForTraces(1)
+
+ expect:
+ res == "TESTVAL"
+ TEST_WRITER.size() == 1
+
+ def trace = TEST_WRITER.firstTrace()
+ trace.size() == 1
+ def span = trace[0]
+ span.getServiceName() == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ span.getOperationName() == "redis.query"
+ span.getType() == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ span.getResourceName() == "GET"
+ !span.context().getErrorFlag()
+
+ def tags = span.context().tags
+ tags["span.kind"] == "client"
+ tags["span.type"] == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ tags["db.type"] == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ tags["db.redis.command.args"] == "key<TESTKEY>"
+ tags["component"] == RedisAsyncCommandsInstrumentation.COMPONENT_NAME
+ tags["thread.name"] != null
+ tags["thread.id"] != null
+ }
+
+ def "get non existent key command"() {
+ setup:
+ String res = syncCommands.get("NON_EXISTENT_KEY")
+ TEST_WRITER.waitForTraces(1)
+
+ expect:
+ res == null
+ TEST_WRITER.size() == 1
+
+ def trace = TEST_WRITER.firstTrace()
+ trace.size() == 1
+ def span = trace[0]
+ span.getServiceName() == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ span.getOperationName() == "redis.query"
+ span.getType() == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ span.getResourceName() == "GET"
+ !span.context().getErrorFlag()
+
+ def tags = span.context().tags
+ tags["span.kind"] == "client"
+ tags["span.type"] == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ tags["db.type"] == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ tags["db.redis.command.args"] == "key<NON_EXISTENT_KEY>"
+ tags["component"] == RedisAsyncCommandsInstrumentation.COMPONENT_NAME
+ tags["thread.name"] != null
+ tags["thread.id"] != null
+ }
+
+ def "command with no arguments"() {
+ setup:
+ def keyRetrieved = syncCommands.randomkey()
+ TEST_WRITER.waitForTraces(1)
+
+ expect:
+ keyRetrieved == "TESTKEY"
+ TEST_WRITER.size() == 1
+
+ def trace = TEST_WRITER.firstTrace()
+ trace.size() == 1
+
+ def span = trace[0]
+ span.getServiceName() == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ span.getOperationName() == "redis.query"
+ span.getType() == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ span.getResourceName() == "RANDOMKEY"
+ !span.context().getErrorFlag()
+
+ def tags = span.context().tags
+ tags["span.kind"] == "client"
+ tags["span.type"] == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ tags["db.type"] == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ tags["db.redis.command.args"] == null
+ tags["component"] == RedisAsyncCommandsInstrumentation.COMPONENT_NAME
+ tags["thread.name"] != null
+ tags["thread.id"] != null
+ }
+
+ def "list command"() {
+ setup:
+ long res = syncCommands.lpush("TESTLIST", "TESTLIST ELEMENT")
+ TEST_WRITER.waitForTraces(1)
+
+ expect:
+ res == 1
+ def trace = TEST_WRITER.firstTrace()
+ trace.size() == 1
+
+ def span = trace[0]
+ span.getServiceName() == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ span.getOperationName() == "redis.query"
+ span.getType() == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ span.getResourceName() == "LPUSH"
+ !span.context().getErrorFlag()
+
+ def tags = span.context().tags
+ tags["span.kind"] == "client"
+ tags["span.type"] == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ tags["db.type"] == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ tags["db.redis.command.args"] == "key<TESTLIST> value<TESTLIST ELEMENT>"
+ tags["component"] == RedisAsyncCommandsInstrumentation.COMPONENT_NAME
+ tags["thread.name"] != null
+ tags["thread.id"] != null
+ }
+
+ def "hash set command"() {
+ setup:
+ def res = syncCommands.hmset("user", testHashMap)
+ TEST_WRITER.waitForTraces(1)
+
+ expect:
+ res == "OK"
+ def trace = TEST_WRITER.firstTrace()
+ trace.size() == 1
+
+ def span = trace[0]
+ span.getServiceName() == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ span.getOperationName() == "redis.query"
+ span.getType() == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ span.getResourceName() == "HMSET"
+ !span.context().getErrorFlag()
+
+ def tags = span.context().tags
+ tags["span.kind"] == "client"
+ tags["span.type"] == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ tags["db.type"] == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ tags["db.redis.command.args"] == "key<user> key<firstname> value<John> key<lastname> value<Doe> key<age> value<53>"
+ tags["component"] == RedisAsyncCommandsInstrumentation.COMPONENT_NAME
+ tags["thread.name"] != null
+ tags["thread.id"] != null
+ }
+
+ def "hash getall command"() {
+ setup:
+ Map<String, String> res = syncCommands.hgetall("user")
+ TEST_WRITER.waitForTraces(1)
+
+ expect:
+ res == testHashMap
+ def trace = TEST_WRITER.firstTrace()
+ trace.size() == 1
+
+ def span = trace[0]
+ span.getServiceName() == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ span.getOperationName() == "redis.query"
+ span.getType() == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ span.getResourceName() == "HGETALL"
+ !span.context().getErrorFlag()
+
+ def tags = span.context().tags
+ tags["span.kind"] == "client"
+ tags["span.type"] == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ tags["db.type"] == RedisAsyncCommandsInstrumentation.SERVICE_NAME
+ tags["db.redis.command.args"] == "key<user>"
+ tags["component"] == RedisAsyncCommandsInstrumentation.COMPONENT_NAME
+ tags["thread.name"] != null
+ tags["thread.id"] != null
+ }
+}

MODIFY
settings.gradle
settings.gradle
@@ -33,6 +33,7 @@ include ':dd-java-agent:instrumentation:jms-1'
include ':dd-java-agent:instrumentation:jms-2'
include ':dd-java-agent:instrumentation:kafka-clients-0.11'
include ':dd-java-agent:instrumentation:kafka-streams-0.11'
+include ':dd-java-agent:instrumentation:lettuce-5'
include ':dd-java-agent:instrumentation:mongo-3.1'
include ':dd-java-agent:instrumentation:mongo-async-3.3'
include ':dd-java-agent:instrumentation:okhttp-3'

lettuce redis client sync and async api instrumentation with sync unit tests complete