berater 0.3.0 → 0.6.2

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: d007dbb664fa57c921047a28dd1e5a591cf5c51fba4931dbc367ab552e4b8a8d
4
- data.tar.gz: c710d31a2d8cfdc50af4c2fae87cabad177096d13da52a7f9c48b61aa31b6038
3
+ metadata.gz: dd61012052d49f83a59057ed5ce4cc46b9c25a2d2bf4b5e494f02e590c2e2249
4
+ data.tar.gz: 3add1b481830566088eb7a9dce3540716b55dda9bb658a7b26b93a261b403f3a
5
5
  SHA512:
6
- metadata.gz: 0d9c211bb65e7341b98e637d3f0ce0e31f0a879437c1c309fb01e1efab5e5c6b2d5044158931ab16a67cbbdebd00ddb5f10afc97430230b153d8cc236f7517ab
7
- data.tar.gz: 5d8b38bfd7683df641fd6891cd1a493079453ea6c61b2880581eeb2d963fb171365fb3448443e2548b34297d8dcfa352e67cff5688cd24d25422efb5e55d4afd
6
+ metadata.gz: 476cf552a5fb56498f81a737dba79358d016bedbd76c3c2cc8069c420e82a221c54280781821ef701d25da2551ddc25910818dc400dc7c046cecbe8f6db0ddd2
7
+ data.tar.gz: a92e64f027425ecdc4a39ee5f88a792f4f2625ec4bcf65fac70e34d591fb136e7544caa80138474ae5f5e6158fc19ec8894892585cf858966bfe5fe13a3d011b
data/lib/berater.rb CHANGED
@@ -1,48 +1,40 @@
1
- require 'berater/version'
1
+ require 'berater/limiter'
2
2
  require 'berater/lock'
3
-
3
+ require 'berater/lua_script'
4
+ require 'berater/utils'
5
+ require 'berater/version'
4
6
 
5
7
  module Berater
6
8
  extend self
7
9
 
8
10
  class Overloaded < StandardError; end
9
11
 
10
- MODES = {}
11
-
12
12
  attr_accessor :redis
13
13
 
14
14
  def configure
15
15
  yield self
16
16
  end
17
17
 
18
- def new(key, mode = nil, *args, **opts, &block)
19
- if mode.nil?
20
- unless args.empty?
21
- raise ArgumentError, '0 arguments expected with block'
22
- end
23
-
24
- unless block_given?
25
- raise ArgumentError, 'expected either mode or block'
26
- end
18
+ def reset
19
+ @redis = nil
20
+ end
27
21
 
28
- mode, *args = DSL.eval(&block)
22
+ def new(key, capacity, interval = nil, **opts)
23
+ case capacity
24
+ when :unlimited, Float::INFINITY
25
+ Berater::Unlimiter
26
+ when :inhibited, 0
27
+ Berater::Inhibitor
29
28
  else
30
- if block_given?
31
- raise ArgumentError, 'expected either mode or block, not both'
29
+ if interval
30
+ Berater::RateLimiter
31
+ else
32
+ Berater::ConcurrencyLimiter
32
33
  end
34
+ end.yield_self do |klass|
35
+ args = [ key, capacity, interval ].compact
36
+ klass.new(*args, **opts)
33
37
  end
34
-
35
- klass = MODES[mode.to_sym]
36
-
37
- unless klass
38
- raise ArgumentError, "invalid mode: #{mode}"
39
- end
40
-
41
- klass.new(key, *args, **opts)
42
- end
43
-
44
- def register(mode, klass)
45
- MODES[mode.to_sym] = klass
46
38
  end
47
39
 
48
40
  def expunge
@@ -54,20 +46,17 @@ module Berater
54
46
  end
55
47
 
56
48
  # convenience method
57
- def Berater(key, mode, *args, **opts, &block)
58
- Berater.new(key, mode, *args, **opts).limit(&block)
49
+ def Berater(key, capacity, interval = nil, **opts, &block)
50
+ limiter = Berater.new(key, capacity, interval, **opts)
51
+ if block_given?
52
+ limiter.limit(&block)
53
+ else
54
+ limiter
55
+ end
59
56
  end
60
57
 
61
58
  # load limiters
62
- require 'berater/limiter'
63
59
  require 'berater/concurrency_limiter'
64
60
  require 'berater/inhibitor'
65
61
  require 'berater/rate_limiter'
66
62
  require 'berater/unlimiter'
67
-
68
- Berater.register(:concurrency, Berater::ConcurrencyLimiter)
69
- Berater.register(:inhibited, Berater::Inhibitor)
70
- Berater.register(:rate, Berater::RateLimiter)
71
- Berater.register(:unlimited, Berater::Unlimiter)
72
-
73
- require 'berater/dsl'
@@ -3,42 +3,28 @@ module Berater
3
3
 
4
4
  class Incapacitated < Overloaded; end
5
5
 
6
- attr_reader :capacity, :timeout
6
+ attr_reader :timeout
7
7
 
8
8
  def initialize(key, capacity, **opts)
9
- super(key, **opts)
9
+ super(key, capacity, **opts)
10
10
 
