Introduce ReactiveRedisOperations.executeInSession(…) and session callback interface.

Allows reuse of a bound connection without additional connection acquisition overhead.

Closes: #2110
Original Pull Request: #2129

src/main/java/org/springframework/data/redis/core/ReactiveRedisOperations.java

@@ -63,6 +63,23 @@ public interface ReactiveRedisOperations<K, V> {
 	 */
 	<T> Flux<T> execute(ReactiveRedisCallback<T> action);
 
+	/**
+	 * Executes the given action within a Redis session using the same
+	 * {@link org.springframework.data.redis.connection.ReactiveRedisConnection}. Application exceptions thrown by the
+	 * action object get propagated to the caller (can only be unchecked) whenever possible. Redis exceptions are
+	 * transformed into appropriate DAO ones. Allows for returning a result object, that is a domain object or a
+	 * collection of domain objects. Performs automatic serialization/deserialization for the given objects to and from
+	 * binary data suitable for the Redis storage. Note: Callback code is not supposed to handle transactions itself! Use
+	 * an appropriate transaction manager. Generally, callback code must not touch any Connection lifecycle methods, like
+	 * close, to let the template do its work.
+	 *
+	 * @param <T> return type
+	 * @param action callback object that specifies the Redis action
+	 * @return a result object returned by the action or {@link Flux#empty()}.
+	 * @since 2.6
+	 */
+	<T> Flux<T> executeInSession(ReactiveRedisSessionCallback<K, V, T> action);
+
 	// -------------------------------------------------------------------------
 	// Methods dealing with Redis Pub/Sub
 	// -------------------------------------------------------------------------

src/main/java/org/springframework/data/redis/core/ReactiveRedisSessionCallback.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright 2017-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 org.springframework.data.redis.core;
+
+import org.reactivestreams.Publisher;
+import org.springframework.dao.DataAccessException;
+import org.springframework.data.redis.connection.ReactiveRedisConnection;
+
+/**
+ * Generic callback interface for code that wants to use the same {@link ReactiveRedisConnection} avoiding connection
+ * allocation overhead upon each Template API method call. Allows to execute any number of operations on a single
+ * {@link ReactiveRedisConnection}, using any type and number of commands.
+ * <p>
+ * This is particularly useful for issuing multiple calls on the same connection.
+ *
+ * @param <T>
+ * @author Mark Paluch
+ * @since 2.6
+ * @see ReactiveRedisOperations#executeInSession(ReactiveRedisSessionCallback)
+ */
+public interface ReactiveRedisSessionCallback<K, V, T> {
+
+	/**
+	 * Gets called by {@link ReactiveRedisOperations#executeInSession(ReactiveRedisSessionCallback)} with an active Redis
+	 * connection. Does not need to care about activating or closing the {@link ReactiveRedisConnection}.
+	 * <p>
+	 * Allows for returning a result object created within the callback, i.e. a domain object or a collection of domain
+	 * objects.
+	 *
+	 * @param operations template associated with a connection.
+	 * @return a result object publisher
+	 * @throws DataAccessException in case of custom exceptions
+	 */
+	Publisher<T> doWithOperations(ReactiveRedisOperations<K, V> operations) throws DataAccessException;
+}

src/main/java/org/springframework/data/redis/core/ReactiveRedisTemplate.java

@@ -151,6 +151,14 @@ public class ReactiveRedisTemplate<K, V> implements ReactiveRedisOperations<K, V
 		return Flux.from(doInConnection(action, exposeConnection));
 	}
 
+	@Override
+	public <T> Flux<T> executeInSession(ReactiveRedisSessionCallback<K, V, T> action) {
+
+		Assert.notNull(action, "Callback object must not be null");
+		return Flux
+				.from(doInConnection(connection -> action.doWithOperations(withConnection(connection)), exposeConnection));
+	}
+
 	/**
 	 * Create a reusable Flux for a {@link ReactiveRedisCallback}. Callback is executed within a connection context. The
 	 * connection is released outside the callback.
@@ -188,7 +196,7 @@ public class ReactiveRedisTemplate<K, V> implements ReactiveRedisOperations<K, V
 	 * @param exposeConnection whether to enforce exposure of the native Redis Connection to callback code
 	 * @return object returned by the action
 	 */
