berater 0.1.4 → 0.6.0

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,26 @@
1
+ module Berater
2
+ class Lock
3
+
4
+ attr_reader :capacity, :contention
5
+
6
+ def initialize(capacity, contention, release_fn = nil)
7
+ @capacity = capacity
8
+ @contention = contention
9
+ @locked_at = Time.now
10
+ @release_fn = release_fn
11
+ @released_at = nil
12
+ end
13
+
14
+ def locked?
15
+ @released_at.nil?
16
+ end
17
+
18
+ def release
19
+ raise 'lock already released' unless locked?
20
+
21
+ @released_at = Time.now
22
+ @release_fn ? @release_fn.call : true
23
+ end
24
+
25
+ end
26
+ end
@@ -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,82 +1,105 @@
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(count, interval, **opts)
9
- super(**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
- 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
- def interval=(interval)
26
- @interval = interval.dup
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)
42
+ end
27
43
 
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
44
+ if cost == 0 then
45
+ -- just check limit, ie. for .overlimit?
46
+ allowed = count < capacity
34
47
  else
35
- raise ArgumentError, "unexpected interval type: #{interval.class}"
36
- end
48
+ allowed = (count + cost) <= capacity
37
49
 
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}"
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(**opts, &block)
55
- unless opts.empty?
56
- return self.class.new(
57
- count,
58
- interval,
59
- options.merge(opts)
60
- ).limit(&block)
61
- end
66
+ def limit(capacity: nil, cost: 1, &block)
67
+ capacity ||= @capacity
62
68
 
63
- ts = Time.now.to_i
69
+ # timestamp in milliseconds
70
+ ts = (Time.now.to_f * 10**3).to_i
64
71
 
65
- # bucket into time slot
66
- rkey = "%s:%d" % [ key, ts - ts % @interval ]
72
+ count, allowed = LUA_SCRIPT.eval(
73
+ redis,
74
+ [ cache_key(key), cache_key("#{key}-ts") ],
75
+ [ ts, capacity, @interval_msec, cost ]
76
+ )
67
77
 
68
- count, _ = redis.multi do
69
- redis.incr rkey
70
- redis.expire rkey, @interval * 2
71
- end
78
+ raise Overrated unless allowed
72
79
 
73
- raise Overrated if count > @count
80
+ lock = Lock.new(capacity, count)
81
+ yield_lock(lock, &block)
82
+ end
74
83
 
75
- if block_given?
76
- yield
84
+ def overloaded?
85
+ limit(cost: 0) { false }
86
+ rescue Overrated
87
+ true
88
+ end
89
+ alias overrated? overloaded?
90
+
91
+ def to_s
92
+ msg = if interval.is_a? Numeric
93
+ if interval == 1
94
+ "every second"
95
+ else
96
+ "every #{interval} seconds"
97
+ end
77
98
  else
78
- count
99
+ "per #{interval}"
79
100
  end
101
+
102
+ "#<#{self.class}(#{key}: #{capacity} #{msg})>"
80
103
  end
81
104
 
82
105
  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,83 @@
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
+ begin
14
+ case obj
15
+ when Proc
16
+ # eg. expect { ... }.to be_overrated
17
+ res = obj.call
18
+
19
+ if res.is_a? Berater::Limiter
20
+ # eg. expect { Berater.new(...) }.to be_overloaded
21
+ @limiter = res
22
+ res.overloaded?
23
+ else
24
+ # eg. expect { Berater(...) }.to be_overloaded
25
+ # eg. expect { limiter.limit }.to be_overloaded
26
+ false
27
+ end
28
+ when Berater::Limiter
29
+ # eg. expect(Berater.new(...)).to be_overloaded
30
+ @limiter = obj
31
+ obj.overloaded?
32
+ end
33
+ rescue @type
34
+ true
35
+ end
36
+ end
37
+
38
+ def description
39
+ if @limiter
40
+ "be #{verb}"
41
+ else
42
+ "raise #{@type}"
43
+ end
44
+ end
45
+
46
+ def failure_message
47
+ if @limiter
48
+ "expected to be #{verb}"
49
+ else
50
+ "expected #{@type} to be raised"
51
+ end
52
+ end
53
+
54
+ def failure_message_when_negated
55
+ if @limiter
56
+ "expected not to be #{verb}"
57
+ else
58
+ "did not expect #{@type} to be raised"
59
+ end
60
+ end
61
+
62
+ private def verb
63
+ @type.to_s.split('::')[-1].downcase
64
+ end
65
+ end
66
+
67
+ def be_overloaded
68
+ Overloaded.new(Berater::Overloaded)
69
+ end
70
+
71
+ def be_overrated
72
+ Overloaded.new(Berater::RateLimiter::Overrated)
73
+ end
74
+
75
+ def be_incapacitated
76
+ Overloaded.new(Berater::ConcurrencyLimiter::Incapacitated)
77
+ end
78
+
79
+ def be_inhibited
80
+ Overloaded.new(Berater::Inhibitor::Inhibited)
81
+ end
82
+ end
83
+ end
@@ -0,0 +1,52 @@
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
+
15
+ # overload class methods
16
+ unless Berater::Limiter.singleton_class.ancestors.include?(TestMode)
17
+ Berater::Limiter.singleton_class.prepend(TestMode)
18
+ end
19
+ end
20
+
21
+ module TestMode
22
+ def new(*args, **opts)
23
+ return super unless Berater.test_mode
24
+
25
+ # chose a stub class with desired behavior
26
+ stub_klass = case Berater.test_mode
27
+ when :pass
28
+ Berater::Unlimiter
29
+ when :fail
30
+ Berater::Inhibitor
31
+ end
32
+
33
+ # don't stub self
34
+ return super if self < stub_klass
35
+
36
+ # swap out limit and overloaded? methods with stub
37
+ super.tap do |instance|
38
+ stub = stub_klass.allocate
39
+ stub.send(:initialize, *args, **opts)
40
+
41
+ instance.define_singleton_method(:limit) do |**opts, &block|
42
+ stub.limit(**opts, &block)
43
+ end
44
+
45
+ instance.define_singleton_method(:overloaded?) do
46
+ stub.overloaded?
47
+ end
48
+ end
49
+ end
50
+ end
51
+
52
+ end
@@ -1,18 +1,20 @@
1
1
  module Berater
2
- class Unlimiter < BaseLimiter
2
+ class Unlimiter < Limiter
3
3
 
4
- def initialize(*args, **opts)
5
- super(**opts)
4
+ def initialize(key = :unlimiter, *args, **opts)
5
+ super(key, Float::INFINITY, **opts)
6
6
  end
7
7
 
8
8
  def limit(**opts, &block)
9
- unless opts.empty?
10
- return self.class.new(
11
- **options.merge(opts)
12
- ).limit(&block)
13
- end
9
+ yield_lock(Lock.new(Float::INFINITY, 0), &block)
10
+ end
11
+
12
+ def overloaded?
13
+ false
14
+ end
14
15
 
15
- yield if block_given?
16
+ protected def capacity=(*)
17
+ @capacity = Float::INFINITY
16
18
  end
17
19
 
18
20
  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