lib/simple_rate_limiter.ex

defmodule SimpleRateLimiter do
  @moduledoc """
  A GenServer module that implements token bucket rate limiting.

  To start the rate limiter, specify the maximum number of tokens and the
  refill interval in milliseconds. The token bucket is initially filled to
  the maximum.

  To check if an action can proceed, call `SimpleRateLimiter.can_proceed?/1`.
  It will return `:ok` if there are tokens available, and decrement the
  token count by 1. If there are no tokens available, it will return an
  error tuple with the amount of time remaining until the next token is
  available.
  """
  use GenServer

  @doc """
  Starts the rate limiter.

    ## Options

    - `:interval`: The refill interval in milliseconds (default `1_000`).
    - `:max`: The maximum number of tokens (default `10`).

    ## Example

        iex> {:ok, pid} = SimpleRateLimiter.start_link(interval: 1_000, max: 10)
  """
  def start_link(opts \\ []) do
    GenServer.start_link(__MODULE__, opts, name: __MODULE__)
  end

  @doc false
  def init(opts) do
    interval = Keyword.get(opts, :interval, 1_000)
    max = Keyword.get(opts, :max, 10)
    {:ok, %{interval: interval, max: max, tokens: max, last_token_time: :os.system_time(:millisecond)}}
  end

  @doc """
    Checks if an action can proceed by trying to consume a token.

    ## Example

        iex> SimpleRateLimiter.can_proceed?(pid)
        :ok
  """
  def can_proceed?(server \\ __MODULE__) do
    GenServer.call(server, :can_proceed?)
  end

  @doc false
  def handle_call(:can_proceed?, _from, state) do
    current_time = :os.system_time(:millisecond)
    elapsed_time = current_time - state.last_token_time

    added_tokens = div(elapsed_time, state.interval)
    new_tokens = min(state.tokens + added_tokens, state.max)

    if new_tokens > 0 do
      new_state = %{state | tokens: new_tokens - 1, last_token_time: current_time}
      {:reply, :ok, new_state}
    else
      remaining_time = state.interval - rem(elapsed_time, state.interval)
      {:reply, {:error, :rate_limit_exceeded, remaining_time: remaining_time}, state}
    end
  end

  @doc """
    Waits until an action can proceed, then calls the given function.

    ## Example

        iex> SimpleRateLimiter.wait_and_proceed(pid, fn -> IO.puts("Hello!") end)
        Hello!
  """
  def wait_and_proceed(server \\ __MODULE__, action_fun) when is_function(action_fun, 0) do
    case can_proceed?(server) do
      :ok ->
        action_fun.()
      {:error, :rate_limit_exceeded, remaining_time: remaining_time} ->
        :timer.sleep(remaining_time)
        wait_and_proceed(server, action_fun)
    end
  end
end