RateLimiter.java

import java.util.concurrent.TimeUnit;
import java.util.Map;
import java.util.HashMap;
import java.util.function.Function;

public class RateLimiter<T> {
  
  private final long rateLimit;
  private final Function<T, ?> keyMapper;
  private final Map<?, Long> timestamps;
  
  private RateLimiter(long rateLimit, Function<T, ?> keyMapper) {
    this.rateLimit = rateLimit;
    this.keyMapper = keyMapper;
    this.timestamps = new HashMap<>();
  }
  
  public static <T> RateLimiter<T> create(long rateLimit, TimeUnit timeUnit) {
    return create(Function.identity(), rateLimit, timeUnit);
  }
  
  public static <T> RateLimiter<T> create(Function<T, ?> keyMapper, long rateLimit, TimeUnit timeUnit) {
    return new RateLimiter<>(TimeUnit.MILLISECONDS.convert(rateLimit, timeUnit), keyMapper);
  }
  
  public void coolDown(T t) {
    timestamps.put(keyMapper.apply(t), System.currentTimeMillis());
  }
  
  public boolean isOnCooldown(T t) {
    return System.currentTimeMillis() - timestamps.getOrDefault(keyMapper.apply(t), 0L) >= rateLimit;
  }
  
  public void resetCooldown(T t) {
    timestamps.computeIfPresent(keyMapper.apply(t), (k, v) -> 0L);
  }
}