通过Redis+自定义注解实现接口限流策略

lz 1年前 ⋅ 1202 阅读

1、前言

通过自定义注解+reids+lua实现,接口限流策略,其实质就是对redis的分布式锁的应用。

流程基本如下:

  • 1、Controller接口的方法,实现自定义注解@RateLimiter
  • 2、自定义拦截RateLimiterHandlerInterceptor,拦截包含注解@RateLimiter的接口,进行验证。
  • 3、为了保证在并发请求下的精确性,使用redis+lua脚本进行加锁。
  • 4、如果窗口时间内请求没有达到上限则放行,如果达到了上限,则返回错误。

2、代码实现

2.1 自定义注解

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

    //一个IP下请求的并发限制

    int period() default 1;  //限流时间周期,默认 1S

    int count() default 10; //周期内限制次数,默认  10次

    boolean rateIP() default true; //默认限制IP,设置为false表示只限制接口请求次数
}

rateIP表示更加细致的控制,如果不需要限制ip,而是对接口设置一个统一的访问上限,则将rateIP设置为flase,比如:@RateLimiter(rateIP = false)

2.2 lua脚本配置

为了保证适应高并发,通过reids+lua脚本来实现加锁与解锁,达到一致性的目的。

lua脚本:
resources目录下创建lua文件夹,在lua创建rateLimit.lua文件,内容如下:

local key = KEYS[1]
local limit = tonumber(KEYS[2])
local length = tonumber(KEYS[3])
--redis.log(redis.LOG_NOTICE,' length: '..length)
local current = redis.call('GET', key)
if current == false then
   --redis.log(redis.LOG_NOTICE,key..' is nil ')
   redis.call('SET', key,1)
   redis.call('EXPIRE',key,length)
   --redis.log(redis.LOG_NOTICE,' set expire end')
   return '1'
else
   --redis.log(redis.LOG_NOTICE,key..' value: '..current)
   local num_current = tonumber(current)
   if num_current+1 > limit then
       return '0'
   else
       redis.call('INCRBY',key,1)
       return '1'
   end
end

其中key就是请求的接口信息,length为窗口时间,limit为窗口时间内允许访问的最大并发数量,比如:key=test,length=1,limit=10,就表示接口test1秒内最大的并发为10,脚本中返回1表示成功,0表示失败。

rateLimit.lua逻辑如下:

  • 1、根据key获取当前的请求数量;
  • 2、如果为空,则将key的值设置为1,并且设置过期时间也就是窗口时间,并且返回1;
  • 3、如果不为空,则判断当前请求数+1是否大于limit(最大并发数),如果是则返回0,如果不是则将key的值自加1,并且返回1

RedisLuaConfig:

创建Redis操作Lua的配置类,内容如下:

@Configuration
public class RedisLuaConfig {
    @Resource
    private StringRedisTemplate stringRedisTemplate;

    /**
     * @param keyList redis得Lua脚本中的key列表,我们把参数放在这里面传递
     * @return result 返回 1表示,正常,0表示限制访问
     */
    public String runLuaScript(List<String> keyList) {
        DefaultRedisScript<String> redisScript = new DefaultRedisScript<>();
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("lua/rateLimit.lua")));
        redisScript.setResultType(String.class);
        String args = "none";
        try {
            return stringRedisTemplate.execute(redisScript, keyList);
        } catch (Exception e) {
            e.printStackTrace();
            return "0";
        }
    }
}

2.3 拦截器配置

创建RateLimiterHandlerInterceptor类并且实现HandlerInterceptor接口,在该类中获取含有RateLimiter的请求,并且取出注解的配置,然后通过ip + "@" + httpMethod + "@" + path的形式生成key,再根据窗口时间和最大并发数,请求lua脚本,实现限流判断。

为了防止被暴力请求,默认情况下所有的接口进行限流,默认最大并发的10

@Component
@Slf4j
public class RateLimiterHandlerInterceptor implements HandlerInterceptor {

    @Resource
    private RedisLuaConfig redisLuaConfig;

    private static final int DEFAULT_PERIOD = 1;
    private static final int DEFAULT_COUNT = 10;

    private static final String LIMITER_KEY = "limiter:";
    private static final String LIMITER_IP_KEY = "limiter-ip:";

    private static final String SUCCESS_CODE = "1";

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {

        Method method = ((HandlerMethod) handler).getMethod();

        String path = request.getServletPath();
        String httpMethod = request.getMethod(); //GET、PUT、DELETE、POST

        //获取当前请求IP
        String ip = IpUtils.getIpAddr(request);

        //如果没有开启开启注解,也进行默认限流操作
        int period = DEFAULT_PERIOD;  //限流时间周期,默认 1S
        int count = DEFAULT_COUNT; //周期内限制次数,默认  10次
        boolean rateIP = true;//默认限制IP,设置为false表示只限制接口请求次数

        IpRateLimiter ipRateLimiter = method.getAnnotation(IpRateLimiter.class);
        if (ipRateLimiter != null) {
            period = ipRateLimiter.period();
            count = ipRateLimiter.count();
            //设置当前key,对同一个ip下同一个请求进行限流操作
            rateIP = ipRateLimiter.rateIP();
        }

        //不对ip进行限制
        String key = LIMITER_KEY + httpMethod + "@" + path;
        if (rateIP) {
            //对ip进行次数限制
            key = LIMITER_IP_KEY + ip + "@" + httpMethod + "@" + path;
        }

        String res = "";
        try {
            List<String> keyList = new ArrayList<>();
            keyList.add(key);
            //表示时间周期内运行访问得次数
            keyList.add(String.valueOf(count)); //count
            keyList.add(String.valueOf(period));//period
            res = redisLuaConfig.runLuaScript(keyList);
        } catch (Exception e) {
			e.printStackTrace();
			throw new RuntimeException("你被限流了");
        }
        //正常执行RedisLua
        if (!SUCCESS_CODE.equals(res)) {
           log.error("你被限流了");
           throw new RuntimeException("你被限流了");
        }
        return true;
    }
}

3、测试

Controller中的一个接口,加上RateLimiter注解,为了更好的验证限流,将请求的最大并发数设置为1

@RestController
@RequestMapping("test")
public class TestController {
    @GetMapping("/start")
    @RateLimiter(count = 1)
    public String start() {
        logger.info("成功访问接口");
        return "操作成功";
    }
}

通过浏览器快速刷新访问接口,就会触发限流,结果如下:

 

--end--

 

版权 本着开源共享、共同学习的精神,本文转载自 https://blog.csdn.net/zhuocailing3390/article/details/122992935 , 如果侵权之处,请联系博主进行删除,谢谢~