feat: round robin host selection strategy

crystall-bitquill · crystall-bitquill · commit 3847b5c46cdd · 2023-08-18T20:27:00.000-07:00
diff --git a/docs/using-the-jdbc-driver/using-plugins/UsingTheReadWriteSplittingPlugin.md b/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           |
+|                     | `roundRobinHostWeights`                               | 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
diff --git a/wrapper/src/main/java/software/amazon/jdbc/DataSourceConnectionProvider.java b/wrapper/src/main/java/software/amazon/jdbc/DataSourceConnectionProvider.java
@@ -29,6 +29,7 @@
 import org.checkerframework.checker.nullness.qual.NonNull;
 import software.amazon.jdbc.dialect.Dialect;
 import software.amazon.jdbc.exceptions.SQLLoginException;
+import software.amazon.jdbc.plugin.readwritesplitting.ReadWriteSplittingPlugin;
 import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect;
 import software.amazon.jdbc.util.Messages;
 import software.amazon.jdbc.util.PropertyUtils;
@@ -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;
@@ -114,6 +116,14 @@ public Connection connect(
       final @NonNull HostSpec hostSpec,
       final @NonNull Properties props)
       throws SQLException {
+    final String strategy = ReadWriteSplittingPlugin.READER_HOST_SELECTOR_STRATEGY.getString(props);
+    if (RoundRobinHostSelector.STRATEGY_ROUND_ROBIN.equals(strategy)) {
+      final RoundRobinHostSelector roundRobinHostSelector =
+          (RoundRobinHostSelector) acceptedStrategies.get(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN);
+      if (roundRobinHostSelector.getHostCacheEntry(hostSpec) != null) {
+        roundRobinHostSelector.updateCachePropertiesForHost(hostSpec, props);
+      }
+    }
 
     final Properties copy = PropertyUtils.copyProperties(props);
 
diff --git a/wrapper/src/main/java/software/amazon/jdbc/DriverConnectionProvider.java b/wrapper/src/main/java/software/amazon/jdbc/DriverConnectionProvider.java
@@ -27,6 +27,7 @@
 import org.checkerframework.checker.nullness.qual.NonNull;
 import software.amazon.jdbc.dialect.Dialect;
 import software.amazon.jdbc.exceptions.SQLLoginException;
+import software.amazon.jdbc.plugin.readwritesplitting.ReadWriteSplittingPlugin;
 import software.amazon.jdbc.targetdriverdialect.ConnectInfo;
 import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect;
 import software.amazon.jdbc.util.Messages;
@@ -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());
         }
       });
 
