Redis使用Lua脚本保证原子性

Redis使用同一个Lua解释器来执行所有命令，同时，Redis保证以一种原子性的方式来执行脚本：当lua脚本在执行的时候，不会有其他脚本和命令同时执行，这种语义类似于 MULTI/EXEC。从别的客户端的视角来看，一个lua脚本要么不可见，要么已经执行完。

然而这也意味着，执行一个较慢的lua脚本是不建议的，由于脚本的开销非常低，构造一个快速执行的脚本并非难事。但是你要注意到，当你正在执行一个比较慢的脚本时，所以其他的客户端都无法执行命令。

我们通过lua和java注解实现一个流量限速功能

1.创建一个limit.lua

local c
c = redis.call('get',KEYS[1])
-- 调用不超过最大值，则直接返回
if c and tonumber(c) > tonumber(ARGV[1]) then
    return c;
end
-- 执行计算器自加
c = redis.call('incr',KEYS[1])
if tonumber(c) == 1 then
-- 从第一次调用开始限流，设置对应键值的过期
    redis.call('expire',KEYS[1],ARGV[2])
end
return c;

2.java代码读取lua脚本，自动注入

/**
 * 流量限制脚本
 *
 * @return
 */
@Bean
public DefaultRedisScript<Long> limitScript() {
    DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
    redisScript.setScriptSource(
            new ResourceScriptSource(new ClassPathResource("script/limit.lua")));
    redisScript.setResultType(Long.class);
    return redisScript;
}

3.创建注解LimitPoint

/**
 * 限流注解
 *
 * @author jayce
 */
@Target({ElementType.METHOD, ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
@Inherited
@Documented
public @interface LimitPoint {
    /**
     * 资源的名字 无实际意义，但是可以用于排除异常
     *
     * @return String
     */
    String name() default "";

    /**
     * 资源的key
     *
     * <p>如果下方 limitType 值为自定义，那么全局限流参数来自于此 如果limitType 为ip，则限流条件 为 key+ip
     *
     * @return String
     */
    String key() default "";

    /**
     * Key的prefix redis前缀，可选
     *
     * @return String
     */
    String prefix() default "";

    /**
     * 给定的时间段 单位秒
     *
     * @return int
     */
    int period() default 60;

    /**
     * 最多的访问限制次数
     *
     * @return int
     */
    int limit() default 10;

    /**
     * 类型 ip限制 还是自定义key值限制 建议使用ip，自定义key属于全局限制，ip则是某节点设置，通常情况使用IP
     *
     * @return LimitType
     */
    LimitTypeEnums limitType() default LimitTypeEnums.IP;
}

4.aop拦截操作

@Before("@annotation(limitPointAnnotation)")
public void interceptor(LimitPoint limitPointAnnotation) {
    LimitTypeEnums limitTypeEnums = limitPointAnnotation.limitType();

    String key;
    int limitPeriod = limitPointAnnotation.period();
    int limitCount = limitPointAnnotation.limit();
    if (limitTypeEnums == LimitTypeEnums.CUSTOMER) {
        key = limitPointAnnotation.key();
    } else {
        key =
                limitPointAnnotation.key()
                        + IpUtils.getIpAddress(
                                ((ServletRequestAttributes)
                                                Objects.requireNonNull(
                                                        RequestContextHolder
                                                                .getRequestAttributes()))
                                        .getRequest());
    }
    ImmutableList<String> keys =
            ImmutableList.of(StringUtils.join(limitPointAnnotation.prefix(), key));
    try {
        Number count = redisTemplate.execute(limitScript, keys, limitCount, limitPeriod);
        assert count != null;
        log.info("限制请求{}, 当前请求{},缓存key{}", limitCount, count.intValue(), key);
        // 如果缓存里没有值，或者他的值小于限制频率
        if (count.intValue() >= limitCount) {
            throw new ServiceException(ResultCode.LIMIT_ERROR);
        }
    }
    // 如果从redis中执行都值判定为空，则这里跳过
    catch (NullPointerException ignored) {
    } catch (ServiceException e) {
        throw e;
    } catch (Exception e) {
        log.error("限流异常", e);
        throw new ServiceException(ResultCode.ERROR);
    }
}

如上操作，我们通过下方的代码实现了流量限速，并且保证了原子性

Number count = redisTemplate.execute(limitScript, keys, limitCount, limitPeriod);

下面使我们的调用代码：

@LimitPoint(key = "hello", period = 60, limit = 2,limitType = LimitTypeEnums.IP)
@GetMapping("/test")
public String test() {
    cache.put("test","test");
    return "hello world";
}

然后使我们的效果截图：