@@ -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