Skip to content

Instantly share code, notes, and snippets.

@dhet
Last active March 29, 2022 10:22
Show Gist options
  • Save dhet/4ae6a3da6f2ccf5b81788e0111ad596c to your computer and use it in GitHub Desktop.
Save dhet/4ae6a3da6f2ccf5b81788e0111ad596c to your computer and use it in GitHub Desktop.
Minimal rate limiter in plain Java

Minimal Rate Limiter

A minimal rate limiter in plain Java which...

  • supports sliding windows
  • supports buckets for partitioning a rate limiter across several domains (e.g. one bucket per user)
  • provides information about when the next call will be possible (e.g. for using the Retry-After HTTP header)
  • is thread-safe

Usage

// RateLimiter allows three invocations over a sliding window of 10 seconds.
RateLimiter rateLimiter = new RateLimiter(3, Duration.ofSeconds(10));

rateLimiter.guard("1", () -> {
  // This code will only run if the rate limit has not been exceeded.
  // Otherwise, a RateLimitExceededException is thrown.
});
import java.time.Duration;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.Deque;
import java.util.concurrent.*;
public class RateLimiter {
private final int invocations;
private final Duration window;
private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
private final ConcurrentMap<Object, Deque<Instant>> buckets = new ConcurrentHashMap<>();
public RateLimiter(int invocations, Duration perTimespan) {
this.invocations = invocations;
this.window = perTimespan;
}
public void guard(Object bucketId, Runnable runnable) throws RateLimitExceededException {
Deque<Instant> deque = getOrCreateDeque(bucketId);
tryEnqueue(deque);
scheduleCleanUp(deque);
runnable.run();
}
private Deque<Instant> getOrCreateDeque(Object bucket) {
return buckets.computeIfAbsent(bucket, (k) -> new LinkedBlockingDeque<>(invocations));
}
private void scheduleCleanUp(Deque<Instant> deque) throws IllegalStateException {
scheduler.schedule((Runnable) deque::pollFirst, window.toNanos(), TimeUnit.NANOSECONDS);
}
private void tryEnqueue(Deque<Instant> deque) throws RateLimitExceededException {
try {
deque.addLast(Instant.now());
} catch (IllegalStateException e) {
Instant first = deque.peekFirst();
Instant nextOpportunity = first != null ? first.plus(window.toNanos(), ChronoUnit.NANOS) : Instant.now();
throw new RateLimitExceededException(nextOpportunity, invocations, window);
}
}
public static class RateLimitExceededException extends RuntimeException {
private final Instant nextOpportunity;
public RateLimitExceededException(Instant nextOpportunity, long invocations, Duration window) {
super(String.format("Rate limit exceeded. More than %s invocations during the last %s ms.",
invocations, window.toMillis()));
this.nextOpportunity = nextOpportunity;
}
public Instant getNextOpportunity() {
return nextOpportunity;
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment