[WIP] Implement refactoring changes

benchmarks/src/jmh/java/software/amazon/jdbc/benchmarks/PluginBenchmarks.java

@@ -62,6 +62,8 @@
 import software.amazon.jdbc.dialect.Dialect;
 import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy;
 import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect;
+import software.amazon.jdbc.util.storage.StorageService;
+import software.amazon.jdbc.util.storage.StorageServiceImpl;
 import software.amazon.jdbc.util.telemetry.GaugeCallable;
 import software.amazon.jdbc.util.telemetry.TelemetryContext;
 import software.amazon.jdbc.util.telemetry.TelemetryCounter;
@@ -85,6 +87,7 @@ public class PluginBenchmarks {
       "jdbc:aws-wrapper:postgresql://instance-0.XYZ.us-east-2.rds.amazonaws.com";
   private static final String TEST_HOST = "instance-0";
   private static final int TEST_PORT = 5432;
+  private static final StorageService storageService = new StorageServiceImpl();
   private final HostSpec writerHostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy())
       .host(TEST_HOST).port(TEST_PORT).build();
 
@@ -164,7 +167,8 @@ public ConnectionWrapper initAndReleaseWithExecutionTimePlugin() throws SQLExcep
         mockTelemetryFactory,
         mockPluginService,
         mockHostListProviderService,
-        mockPluginManagerService)) {
+        mockPluginManagerService,
+        storageService)) {
       wrapper.releaseResources();
       return wrapper;
     }
@@ -179,7 +183,8 @@ public ConnectionWrapper initAndReleaseWithAuroraHostListPlugin() throws SQLExce
         mockTelemetryFactory,
         mockPluginService,
         mockHostListProviderService,
-        mockPluginManagerService)) {
+        mockPluginManagerService,
+        storageService)) {
       wrapper.releaseResources();
       return wrapper;
     }
@@ -194,7 +199,8 @@ public ConnectionWrapper initAndReleaseWithExecutionTimeAndAuroraHostListPlugins
         mockTelemetryFactory,
         mockPluginService,
         mockHostListProviderService,
-        mockPluginManagerService)) {
+        mockPluginManagerService,
+        storageService)) {
       wrapper.releaseResources();
       return wrapper;
     }
@@ -209,7 +215,8 @@ public ConnectionWrapper initAndReleaseWithReadWriteSplittingPlugin() throws SQL
         mockTelemetryFactory,
         mockPluginService,
         mockHostListProviderService,
-        mockPluginManagerService)) {
+        mockPluginManagerService,
+        storageService)) {
       wrapper.releaseResources();
       return wrapper;
     }
@@ -225,7 +232,8 @@ public ConnectionWrapper initAndReleaseWithAuroraHostListAndReadWriteSplittingPl
         mockTelemetryFactory,
         mockPluginService,
         mockHostListProviderService,
-        mockPluginManagerService)) {
+        mockPluginManagerService,
+        storageService)) {
       wrapper.releaseResources();
       return wrapper;
     }
@@ -243,7 +251,8 @@ public ConnectionWrapper initAndReleaseWithReadWriteSplittingPlugin_internalConn
         mockTelemetryFactory,
         mockPluginService,
         mockHostListProviderService,
-        mockPluginManagerService)) {
+        mockPluginManagerService,
+        storageService)) {
       wrapper.releaseResources();
       ConnectionProviderManager.releaseResources();
       Driver.resetCustomConnectionProvider();
@@ -264,7 +273,8 @@ public ConnectionWrapper initAndReleaseWithAuroraHostListAndReadWriteSplittingPl
         mockTelemetryFactory,
         mockPluginService,
         mockHostListProviderService,
-        mockPluginManagerService)) {
+        mockPluginManagerService,
+        storageService)) {
       wrapper.releaseResources();
       ConnectionProviderManager.releaseResources();
       Driver.resetCustomConnectionProvider();
@@ -281,7 +291,8 @@ public Statement executeStatementBaseline() throws SQLException {
         mockTelemetryFactory,
         mockPluginService,
         mockHostListProviderService,
-        mockPluginManagerService);
+        mockPluginManagerService,
+        storageService);
          Statement statement = wrapper.createStatement()) {
       return statement;
     }
