traffic_jam 1.0.1 → 1.1.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
  SHA1:
3
- metadata.gz: dea90b9ea6d6cdd3f95e0c80541bf26fce1b3690
4
- data.tar.gz: 6870c79d37669b7997aea051b45bffe041a176b8
3
+ metadata.gz: dc24485cfcc470e68b80dbb53f6b135ab9879bf0
4
+ data.tar.gz: ab8ea5c0a5e1075af9606fae0e4a0cc9bb1fc071
5
5
  SHA512:
6
- metadata.gz: 79106d2df976e9826ecc5a1b8a7ae3882a1091a79c82f2193f0ea4b9a79bd83b0fa3c35e07e822a77609e25c7b2cb29505c204baae5b5248d5bd24bdc54e669d
7
- data.tar.gz: 938bf7f0fcb98d9677015495399078358f50defc35962203f9de833d0cfea646ea0f86a10ac48ef96b6a28f70ba8d486d99f0c84175c662f1f4946d0a4026fc8
6
+ metadata.gz: f428099b9b945e9196a4e54146a79af29d8e8e8c430320f1e9237be2e96df3b9f39e29cad2b78c53694f9c76e85ee698f38af91b23ab715339ab9f578eaa4434
7
+ data.tar.gz: 7c9afc9302a854fa68c042c9329e73a6224858bec71e25fb0ff5ac2a72fdcd6ab7f5aeedba4af936242a20e5fb96bc4fc87dcacdf9a77b3291f582adb5ddd540
@@ -10,5 +10,8 @@ module TrafficJam
10
10
  @limit = limit
11
11
  end
12
12
  end
13
+
14
+ class InvalidKeyError < StandardError; end
15
+ class UnknownReturnValue < StandardError; end
13
16
  end
14
17
  end
@@ -0,0 +1,53 @@
1
+ module TrafficJam
2
+ # This class represents a lifetime limit on an action, value pair. For example, if
3
+ # limiting the amount of money a user can transfer, the action could be
4
+ # +:transfers+ and the value would be the user ID. The class exposes atomic
5
+ # increment operations and allows querying of the current amount used and
6
+ # amount remaining.
7
+ class LifetimeLimit < Limit
8
+ # Constructor takes an action name as a symbol, a maximum cap, and the
9
+ # period of limit. +max+ and +period+ are required keyword arguments.
10
+ #
11
+ # @param action [Symbol] action name
12
+ # @param value [String] limit target value
13
+ # @param max [Integer] required limit maximum
14
+ # @raise [ArgumentError] if max is nil
15
+ def initialize(action, value, max: nil)
16
+ super(action, value, max: max, period: -1)
17
+ end
18
+
19
+ # Increment the amount used by the given number. Does not perform increment
20
+ # if the operation would exceed the limit. Returns whether the operation was
21
+ # successful.
22
+ #
23
+ # @param amount [Integer] amount to increment by
24
+ # @return [Boolean] true if increment succeded and false if incrementing
25
+ # would exceed the limit
26
+ def increment(amount = 1, time: Time.now)
27
+ raise ArgumentError, 'Amount must be an integer' if amount != amount.to_i
28
+ return amount <= 0 if max.zero?
29
+
30
+ !!run_script([amount.to_i, max])
31
+ end
32
+
33
+ # Return amount of limit used
34
+ #
35
+ # @return [Integer] amount used
36
+ def used
37
+ return 0 if max.zero?
38
+ amount = redis.get(key) || 0
39
+ [amount.to_i, max].min
40
+ end
41
+
42
+ private
43
+
44
+ def run_script(argv)
45
+ redis.evalsha(
46
+ Scripts::INCRBY_HASH, keys: [key], argv: argv
47
+ )
48
+ rescue Redis::CommandError => error
49
+ raise error if /ERR Error running script/ =~ error.message
50
+ redis.eval(Scripts::INCRBY, keys: [key], argv: argv)
51
+ end
52
+ end
53
+ end
@@ -33,8 +33,8 @@ module TrafficJam
33
33
  # @param period [Integer] required limit period in seconds
34
34
  # @raise [ArgumentError] if max or period is nil
35
35
  def initialize(action, value, max: nil, period: nil)
