berater 0.2.0 → 0.6.1

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.
@@ -0,0 +1,55 @@
1
+ require 'digest'
2
+
3
+ module Berater
4
+ class LuaScript
5
+
6
+ attr_reader :source
7
+
8
+ def initialize(source)
9
+ @source = source
10
+ end
11
+
12
+ def sha
13
+ @sha ||= Digest::SHA1.hexdigest(minify)
14
+ end
15
+
16
+ def eval(redis, *args)
17
+ redis.evalsha(sha, *args)
18
+ rescue Redis::CommandError => e
19
+ raise unless e.message.include?('NOSCRIPT')
20
+
21
+ # fall back to regular eval, which will trigger
22
+ # script to be cached for next time
23
+ redis.eval(minify, *args)
24
+ end
25
+
26
+ def load(redis)
27
+ redis.script(:load, minify).tap do |sha|
28
+ unless sha == self.sha
29
+ raise "unexpected script SHA: expected #{self.sha}, got #{sha}"
30
+ end
31
+ end
32
+ end
33
+
34
+ def loaded?(redis)
35
+ redis.script(:exists, sha)
36
+ end
37
+
38
+ def to_s
39
+ source
40
+ end
41
+
42
+ private
43
+
44
+ def minify
45
+ # trim comments (whole line and partial)
46
+ # and whitespace (prefix and empty lines)
47
+ @minify ||= source.gsub(/^\s*--.*\n|\s*--.*|^\s*|^$\n/, '').chomp
48
+ end
49
+
50
+ end
51
+
52
+ def LuaScript(source)
53
+ LuaScript.new(source)
54
+ end
55
+ end
@@ -1,80 +1,97 @@
1
1
  module Berater
2
- class RateLimiter < BaseLimiter
2
+ class RateLimiter < Limiter
3
3
 
4
4
  class Overrated < Overloaded; end
5
5
 
6
- attr_accessor :count, :interval
6
+ attr_accessor :interval
7
7
 
8
- def initialize(key, count, interval, **opts)
9
- super(key, **opts)
10
-
11
- self.count = count
8
+ def initialize(key, capacity, interval, **opts)
12
9
  self.interval = interval
10
+ super(key, capacity, @interval_msec, **opts)
13
11
  end
14
12
 
15
- private 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
13
+ private def interval=(interval)
14
+ @interval = interval
15
+ @interval_msec = Berater::Utils.to_msec(interval)
21
16
 
22
- @count = count
17
+ unless @interval_msec > 0
18
+ raise ArgumentError, 'interval must be > 0'
19
+ end
23
20
  end
24
21
 
25
- private 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}"
22
+ LUA_SCRIPT = Berater::LuaScript(<<~LUA
23
+ local key = KEYS[1]
24
+ local ts_key = KEYS[2]
25
+ local ts = tonumber(ARGV[1])
26
+ local capacity = tonumber(ARGV[2])
27
+ local interval_msec = tonumber(ARGV[3])
28
+ local cost = tonumber(ARGV[4])
29
+ local count = 0
30
+ local allowed
31
+ local msec_per_drip = interval_msec / capacity
32
+
33
+ -- timestamp of last update
34
+ local last_ts = tonumber(redis.call('GET', ts_key))
35
+
36
+ if last_ts then
37
+ count = tonumber(redis.call('GET', key)) or 0
38
+
39
+ -- adjust for time passing
40
+ local drips = math.floor((ts - last_ts) / msec_per_drip)
41
+ count = math.max(0, count - drips)
36
42
  end
37
43
 
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}"
44
+ if cost == 0 then
45
+ -- just check limit, ie. for .overlimit?
46
+ allowed = count < capacity
47
+ else
48
+ allowed = (count + cost) <= capacity
49
+
50
+ if allowed then
51
+ count = count + cost
52
+
53
+ -- time for bucket to empty, in milliseconds
54
+ local ttl = math.ceil(count * msec_per_drip)
55
+
56
+ -- update count and last_ts, with expirations
57
+ redis.call('SET', key, count, 'PX', ttl)
58
+ redis.call('SET', ts_key, ts, 'PX', ttl)
48
59
  end
49
60
  end
50
61
 
51
- @interval
52
- end
62
+ return { count, allowed }
63
+ LUA
64
+ )
53
65
 
54
- def limit
55
- ts = Time.now.to_i
66
+ protected def acquire_lock(capacity, cost)
67
+ # timestamp in milliseconds
68
+ ts = (Time.now.to_f * 10**3).to_i
56
69
 
57
- # bucket into time slot
58
- rkey = "%s:%d" % [ cache_key(key), ts - ts % @interval ]
70
+ count, allowed = LUA_SCRIPT.eval(
71
+ redis,
72
+ [ cache_key(key), cache_key("#{key}-ts") ],
73
+ [ ts, capacity, @interval_msec, cost ]
74
+ )
59
75
 
60
- count, _ = redis.multi do
61
- redis.incr rkey
62
- redis.expire rkey, @interval * 2
63
- end
76
+ raise Overrated unless allowed
64
77
 
65
- raise Overrated if count > @count
78
+ Lock.new(capacity, count)
79
+ end
66
80
 
67
- lock = Lock.new(self, count, count)
81
+ alias overrated? overloaded?
68
82
 
69
- if block_given?
70
- begin
71
- yield lock
72
- ensure
73
- lock.release
83
+ def to_s
84
+ msg = if interval.is_a? Numeric
85
+ if interval == 1
86
+ "every second"
87
+ else
88
+ "every #{interval} seconds"
74
89
  end
75
90
  else
76
- lock
91
+ "per #{interval}"
77
92
  end
93
+
94
+ "#<#{self.class}(#{key}: #{capacity} #{msg})>"
78
95
  end
79
96
 
80
97
  end
@@ -0,0 +1,14 @@
1
+ require 'berater'
2
+ require 'berater/rspec/matchers'
3
+ require 'berater/test_mode'
4
+ require 'rspec'
5
+
6
+ RSpec.configure do |config|
7
+ config.include(Berater::Matchers)
8
+
9
+ config.after do
10
+ Berater.expunge rescue nil
11
+ Berater.redis.script(:flush) rescue nil
12
+ Berater.reset
13
+ end
14
+ end
@@ -0,0 +1,81 @@
1
+ module Berater
2
+ module Matchers
3
+ class Overloaded
4
+ def initialize(type)
5
+ @type = type
6
+ end
7
+
8
+ def supports_block_expectations?
9
+ true
10
+ end
11
+
12
+ def matches?(obj)
13
+ case obj
14
+ when Proc
15
+ # eg. expect { ... }.to be_overrated
16
+ res = obj.call
17
+
18
+ if res.is_a? Berater::Limiter
19
+ # eg. expect { Berater.new(...) }.to be_overloaded
20
+ @limiter = res
21
+ res.overloaded?
22
+ else
23
+ # eg. expect { Berater(...) }.to be_overloaded
24
+ # eg. expect { limiter.limit }.to be_overloaded
25
+ false
26
+ end
27
+ when Berater::Limiter
28
+ # eg. expect(Berater.new(...)).to be_overloaded
29
+ @limiter = obj
30
+ obj.overloaded?
31
+ end
32
+ rescue @type
33
+ true
34
+ end
35
+
36
+ def description
37
+ if @limiter
38
+ "be #{verb}"
39
+ else
40
+ "raise #{@type}"
41
+ end
42
+ end
43
+
44
+ def failure_message
45
+ if @limiter
46
+ "expected to be #{verb}"
47
+ else
48
+ "expected #{@type} to be raised"
49
+ end
50
+ end
51
+
52
+ def failure_message_when_negated
53
+ if @limiter
54
+ "expected not to be #{verb}"
55
+ else
56
+ "did not expect #{@type} to be raised"
57
+ end
58
+ end
59
+
60
+ private def verb
61
+ @type.to_s.split('::')[-1].downcase
62
+ end
63
+ end
64
+
65
+ def be_overloaded
66
+ Overloaded.new(Berater::Overloaded)
67
+ end
68
+
69
+ def be_overrated
70
+ Overloaded.new(Berater::RateLimiter::Overrated)
71
+ end
72
+
73
+ def be_incapacitated
74
+ Overloaded.new(Berater::ConcurrencyLimiter::Incapacitated)
75
+ end
76
+
77
+ def be_inhibited
78
+ Overloaded.new(Berater::Inhibitor::Inhibited)
79
+ end
80
+ end
81
+ end
@@ -0,0 +1,43 @@
1
+ require 'berater'
2
+
3
+ module Berater
4
+ extend self
5
+
6
+ attr_reader :test_mode
7
+
8
+ def test_mode=(mode)
9
+ unless [ nil, :pass, :fail ].include?(mode)
10
+ raise ArgumentError, "invalid mode: #{Berater.test_mode}"
11
+ end
12
+
13
+ @test_mode = mode
14
+ end
15
+
16
+ module TestMode
17
+ def acquire_lock(*)
18
+ case Berater.test_mode
19
+ when :pass
20
+ Lock.new(Float::INFINITY, 0)
21
+ when :fail
22
+ # find class specific Overloaded error
23
+ e = self.class.constants.map do |name|
24
+ self.class.const_get(name)
25
+ end.find do |const|
26
+ const < Berater::Overloaded
27
+ end || Berater::Overloaded
28
+
29
+ raise e
30
+ else
31
+ super
32
+ end
33
+ end
34
+ end
35
+
36
+ end
37
+
38
+ # stub each Limiter subclass
39
+ ObjectSpace.each_object(Class).each do |klass|
40
+ next unless klass < Berater::Limiter
41
+
42
+ klass.prepend(Berater::TestMode)
43
+ end
@@ -1,23 +1,18 @@
1
1
  module Berater
