public class DefaultTPSLimiter implements TPSLimiter {
private final ConcurrentMap<String, StatItem> stats = new ConcurrentHashMap<String, StatItem>();
@Override
public boolean isAllowable(URL url, Invocation invocation) {
// 时间间隔内的上限值
int rate = url.getMethodParameter(invocation.getMethodName(), TPS_LIMIT_RATE_KEY, -1);
// 时间间隔
long interval = url.getMethodParameter(invocation.getMethodName(), TPS_LIMIT_INTERVAL_KEY, DEFAULT_TPS_LIMIT_INTERVAL);
String serviceKey = url.getServiceKey();
if (rate > 0) {
StatItem statItem = stats.get(serviceKey);
if (statItem == null) {
stats.putIfAbsent(serviceKey, new StatItem(serviceKey, rate, interval));
statItem = stats.get(serviceKey);
} else {
// 如果配置有修改,则直接覆盖。
// 这里不直接修改字段,而是重新生成一个不可变对象。(值得学习)
if (statItem.getRate() != rate || statItem.getInterval() != interval) {
stats.put(serviceKey, new StatItem(serviceKey, rate, interval));
statItem = stats.get(serviceKey);
}
}
return statItem.isAllowable();
} else {
// 上限值为负数则不限流
StatItem statItem = stats.get(serviceKey);
if (statItem != null) {
// ? 不太理解,为啥删除前还有判断一下
stats.remove(serviceKey);
}
}
return true;
}
}
class StatItem {
private final String name;
private final AtomicLong lastResetTime;
private final long interval;
private final AtomicInteger token;
private final int rate;
StatItem(String name, int rate, long interval) {
this.name = name;
this.rate = rate;
this.interval = interval;
this.lastResetTime = new AtomicLong(System.currentTimeMillis());
this.token = new AtomicInteger(rate);
}
public boolean isAllowable() {
long now = System.currentTimeMillis();
// 重置时间间隔内的计数
if (now > lastResetTime.get() + interval) {
token.set(rate);
lastResetTime.set(now);
}
return token.decrementAndGet() >= 0;
}
public long getInterval() {
return interval;
}
public int getRate() {
return rate;
}
long getLastResetTime() {
return lastResetTime.get();
}
int getToken() {
return token.get();
}
@Override
public String toString() {
return "StatItem " +
"[name=" + name + ", " +
"rate = " + rate + ", " +
"interval = " + interval + ']';
}
}