Use a generic Session MethodAdvice to instrument all synchronous Session methods; also support re-entrant Session methods
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
@@ -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<String, String> contextStore() {
final Map<String, String> 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<Query, Span> contextStore =
- InstrumentationContext.get(Query.class, Span.class);
+ final ContextStore<Query, SessionState> 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);
}
}
}

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
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<String, String> contextStore() {
final Map<String, String> 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<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
final Map<ElementMatcher<? super MethodDescription>, 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<Session, Span> contextStore =
- InstrumentationContext.get(Session.class, Span.class);
- contextStore.putIfAbsent(session, span);
+ final ContextStore<Session, SessionState> 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<Session, Span> contextStore =
- InstrumentationContext.get(Session.class, Span.class);
- final Span span = contextStore.get(session);
-
- if (span == null) {
+ final ContextStore<Session, SessionState> 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<Session, Span> contextStore =
- InstrumentationContext.get(Session.class, Span.class);
- return SessionMethodUtils.startScopeFrom(contextStore, session, "hibernate.save");
+ final ContextStore<Session, SessionState> 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<Session, Span> sessionContextStore =
- InstrumentationContext.get(Session.class, Span.class);
- final ContextStore<Query, Span> queryContextStore =
- InstrumentationContext.get(Query.class, Span.class);
+ final ContextStore<Session, SessionState> sessionContextStore =
+ InstrumentationContext.get(Session.class, SessionState.class);
+ final ContextStore<Query, SessionState> 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<Session, Span> sessionContextStore =
- InstrumentationContext.get(Session.class, Span.class);
- final ContextStore<Transaction, Span> transactionContextStore =
- InstrumentationContext.get(Transaction.class, Span.class);
+ final ContextStore<Session, SessionState> sessionContextStore =
+ InstrumentationContext.get(Session.class, SessionState.class);
+ final ContextStore<Transaction, SessionState> transactionContextStore =
+ InstrumentationContext.get(Transaction.class, SessionState.class);
SessionMethodUtils.attachSpanFromSession(
sessionContextStore, session, transactionContextStore, transaction);

MODIFY
dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/SessionMethodUtils.java
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 <T> Scope startScopeFrom(
- final ContextStore<T, Span> contextStore, final T spanKey, final String operationName) {
+ public static <T> SessionState startScopeFrom(
+ final ContextStore<T, SessionState> contextStore,
+ final T spanKey,
+ final String operationName,
+ final String resourceName) {
+
+ final SessionState sessionState = contextStore.get(spanKey);
- final Span span = 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 <T> SessionState startScopeFrom(
+ final ContextStore<T, SessionState> 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 <T> void attachSpanFromSession(
- final ContextStore<Session, Span> sessionContextStore,
+ final ContextStore<Session, SessionState> sessionContextStore,
final SharedSessionContract session,
- final ContextStore<T, Span> targetContextStore,
+ final ContextStore<T, SessionState> 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);
}
}

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

MODIFY
dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/TransactionInstrumentation.java
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<String, String> contextStore() {
final Map<String, String> 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<Transaction, Span> contextStore =
- InstrumentationContext.get(Transaction.class, Span.class);
+ final ContextStore<Transaction, SessionState> 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);
}
}
}

MODIFY
dd-java-agent/instrumentation/hibernate/src/test/groovy/SessionTest.groovy
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,11 +78,13 @@ 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()
@@ -83,11 +92,97 @@ class SessionTest extends AgentTestRunner {
}
span(3) {
serviceName "hibernate"
- resourceName "hibernate.save"
operationName "hibernate.save"
+ resourceName "Value"
+ spanType DDSpanTypes.HIBERNATE
+ childOf span(0)
+ tags {
+ "$Tags.COMPONENT.key" "hibernate-java"
+ "$DDTags.SPAN_TYPE" DDSpanTypes.HIBERNATE
+ defaultTags()
+ }
+ }
+ span(4) {
+ serviceName "hibernate"
+ operationName "hibernate.save"
+ resourceName "Value"
+ spanType DDSpanTypes.HIBERNATE
+ childOf span(0)
+ tags {
+ "$Tags.COMPONENT.key" "hibernate-java"
+ "$DDTags.SPAN_TYPE" DDSpanTypes.HIBERNATE
+ defaultTags()
+ }
+ }
+ }
+ }
+ }
+
+ 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()
@@ -95,5 +190,24 @@ class SessionTest extends AgentTestRunner {
}
}
}
+
+ 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"))
+ }
}
}

Use a generic Session MethodAdvice to instrument all synchronous Session methods; also support re-entrant Session methods