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

This commit is contained in:
Will Gittoes 2019-02-12 17:01:37 +11:00
parent 724a59c347
commit cd2467c76b
No known key found for this signature in database
GPG Key ID: 521026A02DB0BB42
6 changed files with 262 additions and 65 deletions

View File

@ -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);
}
}
}

View File

@ -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<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);

View File

@ -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 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 <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);
}
}

View File

@ -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;
}

View File

@ -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);
}
}
}

View File

@ -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"))
}
}
}