diff --git a/wrapper/src/main/java/software/amazon/jdbc/LeastConnectionsHostSelector.java b/wrapper/src/main/java/software/amazon/jdbc/LeastConnectionsHostSelector.java index ec5299415..ea9d14cfa 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/LeastConnectionsHostSelector.java +++ b/wrapper/src/main/java/software/amazon/jdbc/LeastConnectionsHostSelector.java @@ -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; @@ -42,7 +43,8 @@ public HostSpec getHost( @NonNull final HostRole role, @Nullable final Properties props) throws SQLException { final List 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()); diff --git a/wrapper/src/main/java/software/amazon/jdbc/RandomHostSelector.java b/wrapper/src/main/java/software/amazon/jdbc/RandomHostSelector.java index 29d3afbaf..0898d4df9 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/RandomHostSelector.java +++ b/wrapper/src/main/java/software/amazon/jdbc/RandomHostSelector.java @@ -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 { @@ -35,7 +36,9 @@ public HostSpec getHost( @NonNull final HostRole role, @Nullable final Properties props) throws SQLException { final List 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})); } diff --git a/wrapper/src/main/java/software/amazon/jdbc/RoundRobinHostSelector.java b/wrapper/src/main/java/software/amazon/jdbc/RoundRobinHostSelector.java index cb489b770..a942da615 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/RoundRobinHostSelector.java +++ b/wrapper/src/main/java/software/amazon/jdbc/RoundRobinHostSelector.java @@ -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; @@ -57,7 +58,8 @@ public synchronized HostSpec getHost( final @NonNull HostRole role, final @Nullable Properties props) throws SQLException { final List 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()); diff --git a/wrapper/src/test/java/RandomHostSelectorTests.java b/wrapper/src/test/java/RandomHostSelectorTests.java new file mode 100644 index 000000000..8eaabe410 --- /dev/null +++ b/wrapper/src/test/java/RandomHostSelectorTests.java @@ -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 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()); + } +}