fix: issues with custom domains during failover

wrapper/src/main/java/software/amazon/jdbc/dialect/DialectManager.java

@@ -29,6 +29,7 @@
 import software.amazon.jdbc.Driver;
 import software.amazon.jdbc.HostSpec;
 import software.amazon.jdbc.PluginService;
+import software.amazon.jdbc.PropertyDefinition;
 import software.amazon.jdbc.util.CacheMap;
 import software.amazon.jdbc.util.ConnectionUrlParser;
 import software.amazon.jdbc.util.Messages;
@@ -81,7 +82,11 @@ public class DialectManager implements DialectProvider {
   private Dialect dialect = null;
   private String dialectCode;
 
-  private PluginService pluginService;
+  private final PluginService pluginService;
+
+  static {
+    PropertyDefinition.registerPluginProperties(DialectManager.class);
+  }
 
   public DialectManager(PluginService pluginService) {
     this.pluginService = pluginService;

wrapper/src/main/java/software/amazon/jdbc/plugin/AuroraConnectionTrackerPlugin.java

@@ -25,6 +25,8 @@
 import java.util.Map;
 import java.util.Properties;
 import java.util.Set;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.logging.Logger;
 import org.checkerframework.checker.nullness.qual.NonNull;
 import software.amazon.jdbc.HostRole;
@@ -43,17 +45,24 @@ public class AuroraConnectionTrackerPlugin extends AbstractConnectionPlugin impl
 
   private static final Logger LOGGER = Logger.getLogger(AuroraConnectionTrackerPlugin.class.getName());
 
+  // Check topology changes 3 min after last failover
+  private static final long TOPOLOGY_CHANGES_EXPECTED_TIME_NANO = TimeUnit.MINUTES.toNanos(3);
+
   static final String METHOD_ABORT = "Connection.abort";
   static final String METHOD_CLOSE = "Connection.close";
   private static final Set<String> subscribedMethods =
       Collections.unmodifiableSet(new HashSet<String>() {
         {
           addAll(SubscribedMethodHelper.NETWORK_BOUND_METHODS);
+          add(METHOD_CLOSE);
+          add(METHOD_ABORT);
           add("connect");
           add("notifyNodeListChanged");
         }
       });
 
+  private static final AtomicLong hostListRefreshEndTimeNano = new AtomicLong(0);
+
   private final PluginService pluginService;
   private final RdsUtils rdsHelper;
   private final OpenedConnectionTracker tracker;
@@ -87,7 +96,7 @@ public Connection connect(final String driverProtocol, final HostSpec hostSpec,
 
     if (conn != null) {
       final RdsUrlType type = this.rdsHelper.identifyRdsType(hostSpec.getHost());
-      if (type.isRdsCluster()) {
+      if (type.isRdsCluster() || type == RdsUrlType.OTHER) {
         hostSpec.resetAliases();
         this.pluginService.fillAliases(conn, hostSpec);
       }
@@ -106,23 +115,50 @@ public <T, E extends Exception> T execute(final Class<T> resultClass, final Clas
     this.rememberWriter();
 
     try {
+      if (!methodName.equals(METHOD_CLOSE) && !methodName.equals(METHOD_ABORT)) {
+        long localHostListRefreshEndTimeNano = hostListRefreshEndTimeNano.get();
+        boolean needRefreshHostLists = false;
+        if (localHostListRefreshEndTimeNano > 0) {
+          if (localHostListRefreshEndTimeNano > System.nanoTime()) {
+            // The time specified in hostListRefreshThresholdTimeNano isn't yet reached.
+            // Need to continue to refresh host list.
+            needRefreshHostLists = true;
+          } else {
+            // The time specified in hostListRefreshThresholdTimeNano is reached, and we can stop further refreshes
+            // of host list. If hostListRefreshThresholdTimeNano has changed while this thread processes the code,
+            // we can't override a new value in hostListRefreshThresholdTimeNano.
+            hostListRefreshEndTimeNano.compareAndSet(localHostListRefreshEndTimeNano, 0);
+          }
+        }
+        if (this.needUpdateCurrentWriter || needRefreshHostLists) {
+          // Calling this method may effectively close/abort a current connection
+          this.checkWriterChanged(needRefreshHostLists);
+        }
+      }
       final T result = jdbcMethodFunc.call();
       if ((methodName.equals(METHOD_CLOSE) || methodName.equals(METHOD_ABORT))) {
-        tracker.invalidateCurrentConnection(currentHostSpec, this.pluginService.getCurrentConnection());
-      } else if (this.needUpdateCurrentWriter) {
-        this.checkWriterChanged();
+        tracker.removeConnectionTracking(currentHostSpec, this.pluginService.getCurrentConnection());
       }
       return result;
 
     } catch (final Exception e) {
       if (e instanceof FailoverSQLException) {
-        this.checkWriterChanged();
+        hostListRefreshEndTimeNano.set(System.nanoTime() + TOPOLOGY_CHANGES_EXPECTED_TIME_NANO);
+        // Calling this method may effectively close/abort a current connection
+        this.checkWriterChanged(true);
       }
       throw e;
     }
   }
 
-  private void checkWriterChanged() {
+  private void checkWriterChanged(boolean needRefreshHostLists) {
+    if (needRefreshHostLists) {
+      try {
+        this.pluginService.refreshHostList();
+      } catch (SQLException ex) {
+        // do nothing
+      }
+    }
     final HostSpec hostSpecAfterFailover = this.getWriter(this.pluginService.getAllHosts());
 
     if (this.currentWriter == null) {
@@ -135,6 +171,7 @@ private void checkWriterChanged() {
       tracker.logOpenedConnections();
       this.currentWriter = hostSpecAfterFailover;
       this.needUpdateCurrentWriter = false;
+      hostListRefreshEndTimeNano.set(0);
     }
   }
 

wrapper/src/main/java/software/amazon/jdbc/plugin/AuroraInitialConnectionStrategyPlugin.java

@@ -19,7 +19,9 @@
 import java.sql.Connection;
 import java.sql.SQLException;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
+import java.util.Map;
 import java.util.Properties;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
@@ -67,16 +69,46 @@
           "1000",
           "Time between each retry of opening a connection.");
 
+  public static final AwsWrapperProperty VERIFY_OPENED_CONNECTION_TYPE =
+      new AwsWrapperProperty(
+          "verifyOpenedConnectionType",
+          null,
+          "Force to verify an opened connection to be either a writer or a reader.");
+
+  private enum VerifyOpenedConnectionType {
+    WRITER,
+    READER;
+
+    private static final Map<String, VerifyOpenedConnectionType> nameToValue =
+        new HashMap<String, VerifyOpenedConnectionType>() {
+          {
+            put("writer", WRITER);
+            put("reader", READER);
+          }
+        };
+
+    public static VerifyOpenedConnectionType fromValue(String value) {
+      if (value == null) {
+        return null;
+      }
+      return nameToValue.get(value.toLowerCase());
+    }
+  }
+
   private final PluginService pluginService;
   private HostListProviderService hostListProviderService;
   private final RdsUtils rdsUtils = new RdsUtils();
 
+  private VerifyOpenedConnectionType verifyOpenedConnectionType = null;
+
   static {
     PropertyDefinition.registerPluginProperties(AuroraInitialConnectionStrategyPlugin.class);
   }
 
   public AuroraInitialConnectionStrategyPlugin(final PluginService pluginService, final Properties properties) {
     this.pluginService = pluginService;
+    this.verifyOpenedConnectionType =
+        VerifyOpenedConnectionType.fromValue(VERIFY_OPENED_CONNECTION_TYPE.getString(properties));
   }
 
   @Override
@@ -110,12 +142,8 @@
 
     final RdsUrlType type = this.rdsUtils.identifyRdsType(hostSpec.getHost());
 
-    if (!type.isRdsCluster()) {
-      // It's not a cluster endpoint. Continue with a normal workflow.
-      return connectFunc.call();
-    }
-
-    if (type == RdsUrlType.RDS_WRITER_CLUSTER) {
+    if (type == RdsUrlType.RDS_WRITER_CLUSTER
+        || isInitialConnection && this.verifyOpenedConnectionType == VerifyOpenedConnectionType.WRITER) {
       Connection writerCandidateConn = this.getVerifiedWriterConnection(props, isInitialConnection, connectFunc);
       if (writerCandidateConn == null) {
         // Can't get writer connection. Continue with a normal workflow.
@@ -124,7 +152,8 @@
       return writerCandidateConn;
     }
 
-    if (type == RdsUrlType.RDS_READER_CLUSTER) {
+    if (type == RdsUrlType.RDS_READER_CLUSTER
+        || isInitialConnection && this.verifyOpenedConnectionType == VerifyOpenedConnectionType.READER) {
       Connection readerCandidateConn = this.getVerifiedReaderConnection(props, isInitialConnection, connectFunc);
       if (readerCandidateConn == null) {
         // Can't get a reader connection. Continue with a normal workflow.

wrapper/src/main/java/software/amazon/jdbc/plugin/OpenedConnectionTracker.java

@@ -20,13 +20,14 @@
 import java.sql.Connection;
 import java.sql.SQLException;
 import java.util.Arrays;
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Objects;
-import java.util.Optional;
 import java.util.Queue;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.logging.Logger;
@@ -35,6 +36,7 @@
 import software.amazon.jdbc.util.Messages;
 import software.amazon.jdbc.util.RdsUtils;
 import software.amazon.jdbc.util.StringUtils;
+import software.amazon.jdbc.util.SynchronousExecutor;
 import software.amazon.jdbc.util.telemetry.TelemetryContext;
 import software.amazon.jdbc.util.telemetry.TelemetryFactory;
 import software.amazon.jdbc.util.telemetry.TelemetryTraceLevel;
@@ -50,16 +52,17 @@ public class OpenedConnectionTracker {
             invalidateThread.setDaemon(true);
             return invalidateThread;
           });
-  private static final ExecutorService abortConnectionExecutorService =
-      Executors.newCachedThreadPool(
-          r -> {
-            final Thread abortThread = new Thread(r);
-            abortThread.setDaemon(true);
-            return abortThread;
-          });
+  private static final Executor abortConnectionExecutor = new SynchronousExecutor();
 
   private static final Logger LOGGER = Logger.getLogger(OpenedConnectionTracker.class.getName());
   private static final RdsUtils rdsUtils = new RdsUtils();
+
+  private static final Set<String> safeToCheckClosedClasses = new HashSet<>(Arrays.asList(
+      "HikariProxyConnection",
+      "org.postgresql.jdbc.PgConnection",
+      "com.mysql.cj.jdbc.ConnectionImpl",
+      "org.mariadb.jdbc.Connection"));
+
   private final PluginService pluginService;
 
   public OpenedConnectionTracker(final PluginService pluginService) {
@@ -72,6 +75,7 @@ public void populateOpenedConnectionQueue(final HostSpec hostSpec, final Connect
     // Check if the connection was established using an instance endpoint
     if (rdsUtils.isRdsInstance(hostSpec.getHost())) {
       trackConnection(hostSpec.getHostAndPort(), conn);
+      logOpenedConnections();
       return;
     }
 
@@ -80,14 +84,17 @@ public void populateOpenedConnectionQueue(final HostSpec hostSpec, final Connect
         .max(String::compareToIgnoreCase)
         .orElse(null);
 
-    if (instanceEndpoint == null) {
-      LOGGER.finest(
-          Messages.get("OpenedConnectionTracker.unableToPopulateOpenedConnectionQueue",
-            new Object[] {hostSpec.getHost()}));
+    if (instanceEndpoint != null) {
+      trackConnection(instanceEndpoint, conn);
+      logOpenedConnections();
       return;
     }
 
-    trackConnection(instanceEndpoint, conn);
+    // It seems there's no RDS instance host found. It might be a custom domain name. Let's track by all aliases
+    for (String alias : aliases) {
+      trackConnection(alias, conn);
+    }
+    logOpenedConnections();
   }
 
   /**
@@ -100,28 +107,27 @@ public void invalidateAllConnections(final HostSpec hostSpec) {
     invalidateAllConnections(hostSpec.getAliases().toArray(new String[] {}));
   }
 
-  public void invalidateAllConnections(final String... node) {
+  public void invalidateAllConnections(final String... keys) {
     TelemetryFactory telemetryFactory = this.pluginService.getTelemetryFactory();
     TelemetryContext telemetryContext = telemetryFactory.openTelemetryContext(
         TELEMETRY_INVALIDATE_CONNECTIONS, TelemetryTraceLevel.NESTED);
 
     try {
-      final Optional<String> instanceEndpoint = Arrays.stream(node)
-          .filter(x -> rdsUtils.isRdsInstance(rdsUtils.removePort(x)))
-          .findFirst();
-      if (!instanceEndpoint.isPresent()) {
-        return;
+      for (String key : keys) {
+        try {
+          final Queue<WeakReference<Connection>> connectionQueue = openedConnections.get(key);
+          logConnectionQueue(key, connectionQueue);
+          invalidateConnections(connectionQueue);
+        } catch (Exception ex) {
+          // ignore and continue
+        }
       }
-      final Queue<WeakReference<Connection>> connectionQueue = openedConnections.get(instanceEndpoint.get());
-      logConnectionQueue(instanceEndpoint.get(), connectionQueue);
-      invalidateConnections(openedConnections.get(instanceEndpoint.get()));
-
     } finally {
       telemetryContext.closeContext();
     }
   }
 
-  public void invalidateCurrentConnection(final HostSpec hostSpec, final Connection connection) {
+  public void removeConnectionTracking(final HostSpec hostSpec, final Connection connection) {
     final String host = rdsUtils.isRdsInstance(hostSpec.getHost())
         ? hostSpec.asAlias()
         : hostSpec.getAliases().stream()
@@ -134,8 +140,11 @@ public void invalidateCurrentConnection(final HostSpec hostSpec, final Connectio
     }
 
     final Queue<WeakReference<Connection>> connectionQueue = openedConnections.get(host);
-    logConnectionQueue(host, connectionQueue);
-    connectionQueue.removeIf(connectionWeakReference -> Objects.equals(connectionWeakReference.get(), connection));
+    if (connectionQueue != null) {
+      logConnectionQueue(host, connectionQueue);
+      connectionQueue.removeIf(
+          connectionWeakReference -> Objects.equals(connectionWeakReference.get(), connection));
+    }
   }
 
   private void trackConnection(final String instanceEndpoint, final Connection connection) {
@@ -144,10 +153,12 @@ private void trackConnection(final String instanceEndpoint, final Connection con
             instanceEndpoint,
             (k) -> new ConcurrentLinkedQueue<>());
     connectionQueue.add(new WeakReference<>(connection));
-    logOpenedConnections();
   }
 
   private void invalidateConnections(final Queue<WeakReference<Connection>> connectionQueue) {
+    if (connectionQueue == null || connectionQueue.isEmpty()) {
+      return;
+    }
     invalidateConnectionsExecutorService.submit(() -> {
       WeakReference<Connection> connReference;
       while ((connReference = connectionQueue.poll()) != null) {
@@ -157,7 +168,7 @@ private void invalidateConnections(final Queue<WeakReference<Connection>> connec
         }
 
         try {
-          conn.abort(abortConnectionExecutorService);
+          conn.abort(abortConnectionExecutor);
         } catch (final SQLException e) {
           // swallow this exception, current connection should be useless anyway.
         }
@@ -204,7 +215,10 @@ public void pruneNullConnections() {
         if (conn == null) {
           return true;
         }
-        if (conn.getClass().getSimpleName().equals("HikariProxyConnection")) {
+        // The following classes do not check connection validity by calling a DB server
+        // so it's safe to check whether connection is already closed.
+        if (safeToCheckClosedClasses.contains(conn.getClass().getSimpleName())
+            || safeToCheckClosedClasses.contains(conn.getClass().getName())) {
           try {
             return conn.isClosed();
           } catch (SQLException ex) {