Move SessionFactory instrumentation into its own class

dd-java-agent/instrumentation/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/SessionFactoryInstrumentation.java

@@ -0,0 +1,89 @@
+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));
+    }
+  }
+}

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")
+                named("org.hibernate.Session")
-                    .or(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 =