Support mapping extensions from different providers into a form that can be used by the Datadog tracer

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

@@ -1,5 +1,7 @@
 package datadog.trace.agent.tooling;
 
+import static datadog.trace.agent.tooling.ExtensionFinder.findExtensions;
+import static datadog.trace.agent.tooling.ExtensionLoader.loadExtensions;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.GlobalIgnoresMatcher.globalIgnoresMatcher;
 import static net.bytebuddy.matcher.ElementMatchers.isDefaultFinalizer;
 
@@ -242,19 +244,18 @@ public void applied(Iterable<String> instrumentationNames) {
   /** Returns an iterable that combines the original sequence with any discovered extensions. */
   private static Iterable<InstrumenterModule> withExtensions(Iterable<InstrumenterModule> initial) {
     String extensionsPath = InstrumenterConfig.get().getTraceExtensionsPath();
-    if (null == extensionsPath) {
-      return initial;
-    }
-    final List<InstrumenterModule> extensions = new ExtensionsLoader(extensionsPath).loadModules();
-    if (extensions.isEmpty()) {
-      return initial;
-    }
-    return new Iterable<InstrumenterModule>() {
-      @Override
-      public Iterator<InstrumenterModule> iterator() {
-        return withExtensions(initial.iterator(), extensions);
+    if (null != extensionsPath) {
+      if (findExtensions(extensionsPath, InstrumenterModule.class)) {
+        final List<InstrumenterModule> extensions = loadExtensions(InstrumenterModule.class);
+        return new Iterable<InstrumenterModule>() {
+          @Override
+          public Iterator<InstrumenterModule> iterator() {
+            return withExtensions(initial.iterator(), extensions);
+          }
+        };
       }
-    };
+    }
+    return initial;
   }
 
   /** Returns an iterator that combines the original sequence with any discovered extensions. */

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

@@ -247,7 +247,7 @@ public void applyAdvice(ElementMatcher<? super MethodDescription> matcher, Strin
     }
     advice.add(
         new AgentBuilder.Transformer.ForAdvice(customMapping)
-            .include(Utils.getBootstrapProxy(), Utils.getAgentClassLoader())
+            .include(Utils.getBootstrapProxy(), Utils.getExtendedClassLoader())
             .withExceptionHandler(ExceptionHandlers.defaultExceptionHandler())
             .advice(not(ignoredMethods).and(matcher), adviceClass));
   }

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

@@ -0,0 +1,165 @@
+package datadog.trace.agent.tooling;
+
+import static datadog.trace.agent.tooling.ExtensionHandler.DATADOG;
+
+import de.thetaphi.forbiddenapis.SuppressForbidden;
+import java.io.File;
+import java.io.FileNotFoundException;
+import java.io.IOException;
+import java.net.MalformedURLException;
+import java.net.URL;
+import java.net.URLClassLoader;
+import java.net.URLConnection;
+import java.net.URLStreamHandler;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.jar.JarEntry;
+import java.util.jar.JarFile;
+import net.bytebuddy.dynamic.loading.MultipleParentClassLoader;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Finds extensions to the Datadog tracer. */
+public final class ExtensionFinder {
+  private static final Logger log = LoggerFactory.getLogger(ExtensionFinder.class);
+
+  private static final ExtensionHandler[] handlers = {DATADOG};
+
+  /**
+   * Discovers extensions on the configured path and creates a classloader for each extension.
+   * Registers the combined classloader with {@link Utils#setExtendedClassLoader(ClassLoader)}.
+   *
+   * @return {@code true} if any extensions were found
+   */
+  public static boolean findExtensions(String extensionsPath, Class<?>... extensionTypes) {
+    List<ClassLoader> classLoaders = new ArrayList<>();
+    List<JarFile> unusedJars = new ArrayList<>();
+
+    ClassLoader parent = Utils.getAgentClassLoader();
+    String[] descriptors = descriptors(extensionTypes);
+
+    for (JarFile jar : findExtensionJars(extensionsPath)) {
+      URL extensionURL = findExtensionURL(jar, descriptors);
+      if (null != extensionURL) {
+        log.debug("Found extension jar {}", jar.getName());
+        classLoaders.add(new URLClassLoader(new URL[] {extensionURL}, parent));
+      } else {
+        unusedJars.add(jar);
+      }
+    }
+
+    close(unusedJars);
+
+    if (classLoaders.size() > 1) {
+      Utils.setExtendedClassLoader(new MultipleParentClassLoader(classLoaders));
+    } else if (!classLoaders.isEmpty()) {
+      Utils.setExtendedClassLoader(classLoaders.get(0));
+    }
+
+    return !classLoaders.isEmpty();
+  }
+
+  /** Closes jar resources from the extension path which did not contain any extensions. */
+  private static void close(List<JarFile> unusedJars) {
+    for (JarFile jar : unusedJars) {
+      try {
+        jar.close();
+      } catch (Exception ignore) {
+        // move onto next jar
+      }
+    }
+  }
+
+  /**
+   * Uses the registered {@link ExtensionHandler}s to find jars containing matching extensions.
+   * Creates a URL that uses the matched extension handler to access content from the extension.
+   */
+  private static URL findExtensionURL(JarFile jar, String[] descriptors) {
+    for (ExtensionHandler handler : handlers) {
+      for (String descriptor : descriptors) {
+        if (null != handler.mapEntry(jar, descriptor)) {
+          return buildExtensionURL(jar, handler);
+        }
+      }
+    }
+    return null;
+  }
+
+  /** 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);
+            }
+          });
+    } 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);
+    }
+    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());
+    }
+  }
+
+  @SuppressForbidden // split on single-character uses fast path
+  private static List<JarFile> findExtensionJars(String extensionsPath) {
+    List<JarFile> extensionJars = new ArrayList<>();
+    for (String entry : extensionsPath.split(",")) {
+      File file = new File(entry);
+      if (file.isDirectory()) {
+        visitDirectory(file, extensionJars);
+      } else if (isJar(file)) {
+        addExtensionJar(file, extensionJars);
+      }
+    }
+    return extensionJars;
+  }
+
+  private static void visitDirectory(File dir, List<JarFile> extensionJars) {
+    File[] files = dir.listFiles(ExtensionFinder::isJar);
+    if (null != files) {
+      for (File file : files) {
+        addExtensionJar(file, extensionJars);
+      }
+    }
+  }
+
+  private static void addExtensionJar(File file, List<JarFile> extensionJars) {
+    try {
+      extensionJars.add(new JarFile(file, false));
+    } catch (Exception e) {
+      log.debug("Problem reading extension jar {}", file, e);
+    }
+  }
+
+  /** The {@code META-INF/service} descriptors to look for. */
+  private static String[] descriptors(Class<?>[] extensionTypes) {
+    String[] descriptors = new String[extensionTypes.length];
+    for (int i = 0; i < extensionTypes.length; i++) {
+      descriptors[i] = "META-INF/services/" + extensionTypes[i].getName();
+    }
+    return descriptors;
+  }
+
+  private static boolean isJar(File file) {
+    return file.getName().endsWith(".jar") && file.isFile();
+  }
+}

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

@@ -0,0 +1,64 @@
+package datadog.trace.agent.tooling;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+import java.io.BufferedReader;
+import java.io.InputStreamReader;
+import java.net.URL;
+import java.util.ArrayList;
+import java.util.Enumeration;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Set;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Loads extensions into the Datadog tracer. */
+public final class ExtensionLoader {
+  private static final Logger log = LoggerFactory.getLogger(ExtensionLoader.class);
+
+  private static final String[] NO_EXTENSIONS = {};
+
+  /** Loads extensions from the extended classloader built by {@link ExtensionFinder}. */
+  public static <T> List<T> loadExtensions(Class<T> extensionType) {
+    return loadExtensions(Utils.getExtendedClassLoader(), extensionType);
+  }
+
+  private static <T> List<T> loadExtensions(ClassLoader classLoader, Class<T> extensionType) {
+    List<T> extensions = new ArrayList<>();
+    for (String className : listExtensionNames(classLoader, extensionType)) {
+      try {
+        @SuppressWarnings("unchecked")
+        Class<T> extensionClass = (Class<T>) classLoader.loadClass(className);
+        extensions.add(extensionClass.getConstructor().newInstance());
+        log.debug("Loaded extension {}", className);
+      } catch (Throwable e) {
+        log.warn("Failed to load extension {}", className, e);
+      }
+    }
+    return extensions;
+  }
+
+  /** Returns the class names listed under the extension's {@code META-INF/services} descriptor. */
+  private static String[] listExtensionNames(ClassLoader classLoader, Class<?> extensionType) {
+    try {
+      Set<String> lines = new LinkedHashSet<>();
+      Enumeration<URL> urls =
+          classLoader.getResources("META-INF/services/" + extensionType.getName());
+      while (urls.hasMoreElements()) {
+        try (BufferedReader reader =
+            new BufferedReader(new InputStreamReader(urls.nextElement().openStream(), UTF_8))) {
+          String line = reader.readLine();
+          while (line != null) {
+            lines.add(line);
+            line = reader.readLine();
+          }
+        }
+      }
+      return lines.toArray(new String[0]);
+    } catch (Throwable e) {
+      log.warn("Problem reading extensions descriptor", e);
+      return NO_EXTENSIONS;
+    }
+  }
+}