Add ReactiveRedisMessageListenerContainer.receiveLater(…) to await subscriptions.

Original Pull Request: #2052
This commit is contained in:
Mark Paluch
2021-04-26 16:02:28 +02:00
committed by Christoph Strobl
parent 328de94d85
commit fe51d78537
2 changed files with 199 additions and 19 deletions

View File

@@ -23,10 +23,12 @@ import java.nio.ByteBuffer;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLongFieldUpdater;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;
@@ -41,9 +43,11 @@ import org.springframework.data.redis.connection.ReactiveSubscription.ChannelMes
import org.springframework.data.redis.connection.ReactiveSubscription.Message;
import org.springframework.data.redis.connection.ReactiveSubscription.PatternMessage;
import org.springframework.data.redis.connection.SubscriptionListener;
import org.springframework.data.redis.connection.util.ByteArrayWrapper;
import org.springframework.data.redis.serializer.RedisElementReader;
import org.springframework.data.redis.serializer.RedisSerializationContext.SerializationPair;
import org.springframework.data.redis.serializer.RedisSerializer;
import org.springframework.data.redis.util.ByteUtils;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
import org.springframework.util.ObjectUtils;
@@ -157,6 +161,28 @@ public class ReactiveRedisMessageListenerContainer implements DisposableBean {
return receive(Arrays.asList(channelTopics), stringSerializationPair, stringSerializationPair);
}
/**
* Subscribe to one or more {@link ChannelTopic}s and receive a stream of {@link ChannelMessage} once the returned
* {@link Mono} completes. Messages and channel names are treated as {@link String}. The message stream subscribes
* lazily to the Redis channels and unsubscribes if the inner {@link org.reactivestreams.Subscription} is
* {@link org.reactivestreams.Subscription#cancel() cancelled}.
* <p/>
* The returned {@link Mono} completes once the connection has been subscribed to the given {@link Topic topics}. Note
* that cancelling the returned {@link Mono} can leave the connection in a subscribed state.
*
* @param channelTopics the channels to subscribe.
* @return the message stream.
* @throws InvalidDataAccessApiUsageException if {@code patternTopics} is empty.
* @since 2.6
*/
public Mono<Flux<Message<String, String>>> receiveLater(ChannelTopic... channelTopics) {
Assert.notNull(channelTopics, "ChannelTopics must not be null!");
Assert.noNullElements(channelTopics, "ChannelTopics must not contain null elements!");
return receiveLater(Arrays.asList(channelTopics), stringSerializationPair, stringSerializationPair);
}
/**
* Subscribe to one or more {@link PatternTopic}s and receive a stream of {@link PatternMessage}. Messages, pattern,
* and channel names are treated as {@link String}. The message stream subscribes lazily to the Redis channels and
@@ -178,6 +204,30 @@ public class ReactiveRedisMessageListenerContainer implements DisposableBean {
.map(m -> (PatternMessage<String, String, String>) m);
}
/**
* Subscribe to one or more {@link PatternTopic}s and receive a stream of {@link PatternMessage} once the returned
* {@link Mono} completes. Messages, pattern, and channel names are treated as {@link String}. The message stream
* subscribes lazily to the Redis channels and unsubscribes if the inner {@link org.reactivestreams.Subscription} is
* {@link org.reactivestreams.Subscription#cancel() cancelled}.
* <p/>
* The returned {@link Mono} completes once the connection has been subscribed to the given {@link Topic topics}. Note
* that cancelling the returned {@link Mono} can leave the connection in a subscribed state.
*
* @param patternTopics the channels to subscribe.
* @return the message stream.
* @throws InvalidDataAccessApiUsageException if {@code patternTopics} is empty.
* @since 2.6
*/
@SuppressWarnings("unchecked")
public Mono<Flux<PatternMessage<String, String, String>>> receiveLater(PatternTopic... patternTopics) {
Assert.notNull(patternTopics, "PatternTopic must not be null!");
Assert.noNullElements(patternTopics, "PatternTopic must not contain null elements!");
return receiveLater(Arrays.asList(patternTopics), stringSerializationPair, stringSerializationPair)
.map(it -> it.map(m -> (PatternMessage<String, String, String>) m));
}
/**
* Subscribe to one or more {@link Topic}s and receive a stream of {@link ChannelMessage}. The stream may contain
* {@link PatternMessage} if subscribed to patterns. Messages, and channel names are serialized/deserialized using the
@@ -281,6 +331,68 @@ public class ReactiveRedisMessageListenerContainer implements DisposableBean {
.map(message -> readMessage(channelSerializer.getReader(), messageSerializer.getReader(), message));
}
/**
* Subscribe to one or more {@link Topic}s and receive a stream of {@link ChannelMessage}. The returned {@link Mono}
* completes once the connection has been subscribed to the given {@link Topic topics}. Note that cancelling the
* returned {@link Mono} can leave the connection in a subscribed state.
*
* @param topics the channels to subscribe.
* @param channelSerializer serialization pair to decode the channel/pattern name.
* @param messageSerializer serialization pair to decode the message body.
* @return the message stream.
* @throws InvalidDataAccessApiUsageException if {@code topics} is empty.
* @since 2.6
*/
private <C, B> Mono<Flux<Message<C, B>>> receiveLater(Iterable<? extends Topic> topics,
SerializationPair<C> channelSerializer, SerializationPair<B> messageSerializer) {
Assert.notNull(topics, "Topics must not be null!");
Assert.notNull(channelSerializer, "Channel serializer must not be null!");
Assert.notNull(messageSerializer, "Message serializer must not be null!");
verifyConnection();
ByteBuffer[] patterns = getTargets(topics, PatternTopic.class);
ByteBuffer[] channels = getTargets(topics, ChannelTopic.class);
if (ObjectUtils.isEmpty(patterns) && ObjectUtils.isEmpty(channels)) {
throw new InvalidDataAccessApiUsageException("No channels or patterns to subscribe to.");
}
return Mono.defer(() -> {
SubscriptionReadyListener readyListener = SubscriptionReadyListener.create(topics, stringSerializationPair);
return doReceiveLater(channelSerializer, messageSerializer,
getRequiredConnection().pubSubCommands().createSubscription(readyListener), patterns, channels)
.delayUntil(it -> readyListener.getTrigger());
});
}
private <C, B> Mono<Flux<Message<C, B>>> doReceiveLater(SerializationPair<C> channelSerializer,
SerializationPair<B> messageSerializer, Mono<ReactiveSubscription> subscription, ByteBuffer[] patterns,
ByteBuffer[] channels) {
return subscription.flatMap(it -> {
Mono<Void> subscribe = subscribe(patterns, channels, it).doOnSuccess(v -> getSubscribers(it).registered());
Sinks.One<Message<ByteBuffer, ByteBuffer>> terminalSink = Sinks.one();
Flux<Message<C, B>> receiver = it.receive().doOnCancel(() -> {
Subscribers subscribers = getSubscribers(it);
if (subscribers.unregister()) {
subscriptions.remove(it);
it.cancel().subscribe(v -> terminalSink.tryEmitEmpty(), terminalSink::tryEmitError);
}
}).mergeWith(terminalSink.asMono())
.map(message -> readMessage(channelSerializer.getReader(), messageSerializer.getReader(), message));
return subscribe.then(Mono.just(receiver));
});
}
private static Mono<Void> subscribe(ByteBuffer[] patterns, ByteBuffer[] channels, ReactiveSubscription it) {
Assert.isTrue(!ObjectUtils.isEmpty(channels) || !ObjectUtils.isEmpty(patterns),
@@ -418,4 +530,54 @@ public class ReactiveRedisMessageListenerContainer implements DisposableBean {
return false;
}
}
static class SubscriptionReadyListener extends AtomicBoolean implements SubscriptionListener {
private final Set<ByteArrayWrapper> toSubscribe;
private final Sinks.Empty<Void> sink = Sinks.empty();
private SubscriptionReadyListener(Set<ByteArrayWrapper> topics) {
this.toSubscribe = topics;
}
public static SubscriptionReadyListener create(Iterable<? extends Topic> topics,
SerializationPair<String> serializationPair) {
Set<ByteArrayWrapper> wrappers = new HashSet<>();
for (Topic topic : topics) {
wrappers.add(new ByteArrayWrapper(ByteUtils.getBytes(serializationPair.getWriter().write(topic.getTopic()))));
}
return new SubscriptionReadyListener(wrappers);
}
@Override
public void onChannelSubscribed(byte[] channel, long count) {
removeRemaining(channel);
}
@Override
public void onPatternSubscribed(byte[] pattern, long count) {
removeRemaining(pattern);
}
private void removeRemaining(byte[] channel) {
boolean done;
synchronized (toSubscribe) {
toSubscribe.remove(new ByteArrayWrapper(channel));
done = toSubscribe.isEmpty();
}
if (done && compareAndSet(false, true)) {
sink.emitEmpty(Sinks.EmitFailureHandler.FAIL_FAST);
}
}
public Mono<Void> getTrigger() {
return sink.asMono();
}
}
}

