-
Notifications
You must be signed in to change notification settings - Fork 3
/
lua.go
129 lines (109 loc) · 3.09 KB
/
lua.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
package redisstore
const luaTemplate = `
local C_EXPIRE = 'EXPIRE'
local C_HGETALL = 'HGETALL'
local C_HSET = 'HSET'
local F_START = 's'
local F_TICK = 't'
local F_INTERVAL = 'i'
local F_TOKENS = 'k'
local F_MAX = 'm'
-- speed up access to next
local next = next
-- get arguments
local key = KEYS[1]
local now = tonumber(ARGV[1]) -- current unix time in nanoseconds
local defmaxtokens = tonumber(ARGV[2]) -- default tokens per interval, only used if no value already exists for the key
local definterval = tonumber(ARGV[3]) -- interval in nanoseconds, only used if no value already exists for the key
-- hgetall gets all the fields as a lua table.
local hgetall = function (key)
local data = redis.call(C_HGETALL, key)
local result = {}
for i = 1, #data, 2 do
result[data[i]] = data[i+1]
end
return result
end
-- availabletokens returns the number of available tokens given the last tick,
-- current tick, max, and fill rate.
local availabletokens = function (last, curr, max, fillrate)
local delta = curr - last
local available = delta * fillrate
if available > max then
available = max
end
return available
end
-- present returns true if the given value is not nil and is not the empty
-- string.
local present = function (val)
return val ~= nil and val ~= ''
end
-- tick returns the total number of times the interval has occurred between
-- start and current.
local tick = function (start, curr, interval)
local val = math.floor((curr - start) / interval)
if val > 0 then
return val
end
return 0
end
-- ttl returns the appropriate ttl in seconds for the given interval, 3x the
-- interval.
local ttl = function (interval)
return 3 * math.floor(interval / 1000000000)
end
--
-- begin exec
--
local data = hgetall(key)
local start = now
if present(data[F_START]) then
start = tonumber(data[F_START])
else
redis.call(C_HSET, key, F_START, now)
redis.call(C_EXPIRE, key, 30)
end
local lasttick = 0
if present(data[F_TICK]) then
lasttick = tonumber(data[F_TICK])
else
redis.call(C_HSET, key, F_TICK, 0)
redis.call(C_EXPIRE, key, 30)
end
local maxtokens = defmaxtokens
if present(data[F_MAX]) then
maxtokens = tonumber(data[F_MAX])
else
redis.call(C_HSET, key, F_MAX, defmaxtokens)
redis.call(C_EXPIRE, key, 30)
end
local tokens = maxtokens
if present(data[F_TOKENS]) then
tokens = tonumber(data[F_TOKENS])
end
local interval = definterval
if present(data[F_INTERVAL]) then
interval = tonumber(data[F_INTERVAL])
end
local currtick = tick(start, now, interval)
local nexttime = start + ((currtick+1) * interval)
if lasttick < currtick then
local rate = interval / tokens
tokens = availabletokens(lasttick, currtick, maxtokens, rate)
lasttick = currtick
redis.call(C_HSET, key,
F_START, start,
F_TICK, lasttick,
F_INTERVAL, interval,
F_TOKENS, tokens)
redis.call(C_EXPIRE, key, ttl(interval))
end
if tokens > 0 then
tokens = tokens - 1
redis.call(C_HSET, key, F_TOKENS, tokens)
redis.call(C_EXPIRE, key, ttl(interval))
return {maxtokens, tokens, nexttime, true}
end
return {maxtokens, tokens, nexttime, false}
`