initial java autoinstrumentation

Author: realark

braintrust-java-agent/bootstrap/build.gradle

@@ -0,0 +1,12 @@
+// Java plugin, toolchain (Java 17 / Adoptium), options.release, and repositories
+// are inherited from the parent's subprojects {} block.
+
+// These are compile-time dependencies so bootstrap classes can reference OTel types.
+// At runtime, the actual OTel JARs are bundled as normal .class files in the agent JAR
+// and placed on the bootstrap classpath via Instrumentation.appendToBootstrapClassLoaderSearch().
+dependencies {
+  compileOnly "io.opentelemetry:opentelemetry-api:${otelVersion}"
+  compileOnly "io.opentelemetry:opentelemetry-sdk:${otelVersion}"
+  compileOnly "io.opentelemetry:opentelemetry-sdk-extension-autoconfigure:${otelVersion}"
+  compileOnly "io.opentelemetry:opentelemetry-sdk-extension-autoconfigure-spi:${otelVersion}"
+}

braintrust-java-agent/bootstrap/src/main/java/dev/braintrust/bootstrap/BraintrustBridge.java

@@ -0,0 +1,31 @@
+package dev.braintrust.bootstrap;
+
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+
+/** Globally available bootstrap classpath resource class */
+public class BraintrustBridge {
+    public static final String INSTRUMENTATION_NAME = "braintrust-java";
+
+    /**
+     * Diagnostic utility tracking the number of times braintrust otel has been installed.
+     *
+     * <p>In a production app, this should be zero until global otel get() is invoked, then it
+     * should remain at 1 for the rest of app's lifetime.
+     */
+    public static final AtomicInteger otelInstallCount = new AtomicInteger(0);
+
+    private static final AtomicReference<BraintrustClassLoader> agentClassLoaderRef =
+            new AtomicReference<>();
+
+    public static BraintrustClassLoader getAgentClassLoader() {
+        return agentClassLoaderRef.get();
+    }
+
+    public static void setAgentClassloaderIfAbsent(BraintrustClassLoader classLoader) {
+        var witness = agentClassLoaderRef.compareAndExchange(null, classLoader);
+        if (null != witness) {
+            throw new IllegalStateException("agent classloader must only be set once");
+        }
+    }
+}

braintrust-java-agent/bootstrap/src/main/java/dev/braintrust/bootstrap/BraintrustClassLoader.java

@@ -0,0 +1,116 @@
+package dev.braintrust.bootstrap;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.net.URL;
+import java.security.CodeSource;
+import java.security.SecureClassLoader;
+import java.security.cert.Certificate;
+import java.util.Collections;
+import java.util.Enumeration;
+import java.util.jar.JarEntry;
+import java.util.jar.JarFile;
+
+/**
+ * A classloader that loads agent-internal classes from {@code .classdata} entries inside the agent
+ * JAR.
+ *
+ * <p>Classes stored under the {@code internal/} prefix with a {@code .classdata} extension are
+ * invisible to the JVM's default classloading mechanism. This classloader knows how to find them,
+ * providing full classloader isolation between the agent's internals and the application's
+ * classpath.
+ */
+public class BraintrustClassLoader extends SecureClassLoader {
+    private static final String ENTRY_PREFIX = "internal/";
+    private static final String CLASS_DATA_SUFFIX = ".classdata";
+
+    private final JarFile agentJarFile;
+    private final CodeSource agentCodeSource;
+    private final String agentResourcePrefix;
+
+    static {
+        registerAsParallelCapable();
+    }
+
+    /**
+     * Creates a new BraintrustClassLoader.
+     *
+     * @param agentJarURL the URL of the agent JAR file (from the -javaagent path)
+     * @param parent the parent classloader (typically the system/platform classloader)
+     */
+    public BraintrustClassLoader(URL agentJarURL, ClassLoader parent) throws Exception {
+        super(parent);
+        this.agentJarFile = new JarFile(new java.io.File(agentJarURL.toURI()), false);
+        this.agentCodeSource = new CodeSource(agentJarURL, (Certificate[]) null);
+        this.agentResourcePrefix = "jar:file:" + agentJarFile.getName() + "!/";
+    }
+
+    @Override
+    protected Class<?> findClass(String name) throws ClassNotFoundException {
+        // Convert "dev.braintrust.agent.internal.BraintrustAgent"
+        //      -> "internal/dev/braintrust/agent/internal/BraintrustAgent.classdata"
+        String entryName = ENTRY_PREFIX + name.replace('.', '/') + CLASS_DATA_SUFFIX;
+        JarEntry entry = agentJarFile.getJarEntry(entryName);
+        if (entry == null) {
+            throw new ClassNotFoundException(name);
+        }
+
+        byte[] classBytes = readEntry(entry, name);
+        return defineClass(name, classBytes, 0, classBytes.length, agentCodeSource);
+    }
+
+    @Override
+    protected URL findResource(String name) {
+        // For .class resource lookups, map to .classdata
+        String entryName;
+        if (name.endsWith(".class")) {
+            entryName =
+                    ENTRY_PREFIX
+                            + name.substring(0, name.length() - ".class".length())
+                            + CLASS_DATA_SUFFIX;
+        } else {
+            entryName = ENTRY_PREFIX + name;
+        }
+
+        JarEntry entry = agentJarFile.getJarEntry(entryName);
+        if (entry != null) {
+            try {
+                return new URL(agentResourcePrefix + entryName);
+            } catch (java.net.MalformedURLException e) {
+                // fall through
+            }
+        }
+        return null;
+    }
+
+    @Override
+    protected Enumeration<URL> findResources(String name) {
+        URL resource = findResource(name);
+        if (resource != null) {
+            return Collections.enumeration(Collections.singletonList(resource));
+        }
+        return Collections.emptyEnumeration();
+    }
+
+    private byte[] readEntry(JarEntry entry, String className) throws ClassNotFoundException {
+        int size = (int) entry.getSize();
+        byte[] buf = new byte[size];
+        try (InputStream in = agentJarFile.getInputStream(entry)) {
+            int offset = 0;
+            while (offset < size) {
+                int bytesRead = in.read(buf, offset, size - offset);
+                if (bytesRead < 0) {
+                    break;
+                }
+                offset += bytesRead;
+            }
+            if (offset != size) {
+                throw new ClassNotFoundException(
+                        className + " (incomplete read: " + offset + "/" + size + " bytes)");
+            }
+            return buf;
+        } catch (IOException e) {
+            throw new ClassNotFoundException(className, e);
+        }
+    }
+}

