feat(rate-limiter): cleanup

Author: THAVEAU Alexis

lua/dshm.lua

@@ -28,7 +28,7 @@ local function escape(key)
     return table.concat(result, ":")
 end
 
-function _M.new(self, opts)
+function _M.new(_, opts)
     local sock, err = tcp()
     if not sock then
         return nil, err
@@ -88,7 +88,8 @@ local function read_response_line(self, data)
         ngx.log(ngx.DEBUG, "PROTOCOL ERROR")
         return nil, _M.PROTOCOL_ERROR
     else
-        local data, err = sock:receive(data)
+        local err
+        data, err = sock:receive(data)
         if not data then
             if err == "timeout" then
                 sock:close()
@@ -97,7 +98,8 @@ local function read_response_line(self, data)
         end
         ngx.log(ngx.DEBUG, "RECEIVE : ", data)
         -- Discard trailing \r\n
-        local trail, err = sock:receive()
+        local trail
+        trail, err = sock:receive()
         if not trail then
             if err == "timeout" then
                 sock:close()
@@ -128,7 +130,7 @@ function _M.get(self, key)
     if resp == "LEN" then
         resp, data = read_response_line(self, data)
         if resp == "DATA" then
-            local resp, _ = read_response_line(self)
+            resp = read_response_line(self)
             if resp == "DONE" then
                 return data, nil
             else
@@ -203,7 +205,7 @@ function _M.incr(self, key, value, init, init_ttl)
     if resp == "LEN" then
         resp, data = read_response_line(self, data)
         if resp == "DATA" then
-            local resp, _ = read_response_line(self)
+            resp = read_response_line(self)
             if resp == "DONE" then
                 return data, nil
             else
@@ -247,7 +249,7 @@ function _M.set(self, key, value, exptime)
     if resp == "LEN" then
         resp, data = read_response_line(self, data)
         if resp == "DATA" then
-            local resp, _ = read_response_line(self)
+            resp = read_response_line(self)
             if resp == "DONE" then
                 return data, nil
             else
@@ -297,6 +299,58 @@ function _M.touch(self, key, exptime)
 
 end
 
+---
+--- Allow to manage rate limit on sliding window. Rate limiter will try to 'consume' a token and return the remaining token available.
+--- If the quota is exceeded, this method will return nil, "rejected"
+--- Otherwhise return the remaining token available in the window
+---
+---@param self self the dshm instance
+---@param self string the key
+---@param capacity number the tokens capacity
+---@param duration number the sliding window in seconds
+---@return number the remaining tokens available or nil if quota is exceeded
+---@return string nil or error. Error code is rejected when quota is exceeded
+---
+function _M.rate_limiter(self, key, capacity, duration)
+
+    ngx.log(ngx.DEBUG, "rate_limiter : ", key, ", capacity : ", capacity, ", duration : ", duration)
+
+    local sock = self.sock
+    if not sock then
+        return nil, "not initialized"
+    end
+
+    local command = "rate_limiter " .. self.escape_key(key) .. " " .. capacity .. " " .. duration .. "\r\n"
+    local bytes, err = sock:send(command)
+    if not bytes then
+        return nil, err
+    end
+
+    local resp, data = read_response_line(self)
+    if resp == "LEN" then
+        resp, data = read_response_line(self, data)
+        if resp == "DATA" then
+            resp, _ = read_response_line(self)
+            if resp == "DONE" then
+                if data == "-1" then
+                    return nil, "rejected"
+                else
+                    return data, nil
+                end
+            else
+                return nil, _M.PROTOCOL_ERROR
+            end
+        else
+            return nil, _M.PROTOCOL_ERROR
+        end
+    elseif resp == "ERROR" then
+        return nil, data
+    else
+        return nil, _M.PROTOCOL_ERROR
+    end
+
+end
+
 function _M.quit(self)
     local sock = self.sock
     if not sock then

src/main/java/io/github/grrolland/hcshm/ShmService.java

@@ -23,7 +23,7 @@
 import io.github.grrolland.hcshm.processor.RateLimiterProcessor;
 import io.github.grrolland.hcshm.processor.TouchProcessor;
 import io.github.grrolland.hcshm.ratelimiter.ConsumptionProbe;
-import io.github.grrolland.hcshm.ratelimiter.RateLimiterValue;
+import io.github.grrolland.hcshm.ratelimiter.RateLimiterShmValue;
 
 import java.time.Duration;
 import java.util.concurrent.TimeUnit;