View File

@@ -20,6 +20,7 @@ import static org.assertj.core.api.Assertions.*;
import reactor.core.Disposable;
import reactor.test.StepVerifier;
import java.nio.ByteBuffer;
import java.time.Duration;
import java.util.Collection;
import java.util.Collections;
@@ -28,6 +29,7 @@ import java.util.concurrent.CompletableFuture;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import java.util.function.Supplier;
import org.awaitility.Awaitility;
@@ -36,6 +38,7 @@ import org.junit.jupiter.api.BeforeEach;
import org.springframework.data.redis.connection.Message;
import org.springframework.data.redis.connection.MessageListener;
import org.springframework.data.redis.connection.ReactiveRedisConnection;
import org.springframework.data.redis.connection.ReactiveSubscription;
import org.springframework.data.redis.connection.ReactiveSubscription.ChannelMessage;
import org.springframework.data.redis.connection.ReactiveSubscription.PatternMessage;
@@ -62,6 +65,7 @@ public class ReactiveRedisMessageListenerContainerIntegrationTests {
private final LettuceConnectionFactory connectionFactory;
private @Nullable RedisConnection connection;
private @Nullable ReactiveRedisConnection reactiveConnection;
/**
* @param connectionFactory
@@ -79,6 +83,7 @@ public class ReactiveRedisMessageListenerContainerIntegrationTests {
@BeforeEach
void before() {
connection = connectionFactory.getConnection();
reactiveConnection = connectionFactory.getReactiveConnection();
}
@AfterEach
@@ -87,16 +92,21 @@ public class ReactiveRedisMessageListenerContainerIntegrationTests {
if (connection != null) {
connection.close();
}
if (reactiveConnection != null) {
reactiveConnection.close();
}
}
@ParameterizedRedisTest // DATAREDIS-612
@ParameterizedRedisTest // DATAREDIS-612, GH-1622
void shouldReceiveChannelMessages() {
ReactiveRedisMessageListenerContainer container = new ReactiveRedisMessageListenerContainer(connectionFactory);
container.receive(ChannelTopic.of(CHANNEL1)).as(StepVerifier::create) //
.then(awaitSubscription(container::getActiveSubscriptions))
.then(() -> connection.publish(CHANNEL1.getBytes(), MESSAGE.getBytes())) //
container.receiveLater(ChannelTopic.of(CHANNEL1)) //
.doOnNext(it -> doPublish(CHANNEL1.getBytes(), MESSAGE.getBytes())) //
.flatMapMany(Function.identity()) //
.as(StepVerifier::create) //
.assertNext(c -> {
assertThat(c.getChannel()).isEqualTo(CHANNEL1);
@@ -136,9 +146,10 @@ public class ReactiveRedisMessageListenerContainerIntegrationTests {
}
};
container.receive(Collections.singletonList(ChannelTopic.of(CHANNEL1)), listener).as(StepVerifier::create) //
container.receive(Collections.singletonList(ChannelTopic.of(CHANNEL1)), listener) //
.as(StepVerifier::create) //
.then(awaitSubscription(container::getActiveSubscriptions))
.then(() -> connection.publish(CHANNEL1.getBytes(), MESSAGE.getBytes())) //
.then(() -> doPublish(CHANNEL1.getBytes(), MESSAGE.getBytes())) //
.assertNext(c -> {
assertThat(c.getChannel()).isEqualTo(CHANNEL1);
@@ -154,14 +165,14 @@ public class ReactiveRedisMessageListenerContainerIntegrationTests {
container.destroy();
}
@ParameterizedRedisTest // DATAREDIS-612
@ParameterizedRedisTest // DATAREDIS-612, GH-1622
void shouldReceivePatternMessages() {
ReactiveRedisMessageListenerContainer container = new ReactiveRedisMessageListenerContainer(connectionFactory);
container.receive(PatternTopic.of(PATTERN1)).as(StepVerifier::create) //
.then(awaitSubscription(container::getActiveSubscriptions))
.then(() -> connection.publish(CHANNEL1.getBytes(), MESSAGE.getBytes())) //
container.receiveLater(PatternTopic.of(PATTERN1)) //
.doOnNext(it -> doPublish(CHANNEL1.getBytes(), MESSAGE.getBytes())).flatMapMany(Function.identity()) //
.as(StepVerifier::create) //
.assertNext(c -> {
assertThat(c.getPattern()).isEqualTo(PATTERN1);
@@ -206,7 +217,7 @@ public class ReactiveRedisMessageListenerContainerIntegrationTests {
.cast(PatternMessage.class) //
.as(StepVerifier::create) //
.then(awaitSubscription(container::getActiveSubscriptions))
.then(() -> connection.publish(CHANNEL1.getBytes(), MESSAGE.getBytes())) //
.then(() -> doPublish(CHANNEL1.getBytes(), MESSAGE.getBytes())) //
.assertNext(c -> {
assertThat(c.getPattern()).isEqualTo(PATTERN1);
@@ -223,19 +234,22 @@ public class ReactiveRedisMessageListenerContainerIntegrationTests {
container.destroy();
}
@ParameterizedRedisTest // DATAREDIS-612
void shouldPublishAndReceiveMessage() throws InterruptedException {
@ParameterizedRedisTest // DATAREDIS-612, GH-1622
void shouldPublishAndReceiveMessage() throws Exception {
ReactiveRedisMessageListenerContainer container = new ReactiveRedisMessageListenerContainer(connectionFactory);
ReactiveRedisTemplate<String, String> template = new ReactiveRedisTemplate<>(connectionFactory,
RedisSerializationContext.string());
BlockingQueue<PatternMessage<String, String, String>> messages = new LinkedBlockingDeque<>();
Disposable subscription = container.receive(PatternTopic.of(PATTERN1)).doOnNext(messages::add).subscribe();
CompletableFuture<Void> subscribed = new CompletableFuture<>();
Disposable subscription = container.receiveLater(PatternTopic.of(PATTERN1))
.doOnNext(it -> subscribed.complete(null)).flatMapMany(Function.identity()).doOnNext(messages::add).subscribe();
StepVerifier.create(template.convertAndSend(CHANNEL1, MESSAGE), 0) //
.then(awaitSubscription(container::getActiveSubscriptions)) //
.thenRequest(1).expectNextCount(1) //
subscribed.get(5, TimeUnit.SECONDS);
template.convertAndSend(CHANNEL1, MESSAGE).as(StepVerifier::create) //
.expectNextCount(1) //
.verifyComplete();
PatternMessage<String, String, String> message = messages.poll(1, TimeUnit.SECONDS);
@@ -257,7 +271,7 @@ public class ReactiveRedisMessageListenerContainerIntegrationTests {
template.listenToChannel(CHANNEL1).as(StepVerifier::create) //
.thenAwait(Duration.ofMillis(100)) // just make sure we the subscription completed
.then(() -> connection.publish(CHANNEL1.getBytes(), MESSAGE.getBytes())) //
.then(() -> doPublish(CHANNEL1.getBytes(), MESSAGE.getBytes())) //
.assertNext(message -> {
assertThat(message).isInstanceOf(ChannelMessage.class);
@@ -276,7 +290,7 @@ public class ReactiveRedisMessageListenerContainerIntegrationTests {
template.listenToPattern(PATTERN1).as(StepVerifier::create) //
.thenAwait(Duration.ofMillis(100)) // just make sure we the subscription completed
.then(() -> connection.publish(CHANNEL1.getBytes(), MESSAGE.getBytes())) //
.then(() -> doPublish(CHANNEL1.getBytes(), MESSAGE.getBytes())) //
.assertNext(message -> {
assertThat(message).isInstanceOf(PatternMessage.class);
@@ -288,6 +302,10 @@ public class ReactiveRedisMessageListenerContainerIntegrationTests {
.verify();
}
private void doPublish(byte[] channel, byte[] message) {
reactiveConnection.pubSubCommands().publish(ByteBuffer.wrap(channel), ByteBuffer.wrap(message)).subscribe();
}
private static Runnable awaitSubscription(Supplier<Collection<ReactiveSubscription>> activeSubscriptions) {
return () -> {