Track subscriptions and unsubscriptions in LettuceReactiveRedisConnection.

We now track subscriptions and unsubscriptions in the reactive API to ensure that we do not prematurely unsubscribe from a channel or pattern if the topic was subscribed multiple times.

Original Pull Request: #2467

src/test/java/org/springframework/data/redis/connection/lettuce/LettuceReactivePubSubCommandsUnitTests.java

@@ -0,0 +1,201 @@
+/*
+ * Copyright 2022 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.connection.lettuce;
+
+import static org.assertj.core.api.Assertions.*;
+import static org.mockito.Mockito.*;
+
+import io.lettuce.core.pubsub.StatefulRedisPubSubConnection;
+import io.lettuce.core.pubsub.api.reactive.RedisPubSubReactiveCommands;
+import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
+
+import java.nio.ByteBuffer;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoSettings;
+import org.mockito.quality.Strictness;
+
+/**
+ * Unit tests for {@link LettuceReactivePubSubCommands}.
+ *
+ * @author Mark Paluch
+ */
+@MockitoSettings(strictness = Strictness.LENIENT)
+class LettuceReactivePubSubCommandsUnitTests {
+
+	LettuceReactivePubSubCommands sut;
+
+	@Mock LettuceReactiveRedisConnection connection;
+
+	@Mock StatefulRedisPubSubConnection<ByteBuffer, ByteBuffer> lettuceConnection;
+	@Mock RedisPubSubReactiveCommands<ByteBuffer, ByteBuffer> reactiveCommands;
+
+	@SuppressWarnings("unchecked")
+	@BeforeEach
+	void setUp() {
+
+		when(connection.getPubSubConnection()).thenReturn(Mono.just(lettuceConnection));
+		when(lettuceConnection.reactive()).thenReturn(reactiveCommands);
+		sut = new LettuceReactivePubSubCommands(connection);
+	}
+
+	@Test // GH-2386
+	void shouldSubscribeChannelMultipleTimes() {
+
+		when(reactiveCommands.subscribe(any())).thenReturn(Mono.empty());
+
+		sut.subscribe(wrap("channel")) //
+				.as(StepVerifier::create) //
+				.verifyComplete();
+
+		sut.subscribe(wrap("channel")) //
+				.as(StepVerifier::create) //
+				.verifyComplete();
+
+		verify(reactiveCommands, times(2)).subscribe(wrap("channel"));
+		assertThat(sut.getChannels()).hasSize(1);
+		assertThat(sut.getPatterns()).isEmpty();
+	}
+
+	@Test // GH-2386
+	void shouldNotUnsubscribeChannelIfUsedMultipleTimes() {
+
+		when(reactiveCommands.subscribe(any())).thenReturn(Mono.empty());
+		when(reactiveCommands.unsubscribe(any())).thenReturn(Mono.empty());
+
+		sut.subscribe(wrap("channel")) //
+				.as(StepVerifier::create) //
+				.verifyComplete();
+
+		sut.subscribe(wrap("channel")) //
+				.as(StepVerifier::create) //
+				.verifyComplete();
+
+		sut.unsubscribe(wrap("channel")) //
+				.as(StepVerifier::create) //
+				.verifyComplete();
+
+		verify(reactiveCommands, times(2)).subscribe(wrap("channel"));
+		verifyNoMoreInteractions(reactiveCommands);
+		assertThat(sut.getChannels()).hasSize(1);
+		assertThat(sut.getPatterns()).isEmpty();
+	}
+
+	@Test // GH-2386
+	void shouldUnsubscribeChannelIfNotUsedAnymore() {
+
+		when(reactiveCommands.subscribe(any())).thenReturn(Mono.empty());
+		when(reactiveCommands.unsubscribe(any())).thenReturn(Mono.empty());
+
+		sut.subscribe(wrap("channel")) //
+				.as(StepVerifier::create) //
+				.verifyComplete();
+
+		sut.subscribe(wrap("channel")) //
+				.as(StepVerifier::create) //
+				.verifyComplete();
+
+		sut.unsubscribe(wrap("channel")) //
+				.as(StepVerifier::create) //
+				.verifyComplete();
+
+		sut.unsubscribe(wrap("channel")) //
+				.as(StepVerifier::create) //
+				.verifyComplete();
+
+		verify(reactiveCommands, times(2)).subscribe(wrap("channel"));
+		verify(reactiveCommands, times(1)).unsubscribe(wrap("channel"));
+		assertThat(sut.getChannels()).isEmpty();
+		assertThat(sut.getPatterns()).isEmpty();
+	}
+
+	@Test // GH-2386
+	void shouldSubscribePatternMultipleTimes() {
+
+		when(reactiveCommands.psubscribe(any())).thenReturn(Mono.empty());
+
+		sut.pSubscribe(wrap("pattern")) //
+				.as(StepVerifier::create) //
+				.verifyComplete();
+
+		sut.pSubscribe(wrap("pattern")) //
+				.as(StepVerifier::create) //
+				.verifyComplete();
+
+		verify(reactiveCommands, times(2)).psubscribe(wrap("pattern"));
+		assertThat(sut.getChannels()).isEmpty();
+		assertThat(sut.getPatterns()).hasSize(1);
+	}
+
+	@Test // GH-2386
+	void shouldNotUnsubscribePatternIfUsedMultipleTimes() {
+
+		when(reactiveCommands.psubscribe(any())).thenReturn(Mono.empty());
+		when(reactiveCommands.punsubscribe(any())).thenReturn(Mono.empty());
+
+		sut.pSubscribe(wrap("pattern")) //
+				.as(StepVerifier::create) //
+				.verifyComplete();
+
+		sut.pSubscribe(wrap("pattern")) //
+				.as(StepVerifier::create) //
+				.verifyComplete();
+
+		sut.pUnsubscribe(wrap("pattern")) //
+				.as(StepVerifier::create) //
+				.verifyComplete();
+
+		verify(reactiveCommands, times(2)).psubscribe(wrap("pattern"));
+		verifyNoMoreInteractions(reactiveCommands);
+		assertThat(sut.getChannels()).isEmpty();
+		assertThat(sut.getPatterns()).hasSize(1);
+	}
+
+	@Test // GH-2386
+	void shouldUnsubscribePatternIfNotUsedAnymore() {
+
+		when(reactiveCommands.psubscribe(any())).thenReturn(Mono.empty());
+		when(reactiveCommands.punsubscribe(any())).thenReturn(Mono.empty());
+
+		sut.pSubscribe(wrap("pattern")) //
+				.as(StepVerifier::create) //
+				.verifyComplete();
+
+		sut.pSubscribe(wrap("pattern")) //
+				.as(StepVerifier::create) //
+				.verifyComplete();
+
+		sut.pUnsubscribe(wrap("pattern")) //
+				.as(StepVerifier::create) //
+				.verifyComplete();
+
+		sut.pUnsubscribe(wrap("pattern")) //
+				.as(StepVerifier::create) //
+				.verifyComplete();
+
+		verify(reactiveCommands, times(2)).psubscribe(wrap("pattern"));
+		verify(reactiveCommands, times(1)).punsubscribe(wrap("pattern"));
+		assertThat(sut.getChannels()).isEmpty();
+		assertThat(sut.getPatterns()).isEmpty();
+	}
+
+	private static ByteBuffer wrap(String content) {
+		return ByteBuffer.wrap(content.getBytes());
+	}
+}

