Skip to content

Commit 03b4063

Browse files
realarkrealark
committed
autoinstrumentation: datadog otel compatibility
1 parent 91263fe commit 03b4063

18 files changed

Lines changed: 2706 additions & 34 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: 31 additions & 30 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 (jvmRunningWithDatadogOtel() && (!isRunningAfterDatadogAgent())) {
51+
log("ERROR: Braintrust agent must run _after_ datadog -javaagent. aborting install.");
5652
return;
5753
}
5854

55+
boolean installOnBootstrap = !jvmRunningWithDatadogOtel();
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-
*/
121+
/** Checks whether the Datadog agent is present and configured for OTel integration */
116122
private static boolean jvmRunningWithDatadogOtel() {
117-
try {
118-
Class.forName("datadog.trace.bootstrap.Agent", false, null);
119-
} catch (ClassNotFoundException e) {
120-
return false;
121-
}
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: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,8 @@ public class BraintrustAgent implements AutoConfigurationCustomizerProvider {
2121
public static void install(String agentArgs, Instrumentation inst) {
2222
if (!(BraintrustAgent.class.getClassLoader() instanceof BraintrustClassLoader)) {
2323
throw new IllegalStateException(
24-
"Braintrust agent can only run on a braintrust classloader");
24+
"Braintrust agent can only run on a braintrust classloader: "
25+
+ BraintrustAgent.class.getClassLoader());
2526
}
2627
log.info(
2728
"invoked on classloader: {}",

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

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
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+
class BTInterceptor implements TraceInterceptor {
20+
private static final AtomicBoolean installed = new AtomicBoolean(false);
21+
22+
static synchronized void install() {
23+
if (!installed.compareAndExchange(false, true)) {
24+
try {
25+
var tbBuilder = SdkTracerProvider.builder();
26+
Braintrust.get()
27+
.openTelemetryEnable(
28+
tbBuilder, SdkLoggerProvider.builder(), SdkMeterProvider.builder());
29+
final var traceProvider = tbBuilder.build();
30+
var interceptor = new BTInterceptor(999, traceProvider);
31+
if (!GlobalTracer.get().addTraceInterceptor(interceptor)) {
32+
log.error(
33+
"trace interceptor install failed due to conflicting priorities."
34+
+ " Braintrust traces will not be reported.");
35+
return;
36+
}
37+
log.info("trace interceptor successfully installed");
38+
} catch (Exception e) {
39+
log.error(
40+
"trace interceptor install failed. Braintrust traces will not be reported.",
41+
e);
42+
// Don't reset the flag. We don't want to try again.
43+
}
44+
}
45+
}
46+
47+
private final int priority;
48+
private final Tracer tracer;
49+
50+
private BTInterceptor(int priority, SdkTracerProvider traceProvider) {
51+
this.priority = priority;
52+
this.tracer = BraintrustTracing.getTracer(traceProvider);
53+
}
54+
55+
@Override
56+
public int priority() {
57+
return priority;
58+
}
59+
60+
@Override
61+
public Collection<? extends MutableSpan> onTraceComplete(
62+
Collection<? extends MutableSpan> trace) {
63+
try {
64+
List<SpanData> spanDataList = DDSpanConverter.convertTrace(List.copyOf(trace));
65+
DDSpanConverter.replayTrace(tracer, spanDataList);
66+
} catch (Exception e) {
67+
log.debug("failed to replay traces", e);
68+
}
69+
return trace;
70+
}
71+
}

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

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
package dev.braintrust.agent.dd;
2+
3+
import static net.bytebuddy.matcher.ElementMatchers.*;
4+
5+
import com.google.auto.service.AutoService;
6+
import dev.braintrust.instrumentation.InstrumentationModule;
7+
import dev.braintrust.instrumentation.TypeInstrumentation;
8+
import dev.braintrust.instrumentation.TypeTransformer;
9+
import java.util.List;
10+
import net.bytebuddy.description.type.TypeDescription;
11+
import net.bytebuddy.matcher.ElementMatcher;
12+
13+
/**
14+
* Replaces {@code GlobalOpenTelemetry.maybeAutoConfigureAndSetGlobal()} so that the Braintrust
15+
* agent controls OTel SDK initialization lazily — triggered on the first call to {@code
16+
* GlobalOpenTelemetry.get()}.
17+
*/
18+
@AutoService(InstrumentationModule.class)
19+
public class DDBridgeInstrumentationModule extends InstrumentationModule {
20+
public DDBridgeInstrumentationModule() {
21+
super("dd-bridge");
22+
}
23+
24+
@Override
25+
public List<TypeInstrumentation> typeInstrumentations() {
26+
return List.of(new GlobalTracerTypeInstrumentation());
27+
}
28+
29+
static class GlobalTracerTypeInstrumentation implements TypeInstrumentation {
30+
@Override
31+
public ElementMatcher<TypeDescription> typeMatcher() {
32+
return named("datadog.trace.api.GlobalTracer")
33+
.and(
34+
typeDefinitions -> {
35+
var ddOtel = jvmRunningWithDatadogOtel();
36+
if (ddOtel) {
37+
BTInterceptor.install();
38+
}
39+
return ddOtel;
40+
});
41+
}
42+
43+
@Override
44+
public void transform(TypeTransformer transformer) {}
45+
}
46+
47+
/** Checks whether the Datadog agent is present and configured for OTel integration */
48+
private static boolean jvmRunningWithDatadogOtel() {
49+
String sysProp = System.getProperty("dd.trace.otel.enabled");
50+
if (sysProp != null) {
51+
return Boolean.parseBoolean(sysProp);
52+
}
53+
String envVar = System.getenv("DD_TRACE_OTEL_ENABLED");
54+
return Boolean.parseBoolean(envVar);
55+
}
56+
}

0 commit comments

Comments
 (0)