1.创建枚举类,表明限流类型

public enum LimitType {

    /**
     * 全局限流
     */
    DEFAULT,

    /**
     * 单ip限流
     */
    IP
}

2.创建注解,标记相关信息

@Documented
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimiter {

    /**
     *redis 限流key的前缀
     */
    String prefix() default "limit:";

    /**
     *限流时间,单位秒
     */
    int time() default 60;

    /**
     * 限流次数
     */
    int count() default 20;

    /**
     * 限流类型
     */
    LimitType type() default LimitType.IP;
}

3.限流Lua脚本

/**
 *  接口限流
 *
 */
@Slf4j
@Component
public class RedisLimitUtil {
    @Autowired
    private StringRedisTemplate redisTemplate;

    /**
     * 限流
     * @param key   键
     * @param count 限流次数
     * @param times 限流时间
     * @return
     */
    public boolean limit(String key, int count, int times) {
        try {
            String script = "local lockKey = KEYS[1]\n" +
                    "local lockCount = KEYS[2]\n" +
                    "local lockExpire = KEYS[3]\n" +
                    "local currentCount = tonumber(redis.call('get', lockKey) or \"0\")\n" +
                    "if currentCount < tonumber(lockCount)\n" +
                    "then\n" +
                    "redis.call(\"INCRBY\", lockKey, \"1\")\n" +
                    "redis.call(\"expire\", lockKey, lockExpire)\n" +
                    "return true\n" +
                    "else\n" +
                    "return false\n" +
                    "end";
            RedisScript<Boolean> redisScript = new DefaultRedisScript<>(script, Boolean.class);
            List<String> keys = Arrays.asList(key, String.valueOf(count), String.valueOf(times));
            log.info("key:{}",key);
            return redisTemplate.execute(redisScript, keys);
        } catch (Exception e) {
            log.error("限流脚本执行失败:{}", e.getMessage());
        }
        return false;
    }
}

4.创建限流切片类,实现限流算法

@Aspect
@Component
public class RateLimiterAspect {
    @Resource
    private RedisLimitUtil redisLimitUtil;

    /**
     * 前置通知,判断是否超出限流次数
     * @param point
     */
    @Before("@annotation(rateLimiter)")
    public void doBefore(JoinPoint point, RateLimiter rateLimiter) {
        try {
            // 拼接key
            String key = getCombineKey(rateLimiter, point);
            // 判断是否超出限流次数
            if (!redisLimitUtil.limit(key, rateLimiter.count(), rateLimiter.time())) {
                throw new FighterRuntimeException(RestResultCode.CODE_500.getMessage(), RestResultCode.CODE_500.getCode(), false);
            }
        }  catch (Exception e) {
            throw new RuntimeException("接口限流异常,请稍候再试");
        }
    }
    /**
     * 根据限流类型拼接key
     */
    public String getCombineKey(RateLimiter limit, JoinPoint point) {
        StringBuilder sb = new StringBuilder(limit.prefix());
        // 按照IP限流
        if (limit.type() == LimitType.IP) {
            sb.append(IPUtil.getIpAddr(HttpRequestUtil.getRequest())).append("-");
        }
        // 拼接类名和方法名
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        Class<?> targetClass = method.getDeclaringClass();
        sb.append(targetClass.getName()).append("-").append(method.getName());
        return sb.toString();
    }
}

5.总结

根据注解中标注的,count(次数),times(时间)在redis中设置一个key,其中key是“limit:IP-类名-方法名”,设置超时时间为times,值为0,每当这个ip访问一次,就通过lua脚本原子操作+1,并判断与count的值,一旦大于等于count,马上返回false,触发限流,如果没有,则在这个key的过期时间内,没访问一次,就加一。