com.lambdaworks.redis.cluster.RedisAdvancedClusterAsyncCommandsImpl Maven / Gradle / Ivy

Go to download
/*
 * Copyright 2011-2016 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 com.lambdaworks.redis.cluster;

import static com.lambdaworks.redis.cluster.ClusterScanSupport.asyncClusterKeyScanCursorMapper;
import static com.lambdaworks.redis.cluster.ClusterScanSupport.asyncClusterStreamScanCursorMapper;
import static com.lambdaworks.redis.cluster.models.partitions.RedisClusterNode.NodeFlag.MASTER;

import java.lang.reflect.Proxy;
import java.util.*;
import java.util.function.BiFunction;
import java.util.function.Function;
import java.util.function.Predicate;

import com.lambdaworks.redis.*;
import com.lambdaworks.redis.api.StatefulRedisConnection;
import com.lambdaworks.redis.api.async.RedisAsyncCommands;
import com.lambdaworks.redis.api.async.RedisKeyAsyncCommands;
import com.lambdaworks.redis.api.async.RedisScriptingAsyncCommands;
import com.lambdaworks.redis.api.async.RedisServerAsyncCommands;
import com.lambdaworks.redis.cluster.ClusterScanSupport.ScanCursorMapper;
import com.lambdaworks.redis.cluster.api.NodeSelectionSupport;
import com.lambdaworks.redis.cluster.api.StatefulRedisClusterConnection;
import com.lambdaworks.redis.cluster.api.async.AsyncNodeSelection;
import com.lambdaworks.redis.cluster.api.async.NodeSelectionAsyncCommands;
import com.lambdaworks.redis.cluster.api.async.RedisAdvancedClusterAsyncCommands;
import com.lambdaworks.redis.cluster.api.async.RedisClusterAsyncCommands;
import com.lambdaworks.redis.cluster.models.partitions.Partitions;
import com.lambdaworks.redis.cluster.models.partitions.RedisClusterNode;
import com.lambdaworks.redis.codec.RedisCodec;
import com.lambdaworks.redis.output.IntegerOutput;
import com.lambdaworks.redis.output.KeyStreamingChannel;
import com.lambdaworks.redis.output.KeyValueStreamingChannel;
import com.lambdaworks.redis.output.ValueStreamingChannel;
import com.lambdaworks.redis.protocol.AsyncCommand;
import com.lambdaworks.redis.protocol.Command;
import com.lambdaworks.redis.protocol.CommandType;

/**
 * An advanced asynchronous and thread-safe API for a Redis Cluster connection.
 * 
 * @author Mark Paluch
 * @since 3.3
 */
