Redis分布式锁的正确实现 --集群版

为了确保分布式锁可用，我们至少要确保锁的实现同时满足以下四个条件：
1、互斥性。在任意时刻，只有一个客户端能持有锁。
2、不会发生死锁。即使有一个客户端在持有锁的期间崩溃而没有主动解锁，也能保证后续其他客户端能加锁。
3、具有容错性。只要大部分的Redis节点正常运行，客户端就可以加锁和解锁。
4、解铃还须系铃人。加锁和解锁必须是同一个客户端，客户端自己不能把别人加的锁给解了。

package com.hz.tgb.data.redis.lock;

import cn.hutool.core.util.IdUtil;
import com.hz.tgb.entity.Book;
import com.hz.tgb.spring.SpringUtils;
import io.lettuce.core.ScriptOutputType;
import io.lettuce.core.SetArgs;
import io.lettuce.core.api.async.RedisAsyncCommands;
import io.lettuce.core.cluster.api.async.RedisAdvancedClusterAsyncCommands;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.data.redis.core.RedisCallback;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.TimeoutUtils;
import org.springframework.data.redis.serializer.RedisSerializer;
import org.springframework.stereotype.Component;

import java.util.Objects;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.LockSupport;
import java.util.function.Supplier;

/**
 * Redis分布式锁 - 集群版
 *
 * @author hezhao on 2019.11.13
 */
@Component
public class RedisClusterLockUtil {

    /*
    为了确保分布式锁可用，我们至少要确保锁的实现同时满足以下四个条件：
    1、互斥性。在任意时刻，只有一个客户端能持有锁。
    2、不会发生死锁。即使有一个客户端在持有锁的期间崩溃而没有主动解锁，也能保证后续其他客户端能加锁。
    3、具有容错性。只要大部分的Redis节点正常运行，客户端就可以加锁和解锁。
    4、解铃还须系铃人。加锁和解锁必须是同一个客户端，客户端自己不能把别人加的锁给解了。
     */

    private static final Logger logger = LoggerFactory.getLogger(RedisLockUtil.class);

    private static RedisTemplate<String, Object> cacheTemplate;

    /** OK: Redis操作是否成功 */
    private static final String REDIS_OK = "OK";
    /** CONN_NOT_FOUND: Redis链接类型不匹配 */
    private static final String REDIS_CONN_NOT_FOUND = "CONN_NOT_FOUND";
    /** 解锁是否成功 */
    private static final Long RELEASE_SUCCESS = 1L;

    /** 解锁Lua脚本 */
    private static final String UNLOCK_LUA_SCRIPT = "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end";

    /**
     * The number of nanoseconds for which it is faster to spin
     * rather than to use timed park. A rough estimate suffices
     * to improve responsiveness with very short timeouts.
     */
    private static final long spinForTimeoutThreshold = 1000000L;

    /**
     * 加锁
     * @param lockKey 锁键
     * @param requestId 请求唯一标识
     * @param expireTime 缓存过期时间
     * @param unit 时间单位
     * @return true: 加锁成功, false: 加锁失败
     */
    @SuppressWarnings("all")
    public static boolean lock(String lockKey, String requestId, long expireTime, TimeUnit unit) {
        // 加锁和设置过期时间必须是原子操作，否则在高并发情况下或者Redis突然崩溃会导致数据错误。
        try {
            // 以毫秒作为过期时间
            long millisecond = TimeoutUtils.toMillis(expireTime, unit);
            String result = execute(connection -> {
                Object nativeConnection = connection.getNativeConnection();
                RedisSerializer<Object> keySerializer = (RedisSerializer<Object>) getRedisTemplate().getKeySerializer();
                RedisSerializer<Object> valueSerializer = (RedisSerializer<Object>) getRedisTemplate().getValueSerializer();
                // springboot 2.0以上的spring-data-redis 包默认使用 lettuce连接包
                // lettuce连接包下序列化键值，否知无法用默认的ByteArrayCodec解析
                byte[] keyByte = keySerializer.serialize(lockKey);
                byte[] valueByte = valueSerializer.serialize(requestId);
                //lettuce连接包，单机模式，ex为秒，px为毫秒
                if (nativeConnection instanceof RedisAsyncCommands) {
                    RedisAsyncCommands commands = (RedisAsyncCommands)nativeConnection;
                    // 同步方法执行、setnx禁止异步
                    return commands.getStatefulConnection().sync().set(keyByte, valueByte, SetArgs.Builder.nx().px(millisecond));
                } else if (nativeConnection instanceof RedisAdvancedClusterAsyncCommands) {
                    // lettuce连接包，集群模式，ex为秒，px为毫秒
                    RedisAdvancedClusterAsyncCommands clusterAsyncCommands = (RedisAdvancedClusterAsyncCommands) nativeConnection;
                    return clusterAsyncCommands.getStatefulConnection().sync().set(keyByte, valueByte, SetArgs.Builder.nx().px(millisecond));
                }
                return REDIS_CONN_NOT_FOUND;
            });

            // 如果链接类型匹配不上，使用默认加锁方法
            if (Objects.equals(result, REDIS_CONN_NOT_FOUND)) {
                return getRedisTemplate().opsForValue().setIfAbsent(lockKey, requestId)
                        && getRedisTemplate().expire(lockKey, expireTime, unit);
            }

            return REDIS_OK.equals(result);
        } catch (Exception e) {
            logger.error("RedisLockUtil lock 加锁失败", e);
        }
        return false;
    }

