Merge pull request #7286 from DataDog/evanchooly/DEBUG-2434

Implement debug context propagation to enable live debugging of java applications

Author: evanchooly

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/span_origin/EntrySpanOriginInfo.java

@@ -67,12 +67,17 @@ public interface ExceptionDebugger {
     void handleException(Throwable t, AgentSpan span);
   }
 
+  public interface SpanDebugger {
+    String captureSnapshot(String signature);
+  }
+
   private static volatile ProbeResolver probeResolver;
   private static volatile ClassFilter classFilter;
   private static volatile MetricForwarder metricForwarder;
   private static volatile Tracer tracer;
   private static volatile ValueSerializer valueSerializer;
   private static volatile ExceptionDebugger exceptionDebugger;
+  private static volatile SpanDebugger spanDebugger;
 
   public static void initProbeResolver(ProbeResolver probeResolver) {
     DebuggerContext.probeResolver = probeResolver;
@@ -98,6 +103,10 @@ public static void initExceptionDebugger(ExceptionDebugger exceptionDebugger) {
     DebuggerContext.exceptionDebugger = exceptionDebugger;
   }
 
+  public static void initSpanDebugger(SpanDebugger spanDebugger) {
+    DebuggerContext.spanDebugger = spanDebugger;
+  }
+
   /**
    * Returns the probe details based on the probe id provided. If no probe is found, try to
    * re-transform the class using the callingClass parameter No-op if no implementation available
@@ -333,6 +342,18 @@ public static void commit(
     }
   }
 
+  public static String captureSnapshot(String signature) {
+    try {
+      SpanDebugger debugger = spanDebugger;
+      if (debugger != null) {
+        return debugger.captureSnapshot(signature);
+      }
+    } catch (Exception ex) {
+      LOGGER.debug("Error in addSnapshot: ", ex);
+    }
+    return null;
+  }
+
   public static void handleException(Throwable t, AgentSpan span) {
     try {
       ExceptionDebugger exDebugger = exceptionDebugger;

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/span_origin/ExitSpanOriginInfo.java

@@ -0,0 +1,27 @@
+package datadog.trace.bootstrap.debugger.spanorigin;
+
+import static datadog.trace.bootstrap.debugger.DebuggerContext.captureSnapshot;
+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;
+
+public class SpanOriginInfo {
+  public static void entry(AgentSpan span, Method method) {
+    if (InstrumenterConfig.get().isSpanOriginEnabled()) {
+      String signature =
+          stream(method.getParameterTypes())
+              .map(Class::getName)
+              .collect(Collectors.joining(",", "(", ")"));
+      captureSnapshot(signature);
+    }
+  }
+
+  public static void exit(AgentSpan span) {
+    if (InstrumenterConfig.get().isSpanOriginEnabled()) {
+      span.getLocalRootSpan().setMetaStruct(captureSnapshot(null), span);
+    }
+  }
+}

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/span_origin/FindFirstStackTraceElement.java

@@ -6,6 +6,7 @@
 public interface ConfigurationAcceptor {
   enum Source {
     REMOTE_CONFIG,
+    SPAN_DEBUG,
     EXCEPTION
   }
 

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/span_origin/LineInfo.java

@@ -9,6 +9,7 @@
 import com.datadog.debugger.sink.ProbeStatusSink;
 import com.datadog.debugger.sink.SnapshotSink;
 import com.datadog.debugger.sink.SymbolSink;
+import com.datadog.debugger.snapshot.DefaultSpanDebugger;
 import com.datadog.debugger.symbol.SymDBEnablement;
 import com.datadog.debugger.symbol.SymbolAggregator;
 import com.datadog.debugger.uploader.BatchUploader;
@@ -98,6 +99,10 @@ public static synchronized void run(
               config.getDebuggerMaxExceptionPerSecond());
       DebuggerContext.initExceptionDebugger(defaultExceptionDebugger);
     }
+    if (config.isDebuggerSpanDebugEnabled()) {
+      DebuggerContext.initSpanDebugger(
+          new DefaultSpanDebugger(configurationUpdater, classNameFiltering));
+    }
     if (config.isDebuggerInstrumentTheWorld()) {
       setupInstrumentTheWorldTransformer(
           config, instrumentation, debuggerSink, statsdMetricForwarder);

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

@@ -7,6 +7,7 @@
 import com.datadog.debugger.instrumentation.InstrumentationResult;
 import com.datadog.debugger.instrumentation.MethodInfo;
 import com.datadog.debugger.probe.ExceptionProbe;
+import com.datadog.debugger.probe.ForceMethodInstrumentation;
 import com.datadog.debugger.probe.LogProbe;
 import com.datadog.debugger.probe.MetricProbe;
 import com.datadog.debugger.probe.ProbeDefinition;
@@ -609,7 +610,7 @@ private ProbeDefinition selectReferenceDefinition(
     ProbeId probeId = capturedContextProbes.get(0).getProbeId();
     for (ProbeDefinition definition : capturedContextProbes) {
       if (definition instanceof LogProbe) {
-        if (definition instanceof ExceptionProbe) {
+        if (definition instanceof ForceMethodInstrumentation) {
           where = Where.convertLineToMethod(where, classFileLines);
         }
         hasLogProbe = true;

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

@@ -79,7 +79,7 @@ public void handleException(Throwable t, AgentSpan span) {
       return;
     }
     if (exceptionProbeManager.isAlreadyInstrumented(fingerprint)) {
-      ThrowableState state = exceptionProbeManager.getSateByThrowable(innerMostException);
+      ThrowableState state = exceptionProbeManager.getStateByThrowable(innerMostException);
       if (state == null) {
         LOGGER.debug("Unable to find state for throwable: {}", innerMostException.toString());
         return;

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/ConfigurationAcceptor.java

@@ -147,7 +147,7 @@ public void addSnapshot(Snapshot snapshot) {
     state.addSnapshot(snapshot);
   }
 
-  public ThrowableState getSateByThrowable(Throwable throwable) {
+  public ThrowableState getStateByThrowable(Throwable throwable) {
     return snapshotsByThrowable.get(throwable);
   }
 

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerAgent.java

@@ -4,30 +4,33 @@
 
 import com.datadog.debugger.util.ClassNameFiltering;
 import java.security.MessageDigest;
-import java.security.NoSuchAlgorithmException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 /** Computes a fingerprint of an exception based on its stacktrace and exception type. */
 public class Fingerprinter {
   private static final Logger LOGGER = LoggerFactory.getLogger(Fingerprinter.class);
+  private static MessageDigest digest;
+
+  static {
+    try {
+      digest = MessageDigest.getInstance("SHA-256");
+    } catch (Throwable e) {
+      LOGGER.debug("Unable to find digest algorithm SHA-256", e);
+    }
+  }
 
   // compute fingerprint of the Throwable based on the stacktrace and exception type
   public static String fingerprint(Throwable t, ClassNameFiltering classNameFiltering) {
+    if (digest == null) {
+      return null;
+    }
     t = getInnerMostThrowable(t);
     if (t == null) {
       LOGGER.debug("Unable to find root cause of exception");
       return null;
     }
     Class<? extends Throwable> clazz = t.getClass();
-    MessageDigest digest;
-    try {
-      // TODO avoid lookup a new instance every time
-      digest = MessageDigest.getInstance("SHA-256");
-    } catch (NoSuchAlgorithmException e) {
-      LOGGER.debug("Unable to find digest algorithm SHA-256", e);
-      return null;
-    }
     String typeName = clazz.getTypeName();
     digest.update(typeName.getBytes());
     StackTraceElement[] stackTrace = t.getStackTrace();
@@ -38,8 +41,15 @@ public static String fingerprint(Throwable t, ClassNameFiltering classNameFilter
       }
       digest.update(stackTraceElement.toString().getBytes());
     }
-    byte[] bytes = digest.digest();
-    return bytesToHex(bytes);
+    return bytesToHex(digest.digest());
+  }
+
+  public static String fingerprint(StackTraceElement element) {
+    if (digest == null) {
+      return null;
+    }
+    digest.update(element.toString().getBytes());
+    return bytesToHex(digest.digest());
   }
 
   // convert byte[] to hex string

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerTransformer.java

@@ -17,7 +17,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-public class ExceptionProbe extends LogProbe {
+public class ExceptionProbe extends LogProbe implements ForceMethodInstrumentation {
   private static final Logger LOGGER = LoggerFactory.getLogger(ExceptionProbe.class);
   private final transient ExceptionProbeManager exceptionProbeManager;
 
@@ -76,7 +76,7 @@ public void evaluate(
       // capture only on uncaught exception matching the fingerprint
       ExceptionProbeStatus exceptionStatus = (ExceptionProbeStatus) status;
       ExceptionProbeManager.ThrowableState state =
-          exceptionProbeManager.getSateByThrowable(innerMostThrowable);
+          exceptionProbeManager.getStateByThrowable(innerMostThrowable);
       if (state != null) {
         // Already unwinding the exception
         if (!state.isSampling()) {

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/exception/DefaultExceptionDebugger.java

@@ -0,0 +1,11 @@
+package com.datadog.debugger.probe;
+
+/**
+ * This marker interface indicates that, regardless of the {@link Where} field, this probe should
+ * always be treated as a method probe for purposes of instrumentation and evaluation. In the cases
+ * of {@link ExceptionProbe} and {@link SpanDebuggerProbe}, e.g., these probes need to be
+ * instrumented as method probes even though specific lines are listed in their {@link Where} fields
+ * because the information they are collecting is very much related to specific lines but the
+ * lifecycle of a line probe is insufficient to gather that data.
+ */
+public interface ForceMethodInstrumentation {}

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/exception/ExceptionProbeManager.java

@@ -584,9 +584,9 @@ public CapturedContext.Status createStatus() {
   }
 
   public static class LogStatus extends CapturedContext.Status {
-    private static final LogStatus EMPTY_LOG_STATUS =
+    public static final LogStatus EMPTY_LOG_STATUS =
         new LogStatus(ProbeImplementation.UNKNOWN, false);
-    private static final LogStatus EMPTY_CAPTURING_LOG_STATUS =
+    public static final LogStatus EMPTY_CAPTURING_LOG_STATUS =
         new LogStatus(ProbeImplementation.UNKNOWN, true);
 
     private boolean condition = true;