Implement internal connection pools for the read-write plugin (#359)

Add functionality to use internal connection pools to establish
connections. The motivation for using pools comes from the read-write
plugin's need to regularly establish new connections. Using pools will
increase performance and decrease demands on the database by allowing
connections to be re-used.

Author: congoamz

.editorconfig

@@ -42,8 +42,9 @@ ij_any_spaces_around_relational_operators = true
 ij_any_spaces_around_shift_operators = true
 ij_continuation_indent_size = 4
 ij_java_align_multiline_parameters = false
+ij_java_doc_indent_on_continuation = true
 ij_java_if_brace_force = always
-ij_java_indent_case_from_switch = false
+ij_java_indent_case_from_switch = true
 ij_java_line_comment_add_space = true
 ij_java_space_after_colon = true
 ij_java_space_before_colon = true

benchmarks/build.gradle.kts

@@ -23,6 +23,7 @@ dependencies {
     implementation("org.postgresql:postgresql:42.5.0")
     implementation("mysql:mysql-connector-java:8.0.31")
     implementation("org.mariadb.jdbc:mariadb-java-client:3.1.0")
+    implementation("com.zaxxer:HikariCP:4.0.3")
 
     testImplementation("org.junit.jupiter:junit-jupiter-api:5.9.0")
     testImplementation("org.mockito:mockito-inline:4.8.0")

benchmarks/src/jmh/java/software/amazon/jdbc/benchmarks/ConnectionPluginManagerBenchmarks.java

@@ -77,7 +77,7 @@ public class ConnectionPluginManagerBenchmarks {
   private static final String WRITER_SESSION_ID = "MASTER_SESSION_ID";
   private static final String FIELD_SERVER_ID = "SERVER_ID";
   private static final String FIELD_SESSION_ID = "SESSION_ID";
-  private final Properties emptyProperties = new Properties();
+  private Properties propertiesWithoutPlugins;
   private Properties propertiesWithPlugins;
   private ConnectionPluginManager pluginManager;
   private ConnectionPluginManager pluginManagerWithNoPlugins;
@@ -125,14 +125,17 @@ public void setUpIteration() throws Exception {
     DriverConfigurationProfiles.addOrReplaceProfile(
         "benchmark",
         pluginFactories);
+    propertiesWithoutPlugins = new Properties();
+    propertiesWithoutPlugins.setProperty(PropertyDefinition.PLUGINS.name, "");
+
     propertiesWithPlugins = new Properties();
     propertiesWithPlugins.setProperty(PropertyDefinition.PROFILE_NAME.name, "benchmark");
 
     pluginManager = new ConnectionPluginManager(mockConnectionProvider, mockConnectionWrapper);
     pluginManager.init(mockPluginService, propertiesWithPlugins, mockPluginManagerService);
 
     pluginManagerWithNoPlugins = new ConnectionPluginManager(mockConnectionProvider, mockConnectionWrapper);
-    pluginManagerWithNoPlugins.init(mockPluginService, emptyProperties, mockPluginManagerService);
+    pluginManagerWithNoPlugins.init(mockPluginService, propertiesWithoutPlugins, mockPluginManagerService);
   }
 
   @TearDown(Level.Iteration)
@@ -143,7 +146,7 @@ public void tearDownIteration() throws Exception {
   @Benchmark
   public ConnectionPluginManager initConnectionPluginManagerWithNoPlugins() throws SQLException {
     final ConnectionPluginManager manager = new ConnectionPluginManager(mockConnectionProvider, mockConnectionWrapper);
-    manager.init(mockPluginService, emptyProperties, mockPluginManagerService);
+    manager.init(mockPluginService, propertiesWithoutPlugins, mockPluginManagerService);
     return manager;
   }
 
@@ -168,7 +171,7 @@ public Connection connectWithNoPlugins() throws SQLException {
     return pluginManagerWithNoPlugins.connect(
         "driverProtocol",
         new HostSpec("host"),
-        emptyProperties,
+        propertiesWithoutPlugins,
         true);
   }
 
@@ -211,7 +214,7 @@ public ConnectionPluginManager initHostProvidersWithNoPlugins() throws SQLExcept
     pluginManagerWithNoPlugins.initHostProvider(
         "protocol",
         "url",
-        emptyProperties,
+        propertiesWithoutPlugins,
         mockHostListProvider);
     return pluginManager;
   }

benchmarks/src/jmh/java/software/amazon/jdbc/benchmarks/PluginBenchmarks.java

@@ -17,10 +17,12 @@
 package software.amazon.jdbc.benchmarks;
 
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.when;
 
+import com.zaxxer.hikari.HikariConfig;
 import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.SQLException;
@@ -46,9 +48,15 @@
 import org.openjdk.jmh.runner.RunnerException;
 import org.openjdk.jmh.runner.options.Options;
 import org.openjdk.jmh.runner.options.OptionsBuilder;
+import software.amazon.jdbc.ConnectionPluginManager;
 import software.amazon.jdbc.ConnectionProvider;
+import software.amazon.jdbc.ConnectionProviderManager;
+import software.amazon.jdbc.HikariPooledConnectionProvider;
+import software.amazon.jdbc.HostListProviderService;
 import software.amazon.jdbc.HostSpec;
+import software.amazon.jdbc.PluginManagerService;
 import software.amazon.jdbc.PluginService;
+import software.amazon.jdbc.benchmarks.testplugin.TestConnectionWrapper;
 import software.amazon.jdbc.dialect.Dialect;
 import software.amazon.jdbc.wrapper.ConnectionWrapper;
 
@@ -63,13 +71,16 @@ public class PluginBenchmarks {
   private static final String WRITER_SESSION_ID = "MASTER_SESSION_ID";
   private static final String FIELD_SERVER_ID = "SERVER_ID";
   private static final String FIELD_SESSION_ID = "SESSION_ID";
-  private static final String CONNECTION_STRING = "driverProtocol://my.domain.com";
+  private static final String CONNECTION_STRING = "jdbc:postgresql://my.domain.com";
   private static final String PG_CONNECTION_STRING =
       "jdbc:aws-wrapper:postgresql://instance-0.XYZ.us-east-2.rds.amazonaws.com";
   private static final int TEST_PORT = 5432;
   private final HostSpec writerHostSpec = new HostSpec("instance-0", TEST_PORT);
 
   @Mock private PluginService mockPluginService;
+  @Mock private ConnectionPluginManager mockConnectionPluginManager;
+  @Mock private HostListProviderService mockHostListProviderService;
+  @Mock private PluginManagerService mockPluginManagerService;
   @Mock ConnectionProvider mockConnectionProvider;
   @Mock Connection mockConnection;
   @Mock Statement mockStatement;
@@ -89,6 +100,11 @@ public static void main(String[] args) throws RunnerException {
   @Setup(Level.Iteration)
   public void setUpIteration() throws Exception {
     closeable = MockitoAnnotations.openMocks(this);
+    when(mockConnectionPluginManager.connect(any(), any(), any(Properties.class), anyBoolean()))
+        .thenReturn(mockConnection);
+    when(mockConnectionPluginManager.execute(
+        any(), any(), any(), eq("Connection.createStatement"), any(), any()))
+        .thenReturn(mockStatement);
     when(mockConnectionProvider.connect(anyString(), any(Properties.class))).thenReturn(
         mockConnection);
     when(mockConnectionProvider.connect(anyString(), any(Dialect.class), any(HostSpec.class), any(Properties.class)))
@@ -101,6 +117,7 @@ public void setUpIteration() throws Exception {
         .thenReturn("instance-0", "instance-1");
     when(mockResultSet.getStatement()).thenReturn(mockStatement);
     when(mockStatement.getConnection()).thenReturn(mockConnection);
+    when(this.mockPluginService.acceptsStrategy(any(), eq("random"))).thenReturn(true);
     when(this.mockPluginService.getCurrentHostSpec()).thenReturn(writerHostSpec);
   }
 
@@ -110,48 +127,60 @@ public void tearDownIteration() throws Exception {
   }
 
   @Benchmark
-  public void initAndReleaseBaseLine() throws SQLException {
+  public void initAndReleaseBaseLine() {
   }
 
   @Benchmark
   public ConnectionWrapper initAndReleaseWithExecutionTimePlugin() throws SQLException {
-    try (ConnectionWrapper wrapper = new ConnectionWrapper(
+    try (ConnectionWrapper wrapper = new TestConnectionWrapper(
         useExecutionTimePlugin(),
         CONNECTION_STRING,
-        mockConnectionProvider)) {
+        mockConnectionPluginManager,
+        mockPluginService,
+        mockHostListProviderService,
+        mockPluginManagerService)) {
       wrapper.releaseResources();
       return wrapper;
     }
   }
 
   @Benchmark
   public ConnectionWrapper initAndReleaseWithAuroraHostListPlugin() throws SQLException {
-    try (ConnectionWrapper wrapper = new ConnectionWrapper(
+    try (ConnectionWrapper wrapper = new TestConnectionWrapper(
         useAuroraHostListPlugin(),
         CONNECTION_STRING,
-        mockConnectionProvider)) {
+        mockConnectionPluginManager,
+        mockPluginService,
+        mockHostListProviderService,
+        mockPluginManagerService)) {
       wrapper.releaseResources();
       return wrapper;
     }
   }
 
   @Benchmark
   public ConnectionWrapper initAndReleaseWithExecutionTimeAndAuroraHostListPlugins() throws SQLException {
-    try (ConnectionWrapper wrapper = new ConnectionWrapper(
+    try (ConnectionWrapper wrapper = new TestConnectionWrapper(
         useExecutionTimeAndAuroraHostListPlugins(),
         CONNECTION_STRING,
-        mockConnectionProvider)) {
+        mockConnectionPluginManager,
+        mockPluginService,
+        mockHostListProviderService,
+        mockPluginManagerService)) {
       wrapper.releaseResources();
       return wrapper;
     }
   }
 
   @Benchmark
   public ConnectionWrapper initAndReleaseWithReadWriteSplittingPlugin() throws SQLException {
-    try (ConnectionWrapper wrapper = new ConnectionWrapper(
+    try (ConnectionWrapper wrapper = new TestConnectionWrapper(
         useReadWriteSplittingPlugin(),
         CONNECTION_STRING,
-        mockConnectionProvider)) {
+        mockConnectionPluginManager,
+        mockPluginService,
+        mockHostListProviderService,
+        mockPluginManagerService)) {
       wrapper.releaseResources();
       return wrapper;
     }
@@ -160,46 +189,66 @@ public ConnectionWrapper initAndReleaseWithReadWriteSplittingPlugin() throws SQL
   @Benchmark
   public ConnectionWrapper initAndReleaseWithAuroraHostListAndReadWriteSplittingPlugin()
       throws SQLException {
-    try (ConnectionWrapper wrapper = new ConnectionWrapper(
+    try (ConnectionWrapper wrapper = new TestConnectionWrapper(
         useAuroraHostListAndReadWriteSplittingPlugin(),
         PG_CONNECTION_STRING,
-        mockConnectionProvider)) {
+        mockConnectionPluginManager,
+        mockPluginService,
+        mockHostListProviderService,
+        mockPluginManagerService)) {
       wrapper.releaseResources();
       return wrapper;
     }
   }
 
   @Benchmark
-  public ConnectionWrapper initAndReleaseWithReadWriteSplittingPluginWithReaderLoadBalancing()
-      throws SQLException {
-    try (ConnectionWrapper wrapper = new ConnectionWrapper(
-        useReadWriteSplittingPluginWithReaderLoadBalancing(),
+  public ConnectionWrapper initAndReleaseWithReadWriteSplittingPlugin_internalConnectionPools() throws SQLException {
+    HikariPooledConnectionProvider provider =
+        new HikariPooledConnectionProvider((hostSpec, props) -> new HikariConfig());
+    ConnectionProviderManager.setConnectionProvider(provider);
+    try (ConnectionWrapper wrapper = new TestConnectionWrapper(
+        useReadWriteSplittingPlugin(),
         CONNECTION_STRING,
-        mockConnectionProvider)) {
+        mockConnectionPluginManager,
+        mockPluginService,
+        mockHostListProviderService,
+        mockPluginManagerService)) {
       wrapper.releaseResources();
+      ConnectionProviderManager.releaseResources();
+      ConnectionProviderManager.resetProvider();
       return wrapper;
     }
   }
 
   @Benchmark
-  public ConnectionWrapper
-  initAndReleaseWithAuroraHostListAndReadWriteSplittingPluginWithReaderLoadBalancing()
+  public ConnectionWrapper initAndReleaseWithAuroraHostListAndReadWriteSplittingPlugin_internalConnectionPools()
       throws SQLException {
-    try (ConnectionWrapper wrapper = new ConnectionWrapper(
-        useAuroraHostListAndReadWriteSplittingPluginWithReaderLoadBalancing(),
+    HikariPooledConnectionProvider provider =
+        new HikariPooledConnectionProvider((hostSpec, props) -> new HikariConfig());
+    ConnectionProviderManager.setConnectionProvider(provider);
+    try (ConnectionWrapper wrapper = new TestConnectionWrapper(
+        useAuroraHostListAndReadWriteSplittingPlugin(),
         PG_CONNECTION_STRING,
-        mockConnectionProvider)) {
+        mockConnectionPluginManager,
+        mockPluginService,
+        mockHostListProviderService,
+        mockPluginManagerService)) {
       wrapper.releaseResources();
+      ConnectionProviderManager.releaseResources();
+      ConnectionProviderManager.resetProvider();
       return wrapper;
     }
   }
 
   @Benchmark
   public Statement executeStatementBaseline() throws SQLException {
-    try (ConnectionWrapper wrapper = new ConnectionWrapper(
+    try (ConnectionWrapper wrapper = new TestConnectionWrapper(
         useExecutionTimePlugin(),
         CONNECTION_STRING,
-        mockConnectionProvider);
+        mockConnectionPluginManager,
+        mockPluginService,
+        mockHostListProviderService,
+        mockPluginManagerService);
          Statement statement = wrapper.createStatement()) {
       return statement;
     }
@@ -208,10 +257,13 @@ public Statement executeStatementBaseline() throws SQLException {
   @Benchmark
   public ResultSet executeStatementWithExecutionTimePlugin() throws SQLException {
     try (
-        ConnectionWrapper wrapper = new ConnectionWrapper(
+        ConnectionWrapper wrapper = new TestConnectionWrapper(
             useExecutionTimePlugin(),
             CONNECTION_STRING,
-            mockConnectionProvider);
+            mockConnectionPluginManager,
+            mockPluginService,
+            mockHostListProviderService,
+            mockPluginManagerService);
         Statement statement = wrapper.createStatement();
         ResultSet resultSet = statement.executeQuery("some sql")) {
       return resultSet;
@@ -247,18 +299,4 @@ Properties useAuroraHostListAndReadWriteSplittingPlugin() {
     properties.setProperty("wrapperPlugins", "auroraHostList,readWriteSplitting");
     return properties;
   }
-
-  Properties useReadWriteSplittingPluginWithReaderLoadBalancing() {
-    final Properties properties = new Properties();
-    properties.setProperty("wrapperPlugins", "readWriteSplitting");
-    properties.setProperty("loadBalanceReadOnlyTraffic", "true");
-    return properties;
-  }
-
-  Properties useAuroraHostListAndReadWriteSplittingPluginWithReaderLoadBalancing() {
-    final Properties properties = new Properties();
-    properties.setProperty("wrapperPlugins", "auroraHostList,readWriteSplitting");
-    properties.setProperty("loadBalanceReadOnlyTraffic", "true");
-    return properties;
-  }
 }

benchmarks/src/jmh/java/software/amazon/jdbc/benchmarks/testplugin/BenchmarkPlugin.java

@@ -26,11 +26,10 @@
 import java.util.Map;
 import java.util.Properties;
 import java.util.Set;
-import java.util.logging.Level;
 import java.util.logging.Logger;
-
 import software.amazon.jdbc.ConnectionPlugin;
 import software.amazon.jdbc.HostListProviderService;
+import software.amazon.jdbc.HostRole;
 import software.amazon.jdbc.HostSpec;
 import software.amazon.jdbc.JdbcCallable;
 import software.amazon.jdbc.NodeChangeOptions;
@@ -60,11 +59,32 @@ public <T, E extends Exception> T execute(Class<T> resultClass, Class<E> excepti
   public Connection connect(String driverProtocol, HostSpec hostSpec, Properties props,
       boolean isInitialConnection, JdbcCallable<Connection, SQLException> connectFunc)
       throws SQLException {
-    LOGGER.finer(() -> String.format("connect=''%s''", driverProtocol));
+    LOGGER.finer(() -> String.format("connect=''%s''", hostSpec.getHost()));
     resources.add("connect");
     return connectFunc.call();
   }
 
+  @Override
+  public Connection forceConnect(String driverProtocol, HostSpec hostSpec, Properties props,
+      boolean isInitialConnection, JdbcCallable<Connection, SQLException> forceConnectFunc)
+      throws SQLException {
+    LOGGER.finer(() -> String.format("forceConnect=''%s''", hostSpec.getHost()));
+    resources.add("forceConnect");
+    return forceConnectFunc.call();
+  }
+
+  @Override
+  public boolean acceptsStrategy(HostRole role, String strategy) {
+    return false;
+  }
+
+  @Override
+  public HostSpec getHostSpecByStrategy(HostRole role, String strategy) {
+    LOGGER.finer(() -> String.format("getHostSpecByStrategy=''%s''", strategy));
+    resources.add("getHostSpecByStrategy");
+    return new HostSpec("host", 1234, role);
+  }
+
   @Override
   public void initHostProvider(String driverProtocol, String initialUrl, Properties props,
       HostListProviderService hostListProviderService,