berater 0.1.3 → 0.5.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.
checksums.yaml CHANGED
@@ -1,7 +1,7 @@
1
1
  ---
2
2
  SHA256:
3
- metadata.gz: 22b4ddfe5c7a48fbb3314743adf40d6c4c04e1c70198c210f34226065bd47f39
4
- data.tar.gz: d4e917aba67a6aebf1e19f1bba80fce705b8e8e7a9f2e4f07045f8fdac36044d
3
+ metadata.gz: fa5d51d6eea9d4691e58e8dd7065a1d596a57abdac1227fa19bd22cd2f42c15d
4
+ data.tar.gz: 55ac15bf6f2f295c525fdea57b3af9b518935a12dd4b644e3c0c9c085429ca50
5
5
  SHA512:
6
- metadata.gz: cd3724e47146912f01e01614341414447ea8cf86ebefbb75a2baf6079c46bc782a2c006df9ddfdf4b402d3f3ab7c55f2e02b9764e60e52994c0d6dd7eda9b804
7
- data.tar.gz: 9b9c77973b16f11399d87924b1c279cb4274c91cdffe8e89a1e1b5d8881e9cefe9f75bf91e86e0160cacf748c137633f33250892ffc047fabe65c6eb4dc9b3aa
6
+ metadata.gz: 52df5781f9c37c7f1c9dde09e579c8da1f2d62f405835b771c6087cda8af046018e85a88629f7c5d3a2be9aafd2680e348d1b2f91bc63fe456e6e51ec5d0fd02
7
+ data.tar.gz: '0928cb6f2e8e07810dec96f5af302ff9f62d519b9f19ef6f28acb5621569f13edb6d784cbd7084b738df8dfc232ace26e3ce0c6c4a7f2ac3e4391c97789bafae'
data/lib/berater.rb CHANGED
@@ -1,46 +1,40 @@
1
+ require 'berater/limiter'
2
+ require 'berater/lock'
3
+ require 'berater/lua_script'
4
+ require 'berater/utils'
1
5
  require 'berater/version'
2
6
 
3
-
4
7
  module Berater
5
8
  extend self
6
9
 
7
10
  class Overloaded < StandardError; end
8
11
 
9
- MODES = {}
10
-
11
- attr_accessor :redis, :mode
12
+ attr_accessor :redis
12
13
 
13
14
  def configure
14
- self.mode = :unlimited # default
15
-
16
15
  yield self
17
16
  end
18
17
 
19
- def new(mode, *args, **opts)
20
- klass = MODES[mode.to_sym]
21
-
22
- unless klass
23
- raise ArgumentError, "invalid mode: #{mode}"
24
- end
25
-
26
- klass.new(*args, **opts)
27
- end
28
-
29
- def register(mode, klass)
30
- MODES[mode.to_sym] = klass
18
+ def reset
19
+ @redis = nil
31
20
  end
32
21
 
33
- def mode=(mode)
34
- unless MODES.include? mode.to_sym
35
- raise ArgumentError, "invalid mode: #{mode}"
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
28
+ else
29
+ if interval
30
+ Berater::RateLimiter
31
+ else
32
+ Berater::ConcurrencyLimiter
33
+ end
34
+ end.yield_self do |klass|
35
+ args = [ key, capacity, interval ].compact
36
+ klass.new(*args, **opts)
36
37
  end
37
-
38
- @mode = mode.to_sym
39
- end
40
-
41
- def limit(*args, **opts, &block)
42
- mode = opts.delete(:mode) { self.mode }
43
- new(mode, *args, **opts).limit(&block)
44
38
  end
45
39
 
46
40
  def expunge
@@ -51,14 +45,18 @@ module Berater
51
45
 
52
46
  end
53
47
 
54
- # load and register limiters
55
- require 'berater/base_limiter'
48
+ # convenience method
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
56
+ end
57
+
58
+ # load limiters
56
59
  require 'berater/concurrency_limiter'
57
60
  require 'berater/inhibitor'
58
61
  require 'berater/rate_limiter'
59
62
  require 'berater/unlimiter'
60
-
61
- Berater.register(:concurrency, Berater::ConcurrencyLimiter)
62
- Berater.register(:inhibited, Berater::Inhibitor)
63
- Berater.register(:rate, Berater::RateLimiter)
64
- Berater.register(:unlimited, Berater::Unlimiter)
@@ -1,153 +1,106 @@
1
1
  module Berater
2
- class ConcurrencyLimiter < BaseLimiter
2
+ class ConcurrencyLimiter < Limiter
3
3
 
