feat: use incremental counter #3

feat: use incremental counter

Author: shreemaan-abhishek

.github/workflows/test.yaml

@@ -8,7 +8,7 @@ on:
 
 jobs:
   test:
-    runs-on: "ubuntu-20.04"
+    runs-on: "buildjet-4vcpu-ubuntu-2004"
     env:
       OPENRESTY_PREFIX: "/usr/local/openresty"
 

lib/resty/limit/count.lua

@@ -55,22 +55,22 @@ local function incoming_new(self, key, commit, cost)
     local limit = self.limit
     local window = self.window
 
-    local remaining, ok, err
+    local consumed, ok, err
 
     if commit then
-        remaining, err = dict:incr(key, -cost, limit, window)
-        if not remaining then
+        consumed, err = dict:incr(key, cost, 0, window)
+        if not consumed then
             return nil, err
         end
     else
-        remaining = (dict:get(key) or limit) - cost
+        consumed = (dict:get(key) or 0) + cost
     end
 
-    if remaining < 0 then
+    if consumed > limit then
         return nil, "rejected"
     end
 
-    return 0, remaining
+    return 0, consumed
 end
 
 -- incoming function using incr and expire
@@ -125,16 +125,16 @@ function _M.uncommit(self, key)
     local dict = self.dict
     local limit = self.limit
 
-    local remaining, err = dict:incr(key, 1)
-    if not remaining then
+    local consumed, err = dict:incr(key, -1)
+    if not consumed then
         if err == "not found" then
-            remaining = limit
+            consumed = limit
         else
             return nil, err
         end
     end
 
-    return remaining
+    return consumed
 end
 
 

t/count.t

@@ -34,14 +34,15 @@ __DATA__
         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 limit = 10
+            local lim = limit_count.new("store", limit, 100)
             local uri = ngx.var.uri
             for i = 1, 12 do
                 local delay, err = lim:incoming(uri, true)
                 if not delay then
                     ngx.say(err)
                 else
-                    local remaining = err
+                    local remaining = limit - err
                     ngx.say("remaining: ", remaining)
                 end
             end
@@ -75,36 +76,37 @@ rejected
         content_by_lua_block {
             local limit_count = require "resty.limit.count"
             ngx.shared.store:flush_all()
-            local lim = limit_count.new("store", 1, 10)
+            local limit = 1
+            local lim = limit_count.new("store", limit, 10)
             local delay1, err1 = lim:incoming("foo", true)
             local delay2, err2 = lim:incoming("foo", true)
             local delay3, err3 = lim:incoming("bar", true)
             local delay4, err4 = lim:incoming("bar", true)
             if not delay1 then
                 ngx.say(err1)
             else
-                local remaining1 = err1
+                local remaining1 = limit - err1
                 ngx.say("remaining1: ", remaining1)
             end
 
             if not delay2 then
                 ngx.say(err2)
             else
-                local remaining2 = err2
+                local remaining2 = limit - err2
                 ngx.say("remaining2: ", remaining2)
             end
 
             if not delay3 then
                 ngx.say(err3)
             else
-                local remaining3 = err3
+                local remaining3 = limit -err3
                 ngx.say("remaining3: ", remaining3)
             end
 
             if not delay4 then
                 ngx.say(err4)
             else
-                local remaining4 = err4
+                local remaining4 = limit - err4
                 ngx.say("remaining4: ", remaining4)
             end
         }
