Add hibernate.Criteria instrumentation
ADD
dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/CriteriaInstrumentation.java
+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);
+ }
+ }
+}

MODIFY
dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/QueryInstrumentation.java
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
- .and(takesArguments(0)),
+ .or(named("scroll"))),
QueryMethodAdvice.class.getName());
return transformers;

MODIFY
dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/SessionInstrumentation.java
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);
+ }
+ }
}

ADD
dd-java-agent/instrumentation/hibernate/src/test/groovy/AbstractHibernateTest.groovy
+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()
+ }
+ }
+}

ADD
dd-java-agent/instrumentation/hibernate/src/test/groovy/CriteriaTest.groovy
+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() }
+ }
+}

MODIFY
dd-java-agent/instrumentation/hibernate/src/test/groovy/QueryTest.groovy
dd-java-agent/instrumentation/hibernate/src/test/groovy/QueryTest.groovy
-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
- 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"() {
+ def "test hibernate query.#queryMethodName single call"() {
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)
+ }
+ }
+ }
+ }
+
}

MODIFY
dd-java-agent/instrumentation/hibernate/src/test/groovy/SessionTest.groovy
dd-java-agent/instrumentation/hibernate/src/test/groovy/SessionTest.groovy
-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 {
-
- @Shared
- private SessionFactory sessionFactory
+class SessionTest extends AbstractHibernateTest {
@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()
- }
}

Add hibernate.Criteria instrumentation