Skip to content

Instantly share code, notes, and snippets.

@JohnnyJayJay
Last active August 27, 2019 16:53
Show Gist options
  • Save JohnnyJayJay/794f30111ae21c96786a09f576f25607 to your computer and use it in GitHub Desktop.
Save JohnnyJayJay/794f30111ae21c96786a09f576f25607 to your computer and use it in GitHub Desktop.
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);
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment