MSSQL R2DBC Instrumentation

instrumentation/r2dbc-mssql/build.gradle

@@ -0,0 +1,18 @@
+dependencies {
+    implementation(project(":agent-bridge"))
+    implementation(project(":agent-bridge-datastore"))
+    implementation("io.r2dbc:r2dbc-mssql:0.9.0.RELEASE")
+}
+
+jar {
+    manifest { attributes 'Implementation-Title': 'com.newrelic.instrumentation.r2dbc-mssql' }
+}
+
+verifyInstrumentation {
+    passesOnly 'io.r2dbc:r2dbc-mssql:[0.8.0,)'
+}
+
+site {
+    title 'MSSQL R2DBC'
+    type 'Datastore'
+}

instrumentation/r2dbc-mssql/src/main/java/io/r2dbc/mssql/ParametrizedMssqlStatement_Instrumentation.java

@@ -0,0 +1,29 @@
+package io.r2dbc.mssql;
+
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.NewField;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+import io.r2dbc.mssql.client.Client;
+import io.r2dbc.mssql.client.R2dbcUtils;
+import reactor.core.publisher.Flux;
+
+@Weave(type = MatchType.ExactClass, originalName = "io.r2dbc.mssql.ParametrizedMssqlStatement")
+final class ParametrizedMssqlStatement_Instrumentation {
+    private final Client client = Weaver.callOriginal();
+
+    @NewField
+    private final String sql;
+
+    public Flux<MssqlResult> execute() {
+        Flux<MssqlResult> request = Weaver.callOriginal();
+        if (request != null && this.client != null) {
+            return R2dbcUtils.wrapRequest(request, sql, client);
+        }
+        return request;
+    }
+
+    ParametrizedMssqlStatement_Instrumentation(Client client, ConnectionOptions connectionOptions, String sql) {
+        this.sql = sql;
+    }
+}

instrumentation/r2dbc-mssql/src/main/java/io/r2dbc/mssql/SimpleMssqlStatement_Instrumentation.java

@@ -0,0 +1,23 @@
+package io.r2dbc.mssql;
+
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+import io.r2dbc.mssql.client.R2dbcUtils;
+import io.r2dbc.mssql.client.Client;
+import reactor.core.publisher.Flux;
+
+@Weave(type = MatchType.ExactClass, originalName = "io.r2dbc.mssql.SimpleMssqlStatement")
+final class SimpleMssqlStatement_Instrumentation {
+    private final Client client = Weaver.callOriginal();
+    private final String sql = Weaver.callOriginal();
+
+    public Flux<MssqlResult> execute() {
+        Flux<MssqlResult> request = Weaver.callOriginal();
+        if(request != null && this.sql != null && this.client != null) {
+            return R2dbcUtils.wrapRequest(request, this.sql, this.client);
+        }
+        return request;
+    }
+}
+

instrumentation/r2dbc-mssql/src/main/java/io/r2dbc/mssql/client/R2dbcUtils.java

@@ -0,0 +1,66 @@
+package io.r2dbc.mssql.client;
+
+import com.newrelic.agent.bridge.NoOpTransaction;
+import com.newrelic.agent.bridge.datastore.DatastoreVendor;
+import com.newrelic.agent.bridge.datastore.OperationAndTableName;
+import com.newrelic.agent.bridge.datastore.R2dbcObfuscator;
+import com.newrelic.agent.bridge.datastore.R2dbcOperation;
+import com.newrelic.api.agent.DatastoreParameters;
+import com.newrelic.api.agent.NewRelic;
+import com.newrelic.api.agent.Segment;
+import com.newrelic.api.agent.Transaction;
+import io.r2dbc.mssql.MssqlResult;
+import org.reactivestreams.Subscription;
+import reactor.core.publisher.Flux;
+import reactor.netty.Connection;
+
+import java.net.InetSocketAddress;
+
+import java.util.function.Consumer;
+
+public class R2dbcUtils {
+    public static Flux<MssqlResult> wrapRequest(Flux<MssqlResult> request, String sql, Client client) {
+        if(request != null) {
+            Transaction transaction = NewRelic.getAgent().getTransaction();
+            if(transaction != null && !(transaction instanceof NoOpTransaction)) {
+                Segment segment = transaction.startSegment("execute");
+                return request
+                        .doOnSubscribe(reportExecution(sql, client, segment))
+                        .doFinally((type) -> segment.end());
+            }
+        }
+        return request;
+    }
+
+    private static Consumer<Subscription> reportExecution(String sql, Client client, Segment segment) {
+        return (subscription) -> {
+            OperationAndTableName sqlOperation = R2dbcOperation.extractFrom(sql);
+            InetSocketAddress socketAddress = extractSocketAddress(client);
+            if (sqlOperation != null && socketAddress != null) {
+                segment.reportAsExternal(DatastoreParameters
+                        .product(DatastoreVendor.MSSQL.name())
+                        .collection(sqlOperation.getTableName())
+                        .operation(sqlOperation.getOperation())
+                        .instance(socketAddress.getHostName(), socketAddress.getPort())
+                        .databaseName(null)
+                        .slowQuery(sql, R2dbcObfuscator.QUERY_CONVERTER)
+                        .build());
+            }
+        };
+    }
+
+    public static InetSocketAddress extractSocketAddress(Client client) {
+        try {
+            if(client instanceof ReactorNettyClient_Instrumentation) {
+                ReactorNettyClient_Instrumentation instrumentedClient = (ReactorNettyClient_Instrumentation) client;
+                Connection clientConnection = instrumentedClient.clientConnection;
+                    if(clientConnection.channel().remoteAddress() != null && clientConnection.channel().remoteAddress() instanceof InetSocketAddress) {
+                        return (InetSocketAddress) clientConnection.channel().remoteAddress();
+                    }
+            }
+            return null;
+        } catch(Exception exception) {
+            return null;
+        }
+    }
+}

instrumentation/r2dbc-mssql/src/main/java/io/r2dbc/mssql/client/ReactorNettyClient_Instrumentation.java

@@ -0,0 +1,16 @@
+package io.r2dbc.mssql.client;
+
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.NewField;
+import com.newrelic.api.agent.weaver.Weave;
+import reactor.netty.Connection;
+
+@Weave(type = MatchType.ExactClass, originalName = "io.r2dbc.mssql.client.ReactorNettyClient")
+public class ReactorNettyClient_Instrumentation {
+    @NewField
+    public final Connection clientConnection;
+
+    private ReactorNettyClient_Instrumentation(Connection connection, TdsEncoder TdsEncoder, ConnectionContext context) {
+        this.clientConnection = connection;
+    }
+}

settings.gradle

@@ -246,6 +246,7 @@ include 'instrumentation:r2dbc-h2'
 include 'instrumentation:r2dbc-mariadb'
 include 'instrumentation:r2dbc-mysql'
 include 'instrumentation:r2dbc-postgresql'
+include 'instrumentation:r2dbc-mssql'
 include 'instrumentation:rabbit-amqp-2.7'
 include 'instrumentation:rabbit-amqp-3.5.0'
 include 'instrumentation:rabbit-amqp-5.0.0'