fix: Limitless plugin set round robin weights to original props used … (#1223)

Author: aaronchung-bitquill

wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessConnectionContext.java

@@ -27,19 +27,22 @@
 public class LimitlessConnectionContext {
   private HostSpec hostSpec;
   private Properties props;
+  private Properties origProps;
   private Connection connection;
   private JdbcCallable<Connection, SQLException> connectFunc;
   private List<HostSpec> limitlessRouters;
 
   public LimitlessConnectionContext(
       final HostSpec hostSpec,
       final Properties props,
+      final Properties origProps,
       final Connection connection,
       final JdbcCallable<Connection, SQLException> connectFunc,
       final List<HostSpec> limitlessRouters
   ) {
     this.hostSpec = hostSpec;
     this.props = props;
+    this.origProps = origProps;
     this.connection = connection;
     this.connectFunc = connectFunc;
     this.limitlessRouters = limitlessRouters;
@@ -53,6 +56,10 @@ public Properties getProps() {
     return this.props;
   }
 
+  public Properties getOrigProps() {
+    return this.origProps;
+  }
+
   public Connection getConnection() {
     return this.connection;
   }

wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessConnectionPlugin.java

@@ -111,13 +111,14 @@ public Connection connect(
 
     final Properties copyProps = PropertyUtils.copyProperties(props);
     copyProps.setProperty(INTERNAL_CONNECT_PROPERTY_NAME, "true");
-    return connectInternal(driverProtocol, hostSpec, copyProps, isInitialConnection, connectFunc);
+    return connectInternal(driverProtocol, hostSpec, props, copyProps, isInitialConnection, connectFunc);
   }
 
   public Connection connectInternal(
       final String driverProtocol,
       final HostSpec hostSpec,
-      final Properties props,
+      final Properties origProps,
+      final Properties copyProps,
       final boolean isInitialConnection,
       final JdbcCallable<Connection, SQLException> connectFunc)
       throws SQLException {
@@ -143,7 +144,8 @@ public Connection connectInternal(
 
     final LimitlessConnectionContext context = new LimitlessConnectionContext(
         hostSpec,
-        props,
+        copyProps,
+        origProps,
         conn,
         connectFunc,
         null);

wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterMonitor.java

@@ -137,7 +137,6 @@ public void run() {
             newLimitlessRouters,
             TimeUnit.MILLISECONDS.toNanos(LimitlessRouterServiceImpl.MONITOR_DISPOSAL_TIME_MS.getLong(props)));
 
-        RoundRobinHostSelector.setRoundRobinHostWeightPairsProperty(this.props, newLimitlessRouters);
         LOGGER.finest(Utils.logTopology(newLimitlessRouters, "[limitlessRouterMonitor] Topology:"));
         TimeUnit.MILLISECONDS.sleep(this.intervalMs); // do not include this in the telemetry
       } catch (final InterruptedException exception) {

wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImpl.java

@@ -131,7 +131,9 @@ public void establishConnection(final LimitlessConnectionContext context) throws
       return;
     }
 
-    RoundRobinHostSelector.setRoundRobinHostWeightPairsProperty(context.getProps(), context.getLimitlessRouters());
+    RoundRobinHostSelector.setRoundRobinHostWeightPairsProperty(
+        context.getOrigProps(),
+        context.getLimitlessRouters());
     HostSpec selectedHostSpec;
     try {
       selectedHostSpec = this.pluginService.getHostSpecByStrategy(
@@ -165,7 +167,7 @@ public void establishConnection(final LimitlessConnectionContext context) throws
     }
   }
 
-  protected List<HostSpec> getLimitlessRouters(final String clusterId, final Properties props) throws SQLException {
+  protected List<HostSpec> getLimitlessRouters(final String clusterId, final Properties props) {
     final long cacheExpirationNano = TimeUnit.MILLISECONDS.toNanos(
         MONITOR_DISPOSAL_TIME_MS.getLong(props));
     return limitlessRouterCache.get(clusterId, cacheExpirationNano);

wrapper/src/test/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImplTest.java

@@ -47,6 +47,7 @@
 import software.amazon.jdbc.RoundRobinHostSelector;
 import software.amazon.jdbc.hostavailability.HostAvailability;
 import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy;
+import software.amazon.jdbc.util.PropertyUtils;
 import software.amazon.jdbc.wrapper.HighestWeightHostSelector;
 
 public class LimitlessRouterServiceImplTest {
@@ -70,6 +71,7 @@ public void init() throws SQLException {
     props = new Properties();
     when(mockConnectFuncLambda.call()).thenReturn(mockConnection);
     when(mockPluginService.getHostListProvider()).thenReturn(mockHostListProvider);
+    when(mockPluginService.getProperties()).thenReturn(props);
     when(mockHostListProvider.getClusterId()).thenReturn(CLUSTER_ID);
   }
 
@@ -85,6 +87,7 @@ void testEstablishConnection_GivenGetEmptyRouterListAndWaitForRouterInfo_ThenThr
 
     final LimitlessConnectionContext inputContext = new LimitlessConnectionContext(
         hostSpec,
+        PropertyUtils.copyProperties(props),
         props,
         null,
         mockConnectFuncLambda,
@@ -103,6 +106,7 @@ void testEstablishConnection_GivenGetEmptyRouterListAndNoWaitForRouterInfo_ThenC
     props.setProperty(LimitlessConnectionPlugin.WAIT_FOR_ROUTER_INFO.name, "false");
     final LimitlessConnectionContext inputContext = new LimitlessConnectionContext(
         hostSpec,
+        PropertyUtils.copyProperties(props),
         props,
         null,
         mockConnectFuncLambda,
@@ -132,6 +136,7 @@ void testEstablishConnection_GivenHostSpecInRouterCache_ThenCallConnectFunc() th
 
     final LimitlessConnectionContext inputContext = new LimitlessConnectionContext(
         routerList.get(1),
+        PropertyUtils.copyProperties(props),
         props,
         null,
         mockConnectFuncLambda,
@@ -163,6 +168,7 @@ void testEstablishConnection_GivenFetchRouterListAndHostSpecInRouterList_ThenCal
 
     final LimitlessConnectionContext inputContext = new LimitlessConnectionContext(
         routerList.get(1),
+        PropertyUtils.copyProperties(props),
         props,
         null,
         mockConnectFuncLambda,
@@ -199,6 +205,7 @@ void testEstablishConnection_GivenRouterCache_ThenSelectsHost() throws SQLExcept
 
     final LimitlessConnectionContext inputContext = new LimitlessConnectionContext(
         hostSpec,
+        PropertyUtils.copyProperties(props),
         props,
         null,
         mockConnectFuncLambda,
@@ -234,6 +241,7 @@ void testEstablishConnection_GivenFetchRouterList_ThenSelectsHost() throws SQLEx
 
     final LimitlessConnectionContext inputContext = new LimitlessConnectionContext(
         hostSpec,
+        PropertyUtils.copyProperties(props),
         props,
         null,
         mockConnectFuncLambda,
@@ -267,6 +275,7 @@ void testEstablishConnection_GivenHostSpecInRouterCacheAndCallConnectFuncThrows_
     final HostSpec selectedRouter = routerList.get(2);
     final LimitlessConnectionContext inputContext = new LimitlessConnectionContext(
         routerList.get(1),
+        PropertyUtils.copyProperties(props),
         props,
         null,
         mockConnectFuncLambda,
@@ -311,6 +320,7 @@ void testEstablishConnection_GivenSelectsHostThrows_ThenRetry() throws SQLExcept
 
     final LimitlessConnectionContext inputContext = new LimitlessConnectionContext(
         hostSpec,
+        PropertyUtils.copyProperties(props),
         props,
         null,
         mockConnectFuncLambda,
@@ -350,6 +360,7 @@ void testEstablishConnection_GivenSelectsHostNull_ThenRetry() throws SQLExceptio
 
     final LimitlessConnectionContext inputContext = new LimitlessConnectionContext(
         hostSpec,
+        PropertyUtils.copyProperties(props),
         props,
         null,
         mockConnectFuncLambda,
@@ -392,6 +403,7 @@ void testEstablishConnection_GivenPluginServiceConnectThrows_ThenRetry() throws
 
     final LimitlessConnectionContext inputContext = new LimitlessConnectionContext(
         hostSpec,
+        PropertyUtils.copyProperties(props),
         props,
         null,
         mockConnectFuncLambda,
@@ -429,6 +441,7 @@ void testEstablishConnection_GivenRetryAndMaxRetriesExceeded_thenThrowSqlExcepti
 
     final LimitlessConnectionContext inputContext = new LimitlessConnectionContext(
         routerList.get(0),
+        PropertyUtils.copyProperties(props),
         props,
         null,
         mockConnectFuncLambda,