@@ -20,54 +20,64 @@ local max               = math.max
2020local  ngx_now            =  ngx .now 
2121local  ngx_null           =  ngx .null 
2222local  tonumber           =  tonumber 
23+ local  core               =  require (" apisix.core" 
2324
2425
2526local  _M  =  {version  =  0.1 }
2627
2728
29+ local  script  =  core .string .compress_script ([=[ 
30+   local state_key  = KEYS[1]             -- state_key (hash), fields: "excess", "last"     
31+   local rate       = tonumber(ARGV[1])   -- req/s 
32+   local now        = tonumber(ARGV[2])   -- ms 
33+   local burst      = tonumber(ARGV[3])   -- req/s 
34+   local commit     = tonumber(ARGV[4])   -- 1/0 
35+ 
36+   local vals = redis.call("HMGET", state_key, "excess", "last") 
37+   local prev_excess = tonumber(vals[1] or "0") 
38+   local prev_last   = tonumber(vals[2] or "0") 
39+ 
40+   local new_excess 
41+   if prev_last > 0 then 
42+     local elapsed = math.abs(now - prev_last) 
43+     new_excess = math.max(prev_excess - rate * (elapsed) / 1000 + 1000, 0) 
44+   else 
45+     new_excess = 0 
46+   end 
47+ 
48+   if new_excess > burst then 
49+     return {0, new_excess} 
50+   end 
51+ 
52+   if commit == 1 then 
53+     redis.call("HMSET", state_key, "excess", new_excess, "last", now) 
54+     redis.call("EXPIRE", state_key, 60) 
55+   end 
56+ 
57+   return {1, new_excess} 
58+ ]=] 
59+ 
60+ 
2861--  the "commit" argument controls whether should we record the event in shm.
2962function  _M .incoming (self , red , key , commit )
3063    local  rate  =  self .rate 
3164    local  now  =  ngx_now () *  1000 
3265
33-     key  =  " limit_req"  ..  " :"  ..  key 
34-     local  excess_key  =  key  ..  " excess" 
35-     local  last_key  =  key  ..  " last" 
66+     local  state_key  =  " limit_req:{"  ..  key  ..  " }:state" 
3667
37-     local  excess , err  =  red :get (excess_key )
38-     if  err  then 
39-         return  nil , err 
40-     end 
41-     local  last , err  =  red :get (last_key )
42-     if  err  then 
68+     local  commit_flag  =  commit  and  " 1"  or  " 0" 
69+ 
70+     local  res , err  =  red :eval (script , 1 , state_key ,
71+                               rate , now , self .burst , commit_flag )
72+     if  not  res  then 
4373        return  nil , err 
4474    end 
4575
46-     if  excess  ~=  ngx_null  and  last  ~=  ngx_null  then 
47-         excess  =  tonumber (excess )
48-         last  =  tonumber (last )
49-         local  elapsed  =  now  -  last 
50-         excess  =  max (excess  -  rate  *  abs (elapsed ) /  1000  +  1000 , 0 )
51- 
52-         if  excess  >  self .burst  then 
53-             return  nil , " rejected" 
54-         end 
55-     else 
56-         excess  =  0 
57-     end 
76+     local  allowed  =  tonumber (res [1 ]) ==  1 
77+     local  excess   =  tonumber (res [2 ]) or  0 
5878
59-     if  commit  then 
60-         local  ok 
61-         local  err 
62-         ok , err  =  red :set (excess_key , excess )
63-         if  not  ok  then 
64-             return  nil , err 
65-         end 
66- 
67-         ok , err  =  red :set (last_key , now )
68-         if  not  ok  then 
69-             return  nil , err 
70-         end 
79+     if  not  allowed  then 
80+         return  nil , " rejected" 
7181    end 
7282
7383    --  return the delay in seconds, as well as excess
0 commit comments