4
4
  class Incapacitated < Overloaded; end
5
5
 
6
- attr_reader :capacity, :timeout
6
+ attr_reader :timeout
7
7
 
8
- def initialize(capacity, **opts)
9
- super(**opts)
8
+ def initialize(key, capacity, **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
- 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
- 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
-
14
+ private def timeout=(timeout)
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
- class Lock
36
- attr_reader :limiter, :id, :contention
37
-
38
- def initialize(limiter, id, contention)
39
- @limiter = limiter
40
- @id = id
41
- @contention = contention
42
- @released = false
43
- @locked_at = Time.now
44
- end
45
-
46
- def release
47
- raise 'lock already released' if released?
48
- raise 'lock expired' if expired?
49
-
50
- @released = limiter.release(self)
51
- end
52
-
53
- def released?
54
- @released
55
- end
56
-
57
- def expired?
58
- @locked_at + limiter.timeout < Time.now
59
- end
60
- end
61
-
62
- LUA_SCRIPT = <<~LUA.gsub(/^\s*(--.*\n)?/, '')
20
+ LUA_SCRIPT = Berater::LuaScript(<<~LUA
63
21
  local key = KEYS[1]
22
+ local lock_key = KEYS[2]
64
23
  local capacity = tonumber(ARGV[1])
65
24
  local ts = tonumber(ARGV[2])
66
25
  local ttl = tonumber(ARGV[3])
26
+ local cost = tonumber(ARGV[4])
27
+ local lock_ids = {}
67
28
 
68
- local exists
69
- local count
70
- local lock
71
-
72
- -- check to see if key already exists
73
- if ttl == 0 then
74
- exists = redis.call('EXISTS', key)
75
- else
76
- -- and refresh TTL while we're at it
77
- exists = redis.call('EXPIRE', key, ttl * 2)
29
+ -- purge stale hosts
30
+ if ttl > 0 then
31
+ redis.call('ZREMRANGEBYSCORE', key, '-inf', ts - ttl)
78
32
  end
79
33
 
80
- if exists == 1 then
81
- -- purge stale hosts
82
- if ttl > 0 then
83
- redis.call('ZREMRANGEBYSCORE', key, '-inf', ts - ttl)
84
- end
85
-
86
- -- check capacity (subtract one for next lock entry)
87
- count = redis.call('ZCARD', key) - 1
34
+ -- check capacity
35
+ local count = redis.call('ZCARD', key)
88
36
 
37
+ if cost == 0 then
38
+ -- just check limit, ie. for .overlimit?
89
39
  if count < capacity then
90
- -- yay, grab a lock
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 = {}
91
46
 
92
- -- regenerate next lock entry, which has score inf
93
- lock = unpack(redis.call('ZPOPMAX', key))
94
- redis.call('ZADD', key, 'inf', (lock + 1) % 2^52)
47
+ for i = lock_id - cost + 1, lock_id do
48
+ table.insert(lock_ids, i)
95
49
 
96
- count = count + 1
50
+ table.insert(locks, ts)
51
+ table.insert(locks, i)
97
52
  end
98
- else
99
- count = 1
100
- lock = "1"
101
53
 
102
- -- create structure to track locks and next id
103
- redis.call('ZADD', key, 'inf', lock + 1)
54
+ redis.call('ZADD', key, unpack(locks))
55
+ count = count + cost
104
56
 
105
57
  if ttl > 0 then
106
- redis.call('EXPIRE', key, ttl * 2)
58
+ redis.call('PEXPIRE', key, ttl)
107
59
  end
108
60
  end
109
61
 
110
- if lock then
111
- -- store lock and timestamp
112
- redis.call('ZADD', key, ts, lock)
113
- end
114
-
115
- return { count, lock }
62
+ return { count, unpack(lock_ids) }
116
63
  LUA
64
+ )
117
65
 
118
- def limit(**opts, &block)
119
- unless opts.empty?
120
- return self.class.new(
121
- capacity,
122
- **options.merge(opts)
123
- # **options.merge(timeout: timeout).merge(opts)
124
- ).limit(&block)
125
- end
66
+ def limit(capacity: nil, cost: 1, &block)
67
+ capacity ||= @capacity
68
+ # cost is Integer >= 0
126
69
 
127
- count, lock_id = redis.eval(
128
- LUA_SCRIPT,
129
- [ key ],
130
- [ capacity, Time.now.to_i, timeout ]
131
- )
70
+ # timestamp in microseconds
71
+ ts = (Time.now.to_f * 10**6).to_i
132
72
 
133
- raise Incapacitated unless lock_id
73
+ count, *lock_ids = LUA_SCRIPT.eval(
74
+ redis,
75
+ [ cache_key(key), cache_key('lock_id') ],
76
+ [ capacity, ts, @timeout_usec, cost ]
77
+ )
134
78
 
135
- lock = Lock.new(self, lock_id, count)
79
+ raise Incapacitated if lock_ids.empty?
136
80
 
137
- if block_given?
138
- begin
139
- yield lock
140
- ensure
141
- release(lock)
142
- end
81
+ if cost == 0
82
+ lock = Lock.new(self, nil, count)
143
83
  else
144
- lock
84
+ lock = Lock.new(self, lock_ids[0], count, -> { release(lock_ids) })
145
85
  end
86
+
87
+ yield_lock(lock, &block)
88
+ end
89
+
90
+ def overloaded?
91
+ limit(cost: 0) { false }
92
+ rescue Overloaded
93
+ true
94
+ end
95
+ alias incapacitated? overloaded?
96
+
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
146
100
  end
147
101
 
148
- def release(lock)
149
- res = redis.zrem(key, lock.id)
150
- res == true || res == 1 # depending on which version of Redis
102
+ def to_s
103
+ "#<#{self.class}(#{key}: #{capacity} at a time)>"
151
104
  end
152
105
 
153
106
  end
@@ -0,0 +1,68 @@
1
+ module Berater
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
+
21
+ extend self
22
+
23
+ def eval &block
24
+ @keywords ||= Class.new do
25
+ # create a class where DSL keywords are methods
26
+ KEYWORDS.each do |keyword|
27
+ define_singleton_method(keyword) { keyword }
28
+ end
29
+ end
30
+
31
+ install
32
+ @keywords.class_eval &block
33
+ ensure
34
+ uninstall
35
+ end
36
+
37
+ private
38
+
39
+ KEYWORDS = [
40
+ :second, :minute, :hour,
41
+ :unlimited, :inhibited,
42
+ ].freeze
43
+
44
+ def install
45
+ Integer.class_eval do
46
+ def per(unit)
47
+ [ self, unit ]
48
+ end
49
+ alias every per
50
+
51
+ def at_once
52
+ [ self ]
53
+ end
54
+ alias concurrently at_once
55
+ alias at_a_time at_once
56
+ end
57
+ end
58
+
59
+ def uninstall
60
+ Integer.remove_method :per
61
+ Integer.remove_method :every
62
+
63
+ Integer.remove_method :at_once
64
+ Integer.remove_method :concurrently
65
+ Integer.remove_method :at_a_time
66
+ end
67
+ end
68
+ end
@@ -1,21 +1,20 @@
1
1
  module Berater
2
- class Inhibitor < BaseLimiter
2
+ class Inhibitor < Limiter
3
3
 
4
4
  class Inhibited < Overloaded; end
5
5
 
6
- def initialize(*args, **opts)
7
- super(**opts)
6
+ def initialize(key = :inhibitor, *args, **opts)
7
+ super(key, 0, **opts)
8
8
  end
9
9
 
10
- def limit(**opts, &block)
11
- unless opts.empty?
12
- return self.class.new(
13
- **options.merge(opts)
14
- ).limit(&block)
15
- end
16
-
10
+ def limit(**opts)
17
11
  raise Inhibited
18
12
  end
19
13
 
14
+ def overloaded?
15
+ true
16
+ end
17
+ alias inhibited? overloaded?
18
+
20
19
  end
21
20
  end
@@ -0,0 +1,76 @@
1
+ module Berater
2
+ class Limiter
3
+
4
+ attr_reader :key, :capacity, :options
5
+
6
+ def redis
7
+ options[:redis] || Berater.redis
8
+ end
9
+
10
+ def limit
11
+ raise NotImplementedError
12
+ end
13
+
14
+ def overloaded?
15
+ raise NotImplementedError
16
+ end
17
+
18
+ def to_s
19
+ "#<#{self.class}>"
20
+ end
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
+
38
+ protected
39
+
40
+ attr_reader :args
41
+
42
+ def initialize(key, capacity, *args, **opts)
43
+ @key = key
44
+ self.capacity = capacity
45
+ @args = args
46
+ @options = opts
47
+ end
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
+
59
+ def cache_key(key)
60
+ "#{self.class}:#{key}"
61
+ end
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
+
75
+ end
76
+ end