java节流 限速 速度控制

基于令牌桶的速度控制算法

/**
 * 基于时间的令牌桶算法，结合ring buffer, 实现精确节流策略。效率极高。
 */
public class TokenLimiter {
	private long[] timeTokenArr;
	private int index = 0;
	// 最后一次尝试时间
	private long latestTime;
	public TokenLimiter(int tokens) {
		if (tokens < 1) {
			tokens = 1;
		}
		this.timeTokenArr = new long[tokens];
	}
	
	public boolean acquire(long duration) {
		int capacity = timeTokenArr.length;
		boolean r = tryUse(getIndex(index, capacity), duration);
		if (r) {
			index++;
			if (index > 100000000) {
				index = getIndex(index, capacity);
			}
		}
		return r;
	}
	
	private boolean tryUse(int index, long durationMillis) {
		long t = System.currentTimeMillis();
		latestTime = t;
		// 当前时间小于 下一个 令牌的可开始使用时间
		if (t < timeTokenArr[index]) {
			return false;
		}
		// 允许使用，并设置下次(下个周期)允许获取时间
		timeTokenArr[index] = t + durationMillis;
		return true;
	}
	
	private int getIndex(int index, int capacity) {
		return index % capacity;
	}
	
	public void reset() {
		for (int i = 0; i < timeTokenArr.length; i++) {
			timeTokenArr[i] = 0L;
		}
	}
	
	/**
	 * 如果 当前时间 - 最后尝试获取令牌时间 (无论成功与否) > 令牌周期，则认为是过期
	 * @param durationMillis 令牌周期
	 */
	public boolean isExpiry(long durationMillis) {
		return System.currentTimeMillis() - latestTime > durationMillis;
	}
	
	/** 最后一次尝试获取令牌时间 */
	public long getLatestTime() {
		return latestTime;
	}
}

关联用户信息

import java.time.Duration;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;


/**
 * 基于用户id的节流控制算法
 */
public class TokenLimiterOfUid {
	private Map<String, TokenLimiter> cache = null;
	private int capacity;
	private long duration;
	
	private ScheduledThreadPoolExecutor cleanExecutor = new ScheduledThreadPoolExecutor(1);
	
	/**
	 * @param capacity 数量
	 * @param period 计算周期
	 */
	public TokenLimiterOfUid(int capacity, Duration period) {
		this(1000, capacity, period);
	}
	
	/**
	 * @param initUserAmount 初始用户数量
	 * @param capacity 数量
	 * @param period 计算周期
	 */
	public TokenLimiterOfUid(int initUserAmount, int capacity, Duration period) {
		cache = new ConcurrentHashMap<String, TokenLimiter>(initUserAmount);
		
		this.capacity = capacity;
		this.duration = period.toMillis();
		
		// 1分钟检查一次，清理没有事件的uid
		cleanExecutor.scheduleAtFixedRate(() -> {
			cache.forEach((key, value) -> {
				if (!value.isExpiry(duration)) {
					return;
				}
				cache.remove(key);
			});
		}, 1, 1, TimeUnit.MINUTES);
	}
	
	/**
	 * 非阻塞方法，线程安全
	 * @param uid 用户id
	 */
	public boolean acquire(String uid) {
		return get(uid).acquire(duration);
	}
	
	private TokenLimiter get(String uid) {
		TokenLimiter limiter = cache.get(uid);
		if (limiter != null) {
			return limiter;
		}
		synchronized (this) {
			limiter = cache.get(uid);
			if (limiter != null) {
				return limiter;
			}
			limiter = new TokenLimiter(capacity);
			cache.put(uid, limiter);
		}
		return limiter;
	}
	
	public void reset(String uid) {
		get(uid).reset();
	}
	
	public synchronized void resetTokens(int tokens) {
		cache.forEach((key, value) -> {
			cache.put(key, new TokenLimiter(tokens));
		});
	}
	
	public void clear() {
		cache.clear();
	}
	
}

用法

TokenLimiterOfUid limiter = new TokenLimiterOfUid(10, Duration.ofSeconds(5));
boolean r = limiter.acquire("test");

或

TokenLimiter limiter = new TokenLimiter(10);
long duration = Duration.ofSeconds(5).toMillis();
boolean r = limiter.acquire(duration);