From 988eb0e1918af80670180af0a28af8a7704a2b2e Mon Sep 17 00:00:00 2001
From: Michal Wyrzykowski <mwyrzykowski@atlassian.com>
Date: Tue, 25 Jul 2023 08:44:54 +0200
Subject: [PATCH] Fix connection leak in AuroraClusterDiscovery

---
 CHANGELOG.md                                  |  7 +++-
 .../aurora/AuroraClusterDiscovery.java        | 42 +++++++++++++------
 2 files changed, 35 insertions(+), 14 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1345e0f..df5c4b7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -22,7 +22,12 @@ Dropping a requirement of a major version of a dependency is a new contract.
 
 
 ## [Unreleased]
-[Unreleased]: https://github.com/atlassian-labs/db-replica/compare/release-2.9.4...master
+[Unreleased]: https://github.com/atlassian-labs/db-replica/compare/release-2.10.0...master
+
+### Fix
+- Connection leak in `AuroraClusterDiscovery`
+
+## [2.10.0] - 2023-07-24
 
 ### Added
 - Support for custom connection provider for replica discovery in `AuroraMultiReplicaConsistency`
diff --git a/src/main/java/com/atlassian/db/replica/internal/aurora/AuroraClusterDiscovery.java b/src/main/java/com/atlassian/db/replica/internal/aurora/AuroraClusterDiscovery.java
index 4b25298..1ad9090 100644
--- a/src/main/java/com/atlassian/db/replica/internal/aurora/AuroraClusterDiscovery.java
+++ b/src/main/java/com/atlassian/db/replica/internal/aurora/AuroraClusterDiscovery.java
@@ -14,6 +14,7 @@
 import java.sql.SQLException;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.List;
 import java.util.function.Supplier;
 
 import static java.lang.String.format;
@@ -44,29 +45,44 @@ private AuroraClusterDiscovery(
     }
 
     public Collection<Database> getReplicas(Supplier<Connection> connectionSupplier) {
-        return discoveredReplicasCache.get(() -> fetchReplicas(discovererConnectionSupplier != null ? discovererConnectionSupplier : connectionSupplier)).orElseGet(
-            () -> {
-                lazyLogger.debug(() -> "AuroraClusterDiscovery#getReplicas no replicas cached.");
-                return Collections.emptyList();
-            });
+        return discoveredReplicasCache
+            .get(() -> discovererConnectionSupplier != null ? fetchReplicas() : fetchReplicas(connectionSupplier)).orElseGet(
+                () -> {
+                    lazyLogger.debug(() -> "AuroraClusterDiscovery#getReplicas no replicas cached.");
+                    return Collections.emptyList();
+                });
     }
 
     private Collection<Database> fetchReplicas(Supplier<Connection> connectionSupplier) {
         try {
             final Connection connection = connectionSupplier.get();
-            final AuroraReplicasDiscoverer discoverer = createDiscoverer(connection);
-            return discoverer.fetchReplicasUrls(connection).stream().map(auroraUrl -> {
-                JdbcUrl url = auroraUrl.toJdbcUrl();
-                return new AuroraReplicaNode(
-                    auroraUrl.getEndpoint().getServerId(),
-                    replicaConnectionPerUrlProvider.getReplicaConnectionProvider(url)
-                );
-            }).collect(toList());
+            return getDatabases(connection);
         } catch (SQLException exception) {
             throw new ReadReplicaDiscoveryOperationException(exception);
         }
     }
 
+    private Collection<Database> fetchReplicas() {
+        try {
+            try (final Connection connection = discovererConnectionSupplier.get()) {
+                return getDatabases(connection);
+            }
+        } catch (SQLException exception) {
+            throw new ReadReplicaDiscoveryOperationException(exception);
+        }
+    }
+
+    private List<Database> getDatabases(Connection connection) throws SQLException {
+        final AuroraReplicasDiscoverer discoverer = createDiscoverer(connection);
+        return discoverer.fetchReplicasUrls(connection).stream().map(auroraUrl -> {
+            JdbcUrl url = auroraUrl.toJdbcUrl();
+            return new AuroraReplicaNode(
+                auroraUrl.getEndpoint().getServerId(),
+                replicaConnectionPerUrlProvider.getReplicaConnectionProvider(url)
+            );
+        }).collect(toList());
+    }
+
     private AuroraReplicasDiscoverer createDiscoverer(Connection connection) {
         if (this.clusterUri != null) {
             return createDiscovererFromClusterUri();