berater 0.0.1 → 0.1.4
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- checksums.yaml +5 -5
- data/lib/berater.rb +50 -16
- data/lib/berater/base_limiter.rb +32 -0
- data/lib/berater/concurrency_limiter.rb +154 -0
- data/lib/berater/inhibitor.rb +21 -0
- data/lib/berater/rate_limiter.rb +84 -0
- data/lib/berater/unlimiter.rb +19 -0
- data/lib/berater/version.rb +3 -0
- data/spec/berater_spec.rb +192 -0
- data/spec/concurrency_limiter_spec.rb +202 -0
- data/spec/concurrency_lock_spec.rb +92 -0
- data/spec/inhibitor_spec.rb +37 -0
- data/spec/matcher_spec.rb +118 -0
- data/spec/rate_limiter_spec.rb +166 -0
- data/spec/unlimiter_spec.rb +61 -0
- metadata +81 -20
checksums.yaml
CHANGED
@@ -1,7 +1,7 @@
|
|
1
1
|
---
|
2
|
-
|
3
|
-
metadata.gz:
|
4
|
-
data.tar.gz:
|
2
|
+
SHA256:
|
3
|
+
metadata.gz: c8ca6fba7e38a014eed325df41dda7372b41ff62d9e78b61565f4d915d8b4677
|
4
|
+
data.tar.gz: 1fae864b93de87ea3dc70531f3f8697e80af526aea53d8ddfe2076eac9bbd61d
|
5
5
|
SHA512:
|
6
|
-
metadata.gz:
|
7
|
-
data.tar.gz:
|
6
|
+
metadata.gz: 1b4b37134edece0657ea9be15119fe25b2208d554b9dbf56edd04f8e5acff77569022be025ec428b81ab1971171ae468e389a9568d7ad0163b0966e622a09f2e
|
7
|
+
data.tar.gz: bd927303eb1475e79a26161a184205008926f95c59a06eee43f453d5e464d15c205a91bffd36228ae28bc5d27f0fa32faebc7c5870ec5d1779bd6c57ecdc0fe2
|
data/lib/berater.rb
CHANGED
@@ -1,30 +1,64 @@
|
|
1
|
+
require 'berater/version'
|
2
|
+
|
3
|
+
|
1
4
|
module Berater
|
2
|
-
|
5
|
+
extend self
|
3
6
|
|
4
|
-
class
|
7
|
+
class Overloaded < StandardError; end
|
5
8
|
|
6
|
-
|
7
|
-
|
8
|
-
|
9
|
+
MODES = {}
|
10
|
+
|
11
|
+
attr_accessor :redis, :mode
|
9
12
|
|
13
|
+
def configure
|
14
|
+
self.mode = :unlimited # default
|
10
15
|
|
11
|
-
|
12
|
-
|
16
|
+
yield self
|
17
|
+
end
|
18
|
+
|
19
|
+
def new(mode, *args, **opts)
|
20
|
+
klass = MODES[mode.to_sym]
|
13
21
|
|
14
|
-
|
15
|
-
|
22
|
+
unless klass
|
23
|
+
raise ArgumentError, "invalid mode: #{mode}"
|
24
|
+
end
|
16
25
|
|
17
|
-
|
18
|
-
|
19
|
-
@@redis.expire rkey, seconds * 2
|
20
|
-
end.first
|
26
|
+
klass.new(*args, **opts)
|
27
|
+
end
|
21
28
|
|
22
|
-
|
29
|
+
def register(mode, klass)
|
30
|
+
MODES[mode.to_sym] = klass
|
31
|
+
end
|
23
32
|
|
24
|
-
|
33
|
+
def mode=(mode)
|
34
|
+
unless MODES.include? mode.to_sym
|
35
|
+
raise ArgumentError, "invalid mode: #{mode}"
|
25
36
|
end
|
26
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
|
+
end
|
45
|
+
|
46
|
+
def expunge
|
47
|
+
redis.scan_each(match: "#{self.name}*") do |key|
|
48
|
+
redis.del key
|
49
|
+
end
|
27
50
|
end
|
28
51
|
|
29
|
-
class LimitExceeded < RuntimeError; end
|
30
52
|
end
|
53
|
+
|
54
|
+
# load and register limiters
|
55
|
+
require 'berater/base_limiter'
|
56
|
+
require 'berater/concurrency_limiter'
|
57
|
+
require 'berater/inhibitor'
|
58
|
+
require 'berater/rate_limiter'
|
59
|
+
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)
|
@@ -0,0 +1,32 @@
|
|
1
|
+
module Berater
|
2
|
+
class BaseLimiter
|
3
|
+
|
4
|
+
attr_reader :options
|
5
|
+
|
6
|
+
def initialize(**opts)
|
7
|
+
@options = opts
|
8
|
+
end
|
9
|
+
|
10
|
+
def key
|
11
|
+
if options[:key]
|
12
|
+
"#{self.class}:#{options[:key]}"
|
13
|
+
else
|
14
|
+
# default value
|
15
|
+
self.class.to_s
|
16
|
+
end
|
17
|
+
end
|
18
|
+
|
19
|
+
def redis
|
20
|
+
options[:redis] || Berater.redis
|
21
|
+
end
|
22
|
+
|
23
|
+
def limit(**opts)
|
24
|
+
raise NotImplementedError
|
25
|
+
end
|
26
|
+
|
27
|
+
def self.limit(*args, **opts, &block)
|
28
|
+
self.new(*args, **opts).limit(&block)
|
29
|
+
end
|
30
|
+
|
31
|
+
end
|
32
|
+
end
|
@@ -0,0 +1,154 @@
|
|
1
|
+
module Berater
|
2
|
+
class ConcurrencyLimiter < BaseLimiter
|
3
|
+
|
4
|
+
class Incapacitated < Overloaded; end
|
5
|
+
|
6
|
+
attr_reader :capacity, :timeout
|
7
|
+
|
8
|
+
def initialize(capacity, **opts)
|
9
|
+
super(**opts)
|
10
|
+
|
11
|
+
self.capacity = capacity
|
12
|
+
self.timeout = opts[:timeout] || 0
|
13
|
+
end
|
14
|
+
|
15
|
+
def capacity=(capacity)
|
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
|
+
|
32
|
+
@timeout = timeout
|
33
|
+
end
|
34
|
+
|
35
|
+
class Lock
|
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
|
+
limiter.timeout > 0 && @locked_at + limiter.timeout < Time.now
|
59
|
+
end
|
60
|
+
end
|
61
|
+
|
62
|
+
LUA_SCRIPT = <<~LUA.gsub(/^\s*(--.*\n)?/, '')
|
63
|
+
local key = KEYS[1]
|
64
|
+
local capacity = tonumber(ARGV[1])
|
65
|
+
local ts = tonumber(ARGV[2])
|
66
|
+
local ttl = tonumber(ARGV[3])
|
67
|
+
|
68
|
+
local exists
|
69
|
+
local count
|
70
|
+
local lock
|
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)
|
78
|
+
end
|
79
|
+
|
80
|
+
if exists == 1 then
|
81
|
+
-- purge stale hosts
|
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
|
88
|
+
|
89
|
+
if count < capacity then
|
90
|
+
-- yay, grab a lock
|
91
|
+
|
92
|
+
-- regenerate next lock entry, which has score inf
|
93
|
+
lock = unpack(redis.call('ZPOPMAX', key))
|
94
|
+
redis.call('ZADD', key, 'inf', (lock + 1) % 2^52)
|
95
|
+
|
96
|
+
count = count + 1
|
97
|
+
end
|
98
|
+
elseif capacity > 0 then
|
99
|
+
count = 1
|
100
|
+
lock = "1"
|
101
|
+
|
102
|
+
-- create structure to track locks and next id
|
103
|
+
redis.call('ZADD', key, 'inf', lock + 1)
|
104
|
+
|
105
|
+
if ttl > 0 then
|
106
|
+
redis.call('EXPIRE', key, ttl * 2)
|
107
|
+
end
|
108
|
+
end
|
109
|
+
|
110
|
+
if lock then
|
111
|
+
-- store lock and timestamp
|
112
|
+
redis.call('ZADD', key, ts, lock)
|
113
|
+
end
|
114
|
+
|
115
|
+
return { count, lock }
|
116
|
+
LUA
|
117
|
+
|
118
|
+
def limit(**opts, &block)
|
119
|
+
unless opts.empty?
|
120
|
+
return self.class.new(
|
121
|
+
capacity,
|
122
|
+
**options.merge(opts)
|
123
|
+
# **options.merge(timeout: timeout).merge(opts)
|
124
|
+
).limit(&block)
|
125
|
+
end
|
126
|
+
|
127
|
+
count, lock_id = redis.eval(
|
128
|
+
LUA_SCRIPT,
|
129
|
+
[ key ],
|
130
|
+
[ capacity, Time.now.to_i, timeout ]
|
131
|
+
)
|
132
|
+
|
133
|
+
raise Incapacitated unless lock_id
|
134
|
+
|
135
|
+
lock = Lock.new(self, lock_id, count)
|
136
|
+
|
137
|
+
if block_given?
|
138
|
+
begin
|
139
|
+
yield lock
|
140
|
+
ensure
|
141
|
+
lock.release
|
142
|
+
end
|
143
|
+
else
|
144
|
+
lock
|
145
|
+
end
|
146
|
+
end
|
147
|
+
|
148
|
+
def release(lock)
|
149
|
+
res = redis.zrem(key, lock.id)
|
150
|
+
res == true || res == 1 # depending on which version of Redis
|
151
|
+
end
|
152
|
+
|
153
|
+
end
|
154
|
+
end
|
@@ -0,0 +1,21 @@
|
|
1
|
+
module Berater
|
2
|
+
class Inhibitor < BaseLimiter
|
3
|
+
|
4
|
+
class Inhibited < Overloaded; end
|
5
|
+
|
6
|
+
def initialize(*args, **opts)
|
7
|
+
super(**opts)
|
8
|
+
end
|
9
|
+
|
10
|
+
def limit(**opts, &block)
|
11
|
+
unless opts.empty?
|
12
|
+
return self.class.new(
|
13
|
+
**options.merge(opts)
|
14
|
+
).limit(&block)
|
15
|
+
end
|
16
|
+
|
17
|
+
raise Inhibited
|
18
|
+
end
|
19
|
+
|
20
|
+
end
|
21
|
+
end
|
@@ -0,0 +1,84 @@
|
|
1
|
+
module Berater
|
2
|
+
class RateLimiter < BaseLimiter
|
3
|
+
|
4
|
+
class Overrated < Overloaded; end
|
5
|
+
|
6
|
+
attr_accessor :count, :interval
|
7
|
+
|
8
|
+
def initialize(count, interval, **opts)
|
9
|
+
super(**opts)
|
10
|
+
|
11
|
+
self.count = count
|
12
|
+
self.interval = interval
|
13
|
+
end
|
14
|
+
|
15
|
+
def count=(count)
|
16
|
+
unless count.is_a? Integer
|
17
|
+
raise ArgumentError, "expected Integer, found #{count.class}"
|
18
|
+
end
|
19
|
+
|
20
|
+
raise ArgumentError, "count must be >= 0" unless count >= 0
|
21
|
+
|
22
|
+
@count = count
|
23
|
+
end
|
24
|
+
|
25
|
+
def interval=(interval)
|
26
|
+
@interval = interval.dup
|
27
|
+
|
28
|
+
case @interval
|
29
|
+
when Integer
|
30
|
+
raise ArgumentError, "interval must be >= 0" unless @interval >= 0
|
31
|
+
when String
|
32
|
+
@interval = @interval.to_sym
|
33
|
+
when Symbol
|
34
|
+
else
|
35
|
+
raise ArgumentError, "unexpected interval type: #{interval.class}"
|
36
|
+
end
|
37
|
+
|
38
|
+
if @interval.is_a? Symbol
|
39
|
+
case @interval
|
40
|
+
when :sec, :second, :seconds
|
41
|
+
@interval = 1
|
42
|
+
when :min, :minute, :minutes
|
43
|
+
@interval = 60
|
44
|
+
when :hour, :hours
|
45
|
+
@interval = 60 * 60
|
46
|
+
else
|
47
|
+
raise ArgumentError, "unexpected interval value: #{interval}"
|
48
|
+
end
|
49
|
+
end
|
50
|
+
|
51
|
+
@interval
|
52
|
+
end
|
53
|
+
|
54
|
+
def limit(**opts, &block)
|
55
|
+
unless opts.empty?
|
56
|
+
return self.class.new(
|
57
|
+
count,
|
58
|
+
interval,
|
59
|
+
options.merge(opts)
|
60
|
+
).limit(&block)
|
61
|
+
end
|
62
|
+
|
63
|
+
ts = Time.now.to_i
|
64
|
+
|
65
|
+
# bucket into time slot
|
66
|
+
rkey = "%s:%d" % [ key, ts - ts % @interval ]
|
67
|
+
|
68
|
+
count, _ = redis.multi do
|
69
|
+
redis.incr rkey
|
70
|
+
redis.expire rkey, @interval * 2
|
71
|
+
end
|
72
|
+
|
73
|
+
raise Overrated if count > @count
|
74
|
+
|
75
|
+
if block_given?
|
76
|
+
yield
|
77
|
+
else
|
78
|
+
count
|
79
|
+
end
|
80
|
+
end
|
81
|
+
|
82
|
+
end
|
83
|
+
end
|
84
|
+
|
@@ -0,0 +1,19 @@
|
|
1
|
+
module Berater
|
2
|
+
class Unlimiter < BaseLimiter
|
3
|
+
|
4
|
+
def initialize(*args, **opts)
|
5
|
+
super(**opts)
|
6
|
+
end
|
7
|
+
|
8
|
+
def limit(**opts, &block)
|
9
|
+
unless opts.empty?
|
10
|
+
return self.class.new(
|
11
|
+
**options.merge(opts)
|
12
|
+
).limit(&block)
|
13
|
+
end
|
14
|
+
|
15
|
+
yield if block_given?
|
16
|
+
end
|
17
|
+
|
18
|
+
end
|
19
|
+
end
|
@@ -0,0 +1,192 @@
|
|
1
|
+
describe Berater do
|
2
|
+
|
3
|
+
it 'is connected to Redis' do
|
4
|
+
expect(Berater.redis.ping).to eq 'PONG'
|
5
|
+
end
|
6
|
+
|
7
|
+
it { is_expected.to respond_to :configure }
|
8
|
+
|
9
|
+
describe '.configure' do
|
10
|
+
it 'can be set via configure' do
|
11
|
+
Berater.configure do |c|
|
12
|
+
c.mode = :rate
|
13
|
+
end
|
14
|
+
|
15
|
+
expect(Berater.mode).to eq :rate
|
16
|
+
end
|
17
|
+
end
|
18
|
+
|
19
|
+
describe '.redis' do
|
20
|
+
it 'can be reset' do
|
21
|
+
expect(Berater.redis).not_to be_nil
|
22
|
+
Berater.redis = nil
|
23
|
+
expect(Berater.redis).to be_nil
|
24
|
+
end
|
25
|
+
end
|
26
|
+
|
27
|
+
describe '.mode' do
|
28
|
+
it 'validates inputs' do
|
29
|
+
expect { Berater.mode = :foo }.to raise_error(ArgumentError)
|
30
|
+
end
|
31
|
+
end
|
32
|
+
|
33
|
+
context 'unlimited mode' do
|
34
|
+
before { Berater.mode = :unlimited }
|
35
|
+
|
36
|
+
describe '.new' do
|
37
|
+
let(:limiter) { Berater.new(:unlimited) }
|
38
|
+
|
39
|
+
it 'instantiates an Unlimiter' do
|
40
|
+
expect(limiter).to be_a Berater::Unlimiter
|
41
|
+
end
|
42
|
+
|
43
|
+
it 'inherits redis' do
|
44
|
+
expect(limiter.redis).to be Berater.redis
|
45
|
+
end
|
46
|
+
|
47
|
+
it 'accepts options' do
|
48
|
+
redis = double('Redis')
|
49
|
+
limiter = Berater.new(:unlimited, key: 'key', redis: redis)
|
50
|
+
expect(limiter.key).to match(/key/)
|
51
|
+
expect(limiter.redis).to be redis
|
52
|
+
end
|
53
|
+
end
|
54
|
+
|
55
|
+
describe '.limit' do
|
56
|
+
it 'works' do
|
57
|
+
expect(Berater.limit).to be_nil
|
58
|
+
end
|
59
|
+
|
60
|
+
it 'yields' do
|
61
|
+
expect {|b| Berater.limit(&b) }.to yield_control
|
62
|
+
end
|
63
|
+
|
64
|
+
it 'never limits' do
|
65
|
+
10.times { expect(Berater.limit { 123 } ).to eq 123 }
|
66
|
+
end
|
67
|
+
end
|
68
|
+
end
|
69
|
+
|
70
|
+
context 'inhibited mode' do
|
71
|
+
before { Berater.mode = :inhibited }
|
72
|
+
|
73
|
+
describe '.new' do
|
74
|
+
let(:limiter) { Berater.new(:inhibited) }
|
75
|
+
|
76
|
+
it 'instantiates an Inhibitor' do
|
77
|
+
expect(limiter).to be_a Berater::Inhibitor
|
78
|
+
end
|
79
|
+
|
80
|
+
it 'inherits redis' do
|
81
|
+
expect(limiter.redis).to be Berater.redis
|
82
|
+
end
|
83
|
+
|
84
|
+
it 'accepts options' do
|
85
|
+
redis = double('Redis')
|
86
|
+
limiter = Berater.new(:inhibited, key: 'key', redis: redis)
|
87
|
+
expect(limiter.key).to match(/key/)
|
88
|
+
expect(limiter.redis).to be redis
|
89
|
+
end
|
90
|
+
end
|
91
|
+
|
92
|
+
describe '.limit' do
|
93
|
+
it 'always limits' do
|
94
|
+
expect { Berater.limit }.to be_inhibited
|
95
|
+
end
|
96
|
+
end
|
97
|
+
end
|
98
|
+
|
99
|
+
context 'rate mode' do
|
100
|
+
before { Berater.mode = :rate }
|
101
|
+
|
102
|
+
describe '.limiter' do
|
103
|
+
let(:limiter) { Berater.new(:rate, 1, :second) }
|
104
|
+
|
105
|
+
it 'instantiates a RateLimiter' do
|
106
|
+
expect(limiter).to be_a Berater::RateLimiter
|
107
|
+
end
|
108
|
+
|
109
|
+
it 'inherits redis' do
|
110
|
+
expect(limiter.redis).to be Berater.redis
|
111
|
+
end
|
112
|
+
|
113
|
+
it 'accepts options' do
|
114
|
+
redis = double('Redis')
|
115
|
+
limiter = Berater.new(:rate, 1, :second, key: 'key', redis: redis)
|
116
|
+
expect(limiter.key).to match(/key/)
|
117
|
+
expect(limiter.redis).to be redis
|
118
|
+
end
|
119
|
+
end
|
120
|
+
|
121
|
+
describe '.limit' do
|
122
|
+
it 'works' do
|
123
|
+
expect(Berater.limit(1, :second)).to eq 1
|
124
|
+
end
|
125
|
+
|
126
|
+
it 'yields' do
|
127
|
+
expect {|b| Berater.limit(2, :second, &b) }.to yield_control
|
128
|
+
expect(Berater.limit(2, :second) { 123 }).to eq 123
|
129
|
+
end
|
130
|
+
|
131
|
+
it 'limits excessive calls' do
|
132
|
+
expect(Berater.limit(1, :second)).to eq 1
|
133
|
+
expect { Berater.limit(1, :second) }.to be_overrated
|
134
|
+
end
|
135
|
+
|
136
|
+
it 'accepts options' do
|
137
|
+
redis = double('Redis')
|
138
|
+
expect(redis).to receive(:multi)
|
139
|
+
|
140
|
+
Berater.limit(1, :second, redis: redis) rescue nil
|
141
|
+
end
|
142
|
+
end
|
143
|
+
end
|
144
|
+
|
145
|
+
context 'concurrency mode' do
|
146
|
+
before { Berater.mode = :concurrency }
|
147
|
+
|
148
|
+
describe '.limiter' do
|
149
|
+
let(:limiter) { Berater.new(:concurrency, 1) }
|
150
|
+
|
151
|
+
it 'instantiates a ConcurrencyLimiter' do
|
152
|
+
expect(limiter).to be_a Berater::ConcurrencyLimiter
|
153
|
+
end
|
154
|
+
|
155
|
+
it 'inherits redis' do
|
156
|
+
expect(limiter.redis).to be Berater.redis
|
157
|
+
end
|
158
|
+
|
159
|
+
it 'accepts options' do
|
160
|
+
redis = double('Redis')
|
161
|
+
limiter = Berater.new(:concurrency, 1, key: 'key', redis: redis)
|
162
|
+
expect(limiter.key).to match(/key/)
|
163
|
+
expect(limiter.redis).to be redis
|
164
|
+
end
|
165
|
+
end
|
166
|
+
|
167
|
+
describe '.limit' do
|
168
|
+
it 'works (without blocks by returning a lock)' do
|
169
|
+
lock = Berater.limit(1)
|
170
|
+
expect(lock).to be_a Berater::ConcurrencyLimiter::Lock
|
171
|
+
expect(lock.release).to be true
|
172
|
+
end
|
173
|
+
|
174
|
+
it 'yields' do
|
175
|
+
expect {|b| Berater.limit(1, &b) }.to yield_control
|
176
|
+
end
|
177
|
+
|
178
|
+
it 'limits excessive calls' do
|
179
|
+
Berater.limit(1)
|
180
|
+
expect { Berater.limit(1) }.to be_incapacitated
|
181
|
+
end
|
182
|
+
|
183
|
+
it 'accepts options' do
|
184
|
+
redis = double('Redis')
|
185
|
+
expect(redis).to receive(:eval)
|
186
|
+
|
187
|
+
Berater.limit(1, redis: redis) rescue nil
|
188
|
+
end
|
189
|
+
end
|
190
|
+
end
|
191
|
+
|
192
|
+
end
|