org.springframework.data.redis.core.script.DefaultRedisScript Java Examples

static RedisScript getScript(ScriptType type) {
  Resource resource = new ClassPathResource(type.getPath());
  DefaultRedisScript script = new DefaultRedisScript();
  script.setLocation(resource);
  switch (type) {
    case ADD_MESSAGE:
    case MOVE_MESSAGE:
    case PUSH_MESSAGE:
    case MOVE_MESSAGE_LIST_TO_LIST:
    case MOVE_MESSAGE_LIST_TO_ZSET:
    case MOVE_MESSAGE_ZSET_TO_ZSET:
    case MOVE_MESSAGE_ZSET_TO_LIST:
      script.setResultType(Long.class);
      return script;
    case POP_MESSAGE:
      script.setResultType(RqueueMessage.class);
      return script;
    default:
      throw new UnknownSwitchCase(type.toString());
  }
}
 

/**
 * 加锁
 *
 * @param lockKey       key
 * @param requestId     随机请求id
 * @param expireSecond  超时秒
 * @param loopTimes     循环次数
 * @param sleepInterval 等待间隔（毫秒）
 * @return
 */
public static boolean tryGetDistributedLock(String lockKey, String requestId, Integer expireSecond, Long loopTimes, Long sleepInterval) {
    DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>(LOCK_SCRIPT_STR, Long.class);
    while (loopTimes-- >= 0) {
        Object result = redisTemplate.execute(
                (RedisConnection connection) -> connection.eval(
                        redisScript.getScriptAsString().getBytes(),
                        ReturnType.INTEGER,
                        1,
                        lockKey.getBytes(),
                        requestId.getBytes(),
                        String.valueOf(expireSecond).getBytes()
                )
        );
        if (SUCCESS.equals(result)) {
            return true;
        }
        try {
            TimeUnit.MILLISECONDS.sleep(sleepInterval);
        } catch (InterruptedException e) {
            log.error("e message: {}", e.getMessage());
            Thread.currentThread().interrupt();
        }
    }
    return false;
}
 

private void initLuaScript() {
    // lock script
    lockScript = new DefaultRedisScript<Long>();
    lockScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("lock/lock.lua")));
    lockScript.setResultType(Long.class);
    LOGGER.debug("init lock lua script success:{}", lockScript.getScriptAsString());
    // unlock script
    unlockScript = new DefaultRedisScript<Long>();
    unlockScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("lock/unlock.lua")));
    unlockScript.setResultType(Long.class);
    LOGGER.debug("init release lua script success:{}", unlockScript.getScriptAsString());
    // renew script
    renewScript = new DefaultRedisScript<Long>();
    renewScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("lock/renew.lua")));
    renewScript.setResultType(Long.class);
    LOGGER.debug("init renew lua script success:{}", renewScript.getScriptAsString());
}
 

/**
 * get Redis Script
 *
 * @return RedisScript
 */
RedisScript<List<Long>> getRedisScript() {
    DefaultRedisScript redisScript = new DefaultRedisScript<>();
    redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource(QPS_LUA_PATH)));
    redisScript.setResultType(List.class);
    return redisScript;
}
 

@Bean
@SuppressWarnings("unchecked")
public RedisScript redisRequestRateLimiterScript() {
	DefaultRedisScript redisScript = new DefaultRedisScript<>();
	redisScript.setScriptSource(new ResourceScriptSource(
			new ClassPathResource("META-INF/scripts/request_rate_limiter.lua")));
	redisScript.setResultType(List.class);
	return redisScript;
}
 

@Bean
@SuppressWarnings("unchecked")
public RedisScript<Long> limitRedisScript() {
    DefaultRedisScript redisScript = new DefaultRedisScript<>();
    redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("scripts/redis/limit.lua")));
    redisScript.setResultType(Long.class);
    return redisScript;
}
 

