Redis+lua使用切面实现流量控制

对接口实现流量控制，使其在N秒内访问控制在M次

lua脚本实现令牌桶算法（来自网络）：

--令牌秒杀 Lua脚本  tokenTimeLimiter.lua
local key = KEYS[1] 			--限流KEY（一秒一个）
local limit = tonumber(ARGV[1]) --限流大小
local exprie = ARGV[2] 			--过期时间
-- 获取当前计数值
local current = tonumber(redis.call('get', key) or "0")
if current + 1 > limit then --如果超出限流大小
    return 0
else
	current = tonumber(redis.call("INCRBY", key, "1")) --请求数+1
	if current == 1 then  --第一次访问需要设置过期时间
    	redis.call("expire", key,exprie) --设置过期时间
    end
end
return 1  --返回1代表不限流

自定义注解：

package com.suning.shengyu.web.aop;

import java.lang.annotation.*;

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

}

切面：

package com.suning.shengyu.web.aop;

import com.alibaba.fastjson.JSON;
import com.suning.framework.sedis.ShardedJedisAction;
import com.suning.framework.sedis.impl.ShardedJedisClientImpl;
import com.suning.shengyu.common.support.log.ShengyuLogger;
import com.suning.shengyu.common.support.outeropenapi.OpenApiRequest;
import com.suning.shengyu.common.support.outeropenapi.OpenResultCode;
import com.suning.shengyu.common.support.outeropenapi.ResultMap;
import com.suning.shengyu.util.AESCoder;
import com.suning.shengyu.web.util.SignPlugin;
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.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import redis.clients.jedis.Jedis;
import redis.clients.jedis.ShardedJedis;

import java.util.ArrayList;
import java.util.List;

@Aspect
@Component
public class RateLimitAdvice {
    private static final ShengyuLogger LOGGER = ShengyuLogger.getLogger(RateLimitAdvice.class);

    private static final String FLOWCONTROL_CACHE_STRING_KEY = "SHENGYU:MULTI-STRING:RATELIMIT:";

    //时间间隔
    private static final String INTERVAL_SECOND = "600";

    //方法调用次数
    private static final String METHOD_CALL_TIMES = "120";

    @Autowired
    private SignPlugin signPlugin;

    @Autowired
//    private ShardedJedisClientImpl jedisClientForSys;
    private ShardedJedisClientImpl jedisClient;

    @Autowired
    private String luaScript;

    @Pointcut(value = "@annotation(com.suning.shengyu.web.aop.RateLimit)")
    private void pointCut(){
        System.out.println("我是切入点。。。");
    };

    @Around("pointCut()")
    public ResultMap around(ProceedingJoinPoint proceedingJoinPoint) {
        Object[] args = proceedingJoinPoint.getArgs();
        OpenApiRequest openApiRequest = (OpenApiRequest)args[0];
        String mallId = AESCoder.decrypt(openApiRequest.getAppKey(), signPlugin.md5(AESCoder.KEYSTRHEXSTR_APPKEY));

        ResultMap resultMap;
        String method = null;
        try {
            method = openApiRequest.getMethod();
            LOGGER.info("流控切面获取的mallId={}, method={}", mallId, method);
            List<String> values = new ArrayList<>();
            values.add(METHOD_CALL_TIMES);
            values.add(INTERVAL_SECOND);
            String key = FLOWCONTROL_CACHE_STRING_KEY + mallId + method;
            Long executeResult = luaEval(luaScript, key, values);
//            Long executeResult = 0l;
            LOGGER.info("流控令牌数量剩余{}", executeResult);

            if (executeResult == 1) {
                LOGGER.info("当前访问时间段内剩余{}次访问次数", executeResult);
                Object obj = proceedingJoinPoint.proceed();
                System.out.println(JSON.toJSONString(obj));
                resultMap = (ResultMap)obj;
            } else {
                LOGGER.info("当前接口调用超过时间段内限流,key:{}", key);
                return ResultMap.fail(OpenResultCode.REQUEST_LIMITED);
            }
        } catch (Throwable throwable) {
            LOGGER.error("接口限流层异常，mallId={},method={}", mallId, method, throwable);
            return ResultMap.fail(OpenResultCode.SYSTEM_ERROR);
        }
        return resultMap;
    }

    public Long luaEval(final String LUA_SCRIPT, final String key, final List<String> args) {
        return jedisClient.execute(new ShardedJedisAction<Long>() {
            @Override
            public Long doAction(ShardedJedis shardedJedis) {
                Jedis jedis = shardedJedis.getShard(key);
                // 参数KEY
                List<String> keys = new ArrayList<>();
                keys.add(key);
                // 脚本调用
                return (Long) jedis.eval(LUA_SCRIPT, keys, args);
            }
        });
    }

}

读取lua脚本：

@Component
public class RateLimitConfig {
    private static final ShengyuLogger LOGGER = ShengyuLogger.getLogger(RateLimitConfig.class);
    private static final String SCRIPT_FILE_NAME = "ratelimit.lua";

    @Bean({ "luaScript" })
    public String luaScript() {
        String luaScript = "";
        try {
            InputStream in = RateLimitConfig.class.getClassLoader().getResourceAsStream(SCRIPT_FILE_NAME);
            luaScript = new String(ByteStreams.toByteArray(in));
            LOGGER.info("流控脚本：{}", luaScript);
        } catch (IOException e) {
            LOGGER.info("加载lua脚本时异常", e);
        }
        return luaScript;
    }
}

使用举例：

    @ResponseBody
	@RequestMapping("/testApi")
	@RateLimit
	public ResultMap testApi(@RequestBody OpenApiRequest openApiRequest) {
        //业务逻辑
		return ResultMap.ok();
	}