- Make explicit-fallback-only a part of balancer instance

- Cache LoadBalanceProperties instances for given url and properties
yugabyte · Apr 23, 2024 · a7eb111 · a7eb111
1 parent 14ed39f
commit a7eb111
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 12 deletions.
diff --git 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
@@ -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<String, LoadBalancer> CONNECTION_MANAGER_MAP =
       new HashMap<>();
 
+  private static Map<LoadBalancerKey, LoadBalanceProperties> 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
@@ -40,10 +40,12 @@ public class TopologyAwareLoadBalancer implements LoadBalancer {
   private int currentPlacementIndex = 1;
   List<String> 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<String, LoadBalanceManager.NodeInfo> e) {
     Set<LoadBalanceManager.CloudPlacement> 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();