Merge pull request #5809 from DataDog/jpbempel/span-deco-mixed-logs

Merge span decoration and log instrumentation
DataDog · Aug 30, 2023 · 6308c00 · 6308c00
2 parents 360dce2 + 6e5b6b7
commit 6308c00
Show file tree

Hide file tree

Showing 7 changed files with 114 additions and 94 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
@@ -7,6 +7,7 @@
 import com.datadog.debugger.instrumentation.InstrumentationResult;
 import com.datadog.debugger.probe.LogProbe;
 import com.datadog.debugger.probe.ProbeDefinition;
+import com.datadog.debugger.probe.SpanDecorationProbe;
 import com.datadog.debugger.probe.Where;
 import com.datadog.debugger.util.ExceptionHelper;
 import datadog.trace.agent.tooling.AgentStrategies;
@@ -477,21 +478,25 @@ private InstrumentationResult applyInstrumentation(
         preCheckInstrumentation(diagnostics, classLoader, methodNode);
     if (status != InstrumentationResult.Status.ERROR) {
       try {
-        List<ProbeDefinition> logProbes = new ArrayList<>();
+        List<ProbeDefinition> capturedContextProbes = new ArrayList<>();
         for (ProbeDefinition definition : definitions) {
-          if (definition instanceof LogProbe) {
-            logProbes.add(definition);
+          // Log and span decoration probe shared the same instrumentor: CaptureContextInstrumentor
+          // and therefore need to be instrumented once
+          if (definition instanceof LogProbe || definition instanceof SpanDecorationProbe) {
+            capturedContextProbes.add(definition);
           } else {
             List<DiagnosticMessage> probeDiagnostics = diagnostics.get(definition.getProbeId());
             definition.instrument(classLoader, classNode, methodNode, probeDiagnostics);
           }
         }
-        if (logProbes.size() > 0) {
-          List<String> probesIds = logProbes.stream().map(ProbeDefinition::getId).collect(toList());
-          List<DiagnosticMessage> probeDiagnostics = diagnostics.get(logProbes.get(0).getProbeId());
-          logProbes
-              .get(0)
-              .instrument(classLoader, classNode, methodNode, probeDiagnostics, probesIds);
+        if (capturedContextProbes.size() > 0) {
+          List<String> probesIds =
+              capturedContextProbes.stream().map(ProbeDefinition::getId).collect(toList());
+          ProbeDefinition referenceDefinition = selectReferenceDefinition(capturedContextProbes);
+          List<DiagnosticMessage> probeDiagnostics =
+              diagnostics.get(referenceDefinition.getProbeId());
+          referenceDefinition.instrument(
+              classLoader, classNode, methodNode, probeDiagnostics, probesIds);
         }
       } catch (Throwable t) {
         log.warn("Exception during instrumentation: ", t);
@@ -503,6 +508,20 @@ private InstrumentationResult applyInstrumentation(
     return new InstrumentationResult(status, diagnostics, classNode, methodNode);
   }
 
+  // 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.
+  // TODO: handle the conflicting limits for log probes + mixing CaptureSnapshot or not
+  private ProbeDefinition selectReferenceDefinition(List<ProbeDefinition> capturedContextProbes) {
+    ProbeDefinition first = capturedContextProbes.get(0);
+    return capturedContextProbes.stream()
+        .filter(it -> it instanceof LogProbe)
+        .findFirst()
+        .orElse(first);
+  }
+
   private InstrumentationResult.Status preCheckInstrumentation(
       Map<ProbeId, List<DiagnosticMessage>> diagnostics,
       ClassLoader classLoader,

diff --git a/...nt/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/LogInstrumentor.java b/...nt/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/LogInstrumentor.java
diff --git a/...bugger/src/main/java/com/datadog/debugger/instrumentation/SpanDecorationInstrumentor.java b/...bugger/src/main/java/com/datadog/debugger/instrumentation/SpanDecorationInstrumentor.java
diff --git a/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/LogProbe.java b/dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/probe/LogProbe.java
@@ -1,13 +1,15 @@
 package com.datadog.debugger.probe;
 
+import static com.datadog.debugger.probe.LogProbe.Capture.toLimits;
+
 import com.datadog.debugger.agent.DebuggerAgent;
 import com.datadog.debugger.agent.Generated;
 import com.datadog.debugger.agent.LogMessageTemplateBuilder;
 import com.datadog.debugger.el.EvaluationException;
 import com.datadog.debugger.el.ProbeCondition;
 import com.datadog.debugger.el.ValueScript;
+import com.datadog.debugger.instrumentation.CapturedContextInstrumentor;
 import com.datadog.debugger.instrumentation.DiagnosticMessage;
-import com.datadog.debugger.instrumentation.LogInstrumentor;
 import com.datadog.debugger.sink.Sink;
 import com.datadog.debugger.sink.Snapshot;
 import com.squareup.moshi.Json;
@@ -347,7 +349,15 @@ public void instrument(
       MethodNode methodNode,
       List<DiagnosticMessage> diagnostics,
       List<String> probeIds) {
-    new LogInstrumentor(this, classLoader, classNode, methodNode, diagnostics, probeIds)
+    new CapturedContextInstrumentor(
+            this,
+            classLoader,
+            classNode,
+            methodNode,
+            diagnostics,
+            probeIds,
+            isCaptureSnapshot(),
+            toLimits(getCapture()))
         .instrument();
   }
 

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
@@ -5,12 +5,13 @@
 import com.datadog.debugger.agent.LogMessageTemplateBuilder;
 import com.datadog.debugger.el.EvaluationException;
 import com.datadog.debugger.el.ProbeCondition;
+import com.datadog.debugger.instrumentation.CapturedContextInstrumentor;
 import com.datadog.debugger.instrumentation.DiagnosticMessage;
-import com.datadog.debugger.instrumentation.SpanDecorationInstrumentor;
 import com.datadog.debugger.sink.Snapshot;
 import datadog.trace.api.Pair;
 import datadog.trace.bootstrap.debugger.CapturedContext;
 import datadog.trace.bootstrap.debugger.EvaluationError;
+import datadog.trace.bootstrap.debugger.Limits;
 import datadog.trace.bootstrap.debugger.MethodLocation;
 import datadog.trace.bootstrap.debugger.ProbeId;
 import datadog.trace.bootstrap.debugger.ProbeImplementation;
@@ -137,7 +138,8 @@ public void instrument(
       MethodNode methodNode,
       List<DiagnosticMessage> diagnostics,
       List<String> probeIds) {
-    new SpanDecorationInstrumentor(this, classLoader, classNode, methodNode, diagnostics, probeIds)
+    new CapturedContextInstrumentor(
+            this, classLoader, classNode, methodNode, diagnostics, probeIds, false, Limits.DEFAULT)
         .instrument();
   }
 

diff --git a/...gger/src/test/java/com/datadog/debugger/agent/SpanDecorationProbeInstrumentationTest.java b/...gger/src/test/java/com/datadog/debugger/agent/SpanDecorationProbeInstrumentationTest.java
@@ -6,8 +6,10 @@
 import static com.datadog.debugger.probe.SpanDecorationProbe.TargetSpan.ACTIVE;
 import static com.datadog.debugger.probe.SpanDecorationProbe.TargetSpan.ROOT;
 import static com.datadog.debugger.util.LogProbeTestHelper.parseTemplate;
+import static java.util.Collections.singletonList;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -16,12 +18,14 @@
 import com.datadog.debugger.el.DSL;
 import com.datadog.debugger.el.ProbeCondition;
 import com.datadog.debugger.el.expressions.BooleanExpression;
+import com.datadog.debugger.probe.LogProbe;
 import com.datadog.debugger.probe.SpanDecorationProbe;
 import com.datadog.debugger.sink.Snapshot;
 import datadog.trace.agent.tooling.TracerInstaller;
 import datadog.trace.api.Config;
 import datadog.trace.api.interceptor.MutableSpan;
 import datadog.trace.api.interceptor.TraceInterceptor;
+import datadog.trace.bootstrap.debugger.CapturedContext;
 import datadog.trace.bootstrap.debugger.DebuggerContext;
 import datadog.trace.bootstrap.debugger.MethodLocation;
 import datadog.trace.bootstrap.debugger.ProbeId;
@@ -41,6 +45,8 @@
 public class SpanDecorationProbeInstrumentationTest extends ProbeInstrumentationTest {
   private static final String LANGUAGE = "java";
   private static final ProbeId PROBE_ID = new ProbeId("beae1807-f3b0-4ea8-a74f-826790c5e6f8", 0);
+  private static final ProbeId PROBE_ID1 = new ProbeId("beae1807-f3b0-4ea8-a74f-826790c5e6f6", 0);
+  private static final ProbeId PROBE_ID2 = new ProbeId("beae1807-f3b0-4ea8-a74f-826790c5e6f7", 0);
 
   private TestTraceInterceptor traceInterceptor = new TestTraceInterceptor();
 
@@ -59,7 +65,7 @@ public void methodActiveSpanSimpleTag() throws IOException, URISyntaxException {
         CLASS_NAME, ACTIVE, decoration, "process", "int (java.lang.String)");
     Class<?> testClass = compileAndLoadClass(CLASS_NAME);
     int result = Reflect.on(testClass).call("main", "1").get();
-    assertEquals(42, result);
+    assertEquals(84, result);
     MutableSpan span = traceInterceptor.getFirstSpan();
     assertEquals("1", span.getTags().get("tag1"));
     assertEquals(PROBE_ID.getId(), span.getTags().get("_dd.di.tag1.probe_id"));
@@ -81,7 +87,7 @@ public void methodActiveSpanTagList() throws IOException, URISyntaxException {
         "int (java.lang.String)");
     Class<?> testClass = compileAndLoadClass(CLASS_NAME);
     int result = Reflect.on(testClass).call("main", "1").get();
-    assertEquals(42, result);
+    assertEquals(84, result);
     MutableSpan span = traceInterceptor.getFirstSpan();
     assertEquals("1", span.getTags().get("tag1"));
     assertEquals("42", span.getTags().get("tag2"));
@@ -124,7 +130,7 @@ public void methodActiveSpanCondition() throws IOException, URISyntaxException {
     Class<?> testClass = compileAndLoadClass(CLASS_NAME);
     for (int i = 0; i < 10; i++) {
       int result = Reflect.on(testClass).call("main", String.valueOf(i)).get();
-      assertEquals(42, result);
+      assertEquals(84, result);
     }
     assertEquals(10, traceInterceptor.getAllTraces().size());
     MutableSpan span = traceInterceptor.getAllTraces().get(5).get(0);
@@ -139,7 +145,7 @@ public void methodTagEvalError() throws IOException, URISyntaxException {
         CLASS_NAME, ACTIVE, decoration, "process", "int (java.lang.String)");
     Class<?> testClass = compileAndLoadClass(CLASS_NAME);
     int result = Reflect.on(testClass).call("main", "1").get();
-    assertEquals(42, result);
+    assertEquals(84, result);
     MutableSpan span = traceInterceptor.getFirstSpan();
     assertNull(span.getTags().get("tag1"));
     assertEquals("Cannot find symbol: noarg", span.getTags().get("_dd.di.tag1.evaluation_error"));
@@ -154,7 +160,7 @@ public void methodActiveSpanInvalidCondition() throws IOException, URISyntaxExce
         CLASS_NAME, ACTIVE, decoration, "process", "int (java.lang.String)");
     Class<?> testClass = compileAndLoadClass(CLASS_NAME);
     int result = Reflect.on(testClass).call("main", "5").get();
-    assertEquals(42, result);
+    assertEquals(84, result);
     assertFalse(traceInterceptor.getFirstSpan().getTags().containsKey("tag1"));
     assertEquals(1, mockSink.getSnapshots().size());
     Snapshot snapshot = mockSink.getSnapshots().get(0);
@@ -166,10 +172,10 @@ public void methodActiveSpanInvalidCondition() throws IOException, URISyntaxExce
   public void lineActiveSpanSimpleTag() throws IOException, URISyntaxException {
     final String CLASS_NAME = "com.datadog.debugger.CapturedSnapshot20";
     SpanDecorationProbe.Decoration decoration = createDecoration("tag1", "{arg}");
-    installSingleSpanDecoration(CLASS_NAME, ACTIVE, decoration, "CapturedSnapshot20.java", 37);
+    installSingleSpanDecoration(CLASS_NAME, ACTIVE, decoration, "CapturedSnapshot20.java", 38);
     Class<?> testClass = compileAndLoadClass(CLASS_NAME);
     int result = Reflect.on(testClass).call("main", "1").get();
-    assertEquals(42, result);
+    assertEquals(84, result);
     MutableSpan span = traceInterceptor.getFirstSpan();
     assertEquals("1", span.getTags().get("tag1"));
     assertEquals(PROBE_ID.getId(), span.getTags().get("_dd.di.tag1.probe_id"));
@@ -198,11 +204,11 @@ public void lineActiveSpanCondition() throws IOException, URISyntaxException {
     final String CLASS_NAME = "com.datadog.debugger.CapturedSnapshot20";
     SpanDecorationProbe.Decoration decoration =
         createDecoration(eq(ref("arg"), value("5")), "arg == '5'", "tag1", "{arg}");
-    installSingleSpanDecoration(CLASS_NAME, ACTIVE, decoration, "CapturedSnapshot20.java", 37);
+    installSingleSpanDecoration(CLASS_NAME, ACTIVE, decoration, "CapturedSnapshot20.java", 38);
     Class<?> testClass = compileAndLoadClass(CLASS_NAME);
     for (int i = 0; i < 10; i++) {
       int result = Reflect.on(testClass).call("main", String.valueOf(i)).get();
-      assertEquals(42, result);
+      assertEquals(84, result);
     }
     assertEquals(10, traceInterceptor.getAllTraces().size());
     MutableSpan span = traceInterceptor.getAllTraces().get(5).get(0);
@@ -214,10 +220,10 @@ public void lineActiveSpanInvalidCondition() throws IOException, URISyntaxExcept
     final String CLASS_NAME = "com.datadog.debugger.CapturedSnapshot20";
     SpanDecorationProbe.Decoration decoration =
         createDecoration(eq(ref("noarg"), value("5")), "arg == '5'", "tag1", "{arg}");
-    installSingleSpanDecoration(CLASS_NAME, ACTIVE, decoration, "CapturedSnapshot20.java", 37);
+    installSingleSpanDecoration(CLASS_NAME, ACTIVE, decoration, "CapturedSnapshot20.java", 38);
     Class<?> testClass = compileAndLoadClass(CLASS_NAME);
     int result = Reflect.on(testClass).call("main", "5").get();
-    assertEquals(42, result);
+    assertEquals(84, result);
     assertFalse(traceInterceptor.getFirstSpan().getTags().containsKey("tag1"));
     assertEquals(1, mockSink.getSnapshots().size());
     Snapshot snapshot = mockSink.getSnapshots().get(0);
@@ -236,6 +242,48 @@ public void nullActiveSpan() throws IOException, URISyntaxException {
     assertEquals(0, traceInterceptor.getAllTraces().size());
   }
 
+  @Test
+  public void mixedWithLogProbes() throws IOException, URISyntaxException {
+    final String CLASS_NAME = "com.datadog.debugger.CapturedSnapshot20";
+    SpanDecorationProbe.Decoration decoration = createDecoration("tag1", "{intLocal}");
+    SpanDecorationProbe spanDecoProbe =
+        createProbeBuilder(
+                PROBE_ID, ACTIVE, singletonList(decoration), CLASS_NAME, "process", null, null)
+            .evaluateAt(MethodLocation.EXIT)
+            .build();
+    LogProbe logProbe1 =
+        LogProbe.builder()
+            .probeId(PROBE_ID1)
+            .where(CLASS_NAME, "process")
+            .captureSnapshot(true)
+            .build();
+    LogProbe logProbe2 =
+        LogProbe.builder()
+            .probeId(PROBE_ID2)
+            .where(CLASS_NAME, "process")
+            .captureSnapshot(true)
+            .build();
+    Configuration configuration =
+        Configuration.builder()
+            .setService(SERVICE_NAME)
+            .add(logProbe1)
+            .add(logProbe2)
+            .add(spanDecoProbe)
+            .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(PROBE_ID.getId(), span.getTags().get("_dd.di.tag1.probe_id"));
+    List<Snapshot> snapshots = mockSink.getSnapshots();
+    assertEquals(2, snapshots.size());
+    CapturedContext.CapturedValue intLocal =
+        snapshots.get(0).getCaptures().getReturn().getLocals().get("intLocal");
+    assertNotNull(intLocal);
+  }
+
   private SpanDecorationProbe.Decoration createDecoration(String tagName, String valueDsl) {
     List<SpanDecorationProbe.Tag> tags =
         Arrays.asList(
@@ -305,6 +353,7 @@ private static SpanDecorationProbe createProbe(
       String... lines) {
     return createProbeBuilder(
             id, targetSpan, decorationList, typeName, methodName, signature, lines)
+        .evaluateAt(MethodLocation.EXIT)
         .build();
   }
 
@@ -367,19 +416,19 @@ private void installSpanDecorationProbes(String expectedClassName, Configuration
     mockSink = new MockSink();
     DebuggerAgentHelper.injectSink(mockSink);
     DebuggerContext.init(
-        (id, callingClass) ->
-            resolver(id, callingClass, expectedClassName, configuration.getSpanDecorationProbes()),
-        null);
+        (id, callingClass) -> resolver(id, callingClass, expectedClassName, configuration), null);
     DebuggerContext.initClassFilter(new DenyListHelper(null));
   }
 
   private ProbeImplementation resolver(
-      String id,
-      Class<?> callingClass,
-      String expectedClassName,
-      Collection<SpanDecorationProbe> spanDecorationProbes) {
+      String id, Class<?> callingClass, String expectedClassName, Configuration configuration) {
     Assertions.assertEquals(expectedClassName, callingClass.getName());
-    for (SpanDecorationProbe probe : spanDecorationProbes) {
+    for (SpanDecorationProbe probe : configuration.getSpanDecorationProbes()) {
+      if (probe.getId().equals(id)) {
+        return probe;
+      }
+    }
+    for (LogProbe probe : configuration.getLogProbes()) {
       if (probe.getId().equals(id)) {
         return probe;
       }

diff --git a/dd-java-agent/agent-debugger/src/test/resources/com/datadog/debugger/CapturedSnapshot20.java b/dd-java-agent/agent-debugger/src/test/resources/com/datadog/debugger/CapturedSnapshot20.java
@@ -34,6 +34,7 @@ public static int main(String arg) {
   }
 
   private int process(String arg) {
-    return intField;
+    int intLocal = intField + 42;
+    return intLocal;
   }
 }