    /**
     * 解锁
     * @param lockKey 锁键
     * @param requestId 请求唯一标识
     * @return true: 解锁成功, false: 解锁失败
     */
    @SuppressWarnings("all")
    public static boolean unLock(String lockKey, String requestId) {
        try {
            // 使用Lua脚本实现解锁的原子性，如果requestId相等则解锁
            Object result = execute(connection -> {
                Object nativeConnection = connection.getNativeConnection();
                RedisSerializer<Object> keySerializer = (RedisSerializer<Object>) getRedisTemplate().getKeySerializer();
                RedisSerializer<Object> valueSerializer = (RedisSerializer<Object>) getRedisTemplate().getValueSerializer();
                // springboot 2.0以上的spring-data-redis 包默认使用 lettuce连接包
                // lettuce连接包下序列化键值，否知无法用默认的ByteArrayCodec解析
                byte[] keyByte = keySerializer.serialize(lockKey);
                byte[] valueByte = valueSerializer.serialize(requestId);
                //lettuce连接包，单机模式
                if (nativeConnection instanceof RedisAsyncCommands) {
                    RedisAsyncCommands commands = (RedisAsyncCommands)nativeConnection;
                    // 同步方法执行、setnx禁止异步
                    byte[][] keys = {keyByte};
                    byte[][] values = {valueByte};
                    return commands.getStatefulConnection().sync().eval(UNLOCK_LUA_SCRIPT, ScriptOutputType.INTEGER, keys , values);
                } else if (nativeConnection instanceof RedisAdvancedClusterAsyncCommands) {
                    // lettuce连接包，集群模式
                    RedisAdvancedClusterAsyncCommands clusterAsyncCommands = (RedisAdvancedClusterAsyncCommands) nativeConnection;
                    byte[][] keys = {keyByte};
                    byte[][] values = {valueByte};
                    return clusterAsyncCommands.getStatefulConnection().sync().eval(UNLOCK_LUA_SCRIPT, ScriptOutputType.INTEGER, keys , values);
                }
                return REDIS_CONN_NOT_FOUND;
            });

            // 如果链接类型匹配不上，使用默认解锁方法
            if (Objects.equals(result, REDIS_CONN_NOT_FOUND)) {
                return getRedisTemplate().delete(lockKey);
            }

            return Objects.equals(RELEASE_SUCCESS, result);
        } catch (Exception e) {
            logger.error("RedisLockUtil unLock 解锁失败", e);
        }
        return false;
    }

    /**
     * 阻塞锁，拿到锁后执行业务逻辑。注意：超时返回null，程序会继续往下执行
     * @param callback 业务处理逻辑，入参默认为NULL
     * @param lockKey 锁键
     * @param timeout 超时时长, 缓存过期时间默认等于超时时长
     * @param unit 时间单位
     * @return R
     */
    public static <R> R tryLock(Supplier<R> callback, String lockKey, long timeout, TimeUnit unit) {
        return tryLock(callback, lockKey, IdUtil.fastSimpleUUID(), timeout, timeout, unit, TimeOutProcess.DEFAULT);
    }

    /**
     * 阻塞锁，拿到锁后执行业务逻辑。注意：超时会抛出异常
     * @param callback 业务处理逻辑，入参默认为NULL
     * @param lockKey 锁键
     * @param timeout 超时时长, 缓存过期时间默认等于超时时长
     * @param unit 时间单位
     * @return R
     */
    public static <R> R tryLockTimeout(Supplier<R> callback, String lockKey, long timeout, TimeUnit unit) {
        return tryLock(callback, lockKey, IdUtil.fastSimpleUUID(), timeout, timeout, unit, TimeOutProcess.THROW_EXCEPTION);
    }

