Rollback reactive transaction on cancel

This commit introduces a change in reactive transaction semantics for
cancel signals. Canceling a subscription now rolls back a reactive transaction
to achieve a deterministic transaction outcome.

Previously, cancel signals committed a transaction which could
cause partially committed transactions depending on when the cancel happened.

spring-tx/src/main/java/org/springframework/transaction/interceptor/TransactionAspectSupport.java

@@ -878,7 +878,7 @@ public abstract class TransactionAspectSupport implements BeanFactoryAware, Init
 										},
 										this::commitTransactionAfterReturning,
 										(txInfo, err) -> Mono.empty(),
-										this::commitTransactionAfterReturning)
+										this::rollbackTransactionOnCancel)
 										.onErrorResume(ex ->
 												completeTransactionAfterThrowing(it, ex).then(Mono.error(ex)));
 							}
@@ -908,7 +908,7 @@ public abstract class TransactionAspectSupport implements BeanFactoryAware, Init
 											},
 											this::commitTransactionAfterReturning,
 											(txInfo, ex) -> Mono.empty(),
-											this::commitTransactionAfterReturning)
+											this::rollbackTransactionOnCancel)
 									.onErrorResume(ex ->
 											completeTransactionAfterThrowing(it, ex).then(Mono.error(ex)));
 						}
@@ -975,6 +975,16 @@ public abstract class TransactionAspectSupport implements BeanFactoryAware, Init
 			return Mono.empty();
 		}
 
+		private Mono<Void> rollbackTransactionOnCancel(@Nullable ReactiveTransactionInfo txInfo) {
+			if (txInfo != null && txInfo.getReactiveTransaction() != null) {
+				if (logger.isTraceEnabled()) {
+					logger.trace("Rolling back transaction for [" + txInfo.getJoinpointIdentification() + "] after cancellation");
+				}
+				return txInfo.getTransactionManager().rollback(txInfo.getReactiveTransaction());
+			}
+			return Mono.empty();
+		}
+
 		private Mono<Void> completeTransactionAfterThrowing(@Nullable ReactiveTransactionInfo txInfo, Throwable ex) {
 			if (txInfo != null && txInfo.getReactiveTransaction() != null) {
 				if (logger.isTraceEnabled()) {

spring-tx/src/main/java/org/springframework/transaction/reactive/TransactionalOperatorImpl.java

@@ -79,7 +79,7 @@ final class TransactionalOperatorImpl implements TransactionalOperator {
 			// Need re-wrapping of ReactiveTransaction until we get hold of the exception
 			// through usingWhen.
 			return status.flatMap(it -> Mono.usingWhen(Mono.just(it), ignore -> mono,
-					this.transactionManager::commit, (res, err) -> Mono.empty(), this.transactionManager::commit)
+					this.transactionManager::commit, (res, err) -> Mono.empty(), this.transactionManager::rollback)
 					.onErrorResume(ex -> rollbackOnException(it, ex).then(Mono.error(ex))));
 		})
 		.subscriberContext(TransactionContextManager.getOrCreateContext())
@@ -100,7 +100,7 @@ final class TransactionalOperatorImpl implements TransactionalOperator {
 							action::doInTransaction,
 							this.transactionManager::commit,
 							(tx, ex) -> Mono.empty(),
-							this.transactionManager::commit)
+							this.transactionManager::rollback)
 					.onErrorResume(ex ->
 							rollbackOnException(it, ex).then(Mono.error(ex))));
 		})

spring-tx/src/main/kotlin/org/springframework/transaction/reactive/TransactionalOperatorExtensions.kt

@@ -1,9 +1,10 @@
 package org.springframework.transaction.reactive
 
+import java.util.Optional
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.reactive.asFlow
-import kotlinx.coroutines.reactive.awaitFirstOrNull
+import kotlinx.coroutines.reactive.awaitLast
 import kotlinx.coroutines.reactor.asFlux
 import kotlinx.coroutines.reactor.mono
 import org.springframework.transaction.ReactiveTransaction
@@ -22,7 +23,9 @@ fun <T : Any> Flow<T>.transactional(operator: TransactionalOperator): Flow<T> =
 * parameter.
 *
 * @author Sebastien Deleuze
+* @author Mark Paluch
 * @since 5.2
 */
 suspend fun <T : Any> TransactionalOperator.executeAndAwait(f: suspend (ReactiveTransaction) -> T?): T? =
-		execute { status -> mono(Dispatchers.Unconfined) { f(status) } }.awaitFirstOrNull()
+		execute { status -> mono(Dispatchers.Unconfined) { f(status) } }.map { value -> Optional.of(value) }
+				.defaultIfEmpty(Optional.empty()).awaitLast().orElse(null)

spring-tx/src/test/java/org/springframework/transaction/annotation/AnnotationTransactionInterceptorTests.java

@@ -16,13 +16,19 @@
 
 package org.springframework.transaction.annotation;
 
