cassandra bytebuddy

buildSrc/src/main/groovy/VersionScanPlugin.groovy

@@ -39,9 +39,9 @@ class VersionScanPlugin implements Plugin<Project> {
       description = "Queries for all versions of configured modules and finds key classes"
     }
 
-    def hasRelevantTask = project.gradle.startParameter.taskNames.contains('scanVersions')
-    hasRelevantTask |= project.gradle.startParameter.taskNames.contains('scanVersionsReport')
-    hasRelevantTask |= project.gradle.startParameter.taskNames.contains('verifyVersionScan')
+    def hasRelevantTask = project.gradle.startParameter.taskNames.any { it.contains('scanVersions') }
+    hasRelevantTask |= project.gradle.startParameter.taskNames.any { it.contains('scanVersionsReport') }
+    hasRelevantTask |= project.gradle.startParameter.taskNames.any { it.contains('verifyVersionScan') }
 
     if (!hasRelevantTask) {
       return
@@ -80,7 +80,7 @@ class VersionScanPlugin implements Plugin<Project> {
         }
       }
     }
-    if (project.gradle.startParameter.taskNames.contains('scanVersions')) {
+    if (project.gradle.startParameter.taskNames.any { it.contains('scanVersions') }) {
       scanVersions.finalizedBy(scanVersionsReport)
     }
 
@@ -176,7 +176,7 @@ class VersionScanPlugin implements Plugin<Project> {
           }
         }
 
-        if (project.gradle.startParameter.taskNames.contains('scanVersions')) {
+        if (project.gradle.startParameter.taskNames.any { it.contains('scanVersions') }) {
           scanVersions.finalizedBy(verifyVersionScan)
         }
       }

dd-java-agent-ittests/src/test/java/com/datadoghq/agent/integration/CassandraIntegrationTest.java

@@ -2,45 +2,107 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 
+import com.datadoghq.trace.DDBaseSpan;
+import com.datadoghq.trace.DDTracer;
+import com.datadoghq.trace.writer.ListWriter;
 import com.datastax.driver.core.Cluster;
 import com.datastax.driver.core.Session;
-import java.io.IOException;
-import java.util.concurrent.ExecutionException;
-import org.apache.cassandra.exceptions.ConfigurationException;
-import org.apache.thrift.transport.TTransportException;
+import io.opentracing.Tracer;
+import io.opentracing.tag.Tags;
 import org.cassandraunit.utils.EmbeddedCassandraServerHelper;
-import org.junit.After;
-import org.junit.Before;
+import org.junit.AfterClass;
+import org.junit.Assert;
+import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
 
-/** Created by gpolaert on 6/2/17. */
 @Category(ExpensiveTest.class)
 public class CassandraIntegrationTest {
+  private static final ListWriter writer = new ListWriter();
+  private static final Tracer tracer = new DDTracer(writer);
 
-  @Before
-  public void start()
-      throws InterruptedException, TTransportException, ConfigurationException, IOException {
+  @BeforeClass
+  public static void start() throws Exception {
     EmbeddedCassandraServerHelper.startEmbeddedCassandra(40000L);
+    TestUtils.registerOrReplaceGlobalTracer(tracer);
   }
 
-  @After
-  public void stop() {
+  @AfterClass
+  public static void stop() {
     EmbeddedCassandraServerHelper.cleanEmbeddedCassandra();
   }
 
   @Test
-  public void testNewSessionSync() throws ClassNotFoundException {
+  public void testSync() throws ClassNotFoundException {
     final Cluster cluster = EmbeddedCassandraServerHelper.getCluster();
     final Session session = cluster.newSession();
     assertThat(session.getClass().getName()).endsWith("contrib.cassandra.TracingSession");
+    final int origSize = writer.getList().size();
+
+    session.execute("DROP KEYSPACE IF EXISTS sync_test");
+    session.execute(
+        "CREATE KEYSPACE sync_test WITH REPLICATION = {'class':'SimpleStrategy', 'replication_factor':3}");
+    session.execute("CREATE TABLE sync_test.users ( id UUID PRIMARY KEY, name text )");
+    session.execute("INSERT INTO sync_test.users (id, name) values (uuid(), 'alice')");
+    session.execute("SELECT * FROM sync_test.users where name = 'alice' ALLOW FILTERING");
+
+    assertThat(writer.getList().size()).isEqualTo(origSize + 5);
+    DDBaseSpan<?> selectTrace = writer.get(writer.size() - 1).get(0);
+
+    assertThat(selectTrace.getServiceName()).isEqualTo(DDTracer.UNASSIGNED_DEFAULT_SERVICE_NAME);
+    assertThat(selectTrace.getOperationName()).isEqualTo("execute");
+    assertThat(selectTrace.getResourceName()).isEqualTo("execute");
+
+    assertThat(selectTrace.getTags().get(Tags.COMPONENT.getKey())).isEqualTo("java-cassandra");
+    assertThat(selectTrace.getTags().get(Tags.DB_STATEMENT.getKey()))
+        .isEqualTo("SELECT * FROM sync_test.users where name = 'alice' ALLOW FILTERING");
+    assertThat(selectTrace.getTags().get(Tags.DB_TYPE.getKey())).isEqualTo("cassandra");
+    assertThat(selectTrace.getTags().get(Tags.PEER_HOSTNAME.getKey())).isEqualTo("localhost");
+    // More info about IPv4 tag: https://trello.com/c/2el2IwkF/174-mongodb-ot-contrib-provides-a-wrong-peeripv4
+    assertThat(selectTrace.getTags().get(Tags.PEER_HOST_IPV4.getKey())).isEqualTo(2130706433);
+    assertThat(selectTrace.getTags().get(Tags.PEER_PORT.getKey())).isEqualTo(9142);
+    assertThat(selectTrace.getTags().get(Tags.SPAN_KIND.getKey())).isEqualTo("client");
   }
 
   @Test
-  public void testNewSessionAsync()
-      throws ClassNotFoundException, ExecutionException, InterruptedException {
+  public void testAsync() throws Exception {
     final Cluster cluster = EmbeddedCassandraServerHelper.getCluster();
     final Session session = cluster.connectAsync().get();
     assertThat(session.getClass().getName()).endsWith("contrib.cassandra.TracingSession");
+    final int origSize = writer.getList().size();
+
+    session.executeAsync("DROP KEYSPACE IF EXISTS async_test").get();
+    session
+        .executeAsync(
+            "CREATE KEYSPACE async_test WITH REPLICATION = {'class':'SimpleStrategy', 'replication_factor':3}")
+        .get();
+    session.executeAsync("CREATE TABLE async_test.users ( id UUID PRIMARY KEY, name text )").get();
+    session.executeAsync("INSERT INTO async_test.users (id, name) values (uuid(), 'alice')").get();
+    session
+        .executeAsync("SELECT * FROM async_test.users where name = 'alice' ALLOW FILTERING")
+        .get();
+
+    // traces are finished on another thread, so we have some waiting logic
+    for (int timeout = 0; writer.getList().size() < origSize + 5; timeout++) {
+      if (timeout >= 10000) {
+        Assert.fail("Cassandra async test timeout.");
+      }
+      Thread.sleep(1);
+    }
+    DDBaseSpan<?> selectTrace = writer.get(writer.size() - 1).get(0);
+
+    assertThat(selectTrace.getServiceName()).isEqualTo(DDTracer.UNASSIGNED_DEFAULT_SERVICE_NAME);
+    assertThat(selectTrace.getOperationName()).isEqualTo("execute");
+    assertThat(selectTrace.getResourceName()).isEqualTo("execute");
+
+    assertThat(selectTrace.getTags().get(Tags.COMPONENT.getKey())).isEqualTo("java-cassandra");
+    assertThat(selectTrace.getTags().get(Tags.DB_STATEMENT.getKey()))
+        .isEqualTo("SELECT * FROM async_test.users where name = 'alice' ALLOW FILTERING");
+    assertThat(selectTrace.getTags().get(Tags.DB_TYPE.getKey())).isEqualTo("cassandra");
+    assertThat(selectTrace.getTags().get(Tags.PEER_HOSTNAME.getKey())).isEqualTo("localhost");
+    // More info about IPv4 tag: https://trello.com/c/2el2IwkF/174-mongodb-ot-contrib-provides-a-wrong-peeripv4
+    assertThat(selectTrace.getTags().get(Tags.PEER_HOST_IPV4.getKey())).isEqualTo(2130706433);
+    assertThat(selectTrace.getTags().get(Tags.PEER_PORT.getKey())).isEqualTo(9142);
+    assertThat(selectTrace.getTags().get(Tags.SPAN_KIND.getKey())).isEqualTo("client");
   }
 }

dd-java-agent/dd-java-agent.gradle

@@ -24,6 +24,9 @@ dependencies {
   compile(project(':dd-java-agent:integrations:aws-sdk')) {
     transitive = false
   }
+  compile(project(':dd-java-agent:integrations:datastax-cassandra-3.2')) {
+    transitive = false
+  }
   compile(project(':dd-java-agent:integrations:jms-1')) {
     transitive = false
   }

dd-java-agent/integrations/datastax-cassandra-3.2/datastax-cassandra-3.2.gradle

@@ -0,0 +1,44 @@
+// TODO: VersionScan plugin does not report which version failed, which is making it hard to get meaningful results out of this block.
+// Once versionScan can report on which version failed, this can be enabled.
+// The desire is to apply the instrumentation to cassandra-datastax 2.3 and beyond.
+
+apply plugin: 'version-scan'
+
+versionScan {
+  group = "com.datastax.cassandra"
+  module = "cassandra-driver-core"
+  versions = "[3.2.0,)"
+//  verifyPresent = [
+//    // class we're advising
+//    'com.datastax.driver.core.Cluster$Manager'          : null,
+//    // used by TracingSession
+//    'com.datastax.driver.core.BoundStatement'           : null,
+//    'com.datastax.driver.core.BoundStatement'           : null,
+//    'com.datastax.driver.core.CloseFuture'              : null,
+//    'com.datastax.driver.core.Cluster'                  : null,
+//    'com.datastax.driver.core.Host'                     : null,
+//    'com.datastax.driver.core.PreparedStatement'        : null,
+//    'com.datastax.driver.core.RegularStatement'         : null,
+//    'com.datastax.driver.core.ResultSet'                : null,
+//    'com.datastax.driver.core.ResultSetFuture'          : null,
+//    'com.datastax.driver.core.Session'                  : null,
+//    'com.datastax.driver.core.Statement'                : null,
+//    'com.google.common.base.Function'                   : null,
+//    'com.google.common.util.concurrent.Futures'         : null,
+//    'com.google.common.util.concurrent.ListenableFuture': null
+//  ]
+}
+
+apply from: "${rootDir}/gradle/java.gradle"
+
+dependencies {
+  compile group: 'com.datastax.cassandra', name: 'cassandra-driver-core', version: '3.2.0'
+
+  compile project(':dd-trace')
+  // include helpers to pick up opentracing-cassandra-driver helper
+  compile project(':dd-java-agent:integrations:helpers')
+  compile project(':dd-java-agent:tooling')
+
+  compile deps.bytebuddy
+  compile deps.opentracing
+}

dd-java-agent/integrations/datastax-cassandra-3.2/src/main/java/dd/inst/datastax/cassandra/CassandraClientInstrumentation.java

@@ -0,0 +1,70 @@
+package dd.inst.datastax.cassandra;
+
+import static dd.trace.ClassLoaderMatcher.classLoaderHasClasses;
+import static dd.trace.ExceptionHandlers.defaultExceptionHandler;
+import static net.bytebuddy.matcher.ElementMatchers.*;
+
+import com.datastax.driver.core.Session;
+import com.google.auto.service.AutoService;
+import dd.trace.Instrumenter;
+import io.opentracing.Tracer;
+import io.opentracing.util.GlobalTracer;
+import java.lang.reflect.Constructor;
+import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.asm.Advice;
+
+@AutoService(Instrumenter.class)
+public class CassandraClientInstrumentation implements Instrumenter {
+  @Override
+  public AgentBuilder instrument(AgentBuilder agentBuilder) {
+    return agentBuilder
+        .type(
+            named("com.datastax.driver.core.Cluster$Manager"),
+            classLoaderHasClasses(
+                "com.datastax.driver.core.BoundStatement",
+                "com.datastax.driver.core.BoundStatement",
+                "com.datastax.driver.core.CloseFuture",
+                "com.datastax.driver.core.Cluster",
+                "com.datastax.driver.core.Host",
+                "com.datastax.driver.core.PreparedStatement",
+                "com.datastax.driver.core.RegularStatement",
+                "com.datastax.driver.core.ResultSet",
+                "com.datastax.driver.core.ResultSetFuture",
+                "com.datastax.driver.core.Session",
+                "com.datastax.driver.core.Statement",
+                "com.google.common.base.Function",
+                "com.google.common.util.concurrent.Futures",
+                "com.google.common.util.concurrent.ListenableFuture"))
+        .transform(
+            new AgentBuilder.Transformer.ForAdvice()
+                .advice(
+                    isMethod().and(isPrivate()).and(named("newSession")).and(takesArguments(0)),
+                    CassandraClientAdvice.class.getName())
+                .withExceptionHandler(defaultExceptionHandler()))
+        .asDecorator();
+  }
+
+  public static class CassandraClientAdvice {
+    /**
+     * Strategy: each time we build a connection to a Cassandra cluster, the
+     * com.datastax.driver.core.Cluster$Manager.newSession() method is called. The opentracing
+     * contribution is a simple wrapper, so we just have to wrap the new session.
+     *
+     * @param session The fresh session to patch
+     * @return A new tracing session
+     * @throws Exception
+     */
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void injectTracingSession(@Advice.Return(readOnly = false) Session session)
+        throws Exception {
+      if (session.getClass().getName().endsWith("contrib.cassandra.TracingSession")) {
+        return;
+      }
+
+      Class<?> clazz = Class.forName("io.opentracing.contrib.cassandra.TracingSession");
+      Constructor<?> constructor = clazz.getDeclaredConstructor(Session.class, Tracer.class);
+      constructor.setAccessible(true);
+      session = (Session) constructor.newInstance(session, GlobalTracer.get());
+    }
+  }
+}

dd-java-agent/src/main/resources/dd-trace-supported-framework.yaml

@@ -25,14 +25,6 @@ opentracing-aws-sdk:
       com.amazonaws.http.apache.utils.ApacheUtils:
       com.amazonaws.http.request.HttpRequestFactory:
 
-opentracing-cassandra-driver:
-  - artifact: cassandra-driver-core
-    supported_version: 3\.2.*
-    identifying_present_classes:
-      com.datastax.driver.core.utils.MoreObjects:
-      com.datastax.driver.core.RemoteEndpointAwareNettySSLOptions:
-      com.datastax.driver.core.GuavaCompatibility:
-
 opentracing-jms-2_producer:
   - artifact: javax.jms-api
     supported_version: 2\..*

dd-java-agent/src/main/resources/initializer-rules.btm

@@ -13,18 +13,6 @@
 #ENDRULE
 
 
-# Instrument Cassandra client
-# ===========================
-RULE Cluster$Manager-init
-CLASS com.datastax.driver.core.Cluster$Manager
-METHOD <init>
-AT EXIT
-IF TRUE
-DO
-  com.datadoghq.agent.InstrumentationRulesManager.registerClassLoad($0);
-ENDRULE
-
-
 # Instrument OkHttp
 # ===========================
 RULE OkHttpClient$Builder-init

dd-java-agent/src/main/resources/integration-rules.btm

@@ -21,19 +21,6 @@ DO
 ENDRULE
 
 
-# Instrument Cassandra client
-# ===========================
-RULE opentracing-cassandra-driver
-CLASS com.datastax.driver.core.Cluster$Manager
-METHOD newSession()
-HELPER com.datadoghq.agent.integration.CassandraHelper
-AT EXIT
-IF TRUE
-DO
-  $! = patch($!);
-ENDRULE
-
-
 # Instrument OkHttp
 # ===========================
 RULE opentracing-okhttp3

settings.gradle

@@ -15,7 +15,7 @@ include ':dd-trace-annotations'
 include ':dd-java-agent:integrations:helpers'
 include ':dd-java-agent:integrations:apache-httpclient'
 include ':dd-java-agent:integrations:aws-sdk'
-include ':dd-java-agent:integrations:cassandra'
+include ':dd-java-agent:integrations:datastax-cassandra-3.2'
 include ':dd-java-agent:integrations:jms-1'
 include ':dd-java-agent:integrations:jms-2'
 include ':dd-java-agent:integrations:mongo-3.1'