Aerospike 4 instrumentation

.circleci/config.yml

@@ -129,6 +129,8 @@ jobs:
       - image: memcached
         # This is used by rabbitmq instrumentation tests
       - image: rabbitmq
+        # This is used by aerospike instrumentation tests
+      - image: aerospike
 
     parameters:
       testTask:

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

@@ -0,0 +1,26 @@
+
+apply from: "$rootDir/gradle/java.gradle"
+
+apply plugin: 'org.unbroken-dome.test-sets'
+
+testSets {
+  latestDepTest {
+    dirName = 'test'
+  }
+}
+
+muzzle {
+  pass {
+    group = 'com.aerospike'
+    module = 'aerospike-client'
+    versions = "[4,)"
+  }
+}
+
+dependencies {
+  compileOnly group: 'com.aerospike', name: 'aerospike-client', version: '4.0.0'
+
+  testCompile group: 'com.aerospike', name: 'aerospike-client', version: '4.0.0'
+  testCompile group: 'org.testcontainers', name: 'testcontainers', version: '1.15.0'
+  latestDepTestCompile group: 'com.aerospike', name: 'aerospike-client', version: '+'
+}

dd-java-agent/instrumentation/aerospike-4/src/main/java/datadog/trace/instrumentation/aerospike4/AerospikeClientDecorator.java

@@ -0,0 +1,106 @@
+package datadog.trace.instrumentation.aerospike4;
+
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+
+import com.aerospike.client.AerospikeClient;
+import com.aerospike.client.cluster.Cluster;
+import com.aerospike.client.cluster.Node;
+import com.aerospike.client.cluster.Partition;
+import datadog.trace.api.Config;
+import datadog.trace.api.DDTags;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes;
+import datadog.trace.bootstrap.instrumentation.api.Tags;
+import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
+import datadog.trace.bootstrap.instrumentation.decorator.DBTypeProcessingDatabaseClientDecorator;
+
+public class AerospikeClientDecorator extends DBTypeProcessingDatabaseClientDecorator<Node> {
+  public static final UTF8BytesString JAVA_AEROSPIKE = UTF8BytesString.create("java-aerospike");
+  public static final UTF8BytesString AEROSPIKE_COMMAND =
+      UTF8BytesString.create("aerospike.command");
+
+  public static final AerospikeClientDecorator DECORATE = new AerospikeClientDecorator();
+
+  @Override
+  protected String[] instrumentationNames() {
+    return new String[] {"aerospike"};
+  }
+
+  @Override
+  protected String service() {
+    return "aerospike";
+  }
+
+  @Override
+  protected CharSequence component() {
+    return JAVA_AEROSPIKE;
+  }
+
+  @Override
+  protected CharSequence spanType() {
+    return InternalSpanTypes.AEROSPIKE;
+  }
+
+  @Override
+  protected String dbType() {
+    return "aerospike";
+  }
+
+  @Override
+  protected String dbUser(final Node node) {
+    return null;
+  }
+
+  @Override
+  protected String dbInstance(final Node node) {
+    return null;
+  }
+
+  @Override
+  protected String dbHostname(final Node node) {
+    return null;
+  }
+
+  public AgentSpan onConnection(
+      final AgentSpan span, final Node node, final Cluster cluster, final Partition partition) {
+
+    onPeerConnection(span, node.getAddress());
+
+    if (cluster != null && cluster.getUser() != null) {
+      span.setTag(Tags.DB_USER, UTF8BytesString.create(cluster.getUser()));
+    }
+
+    if (partition != null) {
+      String instanceName = partition.toString();
+      final int namespaceEnd = instanceName.indexOf(':');
+      if (namespaceEnd > 0) {
+        instanceName = instanceName.substring(0, namespaceEnd);
+      }
+      span.setTag(Tags.DB_INSTANCE, instanceName);
+      if (Config.get().isDbClientSplitByInstance()) {
+        span.setTag(DDTags.SERVICE_NAME, instanceName);
+      }
+    }
+
+    return span;
+  }
+
+  public void withMethod(final AgentSpan span, final String methodName) {
+    span.setTag(DDTags.RESOURCE_NAME, spanNameForMethod(AerospikeClient.class, methodName));
+  }
+
+  public AgentSpan startAerospikeSpan(final String methodName) {
+    final AgentSpan span = startSpan(AEROSPIKE_COMMAND);
+    afterStart(span);
+    withMethod(span, methodName);
+    return span;
+  }
+
+  public void finishAerospikeSpan(final AgentSpan span, final Throwable error) {
+    if (error != null) {
+      onError(span, error);
+    }
+    beforeFinish(span);
+    span.finish();
+  }
+}

dd-java-agent/instrumentation/aerospike-4/src/main/java/datadog/trace/instrumentation/aerospike4/AerospikeClientInstrumentation.java

