feat: implement autoscaling and the leastConnections strategy (#451)

Author: congoamz

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

@@ -22,6 +22,16 @@
 import static org.mockito.Mockito.when;
 import static org.mockito.MockitoAnnotations.openMocks;
 
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.EnumSet;
+import java.util.List;
+import java.util.Properties;
+import java.util.concurrent.TimeUnit;
 import org.mockito.Mock;
 import org.openjdk.jmh.annotations.Benchmark;
 import org.openjdk.jmh.annotations.BenchmarkMode;
@@ -50,22 +60,11 @@
 import software.amazon.jdbc.PluginManagerService;
 import software.amazon.jdbc.PluginService;
 import software.amazon.jdbc.PropertyDefinition;
+import software.amazon.jdbc.benchmarks.testplugin.BenchmarkPluginFactory;
 import software.amazon.jdbc.dialect.Dialect;
 import software.amazon.jdbc.profile.DriverConfigurationProfiles;
-import software.amazon.jdbc.benchmarks.testplugin.BenchmarkPluginFactory;
 import software.amazon.jdbc.wrapper.ConnectionWrapper;
 
-import java.sql.Connection;
-import java.sql.ResultSet;
-import java.sql.SQLException;
-import java.sql.Statement;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.EnumSet;
-import java.util.List;
-import java.util.Properties;
-import java.util.concurrent.TimeUnit;
-
 @State(Scope.Benchmark)
 @Fork(3)
 @Warmup(iterations = 3)
@@ -108,8 +107,8 @@ public void setUpIteration() throws Exception {
 
     when(mockConnectionProvider.connect(anyString(), any(Properties.class))).thenReturn(
         mockConnection);
-    when(mockConnectionProvider.connect(anyString(), any(Dialect.class), any(HostSpec.class), any(Properties.class)))
-        .thenReturn(mockConnection);
+    when(mockConnectionProvider.connect(anyString(), any(Dialect.class), any(HostSpec.class),
+        any(Properties.class))).thenReturn(mockConnection);
     when(mockConnection.createStatement()).thenReturn(mockStatement);
     when(mockStatement.executeQuery(anyString())).thenReturn(mockResultSet);
     when(mockResultSet.next()).thenReturn(true, true, false);

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

@@ -107,8 +107,8 @@ public void setUpIteration() throws Exception {
         .thenReturn(mockStatement);
     when(mockConnectionProvider.connect(anyString(), any(Properties.class))).thenReturn(
         mockConnection);
-    when(mockConnectionProvider.connect(anyString(), any(Dialect.class), any(HostSpec.class), any(Properties.class)))
-        .thenReturn(mockConnection);
+    when(mockConnectionProvider.connect(anyString(), any(Dialect.class), any(HostSpec.class),
+        any(Properties.class))).thenReturn(mockConnection);
     when(mockConnection.createStatement()).thenReturn(mockStatement);
     when(mockStatement.executeQuery(anyString())).thenReturn(mockResultSet);
     when(mockResultSet.next()).thenReturn(true, true, false);

docs/using-the-jdbc-driver/using-plugins/UsingTheReadWriteSplittingPlugin.md

@@ -39,7 +39,28 @@ The wrapper driver currently uses [Hikari](https://github.com/brettwooldridge/Hi
 - username
 - password
 
-You can optionally pass in a `HikariPoolMapping` function as a second parameter to the `HikariPooledConnectionProvider`. This allows you to decide when new connection pools should be created by defining what is included in the pool map key. A new pool will be created each time a new connection is requested with a unique key. By default, a new pool will be created for each unique instance-user combination. If you would like to define a different key system, you should pass in a `HikariPoolMapping` function defining this logic. Note that the user will always be automatically included in the key for security reasons. Please see [ReadWriteSplittingPostgresExample.java](../../../examples/AWSDriverExample/src/main/java/software/amazon/ReadWriteSplittingPostgresExample.java) for an example of how to configure the pool map key.
+You can optionally pass in a `HikariPoolMapping` function as a second parameter to the `HikariPooledConnectionProvider`. This allows you to decide when new connection pools should be created by defining what is included in the pool map key. A new pool will be created each time a new connection is requested with a unique key. By default, a new pool will be created for each unique instance-user combination. If you would like to define a different key system, you should pass in a `HikariPoolMapping` function defining this logic. A simple example is show below. Please see [ReadWriteSplittingPostgresExample.java](../../../examples/AWSDriverExample/src/main/java/software/amazon/ReadWriteSplittingPostgresExample.java) for the full example.
+
+> :warning: If you do not include the username in your HikariPoolMapping function, connection pools may be shared between different users.
+
+```java
+props.setProperty("somePropertyValue", "1"); // used in getPoolKey
+final HikariPooledConnectionProvider connProvider =
+    new HikariPooledConnectionProvider(
+        ReadWriteSplittingPostgresExample::getHikariConfig,
+        ReadWriteSplittingPostgresExample::getPoolKey
+    );
+ConnectionProviderManager.setConnectionProvider(connProvider);
+
+private static String getPoolKey(HostSpec hostSpec, Properties props) {
+  // Include the URL, user, and somePropertyValue in the connection pool key so that a new
+  // connection pool will be opened for each different instance-user-somePropertyValue
+  // combination.
+  final String user = props.getProperty(PropertyDefinition.USER.name);
+  final String somePropertyValue = props.getProperty("somePropertyValue");
+  return hostSpec.getUrl() + user + somePropertyValue;
+}
+```
 
 2. Call `ConnectionProviderManager.setConnectionProvider`, passing in the `HikariPooledConnectionProvider` you created in step 1.
 

examples/AWSDriverExample/src/main/java/software/amazon/ReadWriteSplittingPostgresExample.java

@@ -164,10 +164,11 @@ private static HikariConfig getHikariConfig(HostSpec hostSpec, Properties props)
   // This method is an optional parameter to `ConnectionProviderManager.setConnectionProvider`.
   // It can be omitted if you do not require it.
   private static String getPoolKey(HostSpec hostSpec, Properties props) {
-    // Include the URL and somePropertyValue in the connection pool key so that a new connection
-    // pool will be opened for each different instance-user-somePropertyValue combination.
-    // (Note that the user will automatically be added to the key).
+    // Include the URL, user, and somePropertyValue in the connection pool key so that a new
+    // connection pool will be opened for each different instance-user-somePropertyValue
+    // combination.
+    final String user = props.getProperty(PropertyDefinition.USER.name);
     final String somePropertyValue = props.getProperty("somePropertyValue");
-    return hostSpec.getUrl() + somePropertyValue;
+    return hostSpec.getUrl() + user + somePropertyValue;
   }
 }

wrapper/build.gradle.kts

@@ -408,3 +408,27 @@ tasks.register<Test>("test-aurora-mysql-performance") {
         systemProperty("test-no-mariadb-engine", "true")
     }
 }
+
+// Autoscaling
+
+tasks.register<Test>("test-autoscaling-only") {
+    group = "verification"
+    filter.includeTestsMatching("integration.refactored.host.TestRunner.runTests")
+    doFirst {
+        systemProperty("test-autoscaling-only", "true")
+        systemProperty("test-no-docker", "true")
+        systemProperty("test-no-performance", "true")
+        systemProperty("test-no-graalvm", "true")
+    }
+}
+
+tasks.register<Test>("debug-autoscaling-only") {
+    group = "verification"
+    filter.includeTestsMatching("integration.refactored.host.TestRunner.debugTests")
+    doFirst {
+        systemProperty("test-autoscaling-only", "true")
+        systemProperty("test-no-docker", "true")
+        systemProperty("test-no-performance", "true")
+        systemProperty("test-no-graalvm", "true")
+    }
+}

wrapper/src/main/java/software/amazon/jdbc/ConnectionProvider.java

@@ -63,8 +63,9 @@ boolean acceptsUrl(
    * @throws SQLException                  if an error occurred while returning the hosts
    * @throws UnsupportedOperationException if the strategy is unsupported by the provider
    */
-  HostSpec getHostSpecByStrategy(@NonNull List<HostSpec> hosts, @NonNull HostRole role,
-      @NonNull String strategy) throws SQLException, UnsupportedOperationException;
+  HostSpec getHostSpecByStrategy(
+      @NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy)
+      throws SQLException, UnsupportedOperationException;
 
   /**
    * Called once per connection that needs to be created.