Skip to content

Commit 58cfdba

Browse files
jordan-wongjordan-wong
committed
feat(r2dbc): add R2DBC instrumentation via r2dbc-proxy listener SPI
Adds tracing for R2DBC (reactive database access) by installing an official r2dbc-proxy ProxyMethodExecutionListener instead of hand-writing advice on Statement.execute()/Batch.execute(). The listener owns the reactive lifecycle (complete/error/cancel) so spans can't leak on cancellation, which a hand-rolled Publisher wrapper is prone to since reactive pipelines cancel constantly (take(1), timeouts, driver-level DiscardOnCancel). - R2dbcInstrumentation hooks ConnectionFactories.find(ConnectionFactoryOptions) via @Advice.AssignReturned-style exit advice and wraps the returned ConnectionFactory with r2dbc-proxy. - TraceProxyExecutionListener implements ProxyMethodExecutionListener, starting/finishing spans in beforeQuery/afterQuery. - R2dbcConnectionCallbackInstrumentation + R2dbcSqlCommentInjector add DBM SQL comment injection, mirroring JDBC's SQLCommenter. - r2dbc-proxy is declared implementation (not compileOnly) and its full class set is listed in helperClassNames() in dependency order: ordinary R2DBC apps don't depend on r2dbc-proxy, so its classes must be injected into the app's classloader (and added to the muzzle validation classpath via extraDependency) or the module is blocked/fails to load at runtime despite passing build-time checks. Verified via a live trace comparison against the OpenTelemetry Java agent on a Spring WebFlux + R2DBC/PostgreSQL sample app: matching INSERT/CREATE TABLE/error-SELECT spans with correct DBM tags (db.instance, db.type, db.user, peer.hostname, _dd.dbm_trace_injected).
1 parent 241817d commit 58cfdba

13 files changed

Lines changed: 1643 additions & 0 deletions

File tree

‎dd-java-agent/instrumentation/build.gradle‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,10 @@ tasks.named('shadowJar', ShadowJar) {
135135
exclude(dependency('com.google.re2j:re2j'))
136136
deps.excludeShared.execute(it)
137137
}
138+
// Redundant metadata file duplicated across sibling io.r2dbc:* artifacts
139+
// (r2dbc-spi, r2dbc-proxy) once a module bundles more than one of them.
140+
// Not required content — drop it instead of failing the aggregate jar.
141+
exclude 'META-INF/CHANGELOG'
138142
}
139143

140144
// temporary config to add slf4j-simple so we get logging from instrumenters while indexing

‎dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/build.gradle‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
plugins {
2+
id 'dd-trace-java.module.instrumentation'
3+
}
4+
5+
muzzle {
6+
pass {
7+
group = "io.r2dbc"
8+
module = "r2dbc-spi"
9+
versions = "[1.0.0.RELEASE,)"
10+
// r2dbc-proxy is compileOnly and referenced directly by the injected
11+
// listener/wrap-helper classes (R2dbcTracingSupport, TraceProxyExecutionListener).
12+
// Muzzle only puts the pinned primary dependency (r2dbc-spi) on the test
13+
// classpath by default, so the compileOnly interception library must be
14+
// added explicitly or muzzle reports its classes as "missing".
15+
extraDependency 'io.r2dbc:r2dbc-proxy:1.1.0.RELEASE'
16+
}
17+
}
18+
19+
addTestSuiteForDir('latestDepTest', 'test')
20+
21+
dependencies {
22+
compileOnly group: 'io.r2dbc', name: 'r2dbc-spi', version: '1.0.0.RELEASE'
23+
// r2dbc-proxy must be bundled into the agent (implementation, NOT compileOnly):
24+
// R2dbcTracingSupport/TraceProxyExecutionListener reference its types directly, and
25+
// ordinary R2DBC apps do not depend on r2dbc-proxy themselves. compileOnly would leave
26+
// those classes absent from the shaded agent jar, and the runtime muzzle safety check
27+
// then blocks the whole module to avoid a NoClassDefFoundError once the target app's
28+
// classloader is checked. r2dbc-spi stays compileOnly — target apps DO provide that one.
29+
implementation group: 'io.r2dbc', name: 'r2dbc-proxy', version: '1.1.0.RELEASE'
30+
31+
testImplementation group: 'io.r2dbc', name: 'r2dbc-spi', version: '1.0.0.RELEASE'
32+
testImplementation group: 'io.r2dbc', name: 'r2dbc-proxy', version: '1.1.0.RELEASE'
33+
// H2 R2DBC driver for in-memory database testing
34+
testImplementation group: 'io.r2dbc', name: 'r2dbc-h2', version: '1.0.0.RELEASE'
35+
// Reactor for blocking on reactive types in tests
36+
testImplementation group: 'io.projectreactor', name: 'reactor-core', version: '3.5.0'
37+
38+
latestDepTestImplementation group: 'io.r2dbc', name: 'r2dbc-spi', version: '1.+'
39+
latestDepTestImplementation group: 'io.r2dbc', name: 'r2dbc-proxy', version: '1.+'
40+
latestDepTestImplementation group: 'io.r2dbc', name: 'r2dbc-h2', version: '1.+'
41+
latestDepTestImplementation group: 'io.projectreactor', name: 'reactor-core', version: '3.+'
42+
}

