Faster type matching

Author: laurit

instrumentation/internal/internal-lambda/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/internal/lambda/InnerClassLambdaMetafactoryInstrumentation.java

@@ -8,8 +8,11 @@
 import static net.bytebuddy.matcher.ElementMatchers.named;
 
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
+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 net.bytebuddy.asm.Advice;
 import net.bytebuddy.asm.AsmVisitorWrapper;
 import net.bytebuddy.description.field.FieldDescription;
 import net.bytebuddy.description.field.FieldList;
@@ -62,6 +65,24 @@ public ClassVisitor wrap(
                         classVisitor, instrumentedType.getInternalName());
                   }
                 }));
+
+    transformer.applyAdviceToMethod(
+        named("spinInnerClass"),
+        InnerClassLambdaMetafactoryInstrumentation.class.getName()
+            + (hasInterfaceClassField() ? "$LambdaJdk17Advice" : "$LambdaAdvice"));
+  }
+
+  @SuppressWarnings("ReturnValueIgnored")
+  private static boolean hasInterfaceClassField() {
+    try {
+      Class<?> clazz = Class.forName("java.lang.invoke.AbstractValidatingLambdaMetafactory");
+      clazz.getDeclaredField("interfaceClass");
+      return true;
+    } catch (NoSuchFieldException exception) {
+      return false;
+    } catch (ClassNotFoundException exception) {
+      throw new IllegalStateException(exception);
+    }
   }
 
   private static class MetaFactoryClassVisitor extends ClassVisitor {
@@ -117,4 +138,34 @@ public void visitMethodInsn(
       return mv;
     }
   }
