Skip to content

Commit

Permalink
Add a provision in driver to try another server on login timeout (#21)
Browse files Browse the repository at this point in the history 
Adds a feature of trying another server when login timeout is specified.
  • Loading branch information
@kneeraj
kneeraj authored Jun 7, 2024
1 parent 8ede139 commit c25ac63
Show file tree
Hide file tree
Showing 2 changed files with 58 additions and 23 deletions.
9 changes: 5 additions & 4 deletions pgjdbc/src/main/java/com/yugabyte/ysql/ClusterAwareLoadBalancer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,7 @@ public boolean hasMorePreferredNode(String chosenHost) {
return false;
}

public synchronized String getLeastLoadedServer(List<String> failedHosts) {
public synchronized String getLeastLoadedServer(List<String> failedHosts, ArrayList<String> timedOutHosts) {
LOGGER.fine("failedHosts: " + failedHosts + ", hostToNumConnMap: " + hostToNumConnMap);
if ((hostToNumConnMap.isEmpty() && currentPublicIps.isEmpty())
|| Boolean.getBoolean(EXPLICIT_FALLBACK_ONLY_KEY)) {
Expand All @@ -110,8 +110,9 @@ public synchronized String getLeastLoadedServer(List<String> failedHosts) {
int min = Integer.MAX_VALUE;
ArrayList<String> minConnectionsHostList = new ArrayList<>();
for (String h : hostToNumConnMap.keySet()) {
if (failedHosts.contains(h)) {
LOGGER.fine("Skipping failed host " + h);
boolean wasTimedOutHost = timedOutHosts != null && timedOutHosts.contains(h);
if (failedHosts.contains(h) || wasTimedOutHost) {
LOGGER.fine("Skipping failed host " + h + "(was timed out host=" + wasTimedOutHost +")");
continue;
}
int currLoad = hostToNumConnMap.get(h);
Expand Down Expand Up @@ -155,7 +156,7 @@ public synchronized String getLeastLoadedServer(List<String> failedHosts) {
}
}
// base condition for this recursive call is useHostColumn != null
return getLeastLoadedServer(failedHosts);
return getLeastLoadedServer(failedHosts, timedOutHosts);
}
}
LOGGER.log(Level.FINE,
Expand Down
72 changes: 53 additions & 19 deletions pgjdbc/src/main/java/org/postgresql/Driver.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -283,15 +283,35 @@ private Properties loadDefaultProperties() throws IOException {
// we managed to establish one after all. See ConnectThread for
// more details.
long timeout = timeout(props);
LoadBalanceProperties lbprops = new LoadBalanceProperties(url, props);

if (timeout <= 0) {
return makeConnection(url, props);
return makeConnection(url, props, lbprops, null);
}

ConnectThread ct = new ConnectThread(url, props);
Thread thread = new Thread(ct, "PostgreSQL JDBC driver connection thread");
thread.setDaemon(true); // Don't prevent the VM from shutting down
thread.start();
return ct.getResult(timeout);
ConnectThread ct;
ArrayList<String> prevTimedOutServers = new ArrayList<>();
int maxRetries = 10;
int tries = 0;
while(true) {
ct = new ConnectThread(url, props, lbprops, prevTimedOutServers);
try {
Thread thread = new Thread(ct, "PostgreSQL JDBC driver connection thread");
thread.setDaemon(true); // Don't prevent the VM from shutting down
thread.start();
return ct.getResult(timeout);
} catch (PSQLException ex1) {
LOGGER.log(Level.INFO, "got exception state: " + ex1.getSQLState());
if (lbprops.hasLoadBalance() && !prevTimedOutServers.isEmpty() && tries++ < maxRetries &&
ex1.getSQLState().equals(PSQLState.CONNECTION_UNABLE_TO_CONNECT.getState())) {
LOGGER.log(Level.INFO, "Connection timeout error occurred with server: "
+ prevTimedOutServers.get(prevTimedOutServers.size()) +
" trying other servers, retryAttempt=" + tries);
} else {
throw ex1;
}
}
}
} catch (PSQLException ex1) {
LOGGER.log(Level.FINE, "Connection error: ", ex1);
// re-throw the exception, otherwise it will be caught next, and a
Expand Down Expand Up @@ -385,17 +405,22 @@ private void setupLoggerFromProperties(final Properties props) {
* while enforcing a login timeout.
*/
private static class ConnectThread implements Runnable {
ConnectThread(String url, Properties props) {
private final ArrayList<String> triedHosts;
private final LoadBalanceProperties lbprops;
ConnectThread(String url, Properties props,
LoadBalanceProperties lbprops, ArrayList<String> prevTimedOutServers) {
this.url = url;
this.props = props;
this.lbprops = lbprops;
triedHosts = prevTimedOutServers;
}

public void run() {
Connection conn;
Throwable error;

try {
conn = makeConnection(url, props);
conn = makeConnection(url, props, lbprops, triedHosts);
error = null;
} catch (Throwable t) {
conn = null;
Expand Down Expand Up @@ -479,32 +504,37 @@ public Connection getResult(long timeout) throws SQLException {
* Create a connection from URL and properties. Always does the connection work in the current
* thread without enforcing a timeout, regardless of any timeout specified in the properties.
*
* @param url the original URL
* @param properties the parsed/defaulted connection properties
* @param url the original URL
* @param properties the parsed/defaulted connection properties
* @param timedOutHosts A list of previously timedout servers passed from Connect thread
* @return a new connection
* @throws SQLException if the connection could not be made
*/
private static Connection makeConnection(String url, Properties properties) throws SQLException {
LoadBalanceProperties lbprops = new LoadBalanceProperties(url, properties);
private static Connection makeConnection(String url, Properties properties,
LoadBalanceProperties lbprops, ArrayList<String> timedOutHosts) throws SQLException {
if (lbprops.hasLoadBalance()) {
Connection conn = getConnectionBalanced(lbprops);
Connection conn = getConnectionBalanced(lbprops, timedOutHosts);
if (conn != null) {
return conn;
}
LOGGER.log(Level.WARNING, "Failed to apply load balance. Trying normal connection");
}
// Make the timedOutHosts empty so that the connect thread does not retry because of failures from
// the original connect attempt.
if (timedOutHosts != null) timedOutHosts.clear();
// Attempt connection with the original properties
return new PgConnection(hostSpecs(properties), user(properties), database(properties), properties, url);
}

private static Connection getConnectionBalanced(LoadBalanceProperties lbprops) {
private static Connection getConnectionBalanced(LoadBalanceProperties lbprops,
ArrayList<String> timedOutHosts) {
LOGGER.log(Level.FINE, "GetConnectionBalanced called");
ClusterAwareLoadBalancer loadBalancer = lbprops.getAppropriateLoadBalancer();
Properties props = lbprops.getStrippedProperties();
String url = lbprops.getStrippedURL();
Set<String> unreachableHosts = loadBalancer.getUnreachableHosts();
List<String> failedHosts = new ArrayList<>(unreachableHosts);
String chosenHost = loadBalancer.getLeastLoadedServer(failedHosts);
String chosenHost = loadBalancer.getLeastLoadedServer(failedHosts, timedOutHosts);
PgConnection newConnection = null;
Connection controlConnection = null;
SQLException firstException = null;
Expand All @@ -525,6 +555,7 @@ hspec, user(lbprops.getOriginalProperties()),
controlConnection.close();
} catch (SQLException ex) {
if (PSQLState.UNDEFINED_FUNCTION.getState().equals(ex.getSQLState())) {
LOGGER.log(Level.WARNING, "yb_servers() is not defined on the server");
return null;
}
gotException = true;
Expand All @@ -537,7 +568,7 @@ hspec, user(lbprops.getOriginalProperties()),
}
}
}
chosenHost = loadBalancer.getLeastLoadedServer(failedHosts);
chosenHost = loadBalancer.getLeastLoadedServer(failedHosts, timedOutHosts);
}
if (chosenHost == null) {
return null;
Expand All @@ -551,6 +582,9 @@ hspec, user(lbprops.getOriginalProperties()),
if (port != null) {
props.setProperty("PGPORT", port);
}
if (timedOutHosts != null) {
timedOutHosts.add(chosenHost);
}
newConnection = new PgConnection(
hostSpecs(props), user(lbprops.getOriginalProperties()), database(props), props, url);
newConnection.setLoadBalancer(loadBalancer);
Expand All @@ -575,7 +609,7 @@ hspec, user(lbprops.getOriginalProperties()),
"A higher priority node than " + chosenHost + " is available");
loadBalancer.decrementHostToNumConnCount(chosenHost);
newConnection.close();
return getConnectionBalanced(lbprops);
return getConnectionBalanced(lbprops, timedOutHosts);
}
return newConnection;
}
Expand All @@ -595,7 +629,7 @@ hspec, user(lbprops.getOriginalProperties()),
firstException = ex;
}
// TODO log exception go to the next one after adding to failed list
LOGGER.log(Level.FINE,
LOGGER.log(Level.INFO,
"couldn't connect to " + chosenHost + ", adding it to failed host list");
loadBalancer.updateFailedHosts(chosenHost);
} else {
Expand All @@ -607,7 +641,7 @@ hspec, user(lbprops.getOriginalProperties()),
"got exception " + ex.getMessage() + ", while connecting to " + chosenHost);
}
}
chosenHost = loadBalancer.getLeastLoadedServer(failedHosts);
chosenHost = loadBalancer.getLeastLoadedServer(failedHosts, timedOutHosts);
}
return null;
}
Expand Down

0 comments on commit c25ac63

Please sign in to comment.