Avoid looking up annotation types during type matching (#5906)

laurit · web-flow · commit 8865ab6ae510 · 2022-05-12T09:30:30.000-07:00
diff --git a/muzzle/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/AgentCachingPoolStrategy.java b/muzzle/src/main/java/io/opentelemetry/javaagent/tooling/muzzle/AgentCachingPoolStrategy.java
@@ -15,11 +15,15 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.description.annotation.AnnotationList;
+import net.bytebuddy.description.annotation.AnnotationValue;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.method.MethodList;
+import net.bytebuddy.description.method.ParameterDescription;
+import net.bytebuddy.description.method.ParameterList;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.description.type.TypeList;
 import net.bytebuddy.dynamic.ClassFileLocator;
@@ -294,6 +298,8 @@ public void clear() {
 
   /** Based on TypePool.Default.WithLazyResolution */
   private class AgentTypePool extends TypePool.Default {
+    // ThreadLocal used for detecting loading of annotation types
+    private final ThreadLocal<Boolean> loadingAnnotations = new ThreadLocal<>();
     private final WeakReference<ClassLoader> classLoaderRef;
 
     public AgentTypePool(
@@ -331,6 +337,18 @@ protected TypePool.Resolution doResolve(String name) {
       return resolution;
     }
 
+    void enterLoadAnnotations() {
+      loadingAnnotations.set(Boolean.TRUE);
+    }
+
+    void exitLoadAnnotations() {
+      loadingAnnotations.set(null);
+    }
+
+    boolean isLoadingAnnotations() {
+      return loadingAnnotations.get() != null;
+    }
+
     /** Based on TypePool.Default.WithLazyResolution.LazyResolution */
     private class LazyResolution implements TypePool.Resolution {
       private final WeakReference<ClassLoader> classLoaderRef;
@@ -343,6 +361,19 @@ private class LazyResolution implements TypePool.Resolution {
 
       @Override
       public boolean isResolved() {
+        // Like jdk, byte-buddy getDeclaredAnnotations does not report annotations whose class is
+        // missing. To do this it needs to locate the bytes for annotation types used in class.
+        // Which means that if we have a matcher that matches methods annotated with @Foo byte-buddy
+        // will end up location bytes for all annotations used on any method in the classes that
+        // this matcher is applied to. From our perspective this is unreasonable, we just want to
+        // match based on annotation name with as little overhead as possible. As we match only
+        // based on annotation name we never need to locate the bytes for the annotation type.
+        // See TypePool.Default.LazyTypeDescription.LazyAnnotationDescription.asList()
+        // When isResolved() is called during loading of annotations delay resolving to avoid
+        // looking up the class bytes.
+        if (isLoadingAnnotations()) {
+          return true;
+        }
         return doResolve(name).isResolved();
       }
 
@@ -354,6 +385,11 @@ public TypeDescription resolve() {
         // super class and interfaces multiple times
         if (cached == null) {
           cached = new AgentTypePool.LazyTypeDescription(classLoaderRef, name);
+          // if we know that an annotation is being loaded wrap the result so that we wouldn't
+          // need to resolve the class bytes to tell whether it is an annotation
+          if (isLoadingAnnotations()) {
+            cached = new AnnotationTypeDescription(cached);
+          }
         }
         return cached;
       }
@@ -376,7 +412,15 @@ protected TypeDescription delegate() {
       @Override
       public AnnotationList getDeclaredAnnotations() {
         if (annotations == null) {
-          annotations = delegate().getDeclaredAnnotations();
+          TypeDescription delegate = delegate();
+          // Run getDeclaredAnnotations with ThreadLocal. ThreadLocal helps us detect types looked
+          // up by getDeclaredAnnotations and treat them specially.
+          enterLoadAnnotations();
+          try {
+            annotations = delegate.getDeclaredAnnotations();
+          } finally {
+            exitLoadAnnotations();
+          }
         }
         return annotations;
       }
@@ -413,12 +457,12 @@ public String getName() {
         return name;
       }
 
-      private volatile Generic cachedSuperClass;
+      private volatile TypeDescription.Generic cachedSuperClass;
 
       @Override
-      public Generic getSuperClass() {
+      public TypeDescription.Generic getSuperClass() {
         if (cachedSuperClass == null) {
-          Generic superClassDescription = delegate().getSuperClass();
+          TypeDescription.Generic superClassDescription = delegate().getSuperClass();
           ClassLoader classLoader = classLoaderRef.get();
           if (canUseFindLoadedClass() && classLoader != null && superClassDescription != null) {
             String superName = superClassDescription.getTypeName();
@@ -445,7 +489,7 @@ public TypeList.Generic getInterfaces() {
             // here we use raw types and loose generic info
             // we don't expect to have matchers that would use the generic info
             List<TypeDescription> result = new ArrayList<>();
-            for (Generic interfaceDescription : interfaces) {
+            for (TypeDescription.Generic interfaceDescription : interfaces) {
               String interfaceName = interfaceDescription.getTypeName();
               Class<?> interfaceClass = findLoadedClass(classLoader, interfaceName);
               if (interfaceClass != null) {
@@ -461,6 +505,45 @@ public TypeList.Generic getInterfaces() {
         }
         return cachedInterfaces;
       }
+
+      private class LazyAnnotationMethodDescription extends DelegatingMethodDescription {
+
+        LazyAnnotationMethodDescription(MethodDescription.InDefinedShape method) {
+          super(method);
+        }
+
+        @Override
+        public AnnotationList getDeclaredAnnotations() {
+          // Run getDeclaredAnnotations with ThreadLocal. ThreadLocal helps us detect types looked
+          // up by getDeclaredAnnotations and treat them specially.
+          enterLoadAnnotations();
+          try {
+            return method.getDeclaredAnnotations();
+          } finally {
+            exitLoadAnnotations();
+          }
+        }
+      }
+
+      @Override
+      public MethodList<MethodDescription.InDefinedShape> getDeclaredMethods() {
+        MethodList<MethodDescription.InDefinedShape> methods = super.getDeclaredMethods();
+
+        class MethodListWrapper extends MethodList.AbstractBase<MethodDescription.InDefinedShape> {
+
+          @Override
+          public MethodDescription.InDefinedShape get(int index) {
+            return new LazyAnnotationMethodDescription(methods.get(index));
+          }
+
+          @Override
+          public int size() {
+            return methods.size();
+          }
+        }
+
+        return new MethodListWrapper();
+      }
     }
 
     private AgentTypePool.LazyTypeDescriptionWithClass newTypeDescription(Class<?> clazz) {
@@ -493,10 +576,10 @@ public String getName() {
         return name;
       }
 
-      private volatile Generic cachedSuperClass;
+      private volatile TypeDescription.Generic cachedSuperClass;
 
       @Override
-      public Generic getSuperClass() {
+      public TypeDescription.Generic getSuperClass() {
         if (cachedSuperClass == null) {
           Class<?> clazz = classRef.get();
           if (clazz == null) {
@@ -556,4 +639,91 @@ private static AgentTypePool.LazyTypeDescriptionWithClass newLazyTypeDescription
     }
     return pool.new LazyTypeDescriptionWithClass(clazz);
   }
+
+  /**
+   * Class descriptor that claims to represent an annotation without checking whether the underlying
+   * type really is an annotation.
+   */
+  private static class AnnotationTypeDescription
+      extends TypeDescription.AbstractBase.OfSimpleType.WithDelegation {
+    private final TypeDescription delegate;
+
+    AnnotationTypeDescription(TypeDescription delegate) {
+      this.delegate = delegate;
+    }
+
+    @Override
+    protected TypeDescription delegate() {
+      return delegate;
+    }
+
+    @Override
+    public String getName() {
+      return delegate.getName();
+    }
+
+    @Override
+    public boolean isAnnotation() {
+      // by default byte-buddy checks whether class modifiers have annotation bit set
+      // as we wish to avoid looking up the class bytes we assume that every that was expected
+      // to be an annotation really is an annotation and return true here
+      // See TypePool.Default.LazyTypeDescription.LazyAnnotationDescription.asList()
+      return true;
+    }
+  }
+
+  private static class DelegatingMethodDescription
+      extends MethodDescription.InDefinedShape.AbstractBase {
+    protected final MethodDescription.InDefinedShape method;
+
+    DelegatingMethodDescription(MethodDescription.InDefinedShape method) {
+      this.method = method;
+    }
+
+    @Nonnull
+    @Override
+    public TypeDescription getDeclaringType() {
+      return method.getDeclaringType();
+    }
+
+    @Override
+    public TypeDescription.Generic getReturnType() {
+      return method.getReturnType();
+    }
+
+    @Override
+    public ParameterList<ParameterDescription.InDefinedShape> getParameters() {
+      return method.getParameters();
+    }
+
+    @Override
+    public TypeList.Generic getExceptionTypes() {
+      return method.getExceptionTypes();
+    }
+
+    @Override
+    public AnnotationValue<?, ?> getDefaultValue() {
+      return method.getDefaultValue();
+    }
+
+    @Override
+    public String getInternalName() {
+      return method.getInternalName();
+    }
+
+    @Override
+    public TypeList.Generic getTypeVariables() {
+      return method.getTypeVariables();
+    }
+
+    @Override
+    public int getModifiers() {
+      return method.getModifiers();
+    }
+
+    @Override
+    public AnnotationList getDeclaredAnnotations() {
+      return method.getDeclaredAnnotations();
+    }
+  }
 }
diff --git a/muzzle/src/test/java/io/opentelemetry/javaagent/tooling/muzzle/AgentCachingPoolStrategyTest.java b/muzzle/src/test/java/io/opentelemetry/javaagent/tooling/muzzle/AgentCachingPoolStrategyTest.java
@@ -0,0 +1,64 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling.muzzle;
+
+import static net.bytebuddy.matcher.ElementMatchers.declaresMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isAnnotatedWith;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import io.opentelemetry.test.AnnotatedTestClass;
+import java.io.IOException;
+import java.io.InputStream;
+import java.net.URL;
+import java.util.Enumeration;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.dynamic.ClassFileLocator;
+import net.bytebuddy.pool.TypePool;
+import org.junit.jupiter.api.Test;
+
+class AgentCachingPoolStrategyTest {
+
+  @Test
+  void testSkipResourceLookupForAnnotations() {
+    ClassLoader classLoader =
+        new ClassLoader(AgentCachingPoolStrategyTest.class.getClassLoader()) {
+
+          private void checkResource(String name) {
+            if (name.contains("TestAnnotation")) {
+              throw new IllegalStateException("Unexpected resource lookup for " + name);
+            }
+          }
+
+          @Override
+          public URL getResource(String name) {
+            checkResource(name);
+            return super.getResource(name);
+          }
+
+          @Override
+          public InputStream getResourceAsStream(String name) {
+            checkResource(name);
+            return super.getResourceAsStream(name);
+          }
+
+          @Override
+          public Enumeration<URL> getResources(String name) throws IOException {
+            checkResource(name);
+            return super.getResources(name);
+          }
+        };
+
+    ClassFileLocator locator = ClassFileLocator.ForClassLoader.of(classLoader);
+    TypePool pool = AgentTooling.poolStrategy().typePool(locator, classLoader);
+    TypePool.Resolution resolution = pool.describe(AnnotatedTestClass.class.getName());
+    TypeDescription typeDescription = resolution.resolve();
+
+    assertTrue(isAnnotatedWith(AnnotatedTestClass.TestAnnotation.class).matches(typeDescription));
+    assertTrue(
+        declaresMethod(isAnnotatedWith(AnnotatedTestClass.TestAnnotation.class))
+            .matches(typeDescription));
+  }
+}
diff --git a/muzzle/src/test/java/io/opentelemetry/test/AnnotatedTestClass.java b/muzzle/src/test/java/io/opentelemetry/test/AnnotatedTestClass.java
@@ -0,0 +1,15 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.test;
+
+@AnnotatedTestClass.TestAnnotation
+public class AnnotatedTestClass {
+
+  @AnnotatedTestClass.TestAnnotation
+  void testMethod() {}
+
+  public @interface TestAnnotation {}
+}
