io.lettuce.core.cluster.StatefulRedisClusterConnectionImpl Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of lettuce-core Show documentation
Advanced and thread-safe Java Redis client for synchronous, asynchronous, and reactive usage. Supports Cluster, Sentinel, Pipelining, Auto-Reconnect, Codecs and much more.
The newest version!
/*
 * Copyright 2011-Present, Redis Ltd. and Contributors
 * All rights reserved.
 *
 * Licensed under the MIT License.
 *
 * This file contains contributions from third-party contributors
 * 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 io.lettuce.core.cluster;

import static io.lettuce.core.ClientOptions.DEFAULT_JSON_PARSER;
import static io.lettuce.core.protocol.CommandType.*;

import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Proxy;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.function.Consumer;
import java.util.stream.Collectors;

import io.lettuce.core.AbstractRedisClient;
import io.lettuce.core.ClientOptions;
import io.lettuce.core.ConnectionState;
import io.lettuce.core.ReadFrom;
import io.lettuce.core.RedisChannelHandler;
import io.lettuce.core.RedisChannelWriter;
import io.lettuce.core.RedisException;
import io.lettuce.core.RedisURI;
import io.lettuce.core.api.StatefulRedisConnection;
import io.lettuce.core.cluster.api.StatefulRedisClusterConnection;
import io.lettuce.core.cluster.api.async.RedisAdvancedClusterAsyncCommands;
import io.lettuce.core.cluster.api.async.RedisClusterAsyncCommands;
import io.lettuce.core.cluster.api.push.RedisClusterPushListener;
import io.lettuce.core.cluster.api.reactive.RedisAdvancedClusterReactiveCommands;
import io.lettuce.core.cluster.api.sync.NodeSelection;
import io.lettuce.core.cluster.api.sync.NodeSelectionCommands;
import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands;
import io.lettuce.core.cluster.models.partitions.Partitions;
import io.lettuce.core.cluster.models.partitions.RedisClusterNode;
import io.lettuce.core.codec.RedisCodec;
import io.lettuce.core.internal.LettuceAssert;
import io.lettuce.core.json.JsonParser;
import io.lettuce.core.protocol.CommandArgsAccessor;
import io.lettuce.core.protocol.CompleteableCommand;
import io.lettuce.core.protocol.ConnectionIntent;
import io.lettuce.core.protocol.ConnectionWatchdog;
import io.lettuce.core.protocol.RedisCommand;
import reactor.core.publisher.Mono;

/**
 * A thread-safe connection to a Redis Cluster. Multiple threads may share one {@link StatefulRedisClusterConnectionImpl}
 *
 * A {@link ConnectionWatchdog} monitors each connection and reconnects automatically until {@link #close} is called. All
 * pending commands will be (re)sent after successful reconnection.
 *
 * @author Mark Paluch
 * @since 4.0
 */