@SuppressWarnings("unchecked")
public class RedisAdvancedClusterAsyncCommandsImpl extends AbstractRedisAsyncCommands implements
         RedisAdvancedClusterAsyncCommands {

    private final Random random = new Random();

    /**
     * Initialize a new connection.
     *
     * @param connection the stateful connection
     * @param codec Codec used to encode/decode keys and values.
     */
    public RedisAdvancedClusterAsyncCommandsImpl(StatefulRedisClusterConnectionImpl connection, RedisCodec codec) {
        super(connection, codec);
    }

    @Override
    public RedisFuture del(K... keys) {
        return del(Arrays.asList(keys));
    }

    @Override
    public RedisFuture del(Iterable keys) {
        Map> partitioned = SlotHash.partition(codec, keys);

        if (partitioned.size() < 2) {
            return super.del(keys);
        }

        Map> executions = new HashMap<>();

        for (Map.Entry> entry : partitioned.entrySet()) {
            RedisFuture del = super.del(entry.getValue());
            executions.put(entry.getKey(), del);
        }

        return MultiNodeExecution.aggregateAsync(executions);
    }

    @Override
    public RedisFuture unlink(K... keys) {
        return unlink(Arrays.asList(keys));
    }

    @Override
    public RedisFuture unlink(Iterable keys) {
        Map> partitioned = SlotHash.partition(codec, keys);

        if (partitioned.size() < 2) {
            return super.unlink(keys);
        }

        Map> executions = new HashMap<>();

        for (Map.Entry> entry : partitioned.entrySet()) {
            RedisFuture unlink = super.unlink(entry.getValue());
            executions.put(entry.getKey(), unlink);
        }

        return MultiNodeExecution.aggregateAsync(executions);
    }

    @Override
    public RedisFuture exists(K... keys) {
        return exists(Arrays.asList(keys));
    }

    public RedisFuture exists(Iterable keys) {
        Map> partitioned = SlotHash.partition(codec, keys);

        if (partitioned.size() < 2) {
            return super.exists(keys);
        }

        Map> executions = new HashMap<>();

        for (Map.Entry> entry : partitioned.entrySet()) {
            RedisFuture exists = super.exists(entry.getValue());
            executions.put(entry.getKey(), exists);
        }

        return MultiNodeExecution.aggregateAsync(executions);
    }

    @Override
    public RedisFuture>> mget(K... keys) {
        return mget(Arrays.asList(keys));
    }

    @Override
    public RedisFuture>> mget(Iterable keys) {
        Map> partitioned = SlotHash.partition(codec, keys);

        if (partitioned.size() < 2) {
            return super.mget(keys);
        }

        Map slots = SlotHash.getSlots(partitioned);
        Map>>> executions = new HashMap<>();

        for (Map.Entry> entry : partitioned.entrySet()) {
            RedisFuture>> mget = super.mget(entry.getValue());
            executions.put(entry.getKey(), mget);
        }

        // restore order of key
        return new PipelinedRedisFuture<>(executions, objectPipelinedRedisFuture -> {
            List> result = new ArrayList<>();
            for (K opKey : keys) {
                int slot = slots.get(opKey);

                int position = partitioned.get(slot).indexOf(opKey);
                RedisFuture>> listRedisFuture = executions.get(slot);
                result.add(MultiNodeExecution.execute(() -> listRedisFuture.get().get(position)));
            }

            return result;
        });
    }

    @Override
    public RedisFuture mget(KeyValueStreamingChannel channel, K... keys) {
       return mget(channel, Arrays.asList(keys));
    }

    @Override
    public RedisFuture mget(KeyValueStreamingChannel channel, Iterable keys) {
        Map> partitioned = SlotHash.partition(codec, keys);

        if (partitioned.size() < 2) {
            return super.mget(channel, keys);
        }

        Map> executions = new HashMap<>();

        for (Map.Entry> entry : partitioned.entrySet()) {
            RedisFuture del = super.mget(channel, entry.getValue());
            executions.put(entry.getKey(), del);
        }

        return MultiNodeExecution.aggregateAsync(executions);
    }

    @Override
    public RedisFuture mset(Map map) {
        Map> partitioned = SlotHash.partition(codec, map.keySet());

        if (partitioned.size() < 2) {
            return super.mset(map);
        }

        Map> executions = new HashMap<>();

        for (Map.Entry> entry : partitioned.entrySet()) {

            Map op = new HashMap<>();
            entry.getValue().forEach(k -> op.put(k, map.get(k)));

            RedisFuture mset = super.mset(op);
            executions.put(entry.getKey(), mset);
        }

        return MultiNodeExecution.firstOfAsync(executions);
    }

    @Override
    public RedisFuture msetnx(Map map) {
        Map> partitioned = SlotHash.partition(codec, map.keySet());

        if (partitioned.size() < 2) {
            return super.msetnx(map);
        }

        Map> executions = new HashMap<>();

        for (Map.Entry> entry : partitioned.entrySet()) {

            Map op = new HashMap<>();
            entry.getValue().forEach(k -> op.put(k, map.get(k)));

            RedisFuture msetnx = super.msetnx(op);
            executions.put(entry.getKey(), msetnx);
        }

        return new PipelinedRedisFuture<>(executions, objectPipelinedRedisFuture -> {
            
            for (RedisFuture listRedisFuture : executions.values()) {
                Boolean b = MultiNodeExecution.execute(() -> listRedisFuture.get());
                if (b == null ||  !b) {
                    return false;
                }
            }

            return !executions.isEmpty();
        });
    }

    @Override
    public RedisFuture clientSetname(K name) {
        Map> executions = new HashMap<>();

        for (RedisClusterNode redisClusterNode : getStatefulConnection().getPartitions()) {
            StatefulRedisConnection byNodeId = getStatefulConnection().getConnection(redisClusterNode.getNodeId());
            if (byNodeId.isOpen()) {
                executions.put("NodeId: " + redisClusterNode.getNodeId(), byNodeId.async().clientSetname(name));
            }

            RedisURI uri = redisClusterNode.getUri();
            StatefulRedisConnection byHost = getStatefulConnection().getConnection(uri.getHost(), uri.getPort());
            if (byHost.isOpen()) {
                executions.put("HostAndPort: " + redisClusterNode.getNodeId(), byHost.async().clientSetname(name));
            }
        }

        return MultiNodeExecution.firstOfAsync(executions);
    }

    @Override
    public RedisFuture> clusterGetKeysInSlot(int slot, int count) {
        RedisClusterAsyncCommands connectionBySlot = findConnectionBySlot(slot);

        if (connectionBySlot != null) {
            return connectionBySlot.clusterGetKeysInSlot(slot, count);
        }

        return super.clusterGetKeysInSlot(slot, count);
    }

    @Override
    public RedisFuture clusterCountKeysInSlot(int slot) {
        RedisClusterAsyncCommands connectionBySlot = findConnectionBySlot(slot);

        if (connectionBySlot != null) {
            return connectionBySlot.clusterCountKeysInSlot(slot);
        }

        return super.clusterCountKeysInSlot(slot);
    }

    @Override
    public RedisFuture dbsize() {
        Map> executions = executeOnMasters(RedisServerAsyncCommands::dbsize);
        return MultiNodeExecution.aggregateAsync(executions);
    }

    @Override
    public RedisFuture flushall() {
        Map> executions = executeOnMasters(RedisServerAsyncCommands::flushall);
        return MultiNodeExecution.firstOfAsync(executions);
    }

    @Override
    public RedisFuture flushdb() {
        Map> executions = executeOnMasters(RedisServerAsyncCommands::flushdb);
        return MultiNodeExecution.firstOfAsync(executions);
    }

    @Override
    public RedisFuture scriptFlush() {
        Map> executions = executeOnNodes(RedisScriptingAsyncCommands::scriptFlush,
                redisClusterNode -> true);
        return MultiNodeExecution.firstOfAsync(executions);
    }

    @Override
    public RedisFuture scriptKill() {
        Map> executions = executeOnNodes(RedisScriptingAsyncCommands::scriptFlush,
                redisClusterNode -> true);
        return MultiNodeExecution.alwaysOkOfAsync(executions);
    }

    @Override
    public RedisFuture randomkey() {

        Partitions partitions = getStatefulConnection().getPartitions();
        int index = random.nextInt(partitions.size());

        RedisClusterAsyncCommands connection = getConnection(partitions.getPartition(index).getNodeId());
        return connection.randomkey();
    }

    @Override
    public RedisFuture> keys(K pattern) {
        Map>> executions = executeOnMasters(commands -> commands.keys(pattern));

        return new PipelinedRedisFuture<>(executions, objectPipelinedRedisFuture -> {
            List result = new ArrayList<>();
            for (RedisFuture> future : executions.values()) {
                result.addAll(MultiNodeExecution.execute(() -> future.get()));
            }
            return result;
        });
    }

    @Override
    public RedisFuture keys(KeyStreamingChannel channel, K pattern) {
        Map> executions = executeOnMasters(commands -> commands.keys(channel, pattern));
        return MultiNodeExecution.aggregateAsync(executions);
    }

    @Override
    public void shutdown(boolean save) {

        executeOnNodes(commands -> {
            commands.shutdown(save);

            Command command = new Command<>(CommandType.SHUTDOWN, new IntegerOutput<>(codec), null);
            AsyncCommand async = new AsyncCommand(command);
            async.complete();
            return async;
        }, redisClusterNode -> true);
    }

    @Override
    public RedisFuture touch(K... keys) {
        return touch(Arrays.asList(keys));
    }

    public RedisFuture touch(Iterable keys) {
        Map> partitioned = SlotHash.partition(codec, keys);

        if (partitioned.size() < 2) {
            return super.touch(keys);
        }

        Map> executions = new HashMap<>();

        for (Map.Entry> entry : partitioned.entrySet()) {
            RedisFuture touch = super.touch(entry.getValue());
            executions.put(entry.getKey(), touch);
        }

        return MultiNodeExecution.aggregateAsync(executions);
    }

    /**
     * Run a command on all available masters,
     *
     * @param function function producing the command
     * @param  result type
     * @return map of a key (counter) and commands.
     */
    protected  Map> executeOnMasters(
            Function, RedisFuture> function) {
        return executeOnNodes(function, redisClusterNode -> redisClusterNode.is(MASTER));
    }

    /**
     * Run a command on all available nodes that match {@code filter}.
     *
     * @param function function producing the command
     * @param filter filter function for the node selection
     * @param  result type
     * @return map of a key (counter) and commands.
     */
    protected  Map> executeOnNodes(
            Function, RedisFuture> function, Function filter) {
        Map> executions = new HashMap<>();

        for (RedisClusterNode redisClusterNode : getStatefulConnection().getPartitions()) {

            if (!filter.apply(redisClusterNode)) {
                continue;
            }

            RedisURI uri = redisClusterNode.getUri();
            StatefulRedisConnection connection = getStatefulConnection().getConnection(uri.getHost(), uri.getPort());
            if (connection.isOpen()) {
                executions.put(redisClusterNode.getNodeId(), function.apply(connection.async()));
            }
        }
        return executions;
    }

    private RedisClusterAsyncCommands findConnectionBySlot(int slot) {
        RedisClusterNode node = getStatefulConnection().getPartitions().getPartitionBySlot(slot);
        if (node != null) {
            return getConnection(node.getUri().getHost(), node.getUri().getPort());
        }

        return null;
    }

    @Override
    public RedisClusterAsyncCommands getConnection(String nodeId) {
        return getStatefulConnection().getConnection(nodeId).async();
    }

    @Override
    public RedisClusterAsyncCommands getConnection(String host, int port) {
        return getStatefulConnection().getConnection(host, port).async();
    }

    @Override
    public StatefulRedisClusterConnection getStatefulConnection() {
        return (StatefulRedisClusterConnection) connection;
    }

    @Override
    public AsyncNodeSelection nodes(Predicate predicate) {
        return nodes(predicate, false);
    }

    @Override
    public AsyncNodeSelection readonly(Predicate predicate) {
        return nodes(predicate, ClusterConnectionProvider.Intent.READ, false);
    }

    @Override
    public AsyncNodeSelection nodes(Predicate predicate, boolean dynamic) {
        return nodes(predicate, ClusterConnectionProvider.Intent.WRITE, dynamic);
    }

    @SuppressWarnings("unchecked")
    protected AsyncNodeSelection nodes(Predicate predicate, ClusterConnectionProvider.Intent intent,
            boolean dynamic) {

        NodeSelectionSupport, ?> selection;

        if (dynamic) {
            selection = new DynamicAsyncNodeSelection<>(getStatefulConnection(), predicate, intent);
        } else {
            selection = new StaticAsyncNodeSelection<>(getStatefulConnection(), predicate, intent);
        }

        NodeSelectionInvocationHandler h = new NodeSelectionInvocationHandler((AbstractNodeSelection) selection);
        return (AsyncNodeSelection) Proxy.newProxyInstance(NodeSelectionSupport.class.getClassLoader(), new Class[] {
                NodeSelectionAsyncCommands.class, AsyncNodeSelection.class }, h);
    }

    @Override
    public RedisFuture> scan() {
        return clusterScan(ScanCursor.INITIAL, (connection, cursor) -> connection.scan(), asyncClusterKeyScanCursorMapper());
    }

    @Override
    public RedisFuture> scan(ScanArgs scanArgs) {
        return clusterScan(ScanCursor.INITIAL, (connection, cursor) -> connection.scan(scanArgs),
                asyncClusterKeyScanCursorMapper());
    }

    @Override
    public RedisFuture> scan(ScanCursor scanCursor, ScanArgs scanArgs) {
        return clusterScan(scanCursor, (connection, cursor) -> connection.scan(cursor, scanArgs),
                asyncClusterKeyScanCursorMapper());
    }

    @Override
    public RedisFuture> scan(ScanCursor scanCursor) {
        return clusterScan(scanCursor, (connection, cursor) -> connection.scan(cursor), asyncClusterKeyScanCursorMapper());
    }

    @Override
    public RedisFuture scan(KeyStreamingChannel channel) {
        return clusterScan(ScanCursor.INITIAL, (connection, cursor) -> connection.scan(channel),
                asyncClusterStreamScanCursorMapper());
    }

    @Override
    public RedisFuture scan(KeyStreamingChannel channel, ScanArgs scanArgs) {
        return clusterScan(ScanCursor.INITIAL, (connection, cursor) -> connection.scan(channel, scanArgs),
                asyncClusterStreamScanCursorMapper());
    }

    @Override
    public RedisFuture scan(KeyStreamingChannel channel, ScanCursor scanCursor, ScanArgs scanArgs) {
        return clusterScan(scanCursor, (connection, cursor) -> connection.scan(channel, cursor, scanArgs),
                asyncClusterStreamScanCursorMapper());
    }

    @Override
    public RedisFuture scan(KeyStreamingChannel channel, ScanCursor scanCursor) {
        return clusterScan(scanCursor, (connection, cursor) -> connection.scan(channel, cursor),
                asyncClusterStreamScanCursorMapper());
    }

    private  RedisFuture clusterScan(ScanCursor cursor,
            BiFunction, ScanCursor, RedisFuture> scanFunction,
            ScanCursorMapper> resultMapper) {

        return clusterScan(getStatefulConnection(), cursor, scanFunction, resultMapper);
    }

    /**
     * Perform a SCAN in the cluster.
     *
     */
    static  RedisFuture clusterScan(StatefulRedisClusterConnection connection,
            ScanCursor cursor, BiFunction, ScanCursor, RedisFuture> scanFunction,
            ScanCursorMapper> mapper) {

        List nodeIds = ClusterScanSupport.getNodeIds(connection, cursor);
        String currentNodeId = ClusterScanSupport.getCurrentNodeId(cursor, nodeIds);
        ScanCursor continuationCursor = ClusterScanSupport.getContinuationCursor(cursor);

        RedisFuture scanCursor = scanFunction.apply(connection.getConnection(currentNodeId).async(), continuationCursor);
        return mapper.map(nodeIds, currentNodeId, scanCursor);
    }

}