分布式限流:redis+lua脚本

分布式的令牌桶算法 redis+lua脚本 自定义注解+aop实现

直接贴代码，不懂清细品

自定义的限流注解,加在controller的方法上即可   limitKey为redis的key  limit为允许的流量

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface DistriLimitAnno {
    public String limitKey() default "limit";
    public int limit() default 1;
}

AOP切面类   

@Aspect
@Component
public class LimitAspect {
    @Autowired
    DistributedLimit distributedLimit;

    @Before("@annotation(com.oyjg.distributedTokenBucket.annotation.DistriLimitAnno)")
    public void beforeLimit(JoinPoint joinPoint) throws  Exception{
　　　　　//获取方法名称
        MethodSignature signature = (MethodSignature)joinPoint.getSignature();
        Method method = signature.getMethod();
        DistriLimitAnno distriLimitAnno = method.getAnnotation(DistriLimitAnno.class);
　　　　　//得到我们自定义注解的两个值
        String key = distriLimitAnno.limitKey();
        int limit = distriLimitAnno.limit();
　　　　 //判断有没有被限流
        Boolean exceededLimit = distributedLimit.distributedLimit(key, String.valueOf(limit));
        if(!exceededLimit){
            throw new RuntimeException("exceedeLimit limit");
        }
    }
}

@Component
public class DistributedLimit {
    @Autowired
    private RedisTemplate<String,String> redisTemplate;
    @Autowired
    private RedisScript<Long> limitScript;

    public Boolean distributedLimit(String key,String limit){
        Long id = 0L;
        try {
            id = redisTemplate.execute(limitScript, Collections.singletonList(key),limit);
        }catch (Exception e){
            System.out.println("error:"+e);
        }
        if(id==0L){
            return false;
        }else {
            return true;
        }
    }
}

Lua脚本

--lua 下标从 1 开始
-- 限流 key
local key = KEYS[1]
-- 限流大小
local limit = tonumber(ARGV[1])

-- 获取当前流量大小
local curentLimit = tonumber(redis.call('get', key) or "0")

if curentLimit + 1 > limit then
  -- 达到限流大小 返回
  return 0;
else
  -- 没有达到阈值 value + 1
  redis.call("INCRBY", key, 1)
  -- EXPIRE后边的单位是秒
  redis.call("EXPIRE", key, 10)
  return curentLimit + 1
end

controller

@Controller
@RequestMapping("/distributedLock")
public class DistributedLockController {
    @Autowired
    DistributedLock lock;

    @PostMapping("/distributedLimit")
    @DistriLimitAnno(limitKey="limit", limit = 100)
    public String distributedLimit(String userId) {
        System.out.println(userId);
        return "ok";
    }
}

测试

@Test
    public void testDistriLimit() {
        String url = "http://localhost:8080/distributedLock/distributedLimit";
        for (int i=0;i<100;i++){
            int finalI = i;
            new Thread(()->{
                MultiValueMap multiValueMap = new LinkedMultiValueMap<>();
                multiValueMap.add("userId",String.valueOf(finalI));
                String result = restTemplate.postForObject(url, multiValueMap, String.class);
                System.out.println("-------"+result);
            }).start();
        }
        try {
            System.in.read();
        } catch (IOException e) {
            //e.printStackTrace();
        }
    }