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

Author: mcculls

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

@@ -17,13 +17,17 @@ 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));
+    // 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);
+    String internalName = name.replace('.', '/');
+    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,29 @@ final class TypeOutline extends WithName {
     this.interfaces = interfaces;
   }
 
+  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-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/TypePoolFacade.java

@@ -4,6 +4,7 @@
 import static datadog.trace.agent.tooling.bytebuddy.outline.TypeFactory.findType;
 import static datadog.trace.agent.tooling.bytebuddy.outline.TypeFactory.typeFactory;
 
+import datadog.instrument.classmatch.ClassFile;
 import datadog.trace.agent.tooling.bytebuddy.SharedTypePools;
 import datadog.trace.agent.tooling.bytebuddy.memoize.Memoizer;
 import datadog.trace.api.InstrumenterConfig;
@@ -36,6 +37,9 @@ public static ClassLoader currentContext() {
 
   @Override
   public void annotationOfInterest(String name) {
+    if (!InstrumenterConfig.get().isVisitorClassParsing()) {
+      ClassFile.annotationOfInterest(name.replace('.', '/'));
+    }
     AnnotationOutline.prepareAnnotationOutline(name);
   }
 

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

@@ -159,6 +159,9 @@ public final class TraceInstrumentationConfig {
   public static final String RESOLVER_USE_URL_CACHES = "resolver.use.url.caches";
   public static final String RESOLVER_RESET_INTERVAL = "resolver.reset.interval";
   public static final String RESOLVER_NAMES_ARE_UNIQUE = "resolver.names.are.unique";
+
+  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";
   public static final String COUCHBASE_INTERNAL_SPANS_ENABLED =

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

@@ -73,6 +73,7 @@
 import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_PEKKO_SCHEDULER_ENABLED;
 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.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;
@@ -161,6 +162,7 @@ public class InstrumenterConfig {
   private final boolean resolverUseLoadClass;
   private final Boolean resolverUseUrlCaches;
   private final int resolverResetInterval;
+  private final boolean visitorClassParsing;
 
   private final boolean runtimeContextFieldInjection;
   private final boolean serialVersionUIDFieldInjection;
@@ -280,6 +282,8 @@ private InstrumenterConfig() {
             ? 0
             : configProvider.getInteger(RESOLVER_RESET_INTERVAL, DEFAULT_RESOLVER_RESET_INTERVAL);
 
+    visitorClassParsing = configProvider.getBoolean(VISITOR_CLASS_PARSING, false);
+
     runtimeContextFieldInjection =
         configProvider.getBoolean(
             RUNTIME_CONTEXT_FIELD_INJECTION, DEFAULT_RUNTIME_CONTEXT_FIELD_INJECTION);
@@ -505,6 +509,10 @@ public String getResolverCacheDir() {
     return resolverCacheDir;
   }
 
+  public boolean isVisitorClassParsing() {
+    return visitorClassParsing;
+  }
+
   public String getInstrumentationConfigId() {
     return instrumentationConfigId;
   }