[reference] R2DBC net-new instrumentation (io.r2dbc:r2dbc-spi 1.0.0, toolkit output)

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) {

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

@@ -0,0 +1,19 @@
+muzzle {
+  pass {
+    group = "io.r2dbc"
+    module = "r2dbc-spi"
+    versions = "[1.0.0.RELEASE,)"
+  }
+}
+
+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.+'
+}

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<TypeDescription> 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();
+      }
+    }
+  }
+}

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<TypeDescription> 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);
+      }
+    }
+  }
+}

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);
+  }
+}

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<Void> {
+  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();
+  }
+}