Faster type matching

instrumentation/external-annotations/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/extannotations/ExternalAnnotationInstrumentation.java

@@ -6,7 +6,6 @@
 package io.opentelemetry.javaagent.instrumentation.extannotations;
 
 import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
-import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasSuperType;
 import static io.opentelemetry.javaagent.instrumentation.extannotations.ExternalAnnotationSingletons.instrumenter;
 import static java.util.logging.Level.WARNING;
 import static net.bytebuddy.matcher.ElementMatchers.declaresMethod;
@@ -100,7 +99,7 @@ public ElementMatcher<ClassLoader> classLoaderOptimization() {
 
   @Override
   public ElementMatcher<TypeDescription> typeMatcher() {
-    return hasSuperType(declaresMethod(isAnnotatedWith(traceAnnotationMatcher)));
+    return declaresMethod(isAnnotatedWith(traceAnnotationMatcher));
   }
 
   @Override

instrumentation/internal/internal-class-loader/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/internal/classloader/ClassLoaderInstrumentationModule.java

@@ -34,6 +34,7 @@ public List<TypeInstrumentation> typeInstrumentations() {
     return asList(
         new BootDelegationInstrumentation(),
         new LoadInjectedClassInstrumentation(),
-        new ResourceInjectionInstrumentation());
+        new ResourceInjectionInstrumentation(),
+        new DefineClassInstrumentation());
   }
 }

