Add Probe Metadata (#9723)

Introducing ProbeMetadata class to help get probe from instrumented
 code. Instead of getting the probe instance from a map based on
probeId we are directly getting it from an array that is filled at
instrumentation time. This way it's cheaper and abstracted the way to
get probe definition to evaluate conditions for example.
In instrumented code, instead of referring the probe by the probeId
hard-coded we are using index in probe metadata array.
ProbeMetadata class encapsulate an AtomicReferenceArray to make sure
every cell as a volatile-style access and provide linear way to remove
or add into this array.
For now linear scanning is fine as we don't expect a lot of probes in
a same location. Could improve on that using bitmaps later if require.
Instrumentation is changed to push int arrays instead of String arrays
and can be optimize later to accept static arguments instead of
varargs/array instead.

Author: jpbempel

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() {