‎dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcConnectionCallbackInstrumentation.java‎

Lines changed: 207 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,207 @@
1+
package datadog.trace.instrumentation.r2dbc;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static datadog.trace.instrumentation.r2dbc.R2dbcDecorator.DECORATE;
5+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
6+
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
7+
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
8+
9+
import com.google.auto.service.AutoService;
10+
import datadog.trace.agent.tooling.Instrumenter;
11+
import datadog.trace.agent.tooling.InstrumenterModule;
12+
import datadog.trace.api.Config;
13+
import io.r2dbc.proxy.core.ConnectionInfo;
14+
import io.r2dbc.spi.ConnectionFactoryOptions;
15+
import java.lang.reflect.Method;
16+
import net.bytebuddy.asm.Advice;
17+
18+
/**
19+
* Instruments {@code io.r2dbc.proxy.callback.ConnectionCallbackHandler} to inject DBM SQL comments
20+
* into queries before they reach the database driver. This is the R2DBC equivalent of JDBC's {@code
21+
* DBMCompatibleConnectionInstrumentation}.
22+
*
23+
* <p>The r2dbc-proxy library uses JDK dynamic proxies for Connection objects, so we cannot
24+
* instrument them with ByteBuddy directly. Instead, we intercept the callback handler's {@code
25+
* invoke} method which is called for every method on the proxied Connection. When {@code
26+
* createStatement(String)} is invoked, we inject the SQL comment into the first argument.
27+
*/
28+
@AutoService(InstrumenterModule.class)
29+
public class R2dbcConnectionCallbackInstrumentation extends InstrumenterModule.Tracing
30+
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
31+
32+
public R2dbcConnectionCallbackInstrumentation() {
33+
super("r2dbc");
34+
}
35+
36+
@Override
37+
public String instrumentedType() {
38+
return "io.r2dbc.proxy.callback.ConnectionCallbackHandler";
39+
}
40+
41+
@Override
42+
public String[] helperClassNames() {
43+
return new String[] {
44+
// See R2dbcInstrumentation#helperClassNames for why the full r2dbc-proxy class set
45+
// (rather than a hand-picked subset) is required, and for the topological ordering
46+
// rationale (supertypes must be injected before their implementing classes).
47+
"io.r2dbc.proxy.callback.AfterQueryCallbackInvoker",
48+
"io.r2dbc.proxy.callback.CallbackHandler",
49+
"io.r2dbc.proxy.callback.CallbackHandlerSupport",
50+
"io.r2dbc.proxy.callback.BatchCallbackHandler",
51+
"io.r2dbc.proxy.callback.CallbackHandlerSupport$MethodInvocationStrategy",
52+
"io.r2dbc.proxy.callback.ConnectionCallbackHandler",
53+
"io.r2dbc.proxy.callback.ConnectionFactoryCallbackHandler",
54+
"io.r2dbc.proxy.callback.MethodInvocationSubscriber",
55+
"io.r2dbc.proxy.callback.ConnectionFactoryCreateMethodInvocationSubscriber",
56+
"io.r2dbc.proxy.callback.ConnectionHolder",
57+
"io.r2dbc.proxy.callback.ConnectionIdManager",
58+
"io.r2dbc.proxy.callback.DefaultConnectionIdManager",
59+
"io.r2dbc.proxy.core.ConnectionInfo",
60+
"io.r2dbc.proxy.callback.DefaultConnectionInfo",
61+
"io.r2dbc.proxy.callback.DelegatingContextView",
62+
"io.r2dbc.proxy.callback.ProxyFactory",
63+
"io.r2dbc.proxy.callback.JdkProxyFactory",
64+
"io.r2dbc.proxy.callback.JdkProxyFactory$CallbackInvocationHandler",
65+
"io.r2dbc.proxy.callback.ProxyFactoryFactory",
66+
"io.r2dbc.proxy.callback.JdkProxyFactoryFactory",
67+
"io.r2dbc.proxy.core.BindInfo",
68+
"io.r2dbc.proxy.callback.MutableBindInfo",
69+
"io.r2dbc.proxy.core.MethodExecutionInfo",
70+
"io.r2dbc.proxy.callback.MutableMethodExecutionInfo",
71+
"io.r2dbc.proxy.core.QueryExecutionInfo",
72+
"io.r2dbc.proxy.callback.MutableQueryExecutionInfo",
73+
"io.r2dbc.proxy.core.StatementInfo",
74+
"io.r2dbc.proxy.callback.MutableStatementInfo",
75+
"io.r2dbc.proxy.callback.ProxyConfig",
76+
"io.r2dbc.proxy.callback.ProxyConfig$1",
77+
"io.r2dbc.proxy.callback.ProxyConfig$Builder",
78+
"io.r2dbc.proxy.callback.ProxyConfigHolder",
79+
"io.r2dbc.proxy.callback.ProxyUtils",
80+
"io.r2dbc.proxy.callback.QueriesExecutionContext",
81+
"io.r2dbc.proxy.callback.QueryInvocationSubscriber",
82+
"io.r2dbc.proxy.callback.ResultCallbackHandler",
83+
"io.r2dbc.proxy.callback.ResultInvocationSubscriber",
84+
"io.r2dbc.proxy.callback.RowCallbackHandler",
85+
"io.r2dbc.proxy.callback.StatementCallbackHandler",
86+
"io.r2dbc.proxy.callback.StopWatch",
87+
"io.r2dbc.proxy.core.Binding",
88+
"io.r2dbc.proxy.core.Bindings",
89+
"io.r2dbc.proxy.core.Bindings$1",
90+
"io.r2dbc.proxy.core.Bindings$IndexBinding",
91+
"io.r2dbc.proxy.core.Bindings$NamedBinding",
92+
"io.r2dbc.proxy.core.BoundValue",
93+
"io.r2dbc.proxy.core.BoundValue$DefaultBoundValue",
94+
"io.r2dbc.proxy.core.ValueStore",
95+
"io.r2dbc.proxy.core.DefaultValueStore",
96+
"io.r2dbc.proxy.core.ExecutionType",
97+
"io.r2dbc.proxy.core.ProxyEventType",
98+
"io.r2dbc.proxy.core.QueryInfo",
99+
"io.r2dbc.proxy.core.R2dbcProxyException",
100+
"io.r2dbc.proxy.listener.BindParameterConverter",
101+
"io.r2dbc.proxy.listener.BindParameterConverter$1",
102+
"io.r2dbc.proxy.listener.BindParameterConverter$BindOperation",
103+
"io.r2dbc.proxy.listener.ProxyExecutionListener",
104+
"io.r2dbc.proxy.listener.CompositeProxyExecutionListener",
105+
"io.r2dbc.proxy.listener.LastExecutionAwareListener",
106+
"io.r2dbc.proxy.listener.ProxyMethodExecutionListener",
107+
"io.r2dbc.proxy.listener.ProxyMethodExecutionListenerAdapter",
108+
"io.r2dbc.proxy.listener.ResultRowConverter",
109+
"io.r2dbc.proxy.listener.ResultRowConverter$GetOperation",
110+
"io.r2dbc.proxy.ProxyConnectionFactory",
111+
"io.r2dbc.proxy.ProxyConnectionFactory$1",
112+
"io.r2dbc.proxy.ProxyConnectionFactory$Builder",
113+
"io.r2dbc.proxy.ProxyConnectionFactory$Builder$1",
114+
"io.r2dbc.proxy.ProxyConnectionFactory$Builder$2",
115+
"io.r2dbc.proxy.ProxyConnectionFactory$Builder$3",
116+
"io.r2dbc.proxy.ProxyConnectionFactory$Builder$4",
117+
"io.r2dbc.proxy.ProxyConnectionFactory$Builder$5",
118+
"io.r2dbc.proxy.ProxyConnectionFactoryProvider",
119+
"io.r2dbc.proxy.support.FormatterUtils",
120+
"io.r2dbc.proxy.support.MethodExecutionInfoFormatter",
121+
"io.r2dbc.proxy.support.QueryExecutionInfoFormatter",
122+
"io.r2dbc.proxy.util.Assert",
123+
packageName + ".R2dbcDecorator",
124+
packageName + ".R2dbcSqlCommentInjector",
125+
packageName + ".R2dbcTracingSupport",
126+
packageName + ".R2dbcTracingSupport$ConnectionMetadataListener",
127+
packageName + ".TraceProxyExecutionListener",
128+
};
129+
}
130+
131+
@Override
132+
public void methodAdvice(MethodTransformer transformer) {
133+
transformer.applyAdvice(
134+
isMethod()
135+
.and(named("invoke"))
136+
.and(takesArguments(3))
137+
.and(takesArgument(0, Object.class))
138+
.and(takesArgument(1, Method.class))
139+
.and(takesArgument(2, Object[].class)),
140+
getClass().getName() + "$InvokeAdvice");
141+
}
142+
143+
public static class InvokeAdvice {
144+
145+
@Advice.OnMethodEnter(suppress = Throwable.class)
146+
public static void onEnter(
147+
@Advice.Argument(1) final Method method,
148+
@Advice.Argument(value = 2, readOnly = false) Object[] args,
149+
@Advice.FieldValue("connectionInfo") final ConnectionInfo connectionInfo) {
150+
if (args == null || args.length == 0) {
151+
return;
152+
}
153+
if (!"createStatement".equals(method.getName())) {
154+
return;
155+
}
156+
if (!(args[0] instanceof String)) {
157+
return;
158+
}
159+
160+
String dbmMode = Config.get().getDbmPropagationMode();
161+
boolean injectComment =
162+
Config.DBM_PROPAGATION_MODE_FULL.equals(dbmMode)
163+
|| Config.DBM_PROPAGATION_MODE_STATIC.equals(dbmMode)
164+
|| Config.DBM_PROPAGATION_MODE_DYNAMIC_SERVICE.equals(dbmMode);
165+
if (!injectComment) {
166+
return;
167+
}
168+
169+
String sql = (String) args[0];
170+
171+
// Look up connection metadata from the map maintained by R2dbcTracingSupport
172+
String hostname = null;
173+
String dbName = null;
174+
String dbService = null;
175+
String dbType = null;
176+
177+
ConnectionFactoryOptions options = R2dbcTracingSupport.CONNECTION_OPTIONS.get(connectionInfo);
178+
if (options != null) {
179+
dbType = DECORATE.extractDbType(options);
180+
dbService = DECORATE.getDbService(options);
181+
CharSequence hostnameSeq = null;
182+
if (options.hasOption(ConnectionFactoryOptions.HOST)) {
183+
Object host = options.getValue(ConnectionFactoryOptions.HOST);
184+
if (host != null) {
185+
hostnameSeq = host.toString();
186+
}
187+
}
188+
hostname = hostnameSeq != null ? hostnameSeq.toString() : null;
189+
if (options.hasOption(ConnectionFactoryOptions.DATABASE)) {
190+
Object db = options.getValue(ConnectionFactoryOptions.DATABASE);
191+
dbName = db != null ? db.toString() : null;
192+
}
193+
}
194+
195+
String injected = R2dbcSqlCommentInjector.inject(sql, dbService, dbType, hostname, dbName);
196+
if (!sql.equals(injected)) {
197+
// Replace the SQL argument with the injected version.
198+
// We must create a new array because ByteBuddy advice cannot mutate the original
199+
// array reference in place for @Advice.Argument(readOnly=false).
200+
Object[] newArgs = new Object[args.length];
201+
System.arraycopy(args, 0, newArgs, 0, args.length);
202+
newArgs[0] = injected;
203+
args = newArgs;
204+
}
205+
}
206+
}
207+
}

