org.springframework.data.redis.connection.lettuce.LettuceReactiveRedisConnection Maven / Gradle / Ivy
/*
* Copyright 2016-2018 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
*
* http://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 io.lettuce.core.api.StatefulConnection;
import io.lettuce.core.api.StatefulRedisConnection;
import io.lettuce.core.api.reactive.BaseRedisReactiveCommands;
import io.lettuce.core.cluster.api.StatefulRedisClusterConnection;
import io.lettuce.core.cluster.api.reactive.RedisClusterReactiveCommands;
import io.lettuce.core.codec.RedisCodec;
import io.lettuce.core.pubsub.StatefulRedisPubSubConnection;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;
import java.nio.ByteBuffer;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import org.reactivestreams.Publisher;
import org.springframework.dao.DataAccessException;
import org.springframework.dao.InvalidDataAccessResourceUsageException;
import org.springframework.data.redis.connection.*;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
/**
* @author Christoph Strobl
* @author Mark Paluch
* @since 2.0
*/
class LettuceReactiveRedisConnection implements ReactiveRedisConnection {
static final RedisCodec CODEC = ByteBufferCodec.INSTANCE;
private final AsyncConnect> dedicatedConnection;
private final AsyncConnect> pubSubConnection;
private @Nullable Mono> sharedConnection;
/**
* Creates new {@link LettuceReactiveRedisConnection}.
*
* @param connectionProvider must not be {@literal null}.
* @throws IllegalArgumentException when {@code client} is {@literal null}.
* @throws InvalidDataAccessResourceUsageException when {@code client} is not suitable for connection.
*/
@SuppressWarnings("unchecked")
LettuceReactiveRedisConnection(LettuceConnectionProvider connectionProvider) {
Assert.notNull(connectionProvider, "LettuceConnectionProvider must not be null!");
this.dedicatedConnection = new AsyncConnect(connectionProvider, StatefulConnection.class);
this.pubSubConnection = new AsyncConnect(connectionProvider, StatefulRedisPubSubConnection.class);
}
/**
* Creates new {@link LettuceReactiveRedisConnection} given a shared {@link StatefulConnection connection}.
*
* @param sharedConnection must not be {@literal null}.
* @param connectionProvider must not be {@literal null}.
* @throws IllegalArgumentException when {@code client} is {@literal null}.
* @throws InvalidDataAccessResourceUsageException when {@code client} is not suitable for connection.
* @since 2.0.1
*/
@SuppressWarnings("unchecked")
LettuceReactiveRedisConnection(StatefulConnection sharedConnection,
LettuceConnectionProvider connectionProvider) {
Assert.notNull(sharedConnection, "Shared StatefulConnection must not be null!");
Assert.notNull(connectionProvider, "LettuceConnectionProvider must not be null!");
this.dedicatedConnection = new AsyncConnect(connectionProvider, StatefulConnection.class);
this.pubSubConnection = new AsyncConnect(connectionProvider, StatefulRedisPubSubConnection.class);
this.sharedConnection = Mono.just(sharedConnection);
}
/*
* (non-Javadoc)
* @see org.springframework.data.redis.connection.ReactiveRedisConnection#keyCommands()
*/
@Override
public ReactiveKeyCommands keyCommands() {
return new LettuceReactiveKeyCommands(this);
}
/*
* (non-Javadoc)
* @see org.springframework.data.redis.connection.ReactiveRedisConnection#stringCommands()
*/
@Override
public ReactiveStringCommands stringCommands() {
return new LettuceReactiveStringCommands(this);
}
/*
* (non-Javadoc)
* @see org.springframework.data.redis.connection.ReactiveRedisConnection#numberCommands()
*/
@Override
public ReactiveNumberCommands numberCommands() {
return new LettuceReactiveNumberCommands(this);
}
/*
* (non-Javadoc)
* @see org.springframework.data.redis.connection.ReactiveRedisConnection#listCommands()
*/
@Override
public ReactiveListCommands listCommands() {
return new LettuceReactiveListCommands(this);
}
/*
* (non-Javadoc)
* @see org.springframework.data.redis.connection.ReactiveRedisConnection#setCommands()
*/
@Override
public ReactiveSetCommands setCommands() {
return new LettuceReactiveSetCommands(this);
}
/*
* (non-Javadoc)
* @see org.springframework.data.redis.connection.ReactiveRedisConnection#zSetCommands()
*/
@Override
public ReactiveZSetCommands zSetCommands() {
return new LettuceReactiveZSetCommands(this);
}
/*
* (non-Javadoc)
* @see org.springframework.data.redis.connection.ReactiveRedisConnection#hashCommands()
*/
@Override
public ReactiveHashCommands hashCommands() {
return new LettuceReactiveHashCommands(this);
}
/*
* (non-Javadoc)
* @see org.springframework.data.redis.connection.ReactiveRedisConnection#geoCommands()
*/
@Override
public ReactiveGeoCommands geoCommands() {
return new LettuceReactiveGeoCommands(this);
}
/*
* (non-Javadoc)
* @see org.springframework.data.redis.connection.ReactiveRedisConnection#hyperLogLogCommands()
*/
@Override
public ReactiveHyperLogLogCommands hyperLogLogCommands() {
return new LettuceReactiveHyperLogLogCommands(this);
}
/*
* (non-Javadoc)
* @see org.springframework.data.redis.connection.ReactiveRedisConnection#pubSubCommands()
*/
@Override
public ReactivePubSubCommands pubSubCommands() {
return new LettuceReactivePubSubCommands(this);
}
/*
* (non-Javadoc)
* @see org.springframework.data.redis.connection.ReactiveRedisConnection#scriptingCommands()
*/
@Override
public ReactiveScriptingCommands scriptingCommands() {
return new LettuceReactiveScriptingCommands(this);
}
/*
* (non-Javadoc)
* @see org.springframework.data.redis.connection.ReactiveRedisConnection#serverCommands()
*/
@Override
public ReactiveServerCommands serverCommands() {
return new LettuceReactiveServerCommands(this);
}
/*
* (non-Javadoc)
* @see org.springframework.data.redis.connection.ReactiveRedisConnection#ping()
*/
@Override
public Mono ping() {
return execute(BaseRedisReactiveCommands::ping).next();
}
/**
* @param callback
* @return
*/
public Flux execute(LettuceReactiveCallback callback) {
return getCommands().flatMapMany(callback::doWithCommands).onErrorMap(translateException());
}
/**
* @param callback
* @return
* @since 2.0.1
*/
public Flux executeDedicated(LettuceReactiveCallback callback) {
return getDedicatedCommands().flatMapMany(callback::doWithCommands).onErrorMap(translateException());
}
/*
* (non-Javadoc)
* @see org.springframework.data.redis.connection.ReactiveRedisConnection#closeLater()
*/
public Mono closeLater() {
return Mono.fromRunnable(() -> dedicatedConnection.close());
}
protected Mono> getConnection() {
if (sharedConnection != null) {
return sharedConnection;
}
return getDedicatedConnection();
}
protected Mono> getDedicatedConnection() {
return dedicatedConnection.getConnection().onErrorMap(translateException());
}
protected Mono> getPubSubConnection() {
return pubSubConnection.getConnection().onErrorMap(translateException());
}
protected Mono> getCommands() {
if (sharedConnection != null) {
return sharedConnection.map(LettuceReactiveRedisConnection::getRedisClusterReactiveCommands);
}
return getDedicatedCommands();
}
protected Mono> getDedicatedCommands() {
return dedicatedConnection.getConnection().map(LettuceReactiveRedisConnection::getRedisClusterReactiveCommands);
}
private static RedisClusterReactiveCommands getRedisClusterReactiveCommands(
StatefulConnection connection) {
if (connection instanceof StatefulRedisConnection) {
return ((StatefulRedisConnection) connection).reactive();
} else if (connection instanceof StatefulRedisClusterConnection) {
return ((StatefulRedisClusterConnection) connection).reactive();
}
throw new IllegalStateException("o.O unknown connection type " + connection);
}
Function translateException() {
return throwable -> {
if (throwable instanceof RuntimeException) {
DataAccessException convertedException = LettuceConverters.exceptionConverter()
.convert((RuntimeException) throwable);
return convertedException != null ? convertedException : throwable;
}
return throwable;
};
}
interface LettuceReactiveCallback {
Publisher doWithCommands(RedisClusterReactiveCommands cmd);
}
enum ByteBufferCodec implements RedisCodec {
INSTANCE;
@Override
public ByteBuffer decodeKey(ByteBuffer bytes) {
ByteBuffer buffer = ByteBuffer.allocate(bytes.remaining());
buffer.put(bytes);
buffer.flip();
return buffer;
}
@Override
public ByteBuffer decodeValue(ByteBuffer bytes) {
return decodeKey(bytes);
}
@Override
public ByteBuffer encodeKey(ByteBuffer key) {
return key.duplicate();
}
@Override
public ByteBuffer encodeValue(ByteBuffer value) {
return value.duplicate();
}
}
/**
* Asynchronous connection utility. This utility has a lifecycle controlled by {@link #getConnection()} and
* {@link #close()} methods:
*
* - Initial state: Not connected. Calling {@link #getConnection()} will transition to the next state.
* - Connection requested: First caller to {@link #getConnection()} initiates an asynchronous connect. Connection is
* accessible through the resulting {@link Mono}.
* - Closing: Call to {@link #close()} initiates connection closing.
* - Closed: Connection is closed.
*
*
* @author Mark Paluch
* @author Christoph Strobl
* @since 2.0.1
*/
static class AsyncConnect> {
private final Mono connectionPublisher;
private final LettuceConnectionProvider connectionProvider;
private AtomicReference state = new AtomicReference<>(State.INITIAL);
private volatile @Nullable CompletableFuture connection;
@SuppressWarnings("unchecked")
AsyncConnect(LettuceConnectionProvider connectionProvider, Class connectionType) {
Assert.notNull(connectionProvider, "LettuceConnectionProvider must not be null!");
this.connectionProvider = connectionProvider;
Mono defer = Mono.defer(() -> Mono. just(connectionProvider.getConnection(connectionType)));
this.connectionPublisher = defer.subscribeOn(Schedulers.elastic());
}
/**
* Obtain a connection publisher. This method connects asynchronously by requesting a connection from
* {@link LettuceConnectionProvider} with non-blocking synchronization if called concurrently.
*
* @return never {@literal null}.
*/
Mono getConnection() {
if (state.get() == State.CLOSED) {
throw new IllegalStateException("Unable to connect. Connection is closed!");
}
CompletableFuture connection = this.connection;
if (connection != null) {
return Mono.fromCompletionStage(connection);
}
if (state.compareAndSet(State.INITIAL, State.CONNECTION_REQUESTED)) {
this.connection = connectionPublisher.toFuture();
}
while (true) {
connection = this.connection;
if (connection != null) {
return Mono.fromCompletionStage(connection);
}
Thread thread = Thread.currentThread();
if (thread.isInterrupted()) {
thread.interrupt();
return Mono.error(new InterruptedException());
}
}
}
/**
* Close connection (blocking call).
*/
void close() {
if (state.compareAndSet(State.CONNECTION_REQUESTED, State.CLOSING)) {
getConnection().doOnSuccess(connectionProvider::release).doOnSuccess(it -> state.set(State.CLOSED)).block();
}
state.compareAndSet(State.INITIAL, State.CLOSED);
}
enum State {
INITIAL, CONNECTION_REQUESTED, CLOSING, CLOSED;
}
}
}