11
- self.capacity = capacity
12
11
  self.timeout = opts[:timeout] || 0
13
12
  end
14
13
 
15
- private 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
14
  private 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
15
  @timeout = timeout
16
+ timeout = 0 if timeout == Float::INFINITY
17
+ @timeout_msec = Berater::Utils.to_msec(timeout)
33
18
  end
34
19
 
35
- LUA_SCRIPT = <<~LUA.gsub(/^\s*(--.*\n)?/, '')
20
+ LUA_SCRIPT = Berater::LuaScript(<<~LUA
36
21
  local key = KEYS[1]
37
22
  local lock_key = KEYS[2]
38
23
  local capacity = tonumber(ARGV[1])
39
24
  local ts = tonumber(ARGV[2])
40
25
  local ttl = tonumber(ARGV[3])
41
- local lock
26
+ local cost = tonumber(ARGV[4])
27
+ local lock_ids = {}
42
28
 
43
29
  -- purge stale hosts
44
30
  if ttl > 0 then
@@ -48,41 +34,63 @@ module Berater
48
34
  -- check capacity
49
35
  local count = redis.call('ZCARD', key)
50
36
 
51
- if count < capacity then
52
- -- grab a lock
53
- lock = redis.call('INCR', lock_key)
54
- redis.call('ZADD', key, ts, lock)
55
- count = count + 1
37
+ if cost == 0 then
38
+ -- just check limit, ie. for .overlimit?
39
+ if count < capacity then
40
+ table.insert(lock_ids, true)
41
+ end
42
+ elseif (count + cost) <= capacity then
43
+ -- grab locks, one per cost
44
+ local lock_id = redis.call('INCRBY', lock_key, cost)
45
+ local locks = {}
46
+
47
+ for i = lock_id - cost + 1, lock_id do
48
+ table.insert(lock_ids, i)
49
+
50
+ table.insert(locks, ts)
51
+ table.insert(locks, i)
52
+ end
53
+
54
+ redis.call('ZADD', key, unpack(locks))
55
+ count = count + cost
56
+
57
+ if ttl > 0 then
58
+ redis.call('PEXPIRE', key, ttl)
59
+ end
56
60
  end
57
61
 
58
- return { count, lock }
62
+ return { count, unpack(lock_ids) }
59
63
  LUA
64
+ )
65
+
66
+ protected def acquire_lock(capacity, cost)
67
+ # fractional cost is not supported, but make it work
68
+ capacity = capacity.to_i
69
+ cost = cost.ceil
60
70
 
