limitless - address PR comments...fix multi cluster lock scenario in limitless monitor service

Author: aaronchung-bitquill

wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessConnectionPlugin.java

@@ -38,6 +38,7 @@
 import software.amazon.jdbc.hostavailability.HostAvailability;
 import software.amazon.jdbc.plugin.AbstractConnectionPlugin;
 import software.amazon.jdbc.util.Messages;
+import software.amazon.jdbc.util.Utils;
 import software.amazon.jdbc.wrapper.HighestWeightHostSelector;
 
 public class LimitlessConnectionPlugin extends AbstractConnectionPlugin {
@@ -154,7 +155,7 @@ private Connection connectInternalWithDialect(
         this.pluginService.getHostListProvider().getClusterId(), props);
 
     Connection conn = null;
-    if (limitlessRouters.isEmpty()) {
+    if (Utils.isNullOrEmpty(limitlessRouters)) {
       conn = connectFunc.call();
       LOGGER.finest(Messages.get("LimitlessConnectionPlugin.limitlessRouterCacheEmpty"));
       final boolean waitForRouterInfo = WAIT_F0R_ROUTER_INFO.getBoolean(props);
@@ -169,7 +170,11 @@ private Connection connectInternalWithDialect(
     if (limitlessRouters.contains(hostSpec)) {
       LOGGER.finest(Messages.get("LimitlessConnectionPlugin.connectWithHost", new Object[] {hostSpec.getHost()}));
       if (conn == null  || conn.isClosed()) {
-        conn = connectFunc.call();
+        try {
+          conn = connectFunc.call();
+        } catch (final SQLException e) {
+          return retryConnectWithLeastLoadedRouters(limitlessRouters, props, conn, hostSpec);
+        }
       }
       return conn;
     }
@@ -184,7 +189,10 @@ private Connection connectInternalWithDialect(
           new Object[] {selectedHostSpec.getHost()}));
     } catch (SQLException e) {
       LOGGER.warning(Messages.get("LimitlessConnectionPlugin.errorSelectingRouter", new Object[] {e.getMessage()}));
-      return retryConnectWithLeastLoadedRouters(limitlessRouters, props, conn, hostSpec, e);
+      if (conn == null  || conn.isClosed()) {
+        conn = connectFunc.call();
+      }
+      return retryConnectWithLeastLoadedRouters(limitlessRouters, props, conn, hostSpec);
     }
 
     try {
@@ -194,8 +202,11 @@ private Connection connectInternalWithDialect(
           "LimitlessConnectionPlugin.failedToConnectToHost",
           new Object[] {selectedHostSpec.getHost()}));
       selectedHostSpec.setAvailability(HostAvailability.NOT_AVAILABLE);
+      if (conn == null  || conn.isClosed()) {
+        conn = connectFunc.call();
+      }
       // Retry connect prioritising healthiest router for best chance of connection over load-balancing with round-robin
-      return retryConnectWithLeastLoadedRouters(limitlessRouters, props, conn, hostSpec, e);
+      return retryConnectWithLeastLoadedRouters(limitlessRouters, props, conn, hostSpec);
     }
   }
 
