berater 0.7.1 → 0.10.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.
checksums.yaml CHANGED
@@ -1,7 +1,7 @@
1
1
  ---
2
2
  SHA256:
3
- metadata.gz: c12ee8f042d6ebb193f114cd093585c14a4a28b802d7a90f79ba5ce41bbd82e6
4
- data.tar.gz: d12ebb1107ea0b24208bfb3ff663a98097cd3367c3972418a139cae9b501e4a1
3
+ metadata.gz: e691d2d3cdc8f002e0e222c24d2e6f9a52ccb1dac7edaa92ea232ecce17fb77e
4
+ data.tar.gz: 5d302fff5fd063f34a94b1aae1b9b2e118b18efbe2315b4591405346f1dbc70c
5
5
  SHA512:
6
- metadata.gz: 491cd616903af917866a41f8b2a4bcc79a469a879f57f1872ce8f4a42c00f1309979c04c142c0083246791aef6cd3847c21197dee0b611ff326ca833cc44ae52
7
- data.tar.gz: 7ba3fd90f8b439e4ad2fe1d5a8008c69978a90c206bc6694e9f470796019d2ac849dddaa85039b9c730e3e20edfec3411cc1275598866c63207c209e688f2a3a
6
+ metadata.gz: 88b8e91557729601336b5235fbcc9710c79aa2f68ff393a09b5c55a1c994cef7371ed4a310c8038b07bdfc5f35d6b367ad7e1fc886c3f7c8579130808af5386b
7
+ data.tar.gz: 2600478d9761b14abbc2ef6c1a38b2e5b5e7476a40c6955f110746d4a90e31baabbda8d6571dcaf7d4e457d0cc19f120e84cba8915ded01b2a680111806d4400
@@ -1,21 +1,22 @@
1
1
  module Berater
2
2
  class ConcurrencyLimiter < Limiter
3
3
 
4
- attr_reader :timeout
5
-
6
4
  def initialize(key, capacity, **opts)
7
5
  super(key, capacity, **opts)
8
6
 
9
- # round fractional capacity
7
+ # truncate fractional capacity
10
8
  self.capacity = capacity.to_i
11
9
 
12
10
  self.timeout = opts[:timeout] || 0
13
11
  end
14
12
 
13
+ def timeout
14
+ options[:timeout]
15
+ end
16
+
15
17
  private def timeout=(timeout)
16
- @timeout = timeout
17
18
  timeout = 0 if timeout == Float::INFINITY
18
- @timeout_msec = Berater::Utils.to_msec(timeout)
19
+ @timeout = Berater::Utils.to_msec(timeout)
19
20
  end
20
21
 
