consolidate cache cleanup in a single place

aws · Jan 10, 2025 · d3daa51 · d3daa51
1 parent 64e3b1b
commit d3daa51
Show file tree

Hide file tree

Showing 30 changed files with 450 additions and 190 deletions.
diff --git a/wrapper/src/main/java/software/amazon/jdbc/Driver.java b/wrapper/src/main/java/software/amazon/jdbc/Driver.java
@@ -33,6 +33,21 @@
 import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import software.amazon.jdbc.authentication.AwsCredentialsManager;
+import software.amazon.jdbc.dialect.DialectManager;
+import software.amazon.jdbc.hostlistprovider.RdsHostListProvider;
+import software.amazon.jdbc.hostlistprovider.monitoring.MonitoringRdsHostListProvider;
+import software.amazon.jdbc.plugin.AwsSecretsManagerCacheHolder;
+import software.amazon.jdbc.plugin.DataCacheConnectionPlugin;
+import software.amazon.jdbc.plugin.OpenedConnectionTracker;
+import software.amazon.jdbc.plugin.customendpoint.CustomEndpointMonitorImpl;
+import software.amazon.jdbc.plugin.customendpoint.CustomEndpointPlugin;
+import software.amazon.jdbc.plugin.efm.MonitorThreadContainer;
+import software.amazon.jdbc.plugin.federatedauth.FederatedAuthCacheHolder;
+import software.amazon.jdbc.plugin.federatedauth.OktaAuthCacheHolder;
+import software.amazon.jdbc.plugin.iam.IamAuthCacheHolder;
+import software.amazon.jdbc.plugin.limitless.LimitlessRouterServiceImpl;
+import software.amazon.jdbc.plugin.strategy.fastestresponse.FastestResponseStrategyPlugin;
+import software.amazon.jdbc.plugin.strategy.fastestresponse.HostResponseTimeServiceImpl;
 import software.amazon.jdbc.profile.ConfigurationProfile;
 import software.amazon.jdbc.profile.DriverConfigurationProfiles;
 import software.amazon.jdbc.states.ResetSessionStateOnCloseCallable;
@@ -277,4 +292,32 @@ public static void setPrepareHostFunc(final Function<String, String> func) {
   public static void resetPrepareHostFunc() {
     RdsUtils.resetPrepareHostFunc();
   }
+
+  public static void clearCaches() {
+    RdsUtils.clearCache();
+    RdsHostListProvider.clearAll();
+    PluginServiceImpl.clearCache();
+    DialectManager.resetEndpointCache();
+    MonitoringRdsHostListProvider.clearCache();
+    CustomEndpointMonitorImpl.clearCache();
+    OpenedConnectionTracker.clearCache();
+    AwsSecretsManagerCacheHolder.clearCache();
+    DataCacheConnectionPlugin.clearCache();
+    FederatedAuthCacheHolder.clearCache();
+    OktaAuthCacheHolder.clearCache();
+    IamAuthCacheHolder.clearCache();
+    LimitlessRouterServiceImpl.clearCache();
+    RoundRobinHostSelector.clearCache();
+    FastestResponseStrategyPlugin.clearCache();
+  }
+
+  public static void releaseResources() {
+    software.amazon.jdbc.plugin.efm2.MonitorServiceImpl.closeAllMonitors();
+    MonitorThreadContainer.releaseInstance();
+    ConnectionProviderManager.releaseResources();
+    CustomEndpointPlugin.closeMonitors();
+    HikariPoolsHolder.closeAllPools();
+    HostResponseTimeServiceImpl.closeAllMonitors();
+    clearCaches();
+  }
 }
diff --git a/wrapper/src/main/java/software/amazon/jdbc/HikariPooledConnectionProvider.java b/wrapper/src/main/java/software/amazon/jdbc/HikariPooledConnectionProvider.java
@@ -39,6 +39,7 @@
 import software.amazon.jdbc.targetdriverdialect.ConnectInfo;
 import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect;
 import software.amazon.jdbc.util.Messages;
+import software.amazon.jdbc.util.Pair;
 import software.amazon.jdbc.util.PropertyUtils;
 import software.amazon.jdbc.util.RdsUrlType;
 import software.amazon.jdbc.util.RdsUtils;
@@ -61,17 +62,22 @@ public class HikariPooledConnectionProvider implements PooledConnectionProvider,
       });
 
   protected static final RdsUtils rdsUtils = new RdsUtils();
-  protected static SlidingExpirationCache<PoolKey, HikariDataSource> databasePools =
-      new SlidingExpirationCache<>(
-          (hikariDataSource) -> hikariDataSource.getHikariPoolMXBean().getActiveConnections() == 0,
-          HikariDataSource::close
-      );
   protected static long poolExpirationCheckNanos = TimeUnit.MINUTES.toNanos(30);
   protected final HikariPoolConfigurator poolConfigurator;
   protected final HikariPoolMapping poolMapping;
   protected final AcceptsUrlFunc acceptsUrlFunc;
   protected final LeastConnectionsHostSelector leastConnectionsHostSelector;
 
+  static {
+    HikariPoolsHolder.databasePools.setShouldDisposeFunc(
+        (hikariDataSource) -> {
+          if (hikariDataSource instanceof HikariDataSource) {
+            return ((HikariDataSource) hikariDataSource).getHikariPoolMXBean().getActiveConnections() == 0;
+          }
+          return true;
+        });
+  }
+
   /**
    * {@link HikariPooledConnectionProvider} constructor. This class can be passed to
    * {@link ConnectionProviderManager#setConnectionProvider} to enable internal connection pools for
@@ -114,7 +120,7 @@ public HikariPooledConnectionProvider(
     this.poolConfigurator = hikariPoolConfigurator;
     this.poolMapping = mapping;
     this.acceptsUrlFunc = null;
-    this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(databasePools);
+    this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(HikariPoolsHolder.databasePools);
   }
 
   /**
@@ -151,8 +157,8 @@ public HikariPooledConnectionProvider(
     this.poolMapping = mapping;
     this.acceptsUrlFunc = null;
     poolExpirationCheckNanos = poolExpirationNanos;
-    databasePools.setCleanupIntervalNanos(poolCleanupNanos);
-    this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(databasePools);
+    HikariPoolsHolder.databasePools.setCleanupIntervalNanos(poolCleanupNanos);
+    this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(HikariPoolsHolder.databasePools);
   }
 
   /**
@@ -193,8 +199,8 @@ public HikariPooledConnectionProvider(
     this.poolMapping = mapping;
     this.acceptsUrlFunc = acceptsUrlFunc;
     poolExpirationCheckNanos = poolExpirationNanos;
-    databasePools.setCleanupIntervalNanos(poolCleanupNanos);
-    this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(databasePools);
+    HikariPoolsHolder.databasePools.setCleanupIntervalNanos(poolCleanupNanos);
+    this.leastConnectionsHostSelector = new LeastConnectionsHostSelector(HikariPoolsHolder.databasePools);
   }
 
 
@@ -272,8 +278,8 @@ public Connection connect(
     final HostSpec finalHostSpec = connectionHostSpec;
     dialect.prepareConnectProperties(copy, protocol, finalHostSpec);
 
-    final HikariDataSource ds = databasePools.computeIfAbsent(
-        new PoolKey(hostSpec.getUrl(), getPoolKey(finalHostSpec, copy)),
+    final HikariDataSource ds = (HikariDataSource) HikariPoolsHolder.databasePools.computeIfAbsent(
+        Pair.create(hostSpec.getUrl(), getPoolKey(finalHostSpec, copy)),
         (lambdaPoolKey) -> createHikariDataSource(protocol, finalHostSpec, copy, targetDriverDialect),
         poolExpirationCheckNanos
     );
@@ -297,22 +303,7 @@ protected String getPoolKey(HostSpec hostSpec, Properties props) {
 
   @Override
   public void releaseResources() {
-    databasePools.getEntries().forEach((poolKey, pool) -> {
-      if (!pool.isClosed()) {
-        pool.close();
-      }
-    });
-    databasePools.clear();
-  }
-
-  // For testing purposes
-  public static void clearCache() {
-    databasePools.getEntries().forEach((poolKey, pool) -> {
-      if (!pool.isClosed()) {
-        pool.close();
-      }
-    });
-    databasePools.clear();
+    HikariPoolsHolder.closeAllPools();
   }
 
   /**
@@ -378,7 +369,7 @@ protected void configurePool(
    * @return the number of active connection pools
    */
   public int getHostCount() {
-    return databasePools.size();
+    return HikariPoolsHolder.databasePools.size();
   }
 
   /**
@@ -388,8 +379,8 @@ public int getHostCount() {
    */
   public Set<String> getHosts() {
     return Collections.unmodifiableSet(
-        databasePools.getEntries().keySet().stream()
-            .map(poolKey -> poolKey.url)
+        HikariPoolsHolder.databasePools.getEntries().keySet().stream()
+            .map(poolKey -> (String) poolKey.getValue1())
             .collect(Collectors.toSet()));
   }
 