+
+  @SuppressWarnings("unused")
+  public static class LambdaAdvice {
+
+    @Advice.OnMethodEnter
+    public static DefineClassContext onEnter(
+        @Advice.FieldValue("samBase") Class<?> lambdaInterface) {
+      return DefineClassHelper.beforeDefineLambdaClass(lambdaInterface);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class)
+    public static void onExit(@Advice.Enter DefineClassContext context) {
+      DefineClassHelper.afterDefineClass(context);
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class LambdaJdk17Advice {
+
+    @Advice.OnMethodEnter
+    public static DefineClassContext onEnter(
+        @Advice.FieldValue("interfaceClass") Class<?> lambdaInterface) {
+      return DefineClassHelper.beforeDefineLambdaClass(lambdaInterface);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class)
+    public static void onExit(@Advice.Enter DefineClassContext context) {
+      DefineClassHelper.afterDefineClass(context);
+    }
+  }
 }

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

@@ -14,6 +14,8 @@ public interface Handler {
     DefineClassContext beforeDefineClass(
         ClassLoader classLoader, String className, byte[] classBytes, int offset, int length);
 
+    DefineClassContext beforeDefineLambdaClass(Class<?> lambdaInterface);
+
     void afterDefineClass(DefineClassContext context);
 
     /** Context returned from {@code beforeDefineClass} and passed to {@code afterDefineClass}. */
@@ -48,6 +50,10 @@ public static Handler.DefineClassContext beforeDefineClass(
     }
   }
 
+  public static Handler.DefineClassContext beforeDefineLambdaClass(Class<?> lambdaInterface) {
+    return handler.beforeDefineLambdaClass(lambdaInterface);
+  }
+
   public static void afterDefineClass(Handler.DefineClassContext context) {
     handler.afterDefineClass(context);
   }

javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/matcher/ClassLoaderHasClassesNamedMatcher.java

@@ -12,6 +12,8 @@
 import java.util.List;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
 import net.bytebuddy.matcher.ElementMatcher;
 
 class ClassLoaderHasClassesNamedMatcher extends ElementMatcher.Junction.AbstractBase<ClassLoader> {
@@ -62,12 +64,14 @@ private static boolean hasResources(ClassLoader cl, String... resources) {
   }
 
   private static class Manager {
-    private static final BitSet EMPTY = new BitSet(0);
     static final Manager INSTANCE = new Manager();
     private final List<ClassLoaderHasClassesNamedMatcher> matchers = new CopyOnWriteArrayList<>();
-    // each matcher gets a bit in BitSet, that bit indicates whether current matcher matched or not
-    // for given class loader
+    // each matcher gets a two bits in BitSet, that first bit indicates whether current matcher has
+    // been run for given class loader and the second whether it matched or not
     private final Cache<ClassLoader, BitSet> enabled = Cache.weak();
+    private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
+    private final Lock readLock = lock.readLock();
+    private final Lock writeLock = lock.writeLock();
     private volatile boolean matchCalled = false;
 
     Manager() {
@@ -83,20 +87,36 @@ void add(ClassLoaderHasClassesNamedMatcher matcher) {
 
     boolean match(ClassLoaderHasClassesNamedMatcher matcher, ClassLoader cl) {
       matchCalled = true;
-      BitSet set = enabled.get(cl);
-      if (set == null) {
-        set = new BitSet(counter.get());
-        for (ClassLoaderHasClassesNamedMatcher m : matchers) {
-          if (hasResources(cl, m.resources)) {
-            // set the bit corresponding to the matcher when it matched
-            set.set(m.index);
+      BitSet set = enabled.computeIfAbsent(cl, (unused) -> new BitSet(counter.get() * 2));
+      int matcherRunBit = 2 * matcher.index;
+      int matchedBit = matcherRunBit + 1;
+      readLock.lock();
+      try {
+        if (!set.get(matcherRunBit)) {
+          // read lock needs to be released before upgrading to write lock
+          readLock.unlock();
+          // we do the resource presence check outside the lock to keep the time we need to hold
+          // the write lock minimal
+          boolean matches = hasResources(cl, matcher.resources);
+          writeLock.lock();
+          try {
+            if (!set.get(matcherRunBit)) {
+              if (matches) {
+                set.set(matchedBit);
+              }
+              set.set(matcherRunBit);
+            }
+          } finally {
+            // downgrading the write lock to the read lock
+            readLock.lock();
+            writeLock.unlock();
           }
         }
-        enabled.put(cl, set.isEmpty() ? EMPTY : set);
-      } else if (set.isEmpty()) {
-        return false;
+
+        return set.get(matchedBit);
+      } finally {
+        readLock.unlock();
       }
-      return set.get(matcher.index);
     }
   }
 }

javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/matcher/SafeErasureMatcher.java

@@ -8,6 +8,7 @@
 import static io.opentelemetry.javaagent.extension.matcher.Utils.safeTypeDefinitionName;
 import static java.util.logging.Level.FINE;
 
+import io.opentelemetry.javaagent.extension.matcher.internal.DelegatingMatcher;
 import java.util.logging.Logger;
 import javax.annotation.Nullable;
 import net.bytebuddy.description.type.TypeDefinition;
@@ -24,7 +25,8 @@
  * @param <T> The type of the matched entity.
  * @see net.bytebuddy.matcher.ErasureMatcher
  */
-class SafeErasureMatcher<T extends TypeDefinition> extends ElementMatcher.Junction.AbstractBase<T> {
+class SafeErasureMatcher<T extends TypeDefinition> extends ElementMatcher.Junction.AbstractBase<T>
+    implements DelegatingMatcher {
 
   private static final Logger logger = Logger.getLogger(SafeErasureMatcher.class.getName());
 
@@ -88,4 +90,9 @@ public boolean equals(@Nullable Object obj) {
   public int hashCode() {
     return matcher.hashCode();
   }
+
+  @Override
+  public ElementMatcher<?> getDelegate() {
+    return matcher;
+  }
 }

javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/matcher/SafeExtendsClassMatcher.java

@@ -7,12 +7,14 @@
 
 import static io.opentelemetry.javaagent.extension.matcher.SafeHasSuperTypeMatcher.safeGetSuperClass;
 
+import io.opentelemetry.javaagent.extension.matcher.internal.DelegatingSuperTypeMatcher;
 import javax.annotation.Nullable;
 import net.bytebuddy.description.type.TypeDefinition;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 
-class SafeExtendsClassMatcher extends ElementMatcher.Junction.AbstractBase<TypeDescription> {
+class SafeExtendsClassMatcher extends ElementMatcher.Junction.AbstractBase<TypeDescription>
+    implements DelegatingSuperTypeMatcher {
 
   private final ElementMatcher<TypeDescription.Generic> matcher;
 
@@ -55,4 +57,9 @@ public boolean equals(@Nullable Object obj) {
   public int hashCode() {
     return matcher.hashCode();
   }
+
+  @Override
+  public ElementMatcher<?> getDelegate() {
+    return matcher;
+  }
 }

javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/matcher/SafeHasSuperTypeMatcher.java

@@ -9,6 +9,7 @@
 import static io.opentelemetry.javaagent.extension.matcher.Utils.safeTypeDefinitionName;
 import static java.util.logging.Level.FINE;
 
+import io.opentelemetry.javaagent.extension.matcher.internal.DelegatingSuperTypeMatcher;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.Set;
@@ -35,7 +36,8 @@
  *
  * @see net.bytebuddy.matcher.HasSuperTypeMatcher
  */
-class SafeHasSuperTypeMatcher extends ElementMatcher.Junction.AbstractBase<TypeDescription> {
+class SafeHasSuperTypeMatcher extends ElementMatcher.Junction.AbstractBase<TypeDescription>
+    implements DelegatingSuperTypeMatcher {
 
   private static final Logger logger = Logger.getLogger(SafeHasSuperTypeMatcher.class.getName());
 
@@ -136,6 +138,11 @@ public int hashCode() {
     return matcher.hashCode();
   }
 
+  @Override
+  public ElementMatcher<?> getDelegate() {
+    return matcher;
+  }
+
   /**
    * TypeDefinition#getInterfaces() produces an iterator which may throw an exception during
    * iteration if an interface is absent from the classpath.

javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/matcher/internal/DelegatingMatcher.java

@@ -0,0 +1,20 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.extension.matcher.internal;
+
+import net.bytebuddy.matcher.ElementMatcher;
+
+/**
+ * Interface for extracting the matcher that the given matcher delegates to.
+ *
+ * <p>This class is internal and is hence not for public use. Its APIs are unstable and can change
+ * at any time.
+ */
+public interface DelegatingMatcher {
+
+  /** Returns the matcher that the current matcher delegates to. */
+  ElementMatcher<?> getDelegate();
+}

javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/matcher/internal/DelegatingSuperTypeMatcher.java

@@ -0,0 +1,14 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.extension.matcher.internal;
+
+/**
+ * Marker interface for delegating matchers that match based on the type hierarchy.
+ *
+ * <p>This class is internal and is hence not for public use. Its APIs are unstable and can change
+ * at any time.
+ */
+public interface DelegatingSuperTypeMatcher extends DelegatingMatcher {}

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

@@ -54,6 +54,7 @@
 import javax.annotation.Nullable;
 import net.bytebuddy.ByteBuddy;
 import net.bytebuddy.agent.builder.AgentBuilder;
+import net.bytebuddy.agent.builder.AgentBuilderUtil;
 import net.bytebuddy.agent.builder.ResettableClassFileTransformer;
 import net.bytebuddy.description.type.TypeDefinition;
 import net.bytebuddy.description.type.TypeDescription;
@@ -184,6 +185,7 @@ private static void installBytebuddyAgent(
     }
     logger.log(FINE, "Installed {0} extension(s)", numberOfLoadedExtensions);
 
+    agentBuilder = AgentBuilderUtil.optimize(agentBuilder);
     ResettableClassFileTransformer resettableClassFileTransformer = agentBuilder.installOn(inst);
     ClassFileTransformerHolder.setClassFileTransformer(resettableClassFileTransformer);