Add support for subscription ready/unsubscription done events.

Original Pull Request: #2052

Author: mp911de

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

@@ -18,11 +18,14 @@
 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 void subscribe(MessageListener listener, byte[]... channels) {
 					"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 void pSubscribe(MessageListener listener, byte[]... patterns) {
 					"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 void pSubscribe(MessageListener listener, byte[]... patterns) {
 
 		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 void subscribe(MessageListener listener, byte[]... channels) {
 
 		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 @@ public void onPMessage(byte[] pattern, byte[] channel, byte[] message) {
 	}
 
 	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 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 @@ public void message(byte[] pattern, byte[] channel, byte[] message) {
 	 * (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 @@ protected LettuceSubscription(MessageListener listener,
 		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 @@ protected StatefulRedisPubSubConnection<byte[], byte[]> getNativeConnection() {
 	@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);
 	}