‎dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcDecorator.java‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
package datadog.trace.instrumentation.r2dbc;
2+
3+
import datadog.trace.api.naming.SpanNaming;
4+
import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes;
5+
import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
6+
import datadog.trace.bootstrap.instrumentation.decorator.DatabaseClientDecorator;
7+
import io.r2dbc.spi.ConnectionFactoryOptions;
8+
9+
public class R2dbcDecorator extends DatabaseClientDecorator<ConnectionFactoryOptions> {
10+
11+
public static final R2dbcDecorator DECORATE = new R2dbcDecorator();
12+
13+
static final CharSequence R2DBC_QUERY =
14+
UTF8BytesString.create(SpanNaming.instance().namingSchema().database().operation("r2dbc"));
15+
private static final CharSequence R2DBC = UTF8BytesString.create("r2dbc");
16+
private static final String DEFAULT_SERVICE_NAME =
17+
SpanNaming.instance().namingSchema().database().service("r2dbc");
18+
19+
@Override
20+
protected String[] instrumentationNames() {
21+
return new String[] {"r2dbc"};
22+
}
23+
24+
@Override
25+
protected String service() {
26+
return DEFAULT_SERVICE_NAME;
27+
}
28+
29+
@Override
30+
protected CharSequence component() {
31+
return R2DBC;
32+
}
33+
34+
@Override
35+
protected CharSequence spanType() {
36+
return InternalSpanTypes.SQL;
37+
}
38+
39+
@Override
40+
protected String dbType() {
41+
return "r2dbc";
42+
}
43+
44+
@Override
45+
protected String dbUser(ConnectionFactoryOptions options) {
46+
if (options == null) {
47+
return null;
48+
}
49+
Object user = options.getValue(ConnectionFactoryOptions.USER);
50+
return user != null ? user.toString() : null;
51+
}
52+
53+
@Override
54+
protected String dbInstance(ConnectionFactoryOptions options) {
55+
if (options == null) {
56+
return null;
57+
}
58+
Object database = options.getValue(ConnectionFactoryOptions.DATABASE);
59+
return database != null ? database.toString() : null;
60+
}
61+
62+
@Override
63+
protected CharSequence dbHostname(ConnectionFactoryOptions options) {
64+
if (options == null) {
65+
return null;
66+
}
67+
Object host = options.getValue(ConnectionFactoryOptions.HOST);
68+
return host != null ? host.toString() : null;
69+
}
70+
71+
public String extractDbType(ConnectionFactoryOptions options) {
72+
if (options != null && options.hasOption(ConnectionFactoryOptions.DRIVER)) {
73+
Object driver = options.getValue(ConnectionFactoryOptions.DRIVER);
74+
if (driver != null) {
75+
return driver.toString();
76+
}
77+
}
78+
return "r2dbc";
79+
}
80+
81+
/** Exposes the protected {@link #processDatabaseType} for use by the listener. */
82+
public void applyDatabaseType(
83+
datadog.trace.bootstrap.instrumentation.api.AgentSpan span, String dbType) {
84+
processDatabaseType(span, dbType);
85+
}
86+
87+
/**
88+
* Returns the database service name derived from the connection options. Used for DBM SQL comment
89+
* injection.
90+
*/
91+
public String getDbService(ConnectionFactoryOptions options) {
92+
String dbType = extractDbType(options);
93+
String instanceName = dbInstance(options);
94+
return dbService(dbType, instanceName);
95+
}
96+
97+
@Override
98+
protected void postProcessServiceAndOperationName(
99+
datadog.trace.bootstrap.instrumentation.api.AgentSpan span, NamingEntry namingEntry) {
100+
if (namingEntry.getService() != null) {
101+
span.setServiceName(namingEntry.getService(), component());
102+
}
103+
span.setOperationName(namingEntry.getOperation());
104+
}
105+
}

0 commit comments

Comments
 (0)