Redis+注解实现限流机制（IP、自定义等）

简介

在项目的使用过程中，限流的场景是很多的，尤其是要提供接口给外部使用的时候，但是自己去封装的话，相对比较耗时。

本方式可以使用默认（方法），ip、自定义参数进行限流，根据时间和次数进行。

整合步骤

依赖

  <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-data-redis</artifactId>
        </dependency>

        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-web</artifactId>
        </dependency>

        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-aop</artifactId>
        </dependency>
        <dependency>
            <groupId>org.projectlombok</groupId>
            <artifactId>lombok</artifactId>
        </dependency>
        <dependency>
            <groupId>cn.hutool</groupId>
            <artifactId>hutool-all</artifactId>
            <version>5.8.27</version>
        </dependency>
        <dependency>
            <groupId>com.googlecode.aviator</groupId>
            <artifactId>aviator</artifactId>
            <version>5.4.1</version>
            <scope>compile</scope>
        </dependency>

限流注解

package com.walker.ratelimiter.annotation;

import com.walker.ratelimiter.enums.LimitType;

import java.lang.annotation.*;

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

    /**
     * 限流key
     */
    String key() default "rate_limit";

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

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

    /**
     * 限流类型
     */
    LimitType limitType() default LimitType.DEFAULT;

    /**
    * 自定义编码
     * 支持SPEL表达式
     * 如果使用多参数，则使用:分割
     *
    */
    String customerCode() default "";

    /**
     * 自定义编码分割符
     */
    String customerCodeSplit() default ":";
}

限流配置：获取限流lua脚本

package com.walker.ratelimiter.config;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;

@Configuration
public class RateLimitConfig {

    @Bean
    public DefaultRedisScript<Long> limitScript() {
        DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("lua/limit.lua")));
        redisScript.setResultType(Long.class);
        return redisScript;
    }

}

基础变量

package com.walker.ratelimiter.constants;

public interface BaseConstants {

    String COLON = ":";
}

枚举类型

package com.walker.ratelimiter.enums;

public enum LimitType {

    /**
     * 默认策略
     */
    DEFAULT,

    /**
     * 根据IP进行限流
     */
    IP,

    /**
    * 自定义
    */
    CUSTOME,

}

lua脚本

local key = KEYS[1]
local count = tonumber(ARGV[1])
local time = tonumber(ARGV[2])
local current = redis.call('get', key)
if current and tonumber(current) > count then
    return tonumber(current)
end
current = redis.call('incr', key)
if tonumber(current) == 1 then
    redis.call('expire', key, time)
end
return tonumber(current)

切面类

package com.walker.ratelimiter.aspect;

import cn.hutool.core.util.StrUtil;
import com.walker.ratelimiter.annotation.RateLimiter;
import com.walker.ratelimiter.constants.BaseConstants;
import com.walker.ratelimiter.enums.LimitType;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.context.expression.MethodBasedEvaluationContext;
import org.springframework.core.DefaultParameterNameDiscoverer;
import org.springframework.core.ParameterNameDiscoverer;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.expression.EvaluationContext;
import org.springframework.expression.Expression;
import org.springframework.expression.TypedValue;
import org.springframework.expression.spel.standard.SpelExpressionParser;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.Collections;
import java.util.List;

@Slf4j
@Aspect
@Component
public class RateLimiterAspect {

    private final RedisTemplate redisTemplate;
    private final RedisScript<Long> limitScript;
    private SpelExpressionParser spelExpressionParser = new SpelExpressionParser();


    public RateLimiterAspect(RedisTemplate redisTemplate, RedisScript<Long> limitScript) {
        this.redisTemplate = redisTemplate;
        this.limitScript = limitScript;
    }

