Use ClassLoaderValues to track class-loader state (#10075)

Author: mcculls

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/Agent.java

@@ -22,6 +22,7 @@
 import datadog.environment.OperatingSystem;
 import datadog.environment.SystemProperties;
 import datadog.instrument.classinject.ClassInjector;
+import datadog.instrument.utils.ClassLoaderValue;
 import datadog.trace.api.Config;
 import datadog.trace.api.Platform;
 import datadog.trace.api.StatsDClientManager;
@@ -102,6 +103,8 @@ public class Agent {
 
   private static final int DEFAULT_JMX_START_DELAY = 15; // seconds
 
+  private static final long CLASSLOADER_CLEAN_FREQUENCY_SECONDS = 30;
+
   private static final Logger log;
 
   private enum AgentFeature {
@@ -327,6 +330,7 @@ public static void start(
       startCrashTracking();
       StaticEventLogger.end("crashtracking");
     }
+
     startDatadogAgent(initTelemetry, inst);
 
     final EnumSet<Library> libraries = detectLibraries(log);
@@ -410,6 +414,15 @@ public static void start(
       StaticEventLogger.end("Profiling");
     }
 
+    // This task removes stale ClassLoaderValue entries where the class-loader is gone
+    // It only runs a couple of times a minute since class-loaders are rarely unloaded
+    AgentTaskScheduler.get()
+        .scheduleAtFixedRate(
+            ClassLoaderValue::removeStaleEntries,
+            CLASSLOADER_CLEAN_FREQUENCY_SECONDS,
+            CLASSLOADER_CLEAN_FREQUENCY_SECONDS,
+            TimeUnit.SECONDS);
+
     StaticEventLogger.end("Agent.start");
   }
 

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterState.java

@@ -1,7 +1,6 @@
 package datadog.trace.agent.tooling;
 
-import datadog.trace.api.cache.DDCache;
-import datadog.trace.api.cache.DDCaches;
+import datadog.instrument.utils.ClassLoaderValue;
 import java.util.Arrays;
 import java.util.concurrent.atomic.AtomicLongArray;
 import org.slf4j.Logger;
@@ -38,8 +37,13 @@ public interface Observer {
   private static long[] defaultState = {};
 
   /** Tracks which instrumentations were applied (per-class-loader) and which were blocked. */
-  private static final DDCache<ClassLoader, AtomicLongArray> classLoaderStates =
-      DDCaches.newFixedSizeWeakKeyCache(512);
+  private static final ClassLoaderValue<AtomicLongArray> classLoaderStates =
+      new ClassLoaderValue<AtomicLongArray>() {
+        @Override
+        protected AtomicLongArray computeValue(ClassLoader cl) {
+          return new AtomicLongArray(defaultState);
+        }
+      };
 
   private static Observer observer;
 
@@ -146,13 +150,7 @@ private static void updateState(ClassLoader classLoader, int instrumentationId,
   }
 
   private static AtomicLongArray classLoaderState(ClassLoader classLoader) {
-    return classLoaderStates.computeIfAbsent(
-        null != classLoader ? classLoader : Utils.getBootstrapProxy(),
-        InstrumenterState::buildClassLoaderState);
-  }
-
-  private static AtomicLongArray buildClassLoaderState(ClassLoader classLoader) {
-    return new AtomicLongArray(defaultState);
+    return classLoaderStates.get(classLoader);
   }
 
   private static int currentState(AtomicLongArray state, int bitIndex) {
@@ -173,7 +171,7 @@ public static String summary() {
     StringBuilder summary = new StringBuilder();
     classLoaderStates.visit(
         (loader, state) -> {
-          summary.append(loader.getClass().getName());
+          summary.append(loader != null ? loader.getClass().getName() : "<bootstrap>");
           summarizeState(summary, state);
           summary.append("\n\n");
         });

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

@@ -3,11 +3,10 @@
 import static datadog.trace.bootstrap.AgentClassLoading.LOCATING_CLASS;
 import static datadog.trace.util.Strings.getResourceName;
 
+import datadog.instrument.utils.ClassLoaderValue;
 import datadog.trace.agent.tooling.InstrumenterMetrics;
 import datadog.trace.agent.tooling.Utils;
 import datadog.trace.api.InstrumenterConfig;
-import datadog.trace.api.cache.DDCache;
-import datadog.trace.api.cache.DDCaches;
 import java.io.IOException;
 import java.io.InputStream;
 import java.lang.ref.WeakReference;
@@ -23,8 +22,13 @@
  * cannot find the desired resource, check up the classloader hierarchy until a resource is found.
  */
 public final class ClassFileLocators {
-  private static final DDCache<ClassLoader, DDClassFileLocator> classFileLocators =
-      DDCaches.newFixedSizeWeakKeyCache(64);
+  private static final ClassLoaderValue<DDClassFileLocator> classFileLocators =
+      new ClassLoaderValue<DDClassFileLocator>() {
+        @Override
+        protected DDClassFileLocator computeValue(ClassLoader cl) {
+          return new DDClassFileLocator(cl);
+        }
+      };
 
   private static final ClassFileLocator bootClassFileLocator =
       new ClassFileLocator() {
@@ -49,9 +53,7 @@ public void close() {
       };
 
   public static ClassFileLocator classFileLocator(final ClassLoader classLoader) {
-    return null != classLoader
-        ? classFileLocators.computeIfAbsent(classLoader, DDClassFileLocator::new)
-        : bootClassFileLocator;
+    return null != classLoader ? classFileLocators.get(classLoader) : bootClassFileLocator;
   }
 
   static final class DDClassFileLocator extends WeakReference<ClassLoader>

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

@@ -3,10 +3,9 @@
 import static datadog.trace.bootstrap.AgentClassLoading.PROBING_CLASSLOADER;
 import static net.bytebuddy.matcher.ElementMatchers.any;
 
+import datadog.instrument.utils.ClassLoaderValue;
 import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.api.Tracer;
-import datadog.trace.api.cache.DDCache;
-import datadog.trace.api.cache.DDCaches;
 import datadog.trace.bootstrap.PatchLogger;
 import datadog.trace.util.Strings;
 import java.util.ArrayList;
@@ -140,15 +139,21 @@ private static boolean loadsExpectedClass(
   static final List<String> hasClassResourceNames = new ArrayList<>();
 
   /** Cache of classloader-instance -> has-class mask. */
-  static final DDCache<ClassLoader, BitSet> hasClassCache = DDCaches.newFixedSizeWeakKeyCache(512);
+  static final ClassLoaderValue<BitSet> hasClassCache =
+      new ClassLoaderValue<BitSet>() {
+        @Override
+        protected BitSet computeValue(ClassLoader cl) {
+          return buildHasClassMask(cl);
+        }
+      };
 
   /** Distinct result used to mark an incompatible classloader that the tracer should skip. */
   static final BitSet INCOMPATIBLE_CLASS_LOADER = new BitSet();
 
   static final BitSet NO_CLASS_NAME_MATCHES = new BitSet();
 
   static BitSet hasClassMask(ClassLoader loader) {
-    return hasClassCache.computeIfAbsent(loader, ClassLoaderMatchers::buildHasClassMask);
+    return hasClassCache.get(loader);
   }
 
   static BitSet buildHasClassMask(ClassLoader loader) {