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

#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

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



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

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