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

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 =
+      final ContextStore<Query, SessionState> contextStore =
-          InstrumentationContext.get(Query.class, Span.class);
+          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);
     }
   }
 }

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

@@ -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.Session", SessionState.class.getName());
-    map.put("org.hibernate.query.Query", Span.class.getName());
+    map.put("org.hibernate.query.Query", SessionState.class.getName());
-    map.put("org.hibernate.Transaction", Span.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)),
+            .and(returns(named("java.lang.Object")))
-        SessionSaveAdvice.class.getName());
+            .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(isDeclaredBy(safeHasSuperType(named("org.hibernate.SharedSessionContract")))),
-            .and(takesArguments(1)),
         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 =
+      final ContextStore<Session, SessionState> contextStore =
-          InstrumentationContext.get(Session.class, Span.class);
+          InstrumentationContext.get(Session.class, SessionState.class);
-      contextStore.putIfAbsent(session, span);
+      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 =
+      final ContextStore<Session, SessionState> contextStore =
-          InstrumentationContext.get(Session.class, Span.class);
+          InstrumentationContext.get(Session.class, SessionState.class);
-      final Span span = contextStore.get(session);
+      final SessionState state = contextStore.get(session);
-
+      if (state == null || state.getSessionSpan() == null) {
-      if (span == 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(
+    public static SessionState startSave(
-        @Advice.This final Session session, @Advice.Argument(0) final Object entity) {
+        @Advice.This final Session session,
+        @Advice.Origin("#m") final String name,
+        @Advice.Argument(0) final Object entity) {
 
-      final ContextStore<Session, Span> contextStore =
+      final ContextStore<Session, SessionState> contextStore =
-          InstrumentationContext.get(Session.class, Span.class);
+          InstrumentationContext.get(Session.class, SessionState.class);
-      return SessionMethodUtils.startScopeFrom(contextStore, session, "hibernate.save");
+      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 =
+      final ContextStore<Session, SessionState> sessionContextStore =
-          InstrumentationContext.get(Session.class, Span.class);
+          InstrumentationContext.get(Session.class, SessionState.class);
-      final ContextStore<Query, Span> queryContextStore =
+      final ContextStore<Query, SessionState> queryContextStore =
-          InstrumentationContext.get(Query.class, Span.class);
+          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 =
+      final ContextStore<Session, SessionState> sessionContextStore =
-          InstrumentationContext.get(Session.class, Span.class);
+          InstrumentationContext.get(Session.class, SessionState.class);
-      final ContextStore<Transaction, Span> transactionContextStore =
+      final ContextStore<Transaction, SessionState> transactionContextStore =
-          InstrumentationContext.get(Transaction.class, Span.class);
+          InstrumentationContext.get(Transaction.class, SessionState.class);
 
       SessionMethodUtils.attachSpanFromSession(
           sessionContextStore, session, transactionContextStore, transaction);

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(
+  public static <T> SessionState startScopeFrom(
-      final ContextStore<T, Span> contextStore, final T spanKey, final String operationName) {
+      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);
+    final SessionState state = sessionContextStore.get((Session) session);
-    if (span == null) {
+    if (state == null) {
       return;
     }
 
-    targetContextStore.putIfAbsent(target, span);
+    targetContextStore.putIfAbsent(target, state);
   }
 }

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

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

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 =
+      final ContextStore<Transaction, SessionState> contextStore =
-          InstrumentationContext.get(Transaction.class, Span.class);
+          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);
     }
   }
 }

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,8 +92,20 @@ 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 {
@@ -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"))
+    }
+  }
 }