berater 0.1.3 → 0.5.0
Sign up to get free protection for your applications and to get access to all the features.
- checksums.yaml +4 -4
- data/lib/berater.rb +33 -35
- data/lib/berater/concurrency_limiter.rb +63 -110
- data/lib/berater/dsl.rb +68 -0
- data/lib/berater/inhibitor.rb +9 -10
- data/lib/berater/limiter.rb +76 -0
- data/lib/berater/lock.rb +27 -0
- data/lib/berater/lua_script.rb +55 -0
- data/lib/berater/rate_limiter.rb +74 -55
- data/lib/berater/rspec.rb +14 -0
- data/lib/berater/rspec/matchers.rb +62 -0
- data/lib/berater/test_mode.rb +52 -0
- data/lib/berater/unlimiter.rb +7 -9
- data/lib/berater/utils.rb +46 -0
- data/lib/berater/version.rb +1 -1
- data/spec/berater_spec.rb +43 -101
- data/spec/concurrency_limiter_spec.rb +163 -99
- data/spec/dsl_refinement_spec.rb +46 -0
- data/spec/dsl_spec.rb +72 -0
- data/spec/inhibitor_spec.rb +4 -18
- data/spec/limiter_spec.rb +71 -0
- data/spec/lua_script_spec.rb +97 -0
- data/spec/{matcher_spec.rb → matchers_spec.rb} +16 -4
- data/spec/rate_limiter_spec.rb +124 -102
- data/spec/riddle_spec.rb +102 -0
- data/spec/test_mode_spec.rb +213 -0
- data/spec/unlimiter_spec.rb +6 -37
- data/spec/utils_spec.rb +78 -0
- metadata +41 -8
- data/lib/berater/base_limiter.rb +0 -32
- data/spec/concurrency_lock_spec.rb +0 -50
checksums.yaml
CHANGED
@@ -1,7 +1,7 @@
|
|
1
1
|
---
|
2
2
|
SHA256:
|
3
|
-
metadata.gz:
|
4
|
-
data.tar.gz:
|
3
|
+
metadata.gz: fa5d51d6eea9d4691e58e8dd7065a1d596a57abdac1227fa19bd22cd2f42c15d
|
4
|
+
data.tar.gz: 55ac15bf6f2f295c525fdea57b3af9b518935a12dd4b644e3c0c9c085429ca50
|
5
5
|
SHA512:
|
6
|
-
metadata.gz:
|
7
|
-
data.tar.gz:
|
6
|
+
metadata.gz: 52df5781f9c37c7f1c9dde09e579c8da1f2d62f405835b771c6087cda8af046018e85a88629f7c5d3a2be9aafd2680e348d1b2f91bc63fe456e6e51ec5d0fd02
|
7
|
+
data.tar.gz: '0928cb6f2e8e07810dec96f5af302ff9f62d519b9f19ef6f28acb5621569f13edb6d784cbd7084b738df8dfc232ace26e3ce0c6c4a7f2ac3e4391c97789bafae'
|
data/lib/berater.rb
CHANGED
@@ -1,46 +1,40 @@
|
|
1
|
+
require 'berater/limiter'
|
2
|
+
require 'berater/lock'
|
3
|
+
require 'berater/lua_script'
|
4
|
+
require 'berater/utils'
|
1
5
|
require 'berater/version'
|
2
6
|
|
3
|
-
|
4
7
|
module Berater
|
5
8
|
extend self
|
6
9
|
|
7
10
|
class Overloaded < StandardError; end
|
8
11
|
|
9
|
-
|
10
|
-
|
11
|
-
attr_accessor :redis, :mode
|
12
|
+
attr_accessor :redis
|
12
13
|
|
13
14
|
def configure
|
14
|
-
self.mode = :unlimited # default
|
15
|
-
|
16
15
|
yield self
|
17
16
|
end
|
18
17
|
|
19
|
-
def
|
20
|
-
|
21
|
-
|
22
|
-
unless klass
|
23
|
-
raise ArgumentError, "invalid mode: #{mode}"
|
24
|
-
end
|
25
|
-
|
26
|
-
klass.new(*args, **opts)
|
27
|
-
end
|
28
|
-
|
29
|
-
def register(mode, klass)
|
30
|
-
MODES[mode.to_sym] = klass
|
18
|
+
def reset
|
19
|
+
@redis = nil
|
31
20
|
end
|
32
21
|
|
33
|
-
def
|
34
|
-
|
35
|
-
|
22
|
+
def new(key, capacity, interval = nil, **opts)
|
23
|
+
case capacity
|
24
|
+
when :unlimited, Float::INFINITY
|
25
|
+
Berater::Unlimiter
|
26
|
+
when :inhibited, 0
|
27
|
+
Berater::Inhibitor
|
28
|
+
else
|
29
|
+
if interval
|
30
|
+
Berater::RateLimiter
|
31
|
+
else
|
32
|
+
Berater::ConcurrencyLimiter
|
33
|
+
end
|
34
|
+
end.yield_self do |klass|
|
35
|
+
args = [ key, capacity, interval ].compact
|
36
|
+
klass.new(*args, **opts)
|
36
37
|
end
|
37
|
-
|
38
|
-
@mode = mode.to_sym
|
39
|
-
end
|
40
|
-
|
41
|
-
def limit(*args, **opts, &block)
|
42
|
-
mode = opts.delete(:mode) { self.mode }
|
43
|
-
new(mode, *args, **opts).limit(&block)
|
44
38
|
end
|
45
39
|
|
46
40
|
def expunge
|
@@ -51,14 +45,18 @@ module Berater
|
|
51
45
|
|
52
46
|
end
|
53
47
|
|
54
|
-
#
|
55
|
-
|
48
|
+
# convenience method
|
49
|
+
def Berater(key, capacity, interval = nil, **opts, &block)
|
50
|
+
limiter = Berater.new(key, capacity, interval, **opts)
|
51
|
+
if block_given?
|
52
|
+
limiter.limit(&block)
|
53
|
+
else
|
54
|
+
limiter
|
55
|
+
end
|
56
|
+
end
|
57
|
+
|
58
|
+
# load limiters
|
56
59
|
require 'berater/concurrency_limiter'
|
57
60
|
require 'berater/inhibitor'
|
58
61
|
require 'berater/rate_limiter'
|
59
62
|
require 'berater/unlimiter'
|
60
|
-
|
61
|
-
Berater.register(:concurrency, Berater::ConcurrencyLimiter)
|
62
|
-
Berater.register(:inhibited, Berater::Inhibitor)
|
63
|
-
Berater.register(:rate, Berater::RateLimiter)
|
64
|
-
Berater.register(:unlimited, Berater::Unlimiter)
|
@@ -1,153 +1,106 @@
|
|
1
1
|
module Berater
|
2
|
-
class ConcurrencyLimiter <
|
2
|
+
class ConcurrencyLimiter < Limiter
|
3
3
|
|
4
4
|
class Incapacitated < Overloaded; end
|
5
5
|
|
6
|
-
attr_reader :
|
6
|
+
attr_reader :timeout
|
7
7
|
|
8
|
-
def initialize(capacity, **opts)
|
9
|
-
super(**opts)
|
8
|
+
def initialize(key, capacity, **opts)
|
9
|
+
super(key, capacity, **opts)
|
10
10
|
|
11
|
-
self.capacity = capacity
|
12
11
|
self.timeout = opts[:timeout] || 0
|
13
12
|
end
|
14
13
|
|
15
|
-
def
|
16
|
-
unless capacity.is_a? Integer
|
17
|
-
raise ArgumentError, "expected Integer, found #{capacity.class}"
|
18
|
-
end
|
19
|
-
|
20
|
-
raise ArgumentError, "capacity must be >= 0" unless capacity >= 0
|
21
|
-
|
22
|
-
@capacity = capacity
|
23
|
-
end
|
24
|
-
|
25
|
-
def timeout=(timeout)
|
26
|
-
unless timeout.is_a? Integer
|
27
|
-
raise ArgumentError, "expected Integer, found #{timeout.class}"
|
28
|
-
end
|
29
|
-
|
30
|
-
raise ArgumentError, "timeout must be >= 0" unless timeout >= 0
|
31
|
-
|
14
|
+
private def timeout=(timeout)
|
32
15
|
@timeout = timeout
|
16
|
+
timeout = 0 if timeout == Float::INFINITY
|
17
|
+
@timeout_usec = Berater::Utils.to_usec(timeout)
|
33
18
|
end
|
34
19
|
|
35
|
-
|
36
|
-
attr_reader :limiter, :id, :contention
|
37
|
-
|
38
|
-
def initialize(limiter, id, contention)
|
39
|
-
@limiter = limiter
|
40
|
-
@id = id
|
41
|
-
@contention = contention
|
42
|
-
@released = false
|
43
|
-
@locked_at = Time.now
|
44
|
-
end
|
45
|
-
|
46
|
-
def release
|
47
|
-
raise 'lock already released' if released?
|
48
|
-
raise 'lock expired' if expired?
|
49
|
-
|
50
|
-
@released = limiter.release(self)
|
51
|
-
end
|
52
|
-
|
53
|
-
def released?
|
54
|
-
@released
|
55
|
-
end
|
56
|
-
|
57
|
-
def expired?
|
58
|
-
@locked_at + limiter.timeout < Time.now
|
59
|
-
end
|
60
|
-
end
|
61
|
-
|
62
|
-
LUA_SCRIPT = <<~LUA.gsub(/^\s*(--.*\n)?/, '')
|
20
|
+
LUA_SCRIPT = Berater::LuaScript(<<~LUA
|
63
21
|
local key = KEYS[1]
|
22
|
+
local lock_key = KEYS[2]
|
64
23
|
local capacity = tonumber(ARGV[1])
|
65
24
|
local ts = tonumber(ARGV[2])
|
66
25
|
local ttl = tonumber(ARGV[3])
|
26
|
+
local cost = tonumber(ARGV[4])
|
27
|
+
local lock_ids = {}
|
67
28
|
|
68
|
-
|
69
|
-
|
70
|
-
|
71
|
-
|
72
|
-
-- check to see if key already exists
|
73
|
-
if ttl == 0 then
|
74
|
-
exists = redis.call('EXISTS', key)
|
75
|
-
else
|
76
|
-
-- and refresh TTL while we're at it
|
77
|
-
exists = redis.call('EXPIRE', key, ttl * 2)
|
29
|
+
-- purge stale hosts
|
30
|
+
if ttl > 0 then
|
31
|
+
redis.call('ZREMRANGEBYSCORE', key, '-inf', ts - ttl)
|
78
32
|
end
|
79
33
|
|
80
|
-
|
81
|
-
|
82
|
-
if ttl > 0 then
|
83
|
-
redis.call('ZREMRANGEBYSCORE', key, '-inf', ts - ttl)
|
84
|
-
end
|
85
|
-
|
86
|
-
-- check capacity (subtract one for next lock entry)
|
87
|
-
count = redis.call('ZCARD', key) - 1
|
34
|
+
-- check capacity
|
35
|
+
local count = redis.call('ZCARD', key)
|
88
36
|
|
37
|
+
if cost == 0 then
|
38
|
+
-- just check limit, ie. for .overlimit?
|
89
39
|
if count < capacity then
|
90
|
-
|
40
|
+
table.insert(lock_ids, true)
|
41
|
+
end
|
42
|
+
elseif (count + cost) <= capacity then
|
43
|
+
-- grab locks, one per cost
|
44
|
+
local lock_id = redis.call('INCRBY', lock_key, cost)
|
45
|
+
local locks = {}
|
91
46
|
|
92
|
-
|
93
|
-
|
94
|
-
redis.call('ZADD', key, 'inf', (lock + 1) % 2^52)
|
47
|
+
for i = lock_id - cost + 1, lock_id do
|
48
|
+
table.insert(lock_ids, i)
|
95
49
|
|
96
|
-
|
50
|
+
table.insert(locks, ts)
|
51
|
+
table.insert(locks, i)
|
97
52
|
end
|
98
|
-
else
|
99
|
-
count = 1
|
100
|
-
lock = "1"
|
101
53
|
|
102
|
-
|
103
|
-
|
54
|
+
redis.call('ZADD', key, unpack(locks))
|
55
|
+
count = count + cost
|
104
56
|
|
105
57
|
if ttl > 0 then
|
106
|
-
redis.call('
|
58
|
+
redis.call('PEXPIRE', key, ttl)
|
107
59
|
end
|
108
60
|
end
|
109
61
|
|
110
|
-
|
111
|
-
-- store lock and timestamp
|
112
|
-
redis.call('ZADD', key, ts, lock)
|
113
|
-
end
|
114
|
-
|
115
|
-
return { count, lock }
|
62
|
+
return { count, unpack(lock_ids) }
|
116
63
|
LUA
|
64
|
+
)
|
117
65
|
|
118
|
-
def limit(
|
119
|
-
|
120
|
-
|
121
|
-
capacity,
|
122
|
-
**options.merge(opts)
|
123
|
-
# **options.merge(timeout: timeout).merge(opts)
|
124
|
-
).limit(&block)
|
125
|
-
end
|
66
|
+
def limit(capacity: nil, cost: 1, &block)
|
67
|
+
capacity ||= @capacity
|
68
|
+
# cost is Integer >= 0
|
126
69
|
|
127
|
-
|
128
|
-
|
129
|
-
[ key ],
|
130
|
-
[ capacity, Time.now.to_i, timeout ]
|
131
|
-
)
|
70
|
+
# timestamp in microseconds
|
71
|
+
ts = (Time.now.to_f * 10**6).to_i
|
132
72
|
|
133
|
-
|
73
|
+
count, *lock_ids = LUA_SCRIPT.eval(
|
74
|
+
redis,
|
75
|
+
[ cache_key(key), cache_key('lock_id') ],
|
76
|
+
[ capacity, ts, @timeout_usec, cost ]
|
77
|
+
)
|
134
78
|
|
135
|
-
|
79
|
+
raise Incapacitated if lock_ids.empty?
|
136
80
|
|
137
|
-
if
|
138
|
-
|
139
|
-
yield lock
|
140
|
-
ensure
|
141
|
-
release(lock)
|
142
|
-
end
|
81
|
+
if cost == 0
|
82
|
+
lock = Lock.new(self, nil, count)
|
143
83
|
else
|
144
|
-
lock
|
84
|
+
lock = Lock.new(self, lock_ids[0], count, -> { release(lock_ids) })
|
145
85
|
end
|
86
|
+
|
87
|
+
yield_lock(lock, &block)
|
88
|
+
end
|
89
|
+
|
90
|
+
def overloaded?
|
91
|
+
limit(cost: 0) { false }
|
92
|
+
rescue Overloaded
|
93
|
+
true
|
94
|
+
end
|
95
|
+
alias incapacitated? overloaded?
|
96
|
+
|
97
|
+
private def release(lock_ids)
|
98
|
+
res = redis.zrem(cache_key(key), lock_ids)
|
99
|
+
res == true || res == lock_ids.count # depending on which version of Redis
|
146
100
|
end
|
147
101
|
|
148
|
-
def
|
149
|
-
|
150
|
-
res == true || res == 1 # depending on which version of Redis
|
102
|
+
def to_s
|
103
|
+
"#<#{self.class}(#{key}: #{capacity} at a time)>"
|
151
104
|
end
|
152
105
|
|
153
106
|
end
|
data/lib/berater/dsl.rb
ADDED
@@ -0,0 +1,68 @@
|
|
1
|
+
module Berater
|
2
|
+
module DSL
|
3
|
+
refine Berater.singleton_class do
|
4
|
+
def new(key, mode = nil, *args, **opts, &block)
|
5
|
+
if mode.nil?
|
6
|
+
unless block_given?
|
7
|
+
raise ArgumentError, 'expected either mode or block'
|
8
|
+
end
|
9
|
+
|
10
|
+
mode, *args = DSL.eval(&block)
|
11
|
+
else
|
12
|
+
if block_given?
|
13
|
+
raise ArgumentError, 'expected either mode or block, not both'
|
14
|
+
end
|
15
|
+
end
|
16
|
+
|
17
|
+
super(key, mode, *args, **opts)
|
18
|
+
end
|
19
|
+
end
|
20
|
+
|
21
|
+
extend self
|
22
|
+
|
23
|
+
def eval &block
|
24
|
+
@keywords ||= Class.new do
|
25
|
+
# create a class where DSL keywords are methods
|
26
|
+
KEYWORDS.each do |keyword|
|
27
|
+
define_singleton_method(keyword) { keyword }
|
28
|
+
end
|
29
|
+
end
|
30
|
+
|
31
|
+
install
|
32
|
+
@keywords.class_eval &block
|
33
|
+
ensure
|
34
|
+
uninstall
|
35
|
+
end
|
36
|
+
|
37
|
+
private
|
38
|
+
|
39
|
+
KEYWORDS = [
|
40
|
+
:second, :minute, :hour,
|
41
|
+
:unlimited, :inhibited,
|
42
|
+
].freeze
|
43
|
+
|
44
|
+
def install
|
45
|
+
Integer.class_eval do
|
46
|
+
def per(unit)
|
47
|
+
[ self, unit ]
|
48
|
+
end
|
49
|
+
alias every per
|
50
|
+
|
51
|
+
def at_once
|
52
|
+
[ self ]
|
53
|
+
end
|
54
|
+
alias concurrently at_once
|
55
|
+
alias at_a_time at_once
|
56
|
+
end
|
57
|
+
end
|
58
|
+
|
59
|
+
def uninstall
|
60
|
+
Integer.remove_method :per
|
61
|
+
Integer.remove_method :every
|
62
|
+
|
63
|
+
Integer.remove_method :at_once
|
64
|
+
Integer.remove_method :concurrently
|
65
|
+
Integer.remove_method :at_a_time
|
66
|
+
end
|
67
|
+
end
|
68
|
+
end
|
data/lib/berater/inhibitor.rb
CHANGED
@@ -1,21 +1,20 @@
|
|
1
1
|
module Berater
|
2
|
-
class Inhibitor <
|
2
|
+
class Inhibitor < Limiter
|
3
3
|
|
4
4
|
class Inhibited < Overloaded; end
|
5
5
|
|
6
|
-
def initialize(*args, **opts)
|
7
|
-
super(**opts)
|
6
|
+
def initialize(key = :inhibitor, *args, **opts)
|
7
|
+
super(key, 0, **opts)
|
8
8
|
end
|
9
9
|
|
10
|
-
def limit(**opts
|
11
|
-
unless opts.empty?
|
12
|
-
return self.class.new(
|
13
|
-
**options.merge(opts)
|
14
|
-
).limit(&block)
|
15
|
-
end
|
16
|
-
|
10
|
+
def limit(**opts)
|
17
11
|
raise Inhibited
|
18
12
|
end
|
19
13
|
|
14
|
+
def overloaded?
|
15
|
+
true
|
16
|
+
end
|
17
|
+
alias inhibited? overloaded?
|
18
|
+
|
20
19
|
end
|
21
20
|
end
|
@@ -0,0 +1,76 @@
|
|
1
|
+
module Berater
|
2
|
+
class Limiter
|
3
|
+
|
4
|
+
attr_reader :key, :capacity, :options
|
5
|
+
|
6
|
+
def redis
|
7
|
+
options[:redis] || Berater.redis
|
8
|
+
end
|
9
|
+
|
10
|
+
def limit
|
11
|
+
raise NotImplementedError
|
12
|
+
end
|
13
|
+
|
14
|
+
def overloaded?
|
15
|
+
raise NotImplementedError
|
16
|
+
end
|
17
|
+
|
18
|
+
def to_s
|
19
|
+
"#<#{self.class}>"
|
20
|
+
end
|
21
|
+
|
22
|
+
def ==(other)
|
23
|
+
self.class == other.class &&
|
24
|
+
self.key == other.key &&
|
25
|
+
self.capacity == other.capacity &&
|
26
|
+
self.args == other.args &&
|
27
|
+
self.options == other.options &&
|
28
|
+
self.redis.connection == other.redis.connection
|
29
|
+
end
|
30
|
+
|
31
|
+
def self.new(*)
|
32
|
+
# can only call via subclass
|
33
|
+
raise NotImplementedError if self == Berater::Limiter
|
34
|
+
|
35
|
+
super
|
36
|
+
end
|
37
|
+
|
38
|
+
protected
|
39
|
+
|
40
|
+
attr_reader :args
|
41
|
+
|
42
|
+
def initialize(key, capacity, *args, **opts)
|
43
|
+
@key = key
|
44
|
+
self.capacity = capacity
|
45
|
+
@args = args
|
46
|
+
@options = opts
|
47
|
+
end
|
48
|
+
|
49
|
+
def capacity=(capacity)
|
50
|
+
unless capacity.is_a?(Integer) || capacity == Float::INFINITY
|
51
|
+
raise ArgumentError, "expected Integer, found #{capacity.class}"
|
52
|
+
end
|
53
|
+
|
54
|
+
raise ArgumentError, "capacity must be >= 0" unless capacity >= 0
|
55
|
+
|
56
|
+
@capacity = capacity
|
57
|
+
end
|
58
|
+
|
59
|
+
def cache_key(key)
|
60
|
+
"#{self.class}:#{key}"
|
61
|
+
end
|
62
|
+
|
63
|
+
def yield_lock(lock, &block)
|
64
|
+
if block_given?
|
65
|
+
begin
|
66
|
+
yield lock
|
67
|
+
ensure
|
68
|
+
lock.release
|
69
|
+
end
|
70
|
+
else
|
71
|
+
lock
|
72
|
+
end
|
73
|
+
end
|
74
|
+
|
75
|
+
end
|
76
|
+
end
|