Define packages in exporter class loader (#409)

Author: trask

agent-tooling/src/main/java/io/opentelemetry/auto/tooling/ExporterClassLoader.java

@@ -22,15 +22,19 @@
 import java.net.URL;
 import java.net.URLClassLoader;
 import java.util.Enumeration;
+import java.util.jar.JarFile;
+import java.util.jar.Manifest;
+import lombok.extern.slf4j.Slf4j;
 import net.bytebuddy.jar.asm.ClassReader;
 import net.bytebuddy.jar.asm.ClassWriter;
 import net.bytebuddy.jar.asm.commons.ClassRemapper;
 
+@Slf4j
 public class ExporterClassLoader extends URLClassLoader {
   // We need to prefix the names to prevent the gradle shadowJar relocation rules from touching
   // them. It's possible to do this by excluding this class from shading, but it may cause issue
   // with transitive dependencies down the line.
-  private final ShadingRemapper remapper =
+  private static final ShadingRemapper remapper =
       new ShadingRemapper(
           rule(
               "#io.opentelemetry.OpenTelemetry",
@@ -52,8 +56,11 @@ public class ExporterClassLoader extends URLClassLoader {
           rule("#java.util.logging.Logger", "#io.opentelemetry.auto.bootstrap.PatchLogger"),
           rule("#org.slf4j", "#io.opentelemetry.auto.slf4j"));
 
-  public ExporterClassLoader(final URL[] urls, final ClassLoader parent) {
-    super(urls, parent);
+  private final Manifest manifest;
+
+  public ExporterClassLoader(final URL url, final ClassLoader parent) {
+    super(new URL[] {url}, parent);
+    this.manifest = getManifest(url);
   }
 
   @Override
@@ -69,16 +76,79 @@ public Enumeration<URL> getResources(final String name) throws IOException {
 
   @Override
   protected Class<?> findClass(final String name) throws ClassNotFoundException {
-
     // Use resource loading to get the class as a stream of bytes, then use ASM to transform it.
-    try (final InputStream in = getResourceAsStream(name.replace('.', '/') + ".class")) {
-      final ClassWriter cw = new ClassWriter(0);
-      final ClassReader cr = new ClassReader(in);
-      cr.accept(new ClassRemapper(cw, remapper), ClassReader.EXPAND_FRAMES);
-      final byte[] bytes = cw.toByteArray();
+    InputStream in = getResourceAsStream(name.replace('.', '/') + ".class");
+    if (in == null) {
+      throw new ClassNotFoundException(name);
+    }
+    try {
+      final byte[] bytes = remapClassBytes(in);
+      definePackageIfNeeded(name);
       return defineClass(name, bytes, 0, bytes.length);
     } catch (final IOException e) {
-      throw new ClassNotFoundException(name);
+      throw new ClassNotFoundException(name, e);
+    } finally {
+      try {
+        in.close();
+      } catch (IOException e) {
+        log.debug(e.getMessage(), e);
+      }
+    }
+  }
+
+  private void definePackageIfNeeded(String className) {
+    String packageName = getPackageName(className);
+    if (packageName == null) {
+      // default package
+      return;
+    }
+    if (isPackageDefined(packageName)) {
+      // package has already been defined
+      return;
+    }
+    try {
+      definePackage(packageName);
+    } catch (IllegalArgumentException e) {
+      // this exception is thrown when the package has already been defined, which is possible due
+      // to race condition with the check above
+      if (!isPackageDefined(packageName)) {
+        // this shouldn't happen however
+        log.error(e.getMessage(), e);
+      }
+    }
+  }
+
+  private boolean isPackageDefined(String packageName) {
+    return getPackage(packageName) != null;
+  }
+
+  private void definePackage(String packageName) {
+    if (manifest == null) {
+      definePackage(packageName, null, null, null, null, null, null, null);
+    } else {
+      definePackage(packageName, manifest, null);
+    }
+  }
+
+  private static byte[] remapClassBytes(InputStream in) throws IOException {
+    final ClassWriter cw = new ClassWriter(0);
+    final ClassReader cr = new ClassReader(in);
+    cr.accept(new ClassRemapper(cw, remapper), ClassReader.EXPAND_FRAMES);
+    return cw.toByteArray();
+  }
+
+  private static String getPackageName(String className) {
+    int index = className.lastIndexOf('.');
+    return index == -1 ? null : className.substring(0, index);
+  }
+
+  private static Manifest getManifest(URL url) {
+    try {
+      JarFile jarFile = new JarFile(url.getFile());
+      return jarFile.getManifest();
+    } catch (IOException e) {
+      log.warn(e.getMessage(), e);
     }
+    return null;
   }
 }

agent-tooling/src/main/java/io/opentelemetry/auto/tooling/TracerInstaller.java

@@ -64,7 +64,7 @@ private static synchronized void installExportersFromJar(final String exporterJa
     }
     final DefaultExporterConfig config = new DefaultExporterConfig("exporter");
     final ExporterClassLoader exporterLoader =
-        new ExporterClassLoader(new URL[] {url}, TracerInstaller.class.getClassLoader());
+        new ExporterClassLoader(url, TracerInstaller.class.getClassLoader());
 
     final SpanExporterFactory spanExporterFactory =
         getExporterFactory(SpanExporterFactory.class, exporterLoader);

agent-tooling/src/test/groovy/io/opentelemetry/auto/test/ClassLoaderMatcherTest.groovy

@@ -32,7 +32,8 @@ class ClassLoaderMatcherTest extends AgentSpecification {
 
   def "skips exporter classloader"() {
     setup:
-    final URLClassLoader exporterLoader = new ExporterClassLoader(new URL[0], null)
+    URL url = new URL("file://")
+    final URLClassLoader exporterLoader = new ExporterClassLoader(url, null)
     expect:
     ClassLoaderMatcher.skipClassLoader().matches(exporterLoader)
   }

auto-exporters/src/test/groovy/io/opentelemetry/auto/exporteradapters/ExporterAdaptersTest.groovy

@@ -53,8 +53,7 @@ class ExporterAdaptersTest extends Specification {
     def file = new File(exporter)
     println "Attempting to load ${file.toString()} for ${classname}"
     assert file.exists(): "${file.toString()} does not exist"
-    URL[] urls = [file.toURI().toURL()]
-    def classLoader = new ExporterClassLoader(urls, this.getClass().getClassLoader())
+    def classLoader = new ExporterClassLoader(file.toURI().toURL(), this.getClass().getClassLoader())
     def serviceLoader = ServiceLoader.load(SpanExporterFactory, classLoader)
 
     when: