Move SessionFactory instrumentation into its own class
ADD
dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/SessionFactoryInstrumentation.java
+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 static net.bytebuddy.matcher.ElementMatchers.returns;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.api.DDSpanTypes;
+import datadog.trace.api.DDTags;
+import datadog.trace.bootstrap.ContextStore;
+import datadog.trace.bootstrap.InstrumentationContext;
+import io.opentracing.Span;
+import io.opentracing.tag.Tags;
+import io.opentracing.util.GlobalTracer;
+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.Session;
+
+@AutoService(Instrumenter.class)
+public class SessionFactoryInstrumentation extends Instrumenter.Default {
+
+ public SessionFactoryInstrumentation() {
+ super("hibernate");
+ }
+
+ @Override
+ public Map<String, String> contextStore() {
+ final Map<String, String> map = new HashMap<>();
+ map.put("org.hibernate.Session", SessionState.class.getName());
+ return Collections.unmodifiableMap(map);
+ }
+
+ @Override
+ public String[] helperClassNames() {
+ return new String[] {
+ "datadog.trace.instrumentation.hibernate.SessionState",
+ };
+ }
+
+ @Override
+ public ElementMatcher<TypeDescription> typeMatcher() {
+ return not(isInterface()).and(safeHasSuperType(named("org.hibernate.SessionFactory")));
+ }
+
+ @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"))
+ .and(takesArguments(0))
+ .and(returns(named("org.hibernate.Session"))),
+ SessionFactoryAdvice.class.getName());
+
+ return transformers;
+ }
+
+ public static class SessionFactoryAdvice {
+
+ @Advice.OnMethodExit(suppress = Throwable.class)
+ public static void openSession(@Advice.Return(readOnly = false) final Session session) {
+
+ final Span span =
+ GlobalTracer.get()
+ .buildSpan("hibernate.session")
+ .withTag(DDTags.SERVICE_NAME, "hibernate")
+ .withTag(DDTags.SPAN_TYPE, DDSpanTypes.HIBERNATE)
+ .withTag(Tags.COMPONENT.getKey(), "hibernate-java")
+ .withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CLIENT)
+ .start();
+
+ final ContextStore<Session, SessionState> contextStore =
+ InstrumentationContext.get(Session.class, SessionState.class);
+ contextStore.putIfAbsent(session, new SessionState(span));
+ }
+ }
+}

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
@@ -14,13 +14,10 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
import com.google.auto.service.AutoService;
import datadog.trace.agent.tooling.Instrumenter;
-import datadog.trace.api.DDSpanTypes;
-import datadog.trace.api.DDTags;
import datadog.trace.bootstrap.ContextStore;
import datadog.trace.bootstrap.InstrumentationContext;
import io.opentracing.Span;
import io.opentracing.tag.Tags;
-import io.opentracing.util.GlobalTracer;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
@@ -31,7 +28,6 @@ import net.bytebuddy.matcher.ElementMatcher;
import org.hibernate.Session;
import org.hibernate.SharedSessionContract;
import org.hibernate.Transaction;
-import org.hibernate.internal.AbstractSharedSessionContract;
import org.hibernate.query.Query;
import org.hibernate.query.spi.QueryImplementor;
@@ -64,23 +60,13 @@ public class SessionInstrumentation extends Instrumenter.Default {
return not(isInterface())
.and(
safeHasSuperType(
- named("org.hibernate.SessionFactory")
- .or(named("org.hibernate.Session"))
+ named("org.hibernate.Session")
.or(named("org.hibernate.internal.AbstractSharedSessionContract"))));
}
@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"))
- .and(isDeclaredBy(safeHasSuperType(named("org.hibernate.SessionFactory"))))
- .and(takesArguments(0))
- .and(returns(named("org.hibernate.Session"))),
- SessionFactoryAdvice.class.getName());
transformers.put(
isMethod()
.and(named("close"))
@@ -131,26 +117,6 @@ public class SessionInstrumentation extends Instrumenter.Default {
return transformers;
}
- public static class SessionFactoryAdvice {
-
- @Advice.OnMethodExit(suppress = Throwable.class)
- public static void openSession(@Advice.Return(readOnly = false) final Session session) {
-
- final Span span =
- GlobalTracer.get()
- .buildSpan("hibernate.session")
- .withTag(DDTags.SERVICE_NAME, "hibernate")
- .withTag(DDTags.SPAN_TYPE, DDSpanTypes.HIBERNATE)
- .withTag(Tags.COMPONENT.getKey(), "hibernate-java")
- .withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CLIENT)
- .start();
-
- final ContextStore<Session, SessionState> contextStore =
- InstrumentationContext.get(Session.class, SessionState.class);
- contextStore.putIfAbsent(session, new SessionState(span));
- }
- }
-
public static class SessionCloseAdvice {
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
@@ -225,7 +191,7 @@ public class SessionInstrumentation extends Instrumenter.Default {
@Advice.OnMethodExit(suppress = Throwable.class)
public static void getTransaction(
- @Advice.This final AbstractSharedSessionContract session,
+ @Advice.This final SharedSessionContract session,
@Advice.Return(readOnly = false) final Transaction transaction) {
final ContextStore<Session, SessionState> sessionContextStore =

Move SessionFactory instrumentation into its own class