jnpf-common-v6x/jnpf-boot-common/jnpf-common-ai/target/classes/jnpf/util/AiLimitUtil.java

package jnpf.util;

import cn.hutool.cache.CacheUtil;
import cn.hutool.cache.impl.TimedCache;
import jnpf.config.AiProperties;

import java.util.concurrent.atomic.AtomicInteger;


/**
 * AI接口限流
 *
 * @author JNPF开发平台组
 * @copyright 引迈信息技术有限公司
 * @date 2025/3/7 10:00
 */
public class AiLimitUtil {

    private static AiProperties aiProperties;

    private static TimedCache<String, AtomicInteger> limit = null;

    private static final String TOTAL_KEY = "ai_limit_total";

    public AiLimitUtil(AiProperties aiProperties) {
        AiLimitUtil.aiProperties = aiProperties;
        if (limit == null) {
            limit = CacheUtil.newTimedCache(aiProperties.getUserLimitTime().toMillis());
            // 一分钟清理一次无用数据
            limit.schedulePrune(1000L * 60);
        }
    }

    /**
     * 是否可以请求
     * @param userId
     * @return
     */
    public static boolean tryAcquire(String userId) {
        if(limit == null) {
            return true;
        }
        if (StringUtil.isNotEmpty(userId)) {
            // 按用户限制
            AtomicInteger userCount = limit.get(userId, false, AtomicInteger::new);
            if (userCount.incrementAndGet() > aiProperties.getUserLimitCount()) {
                return false;
            }
        }
        // 所有请求限制
        AtomicInteger totalCount = limit.get(TOTAL_KEY, false, aiProperties.getTotalLimitTime().toMillis(), AtomicInteger::new);
        return totalCount.incrementAndGet() <= aiProperties.getTotalLimitCount();
    }
}