[WIP] feat: query-level load balancer plugin

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

@@ -49,6 +49,7 @@
 import software.amazon.jdbc.plugin.iam.IamAuthConnectionPluginFactory;
 import software.amazon.jdbc.plugin.limitless.LimitlessConnectionPluginFactory;
 import software.amazon.jdbc.plugin.readwritesplitting.GdbReadWriteSplittingPluginFactory;
+import software.amazon.jdbc.plugin.readwritesplitting.QueryLoadBalancingReadWriteSplittingPluginFactory;
 import software.amazon.jdbc.plugin.readwritesplitting.ReadWriteSplittingPluginFactory;
 import software.amazon.jdbc.plugin.srw.SimpleReadWriteSplittingPluginFactory;
 import software.amazon.jdbc.plugin.staledns.AuroraStaleDnsPluginFactory;
@@ -86,6 +87,7 @@ public class ConnectionPluginChainBuilder {
           put("readWriteSplitting", new ReadWriteSplittingPluginFactory());
           put("srw", new SimpleReadWriteSplittingPluginFactory());
           put("gdbReadWriteSplitting", new GdbReadWriteSplittingPluginFactory());
+          put("qlb", new QueryLoadBalancingReadWriteSplittingPluginFactory());
           put("auroraConnectionTracker", new AuroraConnectionTrackerPluginFactory());
           put("driverMetaData", new DriverMetaDataConnectionPluginFactory());
           put("connectTime", new ConnectTimeConnectionPluginFactory());
@@ -118,6 +120,7 @@ public class ConnectionPluginChainBuilder {
           put(ReadWriteSplittingPluginFactory.class, 700);
           put(SimpleReadWriteSplittingPluginFactory.class, 710);
           put(GdbReadWriteSplittingPluginFactory.class, 720);
+          put(QueryLoadBalancingReadWriteSplittingPluginFactory.class, 730);
           put(HostMonitoringConnectionPluginFactory.class, 800);
           put(software.amazon.jdbc.plugin.efm2.HostMonitoringConnectionPluginFactory.class, 810);
           put(FastestResponseStrategyPluginFactory.class, 900);

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

@@ -47,6 +47,7 @@
 import software.amazon.jdbc.plugin.iam.IamAuthConnectionPlugin;
 import software.amazon.jdbc.plugin.limitless.LimitlessConnectionPlugin;
 import software.amazon.jdbc.plugin.readwritesplitting.GdbReadWriteSplittingPlugin;
+import software.amazon.jdbc.plugin.readwritesplitting.QueryLoadBalancingReadWriteSplittingPlugin;
 import software.amazon.jdbc.plugin.readwritesplitting.ReadWriteSplittingPlugin;
 import software.amazon.jdbc.plugin.srw.SimpleReadWriteSplittingPlugin;
 import software.amazon.jdbc.plugin.staledns.AuroraStaleDnsPlugin;
@@ -56,7 +57,6 @@
 import software.amazon.jdbc.util.Messages;
 import software.amazon.jdbc.util.Pair;
 import software.amazon.jdbc.util.StateSnapshotProvider;
-import software.amazon.jdbc.util.Utils;
 import software.amazon.jdbc.util.WrapperUtils;
 import software.amazon.jdbc.util.telemetry.TelemetryContext;
 import software.amazon.jdbc.util.telemetry.TelemetryFactory;
@@ -93,6 +93,7 @@ public class ConnectionPluginManager implements CanReleaseResources, Wrapper, St
           put(ReadWriteSplittingPlugin.class, "plugin:readWriteSplitting");
           put(SimpleReadWriteSplittingPlugin.class, "plugin:srw");
           put(GdbReadWriteSplittingPlugin.class, "plugin:gdbReadWriteSplitting");
+          put(QueryLoadBalancingReadWriteSplittingPlugin.class, "plugin:qlb");
           put(FastestResponseStrategyPlugin.class, "plugin:fastestResponseStrategy");
           put(DefaultConnectionPlugin.class, "plugin:targetDriver");
           put(AuroraInitialConnectionStrategyPlugin.class, "plugin:initialConnection");

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

@@ -28,6 +28,7 @@
 import software.amazon.jdbc.AwsWrapperProperty;
 import software.amazon.jdbc.HostRole;
 import software.amazon.jdbc.HostSpec;
+import software.amazon.jdbc.PropertyDefinition;
 import software.amazon.jdbc.util.ConnectionUrlParser;
 import software.amazon.jdbc.util.Messages;
 
@@ -49,6 +50,10 @@ public class ConnectionStringHostListProvider implements StaticHostListProvider
           "Set to true if you are providing a connection string with multiple comma-delimited hosts and your "
               + "cluster has only one writer. The writer must be the first host in the connection string");
 
+  static {
+    PropertyDefinition.registerPluginProperties(ConnectionStringHostListProvider.class);
+  }
+
   public ConnectionStringHostListProvider(
       final @NonNull Properties properties,
       final String initialUrl,

wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/AbstractReadWriteSplittingPlugin.java

@@ -127,7 +127,7 @@ public OldConnectionSuggestedAction notifyConnectionChanged(
       boolean isCachedConnection =
           currentConnection == this.writerConnection
               || (this.readerCacheItem != null && currentConnection == this.readerCacheItem.get());
-      
+
       if (!isCachedConnection) {
         closeIdleConnections();
       }
@@ -183,6 +183,12 @@ public <T, E extends Exception> T execute(
       } catch (final SQLException e) {
         throw WrapperUtils.wrapExceptionIfNeeded(exceptionClass, e);
       }
+    } else {
+      try {
+        this.switchConnectionIfRequired(methodInvokeOn, methodName, args);
+      } catch (final SQLException e) {
+        throw WrapperUtils.wrapExceptionIfNeeded(exceptionClass, e);
+      }
     }
 
     try {
@@ -237,6 +243,13 @@ protected void setReaderConnection(final Connection conn, final HostSpec host) {
                 host.getHostAndPort()}));
   }
 
