use ExecutorService to manage node monitoring threads

sergiyvamz · sergiyvamz · commit 9533dae2cbf2 · 2025-03-18T08:18:30.000-07:00
diff --git a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/ClusterTopologyMonitorImpl.java b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/ClusterTopologyMonitorImpl.java
@@ -37,6 +37,7 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
@@ -98,7 +99,9 @@ public class ClusterTopologyMonitorImpl implements ClusterTopologyMonitor {
   protected final Object topologyUpdated = new Object();
   protected final AtomicBoolean requestToUpdateTopology = new AtomicBoolean(false);
   protected final AtomicLong ignoreNewTopologyRequestsEndTimeNano = new AtomicLong(-1);
-  protected final ConcurrentHashMap<String /* host */, Thread> nodeThreads = new ConcurrentHashMap<>();
+  protected final ConcurrentHashMap<String, Boolean> submittedNodes = new ConcurrentHashMap<>();
+  protected ExecutorService nodeExecutorService = null;
+  protected final ReentrantLock nodeExecutorLock = new ReentrantLock();
   protected final AtomicBoolean nodeThreadsStop = new AtomicBoolean(false);
   protected final AtomicReference<Connection> nodeThreadsWriterConnection = new AtomicReference<>(null);
   protected final AtomicReference<HostSpec> nodeThreadsWriterHostSpec = new AtomicReference<>(null);
@@ -266,6 +269,7 @@ protected List<HostSpec> waitTillTopologyGetsUpdated(final long timeoutMs) throw
   public void close() throws Exception {
     this.stop.set(true);
     this.nodeThreadsStop.set(true);
+    this.shutdownNodeExecutorService();
 
     // It breaks a waiting/sleeping cycles in monitoring thread
     synchronized (this.requestToUpdateTopology) {
@@ -277,8 +281,6 @@ public void close() throws Exception {
     if (!this.monitorExecutor.awaitTermination(30, TimeUnit.SECONDS)) {
       this.monitorExecutor.shutdownNow();
     }
-
-    this.nodeThreads.clear();
   }
 
   @Override
@@ -292,7 +294,7 @@ public void run() {
 
         if (this.isInPanicMode()) {
 
-          if (this.nodeThreads.isEmpty()) {
+          if (this.submittedNodes.isEmpty()) {
             LOGGER.finest(Messages.get("ClusterTopologyMonitorImpl.startingNodeMonitoringThreads"));
 
             // start node threads
@@ -308,15 +310,19 @@ public void run() {
               hosts = this.openAnyConnectionAndUpdateTopology();
             }
 
+            this.shutdownNodeExecutorService();
+            this.createNodeExecutorService();
+
             if (hosts != null && !this.isVerifiedWriterConnection) {
               for (HostSpec hostSpec : hosts) {
-                this.nodeThreads.computeIfAbsent(hostSpec.getHost(),
+                this.submittedNodes.computeIfAbsent(hostSpec.getHost(),
                     (key) -> {
-                      final Thread thread = this.getNodeMonitoringThread(hostSpec, this.writerHostSpec.get());
-                      thread.start();
-                      return thread;
+                      this.nodeExecutorService.submit(
+                          this.getNodeMonitoringWorker(hostSpec, this.writerHostSpec.get()));
+                      return true;
                     });
               }
+              // It's not possible to call shutdown() on this.nodeExecutorService since more node may be added later.
             }
             // otherwise let's try it again the next round
 
@@ -345,10 +351,8 @@ public void run() {
               }
 
               this.nodeThreadsStop.set(true);
-              for (Thread thread : this.nodeThreads.values()) {
-                thread.interrupt();
-              }
-              this.nodeThreads.clear();
+              this.shutdownNodeExecutorService();
+              this.submittedNodes.clear();
 
               continue;
 
@@ -357,13 +361,14 @@ public void run() {
               List<HostSpec> hosts = this.nodeThreadsLatestTopology.get();
               if (hosts != null && !this.nodeThreadsStop.get()) {
                 for (HostSpec hostSpec : hosts) {
-                  this.nodeThreads.computeIfAbsent(hostSpec.getHost(),
+                  this.submittedNodes.computeIfAbsent(hostSpec.getHost(),
                       (key) -> {
-                        final Thread thread = this.getNodeMonitoringThread(hostSpec, this.writerHostSpec.get());
-                        thread.start();
-                        return thread;
+                        this.nodeExecutorService.submit(
+                            this.getNodeMonitoringWorker(hostSpec, this.writerHostSpec.get()));
+                        return true;
                       });
                 }
+                // It's not possible to call shutdown() on this.nodeExecutorService since more node may be added later.
               }
             }
           }
@@ -373,12 +378,9 @@ public void run() {
         } else {
           // regular mode (not panic mode)
 
-          if (!this.nodeThreads.isEmpty()) {
-            // stop node threads
-            for (Thread thread : this.nodeThreads.values()) {
-              thread.interrupt();
-            }
-            this.nodeThreads.clear();
+          if (!this.submittedNodes.isEmpty()) {
+            this.shutdownNodeExecutorService();
+            this.submittedNodes.clear();
           }
 
           final List<HostSpec> hosts = this.fetchTopologyAndUpdateCache(this.monitoringConnection.get());
@@ -427,6 +429,7 @@ public void run() {
 
     } finally {
       this.stop.set(true);
+      this.shutdownNodeExecutorService();
 
       final Connection conn = this.monitoringConnection.get();
       this.monitoringConnection.set(null);
@@ -438,13 +441,58 @@ public void run() {
     }
   }
 
+  protected void shutdownNodeExecutorService() {
+    if (this.nodeExecutorService != null) {
+
+      this.nodeExecutorLock.lock();
+      try {
+
+        if (this.nodeExecutorService == null) {
+          return;
+        }
+
+        if (!this.nodeExecutorService.isShutdown()) {
+          this.nodeExecutorService.shutdown();
+        }
+
+        try {
+          if (!this.nodeExecutorService.awaitTermination(30, TimeUnit.SECONDS)) {
+            this.nodeExecutorService.shutdownNow();
+          }
+        } catch (InterruptedException e) {
+          // do nothing
+        }
+
+        this.nodeExecutorService = null;
+      } finally {
+        this.nodeExecutorLock.unlock();
+      }
+    }
+  }
+
+  protected void createNodeExecutorService() {
+    this.nodeExecutorLock.lock();
+    try {
+      this.nodeExecutorService = Executors.newCachedThreadPool(runnableTarget -> {
+        final Thread monitoringThread = new Thread(runnableTarget);
+        monitoringThread.setDaemon(true);
+        if (!StringUtils.isNullOrEmpty(monitoringThread.getName())) {
+          monitoringThread.setName(monitoringThread.getName() + "-nm");
+        }
+        return monitoringThread;
+      });
+    } finally {
+      this.nodeExecutorLock.unlock();
+    }
+  }
+
   protected boolean isInPanicMode() {
     return this.monitoringConnection.get() == null
         || !this.isVerifiedWriterConnection;
   }
 
-  protected Thread getNodeMonitoringThread(final HostSpec hostSpec, final @Nullable HostSpec writerHostSpec) {
-    return new NodeMonitoringThread(this, hostSpec, writerHostSpec);
+  protected Runnable getNodeMonitoringWorker(final HostSpec hostSpec, final @Nullable HostSpec writerHostSpec) {
+    return new NodeMonitoringWorker(this, hostSpec, writerHostSpec);
   }
 
   protected List<HostSpec> openAnyConnectionAndUpdateTopology() {
@@ -536,7 +584,7 @@ protected String getNodeId(final Connection connection) {
   }
 
   protected void closeConnection(final @Nullable Connection connection) {
-    this.closeConnection(connection, false);
+    this.closeConnection(connection, true);
   }
 
   protected void closeConnection(final @Nullable Connection connection, final boolean unstableConnection) {
@@ -756,27 +804,23 @@ protected String getHostEndpoint(final String nodeName) {
     return host.replace("?", nodeName);
   }
 
-  private static class NodeMonitoringThread extends Thread {
+  private static class NodeMonitoringWorker implements Runnable {
 
-    private static final Logger LOGGER = Logger.getLogger(NodeMonitoringThread.class.getName());
+    private static final Logger LOGGER = Logger.getLogger(NodeMonitoringWorker.class.getName());
 
     protected final ClusterTopologyMonitorImpl monitor;
     protected final HostSpec hostSpec;
     protected final @Nullable HostSpec writerHostSpec;
     protected boolean writerChanged = false;
 
-    public NodeMonitoringThread(
+    public NodeMonitoringWorker(
         final ClusterTopologyMonitorImpl monitor,
         final HostSpec hostSpec,
         final @Nullable HostSpec writerHostSpec
     ) {
       this.monitor = monitor;
       this.hostSpec = hostSpec;
       this.writerHostSpec = writerHostSpec;
-
-      if (!StringUtils.isNullOrEmpty(this.getName())) {
-        this.setName(this.getName() + "-nm");
-      }
     }
 
     @Override
@@ -837,7 +881,7 @@ public void run() {
                     this.monitor.topologyMap.get(this.monitor.clusterId)));
               }
 
-              // Setting the connection to null here prevents the finally block
+              // Setting the connection to null here prevents the final block
               // from closing nodeThreadsWriterConnection.
               connection = null;
               return;
diff --git a/wrapper/src/main/resources/aws_advanced_jdbc_wrapper_messages.properties b/wrapper/src/main/resources/aws_advanced_jdbc_wrapper_messages.properties
@@ -403,7 +403,7 @@ ClusterTopologyMonitorImpl.exceptionDuringMonitoringStop=Stopping cluster topolo
 ClusterTopologyMonitorImpl.invalidQuery=An error occurred while attempting to obtain the topology because the topology query was invalid. Please ensure you are connecting to an Aurora or RDS Db cluster.
 ClusterTopologyMonitorImpl.errorGettingNetworkTimeout=An error occurred while getting the connection network timeout: {0}
 ClusterTopologyMonitorImpl.invalidTopology=The topology query returned an invalid topology - no writer instance detected.
-ClusterTopologyMonitorImpl.topologyNotUpdated=Topology hasn't been updated after {0} ms.
+ClusterTopologyMonitorImpl.topologyNotUpdated=Topology hasn''t been updated after {0} ms.
 ClusterTopologyMonitorImpl.openedMonitoringConnection=Opened monitoring connection to node ''{0}''.
 ClusterTopologyMonitorImpl.ignoringTopologyRequest=A topology refresh was requested, but the topology was already updated recently. Returning cached hosts:
 ClusterTopologyMonitorImpl.timeoutSetToZero=A topology refresh was requested, but the given timeout for the request was 0ms. Returning cached hosts:
