Add Probe Metadata

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

@@ -294,13 +294,14 @@ public void freeze(TimeoutChecker timeoutChecker) {
   }
 
   public Status evaluate(
-      String encodedProbeId,
       ProbeImplementation probeImplementation,
       String thisClassName,
       long startTimestamp,
       MethodLocation methodLocation) {
     Status status =
-        statusByProbeId.computeIfAbsent(encodedProbeId, key -> probeImplementation.createStatus());
+        statusByProbeId.computeIfAbsent(
+            probeImplementation.getProbeId().getEncodedId(),
+            key -> probeImplementation.createStatus());
     if (methodLocation == MethodLocation.EXIT && startTimestamp > 0) {
       duration = System.nanoTime() - startTimestamp;
       addExtension(
@@ -326,6 +327,18 @@ public Status getStatus(String encodedProbeId) {
     return result;
   }
 
+  public Status getStatus(int probeIndex) {
+    ProbeImplementation probeImplementation = DebuggerContext.resolveProbe(probeIndex);
+    if (probeImplementation != null) {
+      return getStatus(probeImplementation.getProbeId().getEncodedId());
+    }
+    Status result = statusByProbeId.get(ProbeImplementation.UNKNOWN.getProbeId().getEncodedId());
+    if (result == null) {
+      return Status.EMPTY_STATUS;
+    }
+    return result;
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) return true;

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

@@ -51,7 +51,7 @@ public String tag() {
   }
 
   public interface ProbeResolver {
-    ProbeImplementation resolve(String encodedProbeId);
+    ProbeImplementation resolve(int probeIndex);
   }
 
   public interface ClassFilter {
@@ -144,16 +144,13 @@ public static void initCodeOrigin(CodeOriginRecorder codeOriginRecorder) {
     DebuggerContext.codeOriginRecorder = codeOriginRecorder;
   }
 
-  /**
-   * 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
-   */
-  public static ProbeImplementation resolveProbe(String id) {
+  /** Returns the probe details based on the probe idx provided. */
+  public static ProbeImplementation resolveProbe(int probeIndex) {
     ProbeResolver resolver = probeResolver;
     if (resolver == null) {
       return null;
     }
-    return resolver.resolve(id);
+    return resolver.resolve(probeIndex);
   }
 
   /**
@@ -253,7 +250,7 @@ public static DebuggerSpan createSpan(String probeId, String operationName, Stri
    *
    * @return true if can proceed to capture data
    */
-  public static boolean isReadyToCapture(Class<?> callingClass, String... encodedProbeIds) {
+  public static boolean isReadyToCapture(Class<?> callingClass, int... probeIndices) {
     try {
       return checkAndSetInProbe();
     } catch (Exception ex) {
@@ -287,21 +284,17 @@ public static void evalContext(
       Class<?> callingClass,
       long startTimestamp,
       MethodLocation methodLocation,
-      String... encodedProbeIds) {
+      int... probeIndices) {
     try {
       boolean needFreeze = false;
-      for (String encodedProbeId : encodedProbeIds) {
-        ProbeImplementation probeImplementation = resolveProbe(encodedProbeId);
+      for (int probeIndex : probeIndices) {
+        ProbeImplementation probeImplementation = resolveProbe(probeIndex);
         if (probeImplementation == null) {
           continue;
         }
         CapturedContext.Status status =
             context.evaluate(
-                encodedProbeId,
-                probeImplementation,
-                callingClass.getTypeName(),
-                startTimestamp,
-                methodLocation);
+                probeImplementation, callingClass.getTypeName(), startTimestamp, methodLocation);
         needFreeze |= status.shouldFreezeContext();
       }
       // only freeze the context when we have at lest one snapshot probe, and we should send
@@ -321,20 +314,16 @@ public static void evalContext(
    * conditions and commit snapshot to send it if needed. This is for line probes.
    */
   public static void evalContextAndCommit(
-      CapturedContext context, Class<?> callingClass, int line, String... encodedProbeIds) {
+      CapturedContext context, Class<?> callingClass, int line, int... probeIndices) {
     try {
       List<ProbeImplementation> probeImplementations = new ArrayList<>();
-      for (String encodedProbeId : encodedProbeIds) {
-        ProbeImplementation probeImplementation = resolveProbe(encodedProbeId);
+      for (int probeIndex : probeIndices) {
+        ProbeImplementation probeImplementation = resolveProbe(probeIndex);
         if (probeImplementation == null) {
           continue;
         }
         context.evaluate(
-            encodedProbeId,
-            probeImplementation,
-            callingClass.getTypeName(),
-            -1,
-            MethodLocation.DEFAULT);
+            probeImplementation, callingClass.getTypeName(), -1, MethodLocation.DEFAULT);
         probeImplementations.add(probeImplementation);
       }
       for (ProbeImplementation probeImplementation : probeImplementations) {
@@ -345,9 +334,9 @@ public static void evalContextAndCommit(
     }
   }
 
-  public static void codeOrigin(String probeId) {
+  public static void codeOrigin(int probeIndex) {
     try {
-      ProbeImplementation probe = probeResolver.resolve(probeId);
+      ProbeImplementation probe = probeResolver.resolve(probeIndex);
       if (probe != null) {
         probe.commit(
             CapturedContext.EMPTY_CONTEXT, CapturedContext.EMPTY_CONTEXT, Collections.emptyList());
@@ -365,16 +354,16 @@ public static void commit(
       CapturedContext entryContext,
       CapturedContext exitContext,
       List<CapturedContext.CapturedThrowable> caughtExceptions,
-      String... encodedProbeIds) {
+      int... probeIndices) {
     try {
       if (entryContext == CapturedContext.EMPTY_CONTEXT
           && exitContext == CapturedContext.EMPTY_CONTEXT) {
         // rate limited
         return;
       }
-      for (String encodedProbeId : encodedProbeIds) {
-        CapturedContext.Status entryStatus = entryContext.getStatus(encodedProbeId);
-        CapturedContext.Status exitStatus = exitContext.getStatus(encodedProbeId);
+      for (int probeIndex : probeIndices) {
+        CapturedContext.Status entryStatus = entryContext.getStatus(probeIndex);
+        CapturedContext.Status exitStatus = exitContext.getStatus(probeIndex);
         ProbeImplementation probeImplementation;
         if (entryStatus.probeImplementation != ProbeImplementation.UNKNOWN
             && (entryStatus.probeImplementation.getEvaluateAt() == MethodLocation.ENTRY
@@ -383,7 +372,7 @@ public static void commit(
         } else if (exitStatus.probeImplementation.getEvaluateAt() == MethodLocation.EXIT) {
           probeImplementation = exitStatus.probeImplementation;
         } else {
-          throw new IllegalStateException("no probe details for " + encodedProbeId);
+          throw new IllegalStateException("no probe details for " + probeIndex);
         }
         probeImplementation.commit(entryContext, exitContext, caughtExceptions);
       }

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

@@ -47,6 +47,7 @@ DebuggerTransformer supply(
         Config tracerConfig,
         Configuration configuration,
         DebuggerTransformer.InstrumentationListener listener,
+        ProbeMetadata probeMetadata,
         DebuggerSink debuggerSink);
   }
 
@@ -61,6 +62,7 @@ DebuggerTransformer supply(
   private volatile Configuration currentConfiguration;
   private DebuggerTransformer currentTransformer;
   private final Map<String, ProbeDefinition> appliedDefinitions = new ConcurrentHashMap<>();
+  private final ProbeMetadata probeMetadata = new ProbeMetadata();
   private final DebuggerSink sink;
   private final ClassesToRetransformFinder finder;
   private final String serviceName;
@@ -119,6 +121,10 @@ public void handleException(String configId, Exception ex) {
     sink.getProbeStatusSink().addError(probeId, ex);
   }
 
+  ProbeMetadata getProbeMetadata() {
+    return probeMetadata;
+  }
+
   private ProbeId extractPrefix(String prefix, String configId) {
     return new ProbeId(configId.substring(prefix.length()), 0);
   }
@@ -201,7 +207,11 @@ private void installNewDefinitions(Configuration newConfiguration) {
     // install new probe definitions
     DebuggerTransformer newTransformer =
         transformerSupplier.supply(
-            Config.get(), newConfiguration, this::recordInstrumentationProgress, sink);
+            Config.get(),
+            newConfiguration,
+            this::recordInstrumentationProgress,
+            probeMetadata,
+            sink);
     instrumentation.addTransformer(newTransformer, true);
     currentTransformer = newTransformer;
     LOGGER.debug("New transformer installed");
@@ -231,6 +241,7 @@ private void retransformClasses(List<Class<?>> classesToBeTransformed) {
   private void storeDebuggerDefinitions(ConfigurationComparer changes) {
     for (ProbeDefinition definition : changes.getRemovedDefinitions()) {
       appliedDefinitions.remove(definition.getProbeId().getEncodedId());
+      probeMetadata.removeProbe(definition.getProbeId().getEncodedId());
     }
     for (ProbeDefinition definition : changes.getAddedDefinitions()) {
       appliedDefinitions.put(definition.getProbeId().getEncodedId(), definition);
@@ -240,12 +251,12 @@ private void storeDebuggerDefinitions(ConfigurationComparer changes) {
 
   // /!\ This is called potentially by multiple threads from the instrumented code /!\
   @Override
-  public ProbeImplementation resolve(String encodedProbeId) {
-    ProbeDefinition definition = appliedDefinitions.get(encodedProbeId);
-    if (definition == null) {
-      ratelimitedLogger.warn(SEND_TELEMETRY, "Cannot resolve probe id={}", encodedProbeId);
+  public ProbeImplementation resolve(int probeIndex) {
+    ProbeImplementation probeImplementation = probeMetadata.getProbe(probeIndex);
+    if (probeImplementation == null) {
+      ratelimitedLogger.warn(SEND_TELEMETRY, "Cannot resolve probe index={}", probeIndex);
     }
-    return definition;
+    return probeImplementation;
   }
 
   private static void applyRateLimiter(

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

@@ -394,7 +394,8 @@ static ClassFileTransformer setupInstrumentTheWorldTransformer(
       Config config, Instrumentation instrumentation, DebuggerSink debuggerSink) {
     LOGGER.info("install Instrument-The-World transformer");
     DebuggerTransformer transformer =
-        createTransformer(config, Configuration.builder().build(), null, debuggerSink);
+        createTransformer(
+            config, Configuration.builder().build(), null, new ProbeMetadata(), debuggerSink);
     DebuggerContext.initProbeResolver(transformer::instrumentTheWorldResolver);
     instrumentation.addTransformer(transformer);
     return transformer;
@@ -412,8 +413,9 @@ private static DebuggerTransformer createTransformer(
       Config config,
       Configuration configuration,
       DebuggerTransformer.InstrumentationListener listener,
+      ProbeMetadata probeMetadata,
       DebuggerSink debuggerSink) {
-    return new DebuggerTransformer(config, configuration, listener, debuggerSink);
+    return new DebuggerTransformer(config, configuration, listener, probeMetadata, debuggerSink);
   }
 
   static void stop() {