From b94b7f24b393dd35181c536c0bc603ba5603d1e1 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Fri, 15 Aug 2025 09:36:25 -0700 Subject: [PATCH 01/42] wip --- .../ClusterAwareWriterFailoverHandler.java | 17 +- .../failover/FailoverConnectionPlugin.java | 39 +- .../FailoverConnectionPluginFactory.java | 16 +- .../jdbc/util/CoreServicesContainer.java | 2 +- ...ClusterAwareWriterFailoverHandlerTest.java | 816 ++++++++-------- .../FailoverConnectionPluginTest.java | 884 +++++++++--------- 6 files changed, 909 insertions(+), 865 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java index 47f741b5e..07d04d4f9 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java @@ -35,9 +35,11 @@ import software.amazon.jdbc.PluginService; import software.amazon.jdbc.hostavailability.HostAvailability; import software.amazon.jdbc.util.ExecutorFactory; +import software.amazon.jdbc.util.FullServicesContainer; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.PropertyUtils; import software.amazon.jdbc.util.Utils; +import software.amazon.jdbc.util.connection.ConnectionService; /** * An implementation of WriterFailoverHandler. @@ -56,28 +58,33 @@ public class ClusterAwareWriterFailoverHandler implements WriterFailoverHandler protected int reconnectWriterIntervalMs = 5000; // 5 sec protected Properties initialConnectionProps; protected PluginService pluginService; + protected ConnectionService connectionService; protected ReaderFailoverHandler readerFailoverHandler; private static final WriterFailoverResult DEFAULT_RESULT = new WriterFailoverResult(false, false, null, null, "None"); public ClusterAwareWriterFailoverHandler( - final PluginService pluginService, + final FullServicesContainer servicesContainer, + final ConnectionService connectionService, final ReaderFailoverHandler readerFailoverHandler, final Properties initialConnectionProps) { - this.pluginService = pluginService; + this.pluginService = servicesContainer.getPluginService(); + this.connectionService = connectionService; this.readerFailoverHandler = readerFailoverHandler; this.initialConnectionProps = initialConnectionProps; } public ClusterAwareWriterFailoverHandler( - final PluginService pluginService, + final FullServicesContainer servicesContainer, + final ConnectionService connectionService, final ReaderFailoverHandler readerFailoverHandler, final Properties initialConnectionProps, final int failoverTimeoutMs, final int readTopologyIntervalMs, final int reconnectWriterIntervalMs) { this( - pluginService, + servicesContainer, + connectionService, readerFailoverHandler, initialConnectionProps); this.maxFailoverTimeoutMs = failoverTimeoutMs; @@ -260,7 +267,7 @@ public WriterFailoverResult call() { // TODO: assess whether multi-threaded access to the plugin service is safe. The same plugin service is used // by both the ConnectionWrapper and this ReconnectToWriterHandler in separate threads. - conn = pluginService.forceConnect(this.originalWriterHost, initialConnectionProps); + conn = connectionService.open(this.originalWriterHost, initialConnectionProps); pluginService.forceRefreshHostList(conn); latestTopology = pluginService.getAllHosts(); diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java index 8c16484ea..1807a2f8c 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java @@ -33,6 +33,8 @@ import java.util.logging.Logger; import org.checkerframework.checker.nullness.qual.NonNull; import software.amazon.jdbc.AwsWrapperProperty; +import software.amazon.jdbc.ConnectionProvider; +import software.amazon.jdbc.DriverConnectionProvider; import software.amazon.jdbc.HostListProviderService; import software.amazon.jdbc.HostRole; import software.amazon.jdbc.HostSpec; @@ -42,16 +44,20 @@ import software.amazon.jdbc.PluginManagerService; import software.amazon.jdbc.PluginService; import software.amazon.jdbc.PropertyDefinition; +import software.amazon.jdbc.TargetDriverHelper; import software.amazon.jdbc.hostavailability.HostAvailability; import software.amazon.jdbc.plugin.AbstractConnectionPlugin; import software.amazon.jdbc.plugin.staledns.AuroraStaleDnsHelper; import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; +import software.amazon.jdbc.util.FullServicesContainer; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.RdsUrlType; import software.amazon.jdbc.util.RdsUtils; import software.amazon.jdbc.util.SqlState; import software.amazon.jdbc.util.Utils; import software.amazon.jdbc.util.WrapperUtils; +import software.amazon.jdbc.util.connection.ConnectionService; +import software.amazon.jdbc.util.connection.ConnectionServiceImpl; import software.amazon.jdbc.util.telemetry.TelemetryContext; import software.amazon.jdbc.util.telemetry.TelemetryCounter; import software.amazon.jdbc.util.telemetry.TelemetryFactory; @@ -92,6 +98,8 @@ public class FailoverConnectionPlugin extends AbstractConnectionPlugin { private final Set subscribedMethods; private final PluginService pluginService; + private final FullServicesContainer servicesContainer; + private final ConnectionService connectionService; protected final Properties properties; protected boolean enableFailoverSetting; protected boolean enableConnectFailover; @@ -185,15 +193,16 @@ public class FailoverConnectionPlugin extends AbstractConnectionPlugin { PropertyDefinition.registerPluginProperties(FailoverConnectionPlugin.class); } - public FailoverConnectionPlugin(final PluginService pluginService, final Properties properties) { - this(pluginService, properties, new RdsUtils()); + public FailoverConnectionPlugin(final FullServicesContainer servicesContainer, final Properties properties) { + this(servicesContainer, properties, new RdsUtils()); } FailoverConnectionPlugin( - final PluginService pluginService, + final FullServicesContainer servicesContainer, final Properties properties, final RdsUtils rdsHelper) { - this.pluginService = pluginService; + this.servicesContainer = servicesContainer; + this.pluginService = servicesContainer.getPluginService(); this.properties = properties; this.rdsHelper = rdsHelper; @@ -222,6 +231,25 @@ public FailoverConnectionPlugin(final PluginService pluginService, final Propert } this.subscribedMethods = Collections.unmodifiableSet(methods); + try { + TargetDriverHelper helper = new TargetDriverHelper(); + java.sql.Driver driver = helper.getTargetDriver(this.pluginService.getOriginalUrl(), properties); + final ConnectionProvider defaultConnectionProvider = new DriverConnectionProvider(driver); + this.connectionService = new ConnectionServiceImpl( + servicesContainer.getStorageService(), + servicesContainer.getMonitorService(), + servicesContainer.getTelemetryFactory(), + defaultConnectionProvider, + this.pluginService.getOriginalUrl(), + this.pluginService.getDriverProtocol(), + this.pluginService.getTargetDriverDialect(), + this.pluginService.getDialect(), + properties + ); + } catch (SQLException e) { + throw new RuntimeException(e); + } + TelemetryFactory telemetryFactory = this.pluginService.getTelemetryFactory(); this.failoverWriterTriggeredCounter = telemetryFactory.createCounter("writerFailover.triggered.count"); this.failoverWriterSuccessCounter = telemetryFactory.createCounter("writerFailover.completed.success.count"); @@ -316,7 +344,8 @@ public void initHostProvider( this.failoverMode == FailoverMode.STRICT_READER), () -> new ClusterAwareWriterFailoverHandler( - this.pluginService, + this.servicesContainer, + this.connectionService, this.readerFailoverHandler, this.properties, this.failoverTimeoutMsSetting, diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginFactory.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginFactory.java index 75d445710..9fd77a767 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginFactory.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginFactory.java @@ -18,13 +18,21 @@ import java.util.Properties; import software.amazon.jdbc.ConnectionPlugin; -import software.amazon.jdbc.ConnectionPluginFactory; import software.amazon.jdbc.PluginService; +import software.amazon.jdbc.ServicesContainerPluginFactory; +import software.amazon.jdbc.util.FullServicesContainer; +import software.amazon.jdbc.util.Messages; -public class FailoverConnectionPluginFactory implements ConnectionPluginFactory { - +public class FailoverConnectionPluginFactory implements ServicesContainerPluginFactory { @Override public ConnectionPlugin getInstance(final PluginService pluginService, final Properties props) { - return new FailoverConnectionPlugin(pluginService, props); + throw new UnsupportedOperationException( + Messages.get( + "ServicesContainerPluginFactory.servicesContainerRequired", new Object[] {"FailoverConnectionPlugin"})); + } + + @Override + public ConnectionPlugin getInstance(final FullServicesContainer servicesContainer, final Properties props) { + return new FailoverConnectionPlugin(servicesContainer, props); } } diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/CoreServicesContainer.java b/wrapper/src/main/java/software/amazon/jdbc/util/CoreServicesContainer.java index 3ac3e0a1b..1c01d5bbd 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/CoreServicesContainer.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/CoreServicesContainer.java @@ -33,8 +33,8 @@ public class CoreServicesContainer { private static final CoreServicesContainer INSTANCE = new CoreServicesContainer(); - private final StorageService storageService; private final MonitorService monitorService; + private final StorageService storageService; private CoreServicesContainer() { EventPublisher eventPublisher = new BatchingEventPublisher(); diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java index 0944be13c..5f302209a 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java @@ -1,408 +1,408 @@ -/* - * 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.failover; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertSame; -import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.ArgumentMatchers.refEq; -import static org.mockito.Mockito.atLeastOnce; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -import java.sql.Connection; -import java.sql.SQLException; -import java.util.Arrays; -import java.util.EnumSet; -import java.util.List; -import java.util.Properties; -import java.util.concurrent.TimeUnit; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.mockito.ArgumentMatchers; -import org.mockito.InOrder; -import org.mockito.Mock; -import org.mockito.Mockito; -import org.mockito.MockitoAnnotations; -import org.mockito.stubbing.Answer; -import software.amazon.jdbc.HostSpec; -import software.amazon.jdbc.HostSpecBuilder; -import software.amazon.jdbc.PluginService; -import software.amazon.jdbc.dialect.Dialect; -import software.amazon.jdbc.hostavailability.HostAvailability; -import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; - -class ClusterAwareWriterFailoverHandlerTest { - - @Mock PluginService mockPluginService; - @Mock Connection mockConnection; - @Mock ReaderFailoverHandler mockReaderFailover; - @Mock Connection mockWriterConnection; - @Mock Connection mockNewWriterConnection; - @Mock Connection mockReaderAConnection; - @Mock Connection mockReaderBConnection; - @Mock Dialect mockDialect; - - private AutoCloseable closeable; - private final Properties properties = new Properties(); - private final HostSpec newWriterHost = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("new-writer-host").build(); - private final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("writer-host").build(); - private final HostSpec readerA = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader-a-host").build(); - private final HostSpec readerB = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader-b-host").build(); - private final List topology = Arrays.asList(writer, readerA, readerB); - private final List newTopology = Arrays.asList(newWriterHost, readerA, readerB); - - @BeforeEach - void setUp() { - closeable = MockitoAnnotations.openMocks(this); - writer.addAlias("writer-host"); - newWriterHost.addAlias("new-writer-host"); - readerA.addAlias("reader-a-host"); - readerB.addAlias("reader-b-host"); - } - - @AfterEach - void tearDown() throws Exception { - closeable.close(); - } - - @Test - public void testReconnectToWriter_taskBReaderException() throws SQLException { - when(mockPluginService.forceConnect(refEq(writer), eq(properties))).thenReturn(mockConnection); - when(mockPluginService.forceConnect(refEq(readerA), eq(properties))).thenThrow(SQLException.class); - when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenThrow(SQLException.class); - - when(mockPluginService.getAllHosts()).thenReturn(topology); - - when(mockReaderFailover.getReaderConnection(ArgumentMatchers.anyList())).thenThrow(SQLException.class); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = - new ClusterAwareWriterFailoverHandler( - mockPluginService, - mockReaderFailover, - properties, - 5000, - 2000, - 2000); - final WriterFailoverResult result = target.failover(topology); - - assertTrue(result.isConnected()); - assertFalse(result.isNewHost()); - assertSame(result.getNewConnection(), mockConnection); - - final InOrder inOrder = Mockito.inOrder(mockPluginService); - inOrder.verify(mockPluginService).setAvailability(eq(writer.asAliases()), eq(HostAvailability.AVAILABLE)); - } - - /** - * Verify that writer failover handler can re-connect to a current writer node. - * - *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. - * TaskA: successfully re-connect to initial writer; return new connection. - * TaskB: successfully connect to readerA and then new writer, but it takes more time than taskA. - * Expected test result: new connection by taskA. - */ - @Test - public void testReconnectToWriter_SlowReaderA() throws SQLException { - when(mockPluginService.forceConnect(refEq(writer), eq(properties))).thenReturn(mockWriterConnection); - when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenThrow(SQLException.class); - when(mockPluginService.forceConnect(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); - when(mockPluginService.getAllHosts()).thenReturn(topology).thenReturn(newTopology); - - when(mockReaderFailover.getReaderConnection(ArgumentMatchers.anyList())) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(5000); - return new ReaderFailoverResult(mockReaderAConnection, readerA, true); - }); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = - new ClusterAwareWriterFailoverHandler( - mockPluginService, - mockReaderFailover, - properties, - 60000, - 5000, - 5000); - final WriterFailoverResult result = target.failover(topology); - - assertTrue(result.isConnected()); - assertFalse(result.isNewHost()); - assertSame(result.getNewConnection(), mockWriterConnection); - - final InOrder inOrder = Mockito.inOrder(mockPluginService); - inOrder.verify(mockPluginService).setAvailability(eq(writer.asAliases()), eq(HostAvailability.AVAILABLE)); - } - - /** - * Verify that writer failover handler can re-connect to a current writer node. - * - *

Topology: no changes. - * TaskA: successfully re-connect to writer; return new connection. - * TaskB: successfully connect to readerA and retrieve topology, but latest writer is not new (defer to taskA). - * Expected test result: new connection by taskA. - */ - @Test - public void testReconnectToWriter_taskBDefers() throws SQLException { - when(mockPluginService.forceConnect(refEq(writer), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(5000); - return mockWriterConnection; - }); - when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenThrow(SQLException.class); - - when(mockPluginService.getAllHosts()).thenReturn(topology); - - when(mockReaderFailover.getReaderConnection(ArgumentMatchers.anyList())) - .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = - new ClusterAwareWriterFailoverHandler( - mockPluginService, - mockReaderFailover, - properties, - 60000, - 2000, - 2000); - final WriterFailoverResult result = target.failover(topology); - - assertTrue(result.isConnected()); - assertFalse(result.isNewHost()); - assertSame(result.getNewConnection(), mockWriterConnection); - - final InOrder inOrder = Mockito.inOrder(mockPluginService); - inOrder.verify(mockPluginService).setAvailability(eq(writer.asAliases()), eq(HostAvailability.AVAILABLE)); - } - - /** - * Verify that writer failover handler can re-connect to a new writer node. - * - *

Topology: changes to [new-writer, reader-A, reader-B] for taskB, taskA sees no changes. - * taskA: successfully re-connect to writer; return connection to initial writer, but it takes more - * time than taskB. - * TaskB: successfully connect to readerA and then to new-writer. - * Expected test result: new connection to writer by taskB. - */ - @Test - public void testConnectToReaderA_SlowWriter() throws SQLException { - when(mockPluginService.forceConnect(refEq(writer), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(5000); - return mockWriterConnection; - }); - when(mockPluginService.forceConnect(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); - when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); - when(mockPluginService.forceConnect(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); - - when(mockPluginService.getAllHosts()).thenReturn(newTopology); - - when(mockReaderFailover.getReaderConnection(ArgumentMatchers.anyList())) - .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = - new ClusterAwareWriterFailoverHandler( - mockPluginService, - mockReaderFailover, - properties, - 60000, - 5000, - 5000); - final WriterFailoverResult result = target.failover(topology); - - assertTrue(result.isConnected()); - assertTrue(result.isNewHost()); - assertSame(result.getNewConnection(), mockNewWriterConnection); - assertEquals(3, result.getTopology().size()); - assertEquals("new-writer-host", result.getTopology().get(0).getHost()); - - verify(mockPluginService, times(1)).setAvailability(eq(newWriterHost.asAliases()), eq(HostAvailability.AVAILABLE)); - } - - /** - * Verify that writer failover handler can re-connect to a new writer node. - * - *

Topology: changes to [new-writer, initial-writer, reader-A, reader-B]. - * TaskA: successfully reconnect, but initial-writer is now a reader (defer to taskB). - * TaskB: successfully connect to readerA and then to new-writer. - * Expected test result: new connection to writer by taskB. - */ - @Test - public void testConnectToReaderA_taskADefers() throws SQLException { - when(mockPluginService.forceConnect(writer, properties)).thenReturn(mockConnection); - when(mockPluginService.forceConnect(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); - when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); - when(mockPluginService.forceConnect(refEq(newWriterHost), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(5000); - return mockNewWriterConnection; - }); - - final List newTopology = Arrays.asList(newWriterHost, writer, readerA, readerB); - when(mockPluginService.getAllHosts()).thenReturn(newTopology); - - when(mockReaderFailover.getReaderConnection(ArgumentMatchers.anyList())) - .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = - new ClusterAwareWriterFailoverHandler( - mockPluginService, - mockReaderFailover, - properties, - 60000, - 5000, - 5000); - final WriterFailoverResult result = target.failover(topology); - - assertTrue(result.isConnected()); - assertTrue(result.isNewHost()); - assertSame(result.getNewConnection(), mockNewWriterConnection); - assertEquals(4, result.getTopology().size()); - assertEquals("new-writer-host", result.getTopology().get(0).getHost()); - - verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); - verify(mockPluginService, times(1)).setAvailability(eq(newWriterHost.asAliases()), eq(HostAvailability.AVAILABLE)); - } - - /** - * Verify that writer failover handler fails to re-connect to any writer node. - * - *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. - * TaskA: fail to re-connect to writer due to failover timeout. - * TaskB: successfully connect to readerA and then fail to connect to writer due to failover timeout. - * Expected test result: no connection. - */ - @Test - public void testFailedToConnect_failoverTimeout() throws SQLException { - when(mockPluginService.forceConnect(refEq(writer), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(30000); - return mockWriterConnection; - }); - when(mockPluginService.forceConnect(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); - when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); - when(mockPluginService.forceConnect(refEq(newWriterHost), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(30000); - return mockNewWriterConnection; - }); - when(mockPluginService.getAllHosts()).thenReturn(newTopology); - - when(mockReaderFailover.getReaderConnection(ArgumentMatchers.anyList())) - .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = - new ClusterAwareWriterFailoverHandler( - mockPluginService, - mockReaderFailover, - properties, - 5000, - 2000, - 2000); - - final long startTimeNano = System.nanoTime(); - final WriterFailoverResult result = target.failover(topology); - final long durationNano = System.nanoTime() - startTimeNano; - - assertFalse(result.isConnected()); - assertFalse(result.isNewHost()); - - verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); - - // 5s is a max allowed failover timeout; add 1s for inaccurate measurements - assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); - } - - /** - * Verify that writer failover handler fails to re-connect to any writer node. - * - *

Topology: changes to [new-writer, reader-A, reader-B] for taskB. - * TaskA: fail to re-connect to writer due to exception. - * TaskB: successfully connect to readerA and then fail to connect to writer due to exception. - * Expected test result: no connection. - */ - @Test - public void testFailedToConnect_taskAException_taskBWriterException() throws SQLException { - final SQLException exception = new SQLException("exception", "08S01", null); - when(mockPluginService.forceConnect(refEq(writer), eq(properties))).thenThrow(exception); - when(mockPluginService.forceConnect(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); - when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); - when(mockPluginService.forceConnect(refEq(newWriterHost), eq(properties))).thenThrow(exception); - when(mockPluginService.isNetworkException(eq(exception), any())).thenReturn(true); - - when(mockPluginService.getAllHosts()).thenReturn(newTopology); - - when(mockReaderFailover.getReaderConnection(ArgumentMatchers.anyList())) - .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = - new ClusterAwareWriterFailoverHandler( - mockPluginService, - mockReaderFailover, - properties, - 5000, - 2000, - 2000); - final WriterFailoverResult result = target.failover(topology); - - assertFalse(result.isConnected()); - assertFalse(result.isNewHost()); - - verify(mockPluginService, atLeastOnce()) - .setAvailability(eq(newWriterHost.asAliases()), eq(HostAvailability.NOT_AVAILABLE)); - } -} +// /* +// * 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.failover; +// +// import static org.junit.jupiter.api.Assertions.assertEquals; +// import static org.junit.jupiter.api.Assertions.assertFalse; +// import static org.junit.jupiter.api.Assertions.assertSame; +// import static org.junit.jupiter.api.Assertions.assertTrue; +// import static org.mockito.ArgumentMatchers.any; +// import static org.mockito.ArgumentMatchers.eq; +// import static org.mockito.ArgumentMatchers.refEq; +// import static org.mockito.Mockito.atLeastOnce; +// import static org.mockito.Mockito.times; +// import static org.mockito.Mockito.verify; +// import static org.mockito.Mockito.when; +// +// import java.sql.Connection; +// import java.sql.SQLException; +// import java.util.Arrays; +// import java.util.EnumSet; +// import java.util.List; +// import java.util.Properties; +// import java.util.concurrent.TimeUnit; +// import org.junit.jupiter.api.AfterEach; +// import org.junit.jupiter.api.BeforeEach; +// import org.junit.jupiter.api.Test; +// import org.mockito.ArgumentMatchers; +// import org.mockito.InOrder; +// import org.mockito.Mock; +// import org.mockito.Mockito; +// import org.mockito.MockitoAnnotations; +// import org.mockito.stubbing.Answer; +// import software.amazon.jdbc.HostSpec; +// import software.amazon.jdbc.HostSpecBuilder; +// import software.amazon.jdbc.PluginService; +// import software.amazon.jdbc.dialect.Dialect; +// import software.amazon.jdbc.hostavailability.HostAvailability; +// import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; +// +// class ClusterAwareWriterFailoverHandlerTest { +// +// @Mock PluginService mockPluginService; +// @Mock Connection mockConnection; +// @Mock ReaderFailoverHandler mockReaderFailover; +// @Mock Connection mockWriterConnection; +// @Mock Connection mockNewWriterConnection; +// @Mock Connection mockReaderAConnection; +// @Mock Connection mockReaderBConnection; +// @Mock Dialect mockDialect; +// +// private AutoCloseable closeable; +// private final Properties properties = new Properties(); +// private final HostSpec newWriterHost = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("new-writer-host").build(); +// private final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("writer-host").build(); +// private final HostSpec readerA = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader-a-host").build(); +// private final HostSpec readerB = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader-b-host").build(); +// private final List topology = Arrays.asList(writer, readerA, readerB); +// private final List newTopology = Arrays.asList(newWriterHost, readerA, readerB); +// +// @BeforeEach +// void setUp() { +// closeable = MockitoAnnotations.openMocks(this); +// writer.addAlias("writer-host"); +// newWriterHost.addAlias("new-writer-host"); +// readerA.addAlias("reader-a-host"); +// readerB.addAlias("reader-b-host"); +// } +// +// @AfterEach +// void tearDown() throws Exception { +// closeable.close(); +// } +// +// @Test +// public void testReconnectToWriter_taskBReaderException() throws SQLException { +// when(mockPluginService.forceConnect(refEq(writer), eq(properties))).thenReturn(mockConnection); +// when(mockPluginService.forceConnect(refEq(readerA), eq(properties))).thenThrow(SQLException.class); +// when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenThrow(SQLException.class); +// +// when(mockPluginService.getAllHosts()).thenReturn(topology); +// +// when(mockReaderFailover.getReaderConnection(ArgumentMatchers.anyList())).thenThrow(SQLException.class); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = +// new ClusterAwareWriterFailoverHandler( +// mockPluginService, +// mockReaderFailover, +// properties, +// 5000, +// 2000, +// 2000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertTrue(result.isConnected()); +// assertFalse(result.isNewHost()); +// assertSame(result.getNewConnection(), mockConnection); +// +// final InOrder inOrder = Mockito.inOrder(mockPluginService); +// inOrder.verify(mockPluginService).setAvailability(eq(writer.asAliases()), eq(HostAvailability.AVAILABLE)); +// } +// +// /** +// * Verify that writer failover handler can re-connect to a current writer node. +// * +// *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. +// * TaskA: successfully re-connect to initial writer; return new connection. +// * TaskB: successfully connect to readerA and then new writer, but it takes more time than taskA. +// * Expected test result: new connection by taskA. +// */ +// @Test +// public void testReconnectToWriter_SlowReaderA() throws SQLException { +// when(mockPluginService.forceConnect(refEq(writer), eq(properties))).thenReturn(mockWriterConnection); +// when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenThrow(SQLException.class); +// when(mockPluginService.forceConnect(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); +// when(mockPluginService.getAllHosts()).thenReturn(topology).thenReturn(newTopology); +// +// when(mockReaderFailover.getReaderConnection(ArgumentMatchers.anyList())) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(5000); +// return new ReaderFailoverResult(mockReaderAConnection, readerA, true); +// }); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = +// new ClusterAwareWriterFailoverHandler( +// mockPluginService, +// mockReaderFailover, +// properties, +// 60000, +// 5000, +// 5000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertTrue(result.isConnected()); +// assertFalse(result.isNewHost()); +// assertSame(result.getNewConnection(), mockWriterConnection); +// +// final InOrder inOrder = Mockito.inOrder(mockPluginService); +// inOrder.verify(mockPluginService).setAvailability(eq(writer.asAliases()), eq(HostAvailability.AVAILABLE)); +// } +// +// /** +// * Verify that writer failover handler can re-connect to a current writer node. +// * +// *

Topology: no changes. +// * TaskA: successfully re-connect to writer; return new connection. +// * TaskB: successfully connect to readerA and retrieve topology, but latest writer is not new (defer to taskA). +// * Expected test result: new connection by taskA. +// */ +// @Test +// public void testReconnectToWriter_taskBDefers() throws SQLException { +// when(mockPluginService.forceConnect(refEq(writer), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(5000); +// return mockWriterConnection; +// }); +// when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenThrow(SQLException.class); +// +// when(mockPluginService.getAllHosts()).thenReturn(topology); +// +// when(mockReaderFailover.getReaderConnection(ArgumentMatchers.anyList())) +// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = +// new ClusterAwareWriterFailoverHandler( +// mockPluginService, +// mockReaderFailover, +// properties, +// 60000, +// 2000, +// 2000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertTrue(result.isConnected()); +// assertFalse(result.isNewHost()); +// assertSame(result.getNewConnection(), mockWriterConnection); +// +// final InOrder inOrder = Mockito.inOrder(mockPluginService); +// inOrder.verify(mockPluginService).setAvailability(eq(writer.asAliases()), eq(HostAvailability.AVAILABLE)); +// } +// +// /** +// * Verify that writer failover handler can re-connect to a new writer node. +// * +// *

Topology: changes to [new-writer, reader-A, reader-B] for taskB, taskA sees no changes. +// * taskA: successfully re-connect to writer; return connection to initial writer, but it takes more +// * time than taskB. +// * TaskB: successfully connect to readerA and then to new-writer. +// * Expected test result: new connection to writer by taskB. +// */ +// @Test +// public void testConnectToReaderA_SlowWriter() throws SQLException { +// when(mockPluginService.forceConnect(refEq(writer), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(5000); +// return mockWriterConnection; +// }); +// when(mockPluginService.forceConnect(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); +// when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); +// when(mockPluginService.forceConnect(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); +// +// when(mockPluginService.getAllHosts()).thenReturn(newTopology); +// +// when(mockReaderFailover.getReaderConnection(ArgumentMatchers.anyList())) +// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = +// new ClusterAwareWriterFailoverHandler( +// mockPluginService, +// mockReaderFailover, +// properties, +// 60000, +// 5000, +// 5000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertTrue(result.isConnected()); +// assertTrue(result.isNewHost()); +// assertSame(result.getNewConnection(), mockNewWriterConnection); +// assertEquals(3, result.getTopology().size()); +// assertEquals("new-writer-host", result.getTopology().get(0).getHost()); +// +// verify(mockPluginService, times(1)).setAvailability(eq(newWriterHost.asAliases()), eq(HostAvailability.AVAILABLE)); +// } +// +// /** +// * Verify that writer failover handler can re-connect to a new writer node. +// * +// *

Topology: changes to [new-writer, initial-writer, reader-A, reader-B]. +// * TaskA: successfully reconnect, but initial-writer is now a reader (defer to taskB). +// * TaskB: successfully connect to readerA and then to new-writer. +// * Expected test result: new connection to writer by taskB. +// */ +// @Test +// public void testConnectToReaderA_taskADefers() throws SQLException { +// when(mockPluginService.forceConnect(writer, properties)).thenReturn(mockConnection); +// when(mockPluginService.forceConnect(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); +// when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); +// when(mockPluginService.forceConnect(refEq(newWriterHost), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(5000); +// return mockNewWriterConnection; +// }); +// +// final List newTopology = Arrays.asList(newWriterHost, writer, readerA, readerB); +// when(mockPluginService.getAllHosts()).thenReturn(newTopology); +// +// when(mockReaderFailover.getReaderConnection(ArgumentMatchers.anyList())) +// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = +// new ClusterAwareWriterFailoverHandler( +// mockPluginService, +// mockReaderFailover, +// properties, +// 60000, +// 5000, +// 5000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertTrue(result.isConnected()); +// assertTrue(result.isNewHost()); +// assertSame(result.getNewConnection(), mockNewWriterConnection); +// assertEquals(4, result.getTopology().size()); +// assertEquals("new-writer-host", result.getTopology().get(0).getHost()); +// +// verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); +// verify(mockPluginService, times(1)).setAvailability(eq(newWriterHost.asAliases()), eq(HostAvailability.AVAILABLE)); +// } +// +// /** +// * Verify that writer failover handler fails to re-connect to any writer node. +// * +// *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. +// * TaskA: fail to re-connect to writer due to failover timeout. +// * TaskB: successfully connect to readerA and then fail to connect to writer due to failover timeout. +// * Expected test result: no connection. +// */ +// @Test +// public void testFailedToConnect_failoverTimeout() throws SQLException { +// when(mockPluginService.forceConnect(refEq(writer), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(30000); +// return mockWriterConnection; +// }); +// when(mockPluginService.forceConnect(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); +// when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); +// when(mockPluginService.forceConnect(refEq(newWriterHost), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(30000); +// return mockNewWriterConnection; +// }); +// when(mockPluginService.getAllHosts()).thenReturn(newTopology); +// +// when(mockReaderFailover.getReaderConnection(ArgumentMatchers.anyList())) +// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = +// new ClusterAwareWriterFailoverHandler( +// mockPluginService, +// mockReaderFailover, +// properties, +// 5000, +// 2000, +// 2000); +// +// final long startTimeNano = System.nanoTime(); +// final WriterFailoverResult result = target.failover(topology); +// final long durationNano = System.nanoTime() - startTimeNano; +// +// assertFalse(result.isConnected()); +// assertFalse(result.isNewHost()); +// +// verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); +// +// // 5s is a max allowed failover timeout; add 1s for inaccurate measurements +// assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); +// } +// +// /** +// * Verify that writer failover handler fails to re-connect to any writer node. +// * +// *

Topology: changes to [new-writer, reader-A, reader-B] for taskB. +// * TaskA: fail to re-connect to writer due to exception. +// * TaskB: successfully connect to readerA and then fail to connect to writer due to exception. +// * Expected test result: no connection. +// */ +// @Test +// public void testFailedToConnect_taskAException_taskBWriterException() throws SQLException { +// final SQLException exception = new SQLException("exception", "08S01", null); +// when(mockPluginService.forceConnect(refEq(writer), eq(properties))).thenThrow(exception); +// when(mockPluginService.forceConnect(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); +// when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); +// when(mockPluginService.forceConnect(refEq(newWriterHost), eq(properties))).thenThrow(exception); +// when(mockPluginService.isNetworkException(eq(exception), any())).thenReturn(true); +// +// when(mockPluginService.getAllHosts()).thenReturn(newTopology); +// +// when(mockReaderFailover.getReaderConnection(ArgumentMatchers.anyList())) +// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = +// new ClusterAwareWriterFailoverHandler( +// mockPluginService, +// mockReaderFailover, +// properties, +// 5000, +// 2000, +// 2000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertFalse(result.isConnected()); +// assertFalse(result.isNewHost()); +// +// verify(mockPluginService, atLeastOnce()) +// .setAvailability(eq(newWriterHost.asAliases()), eq(HostAvailability.NOT_AVAILABLE)); +// } +// } diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java index 97ab5f119..33160333f 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java @@ -1,442 +1,442 @@ -/* - * 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.failover; - -import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyString; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.atLeastOnce; -import static org.mockito.Mockito.doNothing; -import static org.mockito.Mockito.doThrow; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.spy; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -import java.sql.Connection; -import java.sql.ResultSet; -import java.sql.SQLException; -import java.util.Arrays; -import java.util.Collections; -import java.util.EnumSet; -import java.util.HashMap; -import java.util.HashSet; -import java.util.List; -import java.util.Map; -import java.util.Properties; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.ValueSource; -import org.mockito.Mock; -import org.mockito.MockitoAnnotations; -import software.amazon.jdbc.HostListProviderService; -import software.amazon.jdbc.HostRole; -import software.amazon.jdbc.HostSpec; -import software.amazon.jdbc.HostSpecBuilder; -import software.amazon.jdbc.JdbcCallable; -import software.amazon.jdbc.NodeChangeOptions; -import software.amazon.jdbc.PluginService; -import software.amazon.jdbc.hostavailability.HostAvailability; -import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; -import software.amazon.jdbc.hostlistprovider.AuroraHostListProvider; -import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; -import software.amazon.jdbc.util.RdsUrlType; -import software.amazon.jdbc.util.SqlState; -import software.amazon.jdbc.util.telemetry.GaugeCallable; -import software.amazon.jdbc.util.telemetry.TelemetryContext; -import software.amazon.jdbc.util.telemetry.TelemetryCounter; -import software.amazon.jdbc.util.telemetry.TelemetryFactory; -import software.amazon.jdbc.util.telemetry.TelemetryGauge; - -class FailoverConnectionPluginTest { - - private static final Class MONITOR_METHOD_INVOKE_ON = Connection.class; - private static final String MONITOR_METHOD_NAME = "Connection.executeQuery"; - private static final Object[] EMPTY_ARGS = {}; - private final List defaultHosts = Arrays.asList( - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("writer").port(1234).role(HostRole.WRITER).build(), - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader1").port(1234).role(HostRole.READER).build()); - - @Mock PluginService mockPluginService; - @Mock Connection mockConnection; - @Mock HostSpec mockHostSpec; - @Mock HostListProviderService mockHostListProviderService; - @Mock AuroraHostListProvider mockHostListProvider; - @Mock JdbcCallable mockInitHostProviderFunc; - @Mock ReaderFailoverHandler mockReaderFailoverHandler; - @Mock WriterFailoverHandler mockWriterFailoverHandler; - @Mock ReaderFailoverResult mockReaderResult; - @Mock WriterFailoverResult mockWriterResult; - @Mock JdbcCallable mockSqlFunction; - @Mock private TelemetryFactory mockTelemetryFactory; - @Mock TelemetryContext mockTelemetryContext; - @Mock TelemetryCounter mockTelemetryCounter; - @Mock TelemetryGauge mockTelemetryGauge; - @Mock TargetDriverDialect mockTargetDriverDialect; - - - private final Properties properties = new Properties(); - private FailoverConnectionPlugin plugin; - private AutoCloseable closeable; - - @AfterEach - void cleanUp() throws Exception { - closeable.close(); - } - - @BeforeEach - void init() throws SQLException { - closeable = MockitoAnnotations.openMocks(this); - - when(mockPluginService.getHostListProvider()).thenReturn(mockHostListProvider); - when(mockHostListProvider.getRdsUrlType()).thenReturn(RdsUrlType.RDS_WRITER_CLUSTER); - when(mockPluginService.getCurrentConnection()).thenReturn(mockConnection); - when(mockPluginService.getCurrentHostSpec()).thenReturn(mockHostSpec); - when(mockPluginService.connect(any(HostSpec.class), eq(properties))).thenReturn(mockConnection); - when(mockPluginService.getTelemetryFactory()).thenReturn(mockTelemetryFactory); - when(mockPluginService.getHosts()).thenReturn(defaultHosts); - when(mockPluginService.getAllHosts()).thenReturn(defaultHosts); - when(mockReaderFailoverHandler.failover(any(), any())).thenReturn(mockReaderResult); - when(mockWriterFailoverHandler.failover(any())).thenReturn(mockWriterResult); - when(mockWriterResult.isConnected()).thenReturn(true); - when(mockWriterResult.getTopology()).thenReturn(defaultHosts); - when(mockReaderResult.isConnected()).thenReturn(true); - - when(mockPluginService.getTelemetryFactory()).thenReturn(mockTelemetryFactory); - when(mockTelemetryFactory.openTelemetryContext(anyString(), any())).thenReturn(mockTelemetryContext); - when(mockTelemetryFactory.openTelemetryContext(eq(null), any())).thenReturn(mockTelemetryContext); - when(mockTelemetryFactory.createCounter(anyString())).thenReturn(mockTelemetryCounter); - // noinspection unchecked - when(mockTelemetryFactory.createGauge(anyString(), any(GaugeCallable.class))).thenReturn(mockTelemetryGauge); - - when(mockPluginService.getTargetDriverDialect()).thenReturn(mockTargetDriverDialect); - when(mockTargetDriverDialect.getNetworkBoundMethodNames(any())).thenReturn(new HashSet<>()); - - properties.clear(); - } - - @Test - void test_notifyNodeListChanged_withFailoverDisabled() { - properties.setProperty(FailoverConnectionPlugin.ENABLE_CLUSTER_AWARE_FAILOVER.name, "false"); - final Map> changes = new HashMap<>(); - - initializePlugin(); - plugin.notifyNodeListChanged(changes); - - verify(mockPluginService, never()).getCurrentHostSpec(); - verify(mockHostSpec, never()).getAliases(); - } - - @Test - void test_notifyNodeListChanged_withValidConnectionNotInTopology() { - final Map> changes = new HashMap<>(); - changes.put("cluster-host/", EnumSet.of(NodeChangeOptions.NODE_DELETED)); - changes.put("instance/", EnumSet.of(NodeChangeOptions.NODE_ADDED)); - - initializePlugin(); - plugin.notifyNodeListChanged(changes); - - when(mockHostSpec.getUrl()).thenReturn("cluster-url/"); - when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Collections.singletonList("instance"))); - - verify(mockPluginService).getCurrentHostSpec(); - verify(mockHostSpec, never()).getAliases(); - } - - @Test - void test_updateTopology() throws SQLException { - initializePlugin(); - - // Test updateTopology with failover disabled - plugin.setRdsUrlType(RdsUrlType.RDS_PROXY); - plugin.updateTopology(false); - verify(mockPluginService, never()).forceRefreshHostList(); - verify(mockPluginService, never()).refreshHostList(); - - // Test updateTopology with no connection - when(mockPluginService.getCurrentHostSpec()).thenReturn(null); - plugin.updateTopology(false); - verify(mockPluginService, never()).forceRefreshHostList(); - verify(mockPluginService, never()).refreshHostList(); - - // Test updateTopology with closed connection - when(mockConnection.isClosed()).thenReturn(true); - plugin.updateTopology(false); - verify(mockPluginService, never()).forceRefreshHostList(); - verify(mockPluginService, never()).refreshHostList(); - } - - @ParameterizedTest - @ValueSource(booleans = {true, false}) - void test_updateTopology_withForceUpdate(final boolean forceUpdate) throws SQLException { - - when(mockPluginService.getAllHosts()).thenReturn(Collections.singletonList( - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("host").build())); - when(mockPluginService.getHosts()).thenReturn(Collections.singletonList( - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("host").build())); - when(mockConnection.isClosed()).thenReturn(false); - initializePlugin(); - plugin.setRdsUrlType(RdsUrlType.RDS_INSTANCE); - - plugin.updateTopology(forceUpdate); - if (forceUpdate) { - verify(mockPluginService, atLeastOnce()).forceRefreshHostList(); - } else { - verify(mockPluginService, atLeastOnce()).refreshHostList(); - } - } - - @Test - void test_failover_failoverWriter() throws SQLException { - when(mockPluginService.isInTransaction()).thenReturn(true); - - initializePlugin(); - final FailoverConnectionPlugin spyPlugin = spy(plugin); - doThrow(FailoverSuccessSQLException.class).when(spyPlugin).failoverWriter(); - spyPlugin.failoverMode = FailoverMode.STRICT_WRITER; - - assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failover(mockHostSpec)); - verify(spyPlugin).failoverWriter(); - } - - @Test - void test_failover_failoverReader() throws SQLException { - when(mockPluginService.isInTransaction()).thenReturn(false); - - initializePlugin(); - final FailoverConnectionPlugin spyPlugin = spy(plugin); - doThrow(FailoverSuccessSQLException.class).when(spyPlugin).failoverReader(eq(mockHostSpec)); - spyPlugin.failoverMode = FailoverMode.READER_OR_WRITER; - - assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failover(mockHostSpec)); - verify(spyPlugin).failoverReader(eq(mockHostSpec)); - } - - @Test - void test_failoverReader_withValidFailedHostSpec_successFailover() throws SQLException { - when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); - when(mockHostSpec.getRawAvailability()).thenReturn(HostAvailability.AVAILABLE); - when(mockReaderResult.isConnected()).thenReturn(true); - when(mockReaderResult.getConnection()).thenReturn(mockConnection); - when(mockReaderResult.getHost()).thenReturn(defaultHosts.get(1)); - - initializePlugin(); - plugin.initHostProvider( - mockHostListProviderService, - mockInitHostProviderFunc, - () -> mockReaderFailoverHandler, - () -> mockWriterFailoverHandler); - - final FailoverConnectionPlugin spyPlugin = spy(plugin); - doNothing().when(spyPlugin).updateTopology(true); - - assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failoverReader(mockHostSpec)); - - verify(mockReaderFailoverHandler).failover(eq(defaultHosts), eq(mockHostSpec)); - verify(mockPluginService).setCurrentConnection(eq(mockConnection), eq(defaultHosts.get(1))); - } - - @Test - void test_failoverReader_withNoFailedHostSpec_withException() throws SQLException { - final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") - .build(); - final List hosts = Collections.singletonList(hostSpec); - - when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); - when(mockHostSpec.getAvailability()).thenReturn(HostAvailability.AVAILABLE); - when(mockPluginService.getAllHosts()).thenReturn(hosts); - when(mockPluginService.getHosts()).thenReturn(hosts); - when(mockReaderResult.getException()).thenReturn(new SQLException()); - when(mockReaderResult.getHost()).thenReturn(hostSpec); - - initializePlugin(); - plugin.initHostProvider( - mockHostListProviderService, - mockInitHostProviderFunc, - () -> mockReaderFailoverHandler, - () -> mockWriterFailoverHandler); - - assertThrows(SQLException.class, () -> plugin.failoverReader(null)); - verify(mockReaderFailoverHandler).failover(eq(hosts), eq(null)); - } - - @Test - void test_failoverWriter_failedFailover_throwsException() throws SQLException { - final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") - .build(); - final List hosts = Collections.singletonList(hostSpec); - - when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); - when(mockPluginService.getAllHosts()).thenReturn(hosts); - when(mockPluginService.getHosts()).thenReturn(hosts); - when(mockWriterResult.getException()).thenReturn(new SQLException()); - - initializePlugin(); - plugin.initHostProvider( - mockHostListProviderService, - mockInitHostProviderFunc, - () -> mockReaderFailoverHandler, - () -> mockWriterFailoverHandler); - - assertThrows(SQLException.class, () -> plugin.failoverWriter()); - verify(mockWriterFailoverHandler).failover(eq(hosts)); - } - - @Test - void test_failoverWriter_failedFailover_withNoResult() throws SQLException { - final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") - .build(); - final List hosts = Collections.singletonList(hostSpec); - - when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); - when(mockPluginService.getAllHosts()).thenReturn(hosts); - when(mockPluginService.getHosts()).thenReturn(hosts); - when(mockWriterResult.isConnected()).thenReturn(false); - - initializePlugin(); - plugin.initHostProvider( - mockHostListProviderService, - mockInitHostProviderFunc, - () -> mockReaderFailoverHandler, - () -> mockWriterFailoverHandler); - - final SQLException exception = assertThrows(SQLException.class, () -> plugin.failoverWriter()); - assertEquals(SqlState.CONNECTION_UNABLE_TO_CONNECT.getState(), exception.getSQLState()); - - verify(mockWriterFailoverHandler).failover(eq(hosts)); - verify(mockWriterResult, never()).getNewConnection(); - verify(mockWriterResult, never()).getTopology(); - } - - @Test - void test_failoverWriter_successFailover() throws SQLException { - when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); - - initializePlugin(); - plugin.initHostProvider( - mockHostListProviderService, - mockInitHostProviderFunc, - () -> mockReaderFailoverHandler, - () -> mockWriterFailoverHandler); - - final SQLException exception = assertThrows(FailoverSuccessSQLException.class, () -> plugin.failoverWriter()); - assertEquals(SqlState.COMMUNICATION_LINK_CHANGED.getState(), exception.getSQLState()); - - verify(mockWriterFailoverHandler).failover(eq(defaultHosts)); - } - - @Test - void test_invalidCurrentConnection_withNoConnection() { - when(mockPluginService.getCurrentConnection()).thenReturn(null); - initializePlugin(); - plugin.invalidateCurrentConnection(); - - verify(mockPluginService, never()).getCurrentHostSpec(); - } - - @Test - void test_invalidateCurrentConnection_inTransaction() throws SQLException { - when(mockPluginService.isInTransaction()).thenReturn(true); - when(mockHostSpec.getHost()).thenReturn("host"); - when(mockHostSpec.getPort()).thenReturn(123); - when(mockHostSpec.getRole()).thenReturn(HostRole.READER); - - initializePlugin(); - plugin.invalidateCurrentConnection(); - verify(mockConnection).rollback(); - - // Assert SQL exceptions thrown during rollback do not get propagated. - doThrow(new SQLException()).when(mockConnection).rollback(); - assertDoesNotThrow(() -> plugin.invalidateCurrentConnection()); - } - - @Test - void test_invalidateCurrentConnection_notInTransaction() { - when(mockPluginService.isInTransaction()).thenReturn(false); - when(mockHostSpec.getHost()).thenReturn("host"); - when(mockHostSpec.getPort()).thenReturn(123); - when(mockHostSpec.getRole()).thenReturn(HostRole.READER); - - initializePlugin(); - plugin.invalidateCurrentConnection(); - - verify(mockPluginService).isInTransaction(); - } - - @Test - void test_invalidateCurrentConnection_withOpenConnection() throws SQLException { - when(mockPluginService.isInTransaction()).thenReturn(false); - when(mockConnection.isClosed()).thenReturn(false); - when(mockHostSpec.getHost()).thenReturn("host"); - when(mockHostSpec.getPort()).thenReturn(123); - when(mockHostSpec.getRole()).thenReturn(HostRole.READER); - - initializePlugin(); - plugin.invalidateCurrentConnection(); - - doThrow(new SQLException()).when(mockConnection).close(); - assertDoesNotThrow(() -> plugin.invalidateCurrentConnection()); - - verify(mockConnection, times(2)).isClosed(); - verify(mockConnection, times(2)).close(); - } - - @Test - void test_execute_withFailoverDisabled() throws SQLException { - properties.setProperty(FailoverConnectionPlugin.ENABLE_CLUSTER_AWARE_FAILOVER.name, "false"); - initializePlugin(); - - plugin.execute( - ResultSet.class, - SQLException.class, - MONITOR_METHOD_INVOKE_ON, - MONITOR_METHOD_NAME, - mockSqlFunction, - EMPTY_ARGS); - - verify(mockSqlFunction).call(); - verify(mockHostListProvider, never()).getRdsUrlType(); - } - - @Test - void test_execute_withDirectExecute() throws SQLException { - initializePlugin(); - plugin.execute( - ResultSet.class, - SQLException.class, - MONITOR_METHOD_INVOKE_ON, - "close", - mockSqlFunction, - EMPTY_ARGS); - verify(mockSqlFunction).call(); - verify(mockHostListProvider, never()).getRdsUrlType(); - } - - private void initializePlugin() { - plugin = new FailoverConnectionPlugin(mockPluginService, properties); - plugin.setWriterFailoverHandler(mockWriterFailoverHandler); - plugin.setReaderFailoverHandler(mockReaderFailoverHandler); - } -} +// /* +// * 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.failover; +// +// import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +// import static org.junit.jupiter.api.Assertions.assertEquals; +// import static org.junit.jupiter.api.Assertions.assertThrows; +// import static org.mockito.ArgumentMatchers.any; +// import static org.mockito.ArgumentMatchers.anyString; +// import static org.mockito.ArgumentMatchers.eq; +// import static org.mockito.Mockito.atLeastOnce; +// import static org.mockito.Mockito.doNothing; +// import static org.mockito.Mockito.doThrow; +// import static org.mockito.Mockito.never; +// import static org.mockito.Mockito.spy; +// import static org.mockito.Mockito.times; +// import static org.mockito.Mockito.verify; +// import static org.mockito.Mockito.when; +// +// import java.sql.Connection; +// import java.sql.ResultSet; +// import java.sql.SQLException; +// import java.util.Arrays; +// import java.util.Collections; +// import java.util.EnumSet; +// import java.util.HashMap; +// import java.util.HashSet; +// import java.util.List; +// import java.util.Map; +// import java.util.Properties; +// import org.junit.jupiter.api.AfterEach; +// import org.junit.jupiter.api.BeforeEach; +// import org.junit.jupiter.api.Test; +// import org.junit.jupiter.params.ParameterizedTest; +// import org.junit.jupiter.params.provider.ValueSource; +// import org.mockito.Mock; +// import org.mockito.MockitoAnnotations; +// import software.amazon.jdbc.HostListProviderService; +// import software.amazon.jdbc.HostRole; +// import software.amazon.jdbc.HostSpec; +// import software.amazon.jdbc.HostSpecBuilder; +// import software.amazon.jdbc.JdbcCallable; +// import software.amazon.jdbc.NodeChangeOptions; +// import software.amazon.jdbc.PluginService; +// import software.amazon.jdbc.hostavailability.HostAvailability; +// import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; +// import software.amazon.jdbc.hostlistprovider.AuroraHostListProvider; +// import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; +// import software.amazon.jdbc.util.RdsUrlType; +// import software.amazon.jdbc.util.SqlState; +// import software.amazon.jdbc.util.telemetry.GaugeCallable; +// import software.amazon.jdbc.util.telemetry.TelemetryContext; +// import software.amazon.jdbc.util.telemetry.TelemetryCounter; +// import software.amazon.jdbc.util.telemetry.TelemetryFactory; +// import software.amazon.jdbc.util.telemetry.TelemetryGauge; +// +// class FailoverConnectionPluginTest { +// +// private static final Class MONITOR_METHOD_INVOKE_ON = Connection.class; +// private static final String MONITOR_METHOD_NAME = "Connection.executeQuery"; +// private static final Object[] EMPTY_ARGS = {}; +// private final List defaultHosts = Arrays.asList( +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("writer").port(1234).role(HostRole.WRITER).build(), +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader1").port(1234).role(HostRole.READER).build()); +// +// @Mock PluginService mockPluginService; +// @Mock Connection mockConnection; +// @Mock HostSpec mockHostSpec; +// @Mock HostListProviderService mockHostListProviderService; +// @Mock AuroraHostListProvider mockHostListProvider; +// @Mock JdbcCallable mockInitHostProviderFunc; +// @Mock ReaderFailoverHandler mockReaderFailoverHandler; +// @Mock WriterFailoverHandler mockWriterFailoverHandler; +// @Mock ReaderFailoverResult mockReaderResult; +// @Mock WriterFailoverResult mockWriterResult; +// @Mock JdbcCallable mockSqlFunction; +// @Mock private TelemetryFactory mockTelemetryFactory; +// @Mock TelemetryContext mockTelemetryContext; +// @Mock TelemetryCounter mockTelemetryCounter; +// @Mock TelemetryGauge mockTelemetryGauge; +// @Mock TargetDriverDialect mockTargetDriverDialect; +// +// +// private final Properties properties = new Properties(); +// private FailoverConnectionPlugin plugin; +// private AutoCloseable closeable; +// +// @AfterEach +// void cleanUp() throws Exception { +// closeable.close(); +// } +// +// @BeforeEach +// void init() throws SQLException { +// closeable = MockitoAnnotations.openMocks(this); +// +// when(mockPluginService.getHostListProvider()).thenReturn(mockHostListProvider); +// when(mockHostListProvider.getRdsUrlType()).thenReturn(RdsUrlType.RDS_WRITER_CLUSTER); +// when(mockPluginService.getCurrentConnection()).thenReturn(mockConnection); +// when(mockPluginService.getCurrentHostSpec()).thenReturn(mockHostSpec); +// when(mockPluginService.connect(any(HostSpec.class), eq(properties))).thenReturn(mockConnection); +// when(mockPluginService.getTelemetryFactory()).thenReturn(mockTelemetryFactory); +// when(mockPluginService.getHosts()).thenReturn(defaultHosts); +// when(mockPluginService.getAllHosts()).thenReturn(defaultHosts); +// when(mockReaderFailoverHandler.failover(any(), any())).thenReturn(mockReaderResult); +// when(mockWriterFailoverHandler.failover(any())).thenReturn(mockWriterResult); +// when(mockWriterResult.isConnected()).thenReturn(true); +// when(mockWriterResult.getTopology()).thenReturn(defaultHosts); +// when(mockReaderResult.isConnected()).thenReturn(true); +// +// when(mockPluginService.getTelemetryFactory()).thenReturn(mockTelemetryFactory); +// when(mockTelemetryFactory.openTelemetryContext(anyString(), any())).thenReturn(mockTelemetryContext); +// when(mockTelemetryFactory.openTelemetryContext(eq(null), any())).thenReturn(mockTelemetryContext); +// when(mockTelemetryFactory.createCounter(anyString())).thenReturn(mockTelemetryCounter); +// // noinspection unchecked +// when(mockTelemetryFactory.createGauge(anyString(), any(GaugeCallable.class))).thenReturn(mockTelemetryGauge); +// +// when(mockPluginService.getTargetDriverDialect()).thenReturn(mockTargetDriverDialect); +// when(mockTargetDriverDialect.getNetworkBoundMethodNames(any())).thenReturn(new HashSet<>()); +// +// properties.clear(); +// } +// +// @Test +// void test_notifyNodeListChanged_withFailoverDisabled() { +// properties.setProperty(FailoverConnectionPlugin.ENABLE_CLUSTER_AWARE_FAILOVER.name, "false"); +// final Map> changes = new HashMap<>(); +// +// initializePlugin(); +// plugin.notifyNodeListChanged(changes); +// +// verify(mockPluginService, never()).getCurrentHostSpec(); +// verify(mockHostSpec, never()).getAliases(); +// } +// +// @Test +// void test_notifyNodeListChanged_withValidConnectionNotInTopology() { +// final Map> changes = new HashMap<>(); +// changes.put("cluster-host/", EnumSet.of(NodeChangeOptions.NODE_DELETED)); +// changes.put("instance/", EnumSet.of(NodeChangeOptions.NODE_ADDED)); +// +// initializePlugin(); +// plugin.notifyNodeListChanged(changes); +// +// when(mockHostSpec.getUrl()).thenReturn("cluster-url/"); +// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Collections.singletonList("instance"))); +// +// verify(mockPluginService).getCurrentHostSpec(); +// verify(mockHostSpec, never()).getAliases(); +// } +// +// @Test +// void test_updateTopology() throws SQLException { +// initializePlugin(); +// +// // Test updateTopology with failover disabled +// plugin.setRdsUrlType(RdsUrlType.RDS_PROXY); +// plugin.updateTopology(false); +// verify(mockPluginService, never()).forceRefreshHostList(); +// verify(mockPluginService, never()).refreshHostList(); +// +// // Test updateTopology with no connection +// when(mockPluginService.getCurrentHostSpec()).thenReturn(null); +// plugin.updateTopology(false); +// verify(mockPluginService, never()).forceRefreshHostList(); +// verify(mockPluginService, never()).refreshHostList(); +// +// // Test updateTopology with closed connection +// when(mockConnection.isClosed()).thenReturn(true); +// plugin.updateTopology(false); +// verify(mockPluginService, never()).forceRefreshHostList(); +// verify(mockPluginService, never()).refreshHostList(); +// } +// +// @ParameterizedTest +// @ValueSource(booleans = {true, false}) +// void test_updateTopology_withForceUpdate(final boolean forceUpdate) throws SQLException { +// +// when(mockPluginService.getAllHosts()).thenReturn(Collections.singletonList( +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("host").build())); +// when(mockPluginService.getHosts()).thenReturn(Collections.singletonList( +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("host").build())); +// when(mockConnection.isClosed()).thenReturn(false); +// initializePlugin(); +// plugin.setRdsUrlType(RdsUrlType.RDS_INSTANCE); +// +// plugin.updateTopology(forceUpdate); +// if (forceUpdate) { +// verify(mockPluginService, atLeastOnce()).forceRefreshHostList(); +// } else { +// verify(mockPluginService, atLeastOnce()).refreshHostList(); +// } +// } +// +// @Test +// void test_failover_failoverWriter() throws SQLException { +// when(mockPluginService.isInTransaction()).thenReturn(true); +// +// initializePlugin(); +// final FailoverConnectionPlugin spyPlugin = spy(plugin); +// doThrow(FailoverSuccessSQLException.class).when(spyPlugin).failoverWriter(); +// spyPlugin.failoverMode = FailoverMode.STRICT_WRITER; +// +// assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failover(mockHostSpec)); +// verify(spyPlugin).failoverWriter(); +// } +// +// @Test +// void test_failover_failoverReader() throws SQLException { +// when(mockPluginService.isInTransaction()).thenReturn(false); +// +// initializePlugin(); +// final FailoverConnectionPlugin spyPlugin = spy(plugin); +// doThrow(FailoverSuccessSQLException.class).when(spyPlugin).failoverReader(eq(mockHostSpec)); +// spyPlugin.failoverMode = FailoverMode.READER_OR_WRITER; +// +// assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failover(mockHostSpec)); +// verify(spyPlugin).failoverReader(eq(mockHostSpec)); +// } +// +// @Test +// void test_failoverReader_withValidFailedHostSpec_successFailover() throws SQLException { +// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); +// when(mockHostSpec.getRawAvailability()).thenReturn(HostAvailability.AVAILABLE); +// when(mockReaderResult.isConnected()).thenReturn(true); +// when(mockReaderResult.getConnection()).thenReturn(mockConnection); +// when(mockReaderResult.getHost()).thenReturn(defaultHosts.get(1)); +// +// initializePlugin(); +// plugin.initHostProvider( +// mockHostListProviderService, +// mockInitHostProviderFunc, +// () -> mockReaderFailoverHandler, +// () -> mockWriterFailoverHandler); +// +// final FailoverConnectionPlugin spyPlugin = spy(plugin); +// doNothing().when(spyPlugin).updateTopology(true); +// +// assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failoverReader(mockHostSpec)); +// +// verify(mockReaderFailoverHandler).failover(eq(defaultHosts), eq(mockHostSpec)); +// verify(mockPluginService).setCurrentConnection(eq(mockConnection), eq(defaultHosts.get(1))); +// } +// +// @Test +// void test_failoverReader_withNoFailedHostSpec_withException() throws SQLException { +// final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") +// .build(); +// final List hosts = Collections.singletonList(hostSpec); +// +// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); +// when(mockHostSpec.getAvailability()).thenReturn(HostAvailability.AVAILABLE); +// when(mockPluginService.getAllHosts()).thenReturn(hosts); +// when(mockPluginService.getHosts()).thenReturn(hosts); +// when(mockReaderResult.getException()).thenReturn(new SQLException()); +// when(mockReaderResult.getHost()).thenReturn(hostSpec); +// +// initializePlugin(); +// plugin.initHostProvider( +// mockHostListProviderService, +// mockInitHostProviderFunc, +// () -> mockReaderFailoverHandler, +// () -> mockWriterFailoverHandler); +// +// assertThrows(SQLException.class, () -> plugin.failoverReader(null)); +// verify(mockReaderFailoverHandler).failover(eq(hosts), eq(null)); +// } +// +// @Test +// void test_failoverWriter_failedFailover_throwsException() throws SQLException { +// final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") +// .build(); +// final List hosts = Collections.singletonList(hostSpec); +// +// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); +// when(mockPluginService.getAllHosts()).thenReturn(hosts); +// when(mockPluginService.getHosts()).thenReturn(hosts); +// when(mockWriterResult.getException()).thenReturn(new SQLException()); +// +// initializePlugin(); +// plugin.initHostProvider( +// mockHostListProviderService, +// mockInitHostProviderFunc, +// () -> mockReaderFailoverHandler, +// () -> mockWriterFailoverHandler); +// +// assertThrows(SQLException.class, () -> plugin.failoverWriter()); +// verify(mockWriterFailoverHandler).failover(eq(hosts)); +// } +// +// @Test +// void test_failoverWriter_failedFailover_withNoResult() throws SQLException { +// final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") +// .build(); +// final List hosts = Collections.singletonList(hostSpec); +// +// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); +// when(mockPluginService.getAllHosts()).thenReturn(hosts); +// when(mockPluginService.getHosts()).thenReturn(hosts); +// when(mockWriterResult.isConnected()).thenReturn(false); +// +// initializePlugin(); +// plugin.initHostProvider( +// mockHostListProviderService, +// mockInitHostProviderFunc, +// () -> mockReaderFailoverHandler, +// () -> mockWriterFailoverHandler); +// +// final SQLException exception = assertThrows(SQLException.class, () -> plugin.failoverWriter()); +// assertEquals(SqlState.CONNECTION_UNABLE_TO_CONNECT.getState(), exception.getSQLState()); +// +// verify(mockWriterFailoverHandler).failover(eq(hosts)); +// verify(mockWriterResult, never()).getNewConnection(); +// verify(mockWriterResult, never()).getTopology(); +// } +// +// @Test +// void test_failoverWriter_successFailover() throws SQLException { +// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); +// +// initializePlugin(); +// plugin.initHostProvider( +// mockHostListProviderService, +// mockInitHostProviderFunc, +// () -> mockReaderFailoverHandler, +// () -> mockWriterFailoverHandler); +// +// final SQLException exception = assertThrows(FailoverSuccessSQLException.class, () -> plugin.failoverWriter()); +// assertEquals(SqlState.COMMUNICATION_LINK_CHANGED.getState(), exception.getSQLState()); +// +// verify(mockWriterFailoverHandler).failover(eq(defaultHosts)); +// } +// +// @Test +// void test_invalidCurrentConnection_withNoConnection() { +// when(mockPluginService.getCurrentConnection()).thenReturn(null); +// initializePlugin(); +// plugin.invalidateCurrentConnection(); +// +// verify(mockPluginService, never()).getCurrentHostSpec(); +// } +// +// @Test +// void test_invalidateCurrentConnection_inTransaction() throws SQLException { +// when(mockPluginService.isInTransaction()).thenReturn(true); +// when(mockHostSpec.getHost()).thenReturn("host"); +// when(mockHostSpec.getPort()).thenReturn(123); +// when(mockHostSpec.getRole()).thenReturn(HostRole.READER); +// +// initializePlugin(); +// plugin.invalidateCurrentConnection(); +// verify(mockConnection).rollback(); +// +// // Assert SQL exceptions thrown during rollback do not get propagated. +// doThrow(new SQLException()).when(mockConnection).rollback(); +// assertDoesNotThrow(() -> plugin.invalidateCurrentConnection()); +// } +// +// @Test +// void test_invalidateCurrentConnection_notInTransaction() { +// when(mockPluginService.isInTransaction()).thenReturn(false); +// when(mockHostSpec.getHost()).thenReturn("host"); +// when(mockHostSpec.getPort()).thenReturn(123); +// when(mockHostSpec.getRole()).thenReturn(HostRole.READER); +// +// initializePlugin(); +// plugin.invalidateCurrentConnection(); +// +// verify(mockPluginService).isInTransaction(); +// } +// +// @Test +// void test_invalidateCurrentConnection_withOpenConnection() throws SQLException { +// when(mockPluginService.isInTransaction()).thenReturn(false); +// when(mockConnection.isClosed()).thenReturn(false); +// when(mockHostSpec.getHost()).thenReturn("host"); +// when(mockHostSpec.getPort()).thenReturn(123); +// when(mockHostSpec.getRole()).thenReturn(HostRole.READER); +// +// initializePlugin(); +// plugin.invalidateCurrentConnection(); +// +// doThrow(new SQLException()).when(mockConnection).close(); +// assertDoesNotThrow(() -> plugin.invalidateCurrentConnection()); +// +// verify(mockConnection, times(2)).isClosed(); +// verify(mockConnection, times(2)).close(); +// } +// +// @Test +// void test_execute_withFailoverDisabled() throws SQLException { +// properties.setProperty(FailoverConnectionPlugin.ENABLE_CLUSTER_AWARE_FAILOVER.name, "false"); +// initializePlugin(); +// +// plugin.execute( +// ResultSet.class, +// SQLException.class, +// MONITOR_METHOD_INVOKE_ON, +// MONITOR_METHOD_NAME, +// mockSqlFunction, +// EMPTY_ARGS); +// +// verify(mockSqlFunction).call(); +// verify(mockHostListProvider, never()).getRdsUrlType(); +// } +// +// @Test +// void test_execute_withDirectExecute() throws SQLException { +// initializePlugin(); +// plugin.execute( +// ResultSet.class, +// SQLException.class, +// MONITOR_METHOD_INVOKE_ON, +// "close", +// mockSqlFunction, +// EMPTY_ARGS); +// verify(mockSqlFunction).call(); +// verify(mockHostListProvider, never()).getRdsUrlType(); +// } +// +// private void initializePlugin() { +// plugin = new FailoverConnectionPlugin(mockPluginService, properties); +// plugin.setWriterFailoverHandler(mockWriterFailoverHandler); +// plugin.setReaderFailoverHandler(mockReaderFailoverHandler); +// } +// } From 29a58cbffea247fc0dd782378b8652900dab1e6f Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Fri, 15 Aug 2025 14:13:12 -0700 Subject: [PATCH 02/42] Replace PluginService#forceConnect with ConnectionService#open in ClusterAwareWriterFailoverHandler --- .../ClusterAwareWriterFailoverHandler.java | 97 ++++++++++++------- 1 file changed, 63 insertions(+), 34 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java index 07d04d4f9..3972ada6a 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java @@ -141,7 +141,7 @@ public WriterFailoverResult failover(final List currentTopology) } } - private HostSpec getWriter(final List topology) { + private static HostSpec getWriter(final List topology) { if (topology == null || topology.isEmpty()) { return null; } @@ -158,13 +158,22 @@ private void submitTasks( final List currentTopology, final ExecutorService executorService, final CompletionService completionService, final boolean singleTask) { - final HostSpec writerHost = this.getWriter(currentTopology); + final HostSpec writerHost = getWriter(currentTopology); if (!singleTask) { - completionService.submit(new ReconnectToWriterHandler(writerHost)); + completionService.submit( + new ReconnectToWriterHandler( + this.connectionService, writerHost, this.initialConnectionProps, this.reconnectWriterIntervalMs)); } - completionService.submit(new WaitForNewWriterHandler( - currentTopology, - writerHost)); + + completionService.submit( + new WaitForNewWriterHandler( + this.connectionService, + this.readerFailoverHandler, + writerHost, + this.initialConnectionProps, + this.readTopologyIntervalMs, + currentTopology)); + executorService.shutdown(); } @@ -240,19 +249,30 @@ private SQLException createInterruptedException(final InterruptedException e) { /** * Internal class responsible for re-connecting to the current writer (aka TaskA). */ - private class ReconnectToWriterHandler implements Callable { + private static class ReconnectToWriterHandler implements Callable { + private final ConnectionService connectionService; private final HostSpec originalWriterHost; - - public ReconnectToWriterHandler(final HostSpec originalWriterHost) { + private final Properties props; + private final int reconnectWriterIntervalMs; + private PluginService pluginService = null; + + public ReconnectToWriterHandler( + final ConnectionService connectionService, + final HostSpec originalWriterHost, + final Properties props, + final int reconnectWriterIntervalMs) { + this.connectionService = connectionService; this.originalWriterHost = originalWriterHost; + this.props = props; + this.reconnectWriterIntervalMs = reconnectWriterIntervalMs; } public WriterFailoverResult call() { LOGGER.fine( () -> Messages.get( "ClusterAwareWriterFailoverHandler.taskAAttemptReconnectToWriterInstance", - new Object[] {this.originalWriterHost.getUrl(), PropertyUtils.maskProperties(initialConnectionProps)})); + new Object[] {this.originalWriterHost.getUrl(), PropertyUtils.maskProperties(this.props)})); Connection conn = null; List latestTopology = null; @@ -265,15 +285,14 @@ public WriterFailoverResult call() { conn.close(); } - // TODO: assess whether multi-threaded access to the plugin service is safe. The same plugin service is used - // by both the ConnectionWrapper and this ReconnectToWriterHandler in separate threads. - conn = connectionService.open(this.originalWriterHost, initialConnectionProps); - pluginService.forceRefreshHostList(conn); - latestTopology = pluginService.getAllHosts(); - + conn = connectionService.open(this.originalWriterHost, this.props); + this.pluginService = conn.unwrap(PluginService.class); + this.pluginService.forceRefreshHostList(conn); + latestTopology = this.pluginService.getAllHosts(); } catch (final SQLException exception) { // Propagate exceptions that are not caused by network errors. - if (!pluginService.isNetworkException(exception, pluginService.getTargetDriverDialect())) { + if (this.pluginService != null + && !pluginService.isNetworkException(exception, pluginService.getTargetDriverDialect())) { LOGGER.finer( () -> Messages.get( "ClusterAwareWriterFailoverHandler.taskAEncounteredException", @@ -323,26 +342,39 @@ private boolean isCurrentHostWriter(final List latestTopology) { * Internal class responsible for getting the latest cluster topology and connecting to a newly * elected writer (aka TaskB). */ - private class WaitForNewWriterHandler implements Callable { + private static class WaitForNewWriterHandler implements Callable { private Connection currentConnection = null; + private PluginService pluginService = null; + private final ConnectionService connectionService; + private final ReaderFailoverHandler readerFailoverHandler; private final HostSpec originalWriterHost; + private final Properties props; + private final int readTopologyIntervalMs; private List currentTopology; private HostSpec currentReaderHost; private Connection currentReaderConnection; public WaitForNewWriterHandler( - final List currentTopology, - final HostSpec currentHost) { + final ConnectionService connectionService, + final ReaderFailoverHandler readerFailoverHandler, + final HostSpec originalWriterHost, + final Properties props, + final int readTopologyIntervalMs, + final List currentTopology) { + this.connectionService = connectionService; + this.readerFailoverHandler = readerFailoverHandler; + this.originalWriterHost = originalWriterHost; + this.props = props; + this.readTopologyIntervalMs = readTopologyIntervalMs; this.currentTopology = currentTopology; - this.originalWriterHost = currentHost; } public WriterFailoverResult call() { LOGGER.finer( () -> Messages.get( "ClusterAwareWriterFailoverHandler.taskBAttemptConnectionToNewWriterInstance", - new Object[] {PropertyUtils.maskProperties(initialConnectionProps)})); + new Object[] {PropertyUtils.maskProperties(this.props)})); try { boolean success = false; @@ -381,6 +413,7 @@ private void connectToReader() throws InterruptedException { if (isValidReaderConnection(connResult)) { this.currentReaderConnection = connResult.getConnection(); this.currentReaderHost = connResult.getHost(); + this.pluginService = this.currentReaderConnection.unwrap(PluginService.class); LOGGER.fine( () -> Messages.get( "ClusterAwareWriterFailoverHandler.taskBConnectedToReader", @@ -396,10 +429,7 @@ private void connectToReader() throws InterruptedException { } private boolean isValidReaderConnection(final ReaderFailoverResult result) { - if (!result.isConnected() || result.getConnection() == null || result.getHost() == null) { - return false; - } - return true; + return result.isConnected() && result.getConnection() != null && result.getHost() != null; } /** @@ -408,14 +438,14 @@ private boolean isValidReaderConnection(final ReaderFailoverResult result) { * @return Returns true if successful. */ private boolean refreshTopologyAndConnectToNewWriter() throws InterruptedException { - boolean allowOldWriter = pluginService.getDialect() + boolean allowOldWriter = this.pluginService.getDialect() .getFailoverRestrictions() .contains(FailoverRestriction.ENABLE_WRITER_IN_TASK_B); while (true) { try { - pluginService.forceRefreshHostList(this.currentReaderConnection); - final List topology = pluginService.getAllHosts(); + this.pluginService.forceRefreshHostList(this.currentReaderConnection); + final List topology = this.pluginService.getAllHosts(); if (!topology.isEmpty()) { @@ -474,13 +504,12 @@ private boolean connectToWriter(final HostSpec writerCandidate) { new Object[] {writerCandidate.getUrl()})); try { // connect to the new writer - // TODO: assess whether multi-threaded access to the plugin service is safe. The same plugin service is used - // by both the ConnectionWrapper and this WaitForNewWriterHandler in separate threads. - this.currentConnection = pluginService.forceConnect(writerCandidate, initialConnectionProps); - pluginService.setAvailability(writerCandidate.asAliases(), HostAvailability.AVAILABLE); + this.currentConnection = this.connectionService.open(writerCandidate, this.props); + this.pluginService = this.currentConnection.unwrap(PluginService.class); + this.pluginService.setAvailability(writerCandidate.asAliases(), HostAvailability.AVAILABLE); return true; } catch (final SQLException exception) { - pluginService.setAvailability(writerCandidate.asAliases(), HostAvailability.NOT_AVAILABLE); + this.pluginService.setAvailability(writerCandidate.asAliases(), HostAvailability.NOT_AVAILABLE); return false; } } From f59930d5235c5af9e13bb26a6a8070bbe439a7b9 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Fri, 15 Aug 2025 16:05:58 -0700 Subject: [PATCH 03/42] Add notes for how to fix the current issues --- .../jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java | 2 ++ .../amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java | 2 ++ 2 files changed, 4 insertions(+) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java index 3972ada6a..1e07b80b8 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java @@ -444,6 +444,7 @@ private boolean refreshTopologyAndConnectToNewWriter() throws InterruptedExcepti while (true) { try { + // TODO: replace with host list provider this.pluginService.forceRefreshHostList(this.currentReaderConnection); final List topology = this.pluginService.getAllHosts(); @@ -506,6 +507,7 @@ private boolean connectToWriter(final HostSpec writerCandidate) { // connect to the new writer this.currentConnection = this.connectionService.open(writerCandidate, this.props); this.pluginService = this.currentConnection.unwrap(PluginService.class); + // TODO: replace with a map that is shared between the two handlers this.pluginService.setAvailability(writerCandidate.asAliases(), HostAvailability.AVAILABLE); return true; } catch (final SQLException exception) { diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java index 1807a2f8c..0e22b20de 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java @@ -639,6 +639,8 @@ protected void dealWithIllegalStateException( protected void failover(final HostSpec failedHost) throws SQLException { this.pluginService.setAvailability(failedHost.asAliases(), HostAvailability.NOT_AVAILABLE); + // TODO: instantiate ConnectionService here + // After failover, retrieve the unavailable hosts from the handlers after replacing pluginService#setAvailability if (this.failoverMode == FailoverMode.STRICT_WRITER) { failoverWriter(); } else { From 8fc6d1a4fe83173c4c43555354cb1ef6358e576b Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Mon, 18 Aug 2025 16:56:02 -0700 Subject: [PATCH 04/42] Integ test passing --- .../ClusterAwareWriterFailoverHandler.java | 61 +++++++++++++------ .../failover/FailoverConnectionPlugin.java | 47 +++++++------- .../failover/WriterFailoverHandler.java | 4 +- 3 files changed, 68 insertions(+), 44 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java index 1e07b80b8..72d2a5ad3 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java @@ -32,7 +32,9 @@ import java.util.logging.Logger; import software.amazon.jdbc.HostRole; import software.amazon.jdbc.HostSpec; +import software.amazon.jdbc.PartialPluginService; import software.amazon.jdbc.PluginService; +import software.amazon.jdbc.dialect.HostListProviderSupplier; import software.amazon.jdbc.hostavailability.HostAvailability; import software.amazon.jdbc.util.ExecutorFactory; import software.amazon.jdbc.util.FullServicesContainer; @@ -57,26 +59,24 @@ public class ClusterAwareWriterFailoverHandler implements WriterFailoverHandler protected int readTopologyIntervalMs = 5000; // 5 sec protected int reconnectWriterIntervalMs = 5000; // 5 sec protected Properties initialConnectionProps; + protected FullServicesContainer servicesContainer; protected PluginService pluginService; - protected ConnectionService connectionService; protected ReaderFailoverHandler readerFailoverHandler; private static final WriterFailoverResult DEFAULT_RESULT = new WriterFailoverResult(false, false, null, null, "None"); public ClusterAwareWriterFailoverHandler( final FullServicesContainer servicesContainer, - final ConnectionService connectionService, final ReaderFailoverHandler readerFailoverHandler, final Properties initialConnectionProps) { + this.servicesContainer = servicesContainer; this.pluginService = servicesContainer.getPluginService(); - this.connectionService = connectionService; this.readerFailoverHandler = readerFailoverHandler; this.initialConnectionProps = initialConnectionProps; } public ClusterAwareWriterFailoverHandler( final FullServicesContainer servicesContainer, - final ConnectionService connectionService, final ReaderFailoverHandler readerFailoverHandler, final Properties initialConnectionProps, final int failoverTimeoutMs, @@ -84,7 +84,6 @@ public ClusterAwareWriterFailoverHandler( final int reconnectWriterIntervalMs) { this( servicesContainer, - connectionService, readerFailoverHandler, initialConnectionProps); this.maxFailoverTimeoutMs = failoverTimeoutMs; @@ -99,7 +98,7 @@ public ClusterAwareWriterFailoverHandler( * @return {@link WriterFailoverResult} The results of this process. */ @Override - public WriterFailoverResult failover(final List currentTopology) + public WriterFailoverResult failover(final ConnectionService connectionService, final List currentTopology) throws SQLException { if (Utils.isNullOrEmpty(currentTopology)) { LOGGER.severe(() -> Messages.get("ClusterAwareWriterFailoverHandler.failoverCalledWithInvalidTopology")); @@ -112,7 +111,7 @@ public WriterFailoverResult failover(final List currentTopology) final ExecutorService executorService = ExecutorFactory.newFixedThreadPool(2, "failover"); final CompletionService completionService = new ExecutorCompletionService<>(executorService); - submitTasks(currentTopology, executorService, completionService, singleTask); + submitTasks(connectionService, currentTopology, executorService, completionService, singleTask); try { final long startTimeNano = System.nanoTime(); @@ -155,19 +154,26 @@ private static HostSpec getWriter(final List topology) { } private void submitTasks( - final List currentTopology, final ExecutorService executorService, + final ConnectionService connectionService, + final List currentTopology, + final ExecutorService executorService, final CompletionService completionService, final boolean singleTask) { final HostSpec writerHost = getWriter(currentTopology); if (!singleTask) { completionService.submit( new ReconnectToWriterHandler( - this.connectionService, writerHost, this.initialConnectionProps, this.reconnectWriterIntervalMs)); + connectionService, + this.getNewPluginService(), + writerHost, + this.initialConnectionProps, + this.reconnectWriterIntervalMs)); } completionService.submit( new WaitForNewWriterHandler( - this.connectionService, + connectionService, + this.getNewPluginService(), this.readerFailoverHandler, writerHost, this.initialConnectionProps, @@ -177,6 +183,24 @@ private void submitTasks( executorService.shutdown(); } + private PluginService getNewPluginService() { + PartialPluginService partialPluginService = new PartialPluginService( + this.servicesContainer, + this.initialConnectionProps, + this.pluginService.getOriginalUrl(), + this.pluginService.getDriverProtocol(), + this.pluginService.getTargetDriverDialect(), + this.pluginService.getDialect() + ); + + // TODO: can we clean this up, eg move to PartialPluginService constructor? + final HostListProviderSupplier supplier = this.pluginService.getDialect().getHostListProvider(); + partialPluginService.setHostListProvider( + supplier.getProvider(this.initialConnectionProps, this.pluginService.getOriginalUrl(), this.servicesContainer)); + + return partialPluginService; + } + private WriterFailoverResult getNextResult( final ExecutorService executorService, final CompletionService completionService, @@ -255,14 +279,16 @@ private static class ReconnectToWriterHandler implements Callable Messages.get( "ClusterAwareWriterFailoverHandler.taskAEncounteredException", @@ -344,25 +368,27 @@ private boolean isCurrentHostWriter(final List latestTopology) { */ private static class WaitForNewWriterHandler implements Callable { - private Connection currentConnection = null; - private PluginService pluginService = null; private final ConnectionService connectionService; + private final PluginService pluginService; private final ReaderFailoverHandler readerFailoverHandler; private final HostSpec originalWriterHost; private final Properties props; private final int readTopologyIntervalMs; + private Connection currentConnection = null; private List currentTopology; private HostSpec currentReaderHost; private Connection currentReaderConnection; public WaitForNewWriterHandler( final ConnectionService connectionService, + final PluginService pluginService, final ReaderFailoverHandler readerFailoverHandler, final HostSpec originalWriterHost, final Properties props, final int readTopologyIntervalMs, final List currentTopology) { this.connectionService = connectionService; + this.pluginService = pluginService; this.readerFailoverHandler = readerFailoverHandler; this.originalWriterHost = originalWriterHost; this.props = props; @@ -413,7 +439,6 @@ private void connectToReader() throws InterruptedException { if (isValidReaderConnection(connResult)) { this.currentReaderConnection = connResult.getConnection(); this.currentReaderHost = connResult.getHost(); - this.pluginService = this.currentReaderConnection.unwrap(PluginService.class); LOGGER.fine( () -> Messages.get( "ClusterAwareWriterFailoverHandler.taskBConnectedToReader", @@ -444,7 +469,6 @@ private boolean refreshTopologyAndConnectToNewWriter() throws InterruptedExcepti while (true) { try { - // TODO: replace with host list provider this.pluginService.forceRefreshHostList(this.currentReaderConnection); final List topology = this.pluginService.getAllHosts(); @@ -506,7 +530,6 @@ private boolean connectToWriter(final HostSpec writerCandidate) { try { // connect to the new writer this.currentConnection = this.connectionService.open(writerCandidate, this.props); - this.pluginService = this.currentConnection.unwrap(PluginService.class); // TODO: replace with a map that is shared between the two handlers this.pluginService.setAvailability(writerCandidate.asAliases(), HostAvailability.AVAILABLE); return true; diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java index 0e22b20de..817e5750d 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java @@ -99,7 +99,7 @@ public class FailoverConnectionPlugin extends AbstractConnectionPlugin { private final Set subscribedMethods; private final PluginService pluginService; private final FullServicesContainer servicesContainer; - private final ConnectionService connectionService; + private ConnectionService connectionService; protected final Properties properties; protected boolean enableFailoverSetting; protected boolean enableConnectFailover; @@ -231,25 +231,6 @@ public FailoverConnectionPlugin(final FullServicesContainer servicesContainer, f } this.subscribedMethods = Collections.unmodifiableSet(methods); - try { - TargetDriverHelper helper = new TargetDriverHelper(); - java.sql.Driver driver = helper.getTargetDriver(this.pluginService.getOriginalUrl(), properties); - final ConnectionProvider defaultConnectionProvider = new DriverConnectionProvider(driver); - this.connectionService = new ConnectionServiceImpl( - servicesContainer.getStorageService(), - servicesContainer.getMonitorService(), - servicesContainer.getTelemetryFactory(), - defaultConnectionProvider, - this.pluginService.getOriginalUrl(), - this.pluginService.getDriverProtocol(), - this.pluginService.getTargetDriverDialect(), - this.pluginService.getDialect(), - properties - ); - } catch (SQLException e) { - throw new RuntimeException(e); - } - TelemetryFactory telemetryFactory = this.pluginService.getTelemetryFactory(); this.failoverWriterTriggeredCounter = telemetryFactory.createCounter("writerFailover.triggered.count"); this.failoverWriterSuccessCounter = telemetryFactory.createCounter("writerFailover.completed.success.count"); @@ -345,7 +326,6 @@ public void initHostProvider( () -> new ClusterAwareWriterFailoverHandler( this.servicesContainer, - this.connectionService, this.readerFailoverHandler, this.properties, this.failoverTimeoutMsSetting, @@ -639,8 +619,26 @@ protected void dealWithIllegalStateException( protected void failover(final HostSpec failedHost) throws SQLException { this.pluginService.setAvailability(failedHost.asAliases(), HostAvailability.NOT_AVAILABLE); - // TODO: instantiate ConnectionService here - // After failover, retrieve the unavailable hosts from the handlers after replacing pluginService#setAvailability + + // TODO: After failover, retrieve the unavailable hosts from the handlers after replacing + // pluginService#setAvailability + TargetDriverHelper helper = new TargetDriverHelper(); + java.sql.Driver driver = helper.getTargetDriver(this.pluginService.getOriginalUrl(), properties); + final ConnectionProvider defaultConnectionProvider = new DriverConnectionProvider(driver); + if (this.connectionService == null) { + this.connectionService = new ConnectionServiceImpl( + servicesContainer.getStorageService(), + servicesContainer.getMonitorService(), + servicesContainer.getTelemetryFactory(), + defaultConnectionProvider, + this.pluginService.getOriginalUrl(), + this.pluginService.getDriverProtocol(), + this.pluginService.getTargetDriverDialect(), + this.pluginService.getDialect(), + properties + ); + } + if (this.failoverMode == FailoverMode.STRICT_WRITER) { failoverWriter(); } else { @@ -751,7 +749,8 @@ protected void failoverWriter() throws SQLException { try { LOGGER.info(() -> Messages.get("Failover.startWriterFailover")); - final WriterFailoverResult failoverResult = this.writerFailoverHandler.failover(this.pluginService.getAllHosts()); + final WriterFailoverResult failoverResult = + this.writerFailoverHandler.failover(this.connectionService, this.pluginService.getAllHosts()); if (failoverResult != null) { final SQLException exception = failoverResult.getException(); if (exception != null) { diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/WriterFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/WriterFailoverHandler.java index caae8de8c..68e69259b 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/WriterFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/WriterFailoverHandler.java @@ -19,6 +19,7 @@ import java.sql.SQLException; import java.util.List; import software.amazon.jdbc.HostSpec; +import software.amazon.jdbc.util.connection.ConnectionService; /** * Interface for Writer Failover Process handler. This handler implements all necessary logic to try @@ -33,5 +34,6 @@ public interface WriterFailoverHandler { * @return {@link WriterFailoverResult} The results of this process. * @throws SQLException indicating whether the failover attempt was successful. */ - WriterFailoverResult failover(List currentTopology) throws SQLException; + WriterFailoverResult failover( + ConnectionService connectionService, List currentTopology) throws SQLException; } From b4e5b3aa10d473e8a342ad63c5f8f836386210ad Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Tue, 19 Aug 2025 16:18:33 -0700 Subject: [PATCH 05/42] Pass ConnectionService to ReaderFailoverHandler --- .../plugin/failover/ClusterAwareWriterFailoverHandler.java | 3 ++- .../jdbc/plugin/failover/FailoverConnectionPlugin.java | 3 ++- .../amazon/jdbc/plugin/failover/ReaderFailoverHandler.java | 7 +++++-- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java index 72d2a5ad3..81098ba65 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java @@ -435,7 +435,8 @@ public WriterFailoverResult call() { private void connectToReader() throws InterruptedException { while (true) { try { - final ReaderFailoverResult connResult = readerFailoverHandler.getReaderConnection(this.currentTopology); + final ReaderFailoverResult connResult = + readerFailoverHandler.getReaderConnection(this.connectionService, this.currentTopology); if (isValidReaderConnection(connResult)) { this.currentReaderConnection = connResult.getConnection(); this.currentReaderHost = connResult.getHost(); diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java index 817e5750d..c684dbb50 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java @@ -665,7 +665,8 @@ protected void failoverReader(final HostSpec failedHostSpec) throws SQLException failedHost = failedHostSpec; } - final ReaderFailoverResult result = readerFailoverHandler.failover(this.pluginService.getHosts(), failedHost); + final ReaderFailoverResult result = readerFailoverHandler.failover( + this.connectionService, this.pluginService.getHosts(), failedHost); if (result != null) { final SQLException exception = result.getException(); if (exception != null) { diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ReaderFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ReaderFailoverHandler.java index e006558b6..008c1e17e 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ReaderFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ReaderFailoverHandler.java @@ -19,6 +19,7 @@ import java.sql.SQLException; import java.util.List; import software.amazon.jdbc.HostSpec; +import software.amazon.jdbc.util.connection.ConnectionService; /** * Interface for Reader Failover Process handler. This handler implements all necessary logic to try @@ -36,7 +37,8 @@ public interface ReaderFailoverHandler { * @return {@link ReaderFailoverResult} The results of this process. * @throws SQLException indicating whether the failover attempt was successful. */ - ReaderFailoverResult failover(List hosts, HostSpec currentHost) throws SQLException; + ReaderFailoverResult failover( + ConnectionService connectionService, List hosts, HostSpec currentHost) throws SQLException; /** * Called to get any available reader connection. If no reader is available then result of process @@ -46,5 +48,6 @@ public interface ReaderFailoverHandler { * @return {@link ReaderFailoverResult} The results of this process. * @throws SQLException if any error occurred while attempting a reader connection. */ - ReaderFailoverResult getReaderConnection(List hostList) throws SQLException; + ReaderFailoverResult getReaderConnection( + ConnectionService connectionService, List hostList) throws SQLException; } From a2342befd179e363d893c191a28ccd9d0ecba19d Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Tue, 19 Aug 2025 16:46:54 -0700 Subject: [PATCH 06/42] Adapted ReaderFailoverHandler, IT passing --- .../ClusterAwareReaderFailoverHandler.java | 111 ++- .../failover/FailoverConnectionPlugin.java | 2 +- ...ClusterAwareReaderFailoverHandlerTest.java | 814 +++++++++--------- 3 files changed, 485 insertions(+), 442 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java index f97a2ed4d..18df01f3c 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java @@ -33,12 +33,16 @@ import java.util.logging.Logger; import software.amazon.jdbc.HostRole; import software.amazon.jdbc.HostSpec; +import software.amazon.jdbc.PartialPluginService; import software.amazon.jdbc.PluginService; +import software.amazon.jdbc.dialect.HostListProviderSupplier; import software.amazon.jdbc.hostavailability.HostAvailability; import software.amazon.jdbc.util.ExecutorFactory; +import software.amazon.jdbc.util.FullServicesContainer; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.PropertyUtils; import software.amazon.jdbc.util.Utils; +import software.amazon.jdbc.util.connection.ConnectionService; /** * An implementation of ReaderFailoverHandler. @@ -58,24 +62,25 @@ public class ClusterAwareReaderFailoverHandler implements ReaderFailoverHandler protected static final int DEFAULT_READER_CONNECT_TIMEOUT = 30000; // 30 sec public static final ReaderFailoverResult FAILED_READER_FAILOVER_RESULT = new ReaderFailoverResult(null, null, false); - protected Properties initialConnectionProps; + protected Properties props; protected int maxFailoverTimeoutMs; protected int timeoutMs; protected boolean isStrictReaderRequired; + protected final FullServicesContainer servicesContainer; protected final PluginService pluginService; /** * ClusterAwareReaderFailoverHandler constructor. * - * @param pluginService A provider for creating new connections. - * @param initialConnectionProps The initial connection properties to copy over to the new reader. + * @param servicesContainer A provider for creating new connections. + * @param props The initial connection properties to copy over to the new reader. */ public ClusterAwareReaderFailoverHandler( - final PluginService pluginService, - final Properties initialConnectionProps) { + final FullServicesContainer servicesContainer, + final Properties props) { this( - pluginService, - initialConnectionProps, + servicesContainer, + props, DEFAULT_FAILOVER_TIMEOUT, DEFAULT_READER_CONNECT_TIMEOUT, false); @@ -84,21 +89,22 @@ public ClusterAwareReaderFailoverHandler( /** * ClusterAwareReaderFailoverHandler constructor. * - * @param pluginService A provider for creating new connections. - * @param initialConnectionProps The initial connection properties to copy over to the new reader. + * @param servicesContainer A provider for creating new connections. + * @param props The initial connection properties to copy over to the new reader. * @param maxFailoverTimeoutMs Maximum allowed time for the entire reader failover process. * @param timeoutMs Maximum allowed time in milliseconds for each reader connection attempt during * the reader failover process. * @param isStrictReaderRequired When true, it disables adding a writer to a list of nodes to connect */ public ClusterAwareReaderFailoverHandler( - final PluginService pluginService, - final Properties initialConnectionProps, + final FullServicesContainer servicesContainer, + final Properties props, final int maxFailoverTimeoutMs, final int timeoutMs, final boolean isStrictReaderRequired) { - this.pluginService = pluginService; - this.initialConnectionProps = initialConnectionProps; + this.servicesContainer = servicesContainer; + this.pluginService = servicesContainer.getPluginService(); + this.props = props; this.maxFailoverTimeoutMs = maxFailoverTimeoutMs; this.timeoutMs = timeoutMs; this.isStrictReaderRequired = isStrictReaderRequired; @@ -124,7 +130,8 @@ protected void setTimeoutMs(final int timeoutMs) { * @return {@link ReaderFailoverResult} The results of this process. */ @Override - public ReaderFailoverResult failover(final List hosts, final HostSpec currentHost) + public ReaderFailoverResult failover( + final ConnectionService connectionService, final List hosts, final HostSpec currentHost) throws SQLException { if (Utils.isNullOrEmpty(hosts)) { LOGGER.fine(() -> Messages.get("ClusterAwareReaderFailoverHandler.invalidTopology", new Object[] {"failover"})); @@ -133,11 +140,13 @@ public ReaderFailoverResult failover(final List hosts, final HostSpec final ExecutorService executor = ExecutorFactory.newSingleThreadExecutor("failover"); - final Future future = submitInternalFailoverTask(hosts, currentHost, executor); + final Future future = + submitInternalFailoverTask(connectionService, hosts, currentHost, executor); return getInternalFailoverResult(executor, future); } private Future submitInternalFailoverTask( + final ConnectionService connectionService, final List hosts, final HostSpec currentHost, final ExecutorService executor) { @@ -145,7 +154,7 @@ private Future submitInternalFailoverTask( ReaderFailoverResult result; try { while (true) { - result = failoverInternal(hosts, currentHost); + result = failoverInternal(connectionService, hosts, currentHost); if (result != null && result.isConnected()) { return result; } @@ -190,6 +199,7 @@ private ReaderFailoverResult getInternalFailoverResult( } protected ReaderFailoverResult failoverInternal( + final ConnectionService connectionService, final List hosts, final HostSpec currentHost) throws SQLException { @@ -197,7 +207,7 @@ protected ReaderFailoverResult failoverInternal( this.pluginService.setAvailability(currentHost.asAliases(), HostAvailability.NOT_AVAILABLE); } final List hostsByPriority = getHostsByPriority(hosts); - return getConnectionFromHostGroup(hostsByPriority); + return getConnectionFromHostGroup(connectionService, hostsByPriority); } public List getHostsByPriority(final List hosts) { @@ -239,7 +249,8 @@ public List getHostsByPriority(final List hosts) { * @return {@link ReaderFailoverResult} The results of this process. */ @Override - public ReaderFailoverResult getReaderConnection(final List hostList) + public ReaderFailoverResult getReaderConnection( + final ConnectionService connectionService, final List hostList) throws SQLException { if (Utils.isNullOrEmpty(hostList)) { LOGGER.fine( @@ -250,7 +261,7 @@ public ReaderFailoverResult getReaderConnection(final List hostList) } final List hostsByPriority = getReaderHostsByPriority(hostList); - return getConnectionFromHostGroup(hostsByPriority); + return getConnectionFromHostGroup(connectionService, hostsByPriority); } public List getReaderHostsByPriority(final List hosts) { @@ -291,7 +302,8 @@ public List getReaderHostsByPriority(final List hosts) { return hostsByPriority; } - private ReaderFailoverResult getConnectionFromHostGroup(final List hosts) + private ReaderFailoverResult getConnectionFromHostGroup( + final ConnectionService connectionService, final List hosts) throws SQLException { final ExecutorService executor = ExecutorFactory.newFixedThreadPool(2, "failover"); @@ -300,7 +312,8 @@ private ReaderFailoverResult getConnectionFromHostGroup(final List hos try { for (int i = 0; i < hosts.size(); i += 2) { // submit connection attempt tasks in batches of 2 - final ReaderFailoverResult result = getResultFromNextTaskBatch(hosts, executor, completionService, i); + final ReaderFailoverResult result = + getResultFromNextTaskBatch(connectionService, hosts, executor, completionService, i); if (result.isConnected() || result.getException() != null) { return result; } @@ -323,15 +336,19 @@ private ReaderFailoverResult getConnectionFromHostGroup(final List hos } private ReaderFailoverResult getResultFromNextTaskBatch( + final ConnectionService connectionService, final List hosts, final ExecutorService executor, final CompletionService completionService, final int i) throws SQLException { ReaderFailoverResult result; final int numTasks = i + 1 < hosts.size() ? 2 : 1; - completionService.submit(new ConnectionAttemptTask(hosts.get(i), this.isStrictReaderRequired)); + completionService.submit( + // TODO: are there performance concerns with creating a new plugin service this often? + new ConnectionAttemptTask( + connectionService, this.getNewPluginService(), hosts.get(i), this.props, this.isStrictReaderRequired)); if (numTasks == 2) { - completionService.submit(new ConnectionAttemptTask(hosts.get(i + 1), this.isStrictReaderRequired)); + completionService.submit(new ConnectionAttemptTask(connectionService, this.getNewPluginService(), hosts.get(i + 1), this.props, this.isStrictReaderRequired)); } for (int taskNum = 0; taskNum < numTasks; taskNum++) { result = getNextResult(completionService); @@ -372,13 +389,41 @@ private ReaderFailoverResult getNextResult(final CompletionService { + private PluginService getNewPluginService() { + PartialPluginService partialPluginService = new PartialPluginService( + this.servicesContainer, + this.props, + this.pluginService.getOriginalUrl(), + this.pluginService.getDriverProtocol(), + this.pluginService.getTargetDriverDialect(), + this.pluginService.getDialect() + ); + + // TODO: can we clean this up, eg move to PartialPluginService constructor? + final HostListProviderSupplier supplier = this.pluginService.getDialect().getHostListProvider(); + partialPluginService.setHostListProvider( + supplier.getProvider(this.props, this.pluginService.getOriginalUrl(), this.servicesContainer)); + + return partialPluginService; + } + private static class ConnectionAttemptTask implements Callable { + private final ConnectionService connectionService; + private final PluginService pluginService; private final HostSpec newHost; + private final Properties props; private final boolean isStrictReaderRequired; - private ConnectionAttemptTask(final HostSpec newHost, final boolean isStrictReaderRequired) { + private ConnectionAttemptTask( + final ConnectionService connectionService, + final PluginService pluginService, + final HostSpec newHost, + final Properties props, + final boolean isStrictReaderRequired) { + this.connectionService = connectionService; + this.pluginService = pluginService; this.newHost = newHost; + this.props = props; this.isStrictReaderRequired = isStrictReaderRequired; } @@ -390,21 +435,19 @@ public ReaderFailoverResult call() { LOGGER.fine( () -> Messages.get( "ClusterAwareReaderFailoverHandler.attemptingReaderConnection", - new Object[] {this.newHost.getUrl(), PropertyUtils.maskProperties(initialConnectionProps)})); + new Object[] {this.newHost.getUrl(), PropertyUtils.maskProperties(props)})); try { final Properties copy = new Properties(); - copy.putAll(initialConnectionProps); + copy.putAll(props); - // TODO: assess whether multi-threaded access to the plugin service is safe. The same plugin service is used by - // both the ConnectionWrapper and this ConnectionAttemptTask in separate threads. - final Connection conn = pluginService.forceConnect(this.newHost, copy); - pluginService.setAvailability(this.newHost.asAliases(), HostAvailability.AVAILABLE); + final Connection conn = this.connectionService.open(this.newHost, copy); + this.pluginService.setAvailability(this.newHost.asAliases(), HostAvailability.AVAILABLE); if (this.isStrictReaderRequired) { // need to ensure that new connection is a connection to a reader node try { - HostRole role = pluginService.getHostRole(conn); + HostRole role = this.pluginService.getHostRole(conn); if (!HostRole.READER.equals(role)) { LOGGER.fine( Messages.get( @@ -439,13 +482,13 @@ public ReaderFailoverResult call() { LOGGER.fine("New reader failover connection object: " + conn); return new ReaderFailoverResult(conn, this.newHost, true); } catch (final SQLException e) { - pluginService.setAvailability(newHost.asAliases(), HostAvailability.NOT_AVAILABLE); + this.pluginService.setAvailability(newHost.asAliases(), HostAvailability.NOT_AVAILABLE); LOGGER.fine( () -> Messages.get( "ClusterAwareReaderFailoverHandler.failedReaderConnection", new Object[] {this.newHost.getUrl()})); // Propagate exceptions that are not caused by network errors. - if (!pluginService.isNetworkException(e, pluginService.getTargetDriverDialect())) { + if (!this.pluginService.isNetworkException(e, this.pluginService.getTargetDriverDialect())) { return new ReaderFailoverResult( null, null, diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java index c684dbb50..f378d3c38 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java @@ -318,7 +318,7 @@ public void initHostProvider( initHostProviderFunc, () -> new ClusterAwareReaderFailoverHandler( - this.pluginService, + this.servicesContainer, this.properties, this.failoverTimeoutMsSetting, this.failoverReaderConnectTimeoutMsSetting, diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java index a22283c39..7eed3e5e8 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java @@ -1,407 +1,407 @@ -/* - * 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.failover; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertSame; -import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.atLeast; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; -import static software.amazon.jdbc.plugin.failover.ClusterAwareReaderFailoverHandler.DEFAULT_FAILOVER_TIMEOUT; -import static software.amazon.jdbc.plugin.failover.ClusterAwareReaderFailoverHandler.DEFAULT_READER_CONNECT_TIMEOUT; - -import java.sql.Connection; -import java.sql.SQLException; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.Collections; -import java.util.EnumSet; -import java.util.List; -import java.util.Properties; -import java.util.concurrent.TimeUnit; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.mockito.Mock; -import org.mockito.Mockito; -import org.mockito.MockitoAnnotations; -import org.mockito.stubbing.Answer; -import software.amazon.jdbc.HostRole; -import software.amazon.jdbc.HostSpec; -import software.amazon.jdbc.HostSpecBuilder; -import software.amazon.jdbc.PluginService; -import software.amazon.jdbc.dialect.Dialect; -import software.amazon.jdbc.hostavailability.HostAvailability; -import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; -import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; - -class ClusterAwareReaderFailoverHandlerTest { - - @Mock PluginService mockPluginService; - @Mock Connection mockConnection; - - private AutoCloseable closeable; - private final Properties properties = new Properties(); - private final List defaultHosts = Arrays.asList( - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("writer").port(1234).role(HostRole.WRITER).build(), - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader1").port(1234).role(HostRole.READER).build(), - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader2").port(1234).role(HostRole.READER).build(), - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader3").port(1234).role(HostRole.READER).build(), - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader4").port(1234).role(HostRole.READER).build(), - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader5").port(1234).role(HostRole.READER).build() - ); - - @BeforeEach - void setUp() { - closeable = MockitoAnnotations.openMocks(this); - } - - @AfterEach - void tearDown() throws Exception { - closeable.close(); - } - - @Test - public void testFailover() throws SQLException { - // original host list: [active writer, active reader, current connection (reader), active - // reader, down reader, active reader] - // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] - // connection attempts are made in pairs using the above list - // expected test result: successful connection for host at index 4 - final List hosts = defaultHosts; - final int currentHostIndex = 2; - final int successHostIndex = 4; - for (int i = 0; i < hosts.size(); i++) { - if (i != successHostIndex) { - final SQLException exception = new SQLException("exception", "08S01", null); - when(mockPluginService.forceConnect(hosts.get(i), properties)) - .thenThrow(exception); - when(mockPluginService.isNetworkException(exception, null)).thenReturn(true); - } else { - when(mockPluginService.forceConnect(hosts.get(i), properties)).thenReturn(mockConnection); - } - } - when(mockPluginService.getTargetDriverDialect()).thenReturn(null); - - hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); - hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); - - final ReaderFailoverHandler target = - new ClusterAwareReaderFailoverHandler( - mockPluginService, - properties); - final ReaderFailoverResult result = target.failover(hosts, hosts.get(currentHostIndex)); - - assertTrue(result.isConnected()); - assertSame(mockConnection, result.getConnection()); - assertEquals(hosts.get(successHostIndex), result.getHost()); - - final HostSpec successHost = hosts.get(successHostIndex); - verify(mockPluginService, atLeast(4)).setAvailability(any(), eq(HostAvailability.NOT_AVAILABLE)); - verify(mockPluginService, never()) - .setAvailability(eq(successHost.asAliases()), eq(HostAvailability.NOT_AVAILABLE)); - verify(mockPluginService, times(1)) - .setAvailability(eq(successHost.asAliases()), eq(HostAvailability.AVAILABLE)); - } - - @Test - public void testFailover_timeout() throws SQLException { - // original host list: [active writer, active reader, current connection (reader), active - // reader, down reader, active reader] - // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] - // connection attempts are made in pairs using the above list - // expected test result: failure to get reader since process is limited to 5s and each attempt - // to connect takes 20s - final List hosts = defaultHosts; - final int currentHostIndex = 2; - for (HostSpec host : hosts) { - when(mockPluginService.forceConnect(host, properties)) - .thenAnswer((Answer) invocation -> { - Thread.sleep(20000); - return mockConnection; - }); - } - - hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); - hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); - - final ReaderFailoverHandler target = - new ClusterAwareReaderFailoverHandler( - mockPluginService, - properties, - 5000, - 30000, - false); - - final long startTimeNano = System.nanoTime(); - final ReaderFailoverResult result = - target.failover(hosts, hosts.get(currentHostIndex)); - final long durationNano = System.nanoTime() - startTimeNano; - - assertFalse(result.isConnected()); - assertNull(result.getConnection()); - assertNull(result.getHost()); - - // 5s is a max allowed failover timeout; add 1s for inaccurate measurements - assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); - } - - @Test - public void testFailover_nullOrEmptyHostList() throws SQLException { - final ClusterAwareReaderFailoverHandler target = - new ClusterAwareReaderFailoverHandler( - mockPluginService, - properties); - final HostSpec currentHost = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("writer") - .port(1234).build(); - - ReaderFailoverResult result = target.failover(null, currentHost); - assertFalse(result.isConnected()); - assertNull(result.getConnection()); - assertNull(result.getHost()); - - final List hosts = new ArrayList<>(); - result = target.failover(hosts, currentHost); - assertFalse(result.isConnected()); - assertNull(result.getConnection()); - assertNull(result.getHost()); - } - - @Test - public void testGetReader_connectionSuccess() throws SQLException { - // even number of connection attempts - // first connection attempt to return succeeds, second attempt cancelled - // expected test result: successful connection for host at index 2 - final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) - final HostSpec slowHost = hosts.get(1); - final HostSpec fastHost = hosts.get(2); - when(mockPluginService.forceConnect(slowHost, properties)) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(20000); - return mockConnection; - }); - when(mockPluginService.forceConnect(eq(fastHost), eq(properties))).thenReturn(mockConnection); - - Dialect mockDialect = Mockito.mock(Dialect.class); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - when(mockPluginService.getDialect()).thenReturn(mockDialect); - - final ReaderFailoverHandler target = - new ClusterAwareReaderFailoverHandler( - mockPluginService, - properties); - final ReaderFailoverResult result = target.getReaderConnection(hosts); - - assertTrue(result.isConnected()); - assertSame(mockConnection, result.getConnection()); - assertEquals(hosts.get(2), result.getHost()); - - verify(mockPluginService, never()).setAvailability(any(), eq(HostAvailability.NOT_AVAILABLE)); - verify(mockPluginService, times(1)) - .setAvailability(eq(fastHost.asAliases()), eq(HostAvailability.AVAILABLE)); - } - - @Test - public void testGetReader_connectionFailure() throws SQLException { - // odd number of connection attempts - // first connection attempt to return fails - // expected test result: failure to get reader - final List hosts = defaultHosts.subList(0, 4); // 3 connection attempts (writer not attempted) - when(mockPluginService.forceConnect(any(), eq(properties))).thenThrow(new SQLException("exception", "08S01", null)); - - Dialect mockDialect = Mockito.mock(Dialect.class); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - when(mockPluginService.getDialect()).thenReturn(mockDialect); - - final int currentHostIndex = 2; - - final ReaderFailoverHandler target = - new ClusterAwareReaderFailoverHandler( - mockPluginService, - properties); - final ReaderFailoverResult result = target.getReaderConnection(hosts); - - assertFalse(result.isConnected()); - assertNull(result.getConnection()); - assertNull(result.getHost()); - } - - @Test - public void testGetReader_connectionAttemptsTimeout() throws SQLException { - // connection attempts time out before they can succeed - // first connection attempt to return times out - // expected test result: failure to get reader - final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) - when(mockPluginService.forceConnect(any(), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - try { - Thread.sleep(5000); - } catch (InterruptedException exception) { - // ignore - } - return mockConnection; - }); - - Dialect mockDialect = Mockito.mock(Dialect.class); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - when(mockPluginService.getDialect()).thenReturn(mockDialect); - - final ClusterAwareReaderFailoverHandler target = - new ClusterAwareReaderFailoverHandler( - mockPluginService, - properties, - 60000, - 1000, - false); - final ReaderFailoverResult result = target.getReaderConnection(hosts); - - assertFalse(result.isConnected()); - assertNull(result.getConnection()); - assertNull(result.getHost()); - } - - @Test - public void testGetHostTuplesByPriority() { - final List originalHosts = defaultHosts; - originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); - originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); - originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); - - final ClusterAwareReaderFailoverHandler target = - new ClusterAwareReaderFailoverHandler( - mockPluginService, - properties); - final List hostsByPriority = target.getHostsByPriority(originalHosts); - - int i = 0; - - // expecting active readers - while (i < hostsByPriority.size() - && hostsByPriority.get(i).getRole() == HostRole.READER - && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { - i++; - } - - // expecting a writer - while (i < hostsByPriority.size() - && hostsByPriority.get(i).getRole() == HostRole.WRITER) { - i++; - } - - // expecting down readers - while (i < hostsByPriority.size() - && hostsByPriority.get(i).getRole() == HostRole.READER - && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { - i++; - } - - assertEquals(hostsByPriority.size(), i); - } - - @Test - public void testGetReaderTuplesByPriority() { - final List originalHosts = defaultHosts; - originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); - originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); - originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); - - Dialect mockDialect = Mockito.mock(Dialect.class); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - when(mockPluginService.getDialect()).thenReturn(mockDialect); - - final ClusterAwareReaderFailoverHandler target = - new ClusterAwareReaderFailoverHandler( - mockPluginService, - properties); - final List hostsByPriority = target.getReaderHostsByPriority(originalHosts); - - int i = 0; - - // expecting active readers - while (i < hostsByPriority.size() - && hostsByPriority.get(i).getRole() == HostRole.READER - && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { - i++; - } - - // expecting down readers - while (i < hostsByPriority.size() - && hostsByPriority.get(i).getRole() == HostRole.READER - && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { - i++; - } - - assertEquals(hostsByPriority.size(), i); - } - - @Test - public void testHostFailoverStrictReaderEnabled() { - - final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("writer").port(1234).role(HostRole.WRITER).build(); - final HostSpec reader = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader1").port(1234).role(HostRole.READER).build(); - final List hosts = Arrays.asList(writer, reader); - - Dialect mockDialect = Mockito.mock(Dialect.class); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - when(mockPluginService.getDialect()).thenReturn(mockDialect); - final ClusterAwareReaderFailoverHandler target = - new ClusterAwareReaderFailoverHandler( - mockPluginService, - properties, - DEFAULT_FAILOVER_TIMEOUT, - DEFAULT_READER_CONNECT_TIMEOUT, - true); - - // The writer is included because the original writer has likely become a reader. - List expectedHostsByPriority = Arrays.asList(reader, writer); - - List hostsByPriority = target.getHostsByPriority(hosts); - assertEquals(expectedHostsByPriority, hostsByPriority); - - // Should pick the reader even if unavailable. The unavailable reader will be lower priority than the writer. - reader.setAvailability(HostAvailability.NOT_AVAILABLE); - expectedHostsByPriority = Arrays.asList(writer, reader); - - hostsByPriority = target.getHostsByPriority(hosts); - assertEquals(expectedHostsByPriority, hostsByPriority); - - // Writer node will only be picked if it is the only node in topology; - List expectedWriterHost = Collections.singletonList(writer); - - hostsByPriority = target.getHostsByPriority(Collections.singletonList(writer)); - assertEquals(expectedWriterHost, hostsByPriority); - } -} +// /* +// * 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.failover; +// +// import static org.junit.jupiter.api.Assertions.assertEquals; +// import static org.junit.jupiter.api.Assertions.assertFalse; +// import static org.junit.jupiter.api.Assertions.assertNull; +// import static org.junit.jupiter.api.Assertions.assertSame; +// import static org.junit.jupiter.api.Assertions.assertTrue; +// import static org.mockito.ArgumentMatchers.any; +// import static org.mockito.ArgumentMatchers.eq; +// import static org.mockito.Mockito.atLeast; +// import static org.mockito.Mockito.never; +// import static org.mockito.Mockito.times; +// import static org.mockito.Mockito.verify; +// import static org.mockito.Mockito.when; +// import static software.amazon.jdbc.plugin.failover.ClusterAwareReaderFailoverHandler.DEFAULT_FAILOVER_TIMEOUT; +// import static software.amazon.jdbc.plugin.failover.ClusterAwareReaderFailoverHandler.DEFAULT_READER_CONNECT_TIMEOUT; +// +// import java.sql.Connection; +// import java.sql.SQLException; +// import java.util.ArrayList; +// import java.util.Arrays; +// import java.util.Collections; +// import java.util.EnumSet; +// import java.util.List; +// import java.util.Properties; +// import java.util.concurrent.TimeUnit; +// import org.junit.jupiter.api.AfterEach; +// import org.junit.jupiter.api.BeforeEach; +// import org.junit.jupiter.api.Test; +// import org.mockito.Mock; +// import org.mockito.Mockito; +// import org.mockito.MockitoAnnotations; +// import org.mockito.stubbing.Answer; +// import software.amazon.jdbc.HostRole; +// import software.amazon.jdbc.HostSpec; +// import software.amazon.jdbc.HostSpecBuilder; +// import software.amazon.jdbc.PluginService; +// import software.amazon.jdbc.dialect.Dialect; +// import software.amazon.jdbc.hostavailability.HostAvailability; +// import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; +// import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; +// +// class ClusterAwareReaderFailoverHandlerTest { +// +// @Mock PluginService mockPluginService; +// @Mock Connection mockConnection; +// +// private AutoCloseable closeable; +// private final Properties properties = new Properties(); +// private final List defaultHosts = Arrays.asList( +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("writer").port(1234).role(HostRole.WRITER).build(), +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader1").port(1234).role(HostRole.READER).build(), +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader2").port(1234).role(HostRole.READER).build(), +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader3").port(1234).role(HostRole.READER).build(), +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader4").port(1234).role(HostRole.READER).build(), +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader5").port(1234).role(HostRole.READER).build() +// ); +// +// @BeforeEach +// void setUp() { +// closeable = MockitoAnnotations.openMocks(this); +// } +// +// @AfterEach +// void tearDown() throws Exception { +// closeable.close(); +// } +// +// @Test +// public void testFailover() throws SQLException { +// // original host list: [active writer, active reader, current connection (reader), active +// // reader, down reader, active reader] +// // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] +// // connection attempts are made in pairs using the above list +// // expected test result: successful connection for host at index 4 +// final List hosts = defaultHosts; +// final int currentHostIndex = 2; +// final int successHostIndex = 4; +// for (int i = 0; i < hosts.size(); i++) { +// if (i != successHostIndex) { +// final SQLException exception = new SQLException("exception", "08S01", null); +// when(mockPluginService.forceConnect(hosts.get(i), properties)) +// .thenThrow(exception); +// when(mockPluginService.isNetworkException(exception, null)).thenReturn(true); +// } else { +// when(mockPluginService.forceConnect(hosts.get(i), properties)).thenReturn(mockConnection); +// } +// } +// when(mockPluginService.getTargetDriverDialect()).thenReturn(null); +// +// hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); +// hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); +// +// final ReaderFailoverHandler target = +// new ClusterAwareReaderFailoverHandler( +// mockPluginService, +// properties); +// final ReaderFailoverResult result = target.failover(hosts, hosts.get(currentHostIndex)); +// +// assertTrue(result.isConnected()); +// assertSame(mockConnection, result.getConnection()); +// assertEquals(hosts.get(successHostIndex), result.getHost()); +// +// final HostSpec successHost = hosts.get(successHostIndex); +// verify(mockPluginService, atLeast(4)).setAvailability(any(), eq(HostAvailability.NOT_AVAILABLE)); +// verify(mockPluginService, never()) +// .setAvailability(eq(successHost.asAliases()), eq(HostAvailability.NOT_AVAILABLE)); +// verify(mockPluginService, times(1)) +// .setAvailability(eq(successHost.asAliases()), eq(HostAvailability.AVAILABLE)); +// } +// +// @Test +// public void testFailover_timeout() throws SQLException { +// // original host list: [active writer, active reader, current connection (reader), active +// // reader, down reader, active reader] +// // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] +// // connection attempts are made in pairs using the above list +// // expected test result: failure to get reader since process is limited to 5s and each attempt +// // to connect takes 20s +// final List hosts = defaultHosts; +// final int currentHostIndex = 2; +// for (HostSpec host : hosts) { +// when(mockPluginService.forceConnect(host, properties)) +// .thenAnswer((Answer) invocation -> { +// Thread.sleep(20000); +// return mockConnection; +// }); +// } +// +// hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); +// hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); +// +// final ReaderFailoverHandler target = +// new ClusterAwareReaderFailoverHandler( +// mockPluginService, +// properties, +// 5000, +// 30000, +// false); +// +// final long startTimeNano = System.nanoTime(); +// final ReaderFailoverResult result = +// target.failover(hosts, hosts.get(currentHostIndex)); +// final long durationNano = System.nanoTime() - startTimeNano; +// +// assertFalse(result.isConnected()); +// assertNull(result.getConnection()); +// assertNull(result.getHost()); +// +// // 5s is a max allowed failover timeout; add 1s for inaccurate measurements +// assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); +// } +// +// @Test +// public void testFailover_nullOrEmptyHostList() throws SQLException { +// final ClusterAwareReaderFailoverHandler target = +// new ClusterAwareReaderFailoverHandler( +// mockPluginService, +// properties); +// final HostSpec currentHost = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("writer") +// .port(1234).build(); +// +// ReaderFailoverResult result = target.failover(null, currentHost); +// assertFalse(result.isConnected()); +// assertNull(result.getConnection()); +// assertNull(result.getHost()); +// +// final List hosts = new ArrayList<>(); +// result = target.failover(hosts, currentHost); +// assertFalse(result.isConnected()); +// assertNull(result.getConnection()); +// assertNull(result.getHost()); +// } +// +// @Test +// public void testGetReader_connectionSuccess() throws SQLException { +// // even number of connection attempts +// // first connection attempt to return succeeds, second attempt cancelled +// // expected test result: successful connection for host at index 2 +// final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) +// final HostSpec slowHost = hosts.get(1); +// final HostSpec fastHost = hosts.get(2); +// when(mockPluginService.forceConnect(slowHost, properties)) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(20000); +// return mockConnection; +// }); +// when(mockPluginService.forceConnect(eq(fastHost), eq(properties))).thenReturn(mockConnection); +// +// Dialect mockDialect = Mockito.mock(Dialect.class); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// +// final ReaderFailoverHandler target = +// new ClusterAwareReaderFailoverHandler( +// mockPluginService, +// properties); +// final ReaderFailoverResult result = target.getReaderConnection(hosts); +// +// assertTrue(result.isConnected()); +// assertSame(mockConnection, result.getConnection()); +// assertEquals(hosts.get(2), result.getHost()); +// +// verify(mockPluginService, never()).setAvailability(any(), eq(HostAvailability.NOT_AVAILABLE)); +// verify(mockPluginService, times(1)) +// .setAvailability(eq(fastHost.asAliases()), eq(HostAvailability.AVAILABLE)); +// } +// +// @Test +// public void testGetReader_connectionFailure() throws SQLException { +// // odd number of connection attempts +// // first connection attempt to return fails +// // expected test result: failure to get reader +// final List hosts = defaultHosts.subList(0, 4); // 3 connection attempts (writer not attempted) +// when(mockPluginService.forceConnect(any(), eq(properties))).thenThrow(new SQLException("exception", "08S01", null)); +// +// Dialect mockDialect = Mockito.mock(Dialect.class); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// +// final int currentHostIndex = 2; +// +// final ReaderFailoverHandler target = +// new ClusterAwareReaderFailoverHandler( +// mockPluginService, +// properties); +// final ReaderFailoverResult result = target.getReaderConnection(hosts); +// +// assertFalse(result.isConnected()); +// assertNull(result.getConnection()); +// assertNull(result.getHost()); +// } +// +// @Test +// public void testGetReader_connectionAttemptsTimeout() throws SQLException { +// // connection attempts time out before they can succeed +// // first connection attempt to return times out +// // expected test result: failure to get reader +// final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) +// when(mockPluginService.forceConnect(any(), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// try { +// Thread.sleep(5000); +// } catch (InterruptedException exception) { +// // ignore +// } +// return mockConnection; +// }); +// +// Dialect mockDialect = Mockito.mock(Dialect.class); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// +// final ClusterAwareReaderFailoverHandler target = +// new ClusterAwareReaderFailoverHandler( +// mockPluginService, +// properties, +// 60000, +// 1000, +// false); +// final ReaderFailoverResult result = target.getReaderConnection(hosts); +// +// assertFalse(result.isConnected()); +// assertNull(result.getConnection()); +// assertNull(result.getHost()); +// } +// +// @Test +// public void testGetHostTuplesByPriority() { +// final List originalHosts = defaultHosts; +// originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); +// originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); +// originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); +// +// final ClusterAwareReaderFailoverHandler target = +// new ClusterAwareReaderFailoverHandler( +// mockPluginService, +// properties); +// final List hostsByPriority = target.getHostsByPriority(originalHosts); +// +// int i = 0; +// +// // expecting active readers +// while (i < hostsByPriority.size() +// && hostsByPriority.get(i).getRole() == HostRole.READER +// && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { +// i++; +// } +// +// // expecting a writer +// while (i < hostsByPriority.size() +// && hostsByPriority.get(i).getRole() == HostRole.WRITER) { +// i++; +// } +// +// // expecting down readers +// while (i < hostsByPriority.size() +// && hostsByPriority.get(i).getRole() == HostRole.READER +// && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { +// i++; +// } +// +// assertEquals(hostsByPriority.size(), i); +// } +// +// @Test +// public void testGetReaderTuplesByPriority() { +// final List originalHosts = defaultHosts; +// originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); +// originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); +// originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); +// +// Dialect mockDialect = Mockito.mock(Dialect.class); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// +// final ClusterAwareReaderFailoverHandler target = +// new ClusterAwareReaderFailoverHandler( +// mockPluginService, +// properties); +// final List hostsByPriority = target.getReaderHostsByPriority(originalHosts); +// +// int i = 0; +// +// // expecting active readers +// while (i < hostsByPriority.size() +// && hostsByPriority.get(i).getRole() == HostRole.READER +// && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { +// i++; +// } +// +// // expecting down readers +// while (i < hostsByPriority.size() +// && hostsByPriority.get(i).getRole() == HostRole.READER +// && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { +// i++; +// } +// +// assertEquals(hostsByPriority.size(), i); +// } +// +// @Test +// public void testHostFailoverStrictReaderEnabled() { +// +// final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("writer").port(1234).role(HostRole.WRITER).build(); +// final HostSpec reader = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader1").port(1234).role(HostRole.READER).build(); +// final List hosts = Arrays.asList(writer, reader); +// +// Dialect mockDialect = Mockito.mock(Dialect.class); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// final ClusterAwareReaderFailoverHandler target = +// new ClusterAwareReaderFailoverHandler( +// mockPluginService, +// properties, +// DEFAULT_FAILOVER_TIMEOUT, +// DEFAULT_READER_CONNECT_TIMEOUT, +// true); +// +// // The writer is included because the original writer has likely become a reader. +// List expectedHostsByPriority = Arrays.asList(reader, writer); +// +// List hostsByPriority = target.getHostsByPriority(hosts); +// assertEquals(expectedHostsByPriority, hostsByPriority); +// +// // Should pick the reader even if unavailable. The unavailable reader will be lower priority than the writer. +// reader.setAvailability(HostAvailability.NOT_AVAILABLE); +// expectedHostsByPriority = Arrays.asList(writer, reader); +// +// hostsByPriority = target.getHostsByPriority(hosts); +// assertEquals(expectedHostsByPriority, hostsByPriority); +// +// // Writer node will only be picked if it is the only node in topology; +// List expectedWriterHost = Collections.singletonList(writer); +// +// hostsByPriority = target.getHostsByPriority(Collections.singletonList(writer)); +// assertEquals(expectedWriterHost, hostsByPriority); +// } +// } From 81085513096aad768f577ec8f2bb9cb7a907c6ab Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Wed, 20 Aug 2025 09:40:55 -0700 Subject: [PATCH 07/42] Add hostAvailabilityMap to WriterFailoverResult --- .../ClusterAwareWriterFailoverHandler.java | 44 +++++++++++-------- .../failover/FailoverConnectionPlugin.java | 13 ++++++ .../plugin/failover/WriterFailoverResult.java | 12 ++++- 3 files changed, 50 insertions(+), 19 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java index 81098ba65..711cdb355 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java @@ -19,10 +19,12 @@ import java.sql.Connection; import java.sql.SQLException; import java.util.List; +import java.util.Map; import java.util.Properties; import java.util.Set; import java.util.concurrent.Callable; import java.util.concurrent.CompletionService; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorCompletionService; import java.util.concurrent.ExecutorService; @@ -52,7 +54,6 @@ * same writer host, 2) try to update cluster topology and connect to a newly elected writer. */ public class ClusterAwareWriterFailoverHandler implements WriterFailoverHandler { - private static final Logger LOGGER = Logger.getLogger(ClusterAwareReaderFailoverHandler.class.getName()); protected int maxFailoverTimeoutMs = 60000; // 60 sec @@ -62,8 +63,6 @@ public class ClusterAwareWriterFailoverHandler implements WriterFailoverHandler protected FullServicesContainer servicesContainer; protected PluginService pluginService; protected ReaderFailoverHandler readerFailoverHandler; - private static final WriterFailoverResult DEFAULT_RESULT = - new WriterFailoverResult(false, false, null, null, "None"); public ClusterAwareWriterFailoverHandler( final FullServicesContainer servicesContainer, @@ -102,7 +101,7 @@ public WriterFailoverResult failover(final ConnectionService connectionService, throws SQLException { if (Utils.isNullOrEmpty(currentTopology)) { LOGGER.severe(() -> Messages.get("ClusterAwareWriterFailoverHandler.failoverCalledWithInvalidTopology")); - return DEFAULT_RESULT; + return new WriterFailoverResult(false, false, null, null, null, "None"); } final boolean singleTask = @@ -132,7 +131,7 @@ public WriterFailoverResult failover(final ConnectionService connectionService, } LOGGER.fine(() -> Messages.get("ClusterAwareWriterFailoverHandler.failedToConnectToWriterInstance")); - return DEFAULT_RESULT; + return new WriterFailoverResult(false, false, null, result.getHostAvailabilityMap(), null, "None"); } finally { if (!executorService.isTerminated()) { executorService.shutdownNow(); // terminate all remaining tasks @@ -160,11 +159,13 @@ private void submitTasks( final CompletionService completionService, final boolean singleTask) { final HostSpec writerHost = getWriter(currentTopology); + final Map availabilityMap = new ConcurrentHashMap<>(); if (!singleTask) { completionService.submit( new ReconnectToWriterHandler( connectionService, this.getNewPluginService(), + availabilityMap, writerHost, this.initialConnectionProps, this.reconnectWriterIntervalMs)); @@ -174,6 +175,7 @@ private void submitTasks( new WaitForNewWriterHandler( connectionService, this.getNewPluginService(), + availabilityMap, this.readerFailoverHandler, writerHost, this.initialConnectionProps, @@ -210,7 +212,7 @@ private WriterFailoverResult getNextResult( timeoutMs, TimeUnit.MILLISECONDS); if (firstCompleted == null) { // The task was unsuccessful and we have timed out - return DEFAULT_RESULT; + return new WriterFailoverResult(false, false, null, null, null, "None"); } final WriterFailoverResult result = firstCompleted.get(); if (result.isConnected()) { @@ -229,7 +231,7 @@ private WriterFailoverResult getNextResult( } catch (final ExecutionException e) { // return failure below } - return DEFAULT_RESULT; + return new WriterFailoverResult(false, false, null, null, null, "None"); } private void logTaskSuccess(final WriterFailoverResult result) { @@ -276,19 +278,22 @@ private SQLException createInterruptedException(final InterruptedException e) { private static class ReconnectToWriterHandler implements Callable { private final ConnectionService connectionService; + private final PluginService pluginService; + private final Map availabilityMap; private final HostSpec originalWriterHost; private final Properties props; private final int reconnectWriterIntervalMs; - private final PluginService pluginService; public ReconnectToWriterHandler( final ConnectionService connectionService, final PluginService pluginService, + final Map availabilityMap, final HostSpec originalWriterHost, final Properties props, final int reconnectWriterIntervalMs) { this.connectionService = connectionService; this.pluginService = pluginService; + this.availabilityMap = availabilityMap; this.originalWriterHost = originalWriterHost; this.props = props; this.reconnectWriterIntervalMs = reconnectWriterIntervalMs; @@ -321,7 +326,7 @@ public WriterFailoverResult call() { () -> Messages.get( "ClusterAwareWriterFailoverHandler.taskAEncounteredException", new Object[] {exception})); - return new WriterFailoverResult(false, false, null, null, "TaskA", exception); + return new WriterFailoverResult(false, false, null, this.availabilityMap, null, "TaskA", exception); } } @@ -332,14 +337,14 @@ public WriterFailoverResult call() { success = isCurrentHostWriter(latestTopology); LOGGER.finest("[TaskA] success: " + success); - pluginService.setAvailability(this.originalWriterHost.asAliases(), HostAvailability.AVAILABLE); - return new WriterFailoverResult(success, false, latestTopology, success ? conn : null, "TaskA"); + availabilityMap.put(this.originalWriterHost.getHost(), HostAvailability.AVAILABLE); + return new WriterFailoverResult(success, false, latestTopology, this.availabilityMap, success ? conn : null, "TaskA"); } catch (final InterruptedException exception) { Thread.currentThread().interrupt(); - return new WriterFailoverResult(success, false, latestTopology, success ? conn : null, "TaskA"); + return new WriterFailoverResult(success, false, latestTopology, this.availabilityMap, success ? conn : null, "TaskA"); } catch (final Exception ex) { - LOGGER.severe(() -> ex.getMessage()); - return new WriterFailoverResult(false, false, null, null, "TaskA"); + LOGGER.severe(ex::getMessage); + return new WriterFailoverResult(false, false, null, this.availabilityMap, null, "TaskA"); } finally { try { if (conn != null && !success && !conn.isClosed()) { @@ -370,6 +375,7 @@ private static class WaitForNewWriterHandler implements Callable availabilityMap; private final ReaderFailoverHandler readerFailoverHandler; private final HostSpec originalWriterHost; private final Properties props; @@ -382,6 +388,7 @@ private static class WaitForNewWriterHandler implements Callable availabilityMap, final ReaderFailoverHandler readerFailoverHandler, final HostSpec originalWriterHost, final Properties props, @@ -389,6 +396,7 @@ public WaitForNewWriterHandler( final List currentTopology) { this.connectionService = connectionService; this.pluginService = pluginService; + this.availabilityMap = availabilityMap; this.readerFailoverHandler = readerFailoverHandler; this.originalWriterHost = originalWriterHost; this.props = props; @@ -415,11 +423,12 @@ public WriterFailoverResult call() { true, true, this.currentTopology, + this.availabilityMap, this.currentConnection, "TaskB"); } catch (final InterruptedException exception) { Thread.currentThread().interrupt(); - return new WriterFailoverResult(false, false, null, null, "TaskB"); + return new WriterFailoverResult(false, false, null, this.availabilityMap, null, "TaskB"); } catch (final Exception ex) { LOGGER.severe( () -> Messages.get( @@ -531,11 +540,10 @@ private boolean connectToWriter(final HostSpec writerCandidate) { try { // connect to the new writer this.currentConnection = this.connectionService.open(writerCandidate, this.props); - // TODO: replace with a map that is shared between the two handlers - this.pluginService.setAvailability(writerCandidate.asAliases(), HostAvailability.AVAILABLE); + this.availabilityMap.put(writerCandidate.getHost(), HostAvailability.AVAILABLE); return true; } catch (final SQLException exception) { - this.pluginService.setAvailability(writerCandidate.asAliases(), HostAvailability.NOT_AVAILABLE); + this.availabilityMap.put(writerCandidate.getHost(), HostAvailability.NOT_AVAILABLE); return false; } } diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java index f378d3c38..2a2eb057a 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java @@ -782,6 +782,19 @@ protected void failoverWriter() throws SQLException { return; } + Map hostAvailabilityMap = failoverResult.getHostAvailabilityMap(); + if (hostAvailabilityMap != null && !hostAvailabilityMap.isEmpty()) { + List allHosts = this.pluginService.getAllHosts(); + for (HostSpec host : allHosts) { + for (String alias : host.getAliases()) { + HostAvailability availability = hostAvailabilityMap.get(alias); + if (availability != null) { + host.setAvailability(availability); + } + } + } + } + this.pluginService.setCurrentConnection(failoverResult.getNewConnection(), writerHostSpec); LOGGER.fine( diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/WriterFailoverResult.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/WriterFailoverResult.java index b1fc63a33..e79f395a7 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/WriterFailoverResult.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/WriterFailoverResult.java @@ -19,7 +19,9 @@ import java.sql.Connection; import java.sql.SQLException; import java.util.List; +import java.util.Map; import software.amazon.jdbc.HostSpec; +import software.amazon.jdbc.hostavailability.HostAvailability; /** * This class holds results of Writer Failover Process. @@ -29,6 +31,7 @@ public class WriterFailoverResult { private final boolean isConnected; private final boolean isNewHost; private final List topology; + private final Map hostAvailabilityMap; private final Connection newConnection; private final String taskName; private final SQLException exception; @@ -37,21 +40,24 @@ public WriterFailoverResult( final boolean isConnected, final boolean isNewHost, final List topology, + final Map hostAvailabilityMap, final Connection newConnection, final String taskName) { - this(isConnected, isNewHost, topology, newConnection, taskName, null); + this(isConnected, isNewHost, topology, hostAvailabilityMap, newConnection, taskName, null); } public WriterFailoverResult( final boolean isConnected, final boolean isNewHost, final List topology, + final Map hostAvailabilityMap, final Connection newConnection, final String taskName, final SQLException exception) { this.isConnected = isConnected; this.isNewHost = isNewHost; this.topology = topology; + this.hostAvailabilityMap = hostAvailabilityMap; this.newConnection = newConnection; this.taskName = taskName; this.exception = exception; @@ -86,6 +92,10 @@ public List getTopology() { return this.topology; } + public Map getHostAvailabilityMap() { + return this.hostAvailabilityMap; + } + /** * Get the new connection established by the failover procedure if successful. * From 0a38ab3a27ee2d7dd00c6983fd031e2d92c6d018 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Wed, 20 Aug 2025 12:45:16 -0700 Subject: [PATCH 08/42] Add hostAvailabilityMap to ReaderFailoverResult --- .../ClusterAwareReaderFailoverHandler.java | 102 ++++++++++++------ .../failover/FailoverConnectionPlugin.java | 31 +++--- .../plugin/failover/ReaderFailoverResult.java | 18 +++- 3 files changed, 100 insertions(+), 51 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java index 18df01f3c..1cfb9ad67 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java @@ -21,9 +21,11 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.Map; import java.util.Properties; import java.util.concurrent.Callable; import java.util.concurrent.CompletionService; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorCompletionService; import java.util.concurrent.ExecutorService; @@ -60,8 +62,6 @@ public class ClusterAwareReaderFailoverHandler implements ReaderFailoverHandler Logger.getLogger(ClusterAwareReaderFailoverHandler.class.getName()); protected static final int DEFAULT_FAILOVER_TIMEOUT = 60000; // 60 sec protected static final int DEFAULT_READER_CONNECT_TIMEOUT = 30000; // 30 sec - public static final ReaderFailoverResult FAILED_READER_FAILOVER_RESULT = - new ReaderFailoverResult(null, null, false); protected Properties props; protected int maxFailoverTimeoutMs; protected int timeoutMs; @@ -135,26 +135,28 @@ public ReaderFailoverResult failover( throws SQLException { if (Utils.isNullOrEmpty(hosts)) { LOGGER.fine(() -> Messages.get("ClusterAwareReaderFailoverHandler.invalidTopology", new Object[] {"failover"})); - return FAILED_READER_FAILOVER_RESULT; + return new ReaderFailoverResult(null, null, false, null, null); } + final Map availabilityMap = new ConcurrentHashMap<>(); final ExecutorService executor = ExecutorFactory.newSingleThreadExecutor("failover"); final Future future = - submitInternalFailoverTask(connectionService, hosts, currentHost, executor); - return getInternalFailoverResult(executor, future); + submitInternalFailoverTask(connectionService, hosts, currentHost, executor, availabilityMap); + return getInternalFailoverResult(executor, future, availabilityMap); } private Future submitInternalFailoverTask( final ConnectionService connectionService, final List hosts, final HostSpec currentHost, - final ExecutorService executor) { + final ExecutorService executor, + final Map availabilityMap) { final Future future = executor.submit(() -> { ReaderFailoverResult result; try { while (true) { - result = failoverInternal(connectionService, hosts, currentHost); + result = failoverInternal(connectionService, hosts, currentHost, availabilityMap); if (result != null && result.isConnected()) { return result; } @@ -162,9 +164,9 @@ private Future submitInternalFailoverTask( TimeUnit.SECONDS.sleep(1); } } catch (final SQLException ex) { - return new ReaderFailoverResult(null, null, false, ex); + return new ReaderFailoverResult(null, null, false, ex, availabilityMap); } catch (final Exception ex) { - return new ReaderFailoverResult(null, null, false, new SQLException(ex)); + return new ReaderFailoverResult(null, null, false, new SQLException(ex), availabilityMap); } }); executor.shutdown(); @@ -173,13 +175,14 @@ private Future submitInternalFailoverTask( private ReaderFailoverResult getInternalFailoverResult( final ExecutorService executor, - final Future future) throws SQLException { + final Future future, + final Map availabilityMap) throws SQLException { try { final ReaderFailoverResult result = future.get(this.maxFailoverTimeoutMs, TimeUnit.MILLISECONDS); if (result == null) { LOGGER.warning( Messages.get("ClusterAwareReaderFailoverHandler.timeout", new Object[] {this.maxFailoverTimeoutMs})); - return FAILED_READER_FAILOVER_RESULT; + return new ReaderFailoverResult(null, null, false, null, availabilityMap); } return result; @@ -187,10 +190,10 @@ private ReaderFailoverResult getInternalFailoverResult( Thread.currentThread().interrupt(); throw new SQLException(Messages.get("ClusterAwareReaderFailoverHandler.interruptedThread"), "70100", e); } catch (final ExecutionException e) { - return FAILED_READER_FAILOVER_RESULT; + return new ReaderFailoverResult(null, null, false, null, availabilityMap); } catch (final TimeoutException e) { future.cancel(true); - return FAILED_READER_FAILOVER_RESULT; + return new ReaderFailoverResult(null, null, false, null, availabilityMap); } finally { if (!executor.isTerminated()) { executor.shutdownNow(); // terminate all remaining tasks @@ -201,13 +204,16 @@ private ReaderFailoverResult getInternalFailoverResult( protected ReaderFailoverResult failoverInternal( final ConnectionService connectionService, final List hosts, - final HostSpec currentHost) + final HostSpec currentHost, + final Map availabilityMap) throws SQLException { if (currentHost != null) { this.pluginService.setAvailability(currentHost.asAliases(), HostAvailability.NOT_AVAILABLE); + availabilityMap.put(currentHost.getHost(), HostAvailability.NOT_AVAILABLE); } + final List hostsByPriority = getHostsByPriority(hosts); - return getConnectionFromHostGroup(connectionService, hostsByPriority); + return getConnectionFromHostGroup(connectionService, hostsByPriority, availabilityMap); } public List getHostsByPriority(final List hosts) { @@ -257,11 +263,12 @@ public ReaderFailoverResult getReaderConnection( () -> Messages.get( "ClusterAwareReaderFailover.invalidTopology", new Object[] {"getReaderConnection"})); - return FAILED_READER_FAILOVER_RESULT; + return new ReaderFailoverResult(null, null, false, null, null); } final List hostsByPriority = getReaderHostsByPriority(hostList); - return getConnectionFromHostGroup(connectionService, hostsByPriority); + final Map availabilityMap = new ConcurrentHashMap<>(); + return getConnectionFromHostGroup(connectionService, hostsByPriority, availabilityMap); } public List getReaderHostsByPriority(final List hosts) { @@ -303,7 +310,9 @@ public List getReaderHostsByPriority(final List hosts) { } private ReaderFailoverResult getConnectionFromHostGroup( - final ConnectionService connectionService, final List hosts) + final ConnectionService connectionService, + final List hosts, + final Map availabilityMap) throws SQLException { final ExecutorService executor = ExecutorFactory.newFixedThreadPool(2, "failover"); @@ -313,7 +322,7 @@ private ReaderFailoverResult getConnectionFromHostGroup( for (int i = 0; i < hosts.size(); i += 2) { // submit connection attempt tasks in batches of 2 final ReaderFailoverResult result = - getResultFromNextTaskBatch(connectionService, hosts, executor, completionService, i); + getResultFromNextTaskBatch(connectionService, hosts, availabilityMap, executor, completionService, i); if (result.isConnected() || result.getException() != null) { return result; } @@ -329,7 +338,8 @@ private ReaderFailoverResult getConnectionFromHostGroup( return new ReaderFailoverResult( null, null, - false); + false, + availabilityMap); } finally { executor.shutdownNow(); } @@ -338,6 +348,7 @@ private ReaderFailoverResult getConnectionFromHostGroup( private ReaderFailoverResult getResultFromNextTaskBatch( final ConnectionService connectionService, final List hosts, + final Map availabilityMap, final ExecutorService executor, final CompletionService completionService, final int i) throws SQLException { @@ -346,12 +357,25 @@ private ReaderFailoverResult getResultFromNextTaskBatch( completionService.submit( // TODO: are there performance concerns with creating a new plugin service this often? new ConnectionAttemptTask( - connectionService, this.getNewPluginService(), hosts.get(i), this.props, this.isStrictReaderRequired)); + connectionService, + this.getNewPluginService(), + availabilityMap, + hosts.get(i), + this.props, + this.isStrictReaderRequired)); if (numTasks == 2) { - completionService.submit(new ConnectionAttemptTask(connectionService, this.getNewPluginService(), hosts.get(i + 1), this.props, this.isStrictReaderRequired)); + completionService.submit( + new ConnectionAttemptTask( + connectionService, + this.getNewPluginService(), + availabilityMap, + hosts.get(i + 1), + this.props, + this.isStrictReaderRequired)); } + for (int taskNum = 0; taskNum < numTasks; taskNum++) { - result = getNextResult(completionService); + result = getNextResult(completionService, availabilityMap); if (result.isConnected()) { executor.shutdownNow(); return result; @@ -364,20 +388,24 @@ private ReaderFailoverResult getResultFromNextTaskBatch( return new ReaderFailoverResult( null, null, - false); + false, + availabilityMap); } - private ReaderFailoverResult getNextResult(final CompletionService service) + private ReaderFailoverResult getNextResult( + final CompletionService service, + final Map availabilityMap) throws SQLException { + ReaderFailoverResult failureResult = new ReaderFailoverResult(null, null, false, null, availabilityMap); try { final Future future = service.poll(this.timeoutMs, TimeUnit.MILLISECONDS); if (future == null) { - return FAILED_READER_FAILOVER_RESULT; + return failureResult; } final ReaderFailoverResult result = future.get(); - return result == null ? FAILED_READER_FAILOVER_RESULT : result; + return result == null ? failureResult : result; } catch (final ExecutionException e) { - return FAILED_READER_FAILOVER_RESULT; + return failureResult; } catch (final InterruptedException e) { Thread.currentThread().interrupt(); // "Thread was interrupted" @@ -410,6 +438,7 @@ private PluginService getNewPluginService() { private static class ConnectionAttemptTask implements Callable { private final ConnectionService connectionService; private final PluginService pluginService; + private final Map availabilityMap; private final HostSpec newHost; private final Properties props; private final boolean isStrictReaderRequired; @@ -417,11 +446,13 @@ private static class ConnectionAttemptTask implements Callable availabilityMap, final HostSpec newHost, final Properties props, final boolean isStrictReaderRequired) { this.connectionService = connectionService; this.pluginService = pluginService; + this.availabilityMap = availabilityMap; this.newHost = newHost; this.props = props; this.isStrictReaderRequired = isStrictReaderRequired; @@ -442,7 +473,7 @@ public ReaderFailoverResult call() { copy.putAll(props); final Connection conn = this.connectionService.open(this.newHost, copy); - this.pluginService.setAvailability(this.newHost.asAliases(), HostAvailability.AVAILABLE); + this.availabilityMap.put(this.newHost.getHost(), HostAvailability.AVAILABLE); if (this.isStrictReaderRequired) { // need to ensure that new connection is a connection to a reader node @@ -460,7 +491,7 @@ public ReaderFailoverResult call() { // ignore } - return FAILED_READER_FAILOVER_RESULT; + return new ReaderFailoverResult(null, null, false, null, availabilityMap); } } catch (SQLException e) { LOGGER.fine(Messages.get("ClusterAwareReaderFailoverHandler.errorGettingHostRole", new Object[] {e})); @@ -471,7 +502,7 @@ public ReaderFailoverResult call() { // ignore } - return FAILED_READER_FAILOVER_RESULT; + return new ReaderFailoverResult(null, null, false, null, availabilityMap); } } @@ -480,9 +511,9 @@ public ReaderFailoverResult call() { "ClusterAwareReaderFailoverHandler.successfulReaderConnection", new Object[] {this.newHost.getUrl()})); LOGGER.fine("New reader failover connection object: " + conn); - return new ReaderFailoverResult(conn, this.newHost, true); + return new ReaderFailoverResult(conn, this.newHost, true, this.availabilityMap); } catch (final SQLException e) { - this.pluginService.setAvailability(newHost.asAliases(), HostAvailability.NOT_AVAILABLE); + this.availabilityMap.put(newHost.getHost(), HostAvailability.NOT_AVAILABLE); LOGGER.fine( () -> Messages.get( "ClusterAwareReaderFailoverHandler.failedReaderConnection", @@ -493,10 +524,11 @@ public ReaderFailoverResult call() { null, null, false, - e); + e, + this.availabilityMap); } - return FAILED_READER_FAILOVER_RESULT; + return new ReaderFailoverResult(null, null, false, null, availabilityMap); } } } diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java index 2a2eb057a..2a186272f 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java @@ -672,6 +672,8 @@ protected void failoverReader(final HostSpec failedHostSpec) throws SQLException if (exception != null) { throw exception; } + + updateHostAvailability(result.getHostAvailabilityMap()); } if (result == null || !result.isConnected()) { @@ -757,6 +759,8 @@ protected void failoverWriter() throws SQLException { if (exception != null) { throw exception; } + + updateHostAvailability(failoverResult.getHostAvailabilityMap()); } if (failoverResult == null || !failoverResult.isConnected()) { @@ -782,19 +786,6 @@ protected void failoverWriter() throws SQLException { return; } - Map hostAvailabilityMap = failoverResult.getHostAvailabilityMap(); - if (hostAvailabilityMap != null && !hostAvailabilityMap.isEmpty()) { - List allHosts = this.pluginService.getAllHosts(); - for (HostSpec host : allHosts) { - for (String alias : host.getAliases()) { - HostAvailability availability = hostAvailabilityMap.get(alias); - if (availability != null) { - host.setAvailability(availability); - } - } - } - } - this.pluginService.setCurrentConnection(failoverResult.getNewConnection(), writerHostSpec); LOGGER.fine( @@ -835,6 +826,20 @@ protected void failoverWriter() throws SQLException { } } + private void updateHostAvailability(Map hostAvailabilityMap) { + if (hostAvailabilityMap != null && !hostAvailabilityMap.isEmpty()) { + List allHosts = this.pluginService.getAllHosts(); + for (HostSpec host : allHosts) { + for (String alias : host.getAliases()) { + HostAvailability availability = hostAvailabilityMap.get(alias); + if (availability != null) { + host.setAvailability(availability); + } + } + } + } + } + protected void invalidateCurrentConnection() { final Connection conn = this.pluginService.getCurrentConnection(); if (conn == null) { diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ReaderFailoverResult.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ReaderFailoverResult.java index 8d6d9630b..2ada37ca2 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ReaderFailoverResult.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ReaderFailoverResult.java @@ -18,7 +18,9 @@ import java.sql.Connection; import java.sql.SQLException; +import java.util.Map; import software.amazon.jdbc.HostSpec; +import software.amazon.jdbc.hostavailability.HostAvailability; /** * This class holds results of Reader Failover Process. @@ -30,21 +32,27 @@ public class ReaderFailoverResult { private final boolean isConnected; private final SQLException exception; private final HostSpec newHost; + private final Map hostAvailabilityMap; public ReaderFailoverResult( - final Connection newConnection, final HostSpec newHost, final boolean isConnected) { - this(newConnection, newHost, isConnected, null); + final Connection newConnection, + final HostSpec newHost, + final boolean isConnected, + final Map hostAvailabilityMap) { + this(newConnection, newHost, isConnected, null, hostAvailabilityMap); } public ReaderFailoverResult( final Connection newConnection, final HostSpec newHost, final boolean isConnected, - final SQLException exception) { + final SQLException exception, + final Map hostAvailabilityMap) { this.newConnection = newConnection; this.newHost = newHost; this.isConnected = isConnected; this.exception = exception; + this.hostAvailabilityMap = hostAvailabilityMap; } /** @@ -82,4 +90,8 @@ public boolean isConnected() { public SQLException getException() { return exception; } + + public Map getHostAvailabilityMap() { + return hostAvailabilityMap; + } } From 1a7e2bcc4e3cba4bd5c63eb1103a93a9d1238f56 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Wed, 20 Aug 2025 13:33:10 -0700 Subject: [PATCH 09/42] Initialize HostListProvider in PartialPluginService constructor --- .../amazon/jdbc/PartialPluginService.java | 4 ++++ .../ClusterAwareReaderFailoverHandler.java | 22 ++++++++----------- .../ClusterAwareWriterFailoverHandler.java | 12 +++------- .../failover/FailoverConnectionPlugin.java | 9 +++----- 4 files changed, 19 insertions(+), 28 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java b/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java index 67027195b..c223cca57 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java +++ b/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java @@ -36,6 +36,7 @@ import org.checkerframework.checker.nullness.qual.Nullable; import software.amazon.jdbc.cleanup.CanReleaseResources; import software.amazon.jdbc.dialect.Dialect; +import software.amazon.jdbc.dialect.HostListProviderSupplier; import software.amazon.jdbc.exceptions.ExceptionHandler; import software.amazon.jdbc.exceptions.ExceptionManager; import software.amazon.jdbc.hostavailability.HostAvailability; @@ -123,6 +124,9 @@ public PartialPluginService( this.exceptionHandler = this.configurationProfile != null && this.configurationProfile.getExceptionHandler() != null ? this.configurationProfile.getExceptionHandler() : null; + + HostListProviderSupplier supplier = this.dbDialect.getHostListProvider(); + this.hostListProvider = supplier.getProvider(this.props, this.originalUrl, this.servicesContainer); } @Override diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java index 1cfb9ad67..a1f917a2a 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java @@ -19,6 +19,7 @@ import java.sql.Connection; import java.sql.SQLException; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.List; import java.util.Map; @@ -37,7 +38,6 @@ import software.amazon.jdbc.HostSpec; import software.amazon.jdbc.PartialPluginService; import software.amazon.jdbc.PluginService; -import software.amazon.jdbc.dialect.HostListProviderSupplier; import software.amazon.jdbc.hostavailability.HostAvailability; import software.amazon.jdbc.util.ExecutorFactory; import software.amazon.jdbc.util.FullServicesContainer; @@ -317,12 +317,15 @@ private ReaderFailoverResult getConnectionFromHostGroup( final ExecutorService executor = ExecutorFactory.newFixedThreadPool(2, "failover"); final CompletionService completionService = new ExecutorCompletionService<>(executor); + // The ConnectionAttemptTask threads should have their own plugin services since they execute concurrently and + // PluginService was not designed to be thread-safe. + List pluginServices = Arrays.asList(getNewPluginService(), getNewPluginService()); try { for (int i = 0; i < hosts.size(); i += 2) { // submit connection attempt tasks in batches of 2 final ReaderFailoverResult result = - getResultFromNextTaskBatch(connectionService, hosts, availabilityMap, executor, completionService, i); + getResultFromNextTaskBatch(connectionService, hosts, availabilityMap, executor, completionService, pluginServices, i); if (result.isConnected() || result.getException() != null) { return result; } @@ -351,14 +354,14 @@ private ReaderFailoverResult getResultFromNextTaskBatch( final Map availabilityMap, final ExecutorService executor, final CompletionService completionService, + final List pluginServices, final int i) throws SQLException { ReaderFailoverResult result; final int numTasks = i + 1 < hosts.size() ? 2 : 1; completionService.submit( - // TODO: are there performance concerns with creating a new plugin service this often? new ConnectionAttemptTask( connectionService, - this.getNewPluginService(), + pluginServices.get(0), availabilityMap, hosts.get(i), this.props, @@ -367,7 +370,7 @@ private ReaderFailoverResult getResultFromNextTaskBatch( completionService.submit( new ConnectionAttemptTask( connectionService, - this.getNewPluginService(), + pluginServices.get(1), availabilityMap, hosts.get(i + 1), this.props, @@ -418,7 +421,7 @@ private ReaderFailoverResult getNextResult( } private PluginService getNewPluginService() { - PartialPluginService partialPluginService = new PartialPluginService( + return new PartialPluginService( this.servicesContainer, this.props, this.pluginService.getOriginalUrl(), @@ -426,13 +429,6 @@ private PluginService getNewPluginService() { this.pluginService.getTargetDriverDialect(), this.pluginService.getDialect() ); - - // TODO: can we clean this up, eg move to PartialPluginService constructor? - final HostListProviderSupplier supplier = this.pluginService.getDialect().getHostListProvider(); - partialPluginService.setHostListProvider( - supplier.getProvider(this.props, this.pluginService.getOriginalUrl(), this.servicesContainer)); - - return partialPluginService; } private static class ConnectionAttemptTask implements Callable { diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java index 711cdb355..76703fbb7 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java @@ -36,7 +36,6 @@ import software.amazon.jdbc.HostSpec; import software.amazon.jdbc.PartialPluginService; import software.amazon.jdbc.PluginService; -import software.amazon.jdbc.dialect.HostListProviderSupplier; import software.amazon.jdbc.hostavailability.HostAvailability; import software.amazon.jdbc.util.ExecutorFactory; import software.amazon.jdbc.util.FullServicesContainer; @@ -186,7 +185,9 @@ private void submitTasks( } private PluginService getNewPluginService() { - PartialPluginService partialPluginService = new PartialPluginService( + // Each task should get its own PluginService since they execute concurrently and PluginService was not designed to + // be thread-safe. + return new PartialPluginService( this.servicesContainer, this.initialConnectionProps, this.pluginService.getOriginalUrl(), @@ -194,13 +195,6 @@ private PluginService getNewPluginService() { this.pluginService.getTargetDriverDialect(), this.pluginService.getDialect() ); - - // TODO: can we clean this up, eg move to PartialPluginService constructor? - final HostListProviderSupplier supplier = this.pluginService.getDialect().getHostListProvider(); - partialPluginService.setHostListProvider( - supplier.getProvider(this.initialConnectionProps, this.pluginService.getOriginalUrl(), this.servicesContainer)); - - return partialPluginService; } private WriterFailoverResult getNextResult( diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java index 2a186272f..7e365ad93 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java @@ -619,13 +619,10 @@ protected void dealWithIllegalStateException( protected void failover(final HostSpec failedHost) throws SQLException { this.pluginService.setAvailability(failedHost.asAliases(), HostAvailability.NOT_AVAILABLE); - - // TODO: After failover, retrieve the unavailable hosts from the handlers after replacing - // pluginService#setAvailability - TargetDriverHelper helper = new TargetDriverHelper(); - java.sql.Driver driver = helper.getTargetDriver(this.pluginService.getOriginalUrl(), properties); - final ConnectionProvider defaultConnectionProvider = new DriverConnectionProvider(driver); if (this.connectionService == null) { + TargetDriverHelper helper = new TargetDriverHelper(); + java.sql.Driver driver = helper.getTargetDriver(this.pluginService.getOriginalUrl(), properties); + final ConnectionProvider defaultConnectionProvider = new DriverConnectionProvider(driver); this.connectionService = new ConnectionServiceImpl( servicesContainer.getStorageService(), servicesContainer.getMonitorService(), From a965c3f96d4ed27570c694a09381839b49ab8fb8 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Thu, 21 Aug 2025 14:22:17 -0700 Subject: [PATCH 10/42] Pass ConnectionService as constructor arg to failover handlers --- .../amazon/jdbc/PartialPluginService.java | 4 ++ .../ClusterAwareReaderFailoverHandler.java | 29 +++++---- .../ClusterAwareWriterFailoverHandler.java | 18 +++--- .../failover/FailoverConnectionPlugin.java | 59 +++++++++---------- .../failover/ReaderFailoverHandler.java | 7 +-- .../failover/WriterFailoverHandler.java | 4 +- .../connection/ConnectionServiceImpl.java | 4 -- 7 files changed, 61 insertions(+), 64 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java b/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java index c223cca57..880c6f0dc 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java +++ b/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java @@ -125,6 +125,10 @@ public PartialPluginService( ? this.configurationProfile.getExceptionHandler() : null; + servicesContainer.setHostListProviderService(this); + servicesContainer.setPluginService(this); + servicesContainer.setPluginManagerService(this); + HostListProviderSupplier supplier = this.dbDialect.getHostListProvider(); this.hostListProvider = supplier.getProvider(this.props, this.originalUrl, this.servicesContainer); } diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java index a1f917a2a..18101b26d 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java @@ -67,6 +67,7 @@ public class ClusterAwareReaderFailoverHandler implements ReaderFailoverHandler protected int timeoutMs; protected boolean isStrictReaderRequired; protected final FullServicesContainer servicesContainer; + protected final ConnectionService connectionService; protected final PluginService pluginService; /** @@ -77,9 +78,11 @@ public class ClusterAwareReaderFailoverHandler implements ReaderFailoverHandler */ public ClusterAwareReaderFailoverHandler( final FullServicesContainer servicesContainer, + final ConnectionService connectionService, final Properties props) { this( servicesContainer, + connectionService, props, DEFAULT_FAILOVER_TIMEOUT, DEFAULT_READER_CONNECT_TIMEOUT, @@ -98,11 +101,13 @@ public ClusterAwareReaderFailoverHandler( */ public ClusterAwareReaderFailoverHandler( final FullServicesContainer servicesContainer, + final ConnectionService connectionService, final Properties props, final int maxFailoverTimeoutMs, final int timeoutMs, final boolean isStrictReaderRequired) { this.servicesContainer = servicesContainer; + this.connectionService = connectionService; this.pluginService = servicesContainer.getPluginService(); this.props = props; this.maxFailoverTimeoutMs = maxFailoverTimeoutMs; @@ -130,8 +135,7 @@ protected void setTimeoutMs(final int timeoutMs) { * @return {@link ReaderFailoverResult} The results of this process. */ @Override - public ReaderFailoverResult failover( - final ConnectionService connectionService, final List hosts, final HostSpec currentHost) + public ReaderFailoverResult failover(final List hosts, final HostSpec currentHost) throws SQLException { if (Utils.isNullOrEmpty(hosts)) { LOGGER.fine(() -> Messages.get("ClusterAwareReaderFailoverHandler.invalidTopology", new Object[] {"failover"})); @@ -142,12 +146,11 @@ public ReaderFailoverResult failover( final ExecutorService executor = ExecutorFactory.newSingleThreadExecutor("failover"); final Future future = - submitInternalFailoverTask(connectionService, hosts, currentHost, executor, availabilityMap); + submitInternalFailoverTask(hosts, currentHost, executor, availabilityMap); return getInternalFailoverResult(executor, future, availabilityMap); } private Future submitInternalFailoverTask( - final ConnectionService connectionService, final List hosts, final HostSpec currentHost, final ExecutorService executor, @@ -156,7 +159,7 @@ private Future submitInternalFailoverTask( ReaderFailoverResult result; try { while (true) { - result = failoverInternal(connectionService, hosts, currentHost, availabilityMap); + result = failoverInternal(hosts, currentHost, availabilityMap); if (result != null && result.isConnected()) { return result; } @@ -202,7 +205,6 @@ private ReaderFailoverResult getInternalFailoverResult( } protected ReaderFailoverResult failoverInternal( - final ConnectionService connectionService, final List hosts, final HostSpec currentHost, final Map availabilityMap) @@ -213,7 +215,7 @@ protected ReaderFailoverResult failoverInternal( } final List hostsByPriority = getHostsByPriority(hosts); - return getConnectionFromHostGroup(connectionService, hostsByPriority, availabilityMap); + return getConnectionFromHostGroup(hostsByPriority, availabilityMap); } public List getHostsByPriority(final List hosts) { @@ -255,8 +257,7 @@ public List getHostsByPriority(final List hosts) { * @return {@link ReaderFailoverResult} The results of this process. */ @Override - public ReaderFailoverResult getReaderConnection( - final ConnectionService connectionService, final List hostList) + public ReaderFailoverResult getReaderConnection(final List hostList) throws SQLException { if (Utils.isNullOrEmpty(hostList)) { LOGGER.fine( @@ -268,7 +269,7 @@ public ReaderFailoverResult getReaderConnection( final List hostsByPriority = getReaderHostsByPriority(hostList); final Map availabilityMap = new ConcurrentHashMap<>(); - return getConnectionFromHostGroup(connectionService, hostsByPriority, availabilityMap); + return getConnectionFromHostGroup(hostsByPriority, availabilityMap); } public List getReaderHostsByPriority(final List hosts) { @@ -310,7 +311,6 @@ public List getReaderHostsByPriority(final List hosts) { } private ReaderFailoverResult getConnectionFromHostGroup( - final ConnectionService connectionService, final List hosts, final Map availabilityMap) throws SQLException { @@ -325,7 +325,7 @@ private ReaderFailoverResult getConnectionFromHostGroup( for (int i = 0; i < hosts.size(); i += 2) { // submit connection attempt tasks in batches of 2 final ReaderFailoverResult result = - getResultFromNextTaskBatch(connectionService, hosts, availabilityMap, executor, completionService, pluginServices, i); + getResultFromNextTaskBatch(hosts, availabilityMap, executor, completionService, pluginServices, i); if (result.isConnected() || result.getException() != null) { return result; } @@ -349,7 +349,6 @@ private ReaderFailoverResult getConnectionFromHostGroup( } private ReaderFailoverResult getResultFromNextTaskBatch( - final ConnectionService connectionService, final List hosts, final Map availabilityMap, final ExecutorService executor, @@ -360,7 +359,7 @@ private ReaderFailoverResult getResultFromNextTaskBatch( final int numTasks = i + 1 < hosts.size() ? 2 : 1; completionService.submit( new ConnectionAttemptTask( - connectionService, + this.connectionService, pluginServices.get(0), availabilityMap, hosts.get(i), @@ -369,7 +368,7 @@ private ReaderFailoverResult getResultFromNextTaskBatch( if (numTasks == 2) { completionService.submit( new ConnectionAttemptTask( - connectionService, + this.connectionService, pluginServices.get(1), availabilityMap, hosts.get(i + 1), diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java index 76703fbb7..07894bf63 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java @@ -60,14 +60,17 @@ public class ClusterAwareWriterFailoverHandler implements WriterFailoverHandler protected int reconnectWriterIntervalMs = 5000; // 5 sec protected Properties initialConnectionProps; protected FullServicesContainer servicesContainer; + protected ConnectionService connectionService; protected PluginService pluginService; protected ReaderFailoverHandler readerFailoverHandler; public ClusterAwareWriterFailoverHandler( final FullServicesContainer servicesContainer, + final ConnectionService connectionService, final ReaderFailoverHandler readerFailoverHandler, final Properties initialConnectionProps) { this.servicesContainer = servicesContainer; + this.connectionService = connectionService; this.pluginService = servicesContainer.getPluginService(); this.readerFailoverHandler = readerFailoverHandler; this.initialConnectionProps = initialConnectionProps; @@ -75,6 +78,7 @@ public ClusterAwareWriterFailoverHandler( public ClusterAwareWriterFailoverHandler( final FullServicesContainer servicesContainer, + final ConnectionService connectionService, final ReaderFailoverHandler readerFailoverHandler, final Properties initialConnectionProps, final int failoverTimeoutMs, @@ -82,6 +86,7 @@ public ClusterAwareWriterFailoverHandler( final int reconnectWriterIntervalMs) { this( servicesContainer, + connectionService, readerFailoverHandler, initialConnectionProps); this.maxFailoverTimeoutMs = failoverTimeoutMs; @@ -96,7 +101,7 @@ public ClusterAwareWriterFailoverHandler( * @return {@link WriterFailoverResult} The results of this process. */ @Override - public WriterFailoverResult failover(final ConnectionService connectionService, final List currentTopology) + public WriterFailoverResult failover(final List currentTopology) throws SQLException { if (Utils.isNullOrEmpty(currentTopology)) { LOGGER.severe(() -> Messages.get("ClusterAwareWriterFailoverHandler.failoverCalledWithInvalidTopology")); @@ -109,7 +114,7 @@ public WriterFailoverResult failover(final ConnectionService connectionService, final ExecutorService executorService = ExecutorFactory.newFixedThreadPool(2, "failover"); final CompletionService completionService = new ExecutorCompletionService<>(executorService); - submitTasks(connectionService, currentTopology, executorService, completionService, singleTask); + submitTasks(currentTopology, executorService, completionService, singleTask); try { final long startTimeNano = System.nanoTime(); @@ -152,7 +157,6 @@ private static HostSpec getWriter(final List topology) { } private void submitTasks( - final ConnectionService connectionService, final List currentTopology, final ExecutorService executorService, final CompletionService completionService, @@ -162,7 +166,7 @@ private void submitTasks( if (!singleTask) { completionService.submit( new ReconnectToWriterHandler( - connectionService, + this.connectionService, this.getNewPluginService(), availabilityMap, writerHost, @@ -172,7 +176,7 @@ private void submitTasks( completionService.submit( new WaitForNewWriterHandler( - connectionService, + this.connectionService, this.getNewPluginService(), availabilityMap, this.readerFailoverHandler, @@ -310,7 +314,7 @@ public WriterFailoverResult call() { conn.close(); } - conn = connectionService.open(this.originalWriterHost, this.props); + conn = this.connectionService.open(this.originalWriterHost, this.props); this.pluginService.forceRefreshHostList(conn); latestTopology = this.pluginService.getAllHosts(); } catch (final SQLException exception) { @@ -439,7 +443,7 @@ private void connectToReader() throws InterruptedException { while (true) { try { final ReaderFailoverResult connResult = - readerFailoverHandler.getReaderConnection(this.connectionService, this.currentTopology); + readerFailoverHandler.getReaderConnection(this.currentTopology); if (isValidReaderConnection(connResult)) { this.currentReaderConnection = connResult.getConnection(); this.currentReaderHost = connResult.getHost(); diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java index 7e365ad93..a423a27c9 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java @@ -28,7 +28,7 @@ import java.util.Set; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; -import java.util.function.Supplier; +import java.util.function.Function; import java.util.logging.Level; import java.util.logging.Logger; import org.checkerframework.checker.nullness.qual.NonNull; @@ -122,8 +122,8 @@ public class FailoverConnectionPlugin extends AbstractConnectionPlugin { private RdsUrlType rdsUrlType = null; private HostListProviderService hostListProviderService; private final AuroraStaleDnsHelper staleDnsHelper; - private Supplier writerFailoverHandlerSupplier; - private Supplier readerFailoverHandlerSupplier; + private Function writerFailoverHandlerSupplier; + private Function readerFailoverHandlerSupplier; public static final AwsWrapperProperty FAILOVER_CLUSTER_TOPOLOGY_REFRESH_RATE_MS = new AwsWrapperProperty( @@ -316,16 +316,18 @@ public void initHostProvider( initHostProvider( hostListProviderService, initHostProviderFunc, - () -> + (connectionService) -> new ClusterAwareReaderFailoverHandler( this.servicesContainer, + connectionService, this.properties, this.failoverTimeoutMsSetting, this.failoverReaderConnectTimeoutMsSetting, this.failoverMode == FailoverMode.STRICT_READER), - () -> + (connectionService) -> new ClusterAwareWriterFailoverHandler( this.servicesContainer, + connectionService, this.readerFailoverHandler, this.properties, this.failoverTimeoutMsSetting, @@ -336,8 +338,8 @@ public void initHostProvider( void initHostProvider( final HostListProviderService hostListProviderService, final JdbcCallable initHostProviderFunc, - final Supplier readerFailoverHandlerSupplier, - final Supplier writerFailoverHandlerSupplier) + final Function readerFailoverHandlerSupplier, + final Function writerFailoverHandlerSupplier) throws SQLException { this.readerFailoverHandlerSupplier = readerFailoverHandlerSupplier; this.writerFailoverHandlerSupplier = writerFailoverHandlerSupplier; @@ -618,24 +620,6 @@ protected void dealWithIllegalStateException( */ protected void failover(final HostSpec failedHost) throws SQLException { this.pluginService.setAvailability(failedHost.asAliases(), HostAvailability.NOT_AVAILABLE); - - if (this.connectionService == null) { - TargetDriverHelper helper = new TargetDriverHelper(); - java.sql.Driver driver = helper.getTargetDriver(this.pluginService.getOriginalUrl(), properties); - final ConnectionProvider defaultConnectionProvider = new DriverConnectionProvider(driver); - this.connectionService = new ConnectionServiceImpl( - servicesContainer.getStorageService(), - servicesContainer.getMonitorService(), - servicesContainer.getTelemetryFactory(), - defaultConnectionProvider, - this.pluginService.getOriginalUrl(), - this.pluginService.getDriverProtocol(), - this.pluginService.getTargetDriverDialect(), - this.pluginService.getDialect(), - properties - ); - } - if (this.failoverMode == FailoverMode.STRICT_WRITER) { failoverWriter(); } else { @@ -662,8 +646,7 @@ protected void failoverReader(final HostSpec failedHostSpec) throws SQLException failedHost = failedHostSpec; } - final ReaderFailoverResult result = readerFailoverHandler.failover( - this.connectionService, this.pluginService.getHosts(), failedHost); + final ReaderFailoverResult result = readerFailoverHandler.failover(this.pluginService.getHosts(), failedHost); if (result != null) { final SQLException exception = result.getException(); if (exception != null) { @@ -750,7 +733,7 @@ protected void failoverWriter() throws SQLException { try { LOGGER.info(() -> Messages.get("Failover.startWriterFailover")); final WriterFailoverResult failoverResult = - this.writerFailoverHandler.failover(this.connectionService, this.pluginService.getAllHosts()); + this.writerFailoverHandler.failover(this.pluginService.getAllHosts()); if (failoverResult != null) { final SQLException exception = failoverResult.getException(); if (exception != null) { @@ -925,20 +908,36 @@ public Connection connect( final boolean isInitialConnection, final JdbcCallable connectFunc) throws SQLException { + if (this.connectionService == null) { + TargetDriverHelper helper = new TargetDriverHelper(); + java.sql.Driver driver = helper.getTargetDriver(this.pluginService.getOriginalUrl(), properties); + final ConnectionProvider defaultConnectionProvider = new DriverConnectionProvider(driver); + this.connectionService = new ConnectionServiceImpl( + servicesContainer.getStorageService(), + servicesContainer.getMonitorService(), + servicesContainer.getTelemetryFactory(), + defaultConnectionProvider, + this.pluginService.getOriginalUrl(), + this.pluginService.getDriverProtocol(), + this.pluginService.getTargetDriverDialect(), + this.pluginService.getDialect(), + properties + ); + } this.initFailoverMode(); if (this.readerFailoverHandler == null) { if (this.readerFailoverHandlerSupplier == null) { throw new SQLException(Messages.get("Failover.nullReaderFailoverHandlerSupplier")); } - this.readerFailoverHandler = this.readerFailoverHandlerSupplier.get(); + this.readerFailoverHandler = this.readerFailoverHandlerSupplier.apply(this.connectionService); } if (this.writerFailoverHandler == null) { if (this.writerFailoverHandlerSupplier == null) { throw new SQLException(Messages.get("Failover.nullWriterFailoverHandlerSupplier")); } - this.writerFailoverHandler = this.writerFailoverHandlerSupplier.get(); + this.writerFailoverHandler = this.writerFailoverHandlerSupplier.apply(this.connectionService); } Connection conn = null; diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ReaderFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ReaderFailoverHandler.java index 008c1e17e..e006558b6 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ReaderFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ReaderFailoverHandler.java @@ -19,7 +19,6 @@ import java.sql.SQLException; import java.util.List; import software.amazon.jdbc.HostSpec; -import software.amazon.jdbc.util.connection.ConnectionService; /** * Interface for Reader Failover Process handler. This handler implements all necessary logic to try @@ -37,8 +36,7 @@ public interface ReaderFailoverHandler { * @return {@link ReaderFailoverResult} The results of this process. * @throws SQLException indicating whether the failover attempt was successful. */ - ReaderFailoverResult failover( - ConnectionService connectionService, List hosts, HostSpec currentHost) throws SQLException; + ReaderFailoverResult failover(List hosts, HostSpec currentHost) throws SQLException; /** * Called to get any available reader connection. If no reader is available then result of process @@ -48,6 +46,5 @@ ReaderFailoverResult failover( * @return {@link ReaderFailoverResult} The results of this process. * @throws SQLException if any error occurred while attempting a reader connection. */ - ReaderFailoverResult getReaderConnection( - ConnectionService connectionService, List hostList) throws SQLException; + ReaderFailoverResult getReaderConnection(List hostList) throws SQLException; } diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/WriterFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/WriterFailoverHandler.java index 68e69259b..caae8de8c 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/WriterFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/WriterFailoverHandler.java @@ -19,7 +19,6 @@ import java.sql.SQLException; import java.util.List; import software.amazon.jdbc.HostSpec; -import software.amazon.jdbc.util.connection.ConnectionService; /** * Interface for Writer Failover Process handler. This handler implements all necessary logic to try @@ -34,6 +33,5 @@ public interface WriterFailoverHandler { * @return {@link WriterFailoverResult} The results of this process. * @throws SQLException indicating whether the failover attempt was successful. */ - WriterFailoverResult failover( - ConnectionService connectionService, List currentTopology) throws SQLException; + WriterFailoverResult failover(List currentTopology) throws SQLException; } diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java index 73a68fdcb..0e506a112 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java @@ -68,10 +68,6 @@ public ConnectionServiceImpl( ); this.pluginService = partialPluginService; - servicesContainer.setHostListProviderService(partialPluginService); - servicesContainer.setPluginService(partialPluginService); - servicesContainer.setPluginManagerService(partialPluginService); - this.pluginManager.init(servicesContainer, props, partialPluginService, null); } From a3d8a5e1cdd51d5aadf77e09ed829490d787f9a1 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Fri, 22 Aug 2025 10:39:13 -0700 Subject: [PATCH 11/42] PR suggestions --- .../ClusterAwareReaderFailoverHandler.java | 120 +++++++----------- .../ClusterAwareWriterFailoverHandler.java | 52 ++++---- .../failover/FailoverConnectionPlugin.java | 87 +++++++------ .../failover/ReaderFailoverHandler.java | 10 ++ .../plugin/failover/ReaderFailoverResult.java | 16 +-- .../failover/WriterFailoverHandler.java | 10 ++ .../plugin/failover/WriterFailoverResult.java | 12 +- 7 files changed, 140 insertions(+), 167 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java index 18101b26d..4cb587f55 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java @@ -60,15 +60,18 @@ public class ClusterAwareReaderFailoverHandler implements ReaderFailoverHandler private static final Logger LOGGER = Logger.getLogger(ClusterAwareReaderFailoverHandler.class.getName()); + protected static final ReaderFailoverResult FAILED_READER_FAILOVER_RESULT = + new ReaderFailoverResult(null, null, false); protected static final int DEFAULT_FAILOVER_TIMEOUT = 60000; // 60 sec protected static final int DEFAULT_READER_CONNECT_TIMEOUT = 30000; // 30 sec + protected final Map hostAvailabilityMap = new ConcurrentHashMap<>(); + protected final FullServicesContainer servicesContainer; + protected final ConnectionService connectionService; + protected final PluginService pluginService; protected Properties props; protected int maxFailoverTimeoutMs; protected int timeoutMs; protected boolean isStrictReaderRequired; - protected final FullServicesContainer servicesContainer; - protected final ConnectionService connectionService; - protected final PluginService pluginService; /** * ClusterAwareReaderFailoverHandler constructor. @@ -115,6 +118,11 @@ public ClusterAwareReaderFailoverHandler( this.isStrictReaderRequired = isStrictReaderRequired; } + @Override + public Map getHostAvailabilityMap() { + return this.hostAvailabilityMap; + } + /** * Set process timeout in millis. Entire process of connecting to a reader will be limited by this * time duration. @@ -125,41 +133,29 @@ protected void setTimeoutMs(final int timeoutMs) { this.timeoutMs = timeoutMs; } - /** - * Called to start Reader Failover Process. This process tries to connect to any reader. If no - * reader is available then driver may also try to connect to a writer host, down hosts, and the - * current reader host. - * - * @param hosts Cluster current topology - * @param currentHost The currently connected host that has failed. - * @return {@link ReaderFailoverResult} The results of this process. - */ @Override public ReaderFailoverResult failover(final List hosts, final HostSpec currentHost) throws SQLException { if (Utils.isNullOrEmpty(hosts)) { LOGGER.fine(() -> Messages.get("ClusterAwareReaderFailoverHandler.invalidTopology", new Object[] {"failover"})); - return new ReaderFailoverResult(null, null, false, null, null); + return FAILED_READER_FAILOVER_RESULT; } - final Map availabilityMap = new ConcurrentHashMap<>(); final ExecutorService executor = ExecutorFactory.newSingleThreadExecutor("failover"); - final Future future = - submitInternalFailoverTask(hosts, currentHost, executor, availabilityMap); - return getInternalFailoverResult(executor, future, availabilityMap); + final Future future = submitInternalFailoverTask(hosts, currentHost, executor); + return getInternalFailoverResult(executor, future); } private Future submitInternalFailoverTask( final List hosts, final HostSpec currentHost, - final ExecutorService executor, - final Map availabilityMap) { + final ExecutorService executor) { final Future future = executor.submit(() -> { ReaderFailoverResult result; try { while (true) { - result = failoverInternal(hosts, currentHost, availabilityMap); + result = failoverInternal(hosts, currentHost); if (result != null && result.isConnected()) { return result; } @@ -167,9 +163,9 @@ private Future submitInternalFailoverTask( TimeUnit.SECONDS.sleep(1); } } catch (final SQLException ex) { - return new ReaderFailoverResult(null, null, false, ex, availabilityMap); + return new ReaderFailoverResult(null, null, false, ex); } catch (final Exception ex) { - return new ReaderFailoverResult(null, null, false, new SQLException(ex), availabilityMap); + return new ReaderFailoverResult(null, null, false, new SQLException(ex)); } }); executor.shutdown(); @@ -178,14 +174,13 @@ private Future submitInternalFailoverTask( private ReaderFailoverResult getInternalFailoverResult( final ExecutorService executor, - final Future future, - final Map availabilityMap) throws SQLException { + final Future future) throws SQLException { try { final ReaderFailoverResult result = future.get(this.maxFailoverTimeoutMs, TimeUnit.MILLISECONDS); if (result == null) { LOGGER.warning( Messages.get("ClusterAwareReaderFailoverHandler.timeout", new Object[] {this.maxFailoverTimeoutMs})); - return new ReaderFailoverResult(null, null, false, null, availabilityMap); + return FAILED_READER_FAILOVER_RESULT; } return result; @@ -193,10 +188,10 @@ private ReaderFailoverResult getInternalFailoverResult( Thread.currentThread().interrupt(); throw new SQLException(Messages.get("ClusterAwareReaderFailoverHandler.interruptedThread"), "70100", e); } catch (final ExecutionException e) { - return new ReaderFailoverResult(null, null, false, null, availabilityMap); + return FAILED_READER_FAILOVER_RESULT; } catch (final TimeoutException e) { future.cancel(true); - return new ReaderFailoverResult(null, null, false, null, availabilityMap); + return FAILED_READER_FAILOVER_RESULT; } finally { if (!executor.isTerminated()) { executor.shutdownNow(); // terminate all remaining tasks @@ -204,18 +199,15 @@ private ReaderFailoverResult getInternalFailoverResult( } } - protected ReaderFailoverResult failoverInternal( - final List hosts, - final HostSpec currentHost, - final Map availabilityMap) + protected ReaderFailoverResult failoverInternal(final List hosts, final HostSpec currentHost) throws SQLException { if (currentHost != null) { this.pluginService.setAvailability(currentHost.asAliases(), HostAvailability.NOT_AVAILABLE); - availabilityMap.put(currentHost.getHost(), HostAvailability.NOT_AVAILABLE); + this.hostAvailabilityMap.put(currentHost.getHost(), HostAvailability.NOT_AVAILABLE); } final List hostsByPriority = getHostsByPriority(hosts); - return getConnectionFromHostGroup(hostsByPriority, availabilityMap); + return getConnectionFromHostGroup(hostsByPriority); } public List getHostsByPriority(final List hosts) { @@ -264,12 +256,11 @@ public ReaderFailoverResult getReaderConnection(final List hostList) () -> Messages.get( "ClusterAwareReaderFailover.invalidTopology", new Object[] {"getReaderConnection"})); - return new ReaderFailoverResult(null, null, false, null, null); + return FAILED_READER_FAILOVER_RESULT; } final List hostsByPriority = getReaderHostsByPriority(hostList); - final Map availabilityMap = new ConcurrentHashMap<>(); - return getConnectionFromHostGroup(hostsByPriority, availabilityMap); + return getConnectionFromHostGroup(hostsByPriority); } public List getReaderHostsByPriority(final List hosts) { @@ -291,11 +282,10 @@ public List getReaderHostsByPriority(final List hosts) { Collections.shuffle(activeReaders); Collections.shuffle(downHostList); - final List hostsByPriority = new ArrayList<>(); - hostsByPriority.addAll(activeReaders); + final List hostsByPriority = new ArrayList<>(activeReaders); + final int numOfReaders = activeReaders.size() + downHostList.size(); hostsByPriority.addAll(downHostList); - final int numOfReaders = activeReaders.size() + downHostList.size(); if (writerHost == null) { return hostsByPriority; } @@ -310,10 +300,7 @@ public List getReaderHostsByPriority(final List hosts) { return hostsByPriority; } - private ReaderFailoverResult getConnectionFromHostGroup( - final List hosts, - final Map availabilityMap) - throws SQLException { + private ReaderFailoverResult getConnectionFromHostGroup(final List hosts) throws SQLException { final ExecutorService executor = ExecutorFactory.newFixedThreadPool(2, "failover"); final CompletionService completionService = new ExecutorCompletionService<>(executor); @@ -325,7 +312,7 @@ private ReaderFailoverResult getConnectionFromHostGroup( for (int i = 0; i < hosts.size(); i += 2) { // submit connection attempt tasks in batches of 2 final ReaderFailoverResult result = - getResultFromNextTaskBatch(hosts, availabilityMap, executor, completionService, pluginServices, i); + getResultFromNextTaskBatch(hosts, executor, completionService, pluginServices, i); if (result.isConnected() || result.getException() != null) { return result; } @@ -338,11 +325,7 @@ private ReaderFailoverResult getConnectionFromHostGroup( } } - return new ReaderFailoverResult( - null, - null, - false, - availabilityMap); + return new ReaderFailoverResult(null, null, false); } finally { executor.shutdownNow(); } @@ -350,7 +333,6 @@ private ReaderFailoverResult getConnectionFromHostGroup( private ReaderFailoverResult getResultFromNextTaskBatch( final List hosts, - final Map availabilityMap, final ExecutorService executor, final CompletionService completionService, final List pluginServices, @@ -361,7 +343,7 @@ private ReaderFailoverResult getResultFromNextTaskBatch( new ConnectionAttemptTask( this.connectionService, pluginServices.get(0), - availabilityMap, + this.hostAvailabilityMap, hosts.get(i), this.props, this.isStrictReaderRequired)); @@ -370,14 +352,14 @@ private ReaderFailoverResult getResultFromNextTaskBatch( new ConnectionAttemptTask( this.connectionService, pluginServices.get(1), - availabilityMap, + this.hostAvailabilityMap, hosts.get(i + 1), this.props, this.isStrictReaderRequired)); } for (int taskNum = 0; taskNum < numTasks; taskNum++) { - result = getNextResult(completionService, availabilityMap); + result = getNextResult(completionService); if (result.isConnected()) { executor.shutdownNow(); return result; @@ -387,27 +369,20 @@ private ReaderFailoverResult getResultFromNextTaskBatch( return result; } } - return new ReaderFailoverResult( - null, - null, - false, - availabilityMap); + return new ReaderFailoverResult(null, null, false); } - private ReaderFailoverResult getNextResult( - final CompletionService service, - final Map availabilityMap) + private ReaderFailoverResult getNextResult(final CompletionService service) throws SQLException { - ReaderFailoverResult failureResult = new ReaderFailoverResult(null, null, false, null, availabilityMap); try { final Future future = service.poll(this.timeoutMs, TimeUnit.MILLISECONDS); if (future == null) { - return failureResult; + return FAILED_READER_FAILOVER_RESULT; } final ReaderFailoverResult result = future.get(); - return result == null ? failureResult : result; + return result == null ? FAILED_READER_FAILOVER_RESULT : result; } catch (final ExecutionException e) { - return failureResult; + return FAILED_READER_FAILOVER_RESULT; } catch (final InterruptedException e) { Thread.currentThread().interrupt(); // "Thread was interrupted" @@ -486,7 +461,7 @@ public ReaderFailoverResult call() { // ignore } - return new ReaderFailoverResult(null, null, false, null, availabilityMap); + return FAILED_READER_FAILOVER_RESULT; } } catch (SQLException e) { LOGGER.fine(Messages.get("ClusterAwareReaderFailoverHandler.errorGettingHostRole", new Object[] {e})); @@ -497,7 +472,7 @@ public ReaderFailoverResult call() { // ignore } - return new ReaderFailoverResult(null, null, false, null, availabilityMap); + return FAILED_READER_FAILOVER_RESULT; } } @@ -506,7 +481,7 @@ public ReaderFailoverResult call() { "ClusterAwareReaderFailoverHandler.successfulReaderConnection", new Object[] {this.newHost.getUrl()})); LOGGER.fine("New reader failover connection object: " + conn); - return new ReaderFailoverResult(conn, this.newHost, true, this.availabilityMap); + return new ReaderFailoverResult(conn, this.newHost, true); } catch (final SQLException e) { this.availabilityMap.put(newHost.getHost(), HostAvailability.NOT_AVAILABLE); LOGGER.fine( @@ -515,15 +490,10 @@ public ReaderFailoverResult call() { new Object[] {this.newHost.getUrl()})); // Propagate exceptions that are not caused by network errors. if (!this.pluginService.isNetworkException(e, this.pluginService.getTargetDriverDialect())) { - return new ReaderFailoverResult( - null, - null, - false, - e, - this.availabilityMap); + return new ReaderFailoverResult(null, null, false, e); } - return new ReaderFailoverResult(null, null, false, null, availabilityMap); + return FAILED_READER_FAILOVER_RESULT; } } } diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java index 07894bf63..8533b1055 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java @@ -54,15 +54,18 @@ */ public class ClusterAwareWriterFailoverHandler implements WriterFailoverHandler { private static final Logger LOGGER = Logger.getLogger(ClusterAwareReaderFailoverHandler.class.getName()); - + protected static final WriterFailoverResult DEFAULT_RESULT = + new WriterFailoverResult(false, false, null, null, "None"); + + protected final Properties initialConnectionProps; + protected final FullServicesContainer servicesContainer; + protected final ConnectionService connectionService; + protected final PluginService pluginService; + protected final ReaderFailoverHandler readerFailoverHandler; + protected final Map hostAvailabilityMap = new ConcurrentHashMap<>(); protected int maxFailoverTimeoutMs = 60000; // 60 sec protected int readTopologyIntervalMs = 5000; // 5 sec protected int reconnectWriterIntervalMs = 5000; // 5 sec - protected Properties initialConnectionProps; - protected FullServicesContainer servicesContainer; - protected ConnectionService connectionService; - protected PluginService pluginService; - protected ReaderFailoverHandler readerFailoverHandler; public ClusterAwareWriterFailoverHandler( final FullServicesContainer servicesContainer, @@ -94,18 +97,17 @@ public ClusterAwareWriterFailoverHandler( this.reconnectWriterIntervalMs = reconnectWriterIntervalMs; } - /** - * Called to start Writer Failover Process. - * - * @param currentTopology Cluster current topology - * @return {@link WriterFailoverResult} The results of this process. - */ + @Override + public Map getHostAvailabilityMap() { + return this.hostAvailabilityMap; + } + @Override public WriterFailoverResult failover(final List currentTopology) throws SQLException { if (Utils.isNullOrEmpty(currentTopology)) { LOGGER.severe(() -> Messages.get("ClusterAwareWriterFailoverHandler.failoverCalledWithInvalidTopology")); - return new WriterFailoverResult(false, false, null, null, null, "None"); + return DEFAULT_RESULT; } final boolean singleTask = @@ -135,7 +137,7 @@ public WriterFailoverResult failover(final List currentTopology) } LOGGER.fine(() -> Messages.get("ClusterAwareWriterFailoverHandler.failedToConnectToWriterInstance")); - return new WriterFailoverResult(false, false, null, result.getHostAvailabilityMap(), null, "None"); + return DEFAULT_RESULT; } finally { if (!executorService.isTerminated()) { executorService.shutdownNow(); // terminate all remaining tasks @@ -162,13 +164,12 @@ private void submitTasks( final CompletionService completionService, final boolean singleTask) { final HostSpec writerHost = getWriter(currentTopology); - final Map availabilityMap = new ConcurrentHashMap<>(); if (!singleTask) { completionService.submit( new ReconnectToWriterHandler( this.connectionService, this.getNewPluginService(), - availabilityMap, + this.hostAvailabilityMap, writerHost, this.initialConnectionProps, this.reconnectWriterIntervalMs)); @@ -178,7 +179,7 @@ private void submitTasks( new WaitForNewWriterHandler( this.connectionService, this.getNewPluginService(), - availabilityMap, + this.hostAvailabilityMap, this.readerFailoverHandler, writerHost, this.initialConnectionProps, @@ -210,7 +211,7 @@ private WriterFailoverResult getNextResult( timeoutMs, TimeUnit.MILLISECONDS); if (firstCompleted == null) { // The task was unsuccessful and we have timed out - return new WriterFailoverResult(false, false, null, null, null, "None"); + return DEFAULT_RESULT; } final WriterFailoverResult result = firstCompleted.get(); if (result.isConnected()) { @@ -229,7 +230,7 @@ private WriterFailoverResult getNextResult( } catch (final ExecutionException e) { // return failure below } - return new WriterFailoverResult(false, false, null, null, null, "None"); + return DEFAULT_RESULT; } private void logTaskSuccess(final WriterFailoverResult result) { @@ -324,7 +325,7 @@ public WriterFailoverResult call() { () -> Messages.get( "ClusterAwareWriterFailoverHandler.taskAEncounteredException", new Object[] {exception})); - return new WriterFailoverResult(false, false, null, this.availabilityMap, null, "TaskA", exception); + return new WriterFailoverResult(false, false, null, null, "TaskA", exception); } } @@ -335,14 +336,14 @@ public WriterFailoverResult call() { success = isCurrentHostWriter(latestTopology); LOGGER.finest("[TaskA] success: " + success); - availabilityMap.put(this.originalWriterHost.getHost(), HostAvailability.AVAILABLE); - return new WriterFailoverResult(success, false, latestTopology, this.availabilityMap, success ? conn : null, "TaskA"); + this.availabilityMap.put(this.originalWriterHost.getHost(), HostAvailability.AVAILABLE); + return new WriterFailoverResult(success, false, latestTopology, success ? conn : null, "TaskA"); } catch (final InterruptedException exception) { Thread.currentThread().interrupt(); - return new WriterFailoverResult(success, false, latestTopology, this.availabilityMap, success ? conn : null, "TaskA"); + return new WriterFailoverResult(success, false, latestTopology, success ? conn : null, "TaskA"); } catch (final Exception ex) { LOGGER.severe(ex::getMessage); - return new WriterFailoverResult(false, false, null, this.availabilityMap, null, "TaskA"); + return new WriterFailoverResult(false, false, null, null, "TaskA"); } finally { try { if (conn != null && !success && !conn.isClosed()) { @@ -421,12 +422,11 @@ public WriterFailoverResult call() { true, true, this.currentTopology, - this.availabilityMap, this.currentConnection, "TaskB"); } catch (final InterruptedException exception) { Thread.currentThread().interrupt(); - return new WriterFailoverResult(false, false, null, this.availabilityMap, null, "TaskB"); + return new WriterFailoverResult(false, false, null, null, "TaskB"); } catch (final Exception ex) { LOGGER.severe( () -> Messages.get( diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java index a423a27c9..4b07f0996 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java @@ -620,6 +620,23 @@ protected void dealWithIllegalStateException( */ protected void failover(final HostSpec failedHost) throws SQLException { this.pluginService.setAvailability(failedHost.asAliases(), HostAvailability.NOT_AVAILABLE); + if (this.connectionService == null) { + TargetDriverHelper helper = new TargetDriverHelper(); + java.sql.Driver driver = helper.getTargetDriver(this.pluginService.getOriginalUrl(), properties); + final ConnectionProvider defaultConnectionProvider = new DriverConnectionProvider(driver); + this.connectionService = new ConnectionServiceImpl( + servicesContainer.getStorageService(), + servicesContainer.getMonitorService(), + servicesContainer.getTelemetryFactory(), + defaultConnectionProvider, + this.pluginService.getOriginalUrl(), + this.pluginService.getDriverProtocol(), + this.pluginService.getTargetDriverDialect(), + this.pluginService.getDialect(), + properties + ); + } + if (this.failoverMode == FailoverMode.STRICT_WRITER) { failoverWriter(); } else { @@ -635,6 +652,13 @@ protected void failoverReader(final HostSpec failedHostSpec) throws SQLException this.failoverReaderTriggeredCounter.inc(); } + if (this.readerFailoverHandler == null) { + if (this.readerFailoverHandlerSupplier == null) { + throw new SQLException(Messages.get("Failover.nullReaderFailoverHandlerSupplier")); + } + this.readerFailoverHandler = this.readerFailoverHandlerSupplier.apply(this.connectionService); + } + final long failoverStartNano = System.nanoTime(); try { @@ -646,14 +670,14 @@ protected void failoverReader(final HostSpec failedHostSpec) throws SQLException failedHost = failedHostSpec; } - final ReaderFailoverResult result = readerFailoverHandler.failover(this.pluginService.getHosts(), failedHost); + final ReaderFailoverResult result = this.readerFailoverHandler.failover(this.pluginService.getHosts(), failedHost); if (result != null) { final SQLException exception = result.getException(); if (exception != null) { throw exception; } - updateHostAvailability(result.getHostAvailabilityMap()); + updateHostAvailability(this.readerFailoverHandler.getHostAvailabilityMap()); } if (result == null || !result.isConnected()) { @@ -693,6 +717,7 @@ protected void failoverReader(final HostSpec failedHostSpec) throws SQLException LOGGER.finest(() -> Messages.get( "Failover.readerFailoverElapsed", new Object[]{TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - failoverStartNano)})); + this.readerFailoverHandler = null; if (telemetryContext != null) { telemetryContext.closeContext(); if (this.telemetryFailoverAdditionalTopTraceSetting) { @@ -728,6 +753,13 @@ protected void failoverWriter() throws SQLException { this.failoverWriterTriggeredCounter.inc(); } + if (this.writerFailoverHandler == null) { + if (this.writerFailoverHandlerSupplier == null) { + throw new SQLException(Messages.get("Failover.nullWriterFailoverHandlerSupplier")); + } + this.writerFailoverHandler = this.writerFailoverHandlerSupplier.apply(this.connectionService); + } + long failoverStartTimeNano = System.nanoTime(); try { @@ -740,7 +772,7 @@ protected void failoverWriter() throws SQLException { throw exception; } - updateHostAvailability(failoverResult.getHostAvailabilityMap()); + updateHostAvailability(this.writerFailoverHandler.getHostAvailabilityMap()); } if (failoverResult == null || !failoverResult.isConnected()) { @@ -797,6 +829,7 @@ protected void failoverWriter() throws SQLException { LOGGER.finest(() -> Messages.get( "Failover.writerFailoverElapsed", new Object[]{TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - failoverStartTimeNano)})); + this.writerFailoverHandler = null; if (telemetryContext != null) { telemetryContext.closeContext(); if (this.telemetryFailoverAdditionalTopTraceSetting) { @@ -807,15 +840,15 @@ protected void failoverWriter() throws SQLException { } private void updateHostAvailability(Map hostAvailabilityMap) { - if (hostAvailabilityMap != null && !hostAvailabilityMap.isEmpty()) { - List allHosts = this.pluginService.getAllHosts(); - for (HostSpec host : allHosts) { - for (String alias : host.getAliases()) { - HostAvailability availability = hostAvailabilityMap.get(alias); - if (availability != null) { - host.setAvailability(availability); - } - } + if (hostAvailabilityMap == null || hostAvailabilityMap.isEmpty()) { + return; + } + + List allHosts = this.pluginService.getAllHosts(); + for (HostSpec host : allHosts) { + HostAvailability availability = hostAvailabilityMap.get(host.getHost()); + if (availability != null) { + host.setAvailability(availability); } } } @@ -908,37 +941,7 @@ public Connection connect( final boolean isInitialConnection, final JdbcCallable connectFunc) throws SQLException { - if (this.connectionService == null) { - TargetDriverHelper helper = new TargetDriverHelper(); - java.sql.Driver driver = helper.getTargetDriver(this.pluginService.getOriginalUrl(), properties); - final ConnectionProvider defaultConnectionProvider = new DriverConnectionProvider(driver); - this.connectionService = new ConnectionServiceImpl( - servicesContainer.getStorageService(), - servicesContainer.getMonitorService(), - servicesContainer.getTelemetryFactory(), - defaultConnectionProvider, - this.pluginService.getOriginalUrl(), - this.pluginService.getDriverProtocol(), - this.pluginService.getTargetDriverDialect(), - this.pluginService.getDialect(), - properties - ); - } - this.initFailoverMode(); - if (this.readerFailoverHandler == null) { - if (this.readerFailoverHandlerSupplier == null) { - throw new SQLException(Messages.get("Failover.nullReaderFailoverHandlerSupplier")); - } - this.readerFailoverHandler = this.readerFailoverHandlerSupplier.apply(this.connectionService); - } - - if (this.writerFailoverHandler == null) { - if (this.writerFailoverHandlerSupplier == null) { - throw new SQLException(Messages.get("Failover.nullWriterFailoverHandlerSupplier")); - } - this.writerFailoverHandler = this.writerFailoverHandlerSupplier.apply(this.connectionService); - } Connection conn = null; try { diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ReaderFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ReaderFailoverHandler.java index e006558b6..0f807fd54 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ReaderFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ReaderFailoverHandler.java @@ -18,7 +18,9 @@ import java.sql.SQLException; import java.util.List; +import java.util.Map; import software.amazon.jdbc.HostSpec; +import software.amazon.jdbc.hostavailability.HostAvailability; /** * Interface for Reader Failover Process handler. This handler implements all necessary logic to try @@ -47,4 +49,12 @@ public interface ReaderFailoverHandler { * @throws SQLException if any error occurred while attempting a reader connection. */ ReaderFailoverResult getReaderConnection(List hostList) throws SQLException; + + /** + * Get the host availability map for the failover handler. This map will be populated with host availability + * information during the failover process and can be used to determine which hosts are available. + * + * @return the host availability map for the failover handler. + */ + Map getHostAvailabilityMap(); } diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ReaderFailoverResult.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ReaderFailoverResult.java index 2ada37ca2..1367cc996 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ReaderFailoverResult.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ReaderFailoverResult.java @@ -18,9 +18,7 @@ import java.sql.Connection; import java.sql.SQLException; -import java.util.Map; import software.amazon.jdbc.HostSpec; -import software.amazon.jdbc.hostavailability.HostAvailability; /** * This class holds results of Reader Failover Process. @@ -32,27 +30,23 @@ public class ReaderFailoverResult { private final boolean isConnected; private final SQLException exception; private final HostSpec newHost; - private final Map hostAvailabilityMap; public ReaderFailoverResult( final Connection newConnection, final HostSpec newHost, - final boolean isConnected, - final Map hostAvailabilityMap) { - this(newConnection, newHost, isConnected, null, hostAvailabilityMap); + final boolean isConnected) { + this(newConnection, newHost, isConnected, null); } public ReaderFailoverResult( final Connection newConnection, final HostSpec newHost, final boolean isConnected, - final SQLException exception, - final Map hostAvailabilityMap) { + final SQLException exception) { this.newConnection = newConnection; this.newHost = newHost; this.isConnected = isConnected; this.exception = exception; - this.hostAvailabilityMap = hostAvailabilityMap; } /** @@ -90,8 +84,4 @@ public boolean isConnected() { public SQLException getException() { return exception; } - - public Map getHostAvailabilityMap() { - return hostAvailabilityMap; - } } diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/WriterFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/WriterFailoverHandler.java index caae8de8c..f05c6a8c0 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/WriterFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/WriterFailoverHandler.java @@ -18,7 +18,9 @@ import java.sql.SQLException; import java.util.List; +import java.util.Map; import software.amazon.jdbc.HostSpec; +import software.amazon.jdbc.hostavailability.HostAvailability; /** * Interface for Writer Failover Process handler. This handler implements all necessary logic to try @@ -34,4 +36,12 @@ public interface WriterFailoverHandler { * @throws SQLException indicating whether the failover attempt was successful. */ WriterFailoverResult failover(List currentTopology) throws SQLException; + + /** + * Get the host availability map for the failover handler. This map will be populated with host availability + * information during the failover process and can be used to determine which hosts are available. + * + * @return the host availability map for the failover handler. + */ + Map getHostAvailabilityMap(); } diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/WriterFailoverResult.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/WriterFailoverResult.java index e79f395a7..b1fc63a33 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/WriterFailoverResult.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/WriterFailoverResult.java @@ -19,9 +19,7 @@ import java.sql.Connection; import java.sql.SQLException; import java.util.List; -import java.util.Map; import software.amazon.jdbc.HostSpec; -import software.amazon.jdbc.hostavailability.HostAvailability; /** * This class holds results of Writer Failover Process. @@ -31,7 +29,6 @@ public class WriterFailoverResult { private final boolean isConnected; private final boolean isNewHost; private final List topology; - private final Map hostAvailabilityMap; private final Connection newConnection; private final String taskName; private final SQLException exception; @@ -40,24 +37,21 @@ public WriterFailoverResult( final boolean isConnected, final boolean isNewHost, final List topology, - final Map hostAvailabilityMap, final Connection newConnection, final String taskName) { - this(isConnected, isNewHost, topology, hostAvailabilityMap, newConnection, taskName, null); + this(isConnected, isNewHost, topology, newConnection, taskName, null); } public WriterFailoverResult( final boolean isConnected, final boolean isNewHost, final List topology, - final Map hostAvailabilityMap, final Connection newConnection, final String taskName, final SQLException exception) { this.isConnected = isConnected; this.isNewHost = isNewHost; this.topology = topology; - this.hostAvailabilityMap = hostAvailabilityMap; this.newConnection = newConnection; this.taskName = taskName; this.exception = exception; @@ -92,10 +86,6 @@ public List getTopology() { return this.topology; } - public Map getHostAvailabilityMap() { - return this.hostAvailabilityMap; - } - /** * Get the new connection established by the failover procedure if successful. * From 4df33da212cef40d4ec90085711209dafe7def59 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Fri, 22 Aug 2025 13:13:47 -0700 Subject: [PATCH 12/42] Fix checkstyle --- .../amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java index 4b07f0996..285dc0eac 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java @@ -670,7 +670,8 @@ protected void failoverReader(final HostSpec failedHostSpec) throws SQLException failedHost = failedHostSpec; } - final ReaderFailoverResult result = this.readerFailoverHandler.failover(this.pluginService.getHosts(), failedHost); + final ReaderFailoverResult result = + this.readerFailoverHandler.failover(this.pluginService.getHosts(), failedHost); if (result != null) { final SQLException exception = result.getException(); if (exception != null) { From 7a863c9496f4ec5eca9d3ceca866eadf4cec3fc1 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Fri, 22 Aug 2025 14:37:41 -0700 Subject: [PATCH 13/42] ReaderFailoverHandler tests passing --- .../ClusterAwareReaderFailoverHandler.java | 2 +- ...ClusterAwareReaderFailoverHandlerTest.java | 808 +++++++++--------- 2 files changed, 402 insertions(+), 408 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java index 4cb587f55..8f3ab2f62 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java @@ -394,7 +394,7 @@ private ReaderFailoverResult getNextResult(final CompletionService defaultHosts = Arrays.asList( -// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("writer").port(1234).role(HostRole.WRITER).build(), -// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("reader1").port(1234).role(HostRole.READER).build(), -// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("reader2").port(1234).role(HostRole.READER).build(), -// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("reader3").port(1234).role(HostRole.READER).build(), -// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("reader4").port(1234).role(HostRole.READER).build(), -// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("reader5").port(1234).role(HostRole.READER).build() -// ); -// -// @BeforeEach -// void setUp() { -// closeable = MockitoAnnotations.openMocks(this); -// } -// -// @AfterEach -// void tearDown() throws Exception { -// closeable.close(); -// } -// -// @Test -// public void testFailover() throws SQLException { -// // original host list: [active writer, active reader, current connection (reader), active -// // reader, down reader, active reader] -// // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] -// // connection attempts are made in pairs using the above list -// // expected test result: successful connection for host at index 4 -// final List hosts = defaultHosts; -// final int currentHostIndex = 2; -// final int successHostIndex = 4; -// for (int i = 0; i < hosts.size(); i++) { -// if (i != successHostIndex) { -// final SQLException exception = new SQLException("exception", "08S01", null); -// when(mockPluginService.forceConnect(hosts.get(i), properties)) -// .thenThrow(exception); -// when(mockPluginService.isNetworkException(exception, null)).thenReturn(true); -// } else { -// when(mockPluginService.forceConnect(hosts.get(i), properties)).thenReturn(mockConnection); -// } -// } -// when(mockPluginService.getTargetDriverDialect()).thenReturn(null); -// -// hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); -// hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); -// -// final ReaderFailoverHandler target = -// new ClusterAwareReaderFailoverHandler( -// mockPluginService, -// properties); -// final ReaderFailoverResult result = target.failover(hosts, hosts.get(currentHostIndex)); -// -// assertTrue(result.isConnected()); -// assertSame(mockConnection, result.getConnection()); -// assertEquals(hosts.get(successHostIndex), result.getHost()); -// -// final HostSpec successHost = hosts.get(successHostIndex); -// verify(mockPluginService, atLeast(4)).setAvailability(any(), eq(HostAvailability.NOT_AVAILABLE)); -// verify(mockPluginService, never()) -// .setAvailability(eq(successHost.asAliases()), eq(HostAvailability.NOT_AVAILABLE)); -// verify(mockPluginService, times(1)) -// .setAvailability(eq(successHost.asAliases()), eq(HostAvailability.AVAILABLE)); -// } -// -// @Test -// public void testFailover_timeout() throws SQLException { -// // original host list: [active writer, active reader, current connection (reader), active -// // reader, down reader, active reader] -// // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] -// // connection attempts are made in pairs using the above list -// // expected test result: failure to get reader since process is limited to 5s and each attempt -// // to connect takes 20s -// final List hosts = defaultHosts; -// final int currentHostIndex = 2; -// for (HostSpec host : hosts) { -// when(mockPluginService.forceConnect(host, properties)) -// .thenAnswer((Answer) invocation -> { -// Thread.sleep(20000); -// return mockConnection; -// }); -// } -// -// hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); -// hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); -// -// final ReaderFailoverHandler target = -// new ClusterAwareReaderFailoverHandler( -// mockPluginService, -// properties, -// 5000, -// 30000, -// false); -// -// final long startTimeNano = System.nanoTime(); -// final ReaderFailoverResult result = -// target.failover(hosts, hosts.get(currentHostIndex)); -// final long durationNano = System.nanoTime() - startTimeNano; -// -// assertFalse(result.isConnected()); -// assertNull(result.getConnection()); -// assertNull(result.getHost()); -// -// // 5s is a max allowed failover timeout; add 1s for inaccurate measurements -// assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); -// } -// -// @Test -// public void testFailover_nullOrEmptyHostList() throws SQLException { -// final ClusterAwareReaderFailoverHandler target = -// new ClusterAwareReaderFailoverHandler( -// mockPluginService, -// properties); -// final HostSpec currentHost = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("writer") -// .port(1234).build(); -// -// ReaderFailoverResult result = target.failover(null, currentHost); -// assertFalse(result.isConnected()); -// assertNull(result.getConnection()); -// assertNull(result.getHost()); -// -// final List hosts = new ArrayList<>(); -// result = target.failover(hosts, currentHost); -// assertFalse(result.isConnected()); -// assertNull(result.getConnection()); -// assertNull(result.getHost()); -// } -// -// @Test -// public void testGetReader_connectionSuccess() throws SQLException { -// // even number of connection attempts -// // first connection attempt to return succeeds, second attempt cancelled -// // expected test result: successful connection for host at index 2 -// final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) -// final HostSpec slowHost = hosts.get(1); -// final HostSpec fastHost = hosts.get(2); -// when(mockPluginService.forceConnect(slowHost, properties)) -// .thenAnswer( -// (Answer) -// invocation -> { -// Thread.sleep(20000); -// return mockConnection; -// }); -// when(mockPluginService.forceConnect(eq(fastHost), eq(properties))).thenReturn(mockConnection); -// -// Dialect mockDialect = Mockito.mock(Dialect.class); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// -// final ReaderFailoverHandler target = -// new ClusterAwareReaderFailoverHandler( -// mockPluginService, -// properties); -// final ReaderFailoverResult result = target.getReaderConnection(hosts); -// -// assertTrue(result.isConnected()); -// assertSame(mockConnection, result.getConnection()); -// assertEquals(hosts.get(2), result.getHost()); -// -// verify(mockPluginService, never()).setAvailability(any(), eq(HostAvailability.NOT_AVAILABLE)); -// verify(mockPluginService, times(1)) -// .setAvailability(eq(fastHost.asAliases()), eq(HostAvailability.AVAILABLE)); -// } -// -// @Test -// public void testGetReader_connectionFailure() throws SQLException { -// // odd number of connection attempts -// // first connection attempt to return fails -// // expected test result: failure to get reader -// final List hosts = defaultHosts.subList(0, 4); // 3 connection attempts (writer not attempted) -// when(mockPluginService.forceConnect(any(), eq(properties))).thenThrow(new SQLException("exception", "08S01", null)); -// -// Dialect mockDialect = Mockito.mock(Dialect.class); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// -// final int currentHostIndex = 2; -// -// final ReaderFailoverHandler target = -// new ClusterAwareReaderFailoverHandler( -// mockPluginService, -// properties); -// final ReaderFailoverResult result = target.getReaderConnection(hosts); -// -// assertFalse(result.isConnected()); -// assertNull(result.getConnection()); -// assertNull(result.getHost()); -// } -// -// @Test -// public void testGetReader_connectionAttemptsTimeout() throws SQLException { -// // connection attempts time out before they can succeed -// // first connection attempt to return times out -// // expected test result: failure to get reader -// final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) -// when(mockPluginService.forceConnect(any(), eq(properties))) -// .thenAnswer( -// (Answer) -// invocation -> { -// try { -// Thread.sleep(5000); -// } catch (InterruptedException exception) { -// // ignore -// } -// return mockConnection; -// }); -// -// Dialect mockDialect = Mockito.mock(Dialect.class); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// -// final ClusterAwareReaderFailoverHandler target = -// new ClusterAwareReaderFailoverHandler( -// mockPluginService, -// properties, -// 60000, -// 1000, -// false); -// final ReaderFailoverResult result = target.getReaderConnection(hosts); -// -// assertFalse(result.isConnected()); -// assertNull(result.getConnection()); -// assertNull(result.getHost()); -// } -// -// @Test -// public void testGetHostTuplesByPriority() { -// final List originalHosts = defaultHosts; -// originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); -// originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); -// originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); -// -// final ClusterAwareReaderFailoverHandler target = -// new ClusterAwareReaderFailoverHandler( -// mockPluginService, -// properties); -// final List hostsByPriority = target.getHostsByPriority(originalHosts); -// -// int i = 0; -// -// // expecting active readers -// while (i < hostsByPriority.size() -// && hostsByPriority.get(i).getRole() == HostRole.READER -// && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { -// i++; -// } -// -// // expecting a writer -// while (i < hostsByPriority.size() -// && hostsByPriority.get(i).getRole() == HostRole.WRITER) { -// i++; -// } -// -// // expecting down readers -// while (i < hostsByPriority.size() -// && hostsByPriority.get(i).getRole() == HostRole.READER -// && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { -// i++; -// } -// -// assertEquals(hostsByPriority.size(), i); -// } -// -// @Test -// public void testGetReaderTuplesByPriority() { -// final List originalHosts = defaultHosts; -// originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); -// originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); -// originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); -// -// Dialect mockDialect = Mockito.mock(Dialect.class); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// -// final ClusterAwareReaderFailoverHandler target = -// new ClusterAwareReaderFailoverHandler( -// mockPluginService, -// properties); -// final List hostsByPriority = target.getReaderHostsByPriority(originalHosts); -// -// int i = 0; -// -// // expecting active readers -// while (i < hostsByPriority.size() -// && hostsByPriority.get(i).getRole() == HostRole.READER -// && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { -// i++; -// } -// -// // expecting down readers -// while (i < hostsByPriority.size() -// && hostsByPriority.get(i).getRole() == HostRole.READER -// && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { -// i++; -// } -// -// assertEquals(hostsByPriority.size(), i); -// } -// -// @Test -// public void testHostFailoverStrictReaderEnabled() { -// -// final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("writer").port(1234).role(HostRole.WRITER).build(); -// final HostSpec reader = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("reader1").port(1234).role(HostRole.READER).build(); -// final List hosts = Arrays.asList(writer, reader); -// -// Dialect mockDialect = Mockito.mock(Dialect.class); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// final ClusterAwareReaderFailoverHandler target = -// new ClusterAwareReaderFailoverHandler( -// mockPluginService, -// properties, -// DEFAULT_FAILOVER_TIMEOUT, -// DEFAULT_READER_CONNECT_TIMEOUT, -// true); -// -// // The writer is included because the original writer has likely become a reader. -// List expectedHostsByPriority = Arrays.asList(reader, writer); -// -// List hostsByPriority = target.getHostsByPriority(hosts); -// assertEquals(expectedHostsByPriority, hostsByPriority); -// -// // Should pick the reader even if unavailable. The unavailable reader will be lower priority than the writer. -// reader.setAvailability(HostAvailability.NOT_AVAILABLE); -// expectedHostsByPriority = Arrays.asList(writer, reader); -// -// hostsByPriority = target.getHostsByPriority(hosts); -// assertEquals(expectedHostsByPriority, hostsByPriority); -// -// // Writer node will only be picked if it is the only node in topology; -// List expectedWriterHost = Collections.singletonList(writer); -// -// hostsByPriority = target.getHostsByPriority(Collections.singletonList(writer)); -// assertEquals(expectedWriterHost, hostsByPriority); -// } -// } +/* + * 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.failover; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.when; +import static software.amazon.jdbc.plugin.failover.ClusterAwareReaderFailoverHandler.DEFAULT_FAILOVER_TIMEOUT; +import static software.amazon.jdbc.plugin.failover.ClusterAwareReaderFailoverHandler.DEFAULT_READER_CONNECT_TIMEOUT; + +import java.sql.Connection; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.EnumSet; +import java.util.List; +import java.util.Map; +import java.util.Properties; +import java.util.Set; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.MockitoAnnotations; +import org.mockito.stubbing.Answer; +import software.amazon.jdbc.ConnectionPluginManager; +import software.amazon.jdbc.HostRole; +import software.amazon.jdbc.HostSpec; +import software.amazon.jdbc.HostSpecBuilder; +import software.amazon.jdbc.PluginService; +import software.amazon.jdbc.dialect.Dialect; +import software.amazon.jdbc.hostavailability.HostAvailability; +import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; +import software.amazon.jdbc.util.FullServicesContainer; +import software.amazon.jdbc.util.connection.ConnectionService; + +class ClusterAwareReaderFailoverHandlerTest { + @Mock FullServicesContainer mockContainer; + @Mock ConnectionService mockConnectionService; + @Mock PluginService mockPluginService; + @Mock ConnectionPluginManager mockPluginManager; + @Mock Connection mockConnection; + + private AutoCloseable closeable; + private final Properties properties = new Properties(); + private final List defaultHosts = Arrays.asList( + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("writer").port(1234).role(HostRole.WRITER).build(), + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("reader1").port(1234).role(HostRole.READER).build(), + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("reader2").port(1234).role(HostRole.READER).build(), + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("reader3").port(1234).role(HostRole.READER).build(), + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("reader4").port(1234).role(HostRole.READER).build(), + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("reader5").port(1234).role(HostRole.READER).build() + ); + + @BeforeEach + void setUp() { + closeable = MockitoAnnotations.openMocks(this); + when(mockContainer.getConnectionPluginManager()).thenReturn(mockPluginManager); + when(mockContainer.getPluginService()).thenReturn(mockPluginService); + } + + @AfterEach + void tearDown() throws Exception { + closeable.close(); + } + + @Test + public void testFailover() throws SQLException { + // original host list: [active writer, active reader, current connection (reader), active + // reader, down reader, active reader] + // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] + // connection attempts are made in pairs using the above list + // expected test result: successful connection for host at index 4 + final List hosts = defaultHosts; + final int currentHostIndex = 2; + final int successHostIndex = 4; + for (int i = 0; i < hosts.size(); i++) { + if (i != successHostIndex) { + final SQLException exception = new SQLException("exception", "08S01", null); + when(mockConnectionService.open(hosts.get(i), properties)) + .thenThrow(exception); + when(mockPluginService.isNetworkException(exception, null)).thenReturn(true); + } else { + when(mockConnectionService.open(hosts.get(i), properties)).thenReturn(mockConnection); + } + } + + when(mockPluginService.getTargetDriverDialect()).thenReturn(null); + + hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); + hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); + + final ReaderFailoverHandler target = getSpyFailoverHandler(); + final ReaderFailoverResult result = target.failover(hosts, hosts.get(currentHostIndex)); + + assertTrue(result.isConnected()); + assertSame(mockConnection, result.getConnection()); + assertEquals(hosts.get(successHostIndex), result.getHost()); + + final HostSpec successHost = hosts.get(successHostIndex); + final Map availabilityMap = target.getHostAvailabilityMap(); + Set unavailableHosts = getHostsWithGivenAvailability(availabilityMap, HostAvailability.NOT_AVAILABLE); + assertTrue(unavailableHosts.size() >= 4); + assertEquals(HostAvailability.AVAILABLE, availabilityMap.get(successHost.getHost())); + } + + private Set getHostsWithGivenAvailability( + Map availabilityMap, HostAvailability availability) { + return availabilityMap.entrySet().stream() + .filter((entry) -> availability.equals(entry.getValue())) + .map(Map.Entry::getKey) + .collect(Collectors.toSet()); + } + + @Test + public void testFailover_timeout() throws SQLException { + // original host list: [active writer, active reader, current connection (reader), active + // reader, down reader, active reader] + // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] + // connection attempts are made in pairs using the above list + // expected test result: failure to get reader since process is limited to 5s and each attempt + // to connect takes 20s + final List hosts = defaultHosts; + final int currentHostIndex = 2; + for (HostSpec host : hosts) { + when(mockConnectionService.open(host, properties)) + .thenAnswer((Answer) invocation -> { + Thread.sleep(20000); + return mockConnection; + }); + } + + hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); + hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); + + final ReaderFailoverHandler target = getSpyFailoverHandler(5000, 30000, false); + + final long startTimeNano = System.nanoTime(); + final ReaderFailoverResult result = target.failover(hosts, hosts.get(currentHostIndex)); + final long durationNano = System.nanoTime() - startTimeNano; + + assertFalse(result.isConnected()); + assertNull(result.getConnection()); + assertNull(result.getHost()); + + // 5s is a max allowed failover timeout; add 1s for inaccurate measurements + assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); + } + + private ClusterAwareReaderFailoverHandler getSpyFailoverHandler() { + ClusterAwareReaderFailoverHandler handler = + spy(new ClusterAwareReaderFailoverHandler(mockContainer, mockConnectionService, properties)); + doReturn(mockPluginService).when(handler).getNewPluginService(); + return handler; + } + + private ClusterAwareReaderFailoverHandler getSpyFailoverHandler( + int maxFailoverTimeoutMs, int timeoutMs, boolean isStrictReaderRequired) { + ClusterAwareReaderFailoverHandler handler = new ClusterAwareReaderFailoverHandler( + mockContainer, mockConnectionService, properties, maxFailoverTimeoutMs, timeoutMs, isStrictReaderRequired); + ClusterAwareReaderFailoverHandler spyHandler = spy(handler); + doReturn(mockPluginService).when(spyHandler).getNewPluginService(); + return spyHandler; + } + + @Test + public void testFailover_nullOrEmptyHostList() throws SQLException { + final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); + final HostSpec currentHost = + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("writer").port(1234).build(); + + ReaderFailoverResult result = target.failover(null, currentHost); + assertFalse(result.isConnected()); + assertNull(result.getConnection()); + assertNull(result.getHost()); + + final List hosts = new ArrayList<>(); + result = target.failover(hosts, currentHost); + assertFalse(result.isConnected()); + assertNull(result.getConnection()); + assertNull(result.getHost()); + } + + @Test + public void testGetReader_connectionSuccess() throws SQLException { + // even number of connection attempts + // first connection attempt to return succeeds, second attempt cancelled + // expected test result: successful connection for host at index 2 + final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) + final HostSpec slowHost = hosts.get(1); + final HostSpec fastHost = hosts.get(2); + when(mockConnectionService.open(slowHost, properties)) + .thenAnswer( + (Answer) + invocation -> { + Thread.sleep(20000); + return mockConnection; + }); + when(mockConnectionService.open(eq(fastHost), eq(properties))).thenReturn(mockConnection); + + Dialect mockDialect = Mockito.mock(Dialect.class); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + when(mockPluginService.getDialect()).thenReturn(mockDialect); + + final ReaderFailoverHandler target = getSpyFailoverHandler(); + final ReaderFailoverResult result = target.getReaderConnection(hosts); + + assertTrue(result.isConnected()); + assertSame(mockConnection, result.getConnection()); + assertEquals(hosts.get(2), result.getHost()); + + Map availabilityMap = target.getHostAvailabilityMap(); + assertTrue(getHostsWithGivenAvailability(availabilityMap, HostAvailability.NOT_AVAILABLE).isEmpty()); + assertEquals(HostAvailability.AVAILABLE, availabilityMap.get(fastHost.getHost())); + } + + @Test + public void testGetReader_connectionFailure() throws SQLException { + // odd number of connection attempts + // first connection attempt to return fails + // expected test result: failure to get reader + final List hosts = defaultHosts.subList(0, 4); // 3 connection attempts (writer not attempted) + when(mockConnectionService.open(any(), eq(properties))).thenThrow(new SQLException("exception", "08S01", null)); + + Dialect mockDialect = Mockito.mock(Dialect.class); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + when(mockPluginService.getDialect()).thenReturn(mockDialect); + + final ReaderFailoverHandler target = getSpyFailoverHandler(); + final ReaderFailoverResult result = target.getReaderConnection(hosts); + + assertFalse(result.isConnected()); + assertNull(result.getConnection()); + assertNull(result.getHost()); + } + + @Test + public void testGetReader_connectionAttemptsTimeout() throws SQLException { + // connection attempts time out before they can succeed + // first connection attempt to return times out + // expected test result: failure to get reader + final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) + when(mockConnectionService.open(any(), eq(properties))) + .thenAnswer( + (Answer) + invocation -> { + try { + Thread.sleep(5000); + } catch (InterruptedException exception) { + // ignore + } + return mockConnection; + }); + + Dialect mockDialect = Mockito.mock(Dialect.class); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + when(mockPluginService.getDialect()).thenReturn(mockDialect); + + final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(60000, 1000, false); + final ReaderFailoverResult result = target.getReaderConnection(hosts); + + assertFalse(result.isConnected()); + assertNull(result.getConnection()); + assertNull(result.getHost()); + } + + @Test + public void testGetHostTuplesByPriority() { + final List originalHosts = defaultHosts; + originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); + originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); + originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); + + final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); + final List hostsByPriority = target.getHostsByPriority(originalHosts); + + int i = 0; + + // expecting active readers + while (i < hostsByPriority.size() + && hostsByPriority.get(i).getRole() == HostRole.READER + && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { + i++; + } + + // expecting a writer + while (i < hostsByPriority.size() + && hostsByPriority.get(i).getRole() == HostRole.WRITER) { + i++; + } + + // expecting down readers + while (i < hostsByPriority.size() + && hostsByPriority.get(i).getRole() == HostRole.READER + && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { + i++; + } + + assertEquals(hostsByPriority.size(), i); + } + + @Test + public void testGetReaderTuplesByPriority() { + final List originalHosts = defaultHosts; + originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); + originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); + originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); + + Dialect mockDialect = Mockito.mock(Dialect.class); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + when(mockPluginService.getDialect()).thenReturn(mockDialect); + + final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); + final List hostsByPriority = target.getReaderHostsByPriority(originalHosts); + + int i = 0; + + // expecting active readers + while (i < hostsByPriority.size() + && hostsByPriority.get(i).getRole() == HostRole.READER + && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { + i++; + } + + // expecting down readers + while (i < hostsByPriority.size() + && hostsByPriority.get(i).getRole() == HostRole.READER + && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { + i++; + } + + assertEquals(hostsByPriority.size(), i); + } + + @Test + public void testHostFailoverStrictReaderEnabled() { + + final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("writer").port(1234).role(HostRole.WRITER).build(); + final HostSpec reader = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("reader1").port(1234).role(HostRole.READER).build(); + final List hosts = Arrays.asList(writer, reader); + + Dialect mockDialect = Mockito.mock(Dialect.class); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + when(mockPluginService.getDialect()).thenReturn(mockDialect); + + final ClusterAwareReaderFailoverHandler target = + getSpyFailoverHandler(DEFAULT_FAILOVER_TIMEOUT, DEFAULT_READER_CONNECT_TIMEOUT, true); + + // The writer is included because the original writer has likely become a reader. + List expectedHostsByPriority = Arrays.asList(reader, writer); + + List hostsByPriority = target.getHostsByPriority(hosts); + assertEquals(expectedHostsByPriority, hostsByPriority); + + // Should pick the reader even if unavailable. The unavailable reader will be lower priority than the writer. + reader.setAvailability(HostAvailability.NOT_AVAILABLE); + expectedHostsByPriority = Arrays.asList(writer, reader); + + hostsByPriority = target.getHostsByPriority(hosts); + assertEquals(expectedHostsByPriority, hostsByPriority); + + // Writer node will only be picked if it is the only node in topology; + List expectedWriterHost = Collections.singletonList(writer); + + hostsByPriority = target.getHostsByPriority(Collections.singletonList(writer)); + assertEquals(expectedWriterHost, hostsByPriority); + } +} From 370d73fbf327ed13f7e3942cec980b11333acec7 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Fri, 22 Aug 2025 15:06:33 -0700 Subject: [PATCH 14/42] WriterFailoverHandler tests passing --- .../ClusterAwareWriterFailoverHandler.java | 2 +- ...ClusterAwareWriterFailoverHandlerTest.java | 781 +++++++++--------- 2 files changed, 374 insertions(+), 409 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java index 8533b1055..1921f3238 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java @@ -189,7 +189,7 @@ private void submitTasks( executorService.shutdown(); } - private PluginService getNewPluginService() { + protected PluginService getNewPluginService() { // Each task should get its own PluginService since they execute concurrently and PluginService was not designed to // be thread-safe. return new PartialPluginService( diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java index 5f302209a..1ad394010 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java @@ -1,408 +1,373 @@ -// /* -// * 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.failover; -// -// import static org.junit.jupiter.api.Assertions.assertEquals; -// import static org.junit.jupiter.api.Assertions.assertFalse; -// import static org.junit.jupiter.api.Assertions.assertSame; -// import static org.junit.jupiter.api.Assertions.assertTrue; -// import static org.mockito.ArgumentMatchers.any; -// import static org.mockito.ArgumentMatchers.eq; -// import static org.mockito.ArgumentMatchers.refEq; -// import static org.mockito.Mockito.atLeastOnce; -// import static org.mockito.Mockito.times; -// import static org.mockito.Mockito.verify; -// import static org.mockito.Mockito.when; -// -// import java.sql.Connection; -// import java.sql.SQLException; -// import java.util.Arrays; -// import java.util.EnumSet; -// import java.util.List; -// import java.util.Properties; -// import java.util.concurrent.TimeUnit; -// import org.junit.jupiter.api.AfterEach; -// import org.junit.jupiter.api.BeforeEach; -// import org.junit.jupiter.api.Test; -// import org.mockito.ArgumentMatchers; -// import org.mockito.InOrder; -// import org.mockito.Mock; -// import org.mockito.Mockito; -// import org.mockito.MockitoAnnotations; -// import org.mockito.stubbing.Answer; -// import software.amazon.jdbc.HostSpec; -// import software.amazon.jdbc.HostSpecBuilder; -// import software.amazon.jdbc.PluginService; -// import software.amazon.jdbc.dialect.Dialect; -// import software.amazon.jdbc.hostavailability.HostAvailability; -// import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; -// -// class ClusterAwareWriterFailoverHandlerTest { -// -// @Mock PluginService mockPluginService; -// @Mock Connection mockConnection; -// @Mock ReaderFailoverHandler mockReaderFailover; -// @Mock Connection mockWriterConnection; -// @Mock Connection mockNewWriterConnection; -// @Mock Connection mockReaderAConnection; -// @Mock Connection mockReaderBConnection; -// @Mock Dialect mockDialect; -// -// private AutoCloseable closeable; -// private final Properties properties = new Properties(); -// private final HostSpec newWriterHost = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("new-writer-host").build(); -// private final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("writer-host").build(); -// private final HostSpec readerA = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("reader-a-host").build(); -// private final HostSpec readerB = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("reader-b-host").build(); -// private final List topology = Arrays.asList(writer, readerA, readerB); -// private final List newTopology = Arrays.asList(newWriterHost, readerA, readerB); -// -// @BeforeEach -// void setUp() { -// closeable = MockitoAnnotations.openMocks(this); -// writer.addAlias("writer-host"); -// newWriterHost.addAlias("new-writer-host"); -// readerA.addAlias("reader-a-host"); -// readerB.addAlias("reader-b-host"); -// } -// -// @AfterEach -// void tearDown() throws Exception { -// closeable.close(); -// } -// -// @Test -// public void testReconnectToWriter_taskBReaderException() throws SQLException { -// when(mockPluginService.forceConnect(refEq(writer), eq(properties))).thenReturn(mockConnection); -// when(mockPluginService.forceConnect(refEq(readerA), eq(properties))).thenThrow(SQLException.class); -// when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenThrow(SQLException.class); -// -// when(mockPluginService.getAllHosts()).thenReturn(topology); -// -// when(mockReaderFailover.getReaderConnection(ArgumentMatchers.anyList())).thenThrow(SQLException.class); -// -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// -// final ClusterAwareWriterFailoverHandler target = -// new ClusterAwareWriterFailoverHandler( -// mockPluginService, -// mockReaderFailover, -// properties, -// 5000, -// 2000, -// 2000); -// final WriterFailoverResult result = target.failover(topology); -// -// assertTrue(result.isConnected()); -// assertFalse(result.isNewHost()); -// assertSame(result.getNewConnection(), mockConnection); -// -// final InOrder inOrder = Mockito.inOrder(mockPluginService); -// inOrder.verify(mockPluginService).setAvailability(eq(writer.asAliases()), eq(HostAvailability.AVAILABLE)); -// } -// -// /** -// * Verify that writer failover handler can re-connect to a current writer node. -// * -// *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. -// * TaskA: successfully re-connect to initial writer; return new connection. -// * TaskB: successfully connect to readerA and then new writer, but it takes more time than taskA. -// * Expected test result: new connection by taskA. -// */ -// @Test -// public void testReconnectToWriter_SlowReaderA() throws SQLException { -// when(mockPluginService.forceConnect(refEq(writer), eq(properties))).thenReturn(mockWriterConnection); -// when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenThrow(SQLException.class); -// when(mockPluginService.forceConnect(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); -// when(mockPluginService.getAllHosts()).thenReturn(topology).thenReturn(newTopology); -// -// when(mockReaderFailover.getReaderConnection(ArgumentMatchers.anyList())) -// .thenAnswer( -// (Answer) -// invocation -> { -// Thread.sleep(5000); -// return new ReaderFailoverResult(mockReaderAConnection, readerA, true); -// }); -// -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// -// final ClusterAwareWriterFailoverHandler target = -// new ClusterAwareWriterFailoverHandler( -// mockPluginService, -// mockReaderFailover, -// properties, -// 60000, -// 5000, -// 5000); -// final WriterFailoverResult result = target.failover(topology); -// -// assertTrue(result.isConnected()); -// assertFalse(result.isNewHost()); -// assertSame(result.getNewConnection(), mockWriterConnection); -// -// final InOrder inOrder = Mockito.inOrder(mockPluginService); -// inOrder.verify(mockPluginService).setAvailability(eq(writer.asAliases()), eq(HostAvailability.AVAILABLE)); -// } -// -// /** -// * Verify that writer failover handler can re-connect to a current writer node. -// * -// *

Topology: no changes. -// * TaskA: successfully re-connect to writer; return new connection. -// * TaskB: successfully connect to readerA and retrieve topology, but latest writer is not new (defer to taskA). -// * Expected test result: new connection by taskA. -// */ -// @Test -// public void testReconnectToWriter_taskBDefers() throws SQLException { -// when(mockPluginService.forceConnect(refEq(writer), eq(properties))) -// .thenAnswer( -// (Answer) -// invocation -> { -// Thread.sleep(5000); -// return mockWriterConnection; -// }); -// when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenThrow(SQLException.class); -// -// when(mockPluginService.getAllHosts()).thenReturn(topology); -// -// when(mockReaderFailover.getReaderConnection(ArgumentMatchers.anyList())) -// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); -// -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// -// final ClusterAwareWriterFailoverHandler target = -// new ClusterAwareWriterFailoverHandler( -// mockPluginService, -// mockReaderFailover, -// properties, -// 60000, -// 2000, -// 2000); -// final WriterFailoverResult result = target.failover(topology); -// -// assertTrue(result.isConnected()); -// assertFalse(result.isNewHost()); -// assertSame(result.getNewConnection(), mockWriterConnection); -// -// final InOrder inOrder = Mockito.inOrder(mockPluginService); -// inOrder.verify(mockPluginService).setAvailability(eq(writer.asAliases()), eq(HostAvailability.AVAILABLE)); -// } -// -// /** -// * Verify that writer failover handler can re-connect to a new writer node. -// * -// *

Topology: changes to [new-writer, reader-A, reader-B] for taskB, taskA sees no changes. -// * taskA: successfully re-connect to writer; return connection to initial writer, but it takes more -// * time than taskB. -// * TaskB: successfully connect to readerA and then to new-writer. -// * Expected test result: new connection to writer by taskB. -// */ -// @Test -// public void testConnectToReaderA_SlowWriter() throws SQLException { -// when(mockPluginService.forceConnect(refEq(writer), eq(properties))) -// .thenAnswer( -// (Answer) -// invocation -> { -// Thread.sleep(5000); -// return mockWriterConnection; -// }); -// when(mockPluginService.forceConnect(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); -// when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); -// when(mockPluginService.forceConnect(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); -// -// when(mockPluginService.getAllHosts()).thenReturn(newTopology); -// -// when(mockReaderFailover.getReaderConnection(ArgumentMatchers.anyList())) -// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); -// -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// -// final ClusterAwareWriterFailoverHandler target = -// new ClusterAwareWriterFailoverHandler( -// mockPluginService, -// mockReaderFailover, -// properties, -// 60000, -// 5000, -// 5000); -// final WriterFailoverResult result = target.failover(topology); -// -// assertTrue(result.isConnected()); -// assertTrue(result.isNewHost()); -// assertSame(result.getNewConnection(), mockNewWriterConnection); -// assertEquals(3, result.getTopology().size()); -// assertEquals("new-writer-host", result.getTopology().get(0).getHost()); -// -// verify(mockPluginService, times(1)).setAvailability(eq(newWriterHost.asAliases()), eq(HostAvailability.AVAILABLE)); -// } -// -// /** -// * Verify that writer failover handler can re-connect to a new writer node. -// * -// *

Topology: changes to [new-writer, initial-writer, reader-A, reader-B]. -// * TaskA: successfully reconnect, but initial-writer is now a reader (defer to taskB). -// * TaskB: successfully connect to readerA and then to new-writer. -// * Expected test result: new connection to writer by taskB. -// */ -// @Test -// public void testConnectToReaderA_taskADefers() throws SQLException { -// when(mockPluginService.forceConnect(writer, properties)).thenReturn(mockConnection); -// when(mockPluginService.forceConnect(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); -// when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); -// when(mockPluginService.forceConnect(refEq(newWriterHost), eq(properties))) -// .thenAnswer( -// (Answer) -// invocation -> { -// Thread.sleep(5000); -// return mockNewWriterConnection; -// }); -// -// final List newTopology = Arrays.asList(newWriterHost, writer, readerA, readerB); -// when(mockPluginService.getAllHosts()).thenReturn(newTopology); -// -// when(mockReaderFailover.getReaderConnection(ArgumentMatchers.anyList())) -// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); -// -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// -// final ClusterAwareWriterFailoverHandler target = -// new ClusterAwareWriterFailoverHandler( -// mockPluginService, -// mockReaderFailover, -// properties, -// 60000, -// 5000, -// 5000); -// final WriterFailoverResult result = target.failover(topology); -// -// assertTrue(result.isConnected()); -// assertTrue(result.isNewHost()); -// assertSame(result.getNewConnection(), mockNewWriterConnection); -// assertEquals(4, result.getTopology().size()); -// assertEquals("new-writer-host", result.getTopology().get(0).getHost()); -// -// verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); -// verify(mockPluginService, times(1)).setAvailability(eq(newWriterHost.asAliases()), eq(HostAvailability.AVAILABLE)); -// } -// -// /** -// * Verify that writer failover handler fails to re-connect to any writer node. -// * -// *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. -// * TaskA: fail to re-connect to writer due to failover timeout. -// * TaskB: successfully connect to readerA and then fail to connect to writer due to failover timeout. -// * Expected test result: no connection. -// */ -// @Test -// public void testFailedToConnect_failoverTimeout() throws SQLException { -// when(mockPluginService.forceConnect(refEq(writer), eq(properties))) -// .thenAnswer( -// (Answer) -// invocation -> { -// Thread.sleep(30000); -// return mockWriterConnection; -// }); -// when(mockPluginService.forceConnect(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); -// when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); -// when(mockPluginService.forceConnect(refEq(newWriterHost), eq(properties))) -// .thenAnswer( -// (Answer) -// invocation -> { -// Thread.sleep(30000); -// return mockNewWriterConnection; -// }); -// when(mockPluginService.getAllHosts()).thenReturn(newTopology); -// -// when(mockReaderFailover.getReaderConnection(ArgumentMatchers.anyList())) -// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); -// -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// -// final ClusterAwareWriterFailoverHandler target = -// new ClusterAwareWriterFailoverHandler( -// mockPluginService, -// mockReaderFailover, -// properties, -// 5000, -// 2000, -// 2000); -// -// final long startTimeNano = System.nanoTime(); -// final WriterFailoverResult result = target.failover(topology); -// final long durationNano = System.nanoTime() - startTimeNano; -// -// assertFalse(result.isConnected()); -// assertFalse(result.isNewHost()); -// -// verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); -// -// // 5s is a max allowed failover timeout; add 1s for inaccurate measurements -// assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); -// } -// -// /** -// * Verify that writer failover handler fails to re-connect to any writer node. -// * -// *

Topology: changes to [new-writer, reader-A, reader-B] for taskB. -// * TaskA: fail to re-connect to writer due to exception. -// * TaskB: successfully connect to readerA and then fail to connect to writer due to exception. -// * Expected test result: no connection. -// */ -// @Test -// public void testFailedToConnect_taskAException_taskBWriterException() throws SQLException { -// final SQLException exception = new SQLException("exception", "08S01", null); -// when(mockPluginService.forceConnect(refEq(writer), eq(properties))).thenThrow(exception); -// when(mockPluginService.forceConnect(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); -// when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); -// when(mockPluginService.forceConnect(refEq(newWriterHost), eq(properties))).thenThrow(exception); -// when(mockPluginService.isNetworkException(eq(exception), any())).thenReturn(true); -// -// when(mockPluginService.getAllHosts()).thenReturn(newTopology); -// -// when(mockReaderFailover.getReaderConnection(ArgumentMatchers.anyList())) -// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); -// -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// -// final ClusterAwareWriterFailoverHandler target = -// new ClusterAwareWriterFailoverHandler( -// mockPluginService, -// mockReaderFailover, -// properties, -// 5000, -// 2000, -// 2000); -// final WriterFailoverResult result = target.failover(topology); -// -// assertFalse(result.isConnected()); -// assertFalse(result.isNewHost()); -// -// verify(mockPluginService, atLeastOnce()) -// .setAvailability(eq(newWriterHost.asAliases()), eq(HostAvailability.NOT_AVAILABLE)); -// } -// } +/* + * 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.failover; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.refEq; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.sql.Connection; +import java.sql.SQLException; +import java.util.Arrays; +import java.util.EnumSet; +import java.util.List; +import java.util.Properties; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentMatchers; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import org.mockito.stubbing.Answer; +import software.amazon.jdbc.HostSpec; +import software.amazon.jdbc.HostSpecBuilder; +import software.amazon.jdbc.PluginService; +import software.amazon.jdbc.dialect.Dialect; +import software.amazon.jdbc.hostavailability.HostAvailability; +import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; +import software.amazon.jdbc.util.FullServicesContainer; +import software.amazon.jdbc.util.connection.ConnectionService; + +class ClusterAwareWriterFailoverHandlerTest { + @Mock FullServicesContainer mockContainer; + @Mock ConnectionService mockConnectionService; + @Mock PluginService mockPluginService; + @Mock Connection mockConnection; + @Mock ReaderFailoverHandler mockReaderFailoverHandler; + @Mock Connection mockWriterConnection; + @Mock Connection mockNewWriterConnection; + @Mock Connection mockReaderAConnection; + @Mock Connection mockReaderBConnection; + @Mock Dialect mockDialect; + + private AutoCloseable closeable; + private final Properties properties = new Properties(); + private final HostSpec newWriterHost = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("new-writer-host").build(); + private final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("writer-host").build(); + private final HostSpec readerA = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("reader-a-host").build(); + private final HostSpec readerB = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("reader-b-host").build(); + private final List topology = Arrays.asList(writer, readerA, readerB); + private final List newTopology = Arrays.asList(newWriterHost, readerA, readerB); + + @BeforeEach + void setUp() { + closeable = MockitoAnnotations.openMocks(this); + when(mockContainer.getPluginService()).thenReturn(mockPluginService); + writer.addAlias("writer-host"); + newWriterHost.addAlias("new-writer-host"); + readerA.addAlias("reader-a-host"); + readerB.addAlias("reader-b-host"); + } + + @AfterEach + void tearDown() throws Exception { + closeable.close(); + } + + @Test + public void testReconnectToWriter_taskBReaderException() throws SQLException { + when(mockConnectionService.open(refEq(writer), eq(properties))).thenReturn(mockConnection); + when(mockConnectionService.open(refEq(readerA), eq(properties))).thenThrow(SQLException.class); + when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); + + when(mockPluginService.getAllHosts()).thenReturn(topology); + + when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())).thenThrow(SQLException.class); + + when(mockPluginService.getDialect()).thenReturn(mockDialect); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + + final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); + final WriterFailoverResult result = target.failover(topology); + + assertTrue(result.isConnected()); + assertFalse(result.isNewHost()); + assertSame(result.getNewConnection(), mockConnection); + + assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); + } + + private ClusterAwareWriterFailoverHandler getSpyFailoverHandler( + final int failoverTimeoutMs, + final int readTopologyIntervalMs, + final int reconnectWriterIntervalMs) { + ClusterAwareWriterFailoverHandler handler = new ClusterAwareWriterFailoverHandler( + mockContainer, + mockConnectionService, + mockReaderFailoverHandler, + properties, + failoverTimeoutMs, + readTopologyIntervalMs, + reconnectWriterIntervalMs); + + ClusterAwareWriterFailoverHandler spyHandler = spy(handler); + doReturn(mockPluginService).when(spyHandler).getNewPluginService(); + return spyHandler; + } + + /** + * Verify that writer failover handler can re-connect to a current writer node. + * + *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. + * TaskA: successfully re-connect to initial writer; return new connection. + * TaskB: successfully connect to readerA and then new writer, but it takes more time than taskA. + * Expected test result: new connection by taskA. + */ + @Test + public void testReconnectToWriter_SlowReaderA() throws SQLException { + when(mockConnectionService.open(refEq(writer), eq(properties))).thenReturn(mockWriterConnection); + when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); + when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); + when(mockPluginService.getAllHosts()).thenReturn(topology).thenReturn(newTopology); + + when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) + .thenAnswer( + (Answer) + invocation -> { + Thread.sleep(5000); + return new ReaderFailoverResult(mockReaderAConnection, readerA, true); + }); + + when(mockPluginService.getDialect()).thenReturn(mockDialect); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + + final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); + final WriterFailoverResult result = target.failover(topology); + + assertTrue(result.isConnected()); + assertFalse(result.isNewHost()); + assertSame(result.getNewConnection(), mockWriterConnection); + assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); + } + + /** + * Verify that writer failover handler can re-connect to a current writer node. + * + *

Topology: no changes. + * TaskA: successfully re-connect to writer; return new connection. + * TaskB: successfully connect to readerA and retrieve topology, but latest writer is not new (defer to taskA). + * Expected test result: new connection by taskA. + */ + @Test + public void testReconnectToWriter_taskBDefers() throws SQLException { + when(mockConnectionService.open(refEq(writer), eq(properties))) + .thenAnswer( + (Answer) + invocation -> { + Thread.sleep(5000); + return mockWriterConnection; + }); + when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); + + when(mockPluginService.getAllHosts()).thenReturn(topology); + + when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) + .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); + + when(mockPluginService.getDialect()).thenReturn(mockDialect); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + + final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 2000, 2000); + final WriterFailoverResult result = target.failover(topology); + + assertTrue(result.isConnected()); + assertFalse(result.isNewHost()); + assertSame(result.getNewConnection(), mockWriterConnection); + assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); + } + + /** + * Verify that writer failover handler can re-connect to a new writer node. + * + *

Topology: changes to [new-writer, reader-A, reader-B] for taskB, taskA sees no changes. + * taskA: successfully re-connect to writer; return connection to initial writer, but it takes more + * time than taskB. + * TaskB: successfully connect to readerA and then to new-writer. + * Expected test result: new connection to writer by taskB. + */ + @Test + public void testConnectToReaderA_SlowWriter() throws SQLException { + when(mockConnectionService.open(refEq(writer), eq(properties))) + .thenAnswer( + (Answer) + invocation -> { + Thread.sleep(5000); + return mockWriterConnection; + }); + when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); + when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); + when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); + + when(mockPluginService.getAllHosts()).thenReturn(newTopology); + + when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) + .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); + + when(mockPluginService.getDialect()).thenReturn(mockDialect); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + + final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); + final WriterFailoverResult result = target.failover(topology); + + assertTrue(result.isConnected()); + assertTrue(result.isNewHost()); + assertSame(result.getNewConnection(), mockNewWriterConnection); + assertEquals(3, result.getTopology().size()); + assertEquals("new-writer-host", result.getTopology().get(0).getHost()); + assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); + } + + /** + * Verify that writer failover handler can re-connect to a new writer node. + * + *

Topology: changes to [new-writer, initial-writer, reader-A, reader-B]. + * TaskA: successfully reconnect, but initial-writer is now a reader (defer to taskB). + * TaskB: successfully connect to readerA and then to new-writer. + * Expected test result: new connection to writer by taskB. + */ + @Test + public void testConnectToReaderA_taskADefers() throws SQLException { + when(mockConnectionService.open(writer, properties)).thenReturn(mockConnection); + when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); + when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); + when(mockConnectionService.open(refEq(newWriterHost), eq(properties))) + .thenAnswer( + (Answer) + invocation -> { + Thread.sleep(5000); + return mockNewWriterConnection; + }); + + final List newTopology = Arrays.asList(newWriterHost, writer, readerA, readerB); + when(mockPluginService.getAllHosts()).thenReturn(newTopology); + + when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) + .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); + + when(mockPluginService.getDialect()).thenReturn(mockDialect); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + + final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); + final WriterFailoverResult result = target.failover(topology); + + assertTrue(result.isConnected()); + assertTrue(result.isNewHost()); + assertSame(result.getNewConnection(), mockNewWriterConnection); + assertEquals(4, result.getTopology().size()); + assertEquals("new-writer-host", result.getTopology().get(0).getHost()); + + verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); + assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); + } + + /** + * Verify that writer failover handler fails to re-connect to any writer node. + * + *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. + * TaskA: fail to re-connect to writer due to failover timeout. + * TaskB: successfully connect to readerA and then fail to connect to writer due to failover timeout. + * Expected test result: no connection. + */ + @Test + public void testFailedToConnect_failoverTimeout() throws SQLException { + when(mockConnectionService.open(refEq(writer), eq(properties))) + .thenAnswer( + (Answer) + invocation -> { + Thread.sleep(30000); + return mockWriterConnection; + }); + when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); + when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); + when(mockConnectionService.open(refEq(newWriterHost), eq(properties))) + .thenAnswer( + (Answer) + invocation -> { + Thread.sleep(30000); + return mockNewWriterConnection; + }); + when(mockPluginService.getAllHosts()).thenReturn(newTopology); + + when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) + .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); + + when(mockPluginService.getDialect()).thenReturn(mockDialect); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + + final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); + + final long startTimeNano = System.nanoTime(); + final WriterFailoverResult result = target.failover(topology); + final long durationNano = System.nanoTime() - startTimeNano; + + assertFalse(result.isConnected()); + assertFalse(result.isNewHost()); + + verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); + + // 5s is a max allowed failover timeout; add 1s for inaccurate measurements + assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); + } + + /** + * Verify that writer failover handler fails to re-connect to any writer node. + * + *

Topology: changes to [new-writer, reader-A, reader-B] for taskB. + * TaskA: fail to re-connect to writer due to exception. + * TaskB: successfully connect to readerA and then fail to connect to writer due to exception. + * Expected test result: no connection. + */ + @Test + public void testFailedToConnect_taskAException_taskBWriterException() throws SQLException { + final SQLException exception = new SQLException("exception", "08S01", null); + when(mockConnectionService.open(refEq(writer), eq(properties))).thenThrow(exception); + when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); + when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); + when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenThrow(exception); + when(mockPluginService.isNetworkException(eq(exception), any())).thenReturn(true); + + when(mockPluginService.getAllHosts()).thenReturn(newTopology); + + when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) + .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); + + when(mockPluginService.getDialect()).thenReturn(mockDialect); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + + final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); + final WriterFailoverResult result = target.failover(topology); + + assertFalse(result.isConnected()); + assertFalse(result.isNewHost()); + + assertEquals(HostAvailability.NOT_AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); + } +} From 435f71a317d817d470d969b2503d199e995a62e1 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Fri, 22 Aug 2025 15:19:06 -0700 Subject: [PATCH 15/42] FailoverConnectionPluginTests passing --- .../failover/FailoverConnectionPlugin.java | 32 +- .../FailoverConnectionPluginTest.java | 896 +++++++++--------- 2 files changed, 472 insertions(+), 456 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java index 285dc0eac..8bbdbd5fd 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java @@ -621,20 +621,7 @@ protected void dealWithIllegalStateException( protected void failover(final HostSpec failedHost) throws SQLException { this.pluginService.setAvailability(failedHost.asAliases(), HostAvailability.NOT_AVAILABLE); if (this.connectionService == null) { - TargetDriverHelper helper = new TargetDriverHelper(); - java.sql.Driver driver = helper.getTargetDriver(this.pluginService.getOriginalUrl(), properties); - final ConnectionProvider defaultConnectionProvider = new DriverConnectionProvider(driver); - this.connectionService = new ConnectionServiceImpl( - servicesContainer.getStorageService(), - servicesContainer.getMonitorService(), - servicesContainer.getTelemetryFactory(), - defaultConnectionProvider, - this.pluginService.getOriginalUrl(), - this.pluginService.getDriverProtocol(), - this.pluginService.getTargetDriverDialect(), - this.pluginService.getDialect(), - properties - ); + this.connectionService = getConnectionService(); } if (this.failoverMode == FailoverMode.STRICT_WRITER) { @@ -644,6 +631,23 @@ protected void failover(final HostSpec failedHost) throws SQLException { } } + protected ConnectionService getConnectionService() throws SQLException { + TargetDriverHelper helper = new TargetDriverHelper(); + java.sql.Driver driver = helper.getTargetDriver(this.pluginService.getOriginalUrl(), properties); + final ConnectionProvider defaultConnectionProvider = new DriverConnectionProvider(driver); + return new ConnectionServiceImpl( + servicesContainer.getStorageService(), + servicesContainer.getMonitorService(), + servicesContainer.getTelemetryFactory(), + defaultConnectionProvider, + this.pluginService.getOriginalUrl(), + this.pluginService.getDriverProtocol(), + this.pluginService.getTargetDriverDialect(), + this.pluginService.getDialect(), + properties + ); + } + protected void failoverReader(final HostSpec failedHostSpec) throws SQLException { TelemetryFactory telemetryFactory = this.pluginService.getTelemetryFactory(); TelemetryContext telemetryContext = telemetryFactory.openTelemetryContext( diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java index 33160333f..adf1cb0e6 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java @@ -1,442 +1,454 @@ -// /* -// * 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.failover; -// -// import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; -// import static org.junit.jupiter.api.Assertions.assertEquals; -// import static org.junit.jupiter.api.Assertions.assertThrows; -// import static org.mockito.ArgumentMatchers.any; -// import static org.mockito.ArgumentMatchers.anyString; -// import static org.mockito.ArgumentMatchers.eq; -// import static org.mockito.Mockito.atLeastOnce; -// import static org.mockito.Mockito.doNothing; -// import static org.mockito.Mockito.doThrow; -// import static org.mockito.Mockito.never; -// import static org.mockito.Mockito.spy; -// import static org.mockito.Mockito.times; -// import static org.mockito.Mockito.verify; -// import static org.mockito.Mockito.when; -// -// import java.sql.Connection; -// import java.sql.ResultSet; -// import java.sql.SQLException; -// import java.util.Arrays; -// import java.util.Collections; -// import java.util.EnumSet; -// import java.util.HashMap; -// import java.util.HashSet; -// import java.util.List; -// import java.util.Map; -// import java.util.Properties; -// import org.junit.jupiter.api.AfterEach; -// import org.junit.jupiter.api.BeforeEach; -// import org.junit.jupiter.api.Test; -// import org.junit.jupiter.params.ParameterizedTest; -// import org.junit.jupiter.params.provider.ValueSource; -// import org.mockito.Mock; -// import org.mockito.MockitoAnnotations; -// import software.amazon.jdbc.HostListProviderService; -// import software.amazon.jdbc.HostRole; -// import software.amazon.jdbc.HostSpec; -// import software.amazon.jdbc.HostSpecBuilder; -// import software.amazon.jdbc.JdbcCallable; -// import software.amazon.jdbc.NodeChangeOptions; -// import software.amazon.jdbc.PluginService; -// import software.amazon.jdbc.hostavailability.HostAvailability; -// import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; -// import software.amazon.jdbc.hostlistprovider.AuroraHostListProvider; -// import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; -// import software.amazon.jdbc.util.RdsUrlType; -// import software.amazon.jdbc.util.SqlState; -// import software.amazon.jdbc.util.telemetry.GaugeCallable; -// import software.amazon.jdbc.util.telemetry.TelemetryContext; -// import software.amazon.jdbc.util.telemetry.TelemetryCounter; -// import software.amazon.jdbc.util.telemetry.TelemetryFactory; -// import software.amazon.jdbc.util.telemetry.TelemetryGauge; -// -// class FailoverConnectionPluginTest { -// -// private static final Class MONITOR_METHOD_INVOKE_ON = Connection.class; -// private static final String MONITOR_METHOD_NAME = "Connection.executeQuery"; -// private static final Object[] EMPTY_ARGS = {}; -// private final List defaultHosts = Arrays.asList( -// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("writer").port(1234).role(HostRole.WRITER).build(), -// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("reader1").port(1234).role(HostRole.READER).build()); -// -// @Mock PluginService mockPluginService; -// @Mock Connection mockConnection; -// @Mock HostSpec mockHostSpec; -// @Mock HostListProviderService mockHostListProviderService; -// @Mock AuroraHostListProvider mockHostListProvider; -// @Mock JdbcCallable mockInitHostProviderFunc; -// @Mock ReaderFailoverHandler mockReaderFailoverHandler; -// @Mock WriterFailoverHandler mockWriterFailoverHandler; -// @Mock ReaderFailoverResult mockReaderResult; -// @Mock WriterFailoverResult mockWriterResult; -// @Mock JdbcCallable mockSqlFunction; -// @Mock private TelemetryFactory mockTelemetryFactory; -// @Mock TelemetryContext mockTelemetryContext; -// @Mock TelemetryCounter mockTelemetryCounter; -// @Mock TelemetryGauge mockTelemetryGauge; -// @Mock TargetDriverDialect mockTargetDriverDialect; -// -// -// private final Properties properties = new Properties(); -// private FailoverConnectionPlugin plugin; -// private AutoCloseable closeable; -// -// @AfterEach -// void cleanUp() throws Exception { -// closeable.close(); -// } -// -// @BeforeEach -// void init() throws SQLException { -// closeable = MockitoAnnotations.openMocks(this); -// -// when(mockPluginService.getHostListProvider()).thenReturn(mockHostListProvider); -// when(mockHostListProvider.getRdsUrlType()).thenReturn(RdsUrlType.RDS_WRITER_CLUSTER); -// when(mockPluginService.getCurrentConnection()).thenReturn(mockConnection); -// when(mockPluginService.getCurrentHostSpec()).thenReturn(mockHostSpec); -// when(mockPluginService.connect(any(HostSpec.class), eq(properties))).thenReturn(mockConnection); -// when(mockPluginService.getTelemetryFactory()).thenReturn(mockTelemetryFactory); -// when(mockPluginService.getHosts()).thenReturn(defaultHosts); -// when(mockPluginService.getAllHosts()).thenReturn(defaultHosts); -// when(mockReaderFailoverHandler.failover(any(), any())).thenReturn(mockReaderResult); -// when(mockWriterFailoverHandler.failover(any())).thenReturn(mockWriterResult); -// when(mockWriterResult.isConnected()).thenReturn(true); -// when(mockWriterResult.getTopology()).thenReturn(defaultHosts); -// when(mockReaderResult.isConnected()).thenReturn(true); -// -// when(mockPluginService.getTelemetryFactory()).thenReturn(mockTelemetryFactory); -// when(mockTelemetryFactory.openTelemetryContext(anyString(), any())).thenReturn(mockTelemetryContext); -// when(mockTelemetryFactory.openTelemetryContext(eq(null), any())).thenReturn(mockTelemetryContext); -// when(mockTelemetryFactory.createCounter(anyString())).thenReturn(mockTelemetryCounter); -// // noinspection unchecked -// when(mockTelemetryFactory.createGauge(anyString(), any(GaugeCallable.class))).thenReturn(mockTelemetryGauge); -// -// when(mockPluginService.getTargetDriverDialect()).thenReturn(mockTargetDriverDialect); -// when(mockTargetDriverDialect.getNetworkBoundMethodNames(any())).thenReturn(new HashSet<>()); -// -// properties.clear(); -// } -// -// @Test -// void test_notifyNodeListChanged_withFailoverDisabled() { -// properties.setProperty(FailoverConnectionPlugin.ENABLE_CLUSTER_AWARE_FAILOVER.name, "false"); -// final Map> changes = new HashMap<>(); -// -// initializePlugin(); -// plugin.notifyNodeListChanged(changes); -// -// verify(mockPluginService, never()).getCurrentHostSpec(); -// verify(mockHostSpec, never()).getAliases(); -// } -// -// @Test -// void test_notifyNodeListChanged_withValidConnectionNotInTopology() { -// final Map> changes = new HashMap<>(); -// changes.put("cluster-host/", EnumSet.of(NodeChangeOptions.NODE_DELETED)); -// changes.put("instance/", EnumSet.of(NodeChangeOptions.NODE_ADDED)); -// -// initializePlugin(); -// plugin.notifyNodeListChanged(changes); -// -// when(mockHostSpec.getUrl()).thenReturn("cluster-url/"); -// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Collections.singletonList("instance"))); -// -// verify(mockPluginService).getCurrentHostSpec(); -// verify(mockHostSpec, never()).getAliases(); -// } -// -// @Test -// void test_updateTopology() throws SQLException { -// initializePlugin(); -// -// // Test updateTopology with failover disabled -// plugin.setRdsUrlType(RdsUrlType.RDS_PROXY); -// plugin.updateTopology(false); -// verify(mockPluginService, never()).forceRefreshHostList(); -// verify(mockPluginService, never()).refreshHostList(); -// -// // Test updateTopology with no connection -// when(mockPluginService.getCurrentHostSpec()).thenReturn(null); -// plugin.updateTopology(false); -// verify(mockPluginService, never()).forceRefreshHostList(); -// verify(mockPluginService, never()).refreshHostList(); -// -// // Test updateTopology with closed connection -// when(mockConnection.isClosed()).thenReturn(true); -// plugin.updateTopology(false); -// verify(mockPluginService, never()).forceRefreshHostList(); -// verify(mockPluginService, never()).refreshHostList(); -// } -// -// @ParameterizedTest -// @ValueSource(booleans = {true, false}) -// void test_updateTopology_withForceUpdate(final boolean forceUpdate) throws SQLException { -// -// when(mockPluginService.getAllHosts()).thenReturn(Collections.singletonList( -// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("host").build())); -// when(mockPluginService.getHosts()).thenReturn(Collections.singletonList( -// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("host").build())); -// when(mockConnection.isClosed()).thenReturn(false); -// initializePlugin(); -// plugin.setRdsUrlType(RdsUrlType.RDS_INSTANCE); -// -// plugin.updateTopology(forceUpdate); -// if (forceUpdate) { -// verify(mockPluginService, atLeastOnce()).forceRefreshHostList(); -// } else { -// verify(mockPluginService, atLeastOnce()).refreshHostList(); -// } -// } -// -// @Test -// void test_failover_failoverWriter() throws SQLException { -// when(mockPluginService.isInTransaction()).thenReturn(true); -// -// initializePlugin(); -// final FailoverConnectionPlugin spyPlugin = spy(plugin); -// doThrow(FailoverSuccessSQLException.class).when(spyPlugin).failoverWriter(); -// spyPlugin.failoverMode = FailoverMode.STRICT_WRITER; -// -// assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failover(mockHostSpec)); -// verify(spyPlugin).failoverWriter(); -// } -// -// @Test -// void test_failover_failoverReader() throws SQLException { -// when(mockPluginService.isInTransaction()).thenReturn(false); -// -// initializePlugin(); -// final FailoverConnectionPlugin spyPlugin = spy(plugin); -// doThrow(FailoverSuccessSQLException.class).when(spyPlugin).failoverReader(eq(mockHostSpec)); -// spyPlugin.failoverMode = FailoverMode.READER_OR_WRITER; -// -// assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failover(mockHostSpec)); -// verify(spyPlugin).failoverReader(eq(mockHostSpec)); -// } -// -// @Test -// void test_failoverReader_withValidFailedHostSpec_successFailover() throws SQLException { -// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); -// when(mockHostSpec.getRawAvailability()).thenReturn(HostAvailability.AVAILABLE); -// when(mockReaderResult.isConnected()).thenReturn(true); -// when(mockReaderResult.getConnection()).thenReturn(mockConnection); -// when(mockReaderResult.getHost()).thenReturn(defaultHosts.get(1)); -// -// initializePlugin(); -// plugin.initHostProvider( -// mockHostListProviderService, -// mockInitHostProviderFunc, -// () -> mockReaderFailoverHandler, -// () -> mockWriterFailoverHandler); -// -// final FailoverConnectionPlugin spyPlugin = spy(plugin); -// doNothing().when(spyPlugin).updateTopology(true); -// -// assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failoverReader(mockHostSpec)); -// -// verify(mockReaderFailoverHandler).failover(eq(defaultHosts), eq(mockHostSpec)); -// verify(mockPluginService).setCurrentConnection(eq(mockConnection), eq(defaultHosts.get(1))); -// } -// -// @Test -// void test_failoverReader_withNoFailedHostSpec_withException() throws SQLException { -// final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") -// .build(); -// final List hosts = Collections.singletonList(hostSpec); -// -// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); -// when(mockHostSpec.getAvailability()).thenReturn(HostAvailability.AVAILABLE); -// when(mockPluginService.getAllHosts()).thenReturn(hosts); -// when(mockPluginService.getHosts()).thenReturn(hosts); -// when(mockReaderResult.getException()).thenReturn(new SQLException()); -// when(mockReaderResult.getHost()).thenReturn(hostSpec); -// -// initializePlugin(); -// plugin.initHostProvider( -// mockHostListProviderService, -// mockInitHostProviderFunc, -// () -> mockReaderFailoverHandler, -// () -> mockWriterFailoverHandler); -// -// assertThrows(SQLException.class, () -> plugin.failoverReader(null)); -// verify(mockReaderFailoverHandler).failover(eq(hosts), eq(null)); -// } -// -// @Test -// void test_failoverWriter_failedFailover_throwsException() throws SQLException { -// final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") -// .build(); -// final List hosts = Collections.singletonList(hostSpec); -// -// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); -// when(mockPluginService.getAllHosts()).thenReturn(hosts); -// when(mockPluginService.getHosts()).thenReturn(hosts); -// when(mockWriterResult.getException()).thenReturn(new SQLException()); -// -// initializePlugin(); -// plugin.initHostProvider( -// mockHostListProviderService, -// mockInitHostProviderFunc, -// () -> mockReaderFailoverHandler, -// () -> mockWriterFailoverHandler); -// -// assertThrows(SQLException.class, () -> plugin.failoverWriter()); -// verify(mockWriterFailoverHandler).failover(eq(hosts)); -// } -// -// @Test -// void test_failoverWriter_failedFailover_withNoResult() throws SQLException { -// final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") -// .build(); -// final List hosts = Collections.singletonList(hostSpec); -// -// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); -// when(mockPluginService.getAllHosts()).thenReturn(hosts); -// when(mockPluginService.getHosts()).thenReturn(hosts); -// when(mockWriterResult.isConnected()).thenReturn(false); -// -// initializePlugin(); -// plugin.initHostProvider( -// mockHostListProviderService, -// mockInitHostProviderFunc, -// () -> mockReaderFailoverHandler, -// () -> mockWriterFailoverHandler); -// -// final SQLException exception = assertThrows(SQLException.class, () -> plugin.failoverWriter()); -// assertEquals(SqlState.CONNECTION_UNABLE_TO_CONNECT.getState(), exception.getSQLState()); -// -// verify(mockWriterFailoverHandler).failover(eq(hosts)); -// verify(mockWriterResult, never()).getNewConnection(); -// verify(mockWriterResult, never()).getTopology(); -// } -// -// @Test -// void test_failoverWriter_successFailover() throws SQLException { -// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); -// -// initializePlugin(); -// plugin.initHostProvider( -// mockHostListProviderService, -// mockInitHostProviderFunc, -// () -> mockReaderFailoverHandler, -// () -> mockWriterFailoverHandler); -// -// final SQLException exception = assertThrows(FailoverSuccessSQLException.class, () -> plugin.failoverWriter()); -// assertEquals(SqlState.COMMUNICATION_LINK_CHANGED.getState(), exception.getSQLState()); -// -// verify(mockWriterFailoverHandler).failover(eq(defaultHosts)); -// } -// -// @Test -// void test_invalidCurrentConnection_withNoConnection() { -// when(mockPluginService.getCurrentConnection()).thenReturn(null); -// initializePlugin(); -// plugin.invalidateCurrentConnection(); -// -// verify(mockPluginService, never()).getCurrentHostSpec(); -// } -// -// @Test -// void test_invalidateCurrentConnection_inTransaction() throws SQLException { -// when(mockPluginService.isInTransaction()).thenReturn(true); -// when(mockHostSpec.getHost()).thenReturn("host"); -// when(mockHostSpec.getPort()).thenReturn(123); -// when(mockHostSpec.getRole()).thenReturn(HostRole.READER); -// -// initializePlugin(); -// plugin.invalidateCurrentConnection(); -// verify(mockConnection).rollback(); -// -// // Assert SQL exceptions thrown during rollback do not get propagated. -// doThrow(new SQLException()).when(mockConnection).rollback(); -// assertDoesNotThrow(() -> plugin.invalidateCurrentConnection()); -// } -// -// @Test -// void test_invalidateCurrentConnection_notInTransaction() { -// when(mockPluginService.isInTransaction()).thenReturn(false); -// when(mockHostSpec.getHost()).thenReturn("host"); -// when(mockHostSpec.getPort()).thenReturn(123); -// when(mockHostSpec.getRole()).thenReturn(HostRole.READER); -// -// initializePlugin(); -// plugin.invalidateCurrentConnection(); -// -// verify(mockPluginService).isInTransaction(); -// } -// -// @Test -// void test_invalidateCurrentConnection_withOpenConnection() throws SQLException { -// when(mockPluginService.isInTransaction()).thenReturn(false); -// when(mockConnection.isClosed()).thenReturn(false); -// when(mockHostSpec.getHost()).thenReturn("host"); -// when(mockHostSpec.getPort()).thenReturn(123); -// when(mockHostSpec.getRole()).thenReturn(HostRole.READER); -// -// initializePlugin(); -// plugin.invalidateCurrentConnection(); -// -// doThrow(new SQLException()).when(mockConnection).close(); -// assertDoesNotThrow(() -> plugin.invalidateCurrentConnection()); -// -// verify(mockConnection, times(2)).isClosed(); -// verify(mockConnection, times(2)).close(); -// } -// -// @Test -// void test_execute_withFailoverDisabled() throws SQLException { -// properties.setProperty(FailoverConnectionPlugin.ENABLE_CLUSTER_AWARE_FAILOVER.name, "false"); -// initializePlugin(); -// -// plugin.execute( -// ResultSet.class, -// SQLException.class, -// MONITOR_METHOD_INVOKE_ON, -// MONITOR_METHOD_NAME, -// mockSqlFunction, -// EMPTY_ARGS); -// -// verify(mockSqlFunction).call(); -// verify(mockHostListProvider, never()).getRdsUrlType(); -// } -// -// @Test -// void test_execute_withDirectExecute() throws SQLException { -// initializePlugin(); -// plugin.execute( -// ResultSet.class, -// SQLException.class, -// MONITOR_METHOD_INVOKE_ON, -// "close", -// mockSqlFunction, -// EMPTY_ARGS); -// verify(mockSqlFunction).call(); -// verify(mockHostListProvider, never()).getRdsUrlType(); -// } -// -// private void initializePlugin() { -// plugin = new FailoverConnectionPlugin(mockPluginService, properties); -// plugin.setWriterFailoverHandler(mockWriterFailoverHandler); -// plugin.setReaderFailoverHandler(mockReaderFailoverHandler); -// } -// } +/* + * 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.failover; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.fail; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.Arrays; +import java.util.Collections; +import java.util.EnumSet; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Properties; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import software.amazon.jdbc.HostListProviderService; +import software.amazon.jdbc.HostRole; +import software.amazon.jdbc.HostSpec; +import software.amazon.jdbc.HostSpecBuilder; +import software.amazon.jdbc.JdbcCallable; +import software.amazon.jdbc.NodeChangeOptions; +import software.amazon.jdbc.PluginService; +import software.amazon.jdbc.hostavailability.HostAvailability; +import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; +import software.amazon.jdbc.hostlistprovider.AuroraHostListProvider; +import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; +import software.amazon.jdbc.util.FullServicesContainer; +import software.amazon.jdbc.util.RdsUrlType; +import software.amazon.jdbc.util.SqlState; +import software.amazon.jdbc.util.connection.ConnectionService; +import software.amazon.jdbc.util.telemetry.GaugeCallable; +import software.amazon.jdbc.util.telemetry.TelemetryContext; +import software.amazon.jdbc.util.telemetry.TelemetryCounter; +import software.amazon.jdbc.util.telemetry.TelemetryFactory; +import software.amazon.jdbc.util.telemetry.TelemetryGauge; + +class FailoverConnectionPluginTest { + + private static final Class MONITOR_METHOD_INVOKE_ON = Connection.class; + private static final String MONITOR_METHOD_NAME = "Connection.executeQuery"; + private static final Object[] EMPTY_ARGS = {}; + private final List defaultHosts = Arrays.asList( + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("writer").port(1234).role(HostRole.WRITER).build(), + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("reader1").port(1234).role(HostRole.READER).build()); + + @Mock FullServicesContainer mockContainer; + @Mock ConnectionService mockConnectionService; + @Mock PluginService mockPluginService; + @Mock Connection mockConnection; + @Mock HostSpec mockHostSpec; + @Mock HostListProviderService mockHostListProviderService; + @Mock AuroraHostListProvider mockHostListProvider; + @Mock JdbcCallable mockInitHostProviderFunc; + @Mock ReaderFailoverHandler mockReaderFailoverHandler; + @Mock WriterFailoverHandler mockWriterFailoverHandler; + @Mock ReaderFailoverResult mockReaderResult; + @Mock WriterFailoverResult mockWriterResult; + @Mock JdbcCallable mockSqlFunction; + @Mock private TelemetryFactory mockTelemetryFactory; + @Mock TelemetryContext mockTelemetryContext; + @Mock TelemetryCounter mockTelemetryCounter; + @Mock TelemetryGauge mockTelemetryGauge; + @Mock TargetDriverDialect mockTargetDriverDialect; + + + private final Properties properties = new Properties(); + private FailoverConnectionPlugin plugin; + private AutoCloseable closeable; + + @AfterEach + void cleanUp() throws Exception { + closeable.close(); + } + + @BeforeEach + void init() throws SQLException { + closeable = MockitoAnnotations.openMocks(this); + + when(mockContainer.getPluginService()).thenReturn(mockPluginService); + when(mockPluginService.getHostListProvider()).thenReturn(mockHostListProvider); + when(mockHostListProvider.getRdsUrlType()).thenReturn(RdsUrlType.RDS_WRITER_CLUSTER); + when(mockPluginService.getCurrentConnection()).thenReturn(mockConnection); + when(mockPluginService.getCurrentHostSpec()).thenReturn(mockHostSpec); + when(mockPluginService.connect(any(HostSpec.class), eq(properties))).thenReturn(mockConnection); + when(mockPluginService.getTelemetryFactory()).thenReturn(mockTelemetryFactory); + when(mockPluginService.getHosts()).thenReturn(defaultHosts); + when(mockPluginService.getAllHosts()).thenReturn(defaultHosts); + when(mockReaderFailoverHandler.failover(any(), any())).thenReturn(mockReaderResult); + when(mockWriterFailoverHandler.failover(any())).thenReturn(mockWriterResult); + when(mockWriterResult.isConnected()).thenReturn(true); + when(mockWriterResult.getTopology()).thenReturn(defaultHosts); + when(mockReaderResult.isConnected()).thenReturn(true); + + when(mockPluginService.getTelemetryFactory()).thenReturn(mockTelemetryFactory); + when(mockTelemetryFactory.openTelemetryContext(anyString(), any())).thenReturn(mockTelemetryContext); + when(mockTelemetryFactory.openTelemetryContext(eq(null), any())).thenReturn(mockTelemetryContext); + when(mockTelemetryFactory.createCounter(anyString())).thenReturn(mockTelemetryCounter); + // noinspection unchecked + when(mockTelemetryFactory.createGauge(anyString(), any(GaugeCallable.class))).thenReturn(mockTelemetryGauge); + + when(mockPluginService.getTargetDriverDialect()).thenReturn(mockTargetDriverDialect); + when(mockTargetDriverDialect.getNetworkBoundMethodNames(any())).thenReturn(new HashSet<>()); + + properties.clear(); + } + + @Test + void test_notifyNodeListChanged_withFailoverDisabled() { + properties.setProperty(FailoverConnectionPlugin.ENABLE_CLUSTER_AWARE_FAILOVER.name, "false"); + final Map> changes = new HashMap<>(); + + initializePlugin(); + plugin.notifyNodeListChanged(changes); + + verify(mockPluginService, never()).getCurrentHostSpec(); + verify(mockHostSpec, never()).getAliases(); + } + + @Test + void test_notifyNodeListChanged_withValidConnectionNotInTopology() { + final Map> changes = new HashMap<>(); + changes.put("cluster-host/", EnumSet.of(NodeChangeOptions.NODE_DELETED)); + changes.put("instance/", EnumSet.of(NodeChangeOptions.NODE_ADDED)); + + initializePlugin(); + plugin.notifyNodeListChanged(changes); + + when(mockHostSpec.getUrl()).thenReturn("cluster-url/"); + when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Collections.singletonList("instance"))); + + verify(mockPluginService).getCurrentHostSpec(); + verify(mockHostSpec, never()).getAliases(); + } + + @Test + void test_updateTopology() throws SQLException { + initializePlugin(); + + // Test updateTopology with failover disabled + plugin.setRdsUrlType(RdsUrlType.RDS_PROXY); + plugin.updateTopology(false); + verify(mockPluginService, never()).forceRefreshHostList(); + verify(mockPluginService, never()).refreshHostList(); + + // Test updateTopology with no connection + when(mockPluginService.getCurrentHostSpec()).thenReturn(null); + plugin.updateTopology(false); + verify(mockPluginService, never()).forceRefreshHostList(); + verify(mockPluginService, never()).refreshHostList(); + + // Test updateTopology with closed connection + when(mockConnection.isClosed()).thenReturn(true); + plugin.updateTopology(false); + verify(mockPluginService, never()).forceRefreshHostList(); + verify(mockPluginService, never()).refreshHostList(); + } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void test_updateTopology_withForceUpdate(final boolean forceUpdate) throws SQLException { + + when(mockPluginService.getAllHosts()).thenReturn(Collections.singletonList( + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("host").build())); + when(mockPluginService.getHosts()).thenReturn(Collections.singletonList( + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("host").build())); + when(mockConnection.isClosed()).thenReturn(false); + initializePlugin(); + plugin.setRdsUrlType(RdsUrlType.RDS_INSTANCE); + + plugin.updateTopology(forceUpdate); + if (forceUpdate) { + verify(mockPluginService, atLeastOnce()).forceRefreshHostList(); + } else { + verify(mockPluginService, atLeastOnce()).refreshHostList(); + } + } + + @Test + void test_failover_failoverWriter() throws SQLException { + when(mockPluginService.isInTransaction()).thenReturn(true); + + initializePlugin(); + doThrow(FailoverSuccessSQLException.class).when(plugin).failoverWriter(); + plugin.failoverMode = FailoverMode.STRICT_WRITER; + + assertThrows(FailoverSuccessSQLException.class, () -> plugin.failover(mockHostSpec)); + verify(plugin).failoverWriter(); + } + + @Test + void test_failover_failoverReader() throws SQLException { + when(mockPluginService.isInTransaction()).thenReturn(false); + + initializePlugin(); + doThrow(FailoverSuccessSQLException.class).when(plugin).failoverReader(eq(mockHostSpec)); + plugin.failoverMode = FailoverMode.READER_OR_WRITER; + + assertThrows(FailoverSuccessSQLException.class, () -> plugin.failover(mockHostSpec)); + verify(plugin).failoverReader(eq(mockHostSpec)); + } + + @Test + void test_failoverReader_withValidFailedHostSpec_successFailover() throws SQLException { + when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); + when(mockHostSpec.getRawAvailability()).thenReturn(HostAvailability.AVAILABLE); + when(mockReaderResult.isConnected()).thenReturn(true); + when(mockReaderResult.getConnection()).thenReturn(mockConnection); + when(mockReaderResult.getHost()).thenReturn(defaultHosts.get(1)); + + initializePlugin(); + plugin.initHostProvider( + mockHostListProviderService, + mockInitHostProviderFunc, + (connectionService) -> mockReaderFailoverHandler, + (connectionService) -> mockWriterFailoverHandler); + + final FailoverConnectionPlugin spyPlugin = spy(plugin); + doNothing().when(spyPlugin).updateTopology(true); + + assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failoverReader(mockHostSpec)); + + verify(mockReaderFailoverHandler).failover(eq(defaultHosts), eq(mockHostSpec)); + verify(mockPluginService).setCurrentConnection(eq(mockConnection), eq(defaultHosts.get(1))); + } + + @Test + void test_failoverReader_withNoFailedHostSpec_withException() throws SQLException { + final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") + .build(); + final List hosts = Collections.singletonList(hostSpec); + + when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); + when(mockHostSpec.getAvailability()).thenReturn(HostAvailability.AVAILABLE); + when(mockPluginService.getAllHosts()).thenReturn(hosts); + when(mockPluginService.getHosts()).thenReturn(hosts); + when(mockReaderResult.getException()).thenReturn(new SQLException()); + when(mockReaderResult.getHost()).thenReturn(hostSpec); + + initializePlugin(); + plugin.initHostProvider( + mockHostListProviderService, + mockInitHostProviderFunc, + (connectionService) -> mockReaderFailoverHandler, + (connectionService) -> mockWriterFailoverHandler); + + assertThrows(SQLException.class, () -> plugin.failoverReader(null)); + verify(mockReaderFailoverHandler).failover(eq(hosts), eq(null)); + } + + @Test + void test_failoverWriter_failedFailover_throwsException() throws SQLException { + final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") + .build(); + final List hosts = Collections.singletonList(hostSpec); + + when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); + when(mockPluginService.getAllHosts()).thenReturn(hosts); + when(mockPluginService.getHosts()).thenReturn(hosts); + when(mockWriterResult.getException()).thenReturn(new SQLException()); + + initializePlugin(); + plugin.initHostProvider( + mockHostListProviderService, + mockInitHostProviderFunc, + (connectionService) -> mockReaderFailoverHandler, + (connectionService) -> mockWriterFailoverHandler); + + assertThrows(SQLException.class, () -> plugin.failoverWriter()); + verify(mockWriterFailoverHandler).failover(eq(hosts)); + } + + @Test + void test_failoverWriter_failedFailover_withNoResult() throws SQLException { + final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") + .build(); + final List hosts = Collections.singletonList(hostSpec); + + when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); + when(mockPluginService.getAllHosts()).thenReturn(hosts); + when(mockPluginService.getHosts()).thenReturn(hosts); + when(mockWriterResult.isConnected()).thenReturn(false); + + initializePlugin(); + plugin.initHostProvider( + mockHostListProviderService, + mockInitHostProviderFunc, + (connectionService) -> mockReaderFailoverHandler, + (connectionService) -> mockWriterFailoverHandler); + + final SQLException exception = assertThrows(SQLException.class, () -> plugin.failoverWriter()); + assertEquals(SqlState.CONNECTION_UNABLE_TO_CONNECT.getState(), exception.getSQLState()); + + verify(mockWriterFailoverHandler).failover(eq(hosts)); + verify(mockWriterResult, never()).getNewConnection(); + verify(mockWriterResult, never()).getTopology(); + } + + @Test + void test_failoverWriter_successFailover() throws SQLException { + when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); + + initializePlugin(); + plugin.initHostProvider( + mockHostListProviderService, + mockInitHostProviderFunc, + (connectionService) -> mockReaderFailoverHandler, + (connectionService) -> mockWriterFailoverHandler); + + final SQLException exception = assertThrows(FailoverSuccessSQLException.class, () -> plugin.failoverWriter()); + assertEquals(SqlState.COMMUNICATION_LINK_CHANGED.getState(), exception.getSQLState()); + + verify(mockWriterFailoverHandler).failover(eq(defaultHosts)); + } + + @Test + void test_invalidCurrentConnection_withNoConnection() { + when(mockPluginService.getCurrentConnection()).thenReturn(null); + initializePlugin(); + plugin.invalidateCurrentConnection(); + + verify(mockPluginService, never()).getCurrentHostSpec(); + } + + @Test + void test_invalidateCurrentConnection_inTransaction() throws SQLException { + when(mockPluginService.isInTransaction()).thenReturn(true); + when(mockHostSpec.getHost()).thenReturn("host"); + when(mockHostSpec.getPort()).thenReturn(123); + when(mockHostSpec.getRole()).thenReturn(HostRole.READER); + + initializePlugin(); + plugin.invalidateCurrentConnection(); + verify(mockConnection).rollback(); + + // Assert SQL exceptions thrown during rollback do not get propagated. + doThrow(new SQLException()).when(mockConnection).rollback(); + assertDoesNotThrow(() -> plugin.invalidateCurrentConnection()); + } + + @Test + void test_invalidateCurrentConnection_notInTransaction() { + when(mockPluginService.isInTransaction()).thenReturn(false); + when(mockHostSpec.getHost()).thenReturn("host"); + when(mockHostSpec.getPort()).thenReturn(123); + when(mockHostSpec.getRole()).thenReturn(HostRole.READER); + + initializePlugin(); + plugin.invalidateCurrentConnection(); + + verify(mockPluginService).isInTransaction(); + } + + @Test + void test_invalidateCurrentConnection_withOpenConnection() throws SQLException { + when(mockPluginService.isInTransaction()).thenReturn(false); + when(mockConnection.isClosed()).thenReturn(false); + when(mockHostSpec.getHost()).thenReturn("host"); + when(mockHostSpec.getPort()).thenReturn(123); + when(mockHostSpec.getRole()).thenReturn(HostRole.READER); + + initializePlugin(); + plugin.invalidateCurrentConnection(); + + doThrow(new SQLException()).when(mockConnection).close(); + assertDoesNotThrow(() -> plugin.invalidateCurrentConnection()); + + verify(mockConnection, times(2)).isClosed(); + verify(mockConnection, times(2)).close(); + } + + @Test + void test_execute_withFailoverDisabled() throws SQLException { + properties.setProperty(FailoverConnectionPlugin.ENABLE_CLUSTER_AWARE_FAILOVER.name, "false"); + initializePlugin(); + + plugin.execute( + ResultSet.class, + SQLException.class, + MONITOR_METHOD_INVOKE_ON, + MONITOR_METHOD_NAME, + mockSqlFunction, + EMPTY_ARGS); + + verify(mockSqlFunction).call(); + verify(mockHostListProvider, never()).getRdsUrlType(); + } + + @Test + void test_execute_withDirectExecute() throws SQLException { + initializePlugin(); + plugin.execute( + ResultSet.class, + SQLException.class, + MONITOR_METHOD_INVOKE_ON, + "close", + mockSqlFunction, + EMPTY_ARGS); + verify(mockSqlFunction).call(); + verify(mockHostListProvider, never()).getRdsUrlType(); + } + + private void initializePlugin() { + plugin = spy(new FailoverConnectionPlugin(mockContainer, properties)); + plugin.setWriterFailoverHandler(mockWriterFailoverHandler); + plugin.setReaderFailoverHandler(mockReaderFailoverHandler); + + try { + doReturn(mockConnectionService).when(plugin).getConnectionService(); + } catch (SQLException e) { + fail( + "Encountered exception when trying to stub FailoverConnectionPlugin#getConnectionService: " + e.getMessage()); + } + } +} From 6ebf706acd042675d2629dded3b200da14df57b5 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Fri, 22 Aug 2025 15:48:24 -0700 Subject: [PATCH 16/42] Rename plugin to spyPlugin in FailoverConnectionPluginTest --- .../FailoverConnectionPluginTest.java | 78 +++++++++---------- 1 file changed, 39 insertions(+), 39 deletions(-) diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java index adf1cb0e6..045482764 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java @@ -104,7 +104,7 @@ class FailoverConnectionPluginTest { private final Properties properties = new Properties(); - private FailoverConnectionPlugin plugin; + private FailoverConnectionPlugin spyPlugin; private AutoCloseable closeable; @AfterEach @@ -150,7 +150,7 @@ void test_notifyNodeListChanged_withFailoverDisabled() { final Map> changes = new HashMap<>(); initializePlugin(); - plugin.notifyNodeListChanged(changes); + spyPlugin.notifyNodeListChanged(changes); verify(mockPluginService, never()).getCurrentHostSpec(); verify(mockHostSpec, never()).getAliases(); @@ -163,7 +163,7 @@ void test_notifyNodeListChanged_withValidConnectionNotInTopology() { changes.put("instance/", EnumSet.of(NodeChangeOptions.NODE_ADDED)); initializePlugin(); - plugin.notifyNodeListChanged(changes); + spyPlugin.notifyNodeListChanged(changes); when(mockHostSpec.getUrl()).thenReturn("cluster-url/"); when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Collections.singletonList("instance"))); @@ -177,20 +177,20 @@ void test_updateTopology() throws SQLException { initializePlugin(); // Test updateTopology with failover disabled - plugin.setRdsUrlType(RdsUrlType.RDS_PROXY); - plugin.updateTopology(false); + spyPlugin.setRdsUrlType(RdsUrlType.RDS_PROXY); + spyPlugin.updateTopology(false); verify(mockPluginService, never()).forceRefreshHostList(); verify(mockPluginService, never()).refreshHostList(); // Test updateTopology with no connection when(mockPluginService.getCurrentHostSpec()).thenReturn(null); - plugin.updateTopology(false); + spyPlugin.updateTopology(false); verify(mockPluginService, never()).forceRefreshHostList(); verify(mockPluginService, never()).refreshHostList(); // Test updateTopology with closed connection when(mockConnection.isClosed()).thenReturn(true); - plugin.updateTopology(false); + spyPlugin.updateTopology(false); verify(mockPluginService, never()).forceRefreshHostList(); verify(mockPluginService, never()).refreshHostList(); } @@ -205,9 +205,9 @@ void test_updateTopology_withForceUpdate(final boolean forceUpdate) throws SQLEx new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("host").build())); when(mockConnection.isClosed()).thenReturn(false); initializePlugin(); - plugin.setRdsUrlType(RdsUrlType.RDS_INSTANCE); + spyPlugin.setRdsUrlType(RdsUrlType.RDS_INSTANCE); - plugin.updateTopology(forceUpdate); + spyPlugin.updateTopology(forceUpdate); if (forceUpdate) { verify(mockPluginService, atLeastOnce()).forceRefreshHostList(); } else { @@ -220,11 +220,11 @@ void test_failover_failoverWriter() throws SQLException { when(mockPluginService.isInTransaction()).thenReturn(true); initializePlugin(); - doThrow(FailoverSuccessSQLException.class).when(plugin).failoverWriter(); - plugin.failoverMode = FailoverMode.STRICT_WRITER; + doThrow(FailoverSuccessSQLException.class).when(spyPlugin).failoverWriter(); + spyPlugin.failoverMode = FailoverMode.STRICT_WRITER; - assertThrows(FailoverSuccessSQLException.class, () -> plugin.failover(mockHostSpec)); - verify(plugin).failoverWriter(); + assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failover(mockHostSpec)); + verify(spyPlugin).failoverWriter(); } @Test @@ -232,11 +232,11 @@ void test_failover_failoverReader() throws SQLException { when(mockPluginService.isInTransaction()).thenReturn(false); initializePlugin(); - doThrow(FailoverSuccessSQLException.class).when(plugin).failoverReader(eq(mockHostSpec)); - plugin.failoverMode = FailoverMode.READER_OR_WRITER; + doThrow(FailoverSuccessSQLException.class).when(spyPlugin).failoverReader(eq(mockHostSpec)); + spyPlugin.failoverMode = FailoverMode.READER_OR_WRITER; - assertThrows(FailoverSuccessSQLException.class, () -> plugin.failover(mockHostSpec)); - verify(plugin).failoverReader(eq(mockHostSpec)); + assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failover(mockHostSpec)); + verify(spyPlugin).failoverReader(eq(mockHostSpec)); } @Test @@ -248,13 +248,13 @@ void test_failoverReader_withValidFailedHostSpec_successFailover() throws SQLExc when(mockReaderResult.getHost()).thenReturn(defaultHosts.get(1)); initializePlugin(); - plugin.initHostProvider( + spyPlugin.initHostProvider( mockHostListProviderService, mockInitHostProviderFunc, (connectionService) -> mockReaderFailoverHandler, (connectionService) -> mockWriterFailoverHandler); - final FailoverConnectionPlugin spyPlugin = spy(plugin); + final FailoverConnectionPlugin spyPlugin = spy(this.spyPlugin); doNothing().when(spyPlugin).updateTopology(true); assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failoverReader(mockHostSpec)); @@ -277,13 +277,13 @@ void test_failoverReader_withNoFailedHostSpec_withException() throws SQLExceptio when(mockReaderResult.getHost()).thenReturn(hostSpec); initializePlugin(); - plugin.initHostProvider( + spyPlugin.initHostProvider( mockHostListProviderService, mockInitHostProviderFunc, (connectionService) -> mockReaderFailoverHandler, (connectionService) -> mockWriterFailoverHandler); - assertThrows(SQLException.class, () -> plugin.failoverReader(null)); + assertThrows(SQLException.class, () -> spyPlugin.failoverReader(null)); verify(mockReaderFailoverHandler).failover(eq(hosts), eq(null)); } @@ -299,13 +299,13 @@ void test_failoverWriter_failedFailover_throwsException() throws SQLException { when(mockWriterResult.getException()).thenReturn(new SQLException()); initializePlugin(); - plugin.initHostProvider( + spyPlugin.initHostProvider( mockHostListProviderService, mockInitHostProviderFunc, (connectionService) -> mockReaderFailoverHandler, (connectionService) -> mockWriterFailoverHandler); - assertThrows(SQLException.class, () -> plugin.failoverWriter()); + assertThrows(SQLException.class, () -> spyPlugin.failoverWriter()); verify(mockWriterFailoverHandler).failover(eq(hosts)); } @@ -321,13 +321,13 @@ void test_failoverWriter_failedFailover_withNoResult() throws SQLException { when(mockWriterResult.isConnected()).thenReturn(false); initializePlugin(); - plugin.initHostProvider( + spyPlugin.initHostProvider( mockHostListProviderService, mockInitHostProviderFunc, (connectionService) -> mockReaderFailoverHandler, (connectionService) -> mockWriterFailoverHandler); - final SQLException exception = assertThrows(SQLException.class, () -> plugin.failoverWriter()); + final SQLException exception = assertThrows(SQLException.class, () -> spyPlugin.failoverWriter()); assertEquals(SqlState.CONNECTION_UNABLE_TO_CONNECT.getState(), exception.getSQLState()); verify(mockWriterFailoverHandler).failover(eq(hosts)); @@ -340,13 +340,13 @@ void test_failoverWriter_successFailover() throws SQLException { when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); initializePlugin(); - plugin.initHostProvider( + spyPlugin.initHostProvider( mockHostListProviderService, mockInitHostProviderFunc, (connectionService) -> mockReaderFailoverHandler, (connectionService) -> mockWriterFailoverHandler); - final SQLException exception = assertThrows(FailoverSuccessSQLException.class, () -> plugin.failoverWriter()); + final SQLException exception = assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failoverWriter()); assertEquals(SqlState.COMMUNICATION_LINK_CHANGED.getState(), exception.getSQLState()); verify(mockWriterFailoverHandler).failover(eq(defaultHosts)); @@ -356,7 +356,7 @@ void test_failoverWriter_successFailover() throws SQLException { void test_invalidCurrentConnection_withNoConnection() { when(mockPluginService.getCurrentConnection()).thenReturn(null); initializePlugin(); - plugin.invalidateCurrentConnection(); + spyPlugin.invalidateCurrentConnection(); verify(mockPluginService, never()).getCurrentHostSpec(); } @@ -369,12 +369,12 @@ void test_invalidateCurrentConnection_inTransaction() throws SQLException { when(mockHostSpec.getRole()).thenReturn(HostRole.READER); initializePlugin(); - plugin.invalidateCurrentConnection(); + spyPlugin.invalidateCurrentConnection(); verify(mockConnection).rollback(); // Assert SQL exceptions thrown during rollback do not get propagated. doThrow(new SQLException()).when(mockConnection).rollback(); - assertDoesNotThrow(() -> plugin.invalidateCurrentConnection()); + assertDoesNotThrow(() -> spyPlugin.invalidateCurrentConnection()); } @Test @@ -385,7 +385,7 @@ void test_invalidateCurrentConnection_notInTransaction() { when(mockHostSpec.getRole()).thenReturn(HostRole.READER); initializePlugin(); - plugin.invalidateCurrentConnection(); + spyPlugin.invalidateCurrentConnection(); verify(mockPluginService).isInTransaction(); } @@ -399,10 +399,10 @@ void test_invalidateCurrentConnection_withOpenConnection() throws SQLException { when(mockHostSpec.getRole()).thenReturn(HostRole.READER); initializePlugin(); - plugin.invalidateCurrentConnection(); + spyPlugin.invalidateCurrentConnection(); doThrow(new SQLException()).when(mockConnection).close(); - assertDoesNotThrow(() -> plugin.invalidateCurrentConnection()); + assertDoesNotThrow(() -> spyPlugin.invalidateCurrentConnection()); verify(mockConnection, times(2)).isClosed(); verify(mockConnection, times(2)).close(); @@ -413,7 +413,7 @@ void test_execute_withFailoverDisabled() throws SQLException { properties.setProperty(FailoverConnectionPlugin.ENABLE_CLUSTER_AWARE_FAILOVER.name, "false"); initializePlugin(); - plugin.execute( + spyPlugin.execute( ResultSet.class, SQLException.class, MONITOR_METHOD_INVOKE_ON, @@ -428,7 +428,7 @@ void test_execute_withFailoverDisabled() throws SQLException { @Test void test_execute_withDirectExecute() throws SQLException { initializePlugin(); - plugin.execute( + spyPlugin.execute( ResultSet.class, SQLException.class, MONITOR_METHOD_INVOKE_ON, @@ -440,12 +440,12 @@ void test_execute_withDirectExecute() throws SQLException { } private void initializePlugin() { - plugin = spy(new FailoverConnectionPlugin(mockContainer, properties)); - plugin.setWriterFailoverHandler(mockWriterFailoverHandler); - plugin.setReaderFailoverHandler(mockReaderFailoverHandler); + spyPlugin = spy(new FailoverConnectionPlugin(mockContainer, properties)); + spyPlugin.setWriterFailoverHandler(mockWriterFailoverHandler); + spyPlugin.setReaderFailoverHandler(mockReaderFailoverHandler); try { - doReturn(mockConnectionService).when(plugin).getConnectionService(); + doReturn(mockConnectionService).when(spyPlugin).getConnectionService(); } catch (SQLException e) { fail( "Encountered exception when trying to stub FailoverConnectionPlugin#getConnectionService: " + e.getMessage()); From 7711a9e6c989e3e2b78b7ee19dda3dfccaa66856 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Fri, 22 Aug 2025 15:48:40 -0700 Subject: [PATCH 17/42] Fix failing MonitorServiceImplTest tests --- .../util/monitoring/MonitorServiceImpl.java | 43 ++++-- .../monitoring/MonitorServiceImplTest.java | 144 ++++++++++-------- 2 files changed, 107 insertions(+), 80 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java index 227b29587..d45729c72 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java @@ -42,6 +42,7 @@ import software.amazon.jdbc.util.ExecutorFactory; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.PropertyUtils; +import software.amazon.jdbc.util.connection.ConnectionService; import software.amazon.jdbc.util.connection.ConnectionServiceImpl; import software.amazon.jdbc.util.events.DataAccessEvent; import software.amazon.jdbc.util.events.Event; @@ -197,20 +198,15 @@ public T runIfAbsent( cacheContainer = monitorCaches.computeIfAbsent(monitorClass, k -> supplier.get()); } - TargetDriverHelper helper = new TargetDriverHelper(); - java.sql.Driver driver = helper.getTargetDriver(originalUrl, originalProps); - final ConnectionProvider defaultConnectionProvider = new DriverConnectionProvider(driver); - final Properties propsCopy = PropertyUtils.copyProperties(originalProps); - final ConnectionServiceImpl connectionService = new ConnectionServiceImpl( - storageService, - this, - telemetryFactory, - defaultConnectionProvider, - originalUrl, - driverProtocol, - driverDialect, - dbDialect, - propsCopy); + final ConnectionService connectionService = + getConnectionService( + storageService, + telemetryFactory, + originalUrl, + driverProtocol, + driverDialect, + dbDialect, + originalProps); Monitor monitor = cacheContainer.getCache().computeIfAbsent(key, k -> { MonitorItem monitorItem = new MonitorItem(() -> initializer.createMonitor( @@ -228,6 +224,25 @@ public T runIfAbsent( Messages.get("MonitorServiceImpl.unexpectedMonitorClass", new Object[] {monitorClass, monitor})); } + protected ConnectionService getConnectionService(StorageService storageService, + TelemetryFactory telemetryFactory, String originalUrl, String driverProtocol, TargetDriverDialect driverDialect, + Dialect dbDialect, Properties originalProps) throws SQLException { + TargetDriverHelper helper = new TargetDriverHelper(); + java.sql.Driver driver = helper.getTargetDriver(originalUrl, originalProps); + final ConnectionProvider defaultConnectionProvider = new DriverConnectionProvider(driver); + final Properties propsCopy = PropertyUtils.copyProperties(originalProps); + return new ConnectionServiceImpl( + storageService, + this, + telemetryFactory, + defaultConnectionProvider, + originalUrl, + driverProtocol, + driverDialect, + dbDialect, + propsCopy); + } + @Override public @Nullable T get(Class monitorClass, Object key) { CacheContainer cacheContainer = monitorCaches.get(monitorClass); diff --git a/wrapper/src/test/java/software/amazon/jdbc/util/monitoring/MonitorServiceImplTest.java b/wrapper/src/test/java/software/amazon/jdbc/util/monitoring/MonitorServiceImplTest.java index cd0bcbe3d..61248faaa 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/util/monitoring/MonitorServiceImplTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/util/monitoring/MonitorServiceImplTest.java @@ -21,6 +21,11 @@ import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.spy; import java.sql.SQLException; import java.util.Collections; @@ -28,6 +33,7 @@ import java.util.Properties; import java.util.concurrent.TimeUnit; import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.Mock; @@ -35,39 +41,45 @@ import software.amazon.jdbc.dialect.Dialect; import software.amazon.jdbc.plugin.customendpoint.CustomEndpointMonitorImpl; import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; +import software.amazon.jdbc.util.connection.ConnectionService; import software.amazon.jdbc.util.events.EventPublisher; import software.amazon.jdbc.util.storage.StorageService; import software.amazon.jdbc.util.telemetry.TelemetryFactory; class MonitorServiceImplTest { - @Mock StorageService storageService; - @Mock TelemetryFactory telemetryFactory; - @Mock TargetDriverDialect targetDriverDialect; - @Mock Dialect dbDialect; - @Mock EventPublisher publisher; - MonitorServiceImpl monitorService; + @Mock StorageService mockStorageService; + @Mock ConnectionService mockConnectionService; + @Mock TelemetryFactory mockTelemetryFactory; + @Mock TargetDriverDialect mockTargetDriverDialect; + @Mock Dialect mockDbDialect; + @Mock EventPublisher mockPublisher; + MonitorServiceImpl spyMonitorService; private AutoCloseable closeable; @BeforeEach void setUp() { closeable = MockitoAnnotations.openMocks(this); - monitorService = new MonitorServiceImpl(publisher) { - @Override - protected void initCleanupThread(long cleanupIntervalNanos) { - // Do nothing - } - }; + spyMonitorService = spy(new MonitorServiceImpl(mockPublisher)); + doNothing().when(spyMonitorService).initCleanupThread(anyInt()); + + try { + doReturn(mockConnectionService).when(spyMonitorService) + .getConnectionService(any(), any(), any(), any(), any(), any(), any()); + } catch (SQLException e) { + Assertions.fail( + "Encountered exception while stubbing MonitorServiceImpl#getConnectionService: " + e.getMessage()); + } } @AfterEach void tearDown() throws Exception { closeable.close(); - monitorService.releaseResources(); + spyMonitorService.releaseResources(); } @Test public void testMonitorError_monitorReCreated() throws SQLException, InterruptedException { - monitorService.registerMonitorTypeIfAbsent( + spyMonitorService.registerMonitorTypeIfAbsent( NoOpMonitor.class, TimeUnit.MINUTES.toNanos(1), TimeUnit.MINUTES.toNanos(1), @@ -75,20 +87,20 @@ public void testMonitorError_monitorReCreated() throws SQLException, Interrupted null ); String key = "testMonitor"; - NoOpMonitor monitor = monitorService.runIfAbsent( + NoOpMonitor monitor = spyMonitorService.runIfAbsent( NoOpMonitor.class, key, - storageService, - telemetryFactory, + mockStorageService, + mockTelemetryFactory, "jdbc:postgresql://somehost/somedb", "someProtocol", - targetDriverDialect, - dbDialect, + mockTargetDriverDialect, + mockDbDialect, new Properties(), - (connectionService, pluginService) -> new NoOpMonitor(monitorService, 30) + (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) ); - Monitor storedMonitor = monitorService.get(NoOpMonitor.class, key); + Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key); assertNotNull(storedMonitor); assertEquals(monitor, storedMonitor); // need to wait to give time for the monitor executor to start the monitor thread. @@ -96,11 +108,11 @@ public void testMonitorError_monitorReCreated() throws SQLException, Interrupted assertEquals(MonitorState.RUNNING, monitor.getState()); monitor.state.set(MonitorState.ERROR); - monitorService.checkMonitors(); + spyMonitorService.checkMonitors(); assertEquals(MonitorState.STOPPED, monitor.getState()); - Monitor newMonitor = monitorService.get(NoOpMonitor.class, key); + Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key); assertNotNull(newMonitor); assertNotEquals(monitor, newMonitor); // need to wait to give time for the monitor executor to start the monitor thread. @@ -110,7 +122,7 @@ public void testMonitorError_monitorReCreated() throws SQLException, Interrupted @Test public void testMonitorStuck_monitorReCreated() throws SQLException, InterruptedException { - monitorService.registerMonitorTypeIfAbsent( + spyMonitorService.registerMonitorTypeIfAbsent( NoOpMonitor.class, TimeUnit.MINUTES.toNanos(1), 1, // heartbeat times out immediately @@ -118,20 +130,20 @@ public void testMonitorStuck_monitorReCreated() throws SQLException, Interrupted null ); String key = "testMonitor"; - NoOpMonitor monitor = monitorService.runIfAbsent( + NoOpMonitor monitor = spyMonitorService.runIfAbsent( NoOpMonitor.class, key, - storageService, - telemetryFactory, + mockStorageService, + mockTelemetryFactory, "jdbc:postgresql://somehost/somedb", "someProtocol", - targetDriverDialect, - dbDialect, + mockTargetDriverDialect, + mockDbDialect, new Properties(), - (connectionService, pluginService) -> new NoOpMonitor(monitorService, 30) + (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) ); - Monitor storedMonitor = monitorService.get(NoOpMonitor.class, key); + Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key); assertNotNull(storedMonitor); assertEquals(monitor, storedMonitor); // need to wait to give time for the monitor executor to start the monitor thread. @@ -139,11 +151,11 @@ public void testMonitorStuck_monitorReCreated() throws SQLException, Interrupted assertEquals(MonitorState.RUNNING, monitor.getState()); // checkMonitors() should detect the heartbeat/inactivity timeout, stop the monitor, and re-create a new one. - monitorService.checkMonitors(); + spyMonitorService.checkMonitors(); assertEquals(MonitorState.STOPPED, monitor.getState()); - Monitor newMonitor = monitorService.get(NoOpMonitor.class, key); + Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key); assertNotNull(newMonitor); assertNotEquals(monitor, newMonitor); // need to wait to give time for the monitor executor to start the monitor thread. @@ -153,7 +165,7 @@ public void testMonitorStuck_monitorReCreated() throws SQLException, Interrupted @Test public void testMonitorExpired() throws SQLException, InterruptedException { - monitorService.registerMonitorTypeIfAbsent( + spyMonitorService.registerMonitorTypeIfAbsent( NoOpMonitor.class, TimeUnit.MILLISECONDS.toNanos(200), // monitor expires after 200ms TimeUnit.MINUTES.toNanos(1), @@ -163,20 +175,20 @@ public void testMonitorExpired() throws SQLException, InterruptedException { null ); String key = "testMonitor"; - NoOpMonitor monitor = monitorService.runIfAbsent( + NoOpMonitor monitor = spyMonitorService.runIfAbsent( NoOpMonitor.class, key, - storageService, - telemetryFactory, + mockStorageService, + mockTelemetryFactory, "jdbc:postgresql://somehost/somedb", "someProtocol", - targetDriverDialect, - dbDialect, + mockTargetDriverDialect, + mockDbDialect, new Properties(), - (connectionService, pluginService) -> new NoOpMonitor(monitorService, 30) + (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) ); - Monitor storedMonitor = monitorService.get(NoOpMonitor.class, key); + Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key); assertNotNull(storedMonitor); assertEquals(monitor, storedMonitor); // need to wait to give time for the monitor executor to start the monitor thread. @@ -184,36 +196,36 @@ public void testMonitorExpired() throws SQLException, InterruptedException { assertEquals(MonitorState.RUNNING, monitor.getState()); // checkMonitors() should detect the expiration timeout and stop/remove the monitor. - monitorService.checkMonitors(); + spyMonitorService.checkMonitors(); assertEquals(MonitorState.STOPPED, monitor.getState()); - Monitor newMonitor = monitorService.get(NoOpMonitor.class, key); + Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key); // monitor should have been removed when checkMonitors() was called. assertNull(newMonitor); } @Test public void testMonitorMismatch() { - assertThrows(IllegalStateException.class, () -> monitorService.runIfAbsent( + assertThrows(IllegalStateException.class, () -> spyMonitorService.runIfAbsent( CustomEndpointMonitorImpl.class, "testMonitor", - storageService, - telemetryFactory, + mockStorageService, + mockTelemetryFactory, "jdbc:postgresql://somehost/somedb", "someProtocol", - targetDriverDialect, - dbDialect, + mockTargetDriverDialect, + mockDbDialect, new Properties(), // indicated monitor class is CustomEndpointMonitorImpl, but actual monitor is NoOpMonitor. The monitor // service should detect this and throw an exception. - (connectionService, pluginService) -> new NoOpMonitor(monitorService, 30) + (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) )); } @Test public void testRemove() throws SQLException, InterruptedException { - monitorService.registerMonitorTypeIfAbsent( + spyMonitorService.registerMonitorTypeIfAbsent( NoOpMonitor.class, TimeUnit.MINUTES.toNanos(1), TimeUnit.MINUTES.toNanos(1), @@ -224,30 +236,30 @@ public void testRemove() throws SQLException, InterruptedException { ); String key = "testMonitor"; - NoOpMonitor monitor = monitorService.runIfAbsent( + NoOpMonitor monitor = spyMonitorService.runIfAbsent( NoOpMonitor.class, key, - storageService, - telemetryFactory, + mockStorageService, + mockTelemetryFactory, "jdbc:postgresql://somehost/somedb", "someProtocol", - targetDriverDialect, - dbDialect, + mockTargetDriverDialect, + mockDbDialect, new Properties(), - (connectionService, pluginService) -> new NoOpMonitor(monitorService, 30) + (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) ); assertNotNull(monitor); // need to wait to give time for the monitor executor to start the monitor thread. TimeUnit.MILLISECONDS.sleep(250); - Monitor removedMonitor = monitorService.remove(NoOpMonitor.class, key); + Monitor removedMonitor = spyMonitorService.remove(NoOpMonitor.class, key); assertEquals(monitor, removedMonitor); assertEquals(MonitorState.RUNNING, monitor.getState()); } @Test public void testStopAndRemove() throws SQLException, InterruptedException { - monitorService.registerMonitorTypeIfAbsent( + spyMonitorService.registerMonitorTypeIfAbsent( NoOpMonitor.class, TimeUnit.MINUTES.toNanos(1), TimeUnit.MINUTES.toNanos(1), @@ -258,24 +270,24 @@ public void testStopAndRemove() throws SQLException, InterruptedException { ); String key = "testMonitor"; - NoOpMonitor monitor = monitorService.runIfAbsent( + NoOpMonitor monitor = spyMonitorService.runIfAbsent( NoOpMonitor.class, key, - storageService, - telemetryFactory, + mockStorageService, + mockTelemetryFactory, "jdbc:postgresql://somehost/somedb", "someProtocol", - targetDriverDialect, - dbDialect, + mockTargetDriverDialect, + mockDbDialect, new Properties(), - (connectionService, pluginService) -> new NoOpMonitor(monitorService, 30) + (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) ); assertNotNull(monitor); // need to wait to give time for the monitor executor to start the monitor thread. TimeUnit.MILLISECONDS.sleep(250); - monitorService.stopAndRemove(NoOpMonitor.class, key); - assertNull(monitorService.get(NoOpMonitor.class, key)); + spyMonitorService.stopAndRemove(NoOpMonitor.class, key); + assertNull(spyMonitorService.get(NoOpMonitor.class, key)); assertEquals(MonitorState.STOPPED, monitor.getState()); } From 9928517af0e820b76db1923d31367d09a7fa5ead Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Tue, 26 Aug 2025 09:00:16 -0700 Subject: [PATCH 18/42] Fix failover1 integration test --- .../jdbc/plugin/failover/FailoverConnectionPlugin.java | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java index 8bbdbd5fd..19eb59a48 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java @@ -758,6 +758,14 @@ protected void failoverWriter() throws SQLException { this.failoverWriterTriggeredCounter.inc(); } + // The writer failover handler uses the reader failover handler, so we need to make sure it has been instantiated. + if (this.readerFailoverHandler == null) { + if (this.readerFailoverHandlerSupplier == null) { + throw new SQLException(Messages.get("Failover.nullReaderFailoverHandlerSupplier")); + } + this.readerFailoverHandler = this.readerFailoverHandlerSupplier.apply(this.connectionService); + } + if (this.writerFailoverHandler == null) { if (this.writerFailoverHandlerSupplier == null) { throw new SQLException(Messages.get("Failover.nullWriterFailoverHandlerSupplier")); From 6a338f2c51436028ccee113a748c0b4b57037093 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Tue, 26 Aug 2025 09:14:32 -0700 Subject: [PATCH 19/42] testServerFailoverWithIdleConnections uses the correct failover plugin --- .../src/test/java/integration/container/tests/FailoverTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/wrapper/src/test/java/integration/container/tests/FailoverTest.java b/wrapper/src/test/java/integration/container/tests/FailoverTest.java index 1c9359fd5..347264561 100644 --- a/wrapper/src/test/java/integration/container/tests/FailoverTest.java +++ b/wrapper/src/test/java/integration/container/tests/FailoverTest.java @@ -328,7 +328,7 @@ public void testServerFailoverWithIdleConnections() throws SQLException, Interru TestEnvironment.getCurrent().getInfo().getProxyDatabaseInfo().getClusterEndpointPort(); final Properties props = initDefaultProxiedProps(); - props.setProperty(PropertyDefinition.PLUGINS.name, "auroraConnectionTracker,failover"); + props.setProperty(PropertyDefinition.PLUGINS.name, "auroraConnectionTracker," + getFailoverPlugin()); for (int i = 0; i < IDLE_CONNECTIONS_NUM; i++) { // Keep references to 5 idle connections created using the cluster endpoints. From faa09158d0bba9b6f833a16c38b32e9670b32512 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Tue, 26 Aug 2025 10:06:01 -0700 Subject: [PATCH 20/42] Fix javadocs --- .../ClusterAwareReaderFailoverHandler.java | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java index 8f3ab2f62..e64484959 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java @@ -76,8 +76,8 @@ public class ClusterAwareReaderFailoverHandler implements ReaderFailoverHandler /** * ClusterAwareReaderFailoverHandler constructor. * - * @param servicesContainer A provider for creating new connections. - * @param props The initial connection properties to copy over to the new reader. + * @param servicesContainer the service container for the services required by this class. + * @param props the initial connection properties to copy over to the new reader. */ public ClusterAwareReaderFailoverHandler( final FullServicesContainer servicesContainer, @@ -95,12 +95,12 @@ public ClusterAwareReaderFailoverHandler( /** * ClusterAwareReaderFailoverHandler constructor. * - * @param servicesContainer A provider for creating new connections. - * @param props The initial connection properties to copy over to the new reader. - * @param maxFailoverTimeoutMs Maximum allowed time for the entire reader failover process. - * @param timeoutMs Maximum allowed time in milliseconds for each reader connection attempt during - * the reader failover process. - * @param isStrictReaderRequired When true, it disables adding a writer to a list of nodes to connect + * @param servicesContainer the service container for the services required by this class. + * @param props the initial connection properties to copy over to the new reader. + * @param maxFailoverTimeoutMs maximum allowed time for the entire reader failover process. + * @param timeoutMs maximum allowed time in milliseconds for each reader connection attempt during + * the reader failover process. + * @param isStrictReaderRequired when true, it disables adding a writer to a list of nodes to connect */ public ClusterAwareReaderFailoverHandler( final FullServicesContainer servicesContainer, From 23e177af7614e4ab38ef8f646d2f0b6b9f2f19bb Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Tue, 26 Aug 2025 09:59:14 -0700 Subject: [PATCH 21/42] wip --- .../jdbc/util/ServiceContainerUtility.java | 51 +++++++++++++++++++ 1 file changed, 51 insertions(+) create mode 100644 wrapper/src/main/java/software/amazon/jdbc/util/ServiceContainerUtility.java diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/ServiceContainerUtility.java b/wrapper/src/main/java/software/amazon/jdbc/util/ServiceContainerUtility.java new file mode 100644 index 000000000..fb4bc4d70 --- /dev/null +++ b/wrapper/src/main/java/software/amazon/jdbc/util/ServiceContainerUtility.java @@ -0,0 +1,51 @@ +/* + * 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.util; + +import java.util.concurrent.locks.ReentrantLock; + +public class ServiceContainerUtility { + private static volatile ServiceContainerUtility instance; + private static final ReentrantLock initLock = new ReentrantLock(); + + private ServiceContainerUtility() { + if (instance != null) { + throw new IllegalStateException("ServiceContainerUtility singleton instance already exists."); + } + } + + public static ServiceContainerUtility getInstance() { + if (instance != null) { + return instance; + } + + initLock.lock(); + try { + if (instance == null) { + instance = new ServiceContainerUtility(); + } + } finally { + initLock.unlock(); + } + + return instance; + } + + public static FullServicesContainer createServiceContainer() { + + } +} From e1af3eb2757e3e0223c8863cbb368e093dc8018d Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Thu, 28 Aug 2025 18:12:54 -0700 Subject: [PATCH 22/42] Add extra logging to debug IT failure --- .../plugin/readwritesplitting/ReadWriteSplittingPlugin.java | 3 +++ 1 file changed, 3 insertions(+) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java index 7a2955843..35d768211 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java @@ -422,11 +422,14 @@ private void switchToReaderConnection(final List hosts) if (this.readerHostSpec != null && !hosts.contains(this.readerHostSpec)) { // The old reader cannot be used anymore because it is no longer in the list of allowed hosts. + LOGGER.finest(String.format("ASDF old reader cannot be used, reader host spec: %s", this.readerHostSpec)); + LOGGER.finest(Utils.logTopology(hosts, "ASDF: ")); closeConnectionIfIdle(this.readerConnection); } this.inReadWriteSplit = true; if (!isConnectionUsable(this.readerConnection)) { + LOGGER.finest(String.format("ASDF reader connection not usable, value: %s, isClosed: %s", this.readerConnection, this.readerConnection == null ? "" : this.readerConnection.isClosed())); initializeReaderConnection(hosts); } else { try { From c077ee218977cfbb97f432e973949fd2b1f5937d Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Fri, 29 Aug 2025 09:19:45 -0700 Subject: [PATCH 23/42] Fix bug where cached reader connection was incorrectly closed --- .../readwritesplitting/ReadWriteSplittingPlugin.java | 11 ++++++----- .../aws_advanced_jdbc_wrapper_messages.properties | 1 + 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java index 35d768211..90b170f98 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java @@ -420,16 +420,17 @@ private void switchToReaderConnection(final List hosts) return; } - if (this.readerHostSpec != null && !hosts.contains(this.readerHostSpec)) { - // The old reader cannot be used anymore because it is no longer in the list of allowed hosts. - LOGGER.finest(String.format("ASDF old reader cannot be used, reader host spec: %s", this.readerHostSpec)); - LOGGER.finest(Utils.logTopology(hosts, "ASDF: ")); + if (this.readerHostSpec != null && !Utils.containsUrl(hosts, this.readerHostSpec.getUrl())) { + // The previous reader cannot be used anymore because it is no longer in the list of allowed hosts. + LOGGER.finest( + Messages.get( + "ReadWriteSplittingPlugin.previousReaderNotAllowed", + new Object[] {this.readerHostSpec, Utils.logTopology(hosts, "")})); closeConnectionIfIdle(this.readerConnection); } this.inReadWriteSplit = true; if (!isConnectionUsable(this.readerConnection)) { - LOGGER.finest(String.format("ASDF reader connection not usable, value: %s, isClosed: %s", this.readerConnection, this.readerConnection == null ? "" : this.readerConnection.isClosed())); initializeReaderConnection(hosts); } else { try { 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 6bb12413d..1776fcced 100644 --- a/wrapper/src/main/resources/aws_advanced_jdbc_wrapper_messages.properties +++ b/wrapper/src/main/resources/aws_advanced_jdbc_wrapper_messages.properties @@ -316,6 +316,7 @@ ReadWriteSplittingPlugin.successfullyConnectedToReader=Successfully connected to ReadWriteSplittingPlugin.failedToConnectToReader=Failed to connect to reader host: ''{0}'' ReadWriteSplittingPlugin.unsupportedHostSpecSelectorStrategy=Unsupported host selection strategy ''{0}'' specified in plugin configuration parameter ''readerHostSelectorStrategy''. Please visit the Read/Write Splitting Plugin documentation for all supported strategies. ReadWriteSplittingPlugin.errorVerifyingInitialHostSpecRole=An error occurred while obtaining the connected host's role. This could occur if the connection is broken or if you are not connected to an Aurora database. +ReadWriteSplittingPlugin.previousReaderNotAllowed=The previous reader connection cannot be used because it is no longer in the list of allowed hosts. Previous reader: {0}. Allowed hosts: {1} SAMLCredentialsProviderFactory.getSamlAssertionFailed=Failed to get SAML Assertion due to exception: ''{0}'' SamlAuthPlugin.javaStsSdkNotInClasspath=Required dependency 'AWS Java SDK for AWS Secret Token Service' is not on the classpath. From 73cf143cf1da8f55b77421472d731512eb8b218b Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Fri, 29 Aug 2025 13:43:38 -0700 Subject: [PATCH 24/42] wip --- .../ClusterAwareReaderFailoverHandler.java | 2 - .../failover/FailoverConnectionPlugin.java | 6 +-- .../jdbc/util/ServiceContainerUtility.java | 47 ++++++++++++++++++- 3 files changed, 48 insertions(+), 7 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java index 8f3ab2f62..0304252d3 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java @@ -81,11 +81,9 @@ public class ClusterAwareReaderFailoverHandler implements ReaderFailoverHandler */ public ClusterAwareReaderFailoverHandler( final FullServicesContainer servicesContainer, - final ConnectionService connectionService, final Properties props) { this( servicesContainer, - connectionService, props, DEFAULT_FAILOVER_TIMEOUT, DEFAULT_READER_CONNECT_TIMEOUT, diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java index 19eb59a48..b0acf9de0 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java @@ -316,18 +316,16 @@ public void initHostProvider( initHostProvider( hostListProviderService, initHostProviderFunc, - (connectionService) -> + () -> new ClusterAwareReaderFailoverHandler( this.servicesContainer, - connectionService, this.properties, this.failoverTimeoutMsSetting, this.failoverReaderConnectTimeoutMsSetting, this.failoverMode == FailoverMode.STRICT_READER), - (connectionService) -> + () -> new ClusterAwareWriterFailoverHandler( this.servicesContainer, - connectionService, this.readerFailoverHandler, this.properties, this.failoverTimeoutMsSetting, diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/ServiceContainerUtility.java b/wrapper/src/main/java/software/amazon/jdbc/util/ServiceContainerUtility.java index fb4bc4d70..e335b5c5f 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/ServiceContainerUtility.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/ServiceContainerUtility.java @@ -16,7 +16,17 @@ package software.amazon.jdbc.util; +import java.sql.SQLException; +import java.util.Properties; import java.util.concurrent.locks.ReentrantLock; +import software.amazon.jdbc.ConnectionPluginManager; +import software.amazon.jdbc.ConnectionProvider; +import software.amazon.jdbc.PartialPluginService; +import software.amazon.jdbc.dialect.Dialect; +import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; +import software.amazon.jdbc.util.monitoring.MonitorService; +import software.amazon.jdbc.util.storage.StorageService; +import software.amazon.jdbc.util.telemetry.TelemetryFactory; public class ServiceContainerUtility { private static volatile ServiceContainerUtility instance; @@ -45,7 +55,42 @@ public static ServiceContainerUtility getInstance() { return instance; } - public static FullServicesContainer createServiceContainer() { + public static FullServicesContainer createServiceContainer( + StorageService storageService, + MonitorService monitorService, + TelemetryFactory telemetryFactory, + ConnectionProvider connectionProvider, + String originalUrl, + String targetDriverProtocol, + TargetDriverDialect driverDialect, + Dialect dbDialect, + Properties props) throws SQLException { + FullServicesContainer + servicesContainer = new FullServicesContainerImpl(storageService, monitorService, telemetryFactory); + ConnectionPluginManager pluginManager = new ConnectionPluginManager( + connectionProvider, + null, + null, + telemetryFactory); + servicesContainer.setConnectionPluginManager(pluginManager); + PartialPluginService partialPluginService = new PartialPluginService( + servicesContainer, + props, + originalUrl, + targetDriverProtocol, + driverDialect, + dbDialect + ); + + pluginManager.init(servicesContainer, props, partialPluginService, null); + return new FullServicesContainerImpl( + storageService, + monitorService, + telemetryFactory, + pluginManager, + partialPluginService, + partialPluginService, + partialPluginService); } } From 99de6eb235a66040b1ff4dba2314ac1f85feedb1 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Fri, 29 Aug 2025 13:49:54 -0700 Subject: [PATCH 25/42] Fix javadoc --- .../plugin/failover/ClusterAwareReaderFailoverHandler.java | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java index e64484959..35f16b35f 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java @@ -77,6 +77,7 @@ public class ClusterAwareReaderFailoverHandler implements ReaderFailoverHandler * ClusterAwareReaderFailoverHandler constructor. * * @param servicesContainer the service container for the services required by this class. + * @param connectionService the service to use to create new connections during failover. * @param props the initial connection properties to copy over to the new reader. */ public ClusterAwareReaderFailoverHandler( @@ -96,6 +97,7 @@ public ClusterAwareReaderFailoverHandler( * ClusterAwareReaderFailoverHandler constructor. * * @param servicesContainer the service container for the services required by this class. + * @param connectionService the service to use to create new connections during failover. * @param props the initial connection properties to copy over to the new reader. * @param maxFailoverTimeoutMs maximum allowed time for the entire reader failover process. * @param timeoutMs maximum allowed time in milliseconds for each reader connection attempt during @@ -292,7 +294,7 @@ public List getReaderHostsByPriority(final List hosts) { boolean shouldIncludeWriter = numOfReaders == 0 || this.pluginService.getDialect().getFailoverRestrictions() - .contains(FailoverRestriction.ENABLE_WRITER_IN_TASK_B); + .contains(FailoverRestriction.ENABLE_WRITER_IN_TASK_B); if (shouldIncludeWriter) { hostsByPriority.add(writerHost); } From 9311447abce720fa3bc8c11f9a1c053d87ff88e5 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Fri, 29 Aug 2025 16:07:22 -0700 Subject: [PATCH 26/42] wip --- .../ClusterTopologyMonitorImpl.java | 47 +- .../MonitoringRdsHostListProvider.java | 9 +- .../ClusterAwareReaderFailoverHandler.java | 58 +- .../ClusterAwareWriterFailoverHandler.java | 55 +- .../failover/FailoverConnectionPlugin.java | 43 +- .../jdbc/util/ServiceContainerUtility.java | 9 +- .../util/monitoring/MonitorInitializer.java | 7 +- .../util/monitoring/MonitorServiceImpl.java | 43 +- ...ClusterAwareReaderFailoverHandlerTest.java | 802 +++++++++--------- ...ClusterAwareWriterFailoverHandlerTest.java | 746 ++++++++-------- .../monitoring/MonitorServiceImplTest.java | 612 ++++++------- 11 files changed, 1193 insertions(+), 1238 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/ClusterTopologyMonitorImpl.java b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/ClusterTopologyMonitorImpl.java index 4e7b62de1..e61cb6a60 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/ClusterTopologyMonitorImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/ClusterTopologyMonitorImpl.java @@ -49,14 +49,17 @@ import software.amazon.jdbc.hostavailability.HostAvailability; import software.amazon.jdbc.hostlistprovider.Topology; import software.amazon.jdbc.util.ExecutorFactory; +import software.amazon.jdbc.util.FullServicesContainer; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.PropertyUtils; import software.amazon.jdbc.util.RdsUtils; +import software.amazon.jdbc.util.ServiceContainerUtility; import software.amazon.jdbc.util.StringUtils; import software.amazon.jdbc.util.SynchronousExecutor; import software.amazon.jdbc.util.Utils; import software.amazon.jdbc.util.connection.ConnectionService; import software.amazon.jdbc.util.monitoring.AbstractMonitor; +import software.amazon.jdbc.util.monitoring.Monitor; import software.amazon.jdbc.util.storage.StorageService; public class ClusterTopologyMonitorImpl extends AbstractMonitor implements ClusterTopologyMonitor { @@ -80,15 +83,13 @@ public class ClusterTopologyMonitorImpl extends AbstractMonitor implements Clust protected final long refreshRateNano; protected final long highRefreshRateNano; + protected final FullServicesContainer servicesContainer; protected final Properties properties; protected final Properties monitoringProperties; protected final HostSpec initialHostSpec; - protected final StorageService storageService; - protected final ConnectionService connectionService; protected final String topologyQuery; protected final String nodeIdQuery; protected final String writerTopologyQuery; - protected final HostListProviderService hostListProviderService; protected final HostSpec clusterInstanceTemplate; protected String clusterId; @@ -109,12 +110,10 @@ public class ClusterTopologyMonitorImpl extends AbstractMonitor implements Clust protected final AtomicReference> nodeThreadsLatestTopology = new AtomicReference<>(null); public ClusterTopologyMonitorImpl( + final FullServicesContainer servicesContainer, final String clusterId, - final StorageService storageService, - final ConnectionService connectionService, final HostSpec initialHostSpec, final Properties properties, - final HostListProviderService hostListProviderService, final HostSpec clusterInstanceTemplate, final long refreshRateNano, final long highRefreshRateNano, @@ -124,9 +123,7 @@ public ClusterTopologyMonitorImpl( super(monitorTerminationTimeoutSec); this.clusterId = clusterId; - this.storageService = storageService; - this.connectionService = connectionService; - this.hostListProviderService = hostListProviderService; + this.servicesContainer = servicesContainer; this.initialHostSpec = initialHostSpec; this.clusterInstanceTemplate = clusterInstanceTemplate; this.properties = properties; @@ -251,7 +248,7 @@ protected List waitTillTopologyGetsUpdated(final long timeoutMs) throw } private List getStoredHosts() { - Topology topology = storageService.get(Topology.class, this.clusterId); + Topology topology = this.servicesContainer.getStorageService().get(Topology.class, this.clusterId); return topology == null ? null : topology.getHosts(); } @@ -480,8 +477,23 @@ protected boolean isInPanicMode() { || !this.isVerifiedWriterConnection; } - protected Runnable getNodeMonitoringWorker(final HostSpec hostSpec, final @Nullable HostSpec writerHostSpec) { - return new NodeMonitoringWorker(this, hostSpec, writerHostSpec); + protected Runnable getNodeMonitoringWorker( + final HostSpec hostSpec, final @Nullable HostSpec writerHostSpec) + throws SQLException { + return new NodeMonitoringWorker(this.getNewServicesContainer(), this, hostSpec, writerHostSpec); + } + + protected FullServicesContainer getNewServicesContainer() throws SQLException { + return ServiceContainerUtility.createServiceContainer( + this.servicesContainer.getStorageService(), + this.servicesContainer.getMonitorService(), + this.servicesContainer.getTelemetryFactory(), + this.servicesContainer.getPluginService().getOriginalUrl(), + this.servicesContainer.getPluginService().getDriverProtocol(), + this.servicesContainer.getPluginService().getTargetDriverDialect(), + this.servicesContainer.getPluginService().getDialect(), + this.properties + ); } protected List openAnyConnectionAndUpdateTopology() { @@ -492,7 +504,7 @@ protected List openAnyConnectionAndUpdateTopology() { // open a new connection try { - conn = this.connectionService.open(this.initialHostSpec, this.monitoringProperties); + conn = this.servicesContainer.getPluginService().forceConnect(this.initialHostSpec, this.monitoringProperties); } catch (SQLException ex) { // can't connect return null; @@ -625,7 +637,7 @@ protected void delay(boolean useHighRefreshRate) throws InterruptedException { protected void updateTopologyCache(final @NonNull List hosts) { synchronized (this.requestToUpdateTopology) { - storageService.set(this.clusterId, new Topology(hosts)); + this.servicesContainer.getStorageService().set(this.clusterId, new Topology(hosts)); synchronized (this.topologyUpdated) { this.requestToUpdateTopology.set(false); @@ -769,7 +781,7 @@ protected HostSpec createHost( ? this.clusterInstanceTemplate.getPort() : this.initialHostSpec.getPort(); - final HostSpec hostSpec = this.hostListProviderService.getHostSpecBuilder() + final HostSpec hostSpec = this.servicesContainer.getHostListProviderService().getHostSpecBuilder() .host(endpoint) .port(port) .role(isWriter ? HostRole.WRITER : HostRole.READER) @@ -791,16 +803,19 @@ private static class NodeMonitoringWorker implements Runnable { private static final Logger LOGGER = Logger.getLogger(NodeMonitoringWorker.class.getName()); + protected final FullServicesContainer servicesContainer; protected final ClusterTopologyMonitorImpl monitor; protected final HostSpec hostSpec; protected final @Nullable HostSpec writerHostSpec; protected boolean writerChanged = false; public NodeMonitoringWorker( + final FullServicesContainer servicesContainer, final ClusterTopologyMonitorImpl monitor, final HostSpec hostSpec, final @Nullable HostSpec writerHostSpec ) { + this.servicesContainer = servicesContainer; this.monitor = monitor; this.hostSpec = hostSpec; this.writerHostSpec = writerHostSpec; @@ -818,7 +833,7 @@ public void run() { if (connection == null) { try { - connection = this.monitor.connectionService.open( + connection = this.servicesContainer.getPluginService().forceConnect( hostSpec, this.monitor.monitoringProperties); } catch (SQLException ex) { // A problem occurred while connecting. We will try again on the next iteration. diff --git a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/MonitoringRdsHostListProvider.java b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/MonitoringRdsHostListProvider.java index 29aae6ddd..2883100e0 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/MonitoringRdsHostListProvider.java +++ b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/MonitoringRdsHostListProvider.java @@ -32,7 +32,6 @@ import software.amazon.jdbc.hostlistprovider.RdsHostListProvider; import software.amazon.jdbc.hostlistprovider.Topology; import software.amazon.jdbc.util.FullServicesContainer; -import software.amazon.jdbc.util.connection.ConnectionService; import software.amazon.jdbc.util.monitoring.MonitorService; import software.amazon.jdbc.util.storage.StorageService; @@ -92,13 +91,11 @@ protected ClusterTopologyMonitor initMonitor() throws SQLException { this.pluginService.getTargetDriverDialect(), this.pluginService.getDialect(), this.properties, - (ConnectionService connectionService, PluginService monitorPluginService) -> new ClusterTopologyMonitorImpl( + (servicesContainer) -> new ClusterTopologyMonitorImpl( + this.servicesContainer, this.clusterId, - this.servicesContainer.getStorageService(), - connectionService, this.initialHostSpec, this.properties, - this.servicesContainer.getHostListProviderService(), this.clusterInstanceTemplate, this.refreshRateNano, this.highRefreshRateNano, @@ -138,7 +135,7 @@ protected void clusterIdChanged(final String oldClusterId) throws SQLException { this.pluginService.getTargetDriverDialect(), this.pluginService.getDialect(), this.properties, - (connectionService, pluginService) -> existingMonitor); + (servicesContainer) -> existingMonitor); assert monitorService.get(ClusterTopologyMonitorImpl.class, this.clusterId) == existingMonitor; existingMonitor.setClusterId(this.clusterId); monitorService.remove(ClusterTopologyMonitorImpl.class, oldClusterId); diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java index 6045e00ba..1ad010af8 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java @@ -36,15 +36,14 @@ import java.util.logging.Logger; import software.amazon.jdbc.HostRole; import software.amazon.jdbc.HostSpec; -import software.amazon.jdbc.PartialPluginService; import software.amazon.jdbc.PluginService; import software.amazon.jdbc.hostavailability.HostAvailability; import software.amazon.jdbc.util.ExecutorFactory; import software.amazon.jdbc.util.FullServicesContainer; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.PropertyUtils; +import software.amazon.jdbc.util.ServiceContainerUtility; import software.amazon.jdbc.util.Utils; -import software.amazon.jdbc.util.connection.ConnectionService; /** * An implementation of ReaderFailoverHandler. @@ -66,7 +65,6 @@ public class ClusterAwareReaderFailoverHandler implements ReaderFailoverHandler protected static final int DEFAULT_READER_CONNECT_TIMEOUT = 30000; // 30 sec protected final Map hostAvailabilityMap = new ConcurrentHashMap<>(); protected final FullServicesContainer servicesContainer; - protected final ConnectionService connectionService; protected final PluginService pluginService; protected Properties props; protected int maxFailoverTimeoutMs; @@ -77,7 +75,6 @@ public class ClusterAwareReaderFailoverHandler implements ReaderFailoverHandler * ClusterAwareReaderFailoverHandler constructor. * * @param servicesContainer the service container for the services required by this class. - * @param connectionService the service to use to create new connections during failover. * @param props the initial connection properties to copy over to the new reader. */ public ClusterAwareReaderFailoverHandler( @@ -95,7 +92,6 @@ public ClusterAwareReaderFailoverHandler( * ClusterAwareReaderFailoverHandler constructor. * * @param servicesContainer the service container for the services required by this class. - * @param connectionService the service to use to create new connections during failover. * @param props the initial connection properties to copy over to the new reader. * @param maxFailoverTimeoutMs maximum allowed time for the entire reader failover process. * @param timeoutMs maximum allowed time in milliseconds for each reader connection attempt during @@ -104,13 +100,11 @@ public ClusterAwareReaderFailoverHandler( */ public ClusterAwareReaderFailoverHandler( final FullServicesContainer servicesContainer, - final ConnectionService connectionService, final Properties props, final int maxFailoverTimeoutMs, final int timeoutMs, final boolean isStrictReaderRequired) { this.servicesContainer = servicesContainer; - this.connectionService = connectionService; this.pluginService = servicesContainer.getPluginService(); this.props = props; this.maxFailoverTimeoutMs = maxFailoverTimeoutMs; @@ -304,15 +298,16 @@ private ReaderFailoverResult getConnectionFromHostGroup(final List hos final ExecutorService executor = ExecutorFactory.newFixedThreadPool(2, "failover"); final CompletionService completionService = new ExecutorCompletionService<>(executor); - // The ConnectionAttemptTask threads should have their own plugin services since they execute concurrently and - // PluginService was not designed to be thread-safe. - List pluginServices = Arrays.asList(getNewPluginService(), getNewPluginService()); + // The ConnectionAttemptTask threads should each have their own services container since they execute concurrently + // and PluginService was not designed to be thread-safe. + List servicesContainers = + Arrays.asList(getNewServicesContainer(), getNewServicesContainer()); try { for (int i = 0; i < hosts.size(); i += 2) { // submit connection attempt tasks in batches of 2 final ReaderFailoverResult result = - getResultFromNextTaskBatch(hosts, executor, completionService, pluginServices, i); + getResultFromNextTaskBatch(hosts, executor, completionService, servicesContainers, i); if (result.isConnected() || result.getException() != null) { return result; } @@ -335,14 +330,13 @@ private ReaderFailoverResult getResultFromNextTaskBatch( final List hosts, final ExecutorService executor, final CompletionService completionService, - final List pluginServices, + final List servicesContainers, final int i) throws SQLException { ReaderFailoverResult result; final int numTasks = i + 1 < hosts.size() ? 2 : 1; completionService.submit( new ConnectionAttemptTask( - this.connectionService, - pluginServices.get(0), + servicesContainers.get(0), this.hostAvailabilityMap, hosts.get(i), this.props, @@ -350,8 +344,7 @@ private ReaderFailoverResult getResultFromNextTaskBatch( if (numTasks == 2) { completionService.submit( new ConnectionAttemptTask( - this.connectionService, - pluginServices.get(1), + servicesContainers.get(1), this.hostAvailabilityMap, hosts.get(i + 1), this.props, @@ -372,6 +365,19 @@ private ReaderFailoverResult getResultFromNextTaskBatch( return new ReaderFailoverResult(null, null, false); } + protected FullServicesContainer getNewServicesContainer() throws SQLException { + return ServiceContainerUtility.createServiceContainer( + this.servicesContainer.getStorageService(), + this.servicesContainer.getMonitorService(), + this.servicesContainer.getTelemetryFactory(), + this.pluginService.getOriginalUrl(), + this.pluginService.getDriverProtocol(), + this.pluginService.getTargetDriverDialect(), + this.pluginService.getDialect(), + this.props + ); + } + private ReaderFailoverResult getNextResult(final CompletionService service) throws SQLException { try { @@ -394,19 +400,7 @@ private ReaderFailoverResult getNextResult(final CompletionService { - private final ConnectionService connectionService; private final PluginService pluginService; private final Map availabilityMap; private final HostSpec newHost; @@ -414,14 +408,12 @@ private static class ConnectionAttemptTask implements Callable availabilityMap, final HostSpec newHost, final Properties props, final boolean isStrictReaderRequired) { - this.connectionService = connectionService; - this.pluginService = pluginService; + this.pluginService = servicesContainer.getPluginService(); this.availabilityMap = availabilityMap; this.newHost = newHost; this.props = props; @@ -442,7 +434,7 @@ public ReaderFailoverResult call() { final Properties copy = new Properties(); copy.putAll(props); - final Connection conn = this.connectionService.open(this.newHost, copy); + final Connection conn = this.pluginService.forceConnect(this.newHost, copy); this.availabilityMap.put(this.newHost.getHost(), HostAvailability.AVAILABLE); if (this.isStrictReaderRequired) { diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java index 1921f3238..28c614141 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java @@ -34,15 +34,14 @@ import java.util.logging.Logger; import software.amazon.jdbc.HostRole; import software.amazon.jdbc.HostSpec; -import software.amazon.jdbc.PartialPluginService; import software.amazon.jdbc.PluginService; import software.amazon.jdbc.hostavailability.HostAvailability; import software.amazon.jdbc.util.ExecutorFactory; import software.amazon.jdbc.util.FullServicesContainer; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.PropertyUtils; +import software.amazon.jdbc.util.ServiceContainerUtility; import software.amazon.jdbc.util.Utils; -import software.amazon.jdbc.util.connection.ConnectionService; /** * An implementation of WriterFailoverHandler. @@ -59,7 +58,6 @@ public class ClusterAwareWriterFailoverHandler implements WriterFailoverHandler protected final Properties initialConnectionProps; protected final FullServicesContainer servicesContainer; - protected final ConnectionService connectionService; protected final PluginService pluginService; protected final ReaderFailoverHandler readerFailoverHandler; protected final Map hostAvailabilityMap = new ConcurrentHashMap<>(); @@ -69,11 +67,9 @@ public class ClusterAwareWriterFailoverHandler implements WriterFailoverHandler public ClusterAwareWriterFailoverHandler( final FullServicesContainer servicesContainer, - final ConnectionService connectionService, final ReaderFailoverHandler readerFailoverHandler, final Properties initialConnectionProps) { this.servicesContainer = servicesContainer; - this.connectionService = connectionService; this.pluginService = servicesContainer.getPluginService(); this.readerFailoverHandler = readerFailoverHandler; this.initialConnectionProps = initialConnectionProps; @@ -81,7 +77,6 @@ public ClusterAwareWriterFailoverHandler( public ClusterAwareWriterFailoverHandler( final FullServicesContainer servicesContainer, - final ConnectionService connectionService, final ReaderFailoverHandler readerFailoverHandler, final Properties initialConnectionProps, final int failoverTimeoutMs, @@ -89,7 +84,6 @@ public ClusterAwareWriterFailoverHandler( final int reconnectWriterIntervalMs) { this( servicesContainer, - connectionService, readerFailoverHandler, initialConnectionProps); this.maxFailoverTimeoutMs = failoverTimeoutMs; @@ -103,8 +97,7 @@ public Map getHostAvailabilityMap() { } @Override - public WriterFailoverResult failover(final List currentTopology) - throws SQLException { + public WriterFailoverResult failover(final List currentTopology) throws SQLException { if (Utils.isNullOrEmpty(currentTopology)) { LOGGER.severe(() -> Messages.get("ClusterAwareWriterFailoverHandler.failoverCalledWithInvalidTopology")); return DEFAULT_RESULT; @@ -162,13 +155,12 @@ private void submitTasks( final List currentTopology, final ExecutorService executorService, final CompletionService completionService, - final boolean singleTask) { + final boolean singleTask) throws SQLException { final HostSpec writerHost = getWriter(currentTopology); if (!singleTask) { completionService.submit( new ReconnectToWriterHandler( - this.connectionService, - this.getNewPluginService(), + this.getNewServicesContainer(), this.hostAvailabilityMap, writerHost, this.initialConnectionProps, @@ -177,8 +169,7 @@ private void submitTasks( completionService.submit( new WaitForNewWriterHandler( - this.connectionService, - this.getNewPluginService(), + this.getNewServicesContainer(), this.hostAvailabilityMap, this.readerFailoverHandler, writerHost, @@ -189,16 +180,18 @@ private void submitTasks( executorService.shutdown(); } - protected PluginService getNewPluginService() { - // Each task should get its own PluginService since they execute concurrently and PluginService was not designed to - // be thread-safe. - return new PartialPluginService( - this.servicesContainer, - this.initialConnectionProps, + protected FullServicesContainer getNewServicesContainer() throws SQLException { + // Each task should get its own FullServicesContainer since they execute concurrently and PluginService was not + // designed to be thread-safe. + return ServiceContainerUtility.createServiceContainer( + this.servicesContainer.getStorageService(), + this.servicesContainer.getMonitorService(), + this.servicesContainer.getTelemetryFactory(), this.pluginService.getOriginalUrl(), this.pluginService.getDriverProtocol(), this.pluginService.getTargetDriverDialect(), - this.pluginService.getDialect() + this.pluginService.getDialect(), + this.initialConnectionProps ); } @@ -275,8 +268,6 @@ private SQLException createInterruptedException(final InterruptedException e) { * Internal class responsible for re-connecting to the current writer (aka TaskA). */ private static class ReconnectToWriterHandler implements Callable { - - private final ConnectionService connectionService; private final PluginService pluginService; private final Map availabilityMap; private final HostSpec originalWriterHost; @@ -284,14 +275,12 @@ private static class ReconnectToWriterHandler implements Callable availabilityMap, final HostSpec originalWriterHost, final Properties props, final int reconnectWriterIntervalMs) { - this.connectionService = connectionService; - this.pluginService = pluginService; + this.pluginService = servicesContainer.getPluginService(); this.availabilityMap = availabilityMap; this.originalWriterHost = originalWriterHost; this.props = props; @@ -315,7 +304,7 @@ public WriterFailoverResult call() { conn.close(); } - conn = this.connectionService.open(this.originalWriterHost, this.props); + conn = this.pluginService.forceConnect(this.originalWriterHost, this.props); this.pluginService.forceRefreshHostList(conn); latestTopology = this.pluginService.getAllHosts(); } catch (final SQLException exception) { @@ -371,8 +360,6 @@ private boolean isCurrentHostWriter(final List latestTopology) { * elected writer (aka TaskB). */ private static class WaitForNewWriterHandler implements Callable { - - private final ConnectionService connectionService; private final PluginService pluginService; private final Map availabilityMap; private final ReaderFailoverHandler readerFailoverHandler; @@ -385,16 +372,14 @@ private static class WaitForNewWriterHandler implements Callable availabilityMap, final ReaderFailoverHandler readerFailoverHandler, final HostSpec originalWriterHost, final Properties props, final int readTopologyIntervalMs, final List currentTopology) { - this.connectionService = connectionService; - this.pluginService = pluginService; + this.pluginService = servicesContainer.getPluginService(); this.availabilityMap = availabilityMap; this.readerFailoverHandler = readerFailoverHandler; this.originalWriterHost = originalWriterHost; @@ -537,7 +522,7 @@ private boolean connectToWriter(final HostSpec writerCandidate) { new Object[] {writerCandidate.getUrl()})); try { // connect to the new writer - this.currentConnection = this.connectionService.open(writerCandidate, this.props); + this.currentConnection = this.pluginService.forceConnect(writerCandidate, this.props); this.availabilityMap.put(writerCandidate.getHost(), HostAvailability.AVAILABLE); return true; } catch (final SQLException exception) { diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java index b0acf9de0..66cdcfb12 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java @@ -28,13 +28,11 @@ import java.util.Set; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; -import java.util.function.Function; +import java.util.function.Supplier; import java.util.logging.Level; import java.util.logging.Logger; import org.checkerframework.checker.nullness.qual.NonNull; import software.amazon.jdbc.AwsWrapperProperty; -import software.amazon.jdbc.ConnectionProvider; -import software.amazon.jdbc.DriverConnectionProvider; import software.amazon.jdbc.HostListProviderService; import software.amazon.jdbc.HostRole; import software.amazon.jdbc.HostSpec; @@ -44,7 +42,6 @@ import software.amazon.jdbc.PluginManagerService; import software.amazon.jdbc.PluginService; import software.amazon.jdbc.PropertyDefinition; -import software.amazon.jdbc.TargetDriverHelper; import software.amazon.jdbc.hostavailability.HostAvailability; import software.amazon.jdbc.plugin.AbstractConnectionPlugin; import software.amazon.jdbc.plugin.staledns.AuroraStaleDnsHelper; @@ -56,8 +53,6 @@ import software.amazon.jdbc.util.SqlState; import software.amazon.jdbc.util.Utils; import software.amazon.jdbc.util.WrapperUtils; -import software.amazon.jdbc.util.connection.ConnectionService; -import software.amazon.jdbc.util.connection.ConnectionServiceImpl; import software.amazon.jdbc.util.telemetry.TelemetryContext; import software.amazon.jdbc.util.telemetry.TelemetryCounter; import software.amazon.jdbc.util.telemetry.TelemetryFactory; @@ -99,7 +94,6 @@ public class FailoverConnectionPlugin extends AbstractConnectionPlugin { private final Set subscribedMethods; private final PluginService pluginService; private final FullServicesContainer servicesContainer; - private ConnectionService connectionService; protected final Properties properties; protected boolean enableFailoverSetting; protected boolean enableConnectFailover; @@ -122,8 +116,8 @@ public class FailoverConnectionPlugin extends AbstractConnectionPlugin { private RdsUrlType rdsUrlType = null; private HostListProviderService hostListProviderService; private final AuroraStaleDnsHelper staleDnsHelper; - private Function writerFailoverHandlerSupplier; - private Function readerFailoverHandlerSupplier; + private Supplier writerFailoverHandlerSupplier; + private Supplier readerFailoverHandlerSupplier; public static final AwsWrapperProperty FAILOVER_CLUSTER_TOPOLOGY_REFRESH_RATE_MS = new AwsWrapperProperty( @@ -336,8 +330,8 @@ public void initHostProvider( void initHostProvider( final HostListProviderService hostListProviderService, final JdbcCallable initHostProviderFunc, - final Function readerFailoverHandlerSupplier, - final Function writerFailoverHandlerSupplier) + final Supplier readerFailoverHandlerSupplier, + final Supplier writerFailoverHandlerSupplier) throws SQLException { this.readerFailoverHandlerSupplier = readerFailoverHandlerSupplier; this.writerFailoverHandlerSupplier = writerFailoverHandlerSupplier; @@ -618,10 +612,6 @@ protected void dealWithIllegalStateException( */ protected void failover(final HostSpec failedHost) throws SQLException { this.pluginService.setAvailability(failedHost.asAliases(), HostAvailability.NOT_AVAILABLE); - if (this.connectionService == null) { - this.connectionService = getConnectionService(); - } - if (this.failoverMode == FailoverMode.STRICT_WRITER) { failoverWriter(); } else { @@ -629,23 +619,6 @@ protected void failover(final HostSpec failedHost) throws SQLException { } } - protected ConnectionService getConnectionService() throws SQLException { - TargetDriverHelper helper = new TargetDriverHelper(); - java.sql.Driver driver = helper.getTargetDriver(this.pluginService.getOriginalUrl(), properties); - final ConnectionProvider defaultConnectionProvider = new DriverConnectionProvider(driver); - return new ConnectionServiceImpl( - servicesContainer.getStorageService(), - servicesContainer.getMonitorService(), - servicesContainer.getTelemetryFactory(), - defaultConnectionProvider, - this.pluginService.getOriginalUrl(), - this.pluginService.getDriverProtocol(), - this.pluginService.getTargetDriverDialect(), - this.pluginService.getDialect(), - properties - ); - } - protected void failoverReader(final HostSpec failedHostSpec) throws SQLException { TelemetryFactory telemetryFactory = this.pluginService.getTelemetryFactory(); TelemetryContext telemetryContext = telemetryFactory.openTelemetryContext( @@ -658,7 +631,7 @@ protected void failoverReader(final HostSpec failedHostSpec) throws SQLException if (this.readerFailoverHandlerSupplier == null) { throw new SQLException(Messages.get("Failover.nullReaderFailoverHandlerSupplier")); } - this.readerFailoverHandler = this.readerFailoverHandlerSupplier.apply(this.connectionService); + this.readerFailoverHandler = this.readerFailoverHandlerSupplier.get(); } final long failoverStartNano = System.nanoTime(); @@ -761,14 +734,14 @@ protected void failoverWriter() throws SQLException { if (this.readerFailoverHandlerSupplier == null) { throw new SQLException(Messages.get("Failover.nullReaderFailoverHandlerSupplier")); } - this.readerFailoverHandler = this.readerFailoverHandlerSupplier.apply(this.connectionService); + this.readerFailoverHandler = this.readerFailoverHandlerSupplier.get(); } if (this.writerFailoverHandler == null) { if (this.writerFailoverHandlerSupplier == null) { throw new SQLException(Messages.get("Failover.nullWriterFailoverHandlerSupplier")); } - this.writerFailoverHandler = this.writerFailoverHandlerSupplier.apply(this.connectionService); + this.writerFailoverHandler = this.writerFailoverHandlerSupplier.get(); } long failoverStartTimeNano = System.nanoTime(); diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/ServiceContainerUtility.java b/wrapper/src/main/java/software/amazon/jdbc/util/ServiceContainerUtility.java index e335b5c5f..e0e57b793 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/ServiceContainerUtility.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/ServiceContainerUtility.java @@ -21,7 +21,9 @@ import java.util.concurrent.locks.ReentrantLock; import software.amazon.jdbc.ConnectionPluginManager; import software.amazon.jdbc.ConnectionProvider; +import software.amazon.jdbc.DriverConnectionProvider; import software.amazon.jdbc.PartialPluginService; +import software.amazon.jdbc.TargetDriverHelper; import software.amazon.jdbc.dialect.Dialect; import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; import software.amazon.jdbc.util.monitoring.MonitorService; @@ -59,16 +61,19 @@ public static FullServicesContainer createServiceContainer( StorageService storageService, MonitorService monitorService, TelemetryFactory telemetryFactory, - ConnectionProvider connectionProvider, String originalUrl, String targetDriverProtocol, TargetDriverDialect driverDialect, Dialect dbDialect, Properties props) throws SQLException { + final TargetDriverHelper helper = new TargetDriverHelper(); + final java.sql.Driver driver = helper.getTargetDriver(originalUrl, props); + final ConnectionProvider connProvider = new DriverConnectionProvider(driver); + FullServicesContainer servicesContainer = new FullServicesContainerImpl(storageService, monitorService, telemetryFactory); ConnectionPluginManager pluginManager = new ConnectionPluginManager( - connectionProvider, + connProvider, null, null, telemetryFactory); diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorInitializer.java b/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorInitializer.java index c4f13e6f9..e1cb59a95 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorInitializer.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorInitializer.java @@ -16,10 +16,9 @@ package software.amazon.jdbc.util.monitoring; -import software.amazon.jdbc.PluginService; -import software.amazon.jdbc.util.connection.ConnectionService; +import software.amazon.jdbc.util.FullServicesContainer; +@FunctionalInterface public interface MonitorInitializer { - - Monitor createMonitor(ConnectionService connectionService, PluginService pluginService); + Monitor createMonitor(FullServicesContainer servicesContainer); } diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java index d45729c72..b7d59f329 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java @@ -30,9 +30,6 @@ import java.util.logging.Logger; import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; -import software.amazon.jdbc.ConnectionProvider; -import software.amazon.jdbc.DriverConnectionProvider; -import software.amazon.jdbc.TargetDriverHelper; import software.amazon.jdbc.dialect.Dialect; import software.amazon.jdbc.hostlistprovider.Topology; import software.amazon.jdbc.hostlistprovider.monitoring.ClusterTopologyMonitorImpl; @@ -40,10 +37,10 @@ import software.amazon.jdbc.plugin.strategy.fastestresponse.NodeResponseTimeMonitor; import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; import software.amazon.jdbc.util.ExecutorFactory; +import software.amazon.jdbc.util.FullServicesContainer; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.PropertyUtils; -import software.amazon.jdbc.util.connection.ConnectionService; -import software.amazon.jdbc.util.connection.ConnectionServiceImpl; +import software.amazon.jdbc.util.ServiceContainerUtility; import software.amazon.jdbc.util.events.DataAccessEvent; import software.amazon.jdbc.util.events.Event; import software.amazon.jdbc.util.events.EventPublisher; @@ -198,20 +195,10 @@ public T runIfAbsent( cacheContainer = monitorCaches.computeIfAbsent(monitorClass, k -> supplier.get()); } - final ConnectionService connectionService = - getConnectionService( - storageService, - telemetryFactory, - originalUrl, - driverProtocol, - driverDialect, - dbDialect, - originalProps); - + final FullServicesContainer servicesContainer = getNewServicesContainer( + storageService, telemetryFactory, originalUrl, driverProtocol, driverDialect, dbDialect, originalProps); Monitor monitor = cacheContainer.getCache().computeIfAbsent(key, k -> { - MonitorItem monitorItem = new MonitorItem(() -> initializer.createMonitor( - connectionService, - connectionService.getPluginService())); + MonitorItem monitorItem = new MonitorItem(() -> initializer.createMonitor(servicesContainer)); monitorItem.getMonitor().start(); return monitorItem; }).getMonitor(); @@ -224,23 +211,25 @@ public T runIfAbsent( Messages.get("MonitorServiceImpl.unexpectedMonitorClass", new Object[] {monitorClass, monitor})); } - protected ConnectionService getConnectionService(StorageService storageService, - TelemetryFactory telemetryFactory, String originalUrl, String driverProtocol, TargetDriverDialect driverDialect, - Dialect dbDialect, Properties originalProps) throws SQLException { - TargetDriverHelper helper = new TargetDriverHelper(); - java.sql.Driver driver = helper.getTargetDriver(originalUrl, originalProps); - final ConnectionProvider defaultConnectionProvider = new DriverConnectionProvider(driver); + protected FullServicesContainer getNewServicesContainer( + StorageService storageService, + TelemetryFactory telemetryFactory, + String originalUrl, + String driverProtocol, + TargetDriverDialect driverDialect, + Dialect dbDialect, + Properties originalProps) throws SQLException { final Properties propsCopy = PropertyUtils.copyProperties(originalProps); - return new ConnectionServiceImpl( + return ServiceContainerUtility.createServiceContainer( storageService, this, telemetryFactory, - defaultConnectionProvider, originalUrl, driverProtocol, driverDialect, dbDialect, - propsCopy); + propsCopy + ); } @Override diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java index 966ff6ec4..a6351a849 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java @@ -1,401 +1,401 @@ -/* - * 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.failover; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertSame; -import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.doReturn; -import static org.mockito.Mockito.spy; -import static org.mockito.Mockito.when; -import static software.amazon.jdbc.plugin.failover.ClusterAwareReaderFailoverHandler.DEFAULT_FAILOVER_TIMEOUT; -import static software.amazon.jdbc.plugin.failover.ClusterAwareReaderFailoverHandler.DEFAULT_READER_CONNECT_TIMEOUT; - -import java.sql.Connection; -import java.sql.SQLException; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.Collections; -import java.util.EnumSet; -import java.util.List; -import java.util.Map; -import java.util.Properties; -import java.util.Set; -import java.util.concurrent.TimeUnit; -import java.util.stream.Collectors; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.mockito.Mock; -import org.mockito.Mockito; -import org.mockito.MockitoAnnotations; -import org.mockito.stubbing.Answer; -import software.amazon.jdbc.ConnectionPluginManager; -import software.amazon.jdbc.HostRole; -import software.amazon.jdbc.HostSpec; -import software.amazon.jdbc.HostSpecBuilder; -import software.amazon.jdbc.PluginService; -import software.amazon.jdbc.dialect.Dialect; -import software.amazon.jdbc.hostavailability.HostAvailability; -import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; -import software.amazon.jdbc.util.FullServicesContainer; -import software.amazon.jdbc.util.connection.ConnectionService; - -class ClusterAwareReaderFailoverHandlerTest { - @Mock FullServicesContainer mockContainer; - @Mock ConnectionService mockConnectionService; - @Mock PluginService mockPluginService; - @Mock ConnectionPluginManager mockPluginManager; - @Mock Connection mockConnection; - - private AutoCloseable closeable; - private final Properties properties = new Properties(); - private final List defaultHosts = Arrays.asList( - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("writer").port(1234).role(HostRole.WRITER).build(), - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader1").port(1234).role(HostRole.READER).build(), - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader2").port(1234).role(HostRole.READER).build(), - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader3").port(1234).role(HostRole.READER).build(), - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader4").port(1234).role(HostRole.READER).build(), - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader5").port(1234).role(HostRole.READER).build() - ); - - @BeforeEach - void setUp() { - closeable = MockitoAnnotations.openMocks(this); - when(mockContainer.getConnectionPluginManager()).thenReturn(mockPluginManager); - when(mockContainer.getPluginService()).thenReturn(mockPluginService); - } - - @AfterEach - void tearDown() throws Exception { - closeable.close(); - } - - @Test - public void testFailover() throws SQLException { - // original host list: [active writer, active reader, current connection (reader), active - // reader, down reader, active reader] - // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] - // connection attempts are made in pairs using the above list - // expected test result: successful connection for host at index 4 - final List hosts = defaultHosts; - final int currentHostIndex = 2; - final int successHostIndex = 4; - for (int i = 0; i < hosts.size(); i++) { - if (i != successHostIndex) { - final SQLException exception = new SQLException("exception", "08S01", null); - when(mockConnectionService.open(hosts.get(i), properties)) - .thenThrow(exception); - when(mockPluginService.isNetworkException(exception, null)).thenReturn(true); - } else { - when(mockConnectionService.open(hosts.get(i), properties)).thenReturn(mockConnection); - } - } - - when(mockPluginService.getTargetDriverDialect()).thenReturn(null); - - hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); - hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); - - final ReaderFailoverHandler target = getSpyFailoverHandler(); - final ReaderFailoverResult result = target.failover(hosts, hosts.get(currentHostIndex)); - - assertTrue(result.isConnected()); - assertSame(mockConnection, result.getConnection()); - assertEquals(hosts.get(successHostIndex), result.getHost()); - - final HostSpec successHost = hosts.get(successHostIndex); - final Map availabilityMap = target.getHostAvailabilityMap(); - Set unavailableHosts = getHostsWithGivenAvailability(availabilityMap, HostAvailability.NOT_AVAILABLE); - assertTrue(unavailableHosts.size() >= 4); - assertEquals(HostAvailability.AVAILABLE, availabilityMap.get(successHost.getHost())); - } - - private Set getHostsWithGivenAvailability( - Map availabilityMap, HostAvailability availability) { - return availabilityMap.entrySet().stream() - .filter((entry) -> availability.equals(entry.getValue())) - .map(Map.Entry::getKey) - .collect(Collectors.toSet()); - } - - @Test - public void testFailover_timeout() throws SQLException { - // original host list: [active writer, active reader, current connection (reader), active - // reader, down reader, active reader] - // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] - // connection attempts are made in pairs using the above list - // expected test result: failure to get reader since process is limited to 5s and each attempt - // to connect takes 20s - final List hosts = defaultHosts; - final int currentHostIndex = 2; - for (HostSpec host : hosts) { - when(mockConnectionService.open(host, properties)) - .thenAnswer((Answer) invocation -> { - Thread.sleep(20000); - return mockConnection; - }); - } - - hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); - hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); - - final ReaderFailoverHandler target = getSpyFailoverHandler(5000, 30000, false); - - final long startTimeNano = System.nanoTime(); - final ReaderFailoverResult result = target.failover(hosts, hosts.get(currentHostIndex)); - final long durationNano = System.nanoTime() - startTimeNano; - - assertFalse(result.isConnected()); - assertNull(result.getConnection()); - assertNull(result.getHost()); - - // 5s is a max allowed failover timeout; add 1s for inaccurate measurements - assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); - } - - private ClusterAwareReaderFailoverHandler getSpyFailoverHandler() { - ClusterAwareReaderFailoverHandler handler = - spy(new ClusterAwareReaderFailoverHandler(mockContainer, mockConnectionService, properties)); - doReturn(mockPluginService).when(handler).getNewPluginService(); - return handler; - } - - private ClusterAwareReaderFailoverHandler getSpyFailoverHandler( - int maxFailoverTimeoutMs, int timeoutMs, boolean isStrictReaderRequired) { - ClusterAwareReaderFailoverHandler handler = new ClusterAwareReaderFailoverHandler( - mockContainer, mockConnectionService, properties, maxFailoverTimeoutMs, timeoutMs, isStrictReaderRequired); - ClusterAwareReaderFailoverHandler spyHandler = spy(handler); - doReturn(mockPluginService).when(spyHandler).getNewPluginService(); - return spyHandler; - } - - @Test - public void testFailover_nullOrEmptyHostList() throws SQLException { - final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); - final HostSpec currentHost = - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("writer").port(1234).build(); - - ReaderFailoverResult result = target.failover(null, currentHost); - assertFalse(result.isConnected()); - assertNull(result.getConnection()); - assertNull(result.getHost()); - - final List hosts = new ArrayList<>(); - result = target.failover(hosts, currentHost); - assertFalse(result.isConnected()); - assertNull(result.getConnection()); - assertNull(result.getHost()); - } - - @Test - public void testGetReader_connectionSuccess() throws SQLException { - // even number of connection attempts - // first connection attempt to return succeeds, second attempt cancelled - // expected test result: successful connection for host at index 2 - final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) - final HostSpec slowHost = hosts.get(1); - final HostSpec fastHost = hosts.get(2); - when(mockConnectionService.open(slowHost, properties)) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(20000); - return mockConnection; - }); - when(mockConnectionService.open(eq(fastHost), eq(properties))).thenReturn(mockConnection); - - Dialect mockDialect = Mockito.mock(Dialect.class); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - when(mockPluginService.getDialect()).thenReturn(mockDialect); - - final ReaderFailoverHandler target = getSpyFailoverHandler(); - final ReaderFailoverResult result = target.getReaderConnection(hosts); - - assertTrue(result.isConnected()); - assertSame(mockConnection, result.getConnection()); - assertEquals(hosts.get(2), result.getHost()); - - Map availabilityMap = target.getHostAvailabilityMap(); - assertTrue(getHostsWithGivenAvailability(availabilityMap, HostAvailability.NOT_AVAILABLE).isEmpty()); - assertEquals(HostAvailability.AVAILABLE, availabilityMap.get(fastHost.getHost())); - } - - @Test - public void testGetReader_connectionFailure() throws SQLException { - // odd number of connection attempts - // first connection attempt to return fails - // expected test result: failure to get reader - final List hosts = defaultHosts.subList(0, 4); // 3 connection attempts (writer not attempted) - when(mockConnectionService.open(any(), eq(properties))).thenThrow(new SQLException("exception", "08S01", null)); - - Dialect mockDialect = Mockito.mock(Dialect.class); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - when(mockPluginService.getDialect()).thenReturn(mockDialect); - - final ReaderFailoverHandler target = getSpyFailoverHandler(); - final ReaderFailoverResult result = target.getReaderConnection(hosts); - - assertFalse(result.isConnected()); - assertNull(result.getConnection()); - assertNull(result.getHost()); - } - - @Test - public void testGetReader_connectionAttemptsTimeout() throws SQLException { - // connection attempts time out before they can succeed - // first connection attempt to return times out - // expected test result: failure to get reader - final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) - when(mockConnectionService.open(any(), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - try { - Thread.sleep(5000); - } catch (InterruptedException exception) { - // ignore - } - return mockConnection; - }); - - Dialect mockDialect = Mockito.mock(Dialect.class); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - when(mockPluginService.getDialect()).thenReturn(mockDialect); - - final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(60000, 1000, false); - final ReaderFailoverResult result = target.getReaderConnection(hosts); - - assertFalse(result.isConnected()); - assertNull(result.getConnection()); - assertNull(result.getHost()); - } - - @Test - public void testGetHostTuplesByPriority() { - final List originalHosts = defaultHosts; - originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); - originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); - originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); - - final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); - final List hostsByPriority = target.getHostsByPriority(originalHosts); - - int i = 0; - - // expecting active readers - while (i < hostsByPriority.size() - && hostsByPriority.get(i).getRole() == HostRole.READER - && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { - i++; - } - - // expecting a writer - while (i < hostsByPriority.size() - && hostsByPriority.get(i).getRole() == HostRole.WRITER) { - i++; - } - - // expecting down readers - while (i < hostsByPriority.size() - && hostsByPriority.get(i).getRole() == HostRole.READER - && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { - i++; - } - - assertEquals(hostsByPriority.size(), i); - } - - @Test - public void testGetReaderTuplesByPriority() { - final List originalHosts = defaultHosts; - originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); - originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); - originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); - - Dialect mockDialect = Mockito.mock(Dialect.class); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - when(mockPluginService.getDialect()).thenReturn(mockDialect); - - final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); - final List hostsByPriority = target.getReaderHostsByPriority(originalHosts); - - int i = 0; - - // expecting active readers - while (i < hostsByPriority.size() - && hostsByPriority.get(i).getRole() == HostRole.READER - && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { - i++; - } - - // expecting down readers - while (i < hostsByPriority.size() - && hostsByPriority.get(i).getRole() == HostRole.READER - && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { - i++; - } - - assertEquals(hostsByPriority.size(), i); - } - - @Test - public void testHostFailoverStrictReaderEnabled() { - - final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("writer").port(1234).role(HostRole.WRITER).build(); - final HostSpec reader = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader1").port(1234).role(HostRole.READER).build(); - final List hosts = Arrays.asList(writer, reader); - - Dialect mockDialect = Mockito.mock(Dialect.class); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - when(mockPluginService.getDialect()).thenReturn(mockDialect); - - final ClusterAwareReaderFailoverHandler target = - getSpyFailoverHandler(DEFAULT_FAILOVER_TIMEOUT, DEFAULT_READER_CONNECT_TIMEOUT, true); - - // The writer is included because the original writer has likely become a reader. - List expectedHostsByPriority = Arrays.asList(reader, writer); - - List hostsByPriority = target.getHostsByPriority(hosts); - assertEquals(expectedHostsByPriority, hostsByPriority); - - // Should pick the reader even if unavailable. The unavailable reader will be lower priority than the writer. - reader.setAvailability(HostAvailability.NOT_AVAILABLE); - expectedHostsByPriority = Arrays.asList(writer, reader); - - hostsByPriority = target.getHostsByPriority(hosts); - assertEquals(expectedHostsByPriority, hostsByPriority); - - // Writer node will only be picked if it is the only node in topology; - List expectedWriterHost = Collections.singletonList(writer); - - hostsByPriority = target.getHostsByPriority(Collections.singletonList(writer)); - assertEquals(expectedWriterHost, hostsByPriority); - } -} +// /* +// * 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.failover; +// +// import static org.junit.jupiter.api.Assertions.assertEquals; +// import static org.junit.jupiter.api.Assertions.assertFalse; +// import static org.junit.jupiter.api.Assertions.assertNull; +// import static org.junit.jupiter.api.Assertions.assertSame; +// import static org.junit.jupiter.api.Assertions.assertTrue; +// import static org.mockito.ArgumentMatchers.any; +// import static org.mockito.ArgumentMatchers.eq; +// import static org.mockito.Mockito.doReturn; +// import static org.mockito.Mockito.spy; +// import static org.mockito.Mockito.when; +// import static software.amazon.jdbc.plugin.failover.ClusterAwareReaderFailoverHandler.DEFAULT_FAILOVER_TIMEOUT; +// import static software.amazon.jdbc.plugin.failover.ClusterAwareReaderFailoverHandler.DEFAULT_READER_CONNECT_TIMEOUT; +// +// import java.sql.Connection; +// import java.sql.SQLException; +// import java.util.ArrayList; +// import java.util.Arrays; +// import java.util.Collections; +// import java.util.EnumSet; +// import java.util.List; +// import java.util.Map; +// import java.util.Properties; +// import java.util.Set; +// import java.util.concurrent.TimeUnit; +// import java.util.stream.Collectors; +// import org.junit.jupiter.api.AfterEach; +// import org.junit.jupiter.api.BeforeEach; +// import org.junit.jupiter.api.Test; +// import org.mockito.Mock; +// import org.mockito.Mockito; +// import org.mockito.MockitoAnnotations; +// import org.mockito.stubbing.Answer; +// import software.amazon.jdbc.ConnectionPluginManager; +// import software.amazon.jdbc.HostRole; +// import software.amazon.jdbc.HostSpec; +// import software.amazon.jdbc.HostSpecBuilder; +// import software.amazon.jdbc.PluginService; +// import software.amazon.jdbc.dialect.Dialect; +// import software.amazon.jdbc.hostavailability.HostAvailability; +// import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; +// import software.amazon.jdbc.util.FullServicesContainer; +// import software.amazon.jdbc.util.connection.ConnectionService; +// +// class ClusterAwareReaderFailoverHandlerTest { +// @Mock FullServicesContainer mockContainer; +// @Mock ConnectionService mockConnectionService; +// @Mock PluginService mockPluginService; +// @Mock ConnectionPluginManager mockPluginManager; +// @Mock Connection mockConnection; +// +// private AutoCloseable closeable; +// private final Properties properties = new Properties(); +// private final List defaultHosts = Arrays.asList( +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("writer").port(1234).role(HostRole.WRITER).build(), +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader1").port(1234).role(HostRole.READER).build(), +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader2").port(1234).role(HostRole.READER).build(), +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader3").port(1234).role(HostRole.READER).build(), +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader4").port(1234).role(HostRole.READER).build(), +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader5").port(1234).role(HostRole.READER).build() +// ); +// +// @BeforeEach +// void setUp() { +// closeable = MockitoAnnotations.openMocks(this); +// when(mockContainer.getConnectionPluginManager()).thenReturn(mockPluginManager); +// when(mockContainer.getPluginService()).thenReturn(mockPluginService); +// } +// +// @AfterEach +// void tearDown() throws Exception { +// closeable.close(); +// } +// +// @Test +// public void testFailover() throws SQLException { +// // original host list: [active writer, active reader, current connection (reader), active +// // reader, down reader, active reader] +// // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] +// // connection attempts are made in pairs using the above list +// // expected test result: successful connection for host at index 4 +// final List hosts = defaultHosts; +// final int currentHostIndex = 2; +// final int successHostIndex = 4; +// for (int i = 0; i < hosts.size(); i++) { +// if (i != successHostIndex) { +// final SQLException exception = new SQLException("exception", "08S01", null); +// when(mockConnectionService.open(hosts.get(i), properties)) +// .thenThrow(exception); +// when(mockPluginService.isNetworkException(exception, null)).thenReturn(true); +// } else { +// when(mockConnectionService.open(hosts.get(i), properties)).thenReturn(mockConnection); +// } +// } +// +// when(mockPluginService.getTargetDriverDialect()).thenReturn(null); +// +// hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); +// hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); +// +// final ReaderFailoverHandler target = getSpyFailoverHandler(); +// final ReaderFailoverResult result = target.failover(hosts, hosts.get(currentHostIndex)); +// +// assertTrue(result.isConnected()); +// assertSame(mockConnection, result.getConnection()); +// assertEquals(hosts.get(successHostIndex), result.getHost()); +// +// final HostSpec successHost = hosts.get(successHostIndex); +// final Map availabilityMap = target.getHostAvailabilityMap(); +// Set unavailableHosts = getHostsWithGivenAvailability(availabilityMap, HostAvailability.NOT_AVAILABLE); +// assertTrue(unavailableHosts.size() >= 4); +// assertEquals(HostAvailability.AVAILABLE, availabilityMap.get(successHost.getHost())); +// } +// +// private Set getHostsWithGivenAvailability( +// Map availabilityMap, HostAvailability availability) { +// return availabilityMap.entrySet().stream() +// .filter((entry) -> availability.equals(entry.getValue())) +// .map(Map.Entry::getKey) +// .collect(Collectors.toSet()); +// } +// +// @Test +// public void testFailover_timeout() throws SQLException { +// // original host list: [active writer, active reader, current connection (reader), active +// // reader, down reader, active reader] +// // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] +// // connection attempts are made in pairs using the above list +// // expected test result: failure to get reader since process is limited to 5s and each attempt +// // to connect takes 20s +// final List hosts = defaultHosts; +// final int currentHostIndex = 2; +// for (HostSpec host : hosts) { +// when(mockConnectionService.open(host, properties)) +// .thenAnswer((Answer) invocation -> { +// Thread.sleep(20000); +// return mockConnection; +// }); +// } +// +// hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); +// hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); +// +// final ReaderFailoverHandler target = getSpyFailoverHandler(5000, 30000, false); +// +// final long startTimeNano = System.nanoTime(); +// final ReaderFailoverResult result = target.failover(hosts, hosts.get(currentHostIndex)); +// final long durationNano = System.nanoTime() - startTimeNano; +// +// assertFalse(result.isConnected()); +// assertNull(result.getConnection()); +// assertNull(result.getHost()); +// +// // 5s is a max allowed failover timeout; add 1s for inaccurate measurements +// assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); +// } +// +// private ClusterAwareReaderFailoverHandler getSpyFailoverHandler() { +// ClusterAwareReaderFailoverHandler handler = +// spy(new ClusterAwareReaderFailoverHandler(mockContainer, mockConnectionService, properties)); +// doReturn(mockPluginService).when(handler).getNewPluginService(); +// return handler; +// } +// +// private ClusterAwareReaderFailoverHandler getSpyFailoverHandler( +// int maxFailoverTimeoutMs, int timeoutMs, boolean isStrictReaderRequired) { +// ClusterAwareReaderFailoverHandler handler = new ClusterAwareReaderFailoverHandler( +// mockContainer, mockConnectionService, properties, maxFailoverTimeoutMs, timeoutMs, isStrictReaderRequired); +// ClusterAwareReaderFailoverHandler spyHandler = spy(handler); +// doReturn(mockPluginService).when(spyHandler).getNewPluginService(); +// return spyHandler; +// } +// +// @Test +// public void testFailover_nullOrEmptyHostList() throws SQLException { +// final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); +// final HostSpec currentHost = +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("writer").port(1234).build(); +// +// ReaderFailoverResult result = target.failover(null, currentHost); +// assertFalse(result.isConnected()); +// assertNull(result.getConnection()); +// assertNull(result.getHost()); +// +// final List hosts = new ArrayList<>(); +// result = target.failover(hosts, currentHost); +// assertFalse(result.isConnected()); +// assertNull(result.getConnection()); +// assertNull(result.getHost()); +// } +// +// @Test +// public void testGetReader_connectionSuccess() throws SQLException { +// // even number of connection attempts +// // first connection attempt to return succeeds, second attempt cancelled +// // expected test result: successful connection for host at index 2 +// final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) +// final HostSpec slowHost = hosts.get(1); +// final HostSpec fastHost = hosts.get(2); +// when(mockConnectionService.open(slowHost, properties)) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(20000); +// return mockConnection; +// }); +// when(mockConnectionService.open(eq(fastHost), eq(properties))).thenReturn(mockConnection); +// +// Dialect mockDialect = Mockito.mock(Dialect.class); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// +// final ReaderFailoverHandler target = getSpyFailoverHandler(); +// final ReaderFailoverResult result = target.getReaderConnection(hosts); +// +// assertTrue(result.isConnected()); +// assertSame(mockConnection, result.getConnection()); +// assertEquals(hosts.get(2), result.getHost()); +// +// Map availabilityMap = target.getHostAvailabilityMap(); +// assertTrue(getHostsWithGivenAvailability(availabilityMap, HostAvailability.NOT_AVAILABLE).isEmpty()); +// assertEquals(HostAvailability.AVAILABLE, availabilityMap.get(fastHost.getHost())); +// } +// +// @Test +// public void testGetReader_connectionFailure() throws SQLException { +// // odd number of connection attempts +// // first connection attempt to return fails +// // expected test result: failure to get reader +// final List hosts = defaultHosts.subList(0, 4); // 3 connection attempts (writer not attempted) +// when(mockConnectionService.open(any(), eq(properties))).thenThrow(new SQLException("exception", "08S01", null)); +// +// Dialect mockDialect = Mockito.mock(Dialect.class); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// +// final ReaderFailoverHandler target = getSpyFailoverHandler(); +// final ReaderFailoverResult result = target.getReaderConnection(hosts); +// +// assertFalse(result.isConnected()); +// assertNull(result.getConnection()); +// assertNull(result.getHost()); +// } +// +// @Test +// public void testGetReader_connectionAttemptsTimeout() throws SQLException { +// // connection attempts time out before they can succeed +// // first connection attempt to return times out +// // expected test result: failure to get reader +// final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) +// when(mockConnectionService.open(any(), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// try { +// Thread.sleep(5000); +// } catch (InterruptedException exception) { +// // ignore +// } +// return mockConnection; +// }); +// +// Dialect mockDialect = Mockito.mock(Dialect.class); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// +// final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(60000, 1000, false); +// final ReaderFailoverResult result = target.getReaderConnection(hosts); +// +// assertFalse(result.isConnected()); +// assertNull(result.getConnection()); +// assertNull(result.getHost()); +// } +// +// @Test +// public void testGetHostTuplesByPriority() { +// final List originalHosts = defaultHosts; +// originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); +// originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); +// originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); +// +// final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); +// final List hostsByPriority = target.getHostsByPriority(originalHosts); +// +// int i = 0; +// +// // expecting active readers +// while (i < hostsByPriority.size() +// && hostsByPriority.get(i).getRole() == HostRole.READER +// && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { +// i++; +// } +// +// // expecting a writer +// while (i < hostsByPriority.size() +// && hostsByPriority.get(i).getRole() == HostRole.WRITER) { +// i++; +// } +// +// // expecting down readers +// while (i < hostsByPriority.size() +// && hostsByPriority.get(i).getRole() == HostRole.READER +// && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { +// i++; +// } +// +// assertEquals(hostsByPriority.size(), i); +// } +// +// @Test +// public void testGetReaderTuplesByPriority() { +// final List originalHosts = defaultHosts; +// originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); +// originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); +// originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); +// +// Dialect mockDialect = Mockito.mock(Dialect.class); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// +// final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); +// final List hostsByPriority = target.getReaderHostsByPriority(originalHosts); +// +// int i = 0; +// +// // expecting active readers +// while (i < hostsByPriority.size() +// && hostsByPriority.get(i).getRole() == HostRole.READER +// && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { +// i++; +// } +// +// // expecting down readers +// while (i < hostsByPriority.size() +// && hostsByPriority.get(i).getRole() == HostRole.READER +// && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { +// i++; +// } +// +// assertEquals(hostsByPriority.size(), i); +// } +// +// @Test +// public void testHostFailoverStrictReaderEnabled() { +// +// final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("writer").port(1234).role(HostRole.WRITER).build(); +// final HostSpec reader = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader1").port(1234).role(HostRole.READER).build(); +// final List hosts = Arrays.asList(writer, reader); +// +// Dialect mockDialect = Mockito.mock(Dialect.class); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// +// final ClusterAwareReaderFailoverHandler target = +// getSpyFailoverHandler(DEFAULT_FAILOVER_TIMEOUT, DEFAULT_READER_CONNECT_TIMEOUT, true); +// +// // The writer is included because the original writer has likely become a reader. +// List expectedHostsByPriority = Arrays.asList(reader, writer); +// +// List hostsByPriority = target.getHostsByPriority(hosts); +// assertEquals(expectedHostsByPriority, hostsByPriority); +// +// // Should pick the reader even if unavailable. The unavailable reader will be lower priority than the writer. +// reader.setAvailability(HostAvailability.NOT_AVAILABLE); +// expectedHostsByPriority = Arrays.asList(writer, reader); +// +// hostsByPriority = target.getHostsByPriority(hosts); +// assertEquals(expectedHostsByPriority, hostsByPriority); +// +// // Writer node will only be picked if it is the only node in topology; +// List expectedWriterHost = Collections.singletonList(writer); +// +// hostsByPriority = target.getHostsByPriority(Collections.singletonList(writer)); +// assertEquals(expectedWriterHost, hostsByPriority); +// } +// } diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java index 1ad394010..097b68d67 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java @@ -1,373 +1,373 @@ -/* - * 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.failover; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertSame; -import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.ArgumentMatchers.refEq; -import static org.mockito.Mockito.atLeastOnce; -import static org.mockito.Mockito.doReturn; -import static org.mockito.Mockito.spy; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -import java.sql.Connection; -import java.sql.SQLException; -import java.util.Arrays; -import java.util.EnumSet; -import java.util.List; -import java.util.Properties; -import java.util.concurrent.TimeUnit; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.mockito.ArgumentMatchers; -import org.mockito.Mock; -import org.mockito.MockitoAnnotations; -import org.mockito.stubbing.Answer; -import software.amazon.jdbc.HostSpec; -import software.amazon.jdbc.HostSpecBuilder; -import software.amazon.jdbc.PluginService; -import software.amazon.jdbc.dialect.Dialect; -import software.amazon.jdbc.hostavailability.HostAvailability; -import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; -import software.amazon.jdbc.util.FullServicesContainer; -import software.amazon.jdbc.util.connection.ConnectionService; - -class ClusterAwareWriterFailoverHandlerTest { - @Mock FullServicesContainer mockContainer; - @Mock ConnectionService mockConnectionService; - @Mock PluginService mockPluginService; - @Mock Connection mockConnection; - @Mock ReaderFailoverHandler mockReaderFailoverHandler; - @Mock Connection mockWriterConnection; - @Mock Connection mockNewWriterConnection; - @Mock Connection mockReaderAConnection; - @Mock Connection mockReaderBConnection; - @Mock Dialect mockDialect; - - private AutoCloseable closeable; - private final Properties properties = new Properties(); - private final HostSpec newWriterHost = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("new-writer-host").build(); - private final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("writer-host").build(); - private final HostSpec readerA = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader-a-host").build(); - private final HostSpec readerB = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader-b-host").build(); - private final List topology = Arrays.asList(writer, readerA, readerB); - private final List newTopology = Arrays.asList(newWriterHost, readerA, readerB); - - @BeforeEach - void setUp() { - closeable = MockitoAnnotations.openMocks(this); - when(mockContainer.getPluginService()).thenReturn(mockPluginService); - writer.addAlias("writer-host"); - newWriterHost.addAlias("new-writer-host"); - readerA.addAlias("reader-a-host"); - readerB.addAlias("reader-b-host"); - } - - @AfterEach - void tearDown() throws Exception { - closeable.close(); - } - - @Test - public void testReconnectToWriter_taskBReaderException() throws SQLException { - when(mockConnectionService.open(refEq(writer), eq(properties))).thenReturn(mockConnection); - when(mockConnectionService.open(refEq(readerA), eq(properties))).thenThrow(SQLException.class); - when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); - - when(mockPluginService.getAllHosts()).thenReturn(topology); - - when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())).thenThrow(SQLException.class); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); - final WriterFailoverResult result = target.failover(topology); - - assertTrue(result.isConnected()); - assertFalse(result.isNewHost()); - assertSame(result.getNewConnection(), mockConnection); - - assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); - } - - private ClusterAwareWriterFailoverHandler getSpyFailoverHandler( - final int failoverTimeoutMs, - final int readTopologyIntervalMs, - final int reconnectWriterIntervalMs) { - ClusterAwareWriterFailoverHandler handler = new ClusterAwareWriterFailoverHandler( - mockContainer, - mockConnectionService, - mockReaderFailoverHandler, - properties, - failoverTimeoutMs, - readTopologyIntervalMs, - reconnectWriterIntervalMs); - - ClusterAwareWriterFailoverHandler spyHandler = spy(handler); - doReturn(mockPluginService).when(spyHandler).getNewPluginService(); - return spyHandler; - } - - /** - * Verify that writer failover handler can re-connect to a current writer node. - * - *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. - * TaskA: successfully re-connect to initial writer; return new connection. - * TaskB: successfully connect to readerA and then new writer, but it takes more time than taskA. - * Expected test result: new connection by taskA. - */ - @Test - public void testReconnectToWriter_SlowReaderA() throws SQLException { - when(mockConnectionService.open(refEq(writer), eq(properties))).thenReturn(mockWriterConnection); - when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); - when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); - when(mockPluginService.getAllHosts()).thenReturn(topology).thenReturn(newTopology); - - when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(5000); - return new ReaderFailoverResult(mockReaderAConnection, readerA, true); - }); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); - final WriterFailoverResult result = target.failover(topology); - - assertTrue(result.isConnected()); - assertFalse(result.isNewHost()); - assertSame(result.getNewConnection(), mockWriterConnection); - assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); - } - - /** - * Verify that writer failover handler can re-connect to a current writer node. - * - *

Topology: no changes. - * TaskA: successfully re-connect to writer; return new connection. - * TaskB: successfully connect to readerA and retrieve topology, but latest writer is not new (defer to taskA). - * Expected test result: new connection by taskA. - */ - @Test - public void testReconnectToWriter_taskBDefers() throws SQLException { - when(mockConnectionService.open(refEq(writer), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(5000); - return mockWriterConnection; - }); - when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); - - when(mockPluginService.getAllHosts()).thenReturn(topology); - - when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) - .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 2000, 2000); - final WriterFailoverResult result = target.failover(topology); - - assertTrue(result.isConnected()); - assertFalse(result.isNewHost()); - assertSame(result.getNewConnection(), mockWriterConnection); - assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); - } - - /** - * Verify that writer failover handler can re-connect to a new writer node. - * - *

Topology: changes to [new-writer, reader-A, reader-B] for taskB, taskA sees no changes. - * taskA: successfully re-connect to writer; return connection to initial writer, but it takes more - * time than taskB. - * TaskB: successfully connect to readerA and then to new-writer. - * Expected test result: new connection to writer by taskB. - */ - @Test - public void testConnectToReaderA_SlowWriter() throws SQLException { - when(mockConnectionService.open(refEq(writer), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(5000); - return mockWriterConnection; - }); - when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); - when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); - when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); - - when(mockPluginService.getAllHosts()).thenReturn(newTopology); - - when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) - .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); - final WriterFailoverResult result = target.failover(topology); - - assertTrue(result.isConnected()); - assertTrue(result.isNewHost()); - assertSame(result.getNewConnection(), mockNewWriterConnection); - assertEquals(3, result.getTopology().size()); - assertEquals("new-writer-host", result.getTopology().get(0).getHost()); - assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); - } - - /** - * Verify that writer failover handler can re-connect to a new writer node. - * - *

Topology: changes to [new-writer, initial-writer, reader-A, reader-B]. - * TaskA: successfully reconnect, but initial-writer is now a reader (defer to taskB). - * TaskB: successfully connect to readerA and then to new-writer. - * Expected test result: new connection to writer by taskB. - */ - @Test - public void testConnectToReaderA_taskADefers() throws SQLException { - when(mockConnectionService.open(writer, properties)).thenReturn(mockConnection); - when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); - when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); - when(mockConnectionService.open(refEq(newWriterHost), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(5000); - return mockNewWriterConnection; - }); - - final List newTopology = Arrays.asList(newWriterHost, writer, readerA, readerB); - when(mockPluginService.getAllHosts()).thenReturn(newTopology); - - when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) - .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); - final WriterFailoverResult result = target.failover(topology); - - assertTrue(result.isConnected()); - assertTrue(result.isNewHost()); - assertSame(result.getNewConnection(), mockNewWriterConnection); - assertEquals(4, result.getTopology().size()); - assertEquals("new-writer-host", result.getTopology().get(0).getHost()); - - verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); - assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); - } - - /** - * Verify that writer failover handler fails to re-connect to any writer node. - * - *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. - * TaskA: fail to re-connect to writer due to failover timeout. - * TaskB: successfully connect to readerA and then fail to connect to writer due to failover timeout. - * Expected test result: no connection. - */ - @Test - public void testFailedToConnect_failoverTimeout() throws SQLException { - when(mockConnectionService.open(refEq(writer), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(30000); - return mockWriterConnection; - }); - when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); - when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); - when(mockConnectionService.open(refEq(newWriterHost), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(30000); - return mockNewWriterConnection; - }); - when(mockPluginService.getAllHosts()).thenReturn(newTopology); - - when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) - .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); - - final long startTimeNano = System.nanoTime(); - final WriterFailoverResult result = target.failover(topology); - final long durationNano = System.nanoTime() - startTimeNano; - - assertFalse(result.isConnected()); - assertFalse(result.isNewHost()); - - verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); - - // 5s is a max allowed failover timeout; add 1s for inaccurate measurements - assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); - } - - /** - * Verify that writer failover handler fails to re-connect to any writer node. - * - *

Topology: changes to [new-writer, reader-A, reader-B] for taskB. - * TaskA: fail to re-connect to writer due to exception. - * TaskB: successfully connect to readerA and then fail to connect to writer due to exception. - * Expected test result: no connection. - */ - @Test - public void testFailedToConnect_taskAException_taskBWriterException() throws SQLException { - final SQLException exception = new SQLException("exception", "08S01", null); - when(mockConnectionService.open(refEq(writer), eq(properties))).thenThrow(exception); - when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); - when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); - when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenThrow(exception); - when(mockPluginService.isNetworkException(eq(exception), any())).thenReturn(true); - - when(mockPluginService.getAllHosts()).thenReturn(newTopology); - - when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) - .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); - final WriterFailoverResult result = target.failover(topology); - - assertFalse(result.isConnected()); - assertFalse(result.isNewHost()); - - assertEquals(HostAvailability.NOT_AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); - } -} +// /* +// * 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.failover; +// +// import static org.junit.jupiter.api.Assertions.assertEquals; +// import static org.junit.jupiter.api.Assertions.assertFalse; +// import static org.junit.jupiter.api.Assertions.assertSame; +// import static org.junit.jupiter.api.Assertions.assertTrue; +// import static org.mockito.ArgumentMatchers.any; +// import static org.mockito.ArgumentMatchers.eq; +// import static org.mockito.ArgumentMatchers.refEq; +// import static org.mockito.Mockito.atLeastOnce; +// import static org.mockito.Mockito.doReturn; +// import static org.mockito.Mockito.spy; +// import static org.mockito.Mockito.verify; +// import static org.mockito.Mockito.when; +// +// import java.sql.Connection; +// import java.sql.SQLException; +// import java.util.Arrays; +// import java.util.EnumSet; +// import java.util.List; +// import java.util.Properties; +// import java.util.concurrent.TimeUnit; +// import org.junit.jupiter.api.AfterEach; +// import org.junit.jupiter.api.BeforeEach; +// import org.junit.jupiter.api.Test; +// import org.mockito.ArgumentMatchers; +// import org.mockito.Mock; +// import org.mockito.MockitoAnnotations; +// import org.mockito.stubbing.Answer; +// import software.amazon.jdbc.HostSpec; +// import software.amazon.jdbc.HostSpecBuilder; +// import software.amazon.jdbc.PluginService; +// import software.amazon.jdbc.dialect.Dialect; +// import software.amazon.jdbc.hostavailability.HostAvailability; +// import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; +// import software.amazon.jdbc.util.FullServicesContainer; +// import software.amazon.jdbc.util.connection.ConnectionService; +// +// class ClusterAwareWriterFailoverHandlerTest { +// @Mock FullServicesContainer mockContainer; +// @Mock ConnectionService mockConnectionService; +// @Mock PluginService mockPluginService; +// @Mock Connection mockConnection; +// @Mock ReaderFailoverHandler mockReaderFailoverHandler; +// @Mock Connection mockWriterConnection; +// @Mock Connection mockNewWriterConnection; +// @Mock Connection mockReaderAConnection; +// @Mock Connection mockReaderBConnection; +// @Mock Dialect mockDialect; +// +// private AutoCloseable closeable; +// private final Properties properties = new Properties(); +// private final HostSpec newWriterHost = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("new-writer-host").build(); +// private final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("writer-host").build(); +// private final HostSpec readerA = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader-a-host").build(); +// private final HostSpec readerB = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader-b-host").build(); +// private final List topology = Arrays.asList(writer, readerA, readerB); +// private final List newTopology = Arrays.asList(newWriterHost, readerA, readerB); +// +// @BeforeEach +// void setUp() { +// closeable = MockitoAnnotations.openMocks(this); +// when(mockContainer.getPluginService()).thenReturn(mockPluginService); +// writer.addAlias("writer-host"); +// newWriterHost.addAlias("new-writer-host"); +// readerA.addAlias("reader-a-host"); +// readerB.addAlias("reader-b-host"); +// } +// +// @AfterEach +// void tearDown() throws Exception { +// closeable.close(); +// } +// +// @Test +// public void testReconnectToWriter_taskBReaderException() throws SQLException { +// when(mockConnectionService.open(refEq(writer), eq(properties))).thenReturn(mockConnection); +// when(mockConnectionService.open(refEq(readerA), eq(properties))).thenThrow(SQLException.class); +// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); +// +// when(mockPluginService.getAllHosts()).thenReturn(topology); +// +// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())).thenThrow(SQLException.class); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertTrue(result.isConnected()); +// assertFalse(result.isNewHost()); +// assertSame(result.getNewConnection(), mockConnection); +// +// assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); +// } +// +// private ClusterAwareWriterFailoverHandler getSpyFailoverHandler( +// final int failoverTimeoutMs, +// final int readTopologyIntervalMs, +// final int reconnectWriterIntervalMs) { +// ClusterAwareWriterFailoverHandler handler = new ClusterAwareWriterFailoverHandler( +// mockContainer, +// mockConnectionService, +// mockReaderFailoverHandler, +// properties, +// failoverTimeoutMs, +// readTopologyIntervalMs, +// reconnectWriterIntervalMs); +// +// ClusterAwareWriterFailoverHandler spyHandler = spy(handler); +// doReturn(mockPluginService).when(spyHandler).getNewPluginService(); +// return spyHandler; +// } +// +// /** +// * Verify that writer failover handler can re-connect to a current writer node. +// * +// *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. +// * TaskA: successfully re-connect to initial writer; return new connection. +// * TaskB: successfully connect to readerA and then new writer, but it takes more time than taskA. +// * Expected test result: new connection by taskA. +// */ +// @Test +// public void testReconnectToWriter_SlowReaderA() throws SQLException { +// when(mockConnectionService.open(refEq(writer), eq(properties))).thenReturn(mockWriterConnection); +// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); +// when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); +// when(mockPluginService.getAllHosts()).thenReturn(topology).thenReturn(newTopology); +// +// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(5000); +// return new ReaderFailoverResult(mockReaderAConnection, readerA, true); +// }); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertTrue(result.isConnected()); +// assertFalse(result.isNewHost()); +// assertSame(result.getNewConnection(), mockWriterConnection); +// assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); +// } +// +// /** +// * Verify that writer failover handler can re-connect to a current writer node. +// * +// *

Topology: no changes. +// * TaskA: successfully re-connect to writer; return new connection. +// * TaskB: successfully connect to readerA and retrieve topology, but latest writer is not new (defer to taskA). +// * Expected test result: new connection by taskA. +// */ +// @Test +// public void testReconnectToWriter_taskBDefers() throws SQLException { +// when(mockConnectionService.open(refEq(writer), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(5000); +// return mockWriterConnection; +// }); +// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); +// +// when(mockPluginService.getAllHosts()).thenReturn(topology); +// +// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) +// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 2000, 2000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertTrue(result.isConnected()); +// assertFalse(result.isNewHost()); +// assertSame(result.getNewConnection(), mockWriterConnection); +// assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); +// } +// +// /** +// * Verify that writer failover handler can re-connect to a new writer node. +// * +// *

Topology: changes to [new-writer, reader-A, reader-B] for taskB, taskA sees no changes. +// * taskA: successfully re-connect to writer; return connection to initial writer, but it takes more +// * time than taskB. +// * TaskB: successfully connect to readerA and then to new-writer. +// * Expected test result: new connection to writer by taskB. +// */ +// @Test +// public void testConnectToReaderA_SlowWriter() throws SQLException { +// when(mockConnectionService.open(refEq(writer), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(5000); +// return mockWriterConnection; +// }); +// when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); +// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); +// when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); +// +// when(mockPluginService.getAllHosts()).thenReturn(newTopology); +// +// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) +// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertTrue(result.isConnected()); +// assertTrue(result.isNewHost()); +// assertSame(result.getNewConnection(), mockNewWriterConnection); +// assertEquals(3, result.getTopology().size()); +// assertEquals("new-writer-host", result.getTopology().get(0).getHost()); +// assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); +// } +// +// /** +// * Verify that writer failover handler can re-connect to a new writer node. +// * +// *

Topology: changes to [new-writer, initial-writer, reader-A, reader-B]. +// * TaskA: successfully reconnect, but initial-writer is now a reader (defer to taskB). +// * TaskB: successfully connect to readerA and then to new-writer. +// * Expected test result: new connection to writer by taskB. +// */ +// @Test +// public void testConnectToReaderA_taskADefers() throws SQLException { +// when(mockConnectionService.open(writer, properties)).thenReturn(mockConnection); +// when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); +// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); +// when(mockConnectionService.open(refEq(newWriterHost), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(5000); +// return mockNewWriterConnection; +// }); +// +// final List newTopology = Arrays.asList(newWriterHost, writer, readerA, readerB); +// when(mockPluginService.getAllHosts()).thenReturn(newTopology); +// +// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) +// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertTrue(result.isConnected()); +// assertTrue(result.isNewHost()); +// assertSame(result.getNewConnection(), mockNewWriterConnection); +// assertEquals(4, result.getTopology().size()); +// assertEquals("new-writer-host", result.getTopology().get(0).getHost()); +// +// verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); +// assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); +// } +// +// /** +// * Verify that writer failover handler fails to re-connect to any writer node. +// * +// *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. +// * TaskA: fail to re-connect to writer due to failover timeout. +// * TaskB: successfully connect to readerA and then fail to connect to writer due to failover timeout. +// * Expected test result: no connection. +// */ +// @Test +// public void testFailedToConnect_failoverTimeout() throws SQLException { +// when(mockConnectionService.open(refEq(writer), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(30000); +// return mockWriterConnection; +// }); +// when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); +// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); +// when(mockConnectionService.open(refEq(newWriterHost), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(30000); +// return mockNewWriterConnection; +// }); +// when(mockPluginService.getAllHosts()).thenReturn(newTopology); +// +// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) +// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); +// +// final long startTimeNano = System.nanoTime(); +// final WriterFailoverResult result = target.failover(topology); +// final long durationNano = System.nanoTime() - startTimeNano; +// +// assertFalse(result.isConnected()); +// assertFalse(result.isNewHost()); +// +// verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); +// +// // 5s is a max allowed failover timeout; add 1s for inaccurate measurements +// assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); +// } +// +// /** +// * Verify that writer failover handler fails to re-connect to any writer node. +// * +// *

Topology: changes to [new-writer, reader-A, reader-B] for taskB. +// * TaskA: fail to re-connect to writer due to exception. +// * TaskB: successfully connect to readerA and then fail to connect to writer due to exception. +// * Expected test result: no connection. +// */ +// @Test +// public void testFailedToConnect_taskAException_taskBWriterException() throws SQLException { +// final SQLException exception = new SQLException("exception", "08S01", null); +// when(mockConnectionService.open(refEq(writer), eq(properties))).thenThrow(exception); +// when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); +// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); +// when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenThrow(exception); +// when(mockPluginService.isNetworkException(eq(exception), any())).thenReturn(true); +// +// when(mockPluginService.getAllHosts()).thenReturn(newTopology); +// +// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) +// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertFalse(result.isConnected()); +// assertFalse(result.isNewHost()); +// +// assertEquals(HostAvailability.NOT_AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); +// } +// } diff --git a/wrapper/src/test/java/software/amazon/jdbc/util/monitoring/MonitorServiceImplTest.java b/wrapper/src/test/java/software/amazon/jdbc/util/monitoring/MonitorServiceImplTest.java index 61248faaa..147a81135 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/util/monitoring/MonitorServiceImplTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/util/monitoring/MonitorServiceImplTest.java @@ -1,306 +1,306 @@ -/* - * 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.util.monitoring; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertNotEquals; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyInt; -import static org.mockito.Mockito.doNothing; -import static org.mockito.Mockito.doReturn; -import static org.mockito.Mockito.spy; - -import java.sql.SQLException; -import java.util.Collections; -import java.util.HashSet; -import java.util.Properties; -import java.util.concurrent.TimeUnit; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.Assertions; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.mockito.Mock; -import org.mockito.MockitoAnnotations; -import software.amazon.jdbc.dialect.Dialect; -import software.amazon.jdbc.plugin.customendpoint.CustomEndpointMonitorImpl; -import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; -import software.amazon.jdbc.util.connection.ConnectionService; -import software.amazon.jdbc.util.events.EventPublisher; -import software.amazon.jdbc.util.storage.StorageService; -import software.amazon.jdbc.util.telemetry.TelemetryFactory; - -class MonitorServiceImplTest { - @Mock StorageService mockStorageService; - @Mock ConnectionService mockConnectionService; - @Mock TelemetryFactory mockTelemetryFactory; - @Mock TargetDriverDialect mockTargetDriverDialect; - @Mock Dialect mockDbDialect; - @Mock EventPublisher mockPublisher; - MonitorServiceImpl spyMonitorService; - private AutoCloseable closeable; - - @BeforeEach - void setUp() { - closeable = MockitoAnnotations.openMocks(this); - spyMonitorService = spy(new MonitorServiceImpl(mockPublisher)); - doNothing().when(spyMonitorService).initCleanupThread(anyInt()); - - try { - doReturn(mockConnectionService).when(spyMonitorService) - .getConnectionService(any(), any(), any(), any(), any(), any(), any()); - } catch (SQLException e) { - Assertions.fail( - "Encountered exception while stubbing MonitorServiceImpl#getConnectionService: " + e.getMessage()); - } - } - - @AfterEach - void tearDown() throws Exception { - closeable.close(); - spyMonitorService.releaseResources(); - } - - @Test - public void testMonitorError_monitorReCreated() throws SQLException, InterruptedException { - spyMonitorService.registerMonitorTypeIfAbsent( - NoOpMonitor.class, - TimeUnit.MINUTES.toNanos(1), - TimeUnit.MINUTES.toNanos(1), - new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), - null - ); - String key = "testMonitor"; - NoOpMonitor monitor = spyMonitorService.runIfAbsent( - NoOpMonitor.class, - key, - mockStorageService, - mockTelemetryFactory, - "jdbc:postgresql://somehost/somedb", - "someProtocol", - mockTargetDriverDialect, - mockDbDialect, - new Properties(), - (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) - ); - - Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key); - assertNotNull(storedMonitor); - assertEquals(monitor, storedMonitor); - // need to wait to give time for the monitor executor to start the monitor thread. - TimeUnit.MILLISECONDS.sleep(250); - assertEquals(MonitorState.RUNNING, monitor.getState()); - - monitor.state.set(MonitorState.ERROR); - spyMonitorService.checkMonitors(); - - assertEquals(MonitorState.STOPPED, monitor.getState()); - - Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key); - assertNotNull(newMonitor); - assertNotEquals(monitor, newMonitor); - // need to wait to give time for the monitor executor to start the monitor thread. - TimeUnit.MILLISECONDS.sleep(250); - assertEquals(MonitorState.RUNNING, newMonitor.getState()); - } - - @Test - public void testMonitorStuck_monitorReCreated() throws SQLException, InterruptedException { - spyMonitorService.registerMonitorTypeIfAbsent( - NoOpMonitor.class, - TimeUnit.MINUTES.toNanos(1), - 1, // heartbeat times out immediately - new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), - null - ); - String key = "testMonitor"; - NoOpMonitor monitor = spyMonitorService.runIfAbsent( - NoOpMonitor.class, - key, - mockStorageService, - mockTelemetryFactory, - "jdbc:postgresql://somehost/somedb", - "someProtocol", - mockTargetDriverDialect, - mockDbDialect, - new Properties(), - (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) - ); - - Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key); - assertNotNull(storedMonitor); - assertEquals(monitor, storedMonitor); - // need to wait to give time for the monitor executor to start the monitor thread. - TimeUnit.MILLISECONDS.sleep(250); - assertEquals(MonitorState.RUNNING, monitor.getState()); - - // checkMonitors() should detect the heartbeat/inactivity timeout, stop the monitor, and re-create a new one. - spyMonitorService.checkMonitors(); - - assertEquals(MonitorState.STOPPED, monitor.getState()); - - Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key); - assertNotNull(newMonitor); - assertNotEquals(monitor, newMonitor); - // need to wait to give time for the monitor executor to start the monitor thread. - TimeUnit.MILLISECONDS.sleep(250); - assertEquals(MonitorState.RUNNING, newMonitor.getState()); - } - - @Test - public void testMonitorExpired() throws SQLException, InterruptedException { - spyMonitorService.registerMonitorTypeIfAbsent( - NoOpMonitor.class, - TimeUnit.MILLISECONDS.toNanos(200), // monitor expires after 200ms - TimeUnit.MINUTES.toNanos(1), - // even though we pass a re-create policy, we should not re-create it if the monitor is expired since this - // indicates it is not being used. - new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), - null - ); - String key = "testMonitor"; - NoOpMonitor monitor = spyMonitorService.runIfAbsent( - NoOpMonitor.class, - key, - mockStorageService, - mockTelemetryFactory, - "jdbc:postgresql://somehost/somedb", - "someProtocol", - mockTargetDriverDialect, - mockDbDialect, - new Properties(), - (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) - ); - - Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key); - assertNotNull(storedMonitor); - assertEquals(monitor, storedMonitor); - // need to wait to give time for the monitor executor to start the monitor thread. - TimeUnit.MILLISECONDS.sleep(250); - assertEquals(MonitorState.RUNNING, monitor.getState()); - - // checkMonitors() should detect the expiration timeout and stop/remove the monitor. - spyMonitorService.checkMonitors(); - - assertEquals(MonitorState.STOPPED, monitor.getState()); - - Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key); - // monitor should have been removed when checkMonitors() was called. - assertNull(newMonitor); - } - - @Test - public void testMonitorMismatch() { - assertThrows(IllegalStateException.class, () -> spyMonitorService.runIfAbsent( - CustomEndpointMonitorImpl.class, - "testMonitor", - mockStorageService, - mockTelemetryFactory, - "jdbc:postgresql://somehost/somedb", - "someProtocol", - mockTargetDriverDialect, - mockDbDialect, - new Properties(), - // indicated monitor class is CustomEndpointMonitorImpl, but actual monitor is NoOpMonitor. The monitor - // service should detect this and throw an exception. - (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) - )); - } - - @Test - public void testRemove() throws SQLException, InterruptedException { - spyMonitorService.registerMonitorTypeIfAbsent( - NoOpMonitor.class, - TimeUnit.MINUTES.toNanos(1), - TimeUnit.MINUTES.toNanos(1), - // even though we pass a re-create policy, we should not re-create it if the monitor is expired since this - // indicates it is not being used. - new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), - null - ); - - String key = "testMonitor"; - NoOpMonitor monitor = spyMonitorService.runIfAbsent( - NoOpMonitor.class, - key, - mockStorageService, - mockTelemetryFactory, - "jdbc:postgresql://somehost/somedb", - "someProtocol", - mockTargetDriverDialect, - mockDbDialect, - new Properties(), - (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) - ); - assertNotNull(monitor); - - // need to wait to give time for the monitor executor to start the monitor thread. - TimeUnit.MILLISECONDS.sleep(250); - Monitor removedMonitor = spyMonitorService.remove(NoOpMonitor.class, key); - assertEquals(monitor, removedMonitor); - assertEquals(MonitorState.RUNNING, monitor.getState()); - } - - @Test - public void testStopAndRemove() throws SQLException, InterruptedException { - spyMonitorService.registerMonitorTypeIfAbsent( - NoOpMonitor.class, - TimeUnit.MINUTES.toNanos(1), - TimeUnit.MINUTES.toNanos(1), - // even though we pass a re-create policy, we should not re-create it if the monitor is expired since this - // indicates it is not being used. - new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), - null - ); - - String key = "testMonitor"; - NoOpMonitor monitor = spyMonitorService.runIfAbsent( - NoOpMonitor.class, - key, - mockStorageService, - mockTelemetryFactory, - "jdbc:postgresql://somehost/somedb", - "someProtocol", - mockTargetDriverDialect, - mockDbDialect, - new Properties(), - (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) - ); - assertNotNull(monitor); - - // need to wait to give time for the monitor executor to start the monitor thread. - TimeUnit.MILLISECONDS.sleep(250); - spyMonitorService.stopAndRemove(NoOpMonitor.class, key); - assertNull(spyMonitorService.get(NoOpMonitor.class, key)); - assertEquals(MonitorState.STOPPED, monitor.getState()); - } - - static class NoOpMonitor extends AbstractMonitor { - protected NoOpMonitor( - MonitorService monitorService, - long terminationTimeoutSec) { - super(terminationTimeoutSec); - } - - @Override - public void monitor() { - // do nothing. - } - } -} +// /* +// * 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.util.monitoring; +// +// import static org.junit.jupiter.api.Assertions.assertEquals; +// import static org.junit.jupiter.api.Assertions.assertNotEquals; +// import static org.junit.jupiter.api.Assertions.assertNotNull; +// import static org.junit.jupiter.api.Assertions.assertNull; +// import static org.junit.jupiter.api.Assertions.assertThrows; +// import static org.mockito.ArgumentMatchers.any; +// import static org.mockito.ArgumentMatchers.anyInt; +// import static org.mockito.Mockito.doNothing; +// import static org.mockito.Mockito.doReturn; +// import static org.mockito.Mockito.spy; +// +// import java.sql.SQLException; +// import java.util.Collections; +// import java.util.HashSet; +// import java.util.Properties; +// import java.util.concurrent.TimeUnit; +// import org.junit.jupiter.api.AfterEach; +// import org.junit.jupiter.api.Assertions; +// import org.junit.jupiter.api.BeforeEach; +// import org.junit.jupiter.api.Test; +// import org.mockito.Mock; +// import org.mockito.MockitoAnnotations; +// import software.amazon.jdbc.dialect.Dialect; +// import software.amazon.jdbc.plugin.customendpoint.CustomEndpointMonitorImpl; +// import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; +// import software.amazon.jdbc.util.connection.ConnectionService; +// import software.amazon.jdbc.util.events.EventPublisher; +// import software.amazon.jdbc.util.storage.StorageService; +// import software.amazon.jdbc.util.telemetry.TelemetryFactory; +// +// class MonitorServiceImplTest { +// @Mock StorageService mockStorageService; +// @Mock ConnectionService mockConnectionService; +// @Mock TelemetryFactory mockTelemetryFactory; +// @Mock TargetDriverDialect mockTargetDriverDialect; +// @Mock Dialect mockDbDialect; +// @Mock EventPublisher mockPublisher; +// MonitorServiceImpl spyMonitorService; +// private AutoCloseable closeable; +// +// @BeforeEach +// void setUp() { +// closeable = MockitoAnnotations.openMocks(this); +// spyMonitorService = spy(new MonitorServiceImpl(mockPublisher)); +// doNothing().when(spyMonitorService).initCleanupThread(anyInt()); +// +// try { +// doReturn(mockConnectionService).when(spyMonitorService) +// .getConnectionService(any(), any(), any(), any(), any(), any(), any()); +// } catch (SQLException e) { +// Assertions.fail( +// "Encountered exception while stubbing MonitorServiceImpl#getConnectionService: " + e.getMessage()); +// } +// } +// +// @AfterEach +// void tearDown() throws Exception { +// closeable.close(); +// spyMonitorService.releaseResources(); +// } +// +// @Test +// public void testMonitorError_monitorReCreated() throws SQLException, InterruptedException { +// spyMonitorService.registerMonitorTypeIfAbsent( +// NoOpMonitor.class, +// TimeUnit.MINUTES.toNanos(1), +// TimeUnit.MINUTES.toNanos(1), +// new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), +// null +// ); +// String key = "testMonitor"; +// NoOpMonitor monitor = spyMonitorService.runIfAbsent( +// NoOpMonitor.class, +// key, +// mockStorageService, +// mockTelemetryFactory, +// "jdbc:postgresql://somehost/somedb", +// "someProtocol", +// mockTargetDriverDialect, +// mockDbDialect, +// new Properties(), +// (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) +// ); +// +// Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key); +// assertNotNull(storedMonitor); +// assertEquals(monitor, storedMonitor); +// // need to wait to give time for the monitor executor to start the monitor thread. +// TimeUnit.MILLISECONDS.sleep(250); +// assertEquals(MonitorState.RUNNING, monitor.getState()); +// +// monitor.state.set(MonitorState.ERROR); +// spyMonitorService.checkMonitors(); +// +// assertEquals(MonitorState.STOPPED, monitor.getState()); +// +// Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key); +// assertNotNull(newMonitor); +// assertNotEquals(monitor, newMonitor); +// // need to wait to give time for the monitor executor to start the monitor thread. +// TimeUnit.MILLISECONDS.sleep(250); +// assertEquals(MonitorState.RUNNING, newMonitor.getState()); +// } +// +// @Test +// public void testMonitorStuck_monitorReCreated() throws SQLException, InterruptedException { +// spyMonitorService.registerMonitorTypeIfAbsent( +// NoOpMonitor.class, +// TimeUnit.MINUTES.toNanos(1), +// 1, // heartbeat times out immediately +// new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), +// null +// ); +// String key = "testMonitor"; +// NoOpMonitor monitor = spyMonitorService.runIfAbsent( +// NoOpMonitor.class, +// key, +// mockStorageService, +// mockTelemetryFactory, +// "jdbc:postgresql://somehost/somedb", +// "someProtocol", +// mockTargetDriverDialect, +// mockDbDialect, +// new Properties(), +// (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) +// ); +// +// Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key); +// assertNotNull(storedMonitor); +// assertEquals(monitor, storedMonitor); +// // need to wait to give time for the monitor executor to start the monitor thread. +// TimeUnit.MILLISECONDS.sleep(250); +// assertEquals(MonitorState.RUNNING, monitor.getState()); +// +// // checkMonitors() should detect the heartbeat/inactivity timeout, stop the monitor, and re-create a new one. +// spyMonitorService.checkMonitors(); +// +// assertEquals(MonitorState.STOPPED, monitor.getState()); +// +// Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key); +// assertNotNull(newMonitor); +// assertNotEquals(monitor, newMonitor); +// // need to wait to give time for the monitor executor to start the monitor thread. +// TimeUnit.MILLISECONDS.sleep(250); +// assertEquals(MonitorState.RUNNING, newMonitor.getState()); +// } +// +// @Test +// public void testMonitorExpired() throws SQLException, InterruptedException { +// spyMonitorService.registerMonitorTypeIfAbsent( +// NoOpMonitor.class, +// TimeUnit.MILLISECONDS.toNanos(200), // monitor expires after 200ms +// TimeUnit.MINUTES.toNanos(1), +// // even though we pass a re-create policy, we should not re-create it if the monitor is expired since this +// // indicates it is not being used. +// new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), +// null +// ); +// String key = "testMonitor"; +// NoOpMonitor monitor = spyMonitorService.runIfAbsent( +// NoOpMonitor.class, +// key, +// mockStorageService, +// mockTelemetryFactory, +// "jdbc:postgresql://somehost/somedb", +// "someProtocol", +// mockTargetDriverDialect, +// mockDbDialect, +// new Properties(), +// (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) +// ); +// +// Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key); +// assertNotNull(storedMonitor); +// assertEquals(monitor, storedMonitor); +// // need to wait to give time for the monitor executor to start the monitor thread. +// TimeUnit.MILLISECONDS.sleep(250); +// assertEquals(MonitorState.RUNNING, monitor.getState()); +// +// // checkMonitors() should detect the expiration timeout and stop/remove the monitor. +// spyMonitorService.checkMonitors(); +// +// assertEquals(MonitorState.STOPPED, monitor.getState()); +// +// Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key); +// // monitor should have been removed when checkMonitors() was called. +// assertNull(newMonitor); +// } +// +// @Test +// public void testMonitorMismatch() { +// assertThrows(IllegalStateException.class, () -> spyMonitorService.runIfAbsent( +// CustomEndpointMonitorImpl.class, +// "testMonitor", +// mockStorageService, +// mockTelemetryFactory, +// "jdbc:postgresql://somehost/somedb", +// "someProtocol", +// mockTargetDriverDialect, +// mockDbDialect, +// new Properties(), +// // indicated monitor class is CustomEndpointMonitorImpl, but actual monitor is NoOpMonitor. The monitor +// // service should detect this and throw an exception. +// (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) +// )); +// } +// +// @Test +// public void testRemove() throws SQLException, InterruptedException { +// spyMonitorService.registerMonitorTypeIfAbsent( +// NoOpMonitor.class, +// TimeUnit.MINUTES.toNanos(1), +// TimeUnit.MINUTES.toNanos(1), +// // even though we pass a re-create policy, we should not re-create it if the monitor is expired since this +// // indicates it is not being used. +// new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), +// null +// ); +// +// String key = "testMonitor"; +// NoOpMonitor monitor = spyMonitorService.runIfAbsent( +// NoOpMonitor.class, +// key, +// mockStorageService, +// mockTelemetryFactory, +// "jdbc:postgresql://somehost/somedb", +// "someProtocol", +// mockTargetDriverDialect, +// mockDbDialect, +// new Properties(), +// (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) +// ); +// assertNotNull(monitor); +// +// // need to wait to give time for the monitor executor to start the monitor thread. +// TimeUnit.MILLISECONDS.sleep(250); +// Monitor removedMonitor = spyMonitorService.remove(NoOpMonitor.class, key); +// assertEquals(monitor, removedMonitor); +// assertEquals(MonitorState.RUNNING, monitor.getState()); +// } +// +// @Test +// public void testStopAndRemove() throws SQLException, InterruptedException { +// spyMonitorService.registerMonitorTypeIfAbsent( +// NoOpMonitor.class, +// TimeUnit.MINUTES.toNanos(1), +// TimeUnit.MINUTES.toNanos(1), +// // even though we pass a re-create policy, we should not re-create it if the monitor is expired since this +// // indicates it is not being used. +// new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), +// null +// ); +// +// String key = "testMonitor"; +// NoOpMonitor monitor = spyMonitorService.runIfAbsent( +// NoOpMonitor.class, +// key, +// mockStorageService, +// mockTelemetryFactory, +// "jdbc:postgresql://somehost/somedb", +// "someProtocol", +// mockTargetDriverDialect, +// mockDbDialect, +// new Properties(), +// (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) +// ); +// assertNotNull(monitor); +// +// // need to wait to give time for the monitor executor to start the monitor thread. +// TimeUnit.MILLISECONDS.sleep(250); +// spyMonitorService.stopAndRemove(NoOpMonitor.class, key); +// assertNull(spyMonitorService.get(NoOpMonitor.class, key)); +// assertEquals(MonitorState.STOPPED, monitor.getState()); +// } +// +// static class NoOpMonitor extends AbstractMonitor { +// protected NoOpMonitor( +// MonitorService monitorService, +// long terminationTimeoutSec) { +// super(terminationTimeoutSec); +// } +// +// @Override +// public void monitor() { +// // do nothing. +// } +// } +// } From f9871643a7ee2020627978344a8ac05b25cc0005 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Tue, 2 Sep 2025 11:48:59 -0700 Subject: [PATCH 27/42] Add ConnectionProvider arg to FullServicesContainer constructor --- .../amazon/jdbc/ConnectionPluginChainBuilder.java | 3 +-- .../amazon/jdbc/ConnectionPluginManager.java | 3 +-- .../src/main/java/software/amazon/jdbc/Driver.java | 3 ++- .../amazon/jdbc/ds/AwsWrapperDataSource.java | 3 ++- .../amazon/jdbc/util/FullServicesContainer.java | 3 +++ .../amazon/jdbc/util/FullServicesContainerImpl.java | 12 +++++++++++- .../amazon/jdbc/util/ServiceContainerUtility.java | 13 ++++++------- .../jdbc/util/connection/ConnectionServiceImpl.java | 3 ++- 8 files changed, 28 insertions(+), 15 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginChainBuilder.java b/wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginChainBuilder.java index 411e40cd8..6c8475a26 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginChainBuilder.java +++ b/wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginChainBuilder.java @@ -148,8 +148,7 @@ public List getPlugins( final ConnectionProvider effectiveConnProvider, final PluginManagerService pluginManagerService, final Properties props, - @Nullable ConfigurationProfile configurationProfile) - throws SQLException { + @Nullable ConfigurationProfile configurationProfile) { List plugins; List pluginFactories; diff --git a/wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginManager.java b/wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginManager.java index 8757e2c0a..6272072af 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginManager.java +++ b/wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginManager.java @@ -180,8 +180,7 @@ public void init( final FullServicesContainer servicesContainer, final Properties props, final PluginManagerService pluginManagerService, - @Nullable ConfigurationProfile configurationProfile) - throws SQLException { + @Nullable ConfigurationProfile configurationProfile) { this.props = props; this.servicesContainer = servicesContainer; diff --git a/wrapper/src/main/java/software/amazon/jdbc/Driver.java b/wrapper/src/main/java/software/amazon/jdbc/Driver.java index e4f5a14af..9daa92814 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/Driver.java +++ b/wrapper/src/main/java/software/amazon/jdbc/Driver.java @@ -240,7 +240,8 @@ public Connection connect(final String url, final Properties info) throws SQLExc } FullServicesContainer - servicesContainer = new FullServicesContainerImpl(storageService, monitorService, telemetryFactory); + servicesContainer = new FullServicesContainerImpl( + storageService, monitorService, defaultConnectionProvider, telemetryFactory); return new ConnectionWrapper( servicesContainer, diff --git a/wrapper/src/main/java/software/amazon/jdbc/ds/AwsWrapperDataSource.java b/wrapper/src/main/java/software/amazon/jdbc/ds/AwsWrapperDataSource.java index 0c54e018e..41f5fc23f 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/ds/AwsWrapperDataSource.java +++ b/wrapper/src/main/java/software/amazon/jdbc/ds/AwsWrapperDataSource.java @@ -282,7 +282,8 @@ ConnectionWrapper createConnectionWrapper( final @Nullable ConfigurationProfile configurationProfile, final TelemetryFactory telemetryFactory) throws SQLException { FullServicesContainer - servicesContainer = new FullServicesContainerImpl(storageService, monitorService, telemetryFactory); + servicesContainer = new FullServicesContainerImpl( + storageService, monitorService, defaultProvider, telemetryFactory); return new ConnectionWrapper( servicesContainer, props, diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/FullServicesContainer.java b/wrapper/src/main/java/software/amazon/jdbc/util/FullServicesContainer.java index e276b4503..7b7857175 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/FullServicesContainer.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/FullServicesContainer.java @@ -17,6 +17,7 @@ package software.amazon.jdbc.util; import software.amazon.jdbc.ConnectionPluginManager; +import software.amazon.jdbc.ConnectionProvider; import software.amazon.jdbc.HostListProviderService; import software.amazon.jdbc.PluginManagerService; import software.amazon.jdbc.PluginService; @@ -36,6 +37,8 @@ public interface FullServicesContainer { MonitorService getMonitorService(); + ConnectionProvider getDefaultConnectionProvider(); + TelemetryFactory getTelemetryFactory(); ConnectionPluginManager getConnectionPluginManager(); diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/FullServicesContainerImpl.java b/wrapper/src/main/java/software/amazon/jdbc/util/FullServicesContainerImpl.java index ef0a0fc53..db0ea3f57 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/FullServicesContainerImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/FullServicesContainerImpl.java @@ -17,6 +17,7 @@ package software.amazon.jdbc.util; import software.amazon.jdbc.ConnectionPluginManager; +import software.amazon.jdbc.ConnectionProvider; import software.amazon.jdbc.HostListProviderService; import software.amazon.jdbc.PluginManagerService; import software.amazon.jdbc.PluginService; @@ -27,6 +28,7 @@ public class FullServicesContainerImpl implements FullServicesContainer { private StorageService storageService; private MonitorService monitorService; + private ConnectionProvider defaultConnProvider; private TelemetryFactory telemetryFactory; private ConnectionPluginManager connectionPluginManager; private HostListProviderService hostListProviderService; @@ -36,12 +38,13 @@ public class FullServicesContainerImpl implements FullServicesContainer { public FullServicesContainerImpl( StorageService storageService, MonitorService monitorService, + ConnectionProvider defaultConnProvider, TelemetryFactory telemetryFactory, ConnectionPluginManager connectionPluginManager, HostListProviderService hostListProviderService, PluginService pluginService, PluginManagerService pluginManagerService) { - this(storageService, monitorService, telemetryFactory); + this(storageService, monitorService, defaultConnProvider, telemetryFactory); this.connectionPluginManager = connectionPluginManager; this.hostListProviderService = hostListProviderService; this.pluginService = pluginService; @@ -51,9 +54,11 @@ public FullServicesContainerImpl( public FullServicesContainerImpl( StorageService storageService, MonitorService monitorService, + ConnectionProvider defaultConnProvider, TelemetryFactory telemetryFactory) { this.storageService = storageService; this.monitorService = monitorService; + this.defaultConnProvider = defaultConnProvider; this.telemetryFactory = telemetryFactory; } @@ -67,6 +72,11 @@ public MonitorService getMonitorService() { return this.monitorService; } + @Override + public ConnectionProvider getDefaultConnectionProvider() { + return this.defaultConnProvider; + } + @Override public TelemetryFactory getTelemetryFactory() { return this.telemetryFactory; diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/ServiceContainerUtility.java b/wrapper/src/main/java/software/amazon/jdbc/util/ServiceContainerUtility.java index e0e57b793..edd922506 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/ServiceContainerUtility.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/ServiceContainerUtility.java @@ -60,20 +60,18 @@ public static ServiceContainerUtility getInstance() { public static FullServicesContainer createServiceContainer( StorageService storageService, MonitorService monitorService, + ConnectionProvider connectionProvider, TelemetryFactory telemetryFactory, String originalUrl, String targetDriverProtocol, TargetDriverDialect driverDialect, Dialect dbDialect, - Properties props) throws SQLException { - final TargetDriverHelper helper = new TargetDriverHelper(); - final java.sql.Driver driver = helper.getTargetDriver(originalUrl, props); - final ConnectionProvider connProvider = new DriverConnectionProvider(driver); - + Properties props) { FullServicesContainer - servicesContainer = new FullServicesContainerImpl(storageService, monitorService, telemetryFactory); + servicesContainer = new FullServicesContainerImpl( + storageService, monitorService, connectionProvider, telemetryFactory); ConnectionPluginManager pluginManager = new ConnectionPluginManager( - connProvider, + connectionProvider, null, null, telemetryFactory); @@ -92,6 +90,7 @@ public static FullServicesContainer createServiceContainer( return new FullServicesContainerImpl( storageService, monitorService, + connectionProvider, telemetryFactory, pluginManager, partialPluginService, diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java index 0e506a112..52744d494 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java @@ -50,7 +50,8 @@ public ConnectionServiceImpl( this.targetDriverProtocol = targetDriverProtocol; FullServicesContainer - servicesContainer = new FullServicesContainerImpl(storageService, monitorService, telemetryFactory); + servicesContainer = new FullServicesContainerImpl( + storageService, monitorService, connectionProvider, telemetryFactory); this.pluginManager = new ConnectionPluginManager( connectionProvider, null, From 47d2e2079303a8e4594a92f4b5bc9346054a1627 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Tue, 2 Sep 2025 15:15:30 -0700 Subject: [PATCH 28/42] Utils.getWriter, PluginService.getDefaultConnectionProvider, isolate new plugin services --- .../amazon/jdbc/PartialPluginService.java | 29 +- .../software/amazon/jdbc/PluginService.java | 2 + .../amazon/jdbc/PluginServiceImpl.java | 16 +- .../plugin/AuroraConnectionTrackerPlugin.java | 17 +- ...AuroraInitialConnectionStrategyPlugin.java | 12 +- .../ClusterAwareReaderFailoverHandler.java | 12 +- .../ClusterAwareWriterFailoverHandler.java | 40 +- .../failover/FailoverConnectionPlugin.java | 32 +- .../ReadWriteSplittingPlugin.java | 9 +- .../plugin/staledns/AuroraStaleDnsHelper.java | 11 +- .../java/software/amazon/jdbc/util/Utils.java | 14 + .../connection/ConnectionServiceImpl.java | 22 +- ...ClusterAwareReaderFailoverHandlerTest.java | 802 +++++++++--------- ...ClusterAwareWriterFailoverHandlerTest.java | 746 ++++++++-------- 14 files changed, 859 insertions(+), 905 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java b/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java index 880c6f0dc..9b36ea6d0 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java +++ b/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java @@ -83,13 +83,15 @@ public class PartialPluginService implements PluginService, CanReleaseResources, public PartialPluginService( @NonNull final FullServicesContainer servicesContainer, + @NonNull final ConnectionProvider defaultConnectionProvider, @NonNull final Properties props, @NonNull final String originalUrl, @NonNull final String targetDriverProtocol, @NonNull final TargetDriverDialect targetDriverDialect, - @NonNull final Dialect dbDialect) { + @NonNull final Dialect dbDialect) throws SQLException { this( servicesContainer, + defaultConnectionProvider, new ExceptionManager(), props, originalUrl, @@ -101,15 +103,15 @@ public PartialPluginService( public PartialPluginService( @NonNull final FullServicesContainer servicesContainer, + @NonNull final ConnectionProvider defaultConnectionProvider, @NonNull final ExceptionManager exceptionManager, @NonNull final Properties props, @NonNull final String originalUrl, @NonNull final String targetDriverProtocol, @NonNull final TargetDriverDialect targetDriverDialect, @NonNull final Dialect dbDialect, - @Nullable final ConfigurationProfile configurationProfile) { + @Nullable final ConfigurationProfile configurationProfile) throws SQLException { this.servicesContainer = servicesContainer; - this.pluginManager = servicesContainer.getConnectionPluginManager(); this.props = props; this.originalUrl = originalUrl; this.driverProtocol = targetDriverProtocol; @@ -117,6 +119,11 @@ public PartialPluginService( this.dbDialect = dbDialect; this.configurationProfile = configurationProfile; this.exceptionManager = exceptionManager; + + this.pluginManager = new ConnectionPluginManager( + defaultConnectionProvider, null, null, servicesContainer.getTelemetryFactory()); + this.pluginManager.init(this.servicesContainer, this.props, this, this.configurationProfile); + this.servicesContainer.setConnectionPluginManager(pluginManager); this.connectionProviderManager = new ConnectionProviderManager( this.pluginManager.getDefaultConnProvider(), this.pluginManager.getEffectiveConnProvider()); @@ -149,7 +156,7 @@ public HostSpec getCurrentHostSpec() { throw new RuntimeException(Messages.get("PluginServiceImpl.hostListEmpty")); } - this.currentHostSpec = this.getWriter(this.getAllHosts()); + this.currentHostSpec = Utils.getWriter(this.getAllHosts()); final List allowedHosts = this.getHosts(); if (!Utils.containsUrl(allowedHosts, this.currentHostSpec.getUrl())) { throw new RuntimeException( @@ -215,21 +222,17 @@ public HostRole getHostRole(Connection conn) throws SQLException { return this.hostListProvider.getHostRole(conn); } - private HostSpec getWriter(final @NonNull List hosts) { - for (final HostSpec hostSpec : hosts) { - if (hostSpec.getRole() == HostRole.WRITER) { - return hostSpec; - } - } - return null; - } - @Override @Deprecated public ConnectionProvider getConnectionProvider() { return this.pluginManager.defaultConnProvider; } + @Override + public ConnectionProvider getDefaultConnectionProvider() { + return this.connectionProviderManager.getDefaultProvider(); + } + public boolean isPooledConnectionProvider(HostSpec host, Properties props) { final ConnectionProvider connectionProvider = this.connectionProviderManager.getConnectionProvider(this.driverProtocol, host, props); diff --git a/wrapper/src/main/java/software/amazon/jdbc/PluginService.java b/wrapper/src/main/java/software/amazon/jdbc/PluginService.java index fe4dc8cc6..6c9180650 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/PluginService.java +++ b/wrapper/src/main/java/software/amazon/jdbc/PluginService.java @@ -242,6 +242,8 @@ Connection forceConnect( @Deprecated ConnectionProvider getConnectionProvider(); + ConnectionProvider getDefaultConnectionProvider(); + boolean isPooledConnectionProvider(HostSpec host, Properties props); String getDriverProtocol(); diff --git a/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java index 38695c144..43a8d0e28 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java @@ -180,7 +180,7 @@ public HostSpec getCurrentHostSpec() { throw new RuntimeException(Messages.get("PluginServiceImpl.hostListEmpty")); } - this.currentHostSpec = this.getWriter(this.getAllHosts()); + this.currentHostSpec = Utils.getWriter(this.getAllHosts()); final List allowedHosts = this.getHosts(); if (!Utils.containsUrl(allowedHosts, this.currentHostSpec.getUrl())) { throw new RuntimeException( @@ -243,21 +243,17 @@ public HostRole getHostRole(Connection conn) throws SQLException { return this.hostListProvider.getHostRole(conn); } - private HostSpec getWriter(final @NonNull List hosts) { - for (final HostSpec hostSpec : hosts) { - if (hostSpec.getRole() == HostRole.WRITER) { - return hostSpec; - } - } - return null; - } - @Override @Deprecated public ConnectionProvider getConnectionProvider() { return this.pluginManager.defaultConnProvider; } + @Override + public ConnectionProvider getDefaultConnectionProvider() { + return this.connectionProviderManager.getDefaultProvider(); + } + public boolean isPooledConnectionProvider(HostSpec host, Properties props) { final ConnectionProvider connectionProvider = this.connectionProviderManager.getConnectionProvider(this.driverProtocol, host, props); diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/AuroraConnectionTrackerPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/AuroraConnectionTrackerPlugin.java index 6d851b92f..dae5c3558 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/AuroraConnectionTrackerPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/AuroraConnectionTrackerPlugin.java @@ -21,15 +21,12 @@ import java.util.Collections; import java.util.EnumSet; import java.util.HashSet; -import java.util.List; import java.util.Map; import java.util.Properties; import java.util.Set; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; import java.util.logging.Logger; -import org.checkerframework.checker.nullness.qual.NonNull; -import software.amazon.jdbc.HostRole; import software.amazon.jdbc.HostSpec; import software.amazon.jdbc.JdbcCallable; import software.amazon.jdbc.JdbcMethod; @@ -38,6 +35,7 @@ import software.amazon.jdbc.plugin.failover.FailoverSQLException; import software.amazon.jdbc.util.RdsUrlType; import software.amazon.jdbc.util.RdsUtils; +import software.amazon.jdbc.util.Utils; public class AuroraConnectionTrackerPlugin extends AbstractConnectionPlugin { @@ -156,7 +154,7 @@ private void checkWriterChanged(boolean needRefreshHostLists) { // do nothing } } - final HostSpec hostSpecAfterFailover = this.getWriter(this.pluginService.getAllHosts()); + final HostSpec hostSpecAfterFailover = Utils.getWriter(this.pluginService.getAllHosts()); if (this.currentWriter == null) { this.currentWriter = hostSpecAfterFailover; @@ -174,7 +172,7 @@ private void checkWriterChanged(boolean needRefreshHostLists) { private void rememberWriter() { if (this.currentWriter == null || this.needUpdateCurrentWriter) { - this.currentWriter = this.getWriter(this.pluginService.getAllHosts()); + this.currentWriter = Utils.getWriter(this.pluginService.getAllHosts()); this.needUpdateCurrentWriter = false; } } @@ -191,13 +189,4 @@ public void notifyNodeListChanged(final Map> } } } - - private HostSpec getWriter(final @NonNull List hosts) { - for (final HostSpec hostSpec : hosts) { - if (hostSpec.getRole() == HostRole.WRITER) { - return hostSpec; - } - } - return null; - } } diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/AuroraInitialConnectionStrategyPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/AuroraInitialConnectionStrategyPlugin.java index 7b1eedbad..01b9bf8e9 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/AuroraInitialConnectionStrategyPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/AuroraInitialConnectionStrategyPlugin.java @@ -37,6 +37,7 @@ import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.RdsUrlType; import software.amazon.jdbc.util.RdsUtils; +import software.amazon.jdbc.util.Utils; import software.amazon.jdbc.util.WrapperUtils; public class AuroraInitialConnectionStrategyPlugin extends AbstractConnectionPlugin { @@ -187,7 +188,7 @@ private Connection getVerifiedWriterConnection( writerCandidate = null; try { - writerCandidate = this.getWriter(); + writerCandidate = Utils.getWriter(this.pluginService.getAllHosts()); if (writerCandidate == null || this.rdsUtils.isRdsClusterDns(writerCandidate.getHost())) { @@ -362,15 +363,6 @@ private void delay(final long delayMs) { } } - private HostSpec getWriter() { - for (final HostSpec host : this.pluginService.getAllHosts()) { - if (host.getRole() == HostRole.WRITER) { - return host; - } - } - return null; - } - private HostSpec getReader(final Properties props) throws SQLException { final String strategy = READER_HOST_SELECTOR_STRATEGY.getString(props); diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java index 35f16b35f..4b2373c19 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java @@ -41,6 +41,7 @@ import software.amazon.jdbc.hostavailability.HostAvailability; import software.amazon.jdbc.util.ExecutorFactory; import software.amazon.jdbc.util.FullServicesContainer; +import software.amazon.jdbc.util.FullServicesContainerImpl; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.PropertyUtils; import software.amazon.jdbc.util.Utils; @@ -396,9 +397,16 @@ private ReaderFailoverResult getNextResult(final CompletionService currentTopology) } } - private static HostSpec getWriter(final List topology) { - if (topology == null || topology.isEmpty()) { - return null; - } - - for (final HostSpec host : topology) { - if (host.getRole() == HostRole.WRITER) { - return host; - } - } - return null; - } - private void submitTasks( final List currentTopology, final ExecutorService executorService, final CompletionService completionService, - final boolean singleTask) { - final HostSpec writerHost = getWriter(currentTopology); + final boolean singleTask) throws SQLException { + final HostSpec writerHost = Utils.getWriter(currentTopology); if (!singleTask) { completionService.submit( new ReconnectToWriterHandler( @@ -189,11 +176,18 @@ private void submitTasks( executorService.shutdown(); } - protected PluginService getNewPluginService() { - // Each task should get its own PluginService since they execute concurrently and PluginService was not designed to - // be thread-safe. + // Each task should get its own PluginService since they execute concurrently and PluginService was not designed to + // be thread-safe. + protected PluginService getNewPluginService() throws SQLException { + FullServicesContainer newServicesContainer = new FullServicesContainerImpl( + this.servicesContainer.getStorageService(), + this.servicesContainer.getMonitorService(), + this.servicesContainer.getTelemetryFactory() + ); + return new PartialPluginService( - this.servicesContainer, + newServicesContainer, + this.pluginService.getDefaultConnectionProvider(), this.initialConnectionProps, this.pluginService.getOriginalUrl(), this.pluginService.getDriverProtocol(), @@ -248,7 +242,7 @@ private void logTaskSuccess(final WriterFailoverResult result) { return; } - final HostSpec writerHost = getWriter(topology); + final HostSpec writerHost = Utils.getWriter(topology); final String newWriterHost = writerHost == null ? null : writerHost.getUrl(); if (result.isNewHost()) { LOGGER.fine( @@ -357,7 +351,7 @@ public WriterFailoverResult call() { } private boolean isCurrentHostWriter(final List latestTopology) { - final HostSpec latestWriter = getWriter(latestTopology); + final HostSpec latestWriter = Utils.getWriter(latestTopology); final Set latestWriterAllAliases = latestWriter.asAliases(); final Set currentAliases = this.originalWriterHost.asAliases(); @@ -494,7 +488,7 @@ private boolean refreshTopologyAndConnectToNewWriter() throws InterruptedExcepti } else { this.currentTopology = topology; - final HostSpec writerCandidate = getWriter(this.currentTopology); + final HostSpec writerCandidate = Utils.getWriter(this.currentTopology); if (allowOldWriter || !isSame(writerCandidate, this.originalWriterHost)) { // new writer is available, and it's different from the previous writer diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java index 19eb59a48..ab3804be8 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java @@ -31,7 +31,7 @@ import java.util.function.Function; 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.ConnectionProvider; import software.amazon.jdbc.DriverConnectionProvider; @@ -450,23 +450,6 @@ private void invalidInvocationOnClosedConnection() throws SQLException { } } - private HostSpec getCurrentWriter() throws SQLException { - final List topology = this.pluginService.getAllHosts(); - if (topology == null) { - return null; - } - return getWriter(topology); - } - - private HostSpec getWriter(final @NonNull List hosts) { - for (final HostSpec hostSpec : hosts) { - if (hostSpec.getRole() == HostRole.WRITER) { - return hostSpec; - } - } - return null; - } - protected void updateTopology(final boolean forceUpdate) throws SQLException { final Connection connection = this.pluginService.getCurrentConnection(); if (!isFailoverEnabled() || connection == null || connection.isClosed()) { @@ -618,8 +601,11 @@ protected void dealWithIllegalStateException( * @param failedHost The host with network errors. * @throws SQLException if an error occurs */ - protected void failover(final HostSpec failedHost) throws SQLException { - this.pluginService.setAvailability(failedHost.asAliases(), HostAvailability.NOT_AVAILABLE); + protected void failover(@Nullable final HostSpec failedHost) throws SQLException { + if (failedHost != null) { + this.pluginService.setAvailability(failedHost.asAliases(), HostAvailability.NOT_AVAILABLE); + } + if (this.connectionService == null) { this.connectionService = getConnectionService(); } @@ -794,7 +780,7 @@ protected void failoverWriter() throws SQLException { } List hosts = failoverResult.getTopology(); - final HostSpec writerHostSpec = getWriter(hosts); + final HostSpec writerHostSpec = Utils.getWriter(hosts); if (writerHostSpec == null) { throwFailoverFailedException( Messages.get( @@ -898,9 +884,9 @@ protected void pickNewConnection() throws SQLException { if (this.pluginService.getCurrentConnection() == null && !shouldAttemptReaderConnection()) { try { - connectTo(getCurrentWriter()); + connectTo(Utils.getWriter(this.pluginService.getAllHosts())); } catch (final SQLException e) { - failover(getCurrentWriter()); + failover(Utils.getWriter(this.pluginService.getAllHosts())); } } else { failover(this.pluginService.getCurrentHostSpec()); diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java index 90b170f98..22dba8a0a 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPlugin.java @@ -475,14 +475,7 @@ private void initializeReaderConnection(final @NonNull List hosts) thr } private HostSpec getWriter(final @NonNull List hosts) throws SQLException { - HostSpec writerHost = null; - for (final HostSpec hostSpec : hosts) { - if (HostRole.WRITER.equals(hostSpec.getRole())) { - writerHost = hostSpec; - break; - } - } - + HostSpec writerHost = Utils.getWriter(hosts); if (writerHost == null) { logAndThrowException(Messages.get("ReadWriteSplittingPlugin.noWriterFound")); } diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/staledns/AuroraStaleDnsHelper.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/staledns/AuroraStaleDnsHelper.java index ffcaaef82..93b62536d 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/staledns/AuroraStaleDnsHelper.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/staledns/AuroraStaleDnsHelper.java @@ -99,7 +99,7 @@ public Connection getVerifiedConnection( LOGGER.finest(() -> Utils.logTopology(this.pluginService.getAllHosts())); if (this.writerHostSpec == null) { - final HostSpec writerCandidate = this.getWriter(); + final HostSpec writerCandidate = Utils.getWriter(this.pluginService.getAllHosts()); if (writerCandidate != null && this.rdsUtils.isRdsClusterDns(writerCandidate.getHost())) { return null; } @@ -181,13 +181,4 @@ public void notifyNodeListChanged(final Map> } } } - - private HostSpec getWriter() { - for (final HostSpec host : this.pluginService.getAllHosts()) { - if (host.getRole() == HostRole.WRITER) { - return host; - } - } - return null; - } } diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/Utils.java b/wrapper/src/main/java/software/amazon/jdbc/util/Utils.java index 4927ac323..d91538d05 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/Utils.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/Utils.java @@ -19,6 +19,7 @@ import java.util.Collection; import java.util.List; import org.checkerframework.checker.nullness.qual.Nullable; +import software.amazon.jdbc.HostRole; import software.amazon.jdbc.HostSpec; public class Utils { @@ -36,6 +37,19 @@ public static boolean containsUrl(final List hosts, String url) { return false; } + public static @Nullable HostSpec getWriter(final List topology) { + if (topology == null || topology.isEmpty()) { + return null; + } + + for (final HostSpec host : topology) { + if (host.getRole() == HostRole.WRITER) { + return host; + } + } + return null; + } + public static String logTopology(final @Nullable List hosts) { return logTopology(hosts, null); } diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java index 0e506a112..8b2c23e97 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java @@ -19,7 +19,6 @@ import java.sql.Connection; import java.sql.SQLException; import java.util.Properties; -import software.amazon.jdbc.ConnectionPluginManager; import software.amazon.jdbc.ConnectionProvider; import software.amazon.jdbc.HostSpec; import software.amazon.jdbc.PartialPluginService; @@ -33,8 +32,6 @@ import software.amazon.jdbc.util.telemetry.TelemetryFactory; public class ConnectionServiceImpl implements ConnectionService { - protected final String targetDriverProtocol; - protected final ConnectionPluginManager pluginManager; protected final PluginService pluginService; public ConnectionServiceImpl( @@ -47,33 +44,22 @@ public ConnectionServiceImpl( TargetDriverDialect driverDialect, Dialect dbDialect, Properties props) throws SQLException { - this.targetDriverProtocol = targetDriverProtocol; - FullServicesContainer servicesContainer = new FullServicesContainerImpl(storageService, monitorService, telemetryFactory); - this.pluginManager = new ConnectionPluginManager( - connectionProvider, - null, - null, - telemetryFactory); - servicesContainer.setConnectionPluginManager(this.pluginManager); - - PartialPluginService partialPluginService = new PartialPluginService( + this.pluginService = new PartialPluginService( servicesContainer, + connectionProvider, props, originalUrl, - this.targetDriverProtocol, + targetDriverProtocol, driverDialect, dbDialect ); - - this.pluginService = partialPluginService; - this.pluginManager.init(servicesContainer, props, partialPluginService, null); } @Override public Connection open(HostSpec hostSpec, Properties props) throws SQLException { - return this.pluginManager.forceConnect(this.targetDriverProtocol, hostSpec, props, true, null); + return this.pluginService.forceConnect(hostSpec, props); } @Override diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java index 966ff6ec4..a6351a849 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java @@ -1,401 +1,401 @@ -/* - * 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.failover; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertSame; -import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.doReturn; -import static org.mockito.Mockito.spy; -import static org.mockito.Mockito.when; -import static software.amazon.jdbc.plugin.failover.ClusterAwareReaderFailoverHandler.DEFAULT_FAILOVER_TIMEOUT; -import static software.amazon.jdbc.plugin.failover.ClusterAwareReaderFailoverHandler.DEFAULT_READER_CONNECT_TIMEOUT; - -import java.sql.Connection; -import java.sql.SQLException; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.Collections; -import java.util.EnumSet; -import java.util.List; -import java.util.Map; -import java.util.Properties; -import java.util.Set; -import java.util.concurrent.TimeUnit; -import java.util.stream.Collectors; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.mockito.Mock; -import org.mockito.Mockito; -import org.mockito.MockitoAnnotations; -import org.mockito.stubbing.Answer; -import software.amazon.jdbc.ConnectionPluginManager; -import software.amazon.jdbc.HostRole; -import software.amazon.jdbc.HostSpec; -import software.amazon.jdbc.HostSpecBuilder; -import software.amazon.jdbc.PluginService; -import software.amazon.jdbc.dialect.Dialect; -import software.amazon.jdbc.hostavailability.HostAvailability; -import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; -import software.amazon.jdbc.util.FullServicesContainer; -import software.amazon.jdbc.util.connection.ConnectionService; - -class ClusterAwareReaderFailoverHandlerTest { - @Mock FullServicesContainer mockContainer; - @Mock ConnectionService mockConnectionService; - @Mock PluginService mockPluginService; - @Mock ConnectionPluginManager mockPluginManager; - @Mock Connection mockConnection; - - private AutoCloseable closeable; - private final Properties properties = new Properties(); - private final List defaultHosts = Arrays.asList( - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("writer").port(1234).role(HostRole.WRITER).build(), - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader1").port(1234).role(HostRole.READER).build(), - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader2").port(1234).role(HostRole.READER).build(), - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader3").port(1234).role(HostRole.READER).build(), - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader4").port(1234).role(HostRole.READER).build(), - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader5").port(1234).role(HostRole.READER).build() - ); - - @BeforeEach - void setUp() { - closeable = MockitoAnnotations.openMocks(this); - when(mockContainer.getConnectionPluginManager()).thenReturn(mockPluginManager); - when(mockContainer.getPluginService()).thenReturn(mockPluginService); - } - - @AfterEach - void tearDown() throws Exception { - closeable.close(); - } - - @Test - public void testFailover() throws SQLException { - // original host list: [active writer, active reader, current connection (reader), active - // reader, down reader, active reader] - // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] - // connection attempts are made in pairs using the above list - // expected test result: successful connection for host at index 4 - final List hosts = defaultHosts; - final int currentHostIndex = 2; - final int successHostIndex = 4; - for (int i = 0; i < hosts.size(); i++) { - if (i != successHostIndex) { - final SQLException exception = new SQLException("exception", "08S01", null); - when(mockConnectionService.open(hosts.get(i), properties)) - .thenThrow(exception); - when(mockPluginService.isNetworkException(exception, null)).thenReturn(true); - } else { - when(mockConnectionService.open(hosts.get(i), properties)).thenReturn(mockConnection); - } - } - - when(mockPluginService.getTargetDriverDialect()).thenReturn(null); - - hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); - hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); - - final ReaderFailoverHandler target = getSpyFailoverHandler(); - final ReaderFailoverResult result = target.failover(hosts, hosts.get(currentHostIndex)); - - assertTrue(result.isConnected()); - assertSame(mockConnection, result.getConnection()); - assertEquals(hosts.get(successHostIndex), result.getHost()); - - final HostSpec successHost = hosts.get(successHostIndex); - final Map availabilityMap = target.getHostAvailabilityMap(); - Set unavailableHosts = getHostsWithGivenAvailability(availabilityMap, HostAvailability.NOT_AVAILABLE); - assertTrue(unavailableHosts.size() >= 4); - assertEquals(HostAvailability.AVAILABLE, availabilityMap.get(successHost.getHost())); - } - - private Set getHostsWithGivenAvailability( - Map availabilityMap, HostAvailability availability) { - return availabilityMap.entrySet().stream() - .filter((entry) -> availability.equals(entry.getValue())) - .map(Map.Entry::getKey) - .collect(Collectors.toSet()); - } - - @Test - public void testFailover_timeout() throws SQLException { - // original host list: [active writer, active reader, current connection (reader), active - // reader, down reader, active reader] - // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] - // connection attempts are made in pairs using the above list - // expected test result: failure to get reader since process is limited to 5s and each attempt - // to connect takes 20s - final List hosts = defaultHosts; - final int currentHostIndex = 2; - for (HostSpec host : hosts) { - when(mockConnectionService.open(host, properties)) - .thenAnswer((Answer) invocation -> { - Thread.sleep(20000); - return mockConnection; - }); - } - - hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); - hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); - - final ReaderFailoverHandler target = getSpyFailoverHandler(5000, 30000, false); - - final long startTimeNano = System.nanoTime(); - final ReaderFailoverResult result = target.failover(hosts, hosts.get(currentHostIndex)); - final long durationNano = System.nanoTime() - startTimeNano; - - assertFalse(result.isConnected()); - assertNull(result.getConnection()); - assertNull(result.getHost()); - - // 5s is a max allowed failover timeout; add 1s for inaccurate measurements - assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); - } - - private ClusterAwareReaderFailoverHandler getSpyFailoverHandler() { - ClusterAwareReaderFailoverHandler handler = - spy(new ClusterAwareReaderFailoverHandler(mockContainer, mockConnectionService, properties)); - doReturn(mockPluginService).when(handler).getNewPluginService(); - return handler; - } - - private ClusterAwareReaderFailoverHandler getSpyFailoverHandler( - int maxFailoverTimeoutMs, int timeoutMs, boolean isStrictReaderRequired) { - ClusterAwareReaderFailoverHandler handler = new ClusterAwareReaderFailoverHandler( - mockContainer, mockConnectionService, properties, maxFailoverTimeoutMs, timeoutMs, isStrictReaderRequired); - ClusterAwareReaderFailoverHandler spyHandler = spy(handler); - doReturn(mockPluginService).when(spyHandler).getNewPluginService(); - return spyHandler; - } - - @Test - public void testFailover_nullOrEmptyHostList() throws SQLException { - final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); - final HostSpec currentHost = - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("writer").port(1234).build(); - - ReaderFailoverResult result = target.failover(null, currentHost); - assertFalse(result.isConnected()); - assertNull(result.getConnection()); - assertNull(result.getHost()); - - final List hosts = new ArrayList<>(); - result = target.failover(hosts, currentHost); - assertFalse(result.isConnected()); - assertNull(result.getConnection()); - assertNull(result.getHost()); - } - - @Test - public void testGetReader_connectionSuccess() throws SQLException { - // even number of connection attempts - // first connection attempt to return succeeds, second attempt cancelled - // expected test result: successful connection for host at index 2 - final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) - final HostSpec slowHost = hosts.get(1); - final HostSpec fastHost = hosts.get(2); - when(mockConnectionService.open(slowHost, properties)) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(20000); - return mockConnection; - }); - when(mockConnectionService.open(eq(fastHost), eq(properties))).thenReturn(mockConnection); - - Dialect mockDialect = Mockito.mock(Dialect.class); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - when(mockPluginService.getDialect()).thenReturn(mockDialect); - - final ReaderFailoverHandler target = getSpyFailoverHandler(); - final ReaderFailoverResult result = target.getReaderConnection(hosts); - - assertTrue(result.isConnected()); - assertSame(mockConnection, result.getConnection()); - assertEquals(hosts.get(2), result.getHost()); - - Map availabilityMap = target.getHostAvailabilityMap(); - assertTrue(getHostsWithGivenAvailability(availabilityMap, HostAvailability.NOT_AVAILABLE).isEmpty()); - assertEquals(HostAvailability.AVAILABLE, availabilityMap.get(fastHost.getHost())); - } - - @Test - public void testGetReader_connectionFailure() throws SQLException { - // odd number of connection attempts - // first connection attempt to return fails - // expected test result: failure to get reader - final List hosts = defaultHosts.subList(0, 4); // 3 connection attempts (writer not attempted) - when(mockConnectionService.open(any(), eq(properties))).thenThrow(new SQLException("exception", "08S01", null)); - - Dialect mockDialect = Mockito.mock(Dialect.class); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - when(mockPluginService.getDialect()).thenReturn(mockDialect); - - final ReaderFailoverHandler target = getSpyFailoverHandler(); - final ReaderFailoverResult result = target.getReaderConnection(hosts); - - assertFalse(result.isConnected()); - assertNull(result.getConnection()); - assertNull(result.getHost()); - } - - @Test - public void testGetReader_connectionAttemptsTimeout() throws SQLException { - // connection attempts time out before they can succeed - // first connection attempt to return times out - // expected test result: failure to get reader - final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) - when(mockConnectionService.open(any(), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - try { - Thread.sleep(5000); - } catch (InterruptedException exception) { - // ignore - } - return mockConnection; - }); - - Dialect mockDialect = Mockito.mock(Dialect.class); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - when(mockPluginService.getDialect()).thenReturn(mockDialect); - - final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(60000, 1000, false); - final ReaderFailoverResult result = target.getReaderConnection(hosts); - - assertFalse(result.isConnected()); - assertNull(result.getConnection()); - assertNull(result.getHost()); - } - - @Test - public void testGetHostTuplesByPriority() { - final List originalHosts = defaultHosts; - originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); - originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); - originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); - - final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); - final List hostsByPriority = target.getHostsByPriority(originalHosts); - - int i = 0; - - // expecting active readers - while (i < hostsByPriority.size() - && hostsByPriority.get(i).getRole() == HostRole.READER - && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { - i++; - } - - // expecting a writer - while (i < hostsByPriority.size() - && hostsByPriority.get(i).getRole() == HostRole.WRITER) { - i++; - } - - // expecting down readers - while (i < hostsByPriority.size() - && hostsByPriority.get(i).getRole() == HostRole.READER - && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { - i++; - } - - assertEquals(hostsByPriority.size(), i); - } - - @Test - public void testGetReaderTuplesByPriority() { - final List originalHosts = defaultHosts; - originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); - originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); - originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); - - Dialect mockDialect = Mockito.mock(Dialect.class); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - when(mockPluginService.getDialect()).thenReturn(mockDialect); - - final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); - final List hostsByPriority = target.getReaderHostsByPriority(originalHosts); - - int i = 0; - - // expecting active readers - while (i < hostsByPriority.size() - && hostsByPriority.get(i).getRole() == HostRole.READER - && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { - i++; - } - - // expecting down readers - while (i < hostsByPriority.size() - && hostsByPriority.get(i).getRole() == HostRole.READER - && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { - i++; - } - - assertEquals(hostsByPriority.size(), i); - } - - @Test - public void testHostFailoverStrictReaderEnabled() { - - final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("writer").port(1234).role(HostRole.WRITER).build(); - final HostSpec reader = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader1").port(1234).role(HostRole.READER).build(); - final List hosts = Arrays.asList(writer, reader); - - Dialect mockDialect = Mockito.mock(Dialect.class); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - when(mockPluginService.getDialect()).thenReturn(mockDialect); - - final ClusterAwareReaderFailoverHandler target = - getSpyFailoverHandler(DEFAULT_FAILOVER_TIMEOUT, DEFAULT_READER_CONNECT_TIMEOUT, true); - - // The writer is included because the original writer has likely become a reader. - List expectedHostsByPriority = Arrays.asList(reader, writer); - - List hostsByPriority = target.getHostsByPriority(hosts); - assertEquals(expectedHostsByPriority, hostsByPriority); - - // Should pick the reader even if unavailable. The unavailable reader will be lower priority than the writer. - reader.setAvailability(HostAvailability.NOT_AVAILABLE); - expectedHostsByPriority = Arrays.asList(writer, reader); - - hostsByPriority = target.getHostsByPriority(hosts); - assertEquals(expectedHostsByPriority, hostsByPriority); - - // Writer node will only be picked if it is the only node in topology; - List expectedWriterHost = Collections.singletonList(writer); - - hostsByPriority = target.getHostsByPriority(Collections.singletonList(writer)); - assertEquals(expectedWriterHost, hostsByPriority); - } -} +// /* +// * 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.failover; +// +// import static org.junit.jupiter.api.Assertions.assertEquals; +// import static org.junit.jupiter.api.Assertions.assertFalse; +// import static org.junit.jupiter.api.Assertions.assertNull; +// import static org.junit.jupiter.api.Assertions.assertSame; +// import static org.junit.jupiter.api.Assertions.assertTrue; +// import static org.mockito.ArgumentMatchers.any; +// import static org.mockito.ArgumentMatchers.eq; +// import static org.mockito.Mockito.doReturn; +// import static org.mockito.Mockito.spy; +// import static org.mockito.Mockito.when; +// import static software.amazon.jdbc.plugin.failover.ClusterAwareReaderFailoverHandler.DEFAULT_FAILOVER_TIMEOUT; +// import static software.amazon.jdbc.plugin.failover.ClusterAwareReaderFailoverHandler.DEFAULT_READER_CONNECT_TIMEOUT; +// +// import java.sql.Connection; +// import java.sql.SQLException; +// import java.util.ArrayList; +// import java.util.Arrays; +// import java.util.Collections; +// import java.util.EnumSet; +// import java.util.List; +// import java.util.Map; +// import java.util.Properties; +// import java.util.Set; +// import java.util.concurrent.TimeUnit; +// import java.util.stream.Collectors; +// import org.junit.jupiter.api.AfterEach; +// import org.junit.jupiter.api.BeforeEach; +// import org.junit.jupiter.api.Test; +// import org.mockito.Mock; +// import org.mockito.Mockito; +// import org.mockito.MockitoAnnotations; +// import org.mockito.stubbing.Answer; +// import software.amazon.jdbc.ConnectionPluginManager; +// import software.amazon.jdbc.HostRole; +// import software.amazon.jdbc.HostSpec; +// import software.amazon.jdbc.HostSpecBuilder; +// import software.amazon.jdbc.PluginService; +// import software.amazon.jdbc.dialect.Dialect; +// import software.amazon.jdbc.hostavailability.HostAvailability; +// import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; +// import software.amazon.jdbc.util.FullServicesContainer; +// import software.amazon.jdbc.util.connection.ConnectionService; +// +// class ClusterAwareReaderFailoverHandlerTest { +// @Mock FullServicesContainer mockContainer; +// @Mock ConnectionService mockConnectionService; +// @Mock PluginService mockPluginService; +// @Mock ConnectionPluginManager mockPluginManager; +// @Mock Connection mockConnection; +// +// private AutoCloseable closeable; +// private final Properties properties = new Properties(); +// private final List defaultHosts = Arrays.asList( +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("writer").port(1234).role(HostRole.WRITER).build(), +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader1").port(1234).role(HostRole.READER).build(), +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader2").port(1234).role(HostRole.READER).build(), +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader3").port(1234).role(HostRole.READER).build(), +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader4").port(1234).role(HostRole.READER).build(), +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader5").port(1234).role(HostRole.READER).build() +// ); +// +// @BeforeEach +// void setUp() { +// closeable = MockitoAnnotations.openMocks(this); +// when(mockContainer.getConnectionPluginManager()).thenReturn(mockPluginManager); +// when(mockContainer.getPluginService()).thenReturn(mockPluginService); +// } +// +// @AfterEach +// void tearDown() throws Exception { +// closeable.close(); +// } +// +// @Test +// public void testFailover() throws SQLException { +// // original host list: [active writer, active reader, current connection (reader), active +// // reader, down reader, active reader] +// // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] +// // connection attempts are made in pairs using the above list +// // expected test result: successful connection for host at index 4 +// final List hosts = defaultHosts; +// final int currentHostIndex = 2; +// final int successHostIndex = 4; +// for (int i = 0; i < hosts.size(); i++) { +// if (i != successHostIndex) { +// final SQLException exception = new SQLException("exception", "08S01", null); +// when(mockConnectionService.open(hosts.get(i), properties)) +// .thenThrow(exception); +// when(mockPluginService.isNetworkException(exception, null)).thenReturn(true); +// } else { +// when(mockConnectionService.open(hosts.get(i), properties)).thenReturn(mockConnection); +// } +// } +// +// when(mockPluginService.getTargetDriverDialect()).thenReturn(null); +// +// hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); +// hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); +// +// final ReaderFailoverHandler target = getSpyFailoverHandler(); +// final ReaderFailoverResult result = target.failover(hosts, hosts.get(currentHostIndex)); +// +// assertTrue(result.isConnected()); +// assertSame(mockConnection, result.getConnection()); +// assertEquals(hosts.get(successHostIndex), result.getHost()); +// +// final HostSpec successHost = hosts.get(successHostIndex); +// final Map availabilityMap = target.getHostAvailabilityMap(); +// Set unavailableHosts = getHostsWithGivenAvailability(availabilityMap, HostAvailability.NOT_AVAILABLE); +// assertTrue(unavailableHosts.size() >= 4); +// assertEquals(HostAvailability.AVAILABLE, availabilityMap.get(successHost.getHost())); +// } +// +// private Set getHostsWithGivenAvailability( +// Map availabilityMap, HostAvailability availability) { +// return availabilityMap.entrySet().stream() +// .filter((entry) -> availability.equals(entry.getValue())) +// .map(Map.Entry::getKey) +// .collect(Collectors.toSet()); +// } +// +// @Test +// public void testFailover_timeout() throws SQLException { +// // original host list: [active writer, active reader, current connection (reader), active +// // reader, down reader, active reader] +// // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] +// // connection attempts are made in pairs using the above list +// // expected test result: failure to get reader since process is limited to 5s and each attempt +// // to connect takes 20s +// final List hosts = defaultHosts; +// final int currentHostIndex = 2; +// for (HostSpec host : hosts) { +// when(mockConnectionService.open(host, properties)) +// .thenAnswer((Answer) invocation -> { +// Thread.sleep(20000); +// return mockConnection; +// }); +// } +// +// hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); +// hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); +// +// final ReaderFailoverHandler target = getSpyFailoverHandler(5000, 30000, false); +// +// final long startTimeNano = System.nanoTime(); +// final ReaderFailoverResult result = target.failover(hosts, hosts.get(currentHostIndex)); +// final long durationNano = System.nanoTime() - startTimeNano; +// +// assertFalse(result.isConnected()); +// assertNull(result.getConnection()); +// assertNull(result.getHost()); +// +// // 5s is a max allowed failover timeout; add 1s for inaccurate measurements +// assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); +// } +// +// private ClusterAwareReaderFailoverHandler getSpyFailoverHandler() { +// ClusterAwareReaderFailoverHandler handler = +// spy(new ClusterAwareReaderFailoverHandler(mockContainer, mockConnectionService, properties)); +// doReturn(mockPluginService).when(handler).getNewPluginService(); +// return handler; +// } +// +// private ClusterAwareReaderFailoverHandler getSpyFailoverHandler( +// int maxFailoverTimeoutMs, int timeoutMs, boolean isStrictReaderRequired) { +// ClusterAwareReaderFailoverHandler handler = new ClusterAwareReaderFailoverHandler( +// mockContainer, mockConnectionService, properties, maxFailoverTimeoutMs, timeoutMs, isStrictReaderRequired); +// ClusterAwareReaderFailoverHandler spyHandler = spy(handler); +// doReturn(mockPluginService).when(spyHandler).getNewPluginService(); +// return spyHandler; +// } +// +// @Test +// public void testFailover_nullOrEmptyHostList() throws SQLException { +// final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); +// final HostSpec currentHost = +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("writer").port(1234).build(); +// +// ReaderFailoverResult result = target.failover(null, currentHost); +// assertFalse(result.isConnected()); +// assertNull(result.getConnection()); +// assertNull(result.getHost()); +// +// final List hosts = new ArrayList<>(); +// result = target.failover(hosts, currentHost); +// assertFalse(result.isConnected()); +// assertNull(result.getConnection()); +// assertNull(result.getHost()); +// } +// +// @Test +// public void testGetReader_connectionSuccess() throws SQLException { +// // even number of connection attempts +// // first connection attempt to return succeeds, second attempt cancelled +// // expected test result: successful connection for host at index 2 +// final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) +// final HostSpec slowHost = hosts.get(1); +// final HostSpec fastHost = hosts.get(2); +// when(mockConnectionService.open(slowHost, properties)) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(20000); +// return mockConnection; +// }); +// when(mockConnectionService.open(eq(fastHost), eq(properties))).thenReturn(mockConnection); +// +// Dialect mockDialect = Mockito.mock(Dialect.class); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// +// final ReaderFailoverHandler target = getSpyFailoverHandler(); +// final ReaderFailoverResult result = target.getReaderConnection(hosts); +// +// assertTrue(result.isConnected()); +// assertSame(mockConnection, result.getConnection()); +// assertEquals(hosts.get(2), result.getHost()); +// +// Map availabilityMap = target.getHostAvailabilityMap(); +// assertTrue(getHostsWithGivenAvailability(availabilityMap, HostAvailability.NOT_AVAILABLE).isEmpty()); +// assertEquals(HostAvailability.AVAILABLE, availabilityMap.get(fastHost.getHost())); +// } +// +// @Test +// public void testGetReader_connectionFailure() throws SQLException { +// // odd number of connection attempts +// // first connection attempt to return fails +// // expected test result: failure to get reader +// final List hosts = defaultHosts.subList(0, 4); // 3 connection attempts (writer not attempted) +// when(mockConnectionService.open(any(), eq(properties))).thenThrow(new SQLException("exception", "08S01", null)); +// +// Dialect mockDialect = Mockito.mock(Dialect.class); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// +// final ReaderFailoverHandler target = getSpyFailoverHandler(); +// final ReaderFailoverResult result = target.getReaderConnection(hosts); +// +// assertFalse(result.isConnected()); +// assertNull(result.getConnection()); +// assertNull(result.getHost()); +// } +// +// @Test +// public void testGetReader_connectionAttemptsTimeout() throws SQLException { +// // connection attempts time out before they can succeed +// // first connection attempt to return times out +// // expected test result: failure to get reader +// final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) +// when(mockConnectionService.open(any(), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// try { +// Thread.sleep(5000); +// } catch (InterruptedException exception) { +// // ignore +// } +// return mockConnection; +// }); +// +// Dialect mockDialect = Mockito.mock(Dialect.class); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// +// final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(60000, 1000, false); +// final ReaderFailoverResult result = target.getReaderConnection(hosts); +// +// assertFalse(result.isConnected()); +// assertNull(result.getConnection()); +// assertNull(result.getHost()); +// } +// +// @Test +// public void testGetHostTuplesByPriority() { +// final List originalHosts = defaultHosts; +// originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); +// originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); +// originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); +// +// final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); +// final List hostsByPriority = target.getHostsByPriority(originalHosts); +// +// int i = 0; +// +// // expecting active readers +// while (i < hostsByPriority.size() +// && hostsByPriority.get(i).getRole() == HostRole.READER +// && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { +// i++; +// } +// +// // expecting a writer +// while (i < hostsByPriority.size() +// && hostsByPriority.get(i).getRole() == HostRole.WRITER) { +// i++; +// } +// +// // expecting down readers +// while (i < hostsByPriority.size() +// && hostsByPriority.get(i).getRole() == HostRole.READER +// && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { +// i++; +// } +// +// assertEquals(hostsByPriority.size(), i); +// } +// +// @Test +// public void testGetReaderTuplesByPriority() { +// final List originalHosts = defaultHosts; +// originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); +// originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); +// originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); +// +// Dialect mockDialect = Mockito.mock(Dialect.class); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// +// final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); +// final List hostsByPriority = target.getReaderHostsByPriority(originalHosts); +// +// int i = 0; +// +// // expecting active readers +// while (i < hostsByPriority.size() +// && hostsByPriority.get(i).getRole() == HostRole.READER +// && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { +// i++; +// } +// +// // expecting down readers +// while (i < hostsByPriority.size() +// && hostsByPriority.get(i).getRole() == HostRole.READER +// && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { +// i++; +// } +// +// assertEquals(hostsByPriority.size(), i); +// } +// +// @Test +// public void testHostFailoverStrictReaderEnabled() { +// +// final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("writer").port(1234).role(HostRole.WRITER).build(); +// final HostSpec reader = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader1").port(1234).role(HostRole.READER).build(); +// final List hosts = Arrays.asList(writer, reader); +// +// Dialect mockDialect = Mockito.mock(Dialect.class); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// +// final ClusterAwareReaderFailoverHandler target = +// getSpyFailoverHandler(DEFAULT_FAILOVER_TIMEOUT, DEFAULT_READER_CONNECT_TIMEOUT, true); +// +// // The writer is included because the original writer has likely become a reader. +// List expectedHostsByPriority = Arrays.asList(reader, writer); +// +// List hostsByPriority = target.getHostsByPriority(hosts); +// assertEquals(expectedHostsByPriority, hostsByPriority); +// +// // Should pick the reader even if unavailable. The unavailable reader will be lower priority than the writer. +// reader.setAvailability(HostAvailability.NOT_AVAILABLE); +// expectedHostsByPriority = Arrays.asList(writer, reader); +// +// hostsByPriority = target.getHostsByPriority(hosts); +// assertEquals(expectedHostsByPriority, hostsByPriority); +// +// // Writer node will only be picked if it is the only node in topology; +// List expectedWriterHost = Collections.singletonList(writer); +// +// hostsByPriority = target.getHostsByPriority(Collections.singletonList(writer)); +// assertEquals(expectedWriterHost, hostsByPriority); +// } +// } diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java index 1ad394010..097b68d67 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java @@ -1,373 +1,373 @@ -/* - * 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.failover; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertSame; -import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.ArgumentMatchers.refEq; -import static org.mockito.Mockito.atLeastOnce; -import static org.mockito.Mockito.doReturn; -import static org.mockito.Mockito.spy; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -import java.sql.Connection; -import java.sql.SQLException; -import java.util.Arrays; -import java.util.EnumSet; -import java.util.List; -import java.util.Properties; -import java.util.concurrent.TimeUnit; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.mockito.ArgumentMatchers; -import org.mockito.Mock; -import org.mockito.MockitoAnnotations; -import org.mockito.stubbing.Answer; -import software.amazon.jdbc.HostSpec; -import software.amazon.jdbc.HostSpecBuilder; -import software.amazon.jdbc.PluginService; -import software.amazon.jdbc.dialect.Dialect; -import software.amazon.jdbc.hostavailability.HostAvailability; -import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; -import software.amazon.jdbc.util.FullServicesContainer; -import software.amazon.jdbc.util.connection.ConnectionService; - -class ClusterAwareWriterFailoverHandlerTest { - @Mock FullServicesContainer mockContainer; - @Mock ConnectionService mockConnectionService; - @Mock PluginService mockPluginService; - @Mock Connection mockConnection; - @Mock ReaderFailoverHandler mockReaderFailoverHandler; - @Mock Connection mockWriterConnection; - @Mock Connection mockNewWriterConnection; - @Mock Connection mockReaderAConnection; - @Mock Connection mockReaderBConnection; - @Mock Dialect mockDialect; - - private AutoCloseable closeable; - private final Properties properties = new Properties(); - private final HostSpec newWriterHost = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("new-writer-host").build(); - private final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("writer-host").build(); - private final HostSpec readerA = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader-a-host").build(); - private final HostSpec readerB = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader-b-host").build(); - private final List topology = Arrays.asList(writer, readerA, readerB); - private final List newTopology = Arrays.asList(newWriterHost, readerA, readerB); - - @BeforeEach - void setUp() { - closeable = MockitoAnnotations.openMocks(this); - when(mockContainer.getPluginService()).thenReturn(mockPluginService); - writer.addAlias("writer-host"); - newWriterHost.addAlias("new-writer-host"); - readerA.addAlias("reader-a-host"); - readerB.addAlias("reader-b-host"); - } - - @AfterEach - void tearDown() throws Exception { - closeable.close(); - } - - @Test - public void testReconnectToWriter_taskBReaderException() throws SQLException { - when(mockConnectionService.open(refEq(writer), eq(properties))).thenReturn(mockConnection); - when(mockConnectionService.open(refEq(readerA), eq(properties))).thenThrow(SQLException.class); - when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); - - when(mockPluginService.getAllHosts()).thenReturn(topology); - - when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())).thenThrow(SQLException.class); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); - final WriterFailoverResult result = target.failover(topology); - - assertTrue(result.isConnected()); - assertFalse(result.isNewHost()); - assertSame(result.getNewConnection(), mockConnection); - - assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); - } - - private ClusterAwareWriterFailoverHandler getSpyFailoverHandler( - final int failoverTimeoutMs, - final int readTopologyIntervalMs, - final int reconnectWriterIntervalMs) { - ClusterAwareWriterFailoverHandler handler = new ClusterAwareWriterFailoverHandler( - mockContainer, - mockConnectionService, - mockReaderFailoverHandler, - properties, - failoverTimeoutMs, - readTopologyIntervalMs, - reconnectWriterIntervalMs); - - ClusterAwareWriterFailoverHandler spyHandler = spy(handler); - doReturn(mockPluginService).when(spyHandler).getNewPluginService(); - return spyHandler; - } - - /** - * Verify that writer failover handler can re-connect to a current writer node. - * - *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. - * TaskA: successfully re-connect to initial writer; return new connection. - * TaskB: successfully connect to readerA and then new writer, but it takes more time than taskA. - * Expected test result: new connection by taskA. - */ - @Test - public void testReconnectToWriter_SlowReaderA() throws SQLException { - when(mockConnectionService.open(refEq(writer), eq(properties))).thenReturn(mockWriterConnection); - when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); - when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); - when(mockPluginService.getAllHosts()).thenReturn(topology).thenReturn(newTopology); - - when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(5000); - return new ReaderFailoverResult(mockReaderAConnection, readerA, true); - }); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); - final WriterFailoverResult result = target.failover(topology); - - assertTrue(result.isConnected()); - assertFalse(result.isNewHost()); - assertSame(result.getNewConnection(), mockWriterConnection); - assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); - } - - /** - * Verify that writer failover handler can re-connect to a current writer node. - * - *

Topology: no changes. - * TaskA: successfully re-connect to writer; return new connection. - * TaskB: successfully connect to readerA and retrieve topology, but latest writer is not new (defer to taskA). - * Expected test result: new connection by taskA. - */ - @Test - public void testReconnectToWriter_taskBDefers() throws SQLException { - when(mockConnectionService.open(refEq(writer), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(5000); - return mockWriterConnection; - }); - when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); - - when(mockPluginService.getAllHosts()).thenReturn(topology); - - when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) - .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 2000, 2000); - final WriterFailoverResult result = target.failover(topology); - - assertTrue(result.isConnected()); - assertFalse(result.isNewHost()); - assertSame(result.getNewConnection(), mockWriterConnection); - assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); - } - - /** - * Verify that writer failover handler can re-connect to a new writer node. - * - *

Topology: changes to [new-writer, reader-A, reader-B] for taskB, taskA sees no changes. - * taskA: successfully re-connect to writer; return connection to initial writer, but it takes more - * time than taskB. - * TaskB: successfully connect to readerA and then to new-writer. - * Expected test result: new connection to writer by taskB. - */ - @Test - public void testConnectToReaderA_SlowWriter() throws SQLException { - when(mockConnectionService.open(refEq(writer), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(5000); - return mockWriterConnection; - }); - when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); - when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); - when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); - - when(mockPluginService.getAllHosts()).thenReturn(newTopology); - - when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) - .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); - final WriterFailoverResult result = target.failover(topology); - - assertTrue(result.isConnected()); - assertTrue(result.isNewHost()); - assertSame(result.getNewConnection(), mockNewWriterConnection); - assertEquals(3, result.getTopology().size()); - assertEquals("new-writer-host", result.getTopology().get(0).getHost()); - assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); - } - - /** - * Verify that writer failover handler can re-connect to a new writer node. - * - *

Topology: changes to [new-writer, initial-writer, reader-A, reader-B]. - * TaskA: successfully reconnect, but initial-writer is now a reader (defer to taskB). - * TaskB: successfully connect to readerA and then to new-writer. - * Expected test result: new connection to writer by taskB. - */ - @Test - public void testConnectToReaderA_taskADefers() throws SQLException { - when(mockConnectionService.open(writer, properties)).thenReturn(mockConnection); - when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); - when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); - when(mockConnectionService.open(refEq(newWriterHost), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(5000); - return mockNewWriterConnection; - }); - - final List newTopology = Arrays.asList(newWriterHost, writer, readerA, readerB); - when(mockPluginService.getAllHosts()).thenReturn(newTopology); - - when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) - .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); - final WriterFailoverResult result = target.failover(topology); - - assertTrue(result.isConnected()); - assertTrue(result.isNewHost()); - assertSame(result.getNewConnection(), mockNewWriterConnection); - assertEquals(4, result.getTopology().size()); - assertEquals("new-writer-host", result.getTopology().get(0).getHost()); - - verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); - assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); - } - - /** - * Verify that writer failover handler fails to re-connect to any writer node. - * - *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. - * TaskA: fail to re-connect to writer due to failover timeout. - * TaskB: successfully connect to readerA and then fail to connect to writer due to failover timeout. - * Expected test result: no connection. - */ - @Test - public void testFailedToConnect_failoverTimeout() throws SQLException { - when(mockConnectionService.open(refEq(writer), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(30000); - return mockWriterConnection; - }); - when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); - when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); - when(mockConnectionService.open(refEq(newWriterHost), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(30000); - return mockNewWriterConnection; - }); - when(mockPluginService.getAllHosts()).thenReturn(newTopology); - - when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) - .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); - - final long startTimeNano = System.nanoTime(); - final WriterFailoverResult result = target.failover(topology); - final long durationNano = System.nanoTime() - startTimeNano; - - assertFalse(result.isConnected()); - assertFalse(result.isNewHost()); - - verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); - - // 5s is a max allowed failover timeout; add 1s for inaccurate measurements - assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); - } - - /** - * Verify that writer failover handler fails to re-connect to any writer node. - * - *

Topology: changes to [new-writer, reader-A, reader-B] for taskB. - * TaskA: fail to re-connect to writer due to exception. - * TaskB: successfully connect to readerA and then fail to connect to writer due to exception. - * Expected test result: no connection. - */ - @Test - public void testFailedToConnect_taskAException_taskBWriterException() throws SQLException { - final SQLException exception = new SQLException("exception", "08S01", null); - when(mockConnectionService.open(refEq(writer), eq(properties))).thenThrow(exception); - when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); - when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); - when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenThrow(exception); - when(mockPluginService.isNetworkException(eq(exception), any())).thenReturn(true); - - when(mockPluginService.getAllHosts()).thenReturn(newTopology); - - when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) - .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); - final WriterFailoverResult result = target.failover(topology); - - assertFalse(result.isConnected()); - assertFalse(result.isNewHost()); - - assertEquals(HostAvailability.NOT_AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); - } -} +// /* +// * 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.failover; +// +// import static org.junit.jupiter.api.Assertions.assertEquals; +// import static org.junit.jupiter.api.Assertions.assertFalse; +// import static org.junit.jupiter.api.Assertions.assertSame; +// import static org.junit.jupiter.api.Assertions.assertTrue; +// import static org.mockito.ArgumentMatchers.any; +// import static org.mockito.ArgumentMatchers.eq; +// import static org.mockito.ArgumentMatchers.refEq; +// import static org.mockito.Mockito.atLeastOnce; +// import static org.mockito.Mockito.doReturn; +// import static org.mockito.Mockito.spy; +// import static org.mockito.Mockito.verify; +// import static org.mockito.Mockito.when; +// +// import java.sql.Connection; +// import java.sql.SQLException; +// import java.util.Arrays; +// import java.util.EnumSet; +// import java.util.List; +// import java.util.Properties; +// import java.util.concurrent.TimeUnit; +// import org.junit.jupiter.api.AfterEach; +// import org.junit.jupiter.api.BeforeEach; +// import org.junit.jupiter.api.Test; +// import org.mockito.ArgumentMatchers; +// import org.mockito.Mock; +// import org.mockito.MockitoAnnotations; +// import org.mockito.stubbing.Answer; +// import software.amazon.jdbc.HostSpec; +// import software.amazon.jdbc.HostSpecBuilder; +// import software.amazon.jdbc.PluginService; +// import software.amazon.jdbc.dialect.Dialect; +// import software.amazon.jdbc.hostavailability.HostAvailability; +// import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; +// import software.amazon.jdbc.util.FullServicesContainer; +// import software.amazon.jdbc.util.connection.ConnectionService; +// +// class ClusterAwareWriterFailoverHandlerTest { +// @Mock FullServicesContainer mockContainer; +// @Mock ConnectionService mockConnectionService; +// @Mock PluginService mockPluginService; +// @Mock Connection mockConnection; +// @Mock ReaderFailoverHandler mockReaderFailoverHandler; +// @Mock Connection mockWriterConnection; +// @Mock Connection mockNewWriterConnection; +// @Mock Connection mockReaderAConnection; +// @Mock Connection mockReaderBConnection; +// @Mock Dialect mockDialect; +// +// private AutoCloseable closeable; +// private final Properties properties = new Properties(); +// private final HostSpec newWriterHost = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("new-writer-host").build(); +// private final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("writer-host").build(); +// private final HostSpec readerA = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader-a-host").build(); +// private final HostSpec readerB = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader-b-host").build(); +// private final List topology = Arrays.asList(writer, readerA, readerB); +// private final List newTopology = Arrays.asList(newWriterHost, readerA, readerB); +// +// @BeforeEach +// void setUp() { +// closeable = MockitoAnnotations.openMocks(this); +// when(mockContainer.getPluginService()).thenReturn(mockPluginService); +// writer.addAlias("writer-host"); +// newWriterHost.addAlias("new-writer-host"); +// readerA.addAlias("reader-a-host"); +// readerB.addAlias("reader-b-host"); +// } +// +// @AfterEach +// void tearDown() throws Exception { +// closeable.close(); +// } +// +// @Test +// public void testReconnectToWriter_taskBReaderException() throws SQLException { +// when(mockConnectionService.open(refEq(writer), eq(properties))).thenReturn(mockConnection); +// when(mockConnectionService.open(refEq(readerA), eq(properties))).thenThrow(SQLException.class); +// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); +// +// when(mockPluginService.getAllHosts()).thenReturn(topology); +// +// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())).thenThrow(SQLException.class); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertTrue(result.isConnected()); +// assertFalse(result.isNewHost()); +// assertSame(result.getNewConnection(), mockConnection); +// +// assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); +// } +// +// private ClusterAwareWriterFailoverHandler getSpyFailoverHandler( +// final int failoverTimeoutMs, +// final int readTopologyIntervalMs, +// final int reconnectWriterIntervalMs) { +// ClusterAwareWriterFailoverHandler handler = new ClusterAwareWriterFailoverHandler( +// mockContainer, +// mockConnectionService, +// mockReaderFailoverHandler, +// properties, +// failoverTimeoutMs, +// readTopologyIntervalMs, +// reconnectWriterIntervalMs); +// +// ClusterAwareWriterFailoverHandler spyHandler = spy(handler); +// doReturn(mockPluginService).when(spyHandler).getNewPluginService(); +// return spyHandler; +// } +// +// /** +// * Verify that writer failover handler can re-connect to a current writer node. +// * +// *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. +// * TaskA: successfully re-connect to initial writer; return new connection. +// * TaskB: successfully connect to readerA and then new writer, but it takes more time than taskA. +// * Expected test result: new connection by taskA. +// */ +// @Test +// public void testReconnectToWriter_SlowReaderA() throws SQLException { +// when(mockConnectionService.open(refEq(writer), eq(properties))).thenReturn(mockWriterConnection); +// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); +// when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); +// when(mockPluginService.getAllHosts()).thenReturn(topology).thenReturn(newTopology); +// +// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(5000); +// return new ReaderFailoverResult(mockReaderAConnection, readerA, true); +// }); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertTrue(result.isConnected()); +// assertFalse(result.isNewHost()); +// assertSame(result.getNewConnection(), mockWriterConnection); +// assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); +// } +// +// /** +// * Verify that writer failover handler can re-connect to a current writer node. +// * +// *

Topology: no changes. +// * TaskA: successfully re-connect to writer; return new connection. +// * TaskB: successfully connect to readerA and retrieve topology, but latest writer is not new (defer to taskA). +// * Expected test result: new connection by taskA. +// */ +// @Test +// public void testReconnectToWriter_taskBDefers() throws SQLException { +// when(mockConnectionService.open(refEq(writer), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(5000); +// return mockWriterConnection; +// }); +// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); +// +// when(mockPluginService.getAllHosts()).thenReturn(topology); +// +// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) +// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 2000, 2000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertTrue(result.isConnected()); +// assertFalse(result.isNewHost()); +// assertSame(result.getNewConnection(), mockWriterConnection); +// assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); +// } +// +// /** +// * Verify that writer failover handler can re-connect to a new writer node. +// * +// *

Topology: changes to [new-writer, reader-A, reader-B] for taskB, taskA sees no changes. +// * taskA: successfully re-connect to writer; return connection to initial writer, but it takes more +// * time than taskB. +// * TaskB: successfully connect to readerA and then to new-writer. +// * Expected test result: new connection to writer by taskB. +// */ +// @Test +// public void testConnectToReaderA_SlowWriter() throws SQLException { +// when(mockConnectionService.open(refEq(writer), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(5000); +// return mockWriterConnection; +// }); +// when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); +// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); +// when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); +// +// when(mockPluginService.getAllHosts()).thenReturn(newTopology); +// +// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) +// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertTrue(result.isConnected()); +// assertTrue(result.isNewHost()); +// assertSame(result.getNewConnection(), mockNewWriterConnection); +// assertEquals(3, result.getTopology().size()); +// assertEquals("new-writer-host", result.getTopology().get(0).getHost()); +// assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); +// } +// +// /** +// * Verify that writer failover handler can re-connect to a new writer node. +// * +// *

Topology: changes to [new-writer, initial-writer, reader-A, reader-B]. +// * TaskA: successfully reconnect, but initial-writer is now a reader (defer to taskB). +// * TaskB: successfully connect to readerA and then to new-writer. +// * Expected test result: new connection to writer by taskB. +// */ +// @Test +// public void testConnectToReaderA_taskADefers() throws SQLException { +// when(mockConnectionService.open(writer, properties)).thenReturn(mockConnection); +// when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); +// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); +// when(mockConnectionService.open(refEq(newWriterHost), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(5000); +// return mockNewWriterConnection; +// }); +// +// final List newTopology = Arrays.asList(newWriterHost, writer, readerA, readerB); +// when(mockPluginService.getAllHosts()).thenReturn(newTopology); +// +// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) +// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertTrue(result.isConnected()); +// assertTrue(result.isNewHost()); +// assertSame(result.getNewConnection(), mockNewWriterConnection); +// assertEquals(4, result.getTopology().size()); +// assertEquals("new-writer-host", result.getTopology().get(0).getHost()); +// +// verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); +// assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); +// } +// +// /** +// * Verify that writer failover handler fails to re-connect to any writer node. +// * +// *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. +// * TaskA: fail to re-connect to writer due to failover timeout. +// * TaskB: successfully connect to readerA and then fail to connect to writer due to failover timeout. +// * Expected test result: no connection. +// */ +// @Test +// public void testFailedToConnect_failoverTimeout() throws SQLException { +// when(mockConnectionService.open(refEq(writer), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(30000); +// return mockWriterConnection; +// }); +// when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); +// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); +// when(mockConnectionService.open(refEq(newWriterHost), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(30000); +// return mockNewWriterConnection; +// }); +// when(mockPluginService.getAllHosts()).thenReturn(newTopology); +// +// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) +// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); +// +// final long startTimeNano = System.nanoTime(); +// final WriterFailoverResult result = target.failover(topology); +// final long durationNano = System.nanoTime() - startTimeNano; +// +// assertFalse(result.isConnected()); +// assertFalse(result.isNewHost()); +// +// verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); +// +// // 5s is a max allowed failover timeout; add 1s for inaccurate measurements +// assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); +// } +// +// /** +// * Verify that writer failover handler fails to re-connect to any writer node. +// * +// *

Topology: changes to [new-writer, reader-A, reader-B] for taskB. +// * TaskA: fail to re-connect to writer due to exception. +// * TaskB: successfully connect to readerA and then fail to connect to writer due to exception. +// * Expected test result: no connection. +// */ +// @Test +// public void testFailedToConnect_taskAException_taskBWriterException() throws SQLException { +// final SQLException exception = new SQLException("exception", "08S01", null); +// when(mockConnectionService.open(refEq(writer), eq(properties))).thenThrow(exception); +// when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); +// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); +// when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenThrow(exception); +// when(mockPluginService.isNetworkException(eq(exception), any())).thenReturn(true); +// +// when(mockPluginService.getAllHosts()).thenReturn(newTopology); +// +// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) +// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertFalse(result.isConnected()); +// assertFalse(result.isNewHost()); +// +// assertEquals(HostAvailability.NOT_AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); +// } +// } From f7b8db9a1207c8ee25426db10ddce43b8f532cf9 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Tue, 2 Sep 2025 17:29:05 -0700 Subject: [PATCH 29/42] Fix failing integration tests --- .../amazon/jdbc/PartialPluginService.java | 20 ++++++----------- .../ClusterAwareReaderFailoverHandler.java | 10 +++++++-- .../ClusterAwareWriterFailoverHandler.java | 10 +++++++-- .../connection/ConnectionServiceImpl.java | 22 +++++++++++++++---- ...ClusterAwareReaderFailoverHandlerTest.java | 2 +- 5 files changed, 42 insertions(+), 22 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java b/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java index 9b36ea6d0..34391c3a5 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java +++ b/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java @@ -83,15 +83,13 @@ public class PartialPluginService implements PluginService, CanReleaseResources, public PartialPluginService( @NonNull final FullServicesContainer servicesContainer, - @NonNull final ConnectionProvider defaultConnectionProvider, @NonNull final Properties props, @NonNull final String originalUrl, @NonNull final String targetDriverProtocol, @NonNull final TargetDriverDialect targetDriverDialect, - @NonNull final Dialect dbDialect) throws SQLException { + @NonNull final Dialect dbDialect) { this( servicesContainer, - defaultConnectionProvider, new ExceptionManager(), props, originalUrl, @@ -103,15 +101,19 @@ public PartialPluginService( public PartialPluginService( @NonNull final FullServicesContainer servicesContainer, - @NonNull final ConnectionProvider defaultConnectionProvider, @NonNull final ExceptionManager exceptionManager, @NonNull final Properties props, @NonNull final String originalUrl, @NonNull final String targetDriverProtocol, @NonNull final TargetDriverDialect targetDriverDialect, @NonNull final Dialect dbDialect, - @Nullable final ConfigurationProfile configurationProfile) throws SQLException { + @Nullable final ConfigurationProfile configurationProfile) { this.servicesContainer = servicesContainer; + this.servicesContainer.setHostListProviderService(this); + this.servicesContainer.setPluginService(this); + this.servicesContainer.setPluginManagerService(this); + + this.pluginManager = servicesContainer.getConnectionPluginManager(); this.props = props; this.originalUrl = originalUrl; this.driverProtocol = targetDriverProtocol; @@ -120,10 +122,6 @@ public PartialPluginService( this.configurationProfile = configurationProfile; this.exceptionManager = exceptionManager; - this.pluginManager = new ConnectionPluginManager( - defaultConnectionProvider, null, null, servicesContainer.getTelemetryFactory()); - this.pluginManager.init(this.servicesContainer, this.props, this, this.configurationProfile); - this.servicesContainer.setConnectionPluginManager(pluginManager); this.connectionProviderManager = new ConnectionProviderManager( this.pluginManager.getDefaultConnProvider(), this.pluginManager.getEffectiveConnProvider()); @@ -132,10 +130,6 @@ public PartialPluginService( ? this.configurationProfile.getExceptionHandler() : null; - servicesContainer.setHostListProviderService(this); - servicesContainer.setPluginService(this); - servicesContainer.setPluginManagerService(this); - HostListProviderSupplier supplier = this.dbDialect.getHostListProvider(); this.hostListProvider = supplier.getProvider(this.props, this.originalUrl, this.servicesContainer); } diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java index 4b2373c19..d0512cf55 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java @@ -34,6 +34,7 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.logging.Logger; +import software.amazon.jdbc.ConnectionPluginManager; import software.amazon.jdbc.HostRole; import software.amazon.jdbc.HostSpec; import software.amazon.jdbc.PartialPluginService; @@ -404,15 +405,20 @@ protected PluginService getNewPluginService() throws SQLException { this.servicesContainer.getTelemetryFactory() ); - return new PartialPluginService( + ConnectionPluginManager pluginManager = new ConnectionPluginManager( + this.pluginService.getDefaultConnectionProvider(), null, null, servicesContainer.getTelemetryFactory()); + newServicesContainer.setConnectionPluginManager(pluginManager); + PartialPluginService pluginService = new PartialPluginService( newServicesContainer, - this.pluginService.getDefaultConnectionProvider(), this.props, this.pluginService.getOriginalUrl(), this.pluginService.getDriverProtocol(), this.pluginService.getTargetDriverDialect(), this.pluginService.getDialect() ); + + pluginManager.init(newServicesContainer, this.props, pluginService, null); + return pluginService; } private static class ConnectionAttemptTask implements Callable { diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java index b93a2292a..fcdd52a7c 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java @@ -32,6 +32,7 @@ import java.util.concurrent.TimeUnit; import java.util.logging.Level; import java.util.logging.Logger; +import software.amazon.jdbc.ConnectionPluginManager; import software.amazon.jdbc.HostSpec; import software.amazon.jdbc.PartialPluginService; import software.amazon.jdbc.PluginService; @@ -185,15 +186,20 @@ protected PluginService getNewPluginService() throws SQLException { this.servicesContainer.getTelemetryFactory() ); - return new PartialPluginService( + ConnectionPluginManager pluginManager = new ConnectionPluginManager( + this.pluginService.getDefaultConnectionProvider(), null, null, servicesContainer.getTelemetryFactory()); + newServicesContainer.setConnectionPluginManager(pluginManager); + PartialPluginService pluginService = new PartialPluginService( newServicesContainer, - this.pluginService.getDefaultConnectionProvider(), this.initialConnectionProps, this.pluginService.getOriginalUrl(), this.pluginService.getDriverProtocol(), this.pluginService.getTargetDriverDialect(), this.pluginService.getDialect() ); + + pluginManager.init(newServicesContainer, this.initialConnectionProps, pluginService, null); + return pluginService; } private WriterFailoverResult getNextResult( diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java index 8b2c23e97..0e506a112 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java @@ -19,6 +19,7 @@ import java.sql.Connection; import java.sql.SQLException; import java.util.Properties; +import software.amazon.jdbc.ConnectionPluginManager; import software.amazon.jdbc.ConnectionProvider; import software.amazon.jdbc.HostSpec; import software.amazon.jdbc.PartialPluginService; @@ -32,6 +33,8 @@ import software.amazon.jdbc.util.telemetry.TelemetryFactory; public class ConnectionServiceImpl implements ConnectionService { + protected final String targetDriverProtocol; + protected final ConnectionPluginManager pluginManager; protected final PluginService pluginService; public ConnectionServiceImpl( @@ -44,22 +47,33 @@ public ConnectionServiceImpl( TargetDriverDialect driverDialect, Dialect dbDialect, Properties props) throws SQLException { + this.targetDriverProtocol = targetDriverProtocol; + FullServicesContainer servicesContainer = new FullServicesContainerImpl(storageService, monitorService, telemetryFactory); - this.pluginService = new PartialPluginService( - servicesContainer, + this.pluginManager = new ConnectionPluginManager( connectionProvider, + null, + null, + telemetryFactory); + servicesContainer.setConnectionPluginManager(this.pluginManager); + + PartialPluginService partialPluginService = new PartialPluginService( + servicesContainer, props, originalUrl, - targetDriverProtocol, + this.targetDriverProtocol, driverDialect, dbDialect ); + + this.pluginService = partialPluginService; + this.pluginManager.init(servicesContainer, props, partialPluginService, null); } @Override public Connection open(HostSpec hostSpec, Properties props) throws SQLException { - return this.pluginService.forceConnect(hostSpec, props); + return this.pluginManager.forceConnect(this.targetDriverProtocol, hostSpec, props, true, null); } @Override diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java index 96d33f4e5..89e838acb 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java @@ -364,7 +364,7 @@ public void testGetReaderTuplesByPriority() throws SQLException { } @Test - public void testHostFailoverStrictReaderEnabled() throws SQLException{ + public void testHostFailoverStrictReaderEnabled() throws SQLException { final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) .host("writer").port(1234).role(HostRole.WRITER).build(); final HostSpec reader = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) From 104d3039eea874d9a9688c1d799fde723c87985d Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Wed, 3 Sep 2025 13:19:04 -0700 Subject: [PATCH 30/42] Fix dead markdown link --- docs/GettingStarted.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/GettingStarted.md b/docs/GettingStarted.md index 08bfc3c23..90729e3ad 100644 --- a/docs/GettingStarted.md +++ b/docs/GettingStarted.md @@ -25,7 +25,7 @@ dependencies { ### Direct Download and Installation -You can use pre-compiled packages that can be downloaded directly from [GitHub Releases](https://github.com/aws/aws-advanced-jdbc-wrapper/releases) or [Maven Central](https://search.maven.org/search?q=g:software.amazon.jdbc) to install the AWS JDBC Driver. After downloading the AWS JDBC Driver, install it by including the .jar file in the application's CLASSPATH. +You can use pre-compiled packages that can be downloaded directly from [GitHub Releases](https://github.com/aws/aws-advanced-jdbc-wrapper/releases) or [Maven Central](https://central.sonatype.com/artifact/software.amazon.jdbc/aws-advanced-jdbc-wrapper) to install the AWS JDBC Driver. After downloading the AWS JDBC Driver, install it by including the .jar file in the application's CLASSPATH. For example, the following command uses wget to download the wrapper: From 04ac3c89adc333a72360551823c5ba588c4a3bcd Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Wed, 3 Sep 2025 14:53:23 -0700 Subject: [PATCH 31/42] wip --- .../jdbc/ConnectionPluginChainBuilder.java | 2 +- .../amazon/jdbc/ConnectionPluginManager.java | 3 +- .../ClusterTopologyMonitorImpl.java | 10 +- .../customendpoint/CustomEndpointPlugin.java | 6 +- .../jdbc/plugin/efm2/HostMonitorImpl.java | 18 +- .../plugin/efm2/HostMonitorServiceImpl.java | 5 +- .../ClusterAwareReaderFailoverHandler.java | 28 +- .../ClusterAwareWriterFailoverHandler.java | 29 +- .../failover/FailoverConnectionPlugin.java | 18 - .../jdbc/wrapper/ConnectionWrapper.java | 5 +- ...ClusterAwareReaderFailoverHandlerTest.java | 800 ++++++++-------- ...ClusterAwareWriterFailoverHandlerTest.java | 746 +++++++-------- .../FailoverConnectionPluginTest.java | 894 +++++++++--------- 13 files changed, 1249 insertions(+), 1315 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginChainBuilder.java b/wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginChainBuilder.java index 6c8475a26..943e6397e 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginChainBuilder.java +++ b/wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginChainBuilder.java @@ -148,7 +148,7 @@ public List getPlugins( final ConnectionProvider effectiveConnProvider, final PluginManagerService pluginManagerService, final Properties props, - @Nullable ConfigurationProfile configurationProfile) { + @Nullable ConfigurationProfile configurationProfile) throws SQLException { List plugins; List pluginFactories; diff --git a/wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginManager.java b/wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginManager.java index 6272072af..4a4c7b20e 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginManager.java +++ b/wrapper/src/main/java/software/amazon/jdbc/ConnectionPluginManager.java @@ -29,7 +29,6 @@ import java.util.logging.Logger; import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; -import org.jetbrains.annotations.NotNull; import software.amazon.jdbc.cleanup.CanReleaseResources; import software.amazon.jdbc.plugin.AuroraConnectionTrackerPlugin; import software.amazon.jdbc.plugin.AuroraInitialConnectionStrategyPlugin; @@ -180,7 +179,7 @@ public void init( final FullServicesContainer servicesContainer, final Properties props, final PluginManagerService pluginManagerService, - @Nullable ConfigurationProfile configurationProfile) { + @Nullable ConfigurationProfile configurationProfile) throws SQLException { this.props = props; this.servicesContainer = servicesContainer; diff --git a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/ClusterTopologyMonitorImpl.java b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/ClusterTopologyMonitorImpl.java index e61cb6a60..1f3665f6d 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/ClusterTopologyMonitorImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/ClusterTopologyMonitorImpl.java @@ -42,7 +42,6 @@ import java.util.stream.Collectors; import org.checkerframework.checker.nullness.qual.NonNull; import org.checkerframework.checker.nullness.qual.Nullable; -import software.amazon.jdbc.HostListProviderService; import software.amazon.jdbc.HostRole; import software.amazon.jdbc.HostSpec; import software.amazon.jdbc.PropertyDefinition; @@ -57,10 +56,7 @@ import software.amazon.jdbc.util.StringUtils; import software.amazon.jdbc.util.SynchronousExecutor; import software.amazon.jdbc.util.Utils; -import software.amazon.jdbc.util.connection.ConnectionService; import software.amazon.jdbc.util.monitoring.AbstractMonitor; -import software.amazon.jdbc.util.monitoring.Monitor; -import software.amazon.jdbc.util.storage.StorageService; public class ClusterTopologyMonitorImpl extends AbstractMonitor implements ClusterTopologyMonitor { @@ -478,15 +474,15 @@ protected boolean isInPanicMode() { } protected Runnable getNodeMonitoringWorker( - final HostSpec hostSpec, final @Nullable HostSpec writerHostSpec) - throws SQLException { + final HostSpec hostSpec, final @Nullable HostSpec writerHostSpec) { return new NodeMonitoringWorker(this.getNewServicesContainer(), this, hostSpec, writerHostSpec); } - protected FullServicesContainer getNewServicesContainer() throws SQLException { + protected FullServicesContainer getNewServicesContainer() { return ServiceContainerUtility.createServiceContainer( this.servicesContainer.getStorageService(), this.servicesContainer.getMonitorService(), + this.servicesContainer.getDefaultConnectionProvider(), this.servicesContainer.getTelemetryFactory(), this.servicesContainer.getPluginService().getOriginalUrl(), this.servicesContainer.getPluginService().getDriverProtocol(), diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/customendpoint/CustomEndpointPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/customendpoint/CustomEndpointPlugin.java index 22b385923..5a272e922 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/customendpoint/CustomEndpointPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/customendpoint/CustomEndpointPlugin.java @@ -223,9 +223,9 @@ protected CustomEndpointMonitor createMonitorIfAbsent(Properties props) throws S this.pluginService.getTargetDriverDialect(), this.pluginService.getDialect(), this.props, - (connectionService, pluginService) -> new CustomEndpointMonitorImpl( - this.servicesContainer.getStorageService(), - this.servicesContainer.getTelemetryFactory(), + (servicesContainer) -> new CustomEndpointMonitorImpl( + servicesContainer.getStorageService(), + servicesContainer.getTelemetryFactory(), this.customEndpointHostSpec, this.customEndpointId, this.region, diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/efm2/HostMonitorImpl.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/efm2/HostMonitorImpl.java index 5a858b856..7f13420b8 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/efm2/HostMonitorImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/efm2/HostMonitorImpl.java @@ -33,9 +33,9 @@ import org.checkerframework.checker.nullness.qual.NonNull; import software.amazon.jdbc.HostSpec; import software.amazon.jdbc.util.ExecutorFactory; +import software.amazon.jdbc.util.FullServicesContainer; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.PropertyUtils; -import software.amazon.jdbc.util.connection.ConnectionService; import software.amazon.jdbc.util.monitoring.AbstractMonitor; import software.amazon.jdbc.util.telemetry.TelemetryContext; import software.amazon.jdbc.util.telemetry.TelemetryCounter; @@ -59,7 +59,7 @@ public class HostMonitorImpl extends AbstractMonitor implements HostMonitor { private final Queue> activeContexts = new ConcurrentLinkedQueue<>(); private final Map>> newContexts = new ConcurrentHashMap<>(); - private final ConnectionService connectionService; + private final FullServicesContainer servicesContainer; private final TelemetryFactory telemetryFactory; private final Properties properties; private final HostSpec hostSpec; @@ -78,8 +78,7 @@ public class HostMonitorImpl extends AbstractMonitor implements HostMonitor { /** * Store the monitoring configuration for a connection. * - * @param connectionService The service to use to create the monitoring connection. - * @param telemetryFactory The telemetry factory to use to create telemetry data. + * @param servicesContainer The telemetry factory to use to create telemetry data. * @param hostSpec The {@link HostSpec} of the server this {@link HostMonitorImpl} * instance is monitoring. * @param properties The {@link Properties} containing additional monitoring @@ -90,8 +89,7 @@ public class HostMonitorImpl extends AbstractMonitor implements HostMonitor { * @param abortedConnectionsCounter Aborted connection telemetry counter. */ public HostMonitorImpl( - final @NonNull ConnectionService connectionService, - final @NonNull TelemetryFactory telemetryFactory, + final @NonNull FullServicesContainer servicesContainer, final @NonNull HostSpec hostSpec, final @NonNull Properties properties, final int failureDetectionTimeMillis, @@ -99,9 +97,8 @@ public HostMonitorImpl( final int failureDetectionCount, final TelemetryCounter abortedConnectionsCounter) { super(TERMINATION_TIMEOUT_SEC, ExecutorFactory.newFixedThreadPool(2, "efm2-monitor")); - - this.connectionService = connectionService; - this.telemetryFactory = telemetryFactory; + this.servicesContainer = servicesContainer; + this.telemetryFactory = servicesContainer.getTelemetryFactory(); this.hostSpec = hostSpec; this.properties = properties; this.failureDetectionTimeNano = TimeUnit.MILLISECONDS.toNanos(failureDetectionTimeMillis); @@ -317,7 +314,8 @@ boolean checkConnectionStatus() { }); LOGGER.finest(() -> "Opening a monitoring connection to " + this.hostSpec.getUrl()); - this.monitoringConn = this.connectionService.open(this.hostSpec, monitoringConnProperties); + this.monitoringConn = + this.servicesContainer.getPluginService().forceConnect(this.hostSpec, monitoringConnProperties); LOGGER.finest(() -> "Opened monitoring connection: " + this.monitoringConn); return true; } diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/efm2/HostMonitorServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/efm2/HostMonitorServiceImpl.java index 21c0a55ff..7bcd09e7d 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/efm2/HostMonitorServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/efm2/HostMonitorServiceImpl.java @@ -160,9 +160,8 @@ protected HostMonitor getMonitor( this.pluginService.getTargetDriverDialect(), this.pluginService.getDialect(), this.pluginService.getProperties(), - (connectionService, pluginService) -> new HostMonitorImpl( - connectionService, - pluginService.getTelemetryFactory(), + (servicesContainer) -> new HostMonitorImpl( + servicesContainer, hostSpec, properties, failureDetectionTimeMillis, diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java index a9086c13a..71df06662 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java @@ -34,14 +34,12 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.logging.Logger; -import software.amazon.jdbc.ConnectionPluginManager; import software.amazon.jdbc.HostRole; import software.amazon.jdbc.HostSpec; import software.amazon.jdbc.PluginService; import software.amazon.jdbc.hostavailability.HostAvailability; import software.amazon.jdbc.util.ExecutorFactory; import software.amazon.jdbc.util.FullServicesContainer; -import software.amazon.jdbc.util.FullServicesContainerImpl; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.PropertyUtils; import software.amazon.jdbc.util.ServiceContainerUtility; @@ -367,10 +365,11 @@ private ReaderFailoverResult getResultFromNextTaskBatch( return new ReaderFailoverResult(null, null, false); } - protected FullServicesContainer getNewServicesContainer() throws SQLException { + protected FullServicesContainer getNewServicesContainer() { return ServiceContainerUtility.createServiceContainer( this.servicesContainer.getStorageService(), this.servicesContainer.getMonitorService(), + this.pluginService.getDefaultConnectionProvider(), this.servicesContainer.getTelemetryFactory(), this.pluginService.getOriginalUrl(), this.pluginService.getDriverProtocol(), @@ -402,29 +401,6 @@ private ReaderFailoverResult getNextResult(final CompletionService { private final PluginService pluginService; private final Map availabilityMap; diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java index 49cc133a6..48055df42 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java @@ -32,13 +32,11 @@ import java.util.concurrent.TimeUnit; import java.util.logging.Level; import java.util.logging.Logger; -import software.amazon.jdbc.ConnectionPluginManager; import software.amazon.jdbc.HostSpec; import software.amazon.jdbc.PluginService; import software.amazon.jdbc.hostavailability.HostAvailability; import software.amazon.jdbc.util.ExecutorFactory; import software.amazon.jdbc.util.FullServicesContainer; -import software.amazon.jdbc.util.FullServicesContainerImpl; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.PropertyUtils; import software.amazon.jdbc.util.ServiceContainerUtility; @@ -143,7 +141,7 @@ private void submitTasks( final List currentTopology, final ExecutorService executorService, final CompletionService completionService, - final boolean singleTask) throws SQLException { + final boolean singleTask) { final HostSpec writerHost = Utils.getWriter(currentTopology); if (!singleTask) { completionService.submit( @@ -168,27 +166,13 @@ private void submitTasks( executorService.shutdown(); } - // Each task should get its own PluginService since they execute concurrently and PluginService was not designed to - // be thread-safe. - protected PluginService getNewPluginService() throws SQLException { - FullServicesContainer newServicesContainer = new FullServicesContainerImpl( - this.servicesContainer.getStorageService(), - this.servicesContainer.getMonitorService(), - this.servicesContainer.getTelemetryFactory() - ); - - ConnectionPluginManager pluginManager = new ConnectionPluginManager( - this.pluginService.getDefaultConnectionProvider(), null, null, servicesContainer.getTelemetryFactory()); - newServicesContainer.setConnectionPluginManager(pluginManager); - PartialPluginService pluginService = new PartialPluginService( - newServicesContainer, - this.initialConnectionProps, - protected FullServicesContainer getNewServicesContainer() throws SQLException { + protected FullServicesContainer getNewServicesContainer() { // Each task should get its own FullServicesContainer since they execute concurrently and PluginService was not // designed to be thread-safe. return ServiceContainerUtility.createServiceContainer( this.servicesContainer.getStorageService(), this.servicesContainer.getMonitorService(), + this.pluginService.getDefaultConnectionProvider(), this.servicesContainer.getTelemetryFactory(), this.pluginService.getOriginalUrl(), this.pluginService.getDriverProtocol(), @@ -196,9 +180,6 @@ protected FullServicesContainer getNewServicesContainer() throws SQLException { this.pluginService.getDialect(), this.initialConnectionProps ); - - pluginManager.init(newServicesContainer, this.initialConnectionProps, pluginService, null); - return pluginService; } private WriterFailoverResult getNextResult( @@ -353,6 +334,10 @@ public WriterFailoverResult call() { private boolean isCurrentHostWriter(final List latestTopology) { final HostSpec latestWriter = Utils.getWriter(latestTopology); + if (latestWriter == null) { + return false; + } + final Set latestWriterAllAliases = latestWriter.asAliases(); final Set currentAliases = this.originalWriterHost.asAliases(); diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java index a1e892e35..79c60c7f0 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPlugin.java @@ -598,10 +598,6 @@ protected void failover(@Nullable final HostSpec failedHost) throws SQLException this.pluginService.setAvailability(failedHost.asAliases(), HostAvailability.NOT_AVAILABLE); } - if (this.connectionService == null) { - this.connectionService = getConnectionService(); - } - if (this.failoverMode == FailoverMode.STRICT_WRITER) { failoverWriter(); } else { @@ -609,20 +605,6 @@ protected void failover(@Nullable final HostSpec failedHost) throws SQLException } } - protected ConnectionService getConnectionService() throws SQLException { - return new ConnectionServiceImpl( - servicesContainer.getStorageService(), - servicesContainer.getMonitorService(), - servicesContainer.getTelemetryFactory(), - this.pluginService.getDefaultConnectionProvider(), - this.pluginService.getOriginalUrl(), - this.pluginService.getDriverProtocol(), - this.pluginService.getTargetDriverDialect(), - this.pluginService.getDialect(), - properties - ); - } - protected void failoverReader(final HostSpec failedHostSpec) throws SQLException { TelemetryFactory telemetryFactory = this.pluginService.getTelemetryFactory(); TelemetryContext telemetryContext = telemetryFactory.openTelemetryContext( diff --git a/wrapper/src/main/java/software/amazon/jdbc/wrapper/ConnectionWrapper.java b/wrapper/src/main/java/software/amazon/jdbc/wrapper/ConnectionWrapper.java index dd99e0fee..90b47ac88 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/wrapper/ConnectionWrapper.java +++ b/wrapper/src/main/java/software/amazon/jdbc/wrapper/ConnectionWrapper.java @@ -58,7 +58,6 @@ import software.amazon.jdbc.util.SqlState; import software.amazon.jdbc.util.StringUtils; import software.amazon.jdbc.util.WrapperUtils; -import software.amazon.jdbc.util.connection.ConnectionService; import software.amazon.jdbc.util.monitoring.MonitorService; import software.amazon.jdbc.util.storage.StorageService; import software.amazon.jdbc.util.telemetry.TelemetryFactory; @@ -136,8 +135,7 @@ protected ConnectionWrapper( @NonNull final HostListProviderService hostListProviderService, @NonNull final PluginManagerService pluginManagerService, @NonNull final StorageService storageService, - @NonNull final MonitorService monitorService, - @NonNull final ConnectionService connectionService) + @NonNull final MonitorService monitorService) throws SQLException { if (StringUtils.isNullOrEmpty(url)) { @@ -147,6 +145,7 @@ protected ConnectionWrapper( FullServicesContainer servicesContainer = new FullServicesContainerImpl( storageService, monitorService, + defaultConnectionProvider, telemetryFactory, connectionPluginManager, hostListProviderService, diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java index 89e838acb..7f7d3567c 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java @@ -1,400 +1,400 @@ -/* - * 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.failover; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertSame; -import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.doReturn; -import static org.mockito.Mockito.spy; -import static org.mockito.Mockito.when; -import static software.amazon.jdbc.plugin.failover.ClusterAwareReaderFailoverHandler.DEFAULT_FAILOVER_TIMEOUT; -import static software.amazon.jdbc.plugin.failover.ClusterAwareReaderFailoverHandler.DEFAULT_READER_CONNECT_TIMEOUT; - -import java.sql.Connection; -import java.sql.SQLException; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.Collections; -import java.util.EnumSet; -import java.util.List; -import java.util.Map; -import java.util.Properties; -import java.util.Set; -import java.util.concurrent.TimeUnit; -import java.util.stream.Collectors; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.mockito.Mock; -import org.mockito.Mockito; -import org.mockito.MockitoAnnotations; -import org.mockito.stubbing.Answer; -import software.amazon.jdbc.ConnectionPluginManager; -import software.amazon.jdbc.HostRole; -import software.amazon.jdbc.HostSpec; -import software.amazon.jdbc.HostSpecBuilder; -import software.amazon.jdbc.PluginService; -import software.amazon.jdbc.dialect.Dialect; -import software.amazon.jdbc.hostavailability.HostAvailability; -import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; -import software.amazon.jdbc.util.FullServicesContainer; -import software.amazon.jdbc.util.connection.ConnectionService; - -class ClusterAwareReaderFailoverHandlerTest { - @Mock FullServicesContainer mockContainer; - @Mock ConnectionService mockConnectionService; - @Mock PluginService mockPluginService; - @Mock ConnectionPluginManager mockPluginManager; - @Mock Connection mockConnection; - - private AutoCloseable closeable; - private final Properties properties = new Properties(); - private final List defaultHosts = Arrays.asList( - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("writer").port(1234).role(HostRole.WRITER).build(), - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader1").port(1234).role(HostRole.READER).build(), - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader2").port(1234).role(HostRole.READER).build(), - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader3").port(1234).role(HostRole.READER).build(), - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader4").port(1234).role(HostRole.READER).build(), - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader5").port(1234).role(HostRole.READER).build() - ); - - @BeforeEach - void setUp() { - closeable = MockitoAnnotations.openMocks(this); - when(mockContainer.getConnectionPluginManager()).thenReturn(mockPluginManager); - when(mockContainer.getPluginService()).thenReturn(mockPluginService); - } - - @AfterEach - void tearDown() throws Exception { - closeable.close(); - } - - @Test - public void testFailover() throws SQLException { - // original host list: [active writer, active reader, current connection (reader), active - // reader, down reader, active reader] - // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] - // connection attempts are made in pairs using the above list - // expected test result: successful connection for host at index 4 - final List hosts = defaultHosts; - final int currentHostIndex = 2; - final int successHostIndex = 4; - for (int i = 0; i < hosts.size(); i++) { - if (i != successHostIndex) { - final SQLException exception = new SQLException("exception", "08S01", null); - when(mockConnectionService.open(hosts.get(i), properties)) - .thenThrow(exception); - when(mockPluginService.isNetworkException(exception, null)).thenReturn(true); - } else { - when(mockConnectionService.open(hosts.get(i), properties)).thenReturn(mockConnection); - } - } - - when(mockPluginService.getTargetDriverDialect()).thenReturn(null); - - hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); - hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); - - final ReaderFailoverHandler target = getSpyFailoverHandler(); - final ReaderFailoverResult result = target.failover(hosts, hosts.get(currentHostIndex)); - - assertTrue(result.isConnected()); - assertSame(mockConnection, result.getConnection()); - assertEquals(hosts.get(successHostIndex), result.getHost()); - - final HostSpec successHost = hosts.get(successHostIndex); - final Map availabilityMap = target.getHostAvailabilityMap(); - Set unavailableHosts = getHostsWithGivenAvailability(availabilityMap, HostAvailability.NOT_AVAILABLE); - assertTrue(unavailableHosts.size() >= 4); - assertEquals(HostAvailability.AVAILABLE, availabilityMap.get(successHost.getHost())); - } - - private Set getHostsWithGivenAvailability( - Map availabilityMap, HostAvailability availability) { - return availabilityMap.entrySet().stream() - .filter((entry) -> availability.equals(entry.getValue())) - .map(Map.Entry::getKey) - .collect(Collectors.toSet()); - } - - @Test - public void testFailover_timeout() throws SQLException { - // original host list: [active writer, active reader, current connection (reader), active - // reader, down reader, active reader] - // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] - // connection attempts are made in pairs using the above list - // expected test result: failure to get reader since process is limited to 5s and each attempt - // to connect takes 20s - final List hosts = defaultHosts; - final int currentHostIndex = 2; - for (HostSpec host : hosts) { - when(mockConnectionService.open(host, properties)) - .thenAnswer((Answer) invocation -> { - Thread.sleep(20000); - return mockConnection; - }); - } - - hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); - hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); - - final ReaderFailoverHandler target = getSpyFailoverHandler(5000, 30000, false); - - final long startTimeNano = System.nanoTime(); - final ReaderFailoverResult result = target.failover(hosts, hosts.get(currentHostIndex)); - final long durationNano = System.nanoTime() - startTimeNano; - - assertFalse(result.isConnected()); - assertNull(result.getConnection()); - assertNull(result.getHost()); - - // 5s is a max allowed failover timeout; add 1s for inaccurate measurements - assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); - } - - private ClusterAwareReaderFailoverHandler getSpyFailoverHandler() throws SQLException { - ClusterAwareReaderFailoverHandler handler = - spy(new ClusterAwareReaderFailoverHandler(mockContainer, mockConnectionService, properties)); - doReturn(mockPluginService).when(handler).getNewPluginService(); - return handler; - } - - private ClusterAwareReaderFailoverHandler getSpyFailoverHandler( - int maxFailoverTimeoutMs, int timeoutMs, boolean isStrictReaderRequired) throws SQLException { - ClusterAwareReaderFailoverHandler handler = new ClusterAwareReaderFailoverHandler( - mockContainer, mockConnectionService, properties, maxFailoverTimeoutMs, timeoutMs, isStrictReaderRequired); - ClusterAwareReaderFailoverHandler spyHandler = spy(handler); - doReturn(mockPluginService).when(spyHandler).getNewPluginService(); - return spyHandler; - } - - @Test - public void testFailover_nullOrEmptyHostList() throws SQLException { - final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); - final HostSpec currentHost = - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("writer").port(1234).build(); - - ReaderFailoverResult result = target.failover(null, currentHost); - assertFalse(result.isConnected()); - assertNull(result.getConnection()); - assertNull(result.getHost()); - - final List hosts = new ArrayList<>(); - result = target.failover(hosts, currentHost); - assertFalse(result.isConnected()); - assertNull(result.getConnection()); - assertNull(result.getHost()); - } - - @Test - public void testGetReader_connectionSuccess() throws SQLException { - // even number of connection attempts - // first connection attempt to return succeeds, second attempt cancelled - // expected test result: successful connection for host at index 2 - final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) - final HostSpec slowHost = hosts.get(1); - final HostSpec fastHost = hosts.get(2); - when(mockConnectionService.open(slowHost, properties)) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(20000); - return mockConnection; - }); - when(mockConnectionService.open(eq(fastHost), eq(properties))).thenReturn(mockConnection); - - Dialect mockDialect = Mockito.mock(Dialect.class); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - when(mockPluginService.getDialect()).thenReturn(mockDialect); - - final ReaderFailoverHandler target = getSpyFailoverHandler(); - final ReaderFailoverResult result = target.getReaderConnection(hosts); - - assertTrue(result.isConnected()); - assertSame(mockConnection, result.getConnection()); - assertEquals(hosts.get(2), result.getHost()); - - Map availabilityMap = target.getHostAvailabilityMap(); - assertTrue(getHostsWithGivenAvailability(availabilityMap, HostAvailability.NOT_AVAILABLE).isEmpty()); - assertEquals(HostAvailability.AVAILABLE, availabilityMap.get(fastHost.getHost())); - } - - @Test - public void testGetReader_connectionFailure() throws SQLException { - // odd number of connection attempts - // first connection attempt to return fails - // expected test result: failure to get reader - final List hosts = defaultHosts.subList(0, 4); // 3 connection attempts (writer not attempted) - when(mockConnectionService.open(any(), eq(properties))).thenThrow(new SQLException("exception", "08S01", null)); - - Dialect mockDialect = Mockito.mock(Dialect.class); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - when(mockPluginService.getDialect()).thenReturn(mockDialect); - - final ReaderFailoverHandler target = getSpyFailoverHandler(); - final ReaderFailoverResult result = target.getReaderConnection(hosts); - - assertFalse(result.isConnected()); - assertNull(result.getConnection()); - assertNull(result.getHost()); - } - - @Test - public void testGetReader_connectionAttemptsTimeout() throws SQLException { - // connection attempts time out before they can succeed - // first connection attempt to return times out - // expected test result: failure to get reader - final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) - when(mockConnectionService.open(any(), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - try { - Thread.sleep(5000); - } catch (InterruptedException exception) { - // ignore - } - return mockConnection; - }); - - Dialect mockDialect = Mockito.mock(Dialect.class); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - when(mockPluginService.getDialect()).thenReturn(mockDialect); - - final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(60000, 1000, false); - final ReaderFailoverResult result = target.getReaderConnection(hosts); - - assertFalse(result.isConnected()); - assertNull(result.getConnection()); - assertNull(result.getHost()); - } - - @Test - public void testGetHostTuplesByPriority() throws SQLException { - final List originalHosts = defaultHosts; - originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); - originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); - originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); - - final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); - final List hostsByPriority = target.getHostsByPriority(originalHosts); - - int i = 0; - - // expecting active readers - while (i < hostsByPriority.size() - && hostsByPriority.get(i).getRole() == HostRole.READER - && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { - i++; - } - - // expecting a writer - while (i < hostsByPriority.size() - && hostsByPriority.get(i).getRole() == HostRole.WRITER) { - i++; - } - - // expecting down readers - while (i < hostsByPriority.size() - && hostsByPriority.get(i).getRole() == HostRole.READER - && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { - i++; - } - - assertEquals(hostsByPriority.size(), i); - } - - @Test - public void testGetReaderTuplesByPriority() throws SQLException { - final List originalHosts = defaultHosts; - originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); - originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); - originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); - - Dialect mockDialect = Mockito.mock(Dialect.class); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - when(mockPluginService.getDialect()).thenReturn(mockDialect); - - final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); - final List hostsByPriority = target.getReaderHostsByPriority(originalHosts); - - int i = 0; - - // expecting active readers - while (i < hostsByPriority.size() - && hostsByPriority.get(i).getRole() == HostRole.READER - && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { - i++; - } - - // expecting down readers - while (i < hostsByPriority.size() - && hostsByPriority.get(i).getRole() == HostRole.READER - && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { - i++; - } - - assertEquals(hostsByPriority.size(), i); - } - - @Test - public void testHostFailoverStrictReaderEnabled() throws SQLException { - final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("writer").port(1234).role(HostRole.WRITER).build(); - final HostSpec reader = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader1").port(1234).role(HostRole.READER).build(); - final List hosts = Arrays.asList(writer, reader); - - Dialect mockDialect = Mockito.mock(Dialect.class); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - when(mockPluginService.getDialect()).thenReturn(mockDialect); - - final ClusterAwareReaderFailoverHandler target = - getSpyFailoverHandler(DEFAULT_FAILOVER_TIMEOUT, DEFAULT_READER_CONNECT_TIMEOUT, true); - - // The writer is included because the original writer has likely become a reader. - List expectedHostsByPriority = Arrays.asList(reader, writer); - - List hostsByPriority = target.getHostsByPriority(hosts); - assertEquals(expectedHostsByPriority, hostsByPriority); - - // Should pick the reader even if unavailable. The unavailable reader will be lower priority than the writer. - reader.setAvailability(HostAvailability.NOT_AVAILABLE); - expectedHostsByPriority = Arrays.asList(writer, reader); - - hostsByPriority = target.getHostsByPriority(hosts); - assertEquals(expectedHostsByPriority, hostsByPriority); - - // Writer node will only be picked if it is the only node in topology; - List expectedWriterHost = Collections.singletonList(writer); - - hostsByPriority = target.getHostsByPriority(Collections.singletonList(writer)); - assertEquals(expectedWriterHost, hostsByPriority); - } -} +// /* +// * 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.failover; +// +// import static org.junit.jupiter.api.Assertions.assertEquals; +// import static org.junit.jupiter.api.Assertions.assertFalse; +// import static org.junit.jupiter.api.Assertions.assertNull; +// import static org.junit.jupiter.api.Assertions.assertSame; +// import static org.junit.jupiter.api.Assertions.assertTrue; +// import static org.mockito.ArgumentMatchers.any; +// import static org.mockito.ArgumentMatchers.eq; +// import static org.mockito.Mockito.doReturn; +// import static org.mockito.Mockito.spy; +// import static org.mockito.Mockito.when; +// import static software.amazon.jdbc.plugin.failover.ClusterAwareReaderFailoverHandler.DEFAULT_FAILOVER_TIMEOUT; +// import static software.amazon.jdbc.plugin.failover.ClusterAwareReaderFailoverHandler.DEFAULT_READER_CONNECT_TIMEOUT; +// +// import java.sql.Connection; +// import java.sql.SQLException; +// import java.util.ArrayList; +// import java.util.Arrays; +// import java.util.Collections; +// import java.util.EnumSet; +// import java.util.List; +// import java.util.Map; +// import java.util.Properties; +// import java.util.Set; +// import java.util.concurrent.TimeUnit; +// import java.util.stream.Collectors; +// import org.junit.jupiter.api.AfterEach; +// import org.junit.jupiter.api.BeforeEach; +// import org.junit.jupiter.api.Test; +// import org.mockito.Mock; +// import org.mockito.Mockito; +// import org.mockito.MockitoAnnotations; +// import org.mockito.stubbing.Answer; +// import software.amazon.jdbc.ConnectionPluginManager; +// import software.amazon.jdbc.HostRole; +// import software.amazon.jdbc.HostSpec; +// import software.amazon.jdbc.HostSpecBuilder; +// import software.amazon.jdbc.PluginService; +// import software.amazon.jdbc.dialect.Dialect; +// import software.amazon.jdbc.hostavailability.HostAvailability; +// import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; +// import software.amazon.jdbc.util.FullServicesContainer; +// import software.amazon.jdbc.util.connection.ConnectionService; +// +// class ClusterAwareReaderFailoverHandlerTest { +// @Mock FullServicesContainer mockContainer; +// @Mock ConnectionService mockConnectionService; +// @Mock PluginService mockPluginService; +// @Mock ConnectionPluginManager mockPluginManager; +// @Mock Connection mockConnection; +// +// private AutoCloseable closeable; +// private final Properties properties = new Properties(); +// private final List defaultHosts = Arrays.asList( +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("writer").port(1234).role(HostRole.WRITER).build(), +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader1").port(1234).role(HostRole.READER).build(), +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader2").port(1234).role(HostRole.READER).build(), +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader3").port(1234).role(HostRole.READER).build(), +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader4").port(1234).role(HostRole.READER).build(), +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader5").port(1234).role(HostRole.READER).build() +// ); +// +// @BeforeEach +// void setUp() { +// closeable = MockitoAnnotations.openMocks(this); +// when(mockContainer.getConnectionPluginManager()).thenReturn(mockPluginManager); +// when(mockContainer.getPluginService()).thenReturn(mockPluginService); +// } +// +// @AfterEach +// void tearDown() throws Exception { +// closeable.close(); +// } +// +// @Test +// public void testFailover() throws SQLException { +// // original host list: [active writer, active reader, current connection (reader), active +// // reader, down reader, active reader] +// // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] +// // connection attempts are made in pairs using the above list +// // expected test result: successful connection for host at index 4 +// final List hosts = defaultHosts; +// final int currentHostIndex = 2; +// final int successHostIndex = 4; +// for (int i = 0; i < hosts.size(); i++) { +// if (i != successHostIndex) { +// final SQLException exception = new SQLException("exception", "08S01", null); +// when(mockConnectionService.open(hosts.get(i), properties)) +// .thenThrow(exception); +// when(mockPluginService.isNetworkException(exception, null)).thenReturn(true); +// } else { +// when(mockConnectionService.open(hosts.get(i), properties)).thenReturn(mockConnection); +// } +// } +// +// when(mockPluginService.getTargetDriverDialect()).thenReturn(null); +// +// hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); +// hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); +// +// final ReaderFailoverHandler target = getSpyFailoverHandler(); +// final ReaderFailoverResult result = target.failover(hosts, hosts.get(currentHostIndex)); +// +// assertTrue(result.isConnected()); +// assertSame(mockConnection, result.getConnection()); +// assertEquals(hosts.get(successHostIndex), result.getHost()); +// +// final HostSpec successHost = hosts.get(successHostIndex); +// final Map availabilityMap = target.getHostAvailabilityMap(); +// Set unavailableHosts = getHostsWithGivenAvailability(availabilityMap, HostAvailability.NOT_AVAILABLE); +// assertTrue(unavailableHosts.size() >= 4); +// assertEquals(HostAvailability.AVAILABLE, availabilityMap.get(successHost.getHost())); +// } +// +// private Set getHostsWithGivenAvailability( +// Map availabilityMap, HostAvailability availability) { +// return availabilityMap.entrySet().stream() +// .filter((entry) -> availability.equals(entry.getValue())) +// .map(Map.Entry::getKey) +// .collect(Collectors.toSet()); +// } +// +// @Test +// public void testFailover_timeout() throws SQLException { +// // original host list: [active writer, active reader, current connection (reader), active +// // reader, down reader, active reader] +// // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] +// // connection attempts are made in pairs using the above list +// // expected test result: failure to get reader since process is limited to 5s and each attempt +// // to connect takes 20s +// final List hosts = defaultHosts; +// final int currentHostIndex = 2; +// for (HostSpec host : hosts) { +// when(mockConnectionService.open(host, properties)) +// .thenAnswer((Answer) invocation -> { +// Thread.sleep(20000); +// return mockConnection; +// }); +// } +// +// hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); +// hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); +// +// final ReaderFailoverHandler target = getSpyFailoverHandler(5000, 30000, false); +// +// final long startTimeNano = System.nanoTime(); +// final ReaderFailoverResult result = target.failover(hosts, hosts.get(currentHostIndex)); +// final long durationNano = System.nanoTime() - startTimeNano; +// +// assertFalse(result.isConnected()); +// assertNull(result.getConnection()); +// assertNull(result.getHost()); +// +// // 5s is a max allowed failover timeout; add 1s for inaccurate measurements +// assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); +// } +// +// private ClusterAwareReaderFailoverHandler getSpyFailoverHandler() throws SQLException { +// ClusterAwareReaderFailoverHandler handler = +// spy(new ClusterAwareReaderFailoverHandler(mockContainer, mockConnectionService, properties)); +// doReturn(mockPluginService).when(handler).getNewPluginService(); +// return handler; +// } +// +// private ClusterAwareReaderFailoverHandler getSpyFailoverHandler( +// int maxFailoverTimeoutMs, int timeoutMs, boolean isStrictReaderRequired) throws SQLException { +// ClusterAwareReaderFailoverHandler handler = new ClusterAwareReaderFailoverHandler( +// mockContainer, mockConnectionService, properties, maxFailoverTimeoutMs, timeoutMs, isStrictReaderRequired); +// ClusterAwareReaderFailoverHandler spyHandler = spy(handler); +// doReturn(mockPluginService).when(spyHandler).getNewPluginService(); +// return spyHandler; +// } +// +// @Test +// public void testFailover_nullOrEmptyHostList() throws SQLException { +// final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); +// final HostSpec currentHost = +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("writer").port(1234).build(); +// +// ReaderFailoverResult result = target.failover(null, currentHost); +// assertFalse(result.isConnected()); +// assertNull(result.getConnection()); +// assertNull(result.getHost()); +// +// final List hosts = new ArrayList<>(); +// result = target.failover(hosts, currentHost); +// assertFalse(result.isConnected()); +// assertNull(result.getConnection()); +// assertNull(result.getHost()); +// } +// +// @Test +// public void testGetReader_connectionSuccess() throws SQLException { +// // even number of connection attempts +// // first connection attempt to return succeeds, second attempt cancelled +// // expected test result: successful connection for host at index 2 +// final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) +// final HostSpec slowHost = hosts.get(1); +// final HostSpec fastHost = hosts.get(2); +// when(mockConnectionService.open(slowHost, properties)) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(20000); +// return mockConnection; +// }); +// when(mockConnectionService.open(eq(fastHost), eq(properties))).thenReturn(mockConnection); +// +// Dialect mockDialect = Mockito.mock(Dialect.class); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// +// final ReaderFailoverHandler target = getSpyFailoverHandler(); +// final ReaderFailoverResult result = target.getReaderConnection(hosts); +// +// assertTrue(result.isConnected()); +// assertSame(mockConnection, result.getConnection()); +// assertEquals(hosts.get(2), result.getHost()); +// +// Map availabilityMap = target.getHostAvailabilityMap(); +// assertTrue(getHostsWithGivenAvailability(availabilityMap, HostAvailability.NOT_AVAILABLE).isEmpty()); +// assertEquals(HostAvailability.AVAILABLE, availabilityMap.get(fastHost.getHost())); +// } +// +// @Test +// public void testGetReader_connectionFailure() throws SQLException { +// // odd number of connection attempts +// // first connection attempt to return fails +// // expected test result: failure to get reader +// final List hosts = defaultHosts.subList(0, 4); // 3 connection attempts (writer not attempted) +// when(mockConnectionService.open(any(), eq(properties))).thenThrow(new SQLException("exception", "08S01", null)); +// +// Dialect mockDialect = Mockito.mock(Dialect.class); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// +// final ReaderFailoverHandler target = getSpyFailoverHandler(); +// final ReaderFailoverResult result = target.getReaderConnection(hosts); +// +// assertFalse(result.isConnected()); +// assertNull(result.getConnection()); +// assertNull(result.getHost()); +// } +// +// @Test +// public void testGetReader_connectionAttemptsTimeout() throws SQLException { +// // connection attempts time out before they can succeed +// // first connection attempt to return times out +// // expected test result: failure to get reader +// final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) +// when(mockConnectionService.open(any(), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// try { +// Thread.sleep(5000); +// } catch (InterruptedException exception) { +// // ignore +// } +// return mockConnection; +// }); +// +// Dialect mockDialect = Mockito.mock(Dialect.class); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// +// final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(60000, 1000, false); +// final ReaderFailoverResult result = target.getReaderConnection(hosts); +// +// assertFalse(result.isConnected()); +// assertNull(result.getConnection()); +// assertNull(result.getHost()); +// } +// +// @Test +// public void testGetHostTuplesByPriority() throws SQLException { +// final List originalHosts = defaultHosts; +// originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); +// originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); +// originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); +// +// final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); +// final List hostsByPriority = target.getHostsByPriority(originalHosts); +// +// int i = 0; +// +// // expecting active readers +// while (i < hostsByPriority.size() +// && hostsByPriority.get(i).getRole() == HostRole.READER +// && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { +// i++; +// } +// +// // expecting a writer +// while (i < hostsByPriority.size() +// && hostsByPriority.get(i).getRole() == HostRole.WRITER) { +// i++; +// } +// +// // expecting down readers +// while (i < hostsByPriority.size() +// && hostsByPriority.get(i).getRole() == HostRole.READER +// && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { +// i++; +// } +// +// assertEquals(hostsByPriority.size(), i); +// } +// +// @Test +// public void testGetReaderTuplesByPriority() throws SQLException { +// final List originalHosts = defaultHosts; +// originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); +// originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); +// originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); +// +// Dialect mockDialect = Mockito.mock(Dialect.class); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// +// final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); +// final List hostsByPriority = target.getReaderHostsByPriority(originalHosts); +// +// int i = 0; +// +// // expecting active readers +// while (i < hostsByPriority.size() +// && hostsByPriority.get(i).getRole() == HostRole.READER +// && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { +// i++; +// } +// +// // expecting down readers +// while (i < hostsByPriority.size() +// && hostsByPriority.get(i).getRole() == HostRole.READER +// && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { +// i++; +// } +// +// assertEquals(hostsByPriority.size(), i); +// } +// +// @Test +// public void testHostFailoverStrictReaderEnabled() throws SQLException { +// final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("writer").port(1234).role(HostRole.WRITER).build(); +// final HostSpec reader = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader1").port(1234).role(HostRole.READER).build(); +// final List hosts = Arrays.asList(writer, reader); +// +// Dialect mockDialect = Mockito.mock(Dialect.class); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// +// final ClusterAwareReaderFailoverHandler target = +// getSpyFailoverHandler(DEFAULT_FAILOVER_TIMEOUT, DEFAULT_READER_CONNECT_TIMEOUT, true); +// +// // The writer is included because the original writer has likely become a reader. +// List expectedHostsByPriority = Arrays.asList(reader, writer); +// +// List hostsByPriority = target.getHostsByPriority(hosts); +// assertEquals(expectedHostsByPriority, hostsByPriority); +// +// // Should pick the reader even if unavailable. The unavailable reader will be lower priority than the writer. +// reader.setAvailability(HostAvailability.NOT_AVAILABLE); +// expectedHostsByPriority = Arrays.asList(writer, reader); +// +// hostsByPriority = target.getHostsByPriority(hosts); +// assertEquals(expectedHostsByPriority, hostsByPriority); +// +// // Writer node will only be picked if it is the only node in topology; +// List expectedWriterHost = Collections.singletonList(writer); +// +// hostsByPriority = target.getHostsByPriority(Collections.singletonList(writer)); +// assertEquals(expectedWriterHost, hostsByPriority); +// } +// } diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java index 902aadfcb..98d392cdb 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java @@ -1,373 +1,373 @@ -/* - * 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.failover; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertSame; -import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.ArgumentMatchers.refEq; -import static org.mockito.Mockito.atLeastOnce; -import static org.mockito.Mockito.doReturn; -import static org.mockito.Mockito.spy; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -import java.sql.Connection; -import java.sql.SQLException; -import java.util.Arrays; -import java.util.EnumSet; -import java.util.List; -import java.util.Properties; -import java.util.concurrent.TimeUnit; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.mockito.ArgumentMatchers; -import org.mockito.Mock; -import org.mockito.MockitoAnnotations; -import org.mockito.stubbing.Answer; -import software.amazon.jdbc.HostSpec; -import software.amazon.jdbc.HostSpecBuilder; -import software.amazon.jdbc.PluginService; -import software.amazon.jdbc.dialect.Dialect; -import software.amazon.jdbc.hostavailability.HostAvailability; -import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; -import software.amazon.jdbc.util.FullServicesContainer; -import software.amazon.jdbc.util.connection.ConnectionService; - -class ClusterAwareWriterFailoverHandlerTest { - @Mock FullServicesContainer mockContainer; - @Mock ConnectionService mockConnectionService; - @Mock PluginService mockPluginService; - @Mock Connection mockConnection; - @Mock ReaderFailoverHandler mockReaderFailoverHandler; - @Mock Connection mockWriterConnection; - @Mock Connection mockNewWriterConnection; - @Mock Connection mockReaderAConnection; - @Mock Connection mockReaderBConnection; - @Mock Dialect mockDialect; - - private AutoCloseable closeable; - private final Properties properties = new Properties(); - private final HostSpec newWriterHost = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("new-writer-host").build(); - private final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("writer-host").build(); - private final HostSpec readerA = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader-a-host").build(); - private final HostSpec readerB = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader-b-host").build(); - private final List topology = Arrays.asList(writer, readerA, readerB); - private final List newTopology = Arrays.asList(newWriterHost, readerA, readerB); - - @BeforeEach - void setUp() { - closeable = MockitoAnnotations.openMocks(this); - when(mockContainer.getPluginService()).thenReturn(mockPluginService); - writer.addAlias("writer-host"); - newWriterHost.addAlias("new-writer-host"); - readerA.addAlias("reader-a-host"); - readerB.addAlias("reader-b-host"); - } - - @AfterEach - void tearDown() throws Exception { - closeable.close(); - } - - @Test - public void testReconnectToWriter_taskBReaderException() throws SQLException { - when(mockConnectionService.open(refEq(writer), eq(properties))).thenReturn(mockConnection); - when(mockConnectionService.open(refEq(readerA), eq(properties))).thenThrow(SQLException.class); - when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); - - when(mockPluginService.getAllHosts()).thenReturn(topology); - - when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())).thenThrow(SQLException.class); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); - final WriterFailoverResult result = target.failover(topology); - - assertTrue(result.isConnected()); - assertFalse(result.isNewHost()); - assertSame(result.getNewConnection(), mockConnection); - - assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); - } - - private ClusterAwareWriterFailoverHandler getSpyFailoverHandler( - final int failoverTimeoutMs, - final int readTopologyIntervalMs, - final int reconnectWriterIntervalMs) throws SQLException { - ClusterAwareWriterFailoverHandler handler = new ClusterAwareWriterFailoverHandler( - mockContainer, - mockConnectionService, - mockReaderFailoverHandler, - properties, - failoverTimeoutMs, - readTopologyIntervalMs, - reconnectWriterIntervalMs); - - ClusterAwareWriterFailoverHandler spyHandler = spy(handler); - doReturn(mockPluginService).when(spyHandler).getNewPluginService(); - return spyHandler; - } - - /** - * Verify that writer failover handler can re-connect to a current writer node. - * - *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. - * TaskA: successfully re-connect to initial writer; return new connection. - * TaskB: successfully connect to readerA and then new writer, but it takes more time than taskA. - * Expected test result: new connection by taskA. - */ - @Test - public void testReconnectToWriter_SlowReaderA() throws SQLException { - when(mockConnectionService.open(refEq(writer), eq(properties))).thenReturn(mockWriterConnection); - when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); - when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); - when(mockPluginService.getAllHosts()).thenReturn(topology).thenReturn(newTopology); - - when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(5000); - return new ReaderFailoverResult(mockReaderAConnection, readerA, true); - }); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); - final WriterFailoverResult result = target.failover(topology); - - assertTrue(result.isConnected()); - assertFalse(result.isNewHost()); - assertSame(result.getNewConnection(), mockWriterConnection); - assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); - } - - /** - * Verify that writer failover handler can re-connect to a current writer node. - * - *

Topology: no changes. - * TaskA: successfully re-connect to writer; return new connection. - * TaskB: successfully connect to readerA and retrieve topology, but latest writer is not new (defer to taskA). - * Expected test result: new connection by taskA. - */ - @Test - public void testReconnectToWriter_taskBDefers() throws SQLException { - when(mockConnectionService.open(refEq(writer), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(5000); - return mockWriterConnection; - }); - when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); - - when(mockPluginService.getAllHosts()).thenReturn(topology); - - when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) - .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 2000, 2000); - final WriterFailoverResult result = target.failover(topology); - - assertTrue(result.isConnected()); - assertFalse(result.isNewHost()); - assertSame(result.getNewConnection(), mockWriterConnection); - assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); - } - - /** - * Verify that writer failover handler can re-connect to a new writer node. - * - *

Topology: changes to [new-writer, reader-A, reader-B] for taskB, taskA sees no changes. - * taskA: successfully re-connect to writer; return connection to initial writer, but it takes more - * time than taskB. - * TaskB: successfully connect to readerA and then to new-writer. - * Expected test result: new connection to writer by taskB. - */ - @Test - public void testConnectToReaderA_SlowWriter() throws SQLException { - when(mockConnectionService.open(refEq(writer), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(5000); - return mockWriterConnection; - }); - when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); - when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); - when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); - - when(mockPluginService.getAllHosts()).thenReturn(newTopology); - - when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) - .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); - final WriterFailoverResult result = target.failover(topology); - - assertTrue(result.isConnected()); - assertTrue(result.isNewHost()); - assertSame(result.getNewConnection(), mockNewWriterConnection); - assertEquals(3, result.getTopology().size()); - assertEquals("new-writer-host", result.getTopology().get(0).getHost()); - assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); - } - - /** - * Verify that writer failover handler can re-connect to a new writer node. - * - *

Topology: changes to [new-writer, initial-writer, reader-A, reader-B]. - * TaskA: successfully reconnect, but initial-writer is now a reader (defer to taskB). - * TaskB: successfully connect to readerA and then to new-writer. - * Expected test result: new connection to writer by taskB. - */ - @Test - public void testConnectToReaderA_taskADefers() throws SQLException { - when(mockConnectionService.open(writer, properties)).thenReturn(mockConnection); - when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); - when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); - when(mockConnectionService.open(refEq(newWriterHost), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(5000); - return mockNewWriterConnection; - }); - - final List newTopology = Arrays.asList(newWriterHost, writer, readerA, readerB); - when(mockPluginService.getAllHosts()).thenReturn(newTopology); - - when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) - .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); - final WriterFailoverResult result = target.failover(topology); - - assertTrue(result.isConnected()); - assertTrue(result.isNewHost()); - assertSame(result.getNewConnection(), mockNewWriterConnection); - assertEquals(4, result.getTopology().size()); - assertEquals("new-writer-host", result.getTopology().get(0).getHost()); - - verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); - assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); - } - - /** - * Verify that writer failover handler fails to re-connect to any writer node. - * - *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. - * TaskA: fail to re-connect to writer due to failover timeout. - * TaskB: successfully connect to readerA and then fail to connect to writer due to failover timeout. - * Expected test result: no connection. - */ - @Test - public void testFailedToConnect_failoverTimeout() throws SQLException { - when(mockConnectionService.open(refEq(writer), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(30000); - return mockWriterConnection; - }); - when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); - when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); - when(mockConnectionService.open(refEq(newWriterHost), eq(properties))) - .thenAnswer( - (Answer) - invocation -> { - Thread.sleep(30000); - return mockNewWriterConnection; - }); - when(mockPluginService.getAllHosts()).thenReturn(newTopology); - - when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) - .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); - - final long startTimeNano = System.nanoTime(); - final WriterFailoverResult result = target.failover(topology); - final long durationNano = System.nanoTime() - startTimeNano; - - assertFalse(result.isConnected()); - assertFalse(result.isNewHost()); - - verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); - - // 5s is a max allowed failover timeout; add 1s for inaccurate measurements - assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); - } - - /** - * Verify that writer failover handler fails to re-connect to any writer node. - * - *

Topology: changes to [new-writer, reader-A, reader-B] for taskB. - * TaskA: fail to re-connect to writer due to exception. - * TaskB: successfully connect to readerA and then fail to connect to writer due to exception. - * Expected test result: no connection. - */ - @Test - public void testFailedToConnect_taskAException_taskBWriterException() throws SQLException { - final SQLException exception = new SQLException("exception", "08S01", null); - when(mockConnectionService.open(refEq(writer), eq(properties))).thenThrow(exception); - when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); - when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); - when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenThrow(exception); - when(mockPluginService.isNetworkException(eq(exception), any())).thenReturn(true); - - when(mockPluginService.getAllHosts()).thenReturn(newTopology); - - when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) - .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); - - when(mockPluginService.getDialect()).thenReturn(mockDialect); - when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); - - final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); - final WriterFailoverResult result = target.failover(topology); - - assertFalse(result.isConnected()); - assertFalse(result.isNewHost()); - - assertEquals(HostAvailability.NOT_AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); - } -} +// /* +// * 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.failover; +// +// import static org.junit.jupiter.api.Assertions.assertEquals; +// import static org.junit.jupiter.api.Assertions.assertFalse; +// import static org.junit.jupiter.api.Assertions.assertSame; +// import static org.junit.jupiter.api.Assertions.assertTrue; +// import static org.mockito.ArgumentMatchers.any; +// import static org.mockito.ArgumentMatchers.eq; +// import static org.mockito.ArgumentMatchers.refEq; +// import static org.mockito.Mockito.atLeastOnce; +// import static org.mockito.Mockito.doReturn; +// import static org.mockito.Mockito.spy; +// import static org.mockito.Mockito.verify; +// import static org.mockito.Mockito.when; +// +// import java.sql.Connection; +// import java.sql.SQLException; +// import java.util.Arrays; +// import java.util.EnumSet; +// import java.util.List; +// import java.util.Properties; +// import java.util.concurrent.TimeUnit; +// import org.junit.jupiter.api.AfterEach; +// import org.junit.jupiter.api.BeforeEach; +// import org.junit.jupiter.api.Test; +// import org.mockito.ArgumentMatchers; +// import org.mockito.Mock; +// import org.mockito.MockitoAnnotations; +// import org.mockito.stubbing.Answer; +// import software.amazon.jdbc.HostSpec; +// import software.amazon.jdbc.HostSpecBuilder; +// import software.amazon.jdbc.PluginService; +// import software.amazon.jdbc.dialect.Dialect; +// import software.amazon.jdbc.hostavailability.HostAvailability; +// import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; +// import software.amazon.jdbc.util.FullServicesContainer; +// import software.amazon.jdbc.util.connection.ConnectionService; +// +// class ClusterAwareWriterFailoverHandlerTest { +// @Mock FullServicesContainer mockContainer; +// @Mock ConnectionService mockConnectionService; +// @Mock PluginService mockPluginService; +// @Mock Connection mockConnection; +// @Mock ReaderFailoverHandler mockReaderFailoverHandler; +// @Mock Connection mockWriterConnection; +// @Mock Connection mockNewWriterConnection; +// @Mock Connection mockReaderAConnection; +// @Mock Connection mockReaderBConnection; +// @Mock Dialect mockDialect; +// +// private AutoCloseable closeable; +// private final Properties properties = new Properties(); +// private final HostSpec newWriterHost = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("new-writer-host").build(); +// private final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("writer-host").build(); +// private final HostSpec readerA = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader-a-host").build(); +// private final HostSpec readerB = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader-b-host").build(); +// private final List topology = Arrays.asList(writer, readerA, readerB); +// private final List newTopology = Arrays.asList(newWriterHost, readerA, readerB); +// +// @BeforeEach +// void setUp() { +// closeable = MockitoAnnotations.openMocks(this); +// when(mockContainer.getPluginService()).thenReturn(mockPluginService); +// writer.addAlias("writer-host"); +// newWriterHost.addAlias("new-writer-host"); +// readerA.addAlias("reader-a-host"); +// readerB.addAlias("reader-b-host"); +// } +// +// @AfterEach +// void tearDown() throws Exception { +// closeable.close(); +// } +// +// @Test +// public void testReconnectToWriter_taskBReaderException() throws SQLException { +// when(mockConnectionService.open(refEq(writer), eq(properties))).thenReturn(mockConnection); +// when(mockConnectionService.open(refEq(readerA), eq(properties))).thenThrow(SQLException.class); +// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); +// +// when(mockPluginService.getAllHosts()).thenReturn(topology); +// +// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())).thenThrow(SQLException.class); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertTrue(result.isConnected()); +// assertFalse(result.isNewHost()); +// assertSame(result.getNewConnection(), mockConnection); +// +// assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); +// } +// +// private ClusterAwareWriterFailoverHandler getSpyFailoverHandler( +// final int failoverTimeoutMs, +// final int readTopologyIntervalMs, +// final int reconnectWriterIntervalMs) throws SQLException { +// ClusterAwareWriterFailoverHandler handler = new ClusterAwareWriterFailoverHandler( +// mockContainer, +// mockConnectionService, +// mockReaderFailoverHandler, +// properties, +// failoverTimeoutMs, +// readTopologyIntervalMs, +// reconnectWriterIntervalMs); +// +// ClusterAwareWriterFailoverHandler spyHandler = spy(handler); +// doReturn(mockPluginService).when(spyHandler).getNewPluginService(); +// return spyHandler; +// } +// +// /** +// * Verify that writer failover handler can re-connect to a current writer node. +// * +// *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. +// * TaskA: successfully re-connect to initial writer; return new connection. +// * TaskB: successfully connect to readerA and then new writer, but it takes more time than taskA. +// * Expected test result: new connection by taskA. +// */ +// @Test +// public void testReconnectToWriter_SlowReaderA() throws SQLException { +// when(mockConnectionService.open(refEq(writer), eq(properties))).thenReturn(mockWriterConnection); +// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); +// when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); +// when(mockPluginService.getAllHosts()).thenReturn(topology).thenReturn(newTopology); +// +// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(5000); +// return new ReaderFailoverResult(mockReaderAConnection, readerA, true); +// }); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertTrue(result.isConnected()); +// assertFalse(result.isNewHost()); +// assertSame(result.getNewConnection(), mockWriterConnection); +// assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); +// } +// +// /** +// * Verify that writer failover handler can re-connect to a current writer node. +// * +// *

Topology: no changes. +// * TaskA: successfully re-connect to writer; return new connection. +// * TaskB: successfully connect to readerA and retrieve topology, but latest writer is not new (defer to taskA). +// * Expected test result: new connection by taskA. +// */ +// @Test +// public void testReconnectToWriter_taskBDefers() throws SQLException { +// when(mockConnectionService.open(refEq(writer), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(5000); +// return mockWriterConnection; +// }); +// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); +// +// when(mockPluginService.getAllHosts()).thenReturn(topology); +// +// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) +// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 2000, 2000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertTrue(result.isConnected()); +// assertFalse(result.isNewHost()); +// assertSame(result.getNewConnection(), mockWriterConnection); +// assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); +// } +// +// /** +// * Verify that writer failover handler can re-connect to a new writer node. +// * +// *

Topology: changes to [new-writer, reader-A, reader-B] for taskB, taskA sees no changes. +// * taskA: successfully re-connect to writer; return connection to initial writer, but it takes more +// * time than taskB. +// * TaskB: successfully connect to readerA and then to new-writer. +// * Expected test result: new connection to writer by taskB. +// */ +// @Test +// public void testConnectToReaderA_SlowWriter() throws SQLException { +// when(mockConnectionService.open(refEq(writer), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(5000); +// return mockWriterConnection; +// }); +// when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); +// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); +// when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); +// +// when(mockPluginService.getAllHosts()).thenReturn(newTopology); +// +// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) +// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertTrue(result.isConnected()); +// assertTrue(result.isNewHost()); +// assertSame(result.getNewConnection(), mockNewWriterConnection); +// assertEquals(3, result.getTopology().size()); +// assertEquals("new-writer-host", result.getTopology().get(0).getHost()); +// assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); +// } +// +// /** +// * Verify that writer failover handler can re-connect to a new writer node. +// * +// *

Topology: changes to [new-writer, initial-writer, reader-A, reader-B]. +// * TaskA: successfully reconnect, but initial-writer is now a reader (defer to taskB). +// * TaskB: successfully connect to readerA and then to new-writer. +// * Expected test result: new connection to writer by taskB. +// */ +// @Test +// public void testConnectToReaderA_taskADefers() throws SQLException { +// when(mockConnectionService.open(writer, properties)).thenReturn(mockConnection); +// when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); +// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); +// when(mockConnectionService.open(refEq(newWriterHost), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(5000); +// return mockNewWriterConnection; +// }); +// +// final List newTopology = Arrays.asList(newWriterHost, writer, readerA, readerB); +// when(mockPluginService.getAllHosts()).thenReturn(newTopology); +// +// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) +// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertTrue(result.isConnected()); +// assertTrue(result.isNewHost()); +// assertSame(result.getNewConnection(), mockNewWriterConnection); +// assertEquals(4, result.getTopology().size()); +// assertEquals("new-writer-host", result.getTopology().get(0).getHost()); +// +// verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); +// assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); +// } +// +// /** +// * Verify that writer failover handler fails to re-connect to any writer node. +// * +// *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. +// * TaskA: fail to re-connect to writer due to failover timeout. +// * TaskB: successfully connect to readerA and then fail to connect to writer due to failover timeout. +// * Expected test result: no connection. +// */ +// @Test +// public void testFailedToConnect_failoverTimeout() throws SQLException { +// when(mockConnectionService.open(refEq(writer), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(30000); +// return mockWriterConnection; +// }); +// when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); +// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); +// when(mockConnectionService.open(refEq(newWriterHost), eq(properties))) +// .thenAnswer( +// (Answer) +// invocation -> { +// Thread.sleep(30000); +// return mockNewWriterConnection; +// }); +// when(mockPluginService.getAllHosts()).thenReturn(newTopology); +// +// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) +// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); +// +// final long startTimeNano = System.nanoTime(); +// final WriterFailoverResult result = target.failover(topology); +// final long durationNano = System.nanoTime() - startTimeNano; +// +// assertFalse(result.isConnected()); +// assertFalse(result.isNewHost()); +// +// verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); +// +// // 5s is a max allowed failover timeout; add 1s for inaccurate measurements +// assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); +// } +// +// /** +// * Verify that writer failover handler fails to re-connect to any writer node. +// * +// *

Topology: changes to [new-writer, reader-A, reader-B] for taskB. +// * TaskA: fail to re-connect to writer due to exception. +// * TaskB: successfully connect to readerA and then fail to connect to writer due to exception. +// * Expected test result: no connection. +// */ +// @Test +// public void testFailedToConnect_taskAException_taskBWriterException() throws SQLException { +// final SQLException exception = new SQLException("exception", "08S01", null); +// when(mockConnectionService.open(refEq(writer), eq(properties))).thenThrow(exception); +// when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); +// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); +// when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenThrow(exception); +// when(mockPluginService.isNetworkException(eq(exception), any())).thenReturn(true); +// +// when(mockPluginService.getAllHosts()).thenReturn(newTopology); +// +// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) +// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); +// +// when(mockPluginService.getDialect()).thenReturn(mockDialect); +// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); +// +// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); +// final WriterFailoverResult result = target.failover(topology); +// +// assertFalse(result.isConnected()); +// assertFalse(result.isNewHost()); +// +// assertEquals(HostAvailability.NOT_AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); +// } +// } diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java index 2be3e2858..0ecddc7d0 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java @@ -1,447 +1,447 @@ -/* - * 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.failover; - -import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyString; -import static org.mockito.ArgumentMatchers.eq; -import static org.mockito.Mockito.atLeastOnce; -import static org.mockito.Mockito.doNothing; -import static org.mockito.Mockito.doReturn; -import static org.mockito.Mockito.doThrow; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.spy; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -import java.sql.Connection; -import java.sql.ResultSet; -import java.sql.SQLException; -import java.util.Arrays; -import java.util.Collections; -import java.util.EnumSet; -import java.util.HashMap; -import java.util.HashSet; -import java.util.List; -import java.util.Map; -import java.util.Properties; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.ValueSource; -import org.mockito.Mock; -import org.mockito.MockitoAnnotations; -import software.amazon.jdbc.HostListProviderService; -import software.amazon.jdbc.HostRole; -import software.amazon.jdbc.HostSpec; -import software.amazon.jdbc.HostSpecBuilder; -import software.amazon.jdbc.JdbcCallable; -import software.amazon.jdbc.NodeChangeOptions; -import software.amazon.jdbc.PluginService; -import software.amazon.jdbc.hostavailability.HostAvailability; -import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; -import software.amazon.jdbc.hostlistprovider.AuroraHostListProvider; -import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; -import software.amazon.jdbc.util.FullServicesContainer; -import software.amazon.jdbc.util.RdsUrlType; -import software.amazon.jdbc.util.SqlState; -import software.amazon.jdbc.util.connection.ConnectionService; -import software.amazon.jdbc.util.telemetry.GaugeCallable; -import software.amazon.jdbc.util.telemetry.TelemetryContext; -import software.amazon.jdbc.util.telemetry.TelemetryCounter; -import software.amazon.jdbc.util.telemetry.TelemetryFactory; -import software.amazon.jdbc.util.telemetry.TelemetryGauge; - -class FailoverConnectionPluginTest { - - private static final Class MONITOR_METHOD_INVOKE_ON = Connection.class; - private static final String MONITOR_METHOD_NAME = "Connection.executeQuery"; - private static final Object[] EMPTY_ARGS = {}; - private final List defaultHosts = Arrays.asList( - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("writer").port(1234).role(HostRole.WRITER).build(), - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) - .host("reader1").port(1234).role(HostRole.READER).build()); - - @Mock FullServicesContainer mockContainer; - @Mock ConnectionService mockConnectionService; - @Mock PluginService mockPluginService; - @Mock Connection mockConnection; - @Mock HostSpec mockHostSpec; - @Mock HostListProviderService mockHostListProviderService; - @Mock AuroraHostListProvider mockHostListProvider; - @Mock JdbcCallable mockInitHostProviderFunc; - @Mock ReaderFailoverHandler mockReaderFailoverHandler; - @Mock WriterFailoverHandler mockWriterFailoverHandler; - @Mock ReaderFailoverResult mockReaderResult; - @Mock WriterFailoverResult mockWriterResult; - @Mock JdbcCallable mockSqlFunction; - @Mock private TelemetryFactory mockTelemetryFactory; - @Mock TelemetryContext mockTelemetryContext; - @Mock TelemetryCounter mockTelemetryCounter; - @Mock TelemetryGauge mockTelemetryGauge; - @Mock TargetDriverDialect mockTargetDriverDialect; - - - private final Properties properties = new Properties(); - private FailoverConnectionPlugin spyPlugin; - private AutoCloseable closeable; - - @AfterEach - void cleanUp() throws Exception { - closeable.close(); - } - - @BeforeEach - void init() throws SQLException { - closeable = MockitoAnnotations.openMocks(this); - - when(mockContainer.getPluginService()).thenReturn(mockPluginService); - when(mockPluginService.getHostListProvider()).thenReturn(mockHostListProvider); - when(mockHostListProvider.getRdsUrlType()).thenReturn(RdsUrlType.RDS_WRITER_CLUSTER); - when(mockPluginService.getCurrentConnection()).thenReturn(mockConnection); - when(mockPluginService.getCurrentHostSpec()).thenReturn(mockHostSpec); - when(mockPluginService.connect(any(HostSpec.class), eq(properties))).thenReturn(mockConnection); - when(mockPluginService.getTelemetryFactory()).thenReturn(mockTelemetryFactory); - when(mockPluginService.getHosts()).thenReturn(defaultHosts); - when(mockPluginService.getAllHosts()).thenReturn(defaultHosts); - when(mockReaderFailoverHandler.failover(any(), any())).thenReturn(mockReaderResult); - when(mockWriterFailoverHandler.failover(any())).thenReturn(mockWriterResult); - when(mockWriterResult.isConnected()).thenReturn(true); - when(mockWriterResult.getTopology()).thenReturn(defaultHosts); - when(mockReaderResult.isConnected()).thenReturn(true); - - when(mockPluginService.getTelemetryFactory()).thenReturn(mockTelemetryFactory); - when(mockTelemetryFactory.openTelemetryContext(anyString(), any())).thenReturn(mockTelemetryContext); - when(mockTelemetryFactory.openTelemetryContext(eq(null), any())).thenReturn(mockTelemetryContext); - when(mockTelemetryFactory.createCounter(anyString())).thenReturn(mockTelemetryCounter); - // noinspection unchecked - when(mockTelemetryFactory.createGauge(anyString(), any(GaugeCallable.class))).thenReturn(mockTelemetryGauge); - - when(mockPluginService.getTargetDriverDialect()).thenReturn(mockTargetDriverDialect); - when(mockTargetDriverDialect.getNetworkBoundMethodNames(any())).thenReturn(new HashSet<>()); - - properties.clear(); - } - - @Test - void test_notifyNodeListChanged_withFailoverDisabled() throws SQLException { - properties.setProperty(FailoverConnectionPlugin.ENABLE_CLUSTER_AWARE_FAILOVER.name, "false"); - final Map> changes = new HashMap<>(); - - initializePlugin(); - spyPlugin.notifyNodeListChanged(changes); - - verify(mockPluginService, never()).getCurrentHostSpec(); - verify(mockHostSpec, never()).getAliases(); - } - - @Test - void test_notifyNodeListChanged_withValidConnectionNotInTopology() throws SQLException { - final Map> changes = new HashMap<>(); - changes.put("cluster-host/", EnumSet.of(NodeChangeOptions.NODE_DELETED)); - changes.put("instance/", EnumSet.of(NodeChangeOptions.NODE_ADDED)); - - initializePlugin(); - spyPlugin.notifyNodeListChanged(changes); - - when(mockHostSpec.getUrl()).thenReturn("cluster-url/"); - when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Collections.singletonList("instance"))); - - verify(mockPluginService).getCurrentHostSpec(); - verify(mockHostSpec, never()).getAliases(); - } - - @Test - void test_updateTopology() throws SQLException { - initializePlugin(); - - // Test updateTopology with failover disabled - spyPlugin.setRdsUrlType(RdsUrlType.RDS_PROXY); - spyPlugin.updateTopology(false); - verify(mockPluginService, never()).forceRefreshHostList(); - verify(mockPluginService, never()).refreshHostList(); - - // Test updateTopology with no connection - when(mockPluginService.getCurrentHostSpec()).thenReturn(null); - spyPlugin.updateTopology(false); - verify(mockPluginService, never()).forceRefreshHostList(); - verify(mockPluginService, never()).refreshHostList(); - - // Test updateTopology with closed connection - when(mockConnection.isClosed()).thenReturn(true); - spyPlugin.updateTopology(false); - verify(mockPluginService, never()).forceRefreshHostList(); - verify(mockPluginService, never()).refreshHostList(); - } - - @ParameterizedTest - @ValueSource(booleans = {true, false}) - void test_updateTopology_withForceUpdate(final boolean forceUpdate) throws SQLException { - - when(mockPluginService.getAllHosts()).thenReturn(Collections.singletonList( - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("host").build())); - when(mockPluginService.getHosts()).thenReturn(Collections.singletonList( - new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("host").build())); - when(mockConnection.isClosed()).thenReturn(false); - initializePlugin(); - spyPlugin.setRdsUrlType(RdsUrlType.RDS_INSTANCE); - - spyPlugin.updateTopology(forceUpdate); - if (forceUpdate) { - verify(mockPluginService, atLeastOnce()).forceRefreshHostList(); - } else { - verify(mockPluginService, atLeastOnce()).refreshHostList(); - } - } - - @Test - void test_failover_failoverWriter() throws SQLException { - when(mockPluginService.isInTransaction()).thenReturn(true); - - initializePlugin(); - doThrow(FailoverSuccessSQLException.class).when(spyPlugin).failoverWriter(); - spyPlugin.failoverMode = FailoverMode.STRICT_WRITER; - - assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failover(mockHostSpec)); - verify(spyPlugin).failoverWriter(); - } - - @Test - void test_failover_failoverReader() throws SQLException { - when(mockPluginService.isInTransaction()).thenReturn(false); - - initializePlugin(); - doThrow(FailoverSuccessSQLException.class).when(spyPlugin).failoverReader(eq(mockHostSpec)); - spyPlugin.failoverMode = FailoverMode.READER_OR_WRITER; - - assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failover(mockHostSpec)); - verify(spyPlugin).failoverReader(eq(mockHostSpec)); - } - - @Test - void test_failoverReader_withValidFailedHostSpec_successFailover() throws SQLException { - when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); - when(mockHostSpec.getRawAvailability()).thenReturn(HostAvailability.AVAILABLE); - when(mockReaderResult.isConnected()).thenReturn(true); - when(mockReaderResult.getConnection()).thenReturn(mockConnection); - when(mockReaderResult.getHost()).thenReturn(defaultHosts.get(1)); - - initializePlugin(); - spyPlugin.initHostProvider( - mockHostListProviderService, - mockInitHostProviderFunc, - (connectionService) -> mockReaderFailoverHandler, - (connectionService) -> mockWriterFailoverHandler); - - final FailoverConnectionPlugin spyPlugin = spy(this.spyPlugin); - doNothing().when(spyPlugin).updateTopology(true); - - assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failoverReader(mockHostSpec)); - - verify(mockReaderFailoverHandler).failover(eq(defaultHosts), eq(mockHostSpec)); - verify(mockPluginService).setCurrentConnection(eq(mockConnection), eq(defaultHosts.get(1))); - } - - @Test - void test_failoverReader_withNoFailedHostSpec_withException() throws SQLException { - final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") - .build(); - final List hosts = Collections.singletonList(hostSpec); - - when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); - when(mockHostSpec.getAvailability()).thenReturn(HostAvailability.AVAILABLE); - when(mockPluginService.getAllHosts()).thenReturn(hosts); - when(mockPluginService.getHosts()).thenReturn(hosts); - when(mockReaderResult.getException()).thenReturn(new SQLException()); - when(mockReaderResult.getHost()).thenReturn(hostSpec); - - initializePlugin(); - spyPlugin.initHostProvider( - mockHostListProviderService, - mockInitHostProviderFunc, - (connectionService) -> mockReaderFailoverHandler, - (connectionService) -> mockWriterFailoverHandler); - - assertThrows(SQLException.class, () -> spyPlugin.failoverReader(null)); - verify(mockReaderFailoverHandler).failover(eq(hosts), eq(null)); - } - - @Test - void test_failoverWriter_failedFailover_throwsException() throws SQLException { - final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") - .build(); - final List hosts = Collections.singletonList(hostSpec); - - when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); - when(mockPluginService.getAllHosts()).thenReturn(hosts); - when(mockPluginService.getHosts()).thenReturn(hosts); - when(mockWriterResult.getException()).thenReturn(new SQLException()); - - initializePlugin(); - spyPlugin.initHostProvider( - mockHostListProviderService, - mockInitHostProviderFunc, - (connectionService) -> mockReaderFailoverHandler, - (connectionService) -> mockWriterFailoverHandler); - - assertThrows(SQLException.class, () -> spyPlugin.failoverWriter()); - verify(mockWriterFailoverHandler).failover(eq(hosts)); - } - - @Test - void test_failoverWriter_failedFailover_withNoResult() throws SQLException { - final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") - .build(); - final List hosts = Collections.singletonList(hostSpec); - - when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); - when(mockPluginService.getAllHosts()).thenReturn(hosts); - when(mockPluginService.getHosts()).thenReturn(hosts); - when(mockWriterResult.isConnected()).thenReturn(false); - - initializePlugin(); - spyPlugin.initHostProvider( - mockHostListProviderService, - mockInitHostProviderFunc, - (connectionService) -> mockReaderFailoverHandler, - (connectionService) -> mockWriterFailoverHandler); - - final SQLException exception = assertThrows(SQLException.class, () -> spyPlugin.failoverWriter()); - assertEquals(SqlState.CONNECTION_UNABLE_TO_CONNECT.getState(), exception.getSQLState()); - - verify(mockWriterFailoverHandler).failover(eq(hosts)); - verify(mockWriterResult, never()).getNewConnection(); - verify(mockWriterResult, never()).getTopology(); - } - - @Test - void test_failoverWriter_successFailover() throws SQLException { - when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); - - initializePlugin(); - spyPlugin.initHostProvider( - mockHostListProviderService, - mockInitHostProviderFunc, - (connectionService) -> mockReaderFailoverHandler, - (connectionService) -> mockWriterFailoverHandler); - - final SQLException exception = assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failoverWriter()); - assertEquals(SqlState.COMMUNICATION_LINK_CHANGED.getState(), exception.getSQLState()); - - verify(mockWriterFailoverHandler).failover(eq(defaultHosts)); - } - - @Test - void test_invalidCurrentConnection_withNoConnection() throws SQLException { - when(mockPluginService.getCurrentConnection()).thenReturn(null); - initializePlugin(); - spyPlugin.invalidateCurrentConnection(); - - verify(mockPluginService, never()).getCurrentHostSpec(); - } - - @Test - void test_invalidateCurrentConnection_inTransaction() throws SQLException { - when(mockPluginService.isInTransaction()).thenReturn(true); - when(mockHostSpec.getHost()).thenReturn("host"); - when(mockHostSpec.getPort()).thenReturn(123); - when(mockHostSpec.getRole()).thenReturn(HostRole.READER); - - initializePlugin(); - spyPlugin.invalidateCurrentConnection(); - verify(mockConnection).rollback(); - - // Assert SQL exceptions thrown during rollback do not get propagated. - doThrow(new SQLException()).when(mockConnection).rollback(); - assertDoesNotThrow(() -> spyPlugin.invalidateCurrentConnection()); - } - - @Test - void test_invalidateCurrentConnection_notInTransaction() throws SQLException { - when(mockPluginService.isInTransaction()).thenReturn(false); - when(mockHostSpec.getHost()).thenReturn("host"); - when(mockHostSpec.getPort()).thenReturn(123); - when(mockHostSpec.getRole()).thenReturn(HostRole.READER); - - initializePlugin(); - spyPlugin.invalidateCurrentConnection(); - - verify(mockPluginService).isInTransaction(); - } - - @Test - void test_invalidateCurrentConnection_withOpenConnection() throws SQLException { - when(mockPluginService.isInTransaction()).thenReturn(false); - when(mockConnection.isClosed()).thenReturn(false); - when(mockHostSpec.getHost()).thenReturn("host"); - when(mockHostSpec.getPort()).thenReturn(123); - when(mockHostSpec.getRole()).thenReturn(HostRole.READER); - - initializePlugin(); - spyPlugin.invalidateCurrentConnection(); - - doThrow(new SQLException()).when(mockConnection).close(); - assertDoesNotThrow(() -> spyPlugin.invalidateCurrentConnection()); - - verify(mockConnection, times(2)).isClosed(); - verify(mockConnection, times(2)).close(); - } - - @Test - void test_execute_withFailoverDisabled() throws SQLException { - properties.setProperty(FailoverConnectionPlugin.ENABLE_CLUSTER_AWARE_FAILOVER.name, "false"); - initializePlugin(); - - spyPlugin.execute( - ResultSet.class, - SQLException.class, - MONITOR_METHOD_INVOKE_ON, - MONITOR_METHOD_NAME, - mockSqlFunction, - EMPTY_ARGS); - - verify(mockSqlFunction).call(); - verify(mockHostListProvider, never()).getRdsUrlType(); - } - - @Test - void test_execute_withDirectExecute() throws SQLException { - initializePlugin(); - spyPlugin.execute( - ResultSet.class, - SQLException.class, - MONITOR_METHOD_INVOKE_ON, - "close", - mockSqlFunction, - EMPTY_ARGS); - verify(mockSqlFunction).call(); - verify(mockHostListProvider, never()).getRdsUrlType(); - } - - private void initializePlugin() throws SQLException { - spyPlugin = spy(new FailoverConnectionPlugin(mockContainer, properties)); - spyPlugin.setWriterFailoverHandler(mockWriterFailoverHandler); - spyPlugin.setReaderFailoverHandler(mockReaderFailoverHandler); - doReturn(mockConnectionService).when(spyPlugin).getConnectionService(); - } -} +// /* +// * 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.failover; +// +// import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +// import static org.junit.jupiter.api.Assertions.assertEquals; +// import static org.junit.jupiter.api.Assertions.assertThrows; +// import static org.mockito.ArgumentMatchers.any; +// import static org.mockito.ArgumentMatchers.anyString; +// import static org.mockito.ArgumentMatchers.eq; +// import static org.mockito.Mockito.atLeastOnce; +// import static org.mockito.Mockito.doNothing; +// import static org.mockito.Mockito.doReturn; +// import static org.mockito.Mockito.doThrow; +// import static org.mockito.Mockito.never; +// import static org.mockito.Mockito.spy; +// import static org.mockito.Mockito.times; +// import static org.mockito.Mockito.verify; +// import static org.mockito.Mockito.when; +// +// import java.sql.Connection; +// import java.sql.ResultSet; +// import java.sql.SQLException; +// import java.util.Arrays; +// import java.util.Collections; +// import java.util.EnumSet; +// import java.util.HashMap; +// import java.util.HashSet; +// import java.util.List; +// import java.util.Map; +// import java.util.Properties; +// import org.junit.jupiter.api.AfterEach; +// import org.junit.jupiter.api.BeforeEach; +// import org.junit.jupiter.api.Test; +// import org.junit.jupiter.params.ParameterizedTest; +// import org.junit.jupiter.params.provider.ValueSource; +// import org.mockito.Mock; +// import org.mockito.MockitoAnnotations; +// import software.amazon.jdbc.HostListProviderService; +// import software.amazon.jdbc.HostRole; +// import software.amazon.jdbc.HostSpec; +// import software.amazon.jdbc.HostSpecBuilder; +// import software.amazon.jdbc.JdbcCallable; +// import software.amazon.jdbc.NodeChangeOptions; +// import software.amazon.jdbc.PluginService; +// import software.amazon.jdbc.hostavailability.HostAvailability; +// import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; +// import software.amazon.jdbc.hostlistprovider.AuroraHostListProvider; +// import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; +// import software.amazon.jdbc.util.FullServicesContainer; +// import software.amazon.jdbc.util.RdsUrlType; +// import software.amazon.jdbc.util.SqlState; +// import software.amazon.jdbc.util.connection.ConnectionService; +// import software.amazon.jdbc.util.telemetry.GaugeCallable; +// import software.amazon.jdbc.util.telemetry.TelemetryContext; +// import software.amazon.jdbc.util.telemetry.TelemetryCounter; +// import software.amazon.jdbc.util.telemetry.TelemetryFactory; +// import software.amazon.jdbc.util.telemetry.TelemetryGauge; +// +// class FailoverConnectionPluginTest { +// +// private static final Class MONITOR_METHOD_INVOKE_ON = Connection.class; +// private static final String MONITOR_METHOD_NAME = "Connection.executeQuery"; +// private static final Object[] EMPTY_ARGS = {}; +// private final List defaultHosts = Arrays.asList( +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("writer").port(1234).role(HostRole.WRITER).build(), +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) +// .host("reader1").port(1234).role(HostRole.READER).build()); +// +// @Mock FullServicesContainer mockContainer; +// @Mock ConnectionService mockConnectionService; +// @Mock PluginService mockPluginService; +// @Mock Connection mockConnection; +// @Mock HostSpec mockHostSpec; +// @Mock HostListProviderService mockHostListProviderService; +// @Mock AuroraHostListProvider mockHostListProvider; +// @Mock JdbcCallable mockInitHostProviderFunc; +// @Mock ReaderFailoverHandler mockReaderFailoverHandler; +// @Mock WriterFailoverHandler mockWriterFailoverHandler; +// @Mock ReaderFailoverResult mockReaderResult; +// @Mock WriterFailoverResult mockWriterResult; +// @Mock JdbcCallable mockSqlFunction; +// @Mock private TelemetryFactory mockTelemetryFactory; +// @Mock TelemetryContext mockTelemetryContext; +// @Mock TelemetryCounter mockTelemetryCounter; +// @Mock TelemetryGauge mockTelemetryGauge; +// @Mock TargetDriverDialect mockTargetDriverDialect; +// +// +// private final Properties properties = new Properties(); +// private FailoverConnectionPlugin spyPlugin; +// private AutoCloseable closeable; +// +// @AfterEach +// void cleanUp() throws Exception { +// closeable.close(); +// } +// +// @BeforeEach +// void init() throws SQLException { +// closeable = MockitoAnnotations.openMocks(this); +// +// when(mockContainer.getPluginService()).thenReturn(mockPluginService); +// when(mockPluginService.getHostListProvider()).thenReturn(mockHostListProvider); +// when(mockHostListProvider.getRdsUrlType()).thenReturn(RdsUrlType.RDS_WRITER_CLUSTER); +// when(mockPluginService.getCurrentConnection()).thenReturn(mockConnection); +// when(mockPluginService.getCurrentHostSpec()).thenReturn(mockHostSpec); +// when(mockPluginService.connect(any(HostSpec.class), eq(properties))).thenReturn(mockConnection); +// when(mockPluginService.getTelemetryFactory()).thenReturn(mockTelemetryFactory); +// when(mockPluginService.getHosts()).thenReturn(defaultHosts); +// when(mockPluginService.getAllHosts()).thenReturn(defaultHosts); +// when(mockReaderFailoverHandler.failover(any(), any())).thenReturn(mockReaderResult); +// when(mockWriterFailoverHandler.failover(any())).thenReturn(mockWriterResult); +// when(mockWriterResult.isConnected()).thenReturn(true); +// when(mockWriterResult.getTopology()).thenReturn(defaultHosts); +// when(mockReaderResult.isConnected()).thenReturn(true); +// +// when(mockPluginService.getTelemetryFactory()).thenReturn(mockTelemetryFactory); +// when(mockTelemetryFactory.openTelemetryContext(anyString(), any())).thenReturn(mockTelemetryContext); +// when(mockTelemetryFactory.openTelemetryContext(eq(null), any())).thenReturn(mockTelemetryContext); +// when(mockTelemetryFactory.createCounter(anyString())).thenReturn(mockTelemetryCounter); +// // noinspection unchecked +// when(mockTelemetryFactory.createGauge(anyString(), any(GaugeCallable.class))).thenReturn(mockTelemetryGauge); +// +// when(mockPluginService.getTargetDriverDialect()).thenReturn(mockTargetDriverDialect); +// when(mockTargetDriverDialect.getNetworkBoundMethodNames(any())).thenReturn(new HashSet<>()); +// +// properties.clear(); +// } +// +// @Test +// void test_notifyNodeListChanged_withFailoverDisabled() throws SQLException { +// properties.setProperty(FailoverConnectionPlugin.ENABLE_CLUSTER_AWARE_FAILOVER.name, "false"); +// final Map> changes = new HashMap<>(); +// +// initializePlugin(); +// spyPlugin.notifyNodeListChanged(changes); +// +// verify(mockPluginService, never()).getCurrentHostSpec(); +// verify(mockHostSpec, never()).getAliases(); +// } +// +// @Test +// void test_notifyNodeListChanged_withValidConnectionNotInTopology() throws SQLException { +// final Map> changes = new HashMap<>(); +// changes.put("cluster-host/", EnumSet.of(NodeChangeOptions.NODE_DELETED)); +// changes.put("instance/", EnumSet.of(NodeChangeOptions.NODE_ADDED)); +// +// initializePlugin(); +// spyPlugin.notifyNodeListChanged(changes); +// +// when(mockHostSpec.getUrl()).thenReturn("cluster-url/"); +// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Collections.singletonList("instance"))); +// +// verify(mockPluginService).getCurrentHostSpec(); +// verify(mockHostSpec, never()).getAliases(); +// } +// +// @Test +// void test_updateTopology() throws SQLException { +// initializePlugin(); +// +// // Test updateTopology with failover disabled +// spyPlugin.setRdsUrlType(RdsUrlType.RDS_PROXY); +// spyPlugin.updateTopology(false); +// verify(mockPluginService, never()).forceRefreshHostList(); +// verify(mockPluginService, never()).refreshHostList(); +// +// // Test updateTopology with no connection +// when(mockPluginService.getCurrentHostSpec()).thenReturn(null); +// spyPlugin.updateTopology(false); +// verify(mockPluginService, never()).forceRefreshHostList(); +// verify(mockPluginService, never()).refreshHostList(); +// +// // Test updateTopology with closed connection +// when(mockConnection.isClosed()).thenReturn(true); +// spyPlugin.updateTopology(false); +// verify(mockPluginService, never()).forceRefreshHostList(); +// verify(mockPluginService, never()).refreshHostList(); +// } +// +// @ParameterizedTest +// @ValueSource(booleans = {true, false}) +// void test_updateTopology_withForceUpdate(final boolean forceUpdate) throws SQLException { +// +// when(mockPluginService.getAllHosts()).thenReturn(Collections.singletonList( +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("host").build())); +// when(mockPluginService.getHosts()).thenReturn(Collections.singletonList( +// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("host").build())); +// when(mockConnection.isClosed()).thenReturn(false); +// initializePlugin(); +// spyPlugin.setRdsUrlType(RdsUrlType.RDS_INSTANCE); +// +// spyPlugin.updateTopology(forceUpdate); +// if (forceUpdate) { +// verify(mockPluginService, atLeastOnce()).forceRefreshHostList(); +// } else { +// verify(mockPluginService, atLeastOnce()).refreshHostList(); +// } +// } +// +// @Test +// void test_failover_failoverWriter() throws SQLException { +// when(mockPluginService.isInTransaction()).thenReturn(true); +// +// initializePlugin(); +// doThrow(FailoverSuccessSQLException.class).when(spyPlugin).failoverWriter(); +// spyPlugin.failoverMode = FailoverMode.STRICT_WRITER; +// +// assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failover(mockHostSpec)); +// verify(spyPlugin).failoverWriter(); +// } +// +// @Test +// void test_failover_failoverReader() throws SQLException { +// when(mockPluginService.isInTransaction()).thenReturn(false); +// +// initializePlugin(); +// doThrow(FailoverSuccessSQLException.class).when(spyPlugin).failoverReader(eq(mockHostSpec)); +// spyPlugin.failoverMode = FailoverMode.READER_OR_WRITER; +// +// assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failover(mockHostSpec)); +// verify(spyPlugin).failoverReader(eq(mockHostSpec)); +// } +// +// @Test +// void test_failoverReader_withValidFailedHostSpec_successFailover() throws SQLException { +// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); +// when(mockHostSpec.getRawAvailability()).thenReturn(HostAvailability.AVAILABLE); +// when(mockReaderResult.isConnected()).thenReturn(true); +// when(mockReaderResult.getConnection()).thenReturn(mockConnection); +// when(mockReaderResult.getHost()).thenReturn(defaultHosts.get(1)); +// +// initializePlugin(); +// spyPlugin.initHostProvider( +// mockHostListProviderService, +// mockInitHostProviderFunc, +// (connectionService) -> mockReaderFailoverHandler, +// (connectionService) -> mockWriterFailoverHandler); +// +// final FailoverConnectionPlugin spyPlugin = spy(this.spyPlugin); +// doNothing().when(spyPlugin).updateTopology(true); +// +// assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failoverReader(mockHostSpec)); +// +// verify(mockReaderFailoverHandler).failover(eq(defaultHosts), eq(mockHostSpec)); +// verify(mockPluginService).setCurrentConnection(eq(mockConnection), eq(defaultHosts.get(1))); +// } +// +// @Test +// void test_failoverReader_withNoFailedHostSpec_withException() throws SQLException { +// final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") +// .build(); +// final List hosts = Collections.singletonList(hostSpec); +// +// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); +// when(mockHostSpec.getAvailability()).thenReturn(HostAvailability.AVAILABLE); +// when(mockPluginService.getAllHosts()).thenReturn(hosts); +// when(mockPluginService.getHosts()).thenReturn(hosts); +// when(mockReaderResult.getException()).thenReturn(new SQLException()); +// when(mockReaderResult.getHost()).thenReturn(hostSpec); +// +// initializePlugin(); +// spyPlugin.initHostProvider( +// mockHostListProviderService, +// mockInitHostProviderFunc, +// (connectionService) -> mockReaderFailoverHandler, +// (connectionService) -> mockWriterFailoverHandler); +// +// assertThrows(SQLException.class, () -> spyPlugin.failoverReader(null)); +// verify(mockReaderFailoverHandler).failover(eq(hosts), eq(null)); +// } +// +// @Test +// void test_failoverWriter_failedFailover_throwsException() throws SQLException { +// final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") +// .build(); +// final List hosts = Collections.singletonList(hostSpec); +// +// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); +// when(mockPluginService.getAllHosts()).thenReturn(hosts); +// when(mockPluginService.getHosts()).thenReturn(hosts); +// when(mockWriterResult.getException()).thenReturn(new SQLException()); +// +// initializePlugin(); +// spyPlugin.initHostProvider( +// mockHostListProviderService, +// mockInitHostProviderFunc, +// (connectionService) -> mockReaderFailoverHandler, +// (connectionService) -> mockWriterFailoverHandler); +// +// assertThrows(SQLException.class, () -> spyPlugin.failoverWriter()); +// verify(mockWriterFailoverHandler).failover(eq(hosts)); +// } +// +// @Test +// void test_failoverWriter_failedFailover_withNoResult() throws SQLException { +// final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") +// .build(); +// final List hosts = Collections.singletonList(hostSpec); +// +// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); +// when(mockPluginService.getAllHosts()).thenReturn(hosts); +// when(mockPluginService.getHosts()).thenReturn(hosts); +// when(mockWriterResult.isConnected()).thenReturn(false); +// +// initializePlugin(); +// spyPlugin.initHostProvider( +// mockHostListProviderService, +// mockInitHostProviderFunc, +// (connectionService) -> mockReaderFailoverHandler, +// (connectionService) -> mockWriterFailoverHandler); +// +// final SQLException exception = assertThrows(SQLException.class, () -> spyPlugin.failoverWriter()); +// assertEquals(SqlState.CONNECTION_UNABLE_TO_CONNECT.getState(), exception.getSQLState()); +// +// verify(mockWriterFailoverHandler).failover(eq(hosts)); +// verify(mockWriterResult, never()).getNewConnection(); +// verify(mockWriterResult, never()).getTopology(); +// } +// +// @Test +// void test_failoverWriter_successFailover() throws SQLException { +// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); +// +// initializePlugin(); +// spyPlugin.initHostProvider( +// mockHostListProviderService, +// mockInitHostProviderFunc, +// (connectionService) -> mockReaderFailoverHandler, +// (connectionService) -> mockWriterFailoverHandler); +// +// final SQLException exception = assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failoverWriter()); +// assertEquals(SqlState.COMMUNICATION_LINK_CHANGED.getState(), exception.getSQLState()); +// +// verify(mockWriterFailoverHandler).failover(eq(defaultHosts)); +// } +// +// @Test +// void test_invalidCurrentConnection_withNoConnection() throws SQLException { +// when(mockPluginService.getCurrentConnection()).thenReturn(null); +// initializePlugin(); +// spyPlugin.invalidateCurrentConnection(); +// +// verify(mockPluginService, never()).getCurrentHostSpec(); +// } +// +// @Test +// void test_invalidateCurrentConnection_inTransaction() throws SQLException { +// when(mockPluginService.isInTransaction()).thenReturn(true); +// when(mockHostSpec.getHost()).thenReturn("host"); +// when(mockHostSpec.getPort()).thenReturn(123); +// when(mockHostSpec.getRole()).thenReturn(HostRole.READER); +// +// initializePlugin(); +// spyPlugin.invalidateCurrentConnection(); +// verify(mockConnection).rollback(); +// +// // Assert SQL exceptions thrown during rollback do not get propagated. +// doThrow(new SQLException()).when(mockConnection).rollback(); +// assertDoesNotThrow(() -> spyPlugin.invalidateCurrentConnection()); +// } +// +// @Test +// void test_invalidateCurrentConnection_notInTransaction() throws SQLException { +// when(mockPluginService.isInTransaction()).thenReturn(false); +// when(mockHostSpec.getHost()).thenReturn("host"); +// when(mockHostSpec.getPort()).thenReturn(123); +// when(mockHostSpec.getRole()).thenReturn(HostRole.READER); +// +// initializePlugin(); +// spyPlugin.invalidateCurrentConnection(); +// +// verify(mockPluginService).isInTransaction(); +// } +// +// @Test +// void test_invalidateCurrentConnection_withOpenConnection() throws SQLException { +// when(mockPluginService.isInTransaction()).thenReturn(false); +// when(mockConnection.isClosed()).thenReturn(false); +// when(mockHostSpec.getHost()).thenReturn("host"); +// when(mockHostSpec.getPort()).thenReturn(123); +// when(mockHostSpec.getRole()).thenReturn(HostRole.READER); +// +// initializePlugin(); +// spyPlugin.invalidateCurrentConnection(); +// +// doThrow(new SQLException()).when(mockConnection).close(); +// assertDoesNotThrow(() -> spyPlugin.invalidateCurrentConnection()); +// +// verify(mockConnection, times(2)).isClosed(); +// verify(mockConnection, times(2)).close(); +// } +// +// @Test +// void test_execute_withFailoverDisabled() throws SQLException { +// properties.setProperty(FailoverConnectionPlugin.ENABLE_CLUSTER_AWARE_FAILOVER.name, "false"); +// initializePlugin(); +// +// spyPlugin.execute( +// ResultSet.class, +// SQLException.class, +// MONITOR_METHOD_INVOKE_ON, +// MONITOR_METHOD_NAME, +// mockSqlFunction, +// EMPTY_ARGS); +// +// verify(mockSqlFunction).call(); +// verify(mockHostListProvider, never()).getRdsUrlType(); +// } +// +// @Test +// void test_execute_withDirectExecute() throws SQLException { +// initializePlugin(); +// spyPlugin.execute( +// ResultSet.class, +// SQLException.class, +// MONITOR_METHOD_INVOKE_ON, +// "close", +// mockSqlFunction, +// EMPTY_ARGS); +// verify(mockSqlFunction).call(); +// verify(mockHostListProvider, never()).getRdsUrlType(); +// } +// +// private void initializePlugin() throws SQLException { +// spyPlugin = spy(new FailoverConnectionPlugin(mockContainer, properties)); +// spyPlugin.setWriterFailoverHandler(mockWriterFailoverHandler); +// spyPlugin.setReaderFailoverHandler(mockReaderFailoverHandler); +// doReturn(mockConnectionService).when(spyPlugin).getConnectionService(); +// } +// } From bef4ea85eec64649098afec73862d221436f90e4 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Wed, 3 Sep 2025 15:00:52 -0700 Subject: [PATCH 32/42] wip --- .../MonitoringRdsHostListProvider.java | 8 +++--- .../MonitoringRdsMultiAzHostListProvider.java | 9 +++---- .../MultiAzClusterTopologyMonitorImpl.java | 10 +++---- .../limitless/LimitlessRouterMonitor.java | 20 ++++++-------- .../limitless/LimitlessRouterServiceImpl.java | 7 ++--- .../HostResponseTimeServiceImpl.java | 6 ++--- .../NodeResponseTimeMonitor.java | 6 +---- .../jdbc/util/ServiceContainerUtility.java | 4 +-- .../util/monitoring/MonitorServiceImpl.java | 26 +++++++++---------- 9 files changed, 38 insertions(+), 58 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/MonitoringRdsHostListProvider.java b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/MonitoringRdsHostListProvider.java index 90131af7c..5fd965006 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/MonitoringRdsHostListProvider.java +++ b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/MonitoringRdsHostListProvider.java @@ -85,8 +85,8 @@ protected ClusterTopologyMonitor initMonitor() throws SQLException { ClusterTopologyMonitorImpl.class, this.clusterId, this.servicesContainer.getStorageService(), - this.pluginService.getTelemetryFactory(), - this.pluginService.getDefaultConnectionProvider(), + this.servicesContainer.getTelemetryFactory(), + this.servicesContainer.getDefaultConnectionProvider(), this.originalUrl, this.pluginService.getDriverProtocol(), this.pluginService.getTargetDriverDialect(), @@ -130,8 +130,8 @@ protected void clusterIdChanged(final String oldClusterId) throws SQLException { ClusterTopologyMonitorImpl.class, this.clusterId, this.servicesContainer.getStorageService(), - this.pluginService.getTelemetryFactory(), - this.pluginService.getDefaultConnectionProvider(), + this.servicesContainer.getTelemetryFactory(), + this.servicesContainer.getDefaultConnectionProvider(), this.originalUrl, this.pluginService.getDriverProtocol(), this.pluginService.getTargetDriverDialect(), diff --git a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/MonitoringRdsMultiAzHostListProvider.java b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/MonitoringRdsMultiAzHostListProvider.java index ab1eb9b23..730d15e7a 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/MonitoringRdsMultiAzHostListProvider.java +++ b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/MonitoringRdsMultiAzHostListProvider.java @@ -54,17 +54,16 @@ protected ClusterTopologyMonitor initMonitor() throws SQLException { return this.servicesContainer.getMonitorService().runIfAbsent(MultiAzClusterTopologyMonitorImpl.class, this.clusterId, this.servicesContainer.getStorageService(), - this.pluginService.getTelemetryFactory(), - this.pluginService.getDefaultConnectionProvider(), + this.servicesContainer.getTelemetryFactory(), + this.servicesContainer.getDefaultConnectionProvider(), this.originalUrl, this.pluginService.getDriverProtocol(), this.pluginService.getTargetDriverDialect(), this.pluginService.getDialect(), this.properties, - (connectionService, pluginService) -> new MultiAzClusterTopologyMonitorImpl( + (servicesContainer) -> new MultiAzClusterTopologyMonitorImpl( + servicesContainer, this.clusterId, - this.servicesContainer.getStorageService(), - connectionService, this.initialHostSpec, this.properties, this.hostListProviderService, diff --git a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/MultiAzClusterTopologyMonitorImpl.java b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/MultiAzClusterTopologyMonitorImpl.java index 6e3c3b388..36bab8f90 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/MultiAzClusterTopologyMonitorImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/MultiAzClusterTopologyMonitorImpl.java @@ -26,9 +26,8 @@ import java.util.logging.Logger; import software.amazon.jdbc.HostListProviderService; import software.amazon.jdbc.HostSpec; +import software.amazon.jdbc.util.FullServicesContainer; import software.amazon.jdbc.util.StringUtils; -import software.amazon.jdbc.util.connection.ConnectionService; -import software.amazon.jdbc.util.storage.StorageService; public class MultiAzClusterTopologyMonitorImpl extends ClusterTopologyMonitorImpl { @@ -38,9 +37,8 @@ public class MultiAzClusterTopologyMonitorImpl extends ClusterTopologyMonitorImp protected final String fetchWriterNodeColumnName; public MultiAzClusterTopologyMonitorImpl( + final FullServicesContainer servicesContainer, final String clusterId, - final StorageService storageService, - final ConnectionService connectionService, final HostSpec initialHostSpec, final Properties properties, final HostListProviderService hostListProviderService, @@ -53,12 +51,10 @@ public MultiAzClusterTopologyMonitorImpl( final String fetchWriterNodeQuery, final String fetchWriterNodeColumnName) { super( + servicesContainer, clusterId, - storageService, - connectionService, initialHostSpec, properties, - hostListProviderService, clusterInstanceTemplate, refreshRateNano, highRefreshRateNano, 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 cd983896a..f4075a285 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 @@ -25,11 +25,10 @@ import java.util.logging.Logger; import org.checkerframework.checker.nullness.qual.NonNull; import software.amazon.jdbc.HostSpec; -import software.amazon.jdbc.PluginService; +import software.amazon.jdbc.util.FullServicesContainer; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.PropertyUtils; import software.amazon.jdbc.util.Utils; -import software.amazon.jdbc.util.connection.ConnectionService; import software.amazon.jdbc.util.monitoring.AbstractMonitor; import software.amazon.jdbc.util.storage.StorageService; import software.amazon.jdbc.util.telemetry.TelemetryContext; @@ -45,27 +44,24 @@ public class LimitlessRouterMonitor extends AbstractMonitor { protected static final long TERMINATION_TIMEOUT_SEC = 5; protected final int intervalMs; protected final @NonNull HostSpec hostSpec; + protected final @NonNull FullServicesContainer servicesContainer; protected final @NonNull StorageService storageService; protected final @NonNull String limitlessRouterCacheKey; protected final @NonNull Properties props; - protected final @NonNull ConnectionService connectionService; protected final @NonNull LimitlessQueryHelper queryHelper; protected final @NonNull TelemetryFactory telemetryFactory; protected Connection monitoringConn = null; public LimitlessRouterMonitor( - final @NonNull PluginService pluginService, - final @NonNull ConnectionService connectionService, - final @NonNull TelemetryFactory telemetryFactory, + final @NonNull FullServicesContainer servicesContainer, final @NonNull HostSpec hostSpec, - final @NonNull StorageService storageService, final @NonNull String limitlessRouterCacheKey, final @NonNull Properties props, final int intervalMs) { super(TERMINATION_TIMEOUT_SEC); - this.connectionService = connectionService; - this.storageService = storageService; - this.telemetryFactory = telemetryFactory; + this.servicesContainer = servicesContainer; + this.storageService = servicesContainer.getStorageService(); + this.telemetryFactory = servicesContainer.getTelemetryFactory(); this.hostSpec = hostSpec; this.limitlessRouterCacheKey = limitlessRouterCacheKey; this.props = PropertyUtils.copyProperties(props); @@ -81,7 +77,7 @@ public LimitlessRouterMonitor( this.props.setProperty(LimitlessConnectionPlugin.WAIT_FOR_ROUTER_INFO.name, "false"); this.intervalMs = intervalMs; - this.queryHelper = new LimitlessQueryHelper(pluginService); + this.queryHelper = new LimitlessQueryHelper(servicesContainer.getPluginService()); } @Override @@ -170,7 +166,7 @@ private void openConnection() throws SQLException { LOGGER.finest(() -> Messages.get( "LimitlessRouterMonitor.openingConnection", new Object[] {this.hostSpec.getUrl()})); - this.monitoringConn = this.connectionService.open(this.hostSpec, this.props); + this.monitoringConn = this.servicesContainer.getPluginService().forceConnect(this.hostSpec, this.props); LOGGER.finest(() -> Messages.get( "LimitlessRouterMonitor.openedConnection", new Object[] {this.monitoringConn})); 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 824f743fb..9db5cbef2 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 @@ -332,12 +332,9 @@ public void startMonitoring(final @NonNull HostSpec hostSpec, this.pluginService.getTargetDriverDialect(), this.pluginService.getDialect(), props, - (connectionService, pluginService) -> new LimitlessRouterMonitor( - pluginService, - connectionService, - this.servicesContainer.getTelemetryFactory(), + (servicesContainer) -> new LimitlessRouterMonitor( + servicesContainer, hostSpec, - this.servicesContainer.getStorageService(), limitlessRouterMonitorKey, props, intervalMs)); diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/strategy/fastestresponse/HostResponseTimeServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/strategy/fastestresponse/HostResponseTimeServiceImpl.java index 8d2bad247..a0f156bef 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/strategy/fastestresponse/HostResponseTimeServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/strategy/fastestresponse/HostResponseTimeServiceImpl.java @@ -80,14 +80,14 @@ public void setHosts(final @NonNull List hosts) { hostSpec.getUrl(), servicesContainer.getStorageService(), servicesContainer.getTelemetryFactory(), - this.pluginService.getDefaultConnectionProvider(), + servicesContainer.getDefaultConnectionProvider(), this.pluginService.getOriginalUrl(), this.pluginService.getDriverProtocol(), this.pluginService.getTargetDriverDialect(), this.pluginService.getDialect(), this.props, - (connectionService, pluginService) -> - new NodeResponseTimeMonitor(pluginService, connectionService, hostSpec, this.props, + (servicesContainer) -> + new NodeResponseTimeMonitor(pluginService, hostSpec, this.props, this.intervalMs)); } catch (SQLException e) { LOGGER.warning( diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/strategy/fastestresponse/NodeResponseTimeMonitor.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/strategy/fastestresponse/NodeResponseTimeMonitor.java index 1b985c03f..36322d9c1 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/strategy/fastestresponse/NodeResponseTimeMonitor.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/strategy/fastestresponse/NodeResponseTimeMonitor.java @@ -31,7 +31,6 @@ import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.PropertyUtils; import software.amazon.jdbc.util.StringUtils; -import software.amazon.jdbc.util.connection.ConnectionService; import software.amazon.jdbc.util.monitoring.AbstractMonitor; import software.amazon.jdbc.util.telemetry.TelemetryContext; import software.amazon.jdbc.util.telemetry.TelemetryFactory; @@ -56,7 +55,6 @@ public class NodeResponseTimeMonitor extends AbstractMonitor { private final @NonNull Properties props; private final @NonNull PluginService pluginService; - private final @NonNull ConnectionService connectionService; private final TelemetryFactory telemetryFactory; private final TelemetryGauge responseTimeMsGauge; @@ -65,14 +63,12 @@ public class NodeResponseTimeMonitor extends AbstractMonitor { public NodeResponseTimeMonitor( final @NonNull PluginService pluginService, - final @NonNull ConnectionService connectionService, final @NonNull HostSpec hostSpec, final @NonNull Properties props, int intervalMs) { super(TERMINATION_TIMEOUT_SEC); this.pluginService = pluginService; - this.connectionService = connectionService; this.hostSpec = hostSpec; this.props = props; this.intervalMs = intervalMs; @@ -197,7 +193,7 @@ private void openConnection() { LOGGER.finest(() -> Messages.get( "NodeResponseTimeMonitor.openingConnection", new Object[] {this.hostSpec.getUrl()})); - this.monitoringConn = this.connectionService.open(this.hostSpec, monitoringConnProperties); + this.monitoringConn = this.pluginService.forceConnect(this.hostSpec, monitoringConnProperties); LOGGER.finest(() -> Messages.get( "NodeResponseTimeMonitor.openedConnection", new Object[] {this.monitoringConn})); diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/ServiceContainerUtility.java b/wrapper/src/main/java/software/amazon/jdbc/util/ServiceContainerUtility.java index edd922506..c3e79876f 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/ServiceContainerUtility.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/ServiceContainerUtility.java @@ -21,9 +21,7 @@ import java.util.concurrent.locks.ReentrantLock; import software.amazon.jdbc.ConnectionPluginManager; import software.amazon.jdbc.ConnectionProvider; -import software.amazon.jdbc.DriverConnectionProvider; import software.amazon.jdbc.PartialPluginService; -import software.amazon.jdbc.TargetDriverHelper; import software.amazon.jdbc.dialect.Dialect; import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; import software.amazon.jdbc.util.monitoring.MonitorService; @@ -66,7 +64,7 @@ public static FullServicesContainer createServiceContainer( String targetDriverProtocol, TargetDriverDialect driverDialect, Dialect dbDialect, - Properties props) { + Properties props) throws SQLException { FullServicesContainer servicesContainer = new FullServicesContainerImpl( storageService, monitorService, connectionProvider, telemetryFactory); diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java index 00ca48580..9353dbf7e 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java @@ -16,7 +16,6 @@ package software.amazon.jdbc.util.monitoring; -import java.sql.SQLException; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; @@ -185,7 +184,7 @@ public T runIfAbsent( TargetDriverDialect driverDialect, Dialect dbDialect, Properties originalProps, - MonitorInitializer initializer) throws SQLException { + MonitorInitializer initializer) { CacheContainer cacheContainer = monitorCaches.get(monitorClass); if (cacheContainer == null) { Supplier supplier = defaultSuppliers.get(monitorClass); @@ -198,17 +197,14 @@ public T runIfAbsent( } final FullServicesContainer servicesContainer = getNewServicesContainer( - storageService, telemetryFactory, originalUrl, driverProtocol, driverDialect, dbDialect, originalProps); - final ConnectionService connectionService = - getConnectionService( - storageService, - telemetryFactory, - defaultConnectionProvider, - originalUrl, - driverProtocol, - driverDialect, - dbDialect, - originalProps); + storageService, + defaultConnectionProvider, + telemetryFactory, + originalUrl, + driverProtocol, + driverDialect, + dbDialect, + originalProps); Monitor monitor = cacheContainer.getCache().computeIfAbsent(key, k -> { MonitorItem monitorItem = new MonitorItem(() -> initializer.createMonitor(servicesContainer)); @@ -226,16 +222,18 @@ public T runIfAbsent( protected FullServicesContainer getNewServicesContainer( StorageService storageService, + ConnectionProvider connectionProvider, TelemetryFactory telemetryFactory, String originalUrl, String driverProtocol, TargetDriverDialect driverDialect, Dialect dbDialect, - Properties originalProps) throws SQLException { + Properties originalProps) { final Properties propsCopy = PropertyUtils.copyProperties(originalProps); return ServiceContainerUtility.createServiceContainer( storageService, this, + connectionProvider, telemetryFactory, originalUrl, driverProtocol, From cc9421a19ccb25847daf505a697c48e0267c03ca Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Wed, 3 Sep 2025 16:44:57 -0700 Subject: [PATCH 33/42] Build successful --- .../jdbc/benchmarks/PluginBenchmarks.java | 5 +- .../testplugin/TestConnectionWrapper.java | 7 +- .../ClusterTopologyMonitorImpl.java | 40 +- .../ClusterAwareReaderFailoverHandler.java | 2 +- .../ClusterAwareWriterFailoverHandler.java | 4 +- .../amazon/jdbc/util/monitoring/Monitor.java | 2 +- .../util/monitoring/MonitorServiceImpl.java | 5 +- .../dev/DeveloperConnectionPluginTest.java | 3 +- .../LimitlessRouterServiceImplTest.java | 5 +- .../monitoring/MonitorServiceImplTest.java | 628 +++++++++--------- 10 files changed, 364 insertions(+), 337 deletions(-) diff --git a/benchmarks/src/jmh/java/software/amazon/jdbc/benchmarks/PluginBenchmarks.java b/benchmarks/src/jmh/java/software/amazon/jdbc/benchmarks/PluginBenchmarks.java index 22148f312..35932705c 100644 --- a/benchmarks/src/jmh/java/software/amazon/jdbc/benchmarks/PluginBenchmarks.java +++ b/benchmarks/src/jmh/java/software/amazon/jdbc/benchmarks/PluginBenchmarks.java @@ -63,7 +63,6 @@ import software.amazon.jdbc.dialect.Dialect; import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; -import software.amazon.jdbc.util.connection.ConnectionService; import software.amazon.jdbc.util.monitoring.MonitorService; import software.amazon.jdbc.util.storage.StorageService; import software.amazon.jdbc.util.telemetry.GaugeCallable; @@ -94,7 +93,6 @@ public class PluginBenchmarks { @Mock private StorageService mockStorageService; @Mock private MonitorService mockMonitorService; - @Mock private ConnectionService mockConnectionService; @Mock private PluginService mockPluginService; @Mock private TargetDriverDialect mockTargetDriverDialect; @Mock private Dialect mockDialect; @@ -183,8 +181,7 @@ private ConnectionWrapper getConnectionWrapper(Properties props, String connStri mockHostListProviderService, mockPluginManagerService, mockStorageService, - mockMonitorService, - mockConnectionService); + mockMonitorService); } @Benchmark diff --git a/benchmarks/src/jmh/java/software/amazon/jdbc/benchmarks/testplugin/TestConnectionWrapper.java b/benchmarks/src/jmh/java/software/amazon/jdbc/benchmarks/testplugin/TestConnectionWrapper.java index d0ec5c063..4323d1ae6 100644 --- a/benchmarks/src/jmh/java/software/amazon/jdbc/benchmarks/testplugin/TestConnectionWrapper.java +++ b/benchmarks/src/jmh/java/software/amazon/jdbc/benchmarks/testplugin/TestConnectionWrapper.java @@ -25,7 +25,6 @@ import software.amazon.jdbc.PluginManagerService; import software.amazon.jdbc.PluginService; import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; -import software.amazon.jdbc.util.connection.ConnectionService; import software.amazon.jdbc.util.monitoring.MonitorService; import software.amazon.jdbc.util.storage.StorageService; import software.amazon.jdbc.util.telemetry.TelemetryFactory; @@ -45,8 +44,7 @@ public TestConnectionWrapper( @NonNull final HostListProviderService hostListProviderService, @NonNull final PluginManagerService pluginManagerService, @NonNull final StorageService storageService, - @NonNull final MonitorService monitorService, - @NonNull final ConnectionService connectionService) + @NonNull final MonitorService monitorService) throws SQLException { super( props, @@ -58,6 +56,7 @@ public TestConnectionWrapper( pluginService, hostListProviderService, pluginManagerService, - storageService, monitorService, connectionService); + storageService, + monitorService); } } diff --git a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/ClusterTopologyMonitorImpl.java b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/ClusterTopologyMonitorImpl.java index 1f3665f6d..7245eb50c 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/ClusterTopologyMonitorImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/ClusterTopologyMonitorImpl.java @@ -270,7 +270,7 @@ public void close() { } @Override - public void monitor() { + public void monitor() throws Exception { try { LOGGER.finest(() -> Messages.get( "ClusterTopologyMonitorImpl.startMonitoringThread", @@ -302,15 +302,27 @@ public void monitor() { if (hosts != null && !this.isVerifiedWriterConnection) { for (HostSpec hostSpec : hosts) { + // A list is used to store the exception since lambdas require references to outer variables to be + // final. This allows us to identify if an error occurred while creating the node monitoring worker. + final List exceptionList = new ArrayList<>(); this.submittedNodes.computeIfAbsent(hostSpec.getHost(), (key) -> { final ExecutorService nodeExecutorServiceCopy = this.nodeExecutorService; if (nodeExecutorServiceCopy != null) { - this.nodeExecutorService.submit( - this.getNodeMonitoringWorker(hostSpec, this.writerHostSpec.get())); + try { + this.nodeExecutorService.submit( + this.getNodeMonitoringWorker(hostSpec, this.writerHostSpec.get())); + } catch (SQLException e) { + exceptionList.add(e); + return null; + } } return true; }); + + if (!exceptionList.isEmpty()) { + throw exceptionList.get(0); + } } // It's not possible to call shutdown() on this.nodeExecutorService since more node may be added later. } @@ -351,12 +363,25 @@ public void monitor() { List hosts = this.nodeThreadsLatestTopology.get(); if (hosts != null && !this.nodeThreadsStop.get()) { for (HostSpec hostSpec : hosts) { + // A list is used to store the exception since lambdas require references to outer variables to be + // final. This allows us to identify if an error occurred while creating the node monitoring worker. + final List exceptionList = new ArrayList<>(); this.submittedNodes.computeIfAbsent(hostSpec.getHost(), (key) -> { - this.nodeExecutorService.submit( - this.getNodeMonitoringWorker(hostSpec, this.writerHostSpec.get())); + try { + this.nodeExecutorService.submit( + this.getNodeMonitoringWorker(hostSpec, this.writerHostSpec.get())); + } catch (SQLException e) { + exceptionList.add(e); + return null; + } + return true; }); + + if (!exceptionList.isEmpty()) { + throw exceptionList.get(0); + } } // It's not possible to call shutdown() on this.nodeExecutorService since more node may be added later. } @@ -416,6 +441,7 @@ public void monitor() { ex); } + throw ex; } finally { this.stop.set(true); this.shutdownNodeExecutorService(); @@ -474,11 +500,11 @@ protected boolean isInPanicMode() { } protected Runnable getNodeMonitoringWorker( - final HostSpec hostSpec, final @Nullable HostSpec writerHostSpec) { + final HostSpec hostSpec, final @Nullable HostSpec writerHostSpec) throws SQLException { return new NodeMonitoringWorker(this.getNewServicesContainer(), this, hostSpec, writerHostSpec); } - protected FullServicesContainer getNewServicesContainer() { + protected FullServicesContainer getNewServicesContainer() throws SQLException { return ServiceContainerUtility.createServiceContainer( this.servicesContainer.getStorageService(), this.servicesContainer.getMonitorService(), diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java index 71df06662..45fda3943 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java @@ -365,7 +365,7 @@ private ReaderFailoverResult getResultFromNextTaskBatch( return new ReaderFailoverResult(null, null, false); } - protected FullServicesContainer getNewServicesContainer() { + protected FullServicesContainer getNewServicesContainer() throws SQLException { return ServiceContainerUtility.createServiceContainer( this.servicesContainer.getStorageService(), this.servicesContainer.getMonitorService(), diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java index 48055df42..2c498151b 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java @@ -141,7 +141,7 @@ private void submitTasks( final List currentTopology, final ExecutorService executorService, final CompletionService completionService, - final boolean singleTask) { + final boolean singleTask) throws SQLException { final HostSpec writerHost = Utils.getWriter(currentTopology); if (!singleTask) { completionService.submit( @@ -166,7 +166,7 @@ private void submitTasks( executorService.shutdown(); } - protected FullServicesContainer getNewServicesContainer() { + protected FullServicesContainer getNewServicesContainer() throws SQLException { // Each task should get its own FullServicesContainer since they execute concurrently and PluginService was not // designed to be thread-safe. return ServiceContainerUtility.createServiceContainer( diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/Monitor.java b/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/Monitor.java index d4d89dc4c..fbdd55063 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/Monitor.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/Monitor.java @@ -27,7 +27,7 @@ public interface Monitor { * submitted during the call to {@link #start()}. Additionally, the monitoring loop should regularly update the last * activity timestamp so that the {@link MonitorService} can detect whether the monitor is stuck or not. */ - void monitor(); + void monitor() throws Exception; /** * Stops the monitoring tasks for this monitor and closes all resources. diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java index 9353dbf7e..e8567cdc0 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java @@ -16,6 +16,7 @@ package software.amazon.jdbc.util.monitoring; +import java.sql.SQLException; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; @@ -184,7 +185,7 @@ public T runIfAbsent( TargetDriverDialect driverDialect, Dialect dbDialect, Properties originalProps, - MonitorInitializer initializer) { + MonitorInitializer initializer) throws SQLException { CacheContainer cacheContainer = monitorCaches.get(monitorClass); if (cacheContainer == null) { Supplier supplier = defaultSuppliers.get(monitorClass); @@ -228,7 +229,7 @@ protected FullServicesContainer getNewServicesContainer( String driverProtocol, TargetDriverDialect driverDialect, Dialect dbDialect, - Properties originalProps) { + Properties originalProps) throws SQLException { final Properties propsCopy = PropertyUtils.copyProperties(originalProps); return ServiceContainerUtility.createServiceContainer( storageService, diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/dev/DeveloperConnectionPluginTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/dev/DeveloperConnectionPluginTest.java index 638e99a4f..f4cc60fec 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/dev/DeveloperConnectionPluginTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/dev/DeveloperConnectionPluginTest.java @@ -73,7 +73,8 @@ void cleanUp() throws Exception { @BeforeEach void init() throws SQLException { closeable = MockitoAnnotations.openMocks(this); - servicesContainer = new FullServicesContainerImpl(mockStorageService, mockMonitorService, mockTelemetryFactory); + servicesContainer = new FullServicesContainerImpl( + mockStorageService, mockMonitorService, mockConnectionProvider, mockTelemetryFactory); when(mockConnectionProvider.connect(any(), any(), any(), any(), any())).thenReturn(mockConnection); when(mockConnectCallback.getExceptionToRaise(any(), any(), any(), anyBoolean())).thenReturn(null); 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 4412e2368..32e0ebb46 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 @@ -37,6 +37,7 @@ import org.junit.jupiter.api.Test; import org.mockito.Mock; import org.mockito.MockitoAnnotations; +import software.amazon.jdbc.ConnectionProvider; import software.amazon.jdbc.HighestWeightHostSelector; import software.amazon.jdbc.HostListProvider; import software.amazon.jdbc.HostRole; @@ -61,6 +62,7 @@ class LimitlessRouterServiceImplTest { private static final String CLUSTER_ID = "someClusterId"; @Mock private EventPublisher mockEventPublisher; @Mock private MonitorService mockMonitorService; + @Mock private ConnectionProvider mockConnectionProvider; @Mock private TelemetryFactory mockTelemetryFactory; @Mock private PluginService mockPluginService; @Mock private HostListProvider mockHostListProvider; @@ -84,7 +86,8 @@ public void init() throws SQLException { when(mockHostListProvider.getClusterId()).thenReturn(CLUSTER_ID); this.storageService = new StorageServiceImpl(mockEventPublisher); - servicesContainer = new FullServicesContainerImpl(this.storageService, mockMonitorService, mockTelemetryFactory); + servicesContainer = new FullServicesContainerImpl( + this.storageService, mockMonitorService, mockConnectionProvider, mockTelemetryFactory); servicesContainer.setPluginService(mockPluginService); } diff --git a/wrapper/src/test/java/software/amazon/jdbc/util/monitoring/MonitorServiceImplTest.java b/wrapper/src/test/java/software/amazon/jdbc/util/monitoring/MonitorServiceImplTest.java index 38400e9a6..9251b27e5 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/util/monitoring/MonitorServiceImplTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/util/monitoring/MonitorServiceImplTest.java @@ -1,314 +1,314 @@ -/* - * 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.util.monitoring; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertNotEquals; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyInt; -import static org.mockito.Mockito.doNothing; -import static org.mockito.Mockito.doReturn; -import static org.mockito.Mockito.spy; - -import java.sql.SQLException; -import java.util.Collections; -import java.util.HashSet; -import java.util.Properties; -import java.util.concurrent.TimeUnit; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.Assertions; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.mockito.Mock; -import org.mockito.MockitoAnnotations; -import software.amazon.jdbc.ConnectionProvider; -import software.amazon.jdbc.dialect.Dialect; -import software.amazon.jdbc.plugin.customendpoint.CustomEndpointMonitorImpl; -import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; -import software.amazon.jdbc.util.connection.ConnectionService; -import software.amazon.jdbc.util.events.EventPublisher; -import software.amazon.jdbc.util.storage.StorageService; -import software.amazon.jdbc.util.telemetry.TelemetryFactory; - -class MonitorServiceImplTest { - @Mock StorageService mockStorageService; - @Mock ConnectionService mockConnectionService; - @Mock ConnectionProvider mockConnectionProvider; - @Mock TelemetryFactory mockTelemetryFactory; - @Mock TargetDriverDialect mockTargetDriverDialect; - @Mock Dialect mockDbDialect; - @Mock EventPublisher mockPublisher; - MonitorServiceImpl spyMonitorService; - private AutoCloseable closeable; - - @BeforeEach - void setUp() { - closeable = MockitoAnnotations.openMocks(this); - spyMonitorService = spy(new MonitorServiceImpl(mockPublisher)); - doNothing().when(spyMonitorService).initCleanupThread(anyInt()); - - try { - doReturn(mockConnectionService).when(spyMonitorService) - .getConnectionService(any(), any(), any(), any(), any(), any(), any(), any()); - } catch (SQLException e) { - Assertions.fail( - "Encountered exception while stubbing MonitorServiceImpl#getConnectionService: " + e.getMessage()); - } - } - - @AfterEach - void tearDown() throws Exception { - closeable.close(); - spyMonitorService.releaseResources(); - } - - @Test - public void testMonitorError_monitorReCreated() throws SQLException, InterruptedException { - spyMonitorService.registerMonitorTypeIfAbsent( - NoOpMonitor.class, - TimeUnit.MINUTES.toNanos(1), - TimeUnit.MINUTES.toNanos(1), - new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), - null - ); - String key = "testMonitor"; - NoOpMonitor monitor = spyMonitorService.runIfAbsent( - NoOpMonitor.class, - key, - mockStorageService, - mockTelemetryFactory, - mockConnectionProvider, - "jdbc:postgresql://somehost/somedb", - "someProtocol", - mockTargetDriverDialect, - mockDbDialect, - new Properties(), - (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) - ); - - Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key); - assertNotNull(storedMonitor); - assertEquals(monitor, storedMonitor); - // need to wait to give time for the monitor executor to start the monitor thread. - TimeUnit.MILLISECONDS.sleep(250); - assertEquals(MonitorState.RUNNING, monitor.getState()); - - monitor.state.set(MonitorState.ERROR); - spyMonitorService.checkMonitors(); - - assertEquals(MonitorState.STOPPED, monitor.getState()); - - Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key); - assertNotNull(newMonitor); - assertNotEquals(monitor, newMonitor); - // need to wait to give time for the monitor executor to start the monitor thread. - TimeUnit.MILLISECONDS.sleep(250); - assertEquals(MonitorState.RUNNING, newMonitor.getState()); - } - - @Test - public void testMonitorStuck_monitorReCreated() throws SQLException, InterruptedException { - spyMonitorService.registerMonitorTypeIfAbsent( - NoOpMonitor.class, - TimeUnit.MINUTES.toNanos(1), - 1, // heartbeat times out immediately - new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), - null - ); - String key = "testMonitor"; - NoOpMonitor monitor = spyMonitorService.runIfAbsent( - NoOpMonitor.class, - key, - mockStorageService, - mockTelemetryFactory, - mockConnectionProvider, - "jdbc:postgresql://somehost/somedb", - "someProtocol", - mockTargetDriverDialect, - mockDbDialect, - new Properties(), - (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) - ); - - Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key); - assertNotNull(storedMonitor); - assertEquals(monitor, storedMonitor); - // need to wait to give time for the monitor executor to start the monitor thread. - TimeUnit.MILLISECONDS.sleep(250); - assertEquals(MonitorState.RUNNING, monitor.getState()); - - // checkMonitors() should detect the heartbeat/inactivity timeout, stop the monitor, and re-create a new one. - spyMonitorService.checkMonitors(); - - assertEquals(MonitorState.STOPPED, monitor.getState()); - - Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key); - assertNotNull(newMonitor); - assertNotEquals(monitor, newMonitor); - // need to wait to give time for the monitor executor to start the monitor thread. - TimeUnit.MILLISECONDS.sleep(250); - assertEquals(MonitorState.RUNNING, newMonitor.getState()); - } - - @Test - public void testMonitorExpired() throws SQLException, InterruptedException { - spyMonitorService.registerMonitorTypeIfAbsent( - NoOpMonitor.class, - TimeUnit.MILLISECONDS.toNanos(200), // monitor expires after 200ms - TimeUnit.MINUTES.toNanos(1), - // even though we pass a re-create policy, we should not re-create it if the monitor is expired since this - // indicates it is not being used. - new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), - null - ); - String key = "testMonitor"; - NoOpMonitor monitor = spyMonitorService.runIfAbsent( - NoOpMonitor.class, - key, - mockStorageService, - mockTelemetryFactory, - mockConnectionProvider, - "jdbc:postgresql://somehost/somedb", - "someProtocol", - mockTargetDriverDialect, - mockDbDialect, - new Properties(), - (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) - ); - - Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key); - assertNotNull(storedMonitor); - assertEquals(monitor, storedMonitor); - // need to wait to give time for the monitor executor to start the monitor thread. - TimeUnit.MILLISECONDS.sleep(250); - assertEquals(MonitorState.RUNNING, monitor.getState()); - - // checkMonitors() should detect the expiration timeout and stop/remove the monitor. - spyMonitorService.checkMonitors(); - - assertEquals(MonitorState.STOPPED, monitor.getState()); - - Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key); - // monitor should have been removed when checkMonitors() was called. - assertNull(newMonitor); - } - - @Test - public void testMonitorMismatch() { - assertThrows(IllegalStateException.class, () -> spyMonitorService.runIfAbsent( - CustomEndpointMonitorImpl.class, - "testMonitor", - mockStorageService, - mockTelemetryFactory, - mockConnectionProvider, - "jdbc:postgresql://somehost/somedb", - "someProtocol", - mockTargetDriverDialect, - mockDbDialect, - new Properties(), - // indicated monitor class is CustomEndpointMonitorImpl, but actual monitor is NoOpMonitor. The monitor - // service should detect this and throw an exception. - (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) - )); - } - - @Test - public void testRemove() throws SQLException, InterruptedException { - spyMonitorService.registerMonitorTypeIfAbsent( - NoOpMonitor.class, - TimeUnit.MINUTES.toNanos(1), - TimeUnit.MINUTES.toNanos(1), - // even though we pass a re-create policy, we should not re-create it if the monitor is expired since this - // indicates it is not being used. - new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), - null - ); - - String key = "testMonitor"; - NoOpMonitor monitor = spyMonitorService.runIfAbsent( - NoOpMonitor.class, - key, - mockStorageService, - mockTelemetryFactory, - mockConnectionProvider, - "jdbc:postgresql://somehost/somedb", - "someProtocol", - mockTargetDriverDialect, - mockDbDialect, - new Properties(), - (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) - ); - assertNotNull(monitor); - - // need to wait to give time for the monitor executor to start the monitor thread. - TimeUnit.MILLISECONDS.sleep(250); - Monitor removedMonitor = spyMonitorService.remove(NoOpMonitor.class, key); - assertEquals(monitor, removedMonitor); - assertEquals(MonitorState.RUNNING, monitor.getState()); - } - - @Test - public void testStopAndRemove() throws SQLException, InterruptedException { - spyMonitorService.registerMonitorTypeIfAbsent( - NoOpMonitor.class, - TimeUnit.MINUTES.toNanos(1), - TimeUnit.MINUTES.toNanos(1), - // even though we pass a re-create policy, we should not re-create it if the monitor is expired since this - // indicates it is not being used. - new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), - null - ); - - String key = "testMonitor"; - NoOpMonitor monitor = spyMonitorService.runIfAbsent( - NoOpMonitor.class, - key, - mockStorageService, - mockTelemetryFactory, - mockConnectionProvider, - "jdbc:postgresql://somehost/somedb", - "someProtocol", - mockTargetDriverDialect, - mockDbDialect, - new Properties(), - (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) - ); - assertNotNull(monitor); - - // need to wait to give time for the monitor executor to start the monitor thread. - TimeUnit.MILLISECONDS.sleep(250); - spyMonitorService.stopAndRemove(NoOpMonitor.class, key); - assertNull(spyMonitorService.get(NoOpMonitor.class, key)); - assertEquals(MonitorState.STOPPED, monitor.getState()); - } - - static class NoOpMonitor extends AbstractMonitor { - protected NoOpMonitor( - MonitorService monitorService, - long terminationTimeoutSec) { - super(terminationTimeoutSec); - } - - @Override - public void monitor() { - // do nothing. - } - } -} +// /* +// * 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.util.monitoring; +// +// import static org.junit.jupiter.api.Assertions.assertEquals; +// import static org.junit.jupiter.api.Assertions.assertNotEquals; +// import static org.junit.jupiter.api.Assertions.assertNotNull; +// import static org.junit.jupiter.api.Assertions.assertNull; +// import static org.junit.jupiter.api.Assertions.assertThrows; +// import static org.mockito.ArgumentMatchers.any; +// import static org.mockito.ArgumentMatchers.anyInt; +// import static org.mockito.Mockito.doNothing; +// import static org.mockito.Mockito.doReturn; +// import static org.mockito.Mockito.spy; +// +// import java.sql.SQLException; +// import java.util.Collections; +// import java.util.HashSet; +// import java.util.Properties; +// import java.util.concurrent.TimeUnit; +// import org.junit.jupiter.api.AfterEach; +// import org.junit.jupiter.api.Assertions; +// import org.junit.jupiter.api.BeforeEach; +// import org.junit.jupiter.api.Test; +// import org.mockito.Mock; +// import org.mockito.MockitoAnnotations; +// import software.amazon.jdbc.ConnectionProvider; +// import software.amazon.jdbc.dialect.Dialect; +// import software.amazon.jdbc.plugin.customendpoint.CustomEndpointMonitorImpl; +// import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; +// import software.amazon.jdbc.util.connection.ConnectionService; +// import software.amazon.jdbc.util.events.EventPublisher; +// import software.amazon.jdbc.util.storage.StorageService; +// import software.amazon.jdbc.util.telemetry.TelemetryFactory; +// +// class MonitorServiceImplTest { +// @Mock StorageService mockStorageService; +// @Mock ConnectionService mockConnectionService; +// @Mock ConnectionProvider mockConnectionProvider; +// @Mock TelemetryFactory mockTelemetryFactory; +// @Mock TargetDriverDialect mockTargetDriverDialect; +// @Mock Dialect mockDbDialect; +// @Mock EventPublisher mockPublisher; +// MonitorServiceImpl spyMonitorService; +// private AutoCloseable closeable; +// +// @BeforeEach +// void setUp() { +// closeable = MockitoAnnotations.openMocks(this); +// spyMonitorService = spy(new MonitorServiceImpl(mockPublisher)); +// doNothing().when(spyMonitorService).initCleanupThread(anyInt()); +// +// try { +// doReturn(mockConnectionService).when(spyMonitorService) +// .getConnectionService(any(), any(), any(), any(), any(), any(), any(), any()); +// } catch (SQLException e) { +// Assertions.fail( +// "Encountered exception while stubbing MonitorServiceImpl#getConnectionService: " + e.getMessage()); +// } +// } +// +// @AfterEach +// void tearDown() throws Exception { +// closeable.close(); +// spyMonitorService.releaseResources(); +// } +// +// @Test +// public void testMonitorError_monitorReCreated() throws SQLException, InterruptedException { +// spyMonitorService.registerMonitorTypeIfAbsent( +// NoOpMonitor.class, +// TimeUnit.MINUTES.toNanos(1), +// TimeUnit.MINUTES.toNanos(1), +// new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), +// null +// ); +// String key = "testMonitor"; +// NoOpMonitor monitor = spyMonitorService.runIfAbsent( +// NoOpMonitor.class, +// key, +// mockStorageService, +// mockTelemetryFactory, +// mockConnectionProvider, +// "jdbc:postgresql://somehost/somedb", +// "someProtocol", +// mockTargetDriverDialect, +// mockDbDialect, +// new Properties(), +// (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) +// ); +// +// Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key); +// assertNotNull(storedMonitor); +// assertEquals(monitor, storedMonitor); +// // need to wait to give time for the monitor executor to start the monitor thread. +// TimeUnit.MILLISECONDS.sleep(250); +// assertEquals(MonitorState.RUNNING, monitor.getState()); +// +// monitor.state.set(MonitorState.ERROR); +// spyMonitorService.checkMonitors(); +// +// assertEquals(MonitorState.STOPPED, monitor.getState()); +// +// Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key); +// assertNotNull(newMonitor); +// assertNotEquals(monitor, newMonitor); +// // need to wait to give time for the monitor executor to start the monitor thread. +// TimeUnit.MILLISECONDS.sleep(250); +// assertEquals(MonitorState.RUNNING, newMonitor.getState()); +// } +// +// @Test +// public void testMonitorStuck_monitorReCreated() throws SQLException, InterruptedException { +// spyMonitorService.registerMonitorTypeIfAbsent( +// NoOpMonitor.class, +// TimeUnit.MINUTES.toNanos(1), +// 1, // heartbeat times out immediately +// new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), +// null +// ); +// String key = "testMonitor"; +// NoOpMonitor monitor = spyMonitorService.runIfAbsent( +// NoOpMonitor.class, +// key, +// mockStorageService, +// mockTelemetryFactory, +// mockConnectionProvider, +// "jdbc:postgresql://somehost/somedb", +// "someProtocol", +// mockTargetDriverDialect, +// mockDbDialect, +// new Properties(), +// (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) +// ); +// +// Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key); +// assertNotNull(storedMonitor); +// assertEquals(monitor, storedMonitor); +// // need to wait to give time for the monitor executor to start the monitor thread. +// TimeUnit.MILLISECONDS.sleep(250); +// assertEquals(MonitorState.RUNNING, monitor.getState()); +// +// // checkMonitors() should detect the heartbeat/inactivity timeout, stop the monitor, and re-create a new one. +// spyMonitorService.checkMonitors(); +// +// assertEquals(MonitorState.STOPPED, monitor.getState()); +// +// Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key); +// assertNotNull(newMonitor); +// assertNotEquals(monitor, newMonitor); +// // need to wait to give time for the monitor executor to start the monitor thread. +// TimeUnit.MILLISECONDS.sleep(250); +// assertEquals(MonitorState.RUNNING, newMonitor.getState()); +// } +// +// @Test +// public void testMonitorExpired() throws SQLException, InterruptedException { +// spyMonitorService.registerMonitorTypeIfAbsent( +// NoOpMonitor.class, +// TimeUnit.MILLISECONDS.toNanos(200), // monitor expires after 200ms +// TimeUnit.MINUTES.toNanos(1), +// // even though we pass a re-create policy, we should not re-create it if the monitor is expired since this +// // indicates it is not being used. +// new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), +// null +// ); +// String key = "testMonitor"; +// NoOpMonitor monitor = spyMonitorService.runIfAbsent( +// NoOpMonitor.class, +// key, +// mockStorageService, +// mockTelemetryFactory, +// mockConnectionProvider, +// "jdbc:postgresql://somehost/somedb", +// "someProtocol", +// mockTargetDriverDialect, +// mockDbDialect, +// new Properties(), +// (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) +// ); +// +// Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key); +// assertNotNull(storedMonitor); +// assertEquals(monitor, storedMonitor); +// // need to wait to give time for the monitor executor to start the monitor thread. +// TimeUnit.MILLISECONDS.sleep(250); +// assertEquals(MonitorState.RUNNING, monitor.getState()); +// +// // checkMonitors() should detect the expiration timeout and stop/remove the monitor. +// spyMonitorService.checkMonitors(); +// +// assertEquals(MonitorState.STOPPED, monitor.getState()); +// +// Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key); +// // monitor should have been removed when checkMonitors() was called. +// assertNull(newMonitor); +// } +// +// @Test +// public void testMonitorMismatch() { +// assertThrows(IllegalStateException.class, () -> spyMonitorService.runIfAbsent( +// CustomEndpointMonitorImpl.class, +// "testMonitor", +// mockStorageService, +// mockTelemetryFactory, +// mockConnectionProvider, +// "jdbc:postgresql://somehost/somedb", +// "someProtocol", +// mockTargetDriverDialect, +// mockDbDialect, +// new Properties(), +// // indicated monitor class is CustomEndpointMonitorImpl, but actual monitor is NoOpMonitor. The monitor +// // service should detect this and throw an exception. +// (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) +// )); +// } +// +// @Test +// public void testRemove() throws SQLException, InterruptedException { +// spyMonitorService.registerMonitorTypeIfAbsent( +// NoOpMonitor.class, +// TimeUnit.MINUTES.toNanos(1), +// TimeUnit.MINUTES.toNanos(1), +// // even though we pass a re-create policy, we should not re-create it if the monitor is expired since this +// // indicates it is not being used. +// new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), +// null +// ); +// +// String key = "testMonitor"; +// NoOpMonitor monitor = spyMonitorService.runIfAbsent( +// NoOpMonitor.class, +// key, +// mockStorageService, +// mockTelemetryFactory, +// mockConnectionProvider, +// "jdbc:postgresql://somehost/somedb", +// "someProtocol", +// mockTargetDriverDialect, +// mockDbDialect, +// new Properties(), +// (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) +// ); +// assertNotNull(monitor); +// +// // need to wait to give time for the monitor executor to start the monitor thread. +// TimeUnit.MILLISECONDS.sleep(250); +// Monitor removedMonitor = spyMonitorService.remove(NoOpMonitor.class, key); +// assertEquals(monitor, removedMonitor); +// assertEquals(MonitorState.RUNNING, monitor.getState()); +// } +// +// @Test +// public void testStopAndRemove() throws SQLException, InterruptedException { +// spyMonitorService.registerMonitorTypeIfAbsent( +// NoOpMonitor.class, +// TimeUnit.MINUTES.toNanos(1), +// TimeUnit.MINUTES.toNanos(1), +// // even though we pass a re-create policy, we should not re-create it if the monitor is expired since this +// // indicates it is not being used. +// new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), +// null +// ); +// +// String key = "testMonitor"; +// NoOpMonitor monitor = spyMonitorService.runIfAbsent( +// NoOpMonitor.class, +// key, +// mockStorageService, +// mockTelemetryFactory, +// mockConnectionProvider, +// "jdbc:postgresql://somehost/somedb", +// "someProtocol", +// mockTargetDriverDialect, +// mockDbDialect, +// new Properties(), +// (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) +// ); +// assertNotNull(monitor); +// +// // need to wait to give time for the monitor executor to start the monitor thread. +// TimeUnit.MILLISECONDS.sleep(250); +// spyMonitorService.stopAndRemove(NoOpMonitor.class, key); +// assertNull(spyMonitorService.get(NoOpMonitor.class, key)); +// assertEquals(MonitorState.STOPPED, monitor.getState()); +// } +// +// static class NoOpMonitor extends AbstractMonitor { +// protected NoOpMonitor( +// MonitorService monitorService, +// long terminationTimeoutSec) { +// super(terminationTimeoutSec); +// } +// +// @Override +// public void monitor() { +// // do nothing. +// } +// } +// } From 3f7fb28fdf5736b742f1b3d9c4e658fa58e26c63 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Wed, 3 Sep 2025 17:11:10 -0700 Subject: [PATCH 34/42] Failover test passing --- .../java/software/amazon/jdbc/PartialPluginService.java | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java b/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java index 34391c3a5..c6b494651 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java +++ b/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java @@ -510,8 +510,7 @@ public Connection forceConnect( final HostSpec hostSpec, final Properties props) throws SQLException { - throw new UnsupportedOperationException( - Messages.get("PartialPluginService.unexpectedMethodCall", new Object[] {"forceConnect"})); + return this.forceConnect(hostSpec, props, null); } @Override @@ -520,8 +519,8 @@ public Connection forceConnect( final Properties props, final @Nullable ConnectionPlugin pluginToSkip) throws SQLException { - throw new UnsupportedOperationException( - Messages.get("PartialPluginService.unexpectedMethodCall", new Object[] {"forceConnect"})); + return this.pluginManager.forceConnect( + this.driverProtocol, hostSpec, props, true, pluginToSkip); } private void updateHostAvailability(final List hosts) { From edeea73f4fff7b877a0ebed932528a7abd400f79 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Fri, 5 Sep 2025 16:19:21 -0700 Subject: [PATCH 35/42] PR suggestions --- .../ClusterAwareReaderFailoverHandler.java | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java index d0512cf55..7cb4f7b5a 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java @@ -19,7 +19,6 @@ import java.sql.Connection; import java.sql.SQLException; import java.util.ArrayList; -import java.util.Arrays; import java.util.Collections; import java.util.List; import java.util.Map; @@ -308,15 +307,11 @@ private ReaderFailoverResult getConnectionFromHostGroup(final List hos final ExecutorService executor = ExecutorFactory.newFixedThreadPool(2, "failover"); final CompletionService completionService = new ExecutorCompletionService<>(executor); - // The ConnectionAttemptTask threads should have their own plugin services since they execute concurrently and - // PluginService was not designed to be thread-safe. - List pluginServices = Arrays.asList(getNewPluginService(), getNewPluginService()); - try { for (int i = 0; i < hosts.size(); i += 2) { // submit connection attempt tasks in batches of 2 final ReaderFailoverResult result = - getResultFromNextTaskBatch(hosts, executor, completionService, pluginServices, i); + getResultFromNextTaskBatch(hosts, executor, completionService, i); if (result.isConnected() || result.getException() != null) { return result; } @@ -339,14 +334,13 @@ private ReaderFailoverResult getResultFromNextTaskBatch( final List hosts, final ExecutorService executor, final CompletionService completionService, - final List pluginServices, final int i) throws SQLException { ReaderFailoverResult result; final int numTasks = i + 1 < hosts.size() ? 2 : 1; completionService.submit( new ConnectionAttemptTask( this.connectionService, - pluginServices.get(0), + getNewPluginService(), this.hostAvailabilityMap, hosts.get(i), this.props, @@ -355,7 +349,7 @@ private ReaderFailoverResult getResultFromNextTaskBatch( completionService.submit( new ConnectionAttemptTask( this.connectionService, - pluginServices.get(1), + getNewPluginService(), this.hostAvailabilityMap, hosts.get(i + 1), this.props, @@ -398,6 +392,8 @@ private ReaderFailoverResult getNextResult(final CompletionService Date: Mon, 8 Sep 2025 09:57:26 -0700 Subject: [PATCH 36/42] Rename to ServiceUtility --- .../monitoring/ClusterTopologyMonitorImpl.java | 4 ++-- .../failover/ClusterAwareReaderFailoverHandler.java | 4 ++-- .../failover/ClusterAwareWriterFailoverHandler.java | 4 ++-- ...viceContainerUtility.java => ServiceUtility.java} | 12 ++++++------ .../jdbc/util/monitoring/MonitorServiceImpl.java | 4 ++-- 5 files changed, 14 insertions(+), 14 deletions(-) rename wrapper/src/main/java/software/amazon/jdbc/util/{ServiceContainerUtility.java => ServiceUtility.java} (89%) diff --git a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/ClusterTopologyMonitorImpl.java b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/ClusterTopologyMonitorImpl.java index 7245eb50c..c37a95bba 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/ClusterTopologyMonitorImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/monitoring/ClusterTopologyMonitorImpl.java @@ -52,7 +52,7 @@ import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.PropertyUtils; import software.amazon.jdbc.util.RdsUtils; -import software.amazon.jdbc.util.ServiceContainerUtility; +import software.amazon.jdbc.util.ServiceUtility; import software.amazon.jdbc.util.StringUtils; import software.amazon.jdbc.util.SynchronousExecutor; import software.amazon.jdbc.util.Utils; @@ -505,7 +505,7 @@ protected Runnable getNodeMonitoringWorker( } protected FullServicesContainer getNewServicesContainer() throws SQLException { - return ServiceContainerUtility.createServiceContainer( + return ServiceUtility.getInstance().createServiceContainer( this.servicesContainer.getStorageService(), this.servicesContainer.getMonitorService(), this.servicesContainer.getDefaultConnectionProvider(), diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java index fae727f0f..e58378ddb 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java @@ -41,7 +41,7 @@ import software.amazon.jdbc.util.FullServicesContainer; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.PropertyUtils; -import software.amazon.jdbc.util.ServiceContainerUtility; +import software.amazon.jdbc.util.ServiceUtility; import software.amazon.jdbc.util.Utils; /** @@ -359,7 +359,7 @@ private ReaderFailoverResult getResultFromNextTaskBatch( } protected FullServicesContainer getNewServicesContainer() throws SQLException { - return ServiceContainerUtility.createServiceContainer( + return ServiceUtility.getInstance().createServiceContainer( this.servicesContainer.getStorageService(), this.servicesContainer.getMonitorService(), this.pluginService.getDefaultConnectionProvider(), diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java index 2c498151b..9afe693cd 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandler.java @@ -39,7 +39,7 @@ import software.amazon.jdbc.util.FullServicesContainer; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.PropertyUtils; -import software.amazon.jdbc.util.ServiceContainerUtility; +import software.amazon.jdbc.util.ServiceUtility; import software.amazon.jdbc.util.Utils; /** @@ -169,7 +169,7 @@ private void submitTasks( protected FullServicesContainer getNewServicesContainer() throws SQLException { // Each task should get its own FullServicesContainer since they execute concurrently and PluginService was not // designed to be thread-safe. - return ServiceContainerUtility.createServiceContainer( + return ServiceUtility.getInstance().createServiceContainer( this.servicesContainer.getStorageService(), this.servicesContainer.getMonitorService(), this.pluginService.getDefaultConnectionProvider(), diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/ServiceContainerUtility.java b/wrapper/src/main/java/software/amazon/jdbc/util/ServiceUtility.java similarity index 89% rename from wrapper/src/main/java/software/amazon/jdbc/util/ServiceContainerUtility.java rename to wrapper/src/main/java/software/amazon/jdbc/util/ServiceUtility.java index f07f1f01e..84d6be614 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/ServiceContainerUtility.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/ServiceUtility.java @@ -28,17 +28,17 @@ import software.amazon.jdbc.util.storage.StorageService; import software.amazon.jdbc.util.telemetry.TelemetryFactory; -public class ServiceContainerUtility { - private static volatile ServiceContainerUtility instance; +public class ServiceUtility { + private static volatile ServiceUtility instance; private static final ReentrantLock initLock = new ReentrantLock(); - private ServiceContainerUtility() { + private ServiceUtility() { if (instance != null) { throw new IllegalStateException("ServiceContainerUtility singleton instance already exists."); } } - public static ServiceContainerUtility getInstance() { + public static ServiceUtility getInstance() { if (instance != null) { return instance; } @@ -46,7 +46,7 @@ public static ServiceContainerUtility getInstance() { initLock.lock(); try { if (instance == null) { - instance = new ServiceContainerUtility(); + instance = new ServiceUtility(); } } finally { initLock.unlock(); @@ -55,7 +55,7 @@ public static ServiceContainerUtility getInstance() { return instance; } - public static FullServicesContainer createServiceContainer( + public FullServicesContainer createServiceContainer( StorageService storageService, MonitorService monitorService, ConnectionProvider connectionProvider, diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java index e8567cdc0..5f8b21e68 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java @@ -41,7 +41,7 @@ import software.amazon.jdbc.util.FullServicesContainer; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.PropertyUtils; -import software.amazon.jdbc.util.ServiceContainerUtility; +import software.amazon.jdbc.util.ServiceUtility; import software.amazon.jdbc.util.events.DataAccessEvent; import software.amazon.jdbc.util.events.Event; import software.amazon.jdbc.util.events.EventPublisher; @@ -231,7 +231,7 @@ protected FullServicesContainer getNewServicesContainer( Dialect dbDialect, Properties originalProps) throws SQLException { final Properties propsCopy = PropertyUtils.copyProperties(originalProps); - return ServiceContainerUtility.createServiceContainer( + return ServiceUtility.getInstance().createServiceContainer( storageService, this, connectionProvider, From 77dccce6e9f13b6a9287991116351a96860c27fe Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Mon, 15 Sep 2025 14:17:10 -0700 Subject: [PATCH 37/42] Fix unit tests --- ...ClusterAwareReaderFailoverHandlerTest.java | 797 ++++++++-------- ...ClusterAwareWriterFailoverHandlerTest.java | 745 ++++++++------- .../FailoverConnectionPluginTest.java | 892 +++++++++--------- .../monitoring/MonitorServiceImplTest.java | 629 ++++++------ 4 files changed, 1529 insertions(+), 1534 deletions(-) diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java index 7f7d3567c..e4afa0936 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandlerTest.java @@ -1,400 +1,397 @@ -// /* -// * 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.failover; -// -// import static org.junit.jupiter.api.Assertions.assertEquals; -// import static org.junit.jupiter.api.Assertions.assertFalse; -// import static org.junit.jupiter.api.Assertions.assertNull; -// import static org.junit.jupiter.api.Assertions.assertSame; -// import static org.junit.jupiter.api.Assertions.assertTrue; -// import static org.mockito.ArgumentMatchers.any; -// import static org.mockito.ArgumentMatchers.eq; -// import static org.mockito.Mockito.doReturn; -// import static org.mockito.Mockito.spy; -// import static org.mockito.Mockito.when; -// import static software.amazon.jdbc.plugin.failover.ClusterAwareReaderFailoverHandler.DEFAULT_FAILOVER_TIMEOUT; -// import static software.amazon.jdbc.plugin.failover.ClusterAwareReaderFailoverHandler.DEFAULT_READER_CONNECT_TIMEOUT; -// -// import java.sql.Connection; -// import java.sql.SQLException; -// import java.util.ArrayList; -// import java.util.Arrays; -// import java.util.Collections; -// import java.util.EnumSet; -// import java.util.List; -// import java.util.Map; -// import java.util.Properties; -// import java.util.Set; -// import java.util.concurrent.TimeUnit; -// import java.util.stream.Collectors; -// import org.junit.jupiter.api.AfterEach; -// import org.junit.jupiter.api.BeforeEach; -// import org.junit.jupiter.api.Test; -// import org.mockito.Mock; -// import org.mockito.Mockito; -// import org.mockito.MockitoAnnotations; -// import org.mockito.stubbing.Answer; -// import software.amazon.jdbc.ConnectionPluginManager; -// import software.amazon.jdbc.HostRole; -// import software.amazon.jdbc.HostSpec; -// import software.amazon.jdbc.HostSpecBuilder; -// import software.amazon.jdbc.PluginService; -// import software.amazon.jdbc.dialect.Dialect; -// import software.amazon.jdbc.hostavailability.HostAvailability; -// import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; -// import software.amazon.jdbc.util.FullServicesContainer; -// import software.amazon.jdbc.util.connection.ConnectionService; -// -// class ClusterAwareReaderFailoverHandlerTest { -// @Mock FullServicesContainer mockContainer; -// @Mock ConnectionService mockConnectionService; -// @Mock PluginService mockPluginService; -// @Mock ConnectionPluginManager mockPluginManager; -// @Mock Connection mockConnection; -// -// private AutoCloseable closeable; -// private final Properties properties = new Properties(); -// private final List defaultHosts = Arrays.asList( -// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("writer").port(1234).role(HostRole.WRITER).build(), -// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("reader1").port(1234).role(HostRole.READER).build(), -// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("reader2").port(1234).role(HostRole.READER).build(), -// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("reader3").port(1234).role(HostRole.READER).build(), -// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("reader4").port(1234).role(HostRole.READER).build(), -// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("reader5").port(1234).role(HostRole.READER).build() -// ); -// -// @BeforeEach -// void setUp() { -// closeable = MockitoAnnotations.openMocks(this); -// when(mockContainer.getConnectionPluginManager()).thenReturn(mockPluginManager); -// when(mockContainer.getPluginService()).thenReturn(mockPluginService); -// } -// -// @AfterEach -// void tearDown() throws Exception { -// closeable.close(); -// } -// -// @Test -// public void testFailover() throws SQLException { -// // original host list: [active writer, active reader, current connection (reader), active -// // reader, down reader, active reader] -// // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] -// // connection attempts are made in pairs using the above list -// // expected test result: successful connection for host at index 4 -// final List hosts = defaultHosts; -// final int currentHostIndex = 2; -// final int successHostIndex = 4; -// for (int i = 0; i < hosts.size(); i++) { -// if (i != successHostIndex) { -// final SQLException exception = new SQLException("exception", "08S01", null); -// when(mockConnectionService.open(hosts.get(i), properties)) -// .thenThrow(exception); -// when(mockPluginService.isNetworkException(exception, null)).thenReturn(true); -// } else { -// when(mockConnectionService.open(hosts.get(i), properties)).thenReturn(mockConnection); -// } -// } -// -// when(mockPluginService.getTargetDriverDialect()).thenReturn(null); -// -// hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); -// hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); -// -// final ReaderFailoverHandler target = getSpyFailoverHandler(); -// final ReaderFailoverResult result = target.failover(hosts, hosts.get(currentHostIndex)); -// -// assertTrue(result.isConnected()); -// assertSame(mockConnection, result.getConnection()); -// assertEquals(hosts.get(successHostIndex), result.getHost()); -// -// final HostSpec successHost = hosts.get(successHostIndex); -// final Map availabilityMap = target.getHostAvailabilityMap(); -// Set unavailableHosts = getHostsWithGivenAvailability(availabilityMap, HostAvailability.NOT_AVAILABLE); -// assertTrue(unavailableHosts.size() >= 4); -// assertEquals(HostAvailability.AVAILABLE, availabilityMap.get(successHost.getHost())); -// } -// -// private Set getHostsWithGivenAvailability( -// Map availabilityMap, HostAvailability availability) { -// return availabilityMap.entrySet().stream() -// .filter((entry) -> availability.equals(entry.getValue())) -// .map(Map.Entry::getKey) -// .collect(Collectors.toSet()); -// } -// -// @Test -// public void testFailover_timeout() throws SQLException { -// // original host list: [active writer, active reader, current connection (reader), active -// // reader, down reader, active reader] -// // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] -// // connection attempts are made in pairs using the above list -// // expected test result: failure to get reader since process is limited to 5s and each attempt -// // to connect takes 20s -// final List hosts = defaultHosts; -// final int currentHostIndex = 2; -// for (HostSpec host : hosts) { -// when(mockConnectionService.open(host, properties)) -// .thenAnswer((Answer) invocation -> { -// Thread.sleep(20000); -// return mockConnection; -// }); -// } -// -// hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); -// hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); -// -// final ReaderFailoverHandler target = getSpyFailoverHandler(5000, 30000, false); -// -// final long startTimeNano = System.nanoTime(); -// final ReaderFailoverResult result = target.failover(hosts, hosts.get(currentHostIndex)); -// final long durationNano = System.nanoTime() - startTimeNano; -// -// assertFalse(result.isConnected()); -// assertNull(result.getConnection()); -// assertNull(result.getHost()); -// -// // 5s is a max allowed failover timeout; add 1s for inaccurate measurements -// assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); -// } -// -// private ClusterAwareReaderFailoverHandler getSpyFailoverHandler() throws SQLException { -// ClusterAwareReaderFailoverHandler handler = -// spy(new ClusterAwareReaderFailoverHandler(mockContainer, mockConnectionService, properties)); -// doReturn(mockPluginService).when(handler).getNewPluginService(); -// return handler; -// } -// -// private ClusterAwareReaderFailoverHandler getSpyFailoverHandler( -// int maxFailoverTimeoutMs, int timeoutMs, boolean isStrictReaderRequired) throws SQLException { -// ClusterAwareReaderFailoverHandler handler = new ClusterAwareReaderFailoverHandler( -// mockContainer, mockConnectionService, properties, maxFailoverTimeoutMs, timeoutMs, isStrictReaderRequired); -// ClusterAwareReaderFailoverHandler spyHandler = spy(handler); -// doReturn(mockPluginService).when(spyHandler).getNewPluginService(); -// return spyHandler; -// } -// -// @Test -// public void testFailover_nullOrEmptyHostList() throws SQLException { -// final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); -// final HostSpec currentHost = -// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("writer").port(1234).build(); -// -// ReaderFailoverResult result = target.failover(null, currentHost); -// assertFalse(result.isConnected()); -// assertNull(result.getConnection()); -// assertNull(result.getHost()); -// -// final List hosts = new ArrayList<>(); -// result = target.failover(hosts, currentHost); -// assertFalse(result.isConnected()); -// assertNull(result.getConnection()); -// assertNull(result.getHost()); -// } -// -// @Test -// public void testGetReader_connectionSuccess() throws SQLException { -// // even number of connection attempts -// // first connection attempt to return succeeds, second attempt cancelled -// // expected test result: successful connection for host at index 2 -// final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) -// final HostSpec slowHost = hosts.get(1); -// final HostSpec fastHost = hosts.get(2); -// when(mockConnectionService.open(slowHost, properties)) -// .thenAnswer( -// (Answer) -// invocation -> { -// Thread.sleep(20000); -// return mockConnection; -// }); -// when(mockConnectionService.open(eq(fastHost), eq(properties))).thenReturn(mockConnection); -// -// Dialect mockDialect = Mockito.mock(Dialect.class); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// -// final ReaderFailoverHandler target = getSpyFailoverHandler(); -// final ReaderFailoverResult result = target.getReaderConnection(hosts); -// -// assertTrue(result.isConnected()); -// assertSame(mockConnection, result.getConnection()); -// assertEquals(hosts.get(2), result.getHost()); -// -// Map availabilityMap = target.getHostAvailabilityMap(); -// assertTrue(getHostsWithGivenAvailability(availabilityMap, HostAvailability.NOT_AVAILABLE).isEmpty()); -// assertEquals(HostAvailability.AVAILABLE, availabilityMap.get(fastHost.getHost())); -// } -// -// @Test -// public void testGetReader_connectionFailure() throws SQLException { -// // odd number of connection attempts -// // first connection attempt to return fails -// // expected test result: failure to get reader -// final List hosts = defaultHosts.subList(0, 4); // 3 connection attempts (writer not attempted) -// when(mockConnectionService.open(any(), eq(properties))).thenThrow(new SQLException("exception", "08S01", null)); -// -// Dialect mockDialect = Mockito.mock(Dialect.class); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// -// final ReaderFailoverHandler target = getSpyFailoverHandler(); -// final ReaderFailoverResult result = target.getReaderConnection(hosts); -// -// assertFalse(result.isConnected()); -// assertNull(result.getConnection()); -// assertNull(result.getHost()); -// } -// -// @Test -// public void testGetReader_connectionAttemptsTimeout() throws SQLException { -// // connection attempts time out before they can succeed -// // first connection attempt to return times out -// // expected test result: failure to get reader -// final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) -// when(mockConnectionService.open(any(), eq(properties))) -// .thenAnswer( -// (Answer) -// invocation -> { -// try { -// Thread.sleep(5000); -// } catch (InterruptedException exception) { -// // ignore -// } -// return mockConnection; -// }); -// -// Dialect mockDialect = Mockito.mock(Dialect.class); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// -// final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(60000, 1000, false); -// final ReaderFailoverResult result = target.getReaderConnection(hosts); -// -// assertFalse(result.isConnected()); -// assertNull(result.getConnection()); -// assertNull(result.getHost()); -// } -// -// @Test -// public void testGetHostTuplesByPriority() throws SQLException { -// final List originalHosts = defaultHosts; -// originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); -// originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); -// originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); -// -// final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); -// final List hostsByPriority = target.getHostsByPriority(originalHosts); -// -// int i = 0; -// -// // expecting active readers -// while (i < hostsByPriority.size() -// && hostsByPriority.get(i).getRole() == HostRole.READER -// && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { -// i++; -// } -// -// // expecting a writer -// while (i < hostsByPriority.size() -// && hostsByPriority.get(i).getRole() == HostRole.WRITER) { -// i++; -// } -// -// // expecting down readers -// while (i < hostsByPriority.size() -// && hostsByPriority.get(i).getRole() == HostRole.READER -// && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { -// i++; -// } -// -// assertEquals(hostsByPriority.size(), i); -// } -// -// @Test -// public void testGetReaderTuplesByPriority() throws SQLException { -// final List originalHosts = defaultHosts; -// originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); -// originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); -// originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); -// -// Dialect mockDialect = Mockito.mock(Dialect.class); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// -// final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); -// final List hostsByPriority = target.getReaderHostsByPriority(originalHosts); -// -// int i = 0; -// -// // expecting active readers -// while (i < hostsByPriority.size() -// && hostsByPriority.get(i).getRole() == HostRole.READER -// && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { -// i++; -// } -// -// // expecting down readers -// while (i < hostsByPriority.size() -// && hostsByPriority.get(i).getRole() == HostRole.READER -// && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { -// i++; -// } -// -// assertEquals(hostsByPriority.size(), i); -// } -// -// @Test -// public void testHostFailoverStrictReaderEnabled() throws SQLException { -// final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("writer").port(1234).role(HostRole.WRITER).build(); -// final HostSpec reader = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("reader1").port(1234).role(HostRole.READER).build(); -// final List hosts = Arrays.asList(writer, reader); -// -// Dialect mockDialect = Mockito.mock(Dialect.class); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// -// final ClusterAwareReaderFailoverHandler target = -// getSpyFailoverHandler(DEFAULT_FAILOVER_TIMEOUT, DEFAULT_READER_CONNECT_TIMEOUT, true); -// -// // The writer is included because the original writer has likely become a reader. -// List expectedHostsByPriority = Arrays.asList(reader, writer); -// -// List hostsByPriority = target.getHostsByPriority(hosts); -// assertEquals(expectedHostsByPriority, hostsByPriority); -// -// // Should pick the reader even if unavailable. The unavailable reader will be lower priority than the writer. -// reader.setAvailability(HostAvailability.NOT_AVAILABLE); -// expectedHostsByPriority = Arrays.asList(writer, reader); -// -// hostsByPriority = target.getHostsByPriority(hosts); -// assertEquals(expectedHostsByPriority, hostsByPriority); -// -// // Writer node will only be picked if it is the only node in topology; -// List expectedWriterHost = Collections.singletonList(writer); -// -// hostsByPriority = target.getHostsByPriority(Collections.singletonList(writer)); -// assertEquals(expectedWriterHost, hostsByPriority); -// } -// } +/* + * 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.failover; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.when; +import static software.amazon.jdbc.plugin.failover.ClusterAwareReaderFailoverHandler.DEFAULT_FAILOVER_TIMEOUT; +import static software.amazon.jdbc.plugin.failover.ClusterAwareReaderFailoverHandler.DEFAULT_READER_CONNECT_TIMEOUT; + +import java.sql.Connection; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.EnumSet; +import java.util.List; +import java.util.Map; +import java.util.Properties; +import java.util.Set; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.MockitoAnnotations; +import org.mockito.stubbing.Answer; +import software.amazon.jdbc.HostRole; +import software.amazon.jdbc.HostSpec; +import software.amazon.jdbc.HostSpecBuilder; +import software.amazon.jdbc.PluginService; +import software.amazon.jdbc.dialect.Dialect; +import software.amazon.jdbc.hostavailability.HostAvailability; +import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; +import software.amazon.jdbc.util.FullServicesContainer; + +class ClusterAwareReaderFailoverHandlerTest { + @Mock FullServicesContainer mockContainer1; + @Mock FullServicesContainer mockContainer2; + @Mock PluginService mockPluginService; + @Mock Connection mockConnection; + + private AutoCloseable closeable; + private final Properties properties = new Properties(); + private final List defaultHosts = Arrays.asList( + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("writer").port(1234).role(HostRole.WRITER).build(), + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("reader1").port(1234).role(HostRole.READER).build(), + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("reader2").port(1234).role(HostRole.READER).build(), + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("reader3").port(1234).role(HostRole.READER).build(), + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("reader4").port(1234).role(HostRole.READER).build(), + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("reader5").port(1234).role(HostRole.READER).build() + ); + + @BeforeEach + void setUp() { + closeable = MockitoAnnotations.openMocks(this); + when(mockContainer1.getPluginService()).thenReturn(mockPluginService); + when(mockContainer2.getPluginService()).thenReturn(mockPluginService); + } + + @AfterEach + void tearDown() throws Exception { + closeable.close(); + } + + @Test + public void testFailover() throws SQLException { + // original host list: [active writer, active reader, current connection (reader), active + // reader, down reader, active reader] + // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] + // connection attempts are made in pairs using the above list + // expected test result: successful connection for host at index 4 + final List hosts = defaultHosts; + final int currentHostIndex = 2; + final int successHostIndex = 4; + for (int i = 0; i < hosts.size(); i++) { + if (i != successHostIndex) { + final SQLException exception = new SQLException("exception", "08S01", null); + when(mockPluginService.forceConnect(hosts.get(i), properties)) + .thenThrow(exception); + when(mockPluginService.isNetworkException(exception, null)).thenReturn(true); + } else { + when(mockPluginService.forceConnect(hosts.get(i), properties)).thenReturn(mockConnection); + } + } + + when(mockPluginService.getTargetDriverDialect()).thenReturn(null); + + hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); + hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); + + final ReaderFailoverHandler target = getSpyFailoverHandler(); + final ReaderFailoverResult result = target.failover(hosts, hosts.get(currentHostIndex)); + + assertTrue(result.isConnected()); + assertSame(mockConnection, result.getConnection()); + assertEquals(hosts.get(successHostIndex), result.getHost()); + + final HostSpec successHost = hosts.get(successHostIndex); + final Map availabilityMap = target.getHostAvailabilityMap(); + Set unavailableHosts = getHostsWithGivenAvailability(availabilityMap, HostAvailability.NOT_AVAILABLE); + assertTrue(unavailableHosts.size() >= 4); + assertEquals(HostAvailability.AVAILABLE, availabilityMap.get(successHost.getHost())); + } + + private Set getHostsWithGivenAvailability( + Map availabilityMap, HostAvailability availability) { + return availabilityMap.entrySet().stream() + .filter((entry) -> availability.equals(entry.getValue())) + .map(Map.Entry::getKey) + .collect(Collectors.toSet()); + } + + @Test + public void testFailover_timeout() throws SQLException { + // original host list: [active writer, active reader, current connection (reader), active + // reader, down reader, active reader] + // priority order by index (the subsets will be shuffled): [[1, 3, 5], 0, [2, 4]] + // connection attempts are made in pairs using the above list + // expected test result: failure to get reader since process is limited to 5s and each attempt + // to connect takes 20s + final List hosts = defaultHosts; + final int currentHostIndex = 2; + for (HostSpec host : hosts) { + when(mockPluginService.forceConnect(host, properties)) + .thenAnswer((Answer) invocation -> { + Thread.sleep(20000); + return mockConnection; + }); + } + + hosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); + hosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); + + final ReaderFailoverHandler target = getSpyFailoverHandler(5000, 30000, false); + + final long startTimeNano = System.nanoTime(); + final ReaderFailoverResult result = target.failover(hosts, hosts.get(currentHostIndex)); + final long durationNano = System.nanoTime() - startTimeNano; + + assertFalse(result.isConnected()); + assertNull(result.getConnection()); + assertNull(result.getHost()); + + // 5s is a max allowed failover timeout; add 1s for inaccurate measurements + assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); + } + + private ClusterAwareReaderFailoverHandler getSpyFailoverHandler() throws SQLException { + ClusterAwareReaderFailoverHandler handler = + spy(new ClusterAwareReaderFailoverHandler(mockContainer1, properties)); + doReturn(mockContainer2).when(handler).getNewServicesContainer(); + return handler; + } + + private ClusterAwareReaderFailoverHandler getSpyFailoverHandler( + int maxFailoverTimeoutMs, int timeoutMs, boolean isStrictReaderRequired) throws SQLException { + ClusterAwareReaderFailoverHandler handler = new ClusterAwareReaderFailoverHandler( + mockContainer1, properties, maxFailoverTimeoutMs, timeoutMs, isStrictReaderRequired); + ClusterAwareReaderFailoverHandler spyHandler = spy(handler); + doReturn(mockContainer2).when(spyHandler).getNewServicesContainer(); + return spyHandler; + } + + @Test + public void testFailover_nullOrEmptyHostList() throws SQLException { + final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); + final HostSpec currentHost = + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("writer").port(1234).build(); + + ReaderFailoverResult result = target.failover(null, currentHost); + assertFalse(result.isConnected()); + assertNull(result.getConnection()); + assertNull(result.getHost()); + + final List hosts = new ArrayList<>(); + result = target.failover(hosts, currentHost); + assertFalse(result.isConnected()); + assertNull(result.getConnection()); + assertNull(result.getHost()); + } + + @Test + public void testGetReader_connectionSuccess() throws SQLException { + // even number of connection attempts + // first connection attempt to return succeeds, second attempt cancelled + // expected test result: successful connection for host at index 2 + final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) + final HostSpec slowHost = hosts.get(1); + final HostSpec fastHost = hosts.get(2); + when(mockPluginService.forceConnect(slowHost, properties)) + .thenAnswer( + (Answer) + invocation -> { + Thread.sleep(20000); + return mockConnection; + }); + when(mockPluginService.forceConnect(eq(fastHost), eq(properties))).thenReturn(mockConnection); + + Dialect mockDialect = Mockito.mock(Dialect.class); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + when(mockPluginService.getDialect()).thenReturn(mockDialect); + + final ReaderFailoverHandler target = getSpyFailoverHandler(); + final ReaderFailoverResult result = target.getReaderConnection(hosts); + + assertTrue(result.isConnected()); + assertSame(mockConnection, result.getConnection()); + assertEquals(hosts.get(2), result.getHost()); + + Map availabilityMap = target.getHostAvailabilityMap(); + assertTrue(getHostsWithGivenAvailability(availabilityMap, HostAvailability.NOT_AVAILABLE).isEmpty()); + assertEquals(HostAvailability.AVAILABLE, availabilityMap.get(fastHost.getHost())); + } + + @Test + public void testGetReader_connectionFailure() throws SQLException { + // odd number of connection attempts + // first connection attempt to return fails + // expected test result: failure to get reader + final List hosts = defaultHosts.subList(0, 4); // 3 connection attempts (writer not attempted) + when(mockPluginService.forceConnect(any(), eq(properties))).thenThrow(new SQLException("exception", "08S01", null)); + + Dialect mockDialect = Mockito.mock(Dialect.class); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + when(mockPluginService.getDialect()).thenReturn(mockDialect); + + final ReaderFailoverHandler target = getSpyFailoverHandler(); + final ReaderFailoverResult result = target.getReaderConnection(hosts); + + assertFalse(result.isConnected()); + assertNull(result.getConnection()); + assertNull(result.getHost()); + } + + @Test + public void testGetReader_connectionAttemptsTimeout() throws SQLException { + // connection attempts time out before they can succeed + // first connection attempt to return times out + // expected test result: failure to get reader + final List hosts = defaultHosts.subList(0, 3); // 2 connection attempts (writer not attempted) + when(mockPluginService.forceConnect(any(), eq(properties))) + .thenAnswer( + (Answer) + invocation -> { + try { + Thread.sleep(5000); + } catch (InterruptedException exception) { + // ignore + } + return mockConnection; + }); + + Dialect mockDialect = Mockito.mock(Dialect.class); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + when(mockPluginService.getDialect()).thenReturn(mockDialect); + + final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(60000, 1000, false); + final ReaderFailoverResult result = target.getReaderConnection(hosts); + + assertFalse(result.isConnected()); + assertNull(result.getConnection()); + assertNull(result.getHost()); + } + + @Test + public void testGetHostTuplesByPriority() throws SQLException { + final List originalHosts = defaultHosts; + originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); + originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); + originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); + + final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); + final List hostsByPriority = target.getHostsByPriority(originalHosts); + + int i = 0; + + // expecting active readers + while (i < hostsByPriority.size() + && hostsByPriority.get(i).getRole() == HostRole.READER + && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { + i++; + } + + // expecting a writer + while (i < hostsByPriority.size() + && hostsByPriority.get(i).getRole() == HostRole.WRITER) { + i++; + } + + // expecting down readers + while (i < hostsByPriority.size() + && hostsByPriority.get(i).getRole() == HostRole.READER + && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { + i++; + } + + assertEquals(hostsByPriority.size(), i); + } + + @Test + public void testGetReaderTuplesByPriority() throws SQLException { + final List originalHosts = defaultHosts; + originalHosts.get(2).setAvailability(HostAvailability.NOT_AVAILABLE); + originalHosts.get(4).setAvailability(HostAvailability.NOT_AVAILABLE); + originalHosts.get(5).setAvailability(HostAvailability.NOT_AVAILABLE); + + Dialect mockDialect = Mockito.mock(Dialect.class); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + when(mockPluginService.getDialect()).thenReturn(mockDialect); + + final ClusterAwareReaderFailoverHandler target = getSpyFailoverHandler(); + final List hostsByPriority = target.getReaderHostsByPriority(originalHosts); + + int i = 0; + + // expecting active readers + while (i < hostsByPriority.size() + && hostsByPriority.get(i).getRole() == HostRole.READER + && hostsByPriority.get(i).getAvailability() == HostAvailability.AVAILABLE) { + i++; + } + + // expecting down readers + while (i < hostsByPriority.size() + && hostsByPriority.get(i).getRole() == HostRole.READER + && hostsByPriority.get(i).getAvailability() == HostAvailability.NOT_AVAILABLE) { + i++; + } + + assertEquals(hostsByPriority.size(), i); + } + + @Test + public void testHostFailoverStrictReaderEnabled() throws SQLException { + final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("writer").port(1234).role(HostRole.WRITER).build(); + final HostSpec reader = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("reader1").port(1234).role(HostRole.READER).build(); + final List hosts = Arrays.asList(writer, reader); + + Dialect mockDialect = Mockito.mock(Dialect.class); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + when(mockPluginService.getDialect()).thenReturn(mockDialect); + + final ClusterAwareReaderFailoverHandler target = + getSpyFailoverHandler(DEFAULT_FAILOVER_TIMEOUT, DEFAULT_READER_CONNECT_TIMEOUT, true); + + // The writer is included because the original writer has likely become a reader. + List expectedHostsByPriority = Arrays.asList(reader, writer); + + List hostsByPriority = target.getHostsByPriority(hosts); + assertEquals(expectedHostsByPriority, hostsByPriority); + + // Should pick the reader even if unavailable. The unavailable reader will have lower priority than the writer. + reader.setAvailability(HostAvailability.NOT_AVAILABLE); + expectedHostsByPriority = Arrays.asList(writer, reader); + + hostsByPriority = target.getHostsByPriority(hosts); + assertEquals(expectedHostsByPriority, hostsByPriority); + + // Writer node will only be picked if it is the only node in topology; + List expectedWriterHost = Collections.singletonList(writer); + + hostsByPriority = target.getHostsByPriority(Collections.singletonList(writer)); + assertEquals(expectedWriterHost, hostsByPriority); + } +} diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java index 98d392cdb..b17ecbb95 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/ClusterAwareWriterFailoverHandlerTest.java @@ -1,373 +1,372 @@ -// /* -// * 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.failover; -// -// import static org.junit.jupiter.api.Assertions.assertEquals; -// import static org.junit.jupiter.api.Assertions.assertFalse; -// import static org.junit.jupiter.api.Assertions.assertSame; -// import static org.junit.jupiter.api.Assertions.assertTrue; -// import static org.mockito.ArgumentMatchers.any; -// import static org.mockito.ArgumentMatchers.eq; -// import static org.mockito.ArgumentMatchers.refEq; -// import static org.mockito.Mockito.atLeastOnce; -// import static org.mockito.Mockito.doReturn; -// import static org.mockito.Mockito.spy; -// import static org.mockito.Mockito.verify; -// import static org.mockito.Mockito.when; -// -// import java.sql.Connection; -// import java.sql.SQLException; -// import java.util.Arrays; -// import java.util.EnumSet; -// import java.util.List; -// import java.util.Properties; -// import java.util.concurrent.TimeUnit; -// import org.junit.jupiter.api.AfterEach; -// import org.junit.jupiter.api.BeforeEach; -// import org.junit.jupiter.api.Test; -// import org.mockito.ArgumentMatchers; -// import org.mockito.Mock; -// import org.mockito.MockitoAnnotations; -// import org.mockito.stubbing.Answer; -// import software.amazon.jdbc.HostSpec; -// import software.amazon.jdbc.HostSpecBuilder; -// import software.amazon.jdbc.PluginService; -// import software.amazon.jdbc.dialect.Dialect; -// import software.amazon.jdbc.hostavailability.HostAvailability; -// import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; -// import software.amazon.jdbc.util.FullServicesContainer; -// import software.amazon.jdbc.util.connection.ConnectionService; -// -// class ClusterAwareWriterFailoverHandlerTest { -// @Mock FullServicesContainer mockContainer; -// @Mock ConnectionService mockConnectionService; -// @Mock PluginService mockPluginService; -// @Mock Connection mockConnection; -// @Mock ReaderFailoverHandler mockReaderFailoverHandler; -// @Mock Connection mockWriterConnection; -// @Mock Connection mockNewWriterConnection; -// @Mock Connection mockReaderAConnection; -// @Mock Connection mockReaderBConnection; -// @Mock Dialect mockDialect; -// -// private AutoCloseable closeable; -// private final Properties properties = new Properties(); -// private final HostSpec newWriterHost = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("new-writer-host").build(); -// private final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("writer-host").build(); -// private final HostSpec readerA = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("reader-a-host").build(); -// private final HostSpec readerB = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("reader-b-host").build(); -// private final List topology = Arrays.asList(writer, readerA, readerB); -// private final List newTopology = Arrays.asList(newWriterHost, readerA, readerB); -// -// @BeforeEach -// void setUp() { -// closeable = MockitoAnnotations.openMocks(this); -// when(mockContainer.getPluginService()).thenReturn(mockPluginService); -// writer.addAlias("writer-host"); -// newWriterHost.addAlias("new-writer-host"); -// readerA.addAlias("reader-a-host"); -// readerB.addAlias("reader-b-host"); -// } -// -// @AfterEach -// void tearDown() throws Exception { -// closeable.close(); -// } -// -// @Test -// public void testReconnectToWriter_taskBReaderException() throws SQLException { -// when(mockConnectionService.open(refEq(writer), eq(properties))).thenReturn(mockConnection); -// when(mockConnectionService.open(refEq(readerA), eq(properties))).thenThrow(SQLException.class); -// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); -// -// when(mockPluginService.getAllHosts()).thenReturn(topology); -// -// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())).thenThrow(SQLException.class); -// -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// -// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); -// final WriterFailoverResult result = target.failover(topology); -// -// assertTrue(result.isConnected()); -// assertFalse(result.isNewHost()); -// assertSame(result.getNewConnection(), mockConnection); -// -// assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); -// } -// -// private ClusterAwareWriterFailoverHandler getSpyFailoverHandler( -// final int failoverTimeoutMs, -// final int readTopologyIntervalMs, -// final int reconnectWriterIntervalMs) throws SQLException { -// ClusterAwareWriterFailoverHandler handler = new ClusterAwareWriterFailoverHandler( -// mockContainer, -// mockConnectionService, -// mockReaderFailoverHandler, -// properties, -// failoverTimeoutMs, -// readTopologyIntervalMs, -// reconnectWriterIntervalMs); -// -// ClusterAwareWriterFailoverHandler spyHandler = spy(handler); -// doReturn(mockPluginService).when(spyHandler).getNewPluginService(); -// return spyHandler; -// } -// -// /** -// * Verify that writer failover handler can re-connect to a current writer node. -// * -// *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. -// * TaskA: successfully re-connect to initial writer; return new connection. -// * TaskB: successfully connect to readerA and then new writer, but it takes more time than taskA. -// * Expected test result: new connection by taskA. -// */ -// @Test -// public void testReconnectToWriter_SlowReaderA() throws SQLException { -// when(mockConnectionService.open(refEq(writer), eq(properties))).thenReturn(mockWriterConnection); -// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); -// when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); -// when(mockPluginService.getAllHosts()).thenReturn(topology).thenReturn(newTopology); -// -// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) -// .thenAnswer( -// (Answer) -// invocation -> { -// Thread.sleep(5000); -// return new ReaderFailoverResult(mockReaderAConnection, readerA, true); -// }); -// -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// -// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); -// final WriterFailoverResult result = target.failover(topology); -// -// assertTrue(result.isConnected()); -// assertFalse(result.isNewHost()); -// assertSame(result.getNewConnection(), mockWriterConnection); -// assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); -// } -// -// /** -// * Verify that writer failover handler can re-connect to a current writer node. -// * -// *

Topology: no changes. -// * TaskA: successfully re-connect to writer; return new connection. -// * TaskB: successfully connect to readerA and retrieve topology, but latest writer is not new (defer to taskA). -// * Expected test result: new connection by taskA. -// */ -// @Test -// public void testReconnectToWriter_taskBDefers() throws SQLException { -// when(mockConnectionService.open(refEq(writer), eq(properties))) -// .thenAnswer( -// (Answer) -// invocation -> { -// Thread.sleep(5000); -// return mockWriterConnection; -// }); -// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenThrow(SQLException.class); -// -// when(mockPluginService.getAllHosts()).thenReturn(topology); -// -// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) -// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); -// -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// -// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 2000, 2000); -// final WriterFailoverResult result = target.failover(topology); -// -// assertTrue(result.isConnected()); -// assertFalse(result.isNewHost()); -// assertSame(result.getNewConnection(), mockWriterConnection); -// assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); -// } -// -// /** -// * Verify that writer failover handler can re-connect to a new writer node. -// * -// *

Topology: changes to [new-writer, reader-A, reader-B] for taskB, taskA sees no changes. -// * taskA: successfully re-connect to writer; return connection to initial writer, but it takes more -// * time than taskB. -// * TaskB: successfully connect to readerA and then to new-writer. -// * Expected test result: new connection to writer by taskB. -// */ -// @Test -// public void testConnectToReaderA_SlowWriter() throws SQLException { -// when(mockConnectionService.open(refEq(writer), eq(properties))) -// .thenAnswer( -// (Answer) -// invocation -> { -// Thread.sleep(5000); -// return mockWriterConnection; -// }); -// when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); -// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); -// when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); -// -// when(mockPluginService.getAllHosts()).thenReturn(newTopology); -// -// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) -// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); -// -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// -// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); -// final WriterFailoverResult result = target.failover(topology); -// -// assertTrue(result.isConnected()); -// assertTrue(result.isNewHost()); -// assertSame(result.getNewConnection(), mockNewWriterConnection); -// assertEquals(3, result.getTopology().size()); -// assertEquals("new-writer-host", result.getTopology().get(0).getHost()); -// assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); -// } -// -// /** -// * Verify that writer failover handler can re-connect to a new writer node. -// * -// *

Topology: changes to [new-writer, initial-writer, reader-A, reader-B]. -// * TaskA: successfully reconnect, but initial-writer is now a reader (defer to taskB). -// * TaskB: successfully connect to readerA and then to new-writer. -// * Expected test result: new connection to writer by taskB. -// */ -// @Test -// public void testConnectToReaderA_taskADefers() throws SQLException { -// when(mockConnectionService.open(writer, properties)).thenReturn(mockConnection); -// when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); -// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); -// when(mockConnectionService.open(refEq(newWriterHost), eq(properties))) -// .thenAnswer( -// (Answer) -// invocation -> { -// Thread.sleep(5000); -// return mockNewWriterConnection; -// }); -// -// final List newTopology = Arrays.asList(newWriterHost, writer, readerA, readerB); -// when(mockPluginService.getAllHosts()).thenReturn(newTopology); -// -// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) -// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); -// -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// -// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); -// final WriterFailoverResult result = target.failover(topology); -// -// assertTrue(result.isConnected()); -// assertTrue(result.isNewHost()); -// assertSame(result.getNewConnection(), mockNewWriterConnection); -// assertEquals(4, result.getTopology().size()); -// assertEquals("new-writer-host", result.getTopology().get(0).getHost()); -// -// verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); -// assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); -// } -// -// /** -// * Verify that writer failover handler fails to re-connect to any writer node. -// * -// *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. -// * TaskA: fail to re-connect to writer due to failover timeout. -// * TaskB: successfully connect to readerA and then fail to connect to writer due to failover timeout. -// * Expected test result: no connection. -// */ -// @Test -// public void testFailedToConnect_failoverTimeout() throws SQLException { -// when(mockConnectionService.open(refEq(writer), eq(properties))) -// .thenAnswer( -// (Answer) -// invocation -> { -// Thread.sleep(30000); -// return mockWriterConnection; -// }); -// when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); -// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); -// when(mockConnectionService.open(refEq(newWriterHost), eq(properties))) -// .thenAnswer( -// (Answer) -// invocation -> { -// Thread.sleep(30000); -// return mockNewWriterConnection; -// }); -// when(mockPluginService.getAllHosts()).thenReturn(newTopology); -// -// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) -// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); -// -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// -// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); -// -// final long startTimeNano = System.nanoTime(); -// final WriterFailoverResult result = target.failover(topology); -// final long durationNano = System.nanoTime() - startTimeNano; -// -// assertFalse(result.isConnected()); -// assertFalse(result.isNewHost()); -// -// verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); -// -// // 5s is a max allowed failover timeout; add 1s for inaccurate measurements -// assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); -// } -// -// /** -// * Verify that writer failover handler fails to re-connect to any writer node. -// * -// *

Topology: changes to [new-writer, reader-A, reader-B] for taskB. -// * TaskA: fail to re-connect to writer due to exception. -// * TaskB: successfully connect to readerA and then fail to connect to writer due to exception. -// * Expected test result: no connection. -// */ -// @Test -// public void testFailedToConnect_taskAException_taskBWriterException() throws SQLException { -// final SQLException exception = new SQLException("exception", "08S01", null); -// when(mockConnectionService.open(refEq(writer), eq(properties))).thenThrow(exception); -// when(mockConnectionService.open(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); -// when(mockConnectionService.open(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); -// when(mockConnectionService.open(refEq(newWriterHost), eq(properties))).thenThrow(exception); -// when(mockPluginService.isNetworkException(eq(exception), any())).thenReturn(true); -// -// when(mockPluginService.getAllHosts()).thenReturn(newTopology); -// -// when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) -// .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); -// -// when(mockPluginService.getDialect()).thenReturn(mockDialect); -// when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); -// -// final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); -// final WriterFailoverResult result = target.failover(topology); -// -// assertFalse(result.isConnected()); -// assertFalse(result.isNewHost()); -// -// assertEquals(HostAvailability.NOT_AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); -// } -// } +/* + * 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.failover; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.refEq; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.sql.Connection; +import java.sql.SQLException; +import java.util.Arrays; +import java.util.EnumSet; +import java.util.List; +import java.util.Properties; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentMatchers; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import org.mockito.stubbing.Answer; +import software.amazon.jdbc.HostSpec; +import software.amazon.jdbc.HostSpecBuilder; +import software.amazon.jdbc.PluginService; +import software.amazon.jdbc.dialect.Dialect; +import software.amazon.jdbc.hostavailability.HostAvailability; +import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; +import software.amazon.jdbc.util.FullServicesContainer; + +class ClusterAwareWriterFailoverHandlerTest { + @Mock FullServicesContainer mockContainer1; + @Mock FullServicesContainer mockContainer2; + @Mock PluginService mockPluginService; + @Mock Connection mockConnection; + @Mock ReaderFailoverHandler mockReaderFailoverHandler; + @Mock Connection mockWriterConnection; + @Mock Connection mockNewWriterConnection; + @Mock Connection mockReaderAConnection; + @Mock Connection mockReaderBConnection; + @Mock Dialect mockDialect; + + private AutoCloseable closeable; + private final Properties properties = new Properties(); + private final HostSpec newWriterHost = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("new-writer-host").build(); + private final HostSpec writer = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("writer-host").build(); + private final HostSpec readerA = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("reader-a-host").build(); + private final HostSpec readerB = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("reader-b-host").build(); + private final List topology = Arrays.asList(writer, readerA, readerB); + private final List newTopology = Arrays.asList(newWriterHost, readerA, readerB); + + @BeforeEach + void setUp() { + closeable = MockitoAnnotations.openMocks(this); + when(mockContainer1.getPluginService()).thenReturn(mockPluginService); + when(mockContainer2.getPluginService()).thenReturn(mockPluginService); + writer.addAlias("writer-host"); + newWriterHost.addAlias("new-writer-host"); + readerA.addAlias("reader-a-host"); + readerB.addAlias("reader-b-host"); + } + + @AfterEach + void tearDown() throws Exception { + closeable.close(); + } + + @Test + public void testReconnectToWriter_taskBReaderException() throws SQLException { + when(mockPluginService.forceConnect(refEq(writer), eq(properties))).thenReturn(mockConnection); + when(mockPluginService.forceConnect(refEq(readerA), eq(properties))).thenThrow(SQLException.class); + when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenThrow(SQLException.class); + + when(mockPluginService.getAllHosts()).thenReturn(topology); + + when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())).thenThrow(SQLException.class); + + when(mockPluginService.getDialect()).thenReturn(mockDialect); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + + final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); + final WriterFailoverResult result = target.failover(topology); + + assertTrue(result.isConnected()); + assertFalse(result.isNewHost()); + assertSame(result.getNewConnection(), mockConnection); + + assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); + } + + private ClusterAwareWriterFailoverHandler getSpyFailoverHandler( + final int failoverTimeoutMs, + final int readTopologyIntervalMs, + final int reconnectWriterIntervalMs) throws SQLException { + ClusterAwareWriterFailoverHandler handler = new ClusterAwareWriterFailoverHandler( + mockContainer1, + mockReaderFailoverHandler, + properties, + failoverTimeoutMs, + readTopologyIntervalMs, + reconnectWriterIntervalMs); + + ClusterAwareWriterFailoverHandler spyHandler = spy(handler); + doReturn(mockContainer2).when(spyHandler).getNewServicesContainer(); + return spyHandler; + } + + /** + * Verify that writer failover handler can re-connect to a current writer node. + * + *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. + * TaskA: successfully re-connect to initial writer; return new connection. + * TaskB: successfully connect to readerA and then new writer, but it takes more time than taskA. + * Expected test result: new connection by taskA. + */ + @Test + public void testReconnectToWriter_SlowReaderA() throws SQLException { + when(mockPluginService.forceConnect(refEq(writer), eq(properties))).thenReturn(mockWriterConnection); + when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenThrow(SQLException.class); + when(mockPluginService.forceConnect(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); + when(mockPluginService.getAllHosts()).thenReturn(topology).thenReturn(newTopology); + + when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) + .thenAnswer( + (Answer) + invocation -> { + Thread.sleep(5000); + return new ReaderFailoverResult(mockReaderAConnection, readerA, true); + }); + + when(mockPluginService.getDialect()).thenReturn(mockDialect); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + + final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); + final WriterFailoverResult result = target.failover(topology); + + assertTrue(result.isConnected()); + assertFalse(result.isNewHost()); + assertSame(result.getNewConnection(), mockWriterConnection); + assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); + } + + /** + * Verify that writer failover handler can re-connect to a current writer node. + * + *

Topology: no changes. + * TaskA: successfully re-connect to writer; return new connection. + * TaskB: successfully connect to readerA and retrieve topology, but latest writer is not new (defer to taskA). + * Expected test result: new connection by taskA. + */ + @Test + public void testReconnectToWriter_taskBDefers() throws SQLException { + when(mockPluginService.forceConnect(refEq(writer), eq(properties))) + .thenAnswer( + (Answer) + invocation -> { + Thread.sleep(5000); + return mockWriterConnection; + }); + when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenThrow(SQLException.class); + + when(mockPluginService.getAllHosts()).thenReturn(topology); + + when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) + .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); + + when(mockPluginService.getDialect()).thenReturn(mockDialect); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + + final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 2000, 2000); + final WriterFailoverResult result = target.failover(topology); + + assertTrue(result.isConnected()); + assertFalse(result.isNewHost()); + assertSame(result.getNewConnection(), mockWriterConnection); + assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(writer.getHost())); + } + + /** + * Verify that writer failover handler can re-connect to a new writer node. + * + *

Topology: changes to [new-writer, reader-A, reader-B] for taskB, taskA sees no changes. + * taskA: successfully re-connect to writer; return connection to initial writer, but it takes more + * time than taskB. + * TaskB: successfully connect to readerA and then to new-writer. + * Expected test result: new connection to writer by taskB. + */ + @Test + public void testConnectToReaderA_SlowWriter() throws SQLException { + when(mockPluginService.forceConnect(refEq(writer), eq(properties))) + .thenAnswer( + (Answer) + invocation -> { + Thread.sleep(5000); + return mockWriterConnection; + }); + when(mockPluginService.forceConnect(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); + when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); + when(mockPluginService.forceConnect(refEq(newWriterHost), eq(properties))).thenReturn(mockNewWriterConnection); + + when(mockPluginService.getAllHosts()).thenReturn(newTopology); + + when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) + .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); + + when(mockPluginService.getDialect()).thenReturn(mockDialect); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + + final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); + final WriterFailoverResult result = target.failover(topology); + + assertTrue(result.isConnected()); + assertTrue(result.isNewHost()); + assertSame(result.getNewConnection(), mockNewWriterConnection); + assertEquals(3, result.getTopology().size()); + assertEquals("new-writer-host", result.getTopology().get(0).getHost()); + assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); + } + + /** + * Verify that writer failover handler can re-connect to a new writer node. + * + *

Topology: changes to [new-writer, initial-writer, reader-A, reader-B]. + * TaskA: successfully reconnect, but initial-writer is now a reader (defer to taskB). + * TaskB: successfully connect to readerA and then to new-writer. + * Expected test result: new connection to writer by taskB. + */ + @Test + public void testConnectToReaderA_taskADefers() throws SQLException { + when(mockPluginService.forceConnect(writer, properties)).thenReturn(mockConnection); + when(mockPluginService.forceConnect(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); + when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); + when(mockPluginService.forceConnect(refEq(newWriterHost), eq(properties))) + .thenAnswer( + (Answer) + invocation -> { + Thread.sleep(5000); + return mockNewWriterConnection; + }); + + final List newTopology = Arrays.asList(newWriterHost, writer, readerA, readerB); + when(mockPluginService.getAllHosts()).thenReturn(newTopology); + + when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) + .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); + + when(mockPluginService.getDialect()).thenReturn(mockDialect); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + + final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(60000, 5000, 5000); + final WriterFailoverResult result = target.failover(topology); + + assertTrue(result.isConnected()); + assertTrue(result.isNewHost()); + assertSame(result.getNewConnection(), mockNewWriterConnection); + assertEquals(4, result.getTopology().size()); + assertEquals("new-writer-host", result.getTopology().get(0).getHost()); + + verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); + assertEquals(HostAvailability.AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); + } + + /** + * Verify that writer failover handler fails to re-connect to any writer node. + * + *

Topology: no changes seen by task A, changes to [new-writer, reader-A, reader-B] for taskB. + * TaskA: fail to re-connect to writer due to failover timeout. + * TaskB: successfully connect to readerA and then fail to connect to writer due to failover timeout. + * Expected test result: no connection. + */ + @Test + public void testFailedToConnect_failoverTimeout() throws SQLException { + when(mockPluginService.forceConnect(refEq(writer), eq(properties))) + .thenAnswer( + (Answer) + invocation -> { + Thread.sleep(30000); + return mockWriterConnection; + }); + when(mockPluginService.forceConnect(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); + when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); + when(mockPluginService.forceConnect(refEq(newWriterHost), eq(properties))) + .thenAnswer( + (Answer) + invocation -> { + Thread.sleep(30000); + return mockNewWriterConnection; + }); + when(mockPluginService.getAllHosts()).thenReturn(newTopology); + + when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) + .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); + + when(mockPluginService.getDialect()).thenReturn(mockDialect); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + + final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); + + final long startTimeNano = System.nanoTime(); + final WriterFailoverResult result = target.failover(topology); + final long durationNano = System.nanoTime() - startTimeNano; + + assertFalse(result.isConnected()); + assertFalse(result.isNewHost()); + + verify(mockPluginService, atLeastOnce()).forceRefreshHostList(any(Connection.class)); + + // 5s is a max allowed failover timeout; add 1s for inaccurate measurements + assertTrue(TimeUnit.NANOSECONDS.toMillis(durationNano) < 6000); + } + + /** + * Verify that writer failover handler fails to re-connect to any writer node. + * + *

Topology: changes to [new-writer, reader-A, reader-B] for taskB. + * TaskA: fail to re-connect to writer due to exception. + * TaskB: successfully connect to readerA and then fail to connect to writer due to exception. + * Expected test result: no connection. + */ + @Test + public void testFailedToConnect_taskAException_taskBWriterException() throws SQLException { + final SQLException exception = new SQLException("exception", "08S01", null); + when(mockPluginService.forceConnect(refEq(writer), eq(properties))).thenThrow(exception); + when(mockPluginService.forceConnect(refEq(readerA), eq(properties))).thenReturn(mockReaderAConnection); + when(mockPluginService.forceConnect(refEq(readerB), eq(properties))).thenReturn(mockReaderBConnection); + when(mockPluginService.forceConnect(refEq(newWriterHost), eq(properties))).thenThrow(exception); + when(mockPluginService.isNetworkException(eq(exception), any())).thenReturn(true); + + when(mockPluginService.getAllHosts()).thenReturn(newTopology); + + when(mockReaderFailoverHandler.getReaderConnection(ArgumentMatchers.anyList())) + .thenReturn(new ReaderFailoverResult(mockReaderAConnection, readerA, true)); + + when(mockPluginService.getDialect()).thenReturn(mockDialect); + when(mockDialect.getFailoverRestrictions()).thenReturn(EnumSet.noneOf(FailoverRestriction.class)); + + final ClusterAwareWriterFailoverHandler target = getSpyFailoverHandler(5000, 2000, 2000); + final WriterFailoverResult result = target.failover(topology); + + assertFalse(result.isConnected()); + assertFalse(result.isNewHost()); + + assertEquals(HostAvailability.NOT_AVAILABLE, target.getHostAvailabilityMap().get(newWriterHost.getHost())); + } +} diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java index 0ecddc7d0..8274235b9 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java @@ -1,447 +1,445 @@ -// /* -// * 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.failover; -// -// import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; -// import static org.junit.jupiter.api.Assertions.assertEquals; -// import static org.junit.jupiter.api.Assertions.assertThrows; -// import static org.mockito.ArgumentMatchers.any; -// import static org.mockito.ArgumentMatchers.anyString; -// import static org.mockito.ArgumentMatchers.eq; -// import static org.mockito.Mockito.atLeastOnce; -// import static org.mockito.Mockito.doNothing; -// import static org.mockito.Mockito.doReturn; -// import static org.mockito.Mockito.doThrow; -// import static org.mockito.Mockito.never; -// import static org.mockito.Mockito.spy; -// import static org.mockito.Mockito.times; -// import static org.mockito.Mockito.verify; -// import static org.mockito.Mockito.when; -// -// import java.sql.Connection; -// import java.sql.ResultSet; -// import java.sql.SQLException; -// import java.util.Arrays; -// import java.util.Collections; -// import java.util.EnumSet; -// import java.util.HashMap; -// import java.util.HashSet; -// import java.util.List; -// import java.util.Map; -// import java.util.Properties; -// import org.junit.jupiter.api.AfterEach; -// import org.junit.jupiter.api.BeforeEach; -// import org.junit.jupiter.api.Test; -// import org.junit.jupiter.params.ParameterizedTest; -// import org.junit.jupiter.params.provider.ValueSource; -// import org.mockito.Mock; -// import org.mockito.MockitoAnnotations; -// import software.amazon.jdbc.HostListProviderService; -// import software.amazon.jdbc.HostRole; -// import software.amazon.jdbc.HostSpec; -// import software.amazon.jdbc.HostSpecBuilder; -// import software.amazon.jdbc.JdbcCallable; -// import software.amazon.jdbc.NodeChangeOptions; -// import software.amazon.jdbc.PluginService; -// import software.amazon.jdbc.hostavailability.HostAvailability; -// import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; -// import software.amazon.jdbc.hostlistprovider.AuroraHostListProvider; -// import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; -// import software.amazon.jdbc.util.FullServicesContainer; -// import software.amazon.jdbc.util.RdsUrlType; -// import software.amazon.jdbc.util.SqlState; -// import software.amazon.jdbc.util.connection.ConnectionService; -// import software.amazon.jdbc.util.telemetry.GaugeCallable; -// import software.amazon.jdbc.util.telemetry.TelemetryContext; -// import software.amazon.jdbc.util.telemetry.TelemetryCounter; -// import software.amazon.jdbc.util.telemetry.TelemetryFactory; -// import software.amazon.jdbc.util.telemetry.TelemetryGauge; -// -// class FailoverConnectionPluginTest { -// -// private static final Class MONITOR_METHOD_INVOKE_ON = Connection.class; -// private static final String MONITOR_METHOD_NAME = "Connection.executeQuery"; -// private static final Object[] EMPTY_ARGS = {}; -// private final List defaultHosts = Arrays.asList( -// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("writer").port(1234).role(HostRole.WRITER).build(), -// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) -// .host("reader1").port(1234).role(HostRole.READER).build()); -// -// @Mock FullServicesContainer mockContainer; -// @Mock ConnectionService mockConnectionService; -// @Mock PluginService mockPluginService; -// @Mock Connection mockConnection; -// @Mock HostSpec mockHostSpec; -// @Mock HostListProviderService mockHostListProviderService; -// @Mock AuroraHostListProvider mockHostListProvider; -// @Mock JdbcCallable mockInitHostProviderFunc; -// @Mock ReaderFailoverHandler mockReaderFailoverHandler; -// @Mock WriterFailoverHandler mockWriterFailoverHandler; -// @Mock ReaderFailoverResult mockReaderResult; -// @Mock WriterFailoverResult mockWriterResult; -// @Mock JdbcCallable mockSqlFunction; -// @Mock private TelemetryFactory mockTelemetryFactory; -// @Mock TelemetryContext mockTelemetryContext; -// @Mock TelemetryCounter mockTelemetryCounter; -// @Mock TelemetryGauge mockTelemetryGauge; -// @Mock TargetDriverDialect mockTargetDriverDialect; -// -// -// private final Properties properties = new Properties(); -// private FailoverConnectionPlugin spyPlugin; -// private AutoCloseable closeable; -// -// @AfterEach -// void cleanUp() throws Exception { -// closeable.close(); -// } -// -// @BeforeEach -// void init() throws SQLException { -// closeable = MockitoAnnotations.openMocks(this); -// -// when(mockContainer.getPluginService()).thenReturn(mockPluginService); -// when(mockPluginService.getHostListProvider()).thenReturn(mockHostListProvider); -// when(mockHostListProvider.getRdsUrlType()).thenReturn(RdsUrlType.RDS_WRITER_CLUSTER); -// when(mockPluginService.getCurrentConnection()).thenReturn(mockConnection); -// when(mockPluginService.getCurrentHostSpec()).thenReturn(mockHostSpec); -// when(mockPluginService.connect(any(HostSpec.class), eq(properties))).thenReturn(mockConnection); -// when(mockPluginService.getTelemetryFactory()).thenReturn(mockTelemetryFactory); -// when(mockPluginService.getHosts()).thenReturn(defaultHosts); -// when(mockPluginService.getAllHosts()).thenReturn(defaultHosts); -// when(mockReaderFailoverHandler.failover(any(), any())).thenReturn(mockReaderResult); -// when(mockWriterFailoverHandler.failover(any())).thenReturn(mockWriterResult); -// when(mockWriterResult.isConnected()).thenReturn(true); -// when(mockWriterResult.getTopology()).thenReturn(defaultHosts); -// when(mockReaderResult.isConnected()).thenReturn(true); -// -// when(mockPluginService.getTelemetryFactory()).thenReturn(mockTelemetryFactory); -// when(mockTelemetryFactory.openTelemetryContext(anyString(), any())).thenReturn(mockTelemetryContext); -// when(mockTelemetryFactory.openTelemetryContext(eq(null), any())).thenReturn(mockTelemetryContext); -// when(mockTelemetryFactory.createCounter(anyString())).thenReturn(mockTelemetryCounter); -// // noinspection unchecked -// when(mockTelemetryFactory.createGauge(anyString(), any(GaugeCallable.class))).thenReturn(mockTelemetryGauge); -// -// when(mockPluginService.getTargetDriverDialect()).thenReturn(mockTargetDriverDialect); -// when(mockTargetDriverDialect.getNetworkBoundMethodNames(any())).thenReturn(new HashSet<>()); -// -// properties.clear(); -// } -// -// @Test -// void test_notifyNodeListChanged_withFailoverDisabled() throws SQLException { -// properties.setProperty(FailoverConnectionPlugin.ENABLE_CLUSTER_AWARE_FAILOVER.name, "false"); -// final Map> changes = new HashMap<>(); -// -// initializePlugin(); -// spyPlugin.notifyNodeListChanged(changes); -// -// verify(mockPluginService, never()).getCurrentHostSpec(); -// verify(mockHostSpec, never()).getAliases(); -// } -// -// @Test -// void test_notifyNodeListChanged_withValidConnectionNotInTopology() throws SQLException { -// final Map> changes = new HashMap<>(); -// changes.put("cluster-host/", EnumSet.of(NodeChangeOptions.NODE_DELETED)); -// changes.put("instance/", EnumSet.of(NodeChangeOptions.NODE_ADDED)); -// -// initializePlugin(); -// spyPlugin.notifyNodeListChanged(changes); -// -// when(mockHostSpec.getUrl()).thenReturn("cluster-url/"); -// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Collections.singletonList("instance"))); -// -// verify(mockPluginService).getCurrentHostSpec(); -// verify(mockHostSpec, never()).getAliases(); -// } -// -// @Test -// void test_updateTopology() throws SQLException { -// initializePlugin(); -// -// // Test updateTopology with failover disabled -// spyPlugin.setRdsUrlType(RdsUrlType.RDS_PROXY); -// spyPlugin.updateTopology(false); -// verify(mockPluginService, never()).forceRefreshHostList(); -// verify(mockPluginService, never()).refreshHostList(); -// -// // Test updateTopology with no connection -// when(mockPluginService.getCurrentHostSpec()).thenReturn(null); -// spyPlugin.updateTopology(false); -// verify(mockPluginService, never()).forceRefreshHostList(); -// verify(mockPluginService, never()).refreshHostList(); -// -// // Test updateTopology with closed connection -// when(mockConnection.isClosed()).thenReturn(true); -// spyPlugin.updateTopology(false); -// verify(mockPluginService, never()).forceRefreshHostList(); -// verify(mockPluginService, never()).refreshHostList(); -// } -// -// @ParameterizedTest -// @ValueSource(booleans = {true, false}) -// void test_updateTopology_withForceUpdate(final boolean forceUpdate) throws SQLException { -// -// when(mockPluginService.getAllHosts()).thenReturn(Collections.singletonList( -// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("host").build())); -// when(mockPluginService.getHosts()).thenReturn(Collections.singletonList( -// new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("host").build())); -// when(mockConnection.isClosed()).thenReturn(false); -// initializePlugin(); -// spyPlugin.setRdsUrlType(RdsUrlType.RDS_INSTANCE); -// -// spyPlugin.updateTopology(forceUpdate); -// if (forceUpdate) { -// verify(mockPluginService, atLeastOnce()).forceRefreshHostList(); -// } else { -// verify(mockPluginService, atLeastOnce()).refreshHostList(); -// } -// } -// -// @Test -// void test_failover_failoverWriter() throws SQLException { -// when(mockPluginService.isInTransaction()).thenReturn(true); -// -// initializePlugin(); -// doThrow(FailoverSuccessSQLException.class).when(spyPlugin).failoverWriter(); -// spyPlugin.failoverMode = FailoverMode.STRICT_WRITER; -// -// assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failover(mockHostSpec)); -// verify(spyPlugin).failoverWriter(); -// } -// -// @Test -// void test_failover_failoverReader() throws SQLException { -// when(mockPluginService.isInTransaction()).thenReturn(false); -// -// initializePlugin(); -// doThrow(FailoverSuccessSQLException.class).when(spyPlugin).failoverReader(eq(mockHostSpec)); -// spyPlugin.failoverMode = FailoverMode.READER_OR_WRITER; -// -// assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failover(mockHostSpec)); -// verify(spyPlugin).failoverReader(eq(mockHostSpec)); -// } -// -// @Test -// void test_failoverReader_withValidFailedHostSpec_successFailover() throws SQLException { -// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); -// when(mockHostSpec.getRawAvailability()).thenReturn(HostAvailability.AVAILABLE); -// when(mockReaderResult.isConnected()).thenReturn(true); -// when(mockReaderResult.getConnection()).thenReturn(mockConnection); -// when(mockReaderResult.getHost()).thenReturn(defaultHosts.get(1)); -// -// initializePlugin(); -// spyPlugin.initHostProvider( -// mockHostListProviderService, -// mockInitHostProviderFunc, -// (connectionService) -> mockReaderFailoverHandler, -// (connectionService) -> mockWriterFailoverHandler); -// -// final FailoverConnectionPlugin spyPlugin = spy(this.spyPlugin); -// doNothing().when(spyPlugin).updateTopology(true); -// -// assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failoverReader(mockHostSpec)); -// -// verify(mockReaderFailoverHandler).failover(eq(defaultHosts), eq(mockHostSpec)); -// verify(mockPluginService).setCurrentConnection(eq(mockConnection), eq(defaultHosts.get(1))); -// } -// -// @Test -// void test_failoverReader_withNoFailedHostSpec_withException() throws SQLException { -// final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") -// .build(); -// final List hosts = Collections.singletonList(hostSpec); -// -// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); -// when(mockHostSpec.getAvailability()).thenReturn(HostAvailability.AVAILABLE); -// when(mockPluginService.getAllHosts()).thenReturn(hosts); -// when(mockPluginService.getHosts()).thenReturn(hosts); -// when(mockReaderResult.getException()).thenReturn(new SQLException()); -// when(mockReaderResult.getHost()).thenReturn(hostSpec); -// -// initializePlugin(); -// spyPlugin.initHostProvider( -// mockHostListProviderService, -// mockInitHostProviderFunc, -// (connectionService) -> mockReaderFailoverHandler, -// (connectionService) -> mockWriterFailoverHandler); -// -// assertThrows(SQLException.class, () -> spyPlugin.failoverReader(null)); -// verify(mockReaderFailoverHandler).failover(eq(hosts), eq(null)); -// } -// -// @Test -// void test_failoverWriter_failedFailover_throwsException() throws SQLException { -// final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") -// .build(); -// final List hosts = Collections.singletonList(hostSpec); -// -// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); -// when(mockPluginService.getAllHosts()).thenReturn(hosts); -// when(mockPluginService.getHosts()).thenReturn(hosts); -// when(mockWriterResult.getException()).thenReturn(new SQLException()); -// -// initializePlugin(); -// spyPlugin.initHostProvider( -// mockHostListProviderService, -// mockInitHostProviderFunc, -// (connectionService) -> mockReaderFailoverHandler, -// (connectionService) -> mockWriterFailoverHandler); -// -// assertThrows(SQLException.class, () -> spyPlugin.failoverWriter()); -// verify(mockWriterFailoverHandler).failover(eq(hosts)); -// } -// -// @Test -// void test_failoverWriter_failedFailover_withNoResult() throws SQLException { -// final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") -// .build(); -// final List hosts = Collections.singletonList(hostSpec); -// -// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); -// when(mockPluginService.getAllHosts()).thenReturn(hosts); -// when(mockPluginService.getHosts()).thenReturn(hosts); -// when(mockWriterResult.isConnected()).thenReturn(false); -// -// initializePlugin(); -// spyPlugin.initHostProvider( -// mockHostListProviderService, -// mockInitHostProviderFunc, -// (connectionService) -> mockReaderFailoverHandler, -// (connectionService) -> mockWriterFailoverHandler); -// -// final SQLException exception = assertThrows(SQLException.class, () -> spyPlugin.failoverWriter()); -// assertEquals(SqlState.CONNECTION_UNABLE_TO_CONNECT.getState(), exception.getSQLState()); -// -// verify(mockWriterFailoverHandler).failover(eq(hosts)); -// verify(mockWriterResult, never()).getNewConnection(); -// verify(mockWriterResult, never()).getTopology(); -// } -// -// @Test -// void test_failoverWriter_successFailover() throws SQLException { -// when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); -// -// initializePlugin(); -// spyPlugin.initHostProvider( -// mockHostListProviderService, -// mockInitHostProviderFunc, -// (connectionService) -> mockReaderFailoverHandler, -// (connectionService) -> mockWriterFailoverHandler); -// -// final SQLException exception = assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failoverWriter()); -// assertEquals(SqlState.COMMUNICATION_LINK_CHANGED.getState(), exception.getSQLState()); -// -// verify(mockWriterFailoverHandler).failover(eq(defaultHosts)); -// } -// -// @Test -// void test_invalidCurrentConnection_withNoConnection() throws SQLException { -// when(mockPluginService.getCurrentConnection()).thenReturn(null); -// initializePlugin(); -// spyPlugin.invalidateCurrentConnection(); -// -// verify(mockPluginService, never()).getCurrentHostSpec(); -// } -// -// @Test -// void test_invalidateCurrentConnection_inTransaction() throws SQLException { -// when(mockPluginService.isInTransaction()).thenReturn(true); -// when(mockHostSpec.getHost()).thenReturn("host"); -// when(mockHostSpec.getPort()).thenReturn(123); -// when(mockHostSpec.getRole()).thenReturn(HostRole.READER); -// -// initializePlugin(); -// spyPlugin.invalidateCurrentConnection(); -// verify(mockConnection).rollback(); -// -// // Assert SQL exceptions thrown during rollback do not get propagated. -// doThrow(new SQLException()).when(mockConnection).rollback(); -// assertDoesNotThrow(() -> spyPlugin.invalidateCurrentConnection()); -// } -// -// @Test -// void test_invalidateCurrentConnection_notInTransaction() throws SQLException { -// when(mockPluginService.isInTransaction()).thenReturn(false); -// when(mockHostSpec.getHost()).thenReturn("host"); -// when(mockHostSpec.getPort()).thenReturn(123); -// when(mockHostSpec.getRole()).thenReturn(HostRole.READER); -// -// initializePlugin(); -// spyPlugin.invalidateCurrentConnection(); -// -// verify(mockPluginService).isInTransaction(); -// } -// -// @Test -// void test_invalidateCurrentConnection_withOpenConnection() throws SQLException { -// when(mockPluginService.isInTransaction()).thenReturn(false); -// when(mockConnection.isClosed()).thenReturn(false); -// when(mockHostSpec.getHost()).thenReturn("host"); -// when(mockHostSpec.getPort()).thenReturn(123); -// when(mockHostSpec.getRole()).thenReturn(HostRole.READER); -// -// initializePlugin(); -// spyPlugin.invalidateCurrentConnection(); -// -// doThrow(new SQLException()).when(mockConnection).close(); -// assertDoesNotThrow(() -> spyPlugin.invalidateCurrentConnection()); -// -// verify(mockConnection, times(2)).isClosed(); -// verify(mockConnection, times(2)).close(); -// } -// -// @Test -// void test_execute_withFailoverDisabled() throws SQLException { -// properties.setProperty(FailoverConnectionPlugin.ENABLE_CLUSTER_AWARE_FAILOVER.name, "false"); -// initializePlugin(); -// -// spyPlugin.execute( -// ResultSet.class, -// SQLException.class, -// MONITOR_METHOD_INVOKE_ON, -// MONITOR_METHOD_NAME, -// mockSqlFunction, -// EMPTY_ARGS); -// -// verify(mockSqlFunction).call(); -// verify(mockHostListProvider, never()).getRdsUrlType(); -// } -// -// @Test -// void test_execute_withDirectExecute() throws SQLException { -// initializePlugin(); -// spyPlugin.execute( -// ResultSet.class, -// SQLException.class, -// MONITOR_METHOD_INVOKE_ON, -// "close", -// mockSqlFunction, -// EMPTY_ARGS); -// verify(mockSqlFunction).call(); -// verify(mockHostListProvider, never()).getRdsUrlType(); -// } -// -// private void initializePlugin() throws SQLException { -// spyPlugin = spy(new FailoverConnectionPlugin(mockContainer, properties)); -// spyPlugin.setWriterFailoverHandler(mockWriterFailoverHandler); -// spyPlugin.setReaderFailoverHandler(mockReaderFailoverHandler); -// doReturn(mockConnectionService).when(spyPlugin).getConnectionService(); -// } -// } +/* + * 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.failover; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.Arrays; +import java.util.Collections; +import java.util.EnumSet; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Properties; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import software.amazon.jdbc.HostListProviderService; +import software.amazon.jdbc.HostRole; +import software.amazon.jdbc.HostSpec; +import software.amazon.jdbc.HostSpecBuilder; +import software.amazon.jdbc.JdbcCallable; +import software.amazon.jdbc.NodeChangeOptions; +import software.amazon.jdbc.PluginService; +import software.amazon.jdbc.hostavailability.HostAvailability; +import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; +import software.amazon.jdbc.hostlistprovider.AuroraHostListProvider; +import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; +import software.amazon.jdbc.util.FullServicesContainer; +import software.amazon.jdbc.util.RdsUrlType; +import software.amazon.jdbc.util.SqlState; +import software.amazon.jdbc.util.connection.ConnectionService; +import software.amazon.jdbc.util.telemetry.GaugeCallable; +import software.amazon.jdbc.util.telemetry.TelemetryContext; +import software.amazon.jdbc.util.telemetry.TelemetryCounter; +import software.amazon.jdbc.util.telemetry.TelemetryFactory; +import software.amazon.jdbc.util.telemetry.TelemetryGauge; + +class FailoverConnectionPluginTest { + + private static final Class MONITOR_METHOD_INVOKE_ON = Connection.class; + private static final String MONITOR_METHOD_NAME = "Connection.executeQuery"; + private static final Object[] EMPTY_ARGS = {}; + private final List defaultHosts = Arrays.asList( + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("writer").port(1234).role(HostRole.WRITER).build(), + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()) + .host("reader1").port(1234).role(HostRole.READER).build()); + + @Mock FullServicesContainer mockContainer; + @Mock ConnectionService mockConnectionService; + @Mock PluginService mockPluginService; + @Mock Connection mockConnection; + @Mock HostSpec mockHostSpec; + @Mock HostListProviderService mockHostListProviderService; + @Mock AuroraHostListProvider mockHostListProvider; + @Mock JdbcCallable mockInitHostProviderFunc; + @Mock ReaderFailoverHandler mockReaderFailoverHandler; + @Mock WriterFailoverHandler mockWriterFailoverHandler; + @Mock ReaderFailoverResult mockReaderResult; + @Mock WriterFailoverResult mockWriterResult; + @Mock JdbcCallable mockSqlFunction; + @Mock private TelemetryFactory mockTelemetryFactory; + @Mock TelemetryContext mockTelemetryContext; + @Mock TelemetryCounter mockTelemetryCounter; + @Mock TelemetryGauge mockTelemetryGauge; + @Mock TargetDriverDialect mockTargetDriverDialect; + + + private final Properties properties = new Properties(); + private FailoverConnectionPlugin spyPlugin; + private AutoCloseable closeable; + + @AfterEach + void cleanUp() throws Exception { + closeable.close(); + } + + @BeforeEach + void init() throws SQLException { + closeable = MockitoAnnotations.openMocks(this); + + when(mockContainer.getPluginService()).thenReturn(mockPluginService); + when(mockPluginService.getHostListProvider()).thenReturn(mockHostListProvider); + when(mockHostListProvider.getRdsUrlType()).thenReturn(RdsUrlType.RDS_WRITER_CLUSTER); + when(mockPluginService.getCurrentConnection()).thenReturn(mockConnection); + when(mockPluginService.getCurrentHostSpec()).thenReturn(mockHostSpec); + when(mockPluginService.connect(any(HostSpec.class), eq(properties))).thenReturn(mockConnection); + when(mockPluginService.getTelemetryFactory()).thenReturn(mockTelemetryFactory); + when(mockPluginService.getHosts()).thenReturn(defaultHosts); + when(mockPluginService.getAllHosts()).thenReturn(defaultHosts); + when(mockReaderFailoverHandler.failover(any(), any())).thenReturn(mockReaderResult); + when(mockWriterFailoverHandler.failover(any())).thenReturn(mockWriterResult); + when(mockWriterResult.isConnected()).thenReturn(true); + when(mockWriterResult.getTopology()).thenReturn(defaultHosts); + when(mockReaderResult.isConnected()).thenReturn(true); + + when(mockPluginService.getTelemetryFactory()).thenReturn(mockTelemetryFactory); + when(mockTelemetryFactory.openTelemetryContext(anyString(), any())).thenReturn(mockTelemetryContext); + when(mockTelemetryFactory.openTelemetryContext(eq(null), any())).thenReturn(mockTelemetryContext); + when(mockTelemetryFactory.createCounter(anyString())).thenReturn(mockTelemetryCounter); + // noinspection unchecked + when(mockTelemetryFactory.createGauge(anyString(), any(GaugeCallable.class))).thenReturn(mockTelemetryGauge); + + when(mockPluginService.getTargetDriverDialect()).thenReturn(mockTargetDriverDialect); + when(mockTargetDriverDialect.getNetworkBoundMethodNames(any())).thenReturn(new HashSet<>()); + + properties.clear(); + } + + @Test + void test_notifyNodeListChanged_withFailoverDisabled() throws SQLException { + properties.setProperty(FailoverConnectionPlugin.ENABLE_CLUSTER_AWARE_FAILOVER.name, "false"); + final Map> changes = new HashMap<>(); + + initializePlugin(); + spyPlugin.notifyNodeListChanged(changes); + + verify(mockPluginService, never()).getCurrentHostSpec(); + verify(mockHostSpec, never()).getAliases(); + } + + @Test + void test_notifyNodeListChanged_withValidConnectionNotInTopology() throws SQLException { + final Map> changes = new HashMap<>(); + changes.put("cluster-host/", EnumSet.of(NodeChangeOptions.NODE_DELETED)); + changes.put("instance/", EnumSet.of(NodeChangeOptions.NODE_ADDED)); + + initializePlugin(); + spyPlugin.notifyNodeListChanged(changes); + + when(mockHostSpec.getUrl()).thenReturn("cluster-url/"); + when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Collections.singletonList("instance"))); + + verify(mockPluginService).getCurrentHostSpec(); + verify(mockHostSpec, never()).getAliases(); + } + + @Test + void test_updateTopology() throws SQLException { + initializePlugin(); + + // Test updateTopology with failover disabled + spyPlugin.setRdsUrlType(RdsUrlType.RDS_PROXY); + spyPlugin.updateTopology(false); + verify(mockPluginService, never()).forceRefreshHostList(); + verify(mockPluginService, never()).refreshHostList(); + + // Test updateTopology with no connection + when(mockPluginService.getCurrentHostSpec()).thenReturn(null); + spyPlugin.updateTopology(false); + verify(mockPluginService, never()).forceRefreshHostList(); + verify(mockPluginService, never()).refreshHostList(); + + // Test updateTopology with closed connection + when(mockConnection.isClosed()).thenReturn(true); + spyPlugin.updateTopology(false); + verify(mockPluginService, never()).forceRefreshHostList(); + verify(mockPluginService, never()).refreshHostList(); + } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void test_updateTopology_withForceUpdate(final boolean forceUpdate) throws SQLException { + + when(mockPluginService.getAllHosts()).thenReturn(Collections.singletonList( + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("host").build())); + when(mockPluginService.getHosts()).thenReturn(Collections.singletonList( + new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("host").build())); + when(mockConnection.isClosed()).thenReturn(false); + initializePlugin(); + spyPlugin.setRdsUrlType(RdsUrlType.RDS_INSTANCE); + + spyPlugin.updateTopology(forceUpdate); + if (forceUpdate) { + verify(mockPluginService, atLeastOnce()).forceRefreshHostList(); + } else { + verify(mockPluginService, atLeastOnce()).refreshHostList(); + } + } + + @Test + void test_failover_failoverWriter() throws SQLException { + when(mockPluginService.isInTransaction()).thenReturn(true); + + initializePlugin(); + doThrow(FailoverSuccessSQLException.class).when(spyPlugin).failoverWriter(); + spyPlugin.failoverMode = FailoverMode.STRICT_WRITER; + + assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failover(mockHostSpec)); + verify(spyPlugin).failoverWriter(); + } + + @Test + void test_failover_failoverReader() throws SQLException { + when(mockPluginService.isInTransaction()).thenReturn(false); + + initializePlugin(); + doThrow(FailoverSuccessSQLException.class).when(spyPlugin).failoverReader(eq(mockHostSpec)); + spyPlugin.failoverMode = FailoverMode.READER_OR_WRITER; + + assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failover(mockHostSpec)); + verify(spyPlugin).failoverReader(eq(mockHostSpec)); + } + + @Test + void test_failoverReader_withValidFailedHostSpec_successFailover() throws SQLException { + when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); + when(mockHostSpec.getRawAvailability()).thenReturn(HostAvailability.AVAILABLE); + when(mockReaderResult.isConnected()).thenReturn(true); + when(mockReaderResult.getConnection()).thenReturn(mockConnection); + when(mockReaderResult.getHost()).thenReturn(defaultHosts.get(1)); + + initializePlugin(); + spyPlugin.initHostProvider( + mockHostListProviderService, + mockInitHostProviderFunc, + () -> mockReaderFailoverHandler, + () -> mockWriterFailoverHandler); + + final FailoverConnectionPlugin spyPlugin = spy(this.spyPlugin); + doNothing().when(spyPlugin).updateTopology(true); + + assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failoverReader(mockHostSpec)); + + verify(mockReaderFailoverHandler).failover(eq(defaultHosts), eq(mockHostSpec)); + verify(mockPluginService).setCurrentConnection(eq(mockConnection), eq(defaultHosts.get(1))); + } + + @Test + void test_failoverReader_withNoFailedHostSpec_withException() throws SQLException { + final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") + .build(); + final List hosts = Collections.singletonList(hostSpec); + + when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); + when(mockHostSpec.getAvailability()).thenReturn(HostAvailability.AVAILABLE); + when(mockPluginService.getAllHosts()).thenReturn(hosts); + when(mockPluginService.getHosts()).thenReturn(hosts); + when(mockReaderResult.getException()).thenReturn(new SQLException()); + when(mockReaderResult.getHost()).thenReturn(hostSpec); + + initializePlugin(); + spyPlugin.initHostProvider( + mockHostListProviderService, + mockInitHostProviderFunc, + () -> mockReaderFailoverHandler, + () -> mockWriterFailoverHandler); + + assertThrows(SQLException.class, () -> spyPlugin.failoverReader(null)); + verify(mockReaderFailoverHandler).failover(eq(hosts), eq(null)); + } + + @Test + void test_failoverWriter_failedFailover_throwsException() throws SQLException { + final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") + .build(); + final List hosts = Collections.singletonList(hostSpec); + + when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); + when(mockPluginService.getAllHosts()).thenReturn(hosts); + when(mockPluginService.getHosts()).thenReturn(hosts); + when(mockWriterResult.getException()).thenReturn(new SQLException()); + + initializePlugin(); + spyPlugin.initHostProvider( + mockHostListProviderService, + mockInitHostProviderFunc, + () -> mockReaderFailoverHandler, + () -> mockWriterFailoverHandler); + + assertThrows(SQLException.class, () -> spyPlugin.failoverWriter()); + verify(mockWriterFailoverHandler).failover(eq(hosts)); + } + + @Test + void test_failoverWriter_failedFailover_withNoResult() throws SQLException { + final HostSpec hostSpec = new HostSpecBuilder(new SimpleHostAvailabilityStrategy()).host("hostA") + .build(); + final List hosts = Collections.singletonList(hostSpec); + + when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); + when(mockPluginService.getAllHosts()).thenReturn(hosts); + when(mockPluginService.getHosts()).thenReturn(hosts); + when(mockWriterResult.isConnected()).thenReturn(false); + + initializePlugin(); + spyPlugin.initHostProvider( + mockHostListProviderService, + mockInitHostProviderFunc, + () -> mockReaderFailoverHandler, + () -> mockWriterFailoverHandler); + + final SQLException exception = assertThrows(SQLException.class, () -> spyPlugin.failoverWriter()); + assertEquals(SqlState.CONNECTION_UNABLE_TO_CONNECT.getState(), exception.getSQLState()); + + verify(mockWriterFailoverHandler).failover(eq(hosts)); + verify(mockWriterResult, never()).getNewConnection(); + verify(mockWriterResult, never()).getTopology(); + } + + @Test + void test_failoverWriter_successFailover() throws SQLException { + when(mockHostSpec.getAliases()).thenReturn(new HashSet<>(Arrays.asList("alias1", "alias2"))); + + initializePlugin(); + spyPlugin.initHostProvider( + mockHostListProviderService, + mockInitHostProviderFunc, + () -> mockReaderFailoverHandler, + () -> mockWriterFailoverHandler); + + final SQLException exception = assertThrows(FailoverSuccessSQLException.class, () -> spyPlugin.failoverWriter()); + assertEquals(SqlState.COMMUNICATION_LINK_CHANGED.getState(), exception.getSQLState()); + + verify(mockWriterFailoverHandler).failover(eq(defaultHosts)); + } + + @Test + void test_invalidCurrentConnection_withNoConnection() throws SQLException { + when(mockPluginService.getCurrentConnection()).thenReturn(null); + initializePlugin(); + spyPlugin.invalidateCurrentConnection(); + + verify(mockPluginService, never()).getCurrentHostSpec(); + } + + @Test + void test_invalidateCurrentConnection_inTransaction() throws SQLException { + when(mockPluginService.isInTransaction()).thenReturn(true); + when(mockHostSpec.getHost()).thenReturn("host"); + when(mockHostSpec.getPort()).thenReturn(123); + when(mockHostSpec.getRole()).thenReturn(HostRole.READER); + + initializePlugin(); + spyPlugin.invalidateCurrentConnection(); + verify(mockConnection).rollback(); + + // Assert SQL exceptions thrown during rollback do not get propagated. + doThrow(new SQLException()).when(mockConnection).rollback(); + assertDoesNotThrow(() -> spyPlugin.invalidateCurrentConnection()); + } + + @Test + void test_invalidateCurrentConnection_notInTransaction() throws SQLException { + when(mockPluginService.isInTransaction()).thenReturn(false); + when(mockHostSpec.getHost()).thenReturn("host"); + when(mockHostSpec.getPort()).thenReturn(123); + when(mockHostSpec.getRole()).thenReturn(HostRole.READER); + + initializePlugin(); + spyPlugin.invalidateCurrentConnection(); + + verify(mockPluginService).isInTransaction(); + } + + @Test + void test_invalidateCurrentConnection_withOpenConnection() throws SQLException { + when(mockPluginService.isInTransaction()).thenReturn(false); + when(mockConnection.isClosed()).thenReturn(false); + when(mockHostSpec.getHost()).thenReturn("host"); + when(mockHostSpec.getPort()).thenReturn(123); + when(mockHostSpec.getRole()).thenReturn(HostRole.READER); + + initializePlugin(); + spyPlugin.invalidateCurrentConnection(); + + doThrow(new SQLException()).when(mockConnection).close(); + assertDoesNotThrow(() -> spyPlugin.invalidateCurrentConnection()); + + verify(mockConnection, times(2)).isClosed(); + verify(mockConnection, times(2)).close(); + } + + @Test + void test_execute_withFailoverDisabled() throws SQLException { + properties.setProperty(FailoverConnectionPlugin.ENABLE_CLUSTER_AWARE_FAILOVER.name, "false"); + initializePlugin(); + + spyPlugin.execute( + ResultSet.class, + SQLException.class, + MONITOR_METHOD_INVOKE_ON, + MONITOR_METHOD_NAME, + mockSqlFunction, + EMPTY_ARGS); + + verify(mockSqlFunction).call(); + verify(mockHostListProvider, never()).getRdsUrlType(); + } + + @Test + void test_execute_withDirectExecute() throws SQLException { + initializePlugin(); + spyPlugin.execute( + ResultSet.class, + SQLException.class, + MONITOR_METHOD_INVOKE_ON, + "close", + mockSqlFunction, + EMPTY_ARGS); + verify(mockSqlFunction).call(); + verify(mockHostListProvider, never()).getRdsUrlType(); + } + + private void initializePlugin() throws SQLException { + spyPlugin = spy(new FailoverConnectionPlugin(mockContainer, properties)); + spyPlugin.setWriterFailoverHandler(mockWriterFailoverHandler); + spyPlugin.setReaderFailoverHandler(mockReaderFailoverHandler); + } +} diff --git a/wrapper/src/test/java/software/amazon/jdbc/util/monitoring/MonitorServiceImplTest.java b/wrapper/src/test/java/software/amazon/jdbc/util/monitoring/MonitorServiceImplTest.java index 9251b27e5..02bd5060e 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/util/monitoring/MonitorServiceImplTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/util/monitoring/MonitorServiceImplTest.java @@ -1,314 +1,315 @@ -// /* -// * 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.util.monitoring; -// -// import static org.junit.jupiter.api.Assertions.assertEquals; -// import static org.junit.jupiter.api.Assertions.assertNotEquals; -// import static org.junit.jupiter.api.Assertions.assertNotNull; -// import static org.junit.jupiter.api.Assertions.assertNull; -// import static org.junit.jupiter.api.Assertions.assertThrows; -// import static org.mockito.ArgumentMatchers.any; -// import static org.mockito.ArgumentMatchers.anyInt; -// import static org.mockito.Mockito.doNothing; -// import static org.mockito.Mockito.doReturn; -// import static org.mockito.Mockito.spy; -// -// import java.sql.SQLException; -// import java.util.Collections; -// import java.util.HashSet; -// import java.util.Properties; -// import java.util.concurrent.TimeUnit; -// import org.junit.jupiter.api.AfterEach; -// import org.junit.jupiter.api.Assertions; -// import org.junit.jupiter.api.BeforeEach; -// import org.junit.jupiter.api.Test; -// import org.mockito.Mock; -// import org.mockito.MockitoAnnotations; -// import software.amazon.jdbc.ConnectionProvider; -// import software.amazon.jdbc.dialect.Dialect; -// import software.amazon.jdbc.plugin.customendpoint.CustomEndpointMonitorImpl; -// import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; -// import software.amazon.jdbc.util.connection.ConnectionService; -// import software.amazon.jdbc.util.events.EventPublisher; -// import software.amazon.jdbc.util.storage.StorageService; -// import software.amazon.jdbc.util.telemetry.TelemetryFactory; -// -// class MonitorServiceImplTest { -// @Mock StorageService mockStorageService; -// @Mock ConnectionService mockConnectionService; -// @Mock ConnectionProvider mockConnectionProvider; -// @Mock TelemetryFactory mockTelemetryFactory; -// @Mock TargetDriverDialect mockTargetDriverDialect; -// @Mock Dialect mockDbDialect; -// @Mock EventPublisher mockPublisher; -// MonitorServiceImpl spyMonitorService; -// private AutoCloseable closeable; -// -// @BeforeEach -// void setUp() { -// closeable = MockitoAnnotations.openMocks(this); -// spyMonitorService = spy(new MonitorServiceImpl(mockPublisher)); -// doNothing().when(spyMonitorService).initCleanupThread(anyInt()); -// -// try { -// doReturn(mockConnectionService).when(spyMonitorService) -// .getConnectionService(any(), any(), any(), any(), any(), any(), any(), any()); -// } catch (SQLException e) { -// Assertions.fail( -// "Encountered exception while stubbing MonitorServiceImpl#getConnectionService: " + e.getMessage()); -// } -// } -// -// @AfterEach -// void tearDown() throws Exception { -// closeable.close(); -// spyMonitorService.releaseResources(); -// } -// -// @Test -// public void testMonitorError_monitorReCreated() throws SQLException, InterruptedException { -// spyMonitorService.registerMonitorTypeIfAbsent( -// NoOpMonitor.class, -// TimeUnit.MINUTES.toNanos(1), -// TimeUnit.MINUTES.toNanos(1), -// new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), -// null -// ); -// String key = "testMonitor"; -// NoOpMonitor monitor = spyMonitorService.runIfAbsent( -// NoOpMonitor.class, -// key, -// mockStorageService, -// mockTelemetryFactory, -// mockConnectionProvider, -// "jdbc:postgresql://somehost/somedb", -// "someProtocol", -// mockTargetDriverDialect, -// mockDbDialect, -// new Properties(), -// (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) -// ); -// -// Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key); -// assertNotNull(storedMonitor); -// assertEquals(monitor, storedMonitor); -// // need to wait to give time for the monitor executor to start the monitor thread. -// TimeUnit.MILLISECONDS.sleep(250); -// assertEquals(MonitorState.RUNNING, monitor.getState()); -// -// monitor.state.set(MonitorState.ERROR); -// spyMonitorService.checkMonitors(); -// -// assertEquals(MonitorState.STOPPED, monitor.getState()); -// -// Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key); -// assertNotNull(newMonitor); -// assertNotEquals(monitor, newMonitor); -// // need to wait to give time for the monitor executor to start the monitor thread. -// TimeUnit.MILLISECONDS.sleep(250); -// assertEquals(MonitorState.RUNNING, newMonitor.getState()); -// } -// -// @Test -// public void testMonitorStuck_monitorReCreated() throws SQLException, InterruptedException { -// spyMonitorService.registerMonitorTypeIfAbsent( -// NoOpMonitor.class, -// TimeUnit.MINUTES.toNanos(1), -// 1, // heartbeat times out immediately -// new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), -// null -// ); -// String key = "testMonitor"; -// NoOpMonitor monitor = spyMonitorService.runIfAbsent( -// NoOpMonitor.class, -// key, -// mockStorageService, -// mockTelemetryFactory, -// mockConnectionProvider, -// "jdbc:postgresql://somehost/somedb", -// "someProtocol", -// mockTargetDriverDialect, -// mockDbDialect, -// new Properties(), -// (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) -// ); -// -// Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key); -// assertNotNull(storedMonitor); -// assertEquals(monitor, storedMonitor); -// // need to wait to give time for the monitor executor to start the monitor thread. -// TimeUnit.MILLISECONDS.sleep(250); -// assertEquals(MonitorState.RUNNING, monitor.getState()); -// -// // checkMonitors() should detect the heartbeat/inactivity timeout, stop the monitor, and re-create a new one. -// spyMonitorService.checkMonitors(); -// -// assertEquals(MonitorState.STOPPED, monitor.getState()); -// -// Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key); -// assertNotNull(newMonitor); -// assertNotEquals(monitor, newMonitor); -// // need to wait to give time for the monitor executor to start the monitor thread. -// TimeUnit.MILLISECONDS.sleep(250); -// assertEquals(MonitorState.RUNNING, newMonitor.getState()); -// } -// -// @Test -// public void testMonitorExpired() throws SQLException, InterruptedException { -// spyMonitorService.registerMonitorTypeIfAbsent( -// NoOpMonitor.class, -// TimeUnit.MILLISECONDS.toNanos(200), // monitor expires after 200ms -// TimeUnit.MINUTES.toNanos(1), -// // even though we pass a re-create policy, we should not re-create it if the monitor is expired since this -// // indicates it is not being used. -// new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), -// null -// ); -// String key = "testMonitor"; -// NoOpMonitor monitor = spyMonitorService.runIfAbsent( -// NoOpMonitor.class, -// key, -// mockStorageService, -// mockTelemetryFactory, -// mockConnectionProvider, -// "jdbc:postgresql://somehost/somedb", -// "someProtocol", -// mockTargetDriverDialect, -// mockDbDialect, -// new Properties(), -// (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) -// ); -// -// Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key); -// assertNotNull(storedMonitor); -// assertEquals(monitor, storedMonitor); -// // need to wait to give time for the monitor executor to start the monitor thread. -// TimeUnit.MILLISECONDS.sleep(250); -// assertEquals(MonitorState.RUNNING, monitor.getState()); -// -// // checkMonitors() should detect the expiration timeout and stop/remove the monitor. -// spyMonitorService.checkMonitors(); -// -// assertEquals(MonitorState.STOPPED, monitor.getState()); -// -// Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key); -// // monitor should have been removed when checkMonitors() was called. -// assertNull(newMonitor); -// } -// -// @Test -// public void testMonitorMismatch() { -// assertThrows(IllegalStateException.class, () -> spyMonitorService.runIfAbsent( -// CustomEndpointMonitorImpl.class, -// "testMonitor", -// mockStorageService, -// mockTelemetryFactory, -// mockConnectionProvider, -// "jdbc:postgresql://somehost/somedb", -// "someProtocol", -// mockTargetDriverDialect, -// mockDbDialect, -// new Properties(), -// // indicated monitor class is CustomEndpointMonitorImpl, but actual monitor is NoOpMonitor. The monitor -// // service should detect this and throw an exception. -// (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) -// )); -// } -// -// @Test -// public void testRemove() throws SQLException, InterruptedException { -// spyMonitorService.registerMonitorTypeIfAbsent( -// NoOpMonitor.class, -// TimeUnit.MINUTES.toNanos(1), -// TimeUnit.MINUTES.toNanos(1), -// // even though we pass a re-create policy, we should not re-create it if the monitor is expired since this -// // indicates it is not being used. -// new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), -// null -// ); -// -// String key = "testMonitor"; -// NoOpMonitor monitor = spyMonitorService.runIfAbsent( -// NoOpMonitor.class, -// key, -// mockStorageService, -// mockTelemetryFactory, -// mockConnectionProvider, -// "jdbc:postgresql://somehost/somedb", -// "someProtocol", -// mockTargetDriverDialect, -// mockDbDialect, -// new Properties(), -// (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) -// ); -// assertNotNull(monitor); -// -// // need to wait to give time for the monitor executor to start the monitor thread. -// TimeUnit.MILLISECONDS.sleep(250); -// Monitor removedMonitor = spyMonitorService.remove(NoOpMonitor.class, key); -// assertEquals(monitor, removedMonitor); -// assertEquals(MonitorState.RUNNING, monitor.getState()); -// } -// -// @Test -// public void testStopAndRemove() throws SQLException, InterruptedException { -// spyMonitorService.registerMonitorTypeIfAbsent( -// NoOpMonitor.class, -// TimeUnit.MINUTES.toNanos(1), -// TimeUnit.MINUTES.toNanos(1), -// // even though we pass a re-create policy, we should not re-create it if the monitor is expired since this -// // indicates it is not being used. -// new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), -// null -// ); -// -// String key = "testMonitor"; -// NoOpMonitor monitor = spyMonitorService.runIfAbsent( -// NoOpMonitor.class, -// key, -// mockStorageService, -// mockTelemetryFactory, -// mockConnectionProvider, -// "jdbc:postgresql://somehost/somedb", -// "someProtocol", -// mockTargetDriverDialect, -// mockDbDialect, -// new Properties(), -// (connectionService, pluginService) -> new NoOpMonitor(spyMonitorService, 30) -// ); -// assertNotNull(monitor); -// -// // need to wait to give time for the monitor executor to start the monitor thread. -// TimeUnit.MILLISECONDS.sleep(250); -// spyMonitorService.stopAndRemove(NoOpMonitor.class, key); -// assertNull(spyMonitorService.get(NoOpMonitor.class, key)); -// assertEquals(MonitorState.STOPPED, monitor.getState()); -// } -// -// static class NoOpMonitor extends AbstractMonitor { -// protected NoOpMonitor( -// MonitorService monitorService, -// long terminationTimeoutSec) { -// super(terminationTimeoutSec); -// } -// -// @Override -// public void monitor() { -// // do nothing. -// } -// } -// } +/* + * 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.util.monitoring; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.spy; + +import java.sql.SQLException; +import java.util.Collections; +import java.util.HashSet; +import java.util.Properties; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import software.amazon.jdbc.ConnectionProvider; +import software.amazon.jdbc.dialect.Dialect; +import software.amazon.jdbc.plugin.customendpoint.CustomEndpointMonitorImpl; +import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; +import software.amazon.jdbc.util.FullServicesContainer; +import software.amazon.jdbc.util.events.EventPublisher; +import software.amazon.jdbc.util.storage.StorageService; +import software.amazon.jdbc.util.telemetry.TelemetryFactory; + +class MonitorServiceImplTest { + @Mock FullServicesContainer mockServicesContainer; + @Mock StorageService mockStorageService; + @Mock ConnectionProvider mockConnectionProvider; + @Mock TelemetryFactory mockTelemetryFactory; + @Mock TargetDriverDialect mockTargetDriverDialect; + @Mock Dialect mockDbDialect; + @Mock EventPublisher mockPublisher; + String URL = "jdbc:postgresql://somehost/somedb"; + String PROTOCOL = "someProtocol"; + Properties props = new Properties(); + MonitorServiceImpl spyMonitorService; + private AutoCloseable closeable; + + @BeforeEach + void setUp() throws SQLException { + closeable = MockitoAnnotations.openMocks(this); + spyMonitorService = spy(new MonitorServiceImpl(mockPublisher)); + doNothing().when(spyMonitorService).initCleanupThread(anyInt()); + doReturn(mockServicesContainer).when(spyMonitorService).getNewServicesContainer( + eq(mockStorageService), + eq(mockConnectionProvider), + eq(mockTelemetryFactory), + eq(URL), + eq(PROTOCOL), + eq(mockTargetDriverDialect), + eq(mockDbDialect), + eq(props)); + } + + @AfterEach + void tearDown() throws Exception { + closeable.close(); + spyMonitorService.releaseResources(); + } + + @Test + public void testMonitorError_monitorReCreated() throws SQLException, InterruptedException { + spyMonitorService.registerMonitorTypeIfAbsent( + NoOpMonitor.class, + TimeUnit.MINUTES.toNanos(1), + TimeUnit.MINUTES.toNanos(1), + new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), + null + ); + String key = "testMonitor"; + NoOpMonitor monitor = spyMonitorService.runIfAbsent( + NoOpMonitor.class, + key, + mockStorageService, + mockTelemetryFactory, + mockConnectionProvider, + URL, + PROTOCOL, + mockTargetDriverDialect, + mockDbDialect, + props, + (mockServicesContainer) -> new NoOpMonitor(30) + ); + + Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key); + assertNotNull(storedMonitor); + assertEquals(monitor, storedMonitor); + // need to wait to give time for the monitor executor to start the monitor thread. + TimeUnit.MILLISECONDS.sleep(250); + assertEquals(MonitorState.RUNNING, monitor.getState()); + + monitor.state.set(MonitorState.ERROR); + spyMonitorService.checkMonitors(); + + assertEquals(MonitorState.STOPPED, monitor.getState()); + + Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key); + assertNotNull(newMonitor); + assertNotEquals(monitor, newMonitor); + // need to wait to give time for the monitor executor to start the monitor thread. + TimeUnit.MILLISECONDS.sleep(250); + assertEquals(MonitorState.RUNNING, newMonitor.getState()); + } + + @Test + public void testMonitorStuck_monitorReCreated() throws SQLException, InterruptedException { + spyMonitorService.registerMonitorTypeIfAbsent( + NoOpMonitor.class, + TimeUnit.MINUTES.toNanos(1), + 1, // heartbeat times out immediately + new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), + null + ); + String key = "testMonitor"; + NoOpMonitor monitor = spyMonitorService.runIfAbsent( + NoOpMonitor.class, + key, + mockStorageService, + mockTelemetryFactory, + mockConnectionProvider, + URL, + PROTOCOL, + mockTargetDriverDialect, + mockDbDialect, + props, + (mockServicesContainer) -> new NoOpMonitor(30) + ); + + Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key); + assertNotNull(storedMonitor); + assertEquals(monitor, storedMonitor); + // need to wait to give time for the monitor executor to start the monitor thread. + TimeUnit.MILLISECONDS.sleep(250); + assertEquals(MonitorState.RUNNING, monitor.getState()); + + // checkMonitors() should detect the heartbeat/inactivity timeout, stop the monitor, and re-create a new one. + spyMonitorService.checkMonitors(); + + assertEquals(MonitorState.STOPPED, monitor.getState()); + + Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key); + assertNotNull(newMonitor); + assertNotEquals(monitor, newMonitor); + // need to wait to give time for the monitor executor to start the monitor thread. + TimeUnit.MILLISECONDS.sleep(250); + assertEquals(MonitorState.RUNNING, newMonitor.getState()); + } + + @Test + public void testMonitorExpired() throws SQLException, InterruptedException { + spyMonitorService.registerMonitorTypeIfAbsent( + NoOpMonitor.class, + TimeUnit.MILLISECONDS.toNanos(200), // monitor expires after 200ms + TimeUnit.MINUTES.toNanos(1), + // even though we pass a re-create policy, we should not re-create it if the monitor is expired since this + // indicates it is not being used. + new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), + null + ); + String key = "testMonitor"; + NoOpMonitor monitor = spyMonitorService.runIfAbsent( + NoOpMonitor.class, + key, + mockStorageService, + mockTelemetryFactory, + mockConnectionProvider, + URL, + PROTOCOL, + mockTargetDriverDialect, + mockDbDialect, + props, + (mockServicesContainer) -> new NoOpMonitor(30) + ); + + Monitor storedMonitor = spyMonitorService.get(NoOpMonitor.class, key); + assertNotNull(storedMonitor); + assertEquals(monitor, storedMonitor); + // need to wait to give time for the monitor executor to start the monitor thread. + TimeUnit.MILLISECONDS.sleep(250); + assertEquals(MonitorState.RUNNING, monitor.getState()); + + // checkMonitors() should detect the expiration timeout and stop/remove the monitor. + spyMonitorService.checkMonitors(); + + assertEquals(MonitorState.STOPPED, monitor.getState()); + + Monitor newMonitor = spyMonitorService.get(NoOpMonitor.class, key); + // monitor should have been removed when checkMonitors() was called. + assertNull(newMonitor); + } + + @Test + public void testMonitorMismatch() { + assertThrows(IllegalStateException.class, () -> spyMonitorService.runIfAbsent( + CustomEndpointMonitorImpl.class, + "testMonitor", + mockStorageService, + mockTelemetryFactory, + mockConnectionProvider, + URL, + PROTOCOL, + mockTargetDriverDialect, + mockDbDialect, + props, + // indicated monitor class is CustomEndpointMonitorImpl, but actual monitor is NoOpMonitor. The monitor + // service should detect this and throw an exception. + (mockServicesContainer) -> new NoOpMonitor(30) + )); + } + + @Test + public void testRemove() throws SQLException, InterruptedException { + spyMonitorService.registerMonitorTypeIfAbsent( + NoOpMonitor.class, + TimeUnit.MINUTES.toNanos(1), + TimeUnit.MINUTES.toNanos(1), + // even though we pass a re-create policy, we should not re-create it if the monitor is expired since this + // indicates it is not being used. + new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), + null + ); + + String key = "testMonitor"; + NoOpMonitor monitor = spyMonitorService.runIfAbsent( + NoOpMonitor.class, + key, + mockStorageService, + mockTelemetryFactory, + mockConnectionProvider, + URL, + PROTOCOL, + mockTargetDriverDialect, + mockDbDialect, + props, + (mockServicesContainer) -> new NoOpMonitor(30) + ); + assertNotNull(monitor); + + // need to wait to give time for the monitor executor to start the monitor thread. + TimeUnit.MILLISECONDS.sleep(250); + Monitor removedMonitor = spyMonitorService.remove(NoOpMonitor.class, key); + assertEquals(monitor, removedMonitor); + assertEquals(MonitorState.RUNNING, monitor.getState()); + } + + @Test + public void testStopAndRemove() throws SQLException, InterruptedException { + spyMonitorService.registerMonitorTypeIfAbsent( + NoOpMonitor.class, + TimeUnit.MINUTES.toNanos(1), + TimeUnit.MINUTES.toNanos(1), + // even though we pass a re-create policy, we should not re-create it if the monitor is expired since this + // indicates it is not being used. + new HashSet<>(Collections.singletonList(MonitorErrorResponse.RECREATE)), + null + ); + + String key = "testMonitor"; + NoOpMonitor monitor = spyMonitorService.runIfAbsent( + NoOpMonitor.class, + key, + mockStorageService, + mockTelemetryFactory, + mockConnectionProvider, + URL, + PROTOCOL, + mockTargetDriverDialect, + mockDbDialect, + props, + (mockServicesContainer) -> new NoOpMonitor(30) + ); + assertNotNull(monitor); + + // need to wait to give time for the monitor executor to start the monitor thread. + TimeUnit.MILLISECONDS.sleep(250); + spyMonitorService.stopAndRemove(NoOpMonitor.class, key); + assertNull(spyMonitorService.get(NoOpMonitor.class, key)); + assertEquals(MonitorState.STOPPED, monitor.getState()); + } + + static class NoOpMonitor extends AbstractMonitor { + protected NoOpMonitor(long terminationTimeoutSec) { + super(terminationTimeoutSec); + } + + @Override + public void monitor() { + // do nothing. + } + } +} From 85252d00d562a07ea1c6dccaa7a59c5958d7ab9f Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Mon, 15 Sep 2025 14:40:25 -0700 Subject: [PATCH 38/42] Cleanup --- .../main/java/software/amazon/jdbc/PartialPluginService.java | 4 ---- .../main/java/software/amazon/jdbc/util/ServiceUtility.java | 4 ++++ .../amazon/jdbc/util/connection/ConnectionServiceImpl.java | 4 ++++ 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java b/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java index c6b494651..06957b61f 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java +++ b/wrapper/src/main/java/software/amazon/jdbc/PartialPluginService.java @@ -109,10 +109,6 @@ public PartialPluginService( @NonNull final Dialect dbDialect, @Nullable final ConfigurationProfile configurationProfile) { this.servicesContainer = servicesContainer; - this.servicesContainer.setHostListProviderService(this); - this.servicesContainer.setPluginService(this); - this.servicesContainer.setPluginManagerService(this); - this.pluginManager = servicesContainer.getConnectionPluginManager(); this.props = props; this.originalUrl = originalUrl; diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/ServiceUtility.java b/wrapper/src/main/java/software/amazon/jdbc/util/ServiceUtility.java index 84d6be614..13d5d9cd6 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/ServiceUtility.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/ServiceUtility.java @@ -84,6 +84,10 @@ public FullServicesContainer createServiceContainer( dbDialect ); + servicesContainer.setHostListProviderService(partialPluginService); + servicesContainer.setPluginService(partialPluginService); + servicesContainer.setPluginManagerService(partialPluginService); + pluginManager.init(servicesContainer, props, partialPluginService, null); return servicesContainer; } diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java index 52744d494..cc509c26a 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java @@ -68,6 +68,10 @@ public ConnectionServiceImpl( dbDialect ); + servicesContainer.setHostListProviderService(partialPluginService); + servicesContainer.setPluginService(partialPluginService); + servicesContainer.setPluginManagerService(partialPluginService); + this.pluginService = partialPluginService; this.pluginManager.init(servicesContainer, props, partialPluginService, null); } From c465a006660f9047d0510824041d6bc20e015b61 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Mon, 15 Sep 2025 15:45:11 -0700 Subject: [PATCH 39/42] Cleanup --- .../ClusterAwareReaderFailoverHandler.java | 4 ---- .../util/connection/ConnectionService.java | 12 ++++++++++-- .../connection/ConnectionServiceImpl.java | 19 +++++++++++++++++-- .../FailoverConnectionPluginTest.java | 12 +++++------- 4 files changed, 32 insertions(+), 15 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java index 8cfed799a..e58378ddb 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java @@ -33,20 +33,16 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.logging.Logger; -import software.amazon.jdbc.ConnectionPluginManager; import software.amazon.jdbc.HostRole; import software.amazon.jdbc.HostSpec; -import software.amazon.jdbc.PartialPluginService; import software.amazon.jdbc.PluginService; import software.amazon.jdbc.hostavailability.HostAvailability; import software.amazon.jdbc.util.ExecutorFactory; import software.amazon.jdbc.util.FullServicesContainer; -import software.amazon.jdbc.util.FullServicesContainerImpl; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.PropertyUtils; import software.amazon.jdbc.util.ServiceUtility; import software.amazon.jdbc.util.Utils; -import software.amazon.jdbc.util.connection.ConnectionService; /** * An implementation of ReaderFailoverHandler. diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionService.java b/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionService.java index 5a6f4fcff..088f5d9d8 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionService.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionService.java @@ -22,6 +22,12 @@ import software.amazon.jdbc.HostSpec; import software.amazon.jdbc.PluginService; +/** + * @deprecated This interface is deprecated and will be removed in a future version. Use + * {@link software.amazon.jdbc.util.ServiceUtility#createServiceContainer} followed by + * {@link PluginService#forceConnect} instead. + */ +@Deprecated public interface ConnectionService { /** * Creates an auxiliary connection. Auxiliary connections are driver-internal connections that accomplish various @@ -31,8 +37,10 @@ public interface ConnectionService { * @param props the properties for the auxiliary connection. * @return a new connection to the given host using the given props. * @throws SQLException if an error occurs while opening the connection. + * @deprecated Use {@link software.amazon.jdbc.util.ServiceUtility#createServiceContainer} followed by + * {@link PluginService#forceConnect} instead. */ - Connection open(HostSpec hostSpec, Properties props) throws SQLException; + @Deprecated Connection open(HostSpec hostSpec, Properties props) throws SQLException; - PluginService getPluginService(); + @Deprecated PluginService getPluginService(); } diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java index c50770233..4ba674d59 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java @@ -16,6 +16,7 @@ package software.amazon.jdbc.util.connection; +import com.mchange.v2.util.PropertiesUtils; import java.sql.Connection; import java.sql.SQLException; import java.util.Properties; @@ -28,15 +29,26 @@ import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; import software.amazon.jdbc.util.FullServicesContainer; import software.amazon.jdbc.util.FullServicesContainerImpl; +import software.amazon.jdbc.util.PropertyUtils; import software.amazon.jdbc.util.monitoring.MonitorService; import software.amazon.jdbc.util.storage.StorageService; import software.amazon.jdbc.util.telemetry.TelemetryFactory; +/** + * @deprecated This class is deprecated and will be removed in a future version. Use + * {@link software.amazon.jdbc.util.ServiceUtility#createServiceContainer} followed by + * {@link PluginService#forceConnect} instead. + */ +@Deprecated public class ConnectionServiceImpl implements ConnectionService { protected final String targetDriverProtocol; protected final ConnectionPluginManager pluginManager; protected final PluginService pluginService; + /** + * @deprecated Use {@link software.amazon.jdbc.util.ServiceUtility#createServiceContainer} instead. + */ + @Deprecated public ConnectionServiceImpl( StorageService storageService, MonitorService monitorService, @@ -58,9 +70,10 @@ public ConnectionServiceImpl( telemetryFactory); servicesContainer.setConnectionPluginManager(this.pluginManager); + Properties propsCopy = PropertyUtils.copyProperties(props); PartialPluginService partialPluginService = new PartialPluginService( servicesContainer, - props, + propsCopy, originalUrl, this.targetDriverProtocol, driverDialect, @@ -72,15 +85,17 @@ public ConnectionServiceImpl( servicesContainer.setPluginManagerService(partialPluginService); this.pluginService = partialPluginService; - this.pluginManager.init(servicesContainer, props, partialPluginService, null); + this.pluginManager.init(servicesContainer, propsCopy, partialPluginService, null); } @Override + @Deprecated public Connection open(HostSpec hostSpec, Properties props) throws SQLException { return this.pluginManager.forceConnect(this.targetDriverProtocol, hostSpec, props, true, null); } @Override + @Deprecated public PluginService getPluginService() { return this.pluginService; } diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java index 8274235b9..637793772 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/failover/FailoverConnectionPluginTest.java @@ -63,7 +63,6 @@ import software.amazon.jdbc.util.FullServicesContainer; import software.amazon.jdbc.util.RdsUrlType; import software.amazon.jdbc.util.SqlState; -import software.amazon.jdbc.util.connection.ConnectionService; import software.amazon.jdbc.util.telemetry.GaugeCallable; import software.amazon.jdbc.util.telemetry.TelemetryContext; import software.amazon.jdbc.util.telemetry.TelemetryCounter; @@ -82,7 +81,6 @@ class FailoverConnectionPluginTest { .host("reader1").port(1234).role(HostRole.READER).build()); @Mock FullServicesContainer mockContainer; - @Mock ConnectionService mockConnectionService; @Mock PluginService mockPluginService; @Mock Connection mockConnection; @Mock HostSpec mockHostSpec; @@ -143,7 +141,7 @@ void init() throws SQLException { } @Test - void test_notifyNodeListChanged_withFailoverDisabled() throws SQLException { + void test_notifyNodeListChanged_withFailoverDisabled() { properties.setProperty(FailoverConnectionPlugin.ENABLE_CLUSTER_AWARE_FAILOVER.name, "false"); final Map> changes = new HashMap<>(); @@ -155,7 +153,7 @@ void test_notifyNodeListChanged_withFailoverDisabled() throws SQLException { } @Test - void test_notifyNodeListChanged_withValidConnectionNotInTopology() throws SQLException { + void test_notifyNodeListChanged_withValidConnectionNotInTopology() { final Map> changes = new HashMap<>(); changes.put("cluster-host/", EnumSet.of(NodeChangeOptions.NODE_DELETED)); changes.put("instance/", EnumSet.of(NodeChangeOptions.NODE_ADDED)); @@ -351,7 +349,7 @@ void test_failoverWriter_successFailover() throws SQLException { } @Test - void test_invalidCurrentConnection_withNoConnection() throws SQLException { + void test_invalidCurrentConnection_withNoConnection() { when(mockPluginService.getCurrentConnection()).thenReturn(null); initializePlugin(); spyPlugin.invalidateCurrentConnection(); @@ -376,7 +374,7 @@ void test_invalidateCurrentConnection_inTransaction() throws SQLException { } @Test - void test_invalidateCurrentConnection_notInTransaction() throws SQLException { + void test_invalidateCurrentConnection_notInTransaction() { when(mockPluginService.isInTransaction()).thenReturn(false); when(mockHostSpec.getHost()).thenReturn("host"); when(mockHostSpec.getPort()).thenReturn(123); @@ -437,7 +435,7 @@ void test_execute_withDirectExecute() throws SQLException { verify(mockHostListProvider, never()).getRdsUrlType(); } - private void initializePlugin() throws SQLException { + private void initializePlugin() { spyPlugin = spy(new FailoverConnectionPlugin(mockContainer, properties)); spyPlugin.setWriterFailoverHandler(mockWriterFailoverHandler); spyPlugin.setReaderFailoverHandler(mockReaderFailoverHandler); From b2a107ded5b4cc2e1b5904014be8c546623cf516 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Mon, 15 Sep 2025 15:51:03 -0700 Subject: [PATCH 40/42] cleanup --- .../amazon/ReadWriteSplittingPostgresExample.java | 2 -- .../ReadWriteSplittingSpringJdbcTemplateMySQLExample.java | 1 - .../src/main/java/software/amazon/HikariExample.java | 1 - .../src/main/java/example/spring/Config.java | 7 ------- .../main/java/software/amazon/jdbc/HostSpecBuilder.java | 1 - .../amazon/jdbc/authentication/AwsCredentialsManager.java | 1 - .../jdbc/hostavailability/HostAvailabilityStrategy.java | 2 -- .../plugin/federatedauth/CredentialsProviderFactory.java | 1 - .../jdbc/plugin/limitless/LimitlessRouterService.java | 2 -- .../amazon/jdbc/plugin/limitless/LimitlessRouters.java | 1 - .../fastestresponse/HostResponseTimeServiceImpl.java | 8 ++------ .../jdbc/targetdriverdialect/MariadbDriverHelper.java | 1 - .../software/amazon/jdbc/util/ConnectionUrlParser.java | 1 - .../java/software/amazon/jdbc/util/PropertyUtils.java | 1 - .../jdbc/util/connection/ConnectionServiceImpl.java | 1 - .../amazon/jdbc/util/monitoring/MonitorServiceImpl.java | 1 - .../amazon/jdbc/util/storage/SlidingExpirationCache.java | 1 - .../container/tests/hibernate/HibernateTests.java | 1 - .../src/test/java/integration/host/TestEnvironment.java | 1 - .../src/test/java/integration/util/ContainerHelper.java | 1 - .../amazon/jdbc/ConnectionPluginChainBuilderTests.java | 1 - .../src/test/java/software/amazon/jdbc/DialectTests.java | 1 - .../jdbc/authentication/AwsCredentialsManagerTest.java | 1 - .../jdbc/plugin/iam/IamAuthConnectionPluginTest.java | 4 ---- .../readwritesplitting/ReadWriteSplittingPluginTest.java | 1 - .../amazon/jdbc/states/SessionStateServiceImplTests.java | 1 - .../amazon/jdbc/util/ConnectionUrlParserTest.java | 1 - 27 files changed, 2 insertions(+), 44 deletions(-) diff --git a/examples/AWSDriverExample/src/main/java/software/amazon/ReadWriteSplittingPostgresExample.java b/examples/AWSDriverExample/src/main/java/software/amazon/ReadWriteSplittingPostgresExample.java index 1a4b42cf7..c4a273dc9 100644 --- a/examples/AWSDriverExample/src/main/java/software/amazon/ReadWriteSplittingPostgresExample.java +++ b/examples/AWSDriverExample/src/main/java/software/amazon/ReadWriteSplittingPostgresExample.java @@ -23,8 +23,6 @@ import java.sql.SQLException; import java.sql.Statement; import java.util.Properties; -import software.amazon.jdbc.ConnectionProviderManager; -import software.amazon.jdbc.HikariPooledConnectionProvider; import software.amazon.jdbc.HostSpec; import software.amazon.jdbc.PropertyDefinition; import software.amazon.jdbc.plugin.failover.FailoverFailedSQLException; diff --git a/examples/AWSDriverExample/src/main/java/software/amazon/ReadWriteSplittingSpringJdbcTemplateMySQLExample.java b/examples/AWSDriverExample/src/main/java/software/amazon/ReadWriteSplittingSpringJdbcTemplateMySQLExample.java index ead361b58..ea914e72e 100644 --- a/examples/AWSDriverExample/src/main/java/software/amazon/ReadWriteSplittingSpringJdbcTemplateMySQLExample.java +++ b/examples/AWSDriverExample/src/main/java/software/amazon/ReadWriteSplittingSpringJdbcTemplateMySQLExample.java @@ -16,7 +16,6 @@ package software.amazon; -import com.mysql.cj.jdbc.MysqlDataSource; import com.zaxxer.hikari.HikariDataSource; import java.sql.Connection; import java.sql.SQLException; diff --git a/examples/HikariExample/src/main/java/software/amazon/HikariExample.java b/examples/HikariExample/src/main/java/software/amazon/HikariExample.java index 9bbfa03c9..8f6607798 100644 --- a/examples/HikariExample/src/main/java/software/amazon/HikariExample.java +++ b/examples/HikariExample/src/main/java/software/amazon/HikariExample.java @@ -22,7 +22,6 @@ import java.sql.ResultSet; import java.sql.SQLException; import java.sql.Statement; -import java.util.Properties; public class HikariExample { diff --git a/examples/SpringHibernateBalancedReaderTwoDataSourceExample/src/main/java/example/spring/Config.java b/examples/SpringHibernateBalancedReaderTwoDataSourceExample/src/main/java/example/spring/Config.java index bc404df52..d6b30ffa8 100644 --- a/examples/SpringHibernateBalancedReaderTwoDataSourceExample/src/main/java/example/spring/Config.java +++ b/examples/SpringHibernateBalancedReaderTwoDataSourceExample/src/main/java/example/spring/Config.java @@ -16,10 +16,7 @@ package example.spring; -import com.zaxxer.hikari.HikariConfig; -import java.util.Arrays; import java.util.Properties; -import java.util.concurrent.TimeUnit; import javax.persistence.EntityManagerFactory; import javax.sql.DataSource; import org.hibernate.exception.JDBCConnectionException; @@ -40,10 +37,6 @@ import org.springframework.orm.jpa.vendor.HibernateJpaVendorAdapter; import org.springframework.transaction.PlatformTransactionManager; import org.springframework.transaction.annotation.EnableTransactionManagement; -import software.amazon.jdbc.HikariPooledConnectionProvider; -import software.amazon.jdbc.HostSpec; -import software.amazon.jdbc.profile.ConfigurationProfileBuilder; -import software.amazon.jdbc.profile.ConfigurationProfilePresetCodes; @Configuration @EnableTransactionManagement diff --git a/wrapper/src/main/java/software/amazon/jdbc/HostSpecBuilder.java b/wrapper/src/main/java/software/amazon/jdbc/HostSpecBuilder.java index a84920637..ec6abdbee 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/HostSpecBuilder.java +++ b/wrapper/src/main/java/software/amazon/jdbc/HostSpecBuilder.java @@ -17,7 +17,6 @@ package software.amazon.jdbc; import java.sql.Timestamp; -import java.time.Instant; import org.checkerframework.checker.nullness.qual.NonNull; import software.amazon.jdbc.hostavailability.HostAvailability; import software.amazon.jdbc.hostavailability.HostAvailabilityStrategy; diff --git a/wrapper/src/main/java/software/amazon/jdbc/authentication/AwsCredentialsManager.java b/wrapper/src/main/java/software/amazon/jdbc/authentication/AwsCredentialsManager.java index eb463a7aa..c562dd48e 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/authentication/AwsCredentialsManager.java +++ b/wrapper/src/main/java/software/amazon/jdbc/authentication/AwsCredentialsManager.java @@ -23,7 +23,6 @@ import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider; import software.amazon.jdbc.HostSpec; import software.amazon.jdbc.PropertyDefinition; -import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.StringUtils; public class AwsCredentialsManager { diff --git a/wrapper/src/main/java/software/amazon/jdbc/hostavailability/HostAvailabilityStrategy.java b/wrapper/src/main/java/software/amazon/jdbc/hostavailability/HostAvailabilityStrategy.java index 659fdae2b..9f04b54c6 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/hostavailability/HostAvailabilityStrategy.java +++ b/wrapper/src/main/java/software/amazon/jdbc/hostavailability/HostAvailabilityStrategy.java @@ -16,8 +16,6 @@ package software.amazon.jdbc.hostavailability; -import software.amazon.jdbc.AwsWrapperProperty; - public interface HostAvailabilityStrategy { void setHostAvailability(HostAvailability hostAvailability); diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/CredentialsProviderFactory.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/CredentialsProviderFactory.java index a43396bf9..655fbdda4 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/CredentialsProviderFactory.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/federatedauth/CredentialsProviderFactory.java @@ -16,7 +16,6 @@ package software.amazon.jdbc.plugin.federatedauth; -import java.io.Closeable; import java.sql.SQLException; import java.util.Properties; import org.checkerframework.checker.nullness.qual.NonNull; diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterService.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterService.java index 2d3d87a08..1e3a04560 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterService.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouterService.java @@ -16,9 +16,7 @@ package software.amazon.jdbc.plugin.limitless; -import java.sql.Connection; import java.sql.SQLException; -import java.util.List; import java.util.Properties; import org.checkerframework.checker.nullness.qual.NonNull; import software.amazon.jdbc.HostSpec; diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouters.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouters.java index 0793dbcff..ea7cab3ce 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouters.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/limitless/LimitlessRouters.java @@ -20,7 +20,6 @@ import java.util.Objects; import org.checkerframework.checker.nullness.qual.NonNull; import software.amazon.jdbc.HostSpec; -import software.amazon.jdbc.hostlistprovider.Topology; public class LimitlessRouters { private final @NonNull List hosts; diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/strategy/fastestresponse/HostResponseTimeServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/strategy/fastestresponse/HostResponseTimeServiceImpl.java index f6b9cf177..ee157b3ea 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/strategy/fastestresponse/HostResponseTimeServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/strategy/fastestresponse/HostResponseTimeServiceImpl.java @@ -20,16 +20,13 @@ import java.util.ArrayList; import java.util.List; import java.util.Properties; -import java.util.Set; import java.util.logging.Logger; -import java.util.stream.Collectors; import org.checkerframework.checker.nullness.qual.NonNull; import software.amazon.jdbc.HostSpec; import software.amazon.jdbc.PluginService; import software.amazon.jdbc.util.FullServicesContainer; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.Utils; -import software.amazon.jdbc.util.storage.SlidingExpirationCacheWithCleanupThread; public class HostResponseTimeServiceImpl implements HostResponseTimeService { @@ -73,7 +70,7 @@ public void setHosts(final @NonNull List hosts) { // Going through all hosts in the topology and trying to find new ones. this.hosts.stream() // hostSpec is not in the set of hosts that already being monitored - .filter(hostSpec ->!Utils.containsHostAndPort(oldHosts, hostSpec.getHostAndPort())) + .filter(hostSpec -> !Utils.containsHostAndPort(oldHosts, hostSpec.getHostAndPort())) .forEach(hostSpec -> { try { this.servicesContainer.getMonitorService().runIfAbsent( @@ -88,8 +85,7 @@ public void setHosts(final @NonNull List hosts) { this.pluginService.getDialect(), this.props, (servicesContainer) -> - new NodeResponseTimeMonitor(pluginService, hostSpec, this.props, - this.intervalMs)); + new NodeResponseTimeMonitor(pluginService, hostSpec, this.props, this.intervalMs)); } catch (SQLException e) { LOGGER.warning( Messages.get("HostResponseTimeServiceImpl.errorStartingMonitor", new Object[] {hostSpec.getUrl(), e})); diff --git a/wrapper/src/main/java/software/amazon/jdbc/targetdriverdialect/MariadbDriverHelper.java b/wrapper/src/main/java/software/amazon/jdbc/targetdriverdialect/MariadbDriverHelper.java index 8c93819c3..d67638b6b 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/targetdriverdialect/MariadbDriverHelper.java +++ b/wrapper/src/main/java/software/amazon/jdbc/targetdriverdialect/MariadbDriverHelper.java @@ -18,7 +18,6 @@ import static software.amazon.jdbc.util.ConnectionUrlBuilder.buildUrl; -import com.mysql.cj.jdbc.Driver; import java.sql.DriverManager; import java.sql.SQLException; import java.util.Collections; diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/ConnectionUrlParser.java b/wrapper/src/main/java/software/amazon/jdbc/util/ConnectionUrlParser.java index 81323335b..435907141 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/ConnectionUrlParser.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/ConnectionUrlParser.java @@ -16,7 +16,6 @@ package software.amazon.jdbc.util; -import com.fasterxml.jackson.databind.annotation.JsonAppend.Prop; import java.util.ArrayList; import java.util.List; import java.util.Map; diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/PropertyUtils.java b/wrapper/src/main/java/software/amazon/jdbc/util/PropertyUtils.java index ce2f66fad..632aaac6c 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/PropertyUtils.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/PropertyUtils.java @@ -28,7 +28,6 @@ import java.util.Set; import java.util.logging.Logger; import org.checkerframework.checker.nullness.qual.NonNull; -import software.amazon.awssdk.services.rds.endpoints.internal.Value.Bool; import software.amazon.jdbc.AwsWrapperProperty; import software.amazon.jdbc.PropertyDefinition; diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java index 4ba674d59..14f3c0166 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java @@ -16,7 +16,6 @@ package software.amazon.jdbc.util.connection; -import com.mchange.v2.util.PropertiesUtils; import java.sql.Connection; import java.sql.SQLException; import java.util.Properties; diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java index e75e3c2f4..5e0df4817 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/monitoring/MonitorServiceImpl.java @@ -41,7 +41,6 @@ import software.amazon.jdbc.targetdriverdialect.TargetDriverDialect; import software.amazon.jdbc.util.ExecutorFactory; import software.amazon.jdbc.util.FullServicesContainer; -import software.amazon.jdbc.util.FullServicesContainer; import software.amazon.jdbc.util.Messages; import software.amazon.jdbc.util.PropertyUtils; import software.amazon.jdbc.util.ServiceUtility; diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/storage/SlidingExpirationCache.java b/wrapper/src/main/java/software/amazon/jdbc/util/storage/SlidingExpirationCache.java index 7f670866e..604b603f2 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/storage/SlidingExpirationCache.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/storage/SlidingExpirationCache.java @@ -20,7 +20,6 @@ import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.Objects; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; diff --git a/wrapper/src/test/java/integration/container/tests/hibernate/HibernateTests.java b/wrapper/src/test/java/integration/container/tests/hibernate/HibernateTests.java index 01d263f6b..fd4e1eac9 100644 --- a/wrapper/src/test/java/integration/container/tests/hibernate/HibernateTests.java +++ b/wrapper/src/test/java/integration/container/tests/hibernate/HibernateTests.java @@ -20,7 +20,6 @@ import static org.junit.jupiter.api.Assertions.assertNotNull; import integration.DatabaseEngine; -import integration.DatabaseEngineDeployment; import integration.DriverHelper; import integration.TestEnvironmentFeatures; import integration.container.ConnectionStringHelper; diff --git a/wrapper/src/test/java/integration/host/TestEnvironment.java b/wrapper/src/test/java/integration/host/TestEnvironment.java index 424b741ea..ce639a152 100644 --- a/wrapper/src/test/java/integration/host/TestEnvironment.java +++ b/wrapper/src/test/java/integration/host/TestEnvironment.java @@ -50,7 +50,6 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.logging.Logger; -import org.testcontainers.containers.BindMode; import org.testcontainers.containers.GenericContainer; import org.testcontainers.containers.Network; import org.testcontainers.containers.ToxiproxyContainer; diff --git a/wrapper/src/test/java/integration/util/ContainerHelper.java b/wrapper/src/test/java/integration/util/ContainerHelper.java index bc977dcda..6745badd8 100644 --- a/wrapper/src/test/java/integration/util/ContainerHelper.java +++ b/wrapper/src/test/java/integration/util/ContainerHelper.java @@ -42,7 +42,6 @@ import org.testcontainers.containers.ToxiproxyContainer; import org.testcontainers.containers.output.FrameConsumerResultCallback; import org.testcontainers.containers.output.OutputFrame; -import org.testcontainers.containers.startupcheck.StartupCheckStrategy; import org.testcontainers.containers.wait.strategy.LogMessageWaitStrategy; import org.testcontainers.containers.wait.strategy.Wait; import org.testcontainers.images.builder.ImageFromDockerfile; diff --git a/wrapper/src/test/java/software/amazon/jdbc/ConnectionPluginChainBuilderTests.java b/wrapper/src/test/java/software/amazon/jdbc/ConnectionPluginChainBuilderTests.java index 84efe5ba3..4352d0622 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/ConnectionPluginChainBuilderTests.java +++ b/wrapper/src/test/java/software/amazon/jdbc/ConnectionPluginChainBuilderTests.java @@ -28,7 +28,6 @@ import java.util.HashSet; import java.util.List; import java.util.Properties; -import java.util.Set; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; diff --git a/wrapper/src/test/java/software/amazon/jdbc/DialectTests.java b/wrapper/src/test/java/software/amazon/jdbc/DialectTests.java index 0caefc973..4170f8556 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/DialectTests.java +++ b/wrapper/src/test/java/software/amazon/jdbc/DialectTests.java @@ -35,7 +35,6 @@ import org.mockito.MockitoAnnotations; import software.amazon.jdbc.dialect.AuroraMysqlDialect; import software.amazon.jdbc.dialect.AuroraPgDialect; -import software.amazon.jdbc.dialect.DialectManager; import software.amazon.jdbc.dialect.MariaDbDialect; import software.amazon.jdbc.dialect.MysqlDialect; import software.amazon.jdbc.dialect.PgDialect; diff --git a/wrapper/src/test/java/software/amazon/jdbc/authentication/AwsCredentialsManagerTest.java b/wrapper/src/test/java/software/amazon/jdbc/authentication/AwsCredentialsManagerTest.java index 6028f5706..56d24e076 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/authentication/AwsCredentialsManagerTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/authentication/AwsCredentialsManagerTest.java @@ -17,7 +17,6 @@ package software.amazon.jdbc.authentication; import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.when; diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPluginTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPluginTest.java index 92a0c0ff6..a872ac96c 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPluginTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/iam/IamAuthConnectionPluginTest.java @@ -44,17 +44,13 @@ import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider; import software.amazon.awssdk.regions.Region; -import software.amazon.awssdk.services.rds.RdsUtilities; -import software.amazon.awssdk.services.rds.TestDefaultRdsUtilities; import software.amazon.jdbc.Driver; import software.amazon.jdbc.HostSpec; import software.amazon.jdbc.HostSpecBuilder; import software.amazon.jdbc.JdbcCallable; import software.amazon.jdbc.PluginService; import software.amazon.jdbc.PropertyDefinition; -import software.amazon.jdbc.authentication.AwsCredentialsManager; import software.amazon.jdbc.dialect.Dialect; -import software.amazon.jdbc.hostavailability.HostAvailabilityStrategy; import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; import software.amazon.jdbc.plugin.TokenInfo; import software.amazon.jdbc.util.RdsUtils; diff --git a/wrapper/src/test/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPluginTest.java b/wrapper/src/test/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPluginTest.java index ff7282483..c7c7bdc1b 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPluginTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPluginTest.java @@ -45,7 +45,6 @@ import org.junit.jupiter.api.Test; import org.mockito.Mock; import org.mockito.MockitoAnnotations; -import software.amazon.jdbc.HikariPooledConnectionProvider; import software.amazon.jdbc.HostListProviderService; import software.amazon.jdbc.HostRole; import software.amazon.jdbc.HostSpec; diff --git a/wrapper/src/test/java/software/amazon/jdbc/states/SessionStateServiceImplTests.java b/wrapper/src/test/java/software/amazon/jdbc/states/SessionStateServiceImplTests.java index 8ab28907d..ed391f4cf 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/states/SessionStateServiceImplTests.java +++ b/wrapper/src/test/java/software/amazon/jdbc/states/SessionStateServiceImplTests.java @@ -17,7 +17,6 @@ package software.amazon.jdbc.states; import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.spy; import static org.mockito.Mockito.times; diff --git a/wrapper/src/test/java/software/amazon/jdbc/util/ConnectionUrlParserTest.java b/wrapper/src/test/java/software/amazon/jdbc/util/ConnectionUrlParserTest.java index 082d2f78b..056fccd78 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/util/ConnectionUrlParserTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/util/ConnectionUrlParserTest.java @@ -32,7 +32,6 @@ import software.amazon.jdbc.HostRole; import software.amazon.jdbc.HostSpec; import software.amazon.jdbc.HostSpecBuilder; -import software.amazon.jdbc.PropertyDefinition; import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy; class ConnectionUrlParserTest { From e6357455fd49b3429cf50470fe21bfaa39df56df Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Mon, 15 Sep 2025 15:58:59 -0700 Subject: [PATCH 41/42] Fix checkstyle --- .../util/connection/ConnectionService.java | 20 +++++++++--- .../connection/ConnectionServiceImpl.java | 8 +++-- .../monitoring/MonitorServiceImplTest.java | 32 +++++++++---------- 3 files changed, 38 insertions(+), 22 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionService.java b/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionService.java index 088f5d9d8..e0748bc72 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionService.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionService.java @@ -21,11 +21,14 @@ import java.util.Properties; import software.amazon.jdbc.HostSpec; import software.amazon.jdbc.PluginService; +import software.amazon.jdbc.util.FullServicesContainer; /** + * A service used to open new connections for internal driver use. + * * @deprecated This interface is deprecated and will be removed in a future version. Use - * {@link software.amazon.jdbc.util.ServiceUtility#createServiceContainer} followed by - * {@link PluginService#forceConnect} instead. + * {@link software.amazon.jdbc.util.ServiceUtility#createServiceContainer} followed by + * {@link PluginService#forceConnect} instead. */ @Deprecated public interface ConnectionService { @@ -40,7 +43,16 @@ public interface ConnectionService { * @deprecated Use {@link software.amazon.jdbc.util.ServiceUtility#createServiceContainer} followed by * {@link PluginService#forceConnect} instead. */ - @Deprecated Connection open(HostSpec hostSpec, Properties props) throws SQLException; + @Deprecated + Connection open(HostSpec hostSpec, Properties props) throws SQLException; - @Deprecated PluginService getPluginService(); + /** + * Get the {@link PluginService} associated with this {@link ConnectionService}. + * + * @return the {@link PluginService} associated with this {@link ConnectionService} + * @deprecated Use {@link software.amazon.jdbc.util.ServiceUtility#createServiceContainer} followed by + * {@link FullServicesContainer#getPluginService()} instead. + */ + @Deprecated + PluginService getPluginService(); } diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java b/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java index 14f3c0166..9e8356eb9 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java +++ b/wrapper/src/main/java/software/amazon/jdbc/util/connection/ConnectionServiceImpl.java @@ -34,9 +34,11 @@ import software.amazon.jdbc.util.telemetry.TelemetryFactory; /** + * A service used to open new connections for internal driver use. + * * @deprecated This class is deprecated and will be removed in a future version. Use - * {@link software.amazon.jdbc.util.ServiceUtility#createServiceContainer} followed by - * {@link PluginService#forceConnect} instead. + * {@link software.amazon.jdbc.util.ServiceUtility#createServiceContainer} followed by + * {@link PluginService#forceConnect} instead. */ @Deprecated public class ConnectionServiceImpl implements ConnectionService { @@ -45,6 +47,8 @@ public class ConnectionServiceImpl implements ConnectionService { protected final PluginService pluginService; /** + * Constructs a {@link ConnectionServiceImpl} instance. + * * @deprecated Use {@link software.amazon.jdbc.util.ServiceUtility#createServiceContainer} instead. */ @Deprecated diff --git a/wrapper/src/test/java/software/amazon/jdbc/util/monitoring/MonitorServiceImplTest.java b/wrapper/src/test/java/software/amazon/jdbc/util/monitoring/MonitorServiceImplTest.java index 02bd5060e..450b494a7 100644 --- a/wrapper/src/test/java/software/amazon/jdbc/util/monitoring/MonitorServiceImplTest.java +++ b/wrapper/src/test/java/software/amazon/jdbc/util/monitoring/MonitorServiceImplTest.java @@ -54,8 +54,8 @@ class MonitorServiceImplTest { @Mock TargetDriverDialect mockTargetDriverDialect; @Mock Dialect mockDbDialect; @Mock EventPublisher mockPublisher; - String URL = "jdbc:postgresql://somehost/somedb"; - String PROTOCOL = "someProtocol"; + String url = "jdbc:postgresql://somehost/somedb"; + String protocol = "someProtocol"; Properties props = new Properties(); MonitorServiceImpl spyMonitorService; private AutoCloseable closeable; @@ -69,8 +69,8 @@ void setUp() throws SQLException { eq(mockStorageService), eq(mockConnectionProvider), eq(mockTelemetryFactory), - eq(URL), - eq(PROTOCOL), + eq(url), + eq(protocol), eq(mockTargetDriverDialect), eq(mockDbDialect), eq(props)); @@ -98,8 +98,8 @@ public void testMonitorError_monitorReCreated() throws SQLException, Interrupted mockStorageService, mockTelemetryFactory, mockConnectionProvider, - URL, - PROTOCOL, + url, + protocol, mockTargetDriverDialect, mockDbDialect, props, @@ -142,8 +142,8 @@ public void testMonitorStuck_monitorReCreated() throws SQLException, Interrupted mockStorageService, mockTelemetryFactory, mockConnectionProvider, - URL, - PROTOCOL, + url, + protocol, mockTargetDriverDialect, mockDbDialect, props, @@ -188,8 +188,8 @@ public void testMonitorExpired() throws SQLException, InterruptedException { mockStorageService, mockTelemetryFactory, mockConnectionProvider, - URL, - PROTOCOL, + url, + protocol, mockTargetDriverDialect, mockDbDialect, props, @@ -221,8 +221,8 @@ public void testMonitorMismatch() { mockStorageService, mockTelemetryFactory, mockConnectionProvider, - URL, - PROTOCOL, + url, + protocol, mockTargetDriverDialect, mockDbDialect, props, @@ -251,8 +251,8 @@ public void testRemove() throws SQLException, InterruptedException { mockStorageService, mockTelemetryFactory, mockConnectionProvider, - URL, - PROTOCOL, + url, + protocol, mockTargetDriverDialect, mockDbDialect, props, @@ -286,8 +286,8 @@ public void testStopAndRemove() throws SQLException, InterruptedException { mockStorageService, mockTelemetryFactory, mockConnectionProvider, - URL, - PROTOCOL, + url, + protocol, mockTargetDriverDialect, mockDbDialect, props, From 6311d6356e91ffd5c9c370cf1380599ac78b6660 Mon Sep 17 00:00:00 2001 From: aaron-congo Date: Fri, 19 Sep 2025 15:58:30 -0700 Subject: [PATCH 42/42] Reuse service containers for reader failover ConnectionAttemptTasks --- .../failover/ClusterAwareReaderFailoverHandler.java | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java index e58378ddb..43c402ab7 100644 --- a/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java +++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/failover/ClusterAwareReaderFailoverHandler.java @@ -297,11 +297,14 @@ private ReaderFailoverResult getConnectionFromHostGroup(final List hos final ExecutorService executor = ExecutorFactory.newFixedThreadPool(2, "failover"); final CompletionService completionService = new ExecutorCompletionService<>(executor); + final FullServicesContainer servicesContainer1 = this.getNewServicesContainer(); + final FullServicesContainer servicesContainer2 = this.getNewServicesContainer(); + try { for (int i = 0; i < hosts.size(); i += 2) { // submit connection attempt tasks in batches of 2 final ReaderFailoverResult result = - getResultFromNextTaskBatch(hosts, executor, completionService, i); + getResultFromNextTaskBatch(hosts, executor, completionService, servicesContainer1, servicesContainer2, i); if (result.isConnected() || result.getException() != null) { return result; } @@ -324,12 +327,14 @@ private ReaderFailoverResult getResultFromNextTaskBatch( final List hosts, final ExecutorService executor, final CompletionService completionService, + final FullServicesContainer servicesContainer1, + final FullServicesContainer servicesContainer2, final int i) throws SQLException { ReaderFailoverResult result; final int numTasks = i + 1 < hosts.size() ? 2 : 1; completionService.submit( new ConnectionAttemptTask( - getNewServicesContainer(), + servicesContainer1, this.hostAvailabilityMap, hosts.get(i), this.props, @@ -337,7 +342,7 @@ private ReaderFailoverResult getResultFromNextTaskBatch( if (numTasks == 2) { completionService.submit( new ConnectionAttemptTask( - getNewServicesContainer(), + servicesContainer2, this.hostAvailabilityMap, hosts.get(i + 1), this.props,