From be98914d223c34c6bb404a1d9f2edb8fb5a32a43 Mon Sep 17 00:00:00 2001 From: Jordan Wong Date: Wed, 22 Jul 2026 09:11:37 -0400 Subject: [PATCH 1/4] feat(r2dbc): net-new instrumentation for io.r2dbc:r2dbc-spi 1.0.0 Reference PR from toolkit net-new generation (no prior dd-trace-java module exists for R2DBC, so this is not a blind regen). Generated against io.r2dbc:r2dbc-spi:1.0.0.RELEASE. Co-Authored-By: Claude Sonnet 5 --- .../agent/test/assertions/TagsMatcher.java | 2 + .../r2dbc/r2dbc-1.0/build.gradle | 20 ++ .../r2dbc/BatchInstrumentation.java | 67 ++++ .../r2dbc/ConnectionInstrumentation.java | 77 +++++ .../r2dbc/R2dbcConnectionInfo.java | 53 +++ .../instrumentation/r2dbc/R2dbcDecorator.java | 129 +++++++ .../r2dbc/R2dbcInstrumenterModule.java | 44 +++ .../r2dbc/R2dbcSQLCommenter.java | 65 ++++ .../r2dbc/StatementInstrumentation.java | 96 ++++++ .../r2dbc/TracingPublisher.java | 70 ++++ .../r2dbc/R2dbcDBMForkedTest.java | 253 ++++++++++++++ .../r2dbc/R2dbcInstrumentationTest.java | 314 ++++++++++++++++++ .../src/test/java/r2dbc/test/TestBatch.java | 54 +++ .../test/java/r2dbc/test/TestConnection.java | 141 ++++++++ .../r2dbc/test/TestConnectionFactory.java | 31 ++ .../src/test/java/r2dbc/test/TestResult.java | 67 ++++ .../test/java/r2dbc/test/TestStatement.java | 79 +++++ metadata/supported-configurations.json | 24 ++ settings.gradle.kts | 1 + .../trace/test/junit/utils/assertions/Is.java | 3 + 20 files changed, 1590 insertions(+) create mode 100644 dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/build.gradle create mode 100644 dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/BatchInstrumentation.java create mode 100644 dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/ConnectionInstrumentation.java create mode 100644 dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcConnectionInfo.java create mode 100644 dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcDecorator.java create mode 100644 dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcInstrumenterModule.java create mode 100644 dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcSQLCommenter.java create mode 100644 dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/StatementInstrumentation.java create mode 100644 dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/TracingPublisher.java create mode 100644 dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/datadog/trace/instrumentation/r2dbc/R2dbcDBMForkedTest.java create mode 100644 dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/datadog/trace/instrumentation/r2dbc/R2dbcInstrumentationTest.java create mode 100644 dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestBatch.java create mode 100644 dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestConnection.java create mode 100644 dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestConnectionFactory.java create mode 100644 dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestResult.java create mode 100644 dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestStatement.java diff --git a/dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/assertions/TagsMatcher.java b/dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/assertions/TagsMatcher.java index 7fe7cd0645e..c3a5e338f56 100644 --- a/dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/assertions/TagsMatcher.java +++ b/dd-java-agent/instrumentation-testing/src/main/java/datadog/trace/agent/test/assertions/TagsMatcher.java @@ -2,6 +2,7 @@ import static datadog.trace.api.DDTags.BASE_SERVICE; import static datadog.trace.api.DDTags.DD_INTEGRATION; +import static datadog.trace.api.DDTags.DD_SVC_SRC; import static datadog.trace.api.DDTags.DJM_ENABLED; import static datadog.trace.api.DDTags.DSM_ENABLED; import static datadog.trace.api.DDTags.ERROR_MSG; @@ -57,6 +58,7 @@ public static TagsMatcher defaultTags() { tagMatchers.put(PARENT_ID, any()); tagMatchers.put(SPAN_LINKS, any()); // this is checked by LinksAsserter tagMatchers.put(DD_INTEGRATION, any()); + tagMatchers.put(DD_SVC_SRC, any()); tagMatchers.put(TRACER_HOST, any()); for (String tagName : REQUIRED_CODE_ORIGIN_TAGS) { diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/build.gradle b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/build.gradle new file mode 100644 index 00000000000..6a1deb41fcf --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/build.gradle @@ -0,0 +1,20 @@ +muzzle { + pass { + group = "io.r2dbc" + module = "r2dbc-spi" + versions = "[1.0.0.RELEASE,)" + assertInverse = true + } +} + +apply from: "$rootDir/gradle/java.gradle" + +addTestSuiteForDir('latestDepTest', 'test') + +dependencies { + compileOnly group: 'io.r2dbc', name: 'r2dbc-spi', version: '1.0.0.RELEASE' + + testImplementation group: 'io.r2dbc', name: 'r2dbc-spi', version: '1.0.0.RELEASE' + + latestDepTestImplementation group: 'io.r2dbc', name: 'r2dbc-spi', version: '1.+' +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/BatchInstrumentation.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/BatchInstrumentation.java new file mode 100644 index 00000000000..a21f6efb98e --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/BatchInstrumentation.java @@ -0,0 +1,67 @@ +package datadog.trace.instrumentation.r2dbc; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.instrumentation.r2dbc.R2dbcDecorator.DECORATE; +import static datadog.trace.instrumentation.r2dbc.R2dbcDecorator.R2DBC_BATCH; +import static net.bytebuddy.matcher.ElementMatchers.isMethod; +import static net.bytebuddy.matcher.ElementMatchers.isPublic; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; + +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; +import org.reactivestreams.Publisher; + +public class BatchInstrumentation + implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice { + + @Override + public String hierarchyMarkerType() { + return "io.r2dbc.spi.Batch"; + } + + @Override + public ElementMatcher hierarchyMatcher() { + return implementsInterface(named("io.r2dbc.spi.Batch")); + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + isMethod().and(isPublic()).and(named("execute")).and(takesArguments(0)), + BatchInstrumentation.class.getName() + "$BatchExecuteAdvice"); + } + + public static class BatchExecuteAdvice { + + @Advice.OnMethodEnter(suppress = Throwable.class) + public static AgentScope onEnter() { + AgentSpan span = startSpan("r2dbc", R2DBC_BATCH); + DECORATE.afterStart(span); + return activateSpan(span); + } + + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) + public static void onExit( + @Advice.Enter final AgentScope scope, + @Advice.Return(readOnly = false) Publisher publisher, + @Advice.Thrown final Throwable throwable) { + AgentSpan span = scope.span(); + if (throwable != null) { + DECORATE.onError(span, throwable); + DECORATE.beforeFinish(span); + scope.close(); + span.finish(); + } else { + publisher = new TracingPublisher<>(publisher, span); + scope.close(); + } + } + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/ConnectionInstrumentation.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/ConnectionInstrumentation.java new file mode 100644 index 00000000000..d479d3b6980 --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/ConnectionInstrumentation.java @@ -0,0 +1,77 @@ +package datadog.trace.instrumentation.r2dbc; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.instrumentation.r2dbc.R2dbcDecorator.INJECT_COMMENT; +import static net.bytebuddy.matcher.ElementMatchers.isMethod; +import static net.bytebuddy.matcher.ElementMatchers.isPublic; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; + +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.bootstrap.InstrumentationContext; +import io.r2dbc.spi.Connection; +import io.r2dbc.spi.ConnectionMetadata; +import io.r2dbc.spi.Statement; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; + +public class ConnectionInstrumentation + implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice { + + @Override + public String hierarchyMarkerType() { + return "io.r2dbc.spi.Connection"; + } + + @Override + public ElementMatcher hierarchyMatcher() { + return implementsInterface(named("io.r2dbc.spi.Connection")); + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + isMethod() + .and(isPublic()) + .and(named("createStatement")) + .and(takesArguments(1)) + .and(takesArgument(0, String.class)), + ConnectionInstrumentation.class.getName() + "$CreateStatementAdvice"); + } + + public static class CreateStatementAdvice { + + @Advice.OnMethodEnter(suppress = Throwable.class) + public static R2dbcConnectionInfo onEnter( + @Advice.This final Connection connection, + @Advice.Argument(value = 0, readOnly = false) String sql) { + final String originalSql = sql; + String dbType = null; + try { + ConnectionMetadata metadata = connection.getMetadata(); + if (metadata != null) { + String productName = metadata.getDatabaseProductName(); + if (productName != null) { + dbType = productName.toLowerCase(); + } + } + } catch (Throwable ignored) { + // Connection may be closed or metadata unavailable + } + if (INJECT_COMMENT) { + sql = R2dbcSQLCommenter.inject(sql, null, dbType, null, null, null); + } + return R2dbcConnectionInfo.of(originalSql, dbType, null, null, null); + } + + @Advice.OnMethodExit(suppress = Throwable.class) + public static void onExit( + @Advice.Enter final R2dbcConnectionInfo info, @Advice.Return final Statement statement) { + if (statement != null && info != null) { + InstrumentationContext.get(Statement.class, R2dbcConnectionInfo.class).put(statement, info); + } + } + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcConnectionInfo.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcConnectionInfo.java new file mode 100644 index 00000000000..02825ea483f --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcConnectionInfo.java @@ -0,0 +1,53 @@ +package datadog.trace.instrumentation.r2dbc; + +/** + * Holds connection metadata and the original SQL string associated with an R2DBC Statement. Used + * for span tagging and DBM SQL comment injection. + */ +public final class R2dbcConnectionInfo { + private final String sql; + private final String dbType; + private final String dbInstance; + private final String dbUser; + private final String dbHostname; + + private R2dbcConnectionInfo( + String sql, String dbType, String dbInstance, String dbUser, String dbHostname) { + this.sql = sql; + this.dbType = dbType; + this.dbInstance = dbInstance; + this.dbUser = dbUser; + this.dbHostname = dbHostname; + } + + public String getSql() { + return sql; + } + + public String getDbType() { + return dbType; + } + + public String getDbInstance() { + return dbInstance; + } + + public String getDbUser() { + return dbUser; + } + + public String getDbHostname() { + return dbHostname; + } + + /** Creates an info object with only the SQL string (no connection metadata). */ + public static R2dbcConnectionInfo ofSql(String sql) { + return new R2dbcConnectionInfo(sql, null, null, null, null); + } + + /** Creates an info object with SQL string and connection metadata. */ + public static R2dbcConnectionInfo of( + String sql, String dbType, String dbInstance, String dbUser, String dbHostname) { + return new R2dbcConnectionInfo(sql, dbType, dbInstance, dbUser, dbHostname); + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcDecorator.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcDecorator.java new file mode 100644 index 00000000000..ac3c5bf1214 --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcDecorator.java @@ -0,0 +1,129 @@ +package datadog.trace.instrumentation.r2dbc; + +import datadog.trace.api.Config; +import datadog.trace.api.naming.SpanNaming; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes; +import datadog.trace.bootstrap.instrumentation.api.Tags; +import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString; +import datadog.trace.bootstrap.instrumentation.decorator.DBTypeProcessingDatabaseClientDecorator; + +public class R2dbcDecorator extends DBTypeProcessingDatabaseClientDecorator { + public static final R2dbcDecorator DECORATE = new R2dbcDecorator(); + + private static final String R2DBC = "r2dbc"; + public static final CharSequence R2DBC_QUERY = + UTF8BytesString.create(SpanNaming.instance().namingSchema().database().operation(R2DBC)); + public static final CharSequence R2DBC_BATCH = UTF8BytesString.create("r2dbc.batch"); + private static final String SERVICE_NAME = + SpanNaming.instance().namingSchema().database().service(R2DBC); + private static final CharSequence COMPONENT_NAME = UTF8BytesString.create(R2DBC); + + public static final String DBM_PROPAGATION_MODE = Config.get().getDbmPropagationMode(); + + public static final boolean INJECT_COMMENT = + DBM_PROPAGATION_MODE.equals(Config.DBM_PROPAGATION_MODE_FULL) + || DBM_PROPAGATION_MODE.equals(Config.DBM_PROPAGATION_MODE_STATIC) + || DBM_PROPAGATION_MODE.equals(Config.DBM_PROPAGATION_MODE_DYNAMIC_SERVICE); + + public static final boolean INJECT_TRACE_CONTEXT = + DBM_PROPAGATION_MODE.equals(Config.DBM_PROPAGATION_MODE_FULL); + + @Override + protected String[] instrumentationNames() { + return new String[] {R2DBC}; + } + + @Override + protected String service() { + return SERVICE_NAME; + } + + @Override + protected CharSequence component() { + return COMPONENT_NAME; + } + + @Override + protected CharSequence spanType() { + return InternalSpanTypes.SQL; + } + + @Override + protected String dbType() { + return R2DBC; + } + + @Override + protected String dbUser(Void connection) { + return null; + } + + @Override + protected String dbInstance(Void connection) { + return null; + } + + @Override + protected CharSequence dbHostname(Void connection) { + return null; + } + + /** + * Apply connection metadata tags to the span for peer service computation. Overrides db.type with + * the actual database product name when available, and sets peer.hostname and db.instance which + * feed into PeerServiceCalculator. + */ + public void onConnection(AgentSpan span, R2dbcConnectionInfo info) { + if (info != null) { + if (info.getDbType() != null) { + processDatabaseType(span, info.getDbType()); + } + if (info.getDbInstance() != null) { + onInstance(span, info.getDbInstance()); + } + if (info.getDbUser() != null) { + span.setTag(Tags.DB_USER, info.getDbUser()); + } + if (info.getDbHostname() != null) { + span.setTag(Tags.PEER_HOSTNAME, info.getDbHostname()); + } + } + } + + /** Extracts the first SQL keyword (SELECT, INSERT, UPDATE, DELETE, etc.) from a SQL string. */ + public static String extractDbOperation(String sql) { + if (sql == null || sql.isEmpty()) { + return null; + } + // Skip leading whitespace and SQL comments (e.g., /* DBM comment */) + int start = 0; + int len = sql.length(); + while (start < len) { + // Skip whitespace + if (Character.isWhitespace(sql.charAt(start))) { + start++; + continue; + } + // Skip block comments /* ... */ + if (start + 1 < len && sql.charAt(start) == '/' && sql.charAt(start + 1) == '*') { + int endComment = sql.indexOf("*/", start + 2); + if (endComment == -1) { + return null; + } + start = endComment + 2; + continue; + } + break; + } + // Find the end of the first word + int end = start; + while (end < len && !Character.isWhitespace(sql.charAt(end))) { + end++; + } + if (start == end) { + return null; + } + return sql.substring(start, end).toUpperCase(); + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcInstrumenterModule.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcInstrumenterModule.java new file mode 100644 index 00000000000..ab9110fc0e7 --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcInstrumenterModule.java @@ -0,0 +1,44 @@ +package datadog.trace.instrumentation.r2dbc; + +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +@AutoService(InstrumenterModule.class) +public class R2dbcInstrumenterModule extends InstrumenterModule.Tracing { + + public R2dbcInstrumenterModule() { + super("r2dbc"); + } + + @Override + public String[] helperClassNames() { + return new String[] { + packageName + ".R2dbcConnectionInfo", + packageName + ".R2dbcDecorator", + packageName + ".R2dbcSQLCommenter", + packageName + ".TracingPublisher", + packageName + ".TracingPublisher$TracingSubscriber", + }; + } + + @Override + public Map contextStore() { + Map contextStore = new HashMap<>(); + contextStore.put("io.r2dbc.spi.Statement", packageName + ".R2dbcConnectionInfo"); + return contextStore; + } + + @Override + public List typeInstrumentations() { + List instrumenters = new ArrayList<>(3); + instrumenters.add(new ConnectionInstrumentation()); + instrumenters.add(new StatementInstrumentation()); + instrumenters.add(new BatchInstrumentation()); + return instrumenters; + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcSQLCommenter.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcSQLCommenter.java new file mode 100644 index 00000000000..c6c5e2d3123 --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcSQLCommenter.java @@ -0,0 +1,65 @@ +package datadog.trace.instrumentation.r2dbc; + +import datadog.trace.bootstrap.instrumentation.dbm.SharedDBCommenter; + +/** + * Injects DBM trace context as SQL comments for R2DBC queries. Delegates comment content building + * to {@link SharedDBCommenter} and wraps the result in SQL block comment delimiters. + */ +public final class R2dbcSQLCommenter { + private static final String OPEN_COMMENT = "/*"; + private static final String CLOSE_COMMENT = "*/"; + + private R2dbcSQLCommenter() {} + + /** + * Injects a DBM SQL comment into the given query. By default the comment is prepended. + * + * @param sql the original SQL string + * @param dbService the database service name + * @param dbType the database type (e.g. "postgresql", "mysql") + * @param hostname the database hostname + * @param dbName the database name + * @param traceParent the W3C traceparent string (null for service-only mode) + * @return the SQL with injected comment, or the original if injection is not applicable + */ + public static String inject( + String sql, + String dbService, + String dbType, + String hostname, + String dbName, + String traceParent) { + if (sql == null || sql.isEmpty()) { + return sql; + } + // Check if a DD comment is already present + if (hasDDComment(sql)) { + return sql; + } + String commentContent = + SharedDBCommenter.buildComment(dbService, dbType, hostname, dbName, traceParent); + if (commentContent == null) { + return sql; + } + // Prepend the comment to the SQL + StringBuilder sb = new StringBuilder(sql.length() + commentContent.length() + 5); + sb.append(OPEN_COMMENT); + sb.append(commentContent); + sb.append(CLOSE_COMMENT); + sb.append(' '); + sb.append(sql); + return sb.toString(); + } + + private static boolean hasDDComment(String sql) { + if (!sql.startsWith(OPEN_COMMENT)) { + return false; + } + int endIdx = sql.indexOf(CLOSE_COMMENT); + if (endIdx > 2) { + return SharedDBCommenter.containsTraceComment(sql, 2, endIdx); + } + return false; + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/StatementInstrumentation.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/StatementInstrumentation.java new file mode 100644 index 00000000000..b1d0359060e --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/StatementInstrumentation.java @@ -0,0 +1,96 @@ +package datadog.trace.instrumentation.r2dbc; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static datadog.trace.instrumentation.r2dbc.R2dbcDecorator.DECORATE; +import static datadog.trace.instrumentation.r2dbc.R2dbcDecorator.INJECT_COMMENT; +import static datadog.trace.instrumentation.r2dbc.R2dbcDecorator.INJECT_TRACE_CONTEXT; +import static datadog.trace.instrumentation.r2dbc.R2dbcDecorator.R2DBC_QUERY; +import static net.bytebuddy.matcher.ElementMatchers.isMethod; +import static net.bytebuddy.matcher.ElementMatchers.isPublic; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; + +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.bootstrap.InstrumentationContext; +import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags; +import datadog.trace.bootstrap.instrumentation.api.Tags; +import io.r2dbc.spi.Statement; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; +import org.reactivestreams.Publisher; + +public class StatementInstrumentation + implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice { + + @Override + public String hierarchyMarkerType() { + return "io.r2dbc.spi.Statement"; + } + + @Override + public ElementMatcher hierarchyMatcher() { + return implementsInterface(named("io.r2dbc.spi.Statement")); + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + isMethod().and(isPublic()).and(named("execute")).and(takesArguments(0)), + StatementInstrumentation.class.getName() + "$StatementExecuteAdvice"); + } + + public static class StatementExecuteAdvice { + + @Advice.OnMethodEnter(suppress = Throwable.class) + public static AgentScope onEnter(@Advice.This final Statement statement) { + R2dbcConnectionInfo info = + InstrumentationContext.get(Statement.class, R2dbcConnectionInfo.class).get(statement); + + String sql = info != null ? info.getSql() : null; + + AgentSpan span = startSpan("r2dbc", R2DBC_QUERY); + DECORATE.afterStart(span); + + if (sql != null) { + DECORATE.onStatement(span, sql); + String dbOperation = R2dbcDecorator.extractDbOperation(sql); + if (dbOperation != null) { + span.setTag(Tags.DB_OPERATION, dbOperation); + } + } + + if (info != null) { + DECORATE.onConnection(span, info); + } + + if (INJECT_COMMENT && INJECT_TRACE_CONTEXT && info != null) { + span.setTag(InstrumentationTags.DBM_TRACE_INJECTED, true); + } + + return activateSpan(span); + } + + @SuppressWarnings("unchecked") + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) + public static void onExit( + @Advice.Enter final AgentScope scope, + @Advice.Return(readOnly = false) Publisher publisher, + @Advice.Thrown final Throwable throwable) { + AgentSpan span = scope.span(); + if (throwable != null) { + DECORATE.onError(span, throwable); + DECORATE.beforeFinish(span); + scope.close(); + span.finish(); + } else { + publisher = new TracingPublisher<>(publisher, span); + scope.close(); + } + } + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/TracingPublisher.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/TracingPublisher.java new file mode 100644 index 00000000000..1ead1440f1c --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/TracingPublisher.java @@ -0,0 +1,70 @@ +package datadog.trace.instrumentation.r2dbc; + +import static datadog.trace.instrumentation.r2dbc.R2dbcDecorator.DECORATE; + +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import org.reactivestreams.Publisher; +import org.reactivestreams.Subscriber; +import org.reactivestreams.Subscription; + +/** + * A wrapping Publisher that finishes the associated span when the downstream subscriber receives + * onComplete or onError. + */ +public final class TracingPublisher implements Publisher { + + private final Publisher delegate; + private final AgentSpan span; + + public TracingPublisher(Publisher delegate, AgentSpan span) { + this.delegate = delegate; + this.span = span; + } + + @Override + public void subscribe(Subscriber subscriber) { + delegate.subscribe(new TracingSubscriber<>(subscriber, span)); + } + + static final class TracingSubscriber implements Subscriber { + + private final Subscriber delegate; + private final AgentSpan span; + + TracingSubscriber(Subscriber delegate, AgentSpan span) { + this.delegate = delegate; + this.span = span; + } + + @Override + public void onSubscribe(Subscription s) { + delegate.onSubscribe(s); + } + + @Override + public void onNext(T t) { + delegate.onNext(t); + } + + @Override + public void onError(Throwable t) { + try { + DECORATE.onError(span, t); + DECORATE.beforeFinish(span); + span.finish(); + } finally { + delegate.onError(t); + } + } + + @Override + public void onComplete() { + try { + DECORATE.beforeFinish(span); + span.finish(); + } finally { + delegate.onComplete(); + } + } + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/datadog/trace/instrumentation/r2dbc/R2dbcDBMForkedTest.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/datadog/trace/instrumentation/r2dbc/R2dbcDBMForkedTest.java new file mode 100644 index 00000000000..56ecebd2791 --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/datadog/trace/instrumentation/r2dbc/R2dbcDBMForkedTest.java @@ -0,0 +1,253 @@ +package datadog.trace.instrumentation.r2dbc; + +import static datadog.trace.agent.test.assertions.SpanMatcher.span; +import static datadog.trace.agent.test.assertions.TagsMatcher.defaultTags; +import static datadog.trace.agent.test.assertions.TagsMatcher.error; +import static datadog.trace.agent.test.assertions.TagsMatcher.tag; +import static datadog.trace.agent.test.assertions.TraceMatcher.trace; +import static datadog.trace.agent.test.utils.TraceUtils.runnableUnderTrace; +import static datadog.trace.bootstrap.instrumentation.api.Tags.COMPONENT; +import static datadog.trace.bootstrap.instrumentation.api.Tags.DB_OPERATION; +import static datadog.trace.bootstrap.instrumentation.api.Tags.DB_TYPE; +import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND; +import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND_CLIENT; +import static datadog.trace.test.junit.utils.assertions.Matchers.is; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.trace.agent.test.AbstractInstrumentationTest; +import datadog.trace.api.DDSpanTypes; +import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags; +import io.r2dbc.spi.Result; +import io.r2dbc.spi.Statement; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.Test; +import org.reactivestreams.Publisher; +import org.reactivestreams.Subscriber; +import org.reactivestreams.Subscription; +import r2dbc.test.TestConnection; +import r2dbc.test.TestStatement; + +/** + * Forked test that runs with DBM propagation mode set to "full". This enables SQL comment injection + * and the _dd.dbm_trace_injected tag on database spans. + * + *

Uses ForkedTest suffix so the test runs in a separate JVM where the DBM config is applied + * before R2dbcDecorator's static fields are initialized. + */ +@datadog.trace.test.junit.utils.config.WithConfig(key = "dbm.propagation.mode", value = "full") +class R2dbcDBMForkedTest extends AbstractInstrumentationTest { + + @Test + void statementExecuteInjectsDBMCommentAndSetsTraceTags() { + TestConnection connection = new TestConnection(); + Statement statement = connection.createStatement("SELECT * FROM users"); + + // Verify SQL comment was injected into the statement received by the driver + String injectedSql = ((TestStatement) statement).getSql(); + assertNotNull(injectedSql, "SQL should not be null"); + assertTrue( + injectedSql.startsWith("/*"), "SQL should start with DBM comment, got: " + injectedSql); + assertTrue( + injectedSql.contains("ddps="), + "DBM comment should contain parent service tag, got: " + injectedSql); + assertTrue( + injectedSql.endsWith("SELECT * FROM users"), + "SQL should end with original query, got: " + injectedSql); + + runnableUnderTrace( + "parent", + () -> { + Publisher publisher = statement.execute(); + blockOnPublisher(publisher); + }); + + // Verify the span has the original SQL as resource (not the DBM-commented version) + // and includes the _dd.dbm_trace_injected tag + assertTraces( + trace( + span().root().operationName("parent"), + span() + .childOfPrevious() + .operationName("r2dbc.query") + .resourceName("SELECT * FROM users") + .type(DDSpanTypes.SQL) + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("testdb")), + tag(DB_OPERATION, is("SELECT")), + tag(InstrumentationTags.DBM_TRACE_INJECTED, is(true))))); + } + + @Test + void statementExecuteWithInsertInjectsDBMComment() { + TestConnection connection = new TestConnection(); + Statement statement = connection.createStatement("INSERT INTO users (name) VALUES ('Alice')"); + + // Verify SQL comment injection + String injectedSql = ((TestStatement) statement).getSql(); + assertTrue( + injectedSql.startsWith("/*"), "SQL should start with DBM comment, got: " + injectedSql); + assertTrue( + injectedSql.endsWith("INSERT INTO users (name) VALUES ('Alice')"), + "SQL should end with original query, got: " + injectedSql); + + runnableUnderTrace( + "parent", + () -> { + Publisher publisher = statement.execute(); + blockOnPublisher(publisher); + }); + + assertTraces( + trace( + span().root().operationName("parent"), + span() + .childOfPrevious() + .operationName("r2dbc.query") + .resourceName("INSERT INTO users (name) VALUES ('Alice')") + .type(DDSpanTypes.SQL) + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("testdb")), + tag(DB_OPERATION, is("INSERT")), + tag(InstrumentationTags.DBM_TRACE_INJECTED, is(true))))); + } + + @Test + void statementExecuteWithErrorStillHasDBMTags() { + TestConnection connection = new TestConnection(); + Statement statement = connection.createStatement("SELECT * FROM nonexistent"); + ((TestStatement) statement).setFailOnExecute(true); + + try { + runnableUnderTrace( + "parent", + () -> { + Publisher publisher = statement.execute(); + blockOnPublisher(publisher); + }); + } catch (RuntimeException ignored) { + // expected + } + + assertTraces( + trace( + span().root().operationName("parent").error(), + span() + .childOfPrevious() + .operationName("r2dbc.query") + .resourceName("SELECT * FROM nonexistent") + .type(DDSpanTypes.SQL) + .error() + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("testdb")), + tag(DB_OPERATION, is("SELECT")), + tag(InstrumentationTags.DBM_TRACE_INJECTED, is(true)), + error(RuntimeException.class, "query failed"), + tag("error.type", is(RuntimeException.class.getName())), + tag("error.message", is("query failed"))))); + } + + @Test + void dbmCommentIsNotInjectedTwice() { + TestConnection connection = new TestConnection(); + Statement statement = connection.createStatement("SELECT 1"); + + // Get the injected SQL + String injectedSql = ((TestStatement) statement).getSql(); + assertTrue(injectedSql.startsWith("/*"), "Should have DBM comment"); + + // Count how many comment blocks there are + int commentCount = 0; + int idx = 0; + while ((idx = injectedSql.indexOf("/*", idx)) != -1) { + commentCount++; + idx += 2; + } + assertTrue( + commentCount == 1, + "Should have exactly one comment block, found " + commentCount + " in: " + injectedSql); + + runnableUnderTrace( + "parent", + () -> { + Publisher publisher = statement.execute(); + blockOnPublisher(publisher); + }); + + // Span resource should use the original SQL, not the commented one + assertTraces( + trace( + span().root().operationName("parent"), + span() + .childOfPrevious() + .operationName("r2dbc.query") + .resourceName("SELECT 1") + .type(DDSpanTypes.SQL) + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("testdb")), + tag(DB_OPERATION, is("SELECT")), + tag(InstrumentationTags.DBM_TRACE_INJECTED, is(true))))); + } + + /** Subscribes to a Publisher and blocks until completion or error. */ + private static void blockOnPublisher(Publisher publisher) { + List results = new ArrayList<>(); + RuntimeException[] error = new RuntimeException[1]; + CountDownLatch latch = new CountDownLatch(1); + publisher.subscribe( + new Subscriber() { + @Override + public void onSubscribe(Subscription s) { + s.request(Long.MAX_VALUE); + } + + @Override + public void onNext(T t) { + results.add(t); + } + + @Override + public void onError(Throwable t) { + if (t instanceof RuntimeException) { + error[0] = (RuntimeException) t; + } else { + error[0] = new RuntimeException(t); + } + latch.countDown(); + } + + @Override + public void onComplete() { + latch.countDown(); + } + }); + try { + latch.await(5, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Interrupted waiting for publisher", e); + } + if (error[0] != null) { + throw error[0]; + } + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/datadog/trace/instrumentation/r2dbc/R2dbcInstrumentationTest.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/datadog/trace/instrumentation/r2dbc/R2dbcInstrumentationTest.java new file mode 100644 index 00000000000..186e199ce48 --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/datadog/trace/instrumentation/r2dbc/R2dbcInstrumentationTest.java @@ -0,0 +1,314 @@ +package datadog.trace.instrumentation.r2dbc; + +import static datadog.trace.agent.test.assertions.SpanMatcher.span; +import static datadog.trace.agent.test.assertions.TagsMatcher.defaultTags; +import static datadog.trace.agent.test.assertions.TagsMatcher.error; +import static datadog.trace.agent.test.assertions.TagsMatcher.tag; +import static datadog.trace.agent.test.assertions.TraceMatcher.trace; +import static datadog.trace.agent.test.utils.TraceUtils.runnableUnderTrace; +import static datadog.trace.bootstrap.instrumentation.api.Tags.COMPONENT; +import static datadog.trace.bootstrap.instrumentation.api.Tags.DB_OPERATION; +import static datadog.trace.bootstrap.instrumentation.api.Tags.DB_TYPE; +import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND; +import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND_CLIENT; +import static datadog.trace.test.junit.utils.assertions.Matchers.is; +import static org.junit.jupiter.api.Assertions.assertFalse; + +import datadog.trace.agent.test.AbstractInstrumentationTest; +import datadog.trace.agent.test.assertions.TraceMatcher; +import datadog.trace.api.DDSpanTypes; +import io.r2dbc.spi.Batch; +import io.r2dbc.spi.Result; +import io.r2dbc.spi.Statement; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.Test; +import org.reactivestreams.Publisher; +import org.reactivestreams.Subscriber; +import org.reactivestreams.Subscription; +import r2dbc.test.TestBatch; +import r2dbc.test.TestConnection; +import r2dbc.test.TestStatement; + +class R2dbcInstrumentationTest extends AbstractInstrumentationTest { + + @Test + void statementExecuteCreatesSpan() { + TestConnection connection = new TestConnection(); + Statement statement = connection.createStatement("SELECT * FROM users"); + + List[] results = new List[1]; + runnableUnderTrace( + "parent", + () -> { + Publisher publisher = statement.execute(); + results[0] = blockOnPublisher(publisher); + }); + assertFalse(results[0].isEmpty(), "Statement execute should return results"); + + assertTraces( + trace( + span().root().operationName("parent"), + span() + .childOfPrevious() + .operationName("r2dbc.query") + .resourceName("SELECT * FROM users") + .type(DDSpanTypes.SQL) + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("testdb")), + tag(DB_OPERATION, is("SELECT"))))); + } + + @Test + void statementExecuteWithInsert() { + TestConnection connection = new TestConnection(); + Statement statement = connection.createStatement("INSERT INTO users (name) VALUES ($1)"); + + List[] results = new List[1]; + runnableUnderTrace( + "parent", + () -> { + Publisher publisher = statement.execute(); + results[0] = blockOnPublisher(publisher); + }); + assertFalse(results[0].isEmpty(), "Statement execute should return results"); + + assertTraces( + trace( + span().root().operationName("parent"), + span() + .childOfPrevious() + .operationName("r2dbc.query") + .resourceName("INSERT INTO users (name) VALUES ($1)") + .type(DDSpanTypes.SQL) + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("testdb")), + tag(DB_OPERATION, is("INSERT"))))); + } + + @Test + void statementExecuteWithError() { + TestConnection connection = new TestConnection(); + Statement statement = connection.createStatement("SELECT * FROM nonexistent"); + ((TestStatement) statement).setFailOnExecute(true); + + try { + runnableUnderTrace( + "parent", + () -> { + Publisher publisher = statement.execute(); + blockOnPublisher(publisher); + }); + } catch (RuntimeException ignored) { + // expected + } + + assertTraces( + trace( + span().root().operationName("parent").error(), + span() + .childOfPrevious() + .operationName("r2dbc.query") + .resourceName("SELECT * FROM nonexistent") + .type(DDSpanTypes.SQL) + .error() + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("testdb")), + tag(DB_OPERATION, is("SELECT")), + error(RuntimeException.class, "query failed")))); + } + + @Test + void batchExecuteCreatesSpan() { + TestConnection connection = new TestConnection(); + Batch batch = connection.createBatch(); + batch.add("INSERT INTO users (name) VALUES ('Alice')"); + batch.add("INSERT INTO users (name) VALUES ('Bob')"); + + List[] results = new List[1]; + runnableUnderTrace( + "parent", + () -> { + Publisher publisher = batch.execute(); + results[0] = blockOnPublisher(publisher); + }); + assertFalse(results[0].isEmpty(), "Batch execute should return results"); + + assertTraces( + trace( + span().root().operationName("parent"), + span() + .childOfPrevious() + .operationName("r2dbc.batch") + .resourceName("r2dbc.batch") + .type(DDSpanTypes.SQL) + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("r2dbc"))))); + } + + @Test + void batchExecuteWithError() { + TestConnection connection = new TestConnection(); + TestBatch batch = (TestBatch) connection.createBatch(); + batch.add("INSERT INTO users (name) VALUES ('Alice')"); + batch.setFailOnExecute(true); + + try { + runnableUnderTrace( + "parent", + () -> { + Publisher publisher = batch.execute(); + blockOnPublisher(publisher); + }); + } catch (RuntimeException ignored) { + // expected + } + + assertTraces( + trace( + span().root().operationName("parent").error(), + span() + .childOfPrevious() + .operationName("r2dbc.batch") + .resourceName("r2dbc.batch") + .type(DDSpanTypes.SQL) + .error() + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("r2dbc")), + error(RuntimeException.class, "batch failed")))); + } + + @Test + void multipleStatementExecutesCreateSeparateSpans() { + TestConnection connection = new TestConnection(); + Statement statement1 = connection.createStatement("SELECT * FROM users"); + Statement statement2 = connection.createStatement("UPDATE users SET name = 'test'"); + + runnableUnderTrace( + "parent", + () -> { + blockOnPublisher(statement1.execute()); + blockOnPublisher(statement2.execute()); + }); + + assertTraces( + trace( + TraceMatcher.SORT_BY_ANCESTRY, + span().root().operationName("parent"), + span() + .childOfPrevious() + .operationName("r2dbc.query") + .resourceName("SELECT * FROM users") + .type(DDSpanTypes.SQL) + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("testdb")), + tag(DB_OPERATION, is("SELECT"))), + span() + .childOfIndex(0) + .operationName("r2dbc.query") + .resourceName("UPDATE users SET name = 'test'") + .type(DDSpanTypes.SQL) + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("testdb")), + tag(DB_OPERATION, is("UPDATE"))))); + } + + @Test + void statementExecuteWithNoParentSpan() { + TestConnection connection = new TestConnection(); + Statement statement = connection.createStatement("SELECT 1"); + + Publisher publisher = statement.execute(); + List results = blockOnPublisher(publisher); + assertFalse(results.isEmpty(), "Statement execute should return results"); + + assertTraces( + trace( + span() + .root() + .operationName("r2dbc.query") + .resourceName("SELECT 1") + .type(DDSpanTypes.SQL) + .measured() + .tags( + defaultTags(), + tag(SPAN_KIND, is(SPAN_KIND_CLIENT)), + tag(COMPONENT, is("r2dbc")), + tag(DB_TYPE, is("testdb")), + tag(DB_OPERATION, is("SELECT"))))); + } + + /** Subscribes to a Publisher and blocks until completion or error. Returns collected results. */ + private static List blockOnPublisher(Publisher publisher) { + List results = new ArrayList<>(); + RuntimeException[] error = new RuntimeException[1]; + CountDownLatch latch = new CountDownLatch(1); + publisher.subscribe( + new Subscriber() { + @Override + public void onSubscribe(Subscription s) { + s.request(Long.MAX_VALUE); + } + + @Override + public void onNext(T t) { + results.add(t); + } + + @Override + public void onError(Throwable t) { + if (t instanceof RuntimeException) { + error[0] = (RuntimeException) t; + } else { + error[0] = new RuntimeException(t); + } + latch.countDown(); + } + + @Override + public void onComplete() { + latch.countDown(); + } + }); + try { + latch.await(5, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Interrupted waiting for publisher", e); + } + if (error[0] != null) { + throw error[0]; + } + return results; + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestBatch.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestBatch.java new file mode 100644 index 00000000000..f62e0bf3f1c --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestBatch.java @@ -0,0 +1,54 @@ +package r2dbc.test; + +import io.r2dbc.spi.Batch; +import io.r2dbc.spi.Result; +import java.util.ArrayList; +import java.util.List; +import org.reactivestreams.Publisher; +import org.reactivestreams.Subscription; + +public class TestBatch implements Batch { + private final List statements = new ArrayList<>(); + private boolean failOnExecute; + + public void setFailOnExecute(boolean fail) { + this.failOnExecute = fail; + } + + @Override + public Batch add(String sql) { + statements.add(sql); + return this; + } + + @Override + public Publisher execute() { + if (failOnExecute) { + return subscriber -> { + subscriber.onSubscribe( + new Subscription() { + @Override + public void request(long n) { + subscriber.onError(new RuntimeException("batch failed")); + } + + @Override + public void cancel() {} + }); + }; + } + return subscriber -> { + subscriber.onSubscribe( + new Subscription() { + @Override + public void request(long n) { + subscriber.onNext(new TestResult()); + subscriber.onComplete(); + } + + @Override + public void cancel() {} + }); + }; + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestConnection.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestConnection.java new file mode 100644 index 00000000000..ae9ba2d510b --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestConnection.java @@ -0,0 +1,141 @@ +package r2dbc.test; + +import io.r2dbc.spi.Batch; +import io.r2dbc.spi.Connection; +import io.r2dbc.spi.ConnectionMetadata; +import io.r2dbc.spi.IsolationLevel; +import io.r2dbc.spi.Statement; +import io.r2dbc.spi.TransactionDefinition; +import io.r2dbc.spi.ValidationDepth; +import java.time.Duration; +import org.reactivestreams.Publisher; +import org.reactivestreams.Subscription; + +public class TestConnection implements Connection { + @Override + public Publisher beginTransaction() { + return emptyPublisher(); + } + + @Override + public Publisher beginTransaction(TransactionDefinition definition) { + return emptyPublisher(); + } + + @Override + public Publisher close() { + return emptyPublisher(); + } + + @Override + public Publisher commitTransaction() { + return emptyPublisher(); + } + + @Override + public Batch createBatch() { + return new TestBatch(); + } + + @Override + public Publisher createSavepoint(String name) { + return emptyPublisher(); + } + + @Override + public Statement createStatement(String sql) { + return new TestStatement(sql); + } + + @Override + public boolean isAutoCommit() { + return true; + } + + @Override + public ConnectionMetadata getMetadata() { + return new ConnectionMetadata() { + @Override + public String getDatabaseProductName() { + return "TestDB"; + } + + @Override + public String getDatabaseVersion() { + return "1.0"; + } + }; + } + + @Override + public IsolationLevel getTransactionIsolationLevel() { + return IsolationLevel.READ_COMMITTED; + } + + @Override + public Publisher releaseSavepoint(String name) { + return emptyPublisher(); + } + + @Override + public Publisher rollbackTransaction() { + return emptyPublisher(); + } + + @Override + public Publisher rollbackTransactionToSavepoint(String name) { + return emptyPublisher(); + } + + @Override + public Publisher setAutoCommit(boolean autoCommit) { + return emptyPublisher(); + } + + @Override + public Publisher setLockWaitTimeout(Duration timeout) { + return emptyPublisher(); + } + + @Override + public Publisher setStatementTimeout(Duration timeout) { + return emptyPublisher(); + } + + @Override + public Publisher setTransactionIsolationLevel(IsolationLevel isolationLevel) { + return emptyPublisher(); + } + + @Override + public Publisher validate(ValidationDepth depth) { + return subscriber -> { + subscriber.onSubscribe( + new Subscription() { + @Override + public void request(long n) { + subscriber.onNext(true); + subscriber.onComplete(); + } + + @Override + public void cancel() {} + }); + }; + } + + private static Publisher emptyPublisher() { + return subscriber -> { + subscriber.onSubscribe( + new Subscription() { + @Override + public void request(long n) { + subscriber.onComplete(); + } + + @Override + public void cancel() {} + }); + }; + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestConnectionFactory.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestConnectionFactory.java new file mode 100644 index 00000000000..0ff3815d8ee --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestConnectionFactory.java @@ -0,0 +1,31 @@ +package r2dbc.test; + +import io.r2dbc.spi.Connection; +import io.r2dbc.spi.ConnectionFactory; +import io.r2dbc.spi.ConnectionFactoryMetadata; +import org.reactivestreams.Publisher; +import org.reactivestreams.Subscription; + +public class TestConnectionFactory implements ConnectionFactory { + @Override + public Publisher create() { + return subscriber -> { + subscriber.onSubscribe( + new Subscription() { + @Override + public void request(long n) { + subscriber.onNext(new TestConnection()); + subscriber.onComplete(); + } + + @Override + public void cancel() {} + }); + }; + } + + @Override + public ConnectionFactoryMetadata getMetadata() { + return () -> "TestDB"; + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestResult.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestResult.java new file mode 100644 index 00000000000..531b5489368 --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestResult.java @@ -0,0 +1,67 @@ +package r2dbc.test; + +import io.r2dbc.spi.Result; +import io.r2dbc.spi.Row; +import io.r2dbc.spi.RowMetadata; +import java.util.function.BiFunction; +import java.util.function.Function; +import java.util.function.Predicate; +import org.reactivestreams.Publisher; +import org.reactivestreams.Subscription; + +public class TestResult implements Result { + @Override + public Publisher getRowsUpdated() { + return subscriber -> { + subscriber.onSubscribe( + new Subscription() { + @Override + public void request(long n) { + subscriber.onNext(0L); + subscriber.onComplete(); + } + + @Override + public void cancel() {} + }); + }; + } + + @Override + public Publisher map(BiFunction mappingFunction) { + return subscriber -> { + subscriber.onSubscribe( + new Subscription() { + @Override + public void request(long n) { + subscriber.onComplete(); + } + + @Override + public void cancel() {} + }); + }; + } + + @Override + public Result filter(Predicate filter) { + return this; + } + + @Override + public Publisher flatMap( + Function> mappingFunction) { + return subscriber -> { + subscriber.onSubscribe( + new Subscription() { + @Override + public void request(long n) { + subscriber.onComplete(); + } + + @Override + public void cancel() {} + }); + }; + } +} diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestStatement.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestStatement.java new file mode 100644 index 00000000000..1ba5136d330 --- /dev/null +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/test/java/r2dbc/test/TestStatement.java @@ -0,0 +1,79 @@ +package r2dbc.test; + +import io.r2dbc.spi.Result; +import io.r2dbc.spi.Statement; +import org.reactivestreams.Publisher; +import org.reactivestreams.Subscription; + +public class TestStatement implements Statement { + private final String sql; + private boolean failOnExecute; + + public TestStatement(String sql) { + this.sql = sql; + } + + public String getSql() { + return sql; + } + + public void setFailOnExecute(boolean fail) { + this.failOnExecute = fail; + } + + @Override + public Statement add() { + return this; + } + + @Override + public Statement bind(int index, Object value) { + return this; + } + + @Override + public Statement bind(String name, Object value) { + return this; + } + + @Override + public Statement bindNull(int index, Class type) { + return this; + } + + @Override + public Statement bindNull(String name, Class type) { + return this; + } + + @Override + public Publisher execute() { + if (failOnExecute) { + return subscriber -> { + subscriber.onSubscribe( + new Subscription() { + @Override + public void request(long n) { + subscriber.onError(new RuntimeException("query failed")); + } + + @Override + public void cancel() {} + }); + }; + } + return subscriber -> { + subscriber.onSubscribe( + new Subscription() { + @Override + public void request(long n) { + subscriber.onNext(new TestResult()); + subscriber.onComplete(); + } + + @Override + public void cancel() {} + }); + }; + } +} diff --git a/metadata/supported-configurations.json b/metadata/supported-configurations.json index 850c6055d58..821faad6f93 100644 --- a/metadata/supported-configurations.json +++ b/metadata/supported-configurations.json @@ -9065,6 +9065,30 @@ "aliases": ["DD_TRACE_INTEGRATION_QUARTZ_ENABLED", "DD_INTEGRATION_QUARTZ_ENABLED"] } ], + "DD_TRACE_R2DBC_ANALYTICS_ENABLED": [ + { + "version": "A", + "type": "boolean", + "default": "false", + "aliases": ["DD_R2DBC_ANALYTICS_ENABLED"] + } + ], + "DD_TRACE_R2DBC_ANALYTICS_SAMPLE_RATE": [ + { + "version": "A", + "type": "decimal", + "default": "1.0", + "aliases": ["DD_R2DBC_ANALYTICS_SAMPLE_RATE"] + } + ], + "DD_TRACE_R2DBC_ENABLED": [ + { + "version": "A", + "type": "boolean", + "default": "true", + "aliases": ["DD_TRACE_INTEGRATION_R2DBC_ENABLED", "DD_INTEGRATION_R2DBC_ENABLED"] + } + ], "DD_TRACE_RABBITMQ_ANALYTICS_ENABLED": [ { "version": "A", diff --git a/settings.gradle.kts b/settings.gradle.kts index 86aa23cc9d4..4df106170de 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -537,6 +537,7 @@ include( ":dd-java-agent:instrumentation:play:play-appsec-common", ":dd-java-agent:instrumentation:protobuf-3.0", ":dd-java-agent:instrumentation:quartz-2.0", + ":dd-java-agent:instrumentation:r2dbc:r2dbc-1.0", ":dd-java-agent:instrumentation:rabbitmq-amqp-2.7", ":dd-java-agent:instrumentation:ratpack-1.5", ":dd-java-agent:instrumentation:reactive-streams-1.0", diff --git a/utils/test-junit-utils/src/main/java/datadog/trace/test/junit/utils/assertions/Is.java b/utils/test-junit-utils/src/main/java/datadog/trace/test/junit/utils/assertions/Is.java index f7a3346e64e..2deeb06f697 100644 --- a/utils/test-junit-utils/src/main/java/datadog/trace/test/junit/utils/assertions/Is.java +++ b/utils/test-junit-utils/src/main/java/datadog/trace/test/junit/utils/assertions/Is.java @@ -27,6 +27,9 @@ public String failureReason() { @Override public boolean test(T t) { + if (this.expected instanceof CharSequence && t instanceof CharSequence) { + return this.expected.toString().equals(t.toString()); + } return this.expected.equals(t); } } From 521b0d52a1d15b408099cdf371a44ee81ee373a5 Mon Sep 17 00:00:00 2001 From: Jordan Wong Date: Wed, 22 Jul 2026 17:06:42 -0400 Subject: [PATCH 2/4] fix(r2dbc): drop assertInverse from muzzle version range The advice hooks io.r2dbc.spi.Statement/Connection/Batch interfaces, which are structurally unchanged across the 0.8.x/0.9.x/1.0.x releases -- assertInverse was asserting that all versions outside [1.0.0.RELEASE,) must fail muzzle, which is false for the older-but- compatible SPI releases and caused muzzle-AssertFail-io.r2dbc-r2dbc-spi-{0.8.6,0.9.1}.RELEASE to fail CI ("unexpectedly passed Muzzle validation"). This does not touch the hook point -- R2DBC's core finding (missing db.name/peer.hostname/db.user/network.destination.port from hooking ConnectionMetadata instead of ConnectionFactoryOptions, see docs/eval-research/hypotheses/r2dbc.md on the toolkit repo) is unrelated and unaddressed here on purpose. Verified locally: `./gradlew :dd-java-agent:instrumentation:r2dbc:r2dbc-1.0:muzzle` BUILD SUCCESSFUL. Co-Authored-By: Claude Sonnet 5 --- dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/build.gradle | 1 - 1 file changed, 1 deletion(-) diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/build.gradle b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/build.gradle index 6a1deb41fcf..3f820b16c72 100644 --- a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/build.gradle +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/build.gradle @@ -3,7 +3,6 @@ muzzle { group = "io.r2dbc" module = "r2dbc-spi" versions = "[1.0.0.RELEASE,)" - assertInverse = true } } From 368d2094ce0a14b51139de59ab38840223ca9858 Mon Sep 17 00:00:00 2001 From: Jordan Wong Date: Wed, 22 Jul 2026 17:07:07 -0400 Subject: [PATCH 3/4] fix(r2dbc): add r2dbc to agent-jar-checks.properties golden file Integration list differs from metadata/agent-jar-checks.properties -- super("r2dbc") registered in R2dbcInstrumenterModule.java but the golden file was never updated. Inserted alphabetically between quartz and ratpack. Hand-edited rather than running the Gradle golden-file-update task directly, but verified via `./gradlew :dd-java-agent:verifyAgentJarIntegrations` BUILD SUCCESSFUL. Co-Authored-By: Claude Sonnet 5 --- metadata/agent-jar-checks.properties | 1 + 1 file changed, 1 insertion(+) diff --git a/metadata/agent-jar-checks.properties b/metadata/agent-jar-checks.properties index fd04d637e35..1b03884d0bb 100644 --- a/metadata/agent-jar-checks.properties +++ b/metadata/agent-jar-checks.properties @@ -144,6 +144,7 @@ expected.integrations = IastInstrumentation,\ play-ws,\ protobuf,\ quartz,\ + r2dbc,\ ratpack,\ ratpack-request-body,\ reactive-streams,\ From ada5972e50f844252c9c459dda511957aaad0b60 Mon Sep 17 00:00:00 2001 From: Jordan Wong Date: Thu, 30 Jul 2026 09:15:34 -0400 Subject: [PATCH 4/4] fix(r2dbc): finish span on subscription cancel to prevent span leak Codex review (PR #12032): when a returned Publisher is cancelled before onComplete/onError (take(1), timeout, disconnected request), TracingSubscriber passed the driver's Subscription straight through, so cancellation was never observed and the span created in execute() could stay open indefinitely. Wrap the Subscription (named TracingSubscription helper) so cancel() finishes the span, guarded by an AtomicBoolean shared with onComplete/onError so a cancel racing a terminal signal can't double-finish. Registered the new helper class in R2dbcInstrumenterModule.helperClassNames() (muzzle requires it). Verified: compileJava + muzzle + module test suite all green. Co-Authored-By: Claude Sonnet 5 --- .../r2dbc/R2dbcInstrumenterModule.java | 1 + .../r2dbc/TracingPublisher.java | 50 ++++++++++++++++--- 2 files changed, 44 insertions(+), 7 deletions(-) diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcInstrumenterModule.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcInstrumenterModule.java index ab9110fc0e7..e6620363456 100644 --- a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcInstrumenterModule.java +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/R2dbcInstrumenterModule.java @@ -23,6 +23,7 @@ public String[] helperClassNames() { packageName + ".R2dbcSQLCommenter", packageName + ".TracingPublisher", packageName + ".TracingPublisher$TracingSubscriber", + packageName + ".TracingPublisher$TracingSubscriber$TracingSubscription", }; } diff --git a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/TracingPublisher.java b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/TracingPublisher.java index 1ead1440f1c..99786becc39 100644 --- a/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/TracingPublisher.java +++ b/dd-java-agent/instrumentation/r2dbc/r2dbc-1.0/src/main/java/datadog/trace/instrumentation/r2dbc/TracingPublisher.java @@ -3,13 +3,14 @@ import static datadog.trace.instrumentation.r2dbc.R2dbcDecorator.DECORATE; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import java.util.concurrent.atomic.AtomicBoolean; import org.reactivestreams.Publisher; import org.reactivestreams.Subscriber; import org.reactivestreams.Subscription; /** * A wrapping Publisher that finishes the associated span when the downstream subscriber receives - * onComplete or onError. + * onComplete, onError, or cancels its subscription. */ public final class TracingPublisher implements Publisher { @@ -30,15 +31,27 @@ static final class TracingSubscriber implements Subscriber { private final Subscriber delegate; private final AgentSpan span; + // Guards against finishing the span more than once when terminal signals (onComplete/onError) + // race with a downstream cancel(). + private final AtomicBoolean finished = new AtomicBoolean(false); TracingSubscriber(Subscriber delegate, AgentSpan span) { this.delegate = delegate; this.span = span; } + private void finishSpan() { + if (finished.compareAndSet(false, true)) { + DECORATE.beforeFinish(span); + span.finish(); + } + } + @Override public void onSubscribe(Subscription s) { - delegate.onSubscribe(s); + // Wrap the subscription so a downstream cancel() (e.g. take(1), timeout, disconnect) finishes + // the span instead of leaving it open indefinitely. + delegate.onSubscribe(new TracingSubscription(s)); } @Override @@ -46,12 +59,36 @@ public void onNext(T t) { delegate.onNext(t); } + final class TracingSubscription implements Subscription { + private final Subscription delegate; + + TracingSubscription(Subscription delegate) { + this.delegate = delegate; + } + + @Override + public void request(long n) { + delegate.request(n); + } + + @Override + public void cancel() { + try { + delegate.cancel(); + } finally { + finishSpan(); + } + } + } + @Override public void onError(Throwable t) { try { - DECORATE.onError(span, t); - DECORATE.beforeFinish(span); - span.finish(); + if (finished.compareAndSet(false, true)) { + DECORATE.onError(span, t); + DECORATE.beforeFinish(span); + span.finish(); + } } finally { delegate.onError(t); } @@ -60,8 +97,7 @@ public void onError(Throwable t) { @Override public void onComplete() { try { - DECORATE.beforeFinish(span); - span.finish(); + finishSpan(); } finally { delegate.onComplete(); }