diff --git a/src/main/java/org/springframework/data/r2dbc/function/DefaultDatabaseClient.java b/src/main/java/org/springframework/data/r2dbc/function/DefaultDatabaseClient.java index 56bb5fe..d158fb6 100644 --- a/src/main/java/org/springframework/data/r2dbc/function/DefaultDatabaseClient.java +++ b/src/main/java/org/springframework/data/r2dbc/function/DefaultDatabaseClient.java @@ -38,6 +38,7 @@ import java.util.LinkedHashSet; import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.BiFunction; import java.util.function.Function; import java.util.function.Supplier; @@ -122,15 +123,16 @@ class DefaultDatabaseClient implements DatabaseClient, ConnectionAccessor { Assert.notNull(action, "Callback object must not be null"); - Mono connectionMono = getConnection(); - // Create close-suppressing Connection proxy, also preparing returned Statements. + Mono connectionMono = getConnection() + .map(it -> new ConnectionCloseHolder(it, this::closeConnection)); return Mono.usingWhen(connectionMono, it -> { - Connection connectionToUse = createConnectionProxy(it); + // Create close-suppressing Connection proxy + Connection connectionToUse = createConnectionProxy(it.connection); return doInConnection(connectionToUse, action); - }, this::closeConnection, this::closeConnection, this::closeConnection) // + }, ConnectionCloseHolder::close, ConnectionCloseHolder::close, ConnectionCloseHolder::close) // .onErrorMap(R2dbcException.class, ex -> translateException("execute", getSql(action), ex)); } @@ -149,15 +151,16 @@ class DefaultDatabaseClient implements DatabaseClient, ConnectionAccessor { Assert.notNull(action, "Callback object must not be null"); - Mono connectionMono = getConnection(); - // Create close-suppressing Connection proxy, also preparing returned Statements. + Mono connectionMono = getConnection() + .map(it -> new ConnectionCloseHolder(it, this::closeConnection)); return Flux.usingWhen(connectionMono, it -> { - Connection connectionToUse = createConnectionProxy(it); + // Create close-suppressing Connection proxy, also preparing returned Statements. + Connection connectionToUse = createConnectionProxy(it.connection); return doInConnectionMany(connectionToUse, action); - }, this::closeConnection, this::closeConnection, this::closeConnection) // + }, ConnectionCloseHolder::close, ConnectionCloseHolder::close, ConnectionCloseHolder::close) // .onErrorMap(R2dbcException.class, ex -> translateException("executeMany", getSql(action), ex)); } @@ -1104,4 +1107,26 @@ class DefaultDatabaseClient implements DatabaseClient, ConnectionAccessor { } } } + + /** + * Holder for a connection that makes sure the close action is invoked atomically only once. + */ + @RequiredArgsConstructor + static class ConnectionCloseHolder extends AtomicBoolean { + + final Connection connection; + final Function> closeFunction; + + Mono close() { + + return Mono.defer(() -> { + + if (compareAndSet(false, true)) { + return Mono.from(closeFunction.apply(connection)); + } + + return Mono.empty(); + }); + } + } } diff --git a/src/test/java/org/springframework/data/r2dbc/function/DefaultDatabaseClientUnitTests.java b/src/test/java/org/springframework/data/r2dbc/function/DefaultDatabaseClientUnitTests.java new file mode 100644 index 0000000..ac1d7ea --- /dev/null +++ b/src/test/java/org/springframework/data/r2dbc/function/DefaultDatabaseClientUnitTests.java @@ -0,0 +1,87 @@ +/* + * Copyright 2019 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 + * + * http://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 org.springframework.data.r2dbc.function; + +import static org.mockito.Mockito.*; + +import io.r2dbc.spi.Connection; +import io.r2dbc.spi.ConnectionFactory; +import reactor.core.CoreSubscriber; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.Mock; +import org.mockito.junit.MockitoJUnitRunner; +import org.reactivestreams.Publisher; +import org.reactivestreams.Subscription; +import org.springframework.data.r2dbc.support.R2dbcExceptionTranslator; + +/** + * Unit tests for {@link DefaultDatabaseClient}. + * + * @author Mark Paluch + */ +@RunWith(MockitoJUnitRunner.class) +public class DefaultDatabaseClientUnitTests { + + @Mock ConnectionFactory connectionFactory; + @Mock Connection connection; + @Mock ReactiveDataAccessStrategy strategy; + @Mock R2dbcExceptionTranslator translator; + + @Before + public void before() { + when(connectionFactory.create()).thenReturn((Publisher) Mono.just(connection)); + when(connection.close()).thenReturn(Mono.empty()); + } + + @Test // gh-48 + public void shouldCloseConnectionOnlyOnce() { + + DefaultDatabaseClient databaseClient = (DefaultDatabaseClient) DatabaseClient.builder() + .connectionFactory(connectionFactory).dataAccessStrategy(strategy).exceptionTranslator(translator).build(); + + Flux flux = databaseClient.inConnectionMany(it -> { + return Flux.empty(); + }); + + flux.subscribe(new CoreSubscriber() { + Subscription subscription; + + @Override + public void onSubscribe(Subscription s) { + s.request(1); + subscription = s; + } + + @Override + public void onNext(Object o) {} + + @Override + public void onError(Throwable t) {} + + @Override + public void onComplete() { + subscription.cancel(); + } + }); + + verify(connection, times(1)).close(); + } +}