berater 0.1.1 → 0.3.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: 3452315467fe2c079071179703283dd9ea92eb8e6abcc959905f40e16f3e2414
4
- data.tar.gz: 90d783190f15f57c3ac44dc2afc89fb5a000a392163ab9a8351ff498e323901a
3
+ metadata.gz: d007dbb664fa57c921047a28dd1e5a591cf5c51fba4931dbc367ab552e4b8a8d
4
+ data.tar.gz: c710d31a2d8cfdc50af4c2fae87cabad177096d13da52a7f9c48b61aa31b6038
5
5
  SHA512:
6
- metadata.gz: 228f67adfb24027a6ea3c7b090c64aecee1a38e64ed00f7272b0c1af273dbcc73bbec8112fbb15595133cd450b45ac10f603e30ad252218d41e66d5206cc549d
7
- data.tar.gz: 9068fd600128bb11f718318fa41f1030d17cd5ad36b9ce5884cbdc5f105fe56f92da73568dc60f12f6e1c7a64402652381259a64524c80889465d4c794a91e69
6
+ metadata.gz: 0d9c211bb65e7341b98e637d3f0ce0e31f0a879437c1c309fb01e1efab5e5c6b2d5044158931ab16a67cbbdebd00ddb5f10afc97430230b153d8cc236f7517ab
7
+ data.tar.gz: 5d8b38bfd7683df641fd6891cd1a493079453ea6c61b2880581eeb2d963fb171365fb3448443e2548b34297d8dcfa352e67cff5688cd24d25422efb5e55d4afd
data/lib/berater.rb CHANGED
@@ -1,53 +1,50 @@
1
- require 'berater/base_limiter'
2
- require 'berater/concurrency_limiter'
3
- require 'berater/inhibitor'
4
- require 'berater/rate_limiter'
5
- require 'berater/unlimiter'
6
1
  require 'berater/version'
2
+ require 'berater/lock'
7
3
 
8
4
 
9
5
  module Berater
10
6
  extend self
11
7
 
12
- Overloaded = BaseLimiter::Overloaded
8
+ class Overloaded < StandardError; end
13
9
 
14
10
  MODES = {}
15
11
 
16
- attr_accessor :redis, :mode
12
+ attr_accessor :redis
17
13
 
18
14
  def configure
19
- self.mode = :unlimited # default
20
-
21
15
  yield self
22
16
  end
23
17
 
24
- def new(mode, *args, **opts)
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
27
+
28
+ mode, *args = DSL.eval(&block)
29
+ else
30
+ if block_given?
31
+ raise ArgumentError, 'expected either mode or block, not both'
32
+ end
33
+ end
34
+
25
35
  klass = MODES[mode.to_sym]
26
36
 
27
37
  unless klass
28
38
  raise ArgumentError, "invalid mode: #{mode}"
29
39
  end
30
40
 
31
- klass.new(*args, **opts)
41
+ klass.new(key, *args, **opts)
32
42
  end
33
43
 
34
44
  def register(mode, klass)
35
45
  MODES[mode.to_sym] = klass
36
46
  end
37
47
 
38
- def mode=(mode)
39
- unless MODES.include? mode.to_sym
40
- raise ArgumentError, "invalid mode: #{mode}"
41
- end
42
-
43
- @mode = mode.to_sym
44
- end
45
-
46
- def limit(*args, **opts, &block)
47
- mode = opts.delete(:mode) { self.mode }
48
- new(mode, *args, **opts).limit(&block)
49
- end
50
-
51
48
  def expunge
52
49
  redis.scan_each(match: "#{self.name}*") do |key|
53
50
  redis.del key
@@ -56,7 +53,21 @@ module Berater
56
53
 
57
54
  end
58
55
 
56
+ # convenience method
57
+ def Berater(key, mode, *args, **opts, &block)
58
+ Berater.new(key, mode, *args, **opts).limit(&block)
59
+ end
60
+
61
+ # load limiters
62
+ require 'berater/limiter'
63
+ require 'berater/concurrency_limiter'
64
+ require 'berater/inhibitor'
65
+ require 'berater/rate_limiter'
66
+ require 'berater/unlimiter'
67
+
59
68
  Berater.register(:concurrency, Berater::ConcurrencyLimiter)
60
69
  Berater.register(:inhibited, Berater::Inhibitor)
61
70
  Berater.register(:rate, Berater::RateLimiter)
62
71
  Berater.register(:unlimited, Berater::Unlimiter)
72
+
73
+ require 'berater/dsl'
@@ -1,18 +1,18 @@
1
1
  module Berater
2
- class ConcurrencyLimiter < BaseLimiter
2
+ class ConcurrencyLimiter < Limiter
3
3
 
4
4
  class Incapacitated < Overloaded; end
5
5
 
6
6
  attr_reader :capacity, :timeout
7
7
 
8
- def initialize(capacity, **opts)
9
- super(**opts)
8
+ def initialize(key, capacity, **opts)
9
+ super(key, **opts)
10
10
 
11
11
  self.capacity = capacity
12
12
  self.timeout = opts[:timeout] || 0
13
13
  end
14
14
 
15
- def capacity=(capacity)
15
+ private def capacity=(capacity)
16
16
  unless capacity.is_a? Integer
17
17
  raise ArgumentError, "expected Integer, found #{capacity.class}"
18
18
  end
@@ -22,7 +22,7 @@ module Berater
22
22
  @capacity = capacity
23
23
  end
24
24
 
25
- def timeout=(timeout)
25
+ private def timeout=(timeout)
26
26
  unless timeout.is_a? Integer
27
27
  raise ArgumentError, "expected Integer, found #{timeout.class}"
28
28
  end
@@ -32,119 +32,62 @@ module Berater
32
32
  @timeout = timeout
33
33
  end
34
34
 
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
35
  LUA_SCRIPT = <<~LUA.gsub(/^\s*(--.*\n)?/, '')
63
36
  local key = KEYS[1]
37
+ local lock_key = KEYS[2]
64
38
  local capacity = tonumber(ARGV[1])
65
- local ttl = tonumber(ARGV[2])
66
-
67
- local exists
68
- local count
39
+ local ts = tonumber(ARGV[2])
40
+ local ttl = tonumber(ARGV[3])
69
41
  local lock
70
- local ts = unpack(redis.call('TIME'))
71
42
 
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)
43
+ -- purge stale hosts
44
+ if ttl > 0 then
45
+ redis.call('ZREMRANGEBYSCORE', key, '-inf', ts - ttl)
78
46
  end
79
47
 
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
88
-
89
- if count < capacity then
90
- -- yay, grab a lock
91
-
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)
95
-
96
- count = count + 1
97
- end
98
- else
99
- count = 1
100
- lock = "1"
101
-
102
- -- create structure to track locks and next id
103
- redis.call('ZADD', key, 'inf', lock + 1)
48
+ -- check capacity
49
+ local count = redis.call('ZCARD', key)
104
50
 
105
- if ttl > 0 then
106
- redis.call('EXPIRE', key, ttl * 2)
107
- end
108
- end
109
-
110
- if lock then
111
- -- store lock and timestamp
51
+ if count < capacity then
52
+ -- grab a lock
53
+ lock = redis.call('INCR', lock_key)
112
54
  redis.call('ZADD', key, ts, lock)
55
+ count = count + 1
113
56
  end
114
57
 
115
58
  return { count, lock }
116
59
  LUA
117
60
 
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
126
-
127
- count, lock_id = redis.eval(LUA_SCRIPT, [ key ], [ capacity, timeout ])
61
+ def limit
62
+ count, lock_id = redis.eval(
63
+ LUA_SCRIPT,
64
+ [ cache_key(key), cache_key('lock_id') ],
65
+ [ capacity, Time.now.to_i, timeout ]
66
+ )
128
67
 
129
68
  raise Incapacitated unless lock_id
130
69
 
131
- lock = Lock.new(self, lock_id, count)
70
+ lock = Lock.new(self, lock_id, count, -> { release(lock_id) })
132
71
 
133
72
  if block_given?
134
73
  begin
135
74
  yield lock
136
75
  ensure
137
- release(lock)
76
+ lock.release
138
77
  end
139
78
  else
140
79
  lock
141
80
  end
142
81
  end
143
82
 
144
- def release(lock)
145
- res = redis.zrem(key, lock.id)
83
+ private def release(lock_id)
84
+ res = redis.zrem(cache_key(key), lock_id)
146
85
  res == true || res == 1 # depending on which version of Redis
147
86
  end
148
87
 
88
+ def to_s
89
+ "#<#{self.class}(#{key}: #{capacity} at a time)>"
90
+ end
91
+
149
92
  end
150
93
  end
@@ -0,0 +1,57 @@
1
+ module Berater
2
+ module DSL
3
+ extend self
4
+
5
+ def eval &block
6
+ @keywords ||= Class.new do
7
+ # create a class where DSL keywords are methods
8
+ KEYWORDS.each do |keyword|
9
+ define_singleton_method(keyword) { keyword }
10
+ end
11
+ end
12
+
13
+ install
14
+ @keywords.class_eval &block
15
+ ensure
16
+ uninstall
17
+ end
18
+
19
+ private
20
+
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
+ KEYWORDS = [
29
+ :second, :minute, :hour,
30
+ :unlimited, :inhibited,
31
+ ].freeze
32
+
33
+ def install
34
+ Integer.class_eval do
35
+ def per(unit)
36
+ [ :rate, self, unit ]
37
+ end
38
+ alias every per
39
+
40
+ def at_once
41
+ [ :concurrency, self ]
42
+ end
43
+ alias concurrently at_once
44
+ alias at_a_time at_once
45
+ end
46
+ end
47
+
48
+ def uninstall
49
+ Integer.remove_method :per
50
+ Integer.remove_method :every
51
+
52
+ Integer.remove_method :at_once
53
+ Integer.remove_method :concurrently
54
+ Integer.remove_method :at_a_time
55
+ end
56
+ end
57
+ end
@@ -1,19 +1,13 @@
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, **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
17
11
  raise Inhibited
18
12
  end
19
13
 
@@ -0,0 +1,30 @@
1
+ module Berater
2
+ class Limiter
3
+
4
+ attr_reader :key, :options
5
+
6
+ def redis
7
+ options[:redis] || Berater.redis
8
+ end
9
+
10
+ def limit
11
+ raise NotImplementedError
12
+ end
13
+
14
+ def to_s
15
+ "#<#{self.class}>"
16
+ end
17
+
18
+ protected
19
+
20
+ def initialize(key, **opts)
21
+ @key = key
22
+ @options = opts
23
+ end
24
+
25
+ def cache_key(key)
26
+ "#{self.class}:#{key}"
27
+ end
28
+
29
+ end
30
+ end
@@ -0,0 +1,36 @@
1
+ module Berater
2
+ class Lock
3
+
4
+ attr_reader :limiter, :id, :contention
5
+
6
+ def initialize(limiter, id, contention, release_fn = nil)
7
+ @limiter = limiter
8
+ @id = id
9
+ @contention = contention
10
+ @locked_at = Time.now
11
+ @release_fn = release_fn
12
+ @released_at = nil
13
+ end
14
+
15
+ def locked?
16
+ @released_at.nil? && !expired?
17
+ end
18
+
19
+ def expired?
20
+ timeout ? @locked_at + timeout < Time.now : false
21
+ end
22
+
23
+ def release
24
+ raise 'lock expired' if expired?
25
+ raise 'lock already released' unless locked?
26
+
27
+ @released_at = Time.now
28
+ @release_fn ? @release_fn.call : true
29
+ end
30
+
31
+ def timeout
32
+ limiter.respond_to?(:timeout) ? limiter.timeout : nil
33
+ end
34
+
35
+ end
36
+ end
@@ -1,18 +1,18 @@
1
1
  module Berater
2
- class RateLimiter < BaseLimiter
2
+ class RateLimiter < Limiter
3
3
 
4
4
  class Overrated < Overloaded; end
5
5
 
6
6
  attr_accessor :count, :interval
7
7
 
8
- def initialize(count, interval, **opts)
9
- super(**opts)
8
+ def initialize(key, count, interval, **opts)
9
+ super(key, **opts)
10
10
 
11
11
  self.count = count
12
12
  self.interval = interval
13
13
  end
14
14
 
15
- def count=(count)
15
+ private def count=(count)
16
16
  unless count.is_a? Integer
17
17
  raise ArgumentError, "expected Integer, found #{count.class}"
18
18
  end
@@ -22,12 +22,13 @@ module Berater
22
22
  @count = count
23
23
  end
24
24
 
25
- def interval=(interval)
25
+ private def interval=(interval)
26
26
  @interval = interval.dup
27
27
 
28
28
  case @interval
29
29
  when Integer
30
30
  raise ArgumentError, "interval must be >= 0" unless @interval >= 0
31
+ @interval_sec = @interval
31
32
  when String
32
33
  @interval = @interval.to_sym
33
34
  when Symbol
@@ -38,45 +39,58 @@ module Berater
38
39
  if @interval.is_a? Symbol
39
40
  case @interval
40
41
  when :sec, :second, :seconds
41
- @interval = 1
42
+ @interval = :second
43
+ @interval_sec = 1
42
44
  when :min, :minute, :minutes
43
- @interval = 60
45
+ @interval = :minute
46
+ @interval_sec = 60
44
47
  when :hour, :hours
45
- @interval = 60 * 60
48
+ @interval = :hour
49
+ @interval_sec = 60 * 60
46
50
  else
47
51
  raise ArgumentError, "unexpected interval value: #{interval}"
48
52
  end
49
53
  end
50
-
51
- @interval
52
54
  end
53
55
 
54
- def limit(**opts, &block)
55
- unless opts.empty?
56
- return self.class.new(
57
- count,
58
- interval,
59
- options.merge(opts)
60
- ).limit(&block)
61
- end
62
-
56
+ def limit
63
57
  ts = Time.now.to_i
64
58
 
65
59
  # bucket into time slot
66
- rkey = "%s:%d" % [ key, ts - ts % @interval ]
60
+ rkey = "%s:%d" % [ cache_key(key), ts - ts % @interval_sec ]
67
61
 
68
62
  count, _ = redis.multi do
69
63
  redis.incr rkey
70
- redis.expire rkey, @interval * 2
64
+ redis.expire rkey, @interval_sec * 2
71
65
  end
72
66
 
73
67
  raise Overrated if count > @count
74
68
 
69
+ lock = Lock.new(self, count, count)
70
+
75
71
  if block_given?
76
- yield
72
+ begin
73
+ yield lock
74
+ ensure
75
+ lock.release
76
+ end
77
+ else
78
+ lock
79
+ end
80
+ end
81
+
82
+ def to_s
83
+ msg = if @interval.is_a? Integer
84
+ if @interval == 1
85
+ "every second"
86
+ else
87
+ "every #{@interval} seconds"
88
+ end
77
89
  else
78
- count
90
+ "per #{@interval}"
79
91
  end
92
+
93
+ "#<#{self.class}(#{key}: #{count} #{msg})>"
80
94
  end
81
95
 
82
96
  end