FrameTour-BE/src/main/java/com/ycwl/basic/utils/ratelimiter/SlidingWindowRateLimiter.java
2025-04-05 16:42:46 +08:00

50 lines
1.6 KiB
Java

package com.ycwl.basic.utils.ratelimiter;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
public class SlidingWindowRateLimiter implements IRateLimiter {
private final Semaphore semaphore;
private final ScheduledExecutorService scheduler = new ScheduledThreadPoolExecutor(1);
public SlidingWindowRateLimiter(int maxRequestsPerSecond) {
this.semaphore = new Semaphore(maxRequestsPerSecond);
scheduler.scheduleAtFixedRate(() -> {
if (semaphore.availablePermits() < maxRequestsPerSecond) {
semaphore.release(maxRequestsPerSecond - semaphore.availablePermits());
}
}, 1, 1, TimeUnit.SECONDS);
}
public SlidingWindowRateLimiter(int maxRequests, int perSecond) {
this.semaphore = new Semaphore(maxRequests);
scheduler.scheduleAtFixedRate(() -> {
if (semaphore.availablePermits() < maxRequests) {
semaphore.release(maxRequests - semaphore.availablePermits());
}
}, perSecond, perSecond, TimeUnit.SECONDS);
}
@Override
public void acquire() throws InterruptedException {
semaphore.acquire();
}
@Override
public void shutdown() {
scheduler.shutdown();
}
@Override
public boolean tryAcquire() {
return semaphore.tryAcquire();
}
@Override
public boolean tryAcquire(long timeout, TimeUnit unit) throws InterruptedException {
return semaphore.tryAcquire(timeout, unit);
}
}