berater 0.4.0 → 0.7.0

Sign up to get free protection for your applications and to get access to all the features.
checksums.yaml CHANGED
@@ -1,7 +1,7 @@
1
1
  ---
2
2
  SHA256:
3
- metadata.gz: f57aeaa9319ae38bdd7d2468443fc5c9ea9ba44cecf47bc5a558f18ecd0ab274
4
- data.tar.gz: ff14ef90856911e9e18cbfbad6f9799727c999697c5e4d96f47b5fbf665c725c
3
+ metadata.gz: 7d25d186bfb9e709986e5c205c9fd2b8107dcf8df1e915316d28a6a59a0d4c0f
4
+ data.tar.gz: '05380448ad96697b5d3753a93584d4bdfcf3b028c649ca11360a1c92d6afadad'
5
5
  SHA512:
6
- metadata.gz: c9d3abeeb9d35c608c194eaf505c5af1350a82583dd12257e70b9d4f96bd3a6cd30db802f57c3bf77d371b5deb5e856e9238eadd3990768205fd37ea3e2e3703
7
- data.tar.gz: c98078a3c7d3a3eb4e2d523577f9d6dd7f1bbae76c9ecbd8d65961cc7aa1bab6e1a4b8fe84cfb30ee60f4d0ce11dcb58e3b94a3566cf50ee60d1f8e21bab82b1
6
+ metadata.gz: 6ddcd488d3d1aa293d621f32ad8cd684daffb0ba9c413e8cf6a7a292ae63a307e2cdf8e62803cf2a8c2b3ea7b9424e4874bb773f49620e768895b193385e847f
7
+ data.tar.gz: cf923f95875da8b7e6c0c0cbe4dccad5a4e3c9f6436da606a03f93c2b9ff86525485c6379ade0f276f30e1aff3bcf0c279fe88d33e6769455818482a5ff1a179
data/lib/berater.rb CHANGED
@@ -1,48 +1,43 @@
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
18
+ def reset
19
+ @redis = nil
20
+ end
23
21
 
24
- unless block_given?
25
- raise ArgumentError, 'expected either mode or block'
26
- end
22
+ def new(key, capacity, **opts)
23
+ args = []
27
24
 
28
- mode, *args = DSL.eval(&block)
25
+ case capacity
26
+ when Float::INFINITY
27
+ Berater::Unlimiter
28
+ when 0
29
+ Berater::Inhibitor
29
30
  else
30
- if block_given?
31
- raise ArgumentError, 'expected either mode or block, not both'
31
+ if opts[:interval]
32
+ args << opts.delete(:interval)
33
+ Berater::RateLimiter
34
+ else
35
+ Berater::ConcurrencyLimiter
32
36
  end
37
+ end.yield_self do |klass|
38
+ args = [ key, capacity, *args ].compact
39
+ klass.new(*args, **opts)
33
40
  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
41
  end
47
42
 
48
43
  def expunge
@@ -54,20 +49,17 @@ module Berater
54
49
  end
55
50
 
56
51
  # convenience method
57
- def Berater(key, mode, *args, **opts, &block)
58
- Berater.new(key, mode, *args, **opts).limit(&block)
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
59
59
  end
60
60
 
61
61
  # load limiters
62
- require 'berater/limiter'
63
62
  require 'berater/concurrency_limiter'
64
63
  require 'berater/inhibitor'
65
64
  require 'berater/rate_limiter'
66
65
  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'
@@ -1,88 +1,93 @@
1
1
  module Berater
2
2
  class ConcurrencyLimiter < Limiter
3
3
 
4
- class Incapacitated < Overloaded; end
5
-
6
- attr_reader :capacity, :timeout
4
+ attr_reader :timeout
7
5
 
8
6
  def initialize(key, capacity, **opts)
