-
-
Save dlsf/49028a527ebc4053049af35e2ca92a21 to your computer and use it in GitHub Desktop.
RateLimiter
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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(Function<T, ?> keyMapper, long rateLimit, TimeUnit timeUnit) { | |
return new RateLimiter<>(TimeUnit.MILLISECONDS.convert(rateLimit, timeUnit), keyMapper); | |
} | |
public 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