基于Redis的分布式服务限流

基于redis缓存的分布式服务限流

1、定于限流注解，限流注解可以加需要限流的业务类上，value为限流业务类型

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface TPSControl {
     String value();
}

2、限流配置类，用于解析限流配置信息存放

@Data
public class TPSControlConfig {
    
    // 机构号
    private String instNo;
    // 时长s
    private Integer time;
    // 请求数
    private Integer maxRequest;
}

3、redis工具类，分布式系统限流基于redis缓存

@Service
public class RedisUtils {
    // 限流redis脚本，如果当前key存在，自增key并返回自增后的值；如果当前key不存在，自增后设置过期时间，返回自增后的值
    private final static String INCREASE_AND_EXPIRE_LUA_SCRIPT = "local count = redis.call('INCR', KEYS[1])\n" +
            "if count == 1 then\n" +
            "  redis.call('expire', KEYS[1], ARGV[1])\n" +
            "end\n" +
            "return redis.call('get', KEYS[1])\n";

    /**
     * TPS限流
     * 实现redisKey++
     * 如果当前的redisKey已过期，则设置过期时长
     *
     * @param redisKey
     * @param timeout
     * @return
     */
    public Object increaseAndExpireKey(String redisKey, String timeout) {

        List<String> keys = new ArrayList<>();
        keys.add(redisKey);

        List<String> argv = new ArrayList<>();
        argv.add(timeout);
        return excuteLuaScript(keys, argv, INCREASE_AND_EXPIRE_LUA_SCRIPT);
    }

    public Object excuteLuaScript(List keys,List args,String luaScript){
        //spring自带的执行脚本方法中，集群模式直接抛出不支持执行脚本异常，此处拿到原redis的connection执行脚本
        Object result = redisTemplate.execute(new RedisCallback() {
            @Override
            public Object doInRedis(RedisConnection connection) throws DataAccessException {
                Object nativeConnection = connection.getNativeConnection();
                // 集群模式和单点模式虽然执行脚本的方法一样，但是没有共同的接口，所以只能分开执行
                // 集群
                if (nativeConnection instanceof JedisCluster) {
                    System.out.println("集群执行");
                    return  ((JedisCluster) nativeConnection).eval(luaScript, keys, args);
                }

                // 单点
                else if (nativeConnection instanceof Jedis) {
                    System.out.println("单点执行");
                    return ((Jedis) nativeConnection).eval(luaScript, keys, args);
                }
                return null;
            }
        });
        return result;
    }
}

4、限流切面控制，根据注解上的业务类型，根据机构的限流配置，对机构的请求进行限流

@Aspect
@Component
@Slf4j
public class TPSControlAspect {

    /** TPS限流配置 */
    private static Map<String, List<TPSControlConfig>> tpsControlConfig = new HashMap<>(8);

    @Autowired
    private RedisUtils redisUtils;

    @Value("${tpsControl}")
    private void setTpsControlMap(String tpsControl) {
        // 初始化TPS限流配置
        initTpsControlConfig(tpsControl);
    }

    // 切面拦截TPSControl这个注解
    @Pointcut("@annotation(com.jfbank.fincloud.loan.order.loan.core.annotation.TPSControl)")
    public void tpsControl() {

    }

    // 环绕拦截
    @Around("tpsControl()")
    public Object doAround(ProceedingJoinPoint pjp) throws Throwable{

        // 解析当前的TPS限流配置
        TPSControlConfig currentTPSConfig = resolveTPSControlConfig(pjp);
        // 判断是否过载
        if (currentTPSConfig == null || !isOverload(currentTPSConfig)) {
            return pjp.proceed();
        }
        throw new ServiceException("500000", "超过最大访问数");
    }

    /**
     * 
     * 拦截方法参数对象中必须要要有机构号（instNo）属性
     * 限流配置是机构维度的限流配置
     *
     * @param pjp
     * @return
     * @throws Exception
     */
    private TPSControlConfig resolveTPSControlConfig(ProceedingJoinPoint pjp) throws Exception{
        // 目标类
        Class<?> targetClass = pjp.getTarget().getClass();
        // 方法签名
        MethodSignature ms =(MethodSignature) pjp.getSignature();
        // 拦截目标方法
        Method targetMethod = targetClass.getDeclaredMethod(ms.getName(), ms.getParameterTypes());
        // 获取机构号
        String instNo = null;
        for (Object arg : pjp.getArgs()) {
            Field field = null;
            try {
                field = arg.getClass().getDeclaredField("instNo");
            } catch (NoSuchFieldException e) {}
            if (!ObjectUtils.isEmpty(field)) {
                field.setAccessible(true);
                instNo = String.valueOf(field.get(arg));
                break;
            }
        }
        // 获取业务类型
        String businessType = targetMethod.getAnnotation(TPSControl.class).value();
        return getCurrentTPSConfig(businessType, instNo);
    }

