Fix Span Decoration probe evaluateAt=ENTRY

DataDog · Jun 25, 2024 · 3554ced · 3554ced
1 parent d19ceac
commit 3554ced
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 4 deletions.
diff --git a/...va-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerTransformer.java b/...va-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerTransformer.java
@@ -19,6 +19,7 @@
 import com.datadog.debugger.util.ExceptionHelper;
 import datadog.trace.agent.tooling.AgentStrategies;
 import datadog.trace.api.Config;
+import datadog.trace.bootstrap.debugger.MethodLocation;
 import datadog.trace.bootstrap.debugger.ProbeId;
 import datadog.trace.bootstrap.debugger.ProbeImplementation;
 import datadog.trace.util.Strings;
@@ -581,12 +582,19 @@ private List<ToInstrumentInfo> filterAndSortDefinitions(List<ProbeDefinition> de
   }
 
   // Log & Span Decoration probes share the same instrumentor so only one definition should be
-  // selected to
-  // generate the instrumentation. Log probes needs capture limits provided by the configuration
-  // so if the list of definition contains at least 1 log probe this is the log probe that need to
-  // be picked.
+  // selected to generate the instrumentation. Log probes needs capture limits provided by the
+  // configuration so if the list of definition contains at least 1 log probe this is the log probe
+  // that need to be picked.
   // TODO: handle the conflicting limits for log probes + mixing CaptureSnapshot or not
   private ProbeDefinition selectReferenceDefinition(List<ProbeDefinition> capturedContextProbes) {
+    ProbeDefinition atEntryProbe =
+        capturedContextProbes.stream()
+            .filter(probeDefinition -> probeDefinition.getEvaluateAt() == MethodLocation.ENTRY)
+            .findAny()
+            .orElse(null);
+    if (atEntryProbe != null) {
+      return atEntryProbe;
+    }
     ProbeDefinition first = capturedContextProbes.get(0);
     return capturedContextProbes.stream()
         .filter(it -> it instanceof LogProbe)

diff --git a/...va-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/SpanDecorationProbe.java b/...va-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/SpanDecorationProbe.java
@@ -324,6 +324,11 @@ public void addTag(String tagName, String tagValue) {
     public List<Pair<String, String>> getTagsToDecorate() {
       return tagsToDecorate;
     }
+
+    @Override
+    public boolean isCapturing() {
+      return true;
+    }
   }
 
   public static SpanDecorationProbe.Builder builder() {

diff --git a/...gger/src/test/java/com/datadog/debugger/agent/SpanDecorationProbeInstrumentationTest.java b/...gger/src/test/java/com/datadog/debugger/agent/SpanDecorationProbeInstrumentationTest.java
@@ -351,6 +351,36 @@ PROBE_ID, ACTIVE, singletonList(decoration), CLASS_NAME, "process", null, null)
     assertNotNull(intLocal);
   }
 
+  @Test
+  public void mixedEntryExit() throws IOException, URISyntaxException {
+    final String CLASS_NAME = "com.datadog.debugger.CapturedSnapshot20";
+    SpanDecorationProbe.Decoration decoration1 = createDecoration("tag1", "{intLocal}");
+    SpanDecorationProbe.Decoration decoration2 = createDecoration("tag2", "{arg}");
+    SpanDecorationProbe spanDecoProbe1 =
+        createProbeBuilder(
+                PROBE_ID1, ACTIVE, singletonList(decoration1), CLASS_NAME, "process", null, null)
+            .evaluateAt(MethodLocation.EXIT)
+            .build();
+    SpanDecorationProbe spanDecoProbe2 =
+        createProbeBuilder(
+                PROBE_ID2, ACTIVE, singletonList(decoration2), CLASS_NAME, "process", null, null)
+            .evaluateAt(MethodLocation.ENTRY)
+            .build();
+    Configuration configuration =
+        Configuration.builder()
+            .setService(SERVICE_NAME)
+            .add(spanDecoProbe1)
+            .add(spanDecoProbe2)
+            .build();
+    installSpanDecorationProbes(CLASS_NAME, configuration);
+    Class<?> testClass = compileAndLoadClass(CLASS_NAME);
+    int result = Reflect.on(testClass).call("main", "1").get();
+    assertEquals(84, result);
+    MutableSpan span = traceInterceptor.getFirstSpan();
+    assertEquals("84", span.getTags().get("tag1"));
+    assertEquals("1", span.getTags().get("tag2"));
+  }
+
   @Test
   public void keywordRedaction() throws IOException, URISyntaxException {
     final String CLASS_NAME = "com.datadog.debugger.CapturedSnapshot28";