36
- raise ArgumentError('Max is required') if max.nil?
37
- raise ArgumentError('Period is required') if period.nil?
36
+ raise ArgumentError.new('Max is required') if max.nil?
37
+ raise ArgumentError.new('Period is required') if period.nil?
38
38
  @action, @value, @max, @period = action, value, max, period
39
39
  end
40
40
 
@@ -71,14 +71,14 @@ module TrafficJam
71
71
  raise ArgumentError.new("Amount must be an integer")
72
72
  end
73
73
 
74
- timestamp = (time.to_f * 1000).round
74
+ timestamp = (time.to_f * 1000).to_i
75
75
  argv = [timestamp, amount.to_i, max, period * 1000]
76
76
 
77
77
  result =
78
78
  begin
79
79
  redis.evalsha(
80
80
  Scripts::INCREMENT_SCRIPT_HASH, keys: [key], argv: argv)
81
- rescue Redis::CommandError => e
81
+ rescue Redis::CommandError
82
82
  redis.eval(Scripts::INCREMENT_SCRIPT, keys: [key], argv: argv)
83
83
  end
84
84
 
@@ -127,9 +127,7 @@ module TrafficJam
127
127
  def used
128
128
  return 0 if max.zero?
129
129
 
130
- obj = redis.hgetall(key)
131
- timestamp = obj['timestamp']
132
- amount = obj['amount']
130
+ timestamp, amount = redis.hmget(key, 'timestamp', 'amount')
133
131
  if timestamp && amount
134
132
  time_passed = Time.now.to_f - timestamp.to_i / 1000.0
135
133
  drift = max * time_passed / period
@@ -161,7 +159,7 @@ module TrafficJam
161
159
  end
162
160
 
163
161
  def key
164
- if @key.nil?
162
+ if !defined?(@key) || @key.nil?
165
163
  converted_value =
166
164
  begin
167
165
  value.to_rate_limit_value
@@ -170,9 +168,13 @@ module TrafficJam
170
168
  end
171
169
  hash = Digest::MD5.base64digest(converted_value.to_s)
172
170
  hash = hash[0...config.hash_length]
173
- @key = "#{config.key_prefix}:#{action}:#{hash}"
171
+ @key = "#{key_prefix}:#{action}:#{hash}"
174
172
  end
175
173
  @key
176
174
  end
175
+
176
+ def key_prefix
177
+ config.key_prefix
178
+ end
177
179
  end
178
180
  end
@@ -0,0 +1,82 @@
1
+ require_relative 'scripts'
2
+
3
+ module TrafficJam
4
+ # This class represents a rolling limit on an action, value pair. For example,
5
+ # if limiting the amount of money a user can transfer in a week, the action
6
+ # could be +:transfers+ and the value would be the user ID. The class exposes
7
+ # atomic increment operations and allows querying of the current amount used
8
+ # and amount remaining.
9
+ #
10
+ # This class also handles 0 for period, where 0 is no period (each
11
+ # request is compared to the max).
12
+ #
13
+ # This class departs from the design of Limit by tracking a sum of the actions
14
+ # in a second, in a hash keyed by the timestamp. Therefore, this limit can put
15
+ # a lot of data size pressure on the Redis storage, so use it wisely.
16
+ class RollingLimit < Limit
17
+ # Constructor takes an action name as a symbol, a maximum cap, and the
18
+ # period of limit. +max+ and +period+ are required keyword arguments.
19
+ #
20
+ # @param action [Symbol] action name
21
+ # @param value [String] limit target value
22
+ # @param max [Integer] required limit maximum
23
+ # @param period [Integer] required limit period in seconds
24
+ # @raise [ArgumentError] if max or period is nil
25
+ def initialize(action, value, max: nil, period: nil)
26
+ super(action, value, max: max, period: period)
27
+ end
28
+
29
+ # Increment the amount used by the given number. Rolls back the increment
30
+ # if the operation exceeds the limit. Returns whether the operation was
31
+ # successful. Time of increment can be specified optionally with a keyword
32
+ # argument, which is not really useful since it be undone by used.
33
+ #
34
+ # @param amount [Integer] amount to increment by
35
+ # @param time [Time] time when increment occurs (ignored)
36
+ # @return [Boolean] true if increment succeded and false if incrementing
37
+ # would exceed the limit
38
+ def increment(amount = 1, time: Time.now)
39
+ raise ArgumentError, 'Amount must be an integer' if amount != amount.to_i
40
+ return amount <= 0 if max.zero?
41
+ return amount <= max if period.zero?
42
+ return true if amount.zero?
43
+ return false if amount > max
44
+
45
+ !run_incr([time.to_i, amount.to_i, max, period]).nil?
46
+ end
47
+
48
+ # Return amount of limit used
49
+ #
50
+ # @return [Integer] amount used
51
+ def used
52
+ return 0 if max.zero? || period.zero?
53
+ [sum, max].min
54
+ end
55
+
56
+ private
57
+
58
+ def sum
59
+ run_sum([Time.now.to_i, period])
60
+ end
61
+
62
+ def clear_before
63
+ Time.now.to_i - period
64
+ end
65
+
66
+ def run_sum(argv)
67
+ redis.evalsha(Scripts::SUM_ROLLING_HASH, keys: [key], argv: argv)
68
+ rescue Redis::CommandError => error
69
+ raise error if /ERR Error running script/ =~ error.message
70
+ redis.eval(Scripts::SUM_ROLLING, keys: [key], argv: argv)
71
+ end
72
+
73
+ def run_incr(argv)
74
+ redis.evalsha(
75
+ Scripts::INCREMENT_ROLLING_HASH, keys: [key], argv: argv
76
+ )
77
+ rescue Redis::CommandError => error
78
+ raise error if /ERR Error running script/ =~ error.message
79
+ redis.eval(Scripts::INCREMENT_ROLLING, keys: [key], argv: argv)
80
+ end
81
+ end
82
+ end
@@ -10,5 +10,13 @@ module TrafficJam
10
10
 
