diff --git a/pgjdbc/src/main/java/com/yugabyte/ysql/LoadBalanceManager.java b/pgjdbc/src/main/java/com/yugabyte/ysql/LoadBalanceManager.java index 27c86212d3..d40793859d 100644 --- a/pgjdbc/src/main/java/com/yugabyte/ysql/LoadBalanceManager.java +++ b/pgjdbc/src/main/java/com/yugabyte/ysql/LoadBalanceManager.java @@ -235,7 +235,8 @@ public static boolean decrementConnectionCount(String host) { public static Connection getConnection(String url, Properties properties, String user, String database) { - LoadBalanceProperties lbProperties = new LoadBalanceProperties(url, properties); + LoadBalanceProperties lbProperties = LoadBalanceProperties.getLoadBalanceProperties(url, + properties); // Cleanup extra properties used for load balancing? if (lbProperties.hasLoadBalance()) { Connection conn = getConnection(lbProperties, user, database); diff --git a/pgjdbc/src/main/java/com/yugabyte/ysql/LoadBalanceProperties.java b/pgjdbc/src/main/java/com/yugabyte/ysql/LoadBalanceProperties.java index cceec39cf8..7be7c20d59 100644 --- a/pgjdbc/src/main/java/com/yugabyte/ysql/LoadBalanceProperties.java +++ b/pgjdbc/src/main/java/com/yugabyte/ysql/LoadBalanceProperties.java @@ -15,6 +15,7 @@ import java.util.HashMap; import java.util.Map; import java.util.Properties; +import java.util.concurrent.ConcurrentHashMap; import java.util.logging.Level; import java.util.logging.Logger; @@ -52,6 +53,8 @@ public class LoadBalanceProperties { public static final Map CONNECTION_MANAGER_MAP = new HashMap<>(); + private static Map loadBalancePropertiesMap = + new ConcurrentHashMap(); private final String originalUrl; private final Properties originalProperties; private boolean hasLoadBalance; @@ -62,7 +65,22 @@ public class LoadBalanceProperties { private boolean refreshIntervalSpecified; private boolean explicitFallbackOnlySpecified; - public LoadBalanceProperties(String origUrl, Properties origProperties) { + public static LoadBalanceProperties getLoadBalanceProperties(String url, Properties properties) { + LoadBalancerKey key = new LoadBalancerKey(url, properties); + LoadBalanceProperties lbp = loadBalancePropertiesMap.get(key); + if (lbp == null) { + synchronized (LoadBalanceProperties.class) { + lbp = loadBalancePropertiesMap.get(key); + if (lbp == null) { + lbp = new LoadBalanceProperties(url, properties); + loadBalancePropertiesMap.put(key, lbp); + } + } + } + return lbp; + } + + private LoadBalanceProperties(String origUrl, Properties origProperties) { originalUrl = origUrl; originalProperties = (Properties) origProperties.clone(); ybURL = processURLAndProperties(); @@ -198,9 +216,6 @@ public LoadBalancer getAppropriateLoadBalancer() { if (refreshIntervalSpecified) { System.setProperty(REFRESH_INTERVAL_KEY, String.valueOf(refreshInterval)); } - if (explicitFallbackOnlySpecified) { - System.setProperty(EXPLICIT_FALLBACK_ONLY_KEY, String.valueOf(explicitFallbackOnly)); - } LoadBalancer ld = null; if (placements == null) { // return base class conn manager. @@ -215,17 +230,35 @@ public LoadBalancer getAppropriateLoadBalancer() { } } } else { - ld = CONNECTION_MANAGER_MAP.get(placements); + String key = placements + "&" + String.valueOf(explicitFallbackOnly).toLowerCase(); + ld = CONNECTION_MANAGER_MAP.get(key); if (ld == null) { synchronized (CONNECTION_MANAGER_MAP) { - ld = CONNECTION_MANAGER_MAP.get(placements); + ld = CONNECTION_MANAGER_MAP.get(key); if (ld == null) { - ld = new TopologyAwareLoadBalancer(placements); - CONNECTION_MANAGER_MAP.put(placements, ld); + ld = new TopologyAwareLoadBalancer(placements, explicitFallbackOnly); + CONNECTION_MANAGER_MAP.put(key, ld); } } } } return ld; } + + private static class LoadBalancerKey { + private String url; + private Properties properties; + + public LoadBalancerKey(String url, Properties properties) { + this.url = url; + this.properties = properties; + } + + public boolean equals(Object other) { + return other instanceof LoadBalancerKey && + url != null && url.equals(((LoadBalancerKey) other).url) && + properties != null && + properties.equals(((LoadBalancerKey) other).properties); + } + } } diff --git a/pgjdbc/src/main/java/com/yugabyte/ysql/TopologyAwareLoadBalancer.java b/pgjdbc/src/main/java/com/yugabyte/ysql/TopologyAwareLoadBalancer.java index fb76de9cfe..ab4cc0aee6 100644 --- a/pgjdbc/src/main/java/com/yugabyte/ysql/TopologyAwareLoadBalancer.java +++ b/pgjdbc/src/main/java/com/yugabyte/ysql/TopologyAwareLoadBalancer.java @@ -40,10 +40,12 @@ public class TopologyAwareLoadBalancer implements LoadBalancer { private int currentPlacementIndex = 1; List attempted = new ArrayList<>(); private int refreshIntervalSeconds; + private boolean explicitFallbackOnly = false; - public TopologyAwareLoadBalancer(String placementValues) { + public TopologyAwareLoadBalancer(String placementValues, boolean onlyExplicitFallback) { placements = placementValues; + explicitFallbackOnly = onlyExplicitFallback; refreshIntervalSeconds = Integer.getInteger(REFRESH_INTERVAL_KEY, DEFAULT_REFRESH_INTERVAL); parseGeoLocations(); } @@ -97,11 +99,14 @@ public int getRefreshListSeconds() { return Integer.getInteger(REFRESH_INTERVAL_KEY, refreshIntervalSeconds); } + public boolean isExplicitFallbackOnly() { + return explicitFallbackOnly; + } + @Override public boolean isHostEligible(Map.Entry e) { Set set = allowedPlacements.get(currentPlacementIndex); - boolean onlyExplicitFallback = Boolean.getBoolean(EXPLICIT_FALLBACK_ONLY_KEY); - boolean found = (currentPlacementIndex == REST_OF_CLUSTER_INDEX && !onlyExplicitFallback) + boolean found = (currentPlacementIndex == REST_OF_CLUSTER_INDEX && !explicitFallbackOnly) || (set != null && e.getValue().getPlacement().isContainedIn(set)); boolean isAttempted = attempted.contains(e.getKey()); boolean isDown = e.getValue().isDown();