Instrument SharedSessionContract instead of Session
This commit is contained in:
parent
180ba619a2
commit
e019fcbda5
|
@ -24,7 +24,7 @@ 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.Session;
|
||||
import org.hibernate.SharedSessionContract;
|
||||
|
||||
@AutoService(Instrumenter.class)
|
||||
public class SessionFactoryInstrumentation extends Instrumenter.Default {
|
||||
|
@ -36,7 +36,7 @@ public class SessionFactoryInstrumentation extends Instrumenter.Default {
|
|||
@Override
|
||||
public Map<String, String> contextStore() {
|
||||
final Map<String, String> map = new HashMap<>();
|
||||
map.put("org.hibernate.Session", SessionState.class.getName());
|
||||
map.put("org.hibernate.SharedSessionContract", SessionState.class.getName());
|
||||
return Collections.unmodifiableMap(map);
|
||||
}
|
||||
|
||||
|
@ -59,9 +59,11 @@ public class SessionFactoryInstrumentation extends Instrumenter.Default {
|
|||
// generated.
|
||||
transformers.put(
|
||||
isMethod()
|
||||
.and(named("openSession"))
|
||||
.and(named("openSession").or(named("openStatelessSession")))
|
||||
.and(takesArguments(0))
|
||||
.and(returns(named("org.hibernate.Session"))),
|
||||
.and(
|
||||
returns(
|
||||
named("org.hibernate.Session").or(named("org.hibernate.StatelessSession")))),
|
||||
SessionFactoryAdvice.class.getName());
|
||||
|
||||
return transformers;
|
||||
|
@ -70,7 +72,7 @@ public class SessionFactoryInstrumentation extends Instrumenter.Default {
|
|||
public static class SessionFactoryAdvice {
|
||||
|
||||
@Advice.OnMethodExit(suppress = Throwable.class)
|
||||
public static void openSession(@Advice.Return(readOnly = false) final Session session) {
|
||||
public static void openSession(@Advice.Return final SharedSessionContract session) {
|
||||
|
||||
final Span span =
|
||||
GlobalTracer.get()
|
||||
|
@ -81,8 +83,8 @@ public class SessionFactoryInstrumentation extends Instrumenter.Default {
|
|||
.withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CLIENT)
|
||||
.start();
|
||||
|
||||
final ContextStore<Session, SessionState> contextStore =
|
||||
InstrumentationContext.get(Session.class, SessionState.class);
|
||||
final ContextStore<SharedSessionContract, SessionState> contextStore =
|
||||
InstrumentationContext.get(SharedSessionContract.class, SessionState.class);
|
||||
contextStore.putIfAbsent(session, new SessionState(span));
|
||||
}
|
||||
}
|
||||
|
|
|
@ -3,7 +3,6 @@ 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;
|
||||
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
|
||||
import static net.bytebuddy.matcher.ElementMatchers.named;
|
||||
|
@ -25,7 +24,6 @@ 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.Session;
|
||||
import org.hibernate.SharedSessionContract;
|
||||
import org.hibernate.Transaction;
|
||||
import org.hibernate.query.Query;
|
||||
|
@ -41,7 +39,7 @@ public class SessionInstrumentation extends Instrumenter.Default {
|
|||
@Override
|
||||
public Map<String, String> contextStore() {
|
||||
final Map<String, String> map = new HashMap<>();
|
||||
map.put("org.hibernate.Session", SessionState.class.getName());
|
||||
map.put("org.hibernate.SharedSessionContract", SessionState.class.getName());
|
||||
map.put("org.hibernate.query.Query", SessionState.class.getName());
|
||||
map.put("org.hibernate.Transaction", SessionState.class.getName());
|
||||
return Collections.unmodifiableMap(map);
|
||||
|
@ -57,22 +55,14 @@ public class SessionInstrumentation extends Instrumenter.Default {
|
|||
|
||||
@Override
|
||||
public ElementMatcher<TypeDescription> typeMatcher() {
|
||||
return not(isInterface())
|
||||
.and(
|
||||
safeHasSuperType(
|
||||
named("org.hibernate.Session")
|
||||
.or(named("org.hibernate.internal.AbstractSharedSessionContract"))));
|
||||
return not(isInterface()).and(safeHasSuperType(named("org.hibernate.SharedSessionContract")));
|
||||
}
|
||||
|
||||
@Override
|
||||
public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
|
||||
final Map<ElementMatcher<? super MethodDescription>, String> transformers = new HashMap<>();
|
||||
transformers.put(
|
||||
isMethod()
|
||||
.and(named("close"))
|
||||
.and(isDeclaredBy(safeHasSuperType(named("org.hibernate.Session"))))
|
||||
.and(takesArguments(0)),
|
||||
SessionCloseAdvice.class.getName());
|
||||
isMethod().and(named("close")).and(takesArguments(0)), SessionCloseAdvice.class.getName());
|
||||
|
||||
// Session synchronous methods we want to instrument.
|
||||
transformers.put(
|
||||
|
@ -86,33 +76,26 @@ public class SessionInstrumentation extends Instrumenter.Default {
|
|||
.or(named("persist"))
|
||||
.or(named("lock"))
|
||||
.or(named("refresh"))
|
||||
.or(named("delete")))
|
||||
.and(isDeclaredBy(safeHasSuperType(named("org.hibernate.Session")))),
|
||||
.or(named("delete"))),
|
||||
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(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.
|
||||
// current Span to the returned object using a ContextStore.
|
||||
transformers.put(
|
||||
isMethod()
|
||||
.and(named("beginTransaction").or(named("getTransaction")))
|
||||
.and(isDeclaredBy(safeHasSuperType(named("org.hibernate.SharedSessionContract"))))
|
||||
.and(takesArguments(0))
|
||||
.and(returns(named("org.hibernate.Transaction"))),
|
||||
GetTransactionAdvice.class.getName());
|
||||
transformers.put(
|
||||
isMethod()
|
||||
.and(named("createQuery"))
|
||||
.and(isDeclaredBy(safeHasSuperType(named("org.hibernate.SharedSessionContract")))),
|
||||
GetQueryAdvice.class.getName());
|
||||
|
||||
transformers.put(isMethod().and(named("createQuery")), GetQueryAdvice.class.getName());
|
||||
|
||||
return transformers;
|
||||
}
|
||||
|
@ -121,10 +104,11 @@ public class SessionInstrumentation extends Instrumenter.Default {
|
|||
|
||||
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
|
||||
public static void closeSession(
|
||||
@Advice.This final Session session, @Advice.Thrown final Throwable throwable) {
|
||||
@Advice.This final SharedSessionContract session,
|
||||
@Advice.Thrown final Throwable throwable) {
|
||||
|
||||
final ContextStore<Session, SessionState> contextStore =
|
||||
InstrumentationContext.get(Session.class, SessionState.class);
|
||||
final ContextStore<SharedSessionContract, SessionState> contextStore =
|
||||
InstrumentationContext.get(SharedSessionContract.class, SessionState.class);
|
||||
final SessionState state = contextStore.get(session);
|
||||
if (state == null || state.getSessionSpan() == null) {
|
||||
return;
|
||||
|
@ -146,19 +130,19 @@ public class SessionInstrumentation extends Instrumenter.Default {
|
|||
|
||||
@Advice.OnMethodEnter(suppress = Throwable.class)
|
||||
public static SessionState startSave(
|
||||
@Advice.This final Session session,
|
||||
@Advice.This final SharedSessionContract session,
|
||||
@Advice.Origin("#m") final String name,
|
||||
@Advice.Argument(0) final Object entity) {
|
||||
|
||||
final ContextStore<Session, SessionState> contextStore =
|
||||
InstrumentationContext.get(Session.class, SessionState.class);
|
||||
final ContextStore<SharedSessionContract, SessionState> contextStore =
|
||||
InstrumentationContext.get(SharedSessionContract.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.This final SharedSessionContract session,
|
||||
@Advice.Enter final SessionState sessionState,
|
||||
@Advice.Thrown final Throwable throwable) {
|
||||
|
||||
|
@ -171,14 +155,14 @@ public class SessionInstrumentation extends Instrumenter.Default {
|
|||
@Advice.OnMethodExit(suppress = Throwable.class)
|
||||
public static void getQuery(
|
||||
@Advice.This final SharedSessionContract session,
|
||||
@Advice.Return(readOnly = false) final QueryImplementor query) {
|
||||
@Advice.Return final QueryImplementor query) {
|
||||
|
||||
if (!(query instanceof Query)) {
|
||||
return;
|
||||
}
|
||||
|
||||
final ContextStore<Session, SessionState> sessionContextStore =
|
||||
InstrumentationContext.get(Session.class, SessionState.class);
|
||||
final ContextStore<SharedSessionContract, SessionState> sessionContextStore =
|
||||
InstrumentationContext.get(SharedSessionContract.class, SessionState.class);
|
||||
final ContextStore<Query, SessionState> queryContextStore =
|
||||
InstrumentationContext.get(Query.class, SessionState.class);
|
||||
|
||||
|
@ -192,10 +176,10 @@ public class SessionInstrumentation extends Instrumenter.Default {
|
|||
@Advice.OnMethodExit(suppress = Throwable.class)
|
||||
public static void getTransaction(
|
||||
@Advice.This final SharedSessionContract session,
|
||||
@Advice.Return(readOnly = false) final Transaction transaction) {
|
||||
@Advice.Return final Transaction transaction) {
|
||||
|
||||
final ContextStore<Session, SessionState> sessionContextStore =
|
||||
InstrumentationContext.get(Session.class, SessionState.class);
|
||||
final ContextStore<SharedSessionContract, SessionState> sessionContextStore =
|
||||
InstrumentationContext.get(SharedSessionContract.class, SessionState.class);
|
||||
final ContextStore<Transaction, SessionState> transactionContextStore =
|
||||
InstrumentationContext.get(Transaction.class, SessionState.class);
|
||||
|
||||
|
|
|
@ -11,7 +11,6 @@ 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 {
|
||||
|
@ -40,6 +39,11 @@ public class SessionMethodUtils {
|
|||
|
||||
final SessionState sessionState = contextStore.get(spanKey);
|
||||
|
||||
if (sessionState == null) {
|
||||
// No state found. Maybe the instrumentation isn't working correctly.
|
||||
return null;
|
||||
}
|
||||
|
||||
if (sessionState.getMethodScope() != null) {
|
||||
// This method call was re-entrant. Do nothing, since it is being traced by the parent/first
|
||||
// call.
|
||||
|
@ -91,18 +95,18 @@ public class SessionMethodUtils {
|
|||
// 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, SessionState> sessionContextStore,
|
||||
final ContextStore<SharedSessionContract, SessionState> sessionContextStore,
|
||||
final SharedSessionContract session,
|
||||
final ContextStore<T, SessionState> targetContextStore,
|
||||
final T target) {
|
||||
|
||||
if (!(session instanceof Session)
|
||||
if (!(session instanceof SharedSessionContract)
|
||||
|| sessionContextStore == null
|
||||
|| targetContextStore == null) {
|
||||
return;
|
||||
}
|
||||
|
||||
final SessionState state = sessionContextStore.get((Session) session);
|
||||
final SessionState state = sessionContextStore.get((SharedSessionContract) session);
|
||||
if (state == null) {
|
||||
return;
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue