Add hibernate.Criteria instrumentation

dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/CriteriaInstrumentation.java

@@ -0,0 +1,81 @@
+package datadog.trace.instrumentation.hibernate;
+
+import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
+import static net.bytebuddy.matcher.ElementMatchers.isInterface;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.InstrumentationContext;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.hibernate.Criteria;
+
+@AutoService(Instrumenter.class)
+public class CriteriaInstrumentation extends Instrumenter.Default {
+
+  public CriteriaInstrumentation() {
+    super("hibernate");
+  }
+
+  @Override
+  public Map<String, String> contextStore() {
+    final Map<String, String> map = new HashMap<>();
+    map.put("org.hibernate.Criteria", SessionState.class.getName());
+    return Collections.unmodifiableMap(map);
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      "datadog.trace.instrumentation.hibernate.SessionMethodUtils",
+      "datadog.trace.instrumentation.hibernate.SessionState",
+    };
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return not(isInterface()).and(safeHasSuperType(named("org.hibernate.Criteria")));
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    final Map<ElementMatcher<? super MethodDescription>, String> transformers = new HashMap<>();
+    transformers.put(
+        isMethod().and(named("list").or(named("uniqueResult")).or(named("scroll"))),
+        CriteriaMethodAdvice.class.getName());
+
+    return transformers;
+  }
+
+  public static class CriteriaMethodAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static SessionState startMethod(
+        @Advice.This final Criteria criteria, @Advice.Origin("#m") final String name) {
+
+      final ContextStore<Criteria, SessionState> contextStore =
+          InstrumentationContext.get(Criteria.class, SessionState.class);
+
+      return SessionMethodUtils.startScopeFrom(
+          contextStore, criteria, "hibernate.criteria." + name);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void endMethod(
+        @Advice.This final Criteria criteria,
+        @Advice.Enter final SessionState state,
+        @Advice.Thrown final Throwable throwable) {
+
+      SessionMethodUtils.closeScope(state, throwable);
+    }
+  }
+}

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

@@ -5,7 +5,6 @@ import static net.bytebuddy.matcher.ElementMatchers.isInterface;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.not;
-import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
@@ -57,8 +56,7 @@ public class QueryInstrumentation extends Instrumenter.Default {
                     .or(named("executeUpdate"))
                     .or(named("uniqueResult"))
                     .or(named("iterate"))
-                    .or(named("scroll"))) // TODO(will): Instrument the ScrollableWhatever returned
+                    .or(named("scroll"))),
-            .and(takesArguments(0)),
         QueryMethodAdvice.class.getName());
 
     return transformers;

dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/SessionInstrumentation.java

@@ -25,6 +25,7 @@ import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.implementation.bytecode.assign.Assigner;
 import net.bytebuddy.matcher.ElementMatcher;
+import org.hibernate.Criteria;
 import org.hibernate.Query;
 import org.hibernate.SharedSessionContract;
 import org.hibernate.Transaction;
@@ -43,6 +44,7 @@ public class SessionInstrumentation extends Instrumenter.Default {
     map.put("org.hibernate.SharedSessionContract", SessionState.class.getName());
     map.put("org.hibernate.Query", SessionState.class.getName());
     map.put("org.hibernate.Transaction", SessionState.class.getName());
+    map.put("org.hibernate.Criteria", SessionState.class.getName());
     map.put("org.hibernate.engine.HibernateIterator", SessionState.class.getName());
     return Collections.unmodifiableMap(map);
   }
@@ -107,6 +109,10 @@ public class SessionInstrumentation extends Instrumenter.Default {
         isMethod().and(returns(safeHasSuperType(named("org.hibernate.Query")))),
         GetQueryAdvice.class.getName());
 
+    transformers.put(
+        isMethod().and(returns(safeHasSuperType(named("org.hibernate.Criteria")))),
+        GetCriteriaAdvice.class.getName());
+
     return transformers;
   }
 
@@ -206,4 +212,20 @@ public class SessionInstrumentation extends Instrumenter.Default {
           sessionContextStore, session, transactionContextStore, transaction);
     }
   }