@@ -398,8 +389,8 @@ public Set<String> getHosts() {
    *
    * @return a set containing every key associated with an active connection pool
    */
-  public Set<PoolKey> getKeys() {
-    return databasePools.getEntries().keySet();
+  public Set<Pair> getKeys() {
+    return HikariPoolsHolder.databasePools.getEntries().keySet();
   }
 
   @Override
@@ -413,7 +404,7 @@ public String getTargetName() {
   public void logConnections() {
     LOGGER.finest(() -> {
       final StringBuilder builder = new StringBuilder();
-      databasePools.getEntries().forEach((key, dataSource) -> {
+      HikariPoolsHolder.databasePools.getEntries().forEach((key, dataSource) -> {
         builder.append("\t[ ");
         builder.append(key).append(":");
         builder.append("\n\t {");
@@ -436,51 +427,8 @@ HikariDataSource createHikariDataSource(
     return new HikariDataSource(config);
   }
 
-  public static class PoolKey {
-    private final @NonNull String url;
-    private final @NonNull String extraKey;
-
-    public PoolKey(final @NonNull String url, final @NonNull String extraKey) {
-      this.url = url;
-      this.extraKey = extraKey;
-    }
-
-    public String getUrl() {
-      return this.url;
-    }
-
-    @Override
-    public int hashCode() {
-      final int prime = 31;
-      int result = 1;
-      result = prime * result + ((url == null) ? 0 : url.hashCode()) + ((extraKey == null) ? 0 : extraKey.hashCode());
-      return result;
-    }
-
-    @Override
-    public boolean equals(final Object obj) {
-      if (this == obj) {
-        return true;
-      }
-      if (obj == null) {
-        return false;
-      }
-      if (getClass() != obj.getClass()) {
-        return false;
-      }
-      final PoolKey other = (PoolKey) obj;
-      return this.url.equals(other.url) && this.extraKey.equals(other.extraKey);
-    }
-
-    @Override
-    public String toString() {
-      return "PoolKey [url=" + url + ", extraKey=" + extraKey + "]";
-    }
-
-  }
-
   // For testing purposes only
-  void setDatabasePools(SlidingExpirationCache<PoolKey, HikariDataSource> connectionPools) {
-    databasePools = connectionPools;
+  void setDatabasePools(SlidingExpirationCache<Pair, AutoCloseable> connectionPools) {
+    HikariPoolsHolder.databasePools = connectionPools;
   }
 }
diff --git a/wrapper/src/main/java/software/amazon/jdbc/HikariPoolsHolder.java b/wrapper/src/main/java/software/amazon/jdbc/HikariPoolsHolder.java
@@ -0,0 +1,46 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package software.amazon.jdbc;
+
+import software.amazon.jdbc.util.Pair;
+import software.amazon.jdbc.util.SlidingExpirationCache;
+
+public class HikariPoolsHolder {
+  static SlidingExpirationCache<Pair, AutoCloseable> databasePools =
+      new SlidingExpirationCache<>(
+          null,
+          (hikariDataSource) -> {
+            try {
+              hikariDataSource.close();
+            } catch (Exception ex) {
+              // ignore
+            }
+          }
+      );
+
+  public static void closeAllPools() {
+    databasePools.getEntries().forEach((poolKey, pool) -> {
+      try {
+        pool.close();
+      } catch (Exception ex) {
+        // ignore
+      }
+    });
+    databasePools.clear();
+
+  }
+}
diff --git a/wrapper/src/main/java/software/amazon/jdbc/LeastConnectionsHostSelector.java b/wrapper/src/main/java/software/amazon/jdbc/LeastConnectionsHostSelector.java
@@ -26,14 +26,15 @@
 import org.checkerframework.checker.nullness.qual.Nullable;
 import software.amazon.jdbc.hostavailability.HostAvailability;
 import software.amazon.jdbc.util.Messages;
+import software.amazon.jdbc.util.Pair;
 import software.amazon.jdbc.util.SlidingExpirationCache;
 
 public class LeastConnectionsHostSelector implements HostSelector {
   public static final String STRATEGY_LEAST_CONNECTIONS = "leastConnections";
-  private final SlidingExpirationCache<HikariPooledConnectionProvider.PoolKey, HikariDataSource> databasePools;
+  private final SlidingExpirationCache<Pair, AutoCloseable> databasePools;
 
   public LeastConnectionsHostSelector(
-      SlidingExpirationCache<HikariPooledConnectionProvider.PoolKey, HikariDataSource> databasePools) {
+      SlidingExpirationCache<Pair, AutoCloseable> databasePools) {
     this.databasePools = databasePools;
   }
 
@@ -58,15 +59,18 @@ public HostSpec getHost(
 
   private int getNumConnections(
       final HostSpec hostSpec,
-      final SlidingExpirationCache<HikariPooledConnectionProvider.PoolKey, HikariDataSource> databasePools) {
+      final SlidingExpirationCache<Pair, AutoCloseable> databasePools) {
     int numConnections = 0;
     final String url = hostSpec.getUrl();
-    for (final Map.Entry<HikariPooledConnectionProvider.PoolKey, HikariDataSource> entry :
+    for (final Map.Entry<Pair, AutoCloseable> entry :
         databasePools.getEntries().entrySet()) {
-      if (!url.equals(entry.getKey().getUrl())) {
+      if (!url.equals(entry.getKey().getValue1())) {
         continue;
       }
-      numConnections += entry.getValue().getHikariPoolMXBean().getActiveConnections();
+      if (!(entry.getValue() instanceof HikariDataSource)) {
+        continue;
+      }
+      numConnections += ((HikariDataSource) entry.getValue()).getHikariPoolMXBean().getActiveConnections();
     }
     return numConnections;
   }

diff --git a/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java
@@ -754,4 +754,8 @@ public <T> T getPlugin(final Class<T> pluginClazz) {
     }
     return null;
   }
+
+  public static void clearCache() {
+    hostAvailabilityExpiringCache.clear();
+  }
 }
diff --git a/wrapper/src/main/java/software/amazon/jdbc/RoundRobinHostSelector.java b/wrapper/src/main/java/software/amazon/jdbc/RoundRobinHostSelector.java
@@ -251,8 +251,7 @@ private void updateCachedHostWeightPairsPropertiesForRoundRobinClusterInfo(
     }
   }
 
-  // For testing purposes only
-  public void clearCache() {
+  public static void clearCache() {
     roundRobinCache.clear();
   }
 

diff --git a/.../java/software/amazon/jdbc/hostlistprovider/monitoring/MonitoringRdsHostListProvider.java b/.../java/software/amazon/jdbc/hostlistprovider/monitoring/MonitoringRdsHostListProvider.java
@@ -93,6 +93,12 @@ public void clear() {
   }
 
   public static void clearCache() {
+    topologyCache.clear();
+    primaryClusterIdCache.clear();
+    suggestedPrimaryClusterIdCache.clear();
+  }
+
+  public static void closeAllMonitors() {
     monitors.getEntries().values().forEach(monitor -> {
       try {
         monitor.close();
@@ -101,9 +107,7 @@ public static void clearCache() {
       }
     });
     monitors.clear();
-    topologyCache.clear();
-    primaryClusterIdCache.clear();
-    suggestedPrimaryClusterIdCache.clear();
+    clearCache();
   }
 
   @Override