61
- def limit
62
- count, lock_id = redis.eval(
63
- LUA_SCRIPT,
71
+ # timestamp in milliseconds
72
+ ts = (Time.now.to_f * 10**3).to_i
73
+
74
+ count, *lock_ids = LUA_SCRIPT.eval(
75
+ redis,
64
76
  [ cache_key(key), cache_key('lock_id') ],
65
- [ capacity, Time.now.to_i, timeout ]
77
+ [ capacity, ts, @timeout_msec, cost ]
66
78
  )
67
79
 
68
- raise Incapacitated unless lock_id
69
-
70
- lock = Lock.new(self, lock_id, count, -> { release(lock_id) })
80
+ raise Incapacitated if lock_ids.empty?
71
81
 
72
- if block_given?
73
- begin
74
- yield lock
75
- ensure
76
- lock.release
77
- end
78
- else
79
- lock
82
+ release_fn = if cost > 0
83
+ proc { release(lock_ids) }
80
84
  end
85
+
86
+ Lock.new(capacity, count, release_fn)
81
87
  end
82
88
 
83
- private def release(lock_id)
84
- res = redis.zrem(cache_key(key), lock_id)
85
- res == true || res == 1 # depending on which version of Redis
89
+ alias incapacitated? overloaded?
90
+
91
+ private def release(lock_ids)
92
+ res = redis.zrem(cache_key(key), lock_ids)
93
+ res == true || res == lock_ids.count # depending on which version of Redis
86
94
  end
87
95
 
88
96
  def to_s
data/lib/berater/dsl.rb CHANGED
@@ -1,5 +1,23 @@
1
1
  module Berater
2
2
  module DSL
3
+ refine Berater.singleton_class do
4
+ def new(key, mode = nil, *args, **opts, &block)
5
+ if mode.nil?
6
+ unless block_given?
7
+ raise ArgumentError, 'expected either mode or block'
8
+ end
9
+
10
+ mode, *args = DSL.eval(&block)
11
+ else
12
+ if block_given?
13
+ raise ArgumentError, 'expected either mode or block, not both'
14
+ end
15
+ end
16
+
17
+ super(key, mode, *args, **opts)
18
+ end
19
+ end
20
+
3
21
  extend self
4
22
 
5
23
  def eval &block
@@ -18,13 +36,6 @@ module Berater
18
36
 
19
37
  private
20
38
 
21
- def each &block
22
- Berater::MODES.map do |mode, limiter|
23
- next unless limiter.const_defined?(:DSL, false)
24
- limiter.const_get(:DSL)
25
- end.compact.each(&block)
26
- end
27
-
28
39
  KEYWORDS = [
29
40
  :second, :minute, :hour,
30
41
  :unlimited, :inhibited,
@@ -33,12 +44,12 @@ module Berater
33
44
  def install
34
45
  Integer.class_eval do
35
46
  def per(unit)
36
- [ :rate, self, unit ]
47
+ [ self, unit ]
37
48
  end
38
49
  alias every per
39
50
 
40
51
  def at_once
41
- [ :concurrency, self ]
52
+ [ self ]
42
53
  end
43
54
  alias concurrently at_once
44
55
  alias at_a_time at_once
@@ -4,10 +4,12 @@ module Berater
4
4
  class Inhibited < Overloaded; end
5
5
 
6
6
  def initialize(key = :inhibitor, *args, **opts)
7
- super(key, **opts)
7
+ super(key, 0, **opts)
8
8
  end
9
9
 
10
- def limit
10
+ alias inhibited? overloaded?
11
+
12
+ protected def acquire_lock(*)
11
13
  raise Inhibited
12
14
  end
13
15
 
@@ -1,27 +1,91 @@
1
1
  module Berater
2
2
  class Limiter
3
3
 
4
- attr_reader :key, :options
4
+ attr_reader :key, :capacity, :options
5
5
 
6
6
  def redis
7
7
  options[:redis] || Berater.redis
8
8
  end
9
9
 
10
- def limit
11
- raise NotImplementedError
10
+ def limit(capacity: nil, cost: 1, &block)
11
+ capacity ||= @capacity
12
+
13
+ unless capacity.is_a?(Numeric)
14
+ raise ArgumentError, "invalid capacity: #{capacity}"
15
+ end
16
+
17
+ unless cost.is_a?(Numeric) && cost >= 0
18
+ raise ArgumentError, "invalid cost: #{cost}"
19
+ end
20
+
21
+ lock = acquire_lock(capacity, cost)
22
+
23
+ if block_given?
24
+ begin
25
+ yield lock
26
+ ensure
27
+ lock.release
28
+ end
29
+ else
30
+ lock
31
+ end
32
+ end
33
+
34
+ def overloaded?
35
+ limit(cost: 0) { false }
36
+ rescue Overloaded
37
+ true
12
38
  end
13
39
 
14
40
  def to_s
15
41
  "#<#{self.class}>"
16
42
  end
17
43
 
44
+ def ==(other)
45
+ self.class == other.class &&
46
+ self.key == other.key &&
47
+ self.capacity == other.capacity &&
48
+ self.args == other.args &&
49
+ self.options == other.options &&
50
+ self.redis.connection == other.redis.connection
51
+ end
52
+
53
+ def self.new(*)
54
+ # can only call via subclass
55
+ raise NoMethodError if self == Berater::Limiter
56
+
57
+ super
58
+ end
59
+
18
60
  protected
19
61
 
20
- def initialize(key, **opts)
62
+ attr_reader :args
63
+
64
+ def initialize(key, capacity, *args, **opts)
21
65
  @key = key
66
+ self.capacity = capacity
67
+ @args = args
22
68
  @options = opts
23
69
  end
24
70
 
71
+ def capacity=(capacity)
72
+ unless capacity.is_a?(Numeric)
73
+ raise ArgumentError, "expected Numeric, found #{capacity.class}"
74
+ end
75
+
76
+ if capacity == Float::INFINITY
77
+ raise ArgumentError, 'infinite capacity not supported, use Unlimiter'
78
+ end
79
+
80
+ raise ArgumentError, 'capacity must be >= 0' unless capacity >= 0
81
+
82
+ @capacity = capacity
83
+ end
84
+
85
+ def acquire_lock(capacity, cost)
86
+ raise NotImplementedError
87
+ end
88
+
25
89
  def cache_key(key)
26
90
  "#{self.class}:#{key}"
27
91
  end
data/lib/berater/lock.rb CHANGED
@@ -1,11 +1,10 @@
1
1
  module Berater
2
2
  class Lock
3
3
 
4
- attr_reader :limiter, :id, :contention
4
+ attr_reader :capacity, :contention
5
5
 
6
- def initialize(limiter, id, contention, release_fn = nil)
7
- @limiter = limiter
8
- @id = id
6
+ def initialize(capacity, contention, release_fn = nil)
7
+ @capacity = capacity
9
8
  @contention = contention
10
9
  @locked_at = Time.now
11
10
  @release_fn = release_fn
@@ -13,24 +12,15 @@ module Berater
13
12
  end
14
13
 
15
14
  def locked?
16
- @released_at.nil? && !expired?
17
- end
18
-
19
- def expired?
20
- timeout ? @locked_at + timeout < Time.now : false
15
+ @released_at.nil?
21
16
  end
22
17
 
23
18
  def release
24
- raise 'lock expired' if expired?
25
19
  raise 'lock already released' unless locked?
26
20
 
27
21
  @released_at = Time.now
28
22
  @release_fn ? @release_fn.call : true
29
23
  end
30
24
 
31
- def timeout
32
- limiter.respond_to?(:timeout) ? limiter.timeout : nil
33
- end
34
-
35
25
  end
36
26
  end