io.github.dengchen2020.ratelimiter.local.RateLimiter Maven / Gradle / Ivy
package io.github.dengchen2020.ratelimiter.local;
import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import com.github.benmanes.caffeine.cache.Scheduler;
import java.time.Duration;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
/**
* 限流工具
*
* @author dengchen
* @since 2024/4/18
*/
public class RateLimiter {
Cache cache;
ScheduledExecutorService scheduledExecutorService = Executors.newScheduledThreadPool(1, Thread.ofVirtual().name("rate-limiter").factory());
/**
* 限流实例化
*
* @param duration 多少时间内限制多少次
*/
public RateLimiter(Duration duration) {
if (duration.getSeconds() < 1) duration = Duration.ofSeconds(1);
cache = Caffeine.newBuilder()
.scheduler(Scheduler.forScheduledExecutorService(scheduledExecutorService))
.maximumSize(50000)
.expireAfterWrite(duration.getSeconds(), TimeUnit.SECONDS)
.build();
}
/**
* 限流实例化
* @param cache Cache实例
*
*/
public RateLimiter(Cache cache) {
this.cache = cache;
}
/**
* 是否被限制
*
* @param limitKey 限制标识符
* @param limitNum 限制的次数
* @return true：被限制 false：未被限制
*/
public boolean limit(String limitKey, int limitNum) {
AtomicInteger time = cache.get(limitKey, key -> new AtomicInteger(1));
return time.getAndIncrement() > limitNum;
}
/**
* 撤销限制
*
* @param limitKey 限制标识符
*/
public void quashLimit(String limitKey) {
AtomicInteger time = cache.getIfPresent(limitKey);
if (time == null) return;
time.decrementAndGet();
}
}