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 CHANGED
@@ -1,7 +1,7 @@
1
1
  ---
2
- SHA1:
3
- metadata.gz: 305ca5d94c6d92dd4b8bf73f2da7a03e15d419c7
4
- data.tar.gz: 83a209c930780a32cf5cda81d1c932d1f8756f88
2
+ SHA256:
3
+ metadata.gz: c8ca6fba7e38a014eed325df41dda7372b41ff62d9e78b61565f4d915d8b4677
4
+ data.tar.gz: 1fae864b93de87ea3dc70531f3f8697e80af526aea53d8ddfe2076eac9bbd61d
5
5
  SHA512:
6
- metadata.gz: b0970141263af6e4b8461c0b974b01abd0873890d295e613e07db5fbb48ba0b736db44b67638e2eb8aa5f3d220d043ad3e7daaf17e58e00cabe999d67bd66f4d
7
- data.tar.gz: d5d02d7331d13c355f3e3d2c56091552507a30e3ae7d978cbaa555984381c04033bd3d93441e8dd58bee9c019de7dd493a983ad28ad5b3bbd3c3fd23b6af8021
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
- VERSION = '0.0.1'
5
+ extend self
3
6
 
4
- class << self
7
+ class Overloaded < StandardError; end
5
8
 
6
- def init redis
7
- @@redis = redis
8
- end
9
+ MODES = {}
10
+
11
+ attr_accessor :redis, :mode
9
12
 
13
+ def configure
14
+ self.mode = :unlimited # default
10
15
 
11
- def incr key, limit, seconds
12
- ts = Time.now.to_i
16
+ yield self
17
+ end
18
+
19
+ def new(mode, *args, **opts)
20
+ klass = MODES[mode.to_sym]
13
21
 
14
- # bucket into time slot
15
- rkey = "%s:%s:%d" % [ self.to_s, key, ts - ts % seconds ]
22
+ unless klass
23
+ raise ArgumentError, "invalid mode: #{mode}"
24
+ end
16
25
 
17
- count = @@redis.multi do
18
- @@redis.incr rkey
19
- @@redis.expire rkey, seconds * 2
20
- end.first
26
+ klass.new(*args, **opts)
27
+ end
21
28
 
22
- raise LimitExceeded if count > limit
29
+ def register(mode, klass)
30
+ MODES[mode.to_sym] = klass
31
+ end
23
32
 
24
- count
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,3 @@
1
+ module Berater
2
+ VERSION = '0.1.4'
3
+ 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