Minimal rate limiter in plain Java

_Readme.md

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.
});

RateLimiter.java

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;
    }
  }

}