Autoclose fetched replica connections when checking consistency (#125)

build.gradle.kts

@@ -34,6 +34,7 @@ dependencies {
     testImplementation("org.junit.jupiter:junit-jupiter")
     testImplementation("org.assertj:assertj-core:3.19.0")
     testImplementation("org.mockito:mockito-core:3.11.0")
+    testImplementation("org.mockito:mockito-junit-jupiter:3.11.0")
     testImplementation("org.threeten:threeten-extra:1.5.0")
     testImplementation("com.github.docker-java:docker-java-core:3.2.6")
     testImplementation("com.github.docker-java:docker-java-transport-httpclient5:3.2.6")

gradle/dependency-locks/testCompileClasspath.lockfile

@@ -29,6 +29,7 @@ org.junit.jupiter:junit-jupiter:5.7.2
 org.junit.platform:junit-platform-commons:1.7.2
 org.junit:junit-bom:5.7.2
 org.mockito:mockito-core:3.11.0
+org.mockito:mockito-junit-jupiter:3.11.0
 org.objenesis:objenesis:3.2
 org.opentest4j:opentest4j:1.2.0
 org.postgresql:postgresql:42.2.18

gradle/dependency-locks/testRuntimeClasspath.lockfile

@@ -32,6 +32,7 @@ org.junit.platform:junit-platform-commons:1.7.2
 org.junit.platform:junit-platform-engine:1.7.2
 org.junit:junit-bom:5.7.2
 org.mockito:mockito-core:3.11.0
+org.mockito:mockito-junit-jupiter:3.11.0
 org.objenesis:objenesis:3.2
 org.opentest4j:opentest4j:1.2.0
 org.postgresql:postgresql:42.2.18

src/main/java/com/atlassian/db/replica/api/AuroraMultiReplicaConsistency.java

@@ -1,12 +1,14 @@
 package com.atlassian.db.replica.api;
 
+import com.atlassian.db.replica.api.exception.ConnectionCouldNotBeClosedException;
 import com.atlassian.db.replica.internal.NoCacheSuppliedCache;
 import com.atlassian.db.replica.internal.aurora.AuroraClusterDiscovery;
 import com.atlassian.db.replica.spi.ReplicaConnectionPerUrlProvider;
 import com.atlassian.db.replica.spi.ReplicaConsistency;
 import com.atlassian.db.replica.spi.SuppliedCache;
 
 import java.sql.Connection;
+import java.sql.SQLException;
 import java.util.Collection;
 import java.util.function.Supplier;
 
@@ -37,14 +39,15 @@ public void write(Connection main) {
 
     @Override
     public boolean isConsistent(Supplier<Connection> replicaSupplier) {
-        Collection<Database> replicas = cluster.getReplicas(replicaSupplier);
-
-        if (replicas.isEmpty()) {
-            return false;
-        }
-
-        return replicas.stream()
-            .allMatch(replica -> replicaConsistency.isConsistent(() -> replica.getConnectionSupplier().get()));
+        return cluster.getReplicas(replicaSupplier)
+            .stream()
+            .allMatch(replica -> {
+                try (Connection connection = replica.getConnectionSupplier().get()) {
+                    return replicaConsistency.isConsistent(() -> connection);
+                } catch (SQLException exception) {
+                    throw new ConnectionCouldNotBeClosedException(exception);
+                }
+            });
     }
 
     public static final class Builder {

src/main/java/com/atlassian/db/replica/api/exception/ConnectionCouldNotBeClosedException.java

@@ -0,0 +1,7 @@
+package com.atlassian.db.replica.api.exception;
+
+public final class ConnectionCouldNotBeClosedException extends RuntimeException {
+    public ConnectionCouldNotBeClosedException(Throwable cause) {
+        super(cause);
+    }
+}

src/test/java/com/atlassian/db/replica/api/AuroraMultiReplicaConsistencyTest.java

@@ -0,0 +1,108 @@
+package com.atlassian.db.replica.api;
+
+import com.atlassian.db.replica.internal.aurora.AuroraReplicaNode;
+import com.atlassian.db.replica.spi.ReplicaConnectionPerUrlProvider;
+import com.atlassian.db.replica.spi.ReplicaConsistency;
+import com.atlassian.db.replica.spi.SuppliedCache;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.util.Collection;
+import java.util.LinkedList;
+import java.util.Optional;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.*;
+
+@ExtendWith(MockitoExtension.class)
+class AuroraMultiReplicaConsistencyTest {
+    @Mock
+    private Connection supplierConnection;
+    @Mock
+    private SuppliedCache<Collection<Database>> discoveredReplicasCache;
+
+    private AuroraMultiReplicaConsistency sut;
+
+    private Collection<Connection> mockReplicaConnections(int count) {
+        Collection<Connection> replicas = new LinkedList<>();
+        Collection<Database> nodes = new LinkedList<>();
+
+        ReplicaConnectionPerUrlProvider replicaConnectionPerUrlProvider = (url) -> () -> {
+            Connection replica = mock(Connection.class);
+            replicas.add(replica);
+            return replica;
+        };
+
+        for (int i = 0; i < count; i++) {
+            AuroraReplicaNode node = new AuroraReplicaNode(
+                String.valueOf(i),
+                replicaConnectionPerUrlProvider.getReplicaConnectionProvider(null)
+            );
+            nodes.add(node);
+        }
+        when(discoveredReplicasCache.get(any())).thenReturn(Optional.of(nodes));
+
+        return replicas;
+    }
+
+    @Test
+    void shouldBeConsistentWhenNoReplicasAreAvailable() {
+        //given
+        mockReplicaConnections(0);
+
+        sut = AuroraMultiReplicaConsistency.builder()
+            .replicaConsistency(mock(ReplicaConsistency.class))
+            .discoveredReplicasCache(discoveredReplicasCache)
+            .build();
+
+        //when
+        boolean consistent = sut.isConsistent(() -> supplierConnection);
+
+        //then
+        assertTrue(consistent);
+    }
+
+    @Test
+    void shouldCloseReplicaConnectionsWhenConsistent() throws SQLException {
+        //given
+        Collection<Connection> replicas = mockReplicaConnections(5);
+
+        sut = AuroraMultiReplicaConsistency.builder()
+            .replicaConsistency(new ReplicaConsistencyMock(true))
+            .discoveredReplicasCache(discoveredReplicasCache)
+            .build();
+
+        //when
+        sut.isConsistent(() -> supplierConnection);
+
+        //then
+        verify(supplierConnection, never()).close();
+        for (Connection replica : replicas) {
+            verify(replica).close();
+        }
+    }
+
+    @Test
+    void shouldCloseReplicaConnectionsWhenNotConsistent() throws SQLException {
+        //given
+        Collection<Connection> replicas = mockReplicaConnections(5);
+
+        sut = AuroraMultiReplicaConsistency.builder()
+            .replicaConsistency(new ReplicaConsistencyMock(false))
+            .discoveredReplicasCache(discoveredReplicasCache)
+            .build();
+
+        //when
+        sut.isConsistent(() -> supplierConnection);
+
+        //then
+        verify(supplierConnection, never()).close();
+        for (Connection replica : replicas) {
+            verify(replica).close();
+        }
+    }
+}

src/test/java/com/atlassian/db/replica/api/ReplicaConsistencyMock.java

@@ -0,0 +1,25 @@
+package com.atlassian.db.replica.api;
+
+import com.atlassian.db.replica.spi.ReplicaConsistency;
+
+import java.sql.Connection;
+import java.util.function.Supplier;
+
+public class ReplicaConsistencyMock implements ReplicaConsistency {
+    private final boolean consistent;
+
+    public ReplicaConsistencyMock(boolean consistent) {
+        this.consistent = consistent;
+    }
+
+    @Override
+    public void write(Connection main) {
+
+    }
+
+    @Override
+    public boolean isConsistent(Supplier<Connection> replica) {
+        Connection connection = replica.get();
+        return consistent;
+    }
+}