Merge pull request #777 from DataDog/tyler/hibernate

Fix NPE in hibernate instrumentation and fix InstrumentationContext+Muzzle

Author: tylerbenson

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java

@@ -63,6 +63,7 @@ public final AgentBuilder instrument(final AgentBuilder parentAgentBuilder) {
         return parentAgentBuilder;
       }
 
+      final MuzzleMatcher muzzleMatcher = new MuzzleMatcher();
       AgentBuilder.Identified.Extendable agentBuilder =
           parentAgentBuilder
               .type(
@@ -73,13 +74,13 @@ public final AgentBuilder instrument(final AgentBuilder parentAgentBuilder) {
                       classLoaderMatcher(),
                       "Instrumentation class loader matcher unexpected exception: "
                           + getClass().getName()))
-              .and(new MuzzleMatcher())
+              .and(muzzleMatcher)
               .and(new PostMatchHook())
               .transform(DDTransformers.defaultTransformers());
       agentBuilder = injectHelperClasses(agentBuilder);
       agentBuilder = contextProvider.instrumentationTransformer(agentBuilder);
       agentBuilder = applyInstrumentationTransformers(agentBuilder);
-      agentBuilder = contextProvider.additionalInstrumentation(agentBuilder);
+      agentBuilder = contextProvider.additionalInstrumentation(agentBuilder, muzzleMatcher);
       return agentBuilder;
     }
 

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/FieldBackedProvider.java

