R2DBC MSSQL Instrumentation - Instrumentation to capture execute call…

…s when using MSSQL via R2DBC.

Instruments the SimpleMssqlStatement and ParametrizedMssqlStatement child classes of the R2DBC Statement interface implementation. Higher level R2DBC Statement interface doesn't provide enough information to instrument more generically.

Related Github Issue
#197

Testing
Manual testing

Checks
[X] Are your contributions backwards compatible with relevant frameworks and APIs? New instrumentation
[X] Does your code contain any breaking changes? Please describe. No
[X] Does your code introduce any new dependencies? Please describe. Yes, new instrumentation dependencies in new project.

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.MySQL.name())
+                        .collection(sqlOperation.getTableName())
+                        .operation(sqlOperation.getOperation())
+                        .instance(socketAddress.getHostName(), socketAddress.getPort())
+                        .databaseName(null)
+                        .slowQuery(sql, R2dbcObfuscator.MYSQL_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'