diff --git a/dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/CriteriaInstrumentation.java b/dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/CriteriaInstrumentation.java new file mode 100644 index 0000000000..3c716371ee --- /dev/null +++ b/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 contextStore() { + final Map 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 typeMatcher() { + return not(isInterface()).and(safeHasSuperType(named("org.hibernate.Criteria"))); + } + + @Override + public Map, String> transformers() { + final Map, 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 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); + } + } +} diff --git a/dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/QueryInstrumentation.java b/dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/QueryInstrumentation.java index b49dbdde3f..bdaacac34f 100644 --- a/dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/QueryInstrumentation.java +++ b/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; diff --git a/dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/SessionInstrumentation.java b/dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/SessionInstrumentation.java index c3f0f7ce87..d264d7352e 100644 --- a/dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/SessionInstrumentation.java +++ b/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 sessionContextStore = + InstrumentationContext.get(SharedSessionContract.class, SessionState.class); + final ContextStore criteriaContextStore = + InstrumentationContext.get(Criteria.class, SessionState.class); + + SessionMethodUtils.attachSpanFromStore( + sessionContextStore, session, criteriaContextStore, criteria); + } + } } diff --git a/dd-java-agent/instrumentation/hibernate/src/test/groovy/AbstractHibernateTest.groovy b/dd-java-agent/instrumentation/hibernate/src/test/groovy/AbstractHibernateTest.groovy new file mode 100644 index 0000000000..85e65b3d35 --- /dev/null +++ b/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 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() + } + } +} diff --git a/dd-java-agent/instrumentation/hibernate/src/test/groovy/CriteriaTest.groovy b/dd-java-agent/instrumentation/hibernate/src/test/groovy/CriteriaTest.groovy new file mode 100644 index 0000000000..30811ae4fa --- /dev/null +++ b/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() } + } +} diff --git a/dd-java-agent/instrumentation/hibernate/src/test/groovy/QueryTest.groovy b/dd-java-agent/instrumentation/hibernate/src/test/groovy/QueryTest.groovy index 86606119fe..b740a08a3c 100644 --- a/dd-java-agent/instrumentation/hibernate/src/test/groovy/QueryTest.groovy +++ b/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 - 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) + } + } + } + } + } diff --git a/dd-java-agent/instrumentation/hibernate/src/test/groovy/SessionTest.groovy b/dd-java-agent/instrumentation/hibernate/src/test/groovy/SessionTest.groovy index a6ffc772d8..7ee6eb07f5 100644 --- a/dd-java-agent/instrumentation/hibernate/src/test/groovy/SessionTest.groovy +++ b/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 { - - @Shared - private SessionFactory sessionFactory +class SessionTest extends AbstractHibernateTest { @Shared private Map sessionBuilders - @Shared - private List 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() - } }