@Around("pointcut()")
public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
	HttpServletRequest request = RequestHolder.getHttpServletRequest();
	MethodSignature signature = (MethodSignature) joinPoint.getSignature();
	Method signatureMethod = signature.getMethod();
	Limit limit = signatureMethod.getAnnotation(Limit.class);
	LimitType limitType = limit.limitType();
	String key = limit.key();
	if (StringUtil.isEmpty(key)) {
		if (limitType == LimitType.IP) {
			key = WebUtil.getIp(request);
		} else {
			key = signatureMethod.getName();
		}
	}

	ImmutableList<Object> keys = ImmutableList.of(StringUtil.join(limit.prefix(), "_", key, "_", request.getRequestURI().replaceAll(StringUtil.SLASH, "_")));

	String luaScript = buildLuaScript();
	RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
	Number count = redisTemplate.execute(redisScript, keys, limit.count(), limit.period());
	if (null != count && count.intValue() <= limit.count()) {
		logger.info("第{}次访问key为 {}，描述为 [{}] 的接口", count, keys, limit.name());
		return joinPoint.proceed();
	} else {
		throw new BadRequestException("访问次数受限制");
	}
}
 

@PostConstruct
public void init() {
    getRedisScript = new DefaultRedisScript<>();
    getRedisScript.setResultType(Long.class);
    getRedisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("rateLimter.lua")));
    LOGGER.info("RateLimterHandler[分布式限流处理器]脚本加载完成");
}
 

@Override
    public LimitResult tryAccess(LimitEntity limitEntity) {
        String identifier = limitEntity.getIdentifier();
//        if (StringUtils.isEmpty( IdentifierThreadLocal.get() )) {
//            identifier = IdentifierThreadLocal.get();
//        }
//        if (StringUtils.isEmpty( identifier )) {
//            throw new LimitException( "identifier cannot be null" );
//        }
        String key =KeyUtil.getKey(limitEntity);
        if (StringUtils.isEmpty(key)) {
            return null;
        }

        int seconds = limitEntity.getSeconds();
        int limitCount = limitEntity.getLimtNum();
       // String compositeKey = KeyUtil.compositeKey( limitEntity.getIdentifier()==null?"nobody":limitEntity.getIdentifier(), limitEntity.getKey() );
        List<String> keys = new ArrayList<>();
        keys.add( key );
        String luaScript = buildLuaScript();
        RedisScript<Long> redisScript = new DefaultRedisScript<>( luaScript, Long.class );
        Long count = stringRedisTemplate.execute( redisScript, keys, "" + limitCount, "" + seconds );
        log.info( "Access try count is {} for key={}", count, key );
//        return count != 0;
        LimitResult result=new LimitResult();
        result.setUrl(key);
        result.setIdenfier(identifier);
        if(count!=0){
           result.setResultType(LimitResult.ResultType.SUCCESS);
        }else {
            result.setResultType(LimitResult.ResultType.FAIL);
        }
        return result;
    }
 

@Around("pointcut()")
public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
    HttpServletRequest request = RequestHolder.getHttpServletRequest();
    MethodSignature signature = (MethodSignature) joinPoint.getSignature();
    Method signatureMethod = signature.getMethod();
    Limit limit = signatureMethod.getAnnotation(Limit.class);
    LimitType limitType = limit.limitType();
    String key = limit.key();
    if (StringUtils.isEmpty(key)) {
        if (limitType == LimitType.IP) {
            key = StringUtils.getIp(request);
        } else {
            key = signatureMethod.getName();
        }
    }

    ImmutableList<Object> keys = ImmutableList.of(StringUtils.join(limit.prefix(), "_", key, "_", request.getRequestURI().replaceAll("/","_")));

    String luaScript = buildLuaScript();
    RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
    Number count = redisTemplate.execute(redisScript, keys, limit.count(), limit.period());
    if (null != count && count.intValue() <= limit.count()) {
        logger.info("第{}次访问key为 {}，描述为 [{}] 的接口", count, keys, limit.name());
        return joinPoint.proceed();
    } else {
        throw new BadRequestException("访问次数受限制");
    }
}
 

/**
 * 解锁
 */
public boolean releaseLock(String key, String value) {
    String script = "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end";

    RedisScript<Long> redisScript = new DefaultRedisScript<>(script, Long.class);
    Long result = redisTemplate.execute(redisScript, Lists.newArrayList(key), value);

    return SUCCESS.equals(result);
}
 

/**
 * 获取锁
 */
