Add support for subscription ready/unsubscription done events.

Original Pull Request: #2052

src/main/java/org/springframework/data/redis/connection/MessageListener.java

@@ -18,11 +18,14 @@ package org.springframework.data.redis.connection;
 import org.springframework.lang.Nullable;
 
 /**
- * Listener of messages published in Redis.
+ * Listener of messages published in Redis. A MessageListener can implement {@link SubscriptionListener} to receive
+ * notifications for subscription states.
  *
  * @author Costin Leau
  * @author Christoph Strobl
+ * @see SubscriptionListener
  */
+@FunctionalInterface
 public interface MessageListener {
 
 	/**

src/main/java/org/springframework/data/redis/connection/SubscriptionListener.java

@@ -0,0 +1,66 @@
+/*
+ * Copyright 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.connection;
+
+/**
+ * Listener for subscription notifications.
+ * <p/>
+ * Subscription notifications are reported by Redis as confirmation for subscribe and unsubscribe operations for
+ * channels and patterns.
+ *
+ * @author Mark Paluch
+ * @since 2.6
+ */
+public interface SubscriptionListener {
+
+	/**
+	 * Empty {@link SubscriptionListener}.
+	 */
+	SubscriptionListener EMPTY = new SubscriptionListener() {};
+
+	/**
+	 * Notification when Redis has confirmed a channel subscription.
+	 *
+	 * @param channel name of the channel.
+	 * @param count subscriber count.
+	 */
+	default void onChannelSubscribed(byte[] channel, long count) {}
+
+	/**
+	 * Notification when Redis has confirmed a channel un-subscription.
+	 *
+	 * @param channel name of the channel.
+	 * @param count subscriber count.
+	 */
+	default void onChannelUnsubscribed(byte[] channel, long count) {}
+
+	/**
+	 * Notification when Redis has confirmed a pattern subscription.
+	 *
+	 * @param pattern the pattern.
+	 * @param count subscriber count.
+	 */
+	default void onPatternSubscribed(byte[] pattern, long count) {}
+
+	/**
+	 * Notification when Redis has confirmed a pattern un-subscription.
+	 *
+	 * @param pattern the pattern.
+	 * @param count subscriber count.
+	 */
+	default void onPatternUnsubscribed(byte[] pattern, long count) {}
+
+}

src/main/java/org/springframework/data/redis/connection/jedis/JedisClusterConnection.java

@@ -450,7 +450,7 @@ public class JedisClusterConnection implements DefaultedRedisClusterConnection {
 					"Connection already subscribed; use the connection Subscription to cancel or add new channels");
 		}
 		try {
-			BinaryJedisPubSub jedisPubSub = new JedisMessageListener(listener);
+			JedisMessageListener jedisPubSub = new JedisMessageListener(listener);
 			subscription = new JedisSubscription(listener, jedisPubSub, channels, null);
 			cluster.subscribe(jedisPubSub, channels);
 		} catch (Exception ex) {
@@ -466,7 +466,7 @@ public class JedisClusterConnection implements DefaultedRedisClusterConnection {
 					"Connection already subscribed; use the connection Subscription to cancel or add new channels");
 		}
 		try {
-			BinaryJedisPubSub jedisPubSub = new JedisMessageListener(listener);
+			JedisMessageListener jedisPubSub = new JedisMessageListener(listener);
 			subscription = new JedisSubscription(listener, jedisPubSub, null, patterns);
 			cluster.psubscribe(jedisPubSub, patterns);
 		} catch (Exception ex) {

src/main/java/org/springframework/data/redis/connection/jedis/JedisConnection.java

@@ -715,7 +715,7 @@ public class JedisConnection extends AbstractRedisConnection {
 
 		doWithJedis(it -> {
 
-			BinaryJedisPubSub jedisPubSub = new JedisMessageListener(listener);
+			JedisMessageListener jedisPubSub = new JedisMessageListener(listener);
 
 			subscription = new JedisSubscription(listener, jedisPubSub, null, patterns);
 			it.psubscribe(jedisPubSub, patterns);
@@ -740,7 +740,7 @@ public class JedisConnection extends AbstractRedisConnection {
 
 		doWithJedis(it -> {
 
-			BinaryJedisPubSub jedisPubSub = new JedisMessageListener(listener);
+			JedisMessageListener jedisPubSub = new JedisMessageListener(listener);
 
 			subscription = new JedisSubscription(listener, jedisPubSub, channels, null);
 			it.subscribe(jedisPubSub, channels);

src/main/java/org/springframework/data/redis/connection/jedis/JedisMessageListener.java

@@ -15,24 +15,29 @@
  */
 package org.springframework.data.redis.connection.jedis;
 
+import redis.clients.jedis.BinaryJedisPubSub;
+
 import org.springframework.data.redis.connection.DefaultMessage;
 import org.springframework.data.redis.connection.MessageListener;
+import org.springframework.data.redis.connection.SubscriptionListener;
 import org.springframework.util.Assert;
 
-import redis.clients.jedis.BinaryJedisPubSub;
-
 /**
  * MessageListener adapter on top of Jedis.
  *
  * @author Costin Leau
+ * @author Mark Paluch
  */
 class JedisMessageListener extends BinaryJedisPubSub {
 
 	private final MessageListener listener;
+	private final SubscriptionListener subscriptionListener;
 
 	JedisMessageListener(MessageListener listener) {
-		Assert.notNull(listener, "message listener is required");
+		Assert.notNull(listener, "MessageListener is required");
 		this.listener = listener;
+		this.subscriptionListener = listener instanceof SubscriptionListener ? (SubscriptionListener) listener
+				: SubscriptionListener.EMPTY;
 	}
 
 	public void onMessage(byte[] channel, byte[] message) {
@@ -44,18 +49,18 @@ class JedisMessageListener extends BinaryJedisPubSub {
 	}
 
 	public void onPSubscribe(byte[] pattern, int subscribedChannels) {
-		// no-op
+		subscriptionListener.onPatternSubscribed(pattern, subscribedChannels);
 	}
 
 	public void onPUnsubscribe(byte[] pattern, int subscribedChannels) {
-		// no-op
+		subscriptionListener.onPatternUnsubscribed(pattern, subscribedChannels);
 	}
 
 	public void onSubscribe(byte[] channel, int subscribedChannels) {
-		// no-op
+		subscriptionListener.onChannelSubscribed(channel, subscribedChannels);
 	}
 
 	public void onUnsubscribe(byte[] channel, int subscribedChannels) {
-		// no-op
+		subscriptionListener.onChannelUnsubscribed(channel, subscribedChannels);
 	}
 }

src/main/java/org/springframework/data/redis/connection/jedis/JedisSubscription.java

@@ -30,7 +30,7 @@ class JedisSubscription extends AbstractSubscription {
 
 	private final BinaryJedisPubSub jedisPubSub;
 
-	JedisSubscription(MessageListener listener, BinaryJedisPubSub jedisPubSub, @Nullable byte[][] channels,
+	JedisSubscription(MessageListener listener, JedisMessageListener jedisPubSub, @Nullable byte[][] channels,
 			@Nullable byte[][] patterns) {
 		super(listener, channels, patterns);
 		this.jedisPubSub = jedisPubSub;

src/main/java/org/springframework/data/redis/connection/lettuce/LettuceMessageListener.java

@@ -19,20 +19,26 @@ import io.lettuce.core.pubsub.RedisPubSubListener;
 
 import org.springframework.data.redis.connection.DefaultMessage;
 import org.springframework.data.redis.connection.MessageListener;
+import org.springframework.data.redis.connection.SubscriptionListener;
 import org.springframework.util.Assert;
 
 /**
  * MessageListener wrapper around Lettuce {@link RedisPubSubListener}.
  *
  * @author Costin Leau
+ * @author Mark Paluch
  */
 class LettuceMessageListener implements RedisPubSubListener<byte[], byte[]> {
 
 	private final MessageListener listener;
+	private final SubscriptionListener subscriptionListener;
 
 	LettuceMessageListener(MessageListener listener) {
 		Assert.notNull(listener, "MessageListener must not be null!");
+
 		this.listener = listener;
+		this.subscriptionListener = listener instanceof SubscriptionListener ? (SubscriptionListener) listener
+				: SubscriptionListener.EMPTY;
 	}
 
 	/*
@@ -55,23 +61,31 @@ class LettuceMessageListener implements RedisPubSubListener<byte[], byte[]> {
 	 * (non-Javadoc)
 	 * @see io.lettuce.core.pubsub.RedisPubSubListener#subscribed(java.lang.Object, long)
 	 */
-	public void subscribed(byte[] channel, long count) {}
+	public void subscribed(byte[] channel, long count) {
+		subscriptionListener.onChannelSubscribed(channel, count);
+	}
 
 	/*
 	 * (non-Javadoc)
 	 * @see io.lettuce.core.pubsub.RedisPubSubListener#psubscribed(java.lang.Object, long)
 	 */
-	public void psubscribed(byte[] pattern, long count) {}
+	public void psubscribed(byte[] pattern, long count) {
+		subscriptionListener.onPatternSubscribed(pattern, count);
+	}
 
 	/*
 	 * (non-Javadoc)
 	 * @see io.lettuce.core.pubsub.RedisPubSubListener#unsubscribed(java.lang.Object, long)
 	 */
-	public void unsubscribed(byte[] channel, long count) {}
+	public void unsubscribed(byte[] channel, long count) {
+		subscriptionListener.onChannelUnsubscribed(channel, count);
+	}
 
 	/*
 	 * (non-Javadoc)
 	 * @see io.lettuce.core.pubsub.RedisPubSubListener#punsubscribed(java.lang.Object, long)
 	 */
-	public void punsubscribed(byte[] pattern, long count) {}
+	public void punsubscribed(byte[] pattern, long count) {
+		subscriptionListener.onPatternUnsubscribed(pattern, count);
+	}
 }

src/main/java/org/springframework/data/redis/connection/lettuce/LettuceSubscription.java

@@ -16,8 +16,13 @@
 package org.springframework.data.redis.connection.lettuce;
 
 import io.lettuce.core.pubsub.StatefulRedisPubSubConnection;
+import io.lettuce.core.pubsub.api.async.RedisPubSubAsyncCommands;
 import io.lettuce.core.pubsub.api.sync.RedisPubSubCommands;
 
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+
 import org.springframework.data.redis.connection.MessageListener;
 import org.springframework.data.redis.connection.util.AbstractSubscription;
 
@@ -37,6 +42,7 @@ public class LettuceSubscription extends AbstractSubscription {
 	private final LettuceMessageListener listener;
 	private final LettuceConnectionProvider connectionProvider;
 	private final RedisPubSubCommands<byte[], byte[]> pubsub;
+	private final RedisPubSubAsyncCommands<byte[], byte[]> pubSubAsync;
 
 	/**
 	 * Creates a new {@link LettuceSubscription} given {@link MessageListener}, {@link StatefulRedisPubSubConnection}, and
@@ -55,6 +61,7 @@ public class LettuceSubscription extends AbstractSubscription {
 		this.listener = new LettuceMessageListener(listener);
 		this.connectionProvider = connectionProvider;
 		this.pubsub = connection.sync();
+		this.pubSubAsync = connection.async();
 
 		this.connection.addListener(this.listener);
 	}
@@ -70,15 +77,29 @@ public class LettuceSubscription extends AbstractSubscription {
 	@Override
 	protected void doClose() {
 
+		List<CompletableFuture<?>> futures = new ArrayList<>();
+
 		if (!getChannels().isEmpty()) {
-			doUnsubscribe(true);
+			futures.add(pubSubAsync.unsubscribe().toCompletableFuture());
 		}
 
 		if (!getPatterns().isEmpty()) {
-			doPUnsubscribe(true);
+			futures.add(pubSubAsync.punsubscribe().toCompletableFuture());
+		}
+
+		if (!futures.isEmpty()) {
+
+			// this is to ensure completion of the futures and result processing. Since we're unsubscribing first, we expect
+			// that we receive pub/sub confirmations before the PING response.
+			futures.add(pubSubAsync.ping().toCompletableFuture());
+
+			CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).whenComplete((v, t) -> {
+				connection.removeListener(listener);
+			});
+		} else {
+			connection.removeListener(listener);
 		}
 
-		connection.removeListener(this.listener);
 		connectionProvider.release(connection);
 	}
 

src/main/java/org/springframework/data/redis/listener/RedisMessageListenerContainer.java

@@ -28,6 +28,7 @@ import java.util.concurrent.TimeUnit;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
+
 import org.springframework.beans.factory.BeanNameAware;
 import org.springframework.beans.factory.DisposableBean;
 import org.springframework.beans.factory.InitializingBean;
@@ -41,6 +42,7 @@ import org.springframework.data.redis.connection.MessageListener;
 import org.springframework.data.redis.connection.RedisConnection;
 import org.springframework.data.redis.connection.RedisConnectionFactory;
 import org.springframework.data.redis.connection.Subscription;
+import org.springframework.data.redis.connection.SubscriptionListener;
 import org.springframework.data.redis.connection.util.ByteArrayWrapper;
 import org.springframework.data.redis.serializer.RedisSerializer;
 import org.springframework.data.redis.serializer.StringRedisSerializer;
@@ -62,13 +64,16 @@ import org.springframework.util.ErrorHandler;
  * configured).
  * <p>
  * Adding and removing listeners at the same time has undefined results. It is strongly recommended to synchronize/order
- * these methods accordingly.
+ * these methods accordingly. {@link MessageListener Listeners} that wish to receive subscription/unsubscription
+ * callbacks in response to subscribe/unsubscribe commands can implement {@link SubscriptionListener}.
  *
  * @author Costin Leau
  * @author Jennifer Hickey
  * @author Way Joke
  * @author Thomas Darimont
  * @author Mark Paluch
+ * @see MessageListener
+ * @see SubscriptionListener
  */
 public class RedisMessageListenerContainer implements InitializingBean, DisposableBean, BeanNameAware, SmartLifecycle {
 
@@ -133,6 +138,7 @@ public class RedisMessageListenerContainer implements InitializingBean, Disposab
 
 	private long maxSubscriptionRegistrationWaitingTime = DEFAULT_SUBSCRIPTION_REGISTRATION_WAIT_TIME;
 
+	@Override
 	public void afterPropertiesSet() {
 		if (taskExecutor == null) {
 			manageExecutor = true;
@@ -958,7 +964,7 @@ public class RedisMessageListenerContainer implements InitializingBean, Disposab
 	 *
 	 * @author Costin Leau
 	 */
-	private class DispatchMessageListener implements MessageListener {
+	private class DispatchMessageListener implements MessageListener, SubscriptionListener {
 
 		@Override
 		public void onMessage(Message message, @Nullable byte[] pattern) {
@@ -977,6 +983,56 @@ public class RedisMessageListenerContainer implements InitializingBean, Disposab
 				dispatchMessage(listeners, message, pattern);
 			}
 		}
+
+		@Override
+		public void onChannelSubscribed(byte[] channel, long count) {
+			dispatchSubscriptionNotification(
+					channelMapping.getOrDefault(new ByteArrayWrapper(channel), Collections.emptyList()), channel, count,
+					SubscriptionListener::onChannelSubscribed);
+		}
+
+		@Override
+		public void onChannelUnsubscribed(byte[] channel, long count) {
+			dispatchSubscriptionNotification(
+					channelMapping.getOrDefault(new ByteArrayWrapper(channel), Collections.emptyList()), channel, count,
+					SubscriptionListener::onChannelUnsubscribed);
+		}
+
+		@Override
+		public void onPatternSubscribed(byte[] pattern, long count) {
+			dispatchSubscriptionNotification(
+					patternMapping.getOrDefault(new ByteArrayWrapper(pattern), Collections.emptyList()), pattern, count,
+					SubscriptionListener::onPatternSubscribed);
+		}
+
+		@Override
+		public void onPatternUnsubscribed(byte[] pattern, long count) {
+			dispatchSubscriptionNotification(
+					patternMapping.getOrDefault(new ByteArrayWrapper(pattern), Collections.emptyList()), pattern, count,
+					SubscriptionListener::onPatternUnsubscribed);
+		}
+	}
+
+	private void dispatchSubscriptionNotification(Collection<MessageListener> listeners, byte[] pattern, long count,
+			SubscriptionConsumer listenerConsumer) {
+
+		if (!CollectionUtils.isEmpty(listeners)) {
+			byte[] source = pattern.clone();
+
+			for (MessageListener messageListener : listeners) {
+				if (messageListener instanceof SubscriptionListener) {
+					taskExecutor.execute(() -> listenerConsumer.accept((SubscriptionListener) messageListener, source, count));
+				}
+			}
+		}
+	}
+
+	/**
+	 * Represents an operation that accepts three input arguments {@link SubscriptionListener},
+	 * {@code channel or pattern}, and {@code count} and returns no result.
+	 */
+	interface SubscriptionConsumer {
+		void accept(SubscriptionListener listener, byte[] channelOrPattern, long count);
 	}
 
 	private void dispatchMessage(Collection<MessageListener> listeners, Message message, @Nullable byte[] pattern) {