public class StatefulRedisClusterConnectionImpl extends RedisChannelHandler
        implements StatefulRedisClusterConnection {

    private final ClusterPushHandler pushHandler;

    protected final RedisCodec codec;

    protected final Mono parser;

    protected final RedisAdvancedClusterCommands sync;

    protected final RedisAdvancedClusterAsyncCommandsImpl async;

    protected final RedisAdvancedClusterReactiveCommandsImpl reactive;

    private final ClusterConnectionState connectionState = new ClusterConnectionState();

    private volatile Partitions partitions;

    /**
     * Initialize a new connection.
     *
     * @param writer the channel writer
     * @param pushHandler the Cluster push handler
     * @param codec Codec used to encode/decode keys and values.
     * @param timeout Maximum time to wait for a response.
     */
    public StatefulRedisClusterConnectionImpl(RedisChannelWriter writer, ClusterPushHandler pushHandler, RedisCodec codec,
            Duration timeout) {

        this(writer, pushHandler, codec, timeout, DEFAULT_JSON_PARSER);
    }

    /**
     * Initialize a new connection.
     *
     * @param writer the channel writer
     * @param pushHandler the Cluster push handler
     * @param codec Codec used to encode/decode keys and values.
     * @param timeout Maximum time to wait for a response.
     * @param parser the JSON parser
     */
    public StatefulRedisClusterConnectionImpl(RedisChannelWriter writer, ClusterPushHandler pushHandler, RedisCodec codec,
            Duration timeout, Mono parser) {

        super(writer, timeout);
        this.pushHandler = pushHandler;
        this.codec = codec;
        this.parser = parser;

        this.async = newRedisAdvancedClusterAsyncCommandsImpl();
        this.sync = newRedisAdvancedClusterCommandsImpl();
        this.reactive = newRedisAdvancedClusterReactiveCommandsImpl();
    }

    protected RedisAdvancedClusterReactiveCommandsImpl newRedisAdvancedClusterReactiveCommandsImpl() {
        return new RedisAdvancedClusterReactiveCommandsImpl<>((StatefulRedisClusterConnection) this, codec, parser);
    }

    protected RedisAdvancedClusterCommands newRedisAdvancedClusterCommandsImpl() {
        return clusterSyncHandler(RedisAdvancedClusterCommands.class);
    }

    protected  T clusterSyncHandler(Class... interfaces) {
        return (T) Proxy.newProxyInstance(AbstractRedisClient.class.getClassLoader(), interfaces, syncInvocationHandler());
    }

    protected RedisAdvancedClusterAsyncCommandsImpl newRedisAdvancedClusterAsyncCommandsImpl() {
        return new RedisAdvancedClusterAsyncCommandsImpl((StatefulRedisClusterConnection) this, codec, parser);
    }

    @Override
    public RedisAdvancedClusterCommands sync() {
        return sync;
    }

    protected InvocationHandler syncInvocationHandler() {
        return new ClusterFutureSyncInvocationHandler<>(this, RedisClusterAsyncCommands.class, NodeSelection.class,
                NodeSelectionCommands.class, async());
    }

    @Override
    public RedisAdvancedClusterAsyncCommands async() {
        return async;
    }

    @Override
    public RedisAdvancedClusterReactiveCommands reactive() {
        return reactive;
    }

    @Override
    public void addListener(RedisClusterPushListener listener) {
        pushHandler.addListener(listener);
    }

    @Override
    public void removeListener(RedisClusterPushListener listener) {
        pushHandler.removeListener(listener);
    }

    private RedisURI lookup(String nodeId) {

        for (RedisClusterNode partition : partitions) {
            if (partition.getNodeId().equals(nodeId)) {
                return partition.getUri();
            }
        }
        return null;
    }

    @Override
    public StatefulRedisConnection getConnection(String nodeId, ConnectionIntent connectionIntent) {

        RedisURI redisURI = lookup(nodeId);

        if (redisURI == null) {
            throw new RedisException("NodeId " + nodeId + " does not belong to the cluster");
        }

        return getClusterDistributionChannelWriter().getClusterConnectionProvider().getConnection(connectionIntent, nodeId);
    }

    @Override
    public CompletableFuture> getConnectionAsync(String nodeId,
            ConnectionIntent connectionIntent) {

        RedisURI redisURI = lookup(nodeId);

        if (redisURI == null) {
            throw new RedisException("NodeId " + nodeId + " does not belong to the cluster");
        }

        AsyncClusterConnectionProvider provider = (AsyncClusterConnectionProvider) getClusterDistributionChannelWriter()
                .getClusterConnectionProvider();

        return provider.getConnectionAsync(connectionIntent, nodeId);
    }

    @Override
    public StatefulRedisConnection getConnection(String host, int port, ConnectionIntent connectionIntent) {

        return getClusterDistributionChannelWriter().getClusterConnectionProvider().getConnection(connectionIntent, host, port);
    }

    @Override
    public CompletableFuture> getConnectionAsync(String host, int port,
            ConnectionIntent connectionIntent) {

        AsyncClusterConnectionProvider provider = (AsyncClusterConnectionProvider) getClusterDistributionChannelWriter()
                .getClusterConnectionProvider();

        return provider.getConnectionAsync(connectionIntent, host, port);
    }

    @Override
    public void activated() {
        super.activated();

        async.clusterMyId().thenAccept(connectionState::setNodeId);
    }

    ClusterDistributionChannelWriter getClusterDistributionChannelWriter() {
        return (ClusterDistributionChannelWriter) super.getChannelWriter();
    }

    @Override
    public  RedisCommand dispatch(RedisCommand command) {
        return super.dispatch(preProcessCommand(command));
    }

    @Override
    public Collection> dispatch(Collection> commands) {

        List> commandsToSend = new ArrayList<>(commands.size());
        for (RedisCommand command : commands) {
            commandsToSend.add(preProcessCommand(command));
        }

        return super.dispatch(commandsToSend);
    }

    private  RedisCommand preProcessCommand(RedisCommand command) {

        RedisCommand local = command;

        if (local.getType().toString().equals(AUTH.name())) {
            local = attachOnComplete(local, status -> {
                if (status.equals("OK")) {
                    List args = CommandArgsAccessor.getCharArrayArguments(command.getArgs());

                    if (!args.isEmpty()) {
                        this.connectionState.setUserNamePassword(args);
                    } else {

                        List stringArgs = CommandArgsAccessor.getStringArguments(command.getArgs());
                        this.connectionState
                                .setUserNamePassword(stringArgs.stream().map(String::toCharArray).collect(Collectors.toList()));
                    }
                }
            });
        }

        if (local.getType().toString().equals(READONLY.name())) {
            local = attachOnComplete(local, status -> {
                if (status.equals("OK")) {
                    this.connectionState.setReadOnly(true);
                }
            });
        }

        if (local.getType().toString().equals(READWRITE.name())) {
            local = attachOnComplete(local, status -> {
                if (status.equals("OK")) {
                    this.connectionState.setReadOnly(false);
                }
            });
        }
        return local;
    }

    private  RedisCommand attachOnComplete(RedisCommand command, Consumer consumer) {

        if (command instanceof CompleteableCommand) {
            CompleteableCommand completeable = (CompleteableCommand) command;
            completeable.onComplete(consumer);
        }
        return command;
    }

    public void setPartitions(Partitions partitions) {

        LettuceAssert.notNull(partitions, "Partitions must not be null");

        this.partitions = partitions;

        String nodeId = connectionState.getNodeId();
        if (nodeId != null && expireStaleConnections()) {

            if (partitions.getPartitionByNodeId(nodeId) == null) {
                getClusterDistributionChannelWriter().disconnectDefaultEndpoint();
            }
        }

        getClusterDistributionChannelWriter().setPartitions(partitions);
    }

    public Partitions getPartitions() {
        return partitions;
    }

    @Override
    public void setReadFrom(ReadFrom readFrom) {
        LettuceAssert.notNull(readFrom, "ReadFrom must not be null");
        getClusterDistributionChannelWriter().setReadFrom(readFrom);
    }

    @Override
    public ReadFrom getReadFrom() {
        return getClusterDistributionChannelWriter().getReadFrom();
    }

    ConnectionState getConnectionState() {
        return connectionState;
    }

    static class ClusterConnectionState extends ConnectionState {

        private volatile String nodeId;

        @Override
        protected void setUserNamePassword(List args) {
            super.setUserNamePassword(args);
        }

        @Override
        protected void setDb(int db) {
            super.setDb(db);
        }

        @Override
        protected void setReadOnly(boolean readOnly) {
            super.setReadOnly(readOnly);
        }

        public String getNodeId() {
            return nodeId;
        }

        public void setNodeId(String nodeId) {
            this.nodeId = nodeId;
        }

    }

    private boolean expireStaleConnections() {

        ClusterClientOptions options = getClusterClientOptions();
        return options == null || options.isCloseStaleConnections();
    }

    private ClusterClientOptions getClusterClientOptions() {

        ClientOptions options = getOptions();
        return options instanceof ClusterClientOptions ? (ClusterClientOptions) options : null;
    }

}