Provide instrumentation support for lettuce-4.0-4.5.0 sync/async APIs

dd-java-agent/instrumentation/lettuce-4/lettuce-4.gradle

@@ -0,0 +1,33 @@
+// Set properties before any plugins get loaded
+ext {
+  minJavaVersionForTests = JavaVersion.VERSION_1_8
+}
+
+muzzle {
+  pass {
+    group = "biz.paluch.redis"
+    module = "lettuce"
+    versions = "[4.0.Final,)"
+    assertInverse = true
+  }
+}
+
+apply from: "${rootDir}/gradle/java.gradle"
+
+apply plugin: 'org.unbroken-dome.test-sets'
+
+testSets {
+  latestDepTest {
+    dirName = 'test'
+  }
+}
+
+dependencies {
+  compileOnly group: 'biz.paluch.redis', name: 'lettuce', version: '4.0.Final'
+  main_java8CompileOnly group: 'biz.paluch.redis', name: 'lettuce', version: '4.0.Final'
+
+  testCompile group: 'com.github.kstyrc', name: 'embedded-redis', version: '0.6'
+  testCompile group: 'biz.paluch.redis', name: 'lettuce', version: '4.0.Final'
+
+  latestDepTestCompile group: 'biz.paluch.redis', name: 'lettuce', version: '4.+'
+}

dd-java-agent/instrumentation/lettuce-4/src/main/java/datadog/trace/instrumentation/lettuce/LettuceAsyncCommandsInstrumentation.java

@@ -0,0 +1,43 @@
+package datadog.trace.instrumentation.lettuce;
+
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import java.util.Map;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(Instrumenter.class)
+public class LettuceAsyncCommandsInstrumentation extends Instrumenter.Default {
+
+  public LettuceAsyncCommandsInstrumentation() {
+    super("lettuce", "lettuce-4-async");
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("com.lambdaworks.redis.AbstractRedisAsyncCommands");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".LettuceClientDecorator", packageName + ".InstrumentationPoints"
+    };
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+        isMethod()
+            .and(named("dispatch"))
+            .and(takesArgument(0, named("com.lambdaworks.redis.protocol.RedisCommand"))),
+        // Cannot reference class directly here because it would lead to class load failure on Java7
+        packageName + ".LettuceAsyncCommandsAdvice");
+  }
+}

dd-java-agent/instrumentation/lettuce-4/src/main/java/datadog/trace/instrumentation/lettuce/LettuceClientInstrumentation.java

@@ -0,0 +1,40 @@
+package datadog.trace.instrumentation.lettuce;
+
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import java.util.Map;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(Instrumenter.class)
+public final class LettuceClientInstrumentation extends Instrumenter.Default {
+
+  public LettuceClientInstrumentation() {
+    super("lettuce");
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("com.lambdaworks.redis.RedisClient");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".LettuceClientDecorator", packageName + ".InstrumentationPoints"
+    };
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+        isMethod().and(named("connectStandalone")),
+        // Cannot reference class directly here because it would lead to class load failure on Java7
+        packageName + ".RedisConnectionAdvice");
+  }
+}

dd-java-agent/instrumentation/lettuce-4/src/main/java8/datadog/trace/instrumentation/lettuce/InstrumentationPoints.java

