Skip to content

Commit

Permalink
feat: update HostSelector implementations to account for HostAvailabi…
Browse files Browse the repository at this point in the history
…lity
  • Loading branch information
aaronchung-bitquill committed Jan 30, 2024
1 parent 3274f81 commit 94fc121
Show file tree
Hide file tree
Showing 4 changed files with 96 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
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;
import software.amazon.jdbc.util.SlidingExpirationCache;

Expand All @@ -42,7 +43,8 @@ public HostSpec getHost(
@NonNull final HostRole role,
@Nullable final Properties props) throws SQLException {
final List<HostSpec> eligibleHosts = hosts.stream()
.filter(hostSpec -> role.equals(hostSpec.getRole()))
.filter(hostSpec ->
role.equals(hostSpec.getRole()) && hostSpec.getAvailability().equals(HostAvailability.AVAILABLE))
.sorted((hostSpec1, hostSpec2) ->
getNumConnections(hostSpec1, this.databasePools) - getNumConnections(hostSpec2, this.databasePools))
.collect(Collectors.toList());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
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 RandomHostSelector implements HostSelector {
Expand All @@ -35,7 +36,9 @@ public HostSpec getHost(
@NonNull final HostRole role,
@Nullable final Properties props) throws SQLException {
final List<HostSpec> eligibleHosts = hosts.stream()
.filter(hostSpec -> role.equals(hostSpec.getRole())).collect(Collectors.toList());
.filter(hostSpec ->
role.equals(hostSpec.getRole()) && hostSpec.getAvailability().equals(HostAvailability.AVAILABLE))
.collect(Collectors.toList());
if (eligibleHosts.size() == 0) {
throw new SQLException(Messages.get("HostSelector.noHostsMatchingRole", new Object[]{role}));
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
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.CacheMap;
import software.amazon.jdbc.util.Messages;
import software.amazon.jdbc.util.StringUtils;
Expand Down Expand Up @@ -57,7 +58,8 @@ public synchronized HostSpec getHost(
final @NonNull HostRole role,
final @Nullable Properties props) throws SQLException {
final List<HostSpec> eligibleHosts = hosts.stream()
.filter(hostSpec -> role.equals(hostSpec.getRole()))
.filter(hostSpec ->
role.equals(hostSpec.getRole()) && hostSpec.getAvailability().equals(HostAvailability.AVAILABLE))
.sorted(Comparator.comparing(HostSpec::getHost))
.collect(Collectors.toList());

Expand Down
86 changes: 86 additions & 0 deletions wrapper/src/test/java/RandomHostSelectorTests.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,86 @@
/*
* 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.
*/

import static org.junit.jupiter.api.Assertions.assertEquals;

import java.sql.SQLException;
import java.util.Arrays;
import java.util.List;
import java.util.Properties;
import org.junit.jupiter.api.RepeatedTest;
import software.amazon.jdbc.HostRole;
import software.amazon.jdbc.HostSpec;
import software.amazon.jdbc.HostSpecBuilder;
import software.amazon.jdbc.RandomHostSelector;
import software.amazon.jdbc.hostavailability.HostAvailability;
import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy;

class RandomHostSelectorTests {

private static final HostRole HOST_ROLE = HostRole.READER;

@RepeatedTest(value = 50)
void testGetHostGivenUnavailbleHost() throws SQLException {

final HostSpec unavailableHost = new HostSpecBuilder(new SimpleHostAvailabilityStrategy())
.host("someUnavailableHost")
.role(HOST_ROLE)
.availability(HostAvailability.NOT_AVAILABLE)
.build();

final HostSpec availableHost = new HostSpecBuilder(new SimpleHostAvailabilityStrategy())
.host("someAvailableHost")
.role(HOST_ROLE)
.availability(HostAvailability.AVAILABLE)
.build();

final RandomHostSelector hostSelector = new RandomHostSelector();
final HostSpec actualHost = hostSelector.getHost(Arrays.asList(unavailableHost, availableHost), HOST_ROLE,
new Properties());

assertEquals(availableHost, actualHost);
}

@RepeatedTest(value = 50)
void testGetHostGivenMultipleUnavailableHosts() throws SQLException {
List<HostSpec> hostSpecTestsList = Arrays.asList(
new HostSpecBuilder(new SimpleHostAvailabilityStrategy())
.host("someUnavailableHost")
.role(HOST_ROLE)
.availability(HostAvailability.NOT_AVAILABLE)
.build(),
new HostSpecBuilder(new SimpleHostAvailabilityStrategy())
.host("someUnavailableHost")
.role(HOST_ROLE)
.availability(HostAvailability.NOT_AVAILABLE)
.build(),
new HostSpecBuilder(new SimpleHostAvailabilityStrategy())
.host("someAvailableHost")
.role(HOST_ROLE)
.availability(HostAvailability.AVAILABLE)
.build(),
new HostSpecBuilder(new SimpleHostAvailabilityStrategy())
.host("someAvailableHost")
.role(HOST_ROLE)
.availability(HostAvailability.AVAILABLE)
.build()
);

final RandomHostSelector hostSelector = new RandomHostSelector();
final HostSpec actualHost = hostSelector.getHost(hostSpecTestsList, HOST_ROLE, new Properties());
assertEquals(HostAvailability.AVAILABLE, actualHost.getAvailability());
}
}

0 comments on commit 94fc121

Please sign in to comment.