feat: round robin host selection strategy

Author: crystall-bitquill

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

@@ -69,11 +69,7 @@ private static String getPoolKey(HostSpec hostSpec, Properties props) {
 
 2. Call `ConnectionProviderManager.setConnectionProvider`, passing in the `HikariPooledConnectionProvider` you created in step 1.
 
-3. By default, the read-write plugin randomly selects a reader instance the first time that `setReadOnly(true)` is called. If you would like the plugin to select a reader based on the instance with the least connections instead, set the following connection property. Note that this strategy is only available when internal connection pools are enabled - if you set the connection property without enabling internal pools, an exception will be thrown.
-
-```java
-props.setProperty(ReadWriteSplittingPlugin.READER_HOST_SELECTOR_STRATEGY.name, "leastConnections");
-```
+3. By default, the read-write plugin randomly selects a reader instance the first time that `setReadOnly(true)` is called. If you would like the plugin to select a reader based on a different connection strategy, please see the [Connection Strategies](#connection-strategies) section for more information.
 
 4. Continue as normal: create connections and use them as needed.
 
@@ -84,6 +80,23 @@ props.setProperty(ReadWriteSplittingPlugin.READER_HOST_SELECTOR_STRATEGY.name, "
 ### Example
 [ReadWriteSplittingPostgresExample.java](../../../examples/AWSDriverExample/src/main/java/software/amazon/ReadWriteSplittingPostgresExample.java) demonstrates how to enable and configure read-write splitting with the Aws Advanced JDBC Driver.
 
+### Connection Strategies
+By default, the read-write plugin randomly selects a reader instance the first time that `setReadOnly(true)` is called. To balance connections to reader instances more evenly, different connection strategies can be used. The following table describes the currently available connection strategies and any relevant configuration parameters for each strategy.
+
+To indicate which connection strategy to use, the `readerHostSelectorStrategy` configuration parameter can be set to one of the connection strategies in the table below. The following is an example of enabling the least connections strategy:
+
+```java
+props.setProperty(ReadWriteSplittingPlugin.READER_HOST_SELECTOR_STRATEGY.name, "leastConnections");
+```
+
+| Connection Strategy | Configuration Parameter                               | Description                                                                                                                                                                                                                                                                                                                                                                                                                                                  | Default Value |
+|---------------------|-------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------|
+| `random`            | This strategy does not have configuration parameters. | The random strategy is the default connection strategy. When switching to a reader connection, the reader instance will be chosen randomly from the available database instances.                                                                                                                                                                                                                                                                            | N/A           |
+| `leastConnections`  | This strategy does not have configuration parameters. | The least connections strategy will select reader instances based on which database instance has the least number of currently active connections. Note that this strategy is only available when internal connection pools are enabled - if you set the connection property without enabling internal pools, an exception will be thrown.                                                                                                                   | N/A           |
+| `roundRobin`        | See the following rows for configuration parameters.  | The round robin strategy will select a reader instance by taking turns with all available database instances in a cycle. A slight addition to the round robin strategy is the weighted round robin strategy, where more connections will be passed to reader instances based on user specified connection properties.                                                                                                                                        | N/A           |
+|                     | `roundRobinHostWeightPairs`                           | This parameter value must be a `string` type comma separated list of database host-weight pairs in the format `<host>:<weight>`. The host represents the database instance name, and the weight represents how many connections should be directed to the host in one cycle through all available hosts. For example, the value `instance-1:1,instance-2:4` means that for every connection to `instance-1`, there will be four connections to `instance-2`. | `null`        |
+|                     | `roundRobinDefaultWeight`                             | This parameter value must be a numeric value in the form of a `string`. This parameter represents the default weight for any hosts that have not been configured with the `roundRobinHostWeightPairs` parameter. For example, if a connection were already established and host weights were set with `roundRobinHostWeightPairs` but a new reader node was added to the database, the new reader node would use the default weight.                         | `1`           |
+
 ### Limitations
 
 #### General plugin limitations

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

@@ -64,7 +64,7 @@ boolean acceptsUrl(
    * @throws UnsupportedOperationException if the strategy is unsupported by the provider
    */
   HostSpec getHostSpecByStrategy(
-      @NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy)
+      @NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy, @NonNull Properties props)
       throws SQLException, UnsupportedOperationException;
 
   /**

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

@@ -141,14 +141,14 @@ public boolean acceptsStrategy(HostRole role, String strategy) {
    * @throws UnsupportedOperationException if the available {@link ConnectionProvider} instances do
    *                                       not support the requested strategy
    */
-  public HostSpec getHostSpecByStrategy(List<HostSpec> hosts, HostRole role, String strategy)
+  public HostSpec getHostSpecByStrategy(List<HostSpec> hosts, HostRole role, String strategy, Properties props)
       throws SQLException, UnsupportedOperationException {
     HostSpec host = null;
     if (connProvider != null) {
       connProviderLock.readLock().lock();
       try {
         if (connProvider != null && connProvider.acceptsStrategy(role, strategy)) {
-          host = connProvider.getHostSpecByStrategy(hosts, role, strategy);
+          host = connProvider.getHostSpecByStrategy(hosts, role, strategy, props);
         }
       } catch (UnsupportedOperationException e) {
         // The custom provider does not support the provided strategy, ignore it and try with the default provider.
@@ -158,7 +158,7 @@ public HostSpec getHostSpecByStrategy(List<HostSpec> hosts, HostRole role, Strin
     }
 
     if (host == null) {
-      host = defaultProvider.getHostSpecByStrategy(hosts, role, strategy);
+      host = defaultProvider.getHostSpecByStrategy(hosts, role, strategy, props);
     }
 
     return host;

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

@@ -46,7 +46,8 @@ public class DataSourceConnectionProvider implements ConnectionProvider {
   private static final Map<String, HostSelector> acceptedStrategies =
       Collections.unmodifiableMap(new HashMap<String, HostSelector>() {
         {
-          put("random", new RandomHostSelector());
+          put(RandomHostSelector.STRATEGY_RANDOM, new RandomHostSelector());
+          put(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN, new RoundRobinHostSelector());
         }
       });
   private final @NonNull DataSource dataSource;
@@ -86,7 +87,7 @@ public boolean acceptsStrategy(@NonNull HostRole role, @NonNull String strategy)
 
   @Override
   public HostSpec getHostSpecByStrategy(
-      @NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy)
+      @NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy, @NonNull Properties props)
       throws SQLException {
     if (!acceptedStrategies.containsKey(strategy)) {
       throw new UnsupportedOperationException(
@@ -95,7 +96,7 @@ public HostSpec getHostSpecByStrategy(
               new Object[] {strategy, DataSourceConnectionProvider.class}));
     }
 
-    return acceptedStrategies.get(strategy).getHost(hosts, role);
+    return acceptedStrategies.get(strategy).getHost(hosts, role, props);
   }
 
   /**

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

@@ -43,7 +43,8 @@ public class DriverConnectionProvider implements ConnectionProvider {
   private static final Map<String, HostSelector> acceptedStrategies =
       Collections.unmodifiableMap(new HashMap<String, HostSelector>() {
         {
-          put("random", new RandomHostSelector());
+          put(RandomHostSelector.STRATEGY_RANDOM, new RandomHostSelector());
+          put(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN, new RoundRobinHostSelector());
         }
       });
 
@@ -81,7 +82,7 @@ public boolean acceptsStrategy(@NonNull HostRole role, @NonNull String strategy)
 
   @Override
   public HostSpec getHostSpecByStrategy(
-      @NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy)
+      @NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy, @NonNull Properties props)
       throws SQLException {
     if (!acceptedStrategies.containsKey(strategy)) {
       throw new UnsupportedOperationException(
@@ -90,7 +91,7 @@ public HostSpec getHostSpecByStrategy(
               new Object[] {strategy, DriverConnectionProvider.class}));
     }
 
-    return acceptedStrategies.get(strategy).getHost(hosts, role);
+    return acceptedStrategies.get(strategy).getHost(hosts, role, props);
   }
 
   /**

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

@@ -21,8 +21,9 @@
 import java.sql.Connection;
 import java.sql.SQLException;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
-import java.util.Map.Entry;
+import java.util.Map;
 import java.util.Properties;
 import java.util.Set;
 import java.util.StringJoiner;
@@ -44,9 +45,12 @@ public class HikariPooledConnectionProvider implements PooledConnectionProvider,
 
   private static final Logger LOGGER =
       Logger.getLogger(HikariPooledConnectionProvider.class.getName());
-
-  private static final String LEAST_CONNECTIONS_STRATEGY = "leastConnections";
-
+  private static final Map<String, HostSelector> acceptedStrategies =
+      new HashMap<String, HostSelector>() {
+        {
+          put(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN, new RoundRobinHostSelector());
+        }
+      };
   private static final RdsUtils rdsUtils = new RdsUtils();
   private static SlidingExpirationCache<PoolKey, HikariDataSource> databasePools =
       new SlidingExpirationCache<>(
@@ -98,6 +102,8 @@ public HikariPooledConnectionProvider(
       HikariPoolConfigurator hikariPoolConfigurator, HikariPoolMapping mapping) {
     this.poolConfigurator = hikariPoolConfigurator;
     this.poolMapping = mapping;
+    final LeastConnectionsHostSelector hostSelector = new LeastConnectionsHostSelector(databasePools);
+    acceptedStrategies.put(LeastConnectionsHostSelector.STRATEGY_LEAST_CONNECTIONS, hostSelector);
   }
 
   /**
@@ -134,6 +140,8 @@ public HikariPooledConnectionProvider(
     this.poolMapping = mapping;
     poolExpirationCheckNanos = poolExpirationNanos;
     databasePools.setCleanupIntervalNanos(poolCleanupNanos);
+    final LeastConnectionsHostSelector hostSelector = new LeastConnectionsHostSelector(databasePools);
+    acceptedStrategies.put(LeastConnectionsHostSelector.STRATEGY_LEAST_CONNECTIONS, hostSelector);
   }
 
   @Override
@@ -145,44 +153,21 @@ public boolean acceptsUrl(
 
   @Override
   public boolean acceptsStrategy(@NonNull HostRole role, @NonNull String strategy) {
-    return LEAST_CONNECTIONS_STRATEGY.equals(strategy);
+    return acceptedStrategies.containsKey(strategy);
   }
 
   @Override
   public HostSpec getHostSpecByStrategy(
-      @NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy)
+      @NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy, @NonNull Properties props)
       throws SQLException {
-    if (!LEAST_CONNECTIONS_STRATEGY.equals(strategy)) {
+    if (!acceptedStrategies.containsKey(strategy)) {
       throw new UnsupportedOperationException(
           Messages.get(
               "ConnectionProvider.unsupportedHostSpecSelectorStrategy",
-              new Object[] {strategy, HikariPooledConnectionProvider.class}));
-    }
-
-    // Remove hosts with the wrong role
-    List<HostSpec> eligibleHosts = hosts.stream()
-        .filter(hostSpec -> role.equals(hostSpec.getRole()))
-        .sorted((hostSpec1, hostSpec2) ->
-            getNumConnections(hostSpec1) - getNumConnections(hostSpec2))
-        .collect(Collectors.toList());
-
-    if (eligibleHosts.size() == 0) {
-      throw new SQLException(Messages.get("HostSelector.noHostsMatchingRole", new Object[]{role}));
+              new Object[] {strategy, DataSourceConnectionProvider.class}));
     }
 
-    return eligibleHosts.get(0);
-  }
-
-  private int getNumConnections(HostSpec hostSpec) {
-    int numConnections = 0;
-    final String url = hostSpec.getUrl();
-    for (Entry<PoolKey, HikariDataSource> entry : databasePools.getEntries().entrySet()) {
-      if (!url.equals(entry.getKey().url)) {
-        continue;
-      }
-      numConnections += entry.getValue().getHikariPoolMXBean().getActiveConnections();
-    }
-    return numConnections;
+    return acceptedStrategies.get(strategy).getHost(hosts, role, props);
   }
 
   @Override
@@ -338,6 +323,10 @@ public PoolKey(final @NonNull String url, final @NonNull String extraKey) {
       this.extraKey = extraKey;
     }
 
+    public String getUrl() {
+      return this.url;
+    }
+
     @Override
     public int hashCode() {
       final int prime = 31;

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

@@ -18,17 +18,19 @@
 
 import java.sql.SQLException;
 import java.util.List;
+import java.util.Properties;
 
 public interface HostSelector {
 
   /**
    * Selects a host with the requested role from the given host list.
    *
-   * @param hosts a list of available hosts to pick from
-   * @param role  the desired host role - either a writer or a reader
+   * @param hosts a list of available hosts to pick from.
+   * @param role  the desired host role - either a writer or a reader.
+   * @param props connection properties that may be needed by the host selector in order to choose a host.
    * @return a host matching the requested role
    * @throws SQLException if the host list does not contain any hosts matching the requested role or
    *                      an error occurs while selecting a host
    */
-  HostSpec getHost(List<HostSpec> hosts, HostRole role) throws SQLException;
+  HostSpec getHost(List<HostSpec> hosts, HostRole role, Properties props) throws SQLException;
 }

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

@@ -0,0 +1,66 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * 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 software.amazon.jdbc;
+
+import com.zaxxer.hikari.HikariDataSource;
+import java.sql.SQLException;
+import java.util.List;
+import java.util.Map;
+import java.util.Properties;
+import java.util.stream.Collectors;
+import software.amazon.jdbc.util.Messages;
+import software.amazon.jdbc.util.SlidingExpirationCache;
+
+public class LeastConnectionsHostSelector implements HostSelector {
+  public static final String STRATEGY_LEAST_CONNECTIONS = "leastConnections";
+  private final SlidingExpirationCache<HikariPooledConnectionProvider.PoolKey, HikariDataSource> databasePools;
+
+  public LeastConnectionsHostSelector(
+      final SlidingExpirationCache<HikariPooledConnectionProvider.PoolKey, HikariDataSource> databasePools) {
+    this.databasePools = databasePools;
+  }
+
+  @Override
+  public HostSpec getHost(final List<HostSpec> hosts, final HostRole role, final Properties props) throws SQLException {
+    final List<HostSpec> eligibleHosts = hosts.stream()
+        .filter(hostSpec -> role.equals(hostSpec.getRole()))
+        .sorted((hostSpec1, hostSpec2) ->
+            getNumConnections(hostSpec1, databasePools) - getNumConnections(hostSpec2, databasePools))
+        .collect(Collectors.toList());
+
+    if (eligibleHosts.size() == 0) {
+      throw new SQLException(Messages.get("HostSelector.noHostsMatchingRole", new Object[]{role}));
+    }
+
+    return eligibleHosts.get(0);
+  }
+
+  private int getNumConnections(
+      final HostSpec hostSpec,
+      final SlidingExpirationCache<HikariPooledConnectionProvider.PoolKey, HikariDataSource> databasePools) {
+    int numConnections = 0;
+    final String url = hostSpec.getUrl();
+    for (final Map.Entry<HikariPooledConnectionProvider.PoolKey, HikariDataSource> entry :
+        databasePools.getEntries().entrySet()) {
+      if (!url.equals(entry.getKey().getUrl())) {
+        continue;
+      }
+      numConnections += entry.getValue().getHikariPoolMXBean().getActiveConnections();
+    }
+    return numConnections;
+  }
+}

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

@@ -162,4 +162,6 @@ HostSpec getHostSpecByStrategy(HostRole role, String strategy)
   ConnectionProvider getConnectionProvider();
 
   String getDriverProtocol();
+
+  Properties getProperties();
 }

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

@@ -553,4 +553,9 @@ public void fillAliases(Connection connection, HostSpec hostSpec) throws SQLExce
   public HostSpecBuilder getHostSpecBuilder() {
     return new HostSpecBuilder(new HostAvailabilityStrategyFactory().create(this.props));
   }
+
+  @Override
+  public Properties getProperties() {
+    return this.props;
+  }
 }