@@ -222,7 +233,7 @@ private Connection connectInternalWithoutDialect(
 
     List<HostSpec> limitlessRouters = this.limitlessRouterService.getLimitlessRouters(
         this.pluginService.getHostListProvider().getClusterId(), props);
-    if (limitlessRouters.isEmpty()) {
+    if (Utils.isNullOrEmpty(limitlessRouters)) {
       LOGGER.finest(Messages.get("LimitlessConnectionPlugin.limitlessRouterCacheEmpty"));
       final boolean waitForRouterInfo = WAIT_F0R_ROUTER_INFO.getBoolean(props);
       if (waitForRouterInfo) {
@@ -240,7 +251,7 @@ private void initLimitlessRouterMonitorService() {
   }
 
   private Connection retryConnectWithLeastLoadedRouters(final List<HostSpec> limitlessRouters, final Properties props,
-      final Connection conn, final HostSpec hostSpec, final SQLException originalException) throws SQLException {
+      final Connection conn, final HostSpec hostSpec) throws SQLException {
 
     List<HostSpec> currentRouters = limitlessRouters;
     int retryCount = 0;
@@ -252,9 +263,8 @@ private Connection retryConnectWithLeastLoadedRouters(final List<HostSpec> limit
         if (currentRouters == null
             || currentRouters.isEmpty()
             || currentRouters.stream().noneMatch(h -> h.getAvailability().equals(HostAvailability.AVAILABLE))) {
-          throw new SQLException(
-              Messages.get("LimitlessConnectionPlugin.noRoutersAvailableForRetry"),
-              originalException);
+          LOGGER.warning(Messages.get("LimitlessConnectionPlugin.noRoutersAvailableForRetry"));
+          return conn;
         }
       }
 
@@ -266,21 +276,25 @@ private Connection retryConnectWithLeastLoadedRouters(final List<HostSpec> limit
         LOGGER.finest(Messages.get(
             "LimitlessConnectionPlugin.selectedHostForRetry",
             new Object[] {selectedHostSpec.getHost()}));
-      } catch (UnsupportedOperationException e) {
+      } catch (final UnsupportedOperationException e) {
         LOGGER.severe(Messages.get("LimitlessConnectionPlugin.incorrectConfiguration"));
         throw e;
+      } catch (final SQLException e) {
+        // error from host selector
+        continue;
       }
 
       try {
         return pluginService.connect(selectedHostSpec, props);
-      } catch (SQLException e) {
+      } catch (final SQLException e) {
         selectedHostSpec.setAvailability(HostAvailability.NOT_AVAILABLE);
         LOGGER.finest(Messages.get(
             "LimitlessConnectionPlugin.failedToConnectToHost",
             new Object[] {selectedHostSpec.getHost()}));
       }
     }
-    throw new SQLException(Messages.get("LimitlessConnectionPlugin.noRoutersAvailableForRetry"), originalException);
+    LOGGER.warning(Messages.get("LimitlessConnectionPlugin.maxRetriesExceeded"));
+    return conn;
   }
 
   private List<HostSpec> synchronouslyGetLimitlessRoutersWithRetry(

wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImpl.java

@@ -18,9 +18,10 @@
 
 import java.sql.Connection;
 import java.sql.SQLException;
-import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.Properties;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.ReentrantLock;
 import org.checkerframework.checker.nullness.qual.NonNull;
@@ -37,7 +38,7 @@ public class LimitlessRouterServiceImpl implements LimitlessRouterService {
           "600000", // 10min
           "Interval in milliseconds for an Limitless router monitor to be considered inactive and to be disposed.");
   protected static final long CACHE_CLEANUP_NANO = TimeUnit.MINUTES.toNanos(1);
-  protected static final ReentrantLock forceGetLimitlessRoutersLock = new ReentrantLock();
+  protected static final Map<String, ReentrantLock> forceGetLimitlessRoutersLockMap = new ConcurrentHashMap<>();
   protected final PluginService pluginService;
   protected final LimitlessQueryHelper queryHelper;
   protected final LimitlessRouterMonitorInitializer limitlessRouterMonitorInitializer;
@@ -80,11 +81,7 @@ public LimitlessRouterServiceImpl(
   public List<HostSpec> getLimitlessRouters(final String clusterId, final Properties props) throws SQLException {
     final long cacheExpirationNano = TimeUnit.MILLISECONDS.toNanos(
         MONITOR_DISPOSAL_TIME_MS.getLong(props));
-    final List<HostSpec> limitlessRouters = limitlessRouterCache.get(clusterId, cacheExpirationNano);
-    if (limitlessRouters == null) {
-      return Collections.emptyList();
-    }
-    return limitlessRouters;
+    return limitlessRouterCache.get(clusterId, cacheExpirationNano);
   }
 
   @Override
@@ -93,7 +90,11 @@ public List<HostSpec> forceGetLimitlessRoutersWithConn(
     final long cacheExpirationNano = TimeUnit.MILLISECONDS.toNanos(
         MONITOR_DISPOSAL_TIME_MS.getLong(props));
 
-    forceGetLimitlessRoutersLock.lock();
+    final ReentrantLock lock = forceGetLimitlessRoutersLockMap.computeIfAbsent(
+        this.pluginService.getHostListProvider().getClusterId(),
+        (key) -> new ReentrantLock()
+    );
+    lock.lock();
     try {
       final List<HostSpec> limitlessRouters =
           limitlessRouterCache.get(this.pluginService.getHostListProvider().getClusterId(), cacheExpirationNano);
@@ -109,7 +110,7 @@ public List<HostSpec> forceGetLimitlessRoutersWithConn(
           newLimitlessRouters, LimitlessRouterServiceImpl.MONITOR_DISPOSAL_TIME_MS.getLong(props));
       return newLimitlessRouters;
     } finally {
-      forceGetLimitlessRoutersLock.unlock();
+      lock.unlock();
     }
   }
 

wrapper/src/main/resources/aws_advanced_jdbc_wrapper_messages.properties

@@ -193,23 +193,18 @@ LimitlessConnectionPlugin.incorrectConfiguration=Limitless Connection Plugin is
 LimitlessConnectionPlugin.interruptedThread=Thread was interrupted.
 LimitlessConnectionPlugin.limitlessRouterCacheEmpty=Limitless Router cache is empty. This normal during application start up when the cache is not yet populated.
 LimitlessConnectionPlugin.maxRetriesExceeded=Max number of connection retries has been exceeded.
-LimitlessConnectionPlugin.noRoutersAvailable=There are no transaction routers available.
-LimitlessConnectionPlugin.noRoutersAvailableForRetry=There are no transaction routers available for connection retry.
+LimitlessConnectionPlugin.noRoutersAvailable=No transaction routers available.
+LimitlessConnectionPlugin.noRoutersAvailableForRetry=No transaction routers available for connection retry. Retrying with original connection.
 LimitlessConnectionPlugin.selectedHost=Host {0} has been selected.
 LimitlessConnectionPlugin.selectedHostForRetry=Host {0} has been selected for connection retry.
 LimitlessConnectionPlugin.synchronouslyGetLimitlessRouters=Fetching Limitless Routers synchronously.
 LimitlessConnectionPlugin.unsupportedDialectOrDatabase=Unsupported dialect ''{0}'' encountered. Please ensure JDBC connection parameters are correct, and refer to the documentation to ensure that the connecting database is compatible with the Limitless Connection Plugin.
 
-# Limitless Router Service Impl
-LimitlessRouterServiceImpl.nulLimitlessRouterMonitor=LimitlessRouterMonitor with cluster ID {0} is null.
-
 # Limitless Query Helper
 LimitlessQueryHelper.unsupportedDialectOrDatabase=Unsupported dialect ''{0}'' encountered. Please ensure JDBC connection parameters are correct, and refer to the documentation to ensure that the connecting database is compatible with the Limitless Connection Plugin.
 
 # Limitless Router Monitor
 LimitlessRouterMonitor.exceptionDuringMonitoringStop=Unhandled exception was thrown in Limitless Router Monitoring thread for node {0}.
-LimitlessRouterMonitor.forceGetLimitlessRouters=Fetching Limitless Transaction Routers synchronously.
-LimitlessRouterMonitor.forceGetLimitlessRoutersFailed=Failed to fetch Limitless Routers due to closed connection.
 LimitlessRouterMonitor.interruptedExceptionDuringMonitoring=Limitless Router Monitoring thread for node {0} was interrupted.
 LimitlessRouterMonitor.invalidQuery=Limitless Connection Plugin has encountered an error obtaining Limitless Router endpoints. Please ensure that you are connecting to an Aurora Limitless Database Shard Group Endpoint URL.
 LimitlessRouterMonitor.invalidRouterLoad=Invalid load metric value of ''{1}''from the transaction router query aurora_limitless_router_endpoints() for transaction router ''{0}''. The load metric value must be a decimal value between 0 and 1. Host weight be assigned a default weight of 1.

wrapper/src/test/java/software/amazon/jdbc/plugin/limitless/LimitlessConnectionPluginTest.java

@@ -302,9 +302,10 @@ void testConnect_givenMaxRetries_throwsSqlException() throws SQLException {
     final int maxRetries = 7;
     propsWithMaxRetries.setProperty(LimitlessConnectionPlugin.MAX_RETRIES.name, String.valueOf(maxRetries));
 
-    assertThrows(
-        SQLException.class,
-        () -> plugin.connect(DRIVER_PROTOCOL, INPUT_HOST_SPEC, propsWithMaxRetries, true, mockConnectFuncLambda));
+
+    final Connection actualConn = plugin
+        .connect(DRIVER_PROTOCOL, INPUT_HOST_SPEC, propsWithMaxRetries, true, mockConnectFuncLambda);
+    assertEquals(mockConnection, actualConn);
 
     verify(mockLimitlessRouterService, times(1)).startMonitoring(INPUT_HOST_SPEC,
         props, Integer.parseInt(LimitlessConnectionPlugin.INTERVAL_MILLIS.defaultValue));