Atomically prepare and validate connections during acquisition.

We now ensure that no cancel signals interrupt preparation and connection acquisition by using the discardOnCancel(…) operator.

[#140]

Signed-off-by: Mark Paluch <[email protected]>

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

@@ -45,6 +45,7 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiPredicate;
 import java.util.function.Consumer;
@@ -91,37 +92,38 @@ public ConnectionPool(ConnectionPoolConfiguration configuration) {
             });
         }
 
-        String acqName = String.format("Connection Acquisition from [%s]", configuration.getConnectionFactory());
-        String timeoutMessage = String.format("Connection Acquisition timed out after %dms", this.maxAcquireTime.toMillis());
+        String acqName = String.format("Connection acquisition from [%s]", configuration.getConnectionFactory());
+        String timeoutMessage = String.format("Connection acquisition timed out after %dms", this.maxAcquireTime.toMillis());
 
         Function<Connection, Mono<Void>> allocateValidation = getValidationFunction(configuration);
 
         Mono<Connection> create = Mono.defer(() -> {
 
             AtomicReference<PooledRef<Connection>> emitted = new AtomicReference<>();
+            AtomicBoolean cancelled = new AtomicBoolean();
 
             Mono<Connection> mono = this.connectionPool.acquire()
                 .doOnNext(emitted::set)
                 .doOnSubscribe(subscription -> {
 
                     if (logger.isDebugEnabled()) {
-                        logger.debug("Obtaining new connection from the driver");
+                        logger.debug("Obtaining new connection from the pool");
                     }
                 })
                 .flatMap(ref -> {
 
-                    PooledConnection connection = new PooledConnection(ref);
-                    return allocateValidation.apply(connection).thenReturn((Connection) connection).onErrorResume(throwable -> {
-                        return ref.invalidate().then(Mono.error(throwable));
-                    });
-                })
-                .doOnCancel(() -> {
+                    Mono<Connection> conn = getValidConnection(allocateValidation, ref);
 
-                    PooledRef<Connection> ref = emitted.get();
-                    if (ref != null && emitted.compareAndSet(ref, null)) {
-                        ref.release().subscribe();
-                    }
-                }).name(acqName);
+                    return conn.onErrorResume(throwable -> {
+                            emitted.set(null); // prevent release on cancel
+                            return ref.invalidate().then(Mono.error(throwable));
+                        })
+                        .doFinally(s -> cleanup(cancelled, emitted))
+                        .as(self -> Operators.discardOnCancel(self, () -> cancelled.set(true)));
+                })
+                .as(self -> Operators.discardOnCancel(self, () -> cancelled.set(true)))
+                .name(acqName)
+                .doOnNext(it -> emitted.set(null));
 
             if (!this.maxAcquireTime.isZero()) {
                 mono = mono.timeout(this.maxAcquireTime).onErrorMap(TimeoutException.class, e -> new R2dbcTimeoutException(timeoutMessage, e));
@@ -131,6 +133,24 @@ public ConnectionPool(ConnectionPoolConfiguration configuration) {
         this.create = configuration.getAcquireRetry() > 0 ? create.retry(configuration.getAcquireRetry()) : create;
     }
 
+    static void cleanup(AtomicBoolean cancelled, AtomicReference<PooledRef<Connection>> emitted) {
+
+        if (cancelled.compareAndSet(true, false)) {
+
+            PooledRef<Connection> savedRef = emitted.get();
+            if (savedRef != null && emitted.compareAndSet(savedRef, null)) {
+                logger.debug("Releasing connection after cancellation");
+                savedRef.release().subscribe(ignore -> {
+                }, e -> logger.warn("Error during release", e));
+            }
+        }
+    }
+
+    private Mono<Connection> getValidConnection(Function<Connection, Mono<Void>> allocateValidation, PooledRef<Connection> ref) {
+        PooledConnection connection = new PooledConnection(ref);
+        return allocateValidation.apply(connection).thenReturn(connection);
+    }
+
     private Function<Connection, Mono<Void>> getValidationFunction(ConnectionPoolConfiguration configuration) {
 
         String timeoutMessage = String.format("Validation timed out after %dms", this.maxAcquireTime.toMillis());

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

@@ -0,0 +1,126 @@
+/*
+ * Copyright 2021 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.r2dbc.pool;
+
+import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
+import reactor.core.CoreSubscriber;
+import reactor.core.publisher.Mono;
+import reactor.core.publisher.MonoOperator;
+import reactor.core.publisher.Operators;
+import reactor.util.Logger;
+import reactor.util.Loggers;
+import reactor.util.context.Context;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * A decorating operator that replays signals from its source to a {@link Subscriber} and drains the source upon {@link Subscription#cancel() cancel} and drops data signals until termination.
+ * Draining data is required to complete a particular request/response window and clear the protocol state as client code expects to start a request/response conversation without any previous
+ * response state.
+ */
+class MonoDiscardOnCancel<T> extends MonoOperator<T, T> {
+
+    private static final Logger logger = Loggers.getLogger(MonoDiscardOnCancel.class);
+
+    private final Runnable cancelConsumer;
+
+    MonoDiscardOnCancel(Mono<? extends T> source, Runnable cancelConsumer) {
+        super(source);
+        this.cancelConsumer = cancelConsumer;
+    }
+
+    @Override
+    public void subscribe(CoreSubscriber<? super T> actual) {
+        this.source.subscribe(new MonoDiscardOnCancelSubscriber<>(actual, this.cancelConsumer));
+    }
+
+    static class MonoDiscardOnCancelSubscriber<T> extends AtomicBoolean implements CoreSubscriber<T>, Subscription {
+
+        final CoreSubscriber<T> actual;
+
+        final Context ctx;
+
+        final Runnable cancelConsumer;
+
+        Subscription s;
+
+        MonoDiscardOnCancelSubscriber(CoreSubscriber<T> actual, Runnable cancelConsumer) {
+
+            this.actual = actual;
+            this.ctx = actual.currentContext();
+            this.cancelConsumer = cancelConsumer;
+        }
+
+        @Override
+        public void onSubscribe(Subscription s) {
+
+            if (Operators.validate(this.s, s)) {
+                this.s = s;
+                this.actual.onSubscribe(this);
+            }
+        }
+
+        @Override
+        public void onNext(T t) {
+
+            if (this.get()) {
+                Operators.onDiscard(t, this.ctx);
+                return;
+            }
+
+            this.actual.onNext(t);
+        }
+
+        @Override
+        public void onError(Throwable t) {
+            if (!this.get()) {
+                this.actual.onError(t);
+            }
+        }
+
+        @Override
+        public void onComplete() {
+            if (!this.get()) {
+                this.actual.onComplete();
+            }
+        }
+
+        @Override
+        public void request(long n) {
+            this.s.request(n);
+        }
+
+        @Override
+        public void cancel() {
+
+            if (compareAndSet(false, true)) {
+                if (logger.isDebugEnabled()) {
+                    logger.debug("received cancel signal");
+                }
+                try {
+                    this.cancelConsumer.run();
+                } catch (Exception e) {
+                    Operators.onErrorDropped(e, this.ctx);
+                }
+                this.s.request(Long.MAX_VALUE);
+            }
+        }
+
+    }
+
+}

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

@@ -0,0 +1,63 @@
+/*
+ * Copyright 2021 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.r2dbc.pool;
+
+import org.reactivestreams.Subscription;
+import reactor.core.publisher.Mono;
+
+/**
+ * Operator utility.
+ *
+ * @since 0.8.8
+ */
+final class Operators {
+
+    private Operators() {
+    }
+
+    /**
+     * Replay signals from {@link Mono the source} until cancellation. Drains the source for data signals if the subscriber cancels the subscription.
+     * <p>
+     * Draining data is required to complete a particular request/response window and clear the protocol state as client code expects to start a request/response conversation without leaving
+     * previous frames on the stack.
+     *
+     * @param source the source to decorate.
+     * @param <T>    The type of values in both source and output sequences.
+     * @return decorated {@link Mono}.
+     */
+    public static <T> Mono<T> discardOnCancel(Mono<? extends T> source) {
+        return new MonoDiscardOnCancel<>(source, () -> {
+        });
+    }
+
+    /**
+     * Replay signals from {@link Mono the source} until cancellation. Drains the source for data signals if the subscriber cancels the subscription.
+     * <p>
+     * Draining data is required to complete a particular request/response window and clear the protocol state as client code expects to start a request/response conversation without leaving
+     * previous frames on the stack.
+     * <p>Propagate the {@link Subscription#cancel()}  signal to a {@link Runnable consumer}.
+     *
+     * @param source         the source to decorate.
+     * @param cancelConsumer {@link Runnable} notified when the resulting {@link Mono} receives a {@link Subscription#cancel() cancel} signal.
+     * @param <T>            The type of values in both source and output sequences.
+     * @return decorated {@link Mono}.
+     */
+    public static <T> Mono<T> discardOnCancel(Mono<? extends T> source, Runnable cancelConsumer) {
+        return new MonoDiscardOnCancel<>(source, cancelConsumer);
+    }
+
+}

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

@@ -38,7 +38,7 @@ static Mono<Void> validate(Connection connection, ValidationDepth depth) {
         return Flux.from(connection.validate(depth)).handle((state, sink) -> {
 
             if (state) {
-                sink.complete();
+                sink.next(state);
                 return;
             }
 

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

@@ -27,6 +27,7 @@
 import org.junit.jupiter.api.Test;
 import org.reactivestreams.Publisher;
 import org.springframework.util.ReflectionUtils;
+import reactor.core.Disposable;
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
@@ -42,8 +43,10 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
@@ -831,6 +834,40 @@ void shouldReportPoolInPool() {
         new ConnectionPool(configuration);
     }
 
+    @Test
+    void cancelDuringAllocationShouldCompleteAtomically() throws InterruptedException {
+
+        ConnectionFactory connectionFactoryMock = mock(ConnectionFactory.class);
+        Connection connectionMock = mock(Connection.class);
+
+        CountDownLatch validateLatch = new CountDownLatch(1);
+        CountDownLatch postLatch = new CountDownLatch(1);
+        AtomicBoolean seenCancel = new AtomicBoolean();
+        Mono<Boolean> validate = Mono.just(true).doOnSuccess(s -> validateLatch.countDown()).delayElement(Duration.ofSeconds(1)).doOnSuccess(s -> postLatch.countDown()).doOnCancel(() -> {
+            seenCancel.set(true);
+        });
+
+        when(connectionFactoryMock.create()).thenAnswer(it -> Mono.just(connectionMock));
+        when(connectionMock.validate(any())).thenReturn(validate);
+
+        ConnectionPoolConfiguration configuration = ConnectionPoolConfiguration.builder(connectionFactoryMock)
+            .build();
+
+        ConnectionPool pool = new ConnectionPool(configuration);
+        Disposable subscribe = pool.create().subscribe();
+        validateLatch.await();
+        subscribe.dispose();
+        postLatch.await();
+
+        PoolMetrics poolMetrics = pool.getMetrics().get();
+        await().atMost(Duration.ofSeconds(1)).until(() -> poolMetrics.idleSize() == 10);
+
+        assertThat(seenCancel).isFalse();
+        assertThat(poolMetrics.pendingAcquireSize()).isEqualTo(0);
+        assertThat(poolMetrics.allocatedSize()).isEqualTo(10);
+        assertThat(poolMetrics.idleSize()).isEqualTo(10);
+    }
+
     private ConnectionPool createConnectionPoolForDisposeTest(Connection connectionMock) {
         ConnectionFactory connectionFactoryMock = mock(ConnectionFactory.class);