From 1db69cad40416f17bb7071b395560d1a1215ba00 Mon Sep 17 00:00:00 2001 From: Craig Baker Date: Sat, 1 Sep 2018 11:01:41 +1000 Subject: [PATCH] Support key matching rules in Hazlecast and Inmemory implementations --- .../limiter/request/RequestLimitRule.java | 9 +++---- .../DefaultRequestLimitRulesSupplierTest.java | 5 ++-- ...elcastSlidingWindowRequestRateLimiter.java | 15 ++++++------ ...MemorySlidingWindowRequestRateLimiter.java | 24 +++++++++---------- .../AbstractSyncRequestRateLimiterTest.java | 19 ++++++++++++++- 5 files changed, 43 insertions(+), 29 deletions(-) diff --git a/ratelimitj-core/src/main/java/es/moki/ratelimitj/core/limiter/request/RequestLimitRule.java b/ratelimitj-core/src/main/java/es/moki/ratelimitj/core/limiter/request/RequestLimitRule.java index edf9673..7a44776 100644 --- a/ratelimitj-core/src/main/java/es/moki/ratelimitj/core/limiter/request/RequestLimitRule.java +++ b/ratelimitj-core/src/main/java/es/moki/ratelimitj/core/limiter/request/RequestLimitRule.java @@ -1,12 +1,12 @@ package es.moki.ratelimitj.core.limiter.request; +import javax.annotation.Nullable; import javax.annotation.ParametersAreNonnullByDefault; import java.time.Duration; import java.util.Arrays; import java.util.HashSet; import java.util.Objects; import java.util.Set; -import java.util.concurrent.TimeUnit; import static java.util.Objects.requireNonNull; @@ -78,9 +78,9 @@ public RequestLimitRule withName(String name) { * @param keys Defines a set of keys to which the rule applies. * @return a limit rule */ - public RequestLimitRule withKeys(String... keys) { + public RequestLimitRule matchingKeys(String... keys) { Set keySet = keys.length > 0 ? new HashSet<>(Arrays.asList(keys)) : null; - return withKeys(keySet); + return matchingKeys(keySet); } /** @@ -89,7 +89,7 @@ public RequestLimitRule withKeys(String... keys) { * @param keys Defines a set of keys to which the rule applies. * @return a limit rule */ - public RequestLimitRule withKeys(Set keys) { + public RequestLimitRule matchingKeys(Set keys) { return new RequestLimitRule(this.durationSeconds, this.limit, this.precision, this.name, keys); } @@ -124,6 +124,7 @@ public long getLimit() { /** * @return The keys. */ + @Nullable public Set getKeys() { return keys; } diff --git a/ratelimitj-core/src/test/java/es/moki/ratelimitj/core/limiter/request/DefaultRequestLimitRulesSupplierTest.java b/ratelimitj-core/src/test/java/es/moki/ratelimitj/core/limiter/request/DefaultRequestLimitRulesSupplierTest.java index 0523f3e..b7ce44d 100644 --- a/ratelimitj-core/src/test/java/es/moki/ratelimitj/core/limiter/request/DefaultRequestLimitRulesSupplierTest.java +++ b/ratelimitj-core/src/test/java/es/moki/ratelimitj/core/limiter/request/DefaultRequestLimitRulesSupplierTest.java @@ -5,7 +5,6 @@ import java.time.Duration; import java.util.HashSet; import java.util.Set; -import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; import static org.assertj.core.api.Assertions.assertThat; @@ -18,9 +17,9 @@ class DefaultRequestLimitRulesSupplierTest { DefaultRequestLimitRulesSupplierTest() { allRules.add(RequestLimitRule.of(Duration.ofSeconds(1), 10).withName("localhostPerSeconds") - .withKeys("localhost", "127.0.0.1")); + .matchingKeys("localhost", "127.0.0.1")); allRules.add(RequestLimitRule.of(Duration.ofHours(1), 2000).withName("localhostPerHours") - .withKeys("localhost", "127.0.0.1")); + .matchingKeys("localhost", "127.0.0.1")); allRules.add(RequestLimitRule.of(Duration.ofSeconds(1), 5).withName("perSeconds")); allRules.add(RequestLimitRule.of(Duration.ofHours(1), 1000).withName("perHours")); requestLimitRulesSupplier = new DefaultRequestLimitRulesSupplier(allRules); diff --git a/ratelimitj-hazelcast/src/main/java/es/moki/ratelimitj/hazelcast/HazelcastSlidingWindowRequestRateLimiter.java b/ratelimitj-hazelcast/src/main/java/es/moki/ratelimitj/hazelcast/HazelcastSlidingWindowRequestRateLimiter.java index 27a03b3..c6b4479 100644 --- a/ratelimitj-hazelcast/src/main/java/es/moki/ratelimitj/hazelcast/HazelcastSlidingWindowRequestRateLimiter.java +++ b/ratelimitj-hazelcast/src/main/java/es/moki/ratelimitj/hazelcast/HazelcastSlidingWindowRequestRateLimiter.java @@ -3,6 +3,7 @@ import com.hazelcast.config.MapConfig; import com.hazelcast.core.HazelcastInstance; import com.hazelcast.core.IMap; +import es.moki.ratelimitj.core.limiter.request.DefaultRequestLimitRulesSupplier; import es.moki.ratelimitj.core.limiter.request.RequestLimitRule; import es.moki.ratelimitj.core.limiter.request.RequestRateLimiter; import es.moki.ratelimitj.core.time.SystemTimeSupplier; @@ -26,7 +27,7 @@ public class HazelcastSlidingWindowRequestRateLimiter implements RequestRateLimi private static final Logger LOG = LoggerFactory.getLogger(HazelcastSlidingWindowRequestRateLimiter.class); private final HazelcastInstance hz; - private final Set rules; + private final DefaultRequestLimitRulesSupplier rulesSupplier; private final TimeSupplier timeSupplier; public HazelcastSlidingWindowRequestRateLimiter(HazelcastInstance hz, Set rules) { @@ -36,9 +37,12 @@ public HazelcastSlidingWindowRequestRateLimiter(HazelcastInstance hz, Set rules, TimeSupplier timeSupplier) { requireNonNull(hz, "hazelcast can not be null"); requireNonNull(rules, "rules can not be null"); + if (rules.isEmpty()) { + throw new IllegalArgumentException("at least one rule must be provided"); + } requireNonNull(rules, "time supplier can not be null"); this.hz = hz; - this.rules = rules; + this.rulesSupplier = new DefaultRequestLimitRulesSupplier(rules); this.timeSupplier = timeSupplier; } @@ -94,13 +98,8 @@ private IMap getMap(String key, int longestDuration) { private boolean eqOrGeLimit(String key, int weight, boolean strictlyGreater) { - requireNonNull(key, "key cannot be null"); - requireNonNull(rules, "rules cannot be null"); - if (rules.isEmpty()) { - throw new IllegalArgumentException("at least one rule must be provided"); - } - final long now = timeSupplier.get(); + final Set rules = rulesSupplier.getRules(key); // TODO implement cleanup final int longestDuration = rules.stream().map(RequestLimitRule::getDurationSeconds).reduce(Integer::max).orElse(0); List savedKeys = new ArrayList<>(rules.size()); diff --git a/ratelimitj-inmemory/src/main/java/es/moki/ratelimitj/inmemory/request/InMemorySlidingWindowRequestRateLimiter.java b/ratelimitj-inmemory/src/main/java/es/moki/ratelimitj/inmemory/request/InMemorySlidingWindowRequestRateLimiter.java index 32cfeb4..53c978b 100644 --- a/ratelimitj-inmemory/src/main/java/es/moki/ratelimitj/inmemory/request/InMemorySlidingWindowRequestRateLimiter.java +++ b/ratelimitj-inmemory/src/main/java/es/moki/ratelimitj/inmemory/request/InMemorySlidingWindowRequestRateLimiter.java @@ -2,7 +2,9 @@ import de.jkeylockmanager.manager.KeyLockManager; import de.jkeylockmanager.manager.KeyLockManagers; +import es.moki.ratelimitj.core.limiter.request.DefaultRequestLimitRulesSupplier; import es.moki.ratelimitj.core.limiter.request.RequestLimitRule; +import es.moki.ratelimitj.core.limiter.request.RequestLimitRulesSupplier; import es.moki.ratelimitj.core.limiter.request.RequestRateLimiter; import es.moki.ratelimitj.core.time.SystemTimeSupplier; import es.moki.ratelimitj.core.time.TimeSupplier; @@ -25,10 +27,10 @@ public class InMemorySlidingWindowRequestRateLimiter implements RequestRateLimit private static final Logger LOG = LoggerFactory.getLogger(InMemorySlidingWindowRequestRateLimiter.class); - private final Set rules; private final TimeSupplier timeSupplier; private final ExpiringMap> expiringKeyMap; private final KeyLockManager lockManager = KeyLockManagers.newLock(); + private final DefaultRequestLimitRulesSupplier rulesSupplier; public InMemorySlidingWindowRequestRateLimiter(RequestLimitRule rule) { this(Collections.singleton(rule), new SystemTimeSupplier()); @@ -39,17 +41,18 @@ public InMemorySlidingWindowRequestRateLimiter(Set rules) { } public InMemorySlidingWindowRequestRateLimiter(Set rules, TimeSupplier timeSupplier) { - requireNonNull(rules, "rules can not be null"); - requireNonNull(rules, "time supplier can not be null"); - this.rules = rules; - this.timeSupplier = timeSupplier; - this.expiringKeyMap = ExpiringMap.builder().variableExpiration().build(); + this(ExpiringMap.builder().variableExpiration().build(), rules, timeSupplier); } InMemorySlidingWindowRequestRateLimiter(ExpiringMap> expiringKeyMap, Set rules, TimeSupplier timeSupplier) { + requireNonNull(rules, "rules can not be null"); + requireNonNull(rules, "time supplier can not be null"); + if (rules.isEmpty()) { + throw new IllegalArgumentException("at least one rule must be provided"); + } this.expiringKeyMap = expiringKeyMap; - this.rules = rules; this.timeSupplier = timeSupplier; + this.rulesSupplier = new DefaultRequestLimitRulesSupplier(rules); } @Override @@ -102,13 +105,8 @@ private ConcurrentMap getMap(String key, int longestDuration) { private boolean eqOrGeLimit(String key, int weight, boolean strictlyGreater) { - requireNonNull(key, "key cannot be null"); - requireNonNull(rules, "rules cannot be null"); - if (rules.isEmpty()) { - throw new IllegalArgumentException("at least one rule must be provided"); - } - final long now = timeSupplier.get(); + final Set rules = rulesSupplier.getRules(key); // TODO implement cleanup final int longestDurationSeconds = rules.stream().map(RequestLimitRule::getDurationSeconds).reduce(Integer::max).orElse(0); List savedKeys = new ArrayList<>(rules.size()); diff --git a/ratelimitj-test/src/main/java/es/moki/ratelimitj/test/limiter/request/AbstractSyncRequestRateLimiterTest.java b/ratelimitj-test/src/main/java/es/moki/ratelimitj/test/limiter/request/AbstractSyncRequestRateLimiterTest.java index 6f2827e..81e4333 100644 --- a/ratelimitj-test/src/main/java/es/moki/ratelimitj/test/limiter/request/AbstractSyncRequestRateLimiterTest.java +++ b/ratelimitj-test/src/main/java/es/moki/ratelimitj/test/limiter/request/AbstractSyncRequestRateLimiterTest.java @@ -9,7 +9,6 @@ import java.time.Duration; import java.util.Set; -import java.util.concurrent.TimeUnit; import java.util.stream.IntStream; import static org.assertj.core.api.Assertions.assertThat; @@ -83,6 +82,24 @@ void shouldLimitSingleWindowSyncWithMultipleKeys() { keySuffix -> assertThat(requestRateLimiter.overLimitWhenIncremented("ip:127.0.0." + keySuffix)).isFalse()); } + @Test + void shouldLimitSingleWindowSyncWithKeySpecificRules() { + + RequestLimitRule rule1 = RequestLimitRule.of(Duration.ofSeconds(10), 5).matchingKeys("ip:127.9.0.0"); + RequestLimitRule rule2 = RequestLimitRule.of(Duration.ofSeconds(10), 10); + + RequestRateLimiter requestRateLimiter = getRateLimiter(ImmutableSet.of(rule1, rule2), timeBandit); + + IntStream.rangeClosed(1, 5).forEach(value -> { + timeBandit.addUnixTimeMilliSeconds(1000L); + assertThat(requestRateLimiter.overLimitWhenIncremented("ip:127.9.0.0")).isFalse(); + }); + assertThat(requestRateLimiter.overLimitWhenIncremented("ip:127.9.0.0")).isTrue(); + + IntStream.rangeClosed(1, 10).forEach(value -> assertThat(requestRateLimiter.overLimitWhenIncremented("ip:127.9.1.0")).isFalse()); + assertThat(requestRateLimiter.overLimitWhenIncremented("ip:127.9.1.0")).isTrue(); + } + @Test void shouldResetLimit() { ImmutableSet rules = ImmutableSet.of(RequestLimitRule.of(Duration.ofSeconds(60), 1));