Merge pull request #777 from DataDog/tyler/hibernate

Fix NPE in hibernate instrumentation and fix InstrumentationContext+Muzzle

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

@@ -63,6 +63,7 @@ public interface Instrumenter {
         return parentAgentBuilder;
       }
 
+      final MuzzleMatcher muzzleMatcher = new MuzzleMatcher();
       AgentBuilder.Identified.Extendable agentBuilder =
           parentAgentBuilder
               .type(
@@ -73,13 +74,13 @@ public interface Instrumenter {
                       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 @@ import net.bytebuddy.utility.JavaModule;
 @Slf4j
 public class FieldBackedProvider implements InstrumentationContextProvider {
 
-  /*
+  /**
-  Note: the value here has to be inside on of the prefixes in
+   * 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')
+   * datadog.trace.agent.tooling.Utils#BOOTSTRAP_PACKAGE_PREFIXES. This ensures that 'isolating' (or
-  classloaders like jboss and osgi see injected classes. This works because we instrument those classloaders
+   * 'module') classloaders like jboss and osgi see injected classes. This works because we
-  to load everything inside bootstrap packages.
+   * 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 class FieldBackedProvider implements InstrumentationContextProvider {
   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
+       * Install transformer that rewrites accesses to context store with specialized bytecode that
-      storage implementation.
+       * invokes appropriate storage implementation.
        */
       builder =
           builder.transform(getTransformerForASMVisitor(getContextStoreReadsRewritingVisitor()));
 
-      /*
+      /**
-      We inject into bootstrap classloader because field accessor interfaces are needed by
+       * We inject into bootstrap classloader because field accessor interfaces are needed by
-      context store implementations. Unfortunately this forces us to remove stored type checking
+       * context store implementations. Unfortunately this forces us to remove stored type checking
-      because actual classes may not be available at this point.
+       * 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
+       * We inject context store implementation into bootstrap classloader because same
-       * may be used by different instrumentations and it has to use same static map in case of
+       * implementation may be used by different instrumentations and it has to use same static map
-       * fallback to map-backed storage.
+       * in case of fallback to map-backed storage.
        */
       builder = builder.transform(bootstrapHelperInjector(contextStoreImplementations.values()));
     }
@@ -329,13 +330,13 @@ public class FieldBackedProvider implements InstrumentationContextProvider {
 
   @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
+         * For each context store defined in a current instrumentation we create an agent builder
-        that injects necessary fields.
+         * that injects necessary fields.
          */
         builder =
             builder
@@ -343,6 +344,13 @@ public class FieldBackedProvider implements InstrumentationContextProvider {
                     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 class FieldBackedProvider implements InstrumentationContextProvider {
           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)
+         * The idea here is that we can add fields if class is just being loaded
-        and we have to add same fields again if class we added fields before is being transformed again.
+         * (classBeingRedefined == null) and we have to add same fields again if class we added
-        Note: here we assume that Class#getInterfaces() returns list of interfaces defined immediately on a
+         * fields before is being transformed again. Note: here we assume that Class#getInterfaces()
-        given class, not inherited from its parents. It looks like current JVM implementation does exactly
+         * returns list of interfaces defined immediately on a given class, not inherited from its
-        this but javadoc is not explicit about that.
+         * 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 @@ public interface InstrumentationContextProvider {
 
   /** 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 class QueryInstrumentation extends AbstractHibernateInstrumentation {
       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 class QueryInstrumentation extends AbstractHibernateInstrumentation {
       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);
+}

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

@@ -0,0 +1,60 @@
+package spring;
+
+import java.util.Properties;
+import javax.sql.DataSource;
+import org.springframework.context.annotation.Bean;
+import org.springframework.data.jpa.repository.config.EnableJpaRepositories;
+import org.springframework.jdbc.datasource.DriverManagerDataSource;
+import org.springframework.orm.jpa.JpaTransactionManager;
+import org.springframework.orm.jpa.LocalContainerEntityManagerFactoryBean;
+import org.springframework.orm.jpa.vendor.Database;
+import org.springframework.orm.jpa.vendor.HibernateJpaVendorAdapter;
+import org.springframework.transaction.PlatformTransactionManager;
+
+@EnableJpaRepositories(basePackages = "spring")
+public class PersistenceConfig {
+
+  @Bean(name = "transactionManager")
+  public PlatformTransactionManager dbTransactionManager() {
+    final JpaTransactionManager transactionManager = new JpaTransactionManager();
+    transactionManager.setEntityManagerFactory(entityManagerFactory().getObject());
+    return transactionManager;
+  }
+
+  @Bean
+  public LocalContainerEntityManagerFactoryBean entityManagerFactory() {
+
+    final HibernateJpaVendorAdapter vendorAdapter = new HibernateJpaVendorAdapter();
+    vendorAdapter.setDatabase(Database.HSQL);
+    vendorAdapter.setGenerateDdl(true);
+
+    final LocalContainerEntityManagerFactoryBean em = new LocalContainerEntityManagerFactoryBean();
+    em.setDataSource(dataSource());
+    em.setPackagesToScan("spring");
+    em.setJpaVendorAdapter(vendorAdapter);
+    em.setJpaProperties(additionalProperties());
+
+    return em;
+  }
+
+  @Bean
+  public DataSource dataSource() {
+    final DriverManagerDataSource dataSource = new DriverManagerDataSource();
+    dataSource.setDriverClassName("org.hsqldb.jdbcDriver");
+    dataSource.setUrl("jdbc:hsqldb:mem:test");
+    dataSource.setUsername("sa");
+    dataSource.setPassword("1");
+    return dataSource;
+  }
+
+  private Properties additionalProperties() {
+    final Properties properties = new Properties();
+    properties.setProperty("hibernate.show_sql", "true");
+    properties.setProperty("hibernate.hbm2ddl.auto", "create");
+    properties.setProperty("hibernate.dialect", "org.hibernate.dialect.HSQLDialect");
+    //    properties.setProperty(
+    //        "hibernate.format_sql",
+    // env.getProperty("spring.jpa.properties.hibernate.format_sql"));
+    return properties;
+  }
+}