Fix closing connection and synchronizations (#3661)

Author: dstepanov

.github/workflows/gradle.yml

@@ -38,6 +38,9 @@ jobs:
         run: |
          sudo rm -rf "/usr/local/share/boost"
          sudo rm -rf "$AGENT_TOOLSDIRECTORY"
+         sudo rm -rf "/opt/ghc"
+         sudo rm -rf "/usr/share/dotnet"
+         sudo rm -rf "/usr/local/lib/android"
          sudo apt-get clean
          df -h
 

data-connection-jdbc/src/test/groovy/io/micronaut/data/connection/reactive/ReactiveConnectionSynchronizationSpec.groovy

@@ -0,0 +1,78 @@
+package io.micronaut.data.connection.reactive
+
+import io.micronaut.data.connection.ConnectionDefinition
+import io.micronaut.data.connection.ConnectionSynchronization
+import org.reactivestreams.Publisher
+import reactor.core.publisher.Mono
+import spock.lang.Specification
+
+import java.util.concurrent.atomic.AtomicInteger
+
+class ReactiveConnectionSynchronizationSpec extends Specification {
+
+    static final class Tracker {
+        final AtomicInteger executionComplete = new AtomicInteger()
+        final AtomicInteger beforeClosed = new AtomicInteger()
+        final AtomicInteger afterClosed = new AtomicInteger()
+    }
+
+    def "reactive: all registered synchronizations execute even if one throws"() {
+        given:
+        def tracker = new Tracker()
+        // new reactive connection status with a dummy connection
+        def status = new DefaultReactiveConnectionStatus<>(new Object(), ConnectionDefinition.DEFAULT, true)
+
+        and: "a normal synchronization that tracks all callbacks"
+        status.registerSynchronization(new ConnectionSynchronization() {
+            @Override
+            void executionComplete() {
+                tracker.executionComplete.incrementAndGet()
+            }
+
+            @Override
+            void beforeClosed() {
+                tracker.beforeClosed.incrementAndGet()
+            }
+
+            @Override
+            void afterClosed() {
+                tracker.afterClosed.incrementAndGet()
+            }
+        })
+
+        and: "a synchronization that throws during executionComplete"
+        status.registerReactiveSynchronization(new ReactiveConnectionSynchronization() {
+            @Override
+            Publisher<Void> onComplete() {
+                return Mono.defer {
+                    throw new RuntimeException("simulated onComplete failure")
+                }
+            }
+
+            @Override
+            Publisher<Void> onClose() {
+                return Mono.empty()
+            }
+
+            @Override
+            Publisher<Void> afterClose() {
+                return Mono.empty()
+            }
+        })
+
+        and: "a close supplier that succeeds"
+        def closeSupplier = { Mono.empty() } as java.util.function.Supplier<Publisher<Void>>
+
+        when: "we signal completion which triggers executionComplete and finally chain"
+        try {
+            Mono.from(status.onComplete(closeSupplier)).block()
+        } catch (Throwable ignored) {
+            // expected: we only care callbacks were invoked despite failure
+        }
+
+        then: "all normal synchronization callbacks should have been invoked at least once"
+        tracker.executionComplete.get() >= 1
+        tracker.beforeClosed.get() >= 1
+        tracker.afterClosed.get() >= 1
+    }
+}

data-connection/src/main/java/io/micronaut/data/connection/reactive/DefaultReactiveConnectionStatus.java

@@ -112,12 +112,29 @@ private Publisher<Void> forEachSynchronizations(Function<ReactiveConnectionSynch
             return Mono.empty();
         }
         Mono<Void> chain = Mono.empty();
+        List<Throwable> exceptions = new ArrayList<>(connectionSynchronizations.size());
         ListIterator<ReactiveConnectionSynchronization> listIterator = connectionSynchronizations.listIterator(connectionSynchronizations.size());
         while (listIterator.hasPrevious()) {
-            Publisher<Void> next = consumer.apply(listIterator.previous());
-            if (next != EMPTY) {
-                chain = chain.then(Mono.from(next));
+            try {
+                Publisher<Void> next = consumer.apply(listIterator.previous());
+                Mono<Void> wrapped = next == EMPTY ? Mono.empty() : Mono.from(next).onErrorResume(e -> {
+                    exceptions.add(e);
+                    return Mono.empty();
+                });
+                chain = chain.then(wrapped);
+            } catch (Exception e) {
+                exceptions.add(e);
+            }
+        }
+        if (!exceptions.isEmpty()) {
+            if (exceptions.size() == 1) {
+                return chain.then(Mono.error(exceptions.get(0)));
+            }
+            IllegalStateException e = new IllegalStateException("Error executing connection synchronizations", exceptions.get(0));
+            for (int i = 1; i < exceptions.size(); i++) {
+                e.addSuppressed(exceptions.get(i));
             }
+            return chain.then(Mono.error(e));
         }
         return chain;
     }

data-connection/src/main/java/io/micronaut/data/connection/support/AbstractConnectionOperations.java

@@ -207,10 +207,13 @@ public void complete(@NonNull ConnectionStatus<C> status) {
             try {
                 connectionStatus.beforeClosed();
             } finally {
-                if (connectionStatus.isNew()) {
-                    closeConnection(status);
+                try {
+                    if (connectionStatus.isNew()) {
+                        closeConnection(status);
+                    }
+                } finally {
+                    connectionStatus.afterClosed();
                 }
-                connectionStatus.afterClosed();
             }
         }
     }

data-connection/src/main/java/io/micronaut/data/connection/support/DefaultConnectionStatus.java

@@ -74,13 +74,34 @@ public void registerSynchronization(ConnectionSynchronization synchronization) {
 
     private void forEachSynchronizations(Consumer<ConnectionSynchronization> consumer) {
         if (connectionSynchronizations != null) {
+            List<Exception> exceptions = new ArrayList<>(connectionSynchronizations.size());
             ListIterator<ConnectionSynchronization> listIterator = connectionSynchronizations.listIterator(connectionSynchronizations.size());
             while (listIterator.hasPrevious()) {
-                consumer.accept(listIterator.previous());
+                try {
+                    consumer.accept(listIterator.previous());
+                } catch (Exception e) {
+                    exceptions.add(e);
+                }
+            }
+            if (!exceptions.isEmpty()) {
+                if (exceptions.size() == 1) {
+                    sneakyThrow(exceptions.get(0));
+                } else {
+                    IllegalStateException e = new IllegalStateException("Error executing connection synchronizations", exceptions.get(0));
+                    for (int i = 1; i < exceptions.size(); i++) {
+                        e.addSuppressed(exceptions.get(i));
+                    }
+                    throw e;
+                }
             }
         }
     }
 
+    @SuppressWarnings("unchecked")
+    private static <T extends Throwable, R> R sneakyThrow(Throwable t) throws T {
+        throw (T) t;
+    }
+
     public void complete() {
         forEachSynchronizations(ConnectionSynchronization::executionComplete);
     }

data-tx-jdbc/src/test/groovy/io/micronaut/transaction/jdbc/BrokenConnectionPropagationSpec.groovy

@@ -0,0 +1,70 @@
+package io.micronaut.transaction.jdbc
+
+import io.micronaut.test.extensions.spock.annotation.MicronautTest
+import io.micronaut.transaction.SynchronousTransactionManager
+import io.micronaut.transaction.TransactionDefinition
+import io.micronaut.transaction.exceptions.TransactionSystemException
+import io.micronaut.transaction.jdbc.mock.MockConnection
+import io.micronaut.transaction.jdbc.mock.MockDataSource
+import io.micronaut.transaction.jdbc.mock.TestSyncTracker
+import io.micronaut.transaction.jdbc.mock.SyncTrackerSynchronization
+import io.micronaut.transaction.jdbc.mock.ThrowingExecutionCompleteSynchronization
+import jakarta.inject.Inject
+
+import java.sql.Connection
+import java.sql.SQLException
+
+@MicronautTest(transactional = false, environments = "broken-conn")
+class BrokenConnectionPropagationSpec extends spock.lang.Specification {
+
+    @Inject
+    SynchronousTransactionManager<Connection> txManager
+
+    @Inject
+    MockDataSource mockDataSource
+
+    @Inject
+    TestSyncTracker tracker
+
+    def "reproducer: closed connection remains cached in PropagatedContext after rollback/reset throws"() {
+        given: "a new transaction is opened using the synchronous transaction API (context-bound connection)"
+        tracker.reset()
+        def status = txManager.getTransaction(TransactionDefinition.DEFAULT)
+        def firstConn = status.getConnection() as MockConnection
+        int firstId = firstConn.id()
+
+        and: "register connection synchronizations including one that throws during executionComplete"
+        status.getConnectionStatus().registerSynchronization(new SyncTrackerSynchronization(tracker, 0))
+        status.getConnectionStatus().registerSynchronization(new ThrowingExecutionCompleteSynchronization(1))
+
+        when: "the underlying connection becomes broken/closed before rollback (e.g. socket timeout)"
+        mockDataSource.getLastConnection().breakAndClose()
+        txManager.rollback(status)
+
+        then: "rollback fails (simulated driver throws due to closed connection or sync failure)"
+        def ex = thrown(Throwable)
+        assert ex instanceof io.micronaut.data.connection.exceptions.ConnectionException || ex instanceof IllegalStateException
+        and: "all connection synchronizations executed even when one throws"
+        assert tracker.executionComplete.get() >= 1
+        assert tracker.beforeClosed.get() >= 1
+        assert tracker.afterClosed.get() >= 1
+
+        when: "on the same thread, a subsequent transaction is started"
+        def status2 = txManager.getTransaction(TransactionDefinition.DEFAULT)
+        def newConn = status2.getConnection() as MockConnection
+        int secondId = newConn.id()
+        // any simple operation should now succeed on a fresh, open connection
+        newConn.getAutoCommit()
+
+        then: "a fresh connection is obtained and is not the previously closed one"
+        secondId != firstId
+        and: "a new connection was created"
+        mockDataSource.getCreatedCount() == 2
+
+        when: "the new transaction is committed to complete the cleanup path"
+        txManager.commit(status2)
+
+        then:
+        noExceptionThrown()
+    }
+}