Skip to content

Commit 6bad5d7

Browse files
realarkrealark
authored
Merge pull request #55 from braintrustdata/ark/autoinstrumentation-dd-compat
autoinstrumentation + dd otel compat mode
2 parents f3a9b15 + 8f55ca2 commit 6bad5d7

16 files changed

Lines changed: 2907 additions & 35 deletions

File tree

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

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package dev.braintrust.bootstrap;
22

3+
import java.net.URL;
34
import java.util.concurrent.atomic.AtomicInteger;
45
import java.util.concurrent.atomic.AtomicReference;
56

@@ -22,10 +23,14 @@ public static BraintrustClassLoader getAgentClassLoader() {
2223
return agentClassLoaderRef.get();
2324
}
2425

25-
public static void setAgentClassLoaderIfAbsent(BraintrustClassLoader classLoader) {
26-
var witness = agentClassLoaderRef.compareAndExchange(null, classLoader);
26+
public static BraintrustClassLoader createBraintrustClassLoader(
27+
URL agentJarURL, ClassLoader btClassLoaderParent) throws Exception {
28+
BraintrustClassLoader btClassLoader =
29+
new BraintrustClassLoader(agentJarURL, btClassLoaderParent);
30+
var witness = agentClassLoaderRef.compareAndExchange(null, btClassLoader);
2731
if (null != witness) {
2832
throw new IllegalStateException("agent classloader must only be set once");
2933
}
34+
return btClassLoader;
3035
}
3136
}

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

Lines changed: 32 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,9 @@
11
package dev.braintrust.system;
22

3-
import dev.braintrust.bootstrap.BraintrustBridge;
4-
import dev.braintrust.bootstrap.BraintrustClassLoader;
53
import java.io.File;
64
import java.lang.instrument.Instrumentation;
75
import java.net.URL;
6+
import java.net.URLClassLoader;
87
import java.util.jar.JarFile;
98