@@ -109,6 +111,14 @@ public Connection connect(
       final @NonNull HostSpec hostSpec,
       final @NonNull Properties props)
       throws SQLException {
+    final String strategy = ReadWriteSplittingPlugin.READER_HOST_SELECTOR_STRATEGY.getString(props);
+    if (RoundRobinHostSelector.STRATEGY_ROUND_ROBIN.equals(strategy)) {
+      final RoundRobinHostSelector roundRobinHostSelector =
+          (RoundRobinHostSelector) acceptedStrategies.get(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN);
+      if (roundRobinHostSelector.getHostCacheEntry(hostSpec) != null) {
+        roundRobinHostSelector.updateCachePropertiesForHost(hostSpec, props);
+      }
+    }
 
     final Properties copy = PropertyUtils.copyProperties(props);
     final ConnectInfo connectInfo =
diff --git a/wrapper/src/main/java/software/amazon/jdbc/HikariPooledConnectionProvider.java b/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;
@@ -32,6 +33,7 @@
 import org.checkerframework.checker.nullness.qual.NonNull;
 import software.amazon.jdbc.cleanup.CanReleaseResources;
 import software.amazon.jdbc.dialect.Dialect;
+import software.amazon.jdbc.plugin.readwritesplitting.ReadWriteSplittingPlugin;
 import software.amazon.jdbc.util.HikariCPSQLException;
 import software.amazon.jdbc.util.Messages;
 import software.amazon.jdbc.util.RdsUrlType;
@@ -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(LeastConnectionsHostSelector.STRATEGY_LEAST_CONNECTIONS, new LeastConnectionsHostSelector());
+          put(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN, new RoundRobinHostSelector());
+        }
+      });
   private static final RdsUtils rdsUtils = new RdsUtils();
   private static SlidingExpirationCache<PoolKey, HikariDataSource> databasePools =
       new SlidingExpirationCache<>(
@@ -98,6 +104,10 @@ public HikariPooledConnectionProvider(
       HikariPoolConfigurator hikariPoolConfigurator, HikariPoolMapping mapping) {
     this.poolConfigurator = hikariPoolConfigurator;
     this.poolMapping = mapping;
+    final LeastConnectionsHostSelector hostSelector =
+        (LeastConnectionsHostSelector) acceptedStrategies
+            .get(LeastConnectionsHostSelector.STRATEGY_LEAST_CONNECTIONS);
+    hostSelector.setDatabasePools(databasePools);
   }
 
   /**
@@ -134,6 +144,10 @@ public HikariPooledConnectionProvider(
     this.poolMapping = mapping;
     poolExpirationCheckNanos = poolExpirationNanos;
     databasePools.setCleanupIntervalNanos(poolCleanupNanos);
+    final LeastConnectionsHostSelector hostSelector =
+            (LeastConnectionsHostSelector) acceptedStrategies
+                .get(LeastConnectionsHostSelector.STRATEGY_LEAST_CONNECTIONS);
+    hostSelector.setDatabasePools(databasePools);
   }
 
   @Override
@@ -145,44 +159,32 @@ public boolean acceptsUrl(
 
   @Override
   public boolean acceptsStrategy(@NonNull HostRole role, @NonNull String strategy) {
-    return LEAST_CONNECTIONS_STRATEGY.equals(strategy);
+    return 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)) {
-      throw new UnsupportedOperationException(
-          Messages.get(
-              "ConnectionProvider.unsupportedHostSpecSelectorStrategy",
-              new Object[] {strategy, HikariPooledConnectionProvider.class}));
+    final HostSpec selectedHost;
+    switch (strategy) {
+      case LeastConnectionsHostSelector.STRATEGY_LEAST_CONNECTIONS:
+        selectedHost =
+            acceptedStrategies.get(LeastConnectionsHostSelector.STRATEGY_LEAST_CONNECTIONS).getHost(hosts, role);
+        break;
+
+      case RoundRobinHostSelector.STRATEGY_ROUND_ROBIN:
+        selectedHost = acceptedStrategies.get(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN).getHost(hosts, role);
+        break;
+
+      default:
+        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}));
-    }
-
-    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 selectedHost;
   }
 
   @Override
@@ -201,6 +203,15 @@ public Connection connect(
 
     ds.setPassword(props.getProperty(PropertyDefinition.PASSWORD.name));
 
+    final String strategy = ReadWriteSplittingPlugin.READER_HOST_SELECTOR_STRATEGY.getString(props);
+    if (RoundRobinHostSelector.STRATEGY_ROUND_ROBIN.equals(strategy)) {
+      final RoundRobinHostSelector roundRobinHostSelector =
+          (RoundRobinHostSelector) acceptedStrategies.get(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN);
+      if (roundRobinHostSelector.getHostCacheEntry(hostSpec) != null) {
+        roundRobinHostSelector.updateCachePropertiesForHost(hostSpec, props);
+      }
+    }
+
     return ds.getConnection();
   }
 
@@ -330,7 +341,7 @@ HikariDataSource createHikariDataSource(String protocol, HostSpec hostSpec, Prop
   }
 
   public static class PoolKey {
-    private final @NonNull String url;
+    final @NonNull String url;
     private final @NonNull String extraKey;
 
     public PoolKey(final @NonNull String url, final @NonNull String extraKey) {
diff --git a/wrapper/src/main/java/software/amazon/jdbc/LeastConnectionsHostSelector.java b/wrapper/src/main/java/software/amazon/jdbc/LeastConnectionsHostSelector.java
@@ -0,0 +1,64 @@
+/*
+ * 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.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 SlidingExpirationCache<HikariPooledConnectionProvider.PoolKey, HikariDataSource> databasePools;
+
+  public void setDatabasePools(
+      SlidingExpirationCache<HikariPooledConnectionProvider.PoolKey, HikariDataSource> databasePools) {
+    this.databasePools = databasePools;
+  }
+
+  public HostSpec getHost(List<HostSpec> hosts, HostRole role) throws SQLException {
+    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(
+      HostSpec hostSpec,
+      SlidingExpirationCache<HikariPooledConnectionProvider.PoolKey, HikariDataSource> databasePools) {
+    int numConnections = 0;
+    final String url = hostSpec.getUrl();
+    for (Map.Entry<HikariPooledConnectionProvider.PoolKey, HikariDataSource> entry :
+        databasePools.getEntries().entrySet()) {
+      if (!url.equals(entry.getKey().url)) {
+        continue;
+      }
+      numConnections += entry.getValue().getHikariPoolMXBean().getActiveConnections();
+    }
+    return numConnections;
+  }
+}
diff --git a/wrapper/src/main/java/software/amazon/jdbc/RandomHostSelector.java b/wrapper/src/main/java/software/amazon/jdbc/RandomHostSelector.java
@@ -24,12 +24,14 @@
 
 public class RandomHostSelector implements HostSelector {
 
+  public static final String STRATEGY_RANDOM = "random";
+
   @Override
   public HostSpec getHost(List<HostSpec> hosts, HostRole role) throws SQLException {
     List<HostSpec> eligibleHosts = hosts.stream()
         .filter(hostSpec -> role.equals(hostSpec.getRole())).collect(Collectors.toList());
     if (eligibleHosts.size() == 0) {
-      throw new SQLException(Messages.get("RandomHostSelector.noHostsMatchingRole", new Object[]{role}));
+      throw new SQLException(Messages.get("HostSelector.noHostsMatchingRole", new Object[]{role}));
     }
 
     int randomIndex = new Random().nextInt(eligibleHosts.size());
diff --git a/wrapper/src/main/java/software/amazon/jdbc/RoundRobinHostSelector.java b/wrapper/src/main/java/software/amazon/jdbc/RoundRobinHostSelector.java
diff --git a/wrapper/src/main/resources/messages.properties b/wrapper/src/main/resources/messages.properties
diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/RoundRobinHostSelectorTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/RoundRobinHostSelectorTest.java
