GenStateMachine Leaky Bucket example

leaky_bucket.ex

defmodule LeakyBucket do
  @moduledoc """
  A state machine implementation of the leaky bucket rate limiting algorithm
  
  The leaky bucket algorithm is a commonly used rate limiting algorithm. You
  start out with a number of drops in the bucket, and drops "drip" out of the
  bucket for each action taken. Drops drip into the bucket based on a fixed
  timer. This allows for some burst capacity.

  This version obviously doesn't scale very highly, but most leaky buckets out
  there don't refill very quickly.
  """

  use GenStateMachine

  @doc """
  Starts the leaky bucket
  """
  def start_link(drops, drip_timer, opts \\ []) do
    GenStateMachine.start_link(__MODULE__, {drops, drip_timer}, opts)
  end

  @doc """
  Blocks the caller until a drop is ready
  """
  def take_drop(server) do
    GenStateMachine.call(server, :request_drop)
  end

  @doc false
  # default state is ready, with max_drops, sending an interval every drip_timer
  def init([drops, drip_timer]) do
    {:ok, _} = :timer.send_interval(drip_timer, {:timer, :drip})
    max_drops = drops
    {:ok, :ready, {max_drops, drops}}
  end
  
  @doc false
  # Timer comes in, we bump our counter, but no larger than max_drops
  def handle_event(:info, {:timer, :drip}, _, {max_drops, drops}) do
    {:next_state, :ready, {max_drops, min(max_drops, drops+1)}}
  end

  @doc false
  # 1 drop left, send it on, but we're empty now
  def handle_event({:call, from}, :request_drop, :ready, {max_drops, 1}) do
    {:next_state, :empty, {max_drops, 0}, [{:reply, from, :ok}]}
  end
  # >1 drop left, take it
  def handle_event({:call, from}, :request_drop, :ready, {max_drops, drops}) do
    {:keep_state, {max_drops, drops-1}, [{:reply, from, :ok}]}
  end
  # no drops left, postpone incoming requests
  def handle_event({:call, _}, :request_drop, :empty, _) do
    {:keep_state_and_data, [:postpone]}
  end

end