    @Around("@annotation(com.walker.ratelimiter.annotation.RateLimiter)")
    public Object doBefore(ProceedingJoinPoint joinPoint) throws Throwable {
        MethodSignature methodSignature = (MethodSignature) joinPoint.getSignature();

        RateLimiter rateLimiter = methodSignature.getMethod().getAnnotation(RateLimiter.class);

        //判断该方法是否存在限流的注解
        if (null != rateLimiter) {
            //获得注解中的配置信息
            int count = rateLimiter.count();
            int time = rateLimiter.time();

            //调用getCombineKey()获得存入redis中的key   key -> 注解中配置的key前缀-ip地址-方法路径-方法名
            String combineKey = getCombineKey(rateLimiter, methodSignature, joinPoint);
            log.info("combineKey->,{}", combineKey);
            //将combineKey放入集合
            List<Object> keys = Collections.singletonList(combineKey);
            log.info("keys->", keys);
            try {
                //执行lua脚本获得返回值
                Long number = (Long) redisTemplate.execute(limitScript, keys, count, time);
                //如果返回null或者返回次数大于配置次数，则限制访问
                if (number == null || number.intValue() > count) {
                    throw new RuntimeException("访问过于频繁，请稍候再试");
                }
                log.info("限制请求'{}',当前请求'{}',缓存key'{}'", count, number.intValue(), combineKey);
            } catch (RuntimeException e) {
                throw e;
            } catch (Exception e) {
                throw new RuntimeException("服务器限流异常，请稍候再试");
            }
        }

        return joinPoint.proceed();
    }

    /**
     * Gets combine key.
     *
     * @param rateLimiter the rate limiter
     * @param signature   the signature
     * @param joinPoint
     * @return the combine key
     */
    public String getCombineKey(RateLimiter rateLimiter, MethodSignature signature, ProceedingJoinPoint joinPoint) throws UnknownHostException {
        StringBuilder stringBuffer = new StringBuilder(rateLimiter.key());
//        ip限流
        if (rateLimiter.limitType() == LimitType.IP) {
            InetAddress ip = InetAddress.getLocalHost();
            log.info("获取ip地址为：{}", ip);
            String hostAddress = ip.getHostAddress();
            stringBuffer.append(hostAddress).append(BaseConstants.COLON);
//        自定义编码限流
        } else if (rateLimiter.limitType() == LimitType.CUSTOME) {
            if (StrUtil.isEmpty(rateLimiter.customerCode())) {
                throw new RuntimeException("自定义编码不能为空");
            }

            String customerCode = rateLimiter.customerCode();
            String split = rateLimiter.customerCodeSplit();
            String[] customerCodes = customerCode.split(split);
            for (String code : customerCodes) {
                ParameterNameDiscoverer parameterNameDiscoverer = new DefaultParameterNameDiscoverer();
                EvaluationContext evaluationContext = new MethodBasedEvaluationContext(TypedValue.NULL, signature.getMethod(), joinPoint.getArgs(), parameterNameDiscoverer);
                Expression expression = spelExpressionParser.parseExpression(code);
                String resolvedCustomerCode =  String.valueOf(expression.getValue(evaluationContext));
                if(StrUtil.isEmpty(resolvedCustomerCode)){
                    throw new RuntimeException("自定义编码不能为空");
                }
                stringBuffer.append(BaseConstants.COLON).append(resolvedCustomerCode);
            }
        }
        Method method = signature.getMethod();
        Class<?> targetClass = method.getDeclaringClass();
        stringBuffer.append(BaseConstants.COLON).append(targetClass.getName()).append(BaseConstants.COLON).append(method.getName());
        return stringBuffer.toString();
    }


}

使用

• 根据ip进行限流

limitType = LimitType.IP

• 默认

limitType = LimitType.DEFAULT

• 自定义参数限流

使用Spel表达式，从参数中获取自定义的code，然后60s限流5次

@RateLimiter(limitType = LimitType.CUSTOME,
             customerCode = "#form.appCode:#toUserInfo.userUid",
             count = 5,time = 60)
public Result<Boolean> message(ImSendMsgForm form, MissuUsers toUserInfo) {
    
}