chore: close reader connection even if it has expired to avoid connection leaks

Author: karenc-bq

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

@@ -276,6 +276,7 @@ private void setWriterConnection(final Connection writerConnection,
   }
 
   private void setReaderConnection(final Connection conn, final HostSpec host) {
+    closeReaderConnectionIfIdle(this.readerConnection);
     this.readerConnection = new CacheItem<>(conn, this.getKeepAliveTimeout(host));
     this.readerHostSpec = host;
     LOGGER.finest(
@@ -380,8 +381,8 @@ private void switchToWriterConnection(
       switchCurrentConnectionTo(this.writerConnection, writerHost);
     }
 
-    if (this.isReaderConnFromInternalPool && this.readerConnection != null) {
-      this.closeConnectionIfIdle(this.readerConnection.get());
+    if (this.isReaderConnFromInternalPool) {
+      this.closeReaderConnectionIfIdle(this.readerConnection);
     }
 
     LOGGER.finer(() -> Messages.get("ReadWriteSplittingPlugin.switchedFromReaderToWriter",
@@ -412,9 +413,9 @@ private void switchToReaderConnection(final List<HostSpec> hosts)
       return;
     }
 
-    if (this.readerConnection != null && this.readerHostSpec != null && !hosts.contains(this.readerHostSpec)) {
+    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.
-      closeConnectionIfIdle(this.readerConnection.get());
+      closeReaderConnectionIfIdle(this.readerConnection);
     }
 
     this.inReadWriteSplit = true;
@@ -435,22 +436,13 @@ private void switchToReaderConnection(final List<HostSpec> hosts)
               new Object[] {this.readerHostSpec.getUrl()}));
         }
 
-        Connection conn = this.readerConnection.get(true);
-        if (isConnectionUsable(conn)) {
-          try {
-            conn.close();
-          } catch (SQLException ex) {
-            // Do nothing
-          }
-        }
-        this.readerConnection = null;
-        this.readerHostSpec = null;
+        closeReaderConnectionIfIdle(this.readerConnection);
         initializeReaderConnection(hosts);
       }
     }
 
     if (this.isWriterConnFromInternalPool) {
-      this.closeConnectionIfIdle(this.writerConnection);
+      this.closeWriterConnectionIfIdle(this.writerConnection);
     }
   }
 
@@ -540,26 +532,37 @@ public void releaseResources() {
 
   private void closeIdleConnections() {
     LOGGER.finest(() -> Messages.get("ReadWriteSplittingPlugin.closingInternalConnections"));
-    if (this.readerConnection != null) {
-      closeConnectionIfIdle(this.readerConnection.get());
+    closeReaderConnectionIfIdle(this.readerConnection);
+    closeWriterConnectionIfIdle(this.writerConnection);
+  }
+
+  void closeReaderConnectionIfIdle(CacheItem<Connection> readerConnection) {
+    if (readerConnection == null) {
+      return;
     }
-    closeConnectionIfIdle(this.writerConnection);
+
+    final Connection currentConnection = this.pluginService.getCurrentConnection();
+    final Connection readerConnectionCache = readerConnection.get(true);
+
+    try {
+      if (isConnectionUsable(readerConnectionCache) && readerConnectionCache != currentConnection) {
+        readerConnectionCache.close();
+      }
+    } catch (SQLException e) {
+      // Do nothing.
+    }
+
+    this.readerConnection = null;
+    this.readerHostSpec = null;
   }
 
-  void closeConnectionIfIdle(final Connection internalConnection) {
+  void closeWriterConnectionIfIdle(final Connection internalConnection) {
     final Connection currentConnection = this.pluginService.getCurrentConnection();
     try {
       if (isConnectionUsable(internalConnection)
           && internalConnection != currentConnection) {
         internalConnection.close();
-        if (internalConnection == writerConnection) {
-          writerConnection = null;
-        }
-
-        if (internalConnection == readerConnection.get()) {
-          readerConnection = null;
-          readerHostSpec = null;
-        }
+        writerConnection = null;
       }
     } catch (final SQLException e) {
       // ignore

wrapper/src/test/java/software/amazon/jdbc/plugin/readwritesplitting/ReadWriteSplittingPluginTest.java

@@ -607,7 +607,7 @@ public void testClosePooledReaderConnectionAfterSetReadOnly() throws SQLExceptio
     spyPlugin.switchConnectionIfRequired(true);
     spyPlugin.switchConnectionIfRequired(false);
 
-    verify(spyPlugin, times(1)).closeConnectionIfIdle(eq(mockReaderConn1));
+    verify(spyPlugin, times(1)).closeReaderConnectionIfIdle(any());
   }
 
   @Test
@@ -634,7 +634,7 @@ public void testClosePooledWriterConnectionAfterSetReadOnly() throws SQLExceptio
     spyPlugin.switchConnectionIfRequired(false);
     spyPlugin.switchConnectionIfRequired(true);
 
-    verify(spyPlugin, times(1)).closeConnectionIfIdle(eq(mockWriterConn));
+    verify(spyPlugin, times(1)).closeWriterConnectionIfIdle(eq(mockWriterConn));
   }
 
   private static HikariConfig getHikariConfig(HostSpec hostSpec, Properties props) {