Add support for non-visitor based class parsing (via dd-instrument-java) (#9894)

Author: mcculls

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/AnnotationOutline.java

@@ -1,5 +1,6 @@
 package datadog.trace.agent.tooling.bytebuddy.outline;
 
+import datadog.instrument.classmatch.ClassFile;
 import java.lang.annotation.Annotation;
 import java.lang.annotation.ElementType;
 import java.lang.annotation.RetentionPolicy;
@@ -17,13 +18,19 @@ final class AnnotationOutline extends WithName implements AnnotationDescription
   private static final Map<String, AnnotationDescription> annotationOutlines = new HashMap<>();
 
   static void prepareAnnotationOutline(String name) {
-    String descriptor = 'L' + name.replace('.', '/') + ';';
-    annotationOutlines.put(descriptor, new AnnotationOutline(name));
+    String internalName = name.replace('.', '/');
+    // flag that we want to parse this annotation
+    ClassFile.annotationOfInterest(internalName);
+    // only a few annotation outlines get prepared - we register them under
+    // both their internal name and descriptor to support different callers
+    AnnotationOutline annotationOutline = new AnnotationOutline(name);
+    annotationOutlines.put('L' + internalName + ';', annotationOutline);
+    annotationOutlines.put(internalName, annotationOutline);
   }
 
   /** Only provide outlines of annotations of interest used for matching. */
-  static AnnotationDescription annotationOutline(String descriptor) {
-    return annotationOutlines.get(descriptor);
+  static AnnotationDescription annotationOutline(String internalNameOrDescriptor) {
+    return annotationOutlines.get(internalNameOrDescriptor);
   }
 
   private AnnotationOutline(String name) {

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/OutlineTypeParser.java

@@ -4,6 +4,8 @@
 import static net.bytebuddy.jar.asm.ClassReader.SKIP_CODE;
 import static net.bytebuddy.jar.asm.ClassReader.SKIP_DEBUG;
 
+import datadog.instrument.classmatch.ClassFile;
+import datadog.trace.api.InstrumenterConfig;
 import java.lang.annotation.Annotation;
 import java.lang.reflect.Field;
 import java.lang.reflect.Method;
@@ -18,13 +20,19 @@
 
 /** Attempts a minimal parse of just the named elements we need for matching. */
 final class OutlineTypeParser implements TypeParser {
+  private static final boolean visitorClassParsing =
+      InstrumenterConfig.get().isVisitorClassParsing();
 
   @Override
   public TypeDescription parse(byte[] bytecode) {
-    ClassReader classReader = OpenedClassReader.of(bytecode);
-    OutlineTypeExtractor typeExtractor = new OutlineTypeExtractor();
-    classReader.accept(typeExtractor, SKIP_CODE | SKIP_DEBUG);
-    return typeExtractor.typeOutline;
+    if (visitorClassParsing) {
+      ClassReader classReader = OpenedClassReader.of(bytecode);
+      OutlineTypeExtractor typeExtractor = new OutlineTypeExtractor();
+      classReader.accept(typeExtractor, SKIP_CODE | SKIP_DEBUG);
+      return typeExtractor.typeOutline;
+    } else {
+      return new TypeOutline(ClassFile.outline(bytecode));
+    }
   }
 
   @Override
@@ -39,7 +47,7 @@ public TypeDescription parse(Class<?> loadedType) {
             extractTypeNames(loadedType.getInterfaces()));
 
     for (Annotation a : loadedType.getDeclaredAnnotations()) {
-      typeOutline.declare(annotationOutline(Type.getDescriptor(a.annotationType())));
+      typeOutline.declare(annotationOutline(Type.getInternalName(a.annotationType())));
     }
 
     for (Field field : loadedType.getDeclaredFields()) {
@@ -50,7 +58,7 @@ public TypeDescription parse(Class<?> loadedType) {
               field.getName(),
               Type.getDescriptor(field.getType()));
       for (Annotation a : field.getDeclaredAnnotations()) {
-        fieldOutline.declare(annotationOutline(Type.getDescriptor(a.annotationType())));
+        fieldOutline.declare(annotationOutline(Type.getInternalName(a.annotationType())));
       }
       typeOutline.declare(fieldOutline);
     }
@@ -63,7 +71,7 @@ public TypeDescription parse(Class<?> loadedType) {
               method.getName(),
               Type.getMethodDescriptor(method));
       for (Annotation a : method.getDeclaredAnnotations()) {
-        methodOutline.declare(annotationOutline(Type.getDescriptor(a.annotationType())));
+        methodOutline.declare(annotationOutline(Type.getInternalName(a.annotationType())));
       }
       typeOutline.declare(methodOutline);
     }

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/TypeOutline.java

@@ -1,7 +1,9 @@
 package datadog.trace.agent.tooling.bytebuddy.outline;
 
+import static datadog.trace.agent.tooling.bytebuddy.outline.AnnotationOutline.annotationOutline;
 import static datadog.trace.agent.tooling.bytebuddy.outline.TypeFactory.findType;
 
+import datadog.instrument.classmatch.ClassOutline;
 import java.util.ArrayList;
 import java.util.List;
 import net.bytebuddy.description.annotation.AnnotationDescription;
@@ -42,6 +44,30 @@ final class TypeOutline extends WithName {
     this.interfaces = interfaces;
   }
 
+  /** Adapts simpler {@link ClassOutline} structure to existing {@link TypeOutline}. */
+  public TypeOutline(ClassOutline outline) {
+    super(outline.className.replace('/', '.'));
+    this.modifiers = outline.access & ALLOWED_TYPE_MODIFIERS;
+    this.superName = outline.superName;
+    this.interfaces = outline.interfaces;
+
+    for (String annotation : outline.annotations) {
+      declare(annotationOutline(annotation));
+    }
+
+    for (datadog.instrument.classmatch.FieldOutline f : outline.fields) {
+      declare(new FieldOutline(this, f.access, f.fieldName, f.descriptor));
+    }
+
+    for (datadog.instrument.classmatch.MethodOutline m : outline.methods) {
+      MethodOutline method = new MethodOutline(this, m.access, m.methodName, m.descriptor);
+      for (String annotation : m.annotations) {
+        method.declare(annotationOutline(annotation));
+      }
+      declare(method);
+    }
+  }
+
   @Override
   protected TypeDescription delegate() {
     throw new IllegalStateException("Not available in outline for " + name);

dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java

@@ -162,6 +162,7 @@ public final class TraceInstrumentationConfig {
   public static final String RESOLVER_NAMES_ARE_UNIQUE = "resolver.names.are.unique";
 
   public static final String UNSAFE_CLASS_INJECTION = "unsafe.class.injection";
+  public static final String VISITOR_CLASS_PARSING = "visitor.class.parsing";
 
   public static final String CASSANDRA_KEYSPACE_STATEMENT_EXTRACTION_ENABLED =
       "trace.cassandra.keyspace.statement.extraction.enabled";

internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java

@@ -74,6 +74,7 @@
 import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_THREAD_POOL_EXECUTORS_EXCLUDE;
 import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_WEBSOCKET_MESSAGES_ENABLED;
 import static datadog.trace.api.config.TraceInstrumentationConfig.UNSAFE_CLASS_INJECTION;
+import static datadog.trace.api.config.TraceInstrumentationConfig.VISITOR_CLASS_PARSING;
 import static datadog.trace.api.config.UsmConfig.USM_ENABLED;
 import static datadog.trace.util.CollectionUtils.tryMakeImmutableList;
 import static datadog.trace.util.CollectionUtils.tryMakeImmutableSet;
@@ -162,6 +163,7 @@ public class InstrumenterConfig {
   private final boolean resolverUseLoadClass;
   private final Boolean resolverUseUrlCaches;
   private final int resolverResetInterval;
+  private final boolean visitorClassParsing;
 
   private final boolean unsafeClassInjection;
 
@@ -284,6 +286,7 @@ private InstrumenterConfig() {
             : configProvider.getInteger(RESOLVER_RESET_INTERVAL, DEFAULT_RESOLVER_RESET_INTERVAL);
 
     unsafeClassInjection = configProvider.getBoolean(UNSAFE_CLASS_INJECTION, false);
+    visitorClassParsing = configProvider.getBoolean(VISITOR_CLASS_PARSING, false);
 
     runtimeContextFieldInjection =
         configProvider.getBoolean(
@@ -514,6 +517,10 @@ public boolean isUnsafeClassInjection() {
     return unsafeClassInjection;
   }
 
+  public boolean isVisitorClassParsing() {
+    return visitorClassParsing;
+  }
+
   public String getInstrumentationConfigId() {
     return instrumentationConfigId;
   }

metadata/supported-configurations.json

@@ -1343,6 +1343,7 @@
     "DD_UNSAFE_CLASS_INJECTION": ["A"],
     "DD_USM_ENABLED": ["A"],
     "DD_VERSION": ["A"],
+    "DD_VISITOR_CLASS_PARSING": ["A"],
     "DD_WRITER_BAGGAGE_INJECT": ["A"],
     "DD_WRITER_TYPE": ["A"],
     "OTEL_INSTRUMENTATION_HTTP_CLIENT_CAPTURE_REQUEST_HEADERS": ["A"],