Add instrumentation for Rediscala library

dd-java-agent/instrumentation/rediscala/rediscala.gradle

@@ -0,0 +1,41 @@
+muzzle {
+  pass {
+    group = "com.github.Ma27"
+    module = "rediscala_2.11"
+    versions = "[1.8.1,)"
+    assertInverse = true
+  }
+
+  pass {
+    group = "com.github.Ma27"
+    module = "rediscala_2.12"
+    versions = "[1.8.1,)"
+    assertInverse = true
+  }
+
+  pass {
+    group = "com.github.Ma27"
+    module = "rediscala_2.13"
+    versions = "[1.9.0,)"
+    assertInverse = true
+  }
+}
+
+apply from: "${rootDir}/gradle/java.gradle"
+
+apply plugin: 'org.unbroken-dome.test-sets'
+
+testSets {
+  latestDepTest {
+    dirName = 'test'
+  }
+}
+
+dependencies {
+  compileOnly group: 'com.github.Ma27', name: 'rediscala_2.11', version: '1.8.1'
+
+  testCompile group: 'com.github.Ma27', name: 'rediscala_2.11', version: '1.8.1'
+  testCompile group: 'com.github.kstyrc', name: 'embedded-redis', version: '0.6'
+
+  latestDepTestCompile group: 'com.github.Ma27', name: 'rediscala_2.11', version: '+'
+}

dd-java-agent/instrumentation/rediscala/src/main/java/datadog/trace/instrumentation/rediscala/RediscalaClientDecorator.java

@@ -0,0 +1,44 @@
+package datadog.trace.instrumentation.rediscala;
+
+import datadog.trace.agent.decorator.DatabaseClientDecorator;
+import datadog.trace.api.DDSpanTypes;
+import redis.RedisCommand;
+
+public class RediscalaClientDecorator extends DatabaseClientDecorator<RedisCommand> {
+  public static final RediscalaClientDecorator DECORATE = new RediscalaClientDecorator();
+
+  @Override
+  protected String[] instrumentationNames() {
+    return new String[] {"rediscala", "redis"};
+  }
+
+  @Override
+  protected String service() {
+    return "redis";
+  }
+
+  @Override
+  protected String component() {
+    return "redis-command";
+  }
+
+  @Override
+  protected String spanType() {
+    return DDSpanTypes.REDIS;
+  }
+
+  @Override
+  protected String dbType() {
+    return "redis";
+  }
+
+  @Override
+  protected String dbUser(final RedisCommand session) {
+    return null;
+  }
+
+  @Override
+  protected String dbInstance(final RedisCommand session) {
+    return null;
+  }
+}

dd-java-agent/instrumentation/rediscala/src/main/java/datadog/trace/instrumentation/rediscala/RediscalaInstrumentation.java

@@ -0,0 +1,125 @@
+package datadog.trace.instrumentation.rediscala;
+
+import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
+import static datadog.trace.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.instrumentation.api.AgentTracer.activeScope;
+import static datadog.trace.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.instrumentation.rediscala.RediscalaClientDecorator.DECORATE;
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.context.TraceScope;
+import datadog.trace.instrumentation.api.AgentScope;
+import datadog.trace.instrumentation.api.AgentSpan;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import redis.RedisCommand;
+import scala.concurrent.ExecutionContext;
+import scala.concurrent.Future;
+import scala.runtime.AbstractFunction1;
+import scala.util.Try;
+
+@AutoService(Instrumenter.class)
+public final class RediscalaInstrumentation extends Instrumenter.Default {
+
+  private static final String SERVICE_NAME = "redis";
+  private static final String COMPONENT_NAME = SERVICE_NAME + "-command";
+
+  public RediscalaInstrumentation() {
+    super("rediscala", "redis");
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return safeHasSuperType(named("redis.ActorRequest"))
+        .or(safeHasSuperType(named("redis.Request")))
+        .or(safeHasSuperType(named("redis.BufferedRequest")))
+        .or(safeHasSuperType(named("redis.RoundRobinPoolRequest")));
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      RediscalaInstrumentation.class.getName() + "$OnCompleteHandler",
+      "datadog.trace.agent.decorator.BaseDecorator",
+      "datadog.trace.agent.decorator.ClientDecorator",
+      "datadog.trace.agent.decorator.DatabaseClientDecorator",
+      packageName + ".RediscalaClientDecorator",
+    };
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+        isMethod()
+            .and(isPublic())
+            .and(named("send"))
+            .and(takesArgument(0, named("redis.RedisCommand")))
+            .and(returns(named("scala.concurrent.Future"))),
+        RediscalaInstrumentation.class.getName() + "$RediscalaAdvice");
+  }
+
+  public static class RediscalaAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentScope onEnter(@Advice.Argument(0) final RedisCommand cmd) {
+      final AgentSpan span = startSpan("redis.command");
+      DECORATE.afterStart(span);
+      DECORATE.onStatement(span, cmd.getClass().getName());
+      return activateSpan(span, true);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void stopSpan(
+        @Advice.Enter final AgentScope scope,
+        @Advice.Thrown final Throwable throwable,
+        @Advice.FieldValue("executionContext") final ExecutionContext ctx,
+        @Advice.Return(readOnly = false) final Future<Object> responseFuture) {
+
+      final AgentSpan span = scope.span();
+
+      if (throwable == null) {
+        responseFuture.onComplete(new OnCompleteHandler(span), ctx);
+      } else {
+        DECORATE.onError(span, throwable);
+        DECORATE.beforeFinish(span);
+        span.finish();
+      }
+      scope.close();
+    }
+  }
+
+  public static class OnCompleteHandler extends AbstractFunction1<Try<Object>, Void> {
+    private final AgentSpan span;
+
+    public OnCompleteHandler(final AgentSpan span) {
+      this.span = span;
+    }
+
+    @Override
+    public Void apply(final Try<Object> result) {
+      try {
+        if (result.isFailure()) {
+          DECORATE.onError(span, result.failed().get());
+        }
+        DECORATE.beforeFinish(span);
+        final TraceScope scope = activeScope();
+        if (scope != null) {
+          scope.setAsyncPropagation(false);
+        }
+      } finally {
+        span.finish();
+      }
+      return null;
+    }
+  }
+}

