Merge pull request #1446 from DataDog/tyler/migrate-hibernate

DataDog · May 12, 2020 · 7be96e3 · 7be96e3
2 parents 98e7991 + 3e2cc28
commit 7be96e3
Show file tree

Hide file tree

Showing 10 changed files with 24 additions and 18 deletions.
diff --git a/.../main/java/datadog/trace/instrumentation/hibernate/core/v3_3/CriteriaInstrumentation.java b/.../main/java/datadog/trace/instrumentation/hibernate/core/v3_3/CriteriaInstrumentation.java
@@ -58,7 +58,7 @@ public static void endMethod(
         @Advice.Thrown final Throwable throwable,
         @Advice.Return(typing = Assigner.Typing.DYNAMIC) final Object entity) {
 
-      SessionMethodUtils.closeScope(state, throwable, entity);
+      SessionMethodUtils.closeScope(state, throwable, entity, true);
     }
   }
 }
diff --git a/...src/main/java/datadog/trace/instrumentation/hibernate/core/v3_3/QueryInstrumentation.java b/...src/main/java/datadog/trace/instrumentation/hibernate/core/v3_3/QueryInstrumentation.java
@@ -79,7 +79,7 @@ public static void endMethod(
         entity = query.getQueryString();
       }
 
-      SessionMethodUtils.closeScope(state, throwable, entity);
+      SessionMethodUtils.closeScope(state, throwable, entity, true);
     }
   }
 }
diff --git a/...c/main/java/datadog/trace/instrumentation/hibernate/core/v3_3/SessionInstrumentation.java b/...c/main/java/datadog/trace/instrumentation/hibernate/core/v3_3/SessionInstrumentation.java
@@ -143,9 +143,10 @@ public static class SessionMethodAdvice extends V3Advice {
     public static SessionState startMethod(
         @Advice.This final Object session,
         @Advice.Origin("#m") final String name,
-        @Advice.Argument(0) final Object entity) {
+        @Advice.Argument(0) final Object entity,
+        @Advice.Local("startSpan") boolean startSpan) {
 
-      final boolean startSpan = !SCOPE_ONLY_METHODS.contains(name);
+      startSpan = !SCOPE_ONLY_METHODS.contains(name);
       if (session instanceof Session) {
         final ContextStore<Session, SessionState> contextStore =
             InstrumentationContext.get(Session.class, SessionState.class);
@@ -164,9 +165,10 @@ public static SessionState startMethod(
     public static void endMethod(
         @Advice.Enter final SessionState sessionState,
         @Advice.Thrown final Throwable throwable,
+        @Advice.Local("startSpan") final boolean startSpan,
         @Advice.Return(typing = Assigner.Typing.DYNAMIC) final Object returned) {
 
-      SessionMethodUtils.closeScope(sessionState, throwable, returned);
+      SessionMethodUtils.closeScope(sessionState, throwable, returned, startSpan);
     }
   }
 

diff --git a/...in/java/datadog/trace/instrumentation/hibernate/core/v3_3/TransactionInstrumentation.java b/...in/java/datadog/trace/instrumentation/hibernate/core/v3_3/TransactionInstrumentation.java
@@ -57,7 +57,7 @@ public static void endCommit(
         @Advice.Enter final SessionState state,
         @Advice.Thrown final Throwable throwable) {
 
-      SessionMethodUtils.closeScope(state, throwable, null);
+      SessionMethodUtils.closeScope(state, throwable, null, true);
     }
   }
 }
diff --git a/.../main/java/datadog/trace/instrumentation/hibernate/core/v4_0/CriteriaInstrumentation.java b/.../main/java/datadog/trace/instrumentation/hibernate/core/v4_0/CriteriaInstrumentation.java
@@ -58,7 +58,7 @@ public static void endMethod(
         @Advice.Thrown final Throwable throwable,
         @Advice.Return(typing = Assigner.Typing.DYNAMIC) final Object entity) {
 
-      SessionMethodUtils.closeScope(state, throwable, entity);
+      SessionMethodUtils.closeScope(state, throwable, entity, true);
     }
   }
 }
diff --git a/...src/main/java/datadog/trace/instrumentation/hibernate/core/v4_0/QueryInstrumentation.java b/...src/main/java/datadog/trace/instrumentation/hibernate/core/v4_0/QueryInstrumentation.java
@@ -79,7 +79,7 @@ public static void endMethod(
         entity = query.getQueryString();
       }
 
-      SessionMethodUtils.closeScope(state, throwable, entity);
+      SessionMethodUtils.closeScope(state, throwable, entity, true);
     }
   }
 }
