Look up helper class bytes when they are needed (#7839)

We can save a bit of heap space by not keeping the class bytes in
memory, we can easily look them up when needed.

Author: laurit

muzzle/src/main/java/io/opentelemetry/javaagent/tooling/HelperInjector.java

@@ -28,6 +28,7 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Supplier;
 import javax.annotation.Nullable;
 import net.bytebuddy.agent.builder.AgentBuilder.Transformer;
 import net.bytebuddy.description.type.TypeDescription;
@@ -88,7 +89,7 @@ Class<?> inject(ClassLoader classLoader, String className) {
   private final List<HelperResource> helperResources;
   @Nullable private final ClassLoader helpersSource;
   @Nullable private final Instrumentation instrumentation;
-  private final Map<String, byte[]> dynamicTypeMap = new LinkedHashMap<>();
+  private final Map<String, Supplier<byte[]>> dynamicTypeMap = new LinkedHashMap<>();
 
   private final Cache<ClassLoader, Boolean> injectedClassLoaders = Cache.weak();
   private final Cache<ClassLoader, Boolean> resourcesInjectedClassLoaders = Cache.weak();
@@ -108,7 +109,6 @@ public HelperInjector(
       String requestingName,
       List<String> helperClassNames,
       List<HelperResource> helperResources,
-      // TODO can this be replaced with the context class loader?
       ClassLoader helpersSource,
       Instrumentation instrumentation) {
     this.requestingName = requestingName;
@@ -120,7 +120,9 @@ public HelperInjector(
   }
 
   private HelperInjector(
-      String requestingName, Map<String, byte[]> helperMap, Instrumentation instrumentation) {
+      String requestingName,
+      Map<String, Supplier<byte[]>> helperMap,
+      Instrumentation instrumentation) {
     this.requestingName = requestingName;
 
     this.helperClassNames = helperMap.keySet();
@@ -135,9 +137,9 @@ public static HelperInjector forDynamicTypes(
       String requestingName,
       Collection<DynamicType.Unloaded<?>> helpers,
       Instrumentation instrumentation) {
-    Map<String, byte[]> bytes = new HashMap<>(helpers.size());
+    Map<String, Supplier<byte[]>> bytes = new HashMap<>(helpers.size());
     for (DynamicType.Unloaded<?> helper : helpers) {
-      bytes.put(helper.getTypeDescription().getName(), helper.getBytes());
+      bytes.put(helper.getTypeDescription().getName(), helper::getBytes);
     }
     return new HelperInjector(requestingName, bytes, instrumentation);
   }
@@ -146,18 +148,27 @@ public static void setHelperInjectorListener(HelperInjectorListener listener) {
     helperInjectorListener = listener;
   }
 
-  private Map<String, byte[]> getHelperMap() throws IOException {
+  private Map<String, Supplier<byte[]>> getHelperMap() {
     if (dynamicTypeMap.isEmpty()) {
-      Map<String, byte[]> classnameToBytes = new LinkedHashMap<>();
-
-      ClassFileLocator locator = ClassFileLocator.ForClassLoader.of(helpersSource);
+      Map<String, Supplier<byte[]>> result = new LinkedHashMap<>();
 
       for (String helperClassName : helperClassNames) {
-        byte[] classBytes = locator.locate(helperClassName).resolve();
-        classnameToBytes.put(helperClassName, classBytes);
+        result.put(
+            helperClassName,
+            () -> {
+              try (ClassFileLocator locator = ClassFileLocator.ForClassLoader.of(helpersSource)) {
+                return locator.locate(helperClassName).resolve();
+              } catch (IOException exception) {
+                if (logger.isLoggable(SEVERE)) {
+                  logger.log(
+                      SEVERE, "Failed to read {0}", new Object[] {helperClassName}, exception);
+                }
+                throw new IllegalStateException("Failed to read " + helperClassName, exception);
+              }
+            });
       }
 
-      return classnameToBytes;
+      return result;
     } else {
       return dynamicTypeMap;
     }
@@ -248,10 +259,10 @@ private void injectHelperClasses(TypeDescription typeDescription, ClassLoader cl
                   new Object[] {cl, helperClassNames});
             }
 
-            Map<String, byte[]> classnameToBytes = getHelperMap();
+            Map<String, Supplier<byte[]>> classnameToBytes = getHelperMap();
             Map<String, HelperClassInjector> map =
                 helperInjectors.computeIfAbsent(cl, (unused) -> new ConcurrentHashMap<>());
-            for (Map.Entry<String, byte[]> entry : classnameToBytes.entrySet()) {
+            for (Map.Entry<String, Supplier<byte[]>> entry : classnameToBytes.entrySet()) {
               // for boot loader we use a placeholder injector, we only need these classes to be
               // in the injected classes map to later tell which of the classes are injected
               HelperClassInjector injector =
@@ -280,9 +291,18 @@ private void injectHelperClasses(TypeDescription typeDescription, ClassLoader cl
         });
   }
 
-  private Map<String, Class<?>> injectBootstrapClassLoader(Map<String, byte[]> classnameToBytes)
+  private static Map<String, byte[]> resolve(Map<String, Supplier<byte[]>> classes) {
+    Map<String, byte[]> result = new LinkedHashMap<>();
+    for (Map.Entry<String, Supplier<byte[]>> entry : classes.entrySet()) {
+      result.put(entry.getKey(), entry.getValue().get());
+    }
+    return result;
+  }
+
+  private Map<String, Class<?>> injectBootstrapClassLoader(Map<String, Supplier<byte[]>> inject)
       throws IOException {
 
+    Map<String, byte[]> classnameToBytes = resolve(inject);
     if (helperInjectorListener != null) {
       helperInjectorListener.onInjection(classnameToBytes);
     }
@@ -358,16 +378,16 @@ public static Class<?> loadHelperClass(ClassLoader classLoader, String className
   }
 
   private static class HelperClassInjector {
-    private final byte[] bytes;
+    private final Supplier<byte[]> bytes;
 
-    HelperClassInjector(byte[] bytes) {
+    HelperClassInjector(Supplier<byte[]> bytes) {
       this.bytes = bytes;
     }
 
     Class<?> inject(ClassLoader classLoader, String className) {
       Map<String, Class<?>> result =
           new ClassInjector.UsingReflection(classLoader)
-              .injectRaw(Collections.singletonMap(className, bytes));
+              .injectRaw(Collections.singletonMap(className, bytes.get()));
       return result.get(className);
     }
   }