diff --git a/.DS_Store b/.DS_Store new file mode 100644 index 000000000..5008ddfcf Binary files /dev/null and b/.DS_Store differ diff --git a/wrapper/src/main/java/software/amazon/jdbc/C3P0PooledConnectionProvider.java b/wrapper/src/main/java/software/amazon/jdbc/C3P0PooledConnectionProvider.java index 07b8570fd..dc4aa70eb 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/C3P0PooledConnectionProvider.java +++ b/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); diff --git a/wrapper/src/main/java/software/amazon/jdbc/DataSourceConnectionProvider.java b/wrapper/src/main/java/software/amazon/jdbc/DataSourceConnectionProvider.java index 0231739d0..6c5afcda3 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/DataSourceConnectionProvider.java +++ b/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; diff --git a/wrapper/src/main/java/software/amazon/jdbc/DriverConnectionProvider.java b/wrapper/src/main/java/software/amazon/jdbc/DriverConnectionProvider.java index 135eabc4b..e2a04e399 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/DriverConnectionProvider.java +++ b/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()); } }); diff --git a/wrapper/src/main/java/software/amazon/jdbc/HikariPooledConnectionProvider.java b/wrapper/src/main/java/software/amazon/jdbc/HikariPooledConnectionProvider.java index 47e1968cc..f0d58c841 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/HikariPooledConnectionProvider.java +++ b/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()); } }); diff --git a/wrapper/src/main/java/software/amazon/jdbc/WeightedRandomHostSelector.java b/wrapper/src/main/java/software/amazon/jdbc/WeightedRandomHostSelector.java new file mode 100644 index 000000000..8ca372ca2 --- /dev/null +++ b/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 `:`."); + public static final String STRATEGY_WEIGHTED_RANDOM = "weightedRandom"; + static final int DEFAULT_WEIGHT = 1; + static final Pattern HOST_WEIGHT_PAIRS_PATTERN = + Pattern.compile("((?[^:/?#]*):(?[0-9]*))"); + + private Map 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 hosts, + @NonNull HostRole role, + @Nullable Properties props) throws SQLException { + + final Map hostWeightMap = + this.getHostWeightPairMap(WEIGHTED_RANDOM_HOST_WEIGHT_PAIRS.getString(props)); + + // Get and check eligible hosts + final List 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 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 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 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; + } + } +} diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessQueryHelper.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessQueryHelper.java index 74ff24291..9867201e6 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessQueryHelper.java +++ b/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 diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImpl.java index 14fee1c67..7b904c57d 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImpl.java +++ b/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,7 +127,7 @@ public void establishConnection(final LimitlessConnectionContext context) throws return; } - RoundRobinHostSelector.setRoundRobinHostWeightPairsProperty( + HostSelectorUtils.setHostWeightPairsProperty(WeightedRandomHostSelector.WEIGHTED_RANDOM_HOST_WEIGHT_PAIRS, context.getProps(), context.getLimitlessRouters()); HostSpec selectedHostSpec; @@ -134,7 +135,7 @@ public void establishConnection(final LimitlessConnectionContext context) throws 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"})); diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/HostSelectorUtils.java b/wrapper/src/main/java/software/amazon/jdbc/util/HostSelectorUtils.java new file mode 100644 index 000000000..adb989705 --- /dev/null +++ b/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 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); + } +} diff --git a/wrapper/src/main/resources/aws_advanced_jdbc_wrapper_messages.properties b/wrapper/src/main/resources/aws_advanced_jdbc_wrapper_messages.properties index 1776fcced..94f361e5f 100644 --- a/wrapper/src/main/resources/aws_advanced_jdbc_wrapper_messages.properties +++ b/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 :. 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 :. 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}'' diff --git a/wrapper/src/test/java/software/amazon/jdbc/RoundRobinHostSelectorTest.java b/wrapper/src/test/java/software/amazon/jdbc/RoundRobinHostSelectorTest.java index f0ea5a473..7a523a3e7 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/RoundRobinHostSelectorTest.java +++ b/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 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); - } } diff --git a/wrapper/src/test/java/software/amazon/jdbc/WeightedRandomHostSelectorTests.java b/wrapper/src/test/java/software/amazon/jdbc/WeightedRandomHostSelectorTests.java new file mode 100644 index 000000000..65609455d --- /dev/null +++ b/wrapper/src/test/java/software/amazon/jdbc/WeightedRandomHostSelectorTests.java @@ -0,0 +1,203 @@ +/* + * 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 static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.Mockito.when; +import static software.amazon.jdbc.WeightedRandomHostSelector.WEIGHTED_RANDOM_HOST_WEIGHT_PAIRS; + +import java.sql.SQLException; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Properties; +import java.util.Random; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import software.amazon.jdbc.hostavailability.HostAvailability; +import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; + +class WeightedRandomHostSelectorTests { + + @Mock Random mockRandom; + private AutoCloseable closeable; + + @BeforeEach + void init() { + closeable = MockitoAnnotations.openMocks(this); + } + + @AfterEach + void cleanUp() throws Exception { + closeable.close(); + } + + @Test + void testGetHost_emptyHostList() { + final HostSelector hostSelector = new WeightedRandomHostSelector(); + final Properties props = new Properties(); + final List emptyHostList = Collections.emptyList(); + assertThrows(SQLException.class, () -> hostSelector.getHost(emptyHostList, HostRole.WRITER, props)); + } + + @Test + void testGetHost_noEligibleHosts() { + final HostSelector hostSelector = new WeightedRandomHostSelector(); + final Properties props = new Properties(); + final List noEligibleHostsList = Arrays.asList( + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-1").role(HostRole.READER).build(), + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-2").role(HostRole.WRITER) + .availability(HostAvailability.NOT_AVAILABLE).build(), + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-3").role(HostRole.READER).build() + ); + assertThrows(SQLException.class, + () -> hostSelector.getHost(noEligibleHostsList, HostRole.WRITER, props)); + } + + @Test + void testGetHost_invalidWeight() { + final HostSelector hostSelector = new WeightedRandomHostSelector(); + final Properties props = new Properties(); + props.setProperty(WEIGHTED_RANDOM_HOST_WEIGHT_PAIRS.name, "instance-1:3,instance-2:2,instance-3:0"); + final List eligibleHostsList = Arrays.asList( + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-1").role(HostRole.WRITER) + .availability(HostAvailability.AVAILABLE).build(), + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-2").role(HostRole.WRITER) + .availability(HostAvailability.AVAILABLE).build(), + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-3").role(HostRole.WRITER) + .availability(HostAvailability.AVAILABLE).build() + ); + assertThrows(SQLException.class, + () -> hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props)); + } + + @Test + void testGetHost_invalidProps() { + final HostSelector hostSelector = new WeightedRandomHostSelector(); + final Properties props = new Properties(); + props.setProperty(WEIGHTED_RANDOM_HOST_WEIGHT_PAIRS.name, "someInvalidString"); + final List eligibleHostsList = Arrays.asList( + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-1").role(HostRole.WRITER) + .availability(HostAvailability.AVAILABLE).build(), + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-2").role(HostRole.WRITER) + .availability(HostAvailability.AVAILABLE).build(), + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-3").role(HostRole.WRITER) + .availability(HostAvailability.AVAILABLE).build() + ); + assertThrows(SQLException.class, + () -> hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props)); + } + + @Test + void testGetHost() throws SQLException { + final WeightedRandomHostSelector hostSelector = new WeightedRandomHostSelector(mockRandom); + final Properties props = new Properties(); + props.setProperty(WEIGHTED_RANDOM_HOST_WEIGHT_PAIRS.name, "instance-1:3,instance-2:2,instance-3:01"); + final List eligibleHostsList = Arrays.asList( + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-1").role(HostRole.WRITER) + .availability(HostAvailability.AVAILABLE).build(), + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-2").role(HostRole.WRITER) + .availability(HostAvailability.AVAILABLE).build(), + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-3").role(HostRole.WRITER) + .availability(HostAvailability.AVAILABLE).build() + ); + + when(mockRandom.nextInt(anyInt())).thenReturn(1, 2, 3, 4, 5, 6); + + final HostSpec actualHost1 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props); + assertEquals(eligibleHostsList.get(0).getHost(), actualHost1.getHost()); + + final HostSpec actualHost2 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props); + assertEquals(eligibleHostsList.get(0).getHost(), actualHost2.getHost()); + + final HostSpec actualHost3 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props); + assertEquals(eligibleHostsList.get(0).getHost(), actualHost3.getHost()); + + final HostSpec actualHost4 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props); + assertEquals(eligibleHostsList.get(1).getHost(), actualHost4.getHost()); + + final HostSpec actualHost5 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props); + assertEquals(eligibleHostsList.get(1).getHost(), actualHost5.getHost()); + + final HostSpec actualHost6 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props); + assertEquals(eligibleHostsList.get(2).getHost(), actualHost6.getHost()); + } + + @Test + void testGetHost_changeWeights() throws SQLException { + final WeightedRandomHostSelector hostSelector = new WeightedRandomHostSelector(mockRandom); + final Properties props = new Properties(); + + props.setProperty(WEIGHTED_RANDOM_HOST_WEIGHT_PAIRS.name, "instance-1:3,instance-2:2,instance-3:01"); + final List eligibleHostsList = Arrays.asList( + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-1").role(HostRole.WRITER) + .availability(HostAvailability.AVAILABLE).build(), + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-2").role(HostRole.WRITER) + .availability(HostAvailability.AVAILABLE).build(), + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-3").role(HostRole.WRITER) + .availability(HostAvailability.AVAILABLE).build() + ); + + when(mockRandom.nextInt(anyInt())).thenReturn(1, 2, 3, 4, 5, 6, 1, 2, 3, 4, 5, 6, 7); + + final HostSpec actualHost1 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props); + assertEquals(eligibleHostsList.get(0).getHost(), actualHost1.getHost()); + + final HostSpec actualHost2 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props); + assertEquals(eligibleHostsList.get(0).getHost(), actualHost2.getHost()); + + final HostSpec actualHost3 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props); + assertEquals(eligibleHostsList.get(0).getHost(), actualHost3.getHost()); + + final HostSpec actualHost4 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props); + assertEquals(eligibleHostsList.get(1).getHost(), actualHost4.getHost()); + + final HostSpec actualHost5 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props); + assertEquals(eligibleHostsList.get(1).getHost(), actualHost5.getHost()); + + final HostSpec actualHost6 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props); + assertEquals(eligibleHostsList.get(2).getHost(), actualHost6.getHost()); + + props.setProperty(WEIGHTED_RANDOM_HOST_WEIGHT_PAIRS.name, "instance-1:1,instance-2:4,instance-3:2"); + + final HostSpec actualHost7 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props); + assertEquals(eligibleHostsList.get(0).getHost(), actualHost7.getHost()); + + final HostSpec actualHost8 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props); + assertEquals(eligibleHostsList.get(1).getHost(), actualHost8.getHost()); + + final HostSpec actualHost9 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props); + assertEquals(eligibleHostsList.get(1).getHost(), actualHost9.getHost()); + + final HostSpec actualHost10 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props); + assertEquals(eligibleHostsList.get(1).getHost(), actualHost10.getHost()); + + final HostSpec actualHost11 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props); + assertEquals(eligibleHostsList.get(1).getHost(), actualHost11.getHost()); + + final HostSpec actualHost12 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props); + assertEquals(eligibleHostsList.get(2).getHost(), actualHost12.getHost()); + + final HostSpec actualHost13 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props); + assertEquals(eligibleHostsList.get(2).getHost(), actualHost13.getHost()); + } +} diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImplTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImplTest.java index 4412e2368..b37df47a8 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImplTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImplTest.java @@ -44,7 +44,7 @@ import software.amazon.jdbc.HostSpecBuilder; import software.amazon.jdbc.JdbcCallable; import software.amazon.jdbc.PluginService; -import software.amazon.jdbc.RoundRobinHostSelector; +import software.amazon.jdbc.WeightedRandomHostSelector; import software.amazon.jdbc.hostavailability.HostAvailability; import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; import software.amazon.jdbc.util.FullServicesContainer; @@ -226,7 +226,7 @@ void testEstablishConnection_GivenRouterCache_ThenSelectsHost() throws SQLExcept assertEquals(mockConnection, inputContext.getConnection()); verify(mockPluginService, times(1)) - .getHostSpecByStrategy(routerList, HostRole.WRITER, RoundRobinHostSelector.STRATEGY_ROUND_ROBIN); + .getHostSpecByStrategy(routerList, HostRole.WRITER, WeightedRandomHostSelector.STRATEGY_WEIGHTED_RANDOM); verify(mockPluginService, times(1)).connect(selectedRouter, inputContext.getProps(), null); verify(mockConnectFuncLambda, times(0)).call(); } @@ -339,7 +339,7 @@ void testEstablishConnection_GivenSelectsHostThrows_ThenRetry() throws SQLExcept assertEquals(routers, this.storageService.get(LimitlessRouters.class, CLUSTER_ID)); verify(mockPluginService, times(2)).getHostSpecByStrategy(any(), any(), any()); verify(mockPluginService, times(1)) - .getHostSpecByStrategy(routerList, HostRole.WRITER, RoundRobinHostSelector.STRATEGY_ROUND_ROBIN); + .getHostSpecByStrategy(routerList, HostRole.WRITER, WeightedRandomHostSelector.STRATEGY_WEIGHTED_RANDOM); verify(mockPluginService, times(1)) .getHostSpecByStrategy(routerList, HostRole.WRITER, HighestWeightHostSelector.STRATEGY_HIGHEST_WEIGHT); verify(mockPluginService, times(1)).connect(selectedRouter, inputContext.getProps(), null); @@ -377,7 +377,7 @@ void testEstablishConnection_GivenSelectsHostNull_ThenRetry() throws SQLExceptio assertEquals(routers, this.storageService.get(LimitlessRouters.class, CLUSTER_ID)); verify(mockPluginService, times(2)).getHostSpecByStrategy(any(), any(), any()); verify(mockPluginService, times(1)) - .getHostSpecByStrategy(routerList, HostRole.WRITER, RoundRobinHostSelector.STRATEGY_ROUND_ROBIN); + .getHostSpecByStrategy(routerList, HostRole.WRITER, WeightedRandomHostSelector.STRATEGY_WEIGHTED_RANDOM); verify(mockPluginService, times(1)) .getHostSpecByStrategy(routerList, HostRole.WRITER, HighestWeightHostSelector.STRATEGY_HIGHEST_WEIGHT); verify(mockPluginService, times(1)).connect(selectedRouter, inputContext.getProps(), null); @@ -419,7 +419,7 @@ void testEstablishConnection_GivenPluginServiceConnectThrows_ThenRetry() throws assertEquals(routers, this.storageService.get(LimitlessRouters.class, CLUSTER_ID)); verify(mockPluginService, times(2)).getHostSpecByStrategy(any(), any(), any()); verify(mockPluginService, times(1)) - .getHostSpecByStrategy(routerList, HostRole.WRITER, RoundRobinHostSelector.STRATEGY_ROUND_ROBIN); + .getHostSpecByStrategy(routerList, HostRole.WRITER, WeightedRandomHostSelector.STRATEGY_WEIGHTED_RANDOM); verify(mockPluginService, times(1)) .getHostSpecByStrategy(routerList, HostRole.WRITER, HighestWeightHostSelector.STRATEGY_HIGHEST_WEIGHT); verify(mockPluginService, times(2)).connect(any(), any(), any()); diff --git a/wrapper/src/test/java/software/amazon/jdbc/util/HostSelectorUtilsTests.java b/wrapper/src/test/java/software/amazon/jdbc/util/HostSelectorUtilsTests.java new file mode 100644 index 000000000..3339050c4 --- /dev/null +++ b/wrapper/src/test/java/software/amazon/jdbc/util/HostSelectorUtilsTests.java @@ -0,0 +1,62 @@ +/* + * 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 static org.junit.jupiter.api.Assertions.assertEquals; + +import java.util.Arrays; +import java.util.List; +import java.util.Properties; +import org.junit.jupiter.api.Test; +import software.amazon.jdbc.HostSpec; +import software.amazon.jdbc.HostSpecBuilder; +import software.amazon.jdbc.RoundRobinHostSelector; +import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; + +public class HostSelectorUtilsTests { + @Test + void testSetHostWeightPairsProperty() { + final String expectedPropertyValue = "instance-1-id:2,instance-2-id:1,instance-3-id:0"; + + final List 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(); + HostSelectorUtils.setHostWeightPairsProperty( + RoundRobinHostSelector.ROUND_ROBIN_HOST_WEIGHT_PAIRS, + properties, hosts); + + final String actualPropertyValue = properties.getProperty( + RoundRobinHostSelector.ROUND_ROBIN_HOST_WEIGHT_PAIRS.name); + + assertEquals(expectedPropertyValue, actualPropertyValue); + } +}