feat: weighted random host selector strategy

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

@@ -50,6 +50,7 @@ public class C3P0PooledConnectionProvider implements PooledConnectionProvider, C
           put(HighestWeightHostSelector.STRATEGY_HIGHEST_WEIGHT, new HighestWeightHostSelector());
           put(RandomHostSelector.STRATEGY_RANDOM, new RandomHostSelector());
           put(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN, new RoundRobinHostSelector());
+          put(WeightedRandomHostSelector.STRATEGY_WEIGHTED_RANDOM, new WeightedRandomHostSelector());
         }
       });
   protected static final long poolExpirationCheckNanos = TimeUnit.MINUTES.toNanos(30);

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

@@ -53,6 +53,7 @@ public class DataSourceConnectionProvider implements ConnectionProvider {
           put(HighestWeightHostSelector.STRATEGY_HIGHEST_WEIGHT, new HighestWeightHostSelector());
           put(RandomHostSelector.STRATEGY_RANDOM, new RandomHostSelector());
           put(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN, new RoundRobinHostSelector());
+          put(WeightedRandomHostSelector.STRATEGY_WEIGHTED_RANDOM, new WeightedRandomHostSelector());
         }
       });
   private final @NonNull DataSource dataSource;

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

@@ -50,6 +50,7 @@ public class DriverConnectionProvider implements ConnectionProvider {
           put(HighestWeightHostSelector.STRATEGY_HIGHEST_WEIGHT, new HighestWeightHostSelector());
           put(RandomHostSelector.STRATEGY_RANDOM, new RandomHostSelector());
           put(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN, new RoundRobinHostSelector());
+          put(WeightedRandomHostSelector.STRATEGY_WEIGHTED_RANDOM, new WeightedRandomHostSelector());
         }
       });
 

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

@@ -57,6 +57,7 @@ public class HikariPooledConnectionProvider implements PooledConnectionProvider,
           put(HighestWeightHostSelector.STRATEGY_HIGHEST_WEIGHT, new HighestWeightHostSelector());
           put(RandomHostSelector.STRATEGY_RANDOM, new RandomHostSelector());
           put(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN, new RoundRobinHostSelector());
+          put(WeightedRandomHostSelector.STRATEGY_WEIGHTED_RANDOM, new WeightedRandomHostSelector());
         }
       });
 

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