@@ -0,0 +1,95 @@
+package datadog.trace.instrumentation.aerospike4;
+
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.instrumentation.aerospike4.AerospikeClientDecorator.DECORATE;
+import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
+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 datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import java.util.HashMap;
+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;
+
+@AutoService(Instrumenter.class)
+public final class AerospikeClientInstrumentation extends Instrumenter.Default {
+  public AerospikeClientInstrumentation() {
+    super("aerospike");
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("com.aerospike.client.AerospikeClient");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".AerospikeClientDecorator", packageName + ".TracingListener",
+    };
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    final Map<ElementMatcher<? super MethodDescription>, String> transformers = new HashMap<>();
+    transformers.put(
+        isMethod()
+            .and(isPublic())
+            .and(takesArgument(0, nameStartsWith("com.aerospike.client.policy"))),
+        getClass().getName() + "$TraceSyncRequestAdvice");
+    transformers.put(
+        isMethod()
+            .and(isPublic())
+            .and(takesArgument(1, nameStartsWith("com.aerospike.client.listener"))),
+        getClass().getName() + "$TraceAsyncRequestAdvice");
+    return transformers;
+  }
+
+  public static final class TraceSyncRequestAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentScope beginRequest(@Advice.Origin("#m") final String methodName) {
+      AgentSpan clientSpan = DECORATE.startAerospikeSpan(methodName);
+      return activateSpan(clientSpan);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void exitRequest(
+        @Advice.Enter final AgentScope scope, @Advice.Thrown final Throwable error) {
+      DECORATE.finishAerospikeSpan(scope.span(), error);
+      scope.close();
+    }
+  }
+
+  public static final class TraceAsyncRequestAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentScope beginRequest(
+        @Advice.Origin("#m") final String methodName,
+        @Advice.Argument(value = 1, readOnly = false, typing = DYNAMIC) Object listener) {
+      AgentSpan clientSpan = DECORATE.startAerospikeSpan(methodName);
+      AgentScope scope = activateSpan(clientSpan);
+      // always want to wrap even when there's no listener so we get the true async time
+      listener = new TracingListener(clientSpan, scope.capture(), listener);
+      return scope;
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void exitRequest(
+        @Advice.Enter final AgentScope scope, @Advice.Thrown final Throwable error) {
+      if (error != null) {
+        DECORATE.finishAerospikeSpan(scope.span(), error);
+      } else {
+        // span will be finished in the traced listener
+      }
+      scope.close();
+    }
+  }
+}

dd-java-agent/instrumentation/aerospike-4/src/main/java/datadog/trace/instrumentation/aerospike4/CommandInstrumentation.java

@@ -0,0 +1,65 @@
+package datadog.trace.instrumentation.aerospike4;
+
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
+import static datadog.trace.instrumentation.aerospike4.AerospikeClientDecorator.DECORATE;
+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.returns;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.aerospike.client.cluster.Cluster;
+import com.aerospike.client.cluster.Node;
+import com.aerospike.client.cluster.Partition;
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.api.DDSpanTypes;
+import datadog.trace.bootstrap.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;
+
+@AutoService(Instrumenter.class)
+public final class CommandInstrumentation extends Instrumenter.Default {
+  public CommandInstrumentation() {
+    super("aerospike");
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("com.aerospike.client.command.Command");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".AerospikeClientDecorator",
+    };
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+        isMethod()
+            .and(named("getNode"))
+            .and(takesArgument(0, named("com.aerospike.client.cluster.Cluster")))
+            .and(returns(named("com.aerospike.client.cluster.Node"))),
+        getClass().getName() + "$GetNodeAdvice");
+  }
+
+  public static final class GetNodeAdvice {
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void getNode(
+        @Advice.Return final Node node,
+        @Advice.Argument(0) final Cluster cluster,
+        @Advice.Argument(value = 1, optional = true) final Partition partition) {
+      final AgentSpan span = activeSpan();
+      // capture the connection details in the active Aerospike span
+      if (span != null && DDSpanTypes.AEROSPIKE.equals(span.getSpanType())) {
+        DECORATE.onConnection(span, node, cluster, partition);
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/aerospike-4/src/main/java/datadog/trace/instrumentation/aerospike4/NioEventLoopInstrumentation.java

@@ -0,0 +1,56 @@
+package datadog.trace.instrumentation.aerospike4;
+
+import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE;
+import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.exclude;
+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 static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import java.util.Map;
+import java.util.concurrent.FutureTask;
+import java.util.concurrent.RunnableFuture;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(Instrumenter.class)
+public final class NioEventLoopInstrumentation extends Instrumenter.Default {
+  public NioEventLoopInstrumentation() {
+    super("aerospike", "java_concurrent");
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("com.aerospike.client.async.NioEventLoop");
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(
+        isMethod()
+            .and(named("execute"))
+            .and(takesArguments(1))
+            .and(takesArgument(0, Runnable.class)),
+        getClass().getName() + "$WrapAsFutureTaskAdvice");
+  }
+
+  public static final class WrapAsFutureTaskAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void enterExecute(@Advice.Argument(value = 0, readOnly = false) Runnable task) {
+      if (task == null || task instanceof RunnableFuture || exclude(RUNNABLE, task)) {
+        return;
+      }
+      // don't wrap Runnables belonging to NioEventLoop(s) as they want to propagate CloseException
+      // outside of the event loop on close() and wrapping them in FutureTask interferes with that
+      if (task.getClass().getName().startsWith("com.aerospike.client.async.NioEventLoop")) {
+        return;
+      }
+      task = new FutureTask<Void>(task, null);
+    }
+  }
+}