@@ -297,7 +308,8 @@ public ResultSet executeStatementWithExecutionTimePlugin() throws SQLException {
             mockTelemetryFactory,
             mockPluginService,
             mockHostListProviderService,
-            mockPluginManagerService);
+            mockPluginManagerService,
+            storageService);
         Statement statement = wrapper.createStatement();
         ResultSet resultSet = statement.executeQuery("some sql")) {
       return resultSet;
@@ -314,7 +326,8 @@ public ResultSet executeStatementWithTelemetryDisabled() throws SQLException {
             mockTelemetryFactory,
             mockPluginService,
             mockHostListProviderService,
-            mockPluginManagerService);
+            mockPluginManagerService,
+            storageService);
         Statement statement = wrapper.createStatement();
         ResultSet resultSet = statement.executeQuery("some sql")) {
       return resultSet;
@@ -331,7 +344,8 @@ public ResultSet executeStatementWithTelemetry() throws SQLException {
             mockTelemetryFactory,
             mockPluginService,
             mockHostListProviderService,
-            mockPluginManagerService);
+            mockPluginManagerService,
+            storageService);
         Statement statement = wrapper.createStatement();
         ResultSet resultSet = statement.executeQuery("some sql")) {
       return resultSet;

benchmarks/src/jmh/java/software/amazon/jdbc/benchmarks/testplugin/TestConnectionWrapper.java

@@ -23,6 +23,7 @@
 import software.amazon.jdbc.HostListProviderService;
 import software.amazon.jdbc.PluginManagerService;
 import software.amazon.jdbc.PluginService;
+import software.amazon.jdbc.util.storage.StorageService;
 import software.amazon.jdbc.util.telemetry.TelemetryFactory;
 import software.amazon.jdbc.wrapper.ConnectionWrapper;
 
@@ -35,9 +36,10 @@ public TestConnectionWrapper(@NonNull Properties props,
       @NonNull final TelemetryFactory telemetryFactory,
       @NonNull PluginService pluginService,
       @NonNull HostListProviderService hostListProviderService,
-      @NonNull PluginManagerService pluginManagerService)
+      @NonNull PluginManagerService pluginManagerService,
+      @NonNull StorageService storageService)
       throws SQLException {
     super(props, url, connectionPluginManager, telemetryFactory, pluginService, hostListProviderService,
-        pluginManagerService);
+        pluginManagerService, storageService);
   }
 }

wrapper/src/main/java/software/amazon/jdbc/Driver.java

@@ -64,6 +64,8 @@
 import software.amazon.jdbc.util.PropertyUtils;
 import software.amazon.jdbc.util.RdsUtils;
 import software.amazon.jdbc.util.StringUtils;
+import software.amazon.jdbc.util.storage.StorageService;
+import software.amazon.jdbc.util.storage.StorageServiceImpl;
 import software.amazon.jdbc.util.telemetry.DefaultTelemetryFactory;
 import software.amazon.jdbc.util.telemetry.TelemetryContext;
 import software.amazon.jdbc.util.telemetry.TelemetryFactory;
