Support key matching rules in Hazlecast and Inmemory implementations

mokies · Sep 1, 2018 · 1db69ca · 1db69ca
1 parent 7f52a9b
commit 1db69ca
Show file tree

Hide file tree

Showing 5 changed files with 43 additions and 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
@@ -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<String> 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<String> keys) {
+    public RequestLimitRule matchingKeys(Set<String> 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<String> getKeys() {
         return keys;
     }

diff --git a/...st/java/es/moki/ratelimitj/core/limiter/request/DefaultRequestLimitRulesSupplierTest.java b/...st/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/.../src/main/java/es/moki/ratelimitj/hazelcast/HazelcastSlidingWindowRequestRateLimiter.java b/.../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<RequestLimitRule> rules;
+    private final DefaultRequestLimitRulesSupplier rulesSupplier;
     private final TimeSupplier timeSupplier;
 
     public HazelcastSlidingWindowRequestRateLimiter(HazelcastInstance hz, Set<RequestLimitRule> rules) {
@@ -36,9 +37,12 @@ public HazelcastSlidingWindowRequestRateLimiter(HazelcastInstance hz, Set<Reques
     public HazelcastSlidingWindowRequestRateLimiter(HazelcastInstance hz, Set<RequestLimitRule> 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<String, Long> 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<RequestLimitRule> rules = rulesSupplier.getRules(key);
         // TODO implement cleanup
         final int longestDuration = rules.stream().map(RequestLimitRule::getDurationSeconds).reduce(Integer::max).orElse(0);
         List<SavedKey> savedKeys = new ArrayList<>(rules.size());

diff --git a/...ain/java/es/moki/ratelimitj/inmemory/request/InMemorySlidingWindowRequestRateLimiter.java b/...ain/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<RequestLimitRule> rules;
     private final TimeSupplier timeSupplier;
     private final ExpiringMap<String, ConcurrentMap<String, Long>> 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<RequestLimitRule> rules) {
     }
 
     public InMemorySlidingWindowRequestRateLimiter(Set<RequestLimitRule> 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<String, ConcurrentMap<String, Long>> expiringKeyMap, Set<RequestLimitRule> 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<String, Long> 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<RequestLimitRule> rules = rulesSupplier.getRules(key);
         // TODO implement cleanup
         final int longestDurationSeconds = rules.stream().map(RequestLimitRule::getDurationSeconds).reduce(Integer::max).orElse(0);
         List<SavedKey> savedKeys = new ArrayList<>(rules.size());

diff --git a/...main/java/es/moki/ratelimitj/test/limiter/request/AbstractSyncRequestRateLimiterTest.java b/...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<RequestLimitRule> rules = ImmutableSet.of(RequestLimitRule.of(Duration.ofSeconds(60), 1));