dd-java-agent/instrumentation/rediscala/src/test/groovy/RediscalaClientTest.groovy

@@ -0,0 +1,138 @@
+import akka.actor.ActorSystem
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.api.Config
+import datadog.trace.api.DDSpanTypes
+import datadog.trace.instrumentation.api.Tags
+import redis.ByteStringSerializerLowPriority
+import redis.ByteStringDeserializerDefault
+import redis.RedisClient
+import redis.RedisDispatcher
+import redis.embedded.RedisServer
+import scala.Option
+import scala.concurrent.Await
+import scala.concurrent.duration.Duration
+import spock.lang.Shared
+
+class RediscalaClientTest extends AgentTestRunner {
+
+  public static final int PORT = 6399
+
+  @Shared
+  RedisServer redisServer = RedisServer.builder()
+  // bind to localhost to avoid firewall popup
+    .setting("bind 127.0.0.1")
+  // set max memory to avoid problems in CI
+    .setting("maxmemory 128M")
+    .port(PORT).build()
+
+  @Shared
+  ActorSystem system
+
+  @Shared
+  RedisClient redisClient
+
+  def setupSpec() {
+    system = ActorSystem.create()
+    redisClient = new RedisClient("localhost",
+      PORT,
+      Option.apply(null),
+      Option.apply(null),
+      "RedisClient",
+      Option.apply(null),
+      system,
+      new RedisDispatcher("rediscala.rediscala-client-worker-dispatcher"))
+
+    println "Using redis: $redisServer.args"
+    redisServer.start()
+
+    // This setting should have no effect since decorator returns null for the instance.
+    System.setProperty(Config.PREFIX + Config.DB_CLIENT_HOST_SPLIT_BY_INSTANCE, "true")
+  }
+
+  def cleanupSpec() {
+    redisServer.stop()
+    system?.terminate()
+    System.clearProperty(Config.PREFIX + Config.DB_CLIENT_HOST_SPLIT_BY_INSTANCE)
+  }
+
+  def setup() {
+    TEST_WRITER.start()
+  }
+
+  def "set command"() {
+    when:
+    def value = redisClient.set("foo",
+      "bar",
+      Option.apply(null),
+      Option.apply(null),
+      false,
+      false,
+      new ByteStringSerializerLowPriority.String$())
+
+
+    then:
+    Await.result(value, Duration.apply("3 second")) == true
+    assertTraces(1) {
+      trace(0, 1) {
+        span(0) {
+          serviceName "redis"
+          operationName "redis.query"
+          resourceName "redis.api.strings.Set"
+          spanType DDSpanTypes.REDIS
+          tags {
+            "$Tags.COMPONENT" "redis-command"
+            "$Tags.DB_TYPE" "redis"
+            "$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT
+            defaultTags()
+          }
+        }
+      }
+    }
+  }
+
+  def "get command"() {
+    when:
+    def write = redisClient.set("bar",
+      "baz",
+      Option.apply(null),
+      Option.apply(null),
+      false,
+      false,
+      new ByteStringSerializerLowPriority.String$())
+    def value = redisClient.get("bar", new ByteStringDeserializerDefault.String$())
+
+    then:
+    Await.result(write, Duration.apply("3 second")) == true
+    Await.result(value, Duration.apply("3 second")) == Option.apply("baz")
+    assertTraces(2) {
+      trace(0, 1) {
+        span(0) {
+          serviceName "redis"
+          operationName "redis.query"
+          resourceName "redis.api.strings.Set"
+          spanType DDSpanTypes.REDIS
+          tags {
+            "$Tags.COMPONENT" "redis-command"
+            "$Tags.DB_TYPE" "redis"
+            "$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT
+            defaultTags()
+          }
+        }
+      }
+      trace(1, 1) {
+        span(0) {
+          serviceName "redis"
+          operationName "redis.query"
+          resourceName "redis.api.strings.Get"
+          spanType DDSpanTypes.REDIS
+          tags {
+            "$Tags.COMPONENT" "redis-command"
+            "$Tags.DB_TYPE" "redis"
+            "$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT
+            defaultTags()
+          }
+        }
+      }
+    }
+  }
+}

settings.gradle

@@ -126,6 +126,7 @@ include ':dd-java-agent:instrumentation:play-ws-2'
 include ':dd-java-agent:instrumentation:play-ws-2.1'
 include ':dd-java-agent:instrumentation:rabbitmq-amqp-2.7'
 include ':dd-java-agent:instrumentation:ratpack-1.4'
+include ':dd-java-agent:instrumentation:rediscala'
 include ':dd-java-agent:instrumentation:reactor-core-3.1'
 include ':dd-java-agent:instrumentation:rmi'
 include ':dd-java-agent:instrumentation:rxjava-1'