move Trace/SpanId capture at commit time (#8184)

capture of traceId/SpanId was directly instrumented.
it means that every execution of the probe try to fetch from the span
the traceId and the spanId with eventual conversion to make it a
string. While in the end we may have a condition that is not met or
the sampling does not generate a snapshot. which is wasteful.
We just move the capture at the commit time when we are sure to send a
 snapshot with these information.

Author: jpbempel

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

@@ -32,8 +32,6 @@ public class CapturedContext implements ValueReferenceResolver {
   private Map<String, CapturedValue> staticFields;
   private Limits limits = Limits.DEFAULT;
   private String thisClassName;
-  private String traceId;
-  private String spanId;
   private long duration;
   private final Map<String, Status> statusByProbeId = new LinkedHashMap<>();
   private Map<String, CapturedValue> watches;
@@ -239,19 +237,6 @@ public void addStaticFields(CapturedValue[] values) {
     }
   }
 
-  public void addTraceId(CapturedValue capturedValue) {
-    traceId = extractStringId(capturedValue);
-  }
-
-  public void addSpanId(CapturedValue capturedValue) {
-    spanId = extractStringId(capturedValue);
-  }
-
-  private String extractStringId(CapturedValue capturedValue) {
-    Object value = capturedValue.getValue();
-    return value instanceof String ? (String) value : null;
-  }
-
   public void setLimits(
       int maxReferenceDepth, int maxCollectionSize, int maxLength, int maxFieldCount) {
     this.limits = new Limits(maxReferenceDepth, maxCollectionSize, maxLength, maxFieldCount);
@@ -285,14 +270,6 @@ public String getThisClassName() {
     return thisClassName;
   }
 
-  public String getTraceId() {
-    return traceId;
-  }
-
-  public String getSpanId() {
-    return spanId;
-  }
-
   /**
    * 'Freeze' the context. The contained arguments, locals and fields are converted from their Java
    * instance representation into the corresponding string value.

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/CapturedContextInstrumentor.java

@@ -14,7 +14,6 @@
 import static com.datadog.debugger.instrumentation.Types.CAPTURED_VALUE;
 import static com.datadog.debugger.instrumentation.Types.CAPTURE_THROWABLE_TYPE;
 import static com.datadog.debugger.instrumentation.Types.CLASS_TYPE;
-import static com.datadog.debugger.instrumentation.Types.CORRELATION_ACCESS_TYPE;
 import static com.datadog.debugger.instrumentation.Types.DEBUGGER_CONTEXT_TYPE;
 import static com.datadog.debugger.instrumentation.Types.METHOD_LOCATION_TYPE;
 import static com.datadog.debugger.instrumentation.Types.OBJECT_TYPE;
@@ -34,7 +33,6 @@
 import com.datadog.debugger.sink.Snapshot;
 import com.datadog.debugger.util.ClassFileLines;
 import datadog.trace.api.Config;
-import datadog.trace.bootstrap.debugger.CorrelationAccess;
 import datadog.trace.bootstrap.debugger.Limits;
 import datadog.trace.bootstrap.debugger.MethodLocation;
 import datadog.trace.bootstrap.debugger.ProbeId;
@@ -778,8 +776,6 @@ private InsnList collectCapturedContext(Snapshot.Kind kind, AbstractInsnNode loc
     // stack: [capturedcontext]
     collectStaticFields(insnList);
     // stack: [capturedcontext]
-    collectCorrelationInfo(insnList);
-    // stack: [capturedcontext]
     /*
      * It makes no sense collecting local variables for exceptions - the ones contributing to the exception
      * are most likely to be outside of the scope in the exception handler block and there is no way to figure
@@ -1096,42 +1092,6 @@ private void collectStaticFields(InsnList insnList) {
     // stack: [capturedcontext]
   }
 
-  private void collectCorrelationInfo(InsnList insnList) {
-    // expected stack top: [capturedcontext]
-    /*
-     * We are cheating a bit with CorrelationAccess - utilizing the knowledge that it is a singleton loaded by the
-     * bootstrap class loader we can assume that the availability will not change during the app life time.
-     * As a side effect, we happen to initialize the access here and not from the injected code.
-     */
-    boolean correlationAvailable = CorrelationAccess.instance().isAvailable();
-    if (isStatic && !correlationAvailable) {
-      // static method and no correlation info, no need to capture fields
-      return;
-    }
-    extractSpecialId(insnList, "dd.trace_id", "getTraceId", "addTraceId");
-    // stack: [capturedcontext]
-    extractSpecialId(insnList, "dd.span_id", "getSpanId", "addSpanId");
-    // stack: [capturedcontext]
-  }
-
-  private void extractSpecialId(
-      InsnList insnList, String fieldName, String getMethodName, String addMethodName) {
-    insnList.add(new InsnNode(Opcodes.DUP));
-    // stack: [capturedcontext, capturedcontext]
-    ldc(insnList, fieldName);
-    // stack: [capturedcontext, capturedcontext, name]
-    ldc(insnList, STRING_TYPE.getClassName());
-    // stack: [capturedcontext, capturedcontext, name, type_name]
-    invokeStatic(insnList, CORRELATION_ACCESS_TYPE, "instance", CORRELATION_ACCESS_TYPE);
-    // stack: [capturedcontext, capturedcontext, name, type_name, access]
-    invokeVirtual(insnList, CORRELATION_ACCESS_TYPE, getMethodName, STRING_TYPE);
-    // stack: [capturedcontext, capturedcontext, name, type_name, id]
-    addCapturedValueOf(insnList, limits);
-    // stack: [capturedcontext, capturedcontext, captured_value]
-    invokeVirtual(insnList, CAPTURED_CONTEXT_TYPE, addMethodName, Type.VOID_TYPE, CAPTURED_VALUE);
-    // stack: [capturedcontext]
-  }
-
   private static boolean isAccessible(FieldNode fieldNode) {
     Object value = fieldNode.value;
     if (value instanceof Field) {
@@ -1140,67 +1100,6 @@ private static boolean isAccessible(FieldNode fieldNode) {
     return true;
   }
 
-  private static List<FieldNode> extractInstanceField(
-      ClassNode classNode, boolean isStatic, ClassLoader classLoader, Limits limits) {
-    List<FieldNode> results = new ArrayList<>();
-    if (CorrelationAccess.instance().isAvailable()) {
-      results.add(
-          new FieldNode(
-              Opcodes.ACC_PRIVATE, "dd.trace_id", STRING_TYPE.getDescriptor(), null, null));
-      results.add(
-          new FieldNode(
-              Opcodes.ACC_PRIVATE, "dd.span_id", STRING_TYPE.getDescriptor(), null, null));
-    }
-    if (isStatic) {
-      return results;
-    }
-    int fieldCount = 0;
-    for (FieldNode fieldNode : classNode.fields) {
-      if (isStaticField(fieldNode)) {
-        continue;
-      }
-      results.add(fieldNode);
-      fieldCount++;
-      if (fieldCount > limits.maxFieldCount) {
-        return results;
-      }
-    }
-    addInheritedFields(classNode, classLoader, limits, results, fieldCount);
-    return results;
-  }
-
-  private static void addInheritedFields(
-      ClassNode classNode,
-      ClassLoader classLoader,
-      Limits limits,
-      List<FieldNode> results,
-      int fieldCount) {
-    String superClassName = extractSuperClass(classNode);
-    while (!superClassName.equals(Object.class.getTypeName())) {
-      Class<?> clazz;
-      try {
-        clazz = Class.forName(superClassName, false, classLoader);
-      } catch (ClassNotFoundException ex) {
-        break;
-      }
-      for (Field field : clazz.getDeclaredFields()) {
-        if (isStaticField(field)) {
-          continue;
-        }
-        String desc = Type.getDescriptor(field.getType());
-        FieldNode fieldNode =
-            new FieldNode(field.getModifiers(), field.getName(), desc, null, field);
-        results.add(fieldNode);
-        fieldCount++;
-        if (fieldCount > limits.maxFieldCount) {
-          return;
-        }
-      }
-      clazz = clazz.getSuperclass();
-      superClassName = clazz.getTypeName();
-    }
-  }
-
   private static List<FieldNode> extractStaticFields(
       ClassNode classNode, ClassLoader classLoader, Limits limits) {
     int fieldCount = 0;

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

@@ -23,6 +23,7 @@
 import com.squareup.moshi.Types;
 import datadog.trace.api.Config;
 import datadog.trace.bootstrap.debugger.CapturedContext;
+import datadog.trace.bootstrap.debugger.CorrelationAccess;
 import datadog.trace.bootstrap.debugger.DebuggerContext;
 import datadog.trace.bootstrap.debugger.EvaluationError;
 import datadog.trace.bootstrap.debugger.Limits;
@@ -549,25 +550,19 @@ protected boolean fillSnapshot(
     LogStatus entryStatus = convertStatus(entryContext.getStatus(probeId.getEncodedId()));
     LogStatus exitStatus = convertStatus(exitContext.getStatus(probeId.getEncodedId()));
     String message = null;
-    String traceId = null;
-    String spanId = null;
     switch (evaluateAt) {
       case ENTRY:
       case DEFAULT:
         message = entryStatus.getMessage();
-        traceId = entryContext.getTraceId();
-        spanId = entryContext.getSpanId();
         break;
       case EXIT:
         message = exitStatus.getMessage();
-        traceId = exitContext.getTraceId();
-        spanId = exitContext.getSpanId();
         break;
     }
     boolean shouldCommit = false;
     if (entryStatus.shouldSend() && exitStatus.shouldSend()) {
-      snapshot.setTraceId(traceId);
-      snapshot.setSpanId(spanId);
+      snapshot.setTraceId(CorrelationAccess.instance().getTraceId());
+      snapshot.setSpanId(CorrelationAccess.instance().getSpanId());
       if (isCaptureSnapshot()) {
         snapshot.setEntry(entryContext);
         snapshot.setExit(exitContext);
@@ -643,8 +638,8 @@ public void commit(CapturedContext lineContext, int line) {
     Snapshot snapshot = createSnapshot();
     boolean shouldCommit = false;
     if (status.shouldSend()) {
-      snapshot.setTraceId(lineContext.getTraceId());
-      snapshot.setSpanId(lineContext.getSpanId());
+      snapshot.setTraceId(CorrelationAccess.instance().getTraceId());
+      snapshot.setSpanId(CorrelationAccess.instance().getSpanId());
       snapshot.setMessage(status.getMessage());
       shouldCommit = true;
     }