    /**
     * 阻塞锁，拿到锁后执行业务逻辑。注意：超时会给予补偿，即处理正常逻辑
     * @param callback 业务处理逻辑，入参默认为NULL
     * @param lockKey 锁键
     * @param timeout 超时时长, 缓存过期时间默认等于超时时长
     * @param unit 时间单位
     * @return R
     */
    public static <R> R tryLockCompensate(Supplier<R> callback, String lockKey, long timeout, TimeUnit unit) {
        return tryLock(callback, lockKey, IdUtil.fastSimpleUUID(), timeout, timeout, unit, TimeOutProcess.CARRY_ON);
    }

    /**
     * 阻塞锁，拿到锁后执行业务逻辑
     * @param callback 业务处理逻辑
     * @param lockKey 锁键
     * @param requestId 请求唯一标识
     * @param timeout 超时时长
     * @param expireTime 缓存过期时间
     * @param unit 时间单位
     * @param timeoutProceed 超时处理逻辑
     * @return R
     */
    public static <R> R tryLock(Supplier<R> callback, String lockKey, String requestId,
                                long timeout, long expireTime, TimeUnit unit, TimeOutProcess timeoutProceed) {
        boolean lockFlag = false;
        try {
            lockFlag = tryLock(lockKey, requestId, timeout, expireTime, unit);
            if(lockFlag){
                return callback.get();
            }
        } finally {
            if (lockFlag){
                unLock(lockKey, requestId);
            }
        }
        if (timeoutProceed == null) {
            return null;
        }
        if (timeoutProceed == TimeOutProcess.THROW_EXCEPTION) {
            throw new RedisLockTimeOutException();
        }
        if (timeoutProceed == TimeOutProcess.CARRY_ON) {
            return callback.get();
        }
        return null;
    }

    /**
     * 阻塞锁
     * @param lockKey 锁键
     * @param requestId 请求唯一标识
     * @param timeout 超时时长, 缓存过期时间默认等于超时时长
     * @param unit 时间单位
     * @return true: 加锁成功, false: 加锁失败
     */
    public static boolean tryLock(String lockKey, String requestId, long timeout, TimeUnit unit) {
        return tryLock(lockKey, requestId, timeout, timeout, unit);
    }

    /**
     * 阻塞锁
     * @param lockKey 锁键
     * @param requestId 请求唯一标识
     * @param timeout 超时时长
     * @param expireTime 缓存过期时间
     * @param unit 时间单位
     * @return true: 加锁成功, false: 加锁失败
     */
    public static boolean tryLock(String lockKey, String requestId, long timeout, long expireTime, TimeUnit unit) {
        long nanosTimeout = unit.toNanos(timeout);
        if (nanosTimeout <= 0L) {
            return false;
        }
        final long deadline = System.nanoTime() + nanosTimeout;
        for (;;) {
            // 获取到锁
            if (lock(lockKey, requestId, expireTime, unit)) {
                return true;
            }
            // 判断是否需要继续阻塞， 如果已超时则返回false
            nanosTimeout = deadline - System.nanoTime();
            if (nanosTimeout <= 0L) {
                return false;
            }
            // 休眠1毫秒
            if (nanosTimeout > spinForTimeoutThreshold) {
                LockSupport.parkNanos(spinForTimeoutThreshold);
            }
        }
    }

    public static <T> T execute(RedisCallback<T> action) {
        return getRedisTemplate().execute(action);
    }

    public static RedisTemplate<String, Object> getRedisTemplate() {
        if (cacheTemplate == null) {
            cacheTemplate = SpringUtils.getBean("redisTemplate", RedisTemplate.class);
        }
        return cacheTemplate;
    }

    public static void main(String[] args) {
        Book param = new Book();
        param.setBookId(1234);
        param.setName("西游记");

        Boolean flag = tryLock(() -> {
            int bookId = param.getBookId();
            System.out.println(bookId);
            // TODO ...
            return true;
        }, "BOOK-" + param.getBookId(), 3, TimeUnit.SECONDS);

        System.out.println(flag);
    }

    /**
     * 超时处理逻辑
     */
    public enum TimeOutProcess {
        /** 默认，超时返回null，程序会继续往下执行 */
        DEFAULT,
        /** 超时会抛出异常 */
        THROW_EXCEPTION,
        /** 超时会给予补偿，即处理正常逻辑 */
        CARRY_ON,

    }

}