9
- super(key, **opts)
10
-
11
- self.capacity = capacity
12
- self.timeout = opts[:timeout] || 0
13
- end
14
-
15
- private def capacity=(capacity)
16
- unless capacity.is_a? Integer
17
- raise ArgumentError, "expected Integer, found #{capacity.class}"
18
- end
7
+ super(key, capacity, **opts)
19
8
 
20
- raise ArgumentError, "capacity must be >= 0" unless capacity >= 0
9
+ # round fractional capacity
10
+ self.capacity = capacity.to_i
21
11
 
22
- @capacity = capacity
12
+ self.timeout = opts[:timeout] || 0
23
13
  end
24
14
 
25
15
  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
16
  @timeout = timeout
17
+ timeout = 0 if timeout == Float::INFINITY
18
+ @timeout_msec = Berater::Utils.to_msec(timeout)
33
19
  end
34
20
 
35
- LUA_SCRIPT = <<~LUA.gsub(/^\s*|\s*--.*/, '')
21
+ LUA_SCRIPT = Berater::LuaScript(<<~LUA
36
22
  local key = KEYS[1]
37
23
  local lock_key = KEYS[2]
38
24
  local capacity = tonumber(ARGV[1])
39
25
  local ts = tonumber(ARGV[2])
40
26
  local ttl = tonumber(ARGV[3])
41
- local lock
27
+ local cost = tonumber(ARGV[4])
28
+ local lock_ids = {}
42
29
 
43
30
  -- purge stale hosts
44
31
  if ttl > 0 then
45
- redis.call('ZREMRANGEBYSCORE', key, '-inf', ts - ttl)
32
+ redis.call('ZREMRANGEBYSCORE', key, 0, ts - ttl)
46
33
  end
47
34
 
48
35
  -- check capacity
49
36
  local count = redis.call('ZCARD', key)
50
37
 
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
38
+ if cost == 0 then
39
+ -- just checking count
40
+ table.insert(lock_ids, true)
41
+ elseif (count + cost) <= capacity then
42
+ -- grab locks, one per cost
43
+ local lock_id = redis.call('INCRBY', lock_key, cost)
44
+ local locks = {}
45
+
46
+ for i = lock_id - cost + 1, lock_id do
47
+ table.insert(lock_ids, i)
48
+
49
+ table.insert(locks, ts)
50
+ table.insert(locks, i)
51
+ end
52
+
53
+ redis.call('ZADD', key, unpack(locks))
54
+ count = count + cost
55
+
56
+ if ttl > 0 then
57
+ redis.call('PEXPIRE', key, ttl)
58
+ end
56
59
  end
57
60
 
58
- return { count, lock }
61
+ return { count, unpack(lock_ids) }
59
62
  LUA
63
+ )
60
64
 
