From c1d840d87dbac6a43bc816dc44f372374ffe8131 Mon Sep 17 00:00:00 2001
From: Christoph Strobl <cstrobl@pivotal.io>
Date: Thu, 14 Jun 2018 14:55:51 +0200
Subject: [PATCH] DATAMONGO-2005 - Use Flux.usingWhen for resource management
 in reactive transactions.

Original pull request: #574.
---
 .../mongodb/core/ReactiveMongoTemplate.java    | 18 ++----------------
 1 file changed, 2 insertions(+), 16 deletions(-)

diff --git a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/ReactiveMongoTemplate.java b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/ReactiveMongoTemplate.java
index 8919dd796..c306e6052 100644
--- a/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/ReactiveMongoTemplate.java
+++ b/spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/ReactiveMongoTemplate.java
@@ -519,22 +519,8 @@ public class ReactiveMongoTemplate implements ReactiveMongoOperations, Applicati
 						session.startTransaction();
 					}
 
-					return ReactiveMongoTemplate.this.withSession(action, session) //
-							.materialize() //
-							.flatMap(signal -> {
-
-								if (session.hasActiveTransaction()) {
-									if (signal.isOnComplete()) {
-										return Mono.from(session.commitTransaction()).thenReturn(signal);
-									}
-									if (signal.isOnError()) {
-										return Mono.from(session.abortTransaction()).thenReturn(signal);
-									}
-								}
-								return Mono.just(signal);
-							}) //
-							.<T> dematerialize() //
-							.doFinally(signalType -> {
+					return Flux.usingWhen(Mono.just(session), s -> ReactiveMongoTemplate.this.withSession(action, s),
+							ClientSession::commitTransaction, ClientSession::abortTransaction).doFinally(signalType -> {
 								doFinally.accept(session);
 							});
 				});