Skip to content

Commit

Permalink
feat: round robin host selection strategy (aws#603)
Browse files Browse the repository at this point in the history
  • Loading branch information
@crystall-bitquill
crystall-bitquill authored Sep 12, 2023
1 parent 61b2698 commit 6754855
Show file tree
Hide file tree
Showing 18 changed files with 728 additions and 59 deletions.
23 changes: 18 additions & 5 deletions docs/using-the-jdbc-driver/using-plugins/UsingTheReadWriteSplittingPlugin.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -69,11 +69,7 @@ private static String getPoolKey(HostSpec hostSpec, Properties props) {

2. Call `ConnectionProviderManager.setConnectionProvider`, passing in the `HikariPooledConnectionProvider` you created in step 1.

3. By default, the read-write plugin randomly selects a reader instance the first time that `setReadOnly(true)` is called. If you would like the plugin to select a reader based on the instance with the least connections instead, set the following connection property. Note that this strategy is only available when internal connection pools are enabled - if you set the connection property without enabling internal pools, an exception will be thrown.

```java
props.setProperty(ReadWriteSplittingPlugin.READER_HOST_SELECTOR_STRATEGY.name, "leastConnections");
```
3. By default, the read-write plugin randomly selects a reader instance the first time that `setReadOnly(true)` is called. If you would like the plugin to select a reader based on a different connection strategy, please see the [Connection Strategies](#connection-strategies) section for more information.

4. Continue as normal: create connections and use them as needed.

Expand All @@ -84,6 +80,23 @@ props.setProperty(ReadWriteSplittingPlugin.READER_HOST_SELECTOR_STRATEGY.name, "
### Example
[ReadWriteSplittingPostgresExample.java](../../../examples/AWSDriverExample/src/main/java/software/amazon/ReadWriteSplittingPostgresExample.java) demonstrates how to enable and configure read-write splitting with the Aws Advanced JDBC Driver.

### Connection Strategies
By default, the read-write plugin randomly selects a reader instance the first time that `setReadOnly(true)` is called. To balance connections to reader instances more evenly, different connection strategies can be used. The following table describes the currently available connection strategies and any relevant configuration parameters for each strategy.

To indicate which connection strategy to use, the `readerHostSelectorStrategy` configuration parameter can be set to one of the connection strategies in the table below. The following is an example of enabling the least connections strategy:

```java
props.setProperty(ReadWriteSplittingPlugin.READER_HOST_SELECTOR_STRATEGY.name, "leastConnections");
```

| Connection Strategy | Configuration Parameter | Description | Default Value |
|---------------------|-------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------|
| `random` | This strategy does not have configuration parameters. | The random strategy is the default connection strategy. When switching to a reader connection, the reader instance will be chosen randomly from the available database instances. | N/A |
| `leastConnections` | This strategy does not have configuration parameters. | The least connections strategy will select reader instances based on which database instance has the least number of currently active connections. Note that this strategy is only available when internal connection pools are enabled - if you set the connection property without enabling internal pools, an exception will be thrown. | N/A |
| `roundRobin` | See the following rows for configuration parameters. | The round robin strategy will select a reader instance by taking turns with all available database instances in a cycle. A slight addition to the round robin strategy is the weighted round robin strategy, where more connections will be passed to reader instances based on user specified connection properties. | N/A |
| | `roundRobinHostWeightPairs` | This parameter value must be a `string` type comma separated list of database host-weight pairs in the format `<host>:<weight>`. The host represents the database instance name, and the weight represents how many connections should be directed to the host in one cycle through all available hosts. For example, the value `instance-1:1,instance-2:4` means that for every connection to `instance-1`, there will be four connections to `instance-2`. <br><br> **Note:** The `<weight>` value in the string must be an integer greater than or equal to 1. | `null` |
| | `roundRobinDefaultWeight` | This parameter value must be an integer value in the form of a `string`. This parameter represents the default weight for any hosts that have not been configured with the `roundRobinHostWeightPairs` parameter. For example, if a connection were already established and host weights were set with `roundRobinHostWeightPairs` but a new reader node was added to the database, the new reader node would use the default weight. <br><br> **Note:** This value must be an integer greater than or equal to 1. | `1` |

### Limitations

#### General plugin limitations
Expand Down
4 changes: 3 additions & 1 deletion wrapper/src/main/java/software/amazon/jdbc/ConnectionProvider.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import java.util.List;
import java.util.Properties;
import org.checkerframework.checker.nullness.qual.NonNull;
import org.checkerframework.checker.nullness.qual.Nullable;
import software.amazon.jdbc.dialect.Dialect;

/**
Expand Down Expand Up @@ -59,12 +60,13 @@ boolean acceptsUrl(
* @param role determines if the connection provider should return a writer or a reader
* @param strategy the strategy determining how the {@link HostSpec} should be selected, e.g.,
* random or round-robin
* @param props any properties that are required by the provided strategy to select a host
* @return the {@link HostSpec} selected using the specified strategy
* @throws SQLException if an error occurred while returning the hosts
* @throws UnsupportedOperationException if the strategy is unsupported by the provider
*/
HostSpec getHostSpecByStrategy(
@NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy)
@NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy, @Nullable Properties props)
throws SQLException, UnsupportedOperationException;

/**
Expand Down
7 changes: 4 additions & 3 deletions wrapper/src/main/java/software/amazon/jdbc/ConnectionProviderManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -134,21 +134,22 @@ public boolean acceptsStrategy(HostRole role, String strategy) {
* @param role the desired role of the host - either a writer or a reader
* @param strategy the strategy that should be used to select a {@link HostSpec} from the host
* list (eg "random")
* @param props any properties that are required by the provided strategy to select a host
* @return a {@link HostSpec} with the requested role
* @throws SQLException if the available {@link ConnectionProvider} instances
* cannot find a host in the host list matching the
* requested role or an error occurs while selecting a host
* @throws UnsupportedOperationException if the available {@link ConnectionProvider} instances do
* not support the requested strategy
*/
public HostSpec getHostSpecByStrategy(List<HostSpec> hosts, HostRole role, String strategy)
public HostSpec getHostSpecByStrategy(List<HostSpec> hosts, HostRole role, String strategy, Properties props)
throws SQLException, UnsupportedOperationException {
HostSpec host = null;
if (connProvider != null) {
connProviderLock.readLock().lock();
try {
if (connProvider != null && connProvider.acceptsStrategy(role, strategy)) {
host = connProvider.getHostSpecByStrategy(hosts, role, strategy);
host = connProvider.getHostSpecByStrategy(hosts, role, strategy, props);
}
} catch (UnsupportedOperationException e) {
// The custom provider does not support the provided strategy, ignore it and try with the default provider.
Expand All @@ -158,7 +159,7 @@ public HostSpec getHostSpecByStrategy(List<HostSpec> hosts, HostRole role, Strin
}

if (host == null) {
host = defaultProvider.getHostSpecByStrategy(hosts, role, strategy);
host = defaultProvider.getHostSpecByStrategy(hosts, role, strategy, props);
}

return host;
Expand Down
8 changes: 5 additions & 3 deletions wrapper/src/main/java/software/amazon/jdbc/DataSourceConnectionProvider.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
import java.util.logging.Logger;
import javax.sql.DataSource;
import org.checkerframework.checker.nullness.qual.NonNull;
import org.checkerframework.checker.nullness.qual.Nullable;
import software.amazon.jdbc.dialect.Dialect;
import software.amazon.jdbc.exceptions.SQLLoginException;
import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect;
Expand All @@ -46,7 +47,8 @@ public class DataSourceConnectionProvider implements ConnectionProvider {
private static final Map<String, HostSelector> acceptedStrategies =
Collections.unmodifiableMap(new HashMap<String, HostSelector>() {
{
put("random", new RandomHostSelector());
put(RandomHostSelector.STRATEGY_RANDOM, new RandomHostSelector());
put(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN, new RoundRobinHostSelector());
}
});
private final @NonNull DataSource dataSource;
Expand Down Expand Up @@ -86,7 +88,7 @@ public boolean acceptsStrategy(@NonNull HostRole role, @NonNull String strategy)

@Override
public HostSpec getHostSpecByStrategy(
@NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy)
@NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy, @Nullable Properties props)
throws SQLException {
if (!acceptedStrategies.containsKey(strategy)) {
throw new UnsupportedOperationException(
Expand All @@ -95,7 +97,7 @@ public HostSpec getHostSpecByStrategy(
new Object[] {strategy, DataSourceConnectionProvider.class}));
}

return acceptedStrategies.get(strategy).getHost(hosts, role);
return acceptedStrategies.get(strategy).getHost(hosts, role, props);
}

/**
Expand Down
8 changes: 5 additions & 3 deletions wrapper/src/main/java/software/amazon/jdbc/DriverConnectionProvider.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
import java.util.Properties;
import java.util.logging.Logger;
import org.checkerframework.checker.nullness.qual.NonNull;
import org.checkerframework.checker.nullness.qual.Nullable;
import software.amazon.jdbc.dialect.Dialect;
import software.amazon.jdbc.exceptions.SQLLoginException;
import software.amazon.jdbc.targetdriverdialect.ConnectInfo;
Expand All @@ -43,7 +44,8 @@ public class DriverConnectionProvider implements ConnectionProvider {
private static final Map<String, HostSelector> acceptedStrategies =
Collections.unmodifiableMap(new HashMap<String, HostSelector>() {
{
put("random", new RandomHostSelector());
put(RandomHostSelector.STRATEGY_RANDOM, new RandomHostSelector());
put(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN, new RoundRobinHostSelector());
}
});

Expand Down Expand Up @@ -81,7 +83,7 @@ public boolean acceptsStrategy(@NonNull HostRole role, @NonNull String strategy)

@Override
public HostSpec getHostSpecByStrategy(
@NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy)
@NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy, @Nullable Properties props)
throws SQLException {
if (!acceptedStrategies.containsKey(strategy)) {
throw new UnsupportedOperationException(
Expand All @@ -90,7 +92,7 @@ public HostSpec getHostSpecByStrategy(
new Object[] {strategy, DriverConnectionProvider.class}));
}

return acceptedStrategies.get(strategy).getHost(hosts, role);
return acceptedStrategies.get(strategy).getHost(hosts, role, props);
}

/**
Expand Down
61 changes: 29 additions & 32 deletions wrapper/src/main/java/software/amazon/jdbc/HikariPooledConnectionProvider.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,15 +21,17 @@
import java.sql.Connection;
import java.sql.SQLException;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map.Entry;
import java.util.Map;
import java.util.Properties;
import java.util.Set;
import java.util.StringJoiner;
import java.util.concurrent.TimeUnit;
import java.util.logging.Logger;
import java.util.stream.Collectors;
import org.checkerframework.checker.nullness.qual.NonNull;
import org.checkerframework.checker.nullness.qual.Nullable;
import software.amazon.jdbc.cleanup.CanReleaseResources;
import software.amazon.jdbc.dialect.Dialect;
import software.amazon.jdbc.util.HikariCPSQLException;
Expand All @@ -44,9 +46,13 @@ public class HikariPooledConnectionProvider implements PooledConnectionProvider,

private static final Logger LOGGER =
Logger.getLogger(HikariPooledConnectionProvider.class.getName());

private static final String LEAST_CONNECTIONS_STRATEGY = "leastConnections";

private static final Map<String, HostSelector> acceptedStrategies =
Collections.unmodifiableMap(new HashMap<String, HostSelector>() {
{
put(RandomHostSelector.STRATEGY_RANDOM, new RandomHostSelector());
put(RoundRobinHostSelector.STRATEGY_ROUND_ROBIN, new RoundRobinHostSelector());
}
});
private static final RdsUtils rdsUtils = new RdsUtils();
private static SlidingExpirationCache<PoolKey, HikariDataSource> databasePools =
new SlidingExpirationCache<>(
Expand All @@ -56,6 +62,7 @@ public class HikariPooledConnectionProvider implements PooledConnectionProvider,
private static long poolExpirationCheckNanos = TimeUnit.MINUTES.toNanos(30);
private final HikariPoolConfigurator poolConfigurator;
private final HikariPoolMapping poolMapping;
private final LeastConnectionsHostSelector leastConnectionsHostSelector;

/**
* {@link HikariPooledConnectionProvider} constructor. This class can be passed to
Expand Down Expand Up @@ -98,6 +105,7 @@ public HikariPooledConnectionProvider(
HikariPoolConfigurator hikariPoolConfigurator, HikariPoolMapping mapping) {
this.poolConfigurator = hikariPoolConfigurator;
this.poolMapping = mapping;
this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(databasePools);
}

/**
Expand Down Expand Up @@ -134,6 +142,7 @@ public HikariPooledConnectionProvider(
this.poolMapping = mapping;
poolExpirationCheckNanos = poolExpirationNanos;
databasePools.setCleanupIntervalNanos(poolCleanupNanos);
this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(databasePools);
}

@Override
Expand All @@ -145,44 +154,28 @@ public boolean acceptsUrl(

@Override
public boolean acceptsStrategy(@NonNull HostRole role, @NonNull String strategy) {
return LEAST_CONNECTIONS_STRATEGY.equals(strategy);
return acceptedStrategies.containsKey(strategy)
|| LeastConnectionsHostSelector.STRATEGY_LEAST_CONNECTIONS.equals(strategy);
}

@Override
public HostSpec getHostSpecByStrategy(
@NonNull List<HostSpec> hosts, @NonNull HostRole role, @NonNull String strategy)
throws SQLException {
if (!LEAST_CONNECTIONS_STRATEGY.equals(strategy)) {
@NonNull List<HostSpec> hosts,
@NonNull HostRole role,
@NonNull String strategy,
@Nullable Properties props) throws SQLException {
if (!acceptsStrategy(role, strategy)) {
throw new UnsupportedOperationException(
Messages.get(
"ConnectionProvider.unsupportedHostSpecSelectorStrategy",
new Object[] {strategy, HikariPooledConnectionProvider.class}));
}

// Remove hosts with the wrong role
List<HostSpec> eligibleHosts = hosts.stream()
.filter(hostSpec -> role.equals(hostSpec.getRole()))
.sorted((hostSpec1, hostSpec2) ->
getNumConnections(hostSpec1) - getNumConnections(hostSpec2))
.collect(Collectors.toList());

if (eligibleHosts.size() == 0) {
throw new SQLException(Messages.get("HostSelector.noHostsMatchingRole", new Object[]{role}));
new Object[] {strategy, DataSourceConnectionProvider.class}));
}

return eligibleHosts.get(0);
}

private int getNumConnections(HostSpec hostSpec) {
int numConnections = 0;
final String url = hostSpec.getUrl();
for (Entry<PoolKey, HikariDataSource> entry : databasePools.getEntries().entrySet()) {
if (!url.equals(entry.getKey().url)) {
continue;
}
numConnections += entry.getValue().getHikariPoolMXBean().getActiveConnections();
if (LeastConnectionsHostSelector.STRATEGY_LEAST_CONNECTIONS.equals(strategy)) {
return this.leastConnectionsHostSelector.getHost(hosts, role, props);
} else {
return acceptedStrategies.get(strategy).getHost(hosts, role, props);
}
return numConnections;
}

@Override
Expand Down Expand Up @@ -338,6 +331,10 @@ public PoolKey(final @NonNull String url, final @NonNull String extraKey) {
this.extraKey = extraKey;
}

public String getUrl() {
return this.url;
}

@Override
public int hashCode() {
final int prime = 31;
Expand Down
11 changes: 8 additions & 3 deletions wrapper/src/main/java/software/amazon/jdbc/HostSelector.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,17 +18,22 @@

import java.sql.SQLException;
import java.util.List;
import java.util.Properties;
import org.checkerframework.checker.nullness.qual.NonNull;
import org.checkerframework.checker.nullness.qual.Nullable;

public interface HostSelector {

/**
* Selects a host with the requested role from the given host list.
*
* @param hosts a list of available hosts to pick from
* @param role the desired host role - either a writer or a reader
* @param hosts a list of available hosts to pick from.
* @param role the desired host role - either a writer or a reader.
* @param props connection properties that may be needed by the host selector in order to choose a host.
* @return a host matching the requested role
* @throws SQLException if the host list does not contain any hosts matching the requested role or
* an error occurs while selecting a host
*/
HostSpec getHost(List<HostSpec> hosts, HostRole role) throws SQLException;
HostSpec getHost(
@NonNull List<HostSpec> hosts, @NonNull HostRole role, @Nullable Properties props) throws SQLException;
}
Loading

0 comments on commit 6754855

Please sign in to comment.