Elixir Token Bucket Rate Limiting

A token bucket is a limiter that spends tokens for each call. A call runs only when enough tokens remain in the bucket.

A full bucket pays for calls from its balance

A bucket with a capacity of 10 starts full. Taking 3 tokens leaves 7.

defmodule Bucket do
  def new(capacity, refill), do: %{tokens: capacity, capacity: capacity, refill: refill}
  def take(bucket, cost) do
    if bucket.tokens >= cost do
      {%{bucket | tokens: bucket.tokens - cost}, :ok}
    else
      {bucket, :denied}
    end
  end
end
{bucket, result} = Bucket.take(Bucket.new(10, 2), 3)
IO.inspect(result)
IO.puts(bucket.tokens)

A short balance denies the call and keeps its tokens

A bucket holding 2 tokens cannot pay a cost of 5. The denied call leaves the balance at 2.

defmodule Bucket do
  def new(capacity, refill), do: %{tokens: capacity, capacity: capacity, refill: refill}
  def take(bucket, cost) do
    if bucket.tokens >= cost do
      {%{bucket | tokens: bucket.tokens - cost}, :ok}
    else
      {bucket, :denied}
    end
  end
end
bucket = Bucket.new(10, 2)
{bucket, _first} = Bucket.take(bucket, 8)
{bucket, result} = Bucket.take(bucket, 5)
IO.inspect(result)
IO.puts(bucket.tokens)

Refill adds tokens up to the capacity cap

Refill adds 2 tokens per counted step, up to the capacity cap. After 5 steps a balance of 2 is held at 10, not 12.

defmodule Bucket do
  def new(capacity, refill), do: %{tokens: capacity, capacity: capacity, refill: refill}
  def take(bucket, cost) do
    if bucket.tokens >= cost do
      {%{bucket | tokens: bucket.tokens - cost}, :ok}
    else
      {bucket, :denied}
    end
  end
  def refill(bucket, steps), do: %{bucket | tokens: min(bucket.capacity, bucket.tokens + bucket.refill * steps)}
end
bucket = Bucket.new(10, 2)
{bucket, _first} = Bucket.take(bucket, 8)
bucket = Bucket.refill(bucket, 5)
IO.puts(bucket.tokens)

Example

defmodule Bucket do
  def new(capacity, refill), do: %{tokens: capacity, capacity: capacity, refill: refill}
  def take(bucket, cost) do
    if bucket.tokens >= cost do
      {%{bucket | tokens: bucket.tokens - cost}, :ok}
    else
      {bucket, :denied}
    end
  end
  def refill(bucket, steps), do: %{bucket | tokens: min(bucket.capacity, bucket.tokens + bucket.refill * steps)}
end
bucket = Bucket.new(5, 1)
{bucket, first} = Bucket.take(bucket, 4)
{bucket, second} = Bucket.take(bucket, 3)
bucket = Bucket.refill(bucket, 3)
{bucket, third} = Bucket.take(bucket, 4)
IO.inspect([first, second, third])
IO.puts(bucket.tokens)

Key points

Token: a unit of permission that one call spends from the bucket.
Capacity: the largest number of tokens the bucket can hold.
Refill: tokens added back to the bucket by a fixed amount per counted step.
Take: a call removes tokens when the balance covers its cost.
Denied: a call is refused and the balance stays unchanged when tokens are short.

In this example, a bucket with a capacity of 5 and a refill of 1 starts full, pays for 4 tokens, denies a cost of 3, refills by 3 steps to 4 tokens, and then pays for 4 tokens to reach 0.

Note

this lesson uses counted steps for refill, so no real waiting is needed to show the balance.

Try it Yourself

Edit the code, then press Run.

Running is simulated

defmodule Bucket do
  def new(capacity, refill), do: %{tokens: capacity, capacity: capacity, refill: refill}
  def take(bucket, cost) do
    if bucket.tokens >= cost do
      {%{bucket | tokens: bucket.tokens - cost}, :ok}
    else
      {bucket, :denied}
    end
  end
  def refill(bucket, steps), do: %{bucket | tokens: min(bucket.capacity, bucket.tokens + bucket.refill * steps)}
end
bucket = Bucket.new(5, 1)
{bucket, first} = Bucket.take(bucket, 4)
{bucket, second} = Bucket.take(bucket, 3)
bucket = Bucket.refill(bucket, 3)
{bucket, third} = Bucket.take(bucket, 4)
IO.inspect([first, second, third])
IO.puts(bucket.tokens)

Exercise

A bucket holds 3 tokens with a capacity of 5 and a refill of 2 per step. What is the balance after 2 refill steps?

7 3 5 4