braintrust-java-agent/bootstrap/src/main/java/dev/braintrust/system/AgentBootstrap.java

@@ -0,0 +1,178 @@
+package dev.braintrust.system;
+
+import dev.braintrust.bootstrap.BraintrustBridge;
+import dev.braintrust.bootstrap.BraintrustClassLoader;
+import io.opentelemetry.api.trace.*;
+import java.io.File;
+import java.lang.instrument.Instrumentation;
+import java.net.URL;
+import java.util.jar.JarFile;
+
+/**
+ * Braintrust Java Agent entry point.
+ *
+ * <p>Minimal code which bootstraps braintrust classloading and instrumentation.
+ */
+public class AgentBootstrap {
+    private static final String AGENT_CLASS = "dev.braintrust.agent.BraintrustAgent";
+    private static final String INSTALLER_METHOD = "install";
+
+    static volatile boolean installed = false;
+
+    /** Entry point when the agent is loaded at JVM startup via {@code -javaagent}. */
+    public static void premain(String agentArgs, Instrumentation inst) {
+        install(agentArgs, inst);
+    }
+
+    /** Entry point when the agent is attached to a running JVM via the Attach API. */
+    public static void agentmain(String agentArgs, Instrumentation inst) {
+        install(agentArgs, inst);
+    }
+
+    private static synchronized void install(String agentArgs, Instrumentation inst) {
+        if (AgentBootstrap.class.getClassLoader() != ClassLoader.getSystemClassLoader()) {
+            log(
+                    "WARNING: install attempted on non-system classloader. aborting. classloader: "
+                            + AgentBootstrap.class.getClassLoader());
+            return;
+        }
+        if (installed) {
+            log("Agent already installed, skipping.");
+            return;
+        }
+
+        log("Braintrust Java Agent starting...");
+
+        if (jvmRunningWithOtelAgent()) {
+            log(
+                    "ERROR: Braintrust agent is not yet compatible with the OTel javaagent -"
+                            + " skipping install.");
+            return;
+        }
+
+        if (jvmRunningWithDatadogOtel()) {
+            log(
+                    "ERROR: Braintrust agent is not yet compatible with datadog javaagent otel -"
+                            + " skipping install.");
+            return;
+        }
+
+        try {
+            // Locate the agent JAR from our own code source
+            URL agentJarURL =
+                    AgentBootstrap.class.getProtectionDomain().getCodeSource().getLocation();
+            File agentJarFile = new File(agentJarURL.toURI());
+            log("Agent JAR: " + agentJarFile);
+
+            // Enable OTel autoconfigure BEFORE adding to bootstrap, so system properties
+            // are set before anything can trigger GlobalOpenTelemetry.get().
+            enableOtelSDKAutoconfiguration();
+
+            inst.appendToBootstrapClassLoaderSearch(new JarFile(agentJarFile, false));
+            log("Added agent JAR to bootstrap classpath.");
+
+            // Create the isolated braintrust classloader.
+            // Parent is the platform classloader so agent internals can see:
+            //   - Bootstrap classes (OTel API/SDK added via appendToBootstrapClassLoaderSearch)
+            //   - JDK platform modules (java.net.http, java.sql, etc.)
+            // but NOT application classes (those are on the system/app classloader).
+            BraintrustClassLoader btClassLoader =
+                    new BraintrustClassLoader(agentJarURL, ClassLoader.getPlatformClassLoader());
+            BraintrustBridge.setAgentClassloaderIfAbsent(btClassLoader);
+
+            // Load and invoke the real agent installer through the isolated classloader.
+            Class<?> installerClass = btClassLoader.loadClass(AGENT_CLASS);
+            installerClass
+                    .getMethod(INSTALLER_METHOD, String.class, Instrumentation.class)
+                    .invoke(null, agentArgs, inst);
+            log("Braintrust Java Agent installed.");
+            installed = true;
+        } catch (Throwable t) {
+            log("ERROR: Failed to install Braintrust Java Agent: " + t.getMessage());
+            log(t);
+        }
+    }
+
+    /**
+     * Checks whether the OpenTelemetry Java agent is present by looking for its premain class on
+     * the system classloader. Since {@code -javaagent} JARs are always on the system classpath,
+     * this works regardless of agent ordering.
+     */
+    private static boolean jvmRunningWithOtelAgent() {
+        try {
+            Class.forName(
+                    "io.opentelemetry.javaagent.OpenTelemetryAgent",
+                    false,
+                    ClassLoader.getSystemClassLoader());
+            return true;
+        } catch (ClassNotFoundException e) {
+            return false;
+        }
+    }
+
+    /**
+     * Checks whether the Datadog agent is present and configured for OTel integration. Must be
+     * callable from the system classloader (no DD compile deps).
+     */
+    private static boolean jvmRunningWithDatadogOtel() {
+        try {
+            Class.forName("datadog.trace.bootstrap.Agent", false, null);
+        } catch (ClassNotFoundException e) {
+            return false;
+        }
+        String sysProp = System.getProperty("dd.trace.otel.enabled");
+        if (sysProp != null) {
+            return Boolean.parseBoolean(sysProp);
+        }
+        String envVar = System.getenv("DD_TRACE_OTEL_ENABLED");
+        return Boolean.parseBoolean(envVar);
+    }
+
+    /**
+     * Returns true if the Datadog agent's premain has already executed, meaning it was listed
+     * before the Braintrust agent in the {@code -javaagent} flags.
+     */
+    static boolean isRunningAfterDatadogAgent() {
+        // DD's premain appends its jars to the bootstrap classpath, making
+        // {@code datadog.trace.bootstrap.Agent} loadable from the bootstrap (null)
+        // classloader. If that class is not found on bootstrap, DD either isn't
+        // present or hasn't run its premain yet (i.e. BT is first).
+        try {
+            Class.forName("datadog.trace.bootstrap.Agent", false, null);
+            return true;
+        } catch (ClassNotFoundException e) {
+            return false;
+        }
+    }
+
+    private static void enableOtelSDKAutoconfiguration() {
+        // Enable OTel SDK autoconfiguration. When anyone first calls
+        // GlobalOpenTelemetry.get(), the SDK will be built using autoconfigure, which
+        // discovers our BraintrustAutoConfigCustomizer via ServiceLoader and injects the
+        // Braintrust span processor/exporter into the tracer provider.
+        setPropertyIfAbsent("otel.java.global-autoconfigure.enabled", "true");
+
+        // Set default exporter config. We don't want autoconfigure to try loading
+        // OTLP exporters from the bootstrap classpath (they live in BraintrustClassLoader).
+        // Our BraintrustAutoConfigCustomizer adds the real span processor/exporter.
+        setPropertyIfAbsent("otel.traces.exporter", "none");
+        setPropertyIfAbsent("otel.metrics.exporter", "none");
+        setPropertyIfAbsent("otel.logs.exporter", "none");
+        log("Enabled OTel SDK autoconfiguration.");
+    }
+
+    /** Sets a system property only if it hasn't been set already (respects user overrides). */
+    private static void setPropertyIfAbsent(String key, String value) {
+        if (System.getProperty(key) == null) {
+            System.setProperty(key, value);
+        }
+    }
+
+    private static void log(String msg) {
+        System.out.println("[braintrust] " + msg);
+    }
+
+    private static void log(Throwable t) {
+        t.printStackTrace(System.err);
+    }
+}