@@ -77,6 +79,8 @@ public class Driver implements java.sql.Driver {
   private static final Logger LOGGER = Logger.getLogger("software.amazon.jdbc.Driver");
   private static @Nullable Driver registeredDriver;
 
+  private static final StorageService storageService = new StorageServiceImpl();
+
   private static final AtomicReference<ResetSessionStateOnCloseCallable> resetSessionStateOnCloseCallable =
       new AtomicReference<>(null);
   private static final AtomicReference<TransferSessionStateOnSwitchCallable> transferSessionStateOnSwitchCallable =
@@ -223,6 +227,7 @@ public Connection connect(final String url, final Properties info) throws SQLExc
           effectiveConnectionProvider,
           targetDriverDialect,
           configurationProfile,
+          storageService,
           telemetryFactory);
 
     } catch (Exception ex) {
@@ -388,6 +393,7 @@ public static void resetConnectionInitFunc() {
   }
 
   public static void clearCaches() {
+    storageService.clearAll();
     RdsUtils.clearCache();
     RdsHostListProvider.clearAll();
     PluginServiceImpl.clearCache();

wrapper/src/main/java/software/amazon/jdbc/PluginService.java

@@ -29,6 +29,7 @@
 import software.amazon.jdbc.hostavailability.HostAvailability;
 import software.amazon.jdbc.states.SessionStateService;
 import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect;
+import software.amazon.jdbc.util.storage.StorageService;
 import software.amazon.jdbc.util.telemetry.TelemetryFactory;
 
 /**
@@ -80,11 +81,16 @@ EnumSet<NodeChangeOptions> setCurrentConnection(
 
   HostSpec getInitialConnectionHostSpec();
 
+  // TODO: evaluate if there is a better way to pass the storage service to the monitors that need it.
+  StorageService getStorageService();
+
   /**
    * Set the collection of hosts that should be allowed and/or blocked for connections.
    *
    * @param allowedAndBlockedHosts An object defining the allowed and blocked sets of hosts.
+   * @deprecated use StorageService#set(ItemCategory.ALLOWED_AND_BLOCKED_HOSTS, key, allowedAndBlockedHosts) instead.
    */
+  @Deprecated
   void setAllowedAndBlockedHosts(AllowedAndBlockedHosts allowedAndBlockedHosts);
 
   /**

wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java

@@ -31,7 +31,6 @@
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
-import java.util.concurrent.atomic.AtomicReference;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
@@ -53,7 +52,10 @@
 import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect;
 import software.amazon.jdbc.util.CacheMap;
 import software.amazon.jdbc.util.Messages;
+import software.amazon.jdbc.util.ServiceContainer;
 import software.amazon.jdbc.util.Utils;
+import software.amazon.jdbc.util.storage.ItemCategory;
+import software.amazon.jdbc.util.storage.StorageService;
 import software.amazon.jdbc.util.telemetry.TelemetryFactory;
 
 public class PluginServiceImpl implements PluginService, CanReleaseResources,
@@ -63,13 +65,14 @@ public class PluginServiceImpl implements PluginService, CanReleaseResources,
   protected static final long DEFAULT_HOST_AVAILABILITY_CACHE_EXPIRE_NANO = TimeUnit.MINUTES.toNanos(5);
 
   protected static final CacheMap<String, HostAvailability> hostAvailabilityExpiringCache = new CacheMap<>();
+  protected final ServiceContainer serviceContainer;
   protected final ConnectionPluginManager pluginManager;
+  protected final StorageService storageService;
   private final Properties props;
   private final String originalUrl;
   private final String driverProtocol;
   protected volatile HostListProvider hostListProvider;
   protected List<HostSpec> allHosts = new ArrayList<>();
-  protected AtomicReference<AllowedAndBlockedHosts> allowedAndBlockedHosts = new AtomicReference<>();
   protected Connection currentConnection;
   protected HostSpec currentHostSpec;
   protected HostSpec initialConnectionHostSpec;
@@ -87,14 +90,15 @@ public class PluginServiceImpl implements PluginService, CanReleaseResources,
   protected final ReentrantLock connectionSwitchLock = new ReentrantLock();
 
   public PluginServiceImpl(
-      @NonNull final ConnectionPluginManager pluginManager,
+      @NonNull final ServiceContainer serviceContainer,
       @NonNull final Properties props,
       @NonNull final String originalUrl,
       @NonNull final String targetDriverProtocol,
       @NonNull final TargetDriverDialect targetDriverDialect)
       throws SQLException {
 
-    this(pluginManager,
+    this(
+        serviceContainer,
         new ExceptionManager(),
         props,
         originalUrl,
@@ -106,13 +110,14 @@ public PluginServiceImpl(
   }
 
   public PluginServiceImpl(
-      @NonNull final ConnectionPluginManager pluginManager,
+      @NonNull final ServiceContainer serviceContainer,
       @NonNull final Properties props,
       @NonNull final String originalUrl,
       @NonNull final String targetDriverProtocol,
       @NonNull final TargetDriverDialect targetDriverDialect,
       @Nullable final ConfigurationProfile configurationProfile) throws SQLException {
-    this(pluginManager,
+    this(
+        serviceContainer,
         new ExceptionManager(),
         props,
         originalUrl,
@@ -124,7 +129,7 @@ public PluginServiceImpl(
   }
 
   public PluginServiceImpl(
-      @NonNull final ConnectionPluginManager pluginManager,
+      @NonNull final ServiceContainer serviceContainer,
       @NonNull final ExceptionManager exceptionManager,
       @NonNull final Properties props,
       @NonNull final String originalUrl,
@@ -133,14 +138,16 @@ public PluginServiceImpl(
       @NonNull final TargetDriverDialect targetDriverDialect,
       @Nullable final ConfigurationProfile configurationProfile,
       @Nullable final SessionStateService sessionStateService) throws SQLException {
-    this.pluginManager = pluginManager;
+    this.serviceContainer = serviceContainer;
+    this.pluginManager = serviceContainer.getConnectionPluginManager();
     this.props = props;
     this.originalUrl = originalUrl;
     this.driverProtocol = targetDriverProtocol;
     this.configurationProfile = configurationProfile;
     this.exceptionManager = exceptionManager;
     this.dialectProvider = dialectProvider != null ? dialectProvider : new DialectManager(this);
     this.targetDriverDialect = targetDriverDialect;
+    this.storageService = serviceContainer.getStorageService();
     this.connectionProviderManager = new ConnectionProviderManager(
         this.pluginManager.getDefaultConnProvider(),
         this.pluginManager.getEffectiveConnProvider());
@@ -206,8 +213,15 @@ public HostSpec getInitialConnectionHostSpec() {
   }
 
   @Override
+  public StorageService getStorageService() {
+    return this.storageService;
+  }
+
+  @Override
+  @Deprecated
   public void setAllowedAndBlockedHosts(AllowedAndBlockedHosts allowedAndBlockedHosts) {
-    this.allowedAndBlockedHosts.set(allowedAndBlockedHosts);
+    this.storageService.set(
+        ItemCategory.ALLOWED_AND_BLOCKED_HOSTS, this.initialConnectionHostSpec.getHost(), allowedAndBlockedHosts);
   }
 
   @Override
@@ -401,7 +415,8 @@ public List<HostSpec> getAllHosts() {
 
   @Override
   public List<HostSpec> getHosts() {
-    AllowedAndBlockedHosts hostPermissions = this.allowedAndBlockedHosts.get();
+    AllowedAndBlockedHosts hostPermissions = this.storageService.get(
+        ItemCategory.ALLOWED_AND_BLOCKED_HOSTS, this.initialConnectionHostSpec.getHost(), AllowedAndBlockedHosts.class);
     if (hostPermissions == null) {
       return this.allHosts;
     }
@@ -705,7 +720,7 @@ public void updateDialect(final @NonNull Connection connection) throws SQLExcept
     }
 
     final HostListProviderSupplier supplier = this.dialect.getHostListProvider();
-    this.setHostListProvider(supplier.getProvider(props, this.originalUrl, this, this));
+    this.setHostListProvider(supplier.getProvider(this.props, this.originalUrl, this.serviceContainer));
   }
 
   @Override

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

@@ -22,6 +22,7 @@
 import java.sql.Statement;
 import java.util.Collections;
 import java.util.List;
+import software.amazon.jdbc.PluginService;
 import software.amazon.jdbc.hostlistprovider.AuroraHostListProvider;
 import software.amazon.jdbc.hostlistprovider.monitoring.MonitoringRdsHostListProvider;
 import software.amazon.jdbc.plugin.failover2.FailoverConnectionPlugin;
@@ -81,25 +82,24 @@ public boolean isDialect(final Connection connection) {
 
   @Override
   public HostListProviderSupplier getHostListProvider() {
-    return (properties, initialUrl, hostListProviderService, pluginService) -> {
-
+    return  (properties, initialUrl, serviceContainer) -> {
+      final PluginService pluginService = serviceContainer.getPluginService();
       final FailoverConnectionPlugin failover2Plugin = pluginService.getPlugin(FailoverConnectionPlugin.class);
 
       if (failover2Plugin != null) {
         return new MonitoringRdsHostListProvider(
             properties,
             initialUrl,
-            hostListProviderService,
+            serviceContainer,
             TOPOLOGY_QUERY,
             NODE_ID_QUERY,
             IS_READER_QUERY,
-            IS_WRITER_QUERY,
-            pluginService);
+            IS_WRITER_QUERY);
       }
       return new AuroraHostListProvider(
           properties,
           initialUrl,
-          hostListProviderService,
+          serviceContainer,
           TOPOLOGY_QUERY,
           NODE_ID_QUERY,
           IS_READER_QUERY);