+import java.time.Duration;
+
 import io.vavr.control.Try;
 import org.junit.jupiter.api.Test;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
 
 import org.springframework.aop.framework.ProxyFactory;
 import org.springframework.transaction.interceptor.TransactionInterceptor;
 import org.springframework.transaction.support.TransactionSynchronizationManager;
 import org.springframework.transaction.testfixture.CallCountingTransactionManager;
+import org.springframework.transaction.testfixture.ReactiveCallCountingTransactionManager;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
@@ -32,11 +38,14 @@ import static org.assertj.core.api.Assertions.assertThatIllegalStateException;
 /**
  * @author Rob Harrop
  * @author Juergen Hoeller
+ * @author Mark Paluch
  */
 public class AnnotationTransactionInterceptorTests {
 
 	private final CallCountingTransactionManager ptm = new CallCountingTransactionManager();
 
+	private final ReactiveCallCountingTransactionManager rtm = new ReactiveCallCountingTransactionManager();
+
 	private final AnnotationTransactionAttributeSource source = new AnnotationTransactionAttributeSource();
 
 	private final TransactionInterceptor ti = new TransactionInterceptor(this.ptm, this.source);
@@ -169,6 +178,78 @@ public class AnnotationTransactionInterceptorTests {
 			.satisfies(ex -> assertGetTransactionAndRollbackCount(1));
 	}
 
+	@Test
+	public void withMonoSuccess() {
+		ProxyFactory proxyFactory = new ProxyFactory();
+		proxyFactory.setTarget(new TestWithReactive());
+		proxyFactory.addAdvice(new TransactionInterceptor(rtm, this.source));
+
+		TestWithReactive proxy = (TestWithReactive) proxyFactory.getProxy();
+
+		StepVerifier.withVirtualTime(proxy::monoSuccess).thenAwait(Duration.ofSeconds(10)).verifyComplete();
+		assertReactiveGetTransactionAndCommitCount(1);
+	}
+
+	@Test
+	public void withMonoFailure() {
+		ProxyFactory proxyFactory = new ProxyFactory();
+		proxyFactory.setTarget(new TestWithReactive());
+		proxyFactory.addAdvice(new TransactionInterceptor(rtm, this.source));
+
+		TestWithReactive proxy = (TestWithReactive) proxyFactory.getProxy();
+
+		proxy.monoFailure().as(StepVerifier::create).verifyError();
+		assertReactiveGetTransactionAndRollbackCount(1);
+	}
+
+	@Test
+	public void withMonoRollback() {
+		ProxyFactory proxyFactory = new ProxyFactory();
+		proxyFactory.setTarget(new TestWithReactive());
+		proxyFactory.addAdvice(new TransactionInterceptor(rtm, this.source));
+
+		TestWithReactive proxy = (TestWithReactive) proxyFactory.getProxy();
+
+		StepVerifier.withVirtualTime(proxy::monoSuccess).thenAwait(Duration.ofSeconds(1)).thenCancel().verify();
+		assertReactiveGetTransactionAndRollbackCount(1);
+	}
+
+	@Test
+	public void withFluxSuccess() {
+		ProxyFactory proxyFactory = new ProxyFactory();
+		proxyFactory.setTarget(new TestWithReactive());
+		proxyFactory.addAdvice(new TransactionInterceptor(rtm, this.source));
+
+		TestWithReactive proxy = (TestWithReactive) proxyFactory.getProxy();
+
+		StepVerifier.withVirtualTime(proxy::fluxSuccess).thenAwait(Duration.ofSeconds(10)).expectNextCount(1).verifyComplete();
+		assertReactiveGetTransactionAndCommitCount(1);
+	}
+
+	@Test
+	public void withFluxFailure() {
+		ProxyFactory proxyFactory = new ProxyFactory();
+		proxyFactory.setTarget(new TestWithReactive());
+		proxyFactory.addAdvice(new TransactionInterceptor(rtm, this.source));
+
+		TestWithReactive proxy = (TestWithReactive) proxyFactory.getProxy();
+
+		proxy.fluxFailure().as(StepVerifier::create).verifyError();
+		assertReactiveGetTransactionAndRollbackCount(1);
+	}
+
+	@Test
+	public void withFluxRollback() {
+		ProxyFactory proxyFactory = new ProxyFactory();
+		proxyFactory.setTarget(new TestWithReactive());
+		proxyFactory.addAdvice(new TransactionInterceptor(rtm, this.source));
+
+		TestWithReactive proxy = (TestWithReactive) proxyFactory.getProxy();
+
+		StepVerifier.withVirtualTime(proxy::fluxSuccess).thenAwait(Duration.ofSeconds(1)).thenCancel().verify();
+		assertReactiveGetTransactionAndRollbackCount(1);
+	}
+
 	@Test
 	public void withVavrTrySuccess() {
 		ProxyFactory proxyFactory = new ProxyFactory();
@@ -342,6 +423,16 @@ public class AnnotationTransactionInterceptorTests {
 		assertThat(this.ptm.rollbacks).isEqualTo(expectedCount);
 	}
 
+	private void assertReactiveGetTransactionAndCommitCount(int expectedCount) {
+		assertThat(this.rtm.begun).isEqualTo(expectedCount);
+		assertThat(this.rtm.commits).isEqualTo(expectedCount);
+	}
+
+	private void assertReactiveGetTransactionAndRollbackCount(int expectedCount) {
+		assertThat(this.rtm.begun).isEqualTo(expectedCount);
+		assertThat(this.rtm.rollbacks).isEqualTo(expectedCount);
+	}
+
 
 	@Transactional
 	public static class TestClassLevelOnly {
@@ -452,6 +543,25 @@ public class AnnotationTransactionInterceptorTests {
 		}
 	}
 
+	@Transactional
+	public static class TestWithReactive {
+
+		public Mono<Void> monoSuccess() {
+			return Mono.delay(Duration.ofSeconds(10)).then();
+		}
+
+		public Mono<Void> monoFailure() {
+			return Mono.error(new IllegalStateException());
+		}
+
+		public Flux<Object> fluxSuccess() {
+			return Flux.just(new Object()).delayElements(Duration.ofSeconds(10));
+		}
+
+		public Flux<Object> fluxFailure() {
+			return Flux.error(new IllegalStateException());
+		}
+	}
 
 	@Transactional
 	public static class TestWithVavrTry {

spring-tx/src/test/java/org/springframework/transaction/reactive/TransactionalOperatorTests.java

@@ -70,8 +70,8 @@ public class TransactionalOperatorTests {
 				.thenAwait()
 				.thenCancel()
 				.verify();
-		assertThat(tm.commit).isTrue();
-		assertThat(tm.rollback).isFalse();
+		assertThat(tm.commit).isFalse();
+		assertThat(tm.rollback).isTrue();
 		assertThat(cancelled).isTrue();
 	}
 
@@ -84,8 +84,8 @@ public class TransactionalOperatorTests {
 				.thenAwait()
 				.thenCancel()
 				.verify();
-		assertThat(tm.commit).isTrue();
-		assertThat(tm.rollback).isFalse();
+		assertThat(tm.commit).isFalse();
+		assertThat(tm.rollback).isTrue();
 		assertThat(cancelled).isTrue();
 	}
 

spring-tx/src/test/kotlin/org/springframework/transaction/reactive/TransactionalOperatorExtensionsTests.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2019 the original author or authors.
+ * Copyright 2002-2020 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.
@@ -42,6 +42,19 @@ class TransactionalOperatorExtensionsTests {
 		assertThat(tm.rollback).isFalse()
 	}
 
+	@Test
+	fun commitWithEmptySuspendingFunction() {
+		val operator = TransactionalOperator.create(tm, DefaultTransactionDefinition())
+		runBlocking {
+			operator.executeAndAwait {
+				delay(1)
+				null
+			}
+		}
+		assertThat(tm.commit).isTrue()
+		assertThat(tm.rollback).isFalse()
+	}
+
 	@Test
 	fun rollbackWithSuspendingFunction() {
 		val operator = TransactionalOperator.create(tm, DefaultTransactionDefinition())

spring-tx/src/testFixtures/java/org/springframework/transaction/testfixture/ReactiveCallCountingTransactionManager.java

@@ -0,0 +1,71 @@
+/*
+ * Copyright 2002-2020 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 org.springframework.transaction.testfixture;
+
+import reactor.core.publisher.Mono;
+
+import org.springframework.transaction.TransactionDefinition;
+import org.springframework.transaction.TransactionException;
+import org.springframework.transaction.reactive.AbstractReactiveTransactionManager;
+import org.springframework.transaction.reactive.GenericReactiveTransaction;
+import org.springframework.transaction.reactive.TransactionSynchronizationManager;
+
+/**
+ * @author Mark Paluch
+ */
+@SuppressWarnings("serial")
+public class ReactiveCallCountingTransactionManager extends AbstractReactiveTransactionManager {
+
+	public TransactionDefinition lastDefinition;
+	public int begun;
+	public int commits;
+	public int rollbacks;
+	public int inflight;
+
+	@Override
+	protected Object doGetTransaction(TransactionSynchronizationManager synchronizationManager) throws TransactionException {
+		return new Object();
+	}
+
+	@Override
+	protected Mono<Void> doBegin(TransactionSynchronizationManager synchronizationManager, Object transaction, TransactionDefinition definition) throws TransactionException {
+		this.lastDefinition = definition;
+		++begun;
+		++inflight;
+		return Mono.empty();
+	}
+
+	@Override
+	protected Mono<Void> doCommit(TransactionSynchronizationManager synchronizationManager, GenericReactiveTransaction status) throws TransactionException {
+		++commits;
+		--inflight;
+		return Mono.empty();
+	}
+
+	@Override
+	protected Mono<Void> doRollback(TransactionSynchronizationManager synchronizationManager, GenericReactiveTransaction status) throws TransactionException {
+		++rollbacks;
+		--inflight;
+		return Mono.empty();
+	}
+
+
+	public void clear() {
+		begun = commits = rollbacks = inflight = 0;
+	}
+
+}