@@ -0,0 +1,169 @@
+/*
+ * 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 java.sql.SQLException;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Properties;
+import java.util.Random;
+import java.util.concurrent.locks.ReentrantLock;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import software.amazon.jdbc.hostavailability.HostAvailability;
+import software.amazon.jdbc.util.Messages;
+
+public class WeightedRandomHostSelector implements HostSelector {
+  public static final AwsWrapperProperty WEIGHTED_RANDOM_HOST_WEIGHT_PAIRS = new AwsWrapperProperty(
+      "weightedRandomHostWeightPairs", null,
+      "Comma separated list of database host-weight pairs in the format of `<host>:<weight>`.");
+  public static final String STRATEGY_WEIGHTED_RANDOM = "weightedRandom";
+  static final int DEFAULT_WEIGHT = 1;
+  static final Pattern HOST_WEIGHT_PAIRS_PATTERN =
+      Pattern.compile("((?<host>[^:/?#]*):(?<weight>[0-9]*))");
+
+  private Map<String, Integer> cachedHostWeightMap;
+  private String cachedHostWeightMapString;
+  private Random random;
+
+  private final ReentrantLock lock = new ReentrantLock();
+
+  public WeightedRandomHostSelector() {
+    this(new Random());
+  }
+
+  public WeightedRandomHostSelector(final Random random) {
+    this.random = random;
+  }
+
+  public HostSpec getHost(
+      @NonNull List<HostSpec> hosts,
+      @NonNull HostRole role,
+      @Nullable Properties props) throws SQLException {
+
+    final Map<String, Integer> hostWeightMap =
+        this.getHostWeightPairMap(WEIGHTED_RANDOM_HOST_WEIGHT_PAIRS.getString(props));
+
+    // Get and check eligible hosts
+    final List<HostSpec> eligibleHosts = hosts.stream()
+        .filter(hostSpec ->
+            role.equals(hostSpec.getRole()) && hostSpec.getAvailability().equals(HostAvailability.AVAILABLE))
+        .sorted(Comparator.comparing(HostSpec::getHost))
+        .collect(Collectors.toList());
+
+    if (eligibleHosts.isEmpty()) {
+      throw new SQLException(Messages.get("HostSelector.noHostsMatchingRole", new Object[] {role}));
+    }
+
+    final Map<String, NumberRange> hostWeightRangeMap = new HashMap<>();
+    int counter = 1;
+    for (HostSpec host : eligibleHosts) {
+      if (!hostWeightMap.containsKey(host.getHost())) {
+        continue;
+      }
+      final int hostWeight = hostWeightMap.get(host.getHost());
+      if (hostWeight > 0) {
+        final int rangeStart = counter;
+        final int rangeEnd = counter + hostWeight - 1;
+        hostWeightRangeMap.put(host.getHost(), new NumberRange(rangeStart, rangeEnd));
+        counter = counter + hostWeight;
+      } else {
+        hostWeightRangeMap.put(host.getHost(), new NumberRange(counter, counter));
+        counter++;
+      }
+    }
+
+    if (this.random == null) {
+      this.random = new Random();
+    }
+    int randomInt = this.random.nextInt(counter);
+
+    // Check random number is in host weight range map
+    for (final HostSpec host : eligibleHosts) {
+      NumberRange range = hostWeightRangeMap.get(host.getHost());
+      if (range != null && range.isInRange(randomInt)) {
+        return host;
+      }
+    }
+
+    throw new SQLException(Messages.get("HostSelector.weightedRandomUnableToGetHost", new Object[] {role}));
+  }
+
+  private Map<String, Integer> getHostWeightPairMap(final String hostWeightMapString) throws SQLException {
+    try {
+      lock.lock();
+      if (this.cachedHostWeightMapString != null
+          && this.cachedHostWeightMapString.trim().equals(hostWeightMapString.trim())
+          && this.cachedHostWeightMap != null
+          && !this.cachedHostWeightMap.isEmpty()) {
+        return this.cachedHostWeightMap;
+      }
+
+      final Map<String, Integer> hostWeightMap = new HashMap<>();
+      if (hostWeightMapString == null || hostWeightMapString.trim().isEmpty()) {
+        return hostWeightMap;
+      }
+      final String[] hostWeightPairs = hostWeightMapString.split(",");
+      for (final String hostWeightPair : hostWeightPairs) {
+        final Matcher matcher = HOST_WEIGHT_PAIRS_PATTERN.matcher(hostWeightPair);
+        if (!matcher.matches()) {
+          throw new SQLException(Messages.get("HostSelector.weightedRandomInvalidHostWeightPairs"));
+        }
+
+        final String hostName = matcher.group("host").trim();
+        final String hostWeight = matcher.group("weight").trim();
+        if (hostName.isEmpty() || hostWeight.isEmpty()) {
+          throw new SQLException(Messages.get("HostSelector.weightedRandomInvalidHostWeightPairs"));
+        }
+
+        try {
+          final int weight = Integer.parseInt(hostWeight);
+          if (weight < DEFAULT_WEIGHT) {
+            throw new SQLException(Messages.get("HostSelector.weightedRandomInvalidHostWeightPairs"));
+          }
+          hostWeightMap.put(hostName, weight);
+        } catch (NumberFormatException e) {
+          throw new SQLException(Messages.get("HostSelector.roundRobinInvalidHostWeightPairs"));
+        }
+      }
+      this.cachedHostWeightMap = hostWeightMap;
+      this.cachedHostWeightMapString = hostWeightMapString;
+      return hostWeightMap;
+    } finally {
+      lock.unlock();
+    }
+  }
+
+  private static class NumberRange {
+    private int start;
+    private int end;
+
+    public NumberRange(int start, int end) {
+      this.start = start;
+      this.end = end;
+    }
+
+    public boolean isInRange(int value) {
+      return start <= value && value <= end;
+    }
+  }
+}

wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessQueryHelper.java

@@ -96,7 +96,7 @@ protected HostSpec createHost(final ResultSet resultSet, final int hostPortToMap
     final String hostName = resultSet.getString(1);
     final float cpu = resultSet.getFloat(2);
 
-    long weight = Math.round(10 - cpu * 10);
+    long weight = (long) (10 - Math.floor(10 * cpu));
 
     if (weight < 1 || weight > 10) {
       weight = 1; // default to 1

wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImpl.java

@@ -34,9 +34,10 @@
 import software.amazon.jdbc.HostSpec;
 import software.amazon.jdbc.PluginService;
 import software.amazon.jdbc.PropertyDefinition;
-import software.amazon.jdbc.RoundRobinHostSelector;
+import software.amazon.jdbc.WeightedRandomHostSelector;
 import software.amazon.jdbc.hostavailability.HostAvailability;
 import software.amazon.jdbc.util.FullServicesContainer;
+import software.amazon.jdbc.util.HostSelectorUtils;
 import software.amazon.jdbc.util.Messages;
 import software.amazon.jdbc.util.Utils;
 import software.amazon.jdbc.util.monitoring.MonitorErrorResponse;
@@ -126,15 +127,15 @@ public void establishConnection(final LimitlessConnectionContext context) throws
       return;
     }
 
-    RoundRobinHostSelector.setRoundRobinHostWeightPairsProperty(
+    HostSelectorUtils.setHostWeightPairsProperty(WeightedRandomHostSelector.WEIGHTED_RANDOM_HOST_WEIGHT_PAIRS,
         context.getProps(),
         context.getLimitlessRouters());
     HostSpec selectedHostSpec;
     try {
       selectedHostSpec = this.pluginService.getHostSpecByStrategy(
           context.getLimitlessRouters(),
           HostRole.WRITER,
-          RoundRobinHostSelector.STRATEGY_ROUND_ROBIN);
+          WeightedRandomHostSelector.STRATEGY_WEIGHTED_RANDOM);
       LOGGER.fine(Messages.get(
           "LimitlessRouterServiceImpl.selectedHost",
           new Object[] {selectedHostSpec != null ? selectedHostSpec.getHost() : "null"}));

wrapper/src/main/java/software/amazon/jdbc/util/HostSelectorUtils.java

@@ -0,0 +1,43 @@
+/*
+ * 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.util;
+
+import java.util.List;
+import java.util.Properties;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import software.amazon.jdbc.AwsWrapperProperty;
+import software.amazon.jdbc.HostSpec;
+
+public class HostSelectorUtils {
+  public static void setHostWeightPairsProperty(
+      final @NonNull AwsWrapperProperty property,
+      final @NonNull Properties properties,
+      final @NonNull List<HostSpec> hosts) {
+    final StringBuilder builder = new StringBuilder();
+    for (int i = 0; i < hosts.size(); i++) {
+      builder
+          .append(hosts.get(i).getHostId())
+          .append(":")
+          .append(hosts.get(i).getWeight());
+      if (i < hosts.size() - 1) {
+        builder.append(",");
+      }
+    }
+    final String hostWeightPairsString = builder.toString();
+    properties.setProperty(property.name, hostWeightPairsString);
+  }
+}

wrapper/src/main/resources/aws_advanced_jdbc_wrapper_messages.properties

@@ -193,6 +193,8 @@ HostMonitoringConnectionPlugin.unableToIdentifyConnection=Unable to identify the
 HostSelector.noHostsMatchingRole=No hosts were found matching the requested ''{0}'' role.
 HostSelector.roundRobinInvalidHostWeightPairs=The provided host weight pairs have not been configured correctly. Please ensure the provided host weight pairs is a comma separated list of pairs, each pair in the format of <host>:<weight>. Weight values must be an integer greater than or equal to the default weight value of 1.
 HostSelector.roundRobinInvalidDefaultWeight=The provided default weight value is not valid. Weight values must be an integer greater than or equal to the default weight value of 1.
+HostSelector.weightedRandomInvalidHostWeightPairs=The provided host weight pairs have not been configured correctly. Please ensure the provided host weight pairs is a comma separated list of pairs, each pair in the format of <host>:<weight>. Weight values must be an integer greater than or equal to the default weight value of 1.
+HostSelector.weightedRandomUnableToGetHost=Weighted Random strategy was unable to select a host.
 
 IamAuthConnectionPlugin.unhandledException=Unhandled exception: ''{0}''
 IamAuthConnectionPlugin.connectException=Error occurred while opening a connection: ''{0}''

wrapper/src/test/java/software/amazon/jdbc/RoundRobinHostSelectorTest.java

@@ -32,6 +32,7 @@
 import software.amazon.jdbc.HostSpecBuilder;
 import software.amazon.jdbc.RoundRobinHostSelector;
 import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy;
+import software.amazon.jdbc.util.HostSelectorUtils;
 
 public class RoundRobinHostSelectorTest {
   private static final int TEST_PORT = 5432;
@@ -524,34 +525,4 @@ void testGetHost_AllHostsChanged() throws SQLException {
         readerHostSpec4.getHost(),
         roundRobinHostSelector.getHost(hostsList14, HostRole.READER, defaultProps).getHost());
   }
-
-  @Test
-  void testSetRoundRobinHostWeightPairsProperty() {
-    final String expectedPropertyValue = "instance-1-id:2,instance-2-id:1,instance-3-id:0";
-
-    final List<HostSpec> hosts = Arrays.asList(
-        new HostSpecBuilder(new SimpleHostAvailabilityStrategy())
-            .host("instance-1")
-            .hostId("instance-1-id")
-            .weight(2)
-            .build(),
-        new HostSpecBuilder(new SimpleHostAvailabilityStrategy())
-            .host("instance-2")
-            .hostId("instance-2-id")
-            .weight(1)
-            .build(),
-        new HostSpecBuilder(new SimpleHostAvailabilityStrategy())
-            .host("instance-3")
-            .hostId("instance-3-id")
-            .weight(0)
-            .build()
-    );
-    final Properties properties = new Properties();
-    RoundRobinHostSelector.setRoundRobinHostWeightPairsProperty(properties, hosts);
-
-    final String actualPropertyValue = properties.getProperty(
-        RoundRobinHostSelector.ROUND_ROBIN_HOST_WEIGHT_PAIRS.name);
-
-    assertEquals(expectedPropertyValue, actualPropertyValue);
-  }
 }