instrumentation/internal/internal-class-loader/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/internal/classloader/DefineClassInstrumentation.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.internal.classloader;
+
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import io.opentelemetry.javaagent.bootstrap.DefineClassHelper;
+import io.opentelemetry.javaagent.bootstrap.DefineClassHelper.Handler.DefineClassContext;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import java.nio.ByteBuffer;
+import java.security.ProtectionDomain;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class DefineClassInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("java.lang.ClassLoader");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        named("defineClass")
+            .and(
+                takesArguments(
+                    String.class, byte[].class, int.class, int.class, ProtectionDomain.class)),
+        DefineClassInstrumentation.class.getName() + "$DefineClassAdvice");
+    transformer.applyAdviceToMethod(
+        named("defineClass")
+            .and(takesArguments(String.class, ByteBuffer.class, ProtectionDomain.class)),
+        DefineClassInstrumentation.class.getName() + "$DefineClassAdvice2");
+  }
+
+  @SuppressWarnings("unused")
+  public static class DefineClassAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static DefineClassContext onEnter(
+        @Advice.This ClassLoader classLoader,
+        @Advice.Argument(0) String className,
+        @Advice.Argument(1) byte[] classBytes,
+        @Advice.Argument(2) int offset,
+        @Advice.Argument(3) int length) {
+      return DefineClassHelper.beforeDefineClass(
+          classLoader, className, classBytes, offset, length);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(@Advice.Enter DefineClassContext context) {
+      DefineClassHelper.afterDefineClass(context);
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class DefineClassAdvice2 {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static DefineClassContext onEnter(
+        @Advice.This ClassLoader classLoader,
+        @Advice.Argument(0) String className,
+        @Advice.Argument(1) ByteBuffer classBytes) {
+      return DefineClassHelper.beforeDefineClass(classLoader, className, classBytes);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(@Advice.Enter DefineClassContext context) {
+      DefineClassHelper.afterDefineClass(context);
+    }
+  }
+}

javaagent-bootstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/DefineClassHelper.java

@@ -0,0 +1,68 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.bootstrap;
+
+import java.nio.ByteBuffer;
+
+public class DefineClassHelper {
+
+  /** Helper class for {@code ClassLoader.defineClass} callbacks. */
+  public interface Handler {
+    DefineClassContext beforeDefineClass(
+        ClassLoader classLoader, String className, byte[] classBytes, int offset, int length);
+
+    void afterDefineClass(DefineClassContext context);
+
+    /** Context returned from {@code beforeDefineClass} and passed to {@code afterDefineClass}. */
+    interface DefineClassContext {
+      void exit();
+    }
+  }
+
+  private static volatile Handler handler;
+
+  public static Handler.DefineClassContext beforeDefineClass(
+      ClassLoader classLoader, String className, byte[] classBytes, int offset, int length) {
+    return handler.beforeDefineClass(classLoader, className, classBytes, offset, length);
+  }
+
+  public static Handler.DefineClassContext beforeDefineClass(
+      ClassLoader classLoader, String className, ByteBuffer byteBuffer) {
+    // see how ClassLoader handles ByteBuffer
+    // https://github.com/openjdk/jdk11u/blob/487c3344fee3502b4843e7e11acceb77ad16100c/src/java.base/share/classes/java/lang/ClassLoader.java#L1095
+    int length = byteBuffer.remaining();
+    if (byteBuffer.hasArray()) {
+      return beforeDefineClass(
+          classLoader,
+          className,
+          byteBuffer.array(),
+          byteBuffer.position() + byteBuffer.arrayOffset(),
+          length);
+    } else {
+      byte[] classBytes = new byte[length];
+      byteBuffer.duplicate().get(classBytes);
+      return beforeDefineClass(classLoader, className, classBytes, 0, length);
+    }
+  }
+
+  public static void afterDefineClass(Handler.DefineClassContext context) {
+    handler.afterDefineClass(context);
+  }
+
+  /**
+   * Sets the {@link Handler} with callbacks to execute when {@code ClassLoader.defineClass} is
+   * called.
+   */
+  public static void internalSetHandler(Handler handler) {
+    if (DefineClassHelper.handler != null) {
+      // Only possible by misuse of this API, just ignore.
+      return;
+    }
+    DefineClassHelper.handler = handler;
+  }
+
+  private DefineClassHelper() {}
+}

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/AgentInstaller.java

@@ -24,6 +24,7 @@
 import io.opentelemetry.javaagent.bootstrap.AgentInitializer;
 import io.opentelemetry.javaagent.bootstrap.BootstrapPackagePrefixesHolder;
 import io.opentelemetry.javaagent.bootstrap.ClassFileTransformerHolder;
+import io.opentelemetry.javaagent.bootstrap.DefineClassHelper;
 import io.opentelemetry.javaagent.extension.AgentListener;
 import io.opentelemetry.javaagent.extension.ignore.IgnoredTypesConfigurer;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
@@ -120,6 +121,7 @@ public static ResettableClassFileTransformer installBytebuddyAgent(
     Config config = Config.get();
 
     setBootstrapPackages(config);
+    setDefineClassHandler();
 
     // If noop OpenTelemetry is enabled, autoConfiguredSdk will be null and AgentListeners are not
     // called
@@ -209,6 +211,10 @@ private static void setBootstrapPackages(Config config) {
     BootstrapPackagePrefixesHolder.setBoostrapPackagePrefixes(builder.build());
   }
 
+  private static void setDefineClassHandler() {
+    DefineClassHelper.internalSetHandler(DefineClassHandler.INSTANCE);
+  }
+
   private static void runBeforeAgentListeners(
       Iterable<AgentListener> agentListeners,
       Config config,

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/DefineClassHandler.java

@@ -0,0 +1,102 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling;
+
+import io.opentelemetry.javaagent.bootstrap.DefineClassHelper.Handler;
+import java.nio.charset.StandardCharsets;
+import org.objectweb.asm.ClassReader;
+
+public class DefineClassHandler implements Handler {
+  public static final DefineClassHandler INSTANCE = new DefineClassHandler();
+  private static final ThreadLocal<DefineClassContextImpl> defineClassContext =
+      ThreadLocal.withInitial(() -> DefineClassContextImpl.NOP);
+
+  private DefineClassHandler() {}
+
+  @Override
+  public DefineClassContext beforeDefineClass(
+      ClassLoader classLoader, String className, byte[] classBytes, int offset, int length) {
+    // with OpenJ9 class data sharing we don't get real class bytes
+    if (classBytes == null
+        || (classBytes.length == 40
+            && new String(classBytes, StandardCharsets.ISO_8859_1)
+                .startsWith("J9ROMCLASSCOOKIE"))) {
+      return null;
+    }
+
+    DefineClassContextImpl context = null;
+    // attempt to load super types of currently loaded class
+    // for a class to be loaded all of its super types must be loaded, here we just change the order
+    // of operations and load super types before transforming the bytes for current class so that
+    // we could use these super types for resolving the advice that needs to be applied to current
+    // class
+    try {
+      ClassReader cr = new ClassReader(classBytes, offset, length);
+      String superName = cr.getSuperName();
+      if (superName != null) {
+        Class.forName(superName.replace('/', '.'), false, classLoader);
+      }
+      String[] interfaces = cr.getInterfaces();
+      for (String interfaceName : interfaces) {
+        Class.forName(interfaceName.replace('/', '.'), false, classLoader);
+      }
+    } catch (Throwable throwable) {
+      // loading of super class or interface failed
+      // mark current class as failed to skip matching and transforming it
+      // we'll let defining the class proceed as usual so that it would throw the same exception as
+      // it does when running without the agent
+      context = DefineClassContextImpl.enter(className);
+    }
+    return context;
+  }
+
+  @Override
+  public void afterDefineClass(DefineClassContext context) {
+    if (context != null) {
+      context.exit();
+    }
+  }
+
+  /**
+   * Detect whether loading the specified class is known to fail.
+   *
+   * @param className class being loaded
+   * @return true if it is known that loading class with given name will fail
+   */
+  public static boolean isFailedClass(String className) {
+    DefineClassContextImpl context = defineClassContext.get();
+    return context.failedClassName != null && context.failedClassName.equals(className);
+  }
+
+  private static class DefineClassContextImpl implements DefineClassContext {
+    private static final DefineClassContextImpl NOP = new DefineClassContextImpl();
+
+    private final DefineClassContextImpl previous;
+    private final String failedClassName;
+
+    private DefineClassContextImpl() {
+      previous = null;
+      failedClassName = null;
+    }
+
+    private DefineClassContextImpl(DefineClassContextImpl previous, String failedClassName) {
+      this.previous = previous;
+      this.failedClassName = failedClassName;
+    }
+
+    static DefineClassContextImpl enter(String failedClassName) {
+      DefineClassContextImpl previous = defineClassContext.get();
+      DefineClassContextImpl context = new DefineClassContextImpl(previous, failedClassName);
+      defineClassContext.set(context);
+      return context;
+    }
+
+    @Override
+    public void exit() {
+      defineClassContext.set(previous);
+    }
+  }
+}

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/bootstrap/BootstrapProxyProviderImpl.java

@@ -0,0 +1,19 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.tooling.bootstrap;
+
+import com.google.auto.service.AutoService;
+import io.opentelemetry.javaagent.tooling.Utils;
+import io.opentelemetry.javaagent.tooling.muzzle.BootstrapProxyProvider;
+
+@AutoService(BootstrapProxyProvider.class)
+public class BootstrapProxyProviderImpl implements BootstrapProxyProvider {
+
+  @Override
+  public ClassLoader getBootstrapProxy() {
+    return Utils.getBootstrapProxy();
+  }
+}

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/field/FieldBackedImplementationInstaller.java

@@ -18,6 +18,8 @@
 import io.opentelemetry.javaagent.tooling.TransformSafeLogger;
 import io.opentelemetry.javaagent.tooling.instrumentation.InstrumentationModuleInstaller;
 import io.opentelemetry.javaagent.tooling.muzzle.VirtualFieldMappings;
+import io.opentelemetry.javaagent.tooling.util.IgnoreFailedTypeMatcher;
+import io.opentelemetry.javaagent.tooling.util.NamedMatcher;
 import java.lang.instrument.Instrumentation;
 import java.util.Collection;
 import java.util.HashSet;
@@ -28,6 +30,7 @@
 import net.bytebuddy.asm.AsmVisitorWrapper;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.dynamic.DynamicType;
+import net.bytebuddy.matcher.ElementMatcher;
 import net.bytebuddy.utility.JavaModule;
 
 /**
@@ -196,9 +199,15 @@ public AgentBuilder.Identified.Extendable injectFields(
            * For each virtual field defined in a current instrumentation we create an agent builder
            * that injects necessary fields.
            */
+          ElementMatcher<TypeDescription> typeMatcher =
+              new NamedMatcher<>(
+                  "VirtualField",
+                  new IgnoreFailedTypeMatcher(
+                      not(isAbstract()).and(hasSuperType(named(entry.getKey())))));
+
           builder =
               builder
-                  .type(not(isAbstract()).and(hasSuperType(named(entry.getKey()))))
+                  .type(typeMatcher)
                   .and(safeToInjectFieldsMatcher())
                   .and(InstrumentationModuleInstaller.NOT_DECORATOR_MATCHER)
                   .transform(NoOpTransformer.INSTANCE);