Fix bug where topology query could execute indefinitely

wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/AuroraHostListProvider.java

@@ -30,7 +30,14 @@
 import java.util.Properties;
 import java.util.Set;
 import java.util.UUID;
+import java.util.concurrent.CompletionService;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorCompletionService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.logging.Logger;
 import org.checkerframework.checker.nullness.qual.NonNull;
@@ -75,6 +82,9 @@ public class AuroraHostListProvider implements DynamicHostListProvider {
               + "This pattern is required to be specified for IP address or custom domain connections to AWS RDS "
               + "clusters. Otherwise, if unspecified, the pattern will be automatically created for AWS RDS clusters.");
 
+  private final Executor NETWORK_TIMEOUT_EXECUTOR = Executors.newSingleThreadExecutor();
+  private final CompletionService NETWORK_TIMEOUT_COMPLETION_SERVICE =
+      new ExecutorCompletionService(NETWORK_TIMEOUT_EXECUTOR);
   private final HostListProviderService hostListProviderService;
   private final String originalUrl;
   private RdsUrlType rdsUrlType;
@@ -93,6 +103,8 @@ public class AuroraHostListProvider implements DynamicHostListProvider {
   public static final CacheMap<String, String> suggestedPrimaryClusterIdCache = new CacheMap<>();
   public static final CacheMap<String, Boolean> primaryClusterIdCache = new CacheMap<>();
 
+  private static final int defaultTopologyQueryTimeoutMs = 5000;
+  private static final int setNetworkTimeoutMs = 10000;
   private final ReentrantLock lock = new ReentrantLock();
   protected String clusterId;
   protected HostSpec clusterInstanceTemplate;
@@ -333,11 +345,42 @@ protected List<HostSpec> queryForTopology(final Connection conn) throws SQLExcep
       this.topologyAwareDialect = (TopologyAwareDatabaseCluster) this.hostListProviderService.getDialect();
     }
 
+    int networkTimeout = -1;
+    try {
+      networkTimeout = conn.getNetworkTimeout();
+      // The topology query is not monitored by the EFM plugin, so it needs a socket timeout
+      if (networkTimeout == 0) {
+        setNetworkTimeout(conn, defaultTopologyQueryTimeoutMs);
+      }
+    } catch (SQLException e) {
+      LOGGER.warning(() -> Messages.get("AuroraHostListProvider.errorGettingNetworkTimeout",
+          new Object[] {e.getMessage()}));
+    }
+
     try (final Statement stmt = conn.createStatement();
         final ResultSet resultSet = stmt.executeQuery(this.topologyAwareDialect.getTopologyQuery())) {
       return processQueryResults(resultSet);
     } catch (final SQLSyntaxErrorException e) {
       throw new SQLException(Messages.get("AuroraHostListProvider.invalidQuery"), e);
+    } finally {
+      if (networkTimeout == 0 && !conn.isClosed()) {
+        setNetworkTimeout(conn, networkTimeout);
+      }
+    }
+  }
+
+  private void setNetworkTimeout(final Connection conn, final int timeoutMs) throws SQLException {
+    Future<Boolean> future = NETWORK_TIMEOUT_COMPLETION_SERVICE.submit(() -> {
+      conn.setNetworkTimeout(NETWORK_TIMEOUT_EXECUTOR, timeoutMs);
+      return true;
+    });
+
+    try {
+      future.get(setNetworkTimeoutMs, TimeUnit.MILLISECONDS);
+    } catch (InterruptedException | ExecutionException | TimeoutException e) {
+      throw new SQLException(
+          Messages.get("AuroraHostListProvider.errorSettingNetworkTimeout",
+          new Object[]{e.getMessage()}));
     }
   }
 

wrapper/src/main/resources/messages.properties

@@ -20,6 +20,7 @@ AuroraHostListConnectionPlugin.providerAlreadySet=Another dynamic host list prov
 # Aurora Host List Provider
 AuroraHostListProvider.clusterInstanceHostPatternRequired=The ''clusterInstanceHostPattern'' configuration property is required when an IP address or custom domain is used to connect to a cluster that provides topology information. If you would instead like to connect without failover functionality, set the 'enableClusterAwareFailover' configuration property to false.
 AuroraHostListProvider.clusterInstanceHostPatternNotSupportedForRDSProxy=An RDS Proxy url can''t be used as the 'clusterInstanceHostPattern' configuration setting.
+AuroraHostListProvider.errorSettingNetworkTimeout=An error occurred while adjusting the network timeout for the topology query: ''{0}''
 AuroraHostListProvider.invalidPattern=Invalid value for the 'clusterInstanceHostPattern' configuration setting - the host pattern must contain a '?' character as a placeholder for the DB instance identifiers of the instances in the cluster.
 AuroraHostListProvider.invalidTopology=The topology query returned an invalid topology - no writer instance detected.
 AuroraHostListProvider.suggestedClusterId=ClusterId ''{0}'' is suggested for url ''{1}''.
@@ -28,6 +29,7 @@ AuroraHostListProvider.invalidQuery=Error obtaining host list. Provided database
 AuroraHostListProvider.invalidDialect=Expecting a dialect that supports a cluster topology.
 AuroraHostListProvider.invalidDialectForGetHostRole=An Aurora dialect is required to analyze a host's role. The detected dialect was ''{0}''
 AuroraHostListProvider.errorGettingHostRole=An error occurred while obtaining the connected host's role. This could occur if the connection is broken or if you are not connected to an Aurora database.
+AuroraHostListProvider.errorGettingNetworkTimeout=An error occurred while getting the connection network timeout: {0}
 
 # AWS Credentials Manager
 AwsCredentialsManager.nullProvider=The configured AwsCredentialsProvider was null. If you have configured the AwsCredentialsManager to use a custom AwsCredentialsProviderHandler, please ensure the handler does not return null.