@@ -129,23 +131,24 @@ rejected
         content_by_lua_block {
             local limit_count = require "resty.limit.count"
             ngx.shared.store:flush_all()
-            local lim = limit_count.new("store", 1, 1)
+            local limit = 1
+            local lim = limit_count.new("store", limit, 1)
 
             local uri = ngx.var.uri
             for i = 1, 2 do
                 local delay, err = lim:incoming(uri, true)
                 if not delay then
                     ngx.say(err)
                 else
-                    local remaining = err
+                    local remaining = limit -  err
                     ngx.say("remaining: ", remaining)
                 end
 
                 local delay, err = lim:incoming(uri, true)
                 if not delay then
                     ngx.say(err)
                 else
-                    local remaining = err
+                    local remaining = limit - err
                     ngx.say("remaining: ", remaining)
                 end
                 ngx.sleep(1)
@@ -172,15 +175,16 @@ rejected
         content_by_lua_block {
             local limit_count = require "resty.limit.count"
             ngx.shared.store:flush_all()
-            local lim = limit_count.new("store", 5, 10)
+            local limit = 5
+            local lim = limit_count.new("store", limit, 10)
             local begin = ngx.time()
 
             for i = 1, 4 do
                 local delay, err = lim:incoming("foo", i < 3)
                 if not delay then
                     ngx.say(err)
                 else
-                    local remaining = err
+                    local remaining = limit - err
                     ngx.say("remaining: ", remaining)
                 end
             end
@@ -205,15 +209,16 @@ remaining: 2
     location = /t {
         content_by_lua_block {
             local limit_count = require "resty.limit.count"
-            local lim = limit_count.new("store", 2, 10)
+            local limit = 2
+            local lim = limit_count.new("store", limit, 10)
             ngx.shared.store:flush_all()
             local key = "foo"
             for i = 1, 3 do
                 local delay, err = lim:incoming(key, true)
                 if not delay then
                     ngx.say("failed to limit count: ", err)
                 else
-                    local remaining = err
+                    local remaining = limit - err
                     ngx.say("remaining: ", remaining)
                 end
                 local ok, err = lim:uncommit(key)
@@ -242,14 +247,15 @@ remaining: 1
         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 limit = 10
+            local lim = limit_count.new("store", limit, 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
+                    local remaining = limit - err
                     ngx.say("remaining: ", remaining)
                 end
             end

t/traffic.t

@@ -74,11 +74,11 @@ $::HttpConfig
 --- request
     GET /t
 --- response_body_like eval
-qr/^1: 0, conn committed: true, states: 0, 0, 1, 9
-2: 0\.5, conn committed: true, states: 1, 1, 2, 8
-3: 1, conn committed: true, states: 2, 2, 3, 7
+qr/^1: 0, conn committed: true, states: 0, 0, 1, 1
+2: 0\.5, conn committed: true, states: 1, 1, 2, 2
+3: 1, conn committed: true, states: 2, 2, 3, 3
 failed to limit traffic: rejected
-5: 0\.(?:4[6-9]|5|5[0-4])\d*, conn committed: true, states: 0, (?:1|1\.0[0-4]\d*|0\.9[6-9]\d*), 4, 6
+5: 0\.(?:4[6-9]|5|5[0-4])\d*, conn committed: true, states: 0, (?:1|1\.0[0-4]\d*|0\.9[6-9]\d*), 4, 4
 6: 2, conn committed: true, states: 1, (?:2|2\.0[0-4]\d*|1\.9[6-9]\d*), 5, 5
 $/s
 --- no_error_log
@@ -193,15 +193,15 @@ $::HttpConfig
     GET /t
 --- response_body_like eval
 qr/^1: 0, conn committed: true, states: 0, 0, 1, 1
-2: 0\.5, conn committed: true, states: 1, 1, 2, 0
+2: 0\.5, conn committed: true, states: 1, 1, 2, 2
 failed to limit traffic: rejected
-states: 1, 1, 2, 0
+states: 1, 1, 2, 2
 failed to limit traffic: rejected
-states: 1, 1, 2, 0
+states: 1, 1, 2, 2
 failed to limit traffic: rejected
-states: 1, 1, 2, 0
+states: 1, 1, 2, 2
 failed to limit traffic: rejected
-states: 1, 1, 2, 0
+states: 1, 1, 2, 2
 $/s
 --- no_error_log
 [error]