Rework the bytecode cache for extensions so it works across connectio…

…ns (#7133)

The shared cache can be small because extension requests are typically localized.

dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/ExtensionFinder.java

@@ -89,34 +89,34 @@ private static URL findExtensionURL(JarFile jar, String[] descriptors) {
   /** Builds a URL that uses an {@link ExtensionHandler} to access the extension. */
   private static URL buildExtensionURL(JarFile jar, ExtensionHandler handler) {
     try {
-      return new URL(
-          "dd-ext",
-          null,
-          -1,
-          "/",
-          new URLStreamHandler() {
-            @Override
-            protected URLConnection openConnection(URL url) throws IOException {
-              return openExtension(url, jar, handler);
-            }
-          });
+      return new URL("dd-ext", null, -1, "/", new StreamMapper(jar, handler));
     } catch (MalformedURLException ignore) {
       return null;
     }
   }
 
-  /** Uses the given {@link ExtensionHandler} to access content from the extension. */
-  static URLConnection openExtension(URL url, JarFile jar, ExtensionHandler handler)
-      throws IOException {
-    String file = url.getFile();
-    if (!file.isEmpty() && file.charAt(0) == '/') {
-      file = file.substring(1);
+  /** Uses a {@link ExtensionHandler} to map and stream content from the extension. */
+  static final class StreamMapper extends URLStreamHandler {
+    private final JarFile jar;
+    private final ExtensionHandler handler;
+
+    StreamMapper(JarFile jar, ExtensionHandler handler) {
+      this.jar = jar;
+      this.handler = handler;
     }
-    JarEntry jarEntry = handler.mapEntry(jar, file);
-    if (null != jarEntry) {
-      return handler.mapContent(url, jar, jarEntry);
-    } else {
-      throw new FileNotFoundException("JAR entry " + file + " not found in " + jar.getName());
+
+    @Override
+    protected URLConnection openConnection(URL url) throws IOException {
+      String file = url.getFile();
+      if (!file.isEmpty() && file.charAt(0) == '/') {
+        file = file.substring(1);
+      }
+      JarEntry jarEntry = handler.mapEntry(jar, file);
+      if (null != jarEntry) {
+        return handler.mapContent(url, jar, jarEntry);
+      } else {
+        throw new FileNotFoundException("JAR entry " + file + " not found in " + jar.getName());
+      }
     }
   }
 

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

@@ -1,8 +1,11 @@
 package datadog.trace.agent.tooling;
 
+import datadog.trace.api.cache.DDCache;
+import datadog.trace.api.cache.DDCaches;
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.UncheckedIOException;
 import java.net.URL;
 import java.net.URLConnection;
 import java.util.function.Function;
@@ -76,8 +79,9 @@ public long getContentLengthLong() {
 
   /** Provides access to bytecode mapped from the extension. */
   protected static class ClassMappingConnection extends JarFileConnection {
+    private static final DDCache<String, byte[]> BYTECODE_CACHE = DDCaches.newFixedSizeCache(32);
+
     private final Function<ClassVisitor, ClassVisitor> mapping;
-    private byte[] cachedBytecode;
 
     public ClassMappingConnection(
         URL url, JarFile jar, JarEntry entry, Function<ClassVisitor, ClassVisitor> mapping) {
@@ -87,28 +91,35 @@ public ClassMappingConnection(
 
     @Override
     public InputStream getInputStream() throws IOException {
-      return new ByteArrayInputStream(mapBytecode());
+      try {
+        return new ByteArrayInputStream(mapBytecode());
+      } catch (UncheckedIOException e) {
+        throw e.getCause();
+      }
     }
 
     @Override
     public long getContentLengthLong() {
       try {
         return mapBytecode().length;
-      } catch (IOException e) {
+      } catch (UncheckedIOException e) {
         return -1;
       }
     }
 
-    private byte[] mapBytecode() throws IOException {
-      if (null == cachedBytecode) {
-        try (InputStream in = super.getInputStream()) {
-          ClassReader cr = new ClassReader(in);
-          ClassWriter cw = new ClassWriter(cr, 0);
-          cr.accept(mapping.apply(cw), 0);
-          cachedBytecode = cw.toByteArray();
-        }
+    private byte[] mapBytecode() {
+      return BYTECODE_CACHE.computeIfAbsent(url.getFile(), this::doMapBytecode);
+    }
+
+    private byte[] doMapBytecode(String unused) {
+      try (InputStream in = super.getInputStream()) {
+        ClassReader cr = new ClassReader(in);
+        ClassWriter cw = new ClassWriter(cr, 0);
+        cr.accept(mapping.apply(cw), 0);
+        return cw.toByteArray();
+      } catch (IOException e) {
+        throw new UncheckedIOException(e);
       }
-      return cachedBytecode;
     }
   }