Add phantom instrumentation

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

@@ -0,0 +1,38 @@
+// Set properties before any plugins get loaded
+ext {
+  minJavaVersionForTests = JavaVersion.VERSION_1_8
+}
+
+apply from: "${rootDir}/gradle/java.gradle"
+apply from: "${rootDir}/gradle/test-with-scala.gradle"
+
+apply plugin: 'org.unbroken-dome.test-sets'
+
+
+dependencies {
+  main_java8CompileOnly group: 'com.outworkers', name: 'phantom-dsl_2.11', version: '2.59.0'
+  main_java8CompileOnly group: 'com.datastax.cassandra', name: 'cassandra-driver-core', version: '3.0.0'
+  main_java8CompileOnly group: 'io.monix', name: 'monix_2.11', version: '3.2.1'
+  compile project(':dd-trace-api')
+  compile project(':dd-java-agent:agent-tooling')
+
+//  compile deps.opentracing
+//  compile deps.autoservice
+//  annotationProcessor deps.autoservice
+  //compile project(':dd-java-agent:instrumentation:hibernate')
+  testCompile project(':dd-java-agent:instrumentation:trace-annotation')
+  testCompile project(':dd-java-agent:instrumentation:java-concurrent')
+  testCompile project(':dd-java-agent:instrumentation:datastax-cassandra-3')
+  testCompile project(':dd-java-agent:instrumentation:phantom')
+
+  testCompile group: 'com.outworkers', name: 'phantom-dsl_2.11', version: '2.59.0'
+  testCompile group: 'com.outworkers', name: 'phantom-monix_2.11', version: '2.59.0'
+  testCompile group: 'com.datastax.cassandra', name: 'cassandra-driver-core', version: '3.0.0'
+  testCompile group: 'com.typesafe.akka', name: 'akka-actor_2.11', version: '2.5.31'
+  testCompile group: 'io.monix', name: 'monix_2.11', version: '3.2.1'
+  testCompile deps.scala
+  testCompile group: 'org.scala-lang', name: 'scala-reflect', version: '2.11.12'
+  testCompile 'org.scala-lang.modules:scala-java8-compat_2.11:0.9.1'
+  testCompile group: 'org.cassandraunit', name: 'cassandra-unit', version: '3.11.2.0'
+}
+

dd-java-agent/instrumentation/phantom/src/main/java/datadog/trace/instrumentation/phantom/FutureCompletionListener.java

@@ -0,0 +1,49 @@
+package datadog.trace.instrumentation.phantom;
+
+import com.outworkers.phantom.ResultSet;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import lombok.extern.slf4j.Slf4j;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import scala.runtime.AbstractFunction1;
+import scala.util.Try;
+
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
+import static datadog.trace.instrumentation.phantom.PhantomDecorator.DECORATE;
+
+public class FutureCompletionListener extends AbstractFunction1<Try<ResultSet>, Object> {
+  private static final Logger log = LoggerFactory.getLogger(FutureCompletionListener.class);
+
+  private final AgentSpan agentSpan;
+  public FutureCompletionListener(AgentSpan agentSpan) {
+    this.agentSpan = agentSpan;
+  }
+
+  @Override
+  public Object apply(final Try<ResultSet> resultSetTry) {
+    log.debug("phantom future completed");
+    final AgentScope scope = activateSpan(agentSpan);
+    try {
+      final ResultSet resultSet = resultSetTry.get();   // TODO: Optimize the potential throw
+      log.debug("Call completed successfully");
+      if (resultSet != null) {
+        String keyspace = resultSet.getExecutionInfo().getStatement().getKeyspace();
+        if (keyspace != null) {
+          agentSpan.setTag("keyspace", keyspace);
+        }
+      }
+      DECORATE.beforeFinish(agentSpan);
+    } catch (final Throwable t) {
+      log.debug("Call completed with error");
+      DECORATE.onError(agentSpan, t);
+    } finally {
+      log.debug("doing finish and close");
+      agentSpan.finish();
+      scope.setAsyncPropagation(false);
+      scope.close();
+    }
+    return null;
+  }
+}

dd-java-agent/instrumentation/phantom/src/main/java/datadog/trace/instrumentation/phantom/GuavaAdapterAdvice.java