21
22
  LUA_SCRIPT = Berater::LuaScript(<<~LUA
@@ -72,24 +73,19 @@ module Berater
72
73
 
73
74
  count, *lock_ids = LUA_SCRIPT.eval(
74
75
  redis,
75
- [ cache_key(key), cache_key('lock_id') ],
76
- [ capacity, ts, @timeout_msec, cost ]
76
+ [ cache_key, self.class.cache_key('lock_id') ],
77
+ [ capacity, ts, @timeout, cost ]
77
78
  )
78
79
 
79
80
  raise Overloaded if lock_ids.empty?
80
81
 
81
82
  release_fn = if cost > 0
82
- proc { release(lock_ids) }
83
+ proc { redis.zrem(cache_key, lock_ids) }
83
84
  end
84
85
 
85
86
  Lock.new(capacity, count, release_fn)
86
87
  end
87
88
 
88
- private def release(lock_ids)
89
- res = redis.zrem(cache_key(key), lock_ids)
90
- res == true || res == lock_ids.count # depending on which version of Redis
91
- end
92
-
93
89
  def to_s
94
90
  "#<#{self.class}(#{key}: #{capacity} at a time)>"
95
91
  end
@@ -5,6 +5,10 @@ module Berater
5
5
  super(key, 0, **opts)
6
6
  end
7
7
 
8
+ def to_s
9
+ "#<#{self.class}>"
10
+ end
11
+
8
12
  protected def acquire_lock(*)
9
13
  raise Overloaded
10
14
  end
@@ -9,17 +9,12 @@ module Berater
9
9
 
10
10
  def limit(capacity: nil, cost: 1, &block)
11
11
  capacity ||= @capacity
12
+ lock = nil
12
13
 
13
- unless capacity.is_a?(Numeric) && capacity >= 0
14
- raise ArgumentError, "invalid capacity: #{capacity}"
14
+ Berater.middleware.call(self, capacity: capacity, cost: cost) do |limiter, **opts|
15
+ lock = limiter.inner_limit(**opts)
15
16
  end
16
17
 
17
- unless cost.is_a?(Numeric) && cost >= 0 && cost < Float::INFINITY
18
- raise ArgumentError, "invalid cost: #{cost}"
19
- end
20
-
21
- lock = acquire_lock(capacity, cost)
22
-
23
18
  if block_given?
24
19
  begin
25
20
  yield lock
@@ -31,6 +26,23 @@ module Berater
31
26
  end
32
27
  end
33
28
 
29
+ protected def inner_limit(capacity:, cost:)
30
+ unless capacity.is_a?(Numeric) && capacity >= 0
31
+ raise ArgumentError, "invalid capacity: #{capacity}"
32
+ end
33
+
34
+ unless cost.is_a?(Numeric) && cost >= 0 && cost < Float::INFINITY
35
+ raise ArgumentError, "invalid cost: #{cost}"
36
+ end
37
+
38
+ acquire_lock(capacity, cost)
39
+ rescue NoMethodError => e
40
+ raise unless e.message.include?("undefined method `evalsha' for")
41
+
42
+ # repackage error so it's easier to understand
43
+ raise RuntimeError, "invalid redis connection: #{redis}"
44
+ end
45
+
34
46
  def utilization
35
47
  lock = limit(cost: 0)
36
48
 
@@ -43,10 +55,6 @@ module Berater
43
55
  1.0
44
56
  end
45
57
 
46
- def to_s
47
- "#<#{self.class}>"
48
- end
49
-
50
58
  def ==(other)
51
59
  self.class == other.class &&
52
60
  self.key == other.key &&
@@ -56,13 +64,6 @@ module Berater
56
64
  self.redis.connection == other.redis.connection
57
65
  end
58
66
 
59
- def self.new(*)
60
- # can only call via subclass
61
- raise NoMethodError if self == Berater::Limiter
62
-
63
- super
64
- end
65
-
66
67
  protected
67
68
 
68
69
  attr_reader :args
@@ -92,8 +93,39 @@ module Berater
92
93
  raise NotImplementedError
93
94
  end
94
95
 
95
- def cache_key(key)
96
- "#{self.class}:#{key}"
96
+ def cache_key(subkey = nil)
97
+ instance_key = subkey.nil? ? key : "#{key}:#{subkey}"
98
+ self.class.cache_key(instance_key)
99
+ end
100
+
101
+ class << self
102
+ def new(*args, **kwargs)
103
+ # can only call via subclass
104
+ raise NoMethodError if self == Berater::Limiter
105
+
106
+ if RUBY_VERSION < '3' && kwargs.empty?
107
+ # avoid ruby 2 problems with empty hashes
108
+ super(*args)
109
+ else
110
+ super
111
+ end
112
+ end
113
+
114
+ def cache_key(key)
115
+ klass = to_s.split(':')[-1]
116
+ "Berater:#{klass}:#{key}"
117
+ end
118
+
119
+ protected
120
+
121
+ def inherited(subclass)
122
+ # automagically create convenience method
123
+ name = subclass.to_s.split(':')[-1]
124
+
125
+ Berater.define_singleton_method(name) do |*args, **opts, &block|
126
+ Berater::Utils.convenience_fn(subclass, *args, **opts, &block)
127
+ end
128
+ end
97
129
  end
98
130
 
99
131
  end
@@ -0,0 +1,66 @@
1
+ module Berater
2
+ private
3
+
4
+ class LimiterSet
5
+ include Enumerable
6
+
7
+ def initialize
8
+ @limiters = {}
9
+ end
10
+
11
+ def each(&block)
12
+ @limiters.each_value(&block)
13
+ end
14
+
15
+ def <<(limiter)
16
+ key = limiter.key if limiter.respond_to?(:key)
17
+ send(:[]=, key, limiter)
18
+ end
19
+
20
+ def []=(key, limiter)
21
+ unless limiter.is_a? Berater::Limiter
22
+ raise ArgumentError, "expected Berater::Limiter, found: #{limiter}"
23
+ end
24
+
25
+ @limiters[key] = limiter
26
+ end
27
+
28
+ def [](key)
29
+ @limiters[key]
30
+ end
31
+
32
+ def fetch(key, val = default = true, &block)
33
+ args = default ? [ key ] : [ key, val ]
34
+ @limiters.fetch(*args, &block)
35
+ end
36
+
37
+ def include?(key)
38
+ if key.is_a? Berater::Limiter
39
+ @limiters.value?(key)
40
+ else
41
+ @limiters.key?(key)
42
+ end
43
+ end
44
+
45
+ def clear
46
+ @limiters.clear
47
+ end
48
+
49
+ def count
50
+ @limiters.count
51
+ end
52
+
53
+ def delete(key)
54
+ if key.is_a? Berater::Limiter
55
+ @limiters.delete(key.key)
56
+ else
57
+ @limiters.delete(key)
58
+ end
59
+ end
60
+ alias remove delete
61
+
62
+ def empty?
63
+ @limiters.empty?
64
+ end
65
+ end
66
+ end
data/lib/berater/lock.rb CHANGED
@@ -19,7 +19,8 @@ module Berater
19
19
  raise 'lock already released' unless locked?
20
20
 
21
21
  @released_at = Time.now
22
- @release_fn ? @release_fn.call : true
22
+ @release_fn&.call
23
+ true
23
24
  end
24
25
 
25
26
  end
@@ -1,4 +1,5 @@
1
1
  require 'digest'
2
+ require 'redis'
2
3
 
3
4
  module Berater
4
5
  class LuaScript
@@ -6,11 +7,11 @@ module Berater
6
7
  attr_reader :source
7
8
 
8
9
  def initialize(source)
9
- @source = source
10
+ @source = source.dup.freeze
10
11
  end
11
12
 
12
13
  def sha
13
- @sha ||= Digest::SHA1.hexdigest(minify)
14
+ @sha ||= Digest::SHA1.hexdigest(minify).freeze
14
15
  end
15
16
 
16
17
  def eval(redis, *args)
@@ -44,7 +45,7 @@ module Berater
44
45
  def minify
45
46
  # trim comments (whole line and partial)
46
47
  # and whitespace (prefix and empty lines)
47
- @minify ||= source.gsub(/^\s*--.*\n|\s*--.*|^\s*|^$\n/, '').chomp
48
+ @minify ||= source.gsub(/^\s*--.*\n|\s*--.*|^\s*|^$\n/, '').chomp.freeze
48
49
  end
49
50
 
50
51
  end
@@ -1,18 +1,19 @@
1
1
  module Berater
2
2
  class RateLimiter < Limiter
3
3
 
4
- attr_accessor :interval
5
-
6
4
  def initialize(key, capacity, interval, **opts)
5
+ super(key, capacity, interval, **opts)
7
6
  self.interval = interval
8
- super(key, capacity, @interval_msec, **opts)
7
+ end
8
+
9
+ def interval
10
+ args[0]
9
11
  end
10
12
 
11
13
  private def interval=(interval)
12
- @interval = interval
13
- @interval_msec = Berater::Utils.to_msec(interval)
14
+ @interval = Berater::Utils.to_msec(interval)
14
15
 
15
- unless @interval_msec > 0
16
+ unless @interval > 0
16
17
  raise ArgumentError, 'interval must be > 0'
17
18
  end
18
19
  end
@@ -76,8 +77,8 @@ module Berater
76
77
 
77
78
  count, allowed = LUA_SCRIPT.eval(
78
79
  redis,
79
- [ cache_key(key) ],
80
- [ ts, capacity, @interval_msec, cost ]
80
+ [ cache_key ],
81
+ [ ts, capacity, @interval, cost ]
81
82
  )
82
83
 
83
84
  count = count.include?('.') ? count.to_f : count.to_i
data/lib/berater/rspec.rb CHANGED
@@ -8,7 +8,6 @@ RSpec.configure do |config|
8
8
 
9
9
  config.after do
10
10
  Berater.expunge rescue nil
11
- Berater.redis.script(:flush) rescue nil
12
11
  Berater.reset
13
12
  end
14
13
  end
@@ -0,0 +1,49 @@
1
+ module Berater
2
+ class StaticLimiter < Limiter
3
+
4
+ LUA_SCRIPT = Berater::LuaScript(<<~LUA
5
+ local key = KEYS[1]
6
+ local capacity = tonumber(ARGV[1])
7
+ local cost = tonumber(ARGV[2])
8
+
9
+ local count = redis.call('GET', key) or 0
10
+ local allowed = (count + cost) <= capacity
11
+
12
+ if allowed then
13
+ count = count + cost
14
+ redis.call('SET', key, count)
15
+ end
16
+
17
+ return { tostring(count), allowed }
18
+ LUA
19
+ )
20
+
21
+ protected def acquire_lock(capacity, cost)
22
+ if cost == 0
23
+ # utilization check
24
+ count = redis.get(cache_key) || "0"
25
+ allowed = true
26
+ else
27
+ count, allowed = LUA_SCRIPT.eval(
28
+ redis, [ cache_key ], [ capacity, cost ],
29
+ )
30
+ end
31
+
32
+ # Redis returns Floats as strings to maintain precision
33
+ count = count.include?('.') ? count.to_f : count.to_i
34
+
35
+ raise Overloaded unless allowed
36
+
37
+ release_fn = if cost > 0
38
+ proc { redis.incrbyfloat(cache_key, -cost) }
39
+ end
40
+
41
+ Lock.new(capacity, count, release_fn)
42
+ end
43
+
44
+ def to_s
45
+ "#<#{self.class}(#{key}: #{capacity})>"
46
+ end
47
+
48
+ end
49
+ end
@@ -1,36 +1,47 @@
1
1
  require 'berater'
2
2
 
3
3
  module Berater
4
- extend self
5
4
 
6
- attr_reader :test_mode
5
+ module TestMode
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
7
12
 
8
- def test_mode=(mode)
9
- unless [ nil, :pass, :fail ].include?(mode)
10
- raise ArgumentError, "invalid mode: #{Berater.test_mode}"
13
+ @test_mode = mode
11
14
  end
12
15
 
13
- @test_mode = mode
16
+ def reset
17
+ super
18
+ @test_mode = nil
19
+ end
14
20
  end
15
21
 
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
- raise Overloaded
23
- else
24
- super
22
+ class Limiter
23
+ module TestMode
24
+ def acquire_lock(*)
25
+ case Berater.test_mode
26
+ when :pass
27
+ Lock.new(Float::INFINITY, 0)
28
+ when :fail
29
+ raise Overloaded
30
+ else
31
+ super
32
+ end
25
33
  end
26
34
  end
27
35
  end
28
36
 
29
37
  end
30
38
 
39
+ # prepend class methods
40
+ Berater.singleton_class.prepend Berater::TestMode
41
+
31
42
  # stub each Limiter subclass
32
43
  ObjectSpace.each_object(Class).each do |klass|
33
44
  next unless klass < Berater::Limiter
34
45
 
35
- klass.prepend(Berater::TestMode)
46
+ klass.prepend Berater::Limiter::TestMode
36
47
  end
@@ -5,6 +5,10 @@ module Berater
5
5
  super(key, Float::INFINITY, **opts)
6
6
  end
7
7
 
8
+ def to_s
9
+ "#<#{self.class}>"
10
+ end
11
+
8
12
  protected
9
13
 
10
14
  def capacity=(*)
data/lib/berater/utils.rb CHANGED
@@ -42,5 +42,9 @@ module Berater
42
42
  (res * 10**3).to_i
43
43
  end
44
44
 
45
+ def convenience_fn(klass, *args, **opts, &block)
46
+ limiter = klass.new(*args, **opts)
47
+ block ? limiter.limit(&block) : limiter
48
+ end
45
49
  end
46
50
  end
@@ -1,3 +1,3 @@
1
1
  module Berater
2
- VERSION = '0.7.1'
2
+ VERSION = "0.10.1"
3
3
  end
data/lib/berater.rb CHANGED
@@ -1,8 +1,10 @@
1
1
  require 'berater/limiter'
2
+ require 'berater/limiter_set'
2
3
  require 'berater/lock'
3
4
  require 'berater/lua_script'
4
5
  require 'berater/utils'
5
6
  require 'berater/version'
7
+ require 'meddleware'
6
8
 
7
9
  module Berater
8
10
  extend self
@@ -15,8 +17,14 @@ module Berater
15
17
  yield self
16
18
  end
17
19
 
18
- def reset
19
- @redis = nil
20
+ def limiters
21
+ @limiters ||= LimiterSet.new
22
+ end
23
+
24
+ def middleware(&block)
25
+ (@middleware ||= Meddleware.new).tap do
26
+ @middleware.instance_eval(&block) if block_given?
27
+ end
20
28
  end
21
29
 
22
30
  def new(key, capacity, **opts)
@@ -31,8 +39,10 @@ module Berater
31
39
  if opts[:interval]
32
40
  args << opts.delete(:interval)
33
41
  Berater::RateLimiter
34
- else
42
+ elsif opts[:timeout]
35
43
  Berater::ConcurrencyLimiter
44
+ else
45
+ Berater::StaticLimiter
36
46
  end
37
47
  end.yield_self do |klass|
38
48
  args = [ key, capacity, *args ].compact
@@ -46,20 +56,21 @@ module Berater
46
56
  end
47
57
  end
48
58
 
59
+ def reset
60
+ @redis = nil
61
+ limiters.clear
62
+ middleware.clear
63
+ end
49
64
  end
50
65
 
51
66
  # convenience method
52
- def Berater(key, capacity, **opts, &block)
53
- limiter = Berater.new(key, capacity, **opts)
54
- if block_given?
55
- limiter.limit(&block)
56
- else
57
- limiter
58
- end
67
+ def Berater(*args, **opts, &block)
68
+ Berater::Utils.convenience_fn(Berater, *args, **opts, &block)
59
69
  end
60
70
 
61
71
  # load limiters
62
72
  require 'berater/concurrency_limiter'
63
73
  require 'berater/inhibitor'
64
74
  require 'berater/rate_limiter'
75
+ require 'berater/static_limiter'
65
76
  require 'berater/unlimiter'