fix: safety close in AsyncConnectionProvider #2935

src/main/java/io/lettuce/core/internal/AsyncConnectionProvider.java

@@ -142,11 +142,14 @@ public CompletableFuture<Void> close() {
 
         List<CompletableFuture<Void>> futures = new ArrayList<>();
 
-        forEach((connectionKey, closeable) -> {
-
-            futures.add(closeable.closeAsync());
-            connections.remove(connectionKey);
-        });
+        for (K k : connections.keySet()) {
+            Sync<K, T, F> remove = connections.remove(k);
+            if (remove != null) {
+                CompletionStage<Void> closeFuture = remove.future.thenAccept(AsyncCloseable::closeAsync);
+                // always synchronously add the future, made it immutably in Futures.allOf()
+                futures.add(closeFuture.toCompletableFuture());
+            }
+        }
 
         return Futures.allOf(futures);
     }
@@ -160,9 +163,8 @@ public void close(K key) {
 
         LettuceAssert.notNull(key, "ConnectionKey must not be null!");
 
-        Sync<K, T, F> sync = connections.get(key);
+        Sync<K, T, F> sync = connections.remove(key);
         if (sync != null) {
-            connections.remove(key);
             sync.doWithConnection(AsyncCloseable::closeAsync);
         }
     }
@@ -217,7 +219,6 @@ static class Sync<K, T extends AsyncCloseable, F extends CompletionStage<T>> {
 
         @SuppressWarnings("unchecked")
         public Sync(K key, F future) {
-
             this.key = key;
             this.future = (F) future.whenComplete((connection, throwable) -> {
 

src/test/java/io/lettuce/core/cluster/AsyncConnectionProviderIntegrationTests.java

@@ -32,7 +32,9 @@
 import javax.inject.Inject;
 
 import org.apache.commons.lang3.time.StopWatch;
+import org.junit.Assert;
 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.junit.jupiter.api.extension.ExtendWith;
@@ -125,7 +127,7 @@ void shouldCloseConnections() throws IOException {
         ConnectionKey connectionKey = new ConnectionKey(ConnectionIntent.READ, TestSettings.host(), TestSettings.port());
 
         sut.getConnection(connectionKey);
-        TestFutures.awaitOrTimeout(sut.close());
+        assertThatThrownBy(() -> TestFutures.awaitOrTimeout(sut.close())).isInstanceOf(IllegalStateException.class);
 
         assertThat(sut.getConnectionCount()).isEqualTo(0);
         TestFutures.awaitOrTimeout(sut.close());

src/test/java/io/lettuce/core/internal/AsyncConnectionProviderTest.java

@@ -0,0 +1,99 @@
+package io.lettuce.core.internal;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+public class AsyncConnectionProviderTest {
+
+    @Test
+    public void testFutureListLength() throws InterruptedException, ExecutionException, TimeoutException {
+
+        CountDownLatch slowCreate = new CountDownLatch(1);
+        CountDownLatch slowShutdown = new CountDownLatch(1);
+
+        // create a provider with a slow connection creation
+        AsyncConnectionProvider<String, AsyncCloseable, CompletableFuture<AsyncCloseable>> provider = new AsyncConnectionProvider<>(
+                key -> {
+                    return countDownFuture(slowCreate, new io.lettuce.core.api.AsyncCloseable() {
+
+                        @Override
+                        public CompletableFuture<Void> closeAsync() {
+                            return CompletableFuture.completedFuture(null);
+                        }
+
+                    });
+                });
+
+        // add slow shutdown connection first
+        SlowCloseFuture slowCloseFuture = new SlowCloseFuture(slowShutdown);
+        provider.register("slowShutdown", new io.lettuce.core.api.AsyncCloseable() {
+
+            @Override
+            public CompletableFuture<Void> closeAsync() {
+                return slowCloseFuture;
+            }
+
+        });
+
+        // add slow creation connection
+        CompletableFuture<AsyncCloseable> createFuture = provider.getConnection("slowCreate");
+
+        // close the connection.
+        CompletableFuture<Void> closeFuture = provider.close();
+
+        // the connection has not been created yet, so the close futures array always has 1 element
+        // we block the iterator on the slowCloseFuture
+        // then we count down the creation, the close future will be added to the list
+        slowCreate.countDown();
+
+        // the close future is added to the list, we unblock the iterator
+        slowShutdown.countDown();
+
+        // assert close future is completed, and no exceptions are thrown
+        closeFuture.get(10, TimeUnit.SECONDS);
+        Assert.assertTrue(createFuture.isDone());
+    }
+
+    private <T> CompletableFuture<T> countDownFuture(CountDownLatch countDownLatch, T value) {
+        return CompletableFuture.runAsync(() -> {
+            try {
+                countDownLatch.await(1, TimeUnit.SECONDS);
+            } catch (InterruptedException e) {
+                throw new RuntimeException(e);
+            }
+        }).thenApply(v -> value);
+    }
+
+    static class SlowCloseFuture extends CompletableFuture<Void> {
+
+        private final CountDownLatch countDownLatch;
+
+        SlowCloseFuture(CountDownLatch countDownLatch) {
+            this.countDownLatch = countDownLatch;
+        }
+
+        @Override
+        public CompletableFuture<Void> toCompletableFuture() {
+            // we block the iterator on here
+            try {
+                countDownLatch.await(1, TimeUnit.SECONDS);
+            } catch (InterruptedException e) {
+                throw new RuntimeException(e);
+            }
+            return super.toCompletableFuture();
+        }
+
+        @Override
+        public Void get() {
+            return null;
+        }
+
+    }
+
+}