@@ -0,0 +1,51 @@
+package datadog.trace.instrumentation.phantom;
+
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.instrumentation.phantom.PhantomDecorator.DECORATE;
+
+import com.datastax.driver.core.Session;
+import com.datastax.driver.core.Statement;
+import com.outworkers.phantom.ResultSet;
+import com.outworkers.phantom.ops.QueryContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.implementation.bytecode.assign.Assigner;
+import scala.concurrent.ExecutionContext;
+import scala.concurrent.ExecutionContext$;
+import scala.concurrent.ExecutionContextExecutor;
+import scala.concurrent.Future;
+
+public class GuavaAdapterAdvice {
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  public static AgentScope enter(@Advice.Argument(value = 0) final Statement statement,
+                                 @Advice.Argument(value = 1) final Session session,
+                                 @Advice.Argument(value = 2) final ExecutionContextExecutor ctx ) {
+    System.out.println("Calling with context " + ctx.toString());
+    final AgentScope scope = startSpanWithScope(statement);
+    scope.setAsyncPropagation(true);
+    return scope;
+  }
+
+  public static AgentScope startSpanWithScope(final Statement statement) {
+    final AgentSpan span = startSpan("phantom.future");
+    DECORATE.afterStart(span);
+    DECORATE.onStatement(span, statement.toString());
+    //log.debug("activating span " + span);
+    return activateSpan(span);
+  }
+
+  @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+  public static void exit(
+    @Advice.Argument(value = 2) final ExecutionContextExecutor ctx,
+    @Advice.Return(readOnly = false, typing = Assigner.Typing.DYNAMIC) Future<ResultSet> resultSetFuture,
+    @Advice.Enter final AgentScope agentScope) {
+    //log.debug("onMethodExit " + agentScope.toString());
+    if (agentScope == null || resultSetFuture == null) {
+      return;
+    }
+
+    resultSetFuture.onComplete(new FutureCompletionListener(agentScope.span()), ctx);
+  }
+}

dd-java-agent/instrumentation/phantom/src/main/java/datadog/trace/instrumentation/phantom/PhantomAdvice.java

@@ -0,0 +1,48 @@
+package datadog.trace.instrumentation.phantom;
+
+import com.datastax.driver.core.Session;
+import com.outworkers.phantom.ResultSet;
+import com.outworkers.phantom.ops.QueryContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.implementation.bytecode.assign.Assigner;
+import scala.concurrent.ExecutionContextExecutor;
+import scala.concurrent.Future;
+
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.instrumentation.phantom.PhantomDecorator.DECORATE;
+
+public class PhantomAdvice {
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  public static AgentScope enter(@Advice.This final QueryContext.RootQueryOps rootQueryOps,
+                                 @Advice.Argument(value = 0) final Session session,
+                                 @Advice.Argument(value = 1) final ExecutionContextExecutor ctx ) {
+    System.out.println("Calling with context " + ctx.toString());
+    final AgentScope scope = startSpanWithScope(rootQueryOps);
+    scope.setAsyncPropagation(true);
+    return scope;
+  }
+
+  public static AgentScope startSpanWithScope(final QueryContext.RootQueryOps queryOps) {
+    final AgentSpan span = startSpan("phantom.future");
+    DECORATE.afterStart(span);
+    DECORATE.onStatement(span, queryOps.query().queryString());
+    //log.debug("activating span " + span);
+    return activateSpan(span);
+  }
+
+  @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+  public static void exit(
+    @Advice.Argument(value = 1) final ExecutionContextExecutor ctx,
+    @Advice.Return(readOnly = false, typing = Assigner.Typing.DYNAMIC) Future<ResultSet> resultSetFuture,
+    @Advice.Enter final AgentScope agentScope) {
+    //log.debug("onMethodExit " + agentScope.toString());
+    if (agentScope == null || resultSetFuture == null) {
+      return;
+    }
+
+    resultSetFuture.onComplete(new FutureCompletionListener(agentScope.span()), ctx);
+  }
+}

dd-java-agent/instrumentation/phantom/src/main/java/datadog/trace/instrumentation/phantom/PhantomDecorator.java

@@ -0,0 +1,49 @@
+package datadog.trace.instrumentation.phantom;
+
+import datadog.trace.api.DDSpanTypes;
+import datadog.trace.bootstrap.instrumentation.decorator.OrmClientDecorator;
+import datadog.trace.bootstrap.instrumentation.decorator.ServerDecorator;
+
+public class PhantomDecorator extends OrmClientDecorator {
+  public static final PhantomDecorator DECORATE = new PhantomDecorator();
+
+  @Override
+  protected String[] instrumentationNames() {
+    return new String[] {"phantom"};
+  }
+
+  @Override
+  protected String spanType() {
+    return "phantom";
+  }
+
+  @Override
+  protected String component() {
+    return "scala-phantom";
+  }
+
+  @Override
+  public String entityName(Object entity) {
+    return null;
+  }
+
+  @Override
+  protected String dbType() {
+    return DDSpanTypes.CASSANDRA;
+  }
+
+  @Override
+  protected String dbUser(Object o) {
+    return null;
+  }
+
+  @Override
+  protected String dbInstance(Object o) {
+    return null;
+  }
+
+  @Override
+  protected String service() {
+    return "phantom";
+  }
+}

dd-java-agent/instrumentation/phantom/src/main/java/datadog/trace/instrumentation/phantom/PhantomInstrumentation.java

@@ -0,0 +1,60 @@
+package datadog.trace.instrumentation.phantom;
+
+import com.datastax.driver.core.Session;
+import com.datastax.driver.core.Statement;
+import com.google.auto.service.AutoService;
+import com.outworkers.phantom.ResultSet;
+import com.outworkers.phantom.ops.QueryContext;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.context.TraceScope;
+import java.util.HashMap;
+import lombok.extern.slf4j.Slf4j;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.implementation.bind.annotation.SuperCall;
+import net.bytebuddy.matcher.ElementMatcher;
+import net.bytebuddy.matcher.ElementMatchers;
+import scala.runtime.AbstractFunction1;
+import scala.concurrent.ExecutionContextExecutor;
+import scala.concurrent.Future;
+import scala.util.Try;
+
+import java.util.Collections;
+import java.util.Map;
+
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.*;
+import static net.bytebuddy.matcher.ElementMatchers.*;
+
+@Slf4j
+@AutoService(Instrumenter.class)
+public class PhantomInstrumentation extends Instrumenter.Default {
+  public PhantomInstrumentation() {
+    super("phantom", "scala-phantom");
+  }
+
+
+  @Override
+  public ElementMatcher<? super TypeDescription> typeMatcher() {
+    return ElementMatchers.nameStartsWith("com.outworkers.phantom.builder.query.execution.PromiseInterface");
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    final Map<ElementMatcher<? super MethodDescription>, String> transformers = new HashMap<>();
+    //transformers.put(isMethod().and(named("future")), packageName + ".PhantomAdvice");
+    transformers.put(isMethod().and(named("fromGuava").and(ElementMatchers.takesArgument(0, named("com.datastax.driver.core.Statement")))),
+      packageName + ".GuavaAdapterAdvice");
+    return transformers;
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[]{
+      packageName + ".FutureCompletionListener",
+      packageName + ".PhantomDecorator"
+    };
+  }
+}

dd-java-agent/instrumentation/phantom/src/main/java/datadog/trace/instrumentation/phantom/TaskCompletionListener.java

@@ -0,0 +1,26 @@
+package datadog.trace.instrumentation.phantom;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import monix.eval.Task;
+import scala.Function1;
+import scala.Option;
+import scala.runtime.AbstractFunction1;
+import scala.runtime.BoxedUnit;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+public class TaskCompletionListener extends AbstractFunction1<Option<Throwable>, Task<BoxedUnit>> {
+  private final AgentSpan agentSpan;
+
+  public TaskCompletionListener(final AgentSpan agentSpan) {
+    this.agentSpan = agentSpan;
+  }
+
+  @Override
+  public Task<BoxedUnit> apply(final Option<Throwable> throwableOption) {
+    final AgentScope scope = activateSpan(agentSpan);
+    agentSpan.finish();
+    scope.setAsyncPropagation(false);
+    scope.close();
+    return null;
+  }
+}