-	private <T> Publisher<T> doInConnection(ReactiveRedisCallback<T> action, boolean exposeConnection) {
+	<T> Publisher<T> doInConnection(ReactiveRedisCallback<T> action, boolean exposeConnection) {
 
 		Assert.notNull(action, "Callback object must not be null");
 
@@ -742,6 +750,31 @@ public class ReactiveRedisTemplate<K, V> implements ReactiveRedisOperations<K, V
 		return serializationContext;
 	}
 
+	private ReactiveRedisOperations<K, V> withConnection(ReactiveRedisConnection connection) {
+		return new BoundConnectionRedisTemplate(connection, connectionFactory, serializationContext);
+	}
+
+	class BoundConnectionRedisTemplate extends ReactiveRedisTemplate<K, V> {
+
+		private final ReactiveRedisConnection connection;
+
+		public BoundConnectionRedisTemplate(ReactiveRedisConnection connection,
+				ReactiveRedisConnectionFactory connectionFactory, RedisSerializationContext<K, V> serializationContext) {
+			super(connectionFactory, serializationContext, true);
+			this.connection = connection;
+		}
+
+		@Override
+		<T> Publisher<T> doInConnection(ReactiveRedisCallback<T> action, boolean exposeConnection) {
+
+			Assert.notNull(action, "Callback object must not be null");
+
+			ReactiveRedisConnection connToUse = ReactiveRedisTemplate.this.preProcessConnection(connection, true);
+			Publisher<T> result = action.doInRedis(connToUse);
+			return ReactiveRedisTemplate.this.postProcessResult(result, connToUse, true);
+		}
+	}
+
 	private ByteBuffer rawKey(K key) {
 		return getSerializationContext().getKeySerializationPair().getWriter().write(key);
 	}

src/main/kotlin/org/springframework/data/redis/core/ReactiveRedisOperationsExtensions.kt

@@ -22,7 +22,7 @@ import kotlinx.coroutines.reactive.awaitFirstOrNull
 import kotlinx.coroutines.reactive.awaitSingle
 import org.springframework.data.redis.connection.DataType
 import org.springframework.data.redis.connection.ReactiveRedisConnection
-import org.springframework.data.redis.connection.ReactiveSubscription.*
+import org.springframework.data.redis.connection.ReactiveSubscription.Message
 import org.springframework.data.redis.core.script.RedisScript
 import org.springframework.data.redis.listener.Topic
 import org.springframework.data.redis.serializer.RedisElementReader
@@ -36,8 +36,20 @@ import java.time.Instant
  * @author Sebastien Deleuze
  * @since 2.2
  */
-fun <K : Any, V : Any, T : Any> ReactiveRedisOperations<K, V>.executeAsFlow(action: (ReactiveRedisConnection) -> Flow<T>): Flow<T> =
-		execute { action(it).asPublisher() }.asFlow()
+fun <K : Any, V : Any, T : Any> ReactiveRedisOperations<K, V>.executeAsFlow(action: (ReactiveRedisConnection) -> Flow<T>): Flow<T> {
+	return execute { action(it).asPublisher() }.asFlow()
+}
+
+/**
+ * Coroutines variant of [ReactiveRedisOperations.execute].
+ *
+ * @author Mark Paluch
+ * @since 2.6
+ */
+fun <K : Any, V : Any, T : Any> ReactiveRedisOperations<K, V>.executeInSessionAsFlow(
+	action: (ReactiveRedisOperations<K, V>) -> Flow<T>
+): Flow<T> =
+	executeInSession { action(it).asPublisher() }.asFlow()
 
 /**
  * Coroutines variant of [ReactiveRedisOperations.execute].
@@ -45,8 +57,12 @@ fun <K : Any, V : Any, T : Any> ReactiveRedisOperations<K, V>.executeAsFlow(acti
  * @author Sebastien Deleuze
  * @since 2.2
  */
-fun <K : Any, V : Any, T : Any> ReactiveRedisOperations<K, V>.executeAsFlow(script: RedisScript<T>, keys: List<K> = emptyList(), args: List<*> = emptyList<Any>()): Flow<T> =
-		execute(script, keys, args).asFlow()
+fun <K : Any, V : Any, T : Any> ReactiveRedisOperations<K, V>.executeAsFlow(
+	script: RedisScript<T>,
+	keys: List<K> = emptyList(),
+	args: List<*> = emptyList<Any>()
+): Flow<T> =
+	execute(script, keys, args).asFlow()
 
 /**
  * Coroutines variant of [ReactiveRedisOperations.execute].

src/test/kotlin/org/springframework/data/redis/core/ReactiveRedisOperationsExtensionsUnitTests.kt

@@ -50,7 +50,8 @@ class ReactiveRedisOperationsExtensionsUnitTests {
 		every { operations.execute(any<ReactiveRedisCallback<*>>()) } returns Flux.just("foo")
 
 		runBlocking {
-			assertThat(operations.executeAsFlow { flow { emit("foo")} }.toList()).contains("foo")
+			assertThat(operations.executeAsFlow { flow { emit("foo") } }
+				.toList()).contains("foo")
 		}
 
 		verify {
@@ -58,12 +59,36 @@ class ReactiveRedisOperationsExtensionsUnitTests {
 		}
 	}
 
+	@Test // GH-2110
+	fun `executeInSession with calllback`() {
+
+		val operations = mockk<ReactiveRedisOperations<String, String>>()
+		every { operations.executeInSession(any<ReactiveRedisSessionCallback<String, String, *>>()) } returns Flux.just(
+			"foo"
+		)
+
+		runBlocking {
+			assertThat(operations.executeInSessionAsFlow { flow { emit("foo") } }
+				.toList()).contains("foo")
+		}
+
+		verify {
+			operations.executeInSession(any<ReactiveRedisSessionCallback<String, String, *>>())
+		}
+	}
+
 	@Test // DATAREDIS-1033
 	fun `execute with script`() {
 
 		val script = RedisScript.of<String>("foo")
 		val operations = mockk<ReactiveRedisOperations<String, String>>()
-		every { operations.execute(any<RedisScript<*>>(), any(), any()) } returns Flux.just("foo")
+		every {
+			operations.execute(
+				any<RedisScript<*>>(),
+				any(),
+				any()
+			)
+		} returns Flux.just("foo")
 
 		runBlocking {
 			assertThat(operations.executeAsFlow(script).toList()).contains("foo")