109
/**
@@ -44,18 +43,16 @@ private static synchronized void install(String agentArgs, Instrumentation inst)
4443

4544
if (jvmRunningWithOtelAgent()) {
4645
log(
47-
"ERROR: Braintrust agent is not yet compatible with the OTel javaagent -"
48-
+ " skipping install.");
46+
"ERROR: Braintrust agent is not yet compatible with the OTel -javaagent."
47+
+ " aborting install.");
4948
return;
5049
}
51-
52-
if (jvmRunningWithDatadogOtel()) {
53-
log(
54-
"ERROR: Braintrust agent is not yet compatible with datadog javaagent otel -"
55-
+ " skipping install.");
50+
if (jvmRunningWithDatadogOtelConfig() && (!isRunningAfterDatadogAgent())) {
51+
log("ERROR: Braintrust agent must run _after_ datadog -javaagent. aborting install.");
5652
return;
5753
}
5854

55+
boolean installOnBootstrap = !jvmRunningWithDatadogOtelConfig();
5956
try {
6057
// Locate the agent JAR from our own code source
6158
URL agentJarURL =
@@ -67,19 +64,21 @@ private static synchronized void install(String agentArgs, Instrumentation inst)
6764
// are set before anything can trigger GlobalOpenTelemetry.get().
6865
enableOtelSDKAutoconfiguration();
6966

70-
inst.appendToBootstrapClassLoaderSearch(new JarFile(agentJarFile, false));
71-
log("Added agent JAR to bootstrap classpath.");
72-
73-
// Create the isolated braintrust classloader.
74-
// Parent is the platform classloader so agent internals can see:
75-
// - Bootstrap classes (OTel API/SDK added via appendToBootstrapClassLoaderSearch)
76-
// - JDK platform modules (java.net.http, java.sql, etc.)
77-
// but NOT application classes (those are on the system/app classloader).
78-
BraintrustClassLoader btClassLoader =
79-
new BraintrustClassLoader(agentJarURL, ClassLoader.getPlatformClassLoader());
80-
BraintrustBridge.setAgentClassLoaderIfAbsent(btClassLoader);
67+
ClassLoader btClassLoaderParent;
68+
if (installOnBootstrap) {
69+
inst.appendToBootstrapClassLoaderSearch(new JarFile(agentJarFile, false));
70+
btClassLoaderParent = ClassLoader.getPlatformClassLoader();
71+
log("Added agent JAR to bootstrap classpath.");
72+
} else {
73+
btClassLoaderParent =
74+
new URLClassLoader(
75+
new URL[] {agentJarFile.toURI().toURL()},
76+
ClassLoader.getPlatformClassLoader());
77+
log("skipping bootstrap classpath setup");
78+
}
8179

8280
// Load and invoke the real agent installer through the isolated classloader.
81+
ClassLoader btClassLoader = createBTClassLoader(agentJarURL, btClassLoaderParent);
8382
Class<?> installerClass = btClassLoader.loadClass(AGENT_CLASS);
8483
installerClass
8584
.getMethod(INSTALLER_METHOD, String.class, Instrumentation.class)
@@ -92,6 +91,16 @@ private static synchronized void install(String agentArgs, Instrumentation inst)
9291
}
9392
}
9493

94+
private static ClassLoader createBTClassLoader(URL agentJarURL, ClassLoader btClassLoaderParent)
95+
throws Exception {
96+
// NOTE: not caching because we only invoke this once
97+
var bridgeClass =
98+
btClassLoaderParent.loadClass("dev.braintrust.bootstrap.BraintrustBridge");
99+
var createMethod =
100+
bridgeClass.getMethod("createBraintrustClassLoader", URL.class, ClassLoader.class);
101+
return (ClassLoader) createMethod.invoke(null, agentJarURL, btClassLoaderParent);
102+
}
103+
95104
/**
96105
* Checks whether the OpenTelemetry Java agent is present by looking for its premain class on
97106
* the system classloader. Since {@code -javaagent} JARs are always on the system classpath,
@@ -109,16 +118,8 @@ private static boolean jvmRunningWithOtelAgent() {
109118
}
110119
}
111120

112-
/**
113-
* Checks whether the Datadog agent is present and configured for OTel integration. Must be
114-
* callable from the system classloader (no DD compile deps).
115-
*/
116-
private static boolean jvmRunningWithDatadogOtel() {
117-
try {
118-
Class.forName("datadog.trace.bootstrap.Agent", false, null);
119-
} catch (ClassNotFoundException e) {
120-
return false;
121-
}
121+
/** Checks whether the Datadog agent is present and configured for OTel integration */
122+
private static boolean jvmRunningWithDatadogOtelConfig() {
122123
String sysProp = System.getProperty("dd.trace.otel.enabled");
123124
if (sysProp != null) {
124125
return Boolean.parseBoolean(sysProp);
@@ -131,7 +132,7 @@ private static boolean jvmRunningWithDatadogOtel() {
131132
* Returns true if the Datadog agent's premain has already executed, meaning it was listed
132133
* before the Braintrust agent in the {@code -javaagent} flags.
133134
*/
134-
static boolean isRunningAfterDatadogAgent() {
135+
private static boolean isRunningAfterDatadogAgent() {
135136
// DD's premain appends its jars to the bootstrap classpath, making
136137
// {@code datadog.trace.bootstrap.Agent} loadable from the bootstrap (null)
137138
// classloader. If that class is not found on bootstrap, DD either isn't

‎braintrust-java-agent/internal/build.gradle‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,9 @@ dependencies {
4040
// These are the heavy deps that stay in BraintrustClassLoader, NOT on bootstrap.
4141
implementation "io.opentelemetry:opentelemetry-exporter-otlp:${otelVersion}"
4242

43+
// for dd compat mode
44+
compileOnly 'com.datadoghq:dd-trace-api:1.60.1'
45+
4346
// Test dependencies
4447
testImplementation project(':braintrust-java-agent:bootstrap')
4548
testImplementation "org.junit.jupiter:junit-jupiter:${junitVersion}"

‎braintrust-java-agent/internal/src/main/java/dev/braintrust/agent/BraintrustAgent.java‎

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import com.google.auto.service.AutoService;
44
import dev.braintrust.Braintrust;
5+
import dev.braintrust.agent.dd.BTInterceptor;
56
import dev.braintrust.bootstrap.BraintrustBridge;
67
import dev.braintrust.bootstrap.BraintrustClassLoader;
78
import dev.braintrust.instrumentation.Instrumenter;
@@ -21,7 +22,8 @@ public class BraintrustAgent implements AutoConfigurationCustomizerProvider {
2122
public static void install(String agentArgs, Instrumentation inst) {
2223
if (!(BraintrustAgent.class.getClassLoader() instanceof BraintrustClassLoader)) {
2324
throw new IllegalStateException(
24-
"Braintrust agent can only run on a braintrust classloader");
25+
"Braintrust agent can only run on a braintrust classloader: "
26+
+ BraintrustAgent.class.getClassLoader());
2527
}
2628
log.info(
2729
"invoked on classloader: {}",
@@ -31,6 +33,9 @@ public static void install(String agentArgs, Instrumentation inst) {
3133
// Fail fast if there are any issues with the Braintrust SDK
3234
Braintrust.get();
3335
Instrumenter.install(inst, BraintrustAgent.class.getClassLoader());
36+
if (jvmRunningWithDatadogOtelConfig() && ddApiOnBootstrapClasspath()) {
37+
BTInterceptor.install();
38+
}
3439
}
3540

3641
@Override
@@ -51,4 +56,24 @@ public void customize(AutoConfigurationCustomizer autoConfiguration) {
5156
return sdkTracerProviderBuilder;
5257
}));
5358
}
59+
60+
/** Checks whether the Datadog agent is present and configured for OTel integration */
61+
private static boolean ddApiOnBootstrapClasspath() {
62+
try {
63+
BraintrustAgent.class.getClassLoader().loadClass("datadog.trace.api.GlobalTracer");
64+
return true;
65+
} catch (ClassNotFoundException e) {
66+
return false;
67+
}
68+
}
69+
70+
/** Checks whether the Datadog agent is present and configured for OTel integration */
71+
private static boolean jvmRunningWithDatadogOtelConfig() {
72+
String sysProp = System.getProperty("dd.trace.otel.enabled");
73+
if (sysProp != null) {
74+
return Boolean.parseBoolean(sysProp);
75+
}
76+
String envVar = System.getenv("DD_TRACE_OTEL_ENABLED");
77+
return Boolean.parseBoolean(envVar);
78+
}
5479
}

‎braintrust-java-agent/internal/src/main/java/dev/braintrust/agent/dd/BTInterceptor.java‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
package dev.braintrust.agent.dd;
2+
3+
import datadog.trace.api.GlobalTracer;
4+
import datadog.trace.api.interceptor.MutableSpan;
5+
import datadog.trace.api.interceptor.TraceInterceptor;
6+
import dev.braintrust.Braintrust;
7+
import dev.braintrust.trace.BraintrustTracing;
8+
import io.opentelemetry.api.trace.Tracer;
9+
import io.opentelemetry.sdk.logs.SdkLoggerProvider;
10+
import io.opentelemetry.sdk.metrics.SdkMeterProvider;
11+
import io.opentelemetry.sdk.trace.SdkTracerProvider;
12+
import io.opentelemetry.sdk.trace.data.SpanData;
13+
import java.util.Collection;
14+
import java.util.List;
15+
import java.util.concurrent.atomic.AtomicBoolean;
16+
import lombok.extern.slf4j.Slf4j;
17+
18+
@Slf4j
19+
public class BTInterceptor implements TraceInterceptor {
20+
private static final AtomicBoolean installed = new AtomicBoolean(false);
21+
22+
public static void install() {
23+
if (!installed.compareAndExchange(false, true)) {
24+
try {
25+
if (!DDSpanConverter.initialize()) {
26+
log.warn(
27+
"failed to initialize DD span converter. Braintrust traces will not be"
28+
+ " reported.");
29+
return;
30+
}
31+
var tbBuilder =
32+
SdkTracerProvider.builder().setIdGenerator(OverridableIdGenerator.INSTANCE);
33+
Braintrust.get()
34+
.openTelemetryEnable(
35+
tbBuilder, SdkLoggerProvider.builder(), SdkMeterProvider.builder());
36+
final var traceProvider = tbBuilder.build();
37+
var interceptor = new BTInterceptor(999, traceProvider);
38+
if (!GlobalTracer.get().addTraceInterceptor(interceptor)) {
39+
log.warn(
40+
"trace interceptor install failed due to conflicting priorities."
41+
+ " Braintrust traces will not be reported.");
42+
return;
43+
}
44+
log.info("trace interceptor successfully installed");
45+
} catch (Exception e) {
46+
log.warn(
47+
"trace interceptor install failed. Braintrust traces will not be reported.",
48+
e);
49+
// Don't reset the flag. We don't want to try again.
50+
}
51+
}
52+
}
53+
54+
private final int priority;
55+
private final Tracer tracer;
56+
57+
private BTInterceptor(int priority, SdkTracerProvider traceProvider) {
58+
this.priority = priority;
59+
this.tracer = BraintrustTracing.getTracer(traceProvider);
60+
}
61+
62+
@Override
63+
public int priority() {
64+
return priority;
65+
}
66+
67+
@Override
68+
public Collection<? extends MutableSpan> onTraceComplete(
69+
Collection<? extends MutableSpan> trace) {
70+
try {
71+
List<SpanData> spanDataList = DDSpanConverter.convertTrace(List.copyOf(trace));
72+
DDSpanConverter.replayTrace(tracer, spanDataList);
73+
} catch (Exception e) {
74+
log.debug("failed to replay traces", e);
75+
}
76+
return trace;
77+
}
78+
}

0 commit comments

Comments
 (0)