feat: round robin host selection strategy (#603)

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`. <br><br> **Note:** The `<weight>` value in the string must be an integer greater than or equal to 1. | `null`        |
+|                     | `roundRobinDefaultWeight`                             | This parameter value must be an integer 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. <br><br> **Note:** This value must be an integer greater than or equal to 1.                                                | `1`           |
+
 ### Limitations
 
 #### General plugin limitations

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

@@ -21,6 +21,7 @@
 import java.util.List;
 import java.util.Properties;
 import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import software.amazon.jdbc.dialect.Dialect;
 
 /**
@@ -59,12 +60,13 @@ boolean acceptsUrl(
    * @param role     determines if the connection provider should return a writer or a reader
    * @param strategy the strategy determining how the {@link HostSpec} should be selected, e.g.,
    *                 random or round-robin
+   * @param props    any properties that are required by the provided strategy to select a host
    * @return the {@link HostSpec} selected using the specified strategy
    * @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)
+      @NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy, @Nullable Properties props)
       throws SQLException, UnsupportedOperationException;
 
   /**

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

@@ -134,21 +134,22 @@ public boolean acceptsStrategy(HostRole role, String strategy) {
    * @param role     the desired role of the host - either a writer or a reader
    * @param strategy the strategy that should be used to select a {@link HostSpec} from the host
    *                 list (eg "random")
+   * @param props    any properties that are required by the provided strategy to select a host
    * @return a {@link HostSpec} with the requested role
    * @throws SQLException                  if the available {@link ConnectionProvider} instances
    *                                       cannot find a host in the host list matching the
    *                                       requested role or an error occurs while selecting a host
    * @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 +159,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

@@ -27,6 +27,7 @@
 import java.util.logging.Logger;
 import javax.sql.DataSource;
 import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import software.amazon.jdbc.dialect.Dialect;
 import software.amazon.jdbc.exceptions.SQLLoginException;
 import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect;
@@ -46,7 +47,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 +88,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, @Nullable Properties props)
       throws SQLException {
     if (!acceptedStrategies.containsKey(strategy)) {
       throw new UnsupportedOperationException(
@@ -95,7 +97,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

@@ -25,6 +25,7 @@
 import java.util.Properties;
 import java.util.logging.Logger;
 import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import software.amazon.jdbc.dialect.Dialect;
 import software.amazon.jdbc.exceptions.SQLLoginException;
 import software.amazon.jdbc.targetdriverdialect.ConnectInfo;
@@ -43,7 +44,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 +83,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, @Nullable Properties props)
       throws SQLException {
     if (!acceptedStrategies.containsKey(strategy)) {
       throw new UnsupportedOperationException(
@@ -90,7 +92,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,15 +21,17 @@
 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;
 import java.util.concurrent.TimeUnit;
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
 import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import software.amazon.jdbc.cleanup.CanReleaseResources;
 import software.amazon.jdbc.dialect.Dialect;
 import software.amazon.jdbc.util.HikariCPSQLException;
@@ -44,9 +46,13 @@ 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 =
+      Collections.unmodifiableMap(new HashMap<String, HostSelector>() {
+        {
+          put(RandomHostSelector.STRATEGY_RANDOM, new RandomHostSelector());
+          put(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN, new RoundRobinHostSelector());
+        }
+      });
   private static final RdsUtils rdsUtils = new RdsUtils();
   private static SlidingExpirationCache<PoolKey, HikariDataSource> databasePools =
       new SlidingExpirationCache<>(
@@ -56,6 +62,7 @@ public class HikariPooledConnectionProvider implements PooledConnectionProvider,
   private static long poolExpirationCheckNanos = TimeUnit.MINUTES.toNanos(30);
   private final HikariPoolConfigurator poolConfigurator;
   private final HikariPoolMapping poolMapping;
+  private final LeastConnectionsHostSelector leastConnectionsHostSelector;
 
   /**
    * {@link HikariPooledConnectionProvider} constructor. This class can be passed to
@@ -98,6 +105,7 @@ public HikariPooledConnectionProvider(
       HikariPoolConfigurator hikariPoolConfigurator, HikariPoolMapping mapping) {
     this.poolConfigurator = hikariPoolConfigurator;
     this.poolMapping = mapping;
+    this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(databasePools);
   }
 
   /**
@@ -134,6 +142,7 @@ public HikariPooledConnectionProvider(
     this.poolMapping = mapping;
     poolExpirationCheckNanos = poolExpirationNanos;
     databasePools.setCleanupIntervalNanos(poolCleanupNanos);
+    this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(databasePools);
   }
 
   @Override
@@ -145,44 +154,28 @@ public boolean acceptsUrl(
 
   @Override
   public boolean acceptsStrategy(@NonNull HostRole role, @NonNull String strategy) {
-    return LEAST_CONNECTIONS_STRATEGY.equals(strategy);
+    return acceptedStrategies.containsKey(strategy)
+        || LeastConnectionsHostSelector.STRATEGY_LEAST_CONNECTIONS.equals(strategy);
   }
 
   @Override
   public HostSpec getHostSpecByStrategy(
-      @NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy)
-      throws SQLException {
-    if (!LEAST_CONNECTIONS_STRATEGY.equals(strategy)) {
+      @NonNull List<HostSpec> hosts,
+      @NonNull HostRole role,
+      @NonNull String strategy,
+      @Nullable Properties props) throws SQLException {
+    if (!acceptsStrategy(role, 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();
+    if (LeastConnectionsHostSelector.STRATEGY_LEAST_CONNECTIONS.equals(strategy)) {
+      return this.leastConnectionsHostSelector.getHost(hosts, role, props);
+    } else {
+      return acceptedStrategies.get(strategy).getHost(hosts, role, props);
     }
-    return numConnections;
   }
 
   @Override
@@ -338,6 +331,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,22 @@
 
 import java.sql.SQLException;
 import java.util.List;
+import java.util.Properties;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 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(
+      @NonNull List<HostSpec> hosts, @NonNull HostRole role, @Nullable Properties props) throws SQLException;
 }