@@ -70,14 +70,15 @@
 @Slf4j
 public class FieldBackedProvider implements InstrumentationContextProvider {
 
-  /*
-  Note: the value here has to be inside on of the prefixes in
-  datadog.trace.agent.tooling.Utils#BOOTSTRAP_PACKAGE_PREFIXES. This ensures that 'isolating' (or 'module')
-  classloaders like jboss and osgi see injected classes. This works because we instrument those classloaders
-  to load everything inside bootstrap packages.
+  /**
+   * Note: the value here has to be inside on of the prefixes in
+   * datadog.trace.agent.tooling.Utils#BOOTSTRAP_PACKAGE_PREFIXES. This ensures that 'isolating' (or
+   * 'module') classloaders like jboss and osgi see injected classes. This works because we
+   * instrument those classloaders to load everything inside bootstrap packages.
    */
   private static final String DYNAMIC_CLASSES_PACKAGE =
       "datadog.trace.bootstrap.instrumentation.context.";
+
   private static final String INJECTED_FIELDS_MARKER_CLASS_NAME =
       Utils.getInternalName(FieldBackedContextStoreAppliedMarker.class.getName());
 
@@ -118,24 +119,24 @@ public FieldBackedProvider(final Instrumenter.Default instrumenter) {
   public AgentBuilder.Identified.Extendable instrumentationTransformer(
       AgentBuilder.Identified.Extendable builder) {
     if (instrumenter.contextStore().size() > 0) {
-      /*
-      Install transformer that rewrites accesses to context store with specialized bytecode that invokes appropriate
-      storage implementation.
+      /**
+       * Install transformer that rewrites accesses to context store with specialized bytecode that
+       * invokes appropriate storage implementation.
        */
       builder =
           builder.transform(getTransformerForASMVisitor(getContextStoreReadsRewritingVisitor()));
 
-      /*
-      We inject into bootstrap classloader because field accessor interfaces are needed by
-      context store implementations. Unfortunately this forces us to remove stored type checking
-      because actual classes may not be available at this point.
+      /**
+       * We inject into bootstrap classloader because field accessor interfaces are needed by
+       * context store implementations. Unfortunately this forces us to remove stored type checking
+       * because actual classes may not be available at this point.
        */
       builder = builder.transform(bootstrapHelperInjector(fieldAccessorInterfaces.values()));
 
-      /*
-       * We inject context store implementation into bootstrap classloader because same implementation
-       * may be used by different instrumentations and it has to use same static map in case of
-       * fallback to map-backed storage.
+      /**
+       * We inject context store implementation into bootstrap classloader because same
+       * implementation may be used by different instrumentations and it has to use same static map
+       * in case of fallback to map-backed storage.
        */
       builder = builder.transform(bootstrapHelperInjector(contextStoreImplementations.values()));
     }
@@ -329,20 +330,27 @@ public void visitLdcInsn(final Object value) {
 
   @Override
   public AgentBuilder.Identified.Extendable additionalInstrumentation(
-      AgentBuilder.Identified.Extendable builder) {
+      AgentBuilder.Identified.Extendable builder, final AgentBuilder.RawMatcher muzzleMatcher) {
 
     if (fieldInjectionEnabled) {
       for (final Map.Entry<String, String> entry : instrumenter.contextStore().entrySet()) {
-        /*
-        For each context store defined in a current instrumentation we create an agent builder
-        that injects necessary fields.
+        /**
+         * For each context store defined in a current instrumentation we create an agent builder
+         * that injects necessary fields.
          */
         builder =
             builder
                 .type(
                     safeHasSuperType(named(entry.getKey())).and(not(isInterface())),
                     instrumenter.classLoaderMatcher())
                 .and(safeToInjectFieldsMatcher())
+                /**
+                 * By adding the muzzleMatcher here, we are adding risk that the rules for injecting
+                 * the classes into the classloader and the rules for adding the field to the class
+                 * might be different. However the consequences are much greater if the class is not
+                 * injected but the field is added, since that results in a NoClassDef error.
+                 */
+                .and(muzzleMatcher)
                 .transform(
                     getTransformerForASMVisitor(
                         getFieldInjectionVisitor(entry.getKey(), entry.getValue())));
@@ -360,13 +368,14 @@ public boolean matches(
           final JavaModule module,
           final Class<?> classBeingRedefined,
           final ProtectionDomain protectionDomain) {
-        /*
-        The idea here is that we can add fields if class is just being loaded (classBeingRedefined == null)
-        and we have to add same fields again if class we added fields before is being transformed again.
-        Note: here we assume that Class#getInterfaces() returns list of interfaces defined immediately on a
-        given class, not inherited from its parents. It looks like current JVM implementation does exactly
-        this but javadoc is not explicit about that.
-        */
+        /**
+         * The idea here is that we can add fields if class is just being loaded
+         * (classBeingRedefined == null) and we have to add same fields again if class we added
+         * fields before is being transformed again. Note: here we assume that Class#getInterfaces()
+         * returns list of interfaces defined immediately on a given class, not inherited from its
+         * parents. It looks like current JVM implementation does exactly this but javadoc is not
+         * explicit about that.
+         */
         return classBeingRedefined == null
             || Arrays.asList(classBeingRedefined.getInterfaces())
                 .contains(FieldBackedContextStoreAppliedMarker.class);

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/InstrumentationContextProvider.java

@@ -13,5 +13,5 @@ AgentBuilder.Identified.Extendable instrumentationTransformer(
 
   /** Hook to define additional instrumentation. Run at instrumentation advice is hooked up. */
   AgentBuilder.Identified.Extendable additionalInstrumentation(
-      AgentBuilder.Identified.Extendable builder);
+      AgentBuilder.Identified.Extendable builder, final AgentBuilder.RawMatcher muzzleMatcher);
 }

dd-java-agent/instrumentation/hibernate/core-3.3/core-3.3.gradle

@@ -44,6 +44,9 @@ dependencies {
 
   testCompile project(':dd-java-agent:testing')
   testCompile project(':dd-java-agent:instrumentation:jdbc')
+  // Added to ensure cross compatibility:
+  testCompile project(':dd-java-agent:instrumentation:hibernate:core-4.0')
+  testCompile project(':dd-java-agent:instrumentation:hibernate:core-4.3')
 
   testCompile group: 'org.hibernate', name: 'hibernate-core', version: '3.3.0.GA'
   testCompile group: 'org.hibernate', name: 'hibernate-annotations', version: '+'

dd-java-agent/instrumentation/hibernate/core-3.3/src/main/java/datadog/trace/instrumentation/hibernate/core/v3_3/QueryInstrumentation.java

@@ -61,7 +61,9 @@ public static SessionState startMethod(
       final SessionState state =
           SessionMethodUtils.startScopeFrom(
               contextStore, query, "hibernate.query." + name, null, true);
-      DECORATOR.onStatement(state.getMethodScope().span(), query.getQueryString());
+      if (state != null) {
+        DECORATOR.onStatement(state.getMethodScope().span(), query.getQueryString());
+      }
       return state;
     }
 

dd-java-agent/instrumentation/hibernate/core-4.0/core-4.0.gradle

@@ -37,7 +37,9 @@ dependencies {
 
   testCompile project(':dd-java-agent:testing')
   testCompile project(':dd-java-agent:instrumentation:jdbc')
-
+  // Added to ensure cross compatibility:
+  testCompile project(':dd-java-agent:instrumentation:hibernate:core-3.3')
+  testCompile project(':dd-java-agent:instrumentation:hibernate:core-4.3')
 
   testCompile group: 'org.hibernate', name: 'hibernate-core', version: '4.0.0.Final'
   testCompile group: 'com.h2database', name: 'h2', version: '1.4.197'

dd-java-agent/instrumentation/hibernate/core-4.0/src/main/java/datadog/trace/instrumentation/hibernate/core/v4_0/QueryInstrumentation.java

@@ -61,7 +61,9 @@ public static SessionState startMethod(
       final SessionState state =
           SessionMethodUtils.startScopeFrom(
               contextStore, query, "hibernate.query." + name, null, true);
-      DECORATOR.onStatement(state.getMethodScope().span(), query.getQueryString());
+      if (state != null) {
+        DECORATOR.onStatement(state.getMethodScope().span(), query.getQueryString());
+      }
       return state;
     }
 

dd-java-agent/instrumentation/hibernate/core-4.3/core-4.3.gradle

@@ -37,11 +37,17 @@ dependencies {
 
   testCompile project(':dd-java-agent:testing')
   testCompile project(':dd-java-agent:instrumentation:jdbc')
+  // Added to ensure cross compatibility:
+  testCompile project(':dd-java-agent:instrumentation:hibernate:core-3.3')
   testCompile project(':dd-java-agent:instrumentation:hibernate:core-4.0')
 
   testCompile group: 'org.hibernate', name: 'hibernate-core', version: '4.3.0.Final'
+  testCompile group: 'org.hibernate', name: 'hibernate-entitymanager', version: '4.3.0.Final'
   testCompile group: 'org.hsqldb', name: 'hsqldb', version: '2.0.0'
+  testCompile group: 'org.springframework.data', name: 'spring-data-jpa', version: '1.5.1.RELEASE'
 
   latestDepTestCompile group: 'org.hibernate', name: 'hibernate-core', version: '+'
+  latestDepTestCompile group: 'org.hibernate', name: 'hibernate-entitymanager', version: '+'
   latestDepTestCompile group: 'org.hsqldb', name: 'hsqldb', version: '2.0.0'
+  latestDepTestCompile group: 'org.springframework.data', name: 'spring-data-jpa', version: '+'
 }

dd-java-agent/instrumentation/hibernate/core-4.3/src/test/groovy/SpringJpaTest.groovy

@@ -0,0 +1,147 @@
+import datadog.trace.agent.test.AgentTestRunner
+import org.springframework.context.annotation.AnnotationConfigApplicationContext
+import spock.lang.Shared
+import spring.Customer
+import spring.CustomerRepository
+import spring.PersistenceConfig
+
+
+/**
+ * Unfortunately this test verifies that our hibernate instrumentation doesn't currently work with Spring Data Repositories.
+ */
+class SpringJpaTest extends AgentTestRunner {
+
+  @Shared
+  def context = new AnnotationConfigApplicationContext(PersistenceConfig)
+
+  @Shared
+  def repo = context.getBean(CustomerRepository)
+
+  def "test CRUD"() {
+    setup:
+    def customer = new Customer("Bob", "Anonymous")
+
+    expect:
+    customer.id == null
+    !repo.findAll().iterator().hasNext()
+
+    assertTraces(1) {
+      trace(0, 2) {
+        span(0) {
+          serviceName "hsqldb"
+          spanType "sql"
+        }
+        span(1) {
+          serviceName "hsqldb"
+          spanType "sql"
+          childOf(span(0))
+        }
+      }
+    }
+    TEST_WRITER.clear()
+
+    when:
+    repo.save(customer)
+    def savedId = customer.id
+
+    then:
+    customer.id != null
+    // Behavior changed in new version:
+    def extraTrace = TEST_WRITER.size() == 2
+    assertTraces(extraTrace ? 2 : 1) {
+      trace(0, 2) {
+        span(0) {
+          serviceName "hsqldb"
+          spanType "sql"
+        }
+        span(1) {
+          serviceName "hsqldb"
+          spanType "sql"
+          childOf(span(0))
+        }
+      }
+      if (extraTrace) {
+        trace(1, 1) {
+          span(0) {
+            serviceName "hsqldb"
+            spanType "sql"
+          }
+        }
+      }
+    }
+    TEST_WRITER.clear()
+
+    when:
+    customer.firstName = "Bill"
+    repo.save(customer)
+
+    then:
+    customer.id == savedId
+    assertTraces(2) {
+      trace(0, 2) {
+        span(0) {
+          serviceName "hsqldb"
+          spanType "sql"
+        }
+        span(1) {
+          serviceName "hsqldb"
+          spanType "sql"
+          childOf(span(0))
+        }
+      }
+      trace(1, 1) {
+        span(0) {
+          serviceName "hsqldb"
+          spanType "sql"
+        }
+      }
+    }
+    TEST_WRITER.clear()
+
+    when:
+    customer = repo.findByLastName("Anonymous")[0]
+
+    then:
+    customer.id == savedId
+    customer.firstName == "Bill"
+    assertTraces(1) {
+      trace(0, 2) {
+        span(0) {
+          serviceName "hsqldb"
+          spanType "sql"
+        }
+        span(1) {
+          serviceName "hsqldb"
+          spanType "sql"
+          childOf(span(0))
+        }
+      }
+    }
+    TEST_WRITER.clear()
+
+    when:
+    repo.delete(customer)
+
+    then:
+    assertTraces(2) {
+      trace(0, 2) {
+        span(0) {
+          serviceName "hsqldb"
+          spanType "sql"
+        }
+        span(1) {
+          serviceName "hsqldb"
+          spanType "sql"
+          childOf(span(0))
+        }
+      }
+      trace(1, 1) {
+        span(0) {
+          serviceName "hsqldb"
+          spanType "sql"
+        }
+      }
+    }
+    TEST_WRITER.clear()
+  }
+}

dd-java-agent/instrumentation/hibernate/core-4.3/src/test/java/spring/Customer.java

@@ -0,0 +1,31 @@
+package spring;
+
+import javax.persistence.Entity;
+import javax.persistence.GeneratedValue;
+import javax.persistence.GenerationType;
+import javax.persistence.Id;
+import lombok.Data;
+
+@Data
+@Entity
+public class Customer {
+
+  @Id
+  @GeneratedValue(strategy = GenerationType.AUTO)
+  private Long id;
+
+  private String firstName;
+  private String lastName;
+
+  protected Customer() {}
+
+  public Customer(final String firstName, final String lastName) {
+    this.firstName = firstName;
+    this.lastName = lastName;
+  }
+
+  @Override
+  public String toString() {
+    return String.format("Customer[id=%d, firstName='%s', lastName='%s']", id, firstName, lastName);
+  }
+}

dd-java-agent/instrumentation/hibernate/core-4.3/src/test/java/spring/CustomerRepository.java

@@ -0,0 +1,9 @@
+package spring;
+
+import java.util.List;
+import org.springframework.data.repository.CrudRepository;
+
+public interface CustomerRepository extends CrudRepository<Customer, Long> {
+
+  List<Customer> findByLastName(String lastName);
+}