src/test/java/org/springframework/data/redis/connection/lettuce/LettuceReactiveSubscriptionUnitTests.java

@@ -37,7 +37,6 @@ import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
-
 import org.springframework.data.redis.RedisSystemException;
 import org.springframework.data.redis.connection.ReactiveSubscription.Message;
 import org.springframework.data.redis.connection.ReactiveSubscription.PatternMessage;
@@ -57,17 +56,19 @@ class LettuceReactiveSubscriptionUnitTests {
 	@Mock StatefulRedisPubSubConnection<ByteBuffer, ByteBuffer> connectionMock;
 	@Mock RedisPubSubReactiveCommands<ByteBuffer, ByteBuffer> commandsMock;
 
+	@Mock LettuceReactivePubSubCommands pubSubMock;
+
 	@BeforeEach
 	void before() {
 		when(connectionMock.reactive()).thenReturn(commandsMock);
-		subscription = new LettuceReactiveSubscription(mock(SubscriptionListener.class), connectionMock,
+		subscription = new LettuceReactiveSubscription(mock(SubscriptionListener.class), connectionMock, pubSubMock,
 				e -> new RedisSystemException(e.getMessage(), e));
 	}
 
 	@Test // DATAREDIS-612
 	void shouldSubscribeChannels() {
 
-		when(commandsMock.subscribe(any())).thenReturn(Mono.empty());
+		when(pubSubMock.subscribe(any())).thenReturn(Mono.empty());
 
 		Mono<Void> subscribe = subscription.subscribe(getByteBuffer("foo"), getByteBuffer("bar"));
 
@@ -82,7 +83,7 @@ class LettuceReactiveSubscriptionUnitTests {
 	@Test // DATAREDIS-612
 	void shouldSubscribeChannelsShouldFail() {
 
-		when(commandsMock.subscribe(any())).thenReturn(Mono.error(new RedisConnectionException("Foo")));
+		when(pubSubMock.subscribe(any())).thenReturn(Mono.error(new RedisConnectionException("Foo")));
 
 		Mono<Void> subscribe = subscription.subscribe(getByteBuffer("foo"), getByteBuffer("bar"));
 
@@ -92,7 +93,7 @@ class LettuceReactiveSubscriptionUnitTests {
 	@Test // DATAREDIS-612
 	void shouldSubscribePatterns() {
 
-		when(commandsMock.psubscribe(any())).thenReturn(Mono.empty());
+		when(pubSubMock.pSubscribe(any())).thenReturn(Mono.empty());
 
 		Mono<Void> subscribe = subscription.pSubscribe(getByteBuffer("foo"), getByteBuffer("bar"));
 
@@ -107,33 +108,33 @@ class LettuceReactiveSubscriptionUnitTests {
 	@Test // DATAREDIS-612
 	void shouldUnsubscribeChannels() {
 
-		when(commandsMock.subscribe(any())).thenReturn(Mono.empty());
-		when(commandsMock.unsubscribe(any())).thenReturn(Mono.empty());
+		when(pubSubMock.subscribe(any())).thenReturn(Mono.empty());
+		when(pubSubMock.unsubscribe(any())).thenReturn(Mono.empty());
 		subscription.subscribe(getByteBuffer("foo"), getByteBuffer("bar")).as(StepVerifier::create).verifyComplete();
 
 		subscription.unsubscribe().as(StepVerifier::create).verifyComplete();
 
 		assertThat(subscription.getChannels()).isEmpty();
-		verify(commandsMock).unsubscribe(any());
+		verify(pubSubMock).unsubscribe(any());
 	}
 
 	@Test // DATAREDIS-612
 	void shouldUnsubscribePatterns() {
 
-		when(commandsMock.psubscribe(any())).thenReturn(Mono.empty());
-		when(commandsMock.punsubscribe(any())).thenReturn(Mono.empty());
+		when(pubSubMock.pSubscribe(any())).thenReturn(Mono.empty());
+		when(pubSubMock.pUnsubscribe(any())).thenReturn(Mono.empty());
 		subscription.pSubscribe(getByteBuffer("foo"), getByteBuffer("bar")).as(StepVerifier::create).verifyComplete();
 
 		subscription.pUnsubscribe().as(StepVerifier::create).verifyComplete();
 
 		assertThat(subscription.getPatterns()).isEmpty();
-		verify(commandsMock).punsubscribe(any());
+		verify(pubSubMock).pUnsubscribe(any());
 	}
 
 	@Test // DATAREDIS-612
 	void shouldEmitChannelMessage() {
 
-		when(commandsMock.subscribe(any())).thenReturn(Mono.empty());
+		when(pubSubMock.subscribe(any())).thenReturn(Mono.empty());
 		subscription.subscribe(getByteBuffer("foo"), getByteBuffer("bar")).as(StepVerifier::create).verifyComplete();
 
 		Sinks.Many<io.lettuce.core.pubsub.api.reactive.ChannelMessage<ByteBuffer, ByteBuffer>> sink = Sinks.many().unicast()
@@ -153,7 +154,7 @@ class LettuceReactiveSubscriptionUnitTests {
 	@Test // DATAREDIS-612
 	void shouldEmitPatternMessage() {
 
-		when(commandsMock.psubscribe(any())).thenReturn(Mono.empty());
+		when(pubSubMock.pSubscribe(any())).thenReturn(Mono.empty());
 		subscription.pSubscribe(getByteBuffer("foo*"), getByteBuffer("bar*")).as(StepVerifier::create).verifyComplete();
 
 		Sinks.Many<io.lettuce.core.pubsub.api.reactive.PatternMessage<ByteBuffer, ByteBuffer>> sink = Sinks.many().unicast()
@@ -175,7 +176,7 @@ class LettuceReactiveSubscriptionUnitTests {
 	@Test // DATAREDIS-612
 	void shouldEmitError() {
 
-		when(commandsMock.subscribe(any())).thenReturn(Mono.empty());
+		when(pubSubMock.subscribe(any())).thenReturn(Mono.empty());
 		subscription.subscribe(getByteBuffer("foo"), getByteBuffer("bar")).as(StepVerifier::create).verifyComplete();
 
 		Sinks.Many<io.lettuce.core.pubsub.api.reactive.ChannelMessage<ByteBuffer, ByteBuffer>> sink = Sinks.many().unicast()
@@ -192,8 +193,8 @@ class LettuceReactiveSubscriptionUnitTests {
 	@Test // DATAREDIS-612
 	void shouldTerminateActiveSubscriptions() {
 
-		when(commandsMock.psubscribe(any())).thenReturn(Mono.empty());
-		when(commandsMock.punsubscribe(any())).thenReturn(Mono.empty());
+		when(pubSubMock.pSubscribe(any())).thenReturn(Mono.empty());
+		when(pubSubMock.pUnsubscribe(any())).thenReturn(Mono.empty());
 		subscription.pSubscribe(getByteBuffer("foo*")).as(StepVerifier::create).verifyComplete();
 
 		when(commandsMock.observeChannels()).thenReturn(Flux.never());
@@ -212,7 +213,7 @@ class LettuceReactiveSubscriptionUnitTests {
 		Sinks.Many<io.lettuce.core.pubsub.api.reactive.PatternMessage<ByteBuffer, ByteBuffer>> sink = Sinks.many().unicast()
 				.onBackpressureBuffer();
 
-		when(commandsMock.psubscribe(any())).thenReturn(Mono.empty());
+		when(pubSubMock.pSubscribe(any())).thenReturn(Mono.empty());
 		subscription.pSubscribe(getByteBuffer("foo*")).as(StepVerifier::create).verifyComplete();
 
 		when(commandsMock.observeChannels()).thenReturn(Flux.never());

src/test/java/org/springframework/data/redis/listener/ReactiveRedisMessageListenerContainerIntegrationTests.java

@@ -18,12 +18,17 @@ package org.springframework.data.redis.listener;
 import static org.assertj.core.api.Assertions.*;
 
 import reactor.core.Disposable;
+import reactor.core.publisher.Flux;
 import reactor.test.StepVerifier;
 
 import java.nio.ByteBuffer;
 import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.List;
+import java.util.Queue;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.LinkedBlockingDeque;
@@ -47,6 +52,8 @@ import org.springframework.data.redis.connection.SubscriptionListener;
 import org.springframework.data.redis.connection.lettuce.LettuceConnectionFactory;
 import org.springframework.data.redis.core.ReactiveRedisTemplate;
 import org.springframework.data.redis.serializer.RedisSerializationContext;
+import org.springframework.data.redis.serializer.RedisSerializationContext.SerializationPair;
+import org.springframework.data.redis.serializer.RedisSerializer;
 import org.springframework.data.redis.test.extension.parametrized.MethodSource;
 import org.springframework.data.redis.test.extension.parametrized.ParameterizedRedisTest;
 import org.springframework.lang.Nullable;
@@ -302,6 +309,49 @@ public class ReactiveRedisMessageListenerContainerIntegrationTests {
 				.verify();
 	}
 
+	@ParameterizedRedisTest // GH-2386
+	void multipleListenShouldTrackSubscriptions() throws Exception {
+
+		ReactiveRedisMessageListenerContainer container = new ReactiveRedisMessageListenerContainer(connectionFactory);
+
+		Flux<? extends ReactiveSubscription.Message<String, String>> c1 = container.receiveLater(ChannelTopic.of(CHANNEL1))
+				.block();
+		Flux<? extends ReactiveSubscription.Message<String, String>> c1p1 = container
+				.receiveLater(Arrays.asList(ChannelTopic.of(CHANNEL1), PatternTopic.of(PATTERN1)),
+						SerializationPair.fromSerializer(RedisSerializer.string()),
+						SerializationPair.fromSerializer(RedisSerializer.string()))
+				.block();
+
+		BlockingQueue<ReactiveSubscription.Message<String, String>> c1Collector = new LinkedBlockingDeque<>();
+		BlockingQueue<ReactiveSubscription.Message<String, String>> c2Collector = new LinkedBlockingDeque<>();
+
+		Disposable c1Subscription = c1.doOnNext(c1Collector::add).subscribe();
+		Disposable c2Subscription = c1p1.doOnNext(c2Collector::add).subscribe();
+
+		doPublish(CHANNEL1.getBytes(), MESSAGE.getBytes());
+
+		assertThat(c1Collector.poll(5, TimeUnit.SECONDS)).isNotNull();
+		assertThat(c2Collector.poll(5, TimeUnit.SECONDS)).isNotNull();
+		c1Collector.clear();
+		c2Collector.clear();
+
+		c2Subscription.dispose();
+
+		Thread.sleep(200);
+
+		doPublish(CHANNEL1.getBytes(), MESSAGE.getBytes());
+
+		assertThat(c1Collector.poll(5, TimeUnit.SECONDS)).isNotNull();
+		assertThat(c2Collector.poll(100, TimeUnit.MILLISECONDS)).isNull();
+
+		c1Subscription.dispose();
+
+		doPublish(CHANNEL1.getBytes(), MESSAGE.getBytes());
+
+		assertThat(c1Collector.poll(100, TimeUnit.MILLISECONDS)).isNull();
+		assertThat(c2Collector.poll(100, TimeUnit.MILLISECONDS)).isNull();
+	}
+
 	private void doPublish(byte[] channel, byte[] message) {
 		reactiveConnection.pubSubCommands().publish(ByteBuffer.wrap(channel), ByteBuffer.wrap(message)).subscribe();
 	}