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!, #keys, #on_reject, #stats, #throttle

Constructor Details

#initialize(rate:, capacity:) ⇒ TokenBucket

Returns a new instance of TokenBucket.



12
13
14
15
16
17
18
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 12

def initialize(rate:, capacity:)
  @rate = rate.to_f
  @capacity = capacity.to_f
  @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) 


20
21
22
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 20

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



28
29
30
31
32
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 28

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)



57
58
59
60
61
62
63
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 57

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



77
78
79
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 77

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

#info(key) ⇒ Object 



65
66
67
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 65

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

#peek(key) ⇒ Object 



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

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

#refund(key, amount: 1) ⇒ Object 



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

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

#remaining(key) ⇒ Object 



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

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

#reset(key) ⇒ Object 



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

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)



102
103
104
105
106
107
108
109
110
111
112
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 102

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



46
47
48
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 46

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)



86
87
88
89
90
91
92
93
94
95
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 86

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