public boolean getLock(String key, String value, int expireTime) {
    String script = "if redis.call('setNX',KEYS[1],ARGV[1]) == 1 then if redis.call('get',KEYS[1])==ARGV[1] then return redis.call('expire',KEYS[1],ARGV[2]) else return 0 end end";

    RedisScript<Long> redisScript = new DefaultRedisScript<>(script, Long.class);
    Long result = redisTemplate.execute(redisScript, Lists.newArrayList(key), value, expireTime);

    return SUCCESS.equals(result);
}
 

/**
 * 读取限流脚本
 *
 * @return
 */
@Bean
public DefaultRedisScript<Number> redisluaScript() {
    DefaultRedisScript<Number> redisScript = new DefaultRedisScript<>();
    redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("redisLimit.lua")));
    //返回类型
    redisScript.setResultType(Number.class);
    return redisScript;
}
 

@Around("pointcut()")
public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
    HttpServletRequest request = RequestHolder.getHttpServletRequest();
    MethodSignature signature = (MethodSignature) joinPoint.getSignature();
    Method signatureMethod = signature.getMethod();
    Limit limit = signatureMethod.getAnnotation(Limit.class);
    LimitType limitType = limit.limitType();
    String key = limit.key();
    if (StringUtils.isEmpty(key)) {
        if (limitType == LimitType.IP) {
            key = StringUtils.getIp(request);
        } else {
            key = signatureMethod.getName();
        }
    }

    ImmutableList<Object> keys = ImmutableList.of(StringUtils.join(limit.prefix(), "_", key, "_", request.getRequestURI().replaceAll("/","_")));

    String luaScript = buildLuaScript();
    RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
    Number count = redisTemplate.execute(redisScript, keys, limit.count(), limit.period());
    if (null != count && count.intValue() <= limit.count()) {
        logger.info("第{}次访问key为 {}，描述为 [{}] 的接口", count, keys, limit.name());
        return joinPoint.proceed();
    } else {
        throw new BadRequestException("访问次数受限制");
    }
}
 

/**
 * 释放锁
 *
 * @param lockKey   key
 * @param requestId 加锁的请求id
 * @return
 */
public static boolean releaseDistributedLock(String lockKey, String requestId) {
    DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>(UNLOCK_SCRIPT_STR, Long.class);
    Object result = redisTemplate.execute(
            (RedisConnection connection) -> connection.eval(
                    redisScript.getScriptAsString().getBytes(),
                    ReturnType.INTEGER,
                    1,
                    lockKey.getBytes(),
                    requestId.getBytes()
            )
    );
    return SUCCESS.equals(result);
}
 

@Bean
public DefaultRedisScript<Boolean> setScript() {
    DefaultRedisScript<Boolean> defaultRedisScript = new DefaultRedisScript<>();
    defaultRedisScript.setScriptText(Script.SET.getScript());
    defaultRedisScript.setResultType(Boolean.class);
    return defaultRedisScript;
}
 

@Bean
public DefaultRedisScript<Boolean> setNxScript() {
    DefaultRedisScript<Boolean> defaultRedisScript = new DefaultRedisScript<>();
    defaultRedisScript.setScriptText(Script.SET_NX.getScript());
    defaultRedisScript.setResultType(Boolean.class);
    return defaultRedisScript;
}
 

ToolsCacheManager(ToolsCache toolsCache, DefaultRedisScript<Boolean> setNxScript, DefaultRedisScript<Boolean> setScript,
                  RedisCache redisCache, CaffeineCache caffeineCache) {
    super();
    this.toolsCache = toolsCache;
    this.dynamic = toolsCache.isDynamic();
    this.cacheNames = toolsCache.getCacheNames();
    this.setNxScript = setNxScript;
    this.setScript = setScript;
    this.redisCache = redisCache;
    this.caffeineCache = caffeineCache;
}
 

