From ed3bbe71074738a6c653a7ece089f099b392b61f Mon Sep 17 00:00:00 2001 From: aaronchung-bitquill <118320132+aaronchung-bitquill@users.noreply.github.com> Date: Fri, 18 Oct 2024 10:25:25 -0700 Subject: [PATCH] fix: limitless plugin - add null checks (#1152) --- .../limitless/LimitlessConnectionPlugin.java | 40 +++++++++++++------ .../limitless/LimitlessRouterMonitor.java | 2 +- .../LimitlessRouterMonitorInitializer.java | 1 - .../limitless/LimitlessRouterServiceImpl.java | 17 +++++++- .../software/amazon/jdbc/util/RdsUrlType.java | 2 +- ..._advanced_jdbc_wrapper_messages.properties | 3 +- .../LimitlessRouterServiceImplTest.java | 4 +- 7 files changed, 49 insertions(+), 20 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessConnectionPlugin.java index 187009625..c9403a9b1 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessConnectionPlugin.java @@ -161,7 +161,7 @@ private Connection connectInternalWithDialect( try { conn = connectFunc.call(); } catch (final SQLException e) { - return retryConnectWithLeastLoadedRouters(limitlessRouters, props, conn, hostSpec); + return retryConnectWithLeastLoadedRouters(limitlessRouters, props, null, hostSpec); } } return conn; @@ -174,7 +174,7 @@ private Connection connectInternalWithDialect( HostRole.WRITER, RoundRobinHostSelector.STRATEGY_ROUND_ROBIN); LOGGER.fine(Messages.get( "LimitlessConnectionPlugin.selectedHost", - new Object[] {selectedHostSpec.getHost()})); + new Object[] {selectedHostSpec != null ? selectedHostSpec.getHost() : "null"})); } catch (SQLException e) { LOGGER.warning(Messages.get("LimitlessConnectionPlugin.errorSelectingRouter", new Object[] {e.getMessage()})); if (conn == null || conn.isClosed()) { @@ -186,10 +186,12 @@ private Connection connectInternalWithDialect( try { return pluginService.connect(selectedHostSpec, props); } catch (SQLException e) { - LOGGER.fine(Messages.get( - "LimitlessConnectionPlugin.failedToConnectToHost", - new Object[] {selectedHostSpec.getHost()})); - selectedHostSpec.setAvailability(HostAvailability.NOT_AVAILABLE); + if (selectedHostSpec != null) { + LOGGER.fine(Messages.get( + "LimitlessConnectionPlugin.failedToConnectToHost", + new Object[] {selectedHostSpec.getHost()})); + selectedHostSpec.setAvailability(HostAvailability.NOT_AVAILABLE); + } if (conn == null || conn.isClosed()) { conn = connectFunc.call(); } @@ -238,8 +240,11 @@ private void initLimitlessRouterMonitorService() { } } - private Connection retryConnectWithLeastLoadedRouters(final List limitlessRouters, final Properties props, - final Connection conn, final HostSpec hostSpec) throws SQLException { + private Connection retryConnectWithLeastLoadedRouters( + final List limitlessRouters, + final Properties props, + final Connection conn, + final HostSpec hostSpec) throws SQLException { List currentRouters = limitlessRouters; int retryCount = 0; @@ -247,12 +252,19 @@ private Connection retryConnectWithLeastLoadedRouters(final List limit while (retryCount++ < maxRetries) { if (currentRouters.stream().noneMatch(h -> h.getAvailability().equals(HostAvailability.AVAILABLE))) { - currentRouters = synchronouslyGetLimitlessRoutersWithRetry(conn, hostSpec.getPort(), props); + if (conn != null && !conn.isClosed()) { + currentRouters = synchronouslyGetLimitlessRoutersWithRetry(conn, hostSpec.getPort(), props); + } + if (currentRouters == null || currentRouters.isEmpty() || currentRouters.stream().noneMatch(h -> h.getAvailability().equals(HostAvailability.AVAILABLE))) { LOGGER.warning(Messages.get("LimitlessConnectionPlugin.noRoutersAvailableForRetry")); - return conn; + if (conn != null && !conn.isClosed()) { + return conn; + } else { + throw new SQLException(Messages.get("LimitlessConnectionPlugin.noRoutersAvailable")); + } } } @@ -281,8 +293,12 @@ private Connection retryConnectWithLeastLoadedRouters(final List limit new Object[] {selectedHostSpec.getHost()})); } } - LOGGER.warning(Messages.get("LimitlessConnectionPlugin.maxRetriesExceeded")); - return conn; + if (conn != null && !conn.isClosed()) { + LOGGER.warning(Messages.get("LimitlessConnectionPlugin.maxRetriesExceeded")); + return conn; + } else { + throw new SQLException(Messages.get("LimitlessConnectionPlugin.maxRetriesExceeded")); + } } private List synchronouslyGetLimitlessRoutersWithRetry( diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterMonitor.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterMonitor.java index 98b880eaa..0fead0ce5 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterMonitor.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterMonitor.java @@ -137,7 +137,7 @@ public void run() { while (!this.stopped.get()) { TelemetryContext telemetryContext = telemetryFactory.openTelemetryContext( - "node response time thread", TelemetryTraceLevel.TOP_LEVEL); + "limitless router monitor thread", TelemetryTraceLevel.TOP_LEVEL); telemetryContext.setAttribute("url", hostSpec.getUrl()); try { this.openConnection(); diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterMonitorInitializer.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterMonitorInitializer.java index 00c2b9a45..d036f2d42 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterMonitorInitializer.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterMonitorInitializer.java @@ -27,7 +27,6 @@ @FunctionalInterface public interface LimitlessRouterMonitorInitializer { LimitlessRouterMonitor createLimitlessRouterMonitor( - final PluginService pluginService, final HostSpec hostSpec, final SlidingExpirationCacheWithCleanupThread> limitlessRouterCache, final String limitlessRouterCacheKey, diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImpl.java index 1cbcd2c52..954d98958 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImpl.java @@ -65,7 +65,21 @@ public class LimitlessRouterServiceImpl implements LimitlessRouterService { ); public LimitlessRouterServiceImpl(final @NonNull PluginService pluginService) { - this(pluginService, LimitlessRouterMonitor::new, new LimitlessQueryHelper(pluginService)); + this( + pluginService, + (hostSpec, + routerCache, + routerCacheKey, + props, + intervalMs) -> + new LimitlessRouterMonitor( + pluginService, + hostSpec, + routerCache, + routerCacheKey, + props, + intervalMs), + new LimitlessQueryHelper(pluginService)); } public LimitlessRouterServiceImpl( @@ -127,7 +141,6 @@ public void startMonitoring(final @NonNull HostSpec hostSpec, limitlessRouterMonitorKey, key -> this.limitlessRouterMonitorInitializer .createLimitlessRouterMonitor( - pluginService, hostSpec, limitlessRouterCache, this.pluginService.getHostListProvider().getClusterId(), diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/RdsUrlType.java b/wrapper/src/main/java/software/amazon/jdbc/util/RdsUrlType.java index 48845c04a..dff1e663a 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/RdsUrlType.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/RdsUrlType.java @@ -23,7 +23,7 @@ public enum RdsUrlType { RDS_CUSTOM_CLUSTER(true, true), RDS_PROXY(true, false), RDS_INSTANCE(true, false), - RDS_AURORA_LIMITLESS_DB_SHARD_GROUP(true, true), + RDS_AURORA_LIMITLESS_DB_SHARD_GROUP(true, false), OTHER(false, false); private final boolean isRds; diff --git a/wrapper/src/main/resources/aws_advanced_jdbc_wrapper_messages.properties b/wrapper/src/main/resources/aws_advanced_jdbc_wrapper_messages.properties index 2b471e686..40bf87093 100644 --- a/wrapper/src/main/resources/aws_advanced_jdbc_wrapper_messages.properties +++ b/wrapper/src/main/resources/aws_advanced_jdbc_wrapper_messages.properties @@ -187,7 +187,7 @@ IamAuthConnectionPlugin.missingRequiredConfigParameter=Configuration parameter ' # Limitless Connection Plugin LimitlessConnectionPlugin.connectWithHost=Connecting to host {0}. -LimitlessConnectionPlugin.usingProvidedConnectUrl=Connecting using provided connection URL. +LimitlessConnectionPlugin.errorSelectingRouter=An error occurred while selecting Limitless Transaction Router: {0} LimitlessConnectionPlugin.failedToConnectToHost=Failed to connect to host {0}. LimitlessConnectionPlugin.incorrectConfiguration=Limitless Connection Plugin is unable to run. Please ensure the connection settings are correct. LimitlessConnectionPlugin.limitlessRouterCacheEmpty=Limitless Router cache is empty. This normal during application start up when the cache is not yet populated. @@ -198,6 +198,7 @@ LimitlessConnectionPlugin.selectedHost=Host {0} has been selected. LimitlessConnectionPlugin.selectedHostForRetry=Host {0} has been selected for connection retry. LimitlessConnectionPlugin.synchronouslyGetLimitlessRouters=Fetching Limitless Routers synchronously. LimitlessConnectionPlugin.unsupportedDialectOrDatabase=Unsupported dialect ''{0}'' encountered. Please ensure JDBC connection parameters are correct, and refer to the documentation to ensure that the connecting database is compatible with the Limitless Connection Plugin. +LimitlessConnectionPlugin.usingProvidedConnectUrl=Connecting using provided connection URL. # Limitless Query Helper LimitlessQueryHelper.unsupportedDialectOrDatabase=Unsupported dialect ''{0}'' encountered. Please ensure JDBC connection parameters are correct, and refer to the documentation to ensure that the connecting database is compatible with the Limitless Connection Plugin. diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImplTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImplTest.java index 1434cb829..1429842f4 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImplTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterServiceImplTest.java @@ -83,7 +83,7 @@ void testGetLimitlessRouters() throws SQLException { final LimitlessRouterService limitlessRouterService = new LimitlessRouterServiceImpl( mockPluginService, - (a, b, c, d, e, f) -> mockLimitlessRouterMonitor, + (a, b, c, d, e) -> mockLimitlessRouterMonitor, mockLimitlessQueryHelper); limitlessRouterService.startMonitoring(hostSpec, props, intervalMs); final List actualEndpointHostSpecList = limitlessRouterService.getLimitlessRouters(CLUSTER_ID, props); @@ -112,7 +112,7 @@ void testForceGetLimitlessRoutersWithConn() throws SQLException { final LimitlessRouterServiceImpl limitlessRouterService = new LimitlessRouterServiceImpl( mockPluginService, - (a, b, c, d, e, f) -> mockLimitlessRouterMonitor, + (a, b, c, d, e) -> mockLimitlessRouterMonitor, mockLimitlessQueryHelper); final List actualHostSpecList = limitlessRouterService