fixes for exit spans

Author: evanchooly

dd-java-agent/agent-debugger/debugger-bootstrap/src/main/java/datadog/trace/bootstrap/debugger/spanorigin/CodeOriginInfo.java

@@ -4,6 +4,7 @@
 import static java.util.Arrays.stream;
 
 import datadog.trace.api.InstrumenterConfig;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import java.lang.reflect.Method;
 import java.util.stream.Collectors;
 
@@ -18,9 +19,12 @@ public static void entry(Method method) {
     }
   }
 
-  public static void exit() {
+  public static void exit(AgentSpan span) {
     if (InstrumenterConfig.get().isCodeOriginEnabled()) {
-      captureCodeOrigin(null);
+      String probeId = captureCodeOrigin(null);
+      if (span != null) {
+        span.getLocalRootSpan().setTag(probeId, span);
+      }
     }
   }
 }

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/codeorigin/CodeOriginProbeManager.java

@@ -88,9 +88,6 @@ public String createProbeForFrame(String signature) {
 
     } else {
       probe = fingerprints.get(fingerprint);
-      if (span != null && !probe.isEntrySpanProbe()) {
-        span.getLocalRootSpan().setTag(probe.getId(), span);
-      }
     }
 
     return probe.getId();

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/CodeOriginProbe.java

@@ -69,7 +69,7 @@ public void commit(
       CapturedContext exitContext,
       List<CapturedContext.CapturedThrowable> caughtExceptions) {
 
-    AgentSpan span = AgentTracer.activeSpan();
+    AgentSpan span = findSpan(AgentTracer.activeSpan());
 
     if (span == null) {
       LOGGER.debug("Could not find the span for probeId {}", id);
@@ -89,9 +89,7 @@ public void commit(
     span.getLocalRootSpan().setTag(getId(), (String) null); // clear possible span reference
   }
 
-  private void applySpanOriginTags(AgentSpan candidate, String snapshotId) {
-    AgentSpan span = findSpan(candidate);
-
+  private void applySpanOriginTags(AgentSpan span, String snapshotId) {
     List<StackTraceElement> entries = getUserStackFrames();
     recordCodeOrigin(span, entries, snapshotId);
     recordStackFrames(span, entries, snapshotId);
@@ -117,9 +115,11 @@ private void recordCodeOrigin(
 
   private void recordStackFrames(
       AgentSpan span, List<StackTraceElement> entries, String snapshotId) {
-    for (AgentSpan s :
-        entrySpanProbe ? asList(span, span.getLocalRootSpan()) : singletonList(span)) {
-      s.setTag(format(DD_STACK_CODE_ORIGIN_TYPE), entrySpanProbe ? "entry" : "exit");
+    List<AgentSpan> agentSpans =
+        entrySpanProbe ? asList(span, span.getLocalRootSpan()) : singletonList(span);
+
+    for (AgentSpan s : agentSpans) {
+      s.setTag(DD_STACK_CODE_ORIGIN_TYPE, entrySpanProbe ? "entry" : "exit");
 
       for (int i = 0; i < entries.size(); i++) {
         StackTraceElement info = entries.get(i);

dd-java-agent/agent-debugger/src/test/resources/com/datadog/debugger/CodeOrigin01.java

@@ -1,11 +1,9 @@
 package com.datadog.debugger;
 
-import datadog.trace.api.Trace;
 import datadog.trace.bootstrap.debugger.spanorigin.CodeOriginInfo;
 import datadog.trace.bootstrap.instrumentation.api.AgentScope;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
-import datadog.trace.bootstrap.instrumentation.api.AgentTracer.SpanBuilder;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer.TracerAPI;
 import datadog.trace.bootstrap.instrumentation.api.ScopeSource;
 import datadog.trace.bootstrap.instrumentation.api.Tags;
@@ -62,7 +60,7 @@ public void entry() throws NoSuchMethodException {
   private int exit() {
     AgentSpan span = newSpan("exit");
     try(AgentScope scope = tracerAPI.activateSpan(span, ScopeSource.MANUAL)) {
-      CodeOriginInfo.exit();
+      CodeOriginInfo.exit(span);
       return 42;
     } finally {
       span.finish();

dd-java-agent/instrumentation/spring-data-1.8/src/main/java/datadog/trace/instrumentation/springdata/SpringDataDecorator.java

@@ -46,6 +46,6 @@ public void onOperation(
       span.setResourceName(spanNameForMethod(method));
     }
 
-    CodeOriginInfo.exit();
+    CodeOriginInfo.exit(span);
   }
 }