基于redis+lua进行限流的方法

Viveka ·

更新时间:2024-09-20

· 1754 次阅读

1，首先我们redis有很多限流的算法(比如：令牌桶，计数器，时间窗口)等，但是都有一定的缺点，令牌桶在单项目中相对来说比较稳定，但是在分布式集群里面缺显的不那么友好，这时候，在分布式里面进行限流的话，我们则可以使用redis+lua脚本进行限流，能抗住亿级并发

2，下面说说lua+redis进行限流的做法
开发环境：idea+redis+lua
第一：
打开idea的插件市场，然后搜索lua，点击右边的安装，然后安装好了，重启即可

第二：写一个自定义限流注解

package com.sport.sportcloudmarathonh5.config;
import java.lang.annotation.*;
/**
 * @author zdj
 * @version 1.0.0
 * @description 自定义注解实现分布式限流
 */
@Target(value = ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface RedisLimitStream {
    /**
     * 请求限制，一秒内可以允许好多个进入(默认一秒可以支持100个)
     * @return
     */
    int reqLimit() default 1000;
    /**
     * 模块名称
     * @return
     */
    String reqName() default "";
}

第三：在指定的方法上面添加该注解

/**
     * 压测接口
     * @return
     */
    @Login(isLogin = false)
    @RedisLimitStream(reqName = "名额秒杀", reqLimit = 1000)
    @ApiOperation(value = "压测接口", notes = "压测接口", httpMethod = "GET")
    @RequestMapping(value = "/pressure", method = RequestMethod.GET)
    public ResultVO<Object> pressure(){
        return ResultVO.success("抢购成功！");
    }

第四：添加一个拦截器对访问的方法在访问之前进行拦截：

package com.sport.sportcloudmarathonh5.config;
import com.alibaba.fastjson.JSONObject;
import com.sport.sportcloudmarathonh5.service.impl.RedisService;
import org.aspectj.lang.ProceedingJoinPoint;
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.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;
import org.springframework.util.ObjectUtils;
import javax.servlet.http.HttpServletResponse;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.List;
/**
 * @author zdj
 * @version 1.0.0
 * @description MyRedisLimiter注解的切面类
 */
@Aspect
@Component
public class RedisLimiterAspect {
    private final Logger logger = LoggerFactory.getLogger(RedisLimitStream.class);
    /**
     * 当前响应请求
     */
    @Autowired
    private HttpServletResponse response;
    /**
     * redis服务
     */
    @Autowired
    private RedisService redisService;
    /**
     * 执行redis的脚本文件
     */
    @Autowired
    private RedisScript<Boolean> rateLimitLua;
    /**
     * 对所有接口进行拦截
     */
    @Pointcut("execution(public * com.sport.sportcloudmarathonh5.controller.*.*(..))")
    public void pointcut(){}
    /**
     * 对切点进行继续处理
     */
    @Around("pointcut()")
    public Object process(ProceedingJoinPoint proceedingJoinPoint) throws Throwable{
        //使用反射获取RedisLimitStream注解
        MethodSignature signature = (MethodSignature) proceedingJoinPoint.getSignature();
        //没有添加限流注解的方法直接放行
        RedisLimitStream redisLimitStream = signature.getMethod().getDeclaredAnnotation(RedisLimitStream.class);
        if(ObjectUtils.isEmpty(redisLimitStream)){
            return proceedingJoinPoint.proceed();
        }
        //List设置Lua的KEYS[1]
        List<String> keyList = new ArrayList<>();
        keyList.add("ip:" + (System.currentTimeMillis() / 1000));
        //获取注解上的参数，获取配置的速率
        //List设置Lua的ARGV[1]
        int value = redisLimitStream.reqLimit();
        // 调用Redis执行lua脚本，未拿到令牌的，直接返回提示
        boolean acquired = redisService.execute(rateLimitLua, keyList, value);
        logger.info("执行lua结果：" + acquired);
        if(!acquired){
            this.limitStreamBackMsg();
            return null;
        }
        //获取到令牌，继续向下执行
        return proceedingJoinPoint.proceed();
    }
    /**
     * 被拦截的人，提示消息
     */
    private void limitStreamBackMsg() {
        response.setHeader("Content-Type", "text/html;charset=UTF8");
        PrintWriter writer = null;
        try {
            writer = response.getWriter();
            writer.println("{\"code\":503,\"message\":\"当前排队人较多，请稍后再试！\",\"data\":\"null\"}");
            writer.flush();
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (writer != null) {
                writer.close();
            }
        }
    }
}

第五：写个配置类，在启动的时候将我们的lua脚本代码加载到redisscript中

package com.sport.sportcloudmarathonh5.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;
/**
 * @author zdj
 * @version 1.0.0
 * @description 实现redis的编码方式
 */
@Configuration
public class RedisConfiguration {
    /**
     * 初始化将lua脚本加载到redis脚本中
     * @return
     */
    @Bean
    public DefaultRedisScript loadRedisScript() {
        DefaultRedisScript redisScript = new DefaultRedisScript();
        redisScript.setLocation(new ClassPathResource("limit.lua"));
        redisScript.setResultType(Boolean.class);
        return redisScript;
    }
}

第六：redis执行lua的方法

  /**
     * 执行lua脚本
     * @param redisScript lua源代码脚本
     * @param keyList
     * @param value
     * @return
     */
    public boolean execute(RedisScript<Boolean> redisScript, List<String> keyList, int value) {
        return redisTemplate.execute(redisScript, keyList, String.valueOf(value));
    }

第七：在resources目录下面新加一个lua脚本文件，将下面代码拷贝进去即可：

local key = KEYS[1] --限流KEY（一秒一个）
local limit = tonumber(ARGV[1]) --限流大小
local current = tonumber(redis.call('get', key) or "0")
if current + 1 > limit then --如果超出限流大小
    return false
else --请求数+1，并设置2秒过期
    redis.call("INCRBY", key, "1")
    redis.call("expire", key, "2")
end
return true