feat: support specify cost when use incoming function (#1)

Author: monkeyDluffy6017

lib/resty/limit/count.lua

@@ -46,20 +46,24 @@ end
 
 -- incoming function using incr with init_ttl
 -- need OpenResty version > v0.10.12rc2
-local function incoming_new(self, key, commit)
+-- parameter "cost" is only supported on the new interface
+local function incoming_new(self, key, commit, cost)
+    cost = cost or 1
+    assert(cost >= 1)
+
     local dict = self.dict
     local limit = self.limit
     local window = self.window
 
     local remaining, ok, err
 
     if commit then
-        remaining, err = dict:incr(key, -1, limit, window)
+        remaining, err = dict:incr(key, -cost, limit, window)
         if not remaining then
             return nil, err
         end
     else
-        remaining = (dict:get(key) or limit) - 1
+        remaining = (dict:get(key) or limit) - cost
     end
 
     if remaining < 0 then
@@ -74,7 +78,6 @@ local function incoming_old(self, key, commit)
     local dict = self.dict
     local limit = self.limit
     local window = self.window
-
     local remaining, ok, err
 
     if commit then

lib/resty/limit/count.md

@@ -113,7 +113,7 @@ This method takes the following arguments:
 
 incoming
 --------
-**syntax:** `delay, err = obj:incoming(key, commit)`
+**syntax:** `delay, err = obj:incoming(key, commit, cost)`
 
 Fires a new request incoming event and calculates the delay needed (if any) for the current request
 upon the specified key or whether the user should reject it immediately.
@@ -130,6 +130,8 @@ key so that we can set a rate for individual user.
 * `commit` is a boolean value. If set to `true`, the object will actually record the event
 in the shm zone backing the current object; otherwise it would just be a "dry run" (which is the default).
 
+* `cost` is the specified number that should be deducted when a new request arrives.
+
 The return values depend on the following cases:
 
 1. If the request does not exceed the `count` value specified in the [new](#new) method, then

t/count.t

@@ -232,3 +232,39 @@ remaining: 1
 --- no_error_log
 [error]
 [lua]
+
+
+
+=== TEST 6: specify a cost
+--- http_config eval: $::HttpConfig
+--- config
+    location = /t {
+        content_by_lua_block {
+            local limit_count = require "resty.limit.count"
+            ngx.shared.store:flush_all()
+            local lim = limit_count.new("store", 10, 100)
+            local uri = ngx.var.uri
+            for i = 1, 7 do
+                local delay, err = lim:incoming(uri, true, 2)
+                if not delay then
+                    ngx.say(err)
+                else
+                    local remaining = err
+                    ngx.say("remaining: ", remaining)
+                end
+            end
+        }
+    }
+--- request
+    GET /t
+--- response_body
+remaining: 8
+remaining: 6
+remaining: 4
+remaining: 2
+remaining: 0
+rejected
+rejected
+--- no_error_log
+[error]
+[lua]