11
11
  INCREMENT_SCRIPT = load('increment')
12
12
  INCREMENT_SCRIPT_HASH = Digest::SHA1.hexdigest(INCREMENT_SCRIPT)
13
+ INCREMENT_SIMPLE = load('increment_simple')
14
+ INCREMENT_SIMPLE_HASH = Digest::SHA1.hexdigest(INCREMENT_SIMPLE)
15
+ INCREMENT_ROLLING = load('increment_rolling')
16
+ INCREMENT_ROLLING_HASH = Digest::SHA1.hexdigest(INCREMENT_ROLLING)
17
+ INCRBY = load('incrby')
18
+ INCRBY_HASH = Digest::SHA1.hexdigest(INCRBY)
19
+ SUM_ROLLING = load('sum_rolling')
20
+ SUM_ROLLING_HASH = Digest::SHA1.hexdigest(SUM_ROLLING)
13
21
  end
14
22
  end
@@ -0,0 +1,94 @@
1
+ require_relative 'limit'
2
+ require_relative 'scripts'
3
+
4
+ module TrafficJam
5
+ # A SimpleLimit is a limit type that is more efficient for increments but does
6
+ # not support decrements or changing the max value without a complete reset.
7
+ # This means that if the period or max value for an action, value key changes,
8
+ # the used and remaining values cannot be preserved.
9
+ #
10
+ # This works by storing a key in Redis with a millisecond-precision expiry
11
+ # representing the time that the limit will be completely reset. Each
12
+ # increment operation converts the increment amount into the number of
13
+ # milliseconds to be added to the expiry.
14
+ #
15
+ # Example: Limit is 5 per 10 seconds.
16
+ # An increment by 1 first sets the key to expire in 2s.
17
+ # Another immediate increment by 4 sets the expiry to 10s.
18
+ # Subsequent increments fail until clock time catches up to expiry
19
+ class SimpleLimit < Limit
20
+ # Increment the amount used by the given number. Does not perform increment
21
+ # if the operation would exceed the limit. Returns whether the operation was
22
+ # successful.
23
+ #
24
+ # @param amount [Integer] amount to increment by
25
+ # @param time [Time] time is ignored
26
+ # @return [Boolean] true if increment succeded and false if incrementing
27
+ # would exceed the limit
28
+ def increment(amount = 1, time: Time.now)
29
+ return true if amount == 0
30
+ return false if max == 0
31
+ raise ArgumentError.new("Amount must be positive") if amount < 0
32
+
33
+ if amount != amount.to_i
34
+ raise ArgumentError.new("Amount must be an integer")
35
+ end
36
+
37
+ return false if amount > max
38
+
39
+ incrby = (period * 1000 * amount / max).to_i
40
+ argv = [incrby, period * 1000]
41
+
42
+ result =
43
+ begin
44
+ redis.evalsha(
45
+ Scripts::INCREMENT_SIMPLE_HASH, keys: [key], argv: argv)
46
+ rescue Redis::CommandError
47
+ redis.eval(Scripts::INCREMENT_SIMPLE, keys: [key], argv: argv)
48
+ end
49
+
50
+ case result
51
+ when 0
52
+ return true
53
+ when -1
54
+ raise Errors::InvalidKeyError, "Redis key #{key} has no expire time set"
55
+ when -2
56
+ return false
57
+ else
58
+ raise Errors::UnknownReturnValue,
59
+ "Received unexpected return value #{result} from " \
60
+ "increment_simple eval"
61
+ end
62
+ end
63
+
64
+ # Decrement the amount used by the given number.
65
+ #
66
+ # @param amount [Integer] amount to decrement by
67
+ # @param time [Time] time is ignored
68
+ # @raise [NotImplementedError] decrement is not defined for SimpleLimit
69
+ def decrement(_amount = 1, time: Time.now)
70
+ raise NotImplementedError, "decrement is not defined for SimpleLimit"
71
+ end
72
+
73
+ # Return amount of limit used, taking time drift into account.
74
+ #
75
+ # @return [Integer] amount used
76
+ def used
77
+ return 0 if max.zero?
78
+
79
+ expiry = redis.pttl(key)
80
+ case expiry
81
+ when -1 # key exists but has no associated expire
82
+ raise Errors::InvalidKeyError, "Redis key #{key} has no expire time set"
83
+ when -2 # key does not exist
84
+ return 0
85
+ end
86
+
87
+ (max * expiry / (period * 1000.0)).ceil
88
+ end
89
+
90
+ def key_prefix
91
+ "#{config.key_prefix}:s"
92
+ end
93
+ end
94
+ end
data/lib/traffic_jam.rb CHANGED
@@ -4,7 +4,9 @@ require_relative 'traffic_jam/errors'
4
4
  require_relative 'traffic_jam/configuration'
5
5
  require_relative 'traffic_jam/limit'
6
6
  require_relative 'traffic_jam/limit_group'
7
-
7
+ require_relative 'traffic_jam/simple_limit'
8
+ require_relative 'traffic_jam/rolling_limit'
9
+ require_relative 'traffic_jam/lifetime_limit'
8
10
 
9
11
  module TrafficJam
10
12
  include Errors
@@ -0,0 +1,18 @@
1
+ local arg_amount = tonumber(ARGV[1])
2
+ local arg_max = tonumber(ARGV[2])
3
+
4
+ local old_amount = tonumber(redis.call("GET", KEYS[1]))
5
+ local new_amount
6
+
7
+ if not old_amount then
8
+ new_amount = arg_amount
9
+ else
10
+ new_amount = old_amount + arg_amount
11
+
12
+ if new_amount > arg_max then
13
+ return false
14
+ end
15
+ end
16
+
17
+ redis.call("INCRBY", KEYS[1], arg_amount)
18
+ return true
@@ -0,0 +1,44 @@
1
+ -- gets all fields from a hash as a dictionary
2
+ local hgetall = function (key)
3
+ local bulk = redis.call('HGETALL', key)
4
+ local result = {}
5
+ local nextkey
6
+ for i, v in ipairs(bulk) do
7
+ if i % 2 == 1 then
8
+ nextkey = v
9
+ else
10
+ result[nextkey] = v
11
+ end
12
+ end
13
+ return result
14
+ end
15
+
16
+ local arg_timestamp = tonumber(ARGV[1])
17
+ local arg_amount = tonumber(ARGV[2])
18
+ local arg_max = tonumber(ARGV[3])
19
+ local arg_period = tonumber(ARGV[4])
20
+
21
+ local sum = arg_amount
22
+
23
+ local clear_before = arg_timestamp - arg_period
24
+ local mytable = hgetall(KEYS[1])
25
+ if mytable ~= false then
26
+ -- print key -> value for mytable
27
+ print('mytable:')
28
+ for key, val in pairs(mytable) do
29
+ print(' ' .. key .. ' -> ' .. val)
30
+ if tonumber(key) < clear_before then
31
+ redis.call('HDEL', KEYS[1], key)
32
+ else
33
+ sum = sum + tonumber(val)
34
+ end
35
+ end
36
+ end
37
+
38
+ if sum > arg_max then
39
+ return false
40
+ end
41
+
42
+ redis.call("HINCRBY", KEYS[1], arg_timestamp, arg_amount)
43
+ redis.call("EXPIRE", KEYS[1], arg_period)
44
+ return true
@@ -0,0 +1,17 @@
1
+ local arg_incrby = tonumber(ARGV[1])
2
+ local arg_max = tonumber(ARGV[2])
3
+
4
+ local old_value = redis.call("PTTL", KEYS[1])
5
+ if old_value == -1 then -- key exists but has no associated expire
6
+ return -1 -- -1 signals key exists but has no associated expire
7
+ elseif old_value == -2 then -- key does not exist
8
+ redis.call("SET", KEYS[1], "", "PX", arg_incrby)
9
+ else
10
+ local new_value = old_value + arg_incrby
11
+ if new_value > arg_max then
12
+ return -2 -- -2 signals increment exceeds max
13
+ end
14
+ redis.call("PEXPIRE", KEYS[1], new_value)
15
+ end
16
+
17
+ return 0 -- 0 signals success
@@ -0,0 +1,32 @@
1
+ -- gets all fields from a hash as a dictionary
2
+ local hgetall = function (key)
3
+ local bulk = redis.call('HGETALL', key)
4
+ local result = {}
5
+ local nextkey
6
+ for i, v in ipairs(bulk) do
7
+ if i % 2 == 1 then
8
+ nextkey = v
9
+ else
10
+ result[nextkey] = v
11
+ end
12
+ end
13
+ return result
14
+ end
15
+
16
+ local arg_timestamp = tonumber(ARGV[1])
17
+ local arg_period = tonumber(ARGV[2])
18
+ local sum = 0
19
+
20
+ local clear_before = arg_timestamp - arg_period
21
+ local mytable = hgetall(KEYS[1])
22
+ if mytable ~= false then
23
+ for key, val in pairs(mytable) do
24
+ if tonumber(key) < clear_before then
25
+ redis.call('HDEL', KEYS[1], key)
26
+ else
27
+ sum = sum + tonumber(val)
28
+ end
29
+ end
30
+ end
31
+
32
+ return sum
metadata CHANGED
@@ -1,14 +1,14 @@
1
1
  --- !ruby/object:Gem::Specification
