Class: Philiprehberger::RateLimiter::TokenBucket

Inherits:
Object
  • Object
show all
Includes:
StatsTracking
Defined in:
lib/philiprehberger/rate_limiter/token_bucket.rb

Instance Attribute Summary collapse

Instance Method Summary collapse

Methods included from StatsTracking

#allow!, #block, #keys, #on_reject, #stats, #throttle

Constructor Details

#initialize(rate:, capacity:, max_keys: nil) ⇒ TokenBucket

Returns a new instance of TokenBucket.

Parameters:

  • rate (Numeric) —

    tokens refilled per second

  • capacity (Numeric) —

    max tokens the bucket can hold

  • max_keys (Integer, nil) (defaults to: nil) —

    cap on tracked keys; least-recently-touched key is evicted (LRU) once the cap is exceeded. nil means unbounded.



16
17
18
19
20
21
22
23
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 16

def initialize(rate:, capacity:, max_keys: nil)
  @rate = rate.to_f
  @capacity = capacity.to_f
  @max_keys = max_keys
  @store = {}
  @mutex = Mutex.new
  init_stats
end

Instance Attribute Details

#capacity ⇒ Object (readonly)

Returns the value of attribute capacity.



10
11
12
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 10

def capacity
  @capacity
end

#rate ⇒ Object (readonly)

Returns the value of attribute rate.



10
11
12
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 10

def rate
  @rate
end

Instance Method Details

#allow?(key, weight: 1) ⇒ Boolean

Returns:

  • (Boolean) 


25
26
27
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 25

def allow?(key, weight: 1)
  @mutex.synchronize { try_acquire(key, weight.to_f) }
end

#allow_batch(keys) ⇒ Hash{Object => Boolean}

Check many keys in a single mutex acquisition.

Parameters:

  • keys (Array<Symbol, String>) —

    the keys to check and consume

Returns:

  • (Hash{Object => Boolean}) —

    mapping of each key to the allow result



33
34
35
36
37
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 33

def allow_batch(keys)
  @mutex.synchronize do
    keys.to_h { |key| [key, try_acquire(key, 1.0)] }
  end
end

#clear ⇒ void

This method returns an undefined value.

Clear state for all keys (resets quotas and stats for every tracked key)



62
63
64
65
66
67
68
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 62

def clear
  @mutex.synchronize do
    @store.clear
    @stats_store.clear
  end
  nil
end

#drain(key = :default) ⇒ Integer

Forcefully consume all remaining tokens for a key.

Parameters:

  • key (Symbol, String) (defaults to: :default) —

    the rate limit key

Returns:

  • (Integer) —

    the integer floor of tokens drained



114
115
116
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 114

def drain(key = :default)
  @mutex.synchronize { drain_tokens(key) }
end

#info(key) ⇒ Object 



70
71
72
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 70

def info(key)
  @mutex.synchronize { build_info(key) }
end

#info_batch(keys) ⇒ Hash{Object => Hash}

Build info hashes for many keys in a single mutex acquisition.

Mirrors #allow_batch for inspection rather than acquisition.

Parameters:

  • keys (Array<Symbol, String>) —

    the keys to inspect

Returns:

  • (Hash{Object => Hash}) —

    mapping of each key to its info hash



80
81
82
83
84
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 80

def info_batch(keys)
  @mutex.synchronize do
    keys.to_h { |key| [key, build_info(key)] }
  end
end

#peek(key) ⇒ Object 



39
40
41
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 39

def peek(key)
  @mutex.synchronize { token_count(key) >= 1.0 }
end

#prune ⇒ Integer

Drop keys whose bucket has fully refilled back to capacity, reclaiming memory for idle keys. Also clears their stats.

Returns:

  • (Integer) —

    the number of keys pruned



94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 94

def prune
  @mutex.synchronize do
    pruned = 0
    tracked = @store.keys
    tracked.each do |key|
      refill(key)
      next unless @store[key][:tokens] >= @capacity

      @store.delete(key)
      @stats_store.delete(key)
      pruned += 1
    end
    pruned
  end
end

#refund(key, amount: 1) ⇒ Object 



86
87
88
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 86

def refund(key, amount: 1)
  @mutex.synchronize { refund_tokens(key, amount.to_f) }
end

#remaining(key) ⇒ Object 



43
44
45
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 43

def remaining(key)
  @mutex.synchronize { token_count(key).to_i }
end

#reset(key) ⇒ Object 



55
56
57
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 55

def reset(key)
  @mutex.synchronize { @store.delete(key.to_s) }
end

#retry_after(key = :default) ⇒ Float

Seconds until the next request would be allowed, suitable for the HTTP Retry-After header. Returns 0.0 when a token is available right now.

Parameters:

  • key (Symbol, String) (defaults to: :default) —

    the rate limit key

Returns:

  • (Float) —

    seconds until 1 token is available (0.0 if allowed now)



139
140
141
142
143
144
145
146
147
148
149
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 139

def retry_after(key = :default)
  @mutex.synchronize do
    refill(key)
    bucket = @store[key.to_s]
    tokens = bucket ? bucket[:tokens] : @capacity
    return 0.0 if tokens >= 1.0

    needed = 1.0 - tokens
    needed / @rate
  end
end

#used(key) ⇒ Integer

Number of currently consumed tokens for a key.

Parameters:

  • key (Symbol, String) —

    the rate limit key

Returns:

  • (Integer) —

    capacity minus remaining tokens



51
52
53
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 51

def used(key)
  @mutex.synchronize { @capacity.to_i - token_count(key).to_i }
end

#wait_time(key = :default, weight: 1) ⇒ Float

Seconds until the next request would be allowed

Parameters:

  • key (Symbol, String) (defaults to: :default) —

    the rate limit key

  • weight (Integer) (defaults to: 1) —

    tokens needed

Returns:

  • (Float) —

    seconds to wait (0 if allowed now)



123
124
125
126
127
128
129
130
131
132
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 123

def wait_time(key = :default, weight: 1)
  @mutex.synchronize do
    refill(key)
    tokens = @store[key.to_s] ? @store[key.to_s][:tokens] : @capacity
    return 0.0 if tokens >= weight

    needed = weight - tokens
    needed / @rate
  end
end