rateLimit 限流使用

需求背景

目前接口有个需求,需要对接口进行限流或者,指定流量可以走指定方法;
例如:我现在写了一个缓存接口,因为缓存涉及接口比较多,害怕突然上线导致崩盘,就设置了两个开关一个总开关直接关闭全部缓存走db,一个限流开关,便于排查问题,可以限制一小部分流量进入

@Component
public class RateLimitHelper {


    /**
     * key : RateLimiterName ,value : 限流器实例
     */
    private final ConcurrentHashMap<String, RateLimiter> limitMap = new ConcurrentHashMap<>();

    /**
     * key : RateLimiterName ,value : 限流器令牌配置
     */
    private final ConcurrentHashMap<String, Integer> limitUserConfigMap = new ConcurrentHashMap<>();

    @PostConstruct
    void init() {
        // 定时线程用于动态更新限流数量,相当于改了限流数量不用重启项目,让其生效,定时线程一分钟执行一次
        ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
        executor.scheduleAtFixedRate(() -> {
            limitMap.forEach((key, rateLimiter) -> {
                Integer currentUserLimitToken = limitUserConfigMap.get(key);
                if (Objects.nonNull(currentUserLimitToken)) {
                    if (rateLimiter.getRate() > 0 && rateLimiter.getRate() != currentUserLimitToken) {
                        limitMap.put(key, RateLimiter.create(currentUserLimitToken));
                    }
                }
            });
        }, 10, 60, TimeUnit.SECONDS);
    }


    /**
     * 获取限流器
     *
     * @param limiterName    限流器名称
     * @param userLimitToken 限流数量
     * @return boolean true:获取成功 false:获取失败
     */
    public boolean tryAcquire(String limiterName, int userLimitToken) {
        boolean result = false;
        if (StringUtils.isBlank(limiterName)) {
            return false;
        }
        RateLimiter rateLimiter = limitMap.get(limiterName);
        if (rateLimiter != null && rateLimiter.tryAcquire()) {
            result = true;
        } else if (rateLimiter == null) {
            limitMap.put(limiterName, RateLimiter.create(userLimitToken));
        }
         limitUserConfigMap.put(limiterName, userLimitToken);
        return result;
    }

    public enum RateLimiterName {
        /**
         * 限流分类
         */
        USER_FIND_BY_IDS("findUserByIds"),

        /**
         * 通过前缀匹配限流枚举
         */
        RESOURCE_ADOPT_FIND_CODE_PREFIX("resourceAdoptFindCodePrefix"),

        /**
         * 通过code和id匹配的限流枚举
         */
        RESOURCE_ADOPT_FIND_CODE_OR_ID("resourceAdoptFindCodeOrId");

        private final String value;

        RateLimiterName(String value) {
            this.value = value;
        }

        public String getValue() {
            return value;
        }
    }

}

业务调用

    /**
     * 缓存总开关
     */
    @Value("${openCache:true}")
    private boolean openResourceCache;

    /**
     * 获取code前缀每秒限制qps数量
     */
    @Value("${codePrefix.rateLimitMaxToken: 100}")
    private int codePrefixLimitToken;

    /**
     * 获取code或id每秒限制qps数量
     */
    @Value("${codeOrId.rateLimitMaxToken: 100}")
    private int codeOrIdLimitToken;

 public List<TestDO> getCodeListPrefix(Collection<String> codeListPrefix, List<ResourceTypeEnum> resourceTypes) {
 		// 获取是否可以使用
        boolean isCache = isOpenResourceCache(RateLimiterName.RESOURCE_ADOPT_FIND_CODE_PREFIX);
        if (isCache) {
            return cache.getCodeListPrefix(codeListPrefix, resourceTypes);
        } else {
            return testDAO.selectListPrefix(codeListPrefix, resourceTypes);
        }
    }

    private boolean isOpenResourceCache(RateLimiterName rateLimiterNameEnum) {
        // 如果总开关关闭直接不走缓存或者限流器为空不走缓存
        if (!openResourceCache || Objects.isNull(rateLimiterNameEnum)) {
            return false;
        }
        // 获取限流次数
        int limitToken = RateLimiterName.RESOURCE_ADOPT_FIND_CODE_PREFIX.equals(rateLimiterNameEnum)
                ? codePrefixLimitToken
                : codeOrIdLimitToken;

        return rateLimitHelper.tryAcquire(rateLimiterNameEnum.getValue(), limitToken);
    }

注意事项

具体的限流策略根据自己的业务需求来