ToolsCacheAdapter(String name, RedisTemplate<Object, Object> redisTemplate, Cache<Object, Object> caffeineCache, ToolsCache toolsCache,
                  DefaultRedisScript<Boolean> setScript, DefaultRedisScript<Boolean> setNxScript, RedisCache redisCache) {
    super(toolsCache.isCacheValueNullable());
    this.name = name;
    this.redisTemplate = redisTemplate;
    this.caffeineCache = caffeineCache;
    this.cachePrefix = toolsCache.getCachePrefix();
    this.expire = redisCache.getExpire();
    this.everyCacheExpire = redisCache.getEveryCacheExpire();
    this.topic = redisCache.getTopic();
    this.setScript = setScript;
    this.setNxScript = setNxScript;
}
 

@Test
public void redisLuaTest() {
    String script = "return redis.call('set',KEYS[1],ARGV[1])";
    DefaultRedisScript<String> redisScript = new DefaultRedisScript<>(script);
    redisScript.setResultType(String.class);
    List<String> keys = Arrays.asList("apple");
    String result = stringRedisTemplate.execute(redisScript, keys, "1");
    Assert.assertEquals("OK", result);
}
 

@Test
public void redisLuaTest() {
    String script = "return redis.call('set',KEYS[1],ARGV[1])";
    DefaultRedisScript<String> redisScript = new DefaultRedisScript<>(script);
    redisScript.setResultType(String.class);
    List<String> keys = Arrays.asList("apple");
    String result = stringRedisTemplate.execute(redisScript, keys, "1");
    Assert.assertEquals("OK", result);
}
 

@SuppressWarnings("unchecked")
private RedisScript<List<Long>> redisRateLimiterScript() {
	DefaultRedisScript redisScript = new DefaultRedisScript<>();
	redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("META-INF/scripts/mica_rate_limiter.lua")));
	redisScript.setResultType(List.class);
	return redisScript;
}
 

@Override
public void onStart(List<BaseRepository> repositoryList) {
    if (repositoryList == null)
        return;

    long startTme = System.currentTimeMillis();
    logger.info("x7-repo/x7-id-generator starting.... \n");

    final String idGeneratorScript = "local hk = KEYS[1] " +
            "local key = KEYS[2] " +
            "local id = ARGV[1] " +
            "local existId = redis.call('hget',hk,key) " +
            "if tonumber(id) > tonumber(existId) " +
            "then " +
            "redis.call('hset',hk,key,tostring(id)) " +
            "return tonumber(id) "+
            "end " +
            "return tonumber(existId)";

    RedisScript<Long> redisScript = new DefaultRedisScript<Long>() {

        @Override
        public String getSha1(){
            return VerifyUtil.toMD5("id_map_key");
        }

        @Override
        public Class<Long> getResultType() {
            return Long.class;
        }

        @Override
        public String getScriptAsString() {
            return  idGeneratorScript;
        }
    };

    for (BaseRepository baseRepository : repositoryList) {
        CriteriaBuilder.ResultMappedBuilder builder = CriteriaBuilder.buildResultMapped();
        Class clzz = baseRepository.getClz();
        Parsed parsed = Parser.get(clzz);
        String key = parsed.getKey(X.KEY_ONE);
        BeanElement be = parsed.getElement(key);
        if (be.clz == String.class)
            continue;
        builder.reduce(ReduceType.MAX, be.property).paged().ignoreTotalRows();
        Criteria.ResultMappedCriteria resultMappedCriteria = builder.get();

        List<Long> idList = baseRepository.listPlainValue(Long.class,resultMappedCriteria);
        Long maxId = idList.stream().filter(id -> id != null).findFirst().orElse(0L);
        String name = baseRepository.getClz().getName();

        logger.info("Db    : " + name + ".maxId = " + maxId);

        List<String> keys = Arrays.asList(IdGeneratorPolicy.ID_MAP_KEY,name);
        long result = this.stringRedisTemplate.execute(redisScript,keys,String.valueOf(maxId));

        logger.info("Redis : " + name + ".maxId = " + result);

    }
    logger.info("..................................................");
    long endTime = System.currentTimeMillis();
    logger.info("x7-repo/x7-id-generator started, cost time: " + (endTime-startTme) +"ms\n\n");
}
 

private RedisScript<Long> getScript() {
    DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
    redisScript.setLocation(new ClassPathResource("/scripts/ratelimit.lua"));
    redisScript.setResultType(Long.class);
    return redisScript;
}