Class: Pecorino::LeakyBucket

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

Overview

This offers just the leaky bucket implementation with fill control, but without the timed lock. It does not raise any exceptions, it just tracks the state of a leaky bucket in the database.

Leak rate is specified directly in tokens per second, instead of specifying the block period. The bucket level is stored and returned as a Float which allows for finer-grained measurement, but more importantly - makes testing from the outside easier.

Note that this implementation has a peculiar property: the bucket is only “full” once it overflows. Due to a leak rate just a few microseconds after that moment the bucket is no longer going to be full anymore as it will have leaked some tokens by then. This means that the information about whether a bucket has become full or not gets returned in the bucket ‘State` struct right after the database update gets executed, and if your code needs to make decisions based on that data it has to use this returned state, not query the leaky bucket again. Specifically:

state = bucket.fillup(1) # Record 1 request
state.full? #=> true, this is timely information

…is the correct way to perform the check. This, however, is not:

bucket.fillup(1)
bucket.state.full? #=> false, some time has passed after the topup and some tokens have already leaked

The storage use is one DB row per leaky bucket you need to manage (likely - one throttled entity such as a combination of an IP address + the URL you need to procect). The ‘key` is an arbitrary string you provide.

Defined Under Namespace

Classes: ConditionalFillupResult, State

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(key:, capacity:, adapter: Pecorino.adapter, leak_rate: nil, over_time: nil) ⇒ LeakyBucket

Creates a new LeakyBucket. The object controls 1 row in the database is specific to the bucket key.

Parameters:

  • key (String) —

    the key for the bucket. The key also gets used to derive locking keys, so that operations on a particular bucket are always serialized.

  • leak_rate (Float) (defaults to: nil) —

    the leak rate of the bucket, in tokens per second. Either ‘leak_rate` or `over_time` can be used, but not both.

  • over_time (#to_f) (defaults to: nil) —

    over how many seconds the bucket will leak out to 0 tokens. The value is assumed to be the number of seconds

    • or a duration which returns the number of seconds from ‘to_f`.

    Either ‘leak_rate` or `over_time` can be used, but not both.

  • capacity (Numeric) —

    how many tokens is the bucket capped at. Filling up the bucket using ‘fillup()` will add to that number, but the bucket contents will then be capped at this value. So with bucket_capacity set to 12 and a `fillup(14)` the bucket will reach the level of 12, and will then immediately start leaking again.

  • adapter (Pecorino::Adapters::BaseAdapter) (defaults to: Pecorino.adapter) —

    a compatible adapter

Raises:

  • (ArgumentError) 


94
95
96
97
98
99
100
101
 
# File 'lib/pecorino/leaky_bucket.rb', line 94

def initialize(key:, capacity:, adapter: Pecorino.adapter, leak_rate: nil, over_time: nil)
  raise ArgumentError, "Either leak_rate: or over_time: must be specified" if leak_rate.nil? && over_time.nil?
  raise ArgumentError, "Either leak_rate: or over_time: may be specified, but not both" if leak_rate && over_time
  @leak_rate = leak_rate || (capacity / over_time.to_f)
  @key = key
  @capacity = capacity.to_f
  @adapter = adapter
end

Instance Attribute Details

#capacity ⇒ Object (readonly)

The capacity of the bucket in tokens

@return [Float]



74
75
76
 
# File 'lib/pecorino/leaky_bucket.rb', line 74

def capacity
  @capacity
end

#key ⇒ Object (readonly)

The key (name) of the leaky bucket

@return [String]



66
67
68
 
# File 'lib/pecorino/leaky_bucket.rb', line 66

def key
  @key
end

#leak_rate ⇒ Object (readonly)

The leak rate (tokens per second) of the bucket

@return [Float]



70
71
72
 
# File 'lib/pecorino/leaky_bucket.rb', line 70

def leak_rate
  @leak_rate
end

Instance Method Details

#able_to_accept?(n_tokens) ⇒ boolean

Tells whether the bucket can accept the amount of tokens without overflowing. Calling this method will not perform any database writes. Note that this call is not race-safe - another caller may still overflow the bucket. Before performing your action, you still need to call ‘fillup()` - but you can preemptively refuse a request if you already know the bucket is full.

Parameters:

  • n_tokens (Float)

Returns:

  • (boolean) 


157
158
159
 
# File 'lib/pecorino/leaky_bucket.rb', line 157

def able_to_accept?(n_tokens)
  (state.level + n_tokens) <= @capacity
end

#fillup(n_tokens) ⇒ State

Places ‘n` tokens in the bucket. If the bucket has less capacity than `n` tokens, the bucket will be filled to capacity. If the bucket has less capacity than `n` tokens, it will be filled to capacity. If the bucket is already full when the fillup is requested, the bucket stays at capacity.

Once tokens are placed, the bucket is set to expire within 2 times the time it would take it to leak to 0, regardless of how many tokens get put in - since the amount of tokens put in the bucket will always be capped to the ‘capacity:` value you pass to the constructor.

Parameters:

  • n_tokens (Float) —

    How many tokens to fillup by

Returns:

  • (State) —

    the state of the bucket after the operation



113
114
115
116
 
# File 'lib/pecorino/leaky_bucket.rb', line 113

def fillup(n_tokens)
  capped_level_after_fillup, is_full = @adapter.add_tokens(capacity: @capacity, key: @key, leak_rate: @leak_rate, n_tokens: n_tokens)
  State.new(capped_level_after_fillup, is_full)
end

#fillup_conditionally(n_tokens) ⇒ ConditionalFillupResult

Places ‘n` tokens in the bucket. If the bucket has less capacity than `n` tokens, the fillup will be rejected. This can be used for “exactly once” semantics or just more precise rate limiting. Note that if the bucket has exactly `n` tokens of capacity the fillup will be accepted.

Once tokens are placed, the bucket is set to expire within 2 times the time it would take it to leak to 0, regardless of how many tokens get put in - since the amount of tokens put in the bucket will always be capped to the ‘capacity:` value you pass to the constructor.

Examples:

withdrawals = LeakyBuket.new(key: "wallet-#{user.id}", capacity: 200, over_time: 1.day)
if withdrawals.fillup_conditionally(amount_to_withdraw).accepted?
  user.wallet.withdraw(amount_to_withdraw)
else
  raise "You need to wait a bit before withdrawing more"
end

Parameters:

  • n_tokens (Float) —

    How many tokens to fillup by

Returns:



135
136
137
138
 
# File 'lib/pecorino/leaky_bucket.rb', line 135

def fillup_conditionally(n_tokens)
  capped_level_after_fillup, is_full, did_accept = @adapter.add_tokens_conditionally(capacity: @capacity, key: @key, leak_rate: @leak_rate, n_tokens: n_tokens)
  ConditionalFillupResult.new(capped_level_after_fillup, is_full, did_accept)
end

#state ⇒ State

Returns the current state of the bucket, containing the level and whether the bucket is full. Calling this method will not perform any database writes.

Returns:

  • (State) —

    the snapshotted state of the bucket at time of query



144
145
146
147
 
# File 'lib/pecorino/leaky_bucket.rb', line 144

def state
  current_level, is_full = @adapter.state(key: @key, capacity: @capacity, leak_rate: @leak_rate)
  State.new(current_level, is_full)
end