Use ClassLoaderIndex to index class-loader related information

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

@@ -1,8 +1,5 @@
 package datadog.trace.agent.tooling.bytebuddy;
 
-import datadog.trace.api.cache.DDCache;
-import datadog.trace.api.cache.DDCaches;
-import java.lang.ref.WeakReference;
 import java.net.URL;
 import java.util.Arrays;
 import java.util.Objects;
@@ -65,12 +62,13 @@ public SharedTypeInfo<T> find(String className) {
    *
    * @return previously shared information for the named type
    */
-  public SharedTypeInfo<T> share(String className, ClassLoader loader, URL classFile, T typeInfo) {
+  public SharedTypeInfo<T> share(
+      String className, int classLoaderKeyId, URL classFile, T typeInfo) {
     SharedTypeInfo<T> newValue;
     if (namesAreUnique) {
       newValue = new SharedTypeInfo<>(className, typeInfo);
     } else {
-      newValue = new DisambiguatingTypeInfo<>(className, typeInfo, loader, classFile);
+      newValue = new DisambiguatingTypeInfo<>(className, typeInfo, classLoaderKeyId, classFile);
     }
 
     int nameHash = className.hashCode();
@@ -117,7 +115,7 @@ public static class SharedTypeInfo<T> {
       this.typeInfo = typeInfo;
     }
 
-    public boolean sameClassLoader(ClassLoader loader) {
+    public boolean sameClassLoader(int classLoaderKeyId) {
       return true;
     }
 
@@ -132,25 +130,18 @@ public final T get() {
 
   /** Includes the classloader and class file resource it originated from. */
   static final class DisambiguatingTypeInfo<T> extends SharedTypeInfo<T> {
-    private static final ClassLoader BOOTSTRAP_LOADER = null;
-    private static final LoaderId BOOTSTRAP_LOADER_ID = null;
-
-    private static final DDCache<ClassLoader, LoaderId> loaderIds =
-        DDCaches.newFixedSizeWeakKeyCache(64);
-
-    private final LoaderId loaderId;
+    private final int classLoaderKeyId;
     private final URL classFile;
 
-    DisambiguatingTypeInfo(String className, T typeInfo, ClassLoader loader, URL classFile) {
+    DisambiguatingTypeInfo(String className, T typeInfo, int classLoaderKeyId, URL classFile) {
       super(className, typeInfo);
-      this.loaderId = loaderId(loader);
+      this.classLoaderKeyId = classLoaderKeyId;
       this.classFile = classFile;
     }
 
-    public boolean sameClassLoader(ClassLoader loader) {
-      return BOOTSTRAP_LOADER_ID == loaderId
-          ? BOOTSTRAP_LOADER == loader
-          : loaderId.sameClassLoader(loader);
+    @Override
+    public boolean sameClassLoader(int classLoaderKeyId) {
+      return this.classLoaderKeyId == classLoaderKeyId;
     }
 
     public boolean sameClassFile(URL classFile) {
@@ -159,12 +150,6 @@ public boolean sameClassFile(URL classFile) {
           && sameClassFile(this.classFile, classFile);
     }
 
-    private static LoaderId loaderId(ClassLoader loader) {
-      return BOOTSTRAP_LOADER == loader
-          ? BOOTSTRAP_LOADER_ID
-          : loaderIds.computeIfAbsent(loader, LoaderId::new);
-    }
-
     /** Matches class file resources without triggering network lookups. */
     private static boolean sameClassFile(URL lhs, URL rhs) {
       return Objects.equals(lhs.getFile(), rhs.getFile())
@@ -173,18 +158,4 @@ private static boolean sameClassFile(URL lhs, URL rhs) {
           && Objects.equals(lhs.getProtocol(), rhs.getProtocol());
     }
   }
-
-  /** Supports classloader comparisons without strongly referencing the classloader. */
-  static final class LoaderId extends WeakReference<ClassLoader> {
-    private final int loaderHash;
-
-    LoaderId(ClassLoader loader) {
-      super(loader);
-      this.loaderHash = System.identityHashCode(loader);
-    }
-
-    boolean sameClassLoader(ClassLoader loader) {
-      return loaderHash == System.identityHashCode(loader) && loader == get();
-    }
-  }
 }

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/memoize/Memoizer.java

@@ -238,10 +238,10 @@ static BitSet doMemoize(TypeDescription type, Map<String, BitSet> localMemos) {
 
     // otherwise share result for this location (other locations may have different results)
     if (namesAreUnique || name.startsWith("java.") || !(type instanceof WithLocation)) {
-      memos.share(name, null, null, memo);
+      memos.share(name, 0, null, memo);
     } else {
       WithLocation origin = (WithLocation) type;
-      memos.share(name, origin.getClassLoader(), origin.getClassFile(), memo);
+      memos.share(name, origin.getClassLoaderKeyId(), origin.getClassFile(), memo);
     }
 
     return memo;
@@ -254,7 +254,7 @@ static boolean potentialMatch(String name) {
 
   private static boolean sameOrigin(TypeDescription type, SharedTypeInfo<BitSet> sharedMemo) {
     return !(type instanceof WithLocation)
-        || sharedMemo.sameClassLoader(((WithLocation) type).getClassLoader())
+        || sharedMemo.sameClassLoader(((WithLocation) type).getClassLoaderKeyId())
         || sharedMemo.sameClassFile(((WithLocation) type).getClassFile());
   }
 

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

@@ -5,6 +5,7 @@
 import static datadog.trace.bootstrap.AgentClassLoading.LOCATING_CLASS;
 import static net.bytebuddy.dynamic.loading.ClassLoadingStrategy.BOOTSTRAP_LOADER;
 
+import datadog.instrument.utils.ClassLoaderIndex;
 import datadog.trace.agent.tooling.InstrumenterMetrics;
 import datadog.trace.agent.tooling.bytebuddy.ClassFileLocators;
 import datadog.trace.agent.tooling.bytebuddy.TypeInfoCache;
@@ -99,7 +100,9 @@ final class TypeFactory {
 
   ClassLoader originalClassLoader;
 
-  ClassLoader classLoader;
+  ClassLoader currentClassLoader;
+
+  int currentClassLoaderKeyId;
 
   ClassFileLocator classFileLocator;
 
@@ -109,16 +112,17 @@ final class TypeFactory {
 
   /** Sets the current class-loader context of this type-factory. */
   void switchContext(ClassLoader classLoader) {
-    if (this.classLoader != classLoader || null == classFileLocator) {
-      this.classLoader = classLoader;
+    if (currentClassLoader != classLoader || null == classFileLocator) {
+      currentClassLoader = classLoader;
+      currentClassLoaderKeyId = -1;
       classFileLocator = classFileLocator(classLoader);
       // clear local type cache whenever the class-loader context changes
       deferredTypes.clear();
     }
   }
 
   ClassLoader currentContext() {
-    return classLoader;
+    return currentClassLoader;
   }
 
   void beginInstall() {
@@ -149,7 +153,7 @@ void beginTransform(String name, byte[] bytecode) {
     targetBytecode = bytecode;
 
     if (installing) {
-      originalClassLoader = classLoader;
+      originalClassLoader = currentClassLoader;
     }
   }
 
@@ -185,7 +189,8 @@ void endTransform() {
 
   private void clearReferences() {
     if (null != classFileLocator) {
-      classLoader = null;
+      currentClassLoader = null;
+      currentClassLoaderKeyId = -1;
       classFileLocator = null;
       deferredTypes.clear();
     }
@@ -247,13 +252,14 @@ TypeDescription resolveType(LazyType request) {
   private TypeDescription lookupType(
       LazyType request, TypeInfoCache<TypeDescription> types, TypeParser typeParser) {
     String name = request.name;
+    int classLoaderKeyId = request.getClassLoaderKeyId();
     boolean isOutline = typeParser == outlineTypeParser;
     long fromTick = InstrumenterMetrics.tick();
 
     // existing type description from same classloader?
     SharedTypeInfo<TypeDescription> sharedType = types.find(name);
     if (null != sharedType
-        && (name.startsWith("java.") || sharedType.sameClassLoader(classLoader))) {
+        && (name.startsWith("java.") || sharedType.sameClassLoader(classLoaderKeyId))) {
       InstrumenterMetrics.reuseTypeDescription(fromTick, isOutline);
       return sharedType.get();
     }
@@ -285,7 +291,7 @@ private TypeDescription lookupType(
     }
 
     // share result, whether we found it or not
-    types.share(name, classLoader, classFile, type);
+    types.share(name, classLoaderKeyId, classFile, type);
 
     return type;
   }
@@ -295,17 +301,17 @@ private TypeDescription loadType(String name, TypeParser typeParser) {
     LOCATING_CLASS.begin();
     try {
       Class<?> loadedType;
-      if (BOOTSTRAP_LOADER == classLoader) {
+      if (BOOTSTRAP_LOADER == currentClassLoader) {
         loadedType = Class.forName(name, false, BOOTSTRAP_LOADER);
-      } else if (skipLoadClass(classLoader.getClass().getName())) {
+      } else if (skipLoadClass(currentClassLoader.getClass().getName())) {
         return null; // avoid known problematic class-loaders
       } else {
-        loadedType = classLoader.loadClass(name);
+        loadedType = currentClassLoader.loadClass(name);
       }
       log.debug(
           "Direct loadClass type resolution of {} from class loader {} bypasses transformation",
           name,
-          classLoader);
+          currentClassLoader);
       return typeParser.parse(loadedType);
     } catch (Throwable ignored) {
       return null;
@@ -331,8 +337,11 @@ final class LazyType extends WithName implements WithLocation {
     }
 
     @Override
-    public ClassLoader getClassLoader() {
-      return classLoader;
+    public int getClassLoaderKeyId() {
+      if (currentClassLoaderKeyId < 0) {
+        currentClassLoaderKeyId = ClassLoaderIndex.getClassLoaderKeyId(currentClassLoader);
+      }
+      return currentClassLoaderKeyId;
     }
 
     @Override

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

@@ -4,7 +4,7 @@
 
 /** Provides details of where the resolved type was defined. */
 public interface WithLocation {
-  ClassLoader getClassLoader();
+  int getClassLoaderKeyId();
 
   URL getClassFile();