1.令牌桶算法
令牌桶中有初始容量,每来一个请求从桶中获取一个令牌,并且在一定时间间隔中可以生成令牌,多余的令牌被丢弃。可以实现限速功能。
2. 代码实现
主要针对不同用户的请求进行限速,如果单独使用google的RateLimiter可以控制请求的速率,如果超过限定的速率则进行等待,但是无法获取用户的请求速率。如果下可以根据不同的用户进行限速。
* TokenBucket
/**
* 用于放令牌的桶
*/
public class TokenBucket {
//上次填充的时间戳,ms
private Long lastRefillTime;
//桶中剩余的token数量
private Long remainingToken;
public TokenBucket(long lastRefillTime,long remainingToken){
this.lastRefillTime = lastRefillTime;
this.remainingToken = remainingToken;
}
public Long getLastRefillTime() {
return lastRefillTime;
}
public Long getRemainingToken() {
return remainingToken;
}
public void setLastRefillTime(Long lastRefillTime) {
this.lastRefillTime = lastRefillTime;
}
public void setRemainingToken(Long remainingToken) {
this.remainingToken = remainingToken;
}
}
- RateLimitService 限速的主要方法
/**
* 限速1s内10个请求
*/
public class RateLimitService {
//每秒的最大次数
private final long permits = 10;
//1s的时间间隔
private final long intervalInMills = 1 * 1000;
//令牌的生成速率 100ms
private final long intervalPerPermit = intervalInMills / permits;
//用作缓存
public static ConcurrentMap<String, TokenBucket> bucketMap = new ConcurrentHashMap<String, TokenBucket>();
//针对某一个用户进行限速
public boolean available(String key) {
//1.如果缓存中没有相应key的tokenBucket,需要初始化一个,填充时间为当前系统时间,令牌个数为最大数减一(因为为减去这一次的请求)
if (!bucketMap.containsKey(key)) {
TokenBucket tokenBucket = new TokenBucket(System.currentTimeMillis(), permits - 1);
bucketMap.put(key, tokenBucket);
System.out.println(String.format("第一次请求完,桶中还剩%s个令牌", tokenBucket.getRemainingToken()));
return true;
} else { //2.缓存中存在
TokenBucket tokenBucket = bucketMap.get(key);
//获取桶的信息--上次填充距当前时间的间隔,用于计算可以生成多少令牌
long lastRefillTime = tokenBucket.getLastRefillTime();
long refillTime = System.currentTimeMillis();
long intervalSinceLast = refillTime - lastRefillTime;
long tokenNum;
//3.如果时间间隔大于1s,也就是要限速的时间粒度,则令牌桶重置
if (intervalSinceLast > intervalInMills) {
tokenNum = permits;
} else { //4.如果时间间隔不大于1s,则需要计算这段时间内生成的令牌个数
long generatedToken = intervalSinceLast / intervalPerPermit;
//有可能只取了一个令牌,然后一直生成,之和大于bucket的最大值
tokenNum = Math.min(permits, tokenBucket.getRemainingToken() + generatedToken);
}
//设置令牌桶中的remainingToken和填充时间
tokenBucket.setLastRefillTime(refillTime);
//桶中没有token了,并且上次获取token完距现在不足以生成一个token
if (tokenNum == 0) {
tokenBucket.setRemainingToken(tokenNum);
System.out.println("桶中的token已经消费完毕,您已超速。");
bucketMap.put(key,tokenBucket);
return false;
} else {
tokenBucket.setRemainingToken(tokenNum - 1);
System.out.println(String.format("桶中还剩%s个令牌", tokenBucket.getRemainingToken()));
bucketMap.put(key,tokenBucket);
return true;
}
}
}
}
- RateLimiteServiceTest
public class RateLimitServiceTest {
public static RateLimitService rateLimitService = new RateLimitService();
public static void main(String[] args) {
while (true) {
rateLimitService.available("bj");
try {
//1.始终有9个,没有超速
TimeUnit.MILLISECONDS.sleep(100);
/*
第一次请求完,桶中还剩9个令牌
桶中还剩9个令牌
桶中还剩9个令牌
桶中还剩9个令牌
桶中还剩9个令牌
桶中还剩9个令牌
桶中还剩9个令牌
*/
//2.超速
//TimeUnit.MILLISECONDS.sleep(10);
/*
第一次请求完,桶中还剩9个令牌
桶中还剩8个令牌
桶中还剩7个令牌
桶中还剩6个令牌
桶中还剩5个令牌
桶中还剩4个令牌
桶中还剩3个令牌
桶中还剩2个令牌
桶中还剩1个令牌
桶中还剩0个令牌
桶中的token已经消费完毕,您已超速。
*/
} catch (InterruptedException e) {
e.printStackTrace();
}
}
}
}
3 主要应用场景
- google的RateLimiter有用到令牌桶算法
- SparkStreaming的backpressure机制中同样用到了令牌桶的算法