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

#clear ⇒ void

This method returns an undefined value.

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



39
40
41
42
43
44
45
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 39

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



59
60
61
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 59

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

#info(key) ⇒ Object 



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

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

#peek(key) ⇒ Object 



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

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

#refund(key, amount: 1) ⇒ Object 



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

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

#remaining(key) ⇒ Object 



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

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

#reset(key) ⇒ Object 



32
33
34
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 32

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)



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

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

#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)



68
69
70
71
72
73
74
75
76
77
 
# File 'lib/philiprehberger/rate_limiter/token_bucket.rb', line 68

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