RedisLock工具类

RedisLock

package test.utils;

import io.lettuce.core.RedisFuture;
import io.lettuce.core.ScriptOutputType;
import io.lettuce.core.SetArgs;
import io.lettuce.core.api.async.RedisAsyncCommands;
import io.lettuce.core.api.async.RedisScriptingAsyncCommands;
import io.lettuce.core.cluster.api.async.RedisAdvancedClusterAsyncCommands;
import lombok.extern.log4j.Log4j2;
import org.springframework.data.redis.connection.RedisConnection;
import org.springframework.data.redis.core.RedisCallback;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;

import java.nio.charset.StandardCharsets;
import java.util.UUID;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;

@Log4j2
public class RedisLock {

	private StringRedisTemplate redisTemplate;

	/**
	 * 当且仅当 key 不存在时设置 value ，等效于 SETNX
	 */
	public static final String NX = "NX";

	/**
	 * 以秒为单位设置 key 的过期时间，等效于 EXPIRE key seconds
	 */
	public static final String EX = "EX";

	/**
	 * 调用 set 后的返回值
	 */
	public static final String OK = "OK";

	/**
	 * 解锁的lua脚本
	 */
	public static final String UNLOCK_LUA = "if redis.call(\"get\",KEYS[1]) == ARGV[1] then return redis.call(\"del\",KEYS[1]) else return 0 end";

	/**
	 * 锁标志对应的key
	 */
	private String lockKey;

	/**
	 * 锁对应的值
	 */
	private String lockValue;

	/**
	 * 锁的有效时间(s)，默认120
	 */
	private int expireTime = 120;

	/**
	 * 请求锁的超时时间(ms)，默认500
	 */
	private long timeOut = 500;

	/**
	 * 锁标记
	 */
	private volatile boolean locked = false;

	private static final String REDIS_LIB_MISMATCH = "Failed to convert nativeConnection. " +
			"Is your SpringBoot main version > 2.0 ? Only lib:lettuce is supported.";

	/**
	 * 使用默认的锁过期时间和请求锁的超时时间
	 *
	 * @param redisTemplate
	 * @param lockKey 锁的key（Redis的Key）
	 */
	public RedisLock(StringRedisTemplate redisTemplate, String lockKey) {
		this.redisTemplate = redisTemplate;
		this.lockKey = lockKey;
	}

	/**
	 * 锁的过期时间和请求锁的超时时间都是用指定的值
	 *
	 * @param redisTemplate
	 * @param lockKey 锁的key（Redis的Key）
	 * @param expireTime 锁的过期时间(单位：秒)
	 * @param timeOut 请求锁的超时时间(单位：毫秒)
	 */
	public RedisLock(StringRedisTemplate redisTemplate, String lockKey, int expireTime, long timeOut) {
		this(redisTemplate, lockKey);
		this.expireTime = expireTime;
		this.timeOut = timeOut;
	}

	/**
	 * 尝试获取锁，直到超时返回
	 *
	 * @return
	 */
	public boolean tryLock() {
		lockValue = UUID.randomUUID().toString();
		// 超时时间转为纳秒
		long timeout = timeOut * 1000000;
		long nowTime = System.nanoTime();
		while ((System.nanoTime() - nowTime) < timeout) {
			if (OK.equalsIgnoreCase(this.set(lockKey, lockValue, expireTime))) {
				locked = true;
				return true;
			}

			// 每次请求等待一段时间
			try {
				TimeUnit.MILLISECONDS.sleep(100);
			} catch (InterruptedException e) {
				log.info("Sleep is interrupted", e);
			}
		}
		return locked;
	}

	/**
	 * 尝试获取锁，立即返回
	 *
	 * @return
	 */
	public boolean lock() {
		lockValue = UUID.randomUUID().toString();
		//不存在则添加 且设置过期时间（单位ms）
		String result = set(lockKey, lockValue, expireTime);
		locked = OK.equalsIgnoreCase(result);
		return locked;
	}

	/**
	 * 以阻塞方式的获取锁，直到成功获得锁才返回
	 *
	 * @return
	 */
	public boolean lockBlock() {
		lockValue = UUID.randomUUID().toString();
		while (true) {
			//不存在则添加 且设置过期时间（单位ms）
			String result = set(lockKey, lockValue, expireTime);
			if (OK.equalsIgnoreCase(result)) {
				locked = true;
				return locked;
			}

			try {
				TimeUnit.MILLISECONDS.sleep(100);
			} catch (InterruptedException e) {
				log.info("Sleep is interrupted", e);
			}
		}
	}

