Skip to content

Commit 0a044e4

Browse files
aaronchung-bitquillaaronchung-bitquill
committed
feat: weighted random host selector + fix: LimitlessQueryHelper rounding error
1 parent 7f6aba2 commit 0a044e4

File tree

4 files changed

+368
-1
lines changed

4 files changed

+368
-1
lines changed

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

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package software.amazon.jdbc;
18+
19+
import java.sql.SQLException;
20+
import java.util.Comparator;
21+
import java.util.HashMap;
22+
import java.util.List;
23+
import java.util.Map;
24+
import java.util.Map.Entry;
25+
import java.util.Properties;
26+
import java.util.Random;
27+
import java.util.concurrent.Callable;
28+
import java.util.regex.Matcher;
29+
import java.util.regex.Pattern;
30+
import java.util.stream.Collectors;
31+
import org.checkerframework.checker.nullness.qual.NonNull;
32+
import org.checkerframework.checker.nullness.qual.Nullable;
33+
import software.amazon.jdbc.hostavailability.HostAvailability;
34+
import software.amazon.jdbc.util.Messages;
35+
36+
public class WeightedRandomHostSelector implements HostSelector {
37+
public static final AwsWrapperProperty WEIGHTED_RANDOM_HOST_WEIGHT_PAIRS = new AwsWrapperProperty(
38+
"weightedRandomHostWeightPairs", null,
39+
"Comma separated list of database host-weight pairs in the format of `<host>:<weight>`.");
40+
public static final String STRATEGY_WEIGHTED_RANDOM = "weightedRandom";
41+
static final int DEFAULT_WEIGHT = 1;
42+
static final Pattern HOST_WEIGHT_PAIRS_PATTERN =
43+
Pattern.compile("((?<host>[^:/?#]*):(?<weight>[0-9]*))");
44+
45+
private Map<String, Integer> cachedHostWeightMap;
46+
private String cachedHostWeightMapString;
47+
private Random random;
48+
private Callable<Integer> randomFunc;
49+
50+
public HostSpec getHost(
51+
@NonNull List<HostSpec> hosts,
52+
@NonNull HostRole role,
53+
@Nullable Properties props) throws SQLException {
54+
55+
final Map<String, Integer> hostWeightMap =
56+
this.getHostWeightPairMap(WEIGHTED_RANDOM_HOST_WEIGHT_PAIRS.getString(props));
57+
58+
// Get and check eligible hosts
59+
final List<HostSpec> eligibleHosts = hosts.stream()
60+
.filter(hostSpec ->
61+
role.equals(hostSpec.getRole()) && hostSpec.getAvailability().equals(HostAvailability.AVAILABLE))
62+
.sorted(Comparator.comparing(HostSpec::getHost))
63+
.collect(Collectors.toList());
64+
65+
if (eligibleHosts.isEmpty()) {
66+
throw new SQLException(Messages.get("HostSelector.noHostsMatchingRole", new Object[] {role}));
67+
}
68+
69+
final Map<HostSpec, NumberRange> hostWeightRangeMap = new HashMap<>();
70+
int counter = 1;
71+
for (HostSpec host : eligibleHosts) {
72+
if (!hostWeightMap.containsKey(host.getHost())) {
73+
continue;
74+
}
75+
final int hostWeight = hostWeightMap.get(host.getHost());
76+
if (hostWeight > 0) {
77+
final int rangeStart = counter;
78+
final int rangeEnd = counter + hostWeight - 1;
79+
hostWeightRangeMap.put(host, new NumberRange(rangeStart, rangeEnd));
80+
counter = counter + hostWeight;
81+
} else {
82+
hostWeightRangeMap.put(host, new NumberRange(counter, counter));
83+
counter++;
84+
}
85+
}
86+
87+
// Check random number is in host weigh range map
88+
if (this.random == null) {
89+
this.random = new Random();
90+
}
91+
int randomInt = this.random.nextInt(counter);
92+
93+
// This block is for testing purposes
94+
if (this.randomFunc != null) {
95+
try {
96+
randomInt = this.randomFunc.call();
97+
} catch (Exception e) {
98+
// This should not happen
99+
}
100+
}
101+
102+
for (final Entry<HostSpec, NumberRange> entry : hostWeightRangeMap.entrySet()) {
103+
if (hostWeightRangeMap.get(entry.getKey()).isInRange(randomInt)) {
104+
return entry.getKey();
105+
}
106+
}
107+
// TODO: proper messaging
108+
throw new SQLException(Messages.get("HostSelector.TODO", new Object[] {role}));
109+
}
110+
111+
private Map<String, Integer> getHostWeightPairMap(final String hostWeightMapString) throws SQLException {
112+
if (this.cachedHostWeightMapString != null
113+
&& this.cachedHostWeightMapString.trim().equals(hostWeightMapString.trim())
114+
&& this.cachedHostWeightMap != null
115+
&& !this.cachedHostWeightMap.isEmpty()) {
116+
return this.cachedHostWeightMap;
117+
}
118+
119+
final Map<String, Integer> hostWeightMap = new HashMap<>();
120+
if (hostWeightMapString == null || hostWeightMapString.trim().isEmpty()) {
121+
return hostWeightMap;
122+
}
123+
final String[] hostWeightPairs = hostWeightMapString.split(",");
124+
for (final String hostWeightPair : hostWeightPairs) {
125+
final Matcher matcher = HOST_WEIGHT_PAIRS_PATTERN.matcher(hostWeightPair);
126+
if (!matcher.matches()) {
127+
// TODO: add this message
128+
throw new SQLException(Messages.get("HostSelector.weightedRandomInvalidHostWeightPairs"));
129+
}
130+
131+
final String hostName = matcher.group("host").trim();
132+
final String hostWeight = matcher.group("weight").trim();
133+
if (hostName.isEmpty() || hostWeight.isEmpty()) {
134+
throw new SQLException(Messages.get("HostSelector.weightedRandomInvalidHostWeightPairs"));
135+
}
136+
137+
try {
138+
final int weight = Integer.parseInt(hostWeight);
139+
if (weight < DEFAULT_WEIGHT) {
140+
throw new SQLException(Messages.get("HostSelector.weightedRandomInvalidHostWeightPairs"));
141+
}
142+
hostWeightMap.put(hostName, weight);
143+
} catch (NumberFormatException e) {
144+
throw new SQLException(Messages.get("HostSelector.roundRobinInvalidHostWeightPairs"));
145+
}
146+
}
147+
this.cachedHostWeightMap = hostWeightMap;
148+
this.cachedHostWeightMapString = hostWeightMapString;
149+
return hostWeightMap;
150+
}
151+
152+
public void setRandomFunc(final Callable<Integer> randomFunc) {
153+
this.randomFunc = randomFunc;
154+
}
155+
156+
private static class NumberRange {
157+
private int start;
158+
private int end;
159+
160+
public NumberRange(int start, int end) {
161+
this.start = start;
162+
this.end = end;
163+
}
164+
165+
public boolean isInRange(int value) {
166+
return start <= value && value <= end;
167+
}
168+
}
169+
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ protected HostSpec createHost(final ResultSet resultSet, final int hostPortToMap
9696
final String hostName = resultSet.getString(1);
9797
final float cpu = resultSet.getFloat(2);
9898

99-
long weight = Math.round(10 - cpu * 10);
99+
long weight = (long) (10 - Math.floor(10 * cpu));
100100

101101
if (weight < 1 || weight > 10) {
102102
weight = 1; // default to 1

‎wrapper/src/main/resources/aws_advanced_jdbc_wrapper_messages.properties‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -193,6 +193,7 @@ HostMonitoringConnectionPlugin.unableToIdentifyConnection=Unable to identify the
193193
HostSelector.noHostsMatchingRole=No hosts were found matching the requested ''{0}'' role.
194194
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.
195195
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.
196+
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.
196197

197198
IamAuthConnectionPlugin.unhandledException=Unhandled exception: ''{0}''
198199
IamAuthConnectionPlugin.connectException=Error occurred while opening a connection: ''{0}''

‎wrapper/src/test/java/software/amazon/jdbc/WeightedRandomHostSelectorTests.java‎

Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package software.amazon.jdbc;
18+
19+
import static org.junit.jupiter.api.Assertions.assertEquals;
20+
import static org.junit.jupiter.api.Assertions.assertThrows;
21+
import static software.amazon.jdbc.WeightedRandomHostSelector.WEIGHTED_RANDOM_HOST_WEIGHT_PAIRS;
22+
23+
import java.sql.SQLException;
24+
import java.util.Arrays;
25+
import java.util.Collections;
26+
import java.util.List;
27+
import java.util.Properties;
28+
import org.junit.jupiter.api.Test;
29+
import software.amazon.jdbc.hostavailability.HostAvailability;
30+
import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy;
31+
32+
class WeightedRandomHostSelectorTests {
33+
34+
@Test
35+
void testGetHost_emptyHostList() {
36+
final HostSelector hostSelector = new WeightedRandomHostSelector();
37+
final Properties props = new Properties();
38+
final List<HostSpec> emptyHostList = Collections.emptyList();
39+
assertThrows(SQLException.class, () -> hostSelector.getHost(emptyHostList, HostRole.WRITER, props));
40+
}
41+
42+
@Test
43+
void testGetHost_noEligibleHosts() {
44+
final HostSelector hostSelector = new WeightedRandomHostSelector();
45+
final Properties props = new Properties();
46+
final List<HostSpec> noEligibleHostsList = Arrays.asList(
47+
new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-1").role(HostRole.READER).build(),
48+
new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-2").role(HostRole.READER).build(),
49+
new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-3").role(HostRole.READER).build()
50+
);
51+
assertThrows(SQLException.class,
52+
() -> hostSelector.getHost(noEligibleHostsList, HostRole.WRITER, props));
53+
}
54+
55+
@Test
56+
void testGetHost_invalidWeight() {
57+
final HostSelector hostSelector = new WeightedRandomHostSelector();
58+
final Properties props = new Properties();
59+
props.setProperty(WEIGHTED_RANDOM_HOST_WEIGHT_PAIRS.name, "instance-1:3,instance-2:2,instance-3:0");
60+
final List<HostSpec> eligibleHostsList = Arrays.asList(
61+
new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-1").role(HostRole.WRITER)
62+
.availability(HostAvailability.AVAILABLE).build(),
63+
new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-2").role(HostRole.WRITER)
64+
.availability(HostAvailability.AVAILABLE).build(),
65+
new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-3").role(HostRole.WRITER)
66+
.availability(HostAvailability.AVAILABLE).build()
67+
);
68+
assertThrows(SQLException.class,
69+
() -> hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props));
70+
}
71+
72+
@Test
73+
void testGetHost_invalidProps() {
74+
final HostSelector hostSelector = new WeightedRandomHostSelector();
75+
final Properties props = new Properties();
76+
props.setProperty(WEIGHTED_RANDOM_HOST_WEIGHT_PAIRS.name, "someInvalidString");
77+
final List<HostSpec> eligibleHostsList = Arrays.asList(
78+
new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-1").role(HostRole.WRITER)
79+
.availability(HostAvailability.AVAILABLE).build(),
80+
new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-2").role(HostRole.WRITER)
81+
.availability(HostAvailability.AVAILABLE).build(),
82+
new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-3").role(HostRole.WRITER)
83+
.availability(HostAvailability.AVAILABLE).build()
84+
);
85+
assertThrows(SQLException.class,
86+
() -> hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props));
87+
}
88+
89+
@Test
90+
void testGetHost() throws SQLException {
91+
final WeightedRandomHostSelector hostSelector = new WeightedRandomHostSelector();
92+
final Properties props = new Properties();
93+
props.setProperty(WEIGHTED_RANDOM_HOST_WEIGHT_PAIRS.name, "instance-1:3,instance-2:2,instance-3:01");
94+
final List<HostSpec> eligibleHostsList = Arrays.asList(
95+
new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-1").role(HostRole.WRITER)
96+
.availability(HostAvailability.AVAILABLE).build(),
97+
new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-2").role(HostRole.WRITER)
98+
.availability(HostAvailability.AVAILABLE).build(),
99+
new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-3").role(HostRole.WRITER)
100+
.availability(HostAvailability.AVAILABLE).build()
101+
);
102+
103+
hostSelector.setRandomFunc(() -> 1);
104+
final HostSpec actualHost1 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props);
105+
assertEquals(eligibleHostsList.get(0).getHost(), actualHost1.getHost());
106+
107+
hostSelector.setRandomFunc(() -> 2);
108+
final HostSpec actualHost2 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props);
109+
assertEquals(eligibleHostsList.get(0).getHost(), actualHost2.getHost());
110+
111+
hostSelector.setRandomFunc(() -> 3);
112+
final HostSpec actualHost3 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props);
113+
assertEquals(eligibleHostsList.get(0).getHost(), actualHost3.getHost());
114+
115+
hostSelector.setRandomFunc(() -> 4);
116+
final HostSpec actualHost4 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props);
117+
assertEquals(eligibleHostsList.get(1).getHost(), actualHost4.getHost());
118+
119+
hostSelector.setRandomFunc(() -> 5);
120+
final HostSpec actualHost5 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props);
121+
assertEquals(eligibleHostsList.get(1).getHost(), actualHost5.getHost());
122+
123+
hostSelector.setRandomFunc(() -> 6);
124+
final HostSpec actualHost6 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props);
125+
assertEquals(eligibleHostsList.get(2).getHost(), actualHost6.getHost());
126+
}
127+
128+
@Test
129+
void testGetHost_changeWeights() throws SQLException {
130+
final WeightedRandomHostSelector hostSelector = new WeightedRandomHostSelector();
131+
final Properties props = new Properties();
132+
133+
props.setProperty(WEIGHTED_RANDOM_HOST_WEIGHT_PAIRS.name, "instance-1:3,instance-2:2,instance-3:01");
134+
final List<HostSpec> eligibleHostsList = Arrays.asList(
135+
new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-1").role(HostRole.WRITER)
136+
.availability(HostAvailability.AVAILABLE).build(),
137+
new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-2").role(HostRole.WRITER)
138+
.availability(HostAvailability.AVAILABLE).build(),
139+
new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("instance-3").role(HostRole.WRITER)
140+
.availability(HostAvailability.AVAILABLE).build()
141+
);
142+
143+
hostSelector.setRandomFunc(() -> 1);
144+
final HostSpec actualHost1 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props);
145+
assertEquals(eligibleHostsList.get(0).getHost(), actualHost1.getHost());
146+
147+
hostSelector.setRandomFunc(() -> 2);
148+
final HostSpec actualHost2 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props);
149+
assertEquals(eligibleHostsList.get(0).getHost(), actualHost2.getHost());
150+
151+
hostSelector.setRandomFunc(() -> 3);
152+
final HostSpec actualHost3 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props);
153+
assertEquals(eligibleHostsList.get(0).getHost(), actualHost3.getHost());
154+
155+
hostSelector.setRandomFunc(() -> 4);
156+
final HostSpec actualHost4 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props);
157+
assertEquals(eligibleHostsList.get(1).getHost(), actualHost4.getHost());
158+
159+
hostSelector.setRandomFunc(() -> 5);
160+
final HostSpec actualHost5 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props);
161+
assertEquals(eligibleHostsList.get(1).getHost(), actualHost5.getHost());
162+
163+
hostSelector.setRandomFunc(() -> 6);
164+
final HostSpec actualHost6 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props);
165+
assertEquals(eligibleHostsList.get(2).getHost(), actualHost6.getHost());
166+
167+
props.setProperty(WEIGHTED_RANDOM_HOST_WEIGHT_PAIRS.name, "instance-1:1,instance-2:4,instance-3:2");
168+
169+
hostSelector.setRandomFunc(() -> 1);
170+
final HostSpec actualHost7 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props);
171+
assertEquals(eligibleHostsList.get(0).getHost(), actualHost7.getHost());
172+
173+
hostSelector.setRandomFunc(() -> 2);
174+
final HostSpec actualHost8 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props);
175+
assertEquals(eligibleHostsList.get(1).getHost(), actualHost8.getHost());
176+
177+
hostSelector.setRandomFunc(() -> 3);
178+
final HostSpec actualHost9 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props);
179+
assertEquals(eligibleHostsList.get(1).getHost(), actualHost9.getHost());
180+
181+
hostSelector.setRandomFunc(() -> 4);
182+
final HostSpec actualHost10 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props);
183+
assertEquals(eligibleHostsList.get(1).getHost(), actualHost10.getHost());
184+
185+
hostSelector.setRandomFunc(() -> 5);
186+
final HostSpec actualHost11 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props);
187+
assertEquals(eligibleHostsList.get(1).getHost(), actualHost11.getHost());
188+
189+
hostSelector.setRandomFunc(() -> 6);
190+
final HostSpec actualHost12 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props);
191+
assertEquals(eligibleHostsList.get(2).getHost(), actualHost12.getHost());
192+
193+
hostSelector.setRandomFunc(() -> 7);
194+
final HostSpec actualHost13 = hostSelector.getHost(eligibleHostsList, HostRole.WRITER, props);
195+
assertEquals(eligibleHostsList.get(2).getHost(), actualHost13.getHost());
196+
}
197+
}

0 commit comments

Comments
 (0)