61
- def limit
62
- count, lock_id = redis.eval(
63
- LUA_SCRIPT,
65
+ protected def acquire_lock(capacity, cost)
66
+ # round fractional capacity and cost
67
+ capacity = capacity.to_i
68
+ cost = cost.ceil
69
+
70
+ # timestamp in milliseconds
71
+ ts = (Time.now.to_f * 10**3).to_i
72
+
73
+ count, *lock_ids = LUA_SCRIPT.eval(
74
+ redis,
64
75
  [ cache_key(key), cache_key('lock_id') ],
65
- [ capacity, Time.now.to_i, timeout ]
76
+ [ capacity, ts, @timeout_msec, cost ]
66
77
  )
67
78
 
68
- raise Incapacitated unless lock_id
69
-
70
- lock = Lock.new(self, lock_id, count, -> { release(lock_id) })
79
+ raise Overloaded if lock_ids.empty?
71
80
 
72
- if block_given?
73
- begin
74
- yield lock
75
- ensure
76
- lock.release
77
- end
78
- else
79
- lock
81
+ release_fn = if cost > 0
82
+ proc { release(lock_ids) }
80
83
  end
84
+
85
+ Lock.new(capacity, count, release_fn)
81
86
  end
82
87
 
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
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
86
91
  end
87
92
 
88
93
  def to_s
data/lib/berater/dsl.rb CHANGED
@@ -1,5 +1,24 @@
1
1
  module Berater
2
2
  module DSL
3
+ refine Berater.singleton_class do
4
+ def new(key, capacity = nil, **opts, &block)
5
+ if capacity.nil?
6
+ unless block_given?
7
+ raise ArgumentError, 'expected either capacity or block'
8
+ end
9
+
10
+ capacity, more_opts = DSL.eval(&block)
11
+ opts.merge!(more_opts) if more_opts
12
+ else
13
+ if block_given?
14
+ raise ArgumentError, 'expected either capacity or block, not both'
15
+ end
16
+ end
17
+
18
+ super(key, capacity, **opts)
19
+ end
20
+ end
21
+
3
22
  extend self
4
23
 
5
24
  def eval &block
@@ -18,27 +37,19 @@ module Berater
18
37
 
19
38
  private
20
39
 
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
40
  KEYWORDS = [
29
41
  :second, :minute, :hour,
30
- :unlimited, :inhibited,
31
42
  ].freeze
32
43
 
33
44
  def install
34
45
  Integer.class_eval do
35
46
  def per(unit)
36
- [ :rate, self, unit ]
47
+ [ self, interval: 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
@@ -1,14 +1,12 @@
1
1
  module Berater
2
2
  class Inhibitor < Limiter
3
3
 
4
- class Inhibited < Overloaded; end
5
-
6
4
  def initialize(key = :inhibitor, *args, **opts)
7
- super(key, **opts)
5
+ super(key, 0, **opts)
8
6
  end
9
7
 
10
- def limit
11
- raise Inhibited
8
+ protected def acquire_lock(*)
9
+ raise Overloaded
12
10
  end
13
11
 
14
12
  end
@@ -1,27 +1,97 @@
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) && capacity >= 0
14
+ raise ArgumentError, "invalid capacity: #{capacity}"
15
+ end
16
+
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
+ 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 utilization
35
+ lock = limit(cost: 0)
36
+
37
+ if lock.capacity == 0
38
+ 1.0
39
+ else
40
+ lock.contention.to_f / lock.capacity
41
+ end
42
+ rescue Berater::Overloaded
43
+ 1.0
12
44
  end
13
45
 
14
46
  def to_s
15
47
  "#<#{self.class}>"
16
48
  end
17
49
 
50
+ def ==(other)
51
+ self.class == other.class &&
52
+ self.key == other.key &&
53
+ self.capacity == other.capacity &&
54
+ self.args == other.args &&
55
+ self.options == other.options &&
56
+ self.redis.connection == other.redis.connection
57
+ end
58
+
59
+ def self.new(*)
60
+ # can only call via subclass
61
+ raise NoMethodError if self == Berater::Limiter
62
+
63
+ super
64
+ end
65
+
18
66
  protected
19
67
 
20
- def initialize(key, **opts)
68
+ attr_reader :args
69
+
70
+ def initialize(key, capacity, *args, **opts)
21
71
  @key = key
72
+ self.capacity = capacity
73
+ @args = args
22
74
  @options = opts
23
75
  end
24
76
 
77
+ def capacity=(capacity)
78
+ unless capacity.is_a?(Numeric)
79
+ raise ArgumentError, "expected Numeric, found #{capacity.class}"
80
+ end
81
+
82
+ if capacity == Float::INFINITY
83
+ raise ArgumentError, 'infinite capacity not supported, use Unlimiter'
84
+ end
85
+
86
+ raise ArgumentError, 'capacity must be >= 0' unless capacity >= 0
87
+
88
+ @capacity = capacity
89
+ end
90
+
91
+ def acquire_lock(capacity, cost)
92
+ raise NotImplementedError
93
+ end
94
+
25
95
  def cache_key(key)
26
96
  "#{self.class}:#{key}"
27
97
  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