+
+  public static class GetCriteriaAdvice {
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void getCriteria(
+        @Advice.This final SharedSessionContract session, @Advice.Return final Criteria criteria) {
+
+      final ContextStore<SharedSessionContract, SessionState> sessionContextStore =
+          InstrumentationContext.get(SharedSessionContract.class, SessionState.class);
+      final ContextStore<Criteria, SessionState> criteriaContextStore =
+          InstrumentationContext.get(Criteria.class, SessionState.class);
+
+      SessionMethodUtils.attachSpanFromStore(
+          sessionContextStore, session, criteriaContextStore, criteria);
+    }
+  }
 }

dd-java-agent/instrumentation/hibernate/src/test/groovy/AbstractHibernateTest.groovy

@@ -0,0 +1,46 @@
+import datadog.trace.agent.test.AgentTestRunner
+import org.hibernate.Session
+import org.hibernate.SessionFactory
+import org.hibernate.boot.MetadataSources
+import org.hibernate.boot.registry.StandardServiceRegistry
+import org.hibernate.boot.registry.StandardServiceRegistryBuilder
+import spock.lang.Shared
+
+abstract class AbstractHibernateTest extends AgentTestRunner {
+
+  @Shared
+  protected SessionFactory sessionFactory
+
+  @Shared
+  protected List<Value> prepopulated
+
+  def setupSpec() {
+    final StandardServiceRegistry registry =
+      new StandardServiceRegistryBuilder()
+        .configure()
+        .build()
+    try {
+      sessionFactory = new MetadataSources(registry).buildMetadata().buildSessionFactory()
+    } catch (Exception e) {
+      StandardServiceRegistryBuilder.destroy(registry)
+      return
+    }
+
+    // Pre-populate the DB, so delete/update can be tested.
+    Session writer = sessionFactory.openSession()
+    writer.beginTransaction()
+    prepopulated = new ArrayList<>()
+    for (int i = 0; i < 2; i++) {
+      prepopulated.add(new Value("Hello :) " + i))
+      writer.save(prepopulated.get(i))
+    }
+    writer.getTransaction().commit()
+    writer.close()
+  }
+
+  def cleanupSpec() {
+    if (sessionFactory != null) {
+      sessionFactory.close()
+    }
+  }
+}

dd-java-agent/instrumentation/hibernate/src/test/groovy/CriteriaTest.groovy

@@ -0,0 +1,110 @@
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.api.DDSpanTypes
+import datadog.trace.api.DDTags
+import io.opentracing.tag.Tags
+import org.hibernate.Criteria
+import org.hibernate.Session
+import org.hibernate.SessionFactory
+import org.hibernate.boot.MetadataSources
+import org.hibernate.boot.registry.StandardServiceRegistry
+import org.hibernate.boot.registry.StandardServiceRegistryBuilder
+import org.hibernate.criterion.Order
+import org.hibernate.criterion.Restrictions
+import spock.lang.Shared
+
+class CriteriaTest extends AgentTestRunner {
+
+  @Shared
+  private SessionFactory sessionFactory
+
+  def setupSpec() {
+    final StandardServiceRegistry registry =
+      new StandardServiceRegistryBuilder()
+        .configure()
+        .build()
+    try {
+      sessionFactory = new MetadataSources(registry).buildMetadata().buildSessionFactory()
+    } catch (Exception e) {
+      StandardServiceRegistryBuilder.destroy(registry)
+      return
+    }
+
+    Session session = sessionFactory.openSession()
+    session.beginTransaction()
+    session.save(new Value("A Hibernate value to be serialized"))
+    session.save(new Value("Another value"))
+    session.getTransaction().commit()
+    session.close()
+  }
+
+  def cleanupSpec() {
+    if (sessionFactory != null) {
+      sessionFactory.close()
+    }
+  }
+
+  def "test criteria.#methodName"() {
+    setup:
+    Session session = sessionFactory.openSession()
+    session.beginTransaction()
+    Criteria criteria = session.createCriteria(Value.class)
+      .add(Restrictions.like("name", "Hello"))
+      .addOrder(Order.desc("name"))
+    interaction.call(criteria)
+    session.getTransaction().commit()
+    session.close()
+
+    expect:
+    assertTraces(1) {
+      trace(0, 4) {
+        span(0) {
+          serviceName "hibernate"
+          resourceName "hibernate.session"
+          operationName "hibernate.session"
+          spanType DDSpanTypes.HIBERNATE
+          parent()
+          tags {
+            "$Tags.COMPONENT.key" "hibernate-java"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
+            "$DDTags.SPAN_TYPE" DDSpanTypes.HIBERNATE
+            defaultTags()
+          }
+        }
+        span(1) {
+          serviceName "hibernate"
+          resourceName "hibernate.transaction.commit"
+          operationName "hibernate.transaction.commit"
+          spanType DDSpanTypes.HIBERNATE
+          childOf span(0)
+          tags {
+            "$Tags.COMPONENT.key" "hibernate-java"
+            "$DDTags.SPAN_TYPE" DDSpanTypes.HIBERNATE
+            defaultTags()
+          }
+        }
+        span(2) {
+          serviceName "hibernate"
+          resourceName "hibernate.criteria.$methodName"
+          operationName "hibernate.criteria.$methodName"
+          spanType DDSpanTypes.HIBERNATE
+          childOf span(0)
+          tags {
+            "$Tags.COMPONENT.key" "hibernate-java"
+            "$DDTags.SPAN_TYPE" DDSpanTypes.HIBERNATE
+            defaultTags()
+          }
+        }
+        span(3) {
+          serviceName "h2"
+          childOf span(2)
+        }
+      }
+    }
+
+    where:
+    methodName     | interaction
+    "list"         | { c -> c.list() }
+    "uniqueResult" | { c -> c.uniqueResult() }
+    "scroll"       | { c -> c.scroll() }
+  }
+}