	/**
	 * 解锁
	 * <p>
	 * 防止持有过期锁的客户端误删现有锁的情况出现，可以通过以下修改：
	 * <p>
	 * 1. 不使用固定的字符串作为 value，而是使用随机的 UUID 作为 value 。
	 * 2. 不使用 DEL 命令来释放锁，而是发送一个 Lua 脚本，这个脚本只在客户端传入的值和键的口令串相匹配时，才对键进行删除。
	 * 注：EVAL 只在 Redis 2.6.0 及以上版本才支持，低版本会自动改用 DEL 删除key
	 */
	public Boolean unlock() {
		// 只有加锁成功并且锁还有效才去释放锁
		if (locked) {
			try {
				return redisTemplate.execute((RedisConnection connection) -> {
					Object nativeConnection = connection.getNativeConnection();
					Long result = 0L;

					byte[] keyBytes = lockKey.getBytes(StandardCharsets.UTF_8);
					byte[] valueBytes = lockValue.getBytes(StandardCharsets.UTF_8);
					Object[] keyParam = new Object[]{keyBytes};

					if (nativeConnection instanceof RedisScriptingAsyncCommands) {
						RedisScriptingAsyncCommands<Object,byte[]> command = (RedisScriptingAsyncCommands<Object,byte[]>) nativeConnection;
						RedisFuture future = command.eval(UNLOCK_LUA, ScriptOutputType.INTEGER, keyParam, valueBytes);
						result = getEvalResult(future,connection);
					}else{
						log.warn(REDIS_LIB_MISMATCH);
					}

					if (result == 0L && !StringUtils.isEmpty(lockKey)) {
						log.debug("Unlock failed! key={}, time={}", lockKey, System.currentTimeMillis());
					}

					locked = result == 0L;
					return result == 1L;
				});
			} catch (Throwable e) {
				if(log.isDebugEnabled()) {
					log.debug(
							"The redis you are using dose NOT support EVAL. Use downgrade method to unlock. {}",
							e.getMessage());
				}
				String value = this.get(lockKey, String.class);
				if (lockValue.equals(value)) {
					redisTemplate.delete(lockKey);
					return true;
				}
				return false;
			}
		}

		return true;
	}

	private Long getEvalResult(RedisFuture future,RedisConnection connection){
		try {
			Object o = future.get();
			return (Long)o;
		} catch (InterruptedException | ExecutionException e) {
			log.error("Future get failed, trying to close connection.", e);
			closeConnection(connection);
			return 0L;
		}
	}

	/**
	 * 获取锁状态
	 *
	 * @return
	 */
	public boolean isLock() {
		return locked;
	}

	/**
	 * 重写redisTemplate的set方法
	 * <p>
	 * 命令 SET resource-name anystring NX EX max-lock-time 是一种在 Redis 中实现锁的简单方法。
	 * <p>
	 * 客户端执行以上的命令：
	 * <p>
	 * 如果服务器返回 OK ，那么这个客户端获得锁。
	 * 如果服务器返回 NIL ，那么客户端获取锁失败，可以在稍后再重试。
	 *
	 * @param key
	 * @param value
	 * @param expireSeconds
	 * @return
	 */
	private String set(final String key, final String value, final long expireSeconds) {
		Assert.isTrue(!StringUtils.isEmpty(key), "Invalid key");
		return redisTemplate.execute((RedisCallback<String>) connection -> {
			Object nativeConnection = connection.getNativeConnection();
			String result = null;
			byte[] keyByte = key.getBytes(StandardCharsets.UTF_8);
			byte[] valueByte = value.getBytes(StandardCharsets.UTF_8);
			if(nativeConnection instanceof RedisAsyncCommands){
				RedisAsyncCommands command = (RedisAsyncCommands) nativeConnection;
				result = command.getStatefulConnection().sync().set(keyByte, valueByte, SetArgs.Builder.nx().ex(expireSeconds));
			}else if(nativeConnection instanceof RedisAdvancedClusterAsyncCommands){
				RedisAdvancedClusterAsyncCommands clusterAsyncCommands = (RedisAdvancedClusterAsyncCommands) nativeConnection;
				result = clusterAsyncCommands.getStatefulConnection().sync().set(keyByte, valueByte, SetArgs.Builder.nx().ex(expireSeconds));
			}else{
				log.error(REDIS_LIB_MISMATCH);
			}
			return result;
		});
	}

	private void closeConnection(RedisConnection connection){
		try{
			connection.close();
		}catch (Exception e2){
			log.error("close connection fail.", e2);
		}
	}

	/**
	 * 获取redis里面的值
	 *
	 * @param key
	 * @param clazz
	 * @return T
	 */
	private <T> T get(final String key, Class<T> clazz) {
		Assert.isTrue(!StringUtils.isEmpty(key), "Invalid key");
		return redisTemplate.execute((RedisConnection connection) -> {
			Object nativeConnection = connection.getNativeConnection();
			Object result = null;
			byte[] keyByte = key.getBytes(StandardCharsets.UTF_8);

			if(nativeConnection instanceof RedisAsyncCommands){
				RedisAsyncCommands command = (RedisAsyncCommands) nativeConnection;
				result = command.getStatefulConnection().sync().get(keyByte);
			}else if(nativeConnection instanceof RedisAdvancedClusterAsyncCommands){
				RedisAdvancedClusterAsyncCommands clusterAsyncCommands = (RedisAdvancedClusterAsyncCommands) nativeConnection;
				result = clusterAsyncCommands.getStatefulConnection().sync().get(keyByte);
			}
			return clazz.cast(result);
		});
	}

}