diff --git a/...c/main/java/datadog/trace/instrumentation/hibernate/core/v4_0/SessionInstrumentation.java b/...c/main/java/datadog/trace/instrumentation/hibernate/core/v4_0/SessionInstrumentation.java
@@ -132,9 +132,10 @@ public static class SessionMethodAdvice extends V4Advice {
     public static SessionState startMethod(
         @Advice.This final SharedSessionContract session,
         @Advice.Origin("#m") final String name,
-        @Advice.Argument(0) final Object entity) {
+        @Advice.Argument(0) final Object entity,
+        @Advice.Local("startSpan") boolean startSpan) {
 
-      final boolean startSpan = !SCOPE_ONLY_METHODS.contains(name);
+      startSpan = !SCOPE_ONLY_METHODS.contains(name);
       final ContextStore<SharedSessionContract, SessionState> contextStore =
           InstrumentationContext.get(SharedSessionContract.class, SessionState.class);
       return SessionMethodUtils.startScopeFrom(
@@ -143,12 +144,12 @@ public static SessionState startMethod(
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void endMethod(
-        @Advice.This final SharedSessionContract session,
         @Advice.Enter final SessionState sessionState,
         @Advice.Thrown final Throwable throwable,
+        @Advice.Local("startSpan") final boolean startSpan,
         @Advice.Return(typing = Assigner.Typing.DYNAMIC) final Object returned) {
 
-      SessionMethodUtils.closeScope(sessionState, throwable, returned);
+      SessionMethodUtils.closeScope(sessionState, throwable, returned, startSpan);
     }
   }
 

diff --git a/...in/java/datadog/trace/instrumentation/hibernate/core/v4_0/TransactionInstrumentation.java b/...in/java/datadog/trace/instrumentation/hibernate/core/v4_0/TransactionInstrumentation.java
@@ -57,7 +57,7 @@ public static void endCommit(
         @Advice.Enter final SessionState state,
         @Advice.Thrown final Throwable throwable) {
 
-      SessionMethodUtils.closeScope(state, throwable, null);
+      SessionMethodUtils.closeScope(state, throwable, null, true);
     }
   }
 }
diff --git a/.../java/datadog/trace/instrumentation/hibernate/core/v4_3/ProcedureCallInstrumentation.java b/.../java/datadog/trace/instrumentation/hibernate/core/v4_3/ProcedureCallInstrumentation.java
@@ -76,7 +76,7 @@ public static SessionState startMethod(
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void endMethod(
         @Advice.Enter final SessionState state, @Advice.Thrown final Throwable throwable) {
-      SessionMethodUtils.closeScope(state, throwable, null);
+      SessionMethodUtils.closeScope(state, throwable, null, true);
     }
   }
 }
diff --git a/...n/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/SessionMethodUtils.java b/...n/hibernate/src/main/java/datadog/trace/instrumentation/hibernate/SessionMethodUtils.java
@@ -42,9 +42,9 @@ public static <TARGET, ENTITY> SessionState startScopeFrom(
       final AgentSpan span = startSpan(operationName, sessionState.getSessionSpan().context());
       DECORATOR.afterStart(span);
       DECORATOR.onOperation(span, entity);
-      scope = activateSpan(span, true);
+      scope = activateSpan(span);
     } else {
-      scope = activateSpan(sessionState.getSessionSpan(), false);
+      scope = activateSpan(sessionState.getSessionSpan());
       sessionState.setHasChildSpan(false);
     }
 
@@ -54,7 +54,10 @@ public static <TARGET, ENTITY> SessionState startScopeFrom(
 
   // Closes a Scope/Span, adding an error tag if the given Throwable is not null.
   public static void closeScope(
-      final SessionState sessionState, final Throwable throwable, final Object entity) {
+      final SessionState sessionState,
+      final Throwable throwable,
+      final Object entity,
+      final boolean closeSpan) {
 
     if (sessionState == null || sessionState.getMethodScope() == null) {
       // This method call was re-entrant. Do nothing, since it is being traced by the parent/first
@@ -65,7 +68,7 @@ public static void closeScope(
     CallDepthThreadLocalMap.reset(SessionMethodUtils.class);
     final AgentScope scope = sessionState.getMethodScope();
     final AgentSpan span = scope.span();
-    if (span != null && sessionState.hasChildSpan) {
+    if (span != null && (sessionState.hasChildSpan || closeSpan)) {
       DECORATOR.onError(span, throwable);
       if (entity != null) {
         DECORATOR.onOperation(span, entity);