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,

    }

}