refactor: store monitoring endpoint

Author: karenc-bq

wrapper/src/main/java/software/amazon/jdbc/plugin/efm/HostMonitoringConnectionPlugin.java

@@ -83,9 +83,9 @@ public class HostMonitoringConnectionPlugin extends AbstractConnectionPlugin
   protected @NonNull Properties properties;
   private final @NonNull Supplier<MonitorService> monitorServiceSupplier;
   private final @NonNull PluginService pluginService;
-  private final @NonNull Set<String> nodeKeys = ConcurrentHashMap.newKeySet(); // Shared with monitor thread
   private MonitorService monitorService;
-  private RdsUtils rdsHelper;
+  private final RdsUtils rdsHelper;
+  private HostSpec monitoringHostSpec;
 
   /**
    * Initialize the node monitoring plugin.
@@ -155,34 +155,17 @@ public <T, E extends Exception> T execute(
     T result;
     MonitorConnectionContext monitorContext = null;
 
-    HostSpec monitoringHostSpec = this.pluginService.getCurrentHostSpec();
-    final RdsUrlType rdsUrlType = this.rdsHelper.identifyRdsType(monitoringHostSpec.getUrl());
-
-    try {
-      if (rdsUrlType.isRdsCluster()) {
-        monitoringHostSpec = this.pluginService.identifyConnection(this.pluginService.getCurrentConnection());
-        monitoringHostSpec.resetAliases();
-        this.pluginService.fillAliases(this.pluginService.getCurrentConnection(), monitoringHostSpec);
-      }
-    } catch (SQLException e) {
-      // Log and ignore
-      LOGGER.finest(Messages.get("HostMonitoringConnectionPlugin.errorIdentifyingConnection", new Object[] {e}));
-    }
-
     try {
       LOGGER.finest(
           () -> Messages.get(
               "HostMonitoringConnectionPlugin.activatedMonitoring",
               new Object[] {methodName}));
 
-      this.nodeKeys.clear();
-      this.nodeKeys.addAll(this.pluginService.getCurrentHostSpec().asAliases());
-
       monitorContext =
           this.monitorService.startMonitoring(
               this.pluginService.getCurrentConnection(), // abort this connection if needed
-              this.nodeKeys,
-              monitoringHostSpec,
+              this.getMonitoringHostSpec().asAliases(),
+              this.getMonitoringHostSpec(),
               this.properties,
               failureDetectionTimeMillis,
               failureDetectionIntervalMillis,
@@ -196,7 +179,7 @@ public <T, E extends Exception> T execute(
           this.monitorService.stopMonitoring(monitorContext);
 
           if (monitorContext.isNodeUnhealthy()) {
-            this.pluginService.setAvailability(this.nodeKeys, HostAvailability.NOT_AVAILABLE);
+            this.pluginService.setAvailability(this.getMonitoringHostSpec().asAliases(), HostAvailability.NOT_AVAILABLE);
 
             final boolean isConnectionClosed;
             try {
@@ -262,16 +245,16 @@ public void releaseResources() {
 
   @Override
   public OldConnectionSuggestedAction notifyConnectionChanged(final EnumSet<NodeChangeOptions> changes) {
-
     if (changes.contains(NodeChangeOptions.WENT_DOWN)
         || changes.contains(NodeChangeOptions.NODE_DELETED)) {
-      if (!this.nodeKeys.isEmpty()) {
-        this.monitorService.stopMonitoringForAllConnections(this.nodeKeys);
+      if (!this.getMonitoringHostSpec().asAliases().isEmpty()) {
+        this.monitorService.stopMonitoringForAllConnections(this.getMonitoringHostSpec().asAliases());
       }
-      this.nodeKeys.clear();
-      this.nodeKeys.addAll(this.pluginService.getCurrentHostSpec().getAliases());
     }
 
+    // Reset monitoring HostSpec since the associated connection has changed.
+    this.monitoringHostSpec = null;
+
     return OldConnectionSuggestedAction.NO_OPINION;
   }
 
@@ -291,8 +274,11 @@ private Connection connectInternal(String driverProtocol, HostSpec hostSpec,
     final Connection conn = connectFunc.call();
 
     if (conn != null) {
-      hostSpec.resetAliases();
-      this.pluginService.fillAliases(conn, hostSpec);
+      final RdsUrlType type = this.rdsHelper.identifyRdsType(hostSpec.getHost());
+      if (type.isRdsCluster()) {
+        hostSpec.resetAliases();
+        this.pluginService.fillAliases(conn, hostSpec);
+      }
     }
 
     return conn;
@@ -308,4 +294,24 @@ public Connection forceConnect(
       throws SQLException {
     return connectInternal(driverProtocol, hostSpec, forceConnectFunc);
   }
+
+  public HostSpec getMonitoringHostSpec() {
+    if (this.monitoringHostSpec == null) {
+      this.monitoringHostSpec = this.pluginService.getCurrentHostSpec();
+      final RdsUrlType rdsUrlType = this.rdsHelper.identifyRdsType(monitoringHostSpec.getUrl());
+
+      try {
+        if (rdsUrlType.isRdsCluster()) {
+          LOGGER.finest("Monitoring HostSpec is associated with a cluster endpoint, "
+              + "plugin needs to identify the cluster connection.");
+          this.monitoringHostSpec = this.pluginService.identifyConnection(this.pluginService.getCurrentConnection());
+          this.pluginService.fillAliases(this.pluginService.getCurrentConnection(), monitoringHostSpec);
+        }
+      } catch (SQLException e) {
+        // Log and ignore
+        LOGGER.finest(Messages.get("HostMonitoringConnectionPlugin.errorIdentifyingConnection", new Object[] {e}));
+      }
+    }
+    return this.monitoringHostSpec;
+  }
 }

wrapper/src/test/java/software/amazon/jdbc/plugin/efm/HostMonitoringConnectionPluginTest.java

@@ -81,6 +81,7 @@ class HostMonitoringConnectionPluginTest {
   @Captor ArgumentCaptor<String> stringArgumentCaptor;
   Properties properties = new Properties();
   @Mock HostSpec hostSpec;
+  @Mock HostSpec hostSpec2;
   @Mock Supplier<MonitorService> supplier;
   @Mock RdsUtils rdsUtils;
   @Mock MonitorConnectionContext context;
@@ -137,6 +138,9 @@ void initDefaultMockReturns() throws Exception {
     when(hostSpec.getHost()).thenReturn("host");
     when(hostSpec.getHost()).thenReturn("port");
     when(hostSpec.getAliases()).thenReturn(new HashSet<>(Collections.singletonList("host:port")));
+    when(hostSpec2.getHost()).thenReturn("host");
+    when(hostSpec2.getHost()).thenReturn("port");
+    when(hostSpec2.getAliases()).thenReturn(new HashSet<>(Collections.singletonList("host:port")));
     when(connection.createStatement()).thenReturn(statement);
     when(statement.executeQuery(any())).thenReturn(resultSet);
     when(rdsUtils.identifyRdsType(any())).thenReturn(RdsUrlType.RDS_INSTANCE);
@@ -293,16 +297,21 @@ void test_notifyConnectionChanged_nodeWentDown(final NodeChangeOptions option) t
         sqlFunction,
         EMPTY_ARGS);
 
-    final Set<String> aliases = new HashSet<>(Arrays.asList("alias1", "alias2"));
-    when(hostSpec.getAliases()).thenReturn(aliases);
-    assertEquals(OldConnectionSuggestedAction.NO_OPINION, plugin.notifyConnectionChanged(EnumSet.of(option)));
-
-    // NodeKeys should be empty at first
-    verify(monitorService, never()).stopMonitoringForAllConnections(any());
+    final Set<String> aliases1 = new HashSet<>(Arrays.asList("alias1", "alias2"));
+    final Set<String> aliases2 = new HashSet<>(Arrays.asList("alias3", "alias4"));
+    when(hostSpec.asAliases()).thenReturn(aliases1);
+    when(hostSpec2.asAliases()).thenReturn(aliases2);
+    when(pluginService.getCurrentHostSpec()).thenReturn(hostSpec);
 
     assertEquals(OldConnectionSuggestedAction.NO_OPINION, plugin.notifyConnectionChanged(EnumSet.of(option)));
     // NodeKeys should contain {"alias1", "alias2"}
-    verify(monitorService).stopMonitoringForAllConnections(aliases);
+    verify(monitorService).stopMonitoringForAllConnections(aliases1);
+
+    when(pluginService.getCurrentHostSpec()).thenReturn(hostSpec2);
+    assertEquals(OldConnectionSuggestedAction.NO_OPINION, plugin.notifyConnectionChanged(EnumSet.of(option)));
+    // NotifyConnectionChanged should reset the monitoringHostSpec.
+    // NodeKeys should contain {"alias3", "alias4"}
+    verify(monitorService).stopMonitoringForAllConnections(aliases2);
   }
 
   @Test