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 cd44c85f70..1bbc8757e8 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 @@ -12,8 +12,6 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.bootstrap.ContextStore; import datadog.trace.bootstrap.InstrumentationContext; -import io.opentracing.Scope; -import io.opentracing.Span; import java.util.Collections; import java.util.HashMap; import java.util.Map; @@ -33,7 +31,7 @@ public class QueryInstrumentation extends Instrumenter.Default { @Override public Map contextStore() { final Map map = new HashMap<>(); - map.put("org.hibernate.query.Query", Span.class.getName()); + map.put("org.hibernate.query.Query", SessionState.class.getName()); return Collections.unmodifiableMap(map); } @@ -41,6 +39,7 @@ public class QueryInstrumentation extends Instrumenter.Default { public String[] helperClassNames() { return new String[] { "datadog.trace.instrumentation.hibernate.SessionMethodUtils", + "datadog.trace.instrumentation.hibernate.SessionState", }; } @@ -65,10 +64,10 @@ public class QueryInstrumentation extends Instrumenter.Default { public static class QueryListAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static Scope startList(@Advice.This final Query query) { + public static SessionState startList(@Advice.This final Query query) { - final ContextStore contextStore = - InstrumentationContext.get(Query.class, Span.class); + final ContextStore contextStore = + InstrumentationContext.get(Query.class, SessionState.class); return SessionMethodUtils.startScopeFrom(contextStore, query, "hibernate.query.list"); } @@ -76,10 +75,10 @@ public class QueryInstrumentation extends Instrumenter.Default { @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void endList( @Advice.This final Query query, - @Advice.Enter final Scope scope, + @Advice.Enter final SessionState state, @Advice.Thrown final Throwable throwable) { - SessionMethodUtils.closeScope(scope, throwable); + SessionMethodUtils.closeScope(state, throwable); } } } 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 683fd65aa6..47549129d7 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 @@ -1,6 +1,7 @@ package datadog.trace.instrumentation.hibernate; import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType; +import static datadog.trace.instrumentation.hibernate.SessionMethodUtils.entityName; import static io.opentracing.log.Fields.ERROR_OBJECT; import static net.bytebuddy.matcher.ElementMatchers.isDeclaredBy; import static net.bytebuddy.matcher.ElementMatchers.isInterface; @@ -8,6 +9,7 @@ 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.returns; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; @@ -16,7 +18,6 @@ import datadog.trace.api.DDSpanTypes; import datadog.trace.api.DDTags; import datadog.trace.bootstrap.ContextStore; import datadog.trace.bootstrap.InstrumentationContext; -import io.opentracing.Scope; import io.opentracing.Span; import io.opentracing.tag.Tags; import io.opentracing.util.GlobalTracer; @@ -44,9 +45,9 @@ public class SessionInstrumentation extends Instrumenter.Default { @Override public Map contextStore() { final Map map = new HashMap<>(); - map.put("org.hibernate.Session", Span.class.getName()); - map.put("org.hibernate.query.Query", Span.class.getName()); - map.put("org.hibernate.Transaction", Span.class.getName()); + map.put("org.hibernate.Session", SessionState.class.getName()); + map.put("org.hibernate.query.Query", SessionState.class.getName()); + map.put("org.hibernate.Transaction", SessionState.class.getName()); return Collections.unmodifiableMap(map); } @@ -54,6 +55,7 @@ public class SessionInstrumentation extends Instrumenter.Default { public String[] helperClassNames() { return new String[] { "datadog.trace.instrumentation.hibernate.SessionMethodUtils", + "datadog.trace.instrumentation.hibernate.SessionState", }; } @@ -70,6 +72,8 @@ public class SessionInstrumentation extends Instrumenter.Default { @Override public Map, String> transformers() { final Map, String> transformers = new HashMap<>(); + // Session lifecycle. A span will cover openSession->Session.close, but no scope will be + // generated. transformers.put( isMethod() .and(named("openSession")) @@ -83,12 +87,34 @@ public class SessionInstrumentation extends Instrumenter.Default { .and(isDeclaredBy(safeHasSuperType(named("org.hibernate.Session")))) .and(takesArguments(0)), SessionCloseAdvice.class.getName()); + + // Session synchronous methods we want to instrument. transformers.put( isMethod() - .and(named("save")) + .and( + named("save") + .or(named("replicate")) + .or(named("saveOrUpdate")) + .or(named("update")) + .or(named("merge")) + .or(named("persist")) + .or(named("lock")) + .or(named("refresh")) + .or(named("delete"))) + .and(isDeclaredBy(safeHasSuperType(named("org.hibernate.Session")))), + SessionMethodAdvice.class.getName()); + // Handle the generic and non-generic 'get' separately. + transformers.put( + isMethod() + .and(named("get")) .and(isDeclaredBy(safeHasSuperType(named("org.hibernate.Session")))) - .and(takesArguments(1)), - SessionSaveAdvice.class.getName()); + .and(returns(named("java.lang.Object"))) + .and(takesArgument(0, named("java.lang.String"))), + SessionMethodAdvice.class.getName()); + + // These methods return some object that we want to instrument, and so the Advice will pin the + // current Span to + // the returned object using a ContextStore. transformers.put( isMethod() .and(named("beginTransaction").or(named("getTransaction"))) @@ -99,8 +125,7 @@ public class SessionInstrumentation extends Instrumenter.Default { transformers.put( isMethod() .and(named("createQuery")) - .and(isDeclaredBy(safeHasSuperType(named("org.hibernate.SharedSessionContract")))) - .and(takesArguments(1)), + .and(isDeclaredBy(safeHasSuperType(named("org.hibernate.SharedSessionContract")))), GetQueryAdvice.class.getName()); return transformers; @@ -120,9 +145,9 @@ public class SessionInstrumentation extends Instrumenter.Default { .withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CLIENT) .start(); - final ContextStore contextStore = - InstrumentationContext.get(Session.class, Span.class); - contextStore.putIfAbsent(session, span); + final ContextStore contextStore = + InstrumentationContext.get(Session.class, SessionState.class); + contextStore.putIfAbsent(session, new SessionState(span)); } } @@ -132,13 +157,17 @@ public class SessionInstrumentation extends Instrumenter.Default { public static void closeSession( @Advice.This final Session session, @Advice.Thrown final Throwable throwable) { - final ContextStore contextStore = - InstrumentationContext.get(Session.class, Span.class); - final Span span = contextStore.get(session); - - if (span == null) { + final ContextStore contextStore = + InstrumentationContext.get(Session.class, SessionState.class); + final SessionState state = contextStore.get(session); + if (state == null || state.getSessionSpan() == null) { return; } + if (state.getMethodScope() != null) { + System.err.println("THIS IS WRONG"); // TODO: proper warning/logging. + } + final Span span = state.getSessionSpan(); + if (throwable != null) { Tags.ERROR.set(span, true); span.log(Collections.singletonMap(ERROR_OBJECT, throwable)); @@ -147,24 +176,27 @@ public class SessionInstrumentation extends Instrumenter.Default { } } - public static class SessionSaveAdvice { + public static class SessionMethodAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static Scope startSave( - @Advice.This final Session session, @Advice.Argument(0) final Object entity) { + public static SessionState startSave( + @Advice.This final Session session, + @Advice.Origin("#m") final String name, + @Advice.Argument(0) final Object entity) { - final ContextStore contextStore = - InstrumentationContext.get(Session.class, Span.class); - return SessionMethodUtils.startScopeFrom(contextStore, session, "hibernate.save"); + final ContextStore contextStore = + InstrumentationContext.get(Session.class, SessionState.class); + return SessionMethodUtils.startScopeFrom( + contextStore, session, "hibernate." + name, entityName(entity)); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void endSave( @Advice.This final Session session, - @Advice.Enter final Scope scope, + @Advice.Enter final SessionState sessionState, @Advice.Thrown final Throwable throwable) { - SessionMethodUtils.closeScope(scope, throwable); + SessionMethodUtils.closeScope(sessionState, throwable); } } @@ -179,10 +211,10 @@ public class SessionInstrumentation extends Instrumenter.Default { return; } - final ContextStore sessionContextStore = - InstrumentationContext.get(Session.class, Span.class); - final ContextStore queryContextStore = - InstrumentationContext.get(Query.class, Span.class); + final ContextStore sessionContextStore = + InstrumentationContext.get(Session.class, SessionState.class); + final ContextStore queryContextStore = + InstrumentationContext.get(Query.class, SessionState.class); SessionMethodUtils.attachSpanFromSession( sessionContextStore, session, queryContextStore, query); @@ -196,10 +228,10 @@ public class SessionInstrumentation extends Instrumenter.Default { @Advice.This final AbstractSharedSessionContract session, @Advice.Return(readOnly = false) final Transaction transaction) { - final ContextStore sessionContextStore = - InstrumentationContext.get(Session.class, Span.class); - final ContextStore transactionContextStore = - InstrumentationContext.get(Transaction.class, Span.class); + final ContextStore sessionContextStore = + InstrumentationContext.get(Session.class, SessionState.class); + final ContextStore transactionContextStore = + InstrumentationContext.get(Transaction.class, SessionState.class); SessionMethodUtils.attachSpanFromSession( sessionContextStore, session, transactionContextStore, transaction); diff --git a/dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/SessionMethodUtils.java b/dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/SessionMethodUtils.java index ddff1308d2..623e7c1906 100644 --- a/dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/SessionMethodUtils.java +++ b/dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/SessionMethodUtils.java @@ -10,33 +10,73 @@ import io.opentracing.Span; import io.opentracing.tag.Tags; import io.opentracing.util.GlobalTracer; import java.util.Collections; +import javax.persistence.Entity; import org.hibernate.Session; import org.hibernate.SharedSessionContract; public class SessionMethodUtils { + public static String entityName(final Object entityOrName) { + String name = "unknown object"; + if (entityOrName == null) { + return name; // Hibernate can internally call save(null, Entity). + } + if (entityOrName instanceof String) { + name = (String) entityOrName; + } else if (entityOrName.getClass().isAnnotationPresent(Entity.class)) { + // This is the object being replicated. + name = entityOrName.getClass().getName(); + } + return name; + } + // Starts a scope as a child from a Span, where the Span is attached to the given spanKey using // the given contextStore. - public static Scope startScopeFrom( - final ContextStore contextStore, final T spanKey, final String operationName) { + public static SessionState startScopeFrom( + final ContextStore contextStore, + final T spanKey, + final String operationName, + final String resourceName) { - final Span span = contextStore.get(spanKey); + final SessionState sessionState = contextStore.get(spanKey); + + if (sessionState.getMethodScope() != null) { + // This method call was re-entrant. Do nothing, since it is being traced by the parent/first + // call. + return sessionState; + } final Scope scope = GlobalTracer.get() .buildSpan(operationName) - .asChildOf(span) // Can be null. + .asChildOf(sessionState.getSessionSpan()) .withTag(DDTags.SERVICE_NAME, "hibernate") .withTag(DDTags.SPAN_TYPE, DDSpanTypes.HIBERNATE) + .withTag(DDTags.RESOURCE_NAME, resourceName) .withTag(Tags.COMPONENT.getKey(), "hibernate-java") .startActive(true); + sessionState.setMethodScope(scope); - return scope; + return sessionState; + } + + public static SessionState startScopeFrom( + final ContextStore contextStore, + final T spanKey, + final String operationName) { + return startScopeFrom(contextStore, spanKey, operationName, operationName); } // Closes a Scope/Span, adding an error tag if the given Throwable is not null. - public static void closeScope(final Scope scope, final Throwable throwable) { + public static void closeScope(final SessionState sessionState, final Throwable throwable) { + if (sessionState == null || sessionState.getMethodScope() == null) { + // This method call was re-entrant. Do nothing, since it is being traced by the parent/first + // call. + return; + } + + final Scope scope = sessionState.getMethodScope(); final Span span = scope.span(); if (throwable != null) { Tags.ERROR.set(span, true); @@ -45,14 +85,15 @@ public class SessionMethodUtils { span.finish(); scope.close(); + sessionState.setMethodScope(null); } // Copies a span from the given Session ContextStore into the targetContextStore. Used to // propagate a Span from a Session to transient Session objects such as Transaction and Query. public static void attachSpanFromSession( - final ContextStore sessionContextStore, + final ContextStore sessionContextStore, final SharedSessionContract session, - final ContextStore targetContextStore, + final ContextStore targetContextStore, final T target) { if (!(session instanceof Session) @@ -61,11 +102,11 @@ public class SessionMethodUtils { return; } - final Span span = sessionContextStore.get((Session) session); - if (span == null) { + final SessionState state = sessionContextStore.get((Session) session); + if (state == null) { return; } - targetContextStore.putIfAbsent(target, span); + targetContextStore.putIfAbsent(target, state); } } diff --git a/dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/SessionState.java b/dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/SessionState.java new file mode 100644 index 0000000000..26cd2a8ee3 --- /dev/null +++ b/dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/SessionState.java @@ -0,0 +1,12 @@ +package datadog.trace.instrumentation.hibernate; + +import io.opentracing.Scope; +import io.opentracing.Span; +import lombok.Data; +import lombok.NonNull; + +@Data +public class SessionState { + @NonNull public Span sessionSpan; + public Scope methodScope; +} diff --git a/dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/TransactionInstrumentation.java b/dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/TransactionInstrumentation.java index 47b2501343..e0a0d7482c 100644 --- a/dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/TransactionInstrumentation.java +++ b/dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/TransactionInstrumentation.java @@ -12,8 +12,6 @@ import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.bootstrap.ContextStore; import datadog.trace.bootstrap.InstrumentationContext; -import io.opentracing.Scope; -import io.opentracing.Span; import java.util.Collections; import java.util.HashMap; import java.util.Map; @@ -33,7 +31,7 @@ public class TransactionInstrumentation extends Instrumenter.Default { @Override public Map contextStore() { final Map map = new HashMap<>(); - map.put("org.hibernate.Transaction", Span.class.getName()); + map.put("org.hibernate.Transaction", SessionState.class.getName()); return Collections.unmodifiableMap(map); } @@ -41,6 +39,7 @@ public class TransactionInstrumentation extends Instrumenter.Default { public String[] helperClassNames() { return new String[] { "datadog.trace.instrumentation.hibernate.SessionMethodUtils", + "datadog.trace.instrumentation.hibernate.SessionState", }; } @@ -64,10 +63,10 @@ public class TransactionInstrumentation extends Instrumenter.Default { public static class TransactionCommitAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static Scope startCommit(@Advice.This final Transaction transaction) { + public static SessionState startCommit(@Advice.This final Transaction transaction) { - final ContextStore contextStore = - InstrumentationContext.get(Transaction.class, Span.class); + final ContextStore contextStore = + InstrumentationContext.get(Transaction.class, SessionState.class); return SessionMethodUtils.startScopeFrom( contextStore, transaction, "hibernate.transaction.commit"); @@ -76,10 +75,10 @@ public class TransactionInstrumentation extends Instrumenter.Default { @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) public static void endCommit( @Advice.This final Transaction transaction, - @Advice.Enter final Scope scope, + @Advice.Enter final SessionState state, @Advice.Thrown final Throwable throwable) { - SessionMethodUtils.closeScope(scope, throwable); + SessionMethodUtils.closeScope(state, throwable); } } } 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 c32e13b012..7d78d28f14 100644 --- a/dd-java-agent/instrumentation/hibernate/src/test/groovy/SessionTest.groovy +++ b/dd-java-agent/instrumentation/hibernate/src/test/groovy/SessionTest.groovy @@ -2,6 +2,8 @@ import datadog.trace.agent.test.AgentTestRunner import datadog.trace.api.DDSpanTypes import datadog.trace.api.DDTags import io.opentracing.tag.Tags +import org.hibernate.MappingException +import org.hibernate.ReplicationMode import org.hibernate.Session import org.hibernate.SessionFactory import org.hibernate.boot.MetadataSources @@ -32,18 +34,23 @@ class SessionTest extends AgentTestRunner { } } - def "test hibernate session"() { + def "test hibernate session save"() { setup: Session session = sessionFactory.openSession() session.beginTransaction() session.save(new Value("A Hibernate value to be serialized")) session.save(new Value("Another value")) + try { + session.save(new Long(123) /* Not a valid Entity, should throw */) + } catch (Exception e) { + // Should set error tag. + } session.getTransaction().commit() session.close() expect: assertTraces(1) { - trace(0, 4) { + trace(0, 5) { span(0) { serviceName "hibernate" resourceName "hibernate.session" @@ -71,8 +78,22 @@ class SessionTest extends AgentTestRunner { } span(2) { serviceName "hibernate" - resourceName "hibernate.save" operationName "hibernate.save" + resourceName "unknown object" + spanType DDSpanTypes.HIBERNATE + childOf span(0) + errored true + tags { + errorTags(MappingException, "Unknown entity: java.lang.Long") + "$Tags.COMPONENT.key" "hibernate-java" + "$DDTags.SPAN_TYPE" DDSpanTypes.HIBERNATE + defaultTags() + } + } + span(3) { + serviceName "hibernate" + operationName "hibernate.save" + resourceName "Value" spanType DDSpanTypes.HIBERNATE childOf span(0) tags { @@ -81,10 +102,10 @@ class SessionTest extends AgentTestRunner { defaultTags() } } - span(3) { + span(4) { serviceName "hibernate" - resourceName "hibernate.save" operationName "hibernate.save" + resourceName "Value" spanType DDSpanTypes.HIBERNATE childOf span(0) tags { @@ -96,4 +117,97 @@ class SessionTest extends AgentTestRunner { } } } + + def "test hibernate session replicate"() { + setup: + Session session = sessionFactory.openSession() + session.beginTransaction() + Value value = new Value("Hello :)") + session.save(value) + session.getTransaction().commit() + session.close() + + TEST_WRITER.waitForTraces(1) + TEST_WRITER.clear() + + session = sessionFactory.openSession() + session.beginTransaction() + + try { + sessionMethodTest.call(session, value) + } catch (Exception e) { + // We expected this, we should see the error field set on the span. + } + + session.getTransaction().commit() + session.close() + + expect: + assertTraces(1) { + trace(0, 3) { + 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 resource + operationName "hibernate.$methodName" + spanType DDSpanTypes.HIBERNATE + childOf span(0) + if (isError) { + errored true + } + tags { + if (isError) { + errorTags(MappingException, "Unknown entity: java.lang.Long") + } + "$Tags.COMPONENT.key" "hibernate-java" + "$DDTags.SPAN_TYPE" DDSpanTypes.HIBERNATE + defaultTags() + } + } + } + } + + where: + methodName | resource | isError | sessionMethodTest + "replicate" | "Value" | false | { sesh, val -> + Value replicated = new Value(val.getName() + " replicated") + replicated.setId(val.getId()) + sesh.replicate(replicated, ReplicationMode.OVERWRITE) + } + "replicate" | "Value" | false | { sesh, val -> + Value replicated = new Value(val.getName() + " replicated") + replicated.setId(val.getId()) + sesh.replicate("Value", replicated, ReplicationMode.OVERWRITE) + } + "replicate" | "unknown object" | true | { sesh, val -> + sesh.replicate(new Long(123) /* Not a valid entity */, ReplicationMode.OVERWRITE) + } + "save" | "Value" | false | { sesh, val -> + sesh.save(new Value("Another value")) + } + } }