Class: Restrainer

Inherits:
Object
  • Object
show all
Defined in:
lib/restrainer.rb

Overview

Redis backed throttling mechanism to ensure that only a limited number of processes can be executed at any one time.

Usage:

Restrainer.new(:foo, 10).throttle do
# Do something
end

If more than the specified number of processes as identified by the name argument is currently running, then the throttle block will raise an error.

Defined Under Namespace

Classes: ThrottledError

Constant Summary collapse

ADD_PROCESS_SCRIPT = 
<<-LUA
  redis.replicate_commands()

  -- Parse arguments
  local sorted_set = KEYS[1]
  local process_id = ARGV[1]
  local limit = tonumber(ARGV[2])
  local ttl = tonumber(ARGV[3])

  -- Use the Redis server clock so all clients share a consistent view of time.
  local time = redis.call('time')
  local now = tonumber(time[1]) + (tonumber(time[2]) / 1000000)

  -- Get count of current processes. If more than the max, check if any of the processes have timed out
  -- and try again.
  local process_count = redis.call('zcard', sorted_set)
  if process_count >= limit then
    local max_score = now - ttl
    local expired_keys = redis.call('zremrangebyscore', sorted_set, '-inf', max_score)
    if expired_keys > 0 then
      process_count = redis.call('zcard', sorted_set)
    end
  end

  -- Success so add to the list and set a global expiration so the list cleans up after itself.
  if process_count < limit then
    redis.call('zadd', sorted_set, now, process_id)
    redis.call('pexpire', sorted_set, math.ceil(ttl * 1000))
  end

  -- Return the number of processes running before the process was added.
  return process_count
LUA
ADD_PROCESS_SCRIPT_SHA1 =

SHA1 digest of the script used to call it with EVALSHA. The script itself is only loaded to the server if it isn't already cached there.

Digest::SHA1.hexdigest(ADD_PROCESS_SCRIPT).freeze

Instance Attribute Summary collapse

Class Method Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(name, limit:, timeout: 60, redis: nil) ⇒ Restrainer

Create a new restrainer. The name is used to identify the Restrainer and group processes together. You can create any number of Restrainers with different names.

The required limit parameter specifies the maximum number of processes that will be allowed to execute the throttle block at any point in time.

The timeout parameter is used for cleaning up internal data structures so that jobs aren't orphaned if their process is killed. Processes will automatically be removed from the running jobs list after the specified number of seconds. Note that the Restrainer will not handle timing out any code itself. This value is just used to insure the integrity of internal data structures.



109
110
111
112
113
114
115
 
# File 'lib/restrainer.rb', line 109

def initialize(name, limit:, timeout: 60, redis: nil)
  @name = name
  @limit = limit
  @timeout = timeout
  @key = "#{self.class.name}.#{name}"
  @redis = redis
end

Instance Attribute Details

#limit ⇒ Object (readonly)

Returns the value of attribute limit.



18
19
20
 
# File 'lib/restrainer.rb', line 18

def limit
  @limit
end

#name ⇒ Object (readonly)

Returns the value of attribute name.



18
19
20
 
# File 'lib/restrainer.rb', line 18

def name
  @name
end

Class Method Details

.redis(&block) ⇒ Redis

Either configure the redis instance using a block or yield the instance. Configuring with a block allows you to use things like connection pools etc. without hard coding a single instance.

Examples:

Restrainer.redis { redis_pool.instance }

Parameters:

  • block (Proc) —

    A block that returns a redis instance.

Returns:

  • (Redis) —

    The redis instance.



73
74
75
76
77
78
79
80
81
82
83
 
# File 'lib/restrainer.rb', line 73

def redis(&block)
  if block
    @redis = block
  else
    @redis ||= begin
      client = Redis.new
      lambda { client }
    end
    @redis.call
  end
end

.redis=(conn) ⇒ void

This method returns an undefined value.

Set the redis instance to a specific instance. It is usually preferable to use the block form for configurating the instance so that it can be evaluated at runtime.

Examples:

Restrainer.redis = Redis.new

Parameters:

  • conn (Redis) 


93
94
95
 
# File 'lib/restrainer.rb', line 93

def redis=(conn)
  @redis = lambda { conn }
end

Instance Method Details

#clear! ⇒ void

This method returns an undefined value.

Clear all locks.



183
184
185
 
# File 'lib/restrainer.rb', line 183

def clear!
  redis.del(key)
end

#current ⇒ Integer

Get the number of processes currently being executed for this restrainer.

Returns:

  • (Integer) —

    The number of processes currently being executed.



176
177
178
 
# File 'lib/restrainer.rb', line 176

def current
  redis.zcard(key).to_i
end

#lock!(process_id = nil, limit: nil) ⇒ String?

Obtain a lock on one the allowed processes. The method returns a process identifier that must be passed to the release! to release the lock. You can pass in a unique identifier if you already have one.

Parameters:

  • process_id (String) (defaults to: nil) —

    A unique identifier for the process. If none is passed, a unique value will be generated.

  • limit (Integer) (defaults to: nil) —

    The maximum number of processes that can be executing at any one time. Defaults to the value passed to the constructor.

Returns:

  • (String, nil) —

    The process identifier, or nil if the limit is negative (i.e. the throttle is wide open and there is nothing to release).

Raises:



150
151
152
153
154
155
156
157
158
159
160
 
# File 'lib/restrainer.rb', line 150

def lock!(process_id = nil, limit: nil)
  process_id ||= SecureRandom.uuid
  limit ||= self.limit

  # limit of less zero is no limit; limit of zero is allow none
  return nil if limit < 0
  raise ThrottledError.new("#{self.class}: #{@name} is not allowing any processing") if limit == 0

  add_process!(redis, process_id, limit)
  process_id
end

#release!(process_id) ⇒ Boolean

Release one of the allowed processes. You must pass in a process id returned by the lock method.

Parameters:

  • process_id (String) —

    The process identifier returned by the lock call.

Returns:

  • (Boolean) —

    True if the process was released, false if it was not found.



167
168
169
170
171
 
# File 'lib/restrainer.rb', line 167

def release!(process_id)
  return false if process_id.nil?

  remove_process!(redis, process_id)
end

#throttle(limit: nil) ⇒ Object

Wrap a block with this method to throttle concurrent execution. If more than the alotted number of processes (as identified by the name) are currently executing, then a Restrainer::ThrottledError will be raised.

Parameters:

  • limit (Integer) (defaults to: nil) —

    The maximum number of processes that can be executing at any one time. Defaults to the value passed to the constructor.

Returns:

  • (Object) —

    The value returned by the block.

Raises:



125
126
127
128
129
130
131
132
133
134
135
136
137
 
# File 'lib/restrainer.rb', line 125

def throttle(limit: nil)
  limit ||= self.limit

  # limit of less zero is no limit; limit of zero is allow none
  return yield if limit < 0

  process_id = lock!(limit: limit)
  begin
    yield
  ensure
    release!(process_id)
  end
end