【基于Redis的分布式限流】

分布式限流的样例

基于 Redis 的分布式服务限流是一种常用的策略，用于控制和限制对某个服务的访问频率，以防止过载。使用 Redis 可以有效地实现分布式环境下的限流。
1.请求实体

@Data
public class CommonRequestDTO {
    private  String instNo;
}

2.请求所执行的服务方法：

@Service
@Slf4j
public class TPSControlTarget {
    @TPSControl("query")
    public Map<String,String> business(CommonRequestDTO commonRequestDTO){
        log.info("开始执行业务逻辑:{}",commonRequestDTO);
        Map<String,String> map=new HashMap<>();
        map.put("businessResult","SUCESS");
        return map;
    }

}

分布式限流注解的实现

请求实体

package com.serviceLimitQps;

import lombok.Data;

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

}

注解实现

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

    /**
     * The point name should applied for.
     *
     * @return action type, default READ
     */
    String value();
}

切面的实现，需要了解AOP代理

package com.serviceLimitQps;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import lombok.extern.slf4j.Slf4j;
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.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;
import org.springframework.util.ObjectUtils;
import org.springframework.util.StringUtils;

import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;

@Aspect
@Component
@Slf4j
public class TPSControlAspect {

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

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

    // 切面拦截TPSControl这个注解
    @Pointcut("@annotation(com.serviceLimitQps.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;
    }
}

用redis来判断当前的流量，这里区分单机或者集群

package com.serviceLimitQps;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.data.redis.RedisProperties;
import org.springframework.dao.DataAccessException;
import org.springframework.data.redis.connection.RedisConnection;
import org.springframework.data.redis.core.RedisCallback;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Service;
import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisCluster;

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

@Service
public class RedisUtil {
    @Autowired
    private RedisTemplate<String, Object> redisTemplate;
    // 限流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 RedisProperties.Jedis) {
                    System.out.println("单点执行");
                    return ((Jedis) nativeConnection).eval(luaScript, keys, args);
                }
                return null;
            }
        });
        return result;
    }
}

自定义的异常类

public class ServiceException extends Exception{
    private String code;
    public ServiceException(String code,String message) {
        super(message);
        this.code=code;
    }

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

    public String getCode() {
        return code;
    }

    public void setCode(String code) {
        this.code = code;
    }
}