2
2
  name: traffic_jam
3
3
  version: !ruby/object:Gem::Version
4
- version: 1.0.1
4
+ version: 1.1.0
5
5
  platform: ruby
6
6
  authors:
7
7
  - Jim Posen
8
8
  autorequire:
9
9
  bindir: bin
10
10
  cert_chain: []
11
- date: 2015-08-13 00:00:00.000000000 Z
11
+ date: 2018-01-09 00:00:00.000000000 Z
12
12
  dependencies:
13
13
  - !ruby/object:Gem::Dependency
14
14
  name: redis
@@ -47,10 +47,17 @@ files:
47
47
  - lib/traffic_jam.rb
48
48
  - lib/traffic_jam/configuration.rb
49
49
  - lib/traffic_jam/errors.rb
50
+ - lib/traffic_jam/lifetime_limit.rb
50
51
  - lib/traffic_jam/limit.rb
51
52
  - lib/traffic_jam/limit_group.rb
53
+ - lib/traffic_jam/rolling_limit.rb
52
54
  - lib/traffic_jam/scripts.rb
55
+ - lib/traffic_jam/simple_limit.rb
56
+ - scripts/incrby.lua
53
57
  - scripts/increment.lua
58
+ - scripts/increment_rolling.lua
59
+ - scripts/increment_simple.lua
60
+ - scripts/sum_rolling.lua
54
61
  homepage: https://github.com/coinbase/traffic_jam
55
62
  licenses:
56
63
  - MIT
@@ -71,9 +78,8 @@ required_rubygems_version: !ruby/object:Gem::Requirement
71
78
  version: '0'
72
79
  requirements: []
73
80
  rubyforge_project:
74
- rubygems_version: 2.2.2
81
+ rubygems_version: 2.6.13
75
82
  signing_key:
76
83
  specification_version: 4
77
84
  summary: Library for time-based rate limiting
78
85
  test_files: []
79
- has_rdoc: