Convert datastax cassandra instrumentation to bytebuddy

Author: Andrew Kent

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/cassandra/cassandra.gradle

@@ -0,0 +1,43 @@
+// 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/datastax-cassandra-3.2.gradle

@@ -0,0 +1,53 @@
+package dd.inst.datastax.cassandra;
+
+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"))
+        .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/integrations/datastax-cassandra-3.2/src/main/java/dd/inst/datastax/cassandra/CassandraClientInstrumentation.java

@@ -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/integrations/helpers/src/main/java/com/datadoghq/agent/integration/CassandraHelper.java

@@ -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/dd-trace-supported-framework.yaml

@@ -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

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

@@ -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'