2
- class Unlimiter < BaseLimiter
2
+ class Unlimiter < Limiter
3
3
 
4
4
  def initialize(key = :unlimiter, *args, **opts)
5
- super(key, **opts)
5
+ super(key, Float::INFINITY, **opts)
6
6
  end
7
7
 
8
- def limit
9
- count = redis.incr(cache_key('count'))
10
- lock = Lock.new(self, count, count)
8
+ protected
11
9
 
12
- if block_given?
13
- begin
14
- yield lock
15
- ensure
16
- lock.release
17
- end
18
- else
19
- lock
20
- end
10
+ def capacity=(*)
11
+ @capacity = Float::INFINITY
12
+ end
13
+
14
+ def acquire_lock(*)
15
+ Lock.new(Float::INFINITY, 0)
21
16
  end
22
17
 
23
18
  end
@@ -0,0 +1,46 @@
1
+ module Berater
2
+ module Utils
3
+ extend self
4
+
5
+ refine Object do
6
+ def to_msec
7
+ Berater::Utils.to_msec(self)
8
+ end
9
+ end
10
+
11
+ def to_msec(val)
12
+ res = val
13
+
14
+ if val.is_a? String
15
+ # naively attempt casting, otherwise maybe it's a keyword
16
+ res = Float(val) rescue val.to_sym
17
+ end
18
+
19
+ if res.is_a? Symbol
20
+ case res
21
+ when :sec, :second, :seconds
22
+ res = 1
23
+ when :min, :minute, :minutes
24
+ res = 60
25
+ when :hour, :hours
26
+ res = 60 * 60
27
+ end
28
+ end
29
+
30
+ unless res.is_a? Numeric
31
+ raise ArgumentError, "unexpected value: #{val}"
32
+ end
33
+
34
+ if res < 0
35
+ raise ArgumentError, "expected value >= 0, found: #{val}"
36
+ end
37
+
38
+ if res == Float::INFINITY
39
+ raise ArgumentError, "infinite values not allowed"
40
+ end
41
+
42
+ (res * 10**3).to_i
43
+ end
44
+
45
+ end
46
+ end
@@ -1,3 +1,3 @@
1
1
  module Berater
2
- VERSION = '0.2.0'
2
+ VERSION = '0.6.1'
3
3
  end
data/spec/berater_spec.rb CHANGED
@@ -7,12 +7,12 @@ describe Berater do
7
7
  it { is_expected.to respond_to :configure }
8
8
 
9
9
  describe '.configure' do
10
- it 'can be set via configure' do
10
+ it 'is used with a block' do
11
11
  Berater.configure do |c|
12
12
  c.redis = :redis
13
13
  end
14
14
 
15
- expect(Berater.redis).to eq :redis
15
+ expect(Berater.redis).to be :redis
16
16
  end
17
17
  end
18
18
 
@@ -26,7 +26,7 @@ describe Berater do
26
26
 
27
27
  describe '.new' do
28
28
  context 'unlimited mode' do
29
- let(:limiter) { Berater.new(:key, :unlimited) }
29
+ let(:limiter) { Berater.new(:key, Float::INFINITY) }
30
30
 
31
31
  it 'instantiates an Unlimiter' do
32
32
  expect(limiter).to be_a Berater::Unlimiter
@@ -39,18 +39,13 @@ describe Berater do
39
39
 
40
40
  it 'accepts options' do
41
41
  redis = double('Redis')
42
- limiter = Berater.new(:key, :unlimited, redis: redis)
42
+ limiter = Berater.new(:key, Float::INFINITY, redis: redis)
43
43
  expect(limiter.redis).to be redis
44
44
  end
45
-
46
- it 'works with convinience' do
47
- expect(Berater).to receive(:new).and_return(limiter)
48
- expect {|b| Berater(:key, :unlimited, &b) }.to yield_control
49
- end
50
45
  end
51
46
 
52
47
  context 'inhibited mode' do
53
- let(:limiter) { Berater.new(:key, :inhibited) }
48
+ let(:limiter) { Berater.new(:key, 0) }
54
49
 
55
50
  it 'instantiates an Inhibitor' do
56
51
  expect(limiter).to be_a Berater::Inhibitor
@@ -63,18 +58,13 @@ describe Berater do
63
58
 
64
59
  it 'accepts options' do
65
60
  redis = double('Redis')
66
- limiter = Berater.new(:key, :inhibited, redis: redis)
61
+ limiter = Berater.new(:key, 0, redis: redis)
67
62
  expect(limiter.redis).to be redis
68
63
  end
69
-
70
- it 'works with convinience' do
71
- expect(Berater).to receive(:new).and_return(limiter)
72
- expect { Berater(:key, :inhibited) }.to be_inhibited
73
- end
74
64
  end
75
65
 
76
66
  context 'rate mode' do
77
- let(:limiter) { Berater.new(:key, :rate, 1, :second) }
67
+ let(:limiter) { Berater.new(:key, 1, :second) }
78
68
 
79
69
  it 'instantiates a RateLimiter' do
80
70
  expect(limiter).to be_a Berater::RateLimiter
@@ -87,18 +77,13 @@ describe Berater do
87
77
 
88
78
  it 'accepts options' do
89
79
  redis = double('Redis')
90
- limiter = Berater.new(:key, :rate, 1, :second, redis: redis)
80
+ limiter = Berater.new(:key, 1, :second, redis: redis)
91
81
  expect(limiter.redis).to be redis
92
82
  end
93
-
94
- it 'works with convinience' do
95
- expect(Berater).to receive(:new).and_return(limiter)
96
- expect {|b| Berater(:key, :rate, 1, :second, &b) }.to yield_control
97
- end
98
83
  end
99
84
 
100
85
  context 'concurrency mode' do
101
- let(:limiter) { Berater.new(:key, :concurrency, 1) }
86
+ let(:limiter) { Berater.new(:key, 1) }
102
87
 
103
88
  it 'instantiates a ConcurrencyLimiter' do
104
89
  expect(limiter).to be_a Berater::ConcurrencyLimiter
@@ -111,15 +96,39 @@ describe Berater do
111
96
 
112
97
  it 'accepts options' do
113
98
  redis = double('Redis')
114
- limiter = Berater.new(:key, :concurrency, 1, redis: redis)
99
+ limiter = Berater.new(:key, 1, redis: redis)
115
100
  expect(limiter.redis).to be redis
116
101
  end
102
+ end
103
+ end
104
+
105
+ describe 'Berater() - convenience method' do
106
+ RSpec.shared_examples 'test convenience' do |klass, *args|
107
+ it 'creates a limiter' do
108
+ limiter = Berater(:key, *args)
109
+ expect(limiter).to be_a klass
110
+ end
111
+
112
+ context 'with a block' do
113
+ it 'creates a limiter and calls limit' do
114
+ limiter = Berater(:key, *args)
115
+ expect(klass).to receive(:new).and_return(limiter)
116
+ expect(limiter).to receive(:limit).and_call_original
117
117
 
118
- it 'works with convinience' do
119
- expect(Berater).to receive(:new).and_return(limiter)
120
- expect {|b| Berater(:key, :concurrency, 1, &b) }.to yield_control
118
+ begin
119
+ res = Berater(:key, *args) { true }
120
+ expect(res).to be true
121
+ rescue Berater::Overloaded
122
+ expect(klass).to be Berater::Inhibitor
123
+ end
124
+ end
121
125
  end
122
126
  end
127
+
128
+ include_examples 'test convenience', Berater::Unlimiter, Float::INFINITY
129
+ include_examples 'test convenience', Berater::Inhibitor, 0
130
+ include_examples 'test convenience', Berater::RateLimiter, 1, :second
131
+ include_examples 'test convenience', Berater::ConcurrencyLimiter, 1
123
132
  end
124
133
 
125
134
  end