berater 0.4.0 → 0.5.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: fa5d51d6eea9d4691e58e8dd7065a1d596a57abdac1227fa19bd22cd2f42c15d
4
+ data.tar.gz: 55ac15bf6f2f295c525fdea57b3af9b518935a12dd4b644e3c0c9c085429ca50
5
5
  SHA512:
6
- metadata.gz: c9d3abeeb9d35c608c194eaf505c5af1350a82583dd12257e70b9d4f96bd3a6cd30db802f57c3bf77d371b5deb5e856e9238eadd3990768205fd37ea3e2e3703
7
- data.tar.gz: c98078a3c7d3a3eb4e2d523577f9d6dd7f1bbae76c9ecbd8d65961cc7aa1bab6e1a4b8fe84cfb30ee60f4d0ce11dcb58e3b94a3566cf50ee60d1f8e21bab82b1
6
+ metadata.gz: 52df5781f9c37c7f1c9dde09e579c8da1f2d62f405835b771c6087cda8af046018e85a88629f7c5d3a2be9aafd2680e348d1b2f91bc63fe456e6e51ec5d0fd02
7
+ data.tar.gz: '0928cb6f2e8e07810dec96f5af302ff9f62d519b9f19ef6f28acb5621569f13edb6d784cbd7084b738df8dfc232ace26e3ce0c6c4a7f2ac3e4391c97789bafae'
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_usec = Berater::Utils.to_usec(timeout)
33
18
  end
34
19
 
35
- LUA_SCRIPT = <<~LUA.gsub(/^\s*|\s*--.*/, '')
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,69 @@ 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
+ def limit(capacity: nil, cost: 1, &block)
67
+ capacity ||= @capacity
68
+ # cost is Integer >= 0
69
+
70
+ # timestamp in microseconds
71
+ ts = (Time.now.to_f * 10**6).to_i
60
72
 
61
- def limit
62
- count, lock_id = redis.eval(
63
- LUA_SCRIPT,
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_usec, cost ]
66
77
  )
67
78
 
68
- raise Incapacitated unless lock_id
79
+ raise Incapacitated if lock_ids.empty?
69
80
 
70
- lock = Lock.new(self, lock_id, count, -> { release(lock_id) })
71
-
72
- if block_given?
73
- begin
74
- yield lock
75
- ensure
76
- lock.release
77
- end
81
+ if cost == 0
82
+ lock = Lock.new(self, nil, count)
78
83
  else
79
- lock
84
+ lock = Lock.new(self, lock_ids[0], count, -> { release(lock_ids) })
80
85
  end
86
+
87
+ yield_lock(lock, &block)
88
+ end
89
+
90
+ def overloaded?
91
+ limit(cost: 0) { false }
92
+ rescue Overloaded
93
+ true
81
94
  end
95
+ alias incapacitated? overloaded?
82
96
 
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
97
+ private def release(lock_ids)
98
+ res = redis.zrem(cache_key(key), lock_ids)
99
+ res == true || res == lock_ids.count # depending on which version of Redis
86
100
  end
87
101
 
88
102
  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,12 +4,17 @@ 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
+ def limit(**opts)
11
11
  raise Inhibited
12
12
  end
13
13
 
14
+ def overloaded?
15
+ true
16
+ end
17
+ alias inhibited? overloaded?
18
+
14
19
  end
15
20
  end
@@ -1,7 +1,7 @@
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
@@ -11,20 +11,66 @@ module Berater
11
11
  raise NotImplementedError
12
12
  end
13
13
 
14
+ def overloaded?
15
+ raise NotImplementedError
16
+ end
17
+
14
18
  def to_s
15
19
  "#<#{self.class}>"
16
20
  end
17
21
 
22
+ def ==(other)
23
+ self.class == other.class &&
24
+ self.key == other.key &&
25
+ self.capacity == other.capacity &&
26
+ self.args == other.args &&
27
+ self.options == other.options &&
28
+ self.redis.connection == other.redis.connection
29
+ end
30
+
31
+ def self.new(*)
32
+ # can only call via subclass
33
+ raise NotImplementedError if self == Berater::Limiter
34
+
35
+ super
36
+ end
37
+
18
38
  protected
19
39
 
20
- def initialize(key, **opts)
40
+ attr_reader :args
41
+
42
+ def initialize(key, capacity, *args, **opts)
21
43
  @key = key
44
+ self.capacity = capacity
45
+ @args = args
22
46
  @options = opts
23
47
  end
24
48
 
49
+ def capacity=(capacity)
50
+ unless capacity.is_a?(Integer) || capacity == Float::INFINITY
51
+ raise ArgumentError, "expected Integer, found #{capacity.class}"
52
+ end
53
+
54
+ raise ArgumentError, "capacity must be >= 0" unless capacity >= 0
55
+
56
+ @capacity = capacity
57
+ end
58
+
25
59
  def cache_key(key)
26
60
  "#{self.class}:#{key}"
27
61
  end
28
62
 
63
+ def yield_lock(lock, &block)
64
+ if block_given?
65
+ begin
66
+ yield lock
67
+ ensure
68
+ lock.release
69
+ end
70
+ else
71
+ lock
72
+ end
73
+ end
74
+
29
75
  end
30
76
  end
data/lib/berater/lock.rb CHANGED
@@ -13,24 +13,15 @@ module Berater
13
13
  end
14
14
 
15
15
  def locked?
16
- @released_at.nil? && !expired?
17
- end
18
-
19
- def expired?
20
- timeout ? @locked_at + timeout < Time.now : false
16
+ @released_at.nil?
21
17
  end
22
18
 
23
19
  def release
24
- raise 'lock expired' if expired?
25
20
  raise 'lock already released' unless locked?
26
21
 
27
22
  @released_at = Time.now
28
23
  @release_fn ? @release_fn.call : true
29
24
  end
30
25
 
31
- def timeout
32
- limiter.respond_to?(:timeout) ? limiter.timeout : nil
33
- end
34
-
35
26
  end
36
27
  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