@@ -158,11 +158,11 @@ public void flushall(String region) {
      * Consume a token
      * @param key  the key
      * @param capacity  the maximum capacity
-     * @param duration the duration of a token
+     * @param duration the duration of a token in seconds
      * @return the number of tokens remaining
      */
     public String rateLimiter(String key, int capacity, long duration) {
-        final IMap<String, RateLimiterValue> map = regionLocator.getMap(hazelcast, key);
+        final IMap<String, RateLimiterShmValue> map = regionLocator.getMap(hazelcast, key);
         RateLimiterProcessor rateLimiterProcessor = new RateLimiterProcessor(capacity, Duration.ofMillis(duration));
         ConsumptionProbe consumptionProbe = (ConsumptionProbe) map.executeOnKey(key, rateLimiterProcessor);
         return String.valueOf(consumptionProbe.getRemainingTokens());

src/main/java/io/github/grrolland/hcshm/commands/RateLimiterCommand.java

@@ -49,8 +49,8 @@ public String execute(String[] commandTokens) {
             assertTokens(commandTokens, 4);
             String key = getKey(commandTokens[1]);
             int capacity = getIncrValue(commandTokens[2]);
-            long expire = getExpire(commandTokens[3]);
-            String value = getService().rateLimiter(key, capacity, expire);
+            long duration = getExpire(commandTokens[3]);
+            String value = getService().rateLimiter(key, capacity, duration);
             writeLen(response, value);
             writeValue(response, value);
             writeDone(response);

src/main/java/io/github/grrolland/hcshm/processor/RateLimiterProcessor.java

@@ -2,7 +2,7 @@
 
 import com.hazelcast.map.EntryProcessor;
 import io.github.grrolland.hcshm.ratelimiter.ConsumptionProbe;
-import io.github.grrolland.hcshm.ratelimiter.RateLimiterValue;
+import io.github.grrolland.hcshm.ratelimiter.RateLimiterShmValue;
 
 import java.io.Serializable;
 import java.time.Duration;
@@ -12,7 +12,7 @@
 /**
  * Processor for the RATE_LIMITER command
  */
-public class RateLimiterProcessor implements EntryProcessor<String, RateLimiterValue, Object>, Serializable {
+public class RateLimiterProcessor implements EntryProcessor<String, RateLimiterShmValue, Object>, Serializable {
     /**
      * Capacity
      */
@@ -37,10 +37,12 @@ public RateLimiterProcessor(int capacity, Duration duration) {
     }
 
     @Override
-    public ConsumptionProbe process(final Map.Entry<String, RateLimiterValue> entry) {
-        RateLimiterValue rateLimiterValue = Optional.ofNullable(getCurrentValue(entry)).orElseGet(this::create);
-        final ConsumptionProbe consumptionProbe = rateLimiterValue.use();
-        entry.setValue(rateLimiterValue);
+    public ConsumptionProbe process(final Map.Entry<String, RateLimiterShmValue> entry) {
+        RateLimiterShmValue rateLimiterShmValue = Optional.ofNullable(getCurrentValue(entry)).orElseGet(this::create);
+        rateLimiterShmValue.setDuration(this.duration);
+        rateLimiterShmValue.setCapacity(this.capacity);
+        final ConsumptionProbe consumptionProbe = rateLimiterShmValue.take();
+        entry.setValue(rateLimiterShmValue);
         return consumptionProbe;
     }
 
@@ -53,16 +55,16 @@ public ConsumptionProbe process(final Map.Entry<String, RateLimiterValue> entry)
      * @throws BadRequestException
      *         exception
      */
-    private RateLimiterValue getCurrentValue(final Map.Entry<String, RateLimiterValue> entry) throws BadRequestException {
+    private RateLimiterShmValue getCurrentValue(final Map.Entry<String, RateLimiterShmValue> entry) throws BadRequestException {
         try {
             return entry.getValue();
         } catch (ClassCastException e) {
             throw new BadRequestException(e);
         }
     }
 
-    private RateLimiterValue create() {
-        return new RateLimiterValue(this.capacity, this.duration);
+    private RateLimiterShmValue create() {
+        return new RateLimiterShmValue(this.capacity, this.duration);
     }
 
 }

src/main/java/io/github/grrolland/hcshm/ratelimiter/ConsumptionProbe.java

@@ -1,28 +1,18 @@
 package io.github.grrolland.hcshm.ratelimiter;
 
 /**
- * Describes token consumed, and tokens remaining
+ * Describes token consumed, and number of tokens remaining
+ * <p>
+ * remainingTokens -1 means no tokens have been consumed
+ * remainingTokens 0 means there is no more tokens to consume
  */
 public class ConsumptionProbe {
-    /**
-     * Token has been consumed
-     */
-    private final boolean consumed;
 
     /**
      * Number of remaining tokens
      */
     private final int remainingTokens;
 
-    /**
-     * return the flag consumed
-     *
-     * @return true if token was consumed
-     */
-    public boolean isConsumed() {
-        return this.consumed;
-    }
-
     /**
      * Return the remaining records
      *
@@ -36,13 +26,10 @@ public int getRemainingTokens() {
     /**
      * Constructor
      *
-     * @param consumed
-     *         true if token has been consumer
      * @param remainingTokens
      *         the number of remaining token
      */
-    ConsumptionProbe(boolean consumed, int remainingTokens) {
-        this.consumed = consumed;
+    ConsumptionProbe(int remainingTokens) {
         this.remainingTokens = remainingTokens;
     }
 }

src/main/java/io/github/grrolland/hcshm/ratelimiter/RateLimiterValue.java renamed to src/main/java/io/github/grrolland/hcshm/ratelimiter/RateLimiterShmValue.java

@@ -7,30 +7,50 @@
 import java.util.List;
 
 /**
- * RateLimiterValue
+ * RateLimiterShmValue store rate limiter data
  */
-public class RateLimiterValue extends AbstractShmValue {
+public class RateLimiterShmValue extends AbstractShmValue {
 
     /**
      * The current consumption records
      */
-    private final List<Record> records;
+    private final List<Token> tokens;
     /**
      * Sliding Window duration
      */
-    private final Duration duration;
+    private Duration duration;
     /**
      * The capacity
      */
-    private final int capacity;
+    private int capacity;
+
+    /**
+     * Set capacity
+     *
+     * @param pCapacity
+     *         capacity
+     */
+    public void setCapacity(final int pCapacity) {
+        this.capacity = pCapacity;
+    }
+
+    /**
+     * Set duration
+     *
+     * @param pDuration
+     *         duration
+     */
+    public void setDuration(final Duration pDuration) {
+        this.duration = pDuration;
+    }
 
     /**
      * Get the remaining records before capacity is exceeded
      *
      * @return the remaining records
      */
     public int getRemaining() {
-        return this.capacity - this.records.size();
+        return Math.max(this.capacity - this.tokens.size(), 0);
     }
 
     @Override
@@ -47,28 +67,28 @@ public String getValue() {
      * @param duration
      *         the sliding window duration
      */
-    public RateLimiterValue(int capacity, Duration duration) {
-        this.records = new ArrayList<>(capacity);
+    public RateLimiterShmValue(int capacity, Duration duration) {
+        this.tokens = new ArrayList<>(capacity);
         this.duration = duration;
         this.capacity = capacity;
     }
 
     /**
-     * Consume and return the ConsumptionProbe
+     * Try to take a token and return the ConsumptionProbe
      *
-     * @return ConsumptionProbe
+     * @return the ConsumptionProbe
      */
-    public ConsumptionProbe use() {
+    public ConsumptionProbe take() {
         // Clear expired tokens
         this.clearExpired();
 
-        boolean consumed = false;
         int remaining = -1;
-        if (canConsume()) {
-            consumed = records.add(new Record());
+        // Try to consume
+        if (this.canConsume()) {
+            tokens.add(new Token());
             remaining = this.getRemaining();
         }
-        return new ConsumptionProbe(consumed, remaining);
+        return new ConsumptionProbe(remaining);
     }
 
     /**
@@ -77,13 +97,14 @@ public ConsumptionProbe use() {
      * @return true if at least one token is available
      */
     private boolean canConsume() {
-        return this.records.size() < this.capacity;
+        return this.tokens.size() < this.capacity;
     }
 
     /**
      * Clear expired tokens
      */
     private void clearExpired() {
-        records.removeIf(pRecord -> pRecord.isExpired(this.duration));
+
+        tokens.removeIf(pToken -> pToken.isExpired(this.duration));
     }
 }