[#9799] Add connection factory to spring r2dbc plugin

agent-testweb/spring-data-r2dbc-plugin-testweb/src/main/java/com/pinpoint/test/plugin/MysqlR2dbcDatabase.java

@@ -18,13 +18,10 @@
 
 import dev.miku.r2dbc.mysql.MySqlConnectionConfiguration;
 import dev.miku.r2dbc.mysql.MySqlConnectionFactory;
-import io.r2dbc.spi.ConnectionFactories;
+import dev.miku.r2dbc.mysql.constant.SslMode;
 import io.r2dbc.spi.ConnectionFactory;
-import io.r2dbc.spi.ConnectionFactoryOptions;
 import org.springframework.beans.factory.annotation.Qualifier;
-import org.springframework.data.r2dbc.core.R2dbcEntityTemplate;
 import org.springframework.stereotype.Component;
-import reactor.core.publisher.Mono;
 
 import javax.annotation.PostConstruct;
 import javax.annotation.PreDestroy;
@@ -39,11 +36,12 @@ public class MysqlR2dbcDatabase implements R2dbcDatabase {
     public void init() throws Exception {
         MySqlConnectionConfiguration connectionConfiguration = MySqlConnectionConfiguration.builder()
                 .host("localhost")
-                .port(8066)
+                .port(49178)
                 .database("test")
                 .user("root")
                 .password("")
                 .connectTimeout(Duration.ofSeconds(5 * 60))
+                .sslMode(SslMode.DISABLED)
                 .build();
 
         connectionFactory = MySqlConnectionFactory.from(connectionConfiguration);

agent-testweb/spring-data-r2dbc-plugin-testweb/src/main/java/com/pinpoint/test/plugin/SpringDataR2dbcPluginController.java

@@ -44,7 +44,7 @@ public class SpringDataR2dbcPluginController {
 
     R2dbcDatabase r2dbcDatabase;
 
-    public SpringDataR2dbcPluginController(@Qualifier("mssql") R2dbcDatabase r2dbcDatabase) {
+    public SpringDataR2dbcPluginController(@Qualifier("mysql") R2dbcDatabase r2dbcDatabase) {
         this.r2dbcDatabase = r2dbcDatabase;
     }
 

agent-testweb/spring-data-r2dbc-plugin-testweb/src/test/java/com/pinpoint/test/plugin/spring/data/r2dbc/R2dbcMysqlTest.java

@@ -24,8 +24,6 @@
 import org.testcontainers.DockerClientFactory;
 import org.testcontainers.containers.MySQLContainer;
 
-import java.util.concurrent.TimeUnit;
-
 public class R2dbcMysqlTest {
     public static final String DATABASE_NAME = "test";
     public static final String USERNAME = "root";

plugins/spring-data-r2dbc/src/main/java/com/navercorp/pinpoint/plugin/spring/r2dbc/SpringDataR2dbcPlugin.java

@@ -45,37 +45,38 @@
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.DefaultGenericExecuteSpecInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.DefaultGenericExecuteSpecResultFunctionApplyInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.DefaultGenericExecuteSpecStatementFunctionGetLambdaInterceptor;
+import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.SetDatabaseInfoConstructorInterceptor;
+import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.StatementBindInterceptor;
+import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.StatementBindNullInterceptor;
+import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.StatementExecuteInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.h2.H2ConnectionConfigurationConstructorInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.h2.H2ConnectionConstructorInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.h2.H2ConnectionFactoryConstructorInterceptor;
-import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.jasync.JasyncConnectionFactoryConstructorInterceptor;
+import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.h2.SessionClientConstructorInterceptor;
+import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.jasync.ConfigurationConstructorInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.jasync.JasyncClientConnectionConstructorInterceptor;
+import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.jasync.JasyncConnectionFactoryConstructorInterceptor;
+import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.jasync.MySQLConnectionFactoryConstructorInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.mariadb.MariadbConnectionConfigurationConstructorInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.mariadb.MariadbConnectionConstructorInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.mariadb.MariadbConnectionFactoryTransformConstructorInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.mssql.MssqlConnectionConfigurationConstructorInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.mssql.MssqlConnectionConfigurationToConnectionOptionsInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.mssql.MssqlConnectionConstructorInterceptor;
-import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.jasync.MySQLConnectionFactoryConstructorInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.mssql.MssqlConnectionFactoryConstructorInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.mssql.MssqlStatementBindInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.mssql.MssqlStatementBindNullInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.mssql.MssqlStatementExecuteInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.mysql.MySqlConnectionConfigurationInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.mysql.MySqlConnectionConstructorInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.mysql.MySqlConnectionFactoryFromInterceptor;
+import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.mysql.QueryFlowLoginInterceptor;
+import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.mysql.ReactorNettyClientConstructorInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.oracle.OracleConnectionFactoryImplConstructorInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.oracle.OracleConnectionFactoryImplLambdaCreateInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.oracle.OracleConnectionImplConstructorInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.postgresql.PostgresqlConnectionConfigurationInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.postgresql.PostgresqlConnectionConstructorInterceptor;
-import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.mysql.QueryFlowLoginInterceptor;
-import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.mysql.ReactorNettyClientConstructorInterceptor;
-import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.h2.SessionClientConstructorInterceptor;
-import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.StatementBindInterceptor;
-import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.StatementBindNullInterceptor;
-import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.StatementExecuteInterceptor;
-import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.jasync.ConfigurationConstructorInterceptor;
 import com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor.postgresql.PostgresqlConnectionFactoryConstructorInterceptor;
 
 import java.security.ProtectionDomain;
@@ -164,6 +165,14 @@ public void setup(ProfilerPluginSetupContext context) {
         transformTemplate.transform("org.springframework.r2dbc.core.DefaultDatabaseClient$DefaultGenericExecuteSpec", DefaultGenericExecuteSpecTransform.class);
         transformTemplate.transform("org.springframework.r2dbc.core.DefaultFetchSpec", DefaultFetchSpecTransform.class);
 
+        // ConnectionFactory
+        transformTemplate.transform("io.r2dbc.pool.ConnectionPoolConfiguration", ConnectionPoolConfigurationTransform.class);
+        transformTemplate.transform("io.r2dbc.pool.ConnectionPool", ConnectionPoolTransform.class);
+        transformTemplate.transform("org.springframework.boot.r2dbc.OptionsCapableConnectionFactory", OptionsCapableConnectionFactoryTransform.class);
+        transformTemplate.transform("org.springframework.r2dbc.connection.DelegatingConnectionFactory", DelegatingConnectionFactoryTransform.class);
+        transformTemplate.transform("org.springframework.r2dbc.connection.SingleConnectionFactory", SingleConnectionFactoryTransform.class);
+        transformTemplate.transform("org.springframework.r2dbc.connection.TransactionAwareConnectionFactoryProxy", TransactionAwareConnectionFactoryProxyTransform.class);
+
         // statementFunction, resultFunction Lambda
         final Matcher defaultGenericExecuteSpecLambdaMatcher = Matchers.newLambdaExpressionMatcher("org.springframework.r2dbc.core.DefaultDatabaseClient$DefaultGenericExecuteSpec", "java.util.function.Function");
         transformTemplate.transform(defaultGenericExecuteSpecLambdaMatcher, DefaultGenericExecuteSpecStatementFunctionTransform.class);
@@ -964,4 +973,117 @@ public byte[] doInTransform(Instrumentor instrumentor, ClassLoader classLoader,
             return target.toBytecode();
         }
     }
+
+    public static class ConnectionPoolConfigurationTransform implements TransformCallback {
+
+        @Override
+        public byte[] doInTransform(Instrumentor instrumentor, ClassLoader classLoader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws InstrumentException {
+            final InstrumentClass target = instrumentor.getInstrumentClass(classLoader, className, classfileBuffer);
+            target.addField(DatabaseInfoAccessor.class);
+
+            for (InstrumentMethod constructorMethod : target.getDeclaredConstructors()) {
+                if (ArrayUtils.hasLength(constructorMethod.getParameterTypes())) {
+                    // ConnectionPoolConfiguration(int acquireRetry, @Nullable Duration backgroundEvictionInterval, ConnectionFactory connectionFactory, ...)
+                    constructorMethod.addInterceptor(SetDatabaseInfoConstructorInterceptor.class, va(2));
+                }
+            }
+
+            return target.toBytecode();
+        }
+    }
+
+    public static class ConnectionPoolTransform implements TransformCallback {
+
+        @Override
+        public byte[] doInTransform(Instrumentor instrumentor, ClassLoader classLoader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws InstrumentException {
+            final InstrumentClass target = instrumentor.getInstrumentClass(classLoader, className, classfileBuffer);
+            target.addField(DatabaseInfoAccessor.class);
+
+            final InstrumentMethod constructorMethod = target.getConstructor("io.r2dbc.pool.ConnectionPoolConfiguration");
+            if (constructorMethod != null) {
+                // ConnectionPool(ConnectionPoolConfiguration configuration)
+                constructorMethod.addInterceptor(SetDatabaseInfoConstructorInterceptor.class, va(0));
+            }
+            final InstrumentMethod createMethod = target.getDeclaredMethod("create");
+            if (createMethod != null) {
+                createMethod.addInterceptor(ConnectionFactoryCreateInterceptor.class);
+            }
+
+            return target.toBytecode();
+        }
+    }
+
+    public static class OptionsCapableConnectionFactoryTransform implements TransformCallback {
+
+        @Override
+        public byte[] doInTransform(Instrumentor instrumentor, ClassLoader classLoader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws InstrumentException {
+            final InstrumentClass target = instrumentor.getInstrumentClass(classLoader, className, classfileBuffer);
+            target.addField(DatabaseInfoAccessor.class);
+
+            final InstrumentMethod constructorMethod = target.getConstructor("io.r2dbc.spi.ConnectionFactoryOptions", "io.r2dbc.spi.ConnectionFactory");
+            if (constructorMethod != null) {
+                // OptionsCapableConnectionFactory(ConnectionFactoryOptions options, ConnectionFactory delegate)
+                constructorMethod.addInterceptor(SetDatabaseInfoConstructorInterceptor.class, va(1));
+            }
+            final InstrumentMethod createMethod = target.getDeclaredMethod("create");
+            if (createMethod != null) {
+                createMethod.addInterceptor(ConnectionFactoryCreateInterceptor.class);
+            }
+
+            return target.toBytecode();
+        }
+    }
+
+    public static class DelegatingConnectionFactoryTransform implements TransformCallback {
+
+        @Override
+        public byte[] doInTransform(Instrumentor instrumentor, ClassLoader classLoader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws InstrumentException {
+            final InstrumentClass target = instrumentor.getInstrumentClass(classLoader, className, classfileBuffer);
+            target.addField(DatabaseInfoAccessor.class);
+
+            final InstrumentMethod constructorMethod = target.getConstructor("io.r2dbc.spi.ConnectionFactory");
+            if (constructorMethod != null) {
+                // DelegatingConnectionFactory(ConnectionFactory targetConnectionFactory)
+                constructorMethod.addInterceptor(SetDatabaseInfoConstructorInterceptor.class, va(0));
+            }
+            final InstrumentMethod createMethod = target.getDeclaredMethod("create");
+            if (createMethod != null) {
+                createMethod.addInterceptor(ConnectionFactoryCreateInterceptor.class);
+            }
+
+            return target.toBytecode();
+        }
+    }
+
+
+    public static class SingleConnectionFactoryTransform implements TransformCallback {
+
+        @Override
+        public byte[] doInTransform(Instrumentor instrumentor, ClassLoader classLoader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws InstrumentException {
+            final InstrumentClass target = instrumentor.getInstrumentClass(classLoader, className, classfileBuffer);
+
+            final InstrumentMethod createMethod = target.getDeclaredMethod("create");
+            if (createMethod != null) {
+                createMethod.addInterceptor(ConnectionFactoryCreateInterceptor.class);
+            }
+
+            return target.toBytecode();
+        }
+    }
+
+
+    public static class TransactionAwareConnectionFactoryProxyTransform implements TransformCallback {
+
+        @Override
+        public byte[] doInTransform(Instrumentor instrumentor, ClassLoader classLoader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws InstrumentException {
+            final InstrumentClass target = instrumentor.getInstrumentClass(classLoader, className, classfileBuffer);
+
+            final InstrumentMethod createMethod = target.getDeclaredMethod("create");
+            if (createMethod != null) {
+                createMethod.addInterceptor(ConnectionFactoryCreateInterceptor.class);
+            }
+
+            return target.toBytecode();
+        }
+    }
 }

plugins/spring-data-r2dbc/src/main/java/com/navercorp/pinpoint/plugin/spring/r2dbc/interceptor/DefaultGenericExecuteSpecStatementFunctionGetLambdaInterceptor.java

@@ -45,7 +45,7 @@ public void after(Object target, Object[] args, Object result, Throwable throwab
 
         final DatabaseInfo databaseInfo = DatabaseInfoAccessorUtils.getDatabaseInfo(args, 0);
         if (databaseInfo != null) {
-            DatabaseInfoAccessorUtils.setDatabaseInfo(databaseInfo, result);
+            DatabaseInfoAccessorUtils.setDatabaseInfo(databaseInfo, target);
         }
 
         if (Boolean.FALSE == target instanceof AsyncContextAccessor) {

plugins/spring-data-r2dbc/src/main/java/com/navercorp/pinpoint/plugin/spring/r2dbc/interceptor/SetDatabaseInfoConstructorInterceptor.java

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2023 NAVER Corp.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.navercorp.pinpoint.plugin.spring.r2dbc.interceptor;
+
+import com.navercorp.pinpoint.bootstrap.context.DatabaseInfo;
+import com.navercorp.pinpoint.bootstrap.plugin.jdbc.DatabaseInfoAccessor;
+import com.navercorp.pinpoint.common.util.ArrayArgumentUtils;
+
+public class SetDatabaseInfoConstructorInterceptor extends SetDatabaseInfoInterceptor {
+    private final int argumentIndex;
+
+    public SetDatabaseInfoConstructorInterceptor(int argumentIndex) {
+        this.argumentIndex = argumentIndex;
+    }
+
+    @Override
+    public DatabaseInfo getDatabaseInfo(Object target, Object[] args, Object result) {
+        final DatabaseInfoAccessor databaseInfoAccessor = ArrayArgumentUtils.getArgument(args, argumentIndex, DatabaseInfoAccessor.class);
+        if (databaseInfoAccessor == null) {
+            return null;
+        }
+        return databaseInfoAccessor._$PINPOINT$_getDatabaseInfo();
+    }
+
+    @Override
+    public boolean setDatabaseInfo(DatabaseInfo databaseInfo, Object target, Object[] args, Object result) {
+        if (target instanceof DatabaseInfoAccessor) {
+            ((DatabaseInfoAccessor) target)._$PINPOINT$_setDatabaseInfo(databaseInfo);
+            return true;
+        }
+        return false;
+    }
+}