Fix r2dbc connection pool returning multiple pooled objects pointing to same database connection

src/main/java/io/r2dbc/pool/ConnectionPool.java

@@ -48,6 +48,7 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.BiPredicate;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -110,7 +111,7 @@ public ConnectionPool(ConnectionPoolConfiguration configuration) {
         Function<Connection, Mono<Void>> allocateValidation = getValidationFunction(configuration);
 
         Mono<Connection> create = Mono.defer(() -> {
-
+            AtomicBoolean emitting = new AtomicBoolean(true);
             Mono<Connection> mono = this.connectionPool.acquire()
                     .flatMap(ref -> {
 
@@ -140,7 +141,10 @@ public ConnectionPool(ConnectionPoolConfiguration configuration) {
                         conn = conn.onErrorResume(throwable -> ref.invalidate().then(Mono.error(throwable)));
 
                         return Operators.discardOnCancel(conn, () -> {
-                            ref.release().subscribe();
+                            if (emitting.compareAndSet(true, false)) {
+                                logger.debug("Discarding pooled reference as creation was disposed");
+                                ref.release().subscribe();
+                            }
                             return false;
                         });
                     })
@@ -167,7 +171,8 @@ public ConnectionPool(ConnectionPoolConfiguration configuration) {
                     return context.put(HOOK_ON_DROPPED, onNextDropped);
                 }).onErrorMap(TimeoutException.class, e -> new R2dbcTimeoutException(timeoutMessage, e));
             }
-            return mono;
+            return mono
+                .doOnTerminate(() -> emitting.compareAndSet(true, false));
         });
         this.create = configuration.getAcquireRetry() > 0 ? create.retry(configuration.getAcquireRetry()) : create;
     }

src/test/java/io/r2dbc/pool/ConnectionPoolUnitTests.java

@@ -31,6 +31,7 @@
 import org.reactivestreams.Subscription;
 import org.springframework.util.ReflectionUtils;
 import reactor.core.Disposable;
+import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.scheduler.Schedulers;
 import reactor.test.StepVerifier;
@@ -47,8 +48,10 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
@@ -60,8 +63,10 @@
 import static java.lang.String.format;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
+import static org.assertj.core.api.Assertions.assertThatNoException;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.awaitility.Awaitility.await;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
@@ -1119,6 +1124,38 @@ public void cancel() {
 
         assertThat(closed).isTrue();
     }
+
+    @Test
+    void shouldNotDuplicateObjectsDuringCancel() {
+        Set<Connection> uniqueConnections = ConcurrentHashMap.newKeySet();
+        ConnectionFactory connectionFactoryMock = mock(ConnectionFactory.class);
+
+        when(connectionFactoryMock.create()).thenReturn((Publisher)Mono.fromCallable(() -> {
+            ConnectionWithLifecycle connectionMock = mock(ConnectionWithLifecycle.class);
+            when(connectionMock.validate(any())).thenReturn(Mono.just(true));
+            when(connectionMock.postAllocate()).thenReturn(Mono.empty());
+            when(connectionMock.preRelease()).thenReturn(Mono.<Void>fromRunnable(() -> 
+                    uniqueConnections.remove(connectionMock))
+              );
+            return connectionMock;
+        }));
+
+        ConnectionPool pool = new ConnectionPool(ConnectionPoolConfiguration.builder(connectionFactoryMock)
+                .build());
+        CompletableFuture<Long> future = Flux.range(0, 32)
+            .flatMap(i -> Mono.fromCallable(() -> true)
+                .then(Mono.defer(() -> Mono.usingWhen(Flux.defer(pool::create), 
+                    c -> {
+                        assertFalse(!uniqueConnections.add(((PooledConnection)c).unwrap()), 
+                            "duplicate connections returned from pool");
+                        return Mono.delay(Duration.ofMillis(1));
+                    }, 
+                    Connection::close))))
+            .ignoreElements()
+            .toFuture();
+
+        assertThatNoException().isThrownBy(future::join);
+    }
 
     interface ConnectionWithLifecycle extends Connection, Lifecycle {