dd-java-agent/instrumentation/hibernate/src/test/groovy/QueryTest.groovy

@@ -1,47 +1,12 @@
-import datadog.trace.agent.test.AgentTestRunner
 import datadog.trace.api.DDSpanTypes
 import datadog.trace.api.DDTags
 import io.opentracing.tag.Tags
 import org.hibernate.Query
 import org.hibernate.Session
-import org.hibernate.SessionFactory
-import org.hibernate.boot.MetadataSources
-import org.hibernate.boot.registry.StandardServiceRegistry
-import org.hibernate.boot.registry.StandardServiceRegistryBuilder
-import spock.lang.Shared
 
-class QueryTest extends AgentTestRunner {
+class QueryTest extends AbstractHibernateTest {
 
-  @Shared
+  def "test hibernate query.#queryMethodName single call"() {
-  private SessionFactory sessionFactory
-
-  def setupSpec() {
-    final StandardServiceRegistry registry =
-      new StandardServiceRegistryBuilder()
-        .configure()
-        .build()
-    try {
-      sessionFactory = new MetadataSources(registry).buildMetadata().buildSessionFactory()
-    } catch (Exception e) {
-      StandardServiceRegistryBuilder.destroy(registry)
-      return
-    }
-
-    Session session = sessionFactory.openSession()
-    session.beginTransaction()
-    session.save(new Value("A Hibernate value to be serialized"))
-    session.save(new Value("Another value"))
-    session.getTransaction().commit()
-    session.close()
-  }
-
-  def cleanupSpec() {
-    if (sessionFactory != null) {
-      sessionFactory.close()
-    }
-  }
-
-  def "test hibernate query.#queryMethodName"() {
     setup:
 
     Session session = sessionFactory.openSession()
@@ -121,4 +86,89 @@ class QueryTest extends AgentTestRunner {
     }
   }
 
+  def "test hibernate query.iterate"() {
+    setup:
+
+    Session session = sessionFactory.openSession()
+    session.beginTransaction()
+    Query q = session.createQuery("from Value")
+    Iterator it = q.iterate()
+    while (it.hasNext()) {
+      it.next()
+    }
+    session.getTransaction().commit()
+    session.close()
+
+    expect:
+    assertTraces(1) {
+      trace(0, 6) {
+        span(0) {
+          serviceName "hibernate"
+          resourceName "hibernate.session"
+          operationName "hibernate.session"
+          spanType DDSpanTypes.HIBERNATE
+          parent()
+          tags {
+            "$Tags.COMPONENT.key" "hibernate-java"
+            "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
+            "$DDTags.SPAN_TYPE" DDSpanTypes.HIBERNATE
+            defaultTags()
+          }
+        }
+        span(1) {
+          serviceName "hibernate"
+          resourceName "hibernate.transaction.commit"
+          operationName "hibernate.transaction.commit"
+          spanType DDSpanTypes.HIBERNATE
+          childOf span(0)
+          tags {
+            "$Tags.COMPONENT.key" "hibernate-java"
+            "$DDTags.SPAN_TYPE" DDSpanTypes.HIBERNATE
+            defaultTags()
+          }
+        }
+        span(2) {
+          serviceName "hibernate"
+          resourceName "hibernate.iterator.next"
+          operationName "hibernate.iterator.next"
+          spanType DDSpanTypes.HIBERNATE
+          childOf span(0)
+          tags {
+            "$Tags.COMPONENT.key" "hibernate-java"
+            "$DDTags.SPAN_TYPE" DDSpanTypes.HIBERNATE
+            defaultTags()
+          }
+        }
+        span(3) {
+          serviceName "hibernate"
+          resourceName "hibernate.iterator.next"
+          operationName "hibernate.iterator.next"
+          spanType DDSpanTypes.HIBERNATE
+          childOf span(0)
+          tags {
+            "$Tags.COMPONENT.key" "hibernate-java"
+            "$DDTags.SPAN_TYPE" DDSpanTypes.HIBERNATE
+            defaultTags()
+          }
+        }
+        span(4) {
+          serviceName "hibernate"
+          resourceName "hibernate.query.iterate"
+          operationName "hibernate.query.iterate"
+          spanType DDSpanTypes.HIBERNATE
+          childOf span(0)
+          tags {
+            "$Tags.COMPONENT.key" "hibernate-java"
+            "$DDTags.SPAN_TYPE" DDSpanTypes.HIBERNATE
+            defaultTags()
+          }
+        }
+        span(5) {
+          serviceName "h2"
+          childOf span(4)
+        }
+      }
+    }
+  }
+
 }

dd-java-agent/instrumentation/hibernate/src/test/groovy/SessionTest.groovy

@@ -1,59 +1,20 @@
-import datadog.trace.agent.test.AgentTestRunner
 import datadog.trace.api.DDSpanTypes
 import datadog.trace.api.DDTags
 import io.opentracing.tag.Tags
 import org.hibernate.*
-import org.hibernate.boot.MetadataSources
-import org.hibernate.boot.registry.StandardServiceRegistry
-import org.hibernate.boot.registry.StandardServiceRegistryBuilder
 import spock.lang.Shared
 
-class SessionTest extends AgentTestRunner {
+class SessionTest extends AbstractHibernateTest {
-
-  @Shared
-  private SessionFactory sessionFactory
 
   @Shared
   private Map<String, Closure> sessionBuilders
 
-  @Shared
-  private List<Value> prepopulated
-
   def setupSpec() {
-    final StandardServiceRegistry registry =
-      new StandardServiceRegistryBuilder()
-        .configure()
-        .build()
-    try {
-      sessionFactory = new MetadataSources(registry).buildMetadata().buildSessionFactory();
-    } catch (Exception e) {
-      StandardServiceRegistryBuilder.destroy(registry)
-    }
-
     // Test two different types of Session. Groovy doesn't allow testing the cross-product/combinations of two data
     // tables, so we get this hack instead.
     sessionBuilders = new HashMap<>();
     sessionBuilders.put("Session", { return sessionFactory.openSession() })
     sessionBuilders.put("StatelessSession", { return sessionFactory.openStatelessSession() })
-
-    // Pre-populate the DB, so delete/update can be tested.
-    Session writer = sessionFactory.openSession()
-    writer.beginTransaction()
-    prepopulated = new ArrayList<>()
-    for (int i = 0; i < 2; i++) {
-      prepopulated.add(new Value("Hello :)"))
-      writer.save(prepopulated.get(i))
-    }
-    writer.getTransaction().commit()
-    writer.close()
-    TEST_WRITER.waitForTraces(1)
-    TEST_WRITER.clear()
-  }
-
-  def cleanupSpec() {
-    if (sessionFactory != null) {
-      sessionFactory.close()
-    }
   }