Various r2dbc updates and changes

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

@@ -53,11 +53,9 @@ 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();
-                    }
-            }
+                if(instrumentedClient.remoteAddress != null && instrumentedClient.remoteAddress instanceof InetSocketAddress) {
+                    return (InetSocketAddress) instrumentedClient.remoteAddress;
+                }            }
             return null;
         } catch(Exception exception) {
             return null;

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

@@ -4,13 +4,15 @@
 import com.newrelic.api.agent.weaver.NewField;
 import com.newrelic.api.agent.weaver.Weave;
 import reactor.netty.Connection;
+import java.net.SocketAddress;
 
 @Weave(type = MatchType.ExactClass, originalName = "io.r2dbc.mssql.client.ReactorNettyClient")
 public class ReactorNettyClient_Instrumentation {
     @NewField
-    public final Connection clientConnection;
+    public final SocketAddress remoteAddress;
 
     private ReactorNettyClient_Instrumentation(Connection connection, TdsEncoder TdsEncoder, ConnectionContext context) {
-        this.clientConnection = connection;
+        this.remoteAddress = connection == null ? null :
+                connection.channel() == null ? null : connection.channel().remoteAddress();
     }
 }

instrumentation/r2dbc-mysql-1.1.3/build.gradle

@@ -0,0 +1,21 @@
+dependencies {
+    implementation(project(":agent-bridge"))
+    implementation(project(":agent-bridge-datastore"))
+    implementation("io.asyncer:r2dbc-mysql:1.1.3")
+    testImplementation("ch.vorburger.mariaDB4j:mariaDB4j:2.2.1")
+}
+
+jar {
+    manifest { attributes 'Implementation-Title': 'com.newrelic.instrumentation.r2dbc-mysql-1.1.3' }
+}
+
+verifyInstrumentation {
+    // note the older instrumentation is for the dev.mik: r2dbc-mysql, which only covers 8.2.0
+    // and this module only covers 1.1.3+, so we currently have a gap from 0.9.0 to 1.1.2
+    passesOnly 'io.asyncer:r2dbc-mysql:[1.1.3,)'
+}
+
+site {
+    title 'MySQL R2DBC'
+    type 'Datastore'
+}

instrumentation/r2dbc-mysql-1.1.3/src/main/java/io/asyncer/r2dbc/mysql/MySqlStatementSupport_Instrumentation.java

@@ -0,0 +1,11 @@
+package io.asyncer.r2dbc.mysql;
+
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+import io.asyncer.r2dbc.mysql.client.Client;
+
+@Weave(type = MatchType.ExactClass, originalName = "io.asyncer.r2dbc.mysql.MySqlStatementSupport")
+abstract class MySqlStatementSupport_Instrumentation {
+    protected final Client client = Weaver.callOriginal();
+}

instrumentation/r2dbc-mysql-1.1.3/src/main/java/io/asyncer/r2dbc/mysql/ParameterizedStatementSupport_Instrumentation.java

@@ -0,0 +1,11 @@
+package io.asyncer.r2dbc.mysql;
+
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+import io.asyncer.r2dbc.mysql.client.Client;
+
+@Weave(type = MatchType.ExactClass, originalName = "io.asyncer.r2dbc.mysql.ParameterizedStatementSupport")
+abstract class ParameterizedStatementSupport_Instrumentation extends MySqlStatementSupport_Instrumentation {
+    protected final Query query = Weaver.callOriginal();
+}

instrumentation/r2dbc-mysql-1.1.3/src/main/java/io/asyncer/r2dbc/mysql/PrepareParameterizedStatement_Instrumentation.java

@@ -0,0 +1,21 @@
+package io.asyncer.r2dbc.mysql;
+
+import io.asyncer.r2dbc.mysql.api.MySqlResult;
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+import io.asyncer.r2dbc.mysql.client.R2dbcUtils;
+import reactor.core.publisher.Flux;
+
+import java.util.List;
+
+@Weave(type = MatchType.ExactClass, originalName = "io.asyncer.r2dbc.mysql.PrepareParameterizedStatement")
+final class PrepareParameterizedStatement_Instrumentation extends ParameterizedStatementSupport_Instrumentation {
+    public Flux<MySqlResult> execute(List<Binding> bindings) {
+        Flux<MySqlResult> request = Weaver.callOriginal();
+        if(request != null && this.query != null && this.client != null) {
+            return R2dbcUtils.wrapRequest(request, query.getFormattedSql(), client);
+        }
+        return request;
+    }
+}

instrumentation/r2dbc-mysql-1.1.3/src/main/java/io/asyncer/r2dbc/mysql/PrepareSimpleStatement_Instrumentation.java

@@ -0,0 +1,19 @@
+package io.asyncer.r2dbc.mysql;
+
+import io.asyncer.r2dbc.mysql.api.MySqlResult;
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+import io.asyncer.r2dbc.mysql.client.R2dbcUtils;
+import reactor.core.publisher.Flux;
+
+@Weave(type = MatchType.ExactClass, originalName = "io.asyncer.r2dbc.mysql.PrepareSimpleStatement")
+final class PrepareSimpleStatement_Instrumentation extends SimpleStatementSupport_Instrumentation {
+    public Flux<MySqlResult> execute() {
+        Flux<MySqlResult> request = Weaver.callOriginal();
+        if(request != null && this.sql != null && this.client != null) {
+            return R2dbcUtils.wrapRequest(request, this.sql, this.client);
+        }
+        return request;
+    }
+}

instrumentation/r2dbc-mysql-1.1.3/src/main/java/io/asyncer/r2dbc/mysql/SimpleStatementSupport_Instrumentation.java

@@ -0,0 +1,11 @@
+package io.asyncer.r2dbc.mysql;
+
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+import io.asyncer.r2dbc.mysql.client.Client;
+
+@Weave(type = MatchType.ExactClass, originalName = "io.asyncer.r2dbc.mysql.SimpleStatementSupport")
+abstract class SimpleStatementSupport_Instrumentation extends MySqlStatementSupport_Instrumentation {
+    protected final String sql = Weaver.callOriginal();
+}

instrumentation/r2dbc-mysql-1.1.3/src/main/java/io/asyncer/r2dbc/mysql/TextParameterizedStatement_Instrumentation.java

@@ -0,0 +1,21 @@
+package io.asyncer.r2dbc.mysql;
+
+import io.asyncer.r2dbc.mysql.api.MySqlResult;
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+import io.asyncer.r2dbc.mysql.client.R2dbcUtils;
+import reactor.core.publisher.Flux;
+
+import java.util.List;
+
+@Weave(type = MatchType.ExactClass, originalName = "io.asyncer.r2dbc.mysql.TextParameterizedStatement")
+final class TextParameterizedStatement_Instrumentation extends ParameterizedStatementSupport_Instrumentation {
+    protected Flux<MySqlResult> execute(List<Binding> bindings) {
+        Flux<MySqlResult> request = Weaver.callOriginal();
+        if(request != null && this.query != null && this.client != null) {
+            return R2dbcUtils.wrapRequest(request, String.join("", query.getFormattedSql()), client);
+        }
+        return request;
+    }
+}

instrumentation/r2dbc-mysql-1.1.3/src/main/java/io/asyncer/r2dbc/mysql/TextSimpleStatement_Instrumentation.java

@@ -0,0 +1,19 @@
+package io.asyncer.r2dbc.mysql;
+
+import io.asyncer.r2dbc.mysql.api.MySqlResult;
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+import io.asyncer.r2dbc.mysql.client.R2dbcUtils;
+import reactor.core.publisher.Flux;
+
+@Weave(type = MatchType.ExactClass, originalName = "io.asyncer.r2dbc.mysql.TextSimpleStatement")
+final class TextSimpleStatement_Instrumentation extends SimpleStatementSupport_Instrumentation {
+    public Flux<MySqlResult> execute() {
+        Flux<MySqlResult> request = Weaver.callOriginal();
+        if(request != null && this.sql != null && this.client != null) {
+            return R2dbcUtils.wrapRequest(request, this.sql, this.client);
+        }
+        return request;
+    }
+}

instrumentation/r2dbc-mysql-1.1.3/src/main/java/io/asyncer/r2dbc/mysql/client/R2dbcUtils.java

@@ -0,0 +1,65 @@
+package io.asyncer.r2dbc.mysql.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.asyncer.r2dbc.mysql.api.MySqlResult;
+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<MySqlResult> wrapRequest(Flux<MySqlResult> 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;
+                if(instrumentedClient.remoteAddress != null && instrumentedClient.remoteAddress instanceof InetSocketAddress) {
+                    return (InetSocketAddress) instrumentedClient.remoteAddress;
+                }
+            }
+            return null;
+        } catch(Exception exception) {
+            return null;
+        }
+    }
+}

