Support loading trace extensions from a comma-separated list of jars,…

… or directories containing jars

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

@@ -22,8 +22,10 @@
 import java.lang.instrument.Instrumentation;
 import java.util.Collections;
 import java.util.EnumSet;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.NoSuchElementException;
 import java.util.Set;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.TimeUnit;
@@ -172,11 +174,14 @@ public static ClassFileTransformer installBytebuddyAgent(
     // pre-size state before registering instrumentations to reduce number of allocations
     InstrumenterState.initialize(instrumenterIndex.instrumentationCount());
 
+    // combine known modules indexed at build-time with extensions contributed at run-time
+    Iterable<InstrumenterModule> instrumenterModules = withExtensions(instrumenterIndex.modules());
+
     // This needs to be a separate loop through all instrumentations before we start adding
     // advice so that we can exclude field injection, since that will try to check exclusion
     // immediately and we don't have the ability to express dependencies between different
     // instrumentations to control the load order.
-    for (InstrumenterModule module : instrumenterIndex.modules()) {
+    for (InstrumenterModule module : instrumenterModules) {
       if (module instanceof ExcludeFilterProvider) {
         ExcludeFilterProvider provider = (ExcludeFilterProvider) module;
         ExcludeFilter.add(provider.excludedClasses());
@@ -191,7 +196,7 @@ public static ClassFileTransformer installBytebuddyAgent(
         new CombiningTransformerBuilder(agentBuilder, instrumenterIndex);
 
     int installedCount = 0;
-    for (InstrumenterModule module : instrumenterIndex.modules()) {
+    for (InstrumenterModule module : instrumenterModules) {
       if (!module.isApplicable(enabledSystems)) {
         if (DEBUG) {
           log.debug("Not applicable - instrumentation.class={}", module.getClass().getName());
@@ -234,6 +239,53 @@ 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);
+      }
+    };
+  }
+
+  /** Returns an iterator that combines the original sequence with any discovered extensions. */
+  private static Iterator<InstrumenterModule> withExtensions(
+      final Iterator<InstrumenterModule> initial, final Iterable<InstrumenterModule> extensions) {
+    return new Iterator<InstrumenterModule>() {
+      private Iterator<InstrumenterModule> delegate = initial;
+
+      @Override
+      public boolean hasNext() {
+        if (delegate.hasNext()) {
+          return true;
+        } else if (delegate == initial) {
+          delegate = extensions.iterator();
+          return delegate.hasNext();
+        } else {
+          return false;
+        }
+      }
+
+      @Override
+      public InstrumenterModule next() {
+        if (hasNext()) {
+          return delegate.next();
+        } else {
+          throw new NoSuchElementException();
+        }
+      }
+    };
+  }
+
   public static Set<InstrumenterModule.TargetSystem> getEnabledSystems() {
     EnumSet<InstrumenterModule.TargetSystem> enabledSystems =
         EnumSet.noneOf(InstrumenterModule.TargetSystem.class);

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

@@ -0,0 +1,110 @@
+package datadog.trace.agent.tooling;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+import de.thetaphi.forbiddenapis.SuppressForbidden;
+import java.io.BufferedReader;
+import java.io.File;
+import java.io.InputStreamReader;
+import java.net.MalformedURLException;
+import java.net.URL;
+import java.net.URLClassLoader;
+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 trace extensions from a comma-separated list of jars, or directories containing jars. */
+public final class ExtensionsLoader {
+  private static final Logger log = LoggerFactory.getLogger(ExtensionsLoader.class);
+
+  private static final String DATADOG_MODULE_EXTENSION_ID =
+      "datadog.trace.agent.tooling.InstrumenterModule";
+
+  private static final String[] NO_EXTENSIONS = {};
+
+  private final ClassLoader extensionLoader;
+
+  public ExtensionsLoader(String extensionsPath) {
+    extensionLoader =
+        new URLClassLoader(toURLs(extensionsPath), Instrumenter.class.getClassLoader());
+  }
+
+  public List<InstrumenterModule> loadModules() {
+    List<InstrumenterModule> modules = new ArrayList<>();
+    for (String className : discoverExtensions(extensionLoader, DATADOG_MODULE_EXTENSION_ID)) {
+      try {
+        modules.add(loadDatadogModule(className));
+      } catch (Throwable e) {
+        log.warn("Failed to load extension module {}", className, e);
+      }
+    }
+    return modules;
+  }
+
+  private InstrumenterModule loadDatadogModule(String className)
+      throws ReflectiveOperationException {
+    Class<?> moduleClass = extensionLoader.loadClass(className);
+    return (InstrumenterModule) moduleClass.getConstructor().newInstance();
+  }
+
+  /** Similar to {@link java.util.ServiceLoader} but doesn't load the discovered extensions. */
+  private static String[] discoverExtensions(ClassLoader loader, String extensionId) {
+    try {
+      Set<String> lines = new LinkedHashSet<>();
+      Enumeration<URL> urls = loader.getResources("META-INF/services/" + extensionId);
+      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;
+    }
+  }
+
+  @SuppressForbidden // split on single-character uses fast path
+  private static URL[] toURLs(String path) {
+    List<URL> urls = new ArrayList<>();
+    for (String entry : path.split(",")) {
+      File file = new File(entry);
+      if (file.isDirectory()) {
+        visitDirectory(file, urls);
+      } else if (isJar(file)) {
+        addExtensionJar(file, urls);
+      }
+    }
+    return urls.toArray(new URL[0]);
+  }
+
+  private static void visitDirectory(File dir, List<URL> urls) {
+    File[] files = dir.listFiles(ExtensionsLoader::isJar);
+    if (null != files) {
+      for (File file : files) {
+        addExtensionJar(file, urls);
+      }
+    }
+  }
+
+  private static void addExtensionJar(File file, List<URL> urls) {
+    try {
+      urls.add(file.toURI().toURL());
+    } catch (MalformedURLException e) {
+      log.debug("Ignoring extension jar {}", file, e);
+    }
+  }
+
+  private static boolean isJar(File file) {
+    return file.getName().endsWith(".jar") && file.isFile();
+  }
+}

dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java

@@ -13,6 +13,8 @@ public final class TraceInstrumentationConfig {
   public static final String TRACE_OTEL_ENABLED = "trace.otel.enabled";
   public static final String INTEGRATIONS_ENABLED = "integrations.enabled";
 
+  public static final String TRACE_EXTENSIONS_PATH = "trace.extensions.path";
+
   public static final String INTEGRATION_SYNAPSE_LEGACY_OPERATION_NAME =
       "integration.synapse.legacy-operation-name";
   public static final String TRACE_ANNOTATIONS = "trace.annotations";

internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java

@@ -58,6 +58,7 @@
 import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_ENABLED;
 import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_EXECUTORS;
 import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_EXECUTORS_ALL;
+import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_EXTENSIONS_PATH;
 import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_METHODS;
 import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_OTEL_ENABLED;
 import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_THREAD_POOL_EXECUTORS_EXCLUDE;
@@ -106,6 +107,8 @@ public class InstrumenterConfig {
   private final boolean usmEnabled;
   private final boolean telemetryEnabled;
 
+  private final String traceExtensionsPath;
+
   private final boolean traceExecutorsAll;
   private final List<String> traceExecutors;
   private final Set<String> traceThreadPoolExecutorsExclude;
@@ -193,6 +196,8 @@ private InstrumenterConfig() {
       usmEnabled = false;
     }
 
+    traceExtensionsPath = configProvider.getString(TRACE_EXTENSIONS_PATH);
+
     traceExecutorsAll = configProvider.getBoolean(TRACE_EXECUTORS_ALL, DEFAULT_TRACE_EXECUTORS_ALL);
     traceExecutors = tryMakeImmutableList(configProvider.getList(TRACE_EXECUTORS));
     traceThreadPoolExecutorsExclude =
@@ -311,6 +316,10 @@ public boolean isTelemetryEnabled() {
     return telemetryEnabled;
   }
 
+  public String getTraceExtensionsPath() {
+    return traceExtensionsPath;
+  }
+
   public boolean isTraceExecutorsAll() {
     return traceExecutorsAll;
   }
@@ -506,6 +515,8 @@ public String toString() {
         + usmEnabled
         + ", telemetryEnabled="
         + telemetryEnabled
+        + ", traceExtensionsPath="
+        + traceExtensionsPath
         + ", traceExecutorsAll="
         + traceExecutorsAll
         + ", traceExecutors="