redis分布式限流

1    新增限流注解 RateLimiter 

/**
 * redis限流注解
 * @author gaoyan
 * @date 2019/2/21 10:05
 */
@Documented
@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
public @interface RateLimiter {

    /**
     * 限流key
     */
    String key();

    /**
     * expire时间内允许通过的请求数
     */
    long limit() default 20;

    /**
     * 过期时间，单位秒
     */
    long expire() default 60;
}

2    注解切面编写

package per.yan.ding.aspect;

import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.Signature;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.stereotype.Component;
import per.yan.ding.model.constant.CacheNameSpace;
import per.yan.ding.util.el.AspectSupportUtils;

import javax.annotation.PostConstruct;
import java.text.MessageFormat;
import java.util.ArrayList;
import java.util.List;

/**
 * @author gaoyan
 * @date 2019/2/21 10:10
 */
@Slf4j
@Aspect
@Component
public class RateLimiterAspect {

    @Autowired
    private RedisTemplate redisTemplate;

    private DefaultRedisScript<Long> getRedisScript;

    @PostConstruct
    public void init() {
        getRedisScript = new DefaultRedisScript<>();
        getRedisScript.setResultType(Long.class);
        getRedisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("/lua/rateLimiter.lua")));
    }

    @Pointcut("@annotation(per.yan.ding.aspect.RateLimiter)")
    public void rateLimiter() {

    }

    @Around("@annotation(rateLimiter)")
    public boolean execute(ProceedingJoinPoint jp, RateLimiter rateLimiter) {
        Signature signature = jp.getSignature();
        if (!(signature instanceof MethodSignature)) {
            throw new IllegalArgumentException("Annotation @RateLimiter must used on method!");
        }
        String tokenKey = rateLimiter.key();
        String token = (String) AspectSupportUtils.getKeyValue(jp, tokenKey);
        long limitTimes = rateLimiter.limit();
        long expireTime = rateLimiter.expire();

        String limitKey = getRateLimiterKey(token);
        Long limiter;
        List<String> keyList = new ArrayList<>();
        // 设置key值为注解中的值
        keyList.add(limitKey);
        limiter = (Long) redisTemplate.execute(getRedisScript, keyList, expireTime, limitTimes);

        if (null == limiter || limiter == 0) {
            String msg = "由于超过单位时间=" + expireTime + "-允许的请求次数=" + limitTimes + "[触发限流]";
            log.info("触发限流 msg:{}", msg);
            return false;
        }
        return true;
    }

    private String getRateLimiterKey(String token) {
        return MessageFormat.format(CacheNameSpace.Notification.LOCK_RATE_LIMITER, token);
    }
}

3    lua脚本放入resources目录下 /lua/rateLimiter.lua

--获取KEY
local key1 = KEYS[1]

local val = redis.call('incr', key1)
local ttl = redis.call('ttl', key1)

--获取ARGV内的参数并打印
local expire = ARGV[1]
local times = ARGV[2]

redis.log(redis.LOG_DEBUG,tostring(times))
redis.log(redis.LOG_DEBUG,tostring(expire))

redis.log(redis.LOG_NOTICE, "incr "..key1.." "..val);
if val == 1 then
    redis.call('expire', key1, tonumber(expire))
else
    if ttl == -1 then
        redis.call('expire', key1, tonumber(expire))
    end
end

if val > tonumber(times) then
    return 0
end

return 1

4    增加一个限流成功与否的判断方法

@RateLimiter(key = "#token")
public boolean isNotLimit(String token) {
    return true;
}

此时，需要限流的方法执行时可以先用该方法判断，true再执行业务处理。

也可以不用增加第四步的限流方法，直接将注解加到业务处理的方法上，限流切面类里面更改 if判断后的处理逻辑