instrumentation/r2dbc-mysql-1.1.3/src/main/java/io/asyncer/r2dbc/mysql/client/ReactorNettyClient_Instrumentation.java

@@ -0,0 +1,21 @@
+package io.asyncer.r2dbc.mysql.client;
+
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.NewField;
+import com.newrelic.api.agent.weaver.Weave;
+import io.asyncer.r2dbc.mysql.ConnectionContext;
+import io.asyncer.r2dbc.mysql.MySqlSslConfiguration;
+import reactor.netty.Connection;
+
+import java.net.SocketAddress;
+
+@Weave(type = MatchType.ExactClass, originalName = "io.asyncer.r2dbc.mysql.client.ReactorNettyClient")
+class ReactorNettyClient_Instrumentation {
+    @NewField
+    public final SocketAddress remoteAddress;
+
+    ReactorNettyClient_Instrumentation(Connection connection, MySqlSslConfiguration ssl, ConnectionContext context) {
+        this.remoteAddress = connection == null ? null :
+                connection.channel() == null ? null : connection.channel().remoteAddress();
+    }
+}

instrumentation/r2dbc-mysql-1.1.3/src/test/java/com/nr/agent/instrumentation/r2dbc/MySQLInstrumentedTest.java

@@ -0,0 +1,99 @@
+package com.nr.agent.instrumentation.r2dbc;
+
+import ch.vorburger.mariadb4j.DB;
+import ch.vorburger.mariadb4j.DBConfigurationBuilder;
+import com.newrelic.agent.introspec.DatastoreHelper;
+import com.newrelic.agent.introspec.InstrumentationTestConfig;
+import com.newrelic.agent.introspec.InstrumentationTestRunner;
+import com.newrelic.agent.introspec.Introspector;
+import io.r2dbc.spi.Connection;
+import io.r2dbc.spi.ConnectionFactories;
+import io.r2dbc.spi.ConnectionFactory;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import reactor.core.publisher.Mono;
+
+import static org.junit.Assert.assertEquals;
+
+@RunWith(InstrumentationTestRunner.class)
+@InstrumentationTestConfig(includePrefixes = "io.asyncer.r2dbc.mysql")
+public class MySQLInstrumentedTest {
+
+    public static DB mariaDb;
+    public Connection connection;
+
+    @Before
+    public void setup() throws Exception {
+        String databaseName = "MySQL" + System.currentTimeMillis();
+        DBConfigurationBuilder builder = DBConfigurationBuilder.newBuilder().setPort(0);
+        mariaDb = DB.newEmbeddedDB(builder.build());
+        mariaDb.start();
+        mariaDb.createDB(databaseName);
+        mariaDb.source("users.sql", "user", "password", databaseName);
+        ConnectionFactory connectionFactory = ConnectionFactories.get(builder.getURL(databaseName).replace("jdbc", "r2dbc").replace("localhost", "user:password@localhost"));
+        connection = Mono.from(connectionFactory.create()).block();
+    }
+
+    @AfterClass
+    public static void teardown() throws Exception {
+        mariaDb.stop();
+    }
+
+    @Test
+    public void testBasicRequests() {
+        //Given
+        Introspector introspector = InstrumentationTestRunner.getIntrospector();
+        DatastoreHelper helper = new DatastoreHelper("MySQL");
+
+        //When
+        R2dbcTestUtils.basicRequests(connection);
+
+        //Then
+        assertEquals(1, introspector.getFinishedTransactionCount(1000));
+        assertEquals(1, introspector.getTransactionNames().size());
+        String transactionName = introspector.getTransactionNames().stream().findFirst().orElse("");
+        helper.assertScopedStatementMetricCount(transactionName, "INSERT", "USERS", 1);
+        helper.assertScopedStatementMetricCount(transactionName, "SELECT", "USERS", 3);
+        helper.assertScopedStatementMetricCount(transactionName, "UPDATE", "USERS", 1);
+        helper.assertScopedStatementMetricCount(transactionName, "DELETE", "USERS", 1);
+        helper.assertAggregateMetrics();
+        helper.assertUnscopedOperationMetricCount("INSERT", 1);
+        helper.assertUnscopedOperationMetricCount("SELECT", 3);
+        helper.assertUnscopedOperationMetricCount("UPDATE", 1);
+        helper.assertUnscopedOperationMetricCount("DELETE", 1);
+        helper.assertUnscopedStatementMetricCount("INSERT", "USERS", 1);
+        helper.assertUnscopedStatementMetricCount("SELECT", "USERS", 3);
+        helper.assertUnscopedStatementMetricCount("UPDATE", "USERS", 1);
+        helper.assertUnscopedStatementMetricCount("DELETE", "USERS", 1);
+    }
+
+    @Test
+    public void testParametrizedRequests() {
+        //Given
+        Introspector introspector = InstrumentationTestRunner.getIntrospector();
+        DatastoreHelper helper = new DatastoreHelper("MySQL");
+
+        //When
+        R2dbcTestUtils.parametrizedRequests(connection);
+
+        //Then
+        assertEquals(1, introspector.getFinishedTransactionCount(1000));
+        assertEquals(1, introspector.getTransactionNames().size());
+        String transactionName = introspector.getTransactionNames().stream().findFirst().orElse("");
+        helper.assertScopedStatementMetricCount(transactionName, "INSERT", "USERS", 1);
+        helper.assertScopedStatementMetricCount(transactionName, "SELECT", "USERS", 3);
+        helper.assertScopedStatementMetricCount(transactionName, "UPDATE", "USERS", 1);
+        helper.assertScopedStatementMetricCount(transactionName, "DELETE", "USERS", 1);
+        helper.assertAggregateMetrics();
+        helper.assertUnscopedOperationMetricCount("INSERT", 1);
+        helper.assertUnscopedOperationMetricCount("SELECT", 3);
+        helper.assertUnscopedOperationMetricCount("UPDATE", 1);
+        helper.assertUnscopedOperationMetricCount("DELETE", 1);
+        helper.assertUnscopedStatementMetricCount("INSERT", "USERS", 1);
+        helper.assertUnscopedStatementMetricCount("SELECT", "USERS", 3);
+        helper.assertUnscopedStatementMetricCount("UPDATE", "USERS", 1);
+        helper.assertUnscopedStatementMetricCount("DELETE", "USERS", 1);
+    }
+}