+  public void switchConnectionIfRequired(
+      final Object methodInvokeOn,
+      final String methodName,
+      final Object[] args) throws SQLException {
+    // do nothing
+  }
+
   public void switchConnectionIfRequired(final boolean readOnly) throws SQLException {
     final Connection currentConnection = this.pluginService.getCurrentConnection();
     if (currentConnection != null && currentConnection.isClosed()) {
@@ -296,7 +309,7 @@ protected void logAndThrowException(final String logMessage, final SqlState sqlS
     throw new ReadWriteSplittingSQLException(logMessage, sqlState.getState());
   }
 
-  private void logAndThrowException(
+  protected void logAndThrowException(
       final String logMessage, final Throwable cause)
       throws SQLException {
     LOGGER.fine(logMessage);
@@ -306,7 +319,7 @@ private void logAndThrowException(
     throw new ReadWriteSplittingSQLException(logMessage, SqlState.CONNECTION_UNABLE_TO_CONNECT.getState(), cause);
   }
 
-  private void switchToWriterConnection()
+  protected void switchToWriterConnection()
       throws SQLException {
     final Connection currentConnection = this.pluginService.getCurrentConnection();
     final HostSpec currentHost = this.pluginService.getCurrentHostSpec();
@@ -315,7 +328,7 @@ private void switchToWriterConnection()
       return;
     }
 
-    this.inReadWriteSplit = true;
+    this.inReadWriteSplit = false;
     if (!isConnectionUsable(this.writerConnection)) {
       initializeWriterConnection();
     } else {
@@ -346,7 +359,7 @@ protected void switchCurrentConnectionTo(
             newConnectionHost.getHostAndPort()}));
   }
 
-  private void switchToReaderConnection()
+  protected void switchToReaderConnection()
       throws SQLException {
     final Connection currentConnection = this.pluginService.getCurrentConnection();
     final HostSpec currentHost = this.pluginService.getCurrentHostSpec();

wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/QueryLoadBalancingReadWriteSplittingPlugin.java

@@ -0,0 +1,179 @@
+/*
+ * 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.plugin.readwritesplitting;
+
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Properties;
+import java.util.Set;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import software.amazon.jdbc.AwsWrapperProperty;
+import software.amazon.jdbc.HostRole;
+import software.amazon.jdbc.HostSpec;
+import software.amazon.jdbc.JdbcMethod;
+import software.amazon.jdbc.PluginService;
+import software.amazon.jdbc.PropertyDefinition;
+import software.amazon.jdbc.util.Messages;
+import software.amazon.jdbc.util.SqlState;
+import software.amazon.jdbc.util.Utils;
+
+public class QueryLoadBalancingReadWriteSplittingPlugin extends ReadWriteSplittingPlugin {
+
+  private static final Logger LOGGER = Logger.getLogger(QueryLoadBalancingReadWriteSplittingPlugin.class.getName());
+
+  private static final Set<String> subscribedMethods =
+      Collections.unmodifiableSet(new HashSet<String>() {
+        {
+          add(JdbcMethod.INITHOSTPROVIDER.methodName);
+          add(JdbcMethod.CONNECT.methodName);
+          add(JdbcMethod.NOTIFYCONNECTIONCHANGED.methodName);
+          add(JdbcMethod.CONNECTION_SETREADONLY.methodName);
+          add(JdbcMethod.CONNECTION_CLEARWARNINGS.methodName);
+          add(JdbcMethod.STATEMENT_EXECUTE.methodName);
+          add(JdbcMethod.STATEMENT_EXECUTEQUERY.methodName);
+          add(JdbcMethod.STATEMENT_EXECUTEBATCH.methodName);
+          add(JdbcMethod.STATEMENT_EXECUTEUPDATE.methodName);
+          add(JdbcMethod.PREPAREDSTATEMENT_EXECUTE.methodName);
+          add(JdbcMethod.PREPAREDSTATEMENT_EXECUTEUPDATE.methodName);
+          add(JdbcMethod.PREPAREDSTATEMENT_EXECUTELARGEUPDATE.methodName);
+          add(JdbcMethod.PREPAREDSTATEMENT_EXECUTEQUERY.methodName);
+          add(JdbcMethod.PREPAREDSTATEMENT_EXECUTEBATCH.methodName);
+          add(JdbcMethod.CALLABLESTATEMENT_EXECUTE.methodName);
+          add(JdbcMethod.CALLABLESTATEMENT_EXECUTEQUERY.methodName);
+          add(JdbcMethod.CALLABLESTATEMENT_EXECUTELARGEUPDATE.methodName);
+          add(JdbcMethod.CALLABLESTATEMENT_EXECUTEBATCH.methodName);
+          add(JdbcMethod.CALLABLESTATEMENT_EXECUTEUPDATE.methodName);
+          add(JdbcMethod.CONNECTION_SETAUTOCOMMIT.methodName);
+
+          // Extra methods to handle
+          add(JdbcMethod.CONNECTION_CREATESTATEMENT.methodName);
+          add(JdbcMethod.CONNECTION_PREPARESTATEMENT.methodName);
+          add(JdbcMethod.CONNECTION_PREPARECALL.methodName);
+        }
+      });
+
+  public static final AwsWrapperProperty READER_HOST_SELECTOR_STRATEGY =
+      new AwsWrapperProperty(
+          "includeWriterToLoadBalancing",
+          "false",
+          "Set to true to allow writer node to be included in query load balancing.");
+
+  private final boolean includeWriterToLoadBalancing;
+
+  static {
+    PropertyDefinition.registerPluginProperties(QueryLoadBalancingReadWriteSplittingPlugin.class);
+  }
+
+  public QueryLoadBalancingReadWriteSplittingPlugin(final PluginService pluginService, final @NonNull Properties properties) {
+    super(pluginService, properties);
+    this.includeWriterToLoadBalancing = READER_HOST_SELECTOR_STRATEGY.getBoolean(properties);
+  }
+
+  @Override
+  public Set<String> getSubscribedMethods() {
+    return subscribedMethods;
+  }
+
+  @Override
+  public void switchConnectionIfRequired(
+      final Object methodInvokeOn,
+      final String methodName,
+      final Object[] args) throws SQLException {
+    final Connection currentConnection = this.pluginService.getCurrentConnection();
+    final HostSpec currentHostSpec = this.pluginService.getCurrentHostSpec();
+    if (currentConnection == null || currentConnection.isClosed()) {
+      // do nothing; let normal workflow to handle it
+      return;
+    }
+
+    // Check if it's safe to switch to a new reader connection
+    if (this.inReadWriteSplit
+        && !pluginService.isInTransaction()
+        && (JdbcMethod.CONNECTION_CREATESTATEMENT.methodName.equals(methodName)
+            || JdbcMethod.CONNECTION_PREPARESTATEMENT.methodName.equals(methodName)
+            || JdbcMethod.CONNECTION_PREPARECALL.methodName.equals(methodName))
+        ) {
+      this.refreshAndStoreTopology(currentConnection);
+
+      try {
+        this.forceCloseReaderConnection();
+        this.initializeReaderConnection();
+        this.switchToReaderConnection();
+
+      } catch (final SQLException e) {
+        if (!isConnectionUsable(currentConnection)) {
+          logAndThrowException(
+              Messages.get("ReadWriteSplittingPlugin.errorSwitchingToReader", new Object[] {e.getMessage()}),
+              e);
+          return;
+        }
+
+        // Failed to switch to a reader. The current connection will be used as a fallback.
+        LOGGER.fine(() -> Messages.get(
+            "ReadWriteSplittingPlugin.fallbackToCurrentConnection",
+            new Object[] {
+                currentHostSpec.getHostAndPort(),
+                e.getMessage()}));
+        this.setReaderConnection(currentConnection, currentHostSpec);
+      }
+    }
+  }
+
+  protected void forceCloseReaderConnection() {
+    if (this.readerCacheItem == null) {
+      return;
+    }
+
+    final Connection readerConnection = this.readerCacheItem.get(true);
+
+    if (readerConnection != null) {
+      try {
+        if (!readerConnection.isClosed()) {
+          // readerConnection is open but is not currently in use, so we close it.
+          readerConnection.close();
+        }
+      } catch (SQLException e) {
+        // Do nothing.
+      }
+      this.readerCacheItem = null;
+    }
+  }
+
+  @Override
+  protected HostSpec getWriterHost(final @NonNull List<HostSpec> hosts) throws SQLException {
+    HostSpec writerHost = Utils.getWriter(hosts);
+    if (writerHost == null) {
+      if (hosts.isEmpty()) {
+        logAndThrowException(Messages.get("ReadWriteSplittingPlugin.noWriterFound"));
+      }
+      writerHost = hosts.get(0);
+    }
+
+    return writerHost;
+  }
+
+  @Override
+  protected @Nullable HostRole getReaderCandidateRole() {
+    return this.includeWriterToLoadBalancing ? null : HostRole.READER;
+  }
+}

wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/QueryLoadBalancingReadWriteSplittingPluginFactory.java

@@ -0,0 +1,30 @@
+/*
+ * 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.plugin.readwritesplitting;
+
+import java.util.Properties;
+import software.amazon.jdbc.ConnectionPlugin;
+import software.amazon.jdbc.ConnectionPluginFactory;
+import software.amazon.jdbc.util.FullServicesContainer;
+
+public class QueryLoadBalancingReadWriteSplittingPluginFactory implements ConnectionPluginFactory {
+
+  @Override
+  public ConnectionPlugin getInstance(final FullServicesContainer servicesContainer, final Properties props) {
+    return new QueryLoadBalancingReadWriteSplittingPlugin(servicesContainer.getPluginService(), props);
+  }
+}

wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java

@@ -20,11 +20,11 @@
 import java.sql.SQLException;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Map;
 import java.util.Properties;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import software.amazon.jdbc.AwsWrapperProperty;
 import software.amazon.jdbc.HostRole;
 import software.amazon.jdbc.HostSpec;
@@ -178,7 +178,7 @@ protected void closeReaderIfNecessary() {
     }
   }
 
-  private HostSpec getWriterHost(final @NonNull List<HostSpec> hosts) throws SQLException {
+  protected HostSpec getWriterHost(final @NonNull List<HostSpec> hosts) throws SQLException {
     HostSpec writerHost = Utils.getWriter(hosts);
     if (writerHost == null) {
       logAndThrowException(Messages.get("ReadWriteSplittingPlugin.noWriterFound"));
@@ -191,6 +191,10 @@ protected List<HostSpec> getReaderHostCandidates() throws SQLException {
     return this.pluginService.getHosts();
   }
 
+  protected @Nullable HostRole getReaderCandidateRole() {
+    return HostRole.READER;
+  }
+
   protected void openNewReaderConnection() throws SQLException {
     Connection conn = null;
     HostSpec readerHost = null;
@@ -199,7 +203,7 @@ protected void openNewReaderConnection() throws SQLException {
     int connAttempts = hostCandidates.size() * 2;
     for (int i = 0; i < connAttempts; i++) {
       HostSpec hostSpec = this.pluginService.getHostSpecByStrategy(
-          hostCandidates, HostRole.READER, this.readerSelectorStrategy);
+          hostCandidates, this.getReaderCandidateRole(), this.readerSelectorStrategy);
       try {
         conn = this.pluginService.connect(hostSpec, this.properties, this);
         this.isReaderConnFromInternalPool = Boolean.TRUE.equals(this.pluginService.isPooledConnection());