    /**
     * 解析注入配置
     * 无法通过apollo自动更新
     *
     * @param tpsConfig
     */
    private void initTpsControlConfig(String tpsConfig) {

        Map<String, Object> configMap = JSON.parseObject(tpsConfig, Map.class);

        if (CollectionUtils.isEmpty(configMap)) {
            return;
        }
        configMap.forEach((key, value) -> {
            tpsControlConfig.put(key, JSONArray.parseArray(JSON.toJSONString(value), TPSControlConfig.class));
        });
        tpsControlConfig.forEach((key, value) -> {
            if (CollectionUtils.isEmpty(value)) {
                return;
            }
            value.forEach(config -> {
                if (StringUtils.isEmpty(config.getRedisKey())) {
                    config.setRedisKey("fincloud:order:loan:tpslimit:" + key + ":" + config.getInstNo());
                }
            });
        });
    }

    /**
     * 获取当前业务的TPS限流配置
     *
     * @param businessType 业务类型
     * @param instNo 机构号
     * @return
     */
    public TPSControlConfig getCurrentTPSConfig(String businessType, String instNo) {

        List<TPSControlConfig> TPSControlConfigList = tpsControlConfig.get(businessType);
        // 获取当前机构的TPS限流配置
        if (CollectionUtils.isEmpty(TPSControlConfigList)) {
            return null;
        }
        Optional<TPSControlConfig> optional = TPSControlConfigList.stream().filter(config -> config.getInstNo().equals(instNo)).findFirst();
        return optional.orElse(null);
    }

    /**
     * 是否过载
     *
     * @param tpsControlConfig
     * @return
     */
    public boolean isOverload(TPSControlConfig tpsControlConfig) {

        Object result = null;
        // 操作Redis失败，不拦截当前访问
        try {
            result = redisUtils.increaseAndExpireKey(tpsControlConfig.getRedisKey(), String.valueOf(tpsControlConfig.getTime()));
        } catch (Exception e) {
            log.error("redisUtils.increaseAndExpireKey failed!", e);
            return false;
        }

        int currentVisitCount = Integer.parseInt(result.toString());
        if (currentVisitCount > tpsControlConfig.getMaxRequest()) {
            log.error("TPSControl currentVisitCount is over limit! currentVisitCount = {} tpsConfig = {}", currentVisitCount, tpsControlConfig);
            return true;
        }
        return false;
    }
}

5、自定义异常类

public class ServiceException extends RuntimeException {
    private static final long serialVersionUID = 1L;
    private String errCode;

    public ServiceException(String errorCode, String message) {
        super(message);
        this.errCode = errorCode;
    }

    public ServiceException(String errorCode, String message, Throwable cause) {
        super(message, cause);
        this.errCode = errorCode;
    }

    public String getErrCode() {
        return this.errCode;
    }

    public void setErrCode(String errCode) {
        this.errCode = errCode;
    }
}

6、限流配置，对业务类型为query的请求进行机构维度的限流，限制机构号10450的机构每分钟最大请求数为5

tpsControl = {'query':[{'instNo':'10450','time':60,'maxRequest':5}]}

7、限流测试

@Data
public class CommonRequestDTO {
    private String instNo;
}


@Service
@Slf4j
public class TPSControlTarget {

    @TPSControl("query")
    public Map<String, String> business(CommonRequestDTO commonRequestDTO) {
        log.info("--- process business ---");
        Map<String, String> result = new HashMap<>();
        result.put("businessResult", "SUCCEED");
        return result;
    }
}

@RunWith(SpringRunner.class)
@SpringBootTest(classes = OrderLoanApplication.class)
@Slf4j
public class TPSControlTest {

    @Autowired
    TPSControlTarget tpsControlTarget;

    @Test
    public void TPSControlTest() throws Exception{
        for (int i = 0; i < 10; i++) {
            CommonRequestDTO commonRequestDTO = new CommonRequestDTO();
            commonRequestDTO.setInstNo("10450");
            new Thread(() -> {
                Object result = tpsControlTarget.business(commonRequestDTO);
                log.info("result -> {}", result);
            }).start();

        }
        Thread.sleep(3000);
    }
}

8、测试结果

只有五个请求可以正常返回，其余请求抛出异常！