@@ -0,0 +1,120 @@
+package datadog.trace.instrumentation.lettuce;
+
+import static com.lambdaworks.redis.protocol.CommandKeyword.SEGFAULT;
+import static com.lambdaworks.redis.protocol.CommandType.CLIENT;
+import static com.lambdaworks.redis.protocol.CommandType.CLUSTER;
+import static com.lambdaworks.redis.protocol.CommandType.COMMAND;
+import static com.lambdaworks.redis.protocol.CommandType.CONFIG;
+import static com.lambdaworks.redis.protocol.CommandType.DEBUG;
+import static com.lambdaworks.redis.protocol.CommandType.SCRIPT;
+import static com.lambdaworks.redis.protocol.CommandType.SHUTDOWN;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.instrumentation.lettuce.LettuceClientDecorator.DECORATE;
+
+import com.lambdaworks.redis.RedisURI;
+import com.lambdaworks.redis.protocol.AsyncCommand;
+import com.lambdaworks.redis.protocol.CommandType;
+import com.lambdaworks.redis.protocol.ProtocolKeyword;
+import com.lambdaworks.redis.protocol.RedisCommand;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+
+import java.util.EnumSet;
+import java.util.Set;
+import java.util.concurrent.CancellationException;
+
+public final class InstrumentationPoints {
+
+  private static final Set<CommandType> NON_INSTRUMENTING_COMMANDS = EnumSet.of(SHUTDOWN, DEBUG);
+
+  private static final Set<CommandType> AGENT_CRASHING_COMMANDS =
+    EnumSet.of(CLIENT, CLUSTER, COMMAND, CONFIG, DEBUG, SCRIPT);
+
+  public static final String AGENT_CRASHING_COMMAND_PREFIX = "COMMAND-NAME:";
+
+  public static AgentScope beforeCommand(RedisCommand<?, ?, ?> command) {
+    AgentSpan span = startSpan("redis.query");
+    DECORATE.afterStart(span);
+    DECORATE.onCommand(span, command);
+    return activateSpan(span, finishSpanEarly(command));
+  }
+
+  public static void afterCommand(RedisCommand<?, ?, ?> command,
+                                  AgentScope scope,
+                                  Throwable throwable,
+                                  AsyncCommand<?, ?, ?> asyncCommand) {
+      AgentSpan span = scope.span();
+      if (throwable != null) {
+        DECORATE.onError(span, throwable);
+        DECORATE.beforeFinish(span);
+        span.finish();
+      } else if (!finishSpanEarly(command)) {
+        asyncCommand.handleAsync((value, ex) -> {
+          if (ex instanceof CancellationException) {
+            span.setTag("db.command.cancelled", true);
+          } else {
+            DECORATE.onError(span, ex);
+          }
+          DECORATE.beforeFinish(span);
+          span.finish();
+          return null;
+        });
+      }
+      scope.close();
+  }
+
+  public static AgentScope beforeConnect(RedisURI redisURI) {
+    AgentSpan span = startSpan("redis.query");
+    DECORATE.afterStart(span);
+    DECORATE.onConnection(span, redisURI);
+    return activateSpan(span, false);
+  }
+
+  public static void afterConnect(AgentScope scope, Throwable throwable) {
+    AgentSpan span = scope.span();
+    if (throwable != null) {
+      DECORATE.onError(span, throwable);
+      DECORATE.beforeFinish(span);
+    }
+    span.finish();
+    scope.close();
+  }
+
+  /**
+   * Determines whether a redis command should finish its relevant span early (as soon as tags are
+   * added and the command is executed) because these commands have no return values/call backs, so
+   * we must close the span early in order to provide info for the users
+   *
+   * @param command
+   * @return true if finish the span early (the command will not have a return value)
+   */
+  public static boolean finishSpanEarly(RedisCommand<?, ?, ?> command) {
+    ProtocolKeyword keyword = command.getType();
+    return isNonInstrumentingCommand(keyword) || isNonInstrumentingKeyword(keyword);
+  }
+
+  private static boolean isNonInstrumentingCommand(ProtocolKeyword keyword) {
+    return keyword instanceof CommandType && NON_INSTRUMENTING_COMMANDS.contains(keyword);
+  }
+
+  private static boolean isNonInstrumentingKeyword(ProtocolKeyword keyword) {
+    return keyword == SEGFAULT;
+  }
+
+  /**
+   * Workaround to keep trace agent from crashing Currently the commands in
+   * AGENT_CRASHING_COMMANDS will crash the trace agent and traces with these commands as the
+   * resource name will not be processed by the trace agent
+   *
+   * @param keyword the actual redis command
+   * @return the redis command with a prefix if it is a command that will crash the trace agent,
+   *     otherwise, the original command is returned.
+   */
+  public static String getCommandResourceName(ProtocolKeyword keyword) {
+    if (keyword instanceof CommandType && AGENT_CRASHING_COMMANDS.contains(keyword)) {
+      return AGENT_CRASHING_COMMAND_PREFIX + keyword.name();
+    }
+    return keyword.name();
+  }
+}

dd-java-agent/instrumentation/lettuce-4/src/main/java8/datadog/trace/instrumentation/lettuce/LettuceAsyncCommandsAdvice.java

@@ -0,0 +1,24 @@
+package datadog.trace.instrumentation.lettuce;
+
+import com.lambdaworks.redis.protocol.AsyncCommand;
+import com.lambdaworks.redis.protocol.RedisCommand;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import net.bytebuddy.asm.Advice;
+
+public class LettuceAsyncCommandsAdvice {
+
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  public static AgentScope onEnter(@Advice.Argument(0) final RedisCommand<?, ?, ?> command) {
+    return InstrumentationPoints.beforeCommand(command);
+  }
+
+  @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+  public static void onExit(
+      @Advice.Argument(0) final RedisCommand<?, ?, ?> command,
+      @Advice.Enter final AgentScope scope,
+      @Advice.Thrown final Throwable throwable,
+      @Advice.Return final AsyncCommand<?, ?, ?> asyncCommand) {
+    InstrumentationPoints.afterCommand(command, scope, throwable, asyncCommand);
+  }
+
+}

dd-java-agent/instrumentation/lettuce-4/src/main/java8/datadog/trace/instrumentation/lettuce/LettuceClientDecorator.java

@@ -0,0 +1,77 @@
+package datadog.trace.instrumentation.lettuce;
+
+import static datadog.trace.instrumentation.lettuce.InstrumentationPoints.getCommandResourceName;
+
+import com.lambdaworks.redis.RedisURI;
+import com.lambdaworks.redis.protocol.RedisCommand;
+import datadog.trace.api.DDSpanTypes;
+import datadog.trace.api.DDTags;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.Tags;
+import datadog.trace.bootstrap.instrumentation.decorator.DatabaseClientDecorator;
+
+public class LettuceClientDecorator extends DatabaseClientDecorator<RedisURI> {
+
+  public static final LettuceClientDecorator DECORATE = new LettuceClientDecorator();
+
+  @Override
+  protected String[] instrumentationNames() {
+    return new String[] {"lettuce"};
+  }
+
+  @Override
+  protected String service() {
+    return "redis";
+  }
+
+  @Override
+  protected String component() {
+    return "redis-client";
+  }
+
+  @Override
+  protected String spanType() {
+    return DDSpanTypes.REDIS;
+  }
+
+  @Override
+  protected String dbType() {
+    return "redis";
+  }
+
+  @Override
+  protected String dbUser(RedisURI connection) {
+    return null;
+  }
+
+  @Override
+  protected String dbInstance(RedisURI connection) {
+    return null;
+  }
+
+  @Override
+  public AgentSpan onConnection(AgentSpan span, RedisURI connection) {
+    if (connection != null) {
+      span.setTag(Tags.PEER_HOSTNAME, connection.getHost());
+      span.setTag(Tags.PEER_PORT, connection.getPort());
+      span.setTag("db.redis.dbIndex", connection.getDatabase());
+      span.setTag(DDTags.RESOURCE_NAME, resourceName(connection));
+    }
+    return super.onConnection(span, connection);
+  }
+
+  public AgentSpan onCommand(AgentSpan span, RedisCommand<?, ?, ?> command) {
+    span.setTag(DDTags.RESOURCE_NAME,
+        null == command ? "Redis Command" : getCommandResourceName(command.getType()));
+    return span;
+  }
+
+  private static String resourceName(RedisURI connection) {
+    return "CONNECT:"
+      + connection.getHost()
+      + ":"
+      + connection.getPort()
+      + "/"
+      + connection.getDatabase();
+  }
+}

dd-java-agent/instrumentation/lettuce-4/src/main/java8/datadog/trace/instrumentation/lettuce/RedisConnectionAdvice.java

@@ -0,0 +1,19 @@
+package datadog.trace.instrumentation.lettuce;
+
+import com.lambdaworks.redis.RedisURI;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import net.bytebuddy.asm.Advice;
+
+public class RedisConnectionAdvice {
+
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  public static AgentScope onEnter(@Advice.Argument(1) final RedisURI redisURI) {
+    return InstrumentationPoints.beforeConnect(redisURI);
+  }
+
+  @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+  public static void onExit(@Advice.Enter final AgentScope scope,
+                            @Advice.Thrown final Throwable